github/ColossalAI
1
0
Fork 0
mirror of https://github.com/hpcaitech/ColossalAI.git synced 2025-09-15 06:00:07 +00:00
Code Issues Packages Projects Releases Wiki Activity

[launch] Support IPv4 host initialization in launch (#5822)

Browse Source
This commit is contained in:
Kai Lv
2024-06-18 19:18:29 +08:00
committed by GitHub
parent 639394b0d4
commit 0adca5b688
1 changed files with 4 additions and 1 deletions
Download Patch File Download Diff File Expand all files Collapse all files

5
colossalai/initialize.py
View File

@@ -45,7 +45,10 @@ def launch(
backend = cur_accelerator.communication_backend
# init default process group
init_method = f"tcp://[{host}]:{port}"
if ":" in host: # IPv6
init_method = f"tcp://[{host}]:{port}"
else: # IPv4
init_method = f"tcp://{host}:{port}"
dist.init_process_group(rank=rank, world_size=world_size, backend=backend, init_method=init_method)
# set cuda device