vLLM Stateful Weight Sync Sender

[hao-aaron]

Weight sync: migrate onto vLLM 0.28's trainer-send abstraction

Deletes SkyRL's WeightTransferStrategy / WeightTransferSender / WeightExtractor
layers and drives all four weight-sync backends through vLLM's own
TrainerWeightTransferEngine, over vLLM's native RLHF routes.
Follow-up to #2114.

---

Why

vLLM 0.26 shipped only the receive half of the weight-transfer abstraction, so SkyRL
owned the send half: a strategy/sender hierarchy, its own extractor + chunk types, and a
/collective_rpc worker extension that wrapped every receive in
set_current_vllm_config because vLLM did not.

0.28 ships the send half (TrainerWeightTransferEngine, WeightSource,
VLLMWeightSyncClient, WeightTransferTrainerFactory) and wraps set_current_vllm_config
itself in GPUWorker.start_weight_update / update_weights / finish_weight_update. All
of SkyRL's parallel machinery becomes redundant.

Shape of the change

Before — worker → SkyRL strategy → SkyRL sender → /collective_rpc → SkyRL worker
extension → vLLM engine.

After — worker builds a WeightSource → vLLM trainer engine → native RLHF routes →
vLLM engine (subclassed only where SkyRL needs a hook).

init_weight_sync_state is now four statements; broadcast_to_inference_engines is one
engine.send_weights() plus a memory bracket.

self._weight_sync_engine = await self._weight_sync_thread(
    build_trainer_engine,
    ie_cfg=inference_engine_cfg,
    colocate_all=self.cfg.placement.colocate_all,
    rank=torch.distributed.get_rank(),
    inference_world_size=inference_world_size,
    source_factory=self._build_weight_source,
    server_urls=..., data_parallel_size=..., base_model_path=...,
)

New modules (weight_sync/)

file	lines	what
sources.py	120	FsdpWeightSource / MegatronWeightSource — vLLM's metadata() + __iter__, nothing else
trainer_engines.py	310	build_trainer_engine: resolve backend → build source → trainer_init
control_plane.py	275	SkyrlWeightSyncClient (blocking HTTP) + the per-server init rewrites
skyrl_engines.py	195	receive side: skyrl_nccl / skyrl_ipc + the drafter-reload proxy
delta_trainer.py	199	DeltaTrainerWeightTransferEngine
patches/vllm/patch_model_runner_registry.py	59	records GPUModelRunner so an engine can reach the drafter

Deleted

transfer_strategy.py, broadcast_strategy.py, cuda_ipc_strategy.py,
delta_strategy.py, sharded_rdt_strategy.py, nccl_trainer_send.py,
weight_extractor.py, weight_extractor_utils.py, layerwise_reload.py,
rdt_control_plane.py, rdt_control_protocol.py.

Also gone: WeightChunk, WeightUpdateRequest, FSDPWeightExtractor,
MegatronWeightExtractor, get_transfer_strategy_cls, for_servers /
to_api_payload / get_vllm_transfer_engine, sender_initializes_receivers,
RemoteInferenceClient.init_weight_update_communicator / update_weights_nccl /
update_weights_ipc / start_weight_update / finish_weight_update.

sharded_rdt_base.py shrank 410 → 207: the vendored ABCs now come from the wheel; what
remains is the two channels a pull backend needs and vLLM has no concept of
(held_names() ownership, groups() / iter_groups()).

Backends

logical	trainer engine	receive engine
nccl	vLLM's NCCLTrainerWeightTransferEngine	skyrl_nccl
ipc	vLLM's IPCTrainerWeightTransferEngine	skyrl_ipc
delta	DeltaTrainerWeightTransferEngine	delta
sharded_rdt	ShardedRDTTrainerWeightTransferEngine	sharded_rdt

The receive side takes new names for NCCL/IPC because SkyRL subclasses vLLM's engines and
register_engine refuses an already-registered name. get_transfer_strategy stays and is
called from both sides — the driver (to configure the servers) and
build_trainer_engine (to pick the trainer engine) — so the two cannot disagree.

The worker extension

NewInferenceWorkerWrap goes from 8 weight-sync methods to 2. Both remaining are limits of
dispatch, not of the engine abstraction:

• fetch_weights — /collective_rpc reaches worker methods by name only, and no native
  route can invoke an engine method.
