[ckpt, sglang, megatron] feat: P2P checkpoint engine for trainer-to-rollout weight sync via Mooncake RDMA

[AkiRusProd]

What does this PR do?

Adds a new checkpoint-engine backend p2p that pushes weights from Megatron trainer ranks directly into SGLang rollout engines over Mooncake RDMA, without building an NCCL process group between trainer and rollout. Each trainer source rank exports its PP-local shard of HF-format weights, stages them in a pinned-CPU replica, and RDMA-writes them point-to-point into the rollout engines' weight buffers. The wire protocol and update lifecycle are interoperable with the Miles P2P weight-update protocol (sgl-project/miles), including its optional --check-weight-update-equal-style correctness check.

Why this is useful:

• No trainer↔rollout NCCL group: sync works across disaggregated placement groups and is robust to rollout elastic scale up/down (rollout metadata is cached and invalidated on membership change).
• PP-local export: each PP stage exports only its own layers (via Megatron-Bridge), so no full-model gather on rank 0.
• Built-in end-to-end correctness check (check_weight_update_equal): rollout weights are snapshotted at startup, corrupted with random values, and the first sync must restore them bit-exactly — verified on 8 replicas (see Test).
• Fast at scale: in production use of the original implementation (internal verl fork), syncing a ~700B MoE model took ~5 s with the p2p backend vs ~6.5 min with the NCCL checkpoint engine on the same setup (~78× faster).

Similar PR check: no open PR implements P2P/RDMA weight sync

Opening as a draft?: the backend is fully functional against the sglang-miles branch (validated e2e on a live cluster, see Test), while two control-plane pieces it needs are still being upstreamed to sglang (see the CI caveat in Test). Happy to keep this in draft until they land in a released sglang, or adjust scope per maintainer guidance.

Test

Unit tests (pass against the sglang-miles build; require sglang importable — see caveat below):

pytest tests/checkpoint_engine/test_p2p_checkpoint_engine.py \
       tests/checkpoint_engine/test_p2p_transfer_utils.py \
       tests/checkpoint_engine/test_weight_checker.py \
       tests/utils/test_megatron_pp_local_export.py \
       tests/workers/rollout/sglang_rollout/test_p2p_bootstrap.py \
       tests/workers/rollout/sglang_rollout/test_p2p_control_smoke.py -x -q

E2E validation

GRPO smoke run on 4 nodes — 2 trainer nodes (Megatron, TP2/PP2/EP2; PP=2 specifically exercises the PP-local export path) + 2 rollout nodes (SGLang, TP2/EP2, 8 replicas), MoE model with MLA, checkpoint_engine.backend=p2p and check_weight_update_equal=True. The code was developed and validated against the sglang-miles branch at commit c855d1b (the branch carrying the Miles weight-update lifecycle endpoints). Two runs, both SUCCEEDED:

• [WeightChecker] compare PASSED on 8 rollout replica(s) (allow_quant_error=False) — rollout weights corrupted at startup were restored bit-exactly by the first p2p sync.
• Megatron+HF checkpoint save: P2P CPU replica is released before save (avoids host-memory pressure) and restored after (RDMA buffer re-registration works).

Beyond the smoke run above, the original implementation of this backend (which this port follows logic-wise) has been validated in production training with a ~700B MoE model — that is where the sync timings quoted in the overview (~5 s vs ~6.5 min over NCCL) come from.

Tested scope: all validation was done with BF16 weights (allow_quant_error=False, bit-exact comparison). The code carries quantization plumbing inherited from the Miles design (fp8/fp4 GEMM config initialization from rl_quant_profile, check_weight_update_allow_quant_error for tolerance-based comparison), but FP8 rollout has not been tested with this backend.

