fix(glm): align and diagnose cross-framework router parity

[yuhezhang-ai]

Summary

This PR has three distinct scopes:

1. Production GLM correctness fixes for GLM 4.x and GLM 5.2 precision/numerical contracts.
2. A general checkpoint-resume fix for Adam/AdamW parameters whose optimizer state has not yet been created.
3. Checkpoint-robustness harness and diagnostic changes used to characterize and gate the remaining cross-framework numerical drift.

It also enables AutoModel-side checkpoint-robustness coverage for GLM 5.2. This PR was originally stacked on #3620; that base is now merged.

Production model fixes

GLM 4.x

• GLM 4.7 Flash (glm4_moe_lite) now uses the upstream fixed 1e-6 epsilon for its query and key/value MLA latent RMSNorms instead of inheriting the decoder-wide epsilon.
• Both GLM MoE implementations (glm4_moe, including the GLM 4.5 Air family, and glm4_moe_lite) now own their upstream router precision policy: the gate projection defaults to FP32 and selected mixture weights remain FP32.
• The GLM 4.7 Flash recipe declares the FP32 gate policy explicitly.

GLM 5.2

• Query and key/value MLA latent RMSNorms use the upstream fixed 1e-6 epsilon, independently of the decoder block's rms_norm_eps=1e-5.
• The model owns the same FP32 router projection/selected-weight policy as upstream GLM.
• e_score_correction_bias is included in the strict FP32 load contract.

Shared MLA latent norm epsilon

• The shared MLA now takes latent_norm_eps as a required keyword-only parameter with no default, so every model states its checkpoint's contract at the call site instead of silently inheriting a shared-layer default. The value cannot come from config: upstream Hugging Face MLA implementations construct these two norms without an epsilon argument, inheriting the reference RMSNorm default of 1e-6, while config.rms_norm_eps only governs the decoder-block norms and the HF config carries no separate field.
• This applies to every MLA reuser, all verified against their upstream implementations: glm4_moe_lite, deepseek_v3, and mistral4 now pass 1e-6 explicitly. For Mistral 4, whose checkpoints ship rms_norm_eps=1e-5, the latent norms previously inherited that decoder epsilon; they now match upstream's 1e-6, the same correction as GLM. DeepSeek V3 is numerically unchanged (rms_norm_eps is already 1e-6); glm_moe_dsa and deepseek_v32 own their norms directly and are unaffected by the shared signature.

General Adam/AdamW resume correctness

Adam creates per-parameter state lazily, after a parameter first receives a gradient. A trainable parameter that is unused in a step can therefore be present in the optimizer parameter groups but lack step, exp_avg, and exp_avg_sq. The flattened DCP save/load path expected a complete optimizer-state skeleton and could fail to restore such a checkpoint.

OptimizerState.state_dict() now materializes zero-valued lazy Adam state before both native and flattened DCP serialization. This is a general Adam/AdamW checkpoint fix, not a GLM-specific workaround, and zero initialization preserves the behavior of the parameter's future first update.

Checkpoint-robustness coverage

GLM 4.7 Flash

• Keeps the full 2,048-token forward and full-sequence metrics/evidence.
• Gates the first 128 tokens of that same forward for the two diagnosed cross-framework comparisons.
• Uses the shared relaxed profile only for HF-source -> AutoModel and HF-export -> AutoModel reload. AutoModel consolidated reload and native resume remain on the standard/exact correctness gates.

GLM 5.2

• Enables the AutoModel-side train, self-repeat, consolidated-reload, and native-resume phases.
• Uses the known-good PP geometry (GBS=512, LBS=4) and 512 samples, enough to form a complete global batch at DP64.
• Skips the two vanilla-HF phases in routine CI because loading the roughly 1.5 TB checkpoint as the rank-0 HF reference takes about 50 minutes while the multi-node allocation idles. Cross-framework parity was checked out of band: mean KL was 0.0049 at 512 tokens (standard-profile pass) and 0.0071 / 0.0065 at 2K / 4K, with p95 within the standard profile and no divergent-token tail.
• Keeps native DCP resume on the standard trajectory-tolerance profile, with exact pre-update model, buffer, and optimizer state verification. Consolidated AutoModel reload alone uses the shared relaxed same-implementation profile (mean_kl <= 0.02, p95_kl <= 0.05) plus the one metric not covered by that profile, a model-specific cosine >= 0.985 override. This does not relax native resume or other same-implementation comparisons.

