[PyTorch] Missing intra-domain ranks list when initializing Userbuffe…

…rs with data parallelism (#1305) added missing list of intra-domain ranks when num_domains > 1 in initialize_ub Signed-off-by: Alp Dener <[email protected]>
NVIDIA · Nov 2, 2024 · a6a9141 · a6a9141
1 parent 4b8ffef
commit a6a9141
Showing 1 changed file with 1 addition and 0 deletions.
diff --git a/transformer_engine/pytorch/module/base.py b/transformer_engine/pytorch/module/base.py
@@ -234,6 +234,7 @@ def initialize_ub(
                 ranks_per_domain_list, backend=bootstrap_backend
             )
             local_rank = torch.distributed.get_rank(intra_domain_group)
+            intra_domain_ranks = torch.distributed.get_process_group_ranks(intra_domain_group)
 
             inter_domain_group, _ = torch.distributed.new_subgroups_by_enumeration(
                 [list(ranks) for ranks in zip(*ranks_per_domain_list)],