Known CI caveat: this backend uses SGLang control-plane endpoints that are being upstreamed from the Miles weight-update lifecycle. Most of the surface has already landed on sglang main: /weights_checker, /update_weight_version, /remote_instance_transfer_engine_info, and the engine-info bootstrap server (engine_info_bootstrap_port in ServerArgs, /get_transfer_engine_info). Two pieces are not on sglang main yet and currently exist only on the sglang-miles branch (reference: c855d1b): the begin/end weight-update session (BeginWeightUpdateReqInput / EndWeightUpdateReqInput) and /get_parallelism_config on the bootstrap server (related upstreaming PRs: sgl-project/sglang#17326 — parameter mapper used by the trainer-side updater, sgl-project/sglang#20907 — parallelism info). The sglang 0.5.8 pinned by CI predates all of the above. Imports of the corresponding *ReqInput classes are lazy, so nothing breaks at import time on stock sglang; test_filter_server_args_dict_drops_non_serializable_fields may need a skipif on stock builds. Happy to adjust markers however maintainers prefer.

API and Usage Example

No breaking changes; new backend is opt-in via config. mooncake-transfer-engine is a lazy runtime dependency of the trainer only when backend=p2p (not added to requirements).

actor_rollout_ref.rollout.checkpoint_engine.backend=p2p \
actor_rollout_ref.rollout.checkpoint_engine.update_weights_bucket_megabytes=1024 \
# optional bit-exact weight-update verification (debug):
actor_rollout_ref.rollout.checkpoint_engine.check_weight_update_equal=True \
# optional transfer tuning:
+actor_rollout_ref.rollout.checkpoint_engine.engine_kwargs.p2p.p2p_transfer_num_workers=4 \
+actor_rollout_ref.rollout.checkpoint_engine.engine_kwargs.p2p.p2p_transfer_timeout=30.0

Note: enable_memory_saver is forced off under backend=p2p (explicit ValueError if the user sets it) — memory saver is fundamentally incompatible with Mooncake TransferEngine (the seed flag is dropped and the EngineInfoBootstrapServer never comes up, so p2p metadata collection gets connection-refused).

Design & Code Changes

Sync lifecycle (Miles P2P order), driven by CheckpointEngineManager._update_weights_p2p:

abort rollout requests → release kv_cache → begin weight-update session
→ (one-time) connect rollout metadata → RDMA push from trainer source ranks
→ bump weight version → end session → resume

New:

• verl/checkpoint_engine/p2p/ — P2PCheckpointEngine (registered as backend=p2p), trainer-side updater (pinned-CPU replica + Mooncake RDMA writes), transfer planning / bucketing utilities, optional weight checker.
• verl/utils/megatron_pp_local_export.py — PP-local HF export via Megatron-Bridge.
• verl/workers/rollout/sglang_rollout/p2p_bootstrap.py — bootstrap enablement and engine_kwargs sanitization.

Modified:

• verl/checkpoint_engine/base.py — p2p update path; replica-level begin/end weight-update sessions (no-op default implementations added on RolloutReplica, so other backends are unaffected); rollout-metadata cache invalidated on elastic scale up/down.
• verl/workers/engine_workers.py — p2p send path (source ranks only, pp_local_export=True); release/restore of P2P CPU buffers around save_checkpoint.
• verl/workers/engine/megatron/transformer_impl.py — pp_local_export mode in get_per_tensor_param.
• SGLang rollout (async_sglang_server.py, http_server_engine.py, sglang_rollout.py) — bootstrap-port plumbing and control-plane endpoints at server / adapter / replica levels.
• verl/workers/config/rollout.py, verl/trainer/config/rollout/rollout.yaml — config knobs.
• verl/experimental/fully_async_policy/ — tolerate None param_version produced by p2p sync.

AI assistance disclosure

AI assistance was used for this contribution. The original P2P implementation was developed in an internal verl 0.7.x fork and validated there in production training; it was then ported onto current upstream main. The code has been verified by a human: I have reviewed every changed line, run the new unit tests, validated the change end-to-end on a live cluster, and can defend it end-to-end. Duplicate-work check: see links in the overview section; test commands and results are listed above.

Checklist Before Submitting

• Read the Contribute Guide.
• Applied pre-commit checks (ruff check + format clean).
• Documentation: <add docs page for the p2p backend? — decide before submit>
• Unit tests added (tests/checkpoint_engine/, tests/workers/rollout/sglang_rollout/);

  e2e in CI not feasible: requires multi-node RDMA (Mooncake) cluster and an SGLang
  build with the Miles weight-update lifecycle — validated manually, see Test.
  

• Send a message in the ci-request Slack channel once ready for CI.
• Not related to recipe submodule.

[gemini-code-assist]

Code Review

This pull request introduces a disaggregated peer-to-peer (P2P) weight transfer mechanism (Miles-style P2P sync) between the trainer and SGLang rollout replicas, bypassing standard NCCL/broadcast pathways for weight updates. It adds a new P2PCheckpointEngine, corresponding tests, and scripts for control-plane smoke testing. The reviewer's feedback highlights several critical bugs and performance optimization opportunities: potential list length mismatches in _do_p2p_write_one_session when a parameter is missing from the remote session's weight info; redundant weight loading operations across shared CPU replicas; sequential await calls in loops when querying remote metadata or starting weight updates (which should be parallelized using asyncio.gather); blocking synchronous HTTP requests (requests.get) inside async Ray actor methods that should be offloaded using asyncio.to_thread; unsafe use of assert statements for critical runtime validations; a potential TypeError in _normalize_param_version when handling non-numeric string weight versions; and a usability improvement to automatically disable SGLang's memory saver when using the P2P backend instead of raising a fatal error.

Important

The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.

[gemini-code-assist]

Sequential await calls in a loop to query metadata from all rollout replicas will block the event loop and significantly slow down the initialization phase. Since these are independent network requests, they should be parallelized using asyncio.gather.

    import asyncio
    async def query_one(engine_ind, engine_rank):
        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()
        return engine_ind, engine_rank, session_id, weights_info, parallelism_info, filter_server_args_dict(server_info)

    results = await asyncio.gather(
        *[query_one(ind, rank) for ind, rank in targets_to_query]
    )

    for engine_ind, engine_rank, session_id, weights_info, parallelism_info, server_args in results:
        session_id_to_server_args[session_id] = server_args
        remote_weight_infos_by_session_id[session_id] = (weights_info, parallelism_info)
        targets_to_session_id[(engine_ind, engine_rank)] = session_id

References

1. Avoid using ray.get() inside async methods of Ray actors. Instead, use await on remote calls or asyncio.gather to parallelize multiple remote calls asynchronously.

[gemini-code-assist]

begin_weight_update_replicas is executed sequentially for each rollout replica. For larger clusters with multiple replicas, this sequential network dispatch blocks the event loop and delays weight synchronization. Parallelize these requests using asyncio.gather.

Suggested change

	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)
	started_replicas = []
	try:
	results = await asyncio.gather(
	*[replica.begin_weight_update(selector) for replica in self.replicas],
	return_exceptions=True,
	)
	for replica, result in zip(self.replicas, results, strict=True):
	self._raise_for_weight_update_control_result("begin_weight_update", replica, result)
	started_replicas.append(replica)
	except Exception:

