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

5 months ago · 0adca5b688
parent 639394b0d4
commit 0adca5b688
1 changed files with 4 additions and 1 deletions
--- a/colossalai/initialize.py
+++ b/colossalai/initialize.py
@ -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