Skip to content

test(checkpoint): expand parity metrics and phase coverage - #3567

Open
yuhezhang-ai wants to merge 33 commits into
mainfrom
yuhez/test/checkpoint-parity-metrics
Open

test(checkpoint): expand parity metrics and phase coverage#3567
yuhezhang-ai wants to merge 33 commits into
mainfrom
yuhez/test/checkpoint-parity-metrics

Conversation

@yuhezhang-ai

@yuhezhang-ai yuhezhang-ai commented Aug 17, 2026

Copy link
Copy Markdown
Contributor

What does this PR do?

Expand checkpoint-robustness coverage and replace per-model max-KL tuning with stable, profile-based full-logit metrics.

Changelog

  • Make phases 0-4 and process isolation the default for enabled LLM/VLM checkpoint-robustness tests; keep cross-TP as optional phase 5.
  • Stop writing an unused final distributed checkpoint after Phase 4 has restored state and completed its validation trajectory.
  • Clarify the lifecycle: source parity, train/save/reference, AutoModel reload, HF export reload, native resume, and optional cross-TP.
  • Add bounded-memory mean, p95, and max KL; cosine similarity; and mean/max absolute logit-difference reporting.
  • Gate mean KL, p95 KL, and cosine with strict, standard, and relaxed profiles. Keep one optional structured parity_threshold_overrides mapping for measured exceptions across source_load, automodel_reload, hf_reload, and cross_tp; omitted metrics inherit the selected profile. Step uses only its AutoModel-reload mean/cosine entries. Max KL remains diagnostic and is not overridable.
  • Run each loaded HF reference twice and report hf_source_self_repeat / hf_export_self_repeat, so an unstable reference is distinguishable from cross-framework drift.
  • Keep phase load/forward smoke and metrics available when only a proven-unreliable numerical gate is informational (skip_source_load_logit_parity, skip_automodel_reload_logit_parity, skip_hf_reload_logit_parity).
  • Use a frozen, SHA-256-protected 6,000+ word document; tokenize 2,048 unique long-form tokens by default and support configurable 1K-4K lengths without repetition.
  • Resize PP stage activation buffers to the configured parity length.
  • Fix source-path resolution and consolidated-reload retargeting for config-based recipes.
  • Make an explicit trust_remote_code: false override a recipe default of true.
  • Support a documented per-config HF-reload waiter timeout for very large CPU-offloaded references.
  • Remove legacy max-KL and flat numeric threshold fields after migrating live recipes to profiles plus the optional structured override mapping.
  • Migrate retrieval checkpoint robustness to the shared standard profile and default-on HF reload/native-resume phases; structured cosine overrides remain available only if a measured exception appears.
  • Add first-time checkpoint-robustness coverage for DeepSeek-V4-Flash, Qwen3.8-27B, Kimi-Linear-48B-A3B, and GLM-4.7-Flash.
  • Add unit coverage and checkpoint-robustness configuration documentation.
  • Make shared-trajectory runs guarantee enough epochs to reach their checkpoint boundary and continuation while suppressing unrelated epoch-boundary saves.
  • Export Nemotron-H PEFT weights and target-module metadata through the checkpoint's remote-code backbone.* or native Transformers model.* namespace.

Phase contract

Phase Name What it loads or does Required metric/check
0 Source load Load the original HF checkpoint with vanilla HF and AutoModel. Full-logit mean KL, p95 KL, and cosine; tied-weight checks.
1 Train/save reference Train AutoModel; save the consolidated HF checkpoint and native DCP checkpoint; capture reference logits. Training/save succeeds; reference logits and resume artifacts are finite and complete.
2 AutoModel reload A new AutoModel process loads the consolidated HF checkpoint. PEFT reloads the AutoModel model/adapter checkpoint. Full-logit mean KL, p95 KL, and cosine against Phase 1; exact PEFT trainable-weight fingerprints.
3 HF reload Vanilla HF loads the consolidated HF checkpoint. PEFT loads the exported adapter. Full-logit mean KL, p95 KL, and cosine against Phase 1; exact adapter fingerprints.
4 Native resume Load the native DCP checkpoint, including model, optimizer, scheduler, RNG, and data position, then resume training for three additional optimizer steps on the same batches as the uninterrupted run. Exact scheduler/RNG/optimizer metadata and first-pre-update per-rank model/optimizer tensors; identical batch and LR; loss within the resume profile for all three steps.
5 Cross-TP reload AutoModel with a different TP size loads the consolidated HF checkpoint. Full-logit mean KL, p95 KL, and cosine against Phase 1.

