From 0adca5b688367a7b6c01cb43ef2507ce29d2150a Mon Sep 17 00:00:00 2001
From: Kai Lv <39761308+KaiLv69@users.noreply.github.com>
Date: Tue, 18 Jun 2024 19:18:29 +0800
Subject: [PATCH] [launch] Support IPv4 host initialization in launch (#5822)

---
 colossalai/initialize.py | 5 ++++-
 1 file changed, 4 insertions(+), 1 deletion(-)

diff --git a/colossalai/initialize.py b/colossalai/initialize.py
index 934555e19..71d42312e 100644
--- 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