DEV Community

Runcrate
Runcrate

Posted on

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

#pytorch #machinelearning #gpu #python

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.

Setup

  • GPUs (one per machine), on-demand price at launch time: RTX 5090 32 GB (Austin, US, $1.15 per hour), L40S 48 GB (Des Moines, US, $0.54), RTX A6000 48 GB (Montreal, CA, $0.65), H100 PCIe 80 GB (US Midwest, $3.86).
  • Model: LlamaForCausalLM from a config, random init: hidden 1024, 24 layers, 16 heads, FFN 4096, vocab 32,000, 468.2M parameters, weights in bf16.
  • Training step: batch 4 x 2,048 tokens (8,192 tokens per step) of random token ids, loss.backward(), fused AdamW (lr 1e-4), zero_grad(set_to_none=True).
  • Four setups, each in its own process: eager or torch.compile(model) (default mode), times attn_implementation="sdpa" or "eager" (the naive matmul-softmax-matmul attention in transformers).
  • Measured: first step time (includes compilation), then 4 more warm-up steps, then the mean of 20 timed steps, tokens per second, and peak memory from torch.cuda.max_memory_allocated.
  • Software: torch 2.8.0+cu128, transformers 4.57.1, Python 3.12, the same on all four machines.

The script 

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.

What the numbers say

  1. With SDPA attention, compile did not speed up training. Step time compiled against eager: RTX 5090 215 vs 196 ms (9% slower), H100 160 vs 152 ms (5% slower), RTX A6000 455 vs 457 ms (same), L40S 236 vs 244 ms (3% faster). SDPA already runs a fused attention kernel, and at this model size the matmuls dominate, so there was little left for the compiler to fuse.
  2. Compile did save memory: 4 GB, 23%, on every GPU. Peak allocation fell from 17.3 to 13.3 GB. Fused elementwise ops keep fewer intermediate activations for the backward pass, so compile can buy a larger batch even when it does not buy speed.
  3. Compilation took 22 to 47 s. (22 to 39 s with SDPA.) For a short experiment that is a real share of the run, and on the RTX 5090 and the H100, where compiled steps were also slower, it never paid back.
  4. Naive attention is where compile pays. On the H100, eager attention ran at 435 ms per step and 56 GB; compiled, 169 ms (2.6x faster) and 39 GB. On the L40S and the RTX A6000 eager attention ran out of memory without compile and fit with it. On the 32 GB RTX 5090 it did not fit either way.
  5. But the right attention kernel beat the compiler. Eager mode with SDPA was faster than compiled naive attention on every GPU (H100 152 vs 169 ms, L40S 244 vs 391 ms) and used less than half the memory.
  6. Per dollar, the L40S trained the most tokens: 231M tokens per dollar compiled with SDPA, against 131M on the RTX 5090, 100M on the RTX A6000 and 50M on the H100. The H100 was the fastest card (54,022 tokens per second), 1.3x the RTX 5090, at 3.4 times its price.

What to take from it

  • Use SDPA (or FlashAttention) first. That is the large win and costs nothing.
  • Then try torch.compile for memory rather than speed at this size: 23% less activation memory can mean a larger batch or a longer context on the same card.
  • Measure on your GPU. The same compiled model was 9% slower on one card and 3% faster on another.

Caveats

  • One model size (468M), one batch shape, one run per setup, on one day. Larger models, longer sequences or mode="max-autotune" can change the result; I tested the default mode only.
  • Pure bf16 weights and optimizer state, no mixed precision with fp32 master weights.
  • The machines differ in CPU and host; with this batch size the step is GPU-bound, but small differences may come from the host.
  • Prices are on-demand at launch time and change daily. Billed cost for the four machines, which also ran an embedding benchmark, was $0.14 (RTX 5090), $0.11 (L40S, which also ran a model startup test), $0.08 (RTX A6000) and $0.43 (H100).

Reproduce it

  1. Create an API key and add an SSH key (POST /api/v1/ssh-keys).
  2. Launch a machine with POST /api/v1/instances ({"name": "...", "gpu_type": "L40S", "gpu_count": 1, "ssh_key_id": "...", "max_runtime_minutes": 40}), install torch and transformers==4.57.1 with uv.
  3. Run python train.py compile sdpa and the other three combinations, then delete the instance.

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.

Top comments (0)