Config controls: Phase 0 — skip_source_load_parity, skip_source_load_logit_parity; Phase 2 — skip_automodel_reload_logit_parity; Phase 3 — skip_hf_reload, skip_hf_reload_logit_parity; Phase 4 — skip_resume; Phase 5 — cross_tp_size.

Phases 0-4 are the default core coverage. A phase is skipped only for a documented capability/reference limitation; an observed mismatch alone remains blocking.

Numerical profiles

Profile Same implementation mean / p95 / cosine Cross-topology mean / p95 / cosine Cross-framework mean / p95 / cosine
strict 1e-7 / 1e-6 / 0.999999 1e-6 / 1e-5 / 0.99999 1e-4 / 1e-3 / 0.9999
standard (default) 3e-3 / 1.2e-2 / 0.999 6e-3 / 3e-2 / 0.998 6e-3 / 3e-2 / 0.998
relaxed 2e-2 / 5e-2 / 0.995 2e-2 / 5e-2 / 0.995 2.5e-2 / 1e-1 / 0.99

The ordering is monotonic: changing topology is never stricter than the same implementation. standard is the default for dense and stable-MoE paths, while relaxed is reserved for measured distributed or low-precision drift. Neither model size nor TP/PP alone justifies a looser profile. A measured exception uses parity_threshold_overrides for only the exceeded gate(s), with every unspecified profile gate retained.

Max KL remains diagnostic because sparse outliers can be large while distribution-level metrics remain stable; for example, Qwen3-MoE HF reload reached max KL 1.406 with mean/p95 0.00228 / 0.00597 and cosine 0.99879.

Scoped-CI calibration

Earlier 20-model calibration established the profile boundaries and kept unexplained correctness/load failures blocking. The rebased active-LLM repair cohort 63498589 then completed 23 jobs: 15 passed and 8 failed for classified reasons. Follow-up pipeline 63504152 passed all four GPT-OSS Customizer variants and both Nemotron Flash variants without skipping a phase or changing a profile definition; pipeline 63510378 passed the repaired Nano 4B PEFT export reload.

Final routed-MoE regression pipeline 63512263 ran eight exact jobs from SHA f041b88e733f9d8162b3387348371d705304e7d5: all four GPT-OSS variants, Qwen3 MoE, Nemotron Nano Customizer full SFT, and Nano 4B PEFT passed. The Nemotron chat variant restored model/buffers/optimizer/scheduler/RNG exactly and passed steps 5-6, but its step-7 loss differed by 0.021995 (0.497%) versus the standard 0.013853 allowance. Its non-chat sibling passed standard with 0.000277 step-7 drift.

The existing relaxed resume profile allows 0.043201 for the measured chat loss, so final SHA 10812cb4102322dac731c55e25142e80cfe75593 selects it only for that routed hybrid-MoE chat config. Logit gates remain standard, no numeric override or skip was added, and the non-chat sibling remains standard. Exact one-job pipeline 63517674 passed all phases. Its step-7 drift was 0.005666, smaller than the calibration run and well inside the shared profile; exact restored state passed again.

Current decisions

