[fix] Cleanup usage of pg_collection in gtp - #6234
Merged
fanshiqing merged 5 commits intoAug 5, 2026
Merged
Conversation
Signed-off-by: Shiqing Fan <shiqingf@nvidia.com>
svcnvidia-nemo-ci
marked this pull request as draft
August 4, 2026 03:04
Contributor
|
This PR has been automatically converted to draft because all PRs must start as drafts. When you are ready for review, click Ready for Review to begin the review process. This will:
See the contribution guide for more details. |
fanshiqing
changed the title
fanshiqing
marked this pull request as ready for review
August 4, 2026 03:13
Member
Author
|
/ok to test 4e87571 |
fanshiqing
commented
Aug 4, 2026
- Thread pg_collection from TransformerLayer through Attention/MLP into the TE and core linear layers. - Assert in finalize_model_grads when the config enables the axis but the group is missing, instead of treating it as "GTP inactive" and training on wrong gradients. - Add test_gtp_custom_pgs.py: two TransformerBlocks, identical degrees (TP=1, CP=1, GTP_remat=2 over world=4), identical weights and input, one from the MPU groups and one from a custom collection with permuted gtp_remat rank membership. Output and gradients must match; verified it fails without the plumbing. Signed-off-by: Shiqing Fan <shiqingf@nvidia.com>
Member
Author
|
/ok to test 2c47a2e |
fanshiqing
force-pushed
the
shiqingf/gtp_refine_usage_of_pg_collection
branch
from
August 5, 2026 02:44
a3943e3 to
07b0066
Compare
Member
Author
|
/ok to test 07b0066 |
Signed-off-by: ykarnati <ykarnati@nvidia.com>
fanshiqing
force-pushed
the
shiqingf/gtp_refine_usage_of_pg_collection
branch
from
August 5, 2026 02:49
07b0066 to
382835b
Compare
Member
Author
|
/ok to test 382835b |
yashaswikarnati
approved these changes
Aug 5, 2026
yaox12
approved these changes
Aug 5, 2026
Member
|
/claude strict-review |
![claude[bot]](https://avatars.githubusercontent.com/in/1236702?s=60&v=4){kind=small}
Contributor
There was a problem hiding this comment.
Strict review passed — no significant issues found. LGTM
Reviewed as a pg_collection plumbing refactor for GTP-backed layers. Verified:
- Correctness:
resolve_gtp_remat_groupcorrectly usesvars()nothasattr(dataclass__getattr__returns None for unset fields, sohasattrwould always be True and defeat the MPU fallback) — covered by the fall-back test. The new fail-loud assertion infinalize_model_gradsturns a silent wrong-gradient skip into a hard error when*_remat_size > 1but the group is missing/mismatched. - Backward compat: all new
pg_collectionargs are optional and MPU-fallback preserves pre-pg_collection callers; the default training path stays safe because_build_default_pg_collectionnow populates the gtp_remat fields; the one internal signature change (_allreduce_replicated_grads_over_gtp_remat_group) has both callers updated. - Unused-arg check: every added
pg_collectionparameter has a real use path (resolve or downstream propagation); the newgtp_remat_rank/egtp_remat_rankseed args are wired and defaulted. - Cleanup: removed
HAVE_GTPimport andMEGATRON_GTP_FORCE_ENABLEenv sets are confirmed dead across the tree. - Process-group guidance: the direct
parallel_state.get_*_group()reads are confined to the existing MPU-fallback branch — a documented compatibility point.
Strong new test coverage: the permuted-pairing custom-group test precisely targets the silent 'module reads global groups instead of the passed collection' failure this refactor prevents.
Phlip79
approved these changes
Aug 5, 2026
kvareddy
approved these changes
Aug 5, 2026
mathemakitten
approved these changes
Aug 5, 2026
deepakn94
reviewed
Aug 5, 2026
| embd_group = parallel_state.get_embedding_group(check_initialized=False) | ||
| pos_emb_group = parallel_state.get_position_embedding_group(check_initialized=False) | ||
| dp_cp_group = parallel_state.get_data_parallel_group(with_context_parallel=True) | ||
| gtp_remat_group = parallel_state.get_gtp_weight_remat_group(check_initialized=False) |
Contributor
There was a problem hiding this comment.
@yashaswikarnati when are we getting rid of all of this fallback code?
Contributor
There was a problem hiding this comment.
Not only for GTP, but everywhere. It's an eyesore :P
deepakn94
approved these changes
Aug 5, 2026
yaoyu-33
approved these changes
Aug 5, 2026
Contributor
|
🔄 Merge queue validation started! You can track the progress here: https://github.com/NVIDIA/Megatron-LM/actions/runs/31029975296 |
Contributor
|
🔄 Merge queue validation started! You can track the progress here: https://github.com/NVIDIA/Megatron-LM/actions/runs/31034228133 |
This was referenced Aug 5, 2026
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
9 participants
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?
Refine usage of
pg_collectionin gtp, which is a planned following up tiny task of #4967Contribution process
Pre-checks