torch.compile vs eager for training a 468M Llama-style model on four GPUs: step time, memory and compile cost

Does torch.compile make a small transformer train faster? I rented four GPUs on Runcrate today (2026-10-29), an RTX 5090, an L40S, an RTX A6000 and an H100 PCIe, trained the same 468M-parameter Llama-style model in eager mode and compiled, with PyTorch's SDPA attention and with the plain "eager" attention, and deleted every machine when the numbers were saved.

Short version: with SDPA attention, torch.compile changed the step time by between 9% slower (RTX 5090) and 3% faster (L40S), cost 22 to 39 s of compile time on the first step, and cut peak memory by 23% on every GPU (17.3 to 13.3 GB). With the plain attention it was a different story: on the H100 it made the step 2.6x faster and fit the model in 39 GB instead of 56 GB, and on the 48 GB cards it was the only way the run fit at all. Even then, eager mode with SDPA beat compiled plain attention on every GPU.

import json, sys, time, torch from transformers import LlamaConfig, LlamaForCausalLM mode, attn = sys. argv [ 1 ], sys. argv [ 2 ] # eager|compile, sdpa|eager B, S = 4, 2048 cfg = LlamaConfig ( vocab_size = 32000, hidden_size = 1024, intermediate_size = 4096, num_hidden_layers = 24, num_attention_heads = 16, num_key_value_heads = 16, max_position_embeddings = S, attn_implementation = attn ) model = LlamaForCausalLM ( cfg ). to ( " cuda ", dtype = torch. bfloat16 ) opt = torch. optim. AdamW ( model. parameters (), lr = 1e-4, fused = True ) fwd = torch. compile ( model ) if mode == " compile " else model x = torch. randint ( 0, 32000, ( B, S ), device = " cuda " ) def step (): loss = fwd ( input_ids = x, labels = x ). loss loss. backward (); opt. step (); opt. zero_grad ( set_to_none = True ) t0 = time. time (); step (); torch. cuda. synchronize (); first = time. time () - t0 for _ in range ( 4 ): step () torch. cuda. synchronize (); torch. cuda. reset_peak_memory_stats () t0 = time. time () for _ in range ( 20 ): step () torch. cuda. synchronize (); dt = ( time. time () - t0 ) / 20 print ( json. dumps ({ " mode ": mode, " attn ": attn, " first_step_s ": first, " step_ms ": dt * 1000, " tokens_per_s ": B * S / dt, " peak_gb ": torch. cuda. max_memory_allocated () / 1e9 })) Enter fullscreen mode Exit fullscreen mode Results GPU $/hour Setup Step (ms) Tokens/s Peak memory First step (s) M tokens per $ RTX 5090 1.15 eager + sdpa attention 196 41,870 17.3 GB 0.6 131 RTX 5090 1.15 compile + sdpa attention 215 38,040 13.3 GB 21.6 119 RTX 5090 1.15 eager + eager attention OOM RTX 5090 1.15 compile + eager attention OOM L40S 0.54 eager + sdpa attention 244 33,581 17.3 GB 0.6 224 L40S 0.54 compile + sdpa attention 236 34,691 13.3 GB 34.2 231 L40S 0.54 eager + eager attention OOM L40S 0.54 compile + eager attention 391 20,973 39.1 GB 38.9 140 RTX A6000 0.65 eager + sdpa attention 457 17,942 17.3 GB 0.7 99 RTX A6000 0.65 compile + sdpa attention 455 18,005 13.3 GB 39.2 100 RTX A6000 0.65 eager + eager attention OOM RTX A6000 0.65 compile + eager attention 546 15,017 39.1 GB 47.1 83 H100 3.86 eager + sdpa attention 152 54,022 17.4 GB 0.4 50 H100 3.86 compile + sdpa attention 160 51,055 13.4 GB 36.4 48 H100 3.86 eager + eager attention 435 18,813 56.0 GB 0.7 18 H100 3.86 compile + eager attention 169 48,469 39.1 GB 43.9 45 "M tokens per $" is training tokens per dollar of GPU time, in millions, at each machine's hourly price. OOM means the run did not fit in the GPU's memory.

If you want to run it yourself: runcrate.ai

Written by the Runcrate team, with AI assistance. The numbers come from the run described above; raw logs are kept with the post.