I am now pretraining a 1B model (Model: ~977M params, 24 layers, 1792 hidden, 16 heads, 112 head dim, SwiGLU, RoPE, weight-tied embeddings) from a RTX 5080 16gb VRAM, I ran it on WSL Ubuntu and also enables gradient checkpoints, BF16 autocast, 8-bit AdamW, Pytorch SDPA, and now my micro batch size is 1, gradient accumulation is 16, context is 4096; I tried to have larger micro batch but it will cause OOM; now I have a 8000-9000 tokens / second in average. Seeking for ways to make the training faster…
You’re at roughly 65 TFLOPS of useful compute (about 6 × params FLOPs per token, plus the recompute from gradient checkpointing), which leaves headroom: the LLMQ paper reports around 51% MFU on consumer GPUs with an optimized stack. In order of expected payoff:
torch.compile on the model. It fuses RMSNorm, RoPE and SwiGLU into fewer kernels and was worth about 1.5x in published single-GPU training benchmarks.
Liger Kernel’s fused linear + cross-entropy (plus its RMSNorm/RoPE/SwiGLU kernels). It never materializes the full 4096 x vocab logits, which is often the biggest memory hog at your size. The freed VRAM is the real win: you can drop gradient checkpointing or only checkpoint some layers, which removes roughly a third of the extra compute, and then try micro batch 2 or more.
FP8 matmuls (torchao float8 or LLMQ). Blackwell consumer cards support it, but compare the loss curve against bf16 for a few thousand steps before trusting it for a long run.
Housekeeping: confirm SDPA is actually using the flash backend, keep the dataset on the WSL ext4 disk instead of /mnt/c, pre-tokenize and pack sequences so there’s no padding, and use a fused optimizer.
Run the PyTorch profiler for a few steps first to see whether you’re bound by memory-bound ops, recompute or data loading, so you apply the right fix.
Also worth knowing: the big matmuls are compute-bound, but everything between them (norms, RoPE, SwiGLU, softmax, logits) is memory-bandwidth-bound, and on a 16 GB card with ~960 GB/s that’s where a lot of your step time goes. That’s why kernel fusion (compile, Liger) gives such large gains, and why FP8 helps beyond just faster matmuls.
Thank you! That breakdown of memory-bandwidth vs compute-bound actually makes perfect sense. It explains why everything feels so bottlenecked.
I actually tried a couple of these on my setup (RTX 5080, WSL, PyTorch 2.14 + CUDA 13). torch.compile basically died because of graph breaks with my custom checkpointing. Liger was even worse—I hit hard CUBLAS and illegal memory access errors on CUDA 13, even on the nightly build. So I just reverted to plain eager mode, sadly. FP8 sounds terrifying mid-run, so I’m definitely skipping that for now.
Data is already on ext4 and pre-packed, but I’ll double-check my SDPA backend right now.
Since Liger is a dead end for my stack, have you messed around with torchao’s fused kernels for norms/SwiGLU on newer Torch versions? Wondering if that’s a safer path. Thanks again for the pro breakdown!
Torchao is probably not that path — its feature set is quantization, sparsity, and FP8 training, not drop-in fused RMSNorm/SwiGLU kernels, so it leads back toward the FP8 route you already ruled out. Two routes stay in stock PyTorch: chunk the loss — run the LM head and cross-entropy over slices of the sequence (512 tokens at a time, each chunk weighted by its token count) so the full 4096 × vocab logits and their gradient are never materialized; and if your checkpointing is the reentrant variant, try use_reentrant=False, the variant torch.compile traces far more cleanly. Check both on one fixed batch: the chunked loss should match your current loss to floating-point tolerance, and the peak-VRAM drop plus a compile run without graph-break warnings tells you whether the freed memory buys back a checkpointed layer.
One thing worth checking is whether your model is running on the GPU or falling back to the CPU. You can also try a smaller batch size, use a lower-precision format like FP16 or BF16 if your hardware supports it, and make sure you’re using an optimized attention implementation. Benchmark with the same prompt and generation length each time so you can tell which change actually improves tokens per second.