Skip to content

fix(ckpt): skip export of parameters no PP rank owns - #5585

Open
pruprakash wants to merge 1 commit into
mainfrom
pruprakash/fix_ckpt_automapping_pp_skip_unowned_param
Open

fix(ckpt): skip export of parameters no PP rank owns#5585
pruprakash wants to merge 1 commit into
mainfrom
pruprakash/fix_ckpt_automapping_pp_skip_unowned_param

Conversation

@pruprakash

Copy link
Copy Markdown
Contributor

What does this PR do ?

Fixes every Megatron-to-HF export at pipeline_model_parallel_size > 1 crashing with ValueError: Object must exist on at least one PP rank whenever a Megatron parameter has no HF counterpart, by letting the PP broadcast return None so the graceful-skip path that was already written can actually run.

Changelog

  • src/megatron/bridge/models/conversion/param_mapping.py: broadcast_obj_from_pp_rank gains an allow_missing flag that returns None instead of raising when no PP rank owns the object; AutoMapping.megatron_to_hf passes it on the receive path and drops its now-redundant local None check.
  • tests/unit_tests/models/test_param_mapping.py: 3 tests covering the patched branch — no owner with allow_missing returns None, no owner without it still raises, and an owned object still broadcasts unchanged.

GitHub Actions CI

See the CI section in the Contributing doc for how to trigger the CI. A Nvidia developer will need to approve and trigger the CI for external contributors.

Before your PR is "Ready for review"

Pre checks:

  • Make sure you read and followed Contributor guidelines
  • Did you write any new necessary tests?
  • Did you add or update any necessary documentation?
  • Does the PR affect components that are optional to install? (Ex: Numba, Pynini, Apex etc)
    • Reviewer: Does the PR have correct import guards for all optional libraries?

Additional Information

  • Root cause: AutoMapping.megatron_to_hf (param_mapping.py:1589-1602) is written to skip a parameter that no PP rank detected a type for — the comment and the if self._detected_type is None: return {} guard are already there. That guard was unreachable at pp_size > 1 because broadcast_obj_from_pp_rank (param_mapping.py:472-479) raised as soon as its all-gather found no owner, before the caller could inspect the result. At pp_size == 1 the function short-circuits with return obj, _detected_type stays None, and the parameter is skipped correctly — which is why the same model exports at PP=1 and crashes at PP=2.

  • Blast radius: any Megatron-to-HF export at pp_size > 1 where a Megatron parameter has no HF counterpart. Reproduced on NemotronH, where a stale backbone.embeddings.weight mapping in the bridge leaves the embedding unmatched under transformers 5.12.1; the same shape is reachable for MoE modules on dense layers created by moe_layer_freq, which is the case the existing comment names.

  • Regression? No — not introduced by any recent PR. The raising broadcast_obj_from_pp_rank and the unreachable guard are both pre-existing. Found while running coverage for perf(ckpt): cache config broadcasts for Mamba and GDN mappings #5367 (perf(ckpt): cache config broadcasts for Mamba and GDN mappings), but that PR is not implicated: the failing call is AutoMapping's own broadcast_obj_from_pp_rank(None, "detected_type"), none of the four Mamba/GDN call sites perf(ckpt): cache config broadcasts for Mamba and GDN mappings #5367 touched, and perf(ckpt): cache config broadcasts for Mamba and GDN mappings #5367's own contract passes in the same run (12/12 qwen3_5 and qwen3_next cases).

  • Verification: red-green in nvcr.io/nvidian/nemo:nightly, EOS job 5850574 — RED 1 failed with the fix reverted, GREEN 16 passed with it applied. The failing e2e that found it is test_pp_config_broadcast_cache_export (2 GPU, pp_size=2, NemotronH hybrid_override_pattern="M-M*", random weights, no Hub download), which goes 6 failed → passing with this change.

  • NVBug: not filed at time of writing; report at bugs/hermes/2026-08-14_test_pp_config_broadcast_cache_export_bug.md. Add the qa_rcca_done label once it is filed.

  • Known adjacent defect, deliberately NOT fixed here: nemotron_h_bridge.py:421 maps embedding.word_embeddings.weight to backbone.embeddings.weight, a key transformers 5.12.1 does not produce, so the embedding is silently dropped from the export at every PP size. That is a wrong checkpoint rather than a crash and wants its own fix; this PR only stops the crash it triggers.

  • Related to perf(ckpt): cache config broadcasts for Mamba and GDN mappings #5367 (context only — the defect is independent of it)

Megatron-to-HF export crashed at pipeline_model_parallel_size > 1 with
"Object must exist on at least one PP rank" whenever a Megatron parameter
had no HF counterpart, so every rank reached AutoMapping.megatron_to_hf
with megatron_module=None. The graceful-skip path for that case was
already written but unreachable: broadcast_obj_from_pp_rank raised before
it could return None, while at PP=1 it short-circuits and the parameter is
skipped correctly.

Adds a unit test covering the change (red-green verified).
Detected by: megatron-bridge QA

Signed-off-by: Pruthviraj Prakash <pruprakash@nvidia.com>
@copy-pr-bot

copy-pr-bot Bot commented Aug 14, 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.

@claude

claude Bot commented Aug 14, 2026

Copy link
Copy Markdown
Contributor

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

Copy link
Copy Markdown
Contributor Author

/ok to test 49d95e2

@pruprakash pruprakash 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 qa_rcca_done labels Aug 14, 2026
@pruprakash
pruprakash requested a review from yaoyu-33 August 14, 2026 19:47
@yaoyu-33 yaoyu-33 added full-test-suite needs-more-tests Requires additional L0 and L1 test coverage before merge labels Aug 15, 2026
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 full-test-suite 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 qa_rcca_done

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants