[tools] Report paired numerical-diff statistics in loss_compare.py

[yuweih205]

Summary

loss_compare.py currently prints per-step differences and separate per-run averages, but it does not report paired statistics over the common steps. This makes it difficult to identify when a comparison first diverges or how large the trajectory drift is. For example, an overall average difference can round to zero while individual steps still mismatch.

This PR adds a reusable compare_series() helper and reports paired statistics for every selected metric, including loss and grad_norm:

• baseline/test step counts and paired step count
• exact-match count
• first divergent step
• MAE and RMSE
• maximum absolute difference and its step
• final paired-step difference

The comparator is injectable so a future tolerance policy can reuse the same statistics path without introducing a second comparison implementation.

Motivation: observed numerical-alignment mismatch

This is motivated by an observed precision/trajectory mismatch, not only by a reporting-formatting gap. In a controlled Qwen3-235B-A22B comparison with the same initial checkpoint, data order, seed=42, and deterministic execution, old TorchTitan and current main both completed 1,000 steps but showed:

• loss divergence beginning at step 2;
• grad_norm divergence beginning at step 19;
• loss MAE of 0.04191409;
• maximum absolute loss difference of 0.12922001;
• final loss values of 7.260471 (main) vs. 7.174619 (old), a difference of about 0.085853.

The controlled ablations show that this is a measurable numerical mismatch and that specific MoE changes account for different parts of the observed trajectory change:

Comparison	Loss	grad_norm
old vs. current main	first divergence: step 2; MAE: 0.04191409; max abs diff: 0.12922001; final diff: about 0.085853	first divergence: step 19
#3666 enabled vs. disabled	first divergence: step 2; max abs diff: 0.10892582 at step 955	first divergence: step 17; max abs diff: 5.3125 at step 720
#3386 before vs. after, with #3666 disabled on both sides	first divergence: step 2; max abs diff: 1.43051e-4 at step 42	first divergence: step 17; max abs diff: 0.125
#3386 and #3666 both disabled vs. old	1000/1000 bitwise exact	1000/1000 bitwise exact

In the combined-ablation run, the final loss was 7.174618721008301 on both sides and the final grad_norm was 1.046875 on both sides; batch fingerprints also matched for all 1,000 steps. These results are from the controlled experiment documented in the linked discussion and are included here to explain the need for paired reporting, not as a claim that this PR diagnoses the underlying model changes or proves that final convergence quality is worse.

The related expert-bias semantics and migration questions are tracked separately in #4394. This is also consistent with the issue discussed in #3996, where concrete step mismatches were visible even though the reported average difference was 0.000000.

Additional evidence: the largest drift is a composite path

The largest loss difference in the experiment log comes from a combination of modern implementation choices, rather than from one isolated change. In a Qwen3-235B-A22B comparison, FlexAttention + fused QKV + ChunkedLoss versus SDPA + unfused QKV + ordinary CE produced the following result over 200 paired steps:

Comparison	Paired steps	Loss MAE	Max absolute loss difference	grad_norm exact
modern composite vs. SDPA/unfused/plain-CE path	200	0.219222	0.665282	1/200

The same experiment notes also contain a 524-step summary for this modern-vs.-SDPA family (0.12834 MAE and 0.66528 maximum difference). Because the paired windows differ, these summaries should not be combined into one statistic. An earlier 100-step Llama3-8B comparison likewise reported 0.04387 MAE and a 0.22112 maximum difference at step 71.

The documented aligned fallback supports the following limited statement: in the aligned dense-model comparison, replacing the three modern choices with SDPA + unfused QKV + ordinary CE restored bitwise overlap. This is a composite reversion, not proof that each of #3571, #2878, and #2937 independently accounts for the full 0.665282 drift. For the Qwen3 MoE 1,000-step bitwise closure shown above, exactness was instead obtained by disabling #3386 and #3666 while already using FA3 + unfused QKV + plain CE; it should not be described as “reverting the three dense-path changes alone fixes the MoE mismatch.”

The available per-feature measurements are useful for sizing individual effects, but they are not a complete 2^3 factorial design or three strict one-at-a-time removals from the 200-step modern baseline:

Matrix configuration	Paired window	Loss MAE	Max absolute loss difference	Max-diff step	RMSE
FA3 + fused QKV + CE	through step 279	0.0001633	0.0008516	209	0.0002239
FA3 + unfused QKV + ChunkedLoss	through step 279	0.0002068	0.0010004	266	0.0002925
FA3 + fused QKV + ChunkedLoss	through step 279	0.0001707	0.0008869	209	0.0002198
old FlexAttention vs. old SDPA	1,000 steps	0.00222	0.03046	—	—

The first three rows are the feature-level matrix reported in the experiment notes (the corresponding runs were not all successful through the same final step), and the last row uses an older TorchTitan comparison. These numbers therefore should be read as per-path diagnostics, not as additive causal contributions to 0.665282.

Example output:

[LOSS_COMPARE] metric: loss
[LOSS_COMPARE] steps: 1000/1000
[LOSS_COMPARE] paired_steps: 1000
[LOSS_COMPARE] exact_match: N/1000
[LOSS_COMPARE] first_divergent_step: 2
[LOSS_COMPARE] mae: 4.191409e-02
[LOSS_COMPARE] rmse: ...
[LOSS_COMPARE] max_abs_diff: 1.292200e-01 at step ...
[LOSS_COMPARE] final_diff: ... (step 1000)

--assert-equal remains exact by default; no training logic or tolerance flag is changed.

Test plan

• python3 -m unittest tests.unit_tests.cpu.test_loss_compare_statistics -v (7 tests)
• python3 -m py_compile scripts/loss_compare.py tests/unit_tests/cpu/test_loss_compare_statistics.py
• git diff --check
• Pre-commit checks for formatting, AST, flake8, pydoclint, codespell, and link checking
• GPU/integration training comparison (not available in this local environment)

Related context

This PR is complementary to the existing relative-tolerance work in #3847: it does not add another --rtol implementation.

[meta-cla]

Hi @yuweih205!

Thank you for your pull request and welcome to our community.

Action Required

In order to merge any pull request (code, docs, etc.), we require contributors to sign our Contributor License Agreement, and we don't seem to have one on file for you.

Process

In order for us to review and merge your suggested changes, please sign at https://code.facebook.com/cla. If you are contributing on behalf of someone else (eg your employer), the individual CLA may not be sufficient and your employer may need to sign the corporate CLA.

Once the CLA is signed, our tooling will perform checks and validations. Afterwards, the pull request will be tagged with CLA signed. The tagging process may take up to 1 hour after signing. Please give it that time before contacting us about it.

If you have received this in error or have any questions, please contact us at cla@meta.com. Thanks!

[meta-cla]

Thank you for signing our Contributor License Agreement. We can now accept your code for this (and any) Meta Open Source project. Thanks!