fix(ckpt): skip export of parameters no PP rank owns - #5585
Open
pruprakash wants to merge 1 commit into
Open
Conversation
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>
Contributor
|
LGTM - clean, well-tested fix. The asymmetric Suggested test cases (all present in this PR):
No perf tests impacted. |
Contributor
Author
|
/ok to test 49d95e2 |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
What does this PR do ?
Fixes every Megatron-to-HF export at
pipeline_model_parallel_size > 1crashing withValueError: Object must exist on at least one PP rankwhenever a Megatron parameter has no HF counterpart, by letting the PP broadcast returnNoneso the graceful-skip path that was already written can actually run.Changelog
src/megatron/bridge/models/conversion/param_mapping.py:broadcast_obj_from_pp_rankgains anallow_missingflag that returnsNoneinstead of raising when no PP rank owns the object;AutoMapping.megatron_to_hfpasses it on the receive path and drops its now-redundant localNonecheck.tests/unit_tests/models/test_param_mapping.py: 3 tests covering the patched branch — no owner withallow_missingreturnsNone, 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:
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 theif self._detected_type is None: return {}guard are already there. That guard was unreachable atpp_size > 1becausebroadcast_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. Atpp_size == 1the function short-circuits withreturn obj,_detected_typestaysNone, 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 > 1where a Megatron parameter has no HF counterpart. Reproduced on NemotronH, where a stalebackbone.embeddings.weightmapping in the bridge leaves the embedding unmatched under transformers 5.12.1; the same shape is reachable for MoE modules on dense layers created bymoe_layer_freq, which is the case the existing comment names.Regression? No — not introduced by any recent PR. The raising
broadcast_obj_from_pp_rankand 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 isAutoMapping's ownbroadcast_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/12qwen3_5andqwen3_nextcases).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 istest_pp_config_broadcast_cache_export(2 GPU,pp_size=2, NemotronHhybrid_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 theqa_rcca_donelabel once it is filed.Known adjacent defect, deliberately NOT fixed here:
nemotron_h_bridge.py:421mapsembedding.word_embeddings.weighttobackbone.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)