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.