Research map / Architectures and efficiency / Quantization
µnit Scaling: Simple and Scalable FP8 LLM Training
Read the paperCheck your idea against 108,967 papers
AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.
Proposes µnit Scaling (µS), a static-scaled FP8 LLM training recipe combining Unit Scaling and µP with Res-Post-LayerNorm and fixed residuals, enabling zero-shot LR/weight-decay transfer and all-hidden-layer FP8 training up to 13B.
Problem
FP8 LLM training is hard due to limited range and resolution; existing methods need dynamic scale factors, many extra hyperparameters, or keep some matmuls in BF16, and hyperparameter transfer is unreliable in low precision.
Approach
Combine Unit Scaling and µP into a simplified scheme: unit-variance init, static 1/sqrt(fan_in) GEMM scaling, Res-Post-LayerNorm placement, fixed-τ variance-preserving residuals, and LR scaling of hidden layers by sqrt(d_base/d_new), so all hidden linear layers run in FP8 without dynamic scaling.
Key results
- µS FP8 at 1B-13B matches or approaches BF16 final train loss (e.g., 13B: BF16 µS 2.108 vs FP8 µS 2.119) while SP FP8 with TransformerEngine diverged at 13B.
- µS hyperparameters (η, λ) transfer from width 256 to width 5120 (20x), with optimal η stable across widths 256-8192 in sweeps.
- FP8 µS training is 25-33% faster than BF16 and 1-6% faster than FP8 TransformerEngine on 64 H100 GPUs.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Train with µS instead of SP in FP8 | 1B–13B LLMs · Final train loss | cross-entropy loss: 2.588 vs 2.590; 2.400 vs 2.390; 2.231 vs 2.226; 2.211 vs 2.119 | Helped |
| Train with µS FP8 instead of SP FP8 | 13B LLM · 13B ARC Easy (3-shot) | accuracy: 35.7% -> 69.7% | Helped |
| Train with µS FP8 instead of SP BF16 | 13B LLM · 13B GSM8K (0-shot) | accuracy: 8.7% -> 10.9% | Helped |
| Use µS static scaling and fused operators instead of BF16 training | 1B, 3B, 7B, and 13B LLMs · Distributed training | throughput: 25-33% higher throughput | Helped |
| Use µS static scaling and fused operators instead of TransformerEngine FP8 | 1B, 3B, 7B, and 13B LLMs · Distributed training | throughput: 1-6% faster | Helped |
| Use fixed residual modification instead of standard residual summation (ablation) | µS deep transformers · Training convergence | training loss | Helped |
| Use fixed residual modification instead of running-mean residual modification (ablation) | µS deep transformers · Training convergence | training loss | Helped |
| Move LayerNorm from the start to the end of each residual branch (ablation) | µS 100-layer transformers · Training convergence | training loss | Helped |
| Use Res-Post-LayerNorm instead of Pre-LayerNorm (ablation) | 100-layer transformer · Training convergence | training loss | No clear effect |
| Apply square root to softmax attention scores (ablation) | Self-attention · Training attention output variance over sequence position | output variance consistency | Mixed |
| Use µS instead of Unit Scaling (ablation) | 7B FP8 LLM · Training convergence | training loss | Helped |
| Use ReLU instead of GELU or SiLU (ablation) | 4-layer FP8 model · Activation underflow | maximum FP8 underflow: GELU: up to 30%; SiLU: up to 7%; ReLU: 0.04% | Helped |
Surprising findings
- SP FP8 models trained with TransformerEngine frequently spiked and diverged at 13B, and lowering weight decay did not fix it.
- Unit Scaling without µS modifications diverged early in 7B FP8 training.
- GELU/SiLU cause high FP8 activation underflow compared with ReLU, yet GELU was used because it gave lower final loss.
Limitations
- Final quality comparisons are based on a single training run per configuration at compute-optimal token budgets, with modest downstream scores near chance for small models.
- Evaluated only on decoder-only LLMs up to 13B parameters with Lion optimizer and a single hardware setup (H100).
- Uses GELU rather than ReLU despite higher FP8 underflow, leaving activation-function choice for FP8 open.
Experimental setup
Datasets: Databricks pretraining corpus (not specified)
Benchmarks: Databricks Model Gauntlet ARC Easy ARC Challenge HellaSwag WinoGrande SQuAD TriviaQA GSM8K LAMBADA PIQA
Metrics: final training cross-entropy loss accuracy training throughput (TFLOPS/tokens per second)
Baselines: Standard parametrization BF16 Standard parametrization FP8 (TransformerEngine) µP u-µP Unit Scaling
Builds on: Unit Scaling µP u-µP TransformerEngine Lion
Closest papers by meaning
- Towards Fully FP8 GEMM LLM Training at Scale (NeurIPS 2025 · similarity 0.91)
- Scaling FP8 training to trillion-token LLMs (ICLR 2025 · similarity 0.91)
- Optimizing Large Language Model Training Using FP4 Quantization (ICML 2025 · similarity 0.89)
- Scaling Laws for Precision (ICLR 2025 · similarity 0.88)
- Scaling Inference-Efficient Language Models (ICML 2025 · similarity 0.88)
- Scaling Laws Meet Model Architecture: Toward Inference-Efficient LLMs (ICLR 2026 · similarity 0.87)
Writing a paper in this area? scime draft is a free LaTeX editor with AI included, real-time co-authors and a check of your references and venue rules before you submit. It also imports OpenAI Prism projects.