Tool: evaluate layer-wise numerical-error propagation

[jlamypoirier]

Summary

• New tools/evaluate_precision.py — inherits PretrainedGPTModelConfig (so model: and pretrained: are real typed Config fields) and adds variants:, output_dir:, and a few optional knobs. Runs a fp32 reference plus one trainer invocation per variant in-process; captures per-layer forward activations and input gradients via the standard tensor-logs pipeline; emits per-tensor RMS / max diffs as a console table + precision_report.json.
• Variants aren't dtype-only: each is a flat dict of dotted-path overrides (same syntax as Fast-LLM CLI key=value args) so a variant can sweep any config knob — attention implementation, optimizer dtype, fused vs unfused, etc.
• Per-variant trainer configs are built with TrainerConfig.get_subclass(...).from_dict(base, fp32_dtypes, variant_updates, tool_overrides). Tuple-keyed updates compose in precedence order: forced fp32 → variant overrides (which can re-override fp32) → tool-required debug-logging overrides (which always win).
• Training, optimizer, and data sections of the trainer config are hardcoded inside the tool (single iteration, no checkpoint save, random tokens, LR 0, fp32 optimization dtype). Only knobs that affect the propagation measurement are user-facing: model, pretrained, variants, output_dir, num_samples, micro_batch_size, sequence_length.
• Moves compare_tensor_logs.py from tests/utils/ into fast_llm/engine/config_utils/ so it's importable from tools/, and factors a _compute_diff helper out of CompareConfig.compare_tensors so the tool can extract numbers for every tensor — not only those that breach a tolerance. Three test callers updated; behaviour unchanged.
• Fills in the HF metadata allowlist (fast_llm/engine/checkpoint/huggingface.py) with the generic PretrainedConfig keys newer transformers versions serialize: generation defaults, encoder-decoder flags, family markers, torchscript, is_decoder, etc. Without this, loading any modern HF Llama checkpoint trips the coverage walker. None are architecture knobs Fast-LLM consumes.

Usage

python -m tools.evaluate_precision -c tool.yaml

pretrained:
  path: /path/to/local/hf/snapshot
  format: llama
output_dir: /tmp/precision_eval
variants:
  bf16:
    model.distributed.compute_dtype: bfloat16
  bf16_sdpa:
    model.distributed.compute_dtype: bfloat16
    model.base_model.decoder.block.mixer.implementation: sdpa

Fast-LLM's HF loader reads weights from a local directory, so HF Hub IDs need to be snapshot_download'd first. model: and pretrained: can also be combined — pretrained provides architecture+weights, model: overrides individual fields.

Test plan

• Cluster smoke test on a real HF checkpoint (SmolLM2-135M, snapshot via huggingface_hub). Reference fp32 + bf16 variant ran end-to-end; per-layer RMS/max table populated for all 30 decoder layers + embeddings + head, fw + bw; JSON artifact round-trips through json.load. Output shows propagated error growing with depth, with sharp jumps at layers where activation magnitude regime changes (e.g. ref_scale 6 → 777 around block 11, bf16 RMS rel 10% → 0.7% → back up to 13% at block 28).
• Existing layer-comparison tests still pass with the moved compare_tensor_logs.py and the refactored compare_tensors.

🤖 Generated with Claude Code

[jlamypoirier]

Sample precision-evaluation runs

Output of python -m tools.evaluate_precision -c examples/evaluate_precision/<config> on SmolLM2-135M, sequence length 128, single forward+backward step. Numbers are RMS relative diff vs the forced-fp32 reference, in %.

smol.yaml — pretrained HF weights

Variant            embeddings  fw mid med  fw mid max  fw logits  fw head     bw head  bw logits  bw mid med  bw mid max  bw decoder.0
bf16               0.000%      1.192%      19.904%     43.910%    4.597%      20.375%  14.710%    17.426%     22.495%     14.725%
bf16_fp32_lm_head  0.000%      1.192%      19.904%     43.901%    4.673%      19.559%  15.259%    16.797%     22.118%     14.375%
bf16_fp32_residual 0.000%      0.260%      4.569%      5.348%     0.568%      5.768%   4.132%     4.298%      8.025%      4.401%
bf16_max_precision 0.000%      0.260%      4.569%      5.347%     0.353%      6.653%   4.909%     4.959%      7.643%      4.381%

