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