Format-Aware Fusion for Fast FP4 Pretraining

• The paper examines native FP4 training.Source: REPORT FOCUS · p. 1
• Format-aware fusion is a central focus.Source: REPORT FOCUS · p. 1
• It reports measured speed–quality trade-offs for the Llama 3 model family.Source: REPORT FOCUS · p. 1

1 Introduction

• 37.9K tokens/s/GPU is the fastest custom route; MXFP4 with row-gradient SR and fixed H32 reaches 37.2K and ends 2.11% above raw BF16 loss.TE-native and four-final-BF16-block trajectories end 1.34% and 0.87% above raw BF16 at 27.6K and 27.1K tokens/s/GPU, respectively.Source: §1 · p. 2
• The per-operand FP4 contract uses upward-rounded MXFP4 scales, adaptive finer scales for 32 × 32 weight tiles, and separate FP32 row- and column-tile factors for CTA-local NVFP4.It applies stochastic rounding only to the data-gradient multiplication and paired fixed-sign Hadamard preconditioning only to the weight-gradient multiplication.Source: §1 · p. 2
• Custom producers finish each route’s scale calculation and emit the layouts and saved state needed by forward and backward; the profitable fusion boundary differs by format.MXFP4 uses a local 32-value scale, global NVFP4 waits for a tensor-wide maximum, and CTA-local NVFP4 calculates two outer scales inside the matrix-multiplication tile.Source: §1 · p. 2
• Evaluation combines matched throughput probes, trajectories through 160B tokens, checkpoint-matched validation loss, and downstream evaluation of complete recipes.Source: §1 · p. 3
• The evaluated hybrids use depth-wise format assignment or operand-wise assignment motivated by data-gradient error analysis on identical saved tensors.Source: §1 · p. 3

2.1 Related work

• The systems contribution treats quantization, physical layouts, consuming matrix multiplication, and retained backward state as one execution boundary.The related work identifies scale format, block geometry, tensor statistics, and contraction as factors in FP4 behavior.Source: §2.1 · p. 4

2.2 The same four-bit values, three different scale contracts

• All three custom routes use E2M1 operands, but different scale domains determine quantization behavior and when producers can release their outputs.MXFP4 completes local 32-value scales, global NVFP4 depends on tensor amax, and CTA-local NVFP4 completes outer scales within its owning tile.Source: §2.2 · p. 4

2.2.1 MXFP4 encode- and decode-centric scale selection

• Reported MXFP4 routes round activation and gradient scales upward, while 2D32 weight tiles may use a finer scale when aT /sT < 0.87.The weight choice halves quantization spacing while accepting bounded upper-tail clipping.Source: §2.2.1 · p. 6, Figure 1 · p. 7

2.2.2 Global and CTA-local NVFP4 outer scales

• Global NVFP4 waits for tensor amax before dependent packing, whereas CTA-local NVFP4 derives row and column factors within each owning CTA.The local factors remain fixed over the K reduction and are applied together in the epilogue.Source: §2.2.2 · pp. 7–8

2.3 Why forward, data gradient, and weight gradient need different views

• Forward, data-gradient, and weight-gradient contractions consume different row- and column-facing representations of shared tensors.A saved gradient therefore needs a row representation for Dgrad and a column representation for Wgrad.Source: §2.3 · p. 9, Figure 4 · p. 10

2.4 Two-dimensional weights: numerical consistency versus implementation

• Fused 2D weight producers emit one numerical encoding in both physical orientations, so Fprop and Dgrad decode the same weight.Independent one-dimensional quantizers can instead produce different numerical weights for the two contractions.Source: §2.4 · pp. 10–11

2.5 Stochastic rounding and signed Hadamard preconditioning

• Custom routes apply stochastic rounding only to row-oriented dY for Dgrad and use paired fixed-sign Hadamard transforms only for explicitly named Wgrad paths.Adjacent-bin SR removes deterministic bias only before saturation; fixed signs do not carry a concentration guarantee requiring independently sampled signs.Source: §2.5 · pp. 12–13