Harness and diagnostic changes

• Computes cosine dot products and norms from FP64 operands, avoiding FP32 multiply error before FP64 accumulation.
• Supports a shorter cross-framework gate prefix without discarding the full-forward metrics; a non-gating standing-shape HF comparison is added automatically when the gate is shortened.
• Adds opt-in router capture/comparison and shape-sensitivity diagnostics. Route-flip counts remain evidence, not pass/fail criteria.
• Hardens distributed capture validation, supported-family checks, patch restoration, finite-value validation, Phase 3 diagnostic isolation, boolean fixture parsing, schema collision checks, and concise CI logging while retaining complete JSON artifacts.
• Reuses the shared bounded-memory per-token parity implementation and tensorizes correction-bias direction analysis.
• Documents and tests YAML numeric parsing; scientific-notation examples that PyYAML reads as strings use decimal spellings instead.
• Normalizes absent lazy Adam state and materialized zero-step Adam state as the same semantic state in the resume snapshot check. This is a harness representation fix; the production serializer still materializes the complete checkpoint state.

GLM 4.7 Flash diagnosis

The identical-input 2,048-token probe found no single discontinuous layer. Every local block remained close to HF (cosine at least about 0.99993), while exact router-set agreement was about 98.4-99.7% per routed layer. The production-fix scoped run measured:

Token group	Tokens	Mean KL	Share of total KL mass
0-5 flipped layers	976	0.0190	2.9%
6-10	572	0.0271	2.4%
11-20	369	0.2085	11.8%
21+	131	4.1249	83.0%

The 70 tokens with no self-route flips had mean KL 0.0101, versus 0.3290 for tokens with one or more flips. This is an empirical floor rather than a pure routing-free control because causal attention can propagate earlier-token route differences.

Route replay separates the remaining error sources:

AutoModel route control against HF	Full mean KL	p95 KL	Cosine
Natural routing	0.3181	0.5963	0.97229
Replay HF expert indices	0.2002	0.2713	0.98649
Replay HF indices and mixture weights	0.1256	0.1027	0.99096

AutoModel replaying its own captured indices and weights was exactly identical. The evidence shows that route bifurcation is the dominant tail amplifier, route-weight differences also matter, and a smaller continuous cross-framework residual accumulates through the routed layers; it does not identify a remaining correctness mismatch.

Changing only vanilla HF's total forward length reproduced the same sensitivity: at 512 tokens, KL(HF@4096[:512] || HF@512) was 0.1498, versus matched-shape HF-to-AutoModel KL 0.1528. Qwen3-30B-A3B had a similar number of near-tie route flips but mean self-KL only 0.0024, showing why flip count itself is diagnostic evidence rather than a correctness threshold.

GLM 5.2 indexer training scope

The native-resume failure exposed by this PR and the desired training behavior are separate issues:

• The current GLM 5.2 forward returns discrete top-k indices from the DSA indexer. The language-model loss cannot differentiate through that selection, so the indexer parameters can receive grad=None; that is what exposed the missing lazy Adam state during checkpoint restore.
• The original DSA recipe trains the indexer with a separate KL objective that matches index scores to the main attention distribution; the language-model loss trains the main model, not the indexer through top-k. The GLM-5 report follows the same warmup/sparse-adaptation recipe and explicitly freezes the indexer during RL.
• The public GLM report does not state clearly whether the auxiliary indexer objective is retained during ordinary SFT. This PR therefore fixes resume correctness only. A follow-up should make AutoModel's SFT policy intentional: either freeze/exclude the indexer, or implement the DSA auxiliary indexer loss. It should not rely on gradients propagating through discrete top-k.

References: DeepSeek-V3.2 DSA training recipe and GLM-5 report.

Scoped CI

GLM 4.7 Flash

The first exact scoped run (parent 64048635, job 407983121) ran the ordinary 50-step recipe and all robustness phases:

Comparison	Gate mean KL	Gate p95 KL	Gate cosine	Result
HF source to AutoModel	0.013405	0.030912	0.999509	Relaxed pass
AutoModel self-repeat	0	0	1	Exact
AutoModel consolidated reload	0	0	1	Exact; standard
HF export self-repeat	0	0	1	Exact
HF export to AutoModel reload	0.012111	0.030585	0.999911	Relaxed pass