smol.yaml — random init (pretrained.model_weights=False)

Variant            embeddings  fw mid med  fw mid max  fw logits  fw head     bw head  bw logits  bw mid med  bw mid max  bw decoder.0
bf16               0.168%      1.739%      2.334%      2.425%     0.160%      0.284%   0.621%     2.120%      2.861%      3.188%
bf16_fp32_lm_head  0.168%      1.739%      2.334%      2.421%     0.160%      0.284%   0.617%     2.139%      2.937%      3.196%
bf16_fp32_residual 0.168%      1.372%      1.603%      1.689%     0.040%      0.295%   0.447%     1.435%      2.179%      2.321%
bf16_max_precision 0.168%      1.372%      1.603%      1.686%     0.041%      0.295%   0.434%     1.437%      2.180%      2.232%

smol_gspo.yaml — pretrained HF weights

Variant            embeddings  fw mid med  fw mid max  fw logits  fw head     bw head  bw logits  bw mid med  bw mid max  bw decoder.0
bf16               0.000%      0.242%      12.672%     11.312%    10.852%     0.107%   0.00095%   0.125%      0.340%      5.407%
bf16_fp32_lm_head  0.000%      0.242%      12.672%     11.296%    10.800%     0.105%   0.00176%   0.123%      0.336%      5.403%
bf16_fp32_residual 0.000%      0.227%      6.190%      6.593%     2.798%      0.042%   0.00573%   0.031%      0.068%      0.994%
bf16_max_precision 0.000%      0.227%      6.190%      6.587%     4.650%      0.049%   0.00730%   0.053%      0.128%      2.123%

smol_gspo.yaml — random init (pretrained.model_weights=False)

Variant            embeddings  fw mid med  fw mid max  fw logits  fw head     bw head  bw logits   bw mid med  bw mid max  bw decoder.0
bf16               0.173%      1.783%      2.222%      2.152%     2.296%      0.0133%  0.000095%   0.023%      0.158%      4.387%
bf16_fp32_lm_head  0.173%      1.783%      2.222%      2.143%     2.301%      0.0134%  0.000095%   0.023%      0.163%      4.626%
bf16_fp32_residual 0.173%      1.356%      1.625%      1.566%     0.860%      0.0048%  0.000044%   0.011%      0.081%      2.210%
bf16_max_precision 0.173%      1.356%      1.625%      1.560%     0.939%      0.0057%  0.000045%   0.012%      0.091%      2.611%

Observations

