Set intra_group_size from env var inside comm.py. #3697
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 D91617889. |
Ali-Tehrani
force-pushed
the
export-D91617889
branch
from
86d7b53 to
5cea44a
Compare
Ali-Tehrani
added a commit
to Ali-Tehrani/torchrec
that referenced
this pull request
Jan 30, 2026
Summary: Context --------- TorchRec comms needs a way to obtain the pod-size (topology-domain-multiple) and the total amount of process groups of the topology group for TWRW/Grid-sharding. We obtain the number of intra-nodes within a pod by obtaining the `TOPOLOGY_DOMAIN_MULTIPLE` from the environment variables (see diff stack). Implementation ------------------ - created `get_intra_group_size` function that obtains the number of intra-node-size from envrionemnt variable, and if not it defaults to usual `get_local_size`. - updated `intra_and_cross_node_pg` to utilize `get_intra_node_size` instead. Differential Revision: D91617889
Ali-Tehrani
added a commit
to Ali-Tehrani/torchrec
that referenced
this pull request
Jan 30, 2026
Summary: Context --------- TorchRec comms needs a way to obtain the pod-size (topology-domain-multiple) and the total amount of process groups of the topology group for TWRW/Grid-sharding. We obtain the number of intra-nodes within a pod by obtaining the `TOPOLOGY_DOMAIN_MULTIPLE` from the environment variables (see diff stack). Implementation ------------------ - created `get_intra_group_size` function that obtains the number of intra-node-size from envrionemnt variable, and if not it defaults to usual `get_local_size`. - updated `intra_and_cross_node_pg` to utilize `get_intra_node_size` instead. Differential Revision: D91617889
Ali-Tehrani
force-pushed
the
export-D91617889
branch
from
5cea44a to
ee29fbd
Compare
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
Summary: Context --------- TorchRec comms needs a way to obtain the pod-size (topology-domain-multiple) and the total amount of process groups of the topology group for TWRW/Grid-sharding. We obtain the number of intra-nodes within a pod by obtaining the `TOPOLOGY_DOMAIN_MULTIPLE` from the environment variables (see diff stack). Implementation ------------------ - created `get_intra_group_size` function that obtains the number of intra-node-size from envrionemnt variable, and if not it defaults to usual `get_local_size`. - updated `intra_and_cross_node_pg` to utilize `get_intra_node_size` instead. Differential Revision: D91617889
Ali-Tehrani
force-pushed
the
export-D91617889
branch
from
ee29fbd to
c2ef70f
Compare
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
TorchRec comms needs a way to obtain the pod-size (topology-domain-multiple) and the total amount of process groups of the topology group for TWRW/Grid-sharding.
We obtain the number of intra-nodes within a pod by obtaining the
TOPOLOGY_DOMAIN_MULTIPLEfrom the environment variables (see diff stack).Implementation
get_intra_group_sizefunction that obtains the number of intra-node-size from envrionemnt variable, and if not it defaults to usualget_local_size.intra_and_cross_node_pgto utilizeget_intra_node_sizeinstead.Differential Revision: D91617889