[Connor-XY]

This PR moves the training loop's I/O onto the model's ProcessGroupCollection. pretrain() builds one collection and passes it to DDP, the schedules, batch loading, the data loaders and checkpoint save and load. The GPT and ModelOpt spec builders take the collection's pipeline rank. A model built on its own grid, such as an RL inference model or a distillation teacher, then loads, builds and checkpoints the right stage on the right ranks. On the standard grid, results are unchanged. It consolidates six draft PRs from #6307 without changes. Each member is one commit, so review commit by commit. It does not depend on any other unit. It lands after the optimizer unit #8889, which also edits training.py and checkpointing.py, and before the FP8/FP4 amax, CUDA graphs and 1F1B unit #8893, which also edits training.py.

Review

The unit's own commits: Connor-XY/Megatron-LM@eb768e2...74906a8

The full PR diff also includes #7487, #6259, #6293, #6257, #6303, the optimizer unit #8889 and the core process-group API unit #8890 until they merge.

Members

• Pass one process-group collection through the pretrain loop #8037: Pass one process-group collection through the pretrain loop. pretrain() builds the collection once and passes it to every DDP wrapper, the optimizer (when Gloo groups are off), schedule selection, the schedules with a P2PCommunicator over its PP group, and the parameter norm. It also adds the all-gather groups to the collection.
• Pass the model's pipeline rank to spec builders outside core #8048: Pass the model's pipeline rank to spec builders outside core. gpt_builder, GPTModelBuilder, the ModelOpt builder, the distillation teacher and the multimodal example pass their collection's pipeline rank to the GPT block spec builders, and default_layer_spec takes an optional pp_rank.
• Pass the pipeline rank to the ModelOpt GPT spec #8726: Pass the pipeline rank to the ModelOpt GPT spec. get_gpt_modelopt_spec takes vp_stage and pp_rank and forwards them to the layer-count and offset helpers. Its builder edits are identical to Pass the model's pipeline rank to spec builders outside core #8048's, so the commit keeps one copy.
• Raise ValueError for missing sequence-packing process groups #8727: Raise ValueError for missing sequence-packing process groups. wrap_data_iterator raises ValueError naming each missing group instead of asserting, and rejects a None collection.
• Save and load checkpoints on the model's process groups #8740: Save and load checkpoints on the model's process groups. save_checkpoint and load_checkpoint take pg_collection, defaulting to the model's own collection, for shard names, RNG state, replica ids, writers and fully parallel save and load groups.
• Use the model's process groups in get_batch and the pretraining loaders #8742: Use the model's process groups in get_batch and the pretraining loaders. get_batch in pretrain_gpt.py, pretrain_hybrid.py and pretrain_vlm.py and the pretraining data loaders read their TP, PP, CP and DP groups from the collection.

Behavior changes for callers

• --create-all-gather-group now takes effect under pretrain(). All-gather groups whose ranks differ from the data-parallel groups are rejected.
• wrap_data_iterator raises ValueError where it raised AssertionError, or AttributeError for a None collection.
• GPT and hybrid get_batch raise ValueError for a collection with cp=None. pretrain() always sets cp.
• With GTP remat, checkpoint RNG state and legacy writers use the replicate dp and expt_dp, as the training loop already did. With expert GTP remat, the fully parallel expert save and load group becomes expt_dp.
• fsdp_dtensor loads read each rank's own rng_state.(pp, tp) instead of (0, 0). A strict load into a different TP/PP layout fails on the missing key.
• Legacy shard names follow the model's collection, not the _MPU_* setters. External converters that emulate stages with those setters must also set their stand-in pipeline group's rank, as tools/checkpoint/loader_base.py now does.
• With a virtual pipeline and account_for_embedding_in_pipeline_split or account_for_loss_in_pipeline_split, the ModelOpt spec no longer gives the first or last chunk one extra layer.

Teams

core-adlr, core-nemo, training, datasets, post-training. Needs the determinism-exempt label: megatron/training/utils/common_utils.py is a registered kernel file, and only its stage predicate changes.

Testing

Each member was GPU-tested against its base before consolidation. The consolidated chain (every unit in landing order, on current main) was run against main on 4 GPUs: 258 test files, and no test fails on the chain that passes on main. Tests in 28 files fail on both, mostly tests written for 8 GPUs, and one file times out on both.

🤖 Generated with Claude Code

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