fix: bracket IPv6 hosts in distributed init methods

miles/backends/experimental/fsdp_utils/update_weight_utils.py

@@ -18,7 +18,7 @@
 from sglang.srt.utils import MultiprocessingSerializer
 
 from miles.utils.distributed_utils import get_gloo_group, init_process_group
-
+from miles.utils.http_utils import _wrap_ipv6
 
 try:
     from sglang.srt.weight_sync.tensor_bucket import FlattenedTensorBucket  # type: ignore[import]
@@ -240,7 +240,7 @@ def connect_rollout_engines(
             ]
             self._model_update_groups = init_process_group(
                 backend="nccl",
-                init_method=f"tcp://{master_address}:{master_port}",
+                init_method=f"tcp://{_wrap_ipv6(master_address)}:{master_port}",
                 world_size=world_size,
                 rank=0,
                 group_name=self._group_name,

miles/backends/megatron_utils/actor.py

@@ -20,6 +20,7 @@
 from miles.utils.distributed_utils import get_gloo_group, init_process_group
 from miles.utils.ft_utils.indep_dp import IndepDPInfo
 from miles.utils.hf_config import load_hf_config
+from miles.utils.http_utils import _wrap_ipv6
 from miles.utils.memory_utils import clear_memory, print_memory
 from miles.utils.processing_utils import load_tokenizer
 from miles.utils.ray_utils import Box
@@ -679,7 +680,7 @@ def connect_actor_critic(
         world_size = 2
         self._actor_critic_groups = init_process_group(
             backend="nccl",
-            init_method=f"tcp://{master_address}:{master_port}",
+            init_method=f"tcp://{_wrap_ipv6(master_address)}:{master_port}",
             world_size=world_size,
             rank=0 if self.role == "actor" else 1,
             group_name=group_name,

miles/backends/megatron_utils/update_weight/update_weight_from_distributed/broadcast.py

@@ -12,6 +12,7 @@
 
 from miles.backends.training_utils.parallel import get_parallel_state
 from miles.utils.distributed_utils import init_process_group
+from miles.utils.http_utils import _wrap_ipv6
 
 from miles.utils.lora import LORA_ADAPTER_NAME
 from ..common import _check_weight_sync_results
@@ -192,7 +193,7 @@ def connect_rollout_engines_from_distributed(
         rank_cursor += engine_gpu_counts[i]
     model_update_groups = init_process_group(
         backend="nccl",
-        init_method=f"tcp://{master_address}:{master_port}",
+        init_method=f"tcp://{_wrap_ipv6(master_address)}:{master_port}",
         world_size=world_size,
         rank=0,
         group_name=group_name,

tests/fast/backends/megatron_utils/test_ipv6_init_method.py

@@ -0,0 +1,29 @@
+from argparse import Namespace
+from unittest.mock import MagicMock, patch
+
+import pytest
+
+from miles.backends.megatron_utils.update_weight.update_weight_from_distributed import broadcast
+
+
+@pytest.mark.parametrize(
+    "ip,expected_prefix",
+    [
+        ("::1", "tcp://[::1]:"),
+        ("127.0.0.1", "tcp://127.0.0.1:"),
+    ],
+)
+def test_connect_rollout_engines_builds_valid_init_method(ip, expected_prefix):
+    engine = MagicMock()
+    with (
+        patch.object(broadcast, "ray") as mock_ray,
+        patch.object(broadcast, "init_process_group") as mock_init_pg,
+    ):
+        mock_ray._private.services.get_node_ip_address.return_value = ip
+        broadcast.connect_rollout_engines_from_distributed(
+            Namespace(rollout_num_gpus_per_engine=1),
+            group_name="test",
+            rollout_engines=[engine],
+        )
+
+    assert mock_init_pg.call_args.kwargs["init_method"].startswith(expected_prefix)