2.6 Reference NVFP4 recipe versus the routes evaluated here

• The NVIDIA NVFP4 recipe is a bundle of choices, while the evaluated TE-native comparator places all 32 body blocks in NVFP4.The separate four-final-BF16-block route is NVIDIA-inspired, not a reproduction of NVIDIA’s per-linear sensitivity list.Source: §2.6 · p. 13

3.1 Fuse the complete producer–consumer boundary

• Fusion removes an HBM round trip only when a producer emits every scale, layout, and saved value required by the next operation.A boundary includes work only when the next stage can consume its output directly; one generic schedule cannot capture the differing MXFP4, global NVFP4, and CTA-local NVFP4 contracts.Source: §3.1 · pp. 13–14

3.2 What is fused in a Llama block

• The gated FFN path pairs W1 and W3, then emits the SiLU-gated output as the FP4 input to W2 instead of writing and rereading BF16 intermediates.Backward combines the W2 data gradient with the SiLU derivative and produces W1/W3 gradient representations while retaining exact BF16 values where autograd requires them.Source: §3.2 · p. 14
• QKV production creates format-specific projection inputs, while backward combines inverse layout, RoPE, QKV gradient production, and adjacent RMSNorm reduction when dependencies align.Residual–RMSNorm retains the exact residual and row statistics while producing the next quantized representations.Source: §3.2 · p. 14

3.3 Depth-wise and operand-wise hybrids

• The 27/5 depth hybrid assigns CTA-local NVFP4 to blocks 1–27 and MXFP4 to blocks 28–32, without mixing formats within a linear.Both formats use orientation-consistent 2D weights and row-only gradient-data SR, with no Hadamard preconditioning.Source: §3.3 · p. 14
• CTA-local NVFP4 had lower Dgrad error than MXFP4 in every measured layer and component, motivating CTA-local Dgrad with MXFP4 retained for Fprop and Wgrad.Each linear’s producer emits both the CTA-local row representation and MXFP4 column representation from one BF16 gradient-tile load.Source: §3.3 · pp. 14, 16
• The operand hybrid compares plain-H16 mixing with fixed-H32 signed transforms on X and dY, leaving the stored weight untransformed and keeping the other numerical choices fixed.The paired Wgrad transform is applied to both column operands; neither variant changes Fprop, CTA-local Dgrad, row SR, output projection, or loss.Source: §3.3 · p. 16

4.1 Model, data, and training horizon

• The main workload is an 8B-parameter Llama-family decoder trained at sequence length 8,192 and global batch 512 for 160B tokens, ending at step 38,147.Unless specified otherwise, each named low-precision route applies to every body linear in all 32 transformer blocks.Source: §4.1 · p. 16

4.4 Training and checkpoint evaluation

• Training losses are compared at matched optimizer steps, while validation uses the same fixed external stream at every checkpoint and downstream scores use exact step-38,000 checkpoints with BF16 inference.The validation stream is called external rather than held out because zero overlap with all historical training trajectories cannot be proven; each long-run recipe has one trajectory, so seed-to-seed variation is not estimated.Source: §4.1 · p. 16, §4.4 · p. 19

4.2 Recipe ledger

• The recipe ledger distinguishes scale, rounding, and operand choices by contraction, while all custom routes retain BF16 parameters, FP32 optimizer state, and a BF16 output projection.“Row SR” applies only to the row-oriented gradient consumed by Dgrad; paired Wgrad Hadamard transforms act on X and dY, not the stored weight.Source: §4.2 · pp. 16, 18, Table 5 · p. 18

4.3 End-to-end performance measurements

• Throughput probes match accelerator type, sequence length, global batch, seed, BF16 output projection, compiled CE, warmup, metric window, and per-GPU timing, but compare complete configurations.Routes retain their selected local batch, activation-checkpoint policy, and distributed schedule; BF16 MFU is effective BF16-equivalent model work divided by peak BF16 throughput, not literal FP4 Tensor Core occupancy.Source: §4.3 · pp. 18–19