Model Profile / gate decision Current outcome or follow-up
Qwen3 MoE 30B LoRA standard Clean source/reloads and exact resume state.
Nemotron Super V3 PEFT relaxed; resume re-enabled Distributed reload/resume drift is inside the shared relaxed envelopes; legacy max override removed.
MiniMax M2.7 LoRA Failed — blocking Very large HF mismatch and later resume drift (AMINT-286).
GPT-OSS 20B PEFT relaxed; resume enabled Final-head source/reloads pass; restored state and all three continuation losses pass, with PEFT and packing trajectories bitwise exact.
DeepSeek V4 Flash relaxed; failed — blocking AutoModel reload fits relaxed; HF layer_types load failure and 40 resume-buffer mismatches remain under AM-202.
Qwen3 MoE 30B TE/DeepEP standard; resume re-enabled Clean mean/p95/cosine; legacy source/HF max overrides removed.
Nemotron Flash 1B standard logits, relaxed resume Full and PEFT variants pass source, AutoModel reload, repaired vanilla-HF export reload, and native resume.
Nemotron Nano V3 standard Rebased-head checkpoint-robustness job passes after the bounded timeout increase.
Nemotron Nano Customizer standard logits; chat-only relaxed resume Non-chat passes standard; chat restores exact state and only its measured third continuation loss needs the shared routed-MoE envelope.
Nemotron Nano 4B PEFT standard; documented Mamba resume skip Repaired remote-code PEFT export passes source, AutoModel reload, and vanilla-HF reload on the final harness SHA.
Nemotron Nano 9B PEFT Failed — blocking TP2 exact PEFT reload changes replicated Mamba LoRA state on TP-rank 1; AMINT-291. No skip or threshold workaround.
Llama 3.2 1B standard Source, reloads, cross-TP, and resume fit.
Customizer Llama 3.1 8B TP standard; no resume skip Current-main rerun passes all phases: restored state, gradients, and post-step state match exactly; AMINT-289 closed as no longer reproducible.
Kimi Linear 48B-A3B Failed — blocking AM reload/resume fit; the remote Transformers OutputRecorder import failure is tracked by AMINT-288.
GPT-OSS 20B standard Rebased-head checkpoint-robustness job passes after the bounded timeout increase.
GLM 4.7 Flash Failed — blocking Large, repeatable HF mismatch (AMINT-287).
ERNIE 4.5 21B-A3B standard Legacy max-only failure; distribution metrics and resume fit.
Step3.7 200B relaxed + Phase 2 mean 0.04 / cosine 0.99; only HF gates informational AutoModel reload and native-resume gates remain enforced; only the repeatable Phase 2 routed-MoE mean/cosine need targeted overrides, while relaxed p95 stays active. Final-head rerun awaits 8-node capacity.
Qwen3.8 27B standard Cross-framework comparisons and resume fit; HF source self-repeat anomaly remains visible.
Qwen3.5 35B-A3B standard Comparisons and resume fit; HF export self-repeat anomaly remains visible.
Qwen3-VL MoE 30B standard Source, both reloads, and resume fit.
Mistral4 relaxed Clean BF16 HF-export drift anchors the relaxed cross-framework profile.
Gemma4 26B-A4B only HF gates informational AutoModel reload and native resume are bitwise exact and remain enforced; HF source/export are demonstrably non-repeatable.

Final routed-MoE regression metrics

Each parity cell is mean / p95 / max KL; cosine. Resume is the absolute loss difference at steps 5/6/7. The chat row records the calibration run; its independent final-SHA rerun passed with resume differences 0.002103 / 0.000127 / 0.005666.

Job Source HF → AM AutoModel reload HF export reload Resume loss differences
GPT-OSS full .007554/.028003/.554645; .999144 0/0/0; 1 .008223/.034005/.790934; .999154 0/.002969/.002026
GPT-OSS full chat .007554/.028003/.554645; .999144 0/0/0; 1 .005849/.016791/1.754057; .999422 0/.004654/.002209
GPT-OSS PEFT .005401/.019350/.611482; .999410 0/0/0; 1 .007292/.021927/2.030980; .999182 0/0/0
GPT-OSS PEFT packing .006386/.018244/.862887; .999233 0/0/0; 1 .006193/.022491/.814642; .999356 0/0/0
Nemotron Customizer full .002541/.010453/.239663; .999287 .002245/.008229/.080814; .999355 .002588/.009124/.249942; .999291 .001315/.001557/.000277
Nemotron Customizer chat .002594/.010720/.232123; .999278 .000576/.002029/.030139; .999849 .002881/.011180/.305111; .999207 .001347/.003246/.021995
Qwen3 MoE .002598/.009486/.271365; .998172 0/0/0; 1 .001871/.005224/.662879; .999021 0/0/0
Nano 4B PEFT .000459/.001849/.009560; .999925 .000495/.001915/.010723; .999923 .000500/.001897/.006943; .999922 documented Mamba skip

Outstanding model issues

  • AMINT-286: MiniMax M2.7 catastrophic HF parity and later resume drift.
  • AMINT-287: GLM-4.7-Flash cross-framework parity.
  • AMINT-288: Kimi Linear / Transformers OutputRecorder compatibility.
  • AMINT-289 (resolved): the prior Customizer Llama optimizer mismatch no longer reproduces on current main; fresh Phase 4 diagnostics are exact.
  • AM-202: DeepSeek V4 evidence updated with missing layer_types and 40 router-buffer mismatches.
  • AMINT-290 (low-priority CI follow-up): wrappers can remain running while Slurm accounting is temporarily unavailable, delaying propagation of terminal job state.
  • AMINT-291: Nano 9B TP2 PEFT reload changes replicated Mamba LoRA state on TP-rank 1.
Full 20-model metric table

Each comparison cell is mean / p95 / max KL; c cosine; d mean / max absolute logit difference. 0 means bitwise-identical logits. Metrics are from pipeline 63229114 except the marked Step source/AM/resume values, which use the newer pipeline 63229651.

