Integrate GTP with heterogeneous MIMO training - #6268
Closed
yashaswikarnati wants to merge 8 commits into
Closed
Conversation
Signed-off-by: Shiqing Fan <shiqingf@nvidia.com>
- 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>
Signed-off-by: Shiqing Fan <shiqingf@nvidia.com>
hasattr() is always True on ProcessGroupCollection: __getattr__ returns None for
declared-but-unset fields. resolve_gtp_remat_group() gated its MPU fallback on
hasattr, so the fallback was unreachable.
* A collection omitting gtp_remat resolved to None -> weights silently built
UNSHARDED under GTP. Check `attr in vars(...)` instead.
* Add test_pg_collection_without_gtp_remat_falls_back_to_mpu. The existing test
always set gtp_remat explicitly, so it could not catch this.
Signed-off-by: Shiqing Fan <shiqingf@nvidia.com>
Signed-off-by: ykarnati <ykarnati@nvidia.com>
Signed-off-by: ykarnati <ykarnati@nvidia.com>
Signed-off-by: ykarnati <ykarnati@nvidia.com>
Signed-off-by: ykarnati <ykarnati@nvidia.com>
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
3 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.
Purpose
Integration-only draft for end-to-end review and validation. The production changes are split into:
The earlier draft #6249 remains untouched as a reference.
Acceptance
Validation