Merge branch 'no_dist_device_id' into 'main'

Don't pass device_id to torch.distributed.init_process_group, it causes hangs See merge request ADLR/megatron-lm!2091
NVIDIA · Sep 12, 2024 · 76f9f48 · 76f9f48
2 parents 028b777 + dcc6634
commit 76f9f48
Showing 1 changed file with 0 additions and 2 deletions.
diff --git a/megatron/training/initialize.py b/megatron/training/initialize.py
@@ -254,8 +254,6 @@ def _initialize_distributed(get_embedding_ranks, get_position_embedding_ranks):
             'rank': args.rank,
             'timeout': timedelta(minutes=args.distributed_timeout_minutes),
         }
-        if packaging.version.Version(torch.__version__) >= packaging.version.Version("2.3.0"):
-            init_process_group_kwargs['device_id'] = device_id
 
         torch.distributed.init_process_group(**init_process_group_kwargs)