Model HF source self Source HF → AM AM reload HF export self HF export → AM Resume
Qwen3 MoE LoRA 0; c 1; d 0 .002598/.009486/.2714; c .99817; d .1293/8.805 0; c 1; d 0 0; c 1; d 0 .001649/.005984/.03904; c .99905; d .08808/8.906 0; state exact
Nemotron Super V3 PEFT 0; c 1; d 0 .005637/.029114/.06446; c .99987; d .2813/16.5 0; c 1; d 0 0; c 1; d 0 .005842/.028533/.08861; c .99986; d .2889/13.25 not run; now enabled
MiniMax M2.7 LoRA 0; c 1; d 0 13.139/22.326/36.013; c .02378; d 4.132/41.5 0; c 1; d 0 0; c 1; d 0 6.248/11.409/18.178; c .01950; d 4.167/43.03 .003000; state exact
GPT-OSS 20B PEFT 0; c 1; d 0 .005909/.019362/1.368; c .99933; d .08169/14.0 .006648/.022910/.8930; c .99930; d .08942/7.398 0; c 1; d 0 .007906/.021799/1.482; c .99910; d .09503/17.69 not run; now enabled
DeepSeek V4 Flash —; HF load failed —; HF load failed .005166/.021667/.3463; c .99816; d .2250/15.44 —; HF load failed .00431; params/optimizer exact; 40 buffers differ
Qwen3 MoE TE/DeepEP 0; c 1; d 0 .002598/.009486/.2714; c .99817; d .1293/8.805 0; c 1; d 0 0; c 1; d 0 .002278/.005972/1.406; c .99879; d .08868/15.52 not run; now enabled
Nemotron Flash 1B 0; c 1; d 0 .000861/.002358/.00789; c .99990; d .03762/.4375 .001095/.003079/.00973; c .99987; d .04348/.5313 —; unsupported HF reload —; unsupported HF reload .008243; state exact
Nemotron Nano V3 0; c 1; d 0 .002570/.010360/.1663; c .99924; d .08407/6.236 —; timeout —; timeout —; timeout —; timeout
Llama 3.2 1B 0; c 1; d 0 .000542/.001872/.01721; c .99990; d .02540/.9766 0; c 1; d 0 0; c 1; d 0 .000515/.001794/.01643; c .99990; d .02484/1.25 .000627; state exact
Customizer Llama 3.1 8B TP 0; c 1; d 0 .000584/.001993/.00969; c .99984; d .02947/1.531 .000518/.001898/.01021; c .99984; d .02802/3.25 0; c 1; d 0 .000565/.001947/.02115; c .99983; d .02932/1.844 steps 5-6 exact; step 7 .001261; all restored/gradient/post-step digests exact
Kimi Linear 48B-A3B —; remote import failed —; remote import failed .002363/.009074/.06528; c .99934; d .09521/8.652 —; remote import failed —; remote import failed .000847; state exact
GPT-OSS 20B 0; c 1; d 0 .004814/.016797/.5545; c .99944; d .08023/12.88 —; timeout —; timeout —; timeout —; timeout
GLM 4.7 Flash 0; c 1; d 0 .2909/.6578/30.334; c .97475; d 4.668/234 0; c 1; d 0 0; c 1; d 0 .05680/.04250/18.737; c .99818; d 2.539/76 0; state exact
ERNIE 4.5 21B-A3B 0; c 1; d 0 .000957/.003744/.07851; c .99967; d .05159/9.156 0; c 1; d 0 0; c 1; d 0 .000913/.003166/.1424; c .99976; d .04806/4.281 .003327; state exact
Step3.7 200B 1.892/5.616/17.622; c .76583; d 1.407/21.46* 2.599/13.498/39.303; c .80205; d 1.308/42.66* .004473/.012207/.9024; c .99579; d .1383/15.63* 1.948/4.726/14.480; c .80724; d 1.188/26.63 .7678/2.704/23.179; c .85263; d 1.109/32.81 .006711; pre-state and first forward exact*
Qwen3.8 27B .04461/.16181/5.390; c .98238; d .3101/14.31 .000498/.001997/.01446; c .99922; d .03549/11.82 0; c 1; d 0 0; c 1; d 0 .000472/.001946/.01296; c .99937; d .03363/11.05 .000460; state exact
Qwen3.5 35B-A3B 0; c 1; d 0 .002063/.007866/.09401; c .99854; d .08786/7.891 0; c 1; d 0 .05982/.23924/6.741; c .98483; d .3052/12.16 .001870/.007141/.1011; c .99832; d .08810/10.66 .000761; state exact
Qwen3-VL MoE 30B 0; c 1; d 0 .002077/.007950/.2456; c .99870; d .1019/9.633 .001397/.004856/.1208; c .99928; d .06816/6.188 0; c 1; d 0 .001393/.004739/.1531; c .99919; d .07135/8.188 .000554; state exact
Mistral4 0; c 1; d 0 .003163/.011946/.6453; c .99851; d .1272/11.0 .001104/.004206/.03141; c .99943; d .07178/13.31 0; c 1; d 0 .02053/.08288/1.519; c .99223; d .3168/15.15 .000414; state exact
Gemma4 26B-A4B 11.133/19.072/27.995; c .46537; d 5.931/43.13 11.846/19.370/28.673; c .45972; d 5.936/44.63 0; c 1; d 0 .5665/2.322/10.673; c .86457; d 2.477/27.94 .07024/.29182/5.345; c .98810; d .6554/22.17 0; state exact

