diff --git a/scripts/p2p_control_smoke.py b/scripts/p2p_control_smoke.py new file mode 100644 index 00000000000..52a7660d5af --- /dev/null +++ b/scripts/p2p_control_smoke.py @@ -0,0 +1,97 @@ +# Copyright 2025 Bytedance Ltd. and/or its affiliates +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from __future__ import annotations + +import logging +from typing import Any, Protocol + +logger = logging.getLogger(__name__) + + +class P2PControlPlaneReplica(Protocol): + async def abort_all_requests(self) -> None: ... + + async def get_remote_instance_transfer_engine_info(self, rank: int) -> Any: ... + + async def get_parallelism_info(self, rank: int) -> Any: ... + + async def begin_weight_update(self, selector: str = "all") -> dict[str, Any]: ... + + async def update_weight_version(self, weight_version: str, abort_all_requests: bool = True) -> dict[str, Any]: ... + + async def end_weight_update(self) -> dict[str, Any]: ... + + async def resume_generation(self) -> None: ... + + +def _assert_success(result: dict[str, Any], step: str) -> None: + if result.get("success") is False: + raise RuntimeError(f"P2P control-plane smoke failed at {step}: {result.get('message')}") + + +async def run_p2p_control_plane_smoke( + replica: P2PControlPlaneReplica, + *, + num_ranks: int, + weight_version: str = "smoke-1", + selector: str = "all", +) -> dict[str, Any]: + """Exercise Verl -> SGLang P2P control API without transferring weights. + + Lifecycle (Miles P2P order): + abort -> metadata queries -> begin -> version bump -> end -> continue + """ + if num_ranks <= 0: + raise ValueError(f"num_ranks must be positive, got {num_ranks}") + + logger.info("P2P control-plane smoke: abort_all_requests") + await replica.abort_all_requests() + + metadata: dict[int, dict[str, Any]] = {} + for rank in range(num_ranks): + logger.info("P2P control-plane smoke: query metadata for rank=%s", rank) + transfer_info = await replica.get_remote_instance_transfer_engine_info(rank) + parallelism_info = await replica.get_parallelism_info(rank) + if transfer_info is None: + raise RuntimeError(f"P2P control-plane smoke: missing transfer engine info for rank {rank}") + if parallelism_info is None: + raise RuntimeError(f"P2P control-plane smoke: missing parallelism config for rank {rank}") + metadata[rank] = { + "transfer_info": transfer_info, + "parallelism_info": parallelism_info, + } + + logger.info("P2P control-plane smoke: begin_weight_update selector=%s", selector) + begin_result = await replica.begin_weight_update(selector) + _assert_success(begin_result, "begin_weight_update") + + logger.info("P2P control-plane smoke: update_weight_version -> %s", weight_version) + version_result = await replica.update_weight_version(weight_version, abort_all_requests=False) + _assert_success(version_result, "update_weight_version") + + logger.info("P2P control-plane smoke: end_weight_update") + end_result = await replica.end_weight_update() + _assert_success(end_result, "end_weight_update") + + logger.info("P2P control-plane smoke: resume_generation") + await replica.resume_generation() + + return { + "metadata": metadata, + "begin": begin_result, + "update_weight_version": version_result, + "end": end_result, + "weight_version": weight_version, + } diff --git a/scripts/run_p2p_live_smoke.py b/scripts/run_p2p_live_smoke.py new file mode 100755 index 00000000000..22d038e06a9 --- /dev/null +++ b/scripts/run_p2p_live_smoke.py @@ -0,0 +1,459 @@ +#!/usr/bin/env python3 +# Copyright 2025 Bytedance Ltd. and/or its affiliates +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Run P2P control-plane live smoke against a running Verl + SGLang rollout job. + +Requires: + - Ray cluster with named actors ``sglang_server_{replica_rank}_{node_rank}`` + - Rollout launched with checkpoint_engine.backend=p2p, e.g.: + actor_rollout_ref.rollout.checkpoint_engine.backend=p2p + +Example: + python scripts/run_p2p_live_smoke.py --replica-rank 0 --num-ranks 1 + python scripts/run_p2p_live_smoke.py --try-all-replicas --num-ranks 1 + python scripts/run_p2p_live_smoke.py --check-bootstrap-only --replica-rank 0 +""" + +from __future__ import annotations + +import argparse +import asyncio +import json +import logging +import sys +from pathlib import Path +from typing import Any + +_REPO_ROOT = Path(__file__).resolve().parents[1] +if str(_REPO_ROOT) not in sys.path: + sys.path.insert(0, str(_REPO_ROOT)) + +import requests # noqa: E402 + +from scripts.p2p_control_smoke import run_p2p_control_plane_smoke # noqa: E402 + +logger = logging.getLogger(__name__) + + +def _sglang_actor_name( + replica_rank: int, + *, + node_rank: int = 0, + name_suffix: str = "", + reward: bool = False, + teacher: bool = False, +) -> str: + if reward: + return f"sglang_server_reward_{replica_rank}_{node_rank}{name_suffix}" + if teacher: + return f"sglang_server_teacher_{replica_rank}_{node_rank}{name_suffix}" + return f"sglang_server_{replica_rank}_{node_rank}{name_suffix}" + + +def _format_ray_error(exc: BaseException) -> str: + try: + from ray.exceptions import RayTaskError + except ImportError: + return str(exc) + if isinstance(exc, RayTaskError): + if exc.cause is not None: + return f"{exc.cause} (ray task: {exc})" + return str(exc) + return str(exc) + + +def _ray_call(server: Any, method: str, **kwargs: Any) -> Any: + import ray + from ray.exceptions import RayTaskError + + try: + return ray.get(getattr(server, method).remote(**kwargs)) + except RayTaskError as exc: + raise RuntimeError(f"{method} failed: {_format_ray_error(exc)}") from exc + + +class RaySGLangReplica: + """SGLangReplica protocol implemented over a named Ray actor.""" + + def __init__(self, server: Any) -> None: + self._server = server + + async def _call(self, method: str, **kwargs: Any) -> Any: + return await asyncio.to_thread(_ray_call, self._server, method, **kwargs) + + async def abort_all_requests(self) -> None: + await self._call("abort_all_requests") + + async def get_remote_instance_transfer_engine_info(self, rank: int) -> Any: + return await self._call("get_remote_instance_transfer_engine_info", rank=rank) + + async def get_parallelism_info(self, rank: int) -> Any: + return await self._call("get_parallelism_info", rank=rank) + + async def begin_weight_update(self, selector: str = "all") -> dict[str, Any]: + return await self._call("begin_weight_update", selector=selector) + + async def update_weight_version(self, weight_version: str, abort_all_requests: bool = True) -> dict[str, Any]: + return await self._call( + "update_weight_version", + weight_version=weight_version, + abort_all_requests=abort_all_requests, + ) + + async def end_weight_update(self) -> dict[str, Any]: + return await self._call("end_weight_update") + + async def resume_generation(self) -> None: + await self._call("resume_generation") + + +def _actor_lookup_kwargs(actor_meta: dict[str, Any]) -> dict[str, Any]: + kwargs: dict[str, Any] = {"name": actor_meta["name"]} + namespace = actor_meta.get("namespace") + if namespace: + kwargs["namespace"] = namespace + return kwargs + + +def _resolve_named_actor(actor_name: str) -> Any: + import ray + from ray.util import list_named_actors + + try: + return ray.get_actor(actor_name) + except ValueError: + pass + + matches = [actor for actor in list_named_actors(all_namespaces=True) if actor.get("name") == actor_name] + if len(matches) == 1: + return ray.get_actor(**_actor_lookup_kwargs(matches[0])) + if len(matches) > 1: + namespaces = sorted({actor.get("namespace", "default") for actor in matches}) + raise RuntimeError( + f"Multiple actors named {actor_name!r} in namespaces {namespaces}. " + "Re-run with --ray-namespace ." + ) + raise ValueError( + f"Actor {actor_name!r} not found in any Ray namespace. " + "Run with --list-sglang-actors to inspect live rollout actors." + ) + + +async def _get_ray_server(actor_name: str) -> Any: + return await asyncio.to_thread(_resolve_named_actor, actor_name) + + +def _list_sglang_actors() -> None: + import ray + from ray.util import list_named_actors + + ray.init(address="auto", ignore_reinit_error=True) + actors = list_named_actors(all_namespaces=True) + matches = sorted( + (actor for actor in actors if "sglang_server" in actor.get("name", "")), + key=lambda actor: (actor.get("namespace", ""), actor.get("name", "")), + ) + if not matches: + print("No sglang_server* named actors found.") + return + namespaces = sorted({actor.get("namespace", "default") for actor in matches}) + for actor in matches: + print( + f"{actor.get('name')}\tnamespace={actor.get('namespace', 'default')}\tstate={actor.get('state', 'unknown')}" + ) + if len(namespaces) == 1: + print( + "\nExample:\n" + f" python scripts/run_p2p_live_smoke.py " + f"--ray-namespace {namespaces[0]} --replica-rank 0 --num-ranks 1 --check-bootstrap-only" + ) + + +async def _check_bootstrap_via_ray(replica: RaySGLangReplica, num_ranks: int) -> dict[str, Any]: + server_info = await replica._call("get_server_info") + metadata: dict[int, dict[str, Any]] = {} + for rank in range(num_ranks): + transfer_info = await replica.get_remote_instance_transfer_engine_info(rank) + parallelism_info = await replica.get_parallelism_info(rank) + if transfer_info is None: + raise RuntimeError(f"missing transfer engine info for rank {rank}") + if parallelism_info is None: + raise RuntimeError(f"missing parallelism config for rank {rank}") + metadata[rank] = { + "transfer_info": transfer_info, + "parallelism_info": parallelism_info, + } + return {"server_info": server_info, "metadata": metadata} + + +def _check_bootstrap_http(host: str, port: int, num_ranks: int, timeout: float) -> dict[str, Any]: + base = f"http://{host}:{port}" + health = requests.get(f"{base}/health", timeout=timeout) + health.raise_for_status() + + metadata: dict[int, dict[str, Any]] = {} + for rank in range(num_ranks): + transfer = requests.get( + f"{base}/get_transfer_engine_info", + params={"rank": rank}, + timeout=timeout, + ) + transfer.raise_for_status() + parallelism = requests.get( + f"{base}/get_parallelism_config", + params={"rank": rank}, + timeout=timeout, + ) + parallelism.raise_for_status() + metadata[rank] = { + "transfer_info": transfer.json()["remote_instance_transfer_engine_info"], + "parallelism_info": parallelism.json(), + } + + return {"health": health.text, "metadata": metadata} + + +async def _run_smoke_for_replica( + *, + actor_name: str, + num_ranks: int, + weight_version: str, + selector: str, + check_bootstrap_only: bool, + bootstrap_host: str, + bootstrap_port: int | None, + http_timeout: float, +) -> dict[str, Any]: + if bootstrap_port is not None: + logger.info( + "Checking bootstrap HTTP directly at %s:%s (no Ray actor lookup)", + bootstrap_host, + bootstrap_port, + ) + return _check_bootstrap_http(bootstrap_host, bootstrap_port, num_ranks, http_timeout) + + server = await _get_ray_server(actor_name) + replica = RaySGLangReplica(server) + if check_bootstrap_only: + logger.info("Checking bootstrap metadata via Ray actor %s", actor_name) + return await _check_bootstrap_via_ray(replica, num_ranks) + + logger.info("Running P2P control-plane smoke via Ray actor %s", actor_name) + return await run_p2p_control_plane_smoke( + replica, + num_ranks=num_ranks, + weight_version=weight_version, + selector=selector, + ) + + +async def _async_main(args: argparse.Namespace) -> int: + if args.list_sglang_actors: + _list_sglang_actors() + return 0 + + import ray + + ray.init( + address=args.ray_address, + namespace=args.ray_namespace, + ignore_reinit_error=True, + ) + + replica_ranks: list[int] + if args.try_all_replicas: + replica_ranks = list(range(args.max_replicas)) + elif args.replica_rank is not None: + replica_ranks = args.replica_rank + else: + replica_ranks = [0] + + failures = 0 + successes = 0 + for replica_rank in replica_ranks: + actor_name = _sglang_actor_name( + replica_rank, + node_rank=args.node_rank, + name_suffix=args.name_suffix, + reward=args.reward_model, + teacher=args.teacher_model, + ) + try: + result = await _run_smoke_for_replica( + actor_name=actor_name, + num_ranks=args.num_ranks, + weight_version=args.weight_version, + selector=args.selector, + check_bootstrap_only=args.check_bootstrap_only, + bootstrap_host=args.bootstrap_host, + bootstrap_port=args.bootstrap_port, + http_timeout=args.http_timeout, + ) + except ValueError as exc: + if args.try_all_replicas: + logger.info("Skip replica_rank=%s (%s)", replica_rank, exc) + continue + logger.error("Failed for %s: %s", actor_name, exc) + failures += 1 + continue + except Exception as exc: + err = str(exc).lower() + if args.try_all_replicas and ("does not exist" in err or "not found" in err): + logger.info("Stop at replica_rank=%s: actor not found", replica_rank) + break + if args.verbose: + logger.exception("Failed for %s", actor_name) + else: + logger.error("Failed for %s: %s", actor_name, exc) + failures += 1 + continue + + successes += 1 + label = actor_name if args.bootstrap_port is None else f"{args.bootstrap_host}:{args.bootstrap_port}" + print(f"OK {label}") + print(json.dumps(result, indent=2, default=str)) + + if successes == 0: + logger.error("No replicas passed smoke") + return 1 + if failures: + return 1 + return 0 + + +def _build_parser() -> argparse.ArgumentParser: + parser = argparse.ArgumentParser( + description="Run P2P control-plane live smoke on a running SGLang rollout.", + formatter_class=argparse.ArgumentDefaultsHelpFormatter, + ) + parser.add_argument( + "--replica-rank", + type=int, + action="append", + dest="replica_rank", + help="Replica rank to smoke. Repeat for multiple replicas.", + ) + parser.add_argument( + "--try-all-replicas", + action="store_true", + help="Try replica ranks 0..max-replicas-1 until an actor is missing.", + ) + parser.add_argument( + "--max-replicas", + type=int, + default=8, + help="Upper bound when --try-all-replicas is set.", + ) + parser.add_argument( + "--node-rank", + type=int, + default=0, + help="SGLang node rank within the replica (control RPCs use node_rank=0).", + ) + parser.add_argument( + "--name-suffix", + default="", + help="Actor name suffix, if rollout was launched with one.", + ) + parser.add_argument("--reward-model", action="store_true", help="Target reward-model actors.") + parser.add_argument("--teacher-model", action="store_true", help="Target teacher-model actors.") + parser.add_argument( + "--num-ranks", + type=int, + default=None, + help="Number of rollout ranks (tp * dp * pp).", + ) + parser.add_argument("--weight-version", default="live-smoke-1") + parser.add_argument("--selector", default="all") + parser.add_argument( + "--check-bootstrap-only", + action="store_true", + help="Only query bootstrap metadata via Ray actor RPCs; skip begin/end lifecycle.", + ) + parser.add_argument( + "--bootstrap-host", + default="127.0.0.1", + help="Host for direct HTTP bootstrap checks (--bootstrap-port).", + ) + parser.add_argument( + "--bootstrap-port", + type=int, + default=None, + help="Direct HTTP bootstrap check without Ray actor lookup (use host/port from job log).", + ) + parser.add_argument("--http-timeout", type=float, default=5.0) + parser.add_argument("--ray-address", default="auto") + parser.add_argument( + "--ray-namespace", + default=None, + help="Ray job namespace. Initializes the driver in that namespace before actor lookup.", + ) + parser.add_argument( + "--list-sglang-actors", + action="store_true", + help="List live sglang_server* named actors and exit.", + ) + parser.add_argument("-v", "--verbose", action="store_true") + return parser + + +def main() -> None: + parser = _build_parser() + args = parser.parse_args() + logging.basicConfig( + level=logging.DEBUG if args.verbose else logging.INFO, + format="%(levelname)s %(message)s", + ) + if args.try_all_replicas and args.replica_rank: + parser.error("Use either --replica-rank or --try-all-replicas, not both.") + if not args.list_sglang_actors and args.num_ranks is None: + parser.error("--num-ranks is required unless --list-sglang-actors is set.") + exit_code = asyncio.run(_async_main(args)) + sys.exit(exit_code) + + +if __name__ == "__main__": + main() + +######### +# Example: +# Посмотреть живые акторы: +# python scripts/run_p2p_live_smoke.py --list-sglang-actors +# +# Повторить smoke (должно заработать): +# python scripts/run_p2p_live_smoke.py \ +# --replica-rank 0 \ +# --num-ranks 1 \ +# --check-bootstrap-only +# +# Если namespace явный (из шага 1): +# python scripts/run_p2p_live_smoke.py \ +# --ray-namespace \ +# --replica-rank 0 \ +# --num-ranks 1 \ +# --check-bootstrap-only +# +# Полный control-plane smoke (без передачи весов): +# python scripts/run_p2p_live_smoke.py \ +# --ray-namespace \ +# --replica-rank 0 \ +# --num-ranks 1 +# +# Прямой HTTP (если ты на worker-ноде 10.224.194.231; порт из лога): +# python scripts/run_p2p_live_smoke.py \ +# --bootstrap-host 10.224.194.231 \ +# --bootstrap-port 62736 \ +# --num-ranks 1 \ +# --check-bootstrap-only diff --git a/tests/checkpoint_engine/test_p2p_checkpoint_engine.py b/tests/checkpoint_engine/test_p2p_checkpoint_engine.py new file mode 100644 index 00000000000..fcf22a90e44 --- /dev/null +++ b/tests/checkpoint_engine/test_p2p_checkpoint_engine.py @@ -0,0 +1,121 @@ +# Copyright 2025 Bytedance Ltd. and/or its affiliates +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from __future__ import annotations + +import asyncio +from unittest.mock import AsyncMock, MagicMock, patch + +import pytest +from omegaconf import OmegaConf + +from verl.checkpoint_engine import CheckpointEngineRegistry, P2PCheckpointEngine +from verl.checkpoint_engine.base import CheckpointEngineManager +from verl.checkpoint_engine.p2p.transfer_utils import resolve_rollout_model_path +from verl.workers.config.rollout import CheckpointEngineConfig + + +def test_p2p_backend_is_registered(): + assert CheckpointEngineRegistry.get("p2p") is P2PCheckpointEngine + + +def test_p2p_send_weights_requires_connect(): + import asyncio + + async def _run() -> None: + engine = P2PCheckpointEngine(bucket_size=1024, is_master=True) + with pytest.raises(RuntimeError, match="not connected"): + await engine.send_weights(iter([])) + + asyncio.run(_run()) + + +def test_p2p_build_topology_is_noop(): + trainer_kwargs, rollout_kwargs = P2PCheckpointEngine.build_topology(2, 4, [{}] * 6) + assert trainer_kwargs["rank"] == [-1, -1] + assert rollout_kwargs["rank"] == [-1, -1, -1, -1] + + +def test_p2p_prepare_finalize_noop(): + engine = P2PCheckpointEngine(bucket_size=1024, is_master=False) + assert engine.prepare()["backend"] == "p2p" + engine.init_process_group(rank=-1, world_size=0, metadata={}) + engine.finalize() + + +def test_resolve_rollout_model_path_from_omegaconf(): + model_config = OmegaConf.create({"path": "/tmp/hf-model", "use_shm": False}) + with patch("verl.utils.fs.copy_to_local", return_value="/tmp/hf-model-local") as copy_mock: + assert resolve_rollout_model_path(model_config) == "/tmp/hf-model-local" + copy_mock.assert_called_once_with("/tmp/hf-model", use_shm=False) + + +def test_update_weights_p2p_connects_rollout_metadata_once(): + replicas = [MagicMock(model_config=OmegaConf.create({"path": "/tmp/model", "use_shm": False}))] + actor_wg = MagicMock() + manager = CheckpointEngineManager( + config=CheckpointEngineConfig(backend="p2p"), + actor_wg=actor_wg, + replicas=replicas, + ) + + async def _run() -> None: + with ( + patch.object(manager, "abort_replicas", new_callable=AsyncMock), + patch.object(manager, "release_kv_cache_replicas", new_callable=AsyncMock), + patch.object(manager, "begin_weight_update_replicas", new_callable=AsyncMock, return_value=[]), + patch.object(manager, "end_weight_update_replicas", new_callable=AsyncMock), + patch.object(manager, "resume_kv_cache_replicas", new_callable=AsyncMock), + patch.object(manager, "resume_generation_replicas", new_callable=AsyncMock), + patch.object(manager, "update_weight_version_replicas", new_callable=AsyncMock), + patch( + "verl.checkpoint_engine.p2p.transfer_utils.collect_p2p_rollout_metadata", + new_callable=AsyncMock, + return_value={"model_path": "/tmp/model"}, + ) as collect_mock, + patch("verl.checkpoint_engine.base.ray.get"), + ): + await manager._update_weights_p2p(global_steps=1) + await manager._update_weights_p2p(global_steps=2) + + assert collect_mock.await_count == 1 + actor_wg.init_p2p_rollout_metadata.assert_called_once() + assert manager._p2p_rollout_metadata_initialized + + asyncio.run(_run()) + + +def test_add_replicas_invalidates_p2p_rollout_metadata(): + manager = CheckpointEngineManager( + config=CheckpointEngineConfig(backend="p2p"), + actor_wg=MagicMock(), + replicas=[], + ) + manager._p2p_rollout_metadata_initialized = True + manager.add_replicas([MagicMock()]) + assert not manager._p2p_rollout_metadata_initialized + + +def test_p2p_release_restore_delegates_to_updater(): + from unittest.mock import MagicMock + + engine = P2PCheckpointEngine(bucket_size=1024, is_master=True) + updater = MagicMock() + engine._updater = updater + + engine.release_for_checkpoint() + updater.release_for_checkpoint.assert_called_once() + + engine.restore_after_checkpoint() + updater.restore_after_checkpoint.assert_called_once() diff --git a/tests/checkpoint_engine/test_p2p_transfer_utils.py b/tests/checkpoint_engine/test_p2p_transfer_utils.py new file mode 100644 index 00000000000..769f210807c --- /dev/null +++ b/tests/checkpoint_engine/test_p2p_transfer_utils.py @@ -0,0 +1,457 @@ +# Copyright 2025 Bytedance Ltd. and/or its affiliates +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from __future__ import annotations + +import pickle +from unittest.mock import MagicMock, patch + +import torch + +from verl.checkpoint_engine.p2p.trainer_updater import P2PTrainerWeightUpdater +from verl.checkpoint_engine.p2p.transfer_utils import ( + P2PRolloutTopology, + P2PTransferManager, + RemoteTransferPlan, + TransferTaskP2PMeta, + build_rollout_topology_from_replicas, + create_transfer_engine, + filter_server_args_dict, + iter_named_tensor_buckets, + resolve_mooncake_transfer_engine_settings, + serialize_p2p_rollout_metadata, + unregister_cpu_memory, +) + + +class _FakeReplica: + def __init__(self, world_size: int) -> None: + self.world_size = world_size + self.config = type("Cfg", (), {"pipeline_model_parallel_size": 1})() + + +def test_iter_named_tensor_buckets_splits_by_size(): + tensors = [ + ("a", torch.zeros(100, dtype=torch.float32)), + ("b", torch.zeros(100, dtype=torch.float32)), + ("c", torch.zeros(100, dtype=torch.float32)), + ] + # Each tensor is 400 bytes; bucket limit 801 fits two tensors, then one. + buckets = list(iter_named_tensor_buckets(iter(tensors), bucket_bytes=801, clone_tensors=False)) + assert len(buckets) == 2 + assert [name for name, _ in buckets[0]] == ["a", "b"] + assert [name for name, _ in buckets[1]] == ["c"] + + +def test_iter_named_tensor_buckets_oversized_tensor_gets_own_bucket(): + tensors = [ + ("big", torch.zeros(200, dtype=torch.float32)), + ("small", torch.zeros(10, dtype=torch.float32)), + ] + buckets = list(iter_named_tensor_buckets(iter(tensors), bucket_bytes=100, clone_tensors=False)) + assert len(buckets) == 2 + assert buckets[0][0][0] == "big" + assert buckets[1][0][0] == "small" + + +def test_iter_named_tensor_buckets_rejects_invalid_size(): + import pytest + + with pytest.raises(ValueError, match="bucket_bytes"): + list(iter_named_tensor_buckets(iter([]), bucket_bytes=0)) + + +def test_build_rollout_topology_from_replicas(): + replicas = [_FakeReplica(8), _FakeReplica(8)] + topology = build_rollout_topology_from_replicas(replicas) + assert topology.engine_count == 2 + assert topology.gpus_per_engine == 8 + assert topology.total_gpus == 16 + + +def test_remote_transfer_plan_round_robin(): + topology = P2PRolloutTopology(engine_count=2, gpus_per_engine=2) + with patch( + "verl.checkpoint_engine.p2p.transfer_utils.resolve_trainer_parallelism", + return_value=(0, 1, 0, 4), + ): + plan = RemoteTransferPlan(topology) + tasks = plan.plan_p2p() + assert tasks == [TransferTaskP2PMeta(engine_ind=0, engine_rank=0)] + + with patch( + "verl.checkpoint_engine.p2p.transfer_utils.resolve_trainer_parallelism", + return_value=(0, 1, 1, 4), + ): + plan = RemoteTransferPlan(topology) + tasks = plan.plan_p2p() + assert tasks == [TransferTaskP2PMeta(engine_ind=0, engine_rank=1)] + + +def test_serialize_p2p_rollout_metadata_roundtrip_keys(): + topology = P2PRolloutTopology(engine_count=1, gpus_per_engine=1) + payload = serialize_p2p_rollout_metadata( + model_path="/tmp/model", + rollout_topology=topology, + engine_kwargs={"p2p_transfer_num_workers": 2}, + remote_weight_infos_by_session_id={"sess": ({"w": (1, 2, 4)}, {"tp_size": 1})}, + targets_to_session_id={(0, 0): "sess"}, + session_id_to_server_args={"sess": {"model_path": "/tmp/model", "trust_remote_code": True}}, + ) + assert payload["targets_to_session_id"] == {(0, 0): "sess"} + assert payload["session_id_to_server_args"]["sess"]["trust_remote_code"] is True + + +def test_filter_server_args_dict_sanitizes_non_serializable_fields(): + class _CustomConfig: + __module__ = "transformers_modules.configuration_deepseek" + + filtered = filter_server_args_dict( + { + "model_path": "/tmp/model", + "trust_remote_code": True, + "internal_states": {"config": _CustomConfig()}, + "custom_sigquit_handler": lambda: None, + } + ) + assert filtered["model_path"] == "/tmp/model" + assert filtered["trust_remote_code"] is True + # Not a ServerArgs field: always dropped. + assert "internal_states" not in filtered + # A non-serializable value never survives as-is: depending on the sglang + # version the key is either absent (no such ServerArgs field) or nullified. + assert filtered.get("custom_sigquit_handler") is None + # The key guarantee: the result crosses Ray/pickle boundaries safely. + pickle.loads(pickle.dumps(filtered)) + + +def test_serialize_p2p_rollout_metadata_is_pickle_safe(): + topology = P2PRolloutTopology(engine_count=1, gpus_per_engine=1) + server_args = filter_server_args_dict( + { + "model_path": "/tmp/model", + "trust_remote_code": True, + "internal_states": object(), + } + ) + payload = serialize_p2p_rollout_metadata( + model_path="/tmp/model", + rollout_topology=topology, + engine_kwargs={}, + remote_weight_infos_by_session_id={"sess": ({"w": (1, 2, 4)}, {"tp_size": 1})}, + targets_to_session_id={(0, 0): "sess"}, + session_id_to_server_args={"sess": server_args}, + ) + restored = pickle.loads(pickle.dumps(payload)) + assert restored["session_id_to_server_args"]["sess"] == { + "model_path": "/tmp/model", + "trust_remote_code": True, + } + + +def test_resolve_mooncake_transfer_engine_settings_defaults(): + with patch.dict("os.environ", {}, clear=True): + assert resolve_mooncake_transfer_engine_settings() == ("rdma", "") + + +def test_resolve_mooncake_transfer_engine_settings_reads_env(): + with patch.dict( + "os.environ", + { + "MOONCAKE_PROTOCOL": "rdma", + "MOONCAKE_DEVICE": "^=mlx5_0,mlx5_1", + }, + clear=True, + ): + assert resolve_mooncake_transfer_engine_settings() == ("rdma", "^=mlx5_0,mlx5_1") + + +def test_resolve_mooncake_transfer_engine_settings_mc_force_tcp_clears_device(): + with patch.dict( + "os.environ", + { + "MOONCAKE_PROTOCOL": "rdma", + "MOONCAKE_DEVICE": "^=mlx5_0", + "MC_FORCE_TCP": "1", + }, + clear=True, + ): + assert resolve_mooncake_transfer_engine_settings() == ("rdma", "") + + +def test_create_transfer_engine_passes_mooncake_env_to_initialize(): + engine = MagicMock() + engine.initialize.return_value = 0 + with ( + patch.dict( + "os.environ", + { + "MOONCAKE_PROTOCOL": "rdma", + "MOONCAKE_DEVICE": "^=mlx5_0", + }, + clear=True, + ), + patch( + "verl.checkpoint_engine.p2p.transfer_utils.ray._private.services.get_node_ip_address", + return_value="10.0.0.1", + ), + patch("mooncake.engine.TransferEngine", return_value=engine), + ): + created = create_transfer_engine() + assert created is engine + engine.initialize.assert_called_once_with("10.0.0.1", "P2PHANDSHAKE", "rdma", "^=mlx5_0") + + +def test_unregister_cpu_memory_calls_unregister(): + engine = MagicMock() + engine.unregister_memory.return_value = 0 + registry = {"weight": (123, 4, 2)} + unregister_cpu_memory(registry, engine) + engine.unregister_memory.assert_called_once_with(123) + assert registry == {} + + +def test_p2p_updater_release_and_restore(): + from unittest.mock import PropertyMock + + topology = P2PRolloutTopology(engine_count=1, gpus_per_engine=1) + with patch( + "verl.checkpoint_engine.p2p.transfer_utils.resolve_trainer_parallelism", + return_value=(0, 1, 0, 1), + ): + updater = P2PTrainerWeightUpdater(model_path="/tmp/model", rollout_topology=topology, is_master=True) + metadata = { + "remote_weight_infos_by_session_id": {"sid": (["ptr"], {"tp": 1})}, + "targets_to_session_id": {(0, 0): "sid"}, + "session_id_to_server_args": {"sid": {"model_path": "/tmp/model"}}, + } + updater._rollout_metadata = metadata + updater._connected = True + updater._model_registered = True + updater._transfer_engine = MagicMock() + updater._weight_memory_registry = {"p": (1, 2, 4)} + updater._shared_params_dict = {"p": torch.zeros(2)} + updater._transfer_engine_meta_list = [(MagicMock(), [])] + + with patch.object(P2PTrainerWeightUpdater, "is_source", new_callable=PropertyMock, return_value=True): + updater.release_for_checkpoint() + + assert not updater._connected + assert updater._transfer_engine is None + assert not updater._shared_params_dict + assert updater._rollout_metadata["targets_to_session_id"] == {(0, 0): "sid"} + + with patch.object(updater, "connect_rollout_metadata") as connect_mock: + updater.restore_after_checkpoint() + connect_mock.assert_called_once_with(metadata) + + +def test_p2p_release_keeps_rollout_metadata_for_restore(): + from unittest.mock import PropertyMock + + topology = P2PRolloutTopology(engine_count=2, gpus_per_engine=1) + with patch( + "verl.checkpoint_engine.p2p.transfer_utils.resolve_trainer_parallelism", + return_value=(0, 1, 0, 1), + ): + updater = P2PTrainerWeightUpdater(model_path="/tmp/model", rollout_topology=topology, is_master=True) + + remote_infos = {"sid0": (["ptr0"], {"tp": 1})} + targets_to_session = {(0, 0): "sid0"} + server_args = {"sid0": {"model_path": "/tmp/model"}} + + with ( + patch.object(P2PTrainerWeightUpdater, "is_source", new_callable=PropertyMock, return_value=True), + patch.object(updater, "_create_cpu_replica", return_value=MagicMock(named_parameters=lambda: [])), + patch( + "verl.checkpoint_engine.p2p.trainer_updater.create_transfer_engine", + return_value=MagicMock(), + ), + patch.object(updater.transfer_plan, "plan_p2p", return_value=[]), + ): + updater.connect_rollout_metadata( + { + "remote_weight_infos_by_session_id": remote_infos, + "targets_to_session_id": targets_to_session, + "session_id_to_server_args": server_args, + } + ) + + with patch.object(P2PTrainerWeightUpdater, "is_source", new_callable=PropertyMock, return_value=True): + updater.release_for_checkpoint() + + assert updater._rollout_metadata["targets_to_session_id"] == {(0, 0): "sid0"} + + +def test_p2p_updater_send_weights_uses_buckets(): + from unittest.mock import PropertyMock + + topology = P2PRolloutTopology(engine_count=1, gpus_per_engine=1) + with patch( + "verl.checkpoint_engine.p2p.transfer_utils.resolve_trainer_parallelism", + return_value=(0, 1, 0, 1), + ): + updater = P2PTrainerWeightUpdater( + model_path="/tmp/model", + rollout_topology=topology, + is_master=True, + bucket_size_bytes=16, + ) + updater._connected = True + updater._model_registered = True + updater._shared_params_dict = {"p": torch.zeros(2)} + updater._transfer_engine = MagicMock() + updater._weight_memory_registry = {"p": (1, 2, 4)} + updater._transfer_engine_meta_list = [(MagicMock(), [MagicMock()])] + + tensors = [ + ("hf.p", torch.ones(2)), + ("hf.q", torch.ones(2)), + ("hf.r", torch.ones(2)), + ] + + with ( + patch.object(P2PTrainerWeightUpdater, "is_source", new_callable=PropertyMock, return_value=True), + patch.object( + updater, + "_get_transfer_ready_params", + side_effect=[ + (["p"], [("hf.p", tensors[0][1])]), + (["q", "r"], [("hf.q", tensors[1][1]), ("hf.r", tensors[2][1])]), + ], + ) as get_ready_mock, + patch.object(updater, "_transfer_ready_bucket", return_value=1) as transfer_mock, + patch.object(updater.transfer_manager, "wait_transfers"), + ): + updater.send_weights(name_tensor for name_tensor in tensors) + + assert get_ready_mock.call_count == 2 + assert transfer_mock.call_count == 2 + + +def test_p2p_transfer_manager_wait_futures(): + from concurrent.futures import Future + + manager = P2PTransferManager() + future = Future() + future.set_result(None) + manager.transfer_futures.append(future) + + manager.wait_futures([future]) + manager.wait_transfers() + assert manager.transfer_futures == [] + + +def test_p2p_send_weights_waits_only_at_end_like_miles(): + from unittest.mock import PropertyMock + + topology = P2PRolloutTopology(engine_count=1, gpus_per_engine=1) + with patch( + "verl.checkpoint_engine.p2p.transfer_utils.resolve_trainer_parallelism", + return_value=(0, 1, 0, 1), + ): + updater = P2PTrainerWeightUpdater( + model_path="/tmp/model", + rollout_topology=topology, + bucket_size_bytes=9, + ) + updater._connected = True + updater._model_registered = True + updater._shared_params_dict = {"p": torch.zeros(2)} + updater._transfer_engine = MagicMock() + updater._weight_memory_registry = {"p": (1, 2, 4)} + updater._transfer_engine_meta_list = [(MagicMock(), [MagicMock()])] + + tensors = [("hf.p", torch.ones(2)), ("hf.q", torch.ones(2))] + + with ( + patch.object(P2PTrainerWeightUpdater, "is_source", new_callable=PropertyMock, return_value=True), + patch.object( + updater, + "_get_transfer_ready_params", + side_effect=[ + (["p"], [tensors[0]]), + (["q"], [tensors[1]]), + ], + ), + patch.object(updater, "_transfer_ready_bucket", return_value=1) as transfer_mock, + patch.object(updater.transfer_manager, "wait_futures") as wait_futures_mock, + patch.object(updater.transfer_manager, "wait_transfers") as wait_transfers_mock, + ): + updater.send_weights(name_tensor for name_tensor in tensors) + + assert transfer_mock.call_count == 2 + wait_futures_mock.assert_not_called() + wait_transfers_mock.assert_called_once() + + +def test_p2p_transfer_ready_bucket_last_engine_is_async(): + from concurrent.futures import Future + from unittest.mock import PropertyMock + + topology = P2PRolloutTopology(engine_count=1, gpus_per_engine=2) + with patch( + "verl.checkpoint_engine.p2p.transfer_utils.resolve_trainer_parallelism", + return_value=(0, 1, 0, 1), + ): + updater = P2PTrainerWeightUpdater(model_path="/tmp/model", rollout_topology=topology) + + replica0 = MagicMock() + replica1 = MagicMock() + session0 = MagicMock() + session1 = MagicMock() + updater._transfer_engine_meta_list = [ + (replica0, [session0]), + (replica1, [session1]), + ] + updater._weight_memory_registry = {"w": (1, 2, 4)} + updater._transfer_engine = MagicMock() + updater._transfer_engine.batch_transfer_sync_write.return_value = 0 + + pending = Future() + pending.set_result(None) + + with ( + patch.object(P2PTrainerWeightUpdater, "is_source", new_callable=PropertyMock, return_value=True), + patch.object(updater.transfer_manager, "submit_returning_future", return_value=pending) as submit_mock, + patch.object(updater.transfer_manager, "submit") as submit_async_mock, + ): + writes = updater._transfer_ready_bucket(["w"], [("hf.w", torch.ones(2))]) + + assert writes == 2 + submit_mock.assert_called_once() + submit_async_mock.assert_called_once() + replica0.load_weights.assert_called_once() + replica1.load_weights.assert_called_once() + + +def test_get_transfer_ready_params_raises_on_unmapped_replica_param(): + import pytest + + topology = P2PRolloutTopology(engine_count=1, gpus_per_engine=1) + with patch( + "verl.checkpoint_engine.p2p.transfer_utils.resolve_trainer_parallelism", + return_value=(0, 1, 0, 1), + ): + updater = P2PTrainerWeightUpdater(model_path="/tmp/model", rollout_topology=topology) + + updater._shared_params_dict = {"known.weight": torch.zeros(2)} + mapped_result = MagicMock() + mapped_result.sglang_name = "missing.weight" + mapped_result.num_shards = 1 + mapped_result.num_local_experts = None + updater._shared_param_mapper = MagicMock(map=MagicMock(return_value=mapped_result)) + + with pytest.raises(RuntimeError, match="missing.weight"): + updater._get_transfer_ready_params([("hf.missing.weight", torch.ones(2))]) diff --git a/tests/checkpoint_engine/test_weight_checker.py b/tests/checkpoint_engine/test_weight_checker.py new file mode 100644 index 00000000000..3d41dc0bc78 --- /dev/null +++ b/tests/checkpoint_engine/test_weight_checker.py @@ -0,0 +1,193 @@ +# Copyright 2025 Bytedance Ltd. and/or its affiliates +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from __future__ import annotations + +import asyncio +from unittest.mock import AsyncMock, patch + +import pytest + +from verl.checkpoint_engine.p2p.weight_checker import ( + P2PWeightChecker, + check_weights_on_replicas, + compare_weight_updates, + setup_weight_checker_snapshots, +) +from verl.workers.config.rollout import CheckpointEngineConfig + + +class _FakeReplica: + def __init__(self, replica_rank: int, *, success: bool = True) -> None: + self.replica_rank = replica_rank + self.check_weights = AsyncMock(return_value={"success": success, "message": "ok", "replica_rank": replica_rank}) + + +def test_setup_weight_checker_snapshots_order(): + replicas = [_FakeReplica(0), _FakeReplica(1)] + + async def _run() -> None: + await setup_weight_checker_snapshots(replicas) + + asyncio.run(_run()) + + assert replicas[0].check_weights.await_count == 2 + assert replicas[1].check_weights.await_count == 2 + assert replicas[0].check_weights.await_args_list[0].kwargs == { + "action": "snapshot", + "allow_quant_error": False, + } + assert replicas[0].check_weights.await_args_list[1].kwargs == { + "action": "reset_tensors", + "allow_quant_error": False, + } + + +def test_compare_weight_updates_passes_allow_quant_error(): + replica = _FakeReplica(0) + + async def _run() -> None: + with patch("builtins.print") as print_mock: + await compare_weight_updates([replica], allow_quant_error=True) + print_mock.assert_called_once() + assert "compare PASSED" in print_mock.call_args[0][0] + + asyncio.run(_run()) + + replica.check_weights.assert_awaited_once_with(action="compare", allow_quant_error=True) + + +def test_check_weights_on_replicas_raises_on_failure(): + replicas = [_FakeReplica(0, success=True), _FakeReplica(1, success=False)] + + async def _run() -> None: + await check_weights_on_replicas(replicas, action="compare") + + with pytest.raises(RuntimeError, match="replica=1"): + asyncio.run(_run()) + + +def test_checkpoint_engine_config_has_weight_checker_flags(): + cfg = CheckpointEngineConfig(check_weight_update_equal=True) + assert cfg.check_weight_update_equal is True + assert cfg.check_weight_update_allow_quant_error is False + + +def test_p2p_weight_checker_runs_once(): + replicas = [_FakeReplica(0)] + checker = P2PWeightChecker(enabled=True) + + async def _run() -> None: + with ( + patch( + "verl.checkpoint_engine.p2p.weight_checker.setup_weight_checker_snapshots", + new_callable=AsyncMock, + ) as setup_mock, + patch( + "verl.checkpoint_engine.p2p.weight_checker.compare_weight_updates", + new_callable=AsyncMock, + ) as compare_mock, + ): + await checker.setup(replicas) + await checker.setup(replicas) + await checker.compare(replicas) + await checker.compare(replicas) + + setup_mock.assert_awaited_once_with(replicas) + compare_mock.assert_awaited_once_with(replicas, allow_quant_error=False) + + asyncio.run(_run()) + assert checker._setup_done + assert checker._compared + + +def test_p2p_weight_checker_disabled_is_noop(): + replicas = [_FakeReplica(0)] + checker = P2PWeightChecker(enabled=False) + + async def _run() -> None: + with ( + patch( + "verl.checkpoint_engine.p2p.weight_checker.setup_weight_checker_snapshots", + new_callable=AsyncMock, + ) as setup_mock, + patch( + "verl.checkpoint_engine.p2p.weight_checker.compare_weight_updates", + new_callable=AsyncMock, + ) as compare_mock, + ): + await checker.setup(replicas) + await checker.compare(replicas) + + setup_mock.assert_not_awaited() + compare_mock.assert_not_awaited() + + asyncio.run(_run()) + assert not checker._setup_done + assert not checker._compared + + +def test_checkpoint_engine_manager_weight_checker_hooks(): + from unittest.mock import MagicMock, patch + + from verl.checkpoint_engine.base import CheckpointEngineManager + from verl.workers.config.rollout import CheckpointEngineConfig + + replicas = [_FakeReplica(0)] + manager = CheckpointEngineManager( + config=CheckpointEngineConfig(backend="p2p", check_weight_update_equal=True), + actor_wg=MagicMock(), + replicas=replicas, + ) + + async def _run() -> None: + with ( + patch( + "verl.checkpoint_engine.p2p.weight_checker.setup_weight_checker_snapshots", + new_callable=AsyncMock, + ) as setup_mock, + patch( + "verl.checkpoint_engine.p2p.weight_checker.compare_weight_updates", + new_callable=AsyncMock, + ) as compare_mock, + patch.object(manager, "_update_weights_p2p", new_callable=AsyncMock), + ): + await manager.update_weights(global_steps=0) + setup_mock.assert_awaited_once_with(replicas) + compare_mock.assert_awaited_once_with(replicas, allow_quant_error=False) + + asyncio.run(_run()) + assert manager._p2p_weight_checker._setup_done + assert manager._p2p_weight_checker._compared + + +def test_checkpoint_engine_manager_skips_weight_checker_for_naive(): + from unittest.mock import MagicMock, patch + + from verl.checkpoint_engine.base import CheckpointEngineManager + from verl.workers.config.rollout import CheckpointEngineConfig + + replicas = [_FakeReplica(0)] + manager = CheckpointEngineManager( + config=CheckpointEngineConfig(backend="naive", check_weight_update_equal=True), + actor_wg=MagicMock(), + replicas=replicas, + ) + + async def _run() -> None: + with patch("verl.checkpoint_engine.base.ray.get"): + await manager.update_weights(global_steps=0) + + asyncio.run(_run()) + assert manager._p2p_weight_checker is None diff --git a/tests/utils/test_megatron_pp_local_export.py b/tests/utils/test_megatron_pp_local_export.py new file mode 100644 index 00000000000..353b3009280 --- /dev/null +++ b/tests/utils/test_megatron_pp_local_export.py @@ -0,0 +1,105 @@ +# Copyright 2025 Bytedance Ltd. and/or its affiliates +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from __future__ import annotations + +import sys +import types +from types import SimpleNamespace +from unittest.mock import MagicMock + +import pytest +import torch + +from verl.utils.megatron_pp_local_export import ( + disable_megatron_bridge_pp_broadcast, + export_pp_local_hf_weights, + filter_local_conversion_tasks, +) + + +def test_filter_local_conversion_tasks_keeps_only_local_params(): + local = SimpleNamespace(megatron_module=object(), param_weight=torch.ones(2)) + remote = SimpleNamespace(megatron_module=None, param_weight=None) + placeholder = None + + assert filter_local_conversion_tasks([local, remote, placeholder]) == [local] + + +def test_disable_megatron_bridge_pp_broadcast_skips_collectives(monkeypatch: pytest.MonkeyPatch): + param_mapping_module = types.ModuleType("megatron.bridge.models.conversion.param_mapping") + + class DummyMapping: + def broadcast_from_pp_rank(self, tensor, cache_key=None): + raise AssertionError("PP tensor broadcast should be disabled") + + def broadcast_obj_from_pp_rank(self, obj, cache_key=None): + raise AssertionError("PP object broadcast should be disabled") + + param_mapping_module.MegatronParamMapping = DummyMapping + monkeypatch.setitem(sys.modules, "megatron", types.ModuleType("megatron")) + monkeypatch.setitem(sys.modules, "megatron.bridge", types.ModuleType("megatron.bridge")) + monkeypatch.setitem(sys.modules, "megatron.bridge.models", types.ModuleType("megatron.bridge.models")) + monkeypatch.setitem( + sys.modules, + "megatron.bridge.models.conversion", + types.ModuleType("megatron.bridge.models.conversion"), + ) + monkeypatch.setitem( + sys.modules, + "megatron.bridge.models.conversion.param_mapping", + param_mapping_module, + ) + + from megatron.bridge.models.conversion.param_mapping import MegatronParamMapping + + mapping = MegatronParamMapping() + with disable_megatron_bridge_pp_broadcast(): + assert torch.equal(mapping.broadcast_from_pp_rank(torch.tensor([1.0])), torch.tensor([1.0])) + assert mapping.broadcast_obj_from_pp_rank({"x": 1}) == {"x": 1} + + +def test_export_pp_local_hf_weights_uses_local_tasks_only(monkeypatch: pytest.MonkeyPatch): + local_task = SimpleNamespace(megatron_module=object(), param_weight=torch.ones(2)) + remote_task = SimpleNamespace(megatron_module=None, param_weight=None) + + model_bridge = MagicMock() + model_bridge.stream_weights_megatron_to_hf.return_value = [("layer.weight", torch.ones(3))] + + bridge = MagicMock() + bridge.hf_pretrained = object() + bridge._model_bridge = model_bridge + bridge.get_conversion_tasks.return_value = [local_task, remote_task] + + monkeypatch.setattr( + "verl.utils.megatron_pp_local_export.disable_megatron_bridge_pp_broadcast", + lambda: _nullcontext(), + ) + + weights = list(export_pp_local_hf_weights(bridge, [MagicMock()])) + assert len(weights) == 1 + assert weights[0][0] == "layer.weight" + assert torch.equal(weights[0][1], torch.ones(3)) + + model_bridge.stream_weights_megatron_to_hf.assert_called_once() + _, kwargs = model_bridge.stream_weights_megatron_to_hf.call_args + assert kwargs["conversion_tasks"] == [local_task] + + +class _nullcontext: + def __enter__(self): + return None + + def __exit__(self, *args): + return False diff --git a/tests/workers/rollout/sglang_rollout/test_p2p_bootstrap.py b/tests/workers/rollout/sglang_rollout/test_p2p_bootstrap.py new file mode 100644 index 00000000000..84e20a8ebf6 --- /dev/null +++ b/tests/workers/rollout/sglang_rollout/test_p2p_bootstrap.py @@ -0,0 +1,55 @@ +# Copyright 2025 Bytedance Ltd. and/or its affiliates +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from verl.workers.config import CheckpointEngineConfig, RolloutConfig +from verl.workers.rollout.sglang_rollout.p2p_bootstrap import ( + sanitize_sglang_engine_kwargs_for_p2p, + should_enable_p2p_weight_transfer_bootstrap, +) + + +def test_should_enable_p2p_bootstrap_from_checkpoint_backend(): + config = RolloutConfig( + name="sglang", + checkpoint_engine=CheckpointEngineConfig(backend="p2p"), + ) + assert should_enable_p2p_weight_transfer_bootstrap(config) is True + + +def test_should_disable_p2p_bootstrap_by_default(): + config = RolloutConfig(name="sglang") + assert should_enable_p2p_weight_transfer_bootstrap(config) is False + + +def test_should_disable_p2p_bootstrap_when_engine_kwargs_set_seed_flag(): + config = RolloutConfig( + name="sglang", + engine_kwargs={"sglang": {"remote_instance_weight_loader_start_seed_via_transfer_engine": True}}, + ) + assert should_enable_p2p_weight_transfer_bootstrap(config) is False + + +def test_sanitize_strips_p2p_bootstrap_flags_from_engine_kwargs(): + config = RolloutConfig( + name="sglang", + engine_kwargs={ + "sglang": { + "engine_info_bootstrap_port": 19001, + "remote_instance_weight_loader_start_seed_via_transfer_engine": True, + } + }, + ) + sanitized = sanitize_sglang_engine_kwargs_for_p2p(config) + assert "engine_info_bootstrap_port" not in sanitized + assert "remote_instance_weight_loader_start_seed_via_transfer_engine" not in sanitized diff --git a/tests/workers/rollout/sglang_rollout/test_p2p_control_smoke.py b/tests/workers/rollout/sglang_rollout/test_p2p_control_smoke.py new file mode 100644 index 00000000000..cf42ed2dac5 --- /dev/null +++ b/tests/workers/rollout/sglang_rollout/test_p2p_control_smoke.py @@ -0,0 +1,105 @@ +# Copyright 2025 Bytedance Ltd. and/or its affiliates +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from __future__ import annotations + +import asyncio +from typing import Any + +import pytest + +from scripts.p2p_control_smoke import run_p2p_control_plane_smoke + + +class _RecordingReplica: + def __init__(self, num_ranks: int) -> None: + self.num_ranks = num_ranks + self.calls: list[tuple[str, tuple, dict]] = [] + + def _record(self, name: str, *args, **kwargs) -> None: + self.calls.append((name, args, kwargs)) + + async def abort_all_requests(self) -> None: + self._record("abort_all_requests") + + async def get_remote_instance_transfer_engine_info(self, rank: int) -> Any: + self._record("get_remote_instance_transfer_engine_info", rank) + return (f"session-{rank}", {f"weight.{rank}": (0, 1, 4)}) + + async def get_parallelism_info(self, rank: int) -> dict[str, Any]: + self._record("get_parallelism_info", rank) + return {"tp_rank": rank, "tp_size": self.num_ranks} + + async def begin_weight_update(self, selector: str = "all") -> dict[str, Any]: + self._record("begin_weight_update", selector) + return {"success": True, "message": "ok"} + + async def update_weight_version(self, weight_version: str, abort_all_requests: bool = True) -> dict[str, Any]: + self._record( + "update_weight_version", + weight_version, + abort_all_requests=abort_all_requests, + ) + return {"success": True, "message": "ok", "new_version": weight_version} + + async def end_weight_update(self) -> dict[str, Any]: + self._record("end_weight_update") + return {"success": True, "message": "ok"} + + async def resume_generation(self) -> None: + self._record("resume_generation") + + +def test_p2p_control_plane_smoke_call_order(): + replica = _RecordingReplica(num_ranks=2) + + async def _run() -> None: + result = await run_p2p_control_plane_smoke(replica, num_ranks=2, weight_version="42") + assert result["weight_version"] == "42" + assert set(result["metadata"]) == {0, 1} + + asyncio.run(_run()) + + call_names = [name for name, _, _ in replica.calls] + assert call_names == [ + "abort_all_requests", + "get_remote_instance_transfer_engine_info", + "get_parallelism_info", + "get_remote_instance_transfer_engine_info", + "get_parallelism_info", + "begin_weight_update", + "update_weight_version", + "end_weight_update", + "resume_generation", + ] + + _, _, update_kwargs = replica.calls[6] + assert update_kwargs["abort_all_requests"] is False + + +def test_p2p_control_plane_smoke_fails_on_missing_metadata(): + class _BrokenReplica(_RecordingReplica): + async def get_parallelism_info(self, rank: int) -> dict[str, Any]: + return None + + async def _run() -> None: + replica = _BrokenReplica(num_ranks=1) + await run_p2p_control_plane_smoke(replica, num_ranks=1) + + with pytest.raises(RuntimeError, match="missing parallelism config"): + asyncio.run(_run()) + + +def test_p2p_control_plane_smoke_sync_entrypoint(): + asyncio.run(run_p2p_control_plane_smoke(_RecordingReplica(1), num_ranks=1)) diff --git a/verl/checkpoint_engine/__init__.py b/verl/checkpoint_engine/__init__.py index 124f2b35307..87be932673d 100644 --- a/verl/checkpoint_engine/__init__.py +++ b/verl/checkpoint_engine/__init__.py @@ -71,3 +71,7 @@ __all__ += ["DeltaShardedCheckpointEngine"] except ImportError: DeltaShardedCheckpointEngine = None + +from .p2p.checkpoint_engine import P2PCheckpointEngine + +__all__ += ["P2PCheckpointEngine"] diff --git a/verl/checkpoint_engine/base.py b/verl/checkpoint_engine/base.py index 3c0d230296c..72346b4dc90 100644 --- a/verl/checkpoint_engine/base.py +++ b/verl/checkpoint_engine/base.py @@ -399,6 +399,16 @@ def __init__( self.backend_cls = CheckpointEngineRegistry.get(config.backend) self.actor_wg = actor_wg self.replicas = replicas + self._p2p_rollout_metadata_initialized = False + if config.backend == "p2p": + from verl.checkpoint_engine.p2p.weight_checker import P2PWeightChecker + + self._p2p_weight_checker = P2PWeightChecker( + enabled=config.check_weight_update_equal, + allow_quant_error=config.check_weight_update_allow_quant_error, + ) + else: + self._p2p_weight_checker = None def build_process_group(self, rollout: RayWorkerGroup): """Build process group for actor worker group and rollout replicas.""" @@ -434,6 +444,7 @@ def add_replicas(self, replicas: list[RolloutReplica]): replicas: The list of rollout replicas to add. """ self.replicas.extend(replicas) + self._invalidate_p2p_rollout_metadata() def remove_replicas(self, replicas: list[RolloutReplica]): """Remove rollout replicas from the manager for elastic scale down, will rebuild process group. @@ -443,6 +454,11 @@ def remove_replicas(self, replicas: list[RolloutReplica]): """ replicas_set = set(replicas) self.replicas = [r for r in self.replicas if r not in replicas_set] + self._invalidate_p2p_rollout_metadata() + + def _invalidate_p2p_rollout_metadata(self) -> None: + """Drop cached P2P rollout metadata so the next sync re-queries rollout engines.""" + self._p2p_rollout_metadata_initialized = False @auto_await async def sleep_replicas(self): @@ -482,6 +498,72 @@ async def resume_kv_cache_replicas(self): """ await asyncio.gather(*[r.resume_kv_cache() for r in self.replicas]) + @auto_await + async def begin_weight_update_replicas(self, selector: str = "all") -> list[RolloutReplica]: + """Open one weight-update session per rollout replica. + + SGLang's begin/end API is a replica-level transaction boundary. It must + not be opened independently by every CheckpointEngineWorker. + """ + started_replicas = [] + try: + for replica in self.replicas: + result = await replica.begin_weight_update(selector) + self._raise_for_weight_update_control_result("begin_weight_update", replica, result) + started_replicas.append(replica) + except Exception: + if started_replicas: + await self.end_weight_update_replicas(started_replicas) + raise + return started_replicas + + @auto_await + async def end_weight_update_replicas(self, replicas: list[RolloutReplica] | None = None): + """Close weight-update sessions opened by begin_weight_update_replicas.""" + target_replicas = self.replicas if replicas is None else replicas + results = await asyncio.gather( + *[replica.end_weight_update() for replica in target_replicas], + return_exceptions=True, + ) + + errors = [] + for replica, result in zip(target_replicas, results, strict=True): + try: + self._raise_for_weight_update_control_result("end_weight_update", replica, result) + except Exception as exc: + errors.append(str(exc)) + if errors: + raise RuntimeError("; ".join(errors)) + + @staticmethod + def _raise_for_weight_update_control_result( + endpoint: str, replica: RolloutReplica, result: dict[str, Any] | BaseException | None + ) -> None: + if isinstance(result, BaseException): + raise RuntimeError( + f"{endpoint} failed on rollout replica {getattr(replica, 'replica_rank', '?')}" + ) from result + if isinstance(result, dict) and result.get("success") is False: + message = result.get("message", "unknown error") + raise RuntimeError( + f"{endpoint} failed on rollout replica {getattr(replica, 'replica_rank', '?')}: {message}" + ) + + @auto_await + async def update_weight_version_replicas(self, weight_version: str): + results = await asyncio.gather( + *[replica.update_weight_version(weight_version, abort_all_requests=False) for replica in self.replicas], + return_exceptions=True, + ) + errors = [] + for replica, result in zip(self.replicas, results, strict=True): + try: + self._raise_for_weight_update_control_result("update_weight_version", replica, result) + except Exception as exc: + errors.append(str(exc)) + if errors: + raise RuntimeError("; ".join(errors)) + @auto_await async def update_weights(self, global_steps: int = None): """Update weights from actor worker group to rollout replicas. @@ -495,6 +577,13 @@ async def update_weights(self, global_steps: int = None): ray.get(self.actor_wg.update_weights(global_steps=global_steps, mode=self.backend)) return {} + # P2P: trainer RDMA push + rollout control-plane lifecycle. + if self.backend == "p2p": + await self._p2p_weight_checker.setup(self.replicas) + await self._update_weights_p2p(global_steps=global_steps) + await self._p2p_weight_checker.compare(self.replicas) + return {} + # 1. abort and save all unfinished requests for partial rollout await self.abort_replicas() @@ -511,11 +600,15 @@ async def update_weights(self, global_steps: int = None): # 4. build process group self.build_process_group(rollout) - # 5. update weights of all workers - results = ray.get( - actor_wg.update_weights(global_steps=global_steps, mode=self.backend) - + rollout.update_weights(global_steps=global_steps) - ) + # 5. update weights of all workers inside one replica-level SGLang session + weight_update_replicas = await self.begin_weight_update_replicas() + try: + results = ray.get( + actor_wg.update_weights(global_steps=global_steps, mode=self.backend) + + rollout.update_weights(global_steps=global_steps) + ) + finally: + await self.end_weight_update_replicas(weight_update_replicas) # The sender workers return the engine's per-sync metrics (empty for # backends that don't track any); merge and hand them to the trainer. sync_metrics: dict = {} @@ -537,6 +630,39 @@ async def update_weights(self, global_steps: int = None): return sync_metrics + @auto_await + async def _update_weights_p2p(self, global_steps: int | None = None) -> None: + """Run Miles-style P2P sync: one-time rollout metadata connect + trainer RDMA.""" + from verl.checkpoint_engine.p2p.transfer_utils import ( + collect_p2p_rollout_metadata, + resolve_rollout_model_path, + ) + + await self.abort_replicas() + await self.release_kv_cache_replicas() + + weight_update_replicas = await self.begin_weight_update_replicas() + try: + if not self._p2p_rollout_metadata_initialized: + model_path = resolve_rollout_model_path(self.replicas[0].model_config) + rollout_metadata = await collect_p2p_rollout_metadata( + self.replicas, + model_path=model_path, + engine_kwargs=self.config.engine_kwargs.get("p2p", {}), + ) + self.actor_wg.init_p2p_rollout_metadata(rollout_metadata) + self._p2p_rollout_metadata_initialized = True + + ray.get(self.actor_wg.update_weights(global_steps=global_steps, mode="p2p")) + + weight_version = str(global_steps) if global_steps is not None else "0" + await self.update_weight_version_replicas(weight_version) + finally: + await self.end_weight_update_replicas(weight_update_replicas) + + await self.resume_kv_cache_replicas() + await self.resume_generation_replicas() + async def split_weight_chunks( weights: Generator[tuple[str, torch.Tensor], None, None], bucket_size: int diff --git a/verl/checkpoint_engine/p2p/__init__.py b/verl/checkpoint_engine/p2p/__init__.py new file mode 100644 index 00000000000..5ad22f727a3 --- /dev/null +++ b/verl/checkpoint_engine/p2p/__init__.py @@ -0,0 +1,45 @@ +# Copyright 2025 Bytedance Ltd. and/or its affiliates +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from .checkpoint_engine import P2PCheckpointEngine +from .trainer_updater import P2PTrainerWeightUpdater +from .transfer_utils import ( + P2PRolloutTopology, + P2PTransferManager, + RemoteTransferPlan, + RemoteWeightInfo, + TransferTaskP2PMeta, + build_rollout_topology_from_replicas, + collect_p2p_rollout_metadata, +) +from .weight_checker import ( + P2PWeightChecker, + compare_weight_updates, + setup_weight_checker_snapshots, +) + +__all__ = [ + "P2PCheckpointEngine", + "P2PTrainerWeightUpdater", + "P2PRolloutTopology", + "P2PTransferManager", + "RemoteTransferPlan", + "RemoteWeightInfo", + "TransferTaskP2PMeta", + "build_rollout_topology_from_replicas", + "collect_p2p_rollout_metadata", + "P2PWeightChecker", + "setup_weight_checker_snapshots", + "compare_weight_updates", +] diff --git a/verl/checkpoint_engine/p2p/checkpoint_engine.py b/verl/checkpoint_engine/p2p/checkpoint_engine.py new file mode 100644 index 00000000000..15e88b457f2 --- /dev/null +++ b/verl/checkpoint_engine/p2p/checkpoint_engine.py @@ -0,0 +1,120 @@ +# Copyright 2025 Bytedance Ltd. and/or its affiliates +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from __future__ import annotations + +from typing import Any, AsyncGenerator, Generator + +import torch + +from verl.checkpoint_engine.base import CheckpointEngine, CheckpointEngineRegistry + +from .trainer_updater import P2PTrainerWeightUpdater +from .transfer_utils import P2PRolloutTopology + + +@CheckpointEngineRegistry.register("p2p") +class P2PCheckpointEngine(CheckpointEngine): + """Miles-style P2P checkpoint engine: trainer CPU replica + Mooncake RDMA writes. + + Rollout metadata is collected on the driver and passed via + :meth:`connect_rollout_metadata` before :meth:`send_weights`. + """ + + def __init__( + self, + bucket_size: int, + is_master: bool = False, + **_: Any, + ) -> None: + self._bucket_size = bucket_size + self.is_master = is_master + self._updater: P2PTrainerWeightUpdater | None = None + + @property + def is_source(self) -> bool: + """Whether this trainer rank pushes weights to rollout via P2P RDMA.""" + return self._updater.is_source + + def connect_rollout_metadata(self, rollout_metadata: dict[str, Any]) -> None: + topology = rollout_metadata.get("rollout_topology") + model_path = rollout_metadata.get("model_path") + if topology is None or model_path is None: + raise RuntimeError("P2PCheckpointEngine requires model_path and rollout_topology in metadata") + if self._updater is None: + topology = P2PRolloutTopology(**topology) + engine_kwargs = rollout_metadata.get("engine_kwargs", {}) + self._updater = P2PTrainerWeightUpdater( + model_path=model_path, + rollout_topology=topology, + num_workers=int(engine_kwargs.get("p2p_transfer_num_workers", 4)), + transfer_timeout=float(engine_kwargs.get("p2p_transfer_timeout", 30.0)), + is_master=self.is_master, + bucket_size_bytes=self._bucket_size, + ) + + self._updater.connect_rollout_metadata(rollout_metadata) + + def prepare(self) -> dict[str, Any]: + return {"backend": "p2p", "is_master": self.is_master} + + @classmethod + def build_topology( + cls, + trainer_world_size: int, + rollout_world_size: int, + metadata: list[dict[str, Any]], + ) -> tuple[dict[str, list[Any]], dict[str, list[Any]]]: + del metadata + trainer_kwargs = { + "rank": [-1] * trainer_world_size, + "world_size": [0] * trainer_world_size, + "metadata": [{}] * trainer_world_size, + } + rollout_kwargs = { + "rank": [-1] * rollout_world_size, + "world_size": [0] * rollout_world_size, + "metadata": [{}] * rollout_world_size, + } + return trainer_kwargs, rollout_kwargs + + def init_process_group(self, rank: int = -1, world_size: int = 0, metadata: dict[str, Any] | None = None): + del rank, world_size, metadata + + def finalize(self) -> None: + return None + + def release_for_checkpoint(self) -> None: + """Drop trainer-side P2P CPU buffers so Megatron save does not OOM/segfault.""" + if self._updater is not None: + self._updater.release_for_checkpoint() + + def restore_after_checkpoint(self) -> None: + """Rebuild trainer-side P2P CPU replica after checkpoint save.""" + if self._updater is not None: + self._updater.restore_after_checkpoint() + + async def send_weights( + self, + weights: Generator[tuple[str, torch.Tensor], None, None], + global_steps: int | None = None, + ) -> None: + if self._updater is None: + raise RuntimeError("P2P backend is not connected; manager must call connect_rollout_metadata first") + + self._updater.send_weights(weights) + + async def receive_weights(self, global_steps: int | None = None) -> AsyncGenerator[tuple[str, torch.Tensor], None]: + if False: + yield ("", torch.empty(0)) diff --git a/verl/checkpoint_engine/p2p/trainer_updater.py b/verl/checkpoint_engine/p2p/trainer_updater.py new file mode 100644 index 00000000000..26bebe5f111 --- /dev/null +++ b/verl/checkpoint_engine/p2p/trainer_updater.py @@ -0,0 +1,370 @@ +# Copyright 2025 Bytedance Ltd. and/or its affiliates +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Trainer-side P2P weight updater following the design of Miles ``UpdateWeightP2P``.""" + +from __future__ import annotations + +import gc +import logging +from collections.abc import Generator, Mapping +from typing import Any + +import torch + +from .transfer_utils import ( + P2PRolloutTopology, + P2PTransferManager, + RemoteTransferPlan, + RemoteWeightInfo, + TransferTaskP2PMeta, + create_server_args_from_dict, + create_transfer_engine, + iter_named_tensor_buckets, + register_cpu_memory, + unregister_cpu_memory, +) + +logger = logging.getLogger(__name__) + + +class P2PTrainerWeightUpdater: + """Miles-style trainer push: CPU pinned replica + Mooncake RDMA writes.""" + + def __init__( + self, + *, + model_path: str, + rollout_topology: P2PRolloutTopology, + num_workers: int = 4, + transfer_timeout: float = 30.0, + is_master: bool = False, + bucket_size_bytes: int = 2 << 30, + ) -> None: + self.model_path = model_path + self.is_master = is_master + self._bucket_size_bytes = bucket_size_bytes + self.transfer_plan = RemoteTransferPlan(rollout_topology) + self.transfer_manager = P2PTransferManager(num_workers=num_workers, transfer_timeout=transfer_timeout) + + self._connected = False + self._model_registered = False + self._tensor_update_pending: dict[str, int] = {} + self._staged_tensors: dict[str, list[tuple[str, torch.Tensor]]] = {} + self._shared_params_dict: dict[str, torch.Tensor] = {} + self._shared_param_mapper = None + self._weight_memory_registry: dict[str, tuple[int, int, int]] = {} + self._transfer_engine = None + self._transfer_engine_meta_list: list[tuple[torch.nn.Module, list[RemoteWeightInfo]]] = [] + self._rollout_metadata: dict[str, Any] | None = None + + @property + def is_source(self) -> bool: + return self.transfer_plan.is_source + + def connect_rollout_metadata(self, rollout_metadata: dict[str, Any]) -> None: + """Initialize transfer engine state from manager-collected rollout metadata.""" + self._rollout_metadata = rollout_metadata + remote_weight_infos_by_session_id = rollout_metadata["remote_weight_infos_by_session_id"] + targets_to_session_id = rollout_metadata["targets_to_session_id"] + session_id_to_server_args = rollout_metadata["session_id_to_server_args"] + if self._connected: + return + if not self.is_source: + self._connected = True + return + + targets = self.transfer_plan.plan_p2p() + targets_grouped_by_engine_rank: dict[int, list[TransferTaskP2PMeta]] = {} + for target in targets: + targets_grouped_by_engine_rank.setdefault(target.engine_rank, []).append(target) + + self._transfer_engine = create_transfer_engine() + first_engine_rank = True + for rank_targets in targets_grouped_by_engine_rank.values(): + first_target = rank_targets[0] + session_id = targets_to_session_id[(first_target.engine_ind, first_target.engine_rank)] + parallelism_config = self._import_parallelism_config(remote_weight_infos_by_session_id[session_id][1]) + server_args = create_server_args_from_dict(session_id_to_server_args[session_id]) + + model_replica = self._create_cpu_replica( + parallelism_config, + self.model_path, + server_args, + first_engine_rank=first_engine_rank, + ) + if first_engine_rank: + self._shared_params_dict = dict(model_replica.named_parameters()) + from sglang.srt.model_loader.parameter_mapper import ParameterMapper + + self._shared_param_mapper = ParameterMapper.from_model(model_replica) + first_engine_rank = False + + remote_infos = [] + for target in rank_targets: + session_id = targets_to_session_id[(target.engine_ind, target.engine_rank)] + weights_info = remote_weight_infos_by_session_id[session_id][0] + remote_infos.append(RemoteWeightInfo(session_id, weights_info)) + self._transfer_engine_meta_list.append((model_replica, remote_infos)) + + self._connected = True + if self.is_master: + logger.info( + "P2P trainer connected: gathered_dp_rank=%s targets=%s", + self.transfer_plan.gathered_dp_rank, + len(targets), + ) + + def release_for_checkpoint(self) -> None: + """Free P2P CPU replica and Mooncake registrations before Megatron checkpoint save.""" + if not self._connected: + return + + if self.is_source: + self.transfer_manager.wait_transfers() + unregister_cpu_memory(self._weight_memory_registry, self._transfer_engine) + self._transfer_engine_meta_list.clear() + self._shared_params_dict.clear() + self._shared_param_mapper = None + self._staged_tensors.clear() + self._tensor_update_pending.clear() + # Keep rollout metadata for restore_after_checkpoint (small, stable ptrs). + self._transfer_engine = None + + self._model_registered = False + self._connected = False + gc.collect() + if self.is_master: + print("[P2PTrainerWeightUpdater] released CPU replica for checkpoint save", flush=True) + + def restore_after_checkpoint(self) -> None: + """Recreate P2P CPU replica after checkpoint save.""" + if self._rollout_metadata is None or self._connected: + return + self.connect_rollout_metadata(self._rollout_metadata) + if self.is_master: + print("[P2PTrainerWeightUpdater] restored CPU replica after checkpoint save", flush=True) + + def send_weights(self, weights: Generator[tuple[str, torch.Tensor], None, None]) -> None: + """Stage HF tensors, load CPU replica, and issue RDMA writes.""" + if not self._connected: + raise RuntimeError("P2P trainer updater is not connected; call connect_rollout_metadata first") + + if not self.is_source: + raise RuntimeError( + "P2P send_weights must not be called on non-source trainer ranks; skip export/update_weights instead." + ) + + if not self._model_registered: + self._weight_memory_registry = register_cpu_memory(self._shared_params_dict, self._transfer_engine) + self._model_registered = True + + rdma_writes = 0 + num_tensors = 0 + total_bytes = 0 + for bucket in iter_named_tensor_buckets(weights, self._bucket_size_bytes): + num_tensors += len(bucket) + total_bytes += sum(tensor.numel() * tensor.element_size() for _, tensor in bucket) + transfer_ready_params, ready_hf_tensors = self._get_transfer_ready_params(bucket) + if transfer_ready_params and ready_hf_tensors: + rdma_writes += self._transfer_ready_bucket(transfer_ready_params, ready_hf_tensors) + + self.transfer_manager.wait_transfers() + assert not self._tensor_update_pending and not self._staged_tensors, ( + f"Pending staged tensors after P2P update: {self._tensor_update_pending}, {self._staged_tensors}" + ) + + if self.is_master: + logger.info( + "P2P send_weights: %s tensors, %.2f MiB, %s RDMA sessions", + num_tensors, + total_bytes / (1024 * 1024), + rdma_writes, + ) + + def _get_transfer_ready_params( + self, converted_named_tensors: list[tuple[str, torch.Tensor]] + ) -> tuple[list[str], list[tuple[str, torch.Tensor]]]: + transfer_ready_params: list[str] = [] + params_dict = self._shared_params_dict + assert self._shared_param_mapper is not None + + for name, tensor in converted_named_tensors: + mapped_result = self._shared_param_mapper.map(name) + mapped, num_shards, num_experts = ( + mapped_result.sglang_name, + mapped_result.num_shards, + mapped_result.num_local_experts, + ) + if mapped not in params_dict: + raise RuntimeError( + f"P2P mapped parameter {mapped!r} (from HF name {name!r}) " + "not found in shared CPU replica; skipping would leave rollout " + "weights stale" + ) + + if num_experts is not None and num_experts > 0: + total_expected = num_experts * num_shards + else: + total_expected = num_shards + + self._staged_tensors.setdefault(mapped, []).append((name, tensor)) + + if total_expected == 1: + transfer_ready_params.append(mapped) + else: + if mapped not in self._tensor_update_pending: + self._tensor_update_pending[mapped] = total_expected - 1 + else: + self._tensor_update_pending[mapped] -= 1 + if self._tensor_update_pending[mapped] == 0: + transfer_ready_params.append(mapped) + + ready_hf_tensors: list[tuple[str, torch.Tensor]] = [] + for param_name in transfer_ready_params: + staged = self._staged_tensors.pop(param_name, []) + ready_hf_tensors.extend(staged) + self._tensor_update_pending.pop(param_name, None) + + return transfer_ready_params, ready_hf_tensors + + def _transfer_ready_bucket( + self, transfer_ready_params: list[str], ready_hf_tensors: list[tuple[str, torch.Tensor]] + ) -> int: + """Load one ready bucket into the CPU replica and issue RDMA writes.""" + rdma_writes = 0 + last_idx = len(self._transfer_engine_meta_list) - 1 + for i, (model_replica, remote_weight_infos) in enumerate(self._transfer_engine_meta_list): + model_replica.load_weights(ready_hf_tensors) + is_last = i == last_idx + if is_last: + for remote_session in remote_weight_infos: + self.transfer_manager.submit( + self._do_p2p_write_one_session, + remote_session, + transfer_ready_params, + ) + rdma_writes += 1 + else: + futures = [ + self.transfer_manager.submit_returning_future( + self._do_p2p_write_one_session, + remote_session, + transfer_ready_params, + ) + for remote_session in remote_weight_infos + ] + self.transfer_manager.wait_futures(futures) + rdma_writes += len(futures) + return rdma_writes + + def _do_p2p_write_one_session(self, remote_session: RemoteWeightInfo, names: list[str]) -> None: + source_ptrs: list[int] = [] + source_lens: list[int] = [] + valid_names: list[str] = [] + + for name in names: + cpu_reg = self._weight_memory_registry.get(name) + if not cpu_reg: + raise RuntimeError( + f"P2P source registration missing for weight {name}; " + "cannot transfer without corrupting rollout weights" + ) + data_ptr, numel, ele_size = cpu_reg + source_ptrs.append(data_ptr) + source_lens.append(numel * ele_size) + valid_names.append(name) + + if not source_ptrs: + return + + session_id = remote_session.session_id + target_ptrs = [ + remote_session.weights_info[name][0] for name in valid_names if name in remote_session.weights_info + ] + if len(target_ptrs) != len(source_ptrs): + raise RuntimeError( + f"P2P pointer count mismatch for session {session_id}: " + f"source={len(source_ptrs)} target={len(target_ptrs)}" + ) + + ret = self._transfer_engine.batch_transfer_sync_write(session_id, source_ptrs, target_ptrs, source_lens) + if ret < 0: + raise RuntimeError(f"P2P transfer failed for session {session_id}, error: {ret}") + + @staticmethod + def _import_parallelism_config(parallelism_info: Mapping[str, Any]): + from sglang.srt.distributed.parallel_state import RankParallelismConfig + + return RankParallelismConfig.from_dict(dict(parallelism_info)) + + def _create_cpu_replica( + self, + parallelism_config, + model_path: str, + server_args, + *, + first_engine_rank: bool = False, + ) -> torch.nn.Module: + from sglang.srt import server_args as server_args_module + from sglang.srt.configs.device_config import DeviceConfig + from sglang.srt.configs.load_config import LoadConfig + from sglang.srt.configs.model_config import ModelConfig + from sglang.srt.distributed.parallel_state import ParallelismContext + from sglang.srt.layers.moe import initialize_moe_config + from sglang.srt.layers.quantization.fp4_utils import initialize_fp4_gemm_config + from sglang.srt.layers.quantization.fp8_utils import initialize_fp8_gemm_config + from sglang.srt.model_loader import get_model + from sglang.srt.model_loader import loader as model_loader_module + + load_config = LoadConfig( + load_format="dummy", + model_loader_extra_config=None, + rl_quant_profile=server_args.rl_quant_profile, + ) + server_args_module._global_server_args = server_args + initialize_moe_config(server_args) + initialize_fp8_gemm_config(server_args) + initialize_fp4_gemm_config(server_args) + + # Monkey-patch loader-level post_load_weights to no-op before get_model. + # Miles sglang uses `_post_load_weights`; sglang-pretrain exposes `post_load_weights`. + post_load_hook_name = ( + "_post_load_weights" if hasattr(model_loader_module, "_post_load_weights") else "post_load_weights" + ) + original_post_load_weights = getattr(model_loader_module, post_load_hook_name) + setattr(model_loader_module, post_load_hook_name, lambda *args, **kwargs: None) + try: + with ParallelismContext(parallelism_config): + model = get_model( + model_config=ModelConfig(model_path), + load_config=load_config, + device_config=DeviceConfig(device="cpu"), + ) + finally: + setattr(model_loader_module, post_load_hook_name, original_post_load_weights) + + if hasattr(model, "post_load_weights"): + model.post_load_weights = lambda *args, **kwargs: None + + if first_engine_rank: + for param in model.parameters(): + param.data = param.data.pin_memory() + else: + for name, param in model.named_parameters(): + if name not in self._shared_params_dict: + raise RuntimeError(f"P2P shared buffer missing parameter {name}") + param.data = self._shared_params_dict[name] + + return model diff --git a/verl/checkpoint_engine/p2p/transfer_utils.py b/verl/checkpoint_engine/p2p/transfer_utils.py new file mode 100644 index 00000000000..28cd124cf4f --- /dev/null +++ b/verl/checkpoint_engine/p2p/transfer_utils.py @@ -0,0 +1,406 @@ +# Copyright 2025 Bytedance Ltd. and/or its affiliates +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""P2P transfer planning utilities, interoperable with the Miles P2P protocol. + +Transfer planning follows Miles heuristics; implementation is built around verl +``RolloutReplica`` handles instead of Miles rollout Ray actors. +""" + +from __future__ import annotations + +import dataclasses +import logging +import os +from collections import defaultdict +from collections.abc import Callable, Generator, Iterator, Sequence +from concurrent.futures import Future, ThreadPoolExecutor +from typing import TYPE_CHECKING, Any, Protocol + +import ray +import torch +import torch.distributed as dist + +if TYPE_CHECKING: + from omegaconf import DictConfig + +logger = logging.getLogger(__name__) + + +@dataclasses.dataclass(frozen=True) +class P2PRolloutTopology: + """Rollout-side topology for P2P transfer planning.""" + + engine_count: int + gpus_per_engine: int + pipeline_parallel_size: int = 1 + + @property + def total_gpus(self) -> int: + return self.engine_count * self.gpus_per_engine + + +@dataclasses.dataclass +class TransferTaskP2PMeta: + """Specifies a rollout engine rank to connect to.""" + + engine_ind: int + engine_rank: int + + +@dataclasses.dataclass +class RemoteWeightInfo: + """Remote weight registration info for one rollout session.""" + + session_id: str + weights_info: dict[str, tuple[int, int, int]] + + +class P2PRolloutReplica(Protocol): + async def get_remote_instance_transfer_engine_info(self, rank: int) -> Any: ... + + async def get_parallelism_info(self, rank: int) -> Any: ... + + async def get_server_info(self) -> dict[str, Any]: ... + + +def resolve_trainer_parallelism() -> tuple[int, int, int, int]: + """Return ``(pp_rank, pp_size, gathered_dp_rank, gathered_dp_size)``.""" + try: + from megatron.core import parallel_state as mpu + except ImportError as exc: + raise RuntimeError("Megatron parallel_state is required for P2P transfer planning") from exc + + if not dist.is_initialized(): + raise RuntimeError("torch.distributed must be initialized for P2P transfer planning") + + pp_rank = mpu.get_pipeline_model_parallel_rank() + pp_size = mpu.get_pipeline_model_parallel_world_size() + world_size = dist.get_world_size() + gathered_dp_size = world_size // pp_size + + pp_group = mpu.get_pipeline_model_parallel_group() + my_column_id = min(dist.get_process_group_ranks(pp_group)) + all_column_ids = [None] * world_size + dist.all_gather_object(all_column_ids, my_column_id) + sorted_columns = sorted(set(all_column_ids)) + gathered_dp_rank = sorted_columns.index(my_column_id) + return pp_rank, pp_size, gathered_dp_rank, gathered_dp_size + + +class RemoteTransferPlan: + """Plans trainer -> rollout P2P connections (Miles-compatible heuristics).""" + + def __init__(self, rollout_topology: P2PRolloutTopology) -> None: + if rollout_topology.pipeline_parallel_size != 1: + raise NotImplementedError("Rollout pipeline parallelism is not supported for P2P yet.") + + _, _, self._gathered_dp_rank, self._gathered_dp_size = resolve_trainer_parallelism() + self._rollout_num_gpu_per_engine = rollout_topology.gpus_per_engine + self._rollout_engine_count = rollout_topology.engine_count + self._rollout_num_gpus = rollout_topology.total_gpus + + @property + def gathered_dp_rank(self) -> int: + return self._gathered_dp_rank + + @property + def is_source(self) -> bool: + return self._gathered_dp_rank < self._rollout_num_gpus + + def plan_p2p(self) -> list[TransferTaskP2PMeta]: + all_targets = [ + (engine_idx, engine_rank) + for engine_idx in range(self._rollout_engine_count) + for engine_rank in range(self._rollout_num_gpu_per_engine) + ] + assignments: dict[int, dict[int, list[int]]] = defaultdict(lambda: defaultdict(list)) + + p2p_count = 0 + for source_rank, (_, target) in zip(range(self._gathered_dp_size), enumerate(all_targets), strict=False): + p2p_count += 1 + engine_idx, engine_rank = target + assignments[source_rank][engine_rank].append(engine_idx) + + def count_engine_index_assignments(engine_rank: int) -> list[int]: + return [len(assignments[source][engine_rank]) for source in range(self._gathered_dp_size)] + + cur_source_index = 0 + if p2p_count < len(all_targets): + for target in all_targets[p2p_count:]: + engine_idx, engine_rank = target + counted = count_engine_index_assignments(engine_rank) + if max(counted) > 0: + _, select_source = min((val, idx) for (idx, val) in enumerate(counted) if val > 0) + else: + select_source = cur_source_index % self._gathered_dp_size + cur_source_index += 1 + assignments[select_source][engine_rank].append(engine_idx) + + transfer_tasks: list[TransferTaskP2PMeta] = [] + for engine_rank, engine_indices in assignments[self._gathered_dp_rank].items(): + for engine_ind in engine_indices: + transfer_tasks.append( + TransferTaskP2PMeta( + engine_ind=engine_ind, + engine_rank=engine_rank, + ) + ) + return transfer_tasks + + +class P2PTransferManager: + """Async thread-pool executor for RDMA writes.""" + + def __init__(self, num_workers: int = 8, transfer_timeout: float = 30.0) -> None: + self.num_workers = num_workers + self.transfer_timeout = transfer_timeout + self.executor: ThreadPoolExecutor | None = None + self.transfer_futures: list[Future] = [] + + def ensure_started(self) -> None: + if self.executor is None: + self.executor = ThreadPoolExecutor(max_workers=self.num_workers) + + def submit(self, fn: Callable, *args) -> None: + self.ensure_started() + self.transfer_futures.append(self.executor.submit(fn, *args)) + + def submit_returning_future(self, fn: Callable, *args) -> Future: + self.ensure_started() + future = self.executor.submit(fn, *args) + self.transfer_futures.append(future) + return future + + def wait_futures(self, futures: Sequence[Future]) -> None: + for future in futures: + future.result(timeout=self.transfer_timeout) + + def wait_transfers(self) -> None: + for future in self.transfer_futures: + try: + future.result(timeout=self.transfer_timeout) + except Exception as exc: + logger.error("P2P transfer future failed: %s", exc) + raise + self.transfer_futures.clear() + + +def _json_safe_value(value: Any) -> Any: + """Coerce rollout metadata values to Ray-pickle-safe primitives.""" + if value is None or isinstance(value, bool | int | float | str): + return value + if isinstance(value, list | tuple): + return [_json_safe_value(item) for item in value] + if isinstance(value, dict): + return {str(key): _json_safe_value(item) for key, item in value.items()} + return None + + +def filter_server_args_dict(data_dict: dict[str, Any]) -> dict[str, Any]: + """Keep only ServerArgs fields with JSON/pickle-safe values for cross-node transfer.""" + import dataclasses as dc + + from sglang.srt.server_args import ServerArgs + + valid_fields = {f.name for f in dc.fields(ServerArgs)} + return {key: _json_safe_value(value) for key, value in data_dict.items() if key in valid_fields} + + +def create_server_args_from_dict(data_dict: dict[str, Any]): + from sglang.srt.server_args import ServerArgs + + return ServerArgs(**data_dict) + + +def register_cpu_memory(params_dict: dict[str, torch.Tensor], transfer_engine) -> dict[str, tuple[int, int, int]]: + weight_dict: dict[str, tuple[int, int, int]] = {} + for name, cpu_tensor in params_dict.items(): + addr = cpu_tensor.data_ptr() + size = cpu_tensor.numel() * cpu_tensor.element_size() + ret = transfer_engine.register_memory(addr, size) + if ret != 0: + raise RuntimeError(f"register CPU memory failed for weight {name}, error: {ret}") + weight_dict[name] = (addr, cpu_tensor.numel(), cpu_tensor.element_size()) + return weight_dict + + +def unregister_cpu_memory(weight_memory_registry: dict[str, tuple[int, int, int]], transfer_engine) -> None: + """Unregister Mooncake-registered CPU buffers before freeing the P2P replica.""" + if not weight_memory_registry: + return + + for name, (addr, _numel, _ele_size) in list(weight_memory_registry.items()): + ret = transfer_engine.unregister_memory(addr) + if ret != 0: + logger.warning("unregister CPU memory failed for weight %s, error: %s", name, ret) + weight_memory_registry.clear() + + +def resolve_mooncake_transfer_engine_settings() -> tuple[str, str]: + """Return Mooncake protocol and IB device filter (aligned with SGLang envs).""" + protocol = os.environ.get("MOONCAKE_PROTOCOL", "rdma") + if os.environ.get("MC_FORCE_TCP") == "1": + return protocol, "" + return protocol, os.environ.get("MOONCAKE_DEVICE", "") + + +def create_transfer_engine(): + from mooncake.engine import TransferEngine + + transfer_engine = TransferEngine() + local_ip = ray._private.services.get_node_ip_address() + protocol, device_name = resolve_mooncake_transfer_engine_settings() + ret = transfer_engine.initialize(local_ip, "P2PHANDSHAKE", protocol, device_name) + if ret != 0: + raise RuntimeError( + f"TransferEngine initialize failed with ret={ret} (protocol={protocol!r}, device_name={device_name!r})" + ) + return transfer_engine + + +async def query_remote_weight_infos( + replicas: Sequence[P2PRolloutReplica], + targets: Sequence[TransferTaskP2PMeta], +) -> tuple[dict[str, tuple[Any, Any]], dict[tuple[int, int], str], dict[str, Any]]: + """Query rollout replicas for transfer metadata (async, Verl replica API).""" + remote_weight_infos_by_session_id: dict[str, tuple[Any, Any]] = {} + targets_to_session_id: dict[tuple[int, int], str] = {} + session_id_to_server_args: dict[str, Any] = {} + targets_to_query = {(target.engine_ind, target.engine_rank) for target in targets} + + for engine_ind, engine_rank in targets_to_query: + replica = replicas[engine_ind] + transfer_info = await replica.get_remote_instance_transfer_engine_info(engine_rank) + parallelism_info = await replica.get_parallelism_info(engine_rank) + if transfer_info is None: + raise RuntimeError(f"missing transfer engine info for replica={engine_ind} rank={engine_rank}") + if parallelism_info is None: + raise RuntimeError(f"missing parallelism config for replica={engine_ind} rank={engine_rank}") + + session_id, weights_info = transfer_info + assert session_id is not None, f"Failed to get session id from rollout replica {engine_ind} rank {engine_rank}" + server_info = await replica.get_server_info() + # Store plain dicts only: ServerArgs objects may pickle references to + # transformers_modules (trust_remote_code) and fail on trainer workers. + session_id_to_server_args[session_id] = filter_server_args_dict(server_info) + remote_weight_infos_by_session_id[session_id] = (weights_info, parallelism_info) + targets_to_session_id[(engine_ind, engine_rank)] = session_id + + return remote_weight_infos_by_session_id, targets_to_session_id, session_id_to_server_args + + +def build_rollout_topology_from_replicas(replicas: Sequence[Any]) -> P2PRolloutTopology: + if not replicas: + raise ValueError("At least one rollout replica is required for P2P") + gpus_per_engine = replicas[0].world_size + pp_size = getattr(replicas[0].config, "pipeline_model_parallel_size", 1) + for replica in replicas[1:]: + if replica.world_size != gpus_per_engine: + raise NotImplementedError("Heterogeneous replica world sizes are not supported for P2P yet") + return P2PRolloutTopology( + engine_count=len(replicas), + gpus_per_engine=gpus_per_engine, + pipeline_parallel_size=pp_size, + ) + + +def serialize_p2p_rollout_metadata( + *, + model_path: str, + rollout_topology: P2PRolloutTopology, + engine_kwargs: dict[str, Any], + remote_weight_infos_by_session_id: dict[str, tuple[Any, Any]], + targets_to_session_id: dict[tuple[int, int], str], + session_id_to_server_args: dict[str, dict[str, Any]], +) -> dict[str, Any]: + return { + "model_path": model_path, + "rollout_topology": dataclasses.asdict(rollout_topology), + "engine_kwargs": engine_kwargs, + "remote_weight_infos_by_session_id": remote_weight_infos_by_session_id, + "targets_to_session_id": targets_to_session_id, + "session_id_to_server_args": session_id_to_server_args, + } + + +def iter_named_tensor_buckets( + weights: Iterator[tuple[str, torch.Tensor]], + bucket_bytes: int, + *, + clone_tensors: bool = False, +) -> Generator[list[tuple[str, torch.Tensor]], None, None]: + """Group HF weight tensors into byte-bounded buckets (Miles-style export batching).""" + if bucket_bytes <= 0: + raise ValueError(f"bucket_bytes must be greater than 0, got {bucket_bytes}") + + current_bucket: list[tuple[str, torch.Tensor]] = [] + current_size = 0 + for name, tensor in weights: + tensor_size = tensor.element_size() * tensor.numel() + if current_size + tensor_size > bucket_bytes: + if current_bucket: + yield current_bucket + stored = tensor.clone() if clone_tensors else tensor + current_bucket = [(name, stored)] + current_size = tensor_size + else: + stored = tensor.clone() if clone_tensors else tensor + current_bucket.append((name, stored)) + current_size += tensor_size + + if current_bucket: + yield current_bucket + + +def resolve_rollout_model_path(model_config: DictConfig) -> str: + """Resolve the rollout HF model path from the driver-side replica config.""" + from omegaconf import DictConfig, OmegaConf + + from verl.utils.fs import copy_to_local + + if not isinstance(model_config, DictConfig): + raise TypeError(f"resolve_rollout_model_path expects DictConfig, got {type(model_config).__name__}") + + path = OmegaConf.select(model_config, "path", default=None) + if not path: + raise RuntimeError("P2P weight sync requires actor_rollout_ref.model.path") + + use_shm = OmegaConf.select(model_config, "use_shm", default=False) + return copy_to_local(path, use_shm=use_shm) + + +async def collect_p2p_rollout_metadata( + replicas: Sequence[P2PRolloutReplica], + *, + model_path: str, + engine_kwargs: dict[str, Any] | None = None, +) -> dict[str, Any]: + topology = build_rollout_topology_from_replicas(replicas) + all_targets = [ + TransferTaskP2PMeta(engine_ind=engine_ind, engine_rank=engine_rank) + for engine_ind, replica in enumerate(replicas) + for engine_rank in range(replica.world_size) + ] + remote_infos, targets_to_session_id, session_id_to_server_args = await query_remote_weight_infos( + replicas, all_targets + ) + return serialize_p2p_rollout_metadata( + model_path=model_path, + rollout_topology=topology, + engine_kwargs=engine_kwargs or {}, + remote_weight_infos_by_session_id=remote_infos, + targets_to_session_id=targets_to_session_id, + session_id_to_server_args=session_id_to_server_args, + ) diff --git a/verl/checkpoint_engine/p2p/weight_checker.py b/verl/checkpoint_engine/p2p/weight_checker.py new file mode 100644 index 00000000000..e2dd274ca32 --- /dev/null +++ b/verl/checkpoint_engine/p2p/weight_checker.py @@ -0,0 +1,128 @@ +# Copyright 2025 Bytedance Ltd. and/or its affiliates +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""SGLang ``/weights_checker`` helpers (Miles ``--check-weight-update-equal`` parity).""" + +from __future__ import annotations + +import asyncio +import logging +from typing import Any + +logger = logging.getLogger(__name__) + + +class P2PWeightChecker: + """Stateful P2P corruption-test controller (Miles ``--check-weight-update-equal``).""" + + def __init__( + self, + *, + enabled: bool = False, + allow_quant_error: bool = False, + ) -> None: + self._enabled = enabled + self._allow_quant_error = allow_quant_error + self._setup_done = False + self._compared = False + + async def setup(self, replicas: list[Any]) -> None: + """Snapshot rollout weights and corrupt them before the first P2P sync.""" + if not self._enabled or self._setup_done: + return + await setup_weight_checker_snapshots(replicas) + self._setup_done = True + + async def compare(self, replicas: list[Any]) -> None: + """Compare rollout weights against the startup snapshot (once, after first sync).""" + if not self._enabled or self._compared or not self._setup_done: + return + await compare_weight_updates(replicas, allow_quant_error=self._allow_quant_error) + self._compared = True + + +def _replica_label(replica: Any) -> str: + return f"replica={getattr(replica, 'replica_rank', '?')}" + + +def _raise_for_check_weights_result(replica: Any, result: Any, *, action: str) -> None: + if isinstance(result, BaseException): + raise RuntimeError(f"weights_checker action={action!r} failed on rollout {_replica_label(replica)}") from result + if not isinstance(result, dict): + raise RuntimeError( + f"weights_checker action={action!r} returned unexpected result on " + f"rollout {_replica_label(replica)}: {result!r}" + ) + if result.get("success") is False: + message = result.get("message", "unknown error") + raise RuntimeError(f"weights_checker action={action!r} failed on rollout {_replica_label(replica)}: {message}") + + +async def check_weights_on_replicas( + replicas: list[Any], + *, + action: str, + allow_quant_error: bool = False, +) -> list[dict[str, Any]]: + """Fan out ``weights_checker`` to every rollout replica (SGLang only).""" + if not replicas: + return [] + + missing = [r for r in replicas if not hasattr(r, "check_weights")] + if missing: + raise RuntimeError( + "check_weight_update_equal requires SGLang rollout replicas with check_weights(); " + f"missing on {len(missing)} replica(s)" + ) + + results = await asyncio.gather( + *[replica.check_weights(action=action, allow_quant_error=allow_quant_error) for replica in replicas], + return_exceptions=True, + ) + outputs: list[dict[str, Any]] = [] + for replica, result in zip(replicas, results, strict=True): + _raise_for_check_weights_result(replica, result, action=action) + outputs.append(result) + return outputs + + +async def setup_weight_checker_snapshots(replicas: list[Any]) -> None: + """Miles startup sequence: snapshot reference weights, then corrupt with random.""" + logger.info("Weight checker: snapshot reference tensors on %s rollout replica(s)", len(replicas)) + await check_weights_on_replicas(replicas, action="snapshot") + logger.info("Weight checker: reset rollout tensors to random values") + await check_weights_on_replicas(replicas, action="reset_tensors") + print( + f"[WeightChecker] snapshot + reset_tensors completed on {len(replicas)} rollout replica(s)", + flush=True, + ) + + +async def compare_weight_updates( + replicas: list[Any], + *, + allow_quant_error: bool = False, +) -> list[dict[str, Any]]: + """Verify post-sync rollout weights match the snapshot taken at setup.""" + logger.info("Weight checker: compare rollout weights (allow_quant_error=%s)", allow_quant_error) + results = await check_weights_on_replicas( + replicas, + action="compare", + allow_quant_error=allow_quant_error, + ) + print( + f"[WeightChecker] compare PASSED on {len(replicas)} rollout replica(s) (allow_quant_error={allow_quant_error})", + flush=True, + ) + return results diff --git a/verl/experimental/fully_async_policy/detach_utils.py b/verl/experimental/fully_async_policy/detach_utils.py index f96ad910bb7..687595c21a5 100644 --- a/verl/experimental/fully_async_policy/detach_utils.py +++ b/verl/experimental/fully_async_policy/detach_utils.py @@ -150,7 +150,14 @@ def assemble_batch_from_rollout_samples( param_version_start = final_batch.non_tensor_batch["min_global_steps"] param_version_end = final_batch.non_tensor_batch["max_global_steps"] - param_version_diff = [abs(a - b) for a, b in zip(param_version_end, param_version_start, strict=False)] + + def _normalize_param_version(value): + return 0 if value is None else value + + param_version_diff = [ + abs(_normalize_param_version(a) - _normalize_param_version(b)) + for a, b in zip(param_version_end, param_version_start, strict=False) + ] num_diff0 = param_version_diff.count(0) partial_stats = { "fully_async/partial/total_partial_num": len(param_version_diff) - num_diff0, @@ -159,10 +166,11 @@ def assemble_batch_from_rollout_samples( } # add meta_info trajectory_param_versions = final_batch.non_tensor_batch["max_global_steps"] + normalized_trajectory_versions = [_normalize_param_version(v) for v in trajectory_param_versions] final_batch.meta_info.update( { - "param_version_diversity": len(set(trajectory_param_versions)), + "param_version_diversity": len(set(normalized_trajectory_versions)), "trajectory_param_versions": trajectory_param_versions, **processing_time_stats, **rollout_status, diff --git a/verl/experimental/fully_async_policy/fully_async_trainer.py b/verl/experimental/fully_async_policy/fully_async_trainer.py index 20ac0330ba9..816a4f4450e 100644 --- a/verl/experimental/fully_async_policy/fully_async_trainer.py +++ b/verl/experimental/fully_async_policy/fully_async_trainer.py @@ -762,7 +762,9 @@ def _collect_metrics_from_samples(self, batch, metrics): """ if hasattr(batch, "meta_info") and batch.meta_info: trajectory_param_versions = batch.meta_info["trajectory_param_versions"] - stale_traj_count = sum(1 for v in trajectory_param_versions if self.current_param_version - v >= 1) + stale_traj_count = sum( + 1 for v in trajectory_param_versions if v is not None and self.current_param_version - v >= 1 + ) self.stale_trajectory_processed += stale_traj_count metrics.update( { diff --git a/verl/trainer/config/_generated_ppo_megatron_trainer.yaml b/verl/trainer/config/_generated_ppo_megatron_trainer.yaml index 533dae11873..40aac342b05 100644 --- a/verl/trainer/config/_generated_ppo_megatron_trainer.yaml +++ b/verl/trainer/config/_generated_ppo_megatron_trainer.yaml @@ -355,6 +355,8 @@ actor_rollout_ref: backend: naive update_weights_bucket_megabytes: 2048 engine_kwargs: {} + check_weight_update_equal: false + check_weight_update_allow_quant_error: false custom_backend_module: null trace: _target_: verl.workers.config.TraceConfig diff --git a/verl/trainer/config/_generated_ppo_torchtitan_trainer.yaml b/verl/trainer/config/_generated_ppo_torchtitan_trainer.yaml index 12ff012683d..3381eca9deb 100644 --- a/verl/trainer/config/_generated_ppo_torchtitan_trainer.yaml +++ b/verl/trainer/config/_generated_ppo_torchtitan_trainer.yaml @@ -313,6 +313,8 @@ actor_rollout_ref: backend: naive update_weights_bucket_megabytes: 2048 engine_kwargs: {} + check_weight_update_equal: false + check_weight_update_allow_quant_error: false custom_backend_module: null trace: _target_: verl.workers.config.TraceConfig diff --git a/verl/trainer/config/_generated_ppo_trainer.yaml b/verl/trainer/config/_generated_ppo_trainer.yaml index 0660ffc1b1a..912c32e81f4 100644 --- a/verl/trainer/config/_generated_ppo_trainer.yaml +++ b/verl/trainer/config/_generated_ppo_trainer.yaml @@ -318,6 +318,8 @@ actor_rollout_ref: backend: naive update_weights_bucket_megabytes: 2048 engine_kwargs: {} + check_weight_update_equal: false + check_weight_update_allow_quant_error: false custom_backend_module: null trace: _target_: verl.workers.config.TraceConfig diff --git a/verl/trainer/config/_generated_ppo_veomni_trainer.yaml b/verl/trainer/config/_generated_ppo_veomni_trainer.yaml index d7aaf55db49..ecd88079139 100644 --- a/verl/trainer/config/_generated_ppo_veomni_trainer.yaml +++ b/verl/trainer/config/_generated_ppo_veomni_trainer.yaml @@ -317,6 +317,8 @@ actor_rollout_ref: backend: naive update_weights_bucket_megabytes: 2048 engine_kwargs: {} + check_weight_update_equal: false + check_weight_update_allow_quant_error: false custom_backend_module: null trace: _target_: verl.workers.config.TraceConfig diff --git a/verl/trainer/config/rollout/rollout.yaml b/verl/trainer/config/rollout/rollout.yaml index e7373ef780e..ed642835111 100644 --- a/verl/trainer/config/rollout/rollout.yaml +++ b/verl/trainer/config/rollout/rollout.yaml @@ -274,7 +274,7 @@ checkpoint_engine: # Target class for checkpoint engine config _target_: verl.workers.config.CheckpointEngineConfig - # Backend for checkpoint engine: naive, nccl, nixl, hccl + # Backend for checkpoint engine: naive, nccl, nixl, hccl, p2p backend: naive # Specifies the tensor bucket size (in megabytes) for batch weight updates during rollout operations. @@ -294,6 +294,11 @@ checkpoint_engine: # Additional keyword arguments to pass to the checkpoint engine constructor engine_kwargs: {} + # Miles --check-weight-update-equal: corruption test via SGLang /weights_checker. + # Snapshots rollout weights at first sync, resets to random, then compares after update. + check_weight_update_equal: false + check_weight_update_allow_quant_error: false + # If set, this Python module is imported on every worker process before the # backend is instantiated, allowing custom backends to register themselves # in CheckpointEngineRegistry. diff --git a/verl/utils/megatron_pp_local_export.py b/verl/utils/megatron_pp_local_export.py new file mode 100644 index 00000000000..4ea3183e7ec --- /dev/null +++ b/verl/utils/megatron_pp_local_export.py @@ -0,0 +1,125 @@ +# Copyright 2025 Bytedance Ltd. and/or its affiliates +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""PP-local HF weight export for disaggregated P2P sync. + +``bridge.export_hf_weights`` broadcasts every parameter across PP ranks and +yields the full model on each rank. For P2P, each PP stage should export and +send only the layers it owns (after TP/EP gather inside bridge mappings). + +This module reuses megatron-bridge conversion tasks/mappings but skips PP +broadcast so non-owning PP ranks never materialize foreign layers. +""" + +from __future__ import annotations + +import logging +from collections.abc import Generator, Iterable +from contextlib import contextmanager +from typing import Any + +import torch + +from verl.utils.megatron_utils import unwrap_model + +logger = logging.getLogger(__name__) + + +def filter_local_conversion_tasks(conversion_tasks: Iterable[Any]) -> list[Any]: + """Keep bridge tasks whose parameter tensor lives on this rank.""" + local_tasks = [] + for task in conversion_tasks: + if task is None: + continue + if task.megatron_module is None or task.param_weight is None: + continue + local_tasks.append(task) + return local_tasks + + +@contextmanager +def disable_megatron_bridge_pp_broadcast(): + """No-op PP broadcast helpers so export stays on the owning PP stage only.""" + try: + from megatron.bridge.models.conversion.param_mapping import MegatronParamMapping + except ImportError: + yield + return + + original_tensor_broadcast = MegatronParamMapping.broadcast_from_pp_rank + original_obj_broadcast = MegatronParamMapping.broadcast_obj_from_pp_rank + + def _tensor_without_pp_broadcast(self, tensor, cache_key=None): + return tensor + + def _obj_without_pp_broadcast(self, obj, cache_key=None): + return obj + + MegatronParamMapping.broadcast_from_pp_rank = _tensor_without_pp_broadcast + MegatronParamMapping.broadcast_obj_from_pp_rank = _obj_without_pp_broadcast + try: + yield + finally: + MegatronParamMapping.broadcast_from_pp_rank = original_tensor_broadcast + MegatronParamMapping.broadcast_obj_from_pp_rank = original_obj_broadcast + + +def _yield_hf_pairs(weights: Iterable[Any]) -> Generator[tuple[str, torch.Tensor], None, None]: + for item in weights: + if isinstance(item, tuple): + if len(item) >= 2: + yield item[0], item[1] + continue + hf_name = getattr(item, "param_name", None) or getattr(item, "hf_param_name", None) + tensor = getattr(item, "weight", None) + if hf_name is not None and tensor is not None: + yield hf_name, tensor + + +@torch.no_grad() +def export_pp_local_hf_weights( + bridge: Any, + megatron_model: torch.nn.Module | list[torch.nn.Module], + *, + merge_adapter_weights: bool = True, + show_progress: bool = False, +) -> Generator[tuple[str, torch.Tensor], None, None]: + """Export HF weights owned by this PP stage (TP/EP gather via bridge, no PP broadcast). + + All ranks in the PP group must still call ``build_conversion_tasks`` (it + all-gathers global param names). Only local tasks are converted and yielded. + """ + model_list = unwrap_model(megatron_model) + if not isinstance(model_list, list): + model_list = [model_list] + + model_bridge = bridge._model_bridge + hf_pretrained = bridge.hf_pretrained + + conversion_tasks = bridge.get_conversion_tasks(model_list) + local_tasks = filter_local_conversion_tasks(conversion_tasks) + if not local_tasks: + logger.warning("PP-local export produced zero local conversion tasks on this rank.") + return + + with disable_megatron_bridge_pp_broadcast(): + streamed = model_bridge.stream_weights_megatron_to_hf( + model_list, + hf_pretrained, + cpu=False, + show_progress=show_progress, + conversion_tasks=local_tasks, + merge_adapter_weights=merge_adapter_weights, + ) + yield from _yield_hf_pairs(streamed) diff --git a/verl/workers/config/rollout.py b/verl/workers/config/rollout.py index 19cf7f5e322..b561ec791da 100644 --- a/verl/workers/config/rollout.py +++ b/verl/workers/config/rollout.py @@ -129,12 +129,18 @@ class CheckpointEngineConfig(BaseConfig): _mutable_fields = {"backend"} - # Backend for checkpoint engine: naive, nccl, nixl, hccl + # Backend for checkpoint engine: naive, nccl, nixl, hccl, p2p backend: Optional[str] = "naive" # Bucket size in MB to transfer multiple weights at one time update_weights_bucket_megabytes: int = 2048 - # Additional keyword arguments for checkpoint engine + # Additional keyword arguments for checkpoint engine. + # For backend ``p2p``: ``p2p_transfer_num_workers``, ``p2p_transfer_timeout``. engine_kwargs: dict = field(default_factory=dict) + # Miles ``--check-weight-update-equal``: snapshot HF weights at startup, corrupt with + # random values, then verify the first weight sync restores them via SGLang weights_checker. + check_weight_update_equal: bool = False + # Allow quantized tensors to differ by up to 1 ULP in dequantized space during compare. + check_weight_update_allow_quant_error: bool = False # If set, this Python module is imported on every worker process before the # backend is instantiated, allowing custom backends to register themselves # in CheckpointEngineRegistry. diff --git a/verl/workers/engine/megatron/transformer_impl.py b/verl/workers/engine/megatron/transformer_impl.py index 9256982565b..43b4dbbc1ea 100644 --- a/verl/workers/engine/megatron/transformer_impl.py +++ b/verl/workers/engine/megatron/transformer_impl.py @@ -807,7 +807,7 @@ def forward_backward_batch(self, data: TensorDict, loss_function: Callable, forw RouterReplay.clear_global_router_replay_action() return output - def get_per_tensor_param(self, base_sync_done=False, **kwargs): + def get_per_tensor_param(self, base_sync_done=False, pp_local_export=False, **kwargs): peft_config = None non_merge_lora_sync = self.peft_cls is not None and not self.model_config.lora.get("merge", False) adapter_only = base_sync_done and non_merge_lora_sync @@ -815,10 +815,24 @@ def get_per_tensor_param(self, base_sync_done=False, **kwargs): peft_config = build_peft_config_for_vllm(self.model_config.lora) # when lora adapter only, we only load adapter weights when base sync is done, otherwise load all weights load_megatron_model_to_gpu(self.module, load_grad=False, load_frozen_params=not adapter_only) - if self.vanilla_bridge: - per_tensor_param = self.bridge.export_weights(self.module) - elif adapter_only: + if adapter_only: per_tensor_param = self.bridge.export_adapter_weights(self.module) + elif pp_local_export: + if self.vanilla_bridge: + raise ValueError( + "P2P weight sync requires PP-local export via Megatron-Bridge. " + "Set actor_rollout_ref.actor.megatron.vanilla_mbridge=False." + ) + from verl.utils.megatron_pp_local_export import export_pp_local_hf_weights + + per_tensor_param = export_pp_local_hf_weights( + self.bridge, + self.module, + merge_adapter_weights=not non_merge_lora_sync, + show_progress=False, + ) + elif self.vanilla_bridge: + per_tensor_param = self.bridge.export_weights(self.module) else: per_tensor_param = ( self.bridge.export_hf_weights(self.module, merge_adapter_weights=False) diff --git a/verl/workers/engine_workers.py b/verl/workers/engine_workers.py index ab9ca0e9ba4..9102cc130ad 100644 --- a/verl/workers/engine_workers.py +++ b/verl/workers/engine_workers.py @@ -665,7 +665,24 @@ def load_checkpoint(self, local_path, hdfs_path=None, del_local_after_load=False @register(dispatch_mode=Dispatch.ONE_TO_ALL) def save_checkpoint(self, local_path, hdfs_path=None, global_step=0, max_ckpt_to_keep=None): assert "actor" in self.role, "save_checkpoint only support actor role" - self.actor.save_checkpoint(local_path, hdfs_path, global_step, max_ckpt_to_keep) + release = getattr(self.checkpoint_engine, "release_for_checkpoint", None) + restore = getattr(self.checkpoint_engine, "restore_after_checkpoint", None) + if release is not None: + release() + try: + self.actor.save_checkpoint(local_path, hdfs_path, global_step, max_ckpt_to_keep) + finally: + if restore is not None: + restore() + + @register(dispatch_mode=Dispatch.ONE_TO_ALL) + def init_p2p_rollout_metadata(self, rollout_metadata: dict): + if "actor" not in self.role: + return + connect = getattr(self.checkpoint_engine, "connect_rollout_metadata", None) + if connect is None: + raise RuntimeError("checkpoint_engine backend does not support P2P rollout metadata") + connect(rollout_metadata) @register(dispatch_mode=Dispatch.ONE_TO_ALL, blocking=False) async def update_weights(self, global_steps: int = None, mode: str = "auto"): @@ -703,6 +720,12 @@ async def update_weights(self, global_steps: int = None, mode: str = "auto"): # so it consumes the sharded param generator instead of the full-tensor one. if effective_mode == "delta_sharded": per_tensor_param, _ = self.actor.engine.get_per_tensor_param_shard() + elif effective_mode == "p2p": + # P2P pushes PP-local HF weights from source ranks only; other + # ranks just participate in the PP-local export collectives. + if not self.checkpoint_engine.is_source: + return {} + per_tensor_param, _ = self.actor.engine.get_per_tensor_param(pp_local_export=True) else: per_tensor_param, _ = self.actor.engine.get_per_tensor_param() metrics = await self.checkpoint_engine.send_weights(per_tensor_param, global_steps=global_steps) diff --git a/verl/workers/rollout/replica.py b/verl/workers/rollout/replica.py index c01cdd79126..0d09a8c7fc2 100644 --- a/verl/workers/rollout/replica.py +++ b/verl/workers/rollout/replica.py @@ -290,6 +290,22 @@ async def resume_kv_cache(self): """Restore the kv_cache GPU memory after a weight sync.""" await asyncio.gather(*[server.resume_kv_cache.remote() for server in self.servers]) + async def begin_weight_update(self, selector: str = "all") -> dict[str, Any]: + """Begin a rollout weight-update session, if the backend requires one.""" + return {"success": True, "message": "No-op begin_weight_update", "selector": selector} + + async def end_weight_update(self) -> dict[str, Any]: + """End a rollout weight-update session, if the backend requires one.""" + return {"success": True, "message": "No-op end_weight_update"} + + async def update_weight_version(self, weight_version: str, abort_all_requests: bool = True) -> dict[str, Any]: + """Bump rollout weight version after a successful weight sync.""" + return { + "success": True, + "message": "No-op update_weight_version", + "new_version": weight_version, + } + async def start_profile(self, **kwargs): """Start profiling on the replica.""" await asyncio.gather(*[server.start_profile.remote(**kwargs) for server in self.servers]) diff --git a/verl/workers/rollout/sglang_rollout/async_sglang_server.py b/verl/workers/rollout/sglang_rollout/async_sglang_server.py index 99e74a21d95..c5dcb222be2 100644 --- a/verl/workers/rollout/sglang_rollout/async_sglang_server.py +++ b/verl/workers/rollout/sglang_rollout/async_sglang_server.py @@ -22,6 +22,7 @@ from typing import Any, Optional import ray +import requests import sglang import sglang.srt.entrypoints.engine import torch @@ -50,6 +51,10 @@ from verl.utils.tracking import RLInsightLogger from verl.workers.config import HFModelConfig, RolloutConfig from verl.workers.rollout.replica import RolloutMode, RolloutReplica, TokenOutput +from verl.workers.rollout.sglang_rollout.p2p_bootstrap import ( + sanitize_sglang_engine_kwargs_for_p2p, + should_enable_p2p_weight_transfer_bootstrap, +) from verl.workers.rollout.sglang_rollout.sglang_rollout import _set_envs_and_config from verl.workers.rollout.sglang_rollout.utils import SGLANG_LORA_NAME from verl.workers.rollout.utils import get_max_position_embeddings, run_uvicorn @@ -110,6 +115,39 @@ def _extract_prompt_logprobs_sglang( result_dict["prompt_logprobs"] = prompt_logprobs_ls +def _resolve_sglang_launch_subprocesses(): + """Resolve SGLang's subprocess launcher across API generations.""" + try: + from sglang.srt.entrypoints.http_server import _launch_subprocesses + + return _launch_subprocesses + except ImportError: + # Newer SGLang moved the implementation behind Engine._launch_subprocesses. + from sglang.srt.entrypoints.engine import Engine + + return Engine._launch_subprocesses + + +def _normalize_launch_subprocesses_result(result: tuple[Any, ...]) -> tuple[Any, Any, dict[str, Any]]: + """Extract (tokenizer_manager, template_manager, scheduler_info) across SGLang API generations.""" + tokenizer_manager, template_manager = result[:2] + + # Stock SGLang <= 0.5.8: (tokenizer_manager, template_manager, scheduler_info, port_args) + if len(result) >= 3 and isinstance(result[2], dict): + return tokenizer_manager, template_manager, result[2] + + # Stock SGLang 0.5.9: (tokenizer_manager, template_manager, scheduler_infos, port_args) + if len(result) >= 3 and isinstance(result[2], list | tuple) and result[2]: + return tokenizer_manager, template_manager, result[2][0] + + # Stock SGLang >= 0.5.10 and sglang-miles: + # (tokenizer_manager, template_manager, port_args, scheduler_init_result[, watchdog]) + if len(result) >= 4 and hasattr(result[3], "scheduler_infos"): + return tokenizer_manager, template_manager, result[3].scheduler_infos[0] + + raise RuntimeError(f"Unsupported SGLang launch_subprocesses result shape: {type(result)} / len={len(result)}") + + class SGLangHttpServer: """SGLang http server in single node, this is equivalent to launch server with command line: ``` @@ -152,6 +190,10 @@ def __init__( self._disaggregation_bootstrap_port = disaggregation_bootstrap_port self.config: RolloutConfig = omega_conf_to_dataclass(config) + self._enable_p2p_weight_transfer_bootstrap = should_enable_p2p_weight_transfer_bootstrap(self.config) + self._engine_info_bootstrap_port = None + self._engine_info_bootstrap_sock = None + self.model_config: HFModelConfig = omega_conf_to_dataclass(model_config, dataclass_type=HFModelConfig) max_position_embeddings = get_max_position_embeddings(self.model_config.hf_config) if self.config.max_model_len is None: @@ -213,6 +255,30 @@ def get_master_address(self): """Get master address and port for init NCCL process group.""" return self._master_address, self._master_port + def reserve_engine_info_bootstrap_port(self) -> int: + """Reserve a unique bootstrap port on the replica head node for P2P metadata.""" + if not self._enable_p2p_weight_transfer_bootstrap: + raise RuntimeError("P2P bootstrap is not enabled for this rollout server") + if self.node_rank != 0: + raise RuntimeError("engine_info_bootstrap_port can only be reserved on node_rank=0") + if self._engine_info_bootstrap_port is not None: + return self._engine_info_bootstrap_port + + port, sock = get_free_port(self._server_address, with_alive_sock=True) + self._engine_info_bootstrap_port = port + self._engine_info_bootstrap_sock = sock + logger.info( + f"SGLangHttpServer replica_rank={self.replica_rank} reserved " + f"engine_info_bootstrap_port={port} on {self._server_address}" + ) + return port + + def set_engine_info_bootstrap_port(self, port: int) -> None: + """Share the replica-wide bootstrap port with non-head nodes before launch.""" + if not self._enable_p2p_weight_transfer_bootstrap: + return + self._engine_info_bootstrap_port = int(port) + def get_server_address(self): """Get http server address and port.""" assert self._server_port is not None, "http server is not launched, port is None" @@ -254,7 +320,7 @@ async def launch_server(self, master_address: str = None, master_port: int = Non self._master_address = master_address self._master_port = master_port - engine_kwargs = self.config.get("engine_kwargs", {}).get("sglang", {}) or {} + engine_kwargs = sanitize_sglang_engine_kwargs_for_p2p(self.config) attention_backend = engine_kwargs.pop("attention_backend", None) mm_attention_backend = engine_kwargs.pop("mm_attention_backend", None) # Delta checkpoint engines apply sparse weight updates in place through SGLang's @@ -297,7 +363,10 @@ async def launch_server(self, master_address: str = None, master_port: int = Non "dtype": self.config.dtype, "mem_fraction_static": self.config.gpu_memory_utilization, "disable_cuda_graph": self.config.enforce_eager, - "enable_memory_saver": True, + # torch_memory_saver is incompatible with Mooncake TransferEngine (RDMA + # registrations don't survive VMM unmap/remap); SGLang refuses to start + # the P2P engine-info bootstrap server when memory saver is on. + "enable_memory_saver": not self._enable_p2p_weight_transfer_bootstrap, "base_gpu_id": self.base_gpu_id, "gpu_id_step": 1, "tp_size": infer_tp, @@ -376,6 +445,33 @@ async def launch_server(self, master_address: str = None, master_port: int = Non if self.config.enable_rollout_routing_replay: args.update({"enable_return_routed_experts": True}) + if self._enable_p2p_weight_transfer_bootstrap: + if self._engine_info_bootstrap_port is None: + raise RuntimeError( + "P2P bootstrap is enabled but engine_info_bootstrap_port is unset. " + "Reserve it on replica head node before launch_server." + ) + if self._engine_info_bootstrap_sock is not None: + self._engine_info_bootstrap_sock.close() + self._engine_info_bootstrap_sock = None + if args.get("enable_memory_saver"): + raise ValueError( + "checkpoint_engine.backend=p2p is incompatible with SGLang's memory saver: " + "Mooncake TransferEngine RDMA registrations do not survive torch_memory_saver " + "VMM unmap/remap, so SGLang would silently skip starting the P2P engine-info " + "bootstrap server. Remove `engine_kwargs.sglang.enable_memory_saver=True` " + "from the rollout config." + ) + args["remote_instance_weight_loader_start_seed_via_transfer_engine"] = True + args["engine_info_bootstrap_port"] = self._engine_info_bootstrap_port + # Multi-node only: SGLang registers to dist_init_addr host; single-node uses 127.0.0.1. + if self.nnodes > 1: + args["host"] = self._server_address + logger.info( + f"SGLangHttpServer replica_rank={self.replica_rank} node_rank={self.node_rank} " + f"enabling P2P bootstrap on port {self._engine_info_bootstrap_port}" + ) + # mtp if self.config.mtp is not None and self.config.mtp.enable and self.config.mtp.enable_rollout: # Enable weights CPU backup for sglang >= 0.5.6 @@ -395,32 +491,19 @@ async def launch_server(self, master_address: str = None, master_port: int = Non sglang.srt.entrypoints.engine._set_envs_and_config = _set_envs_and_config os.environ["SGLANG_BLOCK_NONZERO_RANK_CHILDREN"] = "0" server_args = ServerArgs(**args) - # For SGLang main branch or version >= 0.5.10 - # The latest main branch of SGLang has wrapped the _launch_subprocesses function inside the Engine class - if version.parse(sglang.__version__) >= version.parse("0.5.10"): - from sglang.srt.entrypoints.http_server import Engine - - self.tokenizer_manager, self.template_manager, self.scheduler_info, *_ = Engine._launch_subprocesses( - server_args=server_args, - init_tokenizer_manager_func=sglang.srt.entrypoints.engine.init_tokenizer_manager, - run_scheduler_process_func=sglang.srt.entrypoints.engine.run_scheduler_process, - run_detokenizer_process_func=sglang.srt.entrypoints.engine.run_detokenizer_process, - ) - elif version.parse(sglang.__version__) >= version.parse("0.5.7"): - from sglang.srt.entrypoints.http_server import _launch_subprocesses - - self.tokenizer_manager, self.template_manager, self.scheduler_info, *_ = _launch_subprocesses( + if version.parse(sglang.__version__) >= version.parse("0.5.7"): + launch_result = _resolve_sglang_launch_subprocesses()( server_args=server_args, init_tokenizer_manager_func=sglang.srt.entrypoints.engine.init_tokenizer_manager, run_scheduler_process_func=sglang.srt.entrypoints.engine.run_scheduler_process, run_detokenizer_process_func=sglang.srt.entrypoints.engine.run_detokenizer_process, ) else: - from sglang.srt.entrypoints.http_server import _launch_subprocesses + launch_result = _resolve_sglang_launch_subprocesses()(server_args=server_args) - self.tokenizer_manager, self.template_manager, self.scheduler_info, *_ = _launch_subprocesses( - server_args=server_args - ) + self.tokenizer_manager, self.template_manager, self.scheduler_info = _normalize_launch_subprocesses_result( + launch_result + ) # In multi-node cases, non-zero rank nodes should not launch http server. if self.node_rank > 0: @@ -704,6 +787,106 @@ async def resume_generation(self): return await self.tokenizer_manager.continue_generation(ContinueGenerationReqInput()) + async def get_remote_instance_transfer_engine_info(self, rank: int): + if self.node_rank != 0: + return None + + server_args = self.tokenizer_manager.server_args + resp = requests.get( + f"{server_args.engine_info_bootstrap_url}/get_transfer_engine_info", + params={"rank": rank}, + timeout=5, + ) + resp.raise_for_status() + return resp.json()["remote_instance_transfer_engine_info"] + + async def get_parallelism_info(self, rank: int): + if self.node_rank != 0: + return None + + server_args = self.tokenizer_manager.server_args + resp = requests.get( + f"{server_args.engine_info_bootstrap_url}/get_parallelism_config", + params={"rank": rank}, + timeout=5, + ) + resp.raise_for_status() + return resp.json() + + async def get_server_info(self) -> dict[str, Any]: + if self.node_rank != 0: + return {} + + server_args = self.tokenizer_manager.server_args + internal_states = await self.tokenizer_manager.get_internal_state() + return { + **dataclasses.asdict(server_args), + **self.scheduler_info, + "internal_states": internal_states, + "version": sglang.__version__, + "kv_events": server_args.describe_kv_events_publisher(), + } + + async def begin_weight_update(self, selector: str = "all") -> dict[str, Any]: + if self.node_rank != 0: + return {} + + # Available in SGLang builds with the Miles weight-update lifecycle. + from sglang.srt.managers.io_struct import BeginWeightUpdateReqInput + + success, message = await self.tokenizer_manager.begin_weight_update( + BeginWeightUpdateReqInput(selector=selector), None + ) + return {"success": success, "message": message} + + async def end_weight_update(self) -> dict[str, Any]: + if self.node_rank != 0: + return {} + + from sglang.srt.managers.io_struct import EndWeightUpdateReqInput + + success, message = await self.tokenizer_manager.end_weight_update(EndWeightUpdateReqInput(), None) + return {"success": success, "message": message} + + async def update_weight_version(self, weight_version: str, abort_all_requests: bool = True) -> dict[str, Any]: + if self.node_rank != 0: + return {} + + new_version = str(weight_version) + if abort_all_requests: + self.tokenizer_manager.abort_request(abort_all=True) + self.tokenizer_manager.server_args.weight_version = new_version + try: + self.global_steps = int(new_version) + except (TypeError, ValueError): + logger.debug( + "weight_version %r is not numeric; keeping global_steps=%s", + new_version, + self.global_steps, + ) + return { + "success": True, + "message": f"Weight version updated to {new_version}", + "new_version": new_version, + } + + async def check_weights(self, action: str, allow_quant_error: bool = False) -> dict[str, Any]: + if self.node_rank != 0: + return {} + + from sglang.srt.managers.io_struct import CheckWeightsReqInput + + success, message, ranks, per_engine_checksum = await self.tokenizer_manager.check_weights( + CheckWeightsReqInput(action=action, allow_quant_error=allow_quant_error), + None, + ) + return { + "success": success, + "message": message, + "ranks": ranks, + "per_engine_checksum": per_engine_checksum, + } + async def start_profile(self, **kwargs): if ( self.profiler_controller.check_enable() @@ -769,6 +952,9 @@ async def launch_servers(self): if os.environ.get(f"RAY_EXPERIMENTAL_NOSET_{visible_devices_keyword}", None): logger.warning(f"RAY_EXPERIMENTAL_NOSET_{visible_devices_keyword} is set True!") base_gpu_id = (0 + self.replica_rank * replica_world_size) % self.gpus_per_node + + enable_p2p_bootstrap = should_enable_p2p_weight_transfer_bootstrap(self.config) + # create server actor in each node with node affinity and cuda visible devices for node_rank in range(self.nnodes): workers = self.workers[ @@ -824,6 +1010,16 @@ async def launch_servers(self): ) self.servers.append(server) + if enable_p2p_bootstrap: + bootstrap_port = await self.servers[0].reserve_engine_info_bootstrap_port.remote() + if self.nnodes > 1: + await asyncio.gather( + *[server.set_engine_info_bootstrap_port.remote(bootstrap_port) for server in self.servers[1:]] + ) + logger.info( + f"SGLangReplica replica_rank={self.replica_rank} allocated engine_info_bootstrap_port={bootstrap_port}" + ) + # launch http server in each node master_address, master_port = None, None if self.nnodes > 1: @@ -855,3 +1051,24 @@ async def abort_all_requests(self): async def resume_generation(self): """Resume generation on the primary server after abort_all_requests.""" await self.servers[0].resume_generation.remote() + + async def get_remote_instance_transfer_engine_info(self, rank: int): + return await self.servers[0].get_remote_instance_transfer_engine_info.remote(rank) + + async def get_parallelism_info(self, rank: int): + return await self.servers[0].get_parallelism_info.remote(rank) + + async def get_server_info(self) -> dict[str, Any]: + return await self.servers[0].get_server_info.remote() + + async def begin_weight_update(self, selector: str = "all") -> dict[str, Any]: + return await self.servers[0].begin_weight_update.remote(selector) + + async def end_weight_update(self) -> dict[str, Any]: + return await self.servers[0].end_weight_update.remote() + + async def update_weight_version(self, weight_version: str, abort_all_requests: bool = True) -> dict[str, Any]: + return await self.servers[0].update_weight_version.remote(weight_version, abort_all_requests) + + async def check_weights(self, action: str, allow_quant_error: bool = False) -> dict[str, Any]: + return await self.servers[0].check_weights.remote(action, allow_quant_error) diff --git a/verl/workers/rollout/sglang_rollout/http_server_engine.py b/verl/workers/rollout/sglang_rollout/http_server_engine.py index 63034857ce7..31ec1a234b8 100644 --- a/verl/workers/rollout/sglang_rollout/http_server_engine.py +++ b/verl/workers/rollout/sglang_rollout/http_server_engine.py @@ -690,7 +690,7 @@ async def _make_async_request( try: async with self._get_session() as session: if method.upper() == "GET": - async with session.get(url, timeout=timeout) as response: + async with session.get(url, params=payload or None, timeout=timeout) as response: response.raise_for_status() return await _read_async_response(response) else: @@ -774,6 +774,44 @@ async def update_weights_from_tensor( }, ) + async def get_remote_instance_transfer_engine_info(self, rank: int): + data = await self._make_async_request( + "remote_instance_transfer_engine_info", + {"rank": rank}, + method="GET", + only_master=False, + ) + return data["remote_instance_transfer_engine_info"] + + async def get_parallelism_info(self, rank: int): + return await self._make_async_request( + "parallelism_config", + {"rank": rank}, + method="GET", + only_master=False, + ) + + async def get_server_info(self) -> dict[str, Any]: + return await self._make_async_request("server_info", method="GET", only_master=False) + + async def begin_weight_update(self, selector: str = "all") -> dict[str, Any]: + return await self._make_async_request("begin_weight_update", {"selector": selector}) + + async def end_weight_update(self) -> dict[str, Any]: + return await self._make_async_request("end_weight_update", {}) + + async def update_weight_version(self, weight_version: str, abort_all_requests: bool = True) -> dict[str, Any]: + return await self._make_async_request( + "update_weight_version", + {"new_version": str(weight_version), "abort_all_requests": abort_all_requests}, + ) + + async def check_weights(self, action: str, allow_quant_error: bool = False) -> dict[str, Any]: + return await self._make_async_request( + "weights_checker", + {"action": action, "allow_quant_error": allow_quant_error}, + ) + async def load_lora_adapter_from_tensor(self, req): return await self._make_async_request( "load_lora_adapter_from_tensors", diff --git a/verl/workers/rollout/sglang_rollout/p2p_bootstrap.py b/verl/workers/rollout/sglang_rollout/p2p_bootstrap.py new file mode 100644 index 00000000000..da71ba88b87 --- /dev/null +++ b/verl/workers/rollout/sglang_rollout/p2p_bootstrap.py @@ -0,0 +1,45 @@ +# Copyright 2025 Bytedance Ltd. and/or its affiliates +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from __future__ import annotations + +from typing import Any + +from verl.workers.config import RolloutConfig + + +def _sglang_engine_kwargs(config: RolloutConfig) -> dict[str, Any]: + engine_kwargs = getattr(config, "engine_kwargs", None) or {} + if not isinstance(engine_kwargs, dict): + engine_kwargs = dict(engine_kwargs) + sglang_kwargs = engine_kwargs.get("sglang", {}) or {} + if not isinstance(sglang_kwargs, dict): + sglang_kwargs = dict(sglang_kwargs) + return sglang_kwargs + + +def should_enable_p2p_weight_transfer_bootstrap(config: RolloutConfig) -> bool: + """Return whether rollout should start EngineInfoBootstrapServer for P2P RDMA.""" + checkpoint_engine = getattr(config, "checkpoint_engine", None) + backend = getattr(checkpoint_engine, "backend", None) + return backend == "p2p" + + +def sanitize_sglang_engine_kwargs_for_p2p(config: RolloutConfig) -> dict[str, Any]: + """Return a copy of SGLang engine kwargs with P2P bootstrap flags owned by Verl.""" + kwargs = dict(_sglang_engine_kwargs(config)) + # Verl auto-allocates bootstrap ports and enables P2P seeding when backend=p2p. + kwargs.pop("engine_info_bootstrap_port", None) + kwargs.pop("remote_instance_weight_loader_start_seed_via_transfer_engine", None) + return kwargs diff --git a/verl/workers/rollout/sglang_rollout/sglang_rollout.py b/verl/workers/rollout/sglang_rollout/sglang_rollout.py index 07e30d07c80..bbe124ee2aa 100644 --- a/verl/workers/rollout/sglang_rollout/sglang_rollout.py +++ b/verl/workers/rollout/sglang_rollout/sglang_rollout.py @@ -19,7 +19,7 @@ import multiprocessing as mp import os from dataclasses import asdict -from typing import Generator +from typing import Any, Generator import ray import sglang.srt.entrypoints.engine @@ -288,6 +288,48 @@ async def release(self): tags = ["kv_cache", "weights"] await self._engine.release_memory_occupation(tags=tags) + async def get_remote_instance_transfer_engine_info(self, rank: int): + await self._init_server_adapter() + if self._engine is None or not self._is_server_tp_leader(): + return None + return await self._engine.get_remote_instance_transfer_engine_info(rank) + + async def get_parallelism_info(self, rank: int): + await self._init_server_adapter() + if self._engine is None or not self._is_server_tp_leader(): + return None + return await self._engine.get_parallelism_info(rank) + + async def get_server_info(self) -> dict[str, Any]: + await self._init_server_adapter() + if self._engine is None or not self._is_server_tp_leader(): + return {} + return await self._engine.get_server_info() + + async def begin_weight_update(self, selector: str = "all") -> dict[str, Any]: + await self._init_server_adapter() + if self._engine is None or not self._is_server_tp_leader(): + return {} + return await self._engine.begin_weight_update(selector) + + async def end_weight_update(self) -> dict[str, Any]: + await self._init_server_adapter() + if self._engine is None or not self._is_server_tp_leader(): + return {} + return await self._engine.end_weight_update() + + async def update_weight_version(self, weight_version: str, abort_all_requests: bool = True) -> dict[str, Any]: + await self._init_server_adapter() + if self._engine is None or not self._is_server_tp_leader(): + return {} + return await self._engine.update_weight_version(weight_version, abort_all_requests=abort_all_requests) + + async def check_weights(self, action: str, allow_quant_error: bool = False) -> dict[str, Any]: + await self._init_server_adapter() + if self._engine is None or not self._is_server_tp_leader(): + return {} + return await self._engine.check_weights(action, allow_quant_error=allow_quant_error) + async def update_weights( self, weights: Generator[tuple[str, torch.Tensor], None, None],