[Q] Async prefetch next batch while model is doing forward pass

[GM-git-dotcom]

Hi,

Love the repo for the self-explanatory code and fun comments!

I would appreciate it if you could explain how the next batch is asynchronously being fetched in train.py:

nanoGPT/train.py

Lines 299 to 305 in 325be85

	with ctx:
	logits, loss = model(X, Y)
	loss = loss / gradient_accumulation_steps # scale the loss to account for gradient accumulation
	# immediately async prefetch next batch while model is doing the forward pass on the GPU
	X, Y = get_batch('train')
	# backward pass, with gradient scaling if training in fp16
	scaler.scale(loss).backward()

At the outset, without explicitly using something like asyncio, the execution seemed sequential to me. Apologies for the potential triviality, thanks!

[karpathy]

It's due to this

x, y = x.pin_memory().to(device, non_blocking=True), y.pin_memory().to(device, non_blocking=True)

happens async