From 5028bbebb8a476d73fb14dab2d1ee432a51b2566 Mon Sep 17 00:00:00 2001 From: kaiyuan xie Date: Tue, 19 May 2026 18:42:41 +0800 Subject: [PATCH 1/2] MegatronvLLM native sync: in-process NCCL weight transfer (remove NcclBridge) --- .../update_weight_from_distributed.py | 225 ++-------------- tests/test_update_weight_from_distributed.py | 248 +++++++++++++++--- 2 files changed, 235 insertions(+), 238 deletions(-) diff --git a/slime/backends/megatron_utils/update_weight/update_weight_from_distributed.py b/slime/backends/megatron_utils/update_weight/update_weight_from_distributed.py index 56f53b6e5..8e9152653 100644 --- a/slime/backends/megatron_utils/update_weight/update_weight_from_distributed.py +++ b/slime/backends/megatron_utils/update_weight/update_weight_from_distributed.py @@ -4,7 +4,6 @@ import os import socket import time -import traceback from argparse import Namespace from collections.abc import Callable, Mapping, Sequence from typing import Any @@ -12,7 +11,6 @@ import ray import torch import torch.distributed as dist -import torch.multiprocessing as mp from megatron.core import mpu from ray import ObjectRef from ray.actor import ActorHandle @@ -26,181 +24,6 @@ logger = logging.getLogger(__name__) -# --------------------------------------------------------------------------- -# NcclBridge: isolate vLLM's PyNcclCommunicator in a subprocess so that it -# never coexists with torch.distributed NCCL groups in the Megatron trainer. -# -# vLLM's weight transfer uses raw NCCL (PyNcclCommunicator) which conflicts -# with torch.distributed's NCCL backend when both exist in the same process -# (see https://github.com/vllm-project/vllm/issues/5477). SGLang avoids -# this because it uses torch.distributed process groups for weight sync. -# --------------------------------------------------------------------------- - - -def _nccl_bridge_worker( - conn, - master_address: str, - master_port: int, - world_size: int, - device: int, - cvd: str, - env_snapshot: dict[str, str], -) -> None: - """Subprocess entry-point: creates PyNcclCommunicator and serves requests. - - GPU tensors are shared from the parent via CUDA IPC (torch.multiprocessing - handles this transparently). No GPU→CPU→GPU copies are needed. - - Protocol over *conn* (multiprocessing.Connection): - parent → child: - {"op": "broadcast", "tensors": [gpu_tensor, ...]} - {"op": "send_packed", "named_tensors": [(name, gpu_tensor), ...]} - None → shutdown - child → parent: - "ready" (after init) - "ok" (after each op) - "error: ..." - """ - try: - os.environ.update(env_snapshot) - if cvd: - os.environ["CUDA_VISIBLE_DEVICES"] = cvd - - import torch as _torch # noqa: PLC0415 — subprocess needs fresh import - import torch.multiprocessing # noqa: F401, PLC0415 — register CUDA IPC reducers - - _torch.cuda.set_device(device) - - from vllm.distributed.device_communicators.pynccl import PyNcclCommunicator # noqa: PLC0415 - from vllm.distributed.utils import StatelessProcessGroup # noqa: PLC0415 - - pg = StatelessProcessGroup.create( - host=master_address, - port=master_port, - rank=0, - world_size=world_size, - ) - comm = PyNcclCommunicator(pg, device=device) - - conn.send("ready") - - while True: - cmd = conn.recv() - if cmd is None: - break - - op = cmd["op"] - if op == "broadcast": - for t in cmd["tensors"]: - comm.broadcast(t, src=0, stream=_torch.cuda.current_stream()) - _torch.cuda.synchronize() - conn.send("ok") - - elif op == "send_packed": - # Prefer NCCLWeightTransferEngine.trainer_send_weights (newer vLLM). Some pip builds omit - # NCCLTrainerSendWeightsArgs but still ship packed_broadcast_producer. - try: - from vllm.distributed.weight_transfer.nccl_engine import ( # noqa: PLC0415 - NCCLTrainerSendWeightsArgs, - NCCLWeightTransferEngine, - ) - - trainer_args = NCCLTrainerSendWeightsArgs( - group=comm, - packed=True, - ) - NCCLWeightTransferEngine.trainer_send_weights( - iterator=iter(cmd["named_tensors"]), - trainer_args=trainer_args, - ) - except ImportError: - from vllm.distributed.weight_transfer.packed_tensor import ( # noqa: PLC0415 - DEFAULT_PACKED_BUFFER_SIZE_BYTES, - DEFAULT_PACKED_NUM_BUFFERS, - packed_broadcast_producer, - ) - - packed_broadcast_producer( - iterator=iter(cmd["named_tensors"]), - group=comm, - src=0, - post_iter_func=lambda x: x[1], - buffer_size_bytes=DEFAULT_PACKED_BUFFER_SIZE_BYTES, - num_buffers=DEFAULT_PACKED_NUM_BUFFERS, - ) - _torch.cuda.synchronize() - conn.send("ok") - - except Exception as e: - try: - conn.send(f"error: {e}") - except Exception: - pass - traceback.print_exc() - - -class _NcclBridge: - """Runs vLLM's PyNcclCommunicator in a separate subprocess. - - This prevents NCCL communicator conflicts with torch.distributed groups - that already exist in the Megatron trainer process. GPU tensors are shared - with the subprocess via CUDA IPC (handled transparently by - torch.multiprocessing), avoiding any GPU→CPU→GPU copies. - """ - - def __init__(self, master_address: str, master_port: int, world_size: int, device: int): - ctx = mp.get_context("spawn") - self._parent_conn, child_conn = ctx.Pipe() - - env_snapshot = dict(os.environ) - cvd = os.environ.get("CUDA_VISIBLE_DEVICES", "") - - self._process = ctx.Process( - target=_nccl_bridge_worker, - args=(child_conn, master_address, master_port, world_size, device, cvd, env_snapshot), - daemon=True, - ) - self._process.start() - - msg = self._parent_conn.recv() - if isinstance(msg, str) and msg.startswith("error:"): - raise RuntimeError(f"NcclBridge init failed: {msg}") - if msg != "ready": - raise RuntimeError(f"NcclBridge init unexpected response: {msg}") - logger.info("NcclBridge ready (pid=%d, device=%d)", self._process.pid, device) - - def broadcast_tensors(self, tensors: list[torch.Tensor]) -> None: - """Broadcast a list of tensors (one-by-one) via the bridge subprocess.""" - gpu_tensors = [t.contiguous() for t in tensors] - self._parent_conn.send({"op": "broadcast", "tensors": gpu_tensors}) - self._wait_ok("broadcast_tensors") - - def send_weights_packed(self, named_tensors: list[tuple[str, torch.Tensor]]) -> None: - """Send weights using vLLM's packed broadcast protocol.""" - gpu_pairs = [] - for name, t in named_tensors: - data = t.data if hasattr(t, "data") else t - gpu_pairs.append((name, data.contiguous())) - self._parent_conn.send({"op": "send_packed", "named_tensors": gpu_pairs}) - self._wait_ok("send_weights_packed") - - def _wait_ok(self, label: str, timeout: float = 600.0) -> None: - if not self._parent_conn.poll(timeout): - raise TimeoutError(f"NcclBridge {label} timed out after {timeout}s") - msg = self._parent_conn.recv() - if msg != "ok": - raise RuntimeError(f"NcclBridge {label} failed: {msg}") - - def shutdown(self) -> None: - try: - self._parent_conn.send(None) - self._process.join(timeout=30) - except Exception: - pass - if self._process.is_alive(): - self._process.terminate() - - class UpdateWeightFromDistributed: """ Update distributed engines via NCCL. Each PP rank: group "slime-pp_{pp_rank}", @@ -523,10 +346,8 @@ def connect_rollout_engines_from_distributed( have heterogeneous TP sizes (e.g. prefill TP=2, decode TP=4), each engine occupies a different number of ranks in the NCCL group. - For vLLM backend, the trainer-side NCCL communicator is created inside a - separate subprocess (_NcclBridge) to avoid conflicts between vLLM's raw - NCCL (PyNcclCommunicator) and the torch.distributed NCCL groups that - Megatron already holds in this process. + Trainer rank 0 uses ``NCCLWeightTransferEngine.trainer_init`` + in-process (StatelessProcessGroup + PyNcclCommunicator). """ if engine_gpu_counts is None: engine_gpu_counts = [args.rollout_num_gpus_per_engine] * len(rollout_engines) @@ -556,20 +377,23 @@ def connect_rollout_engines_from_distributed( torch.cuda.synchronize() torch.cuda.empty_cache() + from vllm.distributed.weight_transfer.nccl_engine import NCCLWeightTransferEngine + device = torch.cuda.current_device() logger.info( - "vLLM weight transfer via NcclBridge: addr=%s port=%d world_size=%d device=%d CVD=%s", + "vLLM in-process weight transfer: addr=%s port=%d world_size=%d device=%d CVD=%s", master_address, master_port, world_size, device, os.environ.get("CUDA_VISIBLE_DEVICES", ""), ) - model_update_groups = _NcclBridge( - master_address=master_address, - master_port=master_port, - world_size=world_size, - device=device, + model_update_groups = NCCLWeightTransferEngine.trainer_init( + { + "master_address": master_address, + "master_port": master_port, + "world_size": world_size, + } ) ray.get(refs) @@ -586,14 +410,12 @@ def disconnect_rollout_engines_from_distributed( Destroy NCCL on training and engines. """ refs = [engine.destroy_weights_update_group.remote(group_name) for engine in rollout_engines] - if isinstance(model_update_groups, _NcclBridge): - model_update_groups.shutdown() ray.get(refs) def update_weights_from_distributed( group_name: str, - group: _NcclBridge, + group: Any, weight_version: int, rollout_engines: Sequence[ActorHandle], converted_named_tensors: Sequence[tuple[str, torch.Tensor]], @@ -601,9 +423,10 @@ def update_weights_from_distributed( packed: bool = False, ) -> list[ObjectRef]: """ - Send metadata (Ray), broadcast tensors (NCCL rank 0 → vLLM engines via - the ``_NcclBridge`` subprocess so that raw NCCL never runs inside the - Megatron trainer process). + Send metadata (Ray), broadcast tensors (NCCL rank 0 → engines). + + The *group* is a vLLM ``PyNcclCommunicator`` from ``trainer_init`` + in the Megatron trainer process. """ kwargs: dict[str, Any] = { "names": [name for name, _ in converted_named_tensors], @@ -616,10 +439,18 @@ def update_weights_from_distributed( refs = [engine.update_weights_from_distributed.remote(**kwargs) for engine in rollout_engines] - if packed: - group.send_weights_packed(list(converted_named_tensors)) - else: - group.broadcast_tensors([param.data for _, param in converted_named_tensors]) + from vllm.distributed.weight_transfer.nccl_engine import NCCLWeightTransferEngine + + named_gpu_iter = ( + (name, (param.data if hasattr(param, "data") else param).contiguous()) + for name, param in converted_named_tensors + ) + NCCLWeightTransferEngine.trainer_send_weights( + iterator=named_gpu_iter, + group=group, + packed=packed, + ) + torch.cuda.synchronize() return refs diff --git a/tests/test_update_weight_from_distributed.py b/tests/test_update_weight_from_distributed.py index 763992eef..323481f56 100644 --- a/tests/test_update_weight_from_distributed.py +++ b/tests/test_update_weight_from_distributed.py @@ -4,7 +4,6 @@ import importlib import inspect -from collections.abc import Iterable from dataclasses import dataclass, field import pytest @@ -37,18 +36,15 @@ class RecordingEngine: update_weights_from_distributed: RecordingRemoteMethod = field( default_factory=lambda: RecordingRemoteMethod("ref") ) + init_weights_update_group: RecordingRemoteMethod = field(default_factory=lambda: RecordingRemoteMethod("init_ref")) + destroy_weights_update_group: RecordingRemoteMethod = field( + default_factory=lambda: RecordingRemoteMethod("destroy_ref") + ) @dataclass -class RecordingNcclBridge: - broadcast_calls: list[list[torch.Tensor]] = field(default_factory=list) - packed_calls: list[list[tuple[str, torch.Tensor]]] = field(default_factory=list) - - def broadcast_tensors(self, tensors: Iterable[torch.Tensor]) -> None: - self.broadcast_calls.append(list(tensors)) - - def send_weights_packed(self, named_tensors: Iterable[tuple[str, torch.Tensor]]) -> None: - self.packed_calls.append(list(named_tensors)) +class DummyGroup: + token: str = "dummy" def _real_tensors(n: int = 2): @@ -69,7 +65,7 @@ def test_signature_rejects_legacy_use_vllm_call(upw): with pytest.raises(TypeError, match="use_vllm"): upw.update_weights_from_distributed( "g", - RecordingNcclBridge(), + DummyGroup(), 1, [RecordingEngine()], _real_tensors(), @@ -79,43 +75,89 @@ def test_signature_rejects_legacy_use_vllm_call(upw): @pytest.mark.unit -def test_packed_true_dispatches_send_weights_packed(upw): - group = RecordingNcclBridge() +def test_packed_true_uses_vllm_trainer_send_weights(upw, monkeypatch): + group = DummyGroup() engine = RecordingEngine() tensors = _real_tensors() + seen = [] + + class DummyNCCLWeightTransferEngine: + @staticmethod + def trainer_send_weights(*, iterator, group, packed): + seen.append({"items": list(iterator), "group": group, "packed": packed}) + + monkeypatch.setattr( + upw, + "NCCLWeightTransferEngine", + DummyNCCLWeightTransferEngine, + raising=False, + ) + monkeypatch.setitem( + __import__("sys").modules, + "vllm.distributed.weight_transfer.nccl_engine", + type("M", (), {"NCCLWeightTransferEngine": DummyNCCLWeightTransferEngine}), + ) + monkeypatch.setattr(upw.torch.cuda, "synchronize", lambda: None) refs = upw.update_weights_from_distributed("groupA", group, 7, [engine], tensors, packed=True) - assert len(group.packed_calls) == 1 - assert len(group.broadcast_calls) == 0 - sent = group.packed_calls[0] + assert len(seen) == 1 + sent = seen[0]["items"] assert [n for n, _ in sent] == [n for n, _ in tensors] + assert seen[0]["group"] is group + assert seen[0]["packed"] is True assert refs == ["ref"] @pytest.mark.unit -def test_packed_false_dispatches_broadcast_tensors(upw): - group = RecordingNcclBridge() +def test_packed_false_still_uses_vllm_trainer_send_weights(upw, monkeypatch): + group = DummyGroup() engine = RecordingEngine() tensors = _real_tensors() + seen = [] + + class DummyNCCLWeightTransferEngine: + @staticmethod + def trainer_send_weights(*, iterator, group, packed): + seen.append({"items": list(iterator), "group": group, "packed": packed}) + + monkeypatch.setitem( + __import__("sys").modules, + "vllm.distributed.weight_transfer.nccl_engine", + type("M", (), {"NCCLWeightTransferEngine": DummyNCCLWeightTransferEngine}), + ) + monkeypatch.setattr(upw.torch.cuda, "synchronize", lambda: None) refs = upw.update_weights_from_distributed("groupB", group, 7, [engine], tensors, packed=False) - assert len(group.broadcast_calls) == 1 - assert len(group.packed_calls) == 0 - assert len(group.broadcast_calls[0]) == len(tensors) + assert len(seen) == 1 + assert len(seen[0]["items"]) == len(tensors) + assert seen[0]["packed"] is False assert refs == ["ref"] @pytest.mark.unit -def test_default_packed_is_false(upw): - group = RecordingNcclBridge() +def test_default_packed_is_false(upw, monkeypatch): + group = DummyGroup() engine = RecordingEngine() + seen = [] + + class DummyNCCLWeightTransferEngine: + @staticmethod + def trainer_send_weights(*, iterator, group, packed): + seen.append({"items": list(iterator), "group": group, "packed": packed}) + + monkeypatch.setitem( + __import__("sys").modules, + "vllm.distributed.weight_transfer.nccl_engine", + type("M", (), {"NCCLWeightTransferEngine": DummyNCCLWeightTransferEngine}), + ) + monkeypatch.setattr(upw.torch.cuda, "synchronize", lambda: None) upw.update_weights_from_distributed("g", group, 1, [engine], _real_tensors()) - assert len(group.broadcast_calls) == 1 - assert len(group.packed_calls) == 0 + assert len(seen) == 1 + assert seen[0]["packed"] is False @pytest.mark.unit @@ -123,28 +165,55 @@ def test_no_dist_broadcast_fallback(upw, monkeypatch): import torch.distributed as dist seen_broadcast = [] + seen_send = [] def fake_broadcast(*a, **k): seen_broadcast.append((a, k)) + class DummyNCCLWeightTransferEngine: + @staticmethod + def trainer_send_weights(*, iterator, group, packed): + seen_send.append({"items": list(iterator), "group": group, "packed": packed}) + monkeypatch.setattr(dist, "broadcast", fake_broadcast) + monkeypatch.setitem( + __import__("sys").modules, + "vllm.distributed.weight_transfer.nccl_engine", + type("M", (), {"NCCLWeightTransferEngine": DummyNCCLWeightTransferEngine}), + ) + monkeypatch.setattr(upw.torch.cuda, "synchronize", lambda: None) - group = RecordingNcclBridge() + group = DummyGroup() engine = RecordingEngine() upw.update_weights_from_distributed("g", group, 1, [engine], _real_tensors(), packed=False) assert seen_broadcast == [] - assert group.broadcast_calls + assert len(seen_send) == 1 @pytest.mark.unit -def test_remote_kwargs_include_packed_true(upw): - group = RecordingNcclBridge() +def test_remote_kwargs_include_packed_true(upw, monkeypatch): + group = DummyGroup() engine = RecordingEngine() tensors = _real_tensors(n=1) + seen_send = [] + + class DummyNCCLWeightTransferEngine: + @staticmethod + def trainer_send_weights(*, iterator, group, packed): + seen_send.append({"items": list(iterator), "group": group, "packed": packed}) + + monkeypatch.setitem( + __import__("sys").modules, + "vllm.distributed.weight_transfer.nccl_engine", + type("M", (), {"NCCLWeightTransferEngine": DummyNCCLWeightTransferEngine}), + ) + monkeypatch.setattr(upw.torch.cuda, "synchronize", lambda: None) upw.update_weights_from_distributed("myg", group, 42, [engine], tensors, packed=True) + assert len(seen_send) == 1 + assert seen_send[0]["packed"] is True assert len(engine.update_weights_from_distributed.calls) == 1 kw = engine.update_weights_from_distributed.calls[0].kwargs assert kw["packed"] is True @@ -156,13 +225,28 @@ def test_remote_kwargs_include_packed_true(upw): @pytest.mark.unit -def test_remote_kwargs_include_packed_false(upw): - group = RecordingNcclBridge() +def test_remote_kwargs_include_packed_false(upw, monkeypatch): + group = DummyGroup() engine = RecordingEngine() tensors = _real_tensors(n=2) + seen_send = [] + + class DummyNCCLWeightTransferEngine: + @staticmethod + def trainer_send_weights(*, iterator, group, packed): + seen_send.append({"items": list(iterator), "group": group, "packed": packed}) + + monkeypatch.setitem( + __import__("sys").modules, + "vllm.distributed.weight_transfer.nccl_engine", + type("M", (), {"NCCLWeightTransferEngine": DummyNCCLWeightTransferEngine}), + ) + monkeypatch.setattr(upw.torch.cuda, "synchronize", lambda: None) upw.update_weights_from_distributed("g", group, 99, [engine], tensors, packed=False) + assert len(seen_send) == 1 + assert seen_send[0]["packed"] is False kw = engine.update_weights_from_distributed.calls[0].kwargs assert kw["packed"] is False assert kw["weight_version"] == "99" @@ -170,29 +254,72 @@ def test_remote_kwargs_include_packed_false(upw): @pytest.mark.unit -def test_remote_kwargs_no_use_vllm(upw): - group = RecordingNcclBridge() +def test_remote_kwargs_no_use_vllm(upw, monkeypatch): + group = DummyGroup() engine = RecordingEngine() + seen_send = [] + + class DummyNCCLWeightTransferEngine: + @staticmethod + def trainer_send_weights(*, iterator, group, packed): + seen_send.append({"items": list(iterator), "group": group, "packed": packed}) + + monkeypatch.setitem( + __import__("sys").modules, + "vllm.distributed.weight_transfer.nccl_engine", + type("M", (), {"NCCLWeightTransferEngine": DummyNCCLWeightTransferEngine}), + ) + monkeypatch.setattr(upw.torch.cuda, "synchronize", lambda: None) upw.update_weights_from_distributed("g", group, 1, [engine], _real_tensors(), packed=False) + assert len(seen_send) == 1 kw = engine.update_weights_from_distributed.calls[0].kwargs assert "use_vllm" not in kw @pytest.mark.unit -def test_multiple_engines_each_get_call(upw): - group = RecordingNcclBridge() +def test_multiple_engines_each_get_call(upw, monkeypatch): + group = DummyGroup() engines = [RecordingEngine() for _ in range(3)] + seen_send = [] + + class DummyNCCLWeightTransferEngine: + @staticmethod + def trainer_send_weights(*, iterator, group, packed): + seen_send.append({"items": list(iterator), "group": group, "packed": packed}) + + monkeypatch.setitem( + __import__("sys").modules, + "vllm.distributed.weight_transfer.nccl_engine", + type("M", (), {"NCCLWeightTransferEngine": DummyNCCLWeightTransferEngine}), + ) + monkeypatch.setattr(upw.torch.cuda, "synchronize", lambda: None) + upw.update_weights_from_distributed("g", group, 1, engines, _real_tensors(), packed=True) + assert len(seen_send) == 1 + assert seen_send[0]["packed"] is True for e in engines: assert len(e.update_weights_from_distributed.calls) == 1 @pytest.mark.unit -def test_empty_tensor_list_still_dispatches(upw): - group = RecordingNcclBridge() +def test_empty_tensor_list_still_dispatches(upw, monkeypatch): + group = DummyGroup() engine = RecordingEngine() + seen_send = [] + + class DummyNCCLWeightTransferEngine: + @staticmethod + def trainer_send_weights(*, iterator, group, packed): + seen_send.append({"items": list(iterator), "group": group, "packed": packed}) + + monkeypatch.setitem( + __import__("sys").modules, + "vllm.distributed.weight_transfer.nccl_engine", + type("M", (), {"NCCLWeightTransferEngine": DummyNCCLWeightTransferEngine}), + ) + monkeypatch.setattr(upw.torch.cuda, "synchronize", lambda: None) refs = upw.update_weights_from_distributed("g", group, 1, [engine], [], packed=False) @@ -200,8 +327,8 @@ def test_empty_tensor_list_still_dispatches(upw): kw = engine.update_weights_from_distributed.calls[0].kwargs assert kw["names"] == [] assert kw["shapes"] == [] - assert len(group.broadcast_calls) == 1 - assert group.broadcast_calls[0] == [] + assert len(seen_send) == 1 + assert seen_send[0]["items"] == [] @pytest.mark.unit @@ -213,5 +340,44 @@ def test_source_no_standalone_use_vllm_param(upw): @pytest.mark.unit def test_source_no_sglang_dist_broadcast_fallback(upw): - fn_src = inspect.getsource(upw.update_weights_from_distributed) - assert "dist.broadcast(" not in fn_src + src = inspect.getsource(upw) + assert "dist.broadcast(" not in src + + +@pytest.mark.unit +def test_source_no_materialized_named_gpu_list(upw): + src = inspect.getsource(upw.update_weights_from_distributed) + assert "named_gpu = []" not in src + assert "named_gpu_iter =" in src + + +@pytest.mark.unit +def test_connect_rollout_engines_always_uses_vllm_trainer_init(upw, monkeypatch): + args = type("Args", (), {"rollout_num_gpus_per_engine": 1})() + engines = [RecordingEngine(), RecordingEngine()] + seen = [] + + class DummyNCCLWeightTransferEngine: + @staticmethod + def trainer_init(cfg): + seen.append(cfg) + return DummyGroup("group-from-trainer-init") + + monkeypatch.setitem( + __import__("sys").modules, + "vllm.distributed.weight_transfer.nccl_engine", + type("M", (), {"NCCLWeightTransferEngine": DummyNCCLWeightTransferEngine}), + ) + monkeypatch.setattr(upw.torch.cuda, "synchronize", lambda: None) + monkeypatch.setattr(upw.torch.cuda, "empty_cache", lambda: None) + monkeypatch.setattr(upw.ray, "get", lambda refs: refs) + monkeypatch.setattr(upw.ray._private.services, "get_node_ip_address", lambda: "127.0.0.1") + + group = upw.connect_rollout_engines_from_distributed(args, "g", engines, engine_gpu_counts=[1, 2]) + + assert isinstance(group, DummyGroup) + assert len(seen) == 1 + assert seen[0]["master_address"] == "127.0.0.1" + assert seen[0]["world_size"] == 4 # 1 + (1 + 2) + assert len(engines[0].init_weights_update_group.calls) == 1 + assert len(engines[1].init_weights_update_group.calls) == 1 From b7a767314152f6266a01ce19fe46336b6fca8d37 Mon Sep 17 00:00:00 2001 From: kaiyuan xie Date: Thu, 21 May 2026 10:01:23 +0800 Subject: [PATCH 2/2] update requirements.txt --- requirements.txt | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/requirements.txt b/requirements.txt index 8b016fb03..7e400e7ee 100644 --- a/requirements.txt +++ b/requirements.txt @@ -13,7 +13,7 @@ qwen_vl_utils # for VLM ray[default] ring_flash_attn sglang-router>=0.2.3 -vllm-router>=0.1.14 tensorboard transformers +vllm-router>=0.1.14 wandb