fix(ci): stabilize GPT-OSS and Mamba checkpoint parity

[yuhezhang-ai]

What does this PR do ?

Stabilize two checkpoint-parity comparisons affected by kernel rounding and MoE expert selection. GPT-OSS 20B PEFT uses the existing 512-token cross-framework prefix gate and ordinary relaxed HF profile. AutoModel/HF parity forwards use a fixed Mamba cumsum head tile, preventing the trained and freshly loaded processes from selecting different numerical kernels.

Changelog

• Set cross_framework_gate_sequence_length: 512 only for GPT-OSS 20B PEFT and remove its custom HF mean-KL override. HF returns to the existing relaxed limits: mean KL 0.025, p95 KL 0.1, cosine 0.99. Full 2048-token forwards and diagnostics remain; AutoModel reload retains its standard full-sequence gate.
• Temporarily select Mamba's existing BLOCK_SIZE_H=1 cumsum-forward configuration inside the shared parity-forward context. The single-config Triton path bypasses benchmarking and cached choices. Restore the original configuration list on success or failure; preserve training's autotune cache and other kernels.
• Add CPU regression coverage using Triton's actual autotuner selection with a stub kernel: cached-choice bypass, subsequent ordinary-forward behavior, exception cleanup, absent optional Mamba, and unsupported fixed-tile failure.
• Keep Nemotron's standard thresholds and full 2048-token comparisons unchanged.

Evidence and validation

Nemotron: all 378 parameter entries and 23 buffer entries match exactly on each of eight ranks after export/reload, including dtypes and layouts. Native mean/p95 KL is 0.00282410 / 0.01047837. The first difference occurs inside the first Mamba mixer, then changes an expert selection in the next MoE layer.

The decisive eight-H100 control holds checkpoint, full 2048-token input, and all other Mamba kernel choices fixed:

Cumsum head tile	Mean KL against trained reference	p95 KL	Bitwise-exact logits
1	0	0	Yes
4	0.00283928	0.01032463	No
1 again	0	0	Yes

This identifies a numerical kernel-selection cause locally. The original failing nightly checkpoint was not retained; the reproduction uses matching Nemotron/harness source, synthetic chat training data, and a compatible staged image. It reproduces drift, not the exact nightly threshold overshoot. The independent GPU control passed. The initial capture run crashed during subsequent inline CPU scoring; separate CPU scoring successfully recovered all retained results.

GPT-OSS: replaying only AutoModel's expert selections in native HF reduced mean KL 0.02154 → 0.00209, retaining HF's own scores/operators. All 603 parameter tensors matched. On the retained adapter, the 512-token prefix of the full forward has mean/p95 KL 0.00836464 / 0.04076376 versus 0.02154484 / 0.07683115 over all 2048 tokens. This supports the scoped gate, not a universal numerical bound.

• 239 focused CPU tests passed, covering the new Mamba guard and existing attention, HF-harness, parity-metric, and config-resolver behavior.
• Repository formatting, lint, and diff checks passed.
• The actual harness-only fix passed an independent eight-H100 reload (Slurm 19025874): all 2048 × 131072 logits are bitwise exact against the retained trained reference, with other Mamba kernels left on normal autotuning.
• Prior GPT-OSS scoped CI passed all six phases at 5ad3e74c5c8ca5b59a4ae942325eb22eaa311040: AutoModel reload bitwise exact; HF prefix mean/p95 KL 0.00849597 / 0.03669290. This predates the Mamba addition.
• Combined-head scoped CI 68960536 passed at 8b0bd420d610f68f5509423b6f2f5e0fbd024348: all three eight-H100 jobs and all 18 checkpoint phases, including source parity, AutoModel reload, HF reload, and native resume. Parent and both child pipelines are terminal and successful.

Recipe	Comparison	Tokens	Profile	Mean KL	p95 KL	Cosine
Nemotron full SFT	Source	2048	standard	0.00243821	0.00998228	0.99924617
Nemotron full SFT	AutoModel reload	2048	standard	0.00000000	0.00000000	1.00000000
Nemotron full SFT	HF reload	2048	standard	0.00252052	0.00939148	0.99927586
Nemotron chat full SFT	Source	2048	standard	0.00243821	0.00998228	0.99924617
Nemotron chat full SFT	AutoModel reload	2048	standard	0.00000000	0.00000000	1.00000000
Nemotron chat full SFT	HF reload	2048	standard	0.00262630	0.01021372	0.99927540
GPT-OSS PEFT	Source	512	standard	0.00217624	0.00827674	0.99963192
GPT-OSS PEFT	AutoModel reload	2048	standard	0.00000000	0.00000000	1.00000000
GPT-OSS PEFT	HF reload	512	relaxed	0.00998306	0.04407578	0.99873771

All three AutoModel reloads have zero maximum absolute logit difference over all 2048 tokens. Both Nemotron HF reloads pass the unchanged standard cross-framework limits (mean KL 0.006, p95 KL 0.03, cosine 0.998); GPT-OSS HF uses the ordinary relaxed profile. No numeric threshold overrides or Nemotron sequence reductions were introduced.

Before your PR is "Ready for review"

• Followed contributor guidance and included DCO signoff.
• Added necessary regression tests and ran the focused suite.
• Documented measured numerical evidence and its limits.
• Complete combined-head scoped GPU CI: all 18 phases passed.

Additional Information

Original issues:

• AMINT-363: GPT-OSS HF reload KL divergence
• AMINT-364: Nemotron AutoModel reload KL divergence

Follow-up to #3937; retains the FlexAttention parity fix from #3868. No production model implementation or shared threshold profile changes.

[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 8b0bd42

[athitten]

LG, thank you @yuhezhang-ai !