FP8 Training & Inference

Table of Contents

1. Learning Objectives
2. Notation
3. Core Intuition
4. FP8 Formats: E4M3 vs E5M2
5. FP8 Training Recipe
6. Per-Tensor Scaling for FP8
7. Delayed Scaling
8. FP8 for Inference
9. Quality Analysis
10. Hardware Support
11. Common Pitfalls
12. Summary
13. Exercises

Learning Objectives

1. Explain the FP8 E4M3/E5M2 format choice for forward/backward passes.
2. Describe the per-tensor scaling mechanism and why it's needed.
3. Derive the delayed scaling strategy for overlapping compute and scale calculation.
4. Compare FP8 training quality to BF16 across model sizes.
5. Identify when FP8 training provides meaningful speedup.

Notation

• E4M3: 1 sign + 4 exponent + 3 mantissa
• E5M2: 1 sign + 5 exponent + 2 mantissa
• $s_X$ — per-tensor scale for tensor $X$

Core Intuition

FP8 is the sweet spot between BF16 (safe but no speedup beyond memory) and INT8 (fast but hard to use for training). Being a floating-point format, FP8 naturally handles varying magnitudes without explicit quantization/dequantization — just scale and cast. The H100 provides native FP8 tensor cores at 2x the throughput of BF16, making FP8 training practical with minimal code changes.

FP8 Formats: E4M3 vs E5M2

E4M3 (forward pass — weights and activations):

• Range: $[-448, 448]$ (limited but sufficient for bounded forward values).
• Precision: 3 mantissa bits → 8 values per power-of-2 interval.
• Why forward: Weights and activations have bounded, predictable ranges.

E5M2 (backward pass — gradients):

• Range: $[-57344, 57344]$ (much larger).
• Precision: 2 mantissa bits → 4 values per interval.
• Why backward: Gradients can spike (especially early in training or at loss spikes). Need range more than precision.

Comparison:

FP8 Training Recipe

Standard mixed-precision (BF16):

• Forward: BF16 weights × BF16 activations.
• Backward: BF16 gradients.
• Optimizer: FP32 master weights + FP32 optimizer states.

FP8 mixed-precision:

• Forward: E4M3 weights × E4M3 activations.
• Backward: E5M2 gradients.
• Optimizer: FP32 master weights + FP32 optimizer states (unchanged).

Three GEMMs per transformer layer:

1. Forward: $\mathbf{Y} = \mathbf{X} \times \mathbf{W}$ — both in E4M3.
2. Weight gradient: $d\mathbf{W} = \mathbf{X}^T \times d\mathbf{Y}$ — $\mathbf{X}$ in E4M3, $d\mathbf{Y}$ in E5M2.
3. Activation gradient: $d\mathbf{X} = d\mathbf{Y} \times \mathbf{W}^T$ — $d\mathbf{Y}$ in E5M2, $\mathbf{W}$ in E4M3.

Per-Tensor Scaling for FP8

Problem: FP8 has limited range ($[-448, 448]$ for E4M3). If tensor values exceed this, they overflow.

Solution: Scale each tensor to fit within FP8 range before casting:

During matmul: The scales are factored out:

The matmul runs in FP8; output is rescaled to BF16/FP32.

Delayed Scaling

Problem with just-in-time scaling: Must compute $\max(|\mathbf{X}|)$ before the matmul. This adds latency (extra kernel launch, synchronization).

Delayed scaling: Use the maximum from the PREVIOUS iteration:

Why it works: Tensor statistics change slowly between iterations. The scale from step $t-1$ is almost always valid for step $t$.

Safety: If the scale is too large (values overflow), the matmul produces infinity/NaN → detect and recompute with corrected scale. Overflow is rare (less than 0.01% of steps).

Benefit: Overlaps scale computation with other work; no pipeline bubble.

FP8 for Inference

Simpler than training: Only forward pass; no gradient scaling needed.

W8A8 FP8: Both weights and activations in E4M3.

• Weights: Quantized offline (static scales per tensor or per channel).
• Activations: Dynamic per-tensor scaling (compute max on-the-fly or use calibration).

Speedup vs INT8: Similar throughput on H100 (both 2x over BF16). FP8 is easier to implement (no quantize/dequantize kernels needed — just cast).

Vs INT8 quality: FP8 E4M3 handles small values better (log-spaced levels vs uniform INT8). Slightly better for weight distributions with many near-zero values.

Quality Analysis

Training convergence (GPT-3 scale):

• BF16 baseline: 100% quality.
• FP8 (E4M3/E5M2 with per-tensor scaling): 99.5-100% quality.
• Quality gap decreases with model size (larger models are more robust).

Failure cases:

• Very early training (random weights, unstable gradients): FP8 can diverge. Use BF16 warmup (1000 steps) then switch.
• Fine-tuning with very small learning rates: FP8 gradient precision may be insufficient. Keep BF16 for gradients.

Recommendation: FP8 for pre-training at scale (above 7B). BF16 for fine-tuning and small models.

Hardware Support

Software: NVIDIA Transformer Engine (TE), PyTorch FP8 (via torch.float8), DeepSpeed FP8, Megatron-LM.

Common Pitfalls

Pitfall 1. Using FP8 training on A100 GPUs. A100 doesn't have FP8 tensor cores — casting to FP8 and back adds overhead with no speedup.

Pitfall 2. Applying FP8 without per-tensor scaling. Raw tensor values often exceed the E4M3 range (448), causing widespread overflow and NaN loss.

Pitfall 3. Using E4M3 for gradients. Gradient magnitudes vary wildly; E4M3's limited range causes frequent overflow. Always use E5M2 for gradients.

Summary

• FP8 training: 2x speedup on H100 with near-BF16 quality.
• E4M3 (forward): More precision, bounded range → weights and activations.
• E5M2 (backward): More range, less precision → gradients.
• Per-tensor scaling: Required to fit tensor values into FP8 range.
• Delayed scaling: Use previous iteration's scale to avoid latency.
• Best for: Large-scale pre-training on H100; less benefit for fine-tuning.

Exercises

Exercise 1. Compute the maximum value representable in E4M3 and E5M2. Verify the 128x range difference.

Exercise 2. For a weight matrix with $\max(|W|) = 0.8$: compute the per-tensor scale for E4M3 and the effective precision (step size near 0).

Exercise 3. Estimate the throughput improvement of FP8 training for a 70B model on 8x H100: account for both matmul speedup and the memory bandwidth savings.

Exercise 4. Design the FP8 training recipe for a ViT-H model. Which layers benefit most from FP8? Which should stay in BF16?

Exercise 5. Analyze the failure mode of delayed scaling: if activation magnitude increases 10x between iterations, what fraction of values overflow in E4M3?