test(ci): report Jensen-Shannon checkpoint parity metrics

[yuhezhang-ai]

What does this PR do ?

Stacked on #3567, report Jensen-Shannon divergence alongside the existing checkpoint-parity metrics so we can compare KL and JSD on real long-context model runs before choosing any JSD gate.

Changelog

• Report mean, p95, and max JSD for every full-logit checkpoint-parity comparison.
• Keep JSD report-only; existing mean-KL, p95-KL, and cosine gates are unchanged.
• Bump the structured metric record to schema version 2 and document the new fields.
• Use phase-specific relaxed reload profiles for Mistral4's measured PP4/EP8 routed-MoE cross-process variance while retaining standard source-load parity.

Before your PR is "Ready for review"

Pre checks:

• Make sure you read and followed Contributor guidelines
• Did you write any new necessary tests?
• Did you add or update any necessary documentation?

Validation

• Focused checkpoint-parity/config unit tests and Ruff checks pass.
• Scoped data was collected on the fixed 2,048-token document. KL is directional: source load reports
  KL(HF source || AutoModel), while export reload reports KL(trained AutoModel || exported HF). JSD is symmetric,
  uses natural logarithms, and is bounded by ln(2) = 0.6931.

Model / topology	Comparison (reference -> candidate)	Policy	KL mean / p95 / max	JSD mean / p95 / max	Cosine	Result
Qwen3 MoE 30B / EP8	HF source -> AutoModel	standard	0.00260 / 0.00949 / 0.271	0.000653 / 0.00239 / 0.0790	0.998172	pass
Qwen3 MoE 30B / EP8	trained AutoModel -> HF export	standard	0.00151 / 0.00505 / 0.155	0.000377 / 0.00125 / 0.0370	0.999202	pass
ERNIE 4.5 21B-A3B / EP8	HF source -> AutoModel	standard	0.000957 / 0.00374 / 0.0785	0.000239 / 0.000943 / 0.0190	0.999666	pass
ERNIE 4.5 21B-A3B / EP8	trained AutoModel -> HF export	standard	0.000938 / 0.00334 / 0.112	0.000239 / 0.000834 / 0.0356	0.999742	pass
GPT-OSS 20B / EP8	HF source -> AutoModel	standard	0.00481 / 0.0168 / 0.555	0.00118 / 0.00419 / 0.146	0.999437	pass
GPT-OSS 20B / EP8	trained AutoModel -> HF export	standard	0.00144 / 0.00514 / 0.326	0.000345 / 0.00128 / 0.0587	0.999853	pass
Qwen3-VL MoE 30B / EP8	HF source -> AutoModel	standard	0.00208 / 0.00795 / 0.246	0.000519 / 0.00197 / 0.0655	0.998698	pass
Qwen3-VL MoE 30B / EP8	trained AutoModel -> HF export	standard	0.00151 / 0.00525 / 0.113	0.000376 / 0.00129 / 0.0260	0.999303	pass
Mistral4 / PP4 x EP8	HF source -> AutoModel	standard	0.00299 / 0.0107 / 0.645	0.000703 / 0.00269 / 0.0829	0.998259	pass
Mistral4 / PP4 x EP8	trained AutoModel -> HF export	relaxed	0.0203 / 0.0725 / 1.79	0.00483 / 0.0181 / 0.277	0.992132	pass
GLM 4.7 Flash / EP8	HF source -> AutoModel	standard	0.291 / 0.658 / 30.3	0.0262 / 0.127 / 0.693	0.974748	fail
GLM 4.7 Flash / EP8	trained AutoModel -> HF export	standard	0.0568 / 0.0425 / 18.7	0.00552 / 0.0103 / 0.693	0.998179	fail
Gemma4 26B-A4B / EP8	HF source -> AutoModel	standard, informational	11.8 / 19.4 / 28.7	0.625 / 0.693 / 0.693	0.459719	outside profile; HF-source repeat is also unstable
Gemma4 26B-A4B / EP8	trained AutoModel -> HF export	standard, informational	0.0702 / 0.292 / 5.34	0.0144 / 0.0644 / 0.571	0.988101	outside profile

Metric-selection conclusion

The measured rows show a metric-sensitivity difference, not a correctness classifier. Mistral4's accepted relaxed
reload drift has mean/p95 KL 0.0203 / 0.0725; GLM has 0.0568 / 0.0425. GLM's mean KL is 2.8x larger even
though its p95 is lower because fewer than 5% of its tokens form an extreme directional-KL tail. JSD compresses that
tail: GLM's mean JSD (0.00552) is only 14% above Mistral4 (0.00483), while GLM's p95 JSD (0.0103) is lower
than Mistral4 (0.0181). This demonstrates KL's greater sensitivity to rare, severe probability differences; it does
not by itself establish which result is a correctness bug.

Follow-up GLM diagnosis changes how its failed row should be interpreted:

• All 9,491 tensors shared by the source checkpoint and a zero-step AutoModel HF export are bit-for-bit identical.
• HF and AutoModel are each exactly repeatable at a fixed 2,048-token shape.
• 160 tokens (7.8%) with route differences in at least 21 MoE layers contribute 84% of total source-load KL.
• Tokens with at most five routed-layer differences form a much smaller numerical-drift floor around mean KL 0.018.
  Replacing only the catastrophic tail with that floor predicts mean KL 0.0486, close to the trained-export result
  0.0568.
• Source correction biases contain many quantized near ties. Training changes that router state, consistent with fewer
  sustained route bifurcations after export.

The GLM output divergence is therefore real and deterministic, but it is not evidence of checkpoint load/export
corruption. Small cross-kernel numerical differences are amplified by near-tied MoE routing over a long sequence.
Independent GLM model-semantics discrepancies still warrant focused fixes, but controlled variants show that those
fixes alone do not remove the long-sequence routed tail.

Neither KL nor JSD can distinguish checkpoint corruption, wrong model math, and sensitive numerical routing without
state, component, and router diagnostics. Directional KL assigns more weight to rare severe probability changes; JSD
is symmetric and bounded, making its scale easier to interpret while saturating the most divergent tokens. Mean and
p95 aggregation determine how much of either distribution becomes gating evidence, while cosine guards overall logit
geometry.

For this PR, JSD therefore remains report-only. The existing mean/p95 KL and cosine gates remain unchanged, and max
JSD is not introduced as a universal gate. A follow-up harness change can retain the fixed 2,048-token forward for
diagnostics while using a shorter blocking cross-framework prefix for route-sensitive models and optionally reporting
top-k boundary margins and sustained router bifurcations.

Repeatability remains useful for separating nondeterminism from deterministic cross-implementation divergence, but it
does not identify the root cause by itself. Gemma4's original HF source is unstable (self-repeat KL
11.1 / 19.1 / 28.0 and JSD 0.621 / 0.693 / 0.693), whereas GLM's fixed-shape self-repeats are exact.

Scoped pipelines: Qwen3 MoE,
Mistral4,
LLM cohort, and
VLM cohort.
The phase-specific Mistral4 profile confirmation is queued in
NeMo-CI 63966458.

Additional Information

• Base PR: test(checkpoint): expand parity metrics and phase coverage #3567
• This PR intentionally does not introduce JSD thresholds; the collected distribution data will inform that decision.

[copy-pr-bot]

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.