Runcrate benchmarked torch.compile versus eager mode for training a 468M-parameter Llama-style model (hidden=1024, 24 layers, 16 heads, FFN=4096, vocab=32k, bfloat16) on four single-GPU instances: RTX 5090 (32 GB, $1.15/hr), L40S (48 GB, $0.54/hr), RTX A6000 (48 GB, $0.65/hr), and H100 PCIe (80 GB, $3.86/hr). Each ran PyTorch 2.8.0, transformers 4.57.1, with batch 4×2048 tokens, fused AdamW (lr=1e-4), and zero_grad(set_to_none=True). Four configurations were tested: eager/compile × SDPA/eager attention. With SDPA, compile changed step time by -9% (RTX 5090: 215 vs 196 ms) to +3% (L40S: 236 vs 244 ms), while reducing peak memory 23% universally (17.3→13.3 GB). Compilation overhead ranged 22–39 s (SDPA) and 39–47 s (eager attention). With naive eager attention, compile yielded 2.6× speedup on H100 (169 vs 435 ms) and 30% memory reduction (39 vs 56 GB), and was the only way the model fit on 48 GB cards; RTX 5090 OOMed regardless. However, eager SDPA outperformed compile+eager attention on all GPUs (H100: 152 vs 169 ms) at half the memory. Cost-efficiency favored L40S (231M tokens/$ compile+SDPA) over RTX 5090 (119M), A6000 (100M), and H100 (48M). H100 was fastest in throughput (54k tok/s) but 3.4× pricier than RTX 5090. Conclusions: prioritize SDPA/FlashAttention first; use compile for memory savings (enabling larger batch/context) rather than speed at this scale; measure per-GPU as compile impact varies (±9%).

Read original