◷train_gpt.py▸torch.compile on model◎learner＄no GPU
1/2
Blocks
Quick summary
dynamic=False: assume fixed shapes → Triton can fuse ops and allocate static buffers. fullgraph=True: any Python…
Full explanation below the code →
fba-lab — train_gpt.py · torch.compile on modelexecuting
// block: torch.compile on model · lines 1871–1872$ study train_gpt.py --block model-compiledynamic=False: assume fixed shapes → Triton can fuse ops and allocate static buffers. fullgraph=True… ✓
Explanation
dynamic=False: assume fixed shapes → Triton can fuse ops and allocate static buffers. fullgraph=True: any Python control flow that breaks the FX graph is an error, not a silent fallback to slow eager mode.
Think about
What would happen if you used dynamic=True instead?
Why the first steps are slow — and how 67ms/step is achieved
If you open the Perfetto trace, you'll see spans labeled Torch-Compiled Region: N/0 across each training step.
On the very first call, the compiler traces the forward pass into Triton kernel source, compiles it to a .so binary, and caches it.
After that upfront cost, every step runs the cached kernels — no Python overhead in the hot path.
The profiler schedule skip_first=2, wait=1, warmup=1, active=3 is deliberately tuned to capture only the compiled steady state.
By step 4, all 11 layers are fully compiled and the per-step time has stabilized at 67.83ms.
@torch.compile(dynamic=False, fullgraph=True) # applied to GPT.forward and polar_express defforward(self, inputs, ...):
... # first call: trace → FX graph → Triton codegen → /tmp/torchinductor_root/mi/*.so
... # all subsequent calls: load cached .so, run fused CUDA kernels directly
Fixed shapes → Triton can fuse ops and allocate static buffers. The result is the 2.4× speedup from ~163ms (pre-compile estimate) to 67.83ms/step.
Fused kernels visible as triton__* in the GPU row of the trace
Shape change (e.g. curriculum window resize) → recompile at that step
fullgraph=True: no silent fallbacks
If any Python-level control flow breaks the FX graph — a data-dependent conditional, a print inside forward — raise an error immediately rather than silently falling back to slow eager mode.
Forces all ops to be compileable — no Python loops inside forward
smear_gate, bigram_embed, all 11 attention layers — one compiled graph