LLM FP8 Fine-Tuning in Production: Amax Scaling, Layer Whitelists, and BF16 Fallback for Stable Convergence

Why FP8 Fine-Tuning Can’t Just Be About “Training Ran”

FP8 is often described as a simple precision toggle—swap BF16 for FP8 and get higher throughput with lower memory usage. This understanding misses the most critical difference between training and inference quantization.

Inference weight quantization mainly cares whether compressed static weights still maintain output quality. In training, however, activations, weights, and gradient distributions change continuously, and scaling factors must be updated dynamically to follow the data and optimization process. A single anomalous batch, a few outliers, inconsistent scales across ranks, or the conversion overhead of a small matrix can all cause FP8’s benefits to vanish, or even produce silent degradation that’s not easily caught by the overall loss.

Therefore, the real goal of FP8 fine-tuning in production isn’t “enable FP8 on all layers,” but to establish an auditable low-precision training strategy:

• Which operators use FP8, and which stay in BF16
• How to choose between E4M3, E5M2, and different scaling recipes
• Whether Amax and Scale are within reasonable ranges
• Whether distributed ranks use consistent scales
• How to locally fall back when saturation, underflow, or capability degradation occurs
• Whether FP8 artifacts pass the same release gates as the BF16 baseline

Core Principles: The Key to FP8 Isn’t 8 Bits, It’s Dynamic Range Management

E4M3 and E5M2 Serve Different Roles

The two common FP8 formats are E4M3 and E5M2:

Format	Exponent Bits	Mantissa Bits	Characteristics	Typical Use
E4M3	4	3	Higher precision, smaller dynamic range	Forward weights and activations
E5M2	5	2	Larger dynamic range, lower precision	Backward gradients

In a typical mixed recipe, NVIDIA Transformer Engine defines the E4M3 + E5M2 combination as Format.HYBRID. But this doesn’t mean all tensors should be forced into the same format—Embedding, Normalization, Loss, LM Head, very small Linear layers, and numerically sensitive modules typically need individual validation and may need to stay in BF16.

Scale Determines Whether Values Fall into the Representable Range

Before converting a high-precision tensor to FP8, a scaling factor is computed based on the tensor’s magnitude. The default form of Transformer Engine’s Delayed Scaling can be summarized as:

scale = (FP8_MAX / amax) / (2 ^ margin)

Where:

• amax: The maximum absolute value of the tensor
• FP8_MAX: The maximum representable value of the target format
• margin: A safety margin reserved for anomalous fluctuations

If the scale is too large, outliers saturate easily; if too small, many small values get pushed near zero. Production governance needs to focus on scaling quality, not just the loss curve.

Delayed Scaling Trades History for Stability

Delayed Scaling doesn’t directly use the scale computed from the current iteration. Instead, it records an Amax History, selects a representative value from a historical window, and uses that for the next iteration.

Common strategy comparison:

Strategy	Behavior	Advantage	Risk
max	Uses the maximum Amax in the window	Reduces risk of sudden overflow	Effective precision drops for normal batches
most_recent	Uses the most recent Amax	Faster response	More sensitive to outlier batches
Custom (quantile/trimmed)	Designed as needed	Flexible adaptation	Requires sufficient replay validation

The length of the Amax History window isn’t always better. A window that’s too long can cause a single historical peak to persistently reduce the effective precision of subsequent tensors; a window that’s too short is prone to batch-to-batch jitter. Experiments should be run based on data distribution, gradient accumulation, sequence length, and training phase.

Distributed Training Must Govern Scale Consistency

FSDP, Tensor Parallel, or other distributed training splits weights and activations across multiple ranks. If each rank independently maintains Amax and Scale, the same logical tensor may use different quantization scales, affecting single-card and multi-card semantic consistency.

Transformer Engine can perform distributed reduction on Amax by default. TorchTitan’s Float8 documentation also emphasizes the need to communicate the global max(abs) in FSDP and TP scenarios. If Amax reduction is disabled for performance, the local Scale must be treated as part of the training state: all ranks must be included in checkpoints, and the correct mapping must be maintained upon restoration.

Engineering Implementation: From Global Toggle to Per-Layer Admission

Step 1: Establish a Reproducible BF16 Baseline