Native DCP resume was exact on all eight ranks: steps 5-7 had zero loss difference, and all parameter, buffer, optimizer, gradient, and post-step digests matched. The calibrated rerun (parent 64054805, job 407998642) passed.

GLM 5.2

The first exact-head run (parent 64724314, retry 413471549) reached all requested phases after an initial rendezvous failure:

• Phase 1 self-repeat was exact.
• Phase 2 consolidated reload measured mean_kl=0.0073655, p95_kl=0.0297506, and cosine=0.991423.
• Phase 3 exposed missing lazy Adam state for DSA indexer parameters that received no gradient.

Commit 105704756 added the general optimizer-state fix, initial Phase 2 bounds, and a complete 512-sample global batch. The conflict-free rerun (parent 64884019, leaf job 414485914) completed ordinary training and all robustness work, despite being summarized as an infrastructure issue:

• Phase 1 self-repeat was again exact (mean_kl=0, p95_kl=0, cosine=1).
• Phase 2 consolidated reload measured mean_kl=0.00889297, p95_kl=0.0375951, and cosine=0.989929. Mean and p95 passed the initial explicit bounds; cosine missed 0.99 by about 0.000071.
• Phase 3 successfully loaded the optimizer checkpoint, confirming that the general lazy-Adam production fix works. Its later snapshot assertion compared pre-save missing lazy state with post-load explicit zero state and reported a counter mismatch even though those representations are semantically equivalent.

Commit 2374374bb fixes that harness-only comparison. Merge commit f0772b5de resolves the latest main conflict and simplifies the consolidated-reload policy to the shared relaxed profile plus only the model-specific cosine override (0.985). Native resume remains on the standard profile and requires exact restored state before the first resumed update.

The final single-job scoped rerun passed against the verified conflict-free merge commit: parent pipeline 64902857, generated pipeline 64903663, leaf job 414650566.

• Phase 1 self-repeat was exact (mean_kl=0, p95_kl=0, cosine=1).
• Consolidated AutoModel reload passed with mean_kl=0.00867099, p95_kl=0.03702399, and cosine=0.99011834; reload self-repeat was exact.
• Native resume restored model parameters, buffers, and optimizer state exactly on all 256 ranks. Step 5's first resumed forward was exact on all ranks. After the first resumed update, the worst loss differences were 0.00091209 at step 6 and 0.00099553 at step 7, using at most 17.1% and 19.5% of the standard allowed difference; there were no tolerance violations. Gradient and post-step digests were not bitwise identical under this HybridEP/PP topology, so this establishes exact restore plus bounded standard-profile trajectory continuation, not bitwise-exact post-update execution.

Local validation

• git diff --check and Ruff checks passed on the changed Python files.
• All 9 optimizer-state tests pass, including the never-used Adam parameter case and native multi-rank pipeline/DTensor DCP round trips.
• 244 focused checkpoint-harness, parity-metric, HF-kwargs, resume-trajectory, and config-resolver tests pass.
• The post-format focused regression tests pass (2 passed).
• All 19 focused GLM router/FP32 precision contract tests passed earlier on this branch.
• The resolved GLM 5.2 fixture scopes relaxed to automodel_reload and preserves the cosine override as the float 0.985.

Follow-up intentionally deferred

• Decide and implement the intended GLM 5.2 ordinary-SFT indexer policy: frozen/excluded, or trained with the DSA auxiliary KL objective.
• Normalize all numeric fixture arguments once at the harness boundary; changing only one field would leave the existing string-valued convention inconsistent.

[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.

[yuhezhang-ai]

/ok to test 50b6edd

[yuhezhang-ai]

/ok to test aca8222

[yuhezhang-ai]

/ok to test 8615128

[yuhezhang-ai]

/ok to test 75a7ce0

[copy-pr-bot]

/ok to test 75a7ce0

@yuhezhang-ai, there was an error processing your request: E2

See the following link for more information: https://docs.gha-runners.nvidia.com/cpr/e/2/

[yuhezhang-ai]

/ok to test 75a7ce0

[yuhezhang-ai]

/ok to test 420b9b3

[yuhezhang-ai]

/ok to test 1057047

[yuhezhang-ai]

/ok to test babe3fa

[yuhezhang-ai]

/ok to test f0772b5

[yuhezhang-ai]

/ok to test 8b6c1c1