karpathy / build-nanogpt

Video+code lecture on building nanoGPT from scratch
3.61k stars 502 forks source link

Make example generation work when the model is torch.compile-d #78

Open IggShaman opened 2 months ago

IggShaman commented 2 months ago

Andrej mentions this issue a couple of times in the video. Apparently, when a model is compiled, it will fix its input and output tensor sizes. The training regime works with inputs of shape (B, T), currently configured to (64, 1024). Example generation works with input tensors of shapes (4, 8) ... (4, 32). A simple workaround here is to pad example generation tensors to the (B, T) shape, and ignore extra rows. The same workaround applies to the hellaswag eval. Unfortunately, this increases the computation time linearly to both the B and T, so quadratically overall.