Add intra_group_size to topology #3696
Open
+28
−5
Conversation
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
meta-cla
bot
added
the
CLA Signed
Contributor
|
@Ali-Tehrani has exported this pull request. If you are a Meta employee, you can view the originating Diff in D91617887. |
Summary:
Context
---------
Every GB200 node has 2 B200 GPU attached to it, however allows up to 72 B200 connected via NVlink. The planner needs to know how big the intra topology group size is going to be.
This causes the `local_world_size` to be different from the `intra_group_size`.
Implementation
------------------
- Topology class:
- Adds `pod_size`, and uses that to calculate the `intra_group_size` (maximum number of processes linked with high intra bandwidth) to Topology class. If isn't given, then it defaults to local_world_size.
- `shard_estimators.py`
- The shard estimators now use the `intra_group_size` instead of `local_world_size`, this allows RW/TW/CW to properly account for larger NVlink that comes with the pods.
Reviewed By: isururanawaka
Differential Revision: D91617887
Ali-Tehrani
force-pushed
the
export-D91617887
branch
from
2128c83 to
9d1ba65
Compare
Ali-Tehrani
added a commit
to Ali-Tehrani/torchrec
that referenced
this pull request
Jan 30, 2026
Summary:
Context
---------
Every GB200 node has 2 B200 GPU attached to it, however allows up to 72 B200 connected via NVlink. The planner needs to know how big the intra topology group size is going to be.
This causes the `local_world_size` to be different from the `intra_group_size`.
Implementation
------------------
- Topology class:
- Adds `pod_size`, and uses that to calculate the `intra_group_size` (maximum number of processes linked with high intra bandwidth) to Topology class. If isn't given, then it defaults to local_world_size.
- `shard_estimators.py`
- The shard estimators now use the `intra_group_size` instead of `local_world_size`, this allows RW/TW/CW to properly account for larger NVlink that comes with the pods.
Reviewed By: isururanawaka
Differential Revision: D91617887
Ali-Tehrani
added a commit
to Ali-Tehrani/torchrec
that referenced
this pull request
Jan 30, 2026
Summary:
Context
---------
Every GB200 node has 2 B200 GPU attached to it, however allows up to 72 B200 connected via NVlink. The planner needs to know how big the intra topology group size is going to be.
This causes the `local_world_size` to be different from the `intra_group_size`.
Implementation
------------------
- Topology class:
- Adds `pod_size`, and uses that to calculate the `intra_group_size` (maximum number of processes linked with high intra bandwidth) to Topology class. If isn't given, then it defaults to local_world_size.
- `shard_estimators.py`
- The shard estimators now use the `intra_group_size` instead of `local_world_size`, this allows RW/TW/CW to properly account for larger NVlink that comes with the pods.
Reviewed By: isururanawaka
Differential Revision: D91617887
Ali-Tehrani
added a commit
to Ali-Tehrani/torchrec
that referenced
this pull request
Jan 30, 2026
Summary:
Context
---------
Every GB200 node has 2 B200 GPU attached to it, however allows up to 72 B200 connected via NVlink. The planner needs to know how big the intra topology group size is going to be.
This causes the `local_world_size` to be different from the `intra_group_size`.
Implementation
------------------
- Topology class:
- Adds `pod_size`, and uses that to calculate the `intra_group_size` (maximum number of processes linked with high intra bandwidth) to Topology class. If isn't given, then it defaults to local_world_size.
- `shard_estimators.py`
- The shard estimators now use the `intra_group_size` instead of `local_world_size`, this allows RW/TW/CW to properly account for larger NVlink that comes with the pods.
Reviewed By: isururanawaka
Differential Revision: D91617887
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
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.
Summary:
Context
Every GB200 node has 2 B200 GPU attached to it, however allows up to 72 B200 connected via NVlink. The planner needs to know how big the intra topology group size is going to be.
This causes the
local_world_sizeto be different from theintra_group_size.Implementation
pod_size, and uses that to calculate theintra_group_size(maximum number of processes linked with high intra bandwidth) to Topology class. If isn't given, then it defaults to local_world_size.shard_estimators.pyintra_group_sizeinstead oflocal_world_size, this allows RW/TW/CW to properly account for larger NVlink that comes with the pods.Differential Revision: D91617887