Before enabling FP8, freeze a BF16 baseline:

• Model and tokenizer fingerprints
• Dataset version, sampling order, and chat template
• Optimizer, scheduler, gradient clipping, and random seeds
• Tokens per step, gradient accumulation, and effective batch size
• Loss, gradient norms, throughput, memory, and key capability set results
• Checkpoints and generation samples at fixed steps

Without a reproducible BF16 baseline, FP8’s performance gains and quality changes cannot be attributed.

Step 2: Only Convert Matrix Layers with Clear Benefits

Start with the larger Linear layers in the Transformer, then gradually expand:

1. Attention Q/K/V and output projections
2. MLP up, gate, and down projections
3. After empirical validation, extend to other modules

The initial BF16 retention list should include:

• Embedding and LM Head
• RMSNorm, LayerNorm, and final Loss
• Linear layers too small for conversion overhead to outweigh Tensor Core benefits
• Layers with historically significant Amax fluctuations or prone to gradient anomalies
• Custom operators and unverified fused kernels

The layer whitelist should record the module’s FQN, shape, recipe, and fallback reason as part of the training artifact.

Step 3: Fix the Recipe, Prohibit Silent Drift During Runtime

In production configuration, write the following fields into an immutable manifest:

precision_policy:
  base_dtype: bf16
  fp8_format: hybrid
  scaling: delayed
  amax_history_len: 16
  amax_compute_algo: max
  margin: 0
  reduce_amax: true
  include_modules:
    - "*.self_attn.q_proj"
    - "*.self_attn.k_proj"
    - "*.self_attn.v_proj"
    - "*.mlp.*_proj"
  exclude_modules:
    - "embed_tokens"
    - "lm_head"
    - "*.norm"

When upgrading Transformer Engine, PyTorch, CUDA, cuDNN, GPU model, or parallel topology, you must re-run numerical and throughput replay. Don’t treat the FP8 recipe as a regular hyperparameter independent of the runtime environment.

Step 4: Establish Amax and Low-Precision Anomaly Profiles

Each converted module should output at least the following per-layer metrics:

Metric Category	Specific Metrics
Scaling Status	Current Amax and historical quantiles, Scale and step-to-step change rate
Quantization Quality	Proportion of values near FP8 upper limit, proportion of zero values after quantization
Gradient Health	Input/weight/gradient norms, NaN/Inf/Overflow/skip step counts
Precision Comparison	Relative error between FP8 and BF16 sampled outputs
Kernel Efficiency	FP8 kernel hit rate and fallback rate per layer

It’s recommended to use three levels of alerts:

Level	Trigger Condition	Action
Observe	Scale jitter, rising zero-value ratio, but loss and capability set normal	Log and continuously monitor
Degrade	Saturation across multiple consecutive windows, anomalous gradients, or expanding FP8/BF16 differential	Locally switch to BF16
Stop	NaN, unrecoverable divergence, multi-rank inconsistency, or key capability drops below hard threshold	Pause training and investigate

Step 5: Implement Per-Layer BF16 Fallback, Not a Global Binary Choice

The precision strategy should allow different layers to use different recipes. Transformer Engine supports nesting multiple autocast regions within the same training run. Engineering-wise, maintain a dynamic but auditable fallback table:

precision_overrides:
  "model.layers.17.mlp.down_proj":
    dtype: bf16
    reason: "persistent saturation"
    activated_at_step: 18240
  "model.layers.29.self_attn.o_proj":
    dtype: bf16
    reason: "capability regression isolation"
    activated_at_step: 20100

The fallback action must also record: trigger metrics, thresholds, step, data version, before/after checkpoints, and recovery conclusions. After a local BF16 fallback, observe whether the throughput loss is acceptable and recalculate the overall training ETA.

Step 6: Use Dual-Track Evaluation to Determine “Convergence Equivalence”

FP8 shouldn’t just be compared to its own previous version. At least three tracks are needed:

1. BF16 Control Group: Same data, same token count, and same optimization parameters
2. FP8 Candidate Group: Target recipe and layer whitelist
3. FP8 Fallback Group: Keep high-risk layers in BF16 to pinpoint degradation sources

Comparison dimensions:

• Loss at the same step and same training tokens
• Gradient norm and update norm distributions
• Key capability set, format adherence, refusal, and safety metrics
• Long/short sequences, different languages, and hard-case buckets
• Training throughput, peak memory, communication overhead, and effective tokens/second
• Metric continuity after checkpoint restoration

Don’t treat a single final score as proof of equivalence. FP8 may have a similar overall average score to BF16 but show local degradation on code, multilingual, long-context, or low-frequency format tasks.

Applicable Scenarios

Scenarios suitable for FP8 fine-tuning:

• Platforms with H100, H200, Blackwell, or other hardware capable of low-precision computation
• SFT, continued pre-training, or large-scale alignment training where large Linear GEMMs are the primary cost
• Teams with a stable BF16 baseline, capability evaluation sets, and per-layer observability
• Sufficiently large training token counts to amortize conversion, scale computation, and compilation costs
• Need to combine with FSDP, TP, or other distributed strategies, with the ability to govern scale communication

Scenarios to approach with caution:

• LoRA fine-tuning with very small data volumes and short training times
• Many small matrices, custom operators, or frequent dynamic shapes
• No BF16 baseline, inability to reproduce data order, or lack of capability regression sets
• Frequent changes in GPU, CUDA, framework, and kernel versions in the runtime environment
• Only allowed to observe the final loss, unable to collect per-layer numerical metrics

Common Misconceptions

Misconception 1: Hardware Support for FP8 Automatically Means Acceleration

Hardware support is only a necessary condition. Small matrices, frequent casts, scale kernels, communication, and unfused operators can negate the benefits. Measure end-to-end tokens/second per module, not just a single GEMM micro-benchmark.

Misconception 2: If Loss Hasn’t Diverged, Quality Is Fine

Loss is not sensitive to local capability degradation. Format adherence, multilingual, long-sequence, and code capabilities need independent regression, and per-layer Amax and gradient metrics cannot be omitted.

Misconception 3: Longer Amax History Means More Stability

A longer window can reduce jitter, but historical outliers will continuously affect the scale, causing most normal values to lose effective precision. The window length must be validated based on data and training phase.

Misconception 4: All Ranks Computing Scale Independently Is Faster and Won’t Affect Results

Local scales can break distributed single-device semantics. If Amax reduction is disabled, you must explicitly accept numerical differences and fully save each rank’s local quantization state.

Misconception 5: Anomalies Only Allow a Full Model Fallback to BF16

Per-layer whitelists and local fallbacks can usually retain most of the performance gains. What’s truly needed is a locatable, auditable precision strategy, not a global toggle.

Production Launch Checklist

Environment and Artifacts

• GPU, driver, CUDA, cuDNN, PyTorch, and Transformer Engine versions are fixed
• FP8 recipe, layer whitelist, and exclusion list are under version control
• Model, tokenizer, dataset, and chat template fingerprints are recorded
• FP8-related Scale/Amax state can be correctly restored with checkpoints
• Restoration semantics for different world sizes and parallel topologies are validated

Numerical and Quality

• BF16 control group is reproducible
• Amax, Scale, zero-value rate, saturation rate, and gradient norms are observable per layer
• NaN, Inf, Overflow, and anomalous skip steps have hard gates
• Core capability set has completed FP8/BF16 dual-track differential analysis
• No significant regression in length, language, domain, and difficulty buckets
• Local BF16 fallback process has been rehearsed

Performance and Operations

• Statistics are based on end-to-end effective tokens/second, not single-operator theoretical peak
• Scale computation, cast, compilation, and distributed Amax communication costs are accounted for
• Small matrices and modules with insufficient benefit are automatically or manually filtered
• Post-failure recovery does not reset or misalign Scale history
• Release manifest can reproduce the complete precision strategy
• Training stop, degrade, and recovery thresholds have been approved

References

1. NVIDIA Transformer Engine — Using FP8 and FP4 with Transformer Engine
2. NVIDIA Transformer Engine — Common API: DelayedScaling / Current Scaling / MXFP8
3. PyTorch AO — torchao.float8 API
4. PyTorch TorchTitan — Float8 Training
5. Micikevicius et al. — FP8 Formats for Deep Learning
6. Liang et al. — TorchTitan: One-stop PyTorch Native Solution for Production Ready LLM Pre-training