Use metadata dp_cp_group for the tied-embeddings replica_id

[going-song]

Motivation

MimoModel runs each submodule on its own process groups (per-module
HyperCommGrids with different parallel layouts), so the global MPU may
never be initialized. In that setup, saving a checkpoint of a language
module with share_embeddings_and_output_weights=True and PP >= 2 fails on
the last pipeline stage:

AssertionError: data parallel group with CP (with GTP_remat) is not initialized

The failing line is the tied output-layer copy's replica id in
LanguageModule.tie_embeddings_and_output_weights_state_dict, which reads
the global parallel_state.get_data_parallel_rank(with_context_parallel=True).

Change

In megatron/core/models/common/language_module/language_module.py, compute
the DP replica index from the process group threaded through
metadata['dp_cp_group']:

-            parallel_state.get_data_parallel_rank(with_context_parallel=True),
+            get_pg_rank(metadata['dp_cp_group']),

The key is guaranteed to be present: LanguageModule.sharded_state_dict
calls ensure_metadata_has_dp_cp_group at its entry, before the tying step,
and the make_tp_sharded_tensor_for_checkpoint call registering this same
tensor a few lines below already consumes metadata['dp_cp_group'].
MimoModel.sharded_state_dict fills the key with each module's dp_cp
group (megatron/core/models/mimo/model/base.py), which makes the save
work without the global MPU.

This follows #2053, which moved this file's other reads (the
sharded_state_dict entry guard and the MTP tying path) to the metadata
group — this was the file's remaining direct parallel_state read — and
the guideline to prefer caller-provided process groups over global
parallel_state reads in megatron/core.

Behavior for existing setups

• With the global MPU initialized, metadata['dp_cp_group'] holds the same
  group the removed call reads — ensure_metadata_has_dp_cp_group fills it
  from get_data_parallel_group(with_context_parallel=True), and the
  trainer's _build_sharded_state_dict_metadata
  (megatron/training/checkpointing.py) threads the same group — so the
  computed rank is unchanged.
• The value also cannot affect save/load decisions either way: the tuple's
  leading 1 ("copy of first stage embedding") already makes
  is_main_replica False regardless of the third component, and
  _sharded_tensor_shard_id does not include replica_id.

Why this hasn't surfaced before

• With the global MPU initialized (all standard trainers), both sources
  yield the same value.
• Reaching the line requires tied embeddings, PP >= 2 and a checkpoint
  save; with PP == 1 the function returns early (the stage holding the
  embedding never registers a tied copy).
• The MIMO checkpoint tests currently exercise untied models only
  (share_embeddings_and_output_weights defaults to False), so this
  combination is not covered there yet.

Testing

Verified on 3 ranks with per-module HyperCommGrids and no global MPU
(llm tp1/pp2/dp1 + one encoder rank): unpatched, sharded_state_dict()
raises the assertion on the last language PP stage; with this change, save
(validate_access_integrity=True) and a load round trip pass, and the
loaded tied weight equals the first-stage embedding. Adds a tied-embeddings
case to tests/unit_tests/models/mimo/test_mimo_checkpoint.py.

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

[going-song]

Same source as the first-stage copy: when replica_id is not given, make_tp_sharded_tensor_for_checkpoint (megatron/core/utils.py) computes this component as get_pg_rank(dp_cp_group) from its dp_cp_group argument, which the standard path fills from metadata['dp_cp_group'].

[wujingyue]

/claude fix

[svcnvidia-nemo-ci]

🛠️ Claude fix commit 88163fa775b4 (attempt 1)

⚠️ This explanation is AI-generated and may be inaccurate; the exact commit is authoritative.

What changed
No edits made. The base merge in pr-head／ is already fully staged and resolved: no unmerged paths, no conflict markers, clean worktree. A read-only git merge-tree re-run produced a tree byte-identical to the staged index, confirming the resolution is faithful. Focused checks pass: all 3 changed files compile, the new get＿pg＿rank import resolves, and metadata('dp＿cp＿group') is populated before use.

Files changed by Claude

• No additional file edits; the pinned base was merged.

Why
Nothing to repair. Git's own merge of the two sides is clean and identical to the staged tree, so the overlapping edits in test＿mimo＿1f1b＿schedule.py were disjoint and needed no manual fix. The CI log dir is empty and the job name blank, so no lint or unit failure was available to diagnose; steering file is empty. Unit tests are unrunnable here (no torch／pytest／GPU; suite needs 8).

DCO
@going-song, please fix any DCO failures on your commits before merge. DCO does not block this workflow.

View exact commit

Sanitized and posted by svcnvidia-nemo-ci.

[svcnvidia-nemo-ci]

/ok to test 88163fa

[svcnvidia-nemo-ci]

❌ Claude fix stopped because the pull request head or base changed. View exact-SHA CI.