Skip to content

fix(ckpt): preserve MoE fine-tuning state and refresh Super VL verification - #6220

Merged
cuichenx merged 20 commits into
NVIDIA-NeMo:mainfrom
cuichenx:chcui/sage/nmvl-ga-final-pr
Sep 28, 2026
Merged

cuichenx merged 20 commits into
NVIDIA-NeMo:mainfrom
cuichenx:chcui/sage/nmvl-ga-final-pr

Conversation

@cuichenx

Copy link
Copy Markdown
Contributor

What does this PR do ?

Fix shared MoE checkpoint/inference issues found during Nemotron Super VL verification, and publish refreshed full-VL and text-only verification cards.

Changelog

  • Preserve updated router expert_bias buffers in native PEFT save/load and merge. MCore updates these buffers during gradient finalization even when router parameters are frozen.
  • Remove the obsolete Bridge expert-bias padding-mask shim; MCore already handles padding natively.
  • Make checkpoint inference explicitly dropless while preserving the saved dispatcher. Disable training-only forced routing, paged stashing, and incompatible NCCL-EP zero-copy; support uneven HybridEP inference inputs. Explicit None overrides clear saved capacity limits rather than silently retaining them.
  • Preserve the checkpoint's explicit shared-expert overlap setting and copy reasoning-parser files during Omni/Super-VL HF export.
  • Remove the duplicate newline after Omni media prompts without changing the fixed-resolution training-image policy.
  • Preserve one physical shared MTP block in text-only HF import/export; Super recipes explicitly use that block at two prediction depths. No extra MTP-schema flag is needed.
  • Reuse an explicitly active inference-launcher environment, matching the training launcher.
  • Use nvidia/NVIDIA-Nemotron-3.5-Super-VL-120B-A12B-BF16 in public examples and recipes, refresh the full-VL card, and add a text-only companion card and generated documentation.
  • Narrowly correct card-validator false positives for literal chat delimiters and loopback coordinator addresses, retaining private-address, path, credential, and placeholder checks.

GitHub Actions CI

Request standard CI plus L1 (needs-more-tests). The added L1_Launch_peft_router_checkpoint runs the native one-GPU router checkpoint regression under a distributed launch; ordinary unit execution skips that distributed path. No L2/full-suite label is requested.

Validation

  • Latest MTP simplification: 157 focused model/conversion/recipe tests passed in nvcr.io/nvidia/nemo:26.08.01.
  • Documentation, catalog, and validator coverage: 116 tests passed. Generated-document consistency, whitespace checks, and all-file pre-commit passed.
  • Earlier integration snapshot: 911 focused Bridge tests, the one-GPU native router checkpoint regression, and 11 upstream MTP gradient/forward tests passed. A two-GPU probe covered packed Hybrid/Omni forward/backward with TP2/SP and PP2; it did not cover MTP or full optimizer/pipeline scheduling.
  • The cards retain the exact source commit for each workload. They record exact GPU weight round-trips including MTP, bounded forward/generation checks, H100 pretrain/resume, and SFT/PEFT/export on H100 and GB200. Full-VL SFT covers 200 steps, full-VL PEFT 500 steps, and text-only SFT/PEFT 500 steps. H100 text-SFT uses a fresh-root continuation from step 250, not an uninterrupted run.
  • Full-weight and training verification was not rerun after the final MTP metadata simplification; the new focused tests cover its config/recipe contract. Historical results have not been relabeled as final-head runs.
  • Fresh independent review of the complete proposed PR diff is performed before publication.

Known verification limits

  • GB200 pretrain/resume remains unverified. Recovery reached step 500, but one of 1,125 overlapping loss checks differed by 2.55%, exceeding the unchanged 1% tolerance. Independent replay is pending; successful model reload is not optimizer-resume equivalence.
  • The public HF revision is intentionally blank, so results are not yet bound to an immutable public-release revision. This leaves four known card-validation errors: the missing immutable revision and missing pinned comparison argument in each card. No validator gate has been relaxed to hide them.
  • CPU conversion, long-context training, and additional text-only pretrain/resume verification are deferred. Stronger visual diagnostic failures remain distinguished from passing functional checks.

Compatibility and scope

