diff --git a/miles/backends/experimental/fsdp_utils/update_weight_utils.py b/miles/backends/experimental/fsdp_utils/update_weight_utils.py index 98000e7d34..14f9ab85cd 100644 --- a/miles/backends/experimental/fsdp_utils/update_weight_utils.py +++ b/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, diff --git a/miles/backends/megatron_utils/actor.py b/miles/backends/megatron_utils/actor.py index b922417598..ab56a7d00a 100644 --- a/miles/backends/megatron_utils/actor.py +++ b/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, diff --git a/miles/backends/megatron_utils/update_weight/update_weight_from_distributed/broadcast.py b/miles/backends/megatron_utils/update_weight/update_weight_from_distributed/broadcast.py index 651c99e487..4faecfe6a3 100644 --- a/miles/backends/megatron_utils/update_weight/update_weight_from_distributed/broadcast.py +++ b/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, diff --git a/tests/fast/backends/megatron_utils/test_ipv6_init_method.py b/tests/fast/backends/megatron_utils/test_ipv6_init_method.py new file mode 100644 index 0000000000..6bfbf65461 --- /dev/null +++ b/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)