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

karpathy / nanoGPT

The simplest, fastest repository for training/finetuning medium-sized GPTs.

MIT License

34.48k stars 5.3k forks source link

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

Open GM-git-dotcom opened 1 month ago

GM-git-dotcom commented 1 month ago

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:

https://github.com/karpathy/nanoGPT/blob/325be85d9be8c81b436728a420e85796c57dba7e/train.py#L299-L305

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

karpathy commented 3 weeks ago

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