github
/
ColossalAI
mirror of https://github.com/hpcaitech/ColossalAI
1
0
Fork 0
Code Issues Projects Releases Wiki Activity

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

pull/5835/head
Kai Lv 2024-06-18 19:18:29 +08:00 committed by GitHub
parent 639394b0d4
commit 0adca5b688
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
1 changed files with 4 additions and 1 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

3
colossalai/initialize.py
View File

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