Older bias-enabled PEFT checkpoints may omit trained router-bias state. Native adapter inference/merge rejects missing state instead of silently replacing it with base-model values. A permissive training-resume policy may retain base bias, but that is not exact recovery; missing trained bias cannot be reconstructed from adapter tensors.

Standalone HF adapter export still omits router-bias buffers. The preservation path verified here is native adapter reload, merge into the base, then full-model HF export.

The validator correction is reusable infrastructure behavior: literal ChatML/reasoning tokens are prompt/output data, and loopback is a portable single-node binding. Tests retain rejection of other placeholders, non-loopback addresses, private paths, and explicit deny terms. Numerical verification thresholds are unchanged.

No dependency pins, CI workflow definitions, or MCore code are changed by this PR. Dropless overrides affect inference, not training. Existing copyright years are unchanged.

Before your PR is "Ready for review"

Pre checks:

  • Read and followed the contributor guidelines.
  • Added focused regressions and a native distributed-checkpoint test.
  • Updated documentation and verification cards with explicit incomplete items.
  • Ran focused checks and all-file pre-commit.
  • Obtained a fresh independent review before opening the PR.
  • Does the PR affect components that are optional to install? No new optional dependencies are introduced.
    • Reviewer: confirm existing optional-library guards remain appropriate.

Additional Information

Follow-up to #6100 and #6130, not a new architecture implementation. Primary area: checkpointing. Verification coverage is partial and source-specific as detailed above.

Signed-off-by: Chen Cui <chcui@nvidia.com>
Signed-off-by: Chen Cui <chcui@nvidia.com>
Signed-off-by: Chen Cui <chcui@nvidia.com>
Signed-off-by: Chen Cui <chcui@nvidia.com>
Signed-off-by: Chen Cui <chcui@nvidia.com>
Signed-off-by: Chen Cui <chcui@nvidia.com>
Signed-off-by: Chen Cui <chcui@nvidia.com>
Signed-off-by: Chen Cui <chcui@nvidia.com>
Signed-off-by: Chen Cui <chcui@nvidia.com>
Signed-off-by: Chen Cui <chcui@nvidia.com>
Signed-off-by: Chen Cui <chcui@nvidia.com>
Signed-off-by: Chen Cui <chcui@nvidia.com>
Signed-off-by: Chen Cui <chcui@nvidia.com>
Allow exact ChatML/reasoning delimiters in public inference prompts and output, and loopback literals for single-node coordinators. Preserve checks on enclosed content, other addresses and paths, credentials, URLs, and deny terms. Add focused positive and negative regressions without changing verification scripts or numerical thresholds.

Signed-off-by: Chen Cui <chcui@nvidia.com>
Publish the refreshed cards, companion guide, generated catalog and navigation, with focused consistency and throughput fixtures. Preserve incomplete GB200 pretrain/resume and deferred workflows as unverified. Leave the public HF revision blank as requested; its four validator errors remain acknowledged rather than weakening revision checks.

Signed-off-by: Chen Cui <chcui@nvidia.com>
Signed-off-by: Chen Cui <chcui@nvidia.com>
Signed-off-by: Chen Cui <chcui@nvidia.com>
@cuichenx cuichenx added bug Something isn't working area:ckpt Checkpoint conversion, loading, export, and save paths needs-review PR is ready for code review and waiting on a reviewer needs-more-tests Requires additional L0 and L1 test coverage before merge labels Sep 28, 2026
@cuichenx

Copy link
Copy Markdown
Contributor Author

/ok to test 52d80a3

Signed-off-by: Chen Cui <chcui@nvidia.com>
@cuichenx

Copy link
Copy Markdown
Contributor Author

/ok to test 4d7ab0b

@cuichenx
cuichenx merged commit 1916758 into NVIDIA-NeMo:main Sep 28, 2026
133 checks passed

This branch was successfully deployed

1 active deployment
test — 4d7ab0b2 Deployed Sep 28, 2026 by copy-pr-bot[bot] via cicd-wait-in-queue #22231
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

area:ckpt Checkpoint conversion, loading, export, and save paths bug Something isn't working needs-more-tests Requires additional L0 and L1 test coverage before merge needs-review PR is ready for code review and waiting on a reviewer

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant