fix(ckpt): reject unowned export metadata at every PP size

[pruprakash]

What does this PR do?

Reject missing export metadata with the mapped parameter name at every pipeline-parallel size. An invalid all-ranks-unowned AutoMapping task now raises instead of silently returning an empty export at PP1; PP2 retains its error and identifies the parameter. Both normal and quantized export entry points use the same check.

The original NemotronH failure no longer reproduces on main after #5262 and #6343. This revision replaces the original allow_missing proposal, which would have hidden an invalid PP2 task. Mapping validation is not a universal ownership guarantee: cached parameter names and caller-supplied tasks still make a clear error useful.

Changelog

• Enforce the existing object-broadcast ownership requirement at PP1 and include the parameter name in the PP2 error.
• Remove the empty-export guards and misleading dense-layer/MoE comment.
• Test missing ownership, valid remote ownership and repeated export, both export entry points, valid False/0 metadata, and missing expert-count metadata.

Validation

• 454 focused and adjacent unit tests passed; two MXFP8 tests require newer GPU hardware and were skipped.
• Current-main and candidate NemotronH PP1/PP2 round trips and separate cold exports: complete 29-key output including embedding, correct shapes/dtypes, finite tensors, exact round-trip values, rank agreement, repeat equality, strict HF reload and identical finite/nonconstant HF logits.
• Real distributed controls: valid parameters arrive from either owning rank; deliberately invalid all-unowned tasks fail with the parameter name at PP1 and PP2 in the candidate. The original proposal returned empty output for the invalid PP2 case.
• Full repository pre-commit checks and diff checks pass. Strict local typing retains the same existing errors as main; it is not claimed clean.

The semantic fixture is a deterministic tiny random HF NemotronH checkpoint, not pretrained full-model weights. The forward comparison is HF before/after conversion, not Megatron forward parity. Quantized-entry ownership checks do not qualify quantized payload numerics.

Related: NVBug6612509.

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

[claude]

LGTM - clean, well-tested fix.

The asymmetric allow_missing usage in megatron_to_hf (owner passes it implicitly False, non-owners pass True) is safe: all_gather_object gives every rank the same ownership list, so the allow_missing branch is only reached when no rank owns the object - in which case all ranks are non-owners and all pass allow_missing=True. When any rank owns it, src_rank is found everywhere and the broadcast runs symmetrically. No collective divergence. The change correctly turns the old ValueError into a graceful {} export for Megatron params with no HF counterpart on any PP rank (e.g. MoE modules on dense layers via moe_layer_freq).

Suggested test cases (all present in this PR):

• test_megatron_to_hf_skips_param_no_pp_rank_owns - unowned param exports as {} and nothing is broadcast
• test_megatron_to_hf_uses_type_broadcast_by_owning_pp_rank - non-owning rank still receives the parallelism type from the owner
• test_broadcast_obj_from_pp_rank_raises_when_unowned_by_default - default (allow_missing=False) still raises when no rank owns the object

No perf tests impacted.

[pruprakash]

/ok to test 49d95e2

[cuichenx]

Reassessed against current main, including #6343. The historical NemotronH failure no longer reproduces: PP1/PP2 round trips and independent cold exports preserve all 29 expected keys, and strict HF reload produces identical finite, nonconstant logits on a tiny random checkpoint.

The locally rebased allow_missing change did turn a deliberately invalid PP2 task into an empty export. I replaced it with a parameter-named error at both PP1 and PP2, including the quantized API entry, and removed the misleading MoE comment. #5262 validates mappings and checkpoint keys, but cached names and supplied conversion tasks mean that is not a universal ownership guarantee.

Validation: 454 focused/adjacent tests passed (two hardware-specific MXFP8 skips), real distributed valid-owner and invalid-owner controls passed, 11 final regression cases passed, and repository hooks passed. Independent review approved the final change. This is small-model HF round-trip evidence, not pretrained-model or Megatron-forward parity coverage.

[cuichenx]

/ok to test fb24c1e

[github-actions]

Automatic Claude reviews have been retired. To request a pull-request review, post a comment containing:

/review

Add model=claude to use a Claude reviewer (the default is model=codex). mode=light|strict selects the review depth; for example, /review model=claude mode=strict. Comment /review help for all options.

[qiaochuz-nv]

Retrospective QA-fix tracking issue: #6399. It records the fix authored by @pruprakash and delivered by this merged PR.