• Pretrained weights produce much larger forward-pass errors than random init — particularly visible at fw mid max (single worst intermediate layer) and at head.logits. The CE loss config peaks around 20-44%, GSPO around 11-13%. Random init keeps everything under ~3%.
• full_precision_residual is the dominant stability lever — it cuts the worst forward-pass numbers in roughly half (pretrained CE) or by a smaller fraction (random / GSPO).
• fp32_lm_head (Add fp32_lm_head flag for vLLM precision parity #526) has limited effect on its own; it visibly helps only when combined with full_precision_residual and even there mostly on the absolute head output (not on logits).
• GSPO backward errors are far smaller than CE backward errors (e.g. bw logits at 1e-3-1e-5% vs ~15% for CE), consistent with the GSPO loss producing much smaller logit gradients.

[jlamypoirier]

Within-engine precision: chosen-token log π

Added per-position log_softmax(logits)[label] as a diagnostic signal — the scalar that policy-gradient importance ratios actually depend on. Wired through as a no-grad chosen_logprob LM loss; the tool auto-adds it and surfaces a dedicated summary with bias, correlation, slope, and residual-after-linear-fit.

Bug fix bundled in: Fast-LLM's data.micro_batch_size is the per-sample sequence length, not the batch dim. The tool was passing 1 thinking it controlled batch size — so every previous run in this PR was on 1-token inputs. All numbers I posted earlier in this PR are invalid; the tables below replace them.

Smol (random labels)

Variant	RMS rel	Bias rel	Corr	Slope
bf16	0.473%	-0.005%	0.99950	+0.99848
bf16_fp32_lm_head	0.371%	-0.005%	0.99970	+0.99839
bf16_fp32_residual	0.358%	-0.014%	0.99972	+1.00001
bf16_max_precision	0.216%	-0.004%	0.99990	+0.99990
bf16_in_fp32_out (diagnostic)	0.371%	-0.005%	0.99970	+0.99839
bf16_reduced_reduction (diagnostic)	0.473%	-0.005%	0.99950	+0.99848

Smol_GSPO (RL data)

Variant	RMS rel	Bias rel	Corr	Slope
bf16	0.931%	+0.047%	0.99982	+1.00095
bf16_fp32_lm_head	0.876%	+0.047%	0.99984	+1.00081
bf16_fp32_residual	0.644%	-0.032%	0.99992	+1.00108
bf16_max_precision	0.568%	-0.017%	0.99994	+1.00109
bf16_in_fp32_out (diagnostic)	0.876%	+0.047%	0.99984	+1.00081
bf16_reduced_reduction (diagnostic)	0.931%	+0.047%	0.99982	+1.00095

Reference scale is ~11–12 nats, so 1% relative ≈ 0.1 nats absolute. Per-token bias sits at 0.005–0.014 nats.

Findings

1. Within-engine bf16 vs fp32 is small across the board. All RMS values < 1%, all bias values < 0.05% (≈0.005 nats), all correlations ≥0.9995, all slopes ≈1.000. There's no systematic distortion — just per-token decorrelated noise plus negligible mean shift. The intrinsic precision of a single bf16 engine compared to its fp32 equivalent is not on a scale that would cause RL collapse on its own.

2. fp32_lm_head's entire effect is the output dtype, not the matmul precision. Tested by adding bf16_in_fp32_out (= fp32_lm_head=True + matmul_precision='medium', which routes the matmul back through bf16 Tensor Cores while keeping logits fp32). Result matches standard fp32_lm_head to 5 sig figs on every metric. So the gain is purely from skipping the bf16 round on the output logits, not from running the matmul in fp32. The flag's name is misleading; it could be implemented as a bf16-in / fp32-out kernel with half the matmul memory bandwidth.

3. fp32_lm_head is downstream of the bias source. bf16 and bf16+fp32_lm_head have identical bias on GSPO data; only fp32_residual moves it. So whatever asymmetric rounding is producing the (very small) bias lives upstream of the LM head — bf16 residual stream / weight quantization. Useful structural fact, but the magnitudes involved are too small to matter intrinsically.

4. allow_bf16_reduced_precision_reduction = True has no observable effect at our matmul size (576 × 49152). cuBLAS likely doesn't pick a split-K kernel here, so the flag is moot.

Caveats and future work

This is a single-engine, single-step, small-model measurement (SmolLM2-135M, 1 fwd+bwd). The literature reports vLLM-vs-trainer log-prob mismatches of 2–24 nats per sequence and RL training collapse without precision fixes; our largest within-engine bias is ~0.005 nats. The gap is real, but many things differ between the two settings, any combination of which could be responsible:

• Cross-engine alignment — rollouts in vLLM vs gradients in the trainer go through different kernel implementations (attention, matmul, fused ops). Our single-engine setup can't see this.
• Model scale — different reduction depth, larger logit magnitudes, larger vocab.
• Training dynamics — small per-step biases compound over thousands of steps; we measure one step.
• Algorithm/data distribution — real RL rollouts (high-entropy regions, importance-ratio outliers) stress the head differently from our synthetic GRPO data.
• Software stack — different RL frameworks (TRL, open-instruct, ScaleRL) have different numerical paths even within the trainer side.

The most direct follow-up — and the one closest to what the literature actually measures — is a vLLM-vs-trainer chosen-logprob comparison at matched scale. That would isolate the cross-engine factor specifically; combining it with the within-engine measurements here would let us decompose the literature's 2–24 nat mismatch into engine-mismatch vs other factors. Out of scope for this PR.

[jlamypoirier]

FP16 vs BF16 within-engine: ~8× precision improvement, no bias

Added an FP16 sweep (fp16, fp16_fp32_residual, fp16_fp32_lm_head, fp16_max_precision) and per-variant auto-calibrated gradient scaling. Each variant runs a calibration pass at constant=1.0 to measure max unscaled gradient, then picks the largest power-of-2 scale that keeps scale × max_unscaled × headroom < fp16_max. _compare unscales per variant. Auto-picked scales for this setup: bf16/fp32 variants → 2^11=2048; fp16 base / fp16_fp32_residual → 2^9=512; fp16 + fp32_lm_head variants → 2^11=2048 (the cleaner head reduces the backward-grad max).

Chosen-token log π — smol (random labels)

Variant	RMS rel	Bias rel	Resid rel	Corr
bf16	0.473%	-0.005%	0.472%	0.99950
bf16_fp32_lm_head	0.371%	-0.005%	0.370%	0.99970
bf16_max_precision	0.216%	-0.004%	0.216%	0.99990
fp16	0.057%	+0.0005%	0.056%	0.99999
fp16_fp32_lm_head	0.044%	+0.0007%	0.044%	1.00000
fp16_max_precision	0.028%	-0.0004%	0.027%	1.00000

Chosen-token log π — smol_GSPO (RL data)

Variant	RMS rel	Bias rel	Resid rel	Corr
bf16	0.931%	+0.047%	0.928%	0.99982
bf16_max_precision	0.568%	-0.017%	0.565%	0.99994
fp16	0.113%	-0.004%	0.110%	1.00000
fp16_max_precision	0.104%	-0.003%	0.101%	1.00000

Gradients — smol_GSPO (RL data)

Variant	linear_med	linear_max	norm_med	norm_max	embeddings
bf16	0.000236%	0.0376%	0.0174%	0.292%	0.000483%
bf16_max_precision	0.000186%	0.0346%	0.0153%	0.608%	0.000355%
fp16	0.000028%	0.0061%	0.0023%	0.109%	0.000057%
fp16_max_precision	0.000028%	0.0073%	0.0022%	0.155%	0.000058%

Findings

1. FP16 gives ~8× precision reduction across the board. Per-token chosen_logprob: 0.473%→0.057% (smol) and 0.931%→0.113% (gspo) — both 8.2-8.3× reduction, matching the 7→10 mantissa-bit ratio. Gradients show the same ~8× ratio across linear/norm/embedding weights. Bias collapses too: from ~0.05% in bf16 to ~0.003-0.005% in fp16.

2. FP16 + fp32_lm_head / fp32_residual add little on top. Unlike bf16 where fp32_residual was meaningful (0.93%→0.64% on gspo), in fp16 the residual is already precise enough that adding fp32 to the residual stream barely moves anything (0.113%→0.111%). fp32_lm_head is similarly minor (0.113%→0.107%). The stability features that matter for bf16 are essentially superfluous at fp16.

3. Correlation slope is exactly 1.0 for fp16. Bf16 had slope deviations from 1.0 in the 0.0008-0.0011 range (very small systematic distortion); fp16's slope is 0.99979-1.00000 — effectively unity. No structural distortion at all.

4. The remaining caveats from the prior bf16 analysis still stand. This is single-engine, single-step, SmolLM2-135M. The literature's reported RL collapse on bf16 still isn't visible at this scale, and we can't probe cross-engine alignment with this tool. What this commit does settle: FP16 has the expected ~8× intrinsic precision over BF16 in our setting, which is consistent with the precision-based mechanism that the FP16 paper (Liu et al. 2510.26788) invokes.

The natural next step to actually attribute the literature's RL-stability claim is still a vLLM-vs-trainer log-prob diff — out of scope for this PR.

Committed in 312343e7.