How to improve my tokens per second?

Research
1

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…

Thank you for reading.

2

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:

  1. 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.
  2. 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.
  3. 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.
  4. 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.