5.1 What throughput do the end-to-end routes achieve?

• 37,876.8 tokens/s/GPU for custom MXFP4 with row-SR compares with 18,812.0 for BF16 and 27,622.0 for TE-native NVFP4 in all blocks.The Table 6 panel reports cross-device medians over steps 40–100; every row uses sequence length 8,192, global batch 512, BF16 output projection, and compiled regular CE.Source: §5.1 · p. 19, Table 6 · p. 20

5.2 How closely do the recipes track BF16 training loss?

• At step 38,140, the four-final-BF16-block TE route has 0.87% excess over BF16, followed by TE-native at 1.34%; custom MXFP4 with row-SR and fixed H32 has 2.11% excess.The fixed-H32 operand hybrid ends 1.85% above the pure MXFP4-H32 route and 4.00% above BF16; the H32 operand analysis found lower CTA-local NVFP4 Dgrad error in all 288 layer–component comparisons, with geometric-mean EMX/ElocalCTA ratio 1.292.Source: §5.2 · p. 23, Table 7 · p. 23

5.3 Validation loss

• At step 38,000, MXFP4 with fixed H32 has the lowest observed validation NLL at 2.422804, versus 2.430206 for BF16 and 2.437475 for the four-final-BF16-block TE route.The fixed-H32 operand hybrid records NLL 2.463842, while the 27/5 depth hybrid records 2.538116, or 4.440% higher than matched BF16.Source: §5.3 · p. 23, Figure 9 · p. 24

5.4 How do the routes compare on downstream tasks?

• Downstream task winners differ: BF16 leads MMLU, the fixed-H32 operand hybrid leads WinoGrande, and MXFP4+row-SR+fixed H32 leads HellaSwag and ARC-Challenge.Relative to pure MXFP4-H32, the operand hybrid differs by −1.66, −2.02, +0.79, and −2.39 points on MMLU, HellaSwag, WinoGrande, and ARC-Challenge, respectively; these are observations from one checkpoint per trajectory.Source: §5.4 · p. 25, Table 8 · p. 25, Figure 10 · p. 26

6 Discussion

• 27.6K and 27.1K tokens/s/GPU accompany TE-native and four-final-BF16-block routes, ending 1.34% and 0.87% above raw BF16 loss, respectively.Custom global NVFP4 reaches 32.5K tokens/s/GPU and ends 2.69% above BF16; the panel does not isolate scaling cost from implementation and batch geometry.Source: §6 · p. 27
• 98.34% of row-SR-only MXFP4 throughput accompanies fixed-H32 MXFP4, with raw endpoint excess moving from 6.53% to 2.11% and validation NLL from 2.593146 to 2.422804.Fixed-H16 CTA-local is slower than CTA-local without the transform and has higher endpoint and validation loss; these are complete-trajectory comparisons, not isolated Hadamard tests.Source: §6 · p. 27
• Neither tested hybrid beats MXFP4 with row-SR and fixed H32 in speed, endpoint loss, or validation loss, despite lower CTA-local Dgrad error in all 288 layer–projection comparisons.The completed hybrid trajectories do not preserve that operator-level advantage.Source: §6 · p. 27

7 Conclusion and Further Work

• MXFP4 with row-SR and fixed-H32 Wgrad preconditioning reaches 37.2K tokens/s/GPU, 86.3% BF16 MFU, and 2.11% raw endpoint-loss excess, with the lowest observed validation NLL.Its downstream ranking remains task dependent.Source: §7 · p. 27
• BF16 scores highest on MMLU, the fixed-H32 operand hybrid on WinoGrande, and MXFP4 with row-SR and fixed H32 on HellaSwag and ARC-Challenge.Late training loss and downstream accuracy do not have a monotonic ordering in this panel.Source: §7 · p. 27
• Neither tested hybrid improves the strongest pure custom route’s combined speed–quality point; the conclusion proposes wider unsigned block scales with CTA-local scheduling or delayed tensor scaling as next steps.UE5M3 FP4 pairs E2M1 payloads with an unsigned E5M3 block scale, and a recipe using it was supported on a separate model and training setup [25].Source: §7 · p. 27

