[cuichenx]

• I, the PR author, have personally reviewed every line of this PR.

Depends on #8010 for shared-embedding initialization with meta-device construction. This PR contains no replacement initialization logic. Keep it draft until that dependency is available.

What does this PR do?

Fix Megatron-FSDP v1 mesh construction and checkpoint restart with non-interleaved pipeline parallelism.

With eight ranks and PP=2, the mesh helpers reshape all eight ranks using stage-local dimensions whose product is four. Include the outer PP axis in both FSDP and HSDP construction, then select the current stage and expert rank. Existing process-group consistency checks remain in place.

FSDP DTensor checkpointing also bypasses the block sharded_state_dict() methods that normally give pipeline layers global names. Two stages can consequently save different weights and Adam moments under the same decoder.layers.0.* keys. Add a shared handler that uses each layer's global layer_number for model tensors, optimizer state, and optimizer group metadata. Run it after fused-parameter, expert, and MTP processing on save and load, preserving the original tensor aliases for in-place loading. MTP retains its separate namespace.

Precision-aware v1 optimizers put gradients in decoupled_grad. Recognize that attribute when reducing shared embeddings across pipeline stages. Empty shards remain None for the optimizer, while synchronization retrieves their metadata-bearing buffer tensors so every DP rank participates.

Existing checkpoints affected by stage-local key collisions cannot recover the omitted weights and must be regenerated. The new checkpoint remapping applies only to v1; v2 optimizer keys are already global and must not be renamed again. PP=1 checkpoint names remain unchanged. The existing default missing-weight validation rejects incomplete legacy checkpoints.

Validation

• Previously passed unchanged coverage: 26 mesh and 16 checkpoint cases on each of eight H100 ranks, plus six Bridge checkpoint tests. Includes PP=1/2/4, uneven stages, CP/EP/TP/ETP and HSDP mesh combinations, and preservation of v2 optimizer keys.
• The revised 96-case embedding-gradient test passed on every rank of an eight-H100 run. Its eager fixture isolates the retained gradient fix across ordinary/decoupled gradients, PP=1/2, TP=1/2, tied/MTP/unshared cases, BF16 wrappers, and verified empty shards. Together with unchanged mesh/checkpoint coverage, 138 Core cases passed per rank.
• Small Lightning with Sync tied embedding copies under meta-device init #8010 (5920bd95) plus this PR: eight hybrid layers, MTP=2, PP=2/EP=2/TP=1, BF16 compute and FP32 master weights/Adam moments. Three training steps, save at step two, then restart: all saved weights and optimizer tensors restored bitwise exactly on all eight ranks. The next loss (5.374913215637207) and gradient norm (2.875216245651245) match uninterrupted training exactly. After that step, 281/287 complete model/Adam tensors are bitwise equal; all pass rtol=1e-5/atol=1e-6, maximum absolute difference 4.331e-8.
• Full 52-layer Lightning with Sync tied embedding copies under meta-device init #8010 plus Core fix(distributed): restore FSDP v1 pipeline training and checkpoints #8159: three forward/backward/optimizer steps completed on eight H100s, PP=2/EP=2/TP=1/CP=1, original widths/128 experts/vocabulary, MTP=2 and synthetic sequence length 256. Loss 12.176648 → 10.490705 → 9.382988, no skipped/NaN iterations; peak allocated memory 69.898 GiB. BF16 Adam moments fit this full-size smoke on one node; checkpoint restart was validated with the smaller FP32-moment run above.
• tools/autoformat.sh check mode passed Black/isort/pylint/ruff. Mypy is not installed in the pinned lint group and is optional in that script.

Training probes use function-level patch injection into NeMo 26.08.01, including #8010 and the existing main-branch embedding-gradient handling from #6486. Distributed unit tests use the full proposed Core source based on cab97ebd1.

Training probes set deallocate_pipeline_outputs=False, since v1's backward hooks return output views rejected by pipeline output deallocation. Coverage is non-interleaved PP; these short synthetic-data tests do not establish convergence or performance.

Focused Core command: torchrun --standalone --nproc_per_node=8 -m pytest tests/unit_tests/distributed/mfsdp_v1/test_fsdp_mesh.py tests/unit_tests/distributed/mfsdp_v1/test_pipeline_checkpoint.py tests/unit_tests/distributed/mfsdp_v1/test_pipeline_embedding_grads.py -q.

Independent subagent review of the removal and remaining diff completed; no source blockers.

Contribution checks

• Added distributed unit regressions for mesh group membership, PP=1/2/4 checkpoint round trips, uneven stages, wrapper prefixes, MTP/expert keys, optimizer aliases, legacy checkpoint rejection, and shared embedding gradients with real BF16 wrappers, distinct TP partitions, ordinary/decoupled gradients, actual backward reduction, and verified empty shards.
• Ran tools/autoformat.sh in check mode (Black/isort/pylint/ruff passed; mypy is not installed in the pinned lint group and is optional in that script).
• Documented behavior and compatibility in the helper docstring and this description.
• No GPU kernel changes.

Bridge requires the companion call-site change: NVIDIA-NeMo/Megatron-Bridge#6424.

[copy-pr-bot]

Auto-sync is disabled for draft pull requests in this repository. Workflows must be run manually.

Contributors can view more details about this message here.

[wujingyue]

Can you split? This PR seems to contain several changes. We should be able to separate the checkpointing fix at least. I guess the changes in finalize_model_grads and the changes in adapter are separatable as well.