Repository navigation
Conversation
Contributor
Author
|
/ok to test 614a403 |
|
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. |
pretrain() built the model on a ProcessGroupCollection, but the rest of the loop reached Megatron Core's fallbacks to the global grid. Standard DDP and Torch FSDP2 were wrapped without the collection (only Megatron-FSDP received it), get_megatron_optimizer read the global groups, the schedule was selected from the global pipeline sizes and ran on a default collection and communicator that the schedules build from parallel_state, and the parameter norm used the train() argument, which is None in the stock path. A collection passed to pretrain() was therefore honored only by the model. --create-all-gather-group was also inert: get_model() created the all-gather groups only when it received no collection, and pretrain() always passes one. Build the collection once, right after initialize_megatron, and pass it to every Megatron Core entry point of the loop: all DDP wrappers (Torch FSDP2 shards over the collection's DP x CP group), the optimizer when Gloo groups are off (the collection carries no Gloo groups yet), the schedule selection and the schedules together with a P2PCommunicator over the collection's pipeline group, and the parameter norm. The collection carries the all-gather groups when they are requested, and a check rejects all-gather groups that do not span the data-parallel ranks. On the standard grid every consumer gets the same group objects as before, so training is bit-identical. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Signed-off-by: Yan Xu <yxu1@nvidia.com>
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Signed-off-by: Yan Xu <yxu1@nvidia.com>
Connor-XY
force-pushed
the
yx/pstate-b2-train-one-collection
branch
from
October 26, 2026 22:26
614a403 to
0c6841b
Compare
Contributor
Author
|
/ok to test 0c6841b |
This was referenced Oct 30, 2026
This was referenced Oct 30, 2026
Connor-XY
added a commit
to Connor-XY/Megatron-LM
that referenced
this pull request
Nov 1, 2026
The end-to-end seeding test that follows checks that the collection pretrain builds after initialize_megatron holds the same groups as the collection the RNG was seeded from. It builds that collection with training._build_pg_collection_from_parallel_state, the helper that NVIDIA#8037 adds. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Signed-off-by: Yan Xu <yxu1@nvidia.com>
Connor-XY
added a commit
to Connor-XY/Megatron-LM
that referenced
this pull request
Nov 1, 2026
NVIDIA#8031, NVIDIA#8037 and NVIDIA#8039 each refreshed the allowlist against main. Merging them keeps the total of 195 that NVIDIA#8031 and NVIDIA#8037 share, since NVIDIA#8039's total equals main's, while the merged entries record 194 reads. Regenerate the file so that its total matches its entries. No entry changes. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Signed-off-by: Yan Xu <yxu1@nvidia.com>
Connor-XY
added a commit
to Connor-XY/Megatron-LM
that referenced
this pull request
Nov 1, 2026
The optimizer change follows the one collection that NVIDIA#8037 passes through pretrain into setup_model_and_optimizer, and its test extends NVIDIA#8037's training collection test. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Signed-off-by: Yan Xu <yxu1@nvidia.com>
Connor-XY
added a commit
to Connor-XY/Megatron-LM
that referenced
this pull request
Nov 1, 2026
NVIDIA#8738 passes the training collection to the optimizer with Gloo groups on, on top of NVIDIA#6259, NVIDIA#8037 and NVIDIA#8055 (which carries NVIDIA#8049 and NVIDIA#6293). The training loop's changes that follow call NVIDIA#8037's collection builder, warn through NVIDIA#6293's and NVIDIA#8049's fallback helpers, and extend NVIDIA#8037's training collection test. NVIDIA#8737, merged above, now carries NVIDIA#8037 as well, so the allowlist conflicts on its total: NVIDIA#8737's side records 194 reads and NVIDIA#8738's 158. It takes 158, which matches the merged entries. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Signed-off-by: Yan Xu <yxu1@nvidia.com>
Connor-XY
added a commit
to Connor-XY/Megatron-LM
that referenced
this pull request
Nov 1, 2026
The data-parallel group helper that follows replaces the DP x CP helper that NVIDIA#8741 adds to the training loop, and the training utilities' warnings extend NVIDIA#8741's packed-sequence statistics and GPT builder code. NVIDIA#8741 is based on main, without NVIDIA#8037's and NVIDIA#8048's neighbouring lines, so three files conflict. They are resolved as follows: - training.py: keep NVIDIA#8037's _build_pg_collection_from_parallel_state and _get_forward_backward_func next to NVIDIA#8741's _get_model_pg_collection and _get_data_parallel_world_size. train_step's pg_collection docstring takes NVIDIA#8037's text plus NVIDIA#8741's sentence on the single-grid fallback; its p2p_communicator docstring and evaluate's keep NVIDIA#8037's text. train's docstring also keeps NVIDIA#8037's text, with NVIDIA#8741's sentence that a multi-module collection rejects the GPU sniff test and its Raises section. - models/gpt.py: default_layer_spec takes (config, vp_stage, pp_rank=None, pg_collection=None), and GPTModelBuilder passes both NVIDIA#8048's pp_rank and NVIDIA#8741's pg_collection. - test_gpt_builder.py: keep NVIDIA#8741's ModelOpt CP test and expect default_layer_spec(config, 3, 1, pg_collection=pg). Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Signed-off-by: Yan Xu <yxu1@nvidia.com>
Conflicts: - megatron/training/training.py: Main adds the tensor metric observer, the run_train_loop_without_updates argument, the batch-size transition helper, new train() docstring entries, and moves the RL inference-model setup out of pretrain() (NVIDIA#4264). The PR builds the process-group collection once right after initialization and passes it everywhere. The result keeps both sides. The setup_model_and_optimizer call passes pg_collection directly, without main's use_mpu_process_groups() fallback, and also passes run_train_loop_without_updates. The docstrings combine main's new entries with the PR's pg_collection and p2p_communicator wording. The process-group allowlist records the reads main added since this PR was written. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Signed-off-by: Yan Xu <yxu1@nvidia.com>
Contributor
Author
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.
Why
pretrain()builds the model on aProcessGroupCollection, but the rest of the loop used the global grid: standard DDP and Torch FSDP2 got no collection,get_megatron_optimizerread the globals, the schedule was chosen from the global PP/VP sizes and ran on a default collection and communicator built fromparallel_state, and the parameter norm gotNone. A collection passed topretrain()reached only the model, and--create-all-gather-groupdid nothing underpretrain().What
initialize_megatronand pass it to every DDP wrapper (Torch FSDP2 gets its DP x CP group), the optimizer when Gloo groups are off, schedule selection and the schedules (with aP2PCommunicatorover its PP group), and the parameter norm.Tests
tests/unit_tests/training/test_training_pg_collection.py(8 cases, all fail onmain).main), 4 GPUs:test_training_pg_collection.pypasses all 11 cases it holds there: 6 of this PR's 8, plus 5 added by later [TRACKING] Deprecate parallel_state: explicit process groups in megatron/core #6307 PRs, one of which replaces the two all-gather rank-mismatch cases.mainhas no such file. In the integration run, this PR's test file shows no regressions against main.--deterministic-mode: loss, grad norm, params norm and evaluation loss are bit-identical tomainover 12 iterations for TP2 PP2 SP, PP2 VPP2, MoE EP2 PP2, distributed optimizer (PP2; 2 instances), Megatron-FSDP with--create-all-gather-group(dense, MoE) and Torch FSDP2, via both the model-builder andmodel_providerpaths. The schedules' default builder is no longer reached.finalize_model_gradsread is unchanged.Compatibility
No API change;
train()/evaluate()without a collection behave as before. With Gloo groups (default), the optimizer still builds its groups fromparallel_stateuntil the collection carries them.Synced with main 77391e7 on 2026-11-01 by a merge commit that lists each conflict and its resolution, so existing commits and review links stay valid.
Part of #6307.
🤖 Generated with Claude Code