Generative-AI Disclosure

• Generative AI tools assisted with code development, experiment orchestration, analysis, figure generation, and report drafting and editing.Source: Generative-AI Disclosure · p. 28
• The author reports reviewing and verifying the code, experiments, numerical results, figures, citations, and final text.Source: Generative-AI Disclosure · p. 28
• The author takes responsibility for the report’s content.Source: Generative-AI Disclosure · p. 28

A Implementation and Validation Notes

• Validation checks outputs and gradients against a higher-precision reference; explicit buffer ownership preserves asynchronous lifetimes, while per-GPU timing reflects the slowest worker.Source: A Implementation and Validation Notes · p. 32

A.1 Weight orientation and the fixed sign pattern

• The custom global NVFP4 weight encoding is exposed in both consumer layouts, following the orientation-consistency rule also used by native 2D32 MXFP4.Source: §A.1 · p. 32
• H32 repeats the fixed 16-sign diagonal across two halves; paired transforms preserve the unquantized Wgrad dot product, but the vector’s optimality is untested.Source: §A.1 · p. 33
• Distributed workers checkpoint stochastic-rounding seeds and counter offsets so resumed runs do not reuse random-number draws.Source: §A.1 · p. 33
• Table 9 reports three compiled-CE B300 production snapshots using BF16 output projection; differing batch geometries and measurement windows mean it does not isolate an implementation change.BF16 MFU uses the B300’s 2,250 TFLOP/s/GPU dense BF16 peak [30].Source: §A.1 · p. 33, Table 9 · p. 33

B Low-Precision Output Head: Negative Result

• Both selective low-precision output-head variants were slower than matched BF16 controls and failed to reproduce control loss; MXFP8 Wgrad also changed gradients.The preserved record does not retain a publication-ready error magnitude, so the numerical differences cannot be ranked by size.Source: Table 10 · p. 33, B Low-Precision Output Head: Negative… · p. 34

C.1 Report inputs

• Custom global NVFP4 has 74 exact post-step-19,000 loss observations and a 24,710–33,000 gap; its late LOWESS curve is visual only, not used for reported endpoints or statistics.Source: §C.1 · p. 34

C.2 Checkpoint and evaluator contract

• The earlier evaluator omitted training RoPE scaling and was excluded; corrected validation uses a fixed external stream, not a held-out split, because historical-training overlap cannot be ruled out.Source: §C.2 · pp. 34–35

C.3 Historical execution identity

• Historical identity is run-bound: global NVFP4 uses orientation-consistent 2D16 weights, while completed MXFP4 and resumed CTA-local runs have manifest-verified code or segments.Source: §C.3 · p. 35

C.4 Controlled Dgrad diagnostic

• CTA-local NVFP4 had lower Dgrad error energy in all 288 comparisons, with equal-tensor geometric-mean EMX/ElocalCTA of 1.292.The capture spans three sequence-2,048 batches and QKV, attention-output, and FFN projections in all 32 blocks.Source: §C.4 · p. 36

C.5 Excluded experiments

• The custom global NVFP4/MXFP4 depth hybrid diverged and retained only 16 of 32 terminal shards, while the four-final-BF16-block TE route and fixed-H32 operand hybrid completed 160B-token trajectories.The completed trajectories have resumable final checkpoints and are included in endpoint panels.Source: §C.5 · p. 36
• Diagnostic screens are excluded: output-head configurations showed no stable model-level speedup, and a late CTA-local transform screen under changed runtime was not a causal comparison.Superseded validation records are also excluded because of omitted training-time RoPE scaling or sensitivity to contiguous source-data regions.Source: §C.5 · p. 36

C.6 Code and data availability

• Code, pinned kernels, recipes, evaluation tools, and figure builders are released, but training corpora, tokenizer assets, model checkpoints, and checkpoint metadata are not redistributed.The repository provides hash-bound input schemas and public recipes for binding authorized local copies.Source: §C.6 · p. 36