References

1. Avoid using ray.get() inside async methods of Ray actors. Instead, use await on remote calls or asyncio.gather to parallelize multiple remote calls asynchronously.

[gemini-code-assist]

Calling synchronous/blocking requests.get inside an async method of a Ray actor blocks the event loop, which can cause severe performance degradation or deadlocks. Use asyncio.to_thread to run the blocking request in a separate thread.

Suggested change

	resp = requests.get(
	f"{server_args.engine_info_bootstrap_url}/get_transfer_engine_info",
	params={"rank": rank},
	timeout=5,
	)
	resp = await asyncio.to_thread(
	requests.get,
	f"{server_args.engine_info_bootstrap_url}/get_transfer_engine_info",
	params={"rank": rank},
	timeout=5,
	)

[gemini-code-assist]

Calling synchronous/blocking requests.get inside an async method of a Ray actor blocks the event loop. Use asyncio.to_thread to run the blocking request in a separate thread.

Suggested change

	resp = requests.get(
	f"{server_args.engine_info_bootstrap_url}/get_parallelism_config",
	params={"rank": rank},
	timeout=5,
	)
	resp = await asyncio.to_thread(
	requests.get,
	f"{server_args.engine_info_bootstrap_url}/get_parallelism_config",
	params={"rank": rank},
	timeout=5,
	)

[gemini-code-assist]

Using assert statements for critical runtime validation is unsafe because assertions are stripped out when Python is run with optimization flags (-O). If assertions are disabled, any pending staged tensors will go unnoticed, leading to silent weight mismatch or corruption in the rollout replicas. Use an explicit if check and raise a RuntimeError instead.

Suggested change

	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._tensor_update_pending or self._staged_tensors:
	raise RuntimeError(
	f"Pending staged tensors after P2P update: {self._tensor_update_pending}, {self._staged_tensors}"
	)