• skyrl_sleep_for_weight_sync / skyrl_wake_for_weight_sync — EngineCore.sleep
  hardcodes clear_prefix_cache = level >= 1, and CuMemBackend.suspend cannot express
  "discard weights, offload kv_cache".

Gone with the rest: skyrl_start/finish_weight_update, update_weights_nccl,
update_weights_ipc, init_weight_transfer_engine_rdt, update_weights_rdt, and
LayerwiseReloadWorkerMixin (0.28's get_numel_loaded is character-identical to SkyRL's
patched copy, so the CopyCounter patch is no longer needed).

Behavior changes worth reviewing

1. No bucketing. The Megatron export is one whole-model export_hf_weights call.
   Bucketing never bounded memory (it accumulates a bucket before handing it on); its
   real purposes — IPC handle count and fused-loader grouping — are subsumed by vLLM's
   packed_ipc_producer. One call also satisfies _accumulate_grouped_export's
   "every task of a group_key in one call" by construction.

2. packed=True forced on IPC, overriding vLLM's default. The unpacked path holds a
   contiguous copy of every parameter until past finish_weight_update — the whole model
   resident on the trainer.

3. Packed buffer sized from the source. vLLM's 1 GiB default is smaller than a
   large-vocab embedding (Qwen3-235B: 1.24 GiB), which raises on the IPC path. Pulls the
   Megatron metadata() dry export forward to init.

4. Delta gained a drafter reload it never had. The old
   DeltaWeightTransferEngine.receive_weights was a bare self.model.load_weights(...), so
   MTP + delta silently drafted with stale weights. Now re-streams iter_tensors a second
   time. This is the one place the PR is not behavior-preserving — happy to split it out.

5. Delta owns its layerwise lifecycle. Previously no-ops, with the /collective_rpc
   wrap doing it; now start/finish_weight_update run initialize/finalize_layerwise_reload.

6. _weight_sync_thread selects this rank's CUDA device inside the worker thread.
   send_weights runs off the event loop, the current device is thread-local and defaults
   to 0, and under RAY_EXPERIMENTAL_NOSET_CUDA_VISIBLE_DEVICES this rank's device is not
   0 — so the NCCL communicator, IPC handles and gathers would land on the wrong GPU.

7. Capability probes replace declared flags. getattr(engine, "skyrl_…", default),
   because two of the four trainer engines are vLLM's own classes and cannot carry SkyRL
   attributes. sender_initializes_receivers is gone — vacuous now that every engine
   rendezvouses in its own trainer_init.

Spec-decode drafter

Unchanged in effect, moved in mechanism. spec_decode_utils.py is untouched.

vLLM's native /start_draft_weight_update is a separate session with its own data —
right for EAGLE (separate checkpoint), wrong for MTP, where the head's weights are already
in the policy's stream and share_embeddings unconditionally aliases the target's
embed_tokens/lm_head into the draft (so a draft session would re-send the single
largest tensor). Instead the drafter is reloaded from the same stream in the same session,
via a _LoadWeightsProxy installed only when this process actually has a drafter.

Per backend: nccl proxy → proxy (same); ipc direct call → proxy (same effect);
delta none → second pass (new); sharded_rdt none → none (still a gap).

Known limitation, pre-existing: none of these run initialize_layerwise_reload on the
draft model, so process_weights_after_loading never runs there. Fine for bf16 MTP
heads; wrong for a quantized drafter.

Tests

Deleted with their subjects: test_transfer_strategies.py (969 lines),
test_sharded_rdt_strategy.py, test_weight_chunk.py, test_weight_extractor_utils.py,
test_rdt_control_plane.py.

Added: test_control_plane.py (the init rewrites — a bad NCCL rank_offset hangs in the
rendezvous rather than erroring), test_sources.py (channel agreement + laziness),
test_trainer_engines.py (init-info choices, backend resolution, capability probes).

Rewritten: test_prefix_cache_reset.py (probes, with a fake engine rather than a Mock so
the defaults are actually exercised), GPU test_weight_sync.py (all four backends now go
through build_trainer_engine, i.e. the production path),
test_megatron_extractor_consistency.py → test_megatron_weight_source.py.

suite	result
CPU -m "not vllm"	1,484 passed
CPU -m "vllm"	257 passed
GPU test_weight_sync.py	4 passed — NCCL TP=2, NCCL 1P1D PD, colocated IPC, sharded RDT
GPU test_delta_weight_sync_e2e.py[fsdp]	1 passed
GPU test_offload_kv_weight_sync.py	passed
ruff + black	clean

Not run — needs a real Megatron model: GPU megatron/test_megatron_weight_source.py and
the megatron half of the delta e2e. Apply run_train_gpu_ci and
run_train_megatron_gpu_ci.

The delta e2e needs ~/data/gsm8k/validation.parquet
(python examples/train/gsm8k/gsm8k_dataset.py --output_dir ~/data/gsm8k). Without it the
test fails in get_test_prompts after the sync has already succeeded, which reads like a
weight-sync failure and is not one.

Known gaps

• No GPU run has had a speculative-decoding drafter, so the proxy branch has never
  executed. A speculative_config parametrization on the delta or offload e2e would close it.
• No Megatron source on GPU — MegatronWeightSource is unit-tested against a fake
  bridge only. NCCLTrainerWeightTransferEngine._checked_iter now validates
  metadata() against iteration at runtime, so expect it to surface latent ordering bugs
  on the first real Megatron MoE run; that is the feature working.
• Single trainer rank in every GPU test — is_sender=False paths (IPC handle
  all-gather, delta's result gather) are unexercised.
• PP-collective cache at pp > 1 — megatron-bridge's broadcast_from_pp_rank takes a
  cache_key, and the old extractor rebuilt tasks each sync "so mapping objects start with
  clean PP-collective caches". A single long whole-model export may accumulate per-mapping
  cached tensors. Needs a peak-memory measurement.

Docs

.claude/docs/weight_sync.md rewritten; architecture.md, backends/fsdp.md and
docs/content/docs/getting-started/inference_architecture.mdx updated.

---

Note

High Risk
This rewires the core training→inference weight path (NCCL rank mapping, packed IPC, delta reload, RDT fan-out); mistakes tend to hang or silently stale weights rather than fail fast.

Overview
Replaces SkyRL’s custom weight transfer strategy/sender/extractor stack with vLLM 0.28’s trainer-send model: each training worker builds a WeightSource and a TrainerWeightTransferEngine whose send_weights() drives the native /init_weight_transfer_engine → /start_weight_update → /update_weights → /finish_weight_update round trip via a blocking SkyrlWeightSyncClient.

Removed the old broadcast/IPC/delta/RDT strategy modules, chunk/extractor types, vendored nccl_trainer_send, layerwise /collective_rpc weight paths, and most of RemoteInferenceClient’s weight-sync APIs. Added sources.py, trainer_engines.build_trainer_engine, control_plane.py (NCCL/RDT per-server init rewrites), receive-side skyrl_nccl / skyrl_ipc engines with drafter reload, delta_trainer, and a GPUModelRunner registry patch.

NewInferenceWorkerWrap now only handles fetch_weights (delta, pre-pause) and direct CuMemAllocator sleep/wake for KV offload during sync; receive lifecycle and chunk loading live in vLLM engines. Delta owns layerwise reload on the receive engine and re-streams the checkpoint for spec-decode drafter reload (behavior change vs. prior bare load_weights). Docs and public inference architecture pages are updated to match the four backends (nccl, ipc, delta, sharded_rdt) and get_vllm_receive_backend on the driver.

^{Reviewed by Cursor Bugbot for commit 055bc62. Bugbot is set up for automated code reviews on this repo. Configure here.}

[gemini-code-assist]

Code Review

This pull request refactors the weight synchronization architecture to use vLLM's native trainer-send abstraction, introducing specialized trainer-side engines and weight sources for the nccl, ipc, delta, and sharded_rdt backends, while removing legacy worker-extension mixins. The review feedback identifies a critical runtime bug in _packed_buffer_size_bytes where itemsize is incorrectly accessed on a torch.dtype object, which should be replaced with element_size(). Additionally, it is recommended to add explicit validation to ensure the number of server URLs is perfectly divisible by the data parallel size in both nccl_init_payloads and rdt_init_payloads to prevent truncation errors during rank offset and replica calculations.

[gemini-code-assist]

If len(server_urls) is not perfectly divisible by data_parallel_size (e.g., due to a misconfiguration), the division len(server_urls) // dp will truncate.

In rdt_init_payloads, this can result in a server having a replica_rank (i.e., i // dp) that is greater than or equal to num_replicas, violating the rank boundary and causing out-of-bounds or mapping errors in RDT.

We should explicitly validate that len(server_urls) % dp == 0 before performing these calculations.

def rdt_init_payloads(
    init_info: Dict[str, Any],
    server_urls: Sequence[str],
    data_parallel_size: int,
) -> List[Dict[str, Any]]:
    """Per-server sharded-RDT init payloads, stamped with the deployment ordinal.

    Each deployment has its own self-contained parallel config, so its internal
    worker index restarts at 0 and would collide under the M:N block assignment.
    Stamping each server with its ordinal (``server_index // data_parallel_size``)
    as ``replica_rank``, plus the deployment count as ``num_replicas``, lets the
    engine offset its consumers into a globally distinct range.

    The ordinal divides by ``data_parallel_size`` because the DP servers of one
    deployment share a parallel config, so they must share one ``replica_rank``.
    """
    dp = max(1, int(data_parallel_size))
    num_servers = len(server_urls)
    if num_servers % dp != 0:
        raise ValueError(
            f"Number of servers ({num_servers}) must be divisible by data_parallel_size ({dp})."
        )
    num_replicas = max(1, num_servers // dp)
    return [{**init_info, "replica_rank": i // dp, "num_replicas": num_replicas} for i in range(num_servers)]

[cursor]

Cursor Bugbot has reviewed your changes using default effort and found 1 potential issue.

^{Reviewed by Cursor Bugbot for commit 1c62fb9. Configure here.}

[SumanthRH]

I had an agent review the deletions in this file. It looks like you've deleted a bunch of tests for code that is still present. Can we revert that? We should always keep tests for code that lives in skyrl even if there's something similar on the vLLM side for RDT.

Deleted class	Still-live code it covered
TestPpLocalOwnership (~120 lines)	rdt_send.py:961 held_names, _pp_local, _demoted, walk reorder
TestExpertOwnership (~75 lines)	EP stamping, "expert count not divisible by ep_size raises" (rdt_send.py:1053+)
TestHeldNamesComposition	ownership composition across pp × ep
TestStampedYieldValidation	held-name/tensor agreement checks
TestQkvIndexDeviceCtx	rdt_send._qkv_index_device_ctx (rdt_send.py:53)
TestShardedRdtVllmRegistration	rdt_vllm_register.ensure_registered()
TestLoraLoadRequest	weight_sync/base.py LoraLoadRequest

[SumanthRH]

Can you simple move these RDT specific tests to a separate file and not delete them

[SumanthRH]

This is not called anywhere?

[SumanthRH]

Deleting this means that SkyRL doesn't set the bucket size for packed transfer anymore. weight_transfer_threshold_cuda_ipc_GB becomes unused.

Can we pass this to vLLM in init info?

https://github.com/vllm-project/vllm/blob/8a728663c1c3eeace834a95f5654fa653cc1998c/vllm/distributed/weight_transfer/ipc_engine.py#L60

[SumanthRH]

Also, the current logic has a common weight extractor for all weight sync backends. Which means that weight_transfer_threshold_cuda_ipc_GB also sets the bucket size for NCCL.

So we should use the same field weight_transfer_threshold_cuda_ipc_GB to set the packed transfer size even in the case of NCCL: https://github.com/vllm-project/vllm/blob/8a728663c1c3eeace834a95f5654fa653cc1998c/vllm/distributed/weight_transfer/nccl_engine.py#L80

I'll work on renaming the config in SkyRL meanwhile to avoid confusion

[SumanthRH]

Remove these docstring edits for update_named_weights?

They are wrong

[SumanthRH]

One pattern I see is some weird docstring comments from the agent that references now deleted code.

I have added many guidelines in the .claude/ folder in SkyRL to specifically avoid such docstrings.

Can you point your agent to the guidelines in .claude/docs and specifically the guide here:
asking it to make edits?

SkyRL/.claude/docs/contributing.md

Lines 51 to 157 in c516f3a

	## Comments
	Comments should describe what the code is doing, not what instruction the user provided or what learnings-on-the-journey the agent stumbled into.
	Do:
	"Uses the start/update/finish lifecycle to enable chunked transfers. Per chunk, all tensors are packed into a single contiguous CUDA buffer (one dtype per chunk, guaranteed by the weight extractor) and one IPC handle is created for the packed buffer per rank."
	Don't:
	"Uses the start/update/finish lifecycle to enable chunked transfers. Per chunk, all tensors are packed into a single contiguous CUDA buffer (one dtype per chunk, guaranteed by the weight extractor) and one IPC handle is created for the packed buffer per rank. This avoids the one-handle-per-param ceiling of vLLM's default IPCWeightTransferEngine, which otherwise dominates latency for models with many small parameters."
	Do:
	```python
	def test_freeze_moe_router_two_level_wrap():
	"""
	Under Megatron's ``bf16=True`` path, chunks are wrapped as
	``DDP(Float16Module(GPTModel))``. This tests whether `freeze_moe_router`
	can handle 2 levels of wrapping.
	"""
	inner_model = _Model()
	```
	Don't:
	```python
	def test_freeze_moe_router_two_level_wrap():
	"""Regression: recursive unwrap handles DDP(Float16Module(model)).
	Under Megatron's ``bf16=True`` path, chunks are wrapped as
	``DDP(Float16Module(GPTModel))``. Single-level ``.module`` peel leaves
	``Float16Module`` (which lacks ``.decoder``) and the helper raises. The
	worker uses a ``while hasattr(..., "module")`` loop — this test inlines
	the same loop to guard against regressions.
	"""
	inner_model = _Model()
	```
	Do:
	```bash
	# SFT training with Megatron backend for Qwen2.5-1.5B-Instruct on a
	# tool-calling dataset (Salesforce/APIGen-MT-5k).
	#
	# This script runs supervised fine-tuning using the Megatron backend with
	# pure data parallelism (DP=4) on 4 GPUs.
	#
	# Usage:
	#
	# export DATA_DIR=$HOME/data/apigen-mt-5k-openai
	# export WANDB_API_KEY=<your_key_here>
	# bash examples/train/sft/run_sft_megatron_apigen_mt.sh num_epochs=1 num_steps=<num_steps>
	```
	Don't:
	```bash
	# SFT training with Megatron backend for Qwen2.5-1.5B-Instruct on a
	# tool-calling dataset (Salesforce/APIGen-MT-5k).
	#
	# This script runs supervised fine-tuning using the Megatron backend with
	# pure data parallelism (DP=4) on 4 GPUs. It exercises the tool-calling SFT
	# path: per-row ``tools`` schemas and ``system`` policy are threaded into
	# every ``apply_chat_template`` call, ``tool`` observation tokens are
	# masked to 0, and every assistant turn (including ``tool_calls``)
	# contributes to the loss via ``train_on_what=all_assistant_messages``.
	#
	# APIGen-MT-5k ships in ShareGPT format and contains 5000 rows. The
	# preprocessing step below converts it to OpenAI messages format and writes
	# a parquet shard the SFT trainer can load directly.
	#
	# Usage:
	# bash examples/train/sft/run_sft_megatron_apigen_mt.sh [extra overrides...]
	#
	# Example (default: 10-step smoke run on 4 GPUs):
	# bash examples/train/sft/run_sft_megatron_apigen_mt.sh
	#
	# Example (full epoch over the 5000 rows):
	# bash examples/train/sft/run_sft_megatron_apigen_mt.sh num_epochs=1 num_steps=null
	```
	## Error messages
	The same holds true for error messages:
	Do:
	```python
	if self._callback_handler.callbacks:
	raise NotImplementedError(
	"Callbacks are not yet supported by `FullyAsyncRayPPOTrainer`. "
	)
	```
	Don't:
	```python
	if self._callback_handler.callbacks:
	raise NotImplementedError(
	"Callbacks are not yet supported by `FullyAsyncRayPPOTrainer`. "
	"Track in a follow-up; the sync RayPPOTrainer and SFTTrainer do support them."
	)
	```

[SumanthRH]

Did a first pass! Take a look!

[SumanthRH]

Overall given the large set of changes, we should also run some e2e tests. It would be good to:

1. Run the Qwen 3 235B weight sync benchmark with RDT + Megatron + vLLM, repeating the same settings in the blog: https://vllm.ai/blog/2026-08-22-rdt-weight-transfer
2. Do a before and after comparison for Qwen 3 235B weight sync for the NCCL broadcast backend with like Megatron TP 8 EP 8 and vLLM TP 4. Ideally this is a real E2E training config with GSM8K for like 2 steps. This is to see if any of the weight export changes on the trainer broke something.