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.
def forward(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
schedule(skip_first=2, wait=1, warmup=1, active=3)
# skip_first=2 absorbs compilation; active=3 captures 3 clean steps → 181 MB trace
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
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