Step3.7 and Gemma4 HF-reference evidence

Model / comparison Mean / p95 / max KL Cosine Mean / max abs logit diff Interpretation
Step HF source self-repeat 1.892 / 5.616 / 17.622 0.76583 1.407 / 21.46 HF differs materially from itself in the same loaded reference.
Step source HF → AM 2.599 / 13.498 / 39.303 0.80205 1.308 / 42.66 Cross-framework gap is the same order as HF self-repeat.
Step AM reload 0.004473 / 0.012207 / 0.9024 0.99579 0.1383 / 15.63 Fits relaxed and remains blocking.
Step HF export self-repeat 1.948 / 4.726 / 14.480 0.80724 1.188 / 26.63 Exported HF reference is also non-repeatable.
Gemma HF source self-repeat 11.133 / 19.072 / 27.995 0.46537 5.931 / 43.13 HF source is highly non-repeatable at 2K.
Gemma source HF → AM 11.846 / 19.370 / 28.673 0.45972 5.936 / 44.63 Cross-framework gap closely tracks HF self-repeat.
Gemma AM reload 0 1 0 AutoModel reload is bitwise exact and remains blocking.
Gemma HF export self-repeat 0.5665 / 2.322 / 10.673 0.86457 2.477 / 27.94 Reloaded HF reference remains non-repeatable.

Stable Qwen/ERNIE MoE controls show that MoE/EP/PP alone does not justify a relaxed or skipped gate. Therefore only Step/Gemma's HF-based numerical gates are informational. Their AutoModel reload and native-resume checks remain active. Nemotron Flash's HF export reload is now enabled and passes after applying the same pre-v5 remote-code compatibility setup used for source load.

Validation

Additional information

  • The former base PR test(checkpoint): improve resume correctness diagnostics #3562 is merged; this branch targets main directly and includes the latest main updates through a normal merge commit.
  • Related: AMINT-260, AMINT-268, AMINT-269, AMINT-270, AMINT-278, AMINT-286, AMINT-287, AMINT-288, AMINT-289, AMINT-290, AMINT-291, AM-202

@copy-pr-bot

copy-pr-bot Bot commented Aug 17, 2026

Copy link
Copy Markdown

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.

Base automatically changed from yuhez/test/resume-correctness-diagnostics to main August 18, 2026 17:06
@yuhezhang-ai
yuhezhang-ai force-pushed the yuhez/test/checkpoint-parity-metrics branch from 6e4df85 to 7b0f1bb Compare August 18, 2026 18:13
@yuhezhang-ai

Copy link
Copy Markdown
Contributor Author

/ok to test dea894a

@yuhezhang-ai

Copy link
Copy Markdown
Contributor Author

/ok to test 10812cb

@yuhezhang-ai

Copy link
Copy Markdown
Contributor Author

/ok to test 7cb1978

@yuhezhang-ai

Copy link
Copy Markdown
Contributor Author

/ok to test 65486db

Signed-off-by: Yuhe Zhang <yuhez@nvidia.com>
…rics' into yuhez/test/checkpoint-parity-metrics

Signed-off-by: Yuhe Zhang <yuhez@nvidia.com>
@yuhezhang-ai

Copy link
Copy Markdown
Contributor Author

/ok to test 4611e89

Signed-off-by: Yuhe Zhang <yuhez@nvidia.com>
@yuhezhang-ai

Copy link
Copy Markdown
Contributor Author

/ok to test f5e6870

@akoumpa

akoumpa commented Aug 21, 2026

Copy link
Copy Markdown
Contributor

/ok to test 92bdc0b

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants