[RL] Support checkpoint-coordinate sparse NCCL weight updates (#53751)

Signed-off-by: aoshen02 <aoshen@inferact.ai>
Signed-off-by: Shuolei Wang <shuoleiwang123@gmail.com>
Co-authored-by: aoshen02 <aoshen@inferact.ai>
Co-authored-by: OpenAI Codex <codex@openai.com>

docs/training/weight_transfer/README.md

@@ -33,7 +33,7 @@ engine drives on your behalf:
 | ------- | --------- | -------- |
 | [NCCL](nccl.md) | NCCL broadcast | Separate GPUs for training and inference |
 | [IPC](ipc.md) | CUDA IPC handles | Colocated training and inference on same GPU |
-| [sparse_nccl](nccl.md#sparse-nccl) | NCCL broadcast | Sparse flat-index weight patches (TP=1/PP=1) |
+| [sparse_nccl](nccl.md#sparse-nccl) | NCCL broadcast | Checkpoint-coordinate sparse weight patches |
 | [sharded_rdt](sharded_rdt.md) | NIXL / Ray Direct Transport (pull-based) | Very large models where each worker needs only its own slice (MoE with expert parallelism) |
 
 ## Quickstart

docs/training/weight_transfer/nccl.md

@@ -132,17 +132,19 @@ pass — this is the invariant to test first.
 
 ## Sparse NCCL
 
-Sparse, flat-index weight patches use a separate backend,
-`WeightTransferConfig(backend="sparse_nccl")`. It shares only NCCL process-group
-initialization with the dense engine; patches are applied directly in place to
-existing parameters, with no layerwise reload. The current sparse MVP requires
-`TP=1` and `PP=1`, and assumes a single-rank trainer.
+Sparse, flat-index weight patches use
+`WeightTransferConfig(backend="sparse_nccl")`. Names, full shapes, and flat
+indices are interpreted in **checkpoint/Hugging Face coordinates**. Every
+inference rank receives the same checkpoint-global patch, then the model's native
+`load_weights()` maps it to rank-local TP/EP and packed runtime parameters.
 
-Sparse is a **delta** backend: each round's patches are a fresh set of deltas
-from the latest optimizer step, not a stable stream of the model's parameters. So
-the engine takes **no `WeightSource`** — passing one raises — and each round's
-patches go straight to `send_weights(patches)`. A round with no patches is a
-no-op.
+Tensor parallelism is supported and does not have to match the trainer's layout.
+Pipeline-parallel behavior is not covered by the current sparse NCCL GPU test.
+
+Sparse is a **delta** backend: each call supplies replacement patches rather
+than a stable stream of the model's parameters. The engine therefore takes no
+`WeightSource`; patches go straight to `send_weights(patches)`, and an empty
+patch list is a no-op.
 
 ```python
 from vllm.distributed.weight_transfer import (
@@ -154,6 +156,7 @@ from vllm.distributed.weight_transfer.sparse_nccl_engine import (
     SparseWeightPatch,
 )
 
+client = RayVLLMWeightSyncClient(llm)
 engine = WeightTransferTrainerFactory.trainer_init(
     init_info=SparseNCCLTrainerInitInfo(
         master_address=master_address,
@@ -161,7 +164,7 @@ engine = WeightTransferTrainerFactory.trainer_init(
         world_size=world_size,
         rank=0,
     ),
-    client=RayVLLMWeightSyncClient(llm),
+    client=client,
 )
 
 patches = [
@@ -175,21 +178,24 @@ patches = [
 engine.send_weights(patches)
 ```
 
-`SparseNCCLTrainerInitInfo` takes the same rendezvous fields as the dense
-backend and no packed params — sparse transfers are never packed.
+Each `send_weights(patches)` call owns a complete one-shot `start` / `update` / `finish` lifecycle.
 
-Patches are validated on the trainer *before* `start_weight_update`: int32
-indices, 1-D, and matching `indices`/`values` lengths, with `full_shape` set. The
-worker checks the same invariants when applying, but by then the broadcasts are
-under way, where a size mismatch wedges both sides instead of raising.
+RL infrastructure that owns the generic worker lifecycle can keep one logical update open across bounded chunks without adding trainer-side session state:
 
-A trainer can hold a dense engine and a sparse engine at once — a full resync
-through the dense engine, cheap sparse deltas in between — as
-[`rlhf_sparse_nccl.py`](../../../examples/rl/rlhf_sparse_nccl.py) demonstrates.
+```python
+client.start_weight_update()
+for patches in patch_chunks:
+    engine.send_weight_chunk(patches)
+client.finish_weight_update()
+```
+
+Sparse NCCL sends only `O(nnz)` indices and values. Checkpoint application still uses `O(N)` staging through the native loader, and the caller owns export/diff state and restart or reseed after a partial failure.
+
+[`rlhf_sparse_nccl.py`](../../../examples/rl/rlhf_sparse_nccl.py) demonstrates per-expert checkpoint patches with a Qwen3 MoE TP2/EP2 engine.
 
 ## Examples
 
 - [RLHF with NCCL weight syncing (`vllm serve`, HTTP)](../../../examples/rl/rlhf_http_nccl.py) - **Start here.** Trainer on one GPU, 2x tensor-parallel fp8 server on two others; HTTP control plane, NCCL data plane. Launches and tears down its own server
 - [RLHF with NCCL + FSDP2 and expert parallelism](../../../examples/rl/rlhf_nccl_fsdp_ep.py) - Multi-rank trainer: every FSDP rank builds an engine and joins the `full_tensor()` gather while only rank 0 touches the wire
-- [RLHF with sparse NCCL weight syncing (offline, Ray)](../../../examples/rl/rlhf_sparse_nccl.py) - Dense-vs-sparse equivalence demo with a real model on a 2-GPU trainer/inference setup; sparse patches use `backend="sparse_nccl"` and currently require `TP=1` and `PP=1`
+- [RLHF with sparse NCCL weight syncing (offline, Ray)](../../../examples/rl/rlhf_sparse_nccl.py) - Qwen3 MoE per-expert checkpoint updates with one trainer GPU and two TP2/EP2 inference GPUs
 - [RLHF with async weight syncing (offline, Ray)](../../../examples/rl/rlhf_async_new_apis.py) - Async generation with mid-flight pause, weight sync, resume, and validation against a fresh model; uses an in-process `AsyncLLMEngine` with `RayVLLMWeightSyncClient`

examples/rl/rlhf_sparse_nccl.py

@@ -1,64 +1,45 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""
-Demonstrates dense-vs-sparse NCCL weight syncing with a real model.
+"""Demonstrate checkpoint-coordinate sparse NCCL updates with expert parallelism.
 
-This example mirrors the validation story used for the sparse NCCL MVP:
-both the dense update path and the sparse patch path start from the same real
-checkpoint and apply the same deterministic trainer-side patch. The script then
-checks that greedy 1-token outputs match between the dense and sparse vLLM
-engines after the update.
+The trainer and vLLM start from the same Qwen3 MoE checkpoint. The trainer
+modifies rows from two global experts, converts its fused expert storage back to
+per-expert checkpoint coordinates, and sends both patches through one
+``send_weights()`` lifecycle. Every vLLM rank receives both patches; the native
+loader applies its local expert and skips the foreign expert.
 
-The example performs the following steps:
-* Load a training model on one GPU via a Ray actor.
-* Launch a vLLM engine with the same real model on a second GPU.
-* Verify trainer vs vLLM baseline agreement before any update.
-* Apply a deterministic patch to ``model.embed_tokens.weight`` on the trainer.
-* Run a dense NCCL update into a fresh vLLM engine and collect post-update
-  outputs.
-* Reset the trainer back to the baseline checkpoint.
-* Apply the same deterministic patch again.
-* Run a sparse NCCL update into another fresh vLLM engine and collect
-  post-update outputs.
-* Compare dense vs sparse baseline outputs, dense vs sparse post-update
-  outputs, estimated payload sizes, and trainer-side send times.
-
-Current sparse weight transfer MVP limitations:
-* ``TP=1`` and ``PP=1`` only
-* sparse updates use runtime/kernel-format parameter names
-* sparse updates are not composable with checkpoint-format or packed updates
-
-This example assumes a single-node cluster with two GPUs.
+This example uses three GPUs on one node: one for the full BF16 Hugging Face
+trainer model and two for a TP2/EP2 vLLM inference engine. Unspecified checkpoint
+elements keep their initialized values, so sparse updates require a known shared
+baseline. ``SPARSE_NCCL_MODEL`` may point to a compatible local Qwen3 MoE
+checkpoint.
 """
 
-import hashlib
 import os
-import time
-from collections.abc import Sequence
+from contextlib import suppress
 
 import ray
 import torch
 from ray.util.placement_group import placement_group
 from ray.util.scheduling_strategies import PlacementGroupSchedulingStrategy
-from transformers import AutoModelForCausalLM, AutoTokenizer
+from transformers import AutoModelForCausalLM
 
 from vllm import LLM, SamplingParams
 from vllm.config import WeightTransferConfig
 from vllm.distributed.weight_transfer import (
-    ModuleSource,
     RayVLLMWeightSyncClient,
     WeightTransferTrainerFactory,
 )
-from vllm.distributed.weight_transfer.nccl_engine import NCCLTrainerInitInfo
 from vllm.distributed.weight_transfer.sparse_nccl_engine import (
     SparseNCCLTrainerInitInfo,
     SparseWeightPatch,
 )
 from vllm.utils.network_utils import get_ip, get_open_port
 
-MODEL_NAME = "Qwen/Qwen2.5-0.5B-Instruct"
-PATCHED_PARAM_NAME = "model.embed_tokens.weight"
-MAX_PATCH_ROWS = 32
+MODEL_NAME = os.environ.get("SPARSE_NCCL_MODEL", "Qwen/Qwen3-30B-A3B")
+INFERENCE_TP_SIZE = 2
+PATCHED_LAYER = 0
+PATCHED_ROWS = 2
 PROMPTS = [
     "Hello, my name is",
     "The president of the United States is",
@@ -72,414 +53,201 @@ class MyLLM(LLM):
     """Configure the vLLM worker for Ray placement group execution."""
 
     def __init__(self, *args, **kwargs):
-        os.environ["VLLM_RAY_BUNDLE_INDICES"] = "0"
+        os.environ["VLLM_RAY_BUNDLE_INDICES"] = ",".join(
+            str(index) for index in range(INFERENCE_TP_SIZE)
+        )
         super().__init__(*args, **kwargs)
 
 
 @ray.remote(num_gpus=1)
 class TrainModel:
-    """Ray actor that owns the trainer-side model and deterministic patch state."""
+    """Own the trainer model and sparse NCCL sender on one GPU."""
 
     def __init__(self, model_name: str):
-        self.model_name = model_name
-        self.tokenizer = AutoTokenizer.from_pretrained(model_name)
-        if self.tokenizer.pad_token_id is None:
-            self.tokenizer.pad_token = self.tokenizer.eos_token
-
-        self.model = None
-        self.patched_param = None
-        self.pending_sparse_patches: list[SparseWeightPatch] | None = None
-        self.dense_engine = None
-        self.sparse_engine = None
-        self.master_address = get_ip()
-        self.port = get_open_port()
-        self.reset_model()
-
-    def reset_model(self) -> None:
         self.model = AutoModelForCausalLM.from_pretrained(
-            self.model_name,
+            model_name,
             torch_dtype=torch.bfloat16,
         ).to("cuda:0")
         self.model.eval()
 
+        config = self.model.config
+        if config.model_type != "qwen3_moe" or config.num_experts % INFERENCE_TP_SIZE:
+            raise RuntimeError(
+                "This recipe requires a Qwen3 MoE model whose experts divide "
+                f"evenly across TP{INFERENCE_TP_SIZE}"
+            )
+        self.expert_ids = (0, config.num_experts // 2)
+        self.expert_intermediate_size = config.moe_intermediate_size
+        fused_name = f"model.layers.{PATCHED_LAYER}.mlp.experts.gate_up_proj"
         try:
-            self.patched_param = self.model.get_parameter(PATCHED_PARAM_NAME)
+            self.gate_up_proj = self.model.get_parameter(fused_name)
         except AttributeError as exc:
             raise RuntimeError(
-                f"Expected trainer model to expose `{PATCHED_PARAM_NAME}`"
+                f"Expected trainer model to expose `{fused_name}`"
             ) from exc
 
-        self.pending_sparse_patches = None
-
-    def create_rendezvous(self) -> tuple[str, int]:
-        self.port = get_open_port()
-        return self.master_address, self.port
-
-    def init_dense_engine(self, world_size: int, llm_handle) -> None:
-        """Build the dense NCCL trainer engine"""
-        self.dense_engine = WeightTransferTrainerFactory.trainer_init(
-            init_info=NCCLTrainerInitInfo(
-                master_address=self.master_address,
-                master_port=self.port,
-                world_size=world_size,
-                rank=0,  # single-GPU trainer is the sole (sender) rank
-                packed=True,
-            ),
-            client=RayVLLMWeightSyncClient(llm_handle),
-            source=ModuleSource(self.model),
+        expected_shape = (
+            config.num_experts,
+            2 * self.expert_intermediate_size,
+            config.hidden_size,
         )
-
-    def shutdown_engines(self) -> None:
-        """Drop the NCCL communicators before the vLLM engine they point at
-        goes away, so the next phase rendezvouses from a clean slate."""
-        for attr in ("dense_engine", "sparse_engine"):
-            engine = getattr(self, attr)
-            if engine is not None:
-                engine.shutdown()
-                setattr(self, attr, None)
-
-    def dense_payload_bytes(self) -> int:
-        return sum(
-            param.numel() * param.element_size() for param in self.model.parameters()
-        )
-
-    @torch.inference_mode()
-    def generate(
-        self,
-        prompts: Sequence[str],
-        max_new_tokens: int = 1,
-    ) -> list[dict[str, object]]:
-        generations = []
-        for prompt in prompts:
-            model_inputs = self.tokenizer(prompt, return_tensors="pt").to("cuda:0")
-            output = self.model.generate(
-                **model_inputs,
-                max_new_tokens=max_new_tokens,
-                do_sample=False,
-                pad_token_id=self.tokenizer.pad_token_id,
+        if self.gate_up_proj.shape != expected_shape:
+            raise RuntimeError(
+                f"Unexpected fused expert shape: {self.gate_up_proj.shape} "
+                f"!= {expected_shape}"
             )
-            new_token_ids = output[0, model_inputs["input_ids"].shape[1] :].tolist()
-            generations.append(
-                {
-                    "token_ids": new_token_ids,
-                    "text": self.tokenizer.decode(
-                        new_token_ids,
-                        skip_special_tokens=False,
-                    ),
-                }
-            )
-        return generations
 
-    def prepare_sparse_patch(
-        self,
-        prompts: Sequence[str],
-        max_patch_rows: int = MAX_PATCH_ROWS,
-    ) -> tuple[list[int], str, int]:
-        selected_token_ids: list[int] = []
-        special_ids = set(self.tokenizer.all_special_ids)
-        for prompt in prompts:
-            token_ids = self.tokenizer(prompt, add_special_tokens=False)["input_ids"]
-            for token_id in token_ids:
-                if token_id in special_ids or token_id in selected_token_ids:
-                    continue
-                selected_token_ids.append(token_id)
-                if len(selected_token_ids) == max_patch_rows:
-                    break
-            if len(selected_token_ids) == max_patch_rows:
-                break
-
-        if not selected_token_ids:
-            raise ValueError("Could not derive any non-special token IDs to patch")
-
-        vocab_size = self.patched_param.shape[0]
-        next_token_id = selected_token_ids[-1]
-        while len(selected_token_ids) < max_patch_rows:
-            next_token_id = (next_token_id + 1) % vocab_size
-            if next_token_id in special_ids or next_token_id in selected_token_ids:
-                continue
-            selected_token_ids.append(next_token_id)
-
-        row_ids = torch.tensor(
-            selected_token_ids,
-            device=self.patched_param.device,
-            dtype=torch.long,
-        )
-        hidden_size = self.patched_param.shape[1]
-        column_offsets = torch.arange(
-            hidden_size,
-            device=self.patched_param.device,
-            dtype=torch.long,
-        )
-
-        with torch.no_grad():
-            # Rotate the selected embedding rows instead of zeroing them so the
-            # patch remains deterministic while avoiding a degenerate collapse
-            # to the same special token after the update.
-            replacement_rows = self.patched_param[row_ids].roll(shifts=1, dims=0)
-            self.patched_param[row_ids] = replacement_rows
-
-        flat_indices = (
-            row_ids.unsqueeze(1).mul(hidden_size).add(column_offsets).reshape(-1)
-        )
-        flat_values = self.patched_param[row_ids].detach().reshape(-1).contiguous()
-        self.pending_sparse_patches = [
-            SparseWeightPatch(
-                name=PATCHED_PARAM_NAME,
-                indices=flat_indices.to(torch.int32),
-                values=flat_values,
-                full_shape=tuple(self.patched_param.shape),
-            )
-        ]
-        patch_digest = hashlib.sha256(
-            self.pending_sparse_patches[0].indices.cpu().numpy().tobytes()
-            + self.pending_sparse_patches[0].values.float().cpu().numpy().tobytes()
-        ).hexdigest()
-
-        sparse_payload_bytes = (
-            flat_indices.numel() * torch.tensor([], dtype=torch.int32).element_size()
-            + flat_values.numel() * flat_values.element_size()
-        )
-        return selected_token_ids, patch_digest, sparse_payload_bytes
-
-    def send_dense_weights(self) -> float:
-        """One call drives start/update/finish on the inference side
-        concurrently with the NCCL broadcast."""
-        if self.dense_engine is None:
-            raise RuntimeError("Dense engine is not initialized")
-
-        start = time.perf_counter()
-        self.dense_engine.send_weights()
-        torch.accelerator.synchronize()
-        return (time.perf_counter() - start) * 1000.0
+        self.master_address = get_ip()
+        self.master_port = get_open_port()
+        self.engine = None
 
     def init_sparse_engine(self, world_size: int, llm_handle) -> None:
-        """Build the sparse trainer engine"""
-        self.sparse_engine = WeightTransferTrainerFactory.trainer_init(
+        self.engine = WeightTransferTrainerFactory.trainer_init(
             init_info=SparseNCCLTrainerInitInfo(
                 master_address=self.master_address,
-                master_port=self.port,
+                master_port=self.master_port,
                 world_size=world_size,
-                rank=0,  # single-GPU trainer is the sole (sender) rank
+                rank=0,
             ),
             client=RayVLLMWeightSyncClient(llm_handle),
         )
 
-    def send_pending_sparse_patch(self) -> float:
-        """One call drives start/update/finish on the inference side
-        concurrently with the patch broadcasts. Sparse deltas differ each round,
-        so the patches are passed to `send_weights` rather than fixed at init."""
-        if self.sparse_engine is None:
-            raise RuntimeError("Sparse engine is not initialized")
-        if self.pending_sparse_patches is None:
-            raise RuntimeError("Sparse patch has not been prepared")
+    @torch.no_grad()
+    def patch_and_send(self) -> tuple[list[str], int]:
+        if self.engine is None:
+            raise RuntimeError("Sparse NCCL engine is not initialized")
 
-        start = time.perf_counter()
-        self.sparse_engine.send_weights(self.pending_sparse_patches)
-        torch.accelerator.synchronize()
-        self.pending_sparse_patches = None
-        return (time.perf_counter() - start) * 1000.0
+        patches = []
+        for expert_id in self.expert_ids:
+            checkpoint_weight = self.gate_up_proj[
+                expert_id, : self.expert_intermediate_size
+            ]
+            original_rows = checkpoint_weight[:PATCHED_ROWS].clone()
+            replacement_rows = original_rows.flip(0)
+            if torch.equal(original_rows, replacement_rows):
+                raise RuntimeError(f"Expert {expert_id} patch would be a no-op")
+            checkpoint_weight[:PATCHED_ROWS] = replacement_rows
+
+            hidden_size = checkpoint_weight.shape[1]
+            flat_indices = torch.arange(
+                PATCHED_ROWS * hidden_size,
+                device=checkpoint_weight.device,
+                dtype=torch.int32,
+            )
+            patches.append(
+                SparseWeightPatch(
+                    name=(
+                        f"model.layers.{PATCHED_LAYER}.mlp.experts."
+                        f"{expert_id}.gate_proj.weight"
+                    ),
+                    full_shape=tuple(checkpoint_weight.shape),
+                    indices=flat_indices,
+                    values=replacement_rows.reshape(-1).contiguous(),
+                )
+            )
+
+        self.engine.send_weights(patches)
+        return [patch.name for patch in patches], sum(
+            patch.indices.numel() for patch in patches
+        )
+
+    def shutdown_engine(self) -> None:
+        if self.engine is not None:
+            self.engine.shutdown()
+            self.engine = None
 
 
-def launch_llm(
-    scheduling_inference: PlacementGroupSchedulingStrategy,
-    backend: str = "nccl",
-):
+def launch_llm(scheduling_strategy: PlacementGroupSchedulingStrategy):
     return ray.remote(
         num_cpus=0,
         num_gpus=0,
-        scheduling_strategy=scheduling_inference,
+        scheduling_strategy=scheduling_strategy,
     )(MyLLM).remote(
         model=MODEL_NAME,
         enforce_eager=True,
-        tensor_parallel_size=1,
+        tensor_parallel_size=INFERENCE_TP_SIZE,
+        enable_expert_parallel=True,
+        expert_placement_strategy="linear",
+        moe_backend="triton",
         distributed_executor_backend="ray",
         gpu_memory_utilization=0.7,
-        weight_transfer_config=WeightTransferConfig(backend=backend),
+        enable_prefix_caching=False,
+        weight_transfer_config=WeightTransferConfig(backend="sparse_nccl"),
     )
 
 
-def collect_vllm_generations(llm_handle) -> list[dict[str, object]]:
+def generate(llm_handle) -> list[dict[str, object]]:
     outputs = ray.get(llm_handle.generate.remote(PROMPTS, SAMPLING_PARAMS))
-    generations = []
-    for output in outputs:
-        generations.append(
-            {
-                "token_ids": output.outputs[0].token_ids,
-                "text": output.outputs[0].text,
-            }
-        )
-    return generations
-
-
-def token_sequences_match(
-    left: Sequence[dict[str, object]],
-    right: Sequence[dict[str, object]],
-) -> bool:
-    return [item["token_ids"] for item in left] == [item["token_ids"] for item in right]
-
-
-def print_generations(label: str, prompts: Sequence[str], generations) -> None:
-    print(f"\n{label}")
-    print("-" * 50)
-    for prompt, generation in zip(prompts, generations):
-        print(f"Prompt: {prompt!r}")
-        print(f"Token IDs: {generation['token_ids']}")
-        print(f"Text: {generation['text']!r}")
-        print("-" * 50)
-
-
-def run_dense_phase(
-    train_model,
-    scheduling_inference: PlacementGroupSchedulingStrategy,
-) -> dict[str, object]:
-    ray.get(train_model.reset_model.remote())
-    llm = launch_llm(scheduling_inference, backend="nccl")
-    try:
-        dense_before = collect_vllm_generations(llm)
-
-        ray.get(llm.sleep.remote(level=0))
-        ray.get(train_model.create_rendezvous.remote())
-        world_size = ray.get(llm.get_world_size.remote()) + 1
-        ray.get(train_model.init_dense_engine.remote(world_size, llm))
-
-        dense_payload_bytes = ray.get(train_model.dense_payload_bytes.remote())
-        selected_token_ids, patch_digest, _ = ray.get(
-            train_model.prepare_sparse_patch.remote(PROMPTS)
-        )
-
-        # One call drives start/update/finish + the NCCL broadcast.
-        dense_send_ms = ray.get(train_model.send_dense_weights.remote())
-        ray.get(llm.wake_up.remote(tags=["scheduling"]))
-
-        dense_after = collect_vllm_generations(llm)
-
-        return {
-            "dense_before": dense_before,
-            "dense_after": dense_after,
-            "selected_token_ids": selected_token_ids,
-            "patch_digest": patch_digest,
-            "dense_payload_bytes": dense_payload_bytes,
-            "dense_send_ms": dense_send_ms,
+    return [
+        {
+            "token_ids": output.outputs[0].token_ids,
+            "text": output.outputs[0].text,
         }
-    finally:
-        ray.get(train_model.shutdown_engines.remote())
-        ray.kill(llm)
+        for output in outputs
+    ]
 
 
-def run_sparse_phase(
-    train_model,
-    scheduling_inference: PlacementGroupSchedulingStrategy,
-) -> dict[str, object]:
-    ray.get(train_model.reset_model.remote())
-    llm = launch_llm(scheduling_inference, backend="sparse_nccl")
+def print_generations(label: str, generations: list[dict[str, object]]) -> None:
+    print(f"\n{label}")
+    for prompt, generation in zip(PROMPTS, generations):
+        print(
+            f"  {prompt!r} -> {generation['text']!r} "
+            f"(token_ids={generation['token_ids']})"
+        )
+
+
+def main() -> None:
+    ray.init()
+    train_model = None
+    pg_inference = None
+    llm = None
     try:
-        sparse_before = collect_vllm_generations(llm)
+        train_model = TrainModel.remote(MODEL_NAME)
+        pg_inference = placement_group(
+            [{"GPU": 1, "CPU": 0}] * INFERENCE_TP_SIZE,
+            strategy="STRICT_PACK",
+        )
+        ray.get(pg_inference.ready())
+        scheduling_strategy = PlacementGroupSchedulingStrategy(
+            placement_group=pg_inference,
+            placement_group_capture_child_tasks=True,
+            placement_group_bundle_index=0,
+        )
+        llm = launch_llm(scheduling_strategy)
+
+        before = generate(llm)
+        print_generations("BEFORE sparse update", before)
 
         ray.get(llm.sleep.remote(level=0))
-        ray.get(train_model.create_rendezvous.remote())
         world_size = ray.get(llm.get_world_size.remote()) + 1
         ray.get(train_model.init_sparse_engine.remote(world_size, llm))
-
-        selected_token_ids, patch_digest, sparse_payload_bytes = ray.get(
-            train_model.prepare_sparse_patch.remote(PROMPTS)
-        )
-        sparse_send_ms = ray.get(train_model.send_pending_sparse_patch.remote())
+        patched_names, num_updates = ray.get(train_model.patch_and_send.remote())
         ray.get(llm.wake_up.remote(tags=["scheduling"]))
 
-        sparse_after = collect_vllm_generations(llm)
-
-        return {
-            "sparse_before": sparse_before,
-            "sparse_after": sparse_after,
-            "selected_token_ids": selected_token_ids,
-            "patch_digest": patch_digest,
-            "sparse_payload_bytes": sparse_payload_bytes,
-            "sparse_send_ms": sparse_send_ms,
-        }
+        after = generate(llm)
+        print_generations("AFTER sparse update", after)
+        print(f"patched_checkpoint_names={patched_names}")
+        print(f"num_sparse_values={num_updates}")
+        outputs_changed = any(
+            old["token_ids"] != new["token_ids"]
+            for old, new in zip(before, after, strict=True)
+        )
+        print(f"outputs_changed={outputs_changed}")
     finally:
-        ray.get(train_model.shutdown_engines.remote())
-        ray.kill(llm)
+        if train_model is not None:
+            with suppress(Exception):
+                ray.get(train_model.shutdown_engine.remote())
+        if llm is not None:
+            with suppress(Exception):
+                ray.kill(llm)
+        if train_model is not None:
+            with suppress(Exception):
+                ray.kill(train_model)
+        if pg_inference is not None:
+            with suppress(Exception):
+                ray.util.remove_placement_group(pg_inference)
+        ray.shutdown()
 
 
-ray.init()
-
-try:
-    train_model = TrainModel.remote(MODEL_NAME)
-
-    pg_inference = placement_group([{"GPU": 1, "CPU": 0}])
-    ray.get(pg_inference.ready())
-    scheduling_inference = PlacementGroupSchedulingStrategy(
-        placement_group=pg_inference,
-        placement_group_capture_child_tasks=True,
-        placement_group_bundle_index=0,
-    )
-
-    dense_results = run_dense_phase(train_model, scheduling_inference)
-    sparse_results = run_sparse_phase(train_model, scheduling_inference)
-
-    baseline_equal = token_sequences_match(
-        dense_results["dense_before"],
-        sparse_results["sparse_before"],
-    )
-    patch_selection_equal = (
-        dense_results["selected_token_ids"] == sparse_results["selected_token_ids"]
-    )
-    patch_digest_equal = dense_results["patch_digest"] == sparse_results["patch_digest"]
-    after_equal = token_sequences_match(
-        dense_results["dense_after"],
-        sparse_results["sparse_after"],
-    )
-    any_output_changed = any(
-        before["token_ids"] != after["token_ids"]
-        for before, after in zip(
-            dense_results["dense_before"],
-            dense_results["dense_after"],
-        )
-    )
-    dense_payload_mb = dense_results["dense_payload_bytes"] / (1024 * 1024)
-    sparse_payload_mb = sparse_results["sparse_payload_bytes"] / (1024 * 1024)
-
-    print_generations(
-        "Dense baseline outputs",
-        PROMPTS,
-        dense_results["dense_before"],
-    )
-    print_generations(
-        "Sparse baseline outputs", PROMPTS, sparse_results["sparse_before"]
-    )
-    print_generations(
-        "Dense outputs after update", PROMPTS, dense_results["dense_after"]
-    )
-    print_generations(
-        "Sparse outputs after update",
-        PROMPTS,
-        sparse_results["sparse_after"],
-    )
-
-    print(f"patched_token_ids = {dense_results['selected_token_ids']}")
-    print(f"patch_selection_equal = {patch_selection_equal}")
-    print(f"dense_patch_digest = {dense_results['patch_digest']}")
-    print(f"sparse_patch_digest = {sparse_results['patch_digest']}")
-    print(f"patch_digest_equal = {patch_digest_equal}")
-    print(f"baseline_equal = {baseline_equal}")
-    print(f"after_equal = {after_equal}")
-    print(f"any_output_changed = {any_output_changed}")
-    print(f"dense_payload_mb = {dense_payload_mb:.2f}")
-    print(f"sparse_payload_mb = {sparse_payload_mb:.2f}")
-    print(f"dense_send_ms = {dense_results['dense_send_ms']:.2f}")
-    print(f"sparse_send_ms = {sparse_results['sparse_send_ms']:.2f}")
-
-    if not baseline_equal:
-        raise RuntimeError(
-            "Dense and sparse phases did not start from the same baseline"
-        )
-    if not patch_selection_equal:
-        raise RuntimeError("Dense and sparse phases used different sparse patches")
-    if not patch_digest_equal:
-        raise RuntimeError("Dense and sparse phases produced different patch values")
-    if not after_equal:
-        raise RuntimeError("Dense and sparse updates produced different outputs")
-    if not any_output_changed:
-        raise RuntimeError("Patch did not change the observed outputs")
-finally:
-    ray.shutdown()
+if __name__ == "__main__":
+    main()

tests/distributed/test_weight_transfer.py

@@ -9,6 +9,7 @@ Integration tests for NCCL and IPC weight transfer between processes using Ray.
 import pickle
 import threading
 import time
+from contextlib import nullcontext
 from unittest.mock import MagicMock
 
 import pybase64 as base64
@@ -313,103 +314,6 @@ class TestEngineRegistry:
             )
 
 
-# --- Unit Tests: Sparse patch application (CPU) ---
-
-
-class TestSparseNCCLPatchApplication:
-    """Test SparseNCCLWeightTransferEngine._apply_patch on a real param."""
-
-    def _make_engine(self, model):
-        config = WeightTransferConfig(backend="sparse_nccl")
-        return SparseNCCLWeightTransferEngine(
-            config, create_mock_vllm_config(), torch.device("cpu"), model
-        )
-
-    def _make_model(self, numel: int = 8):
-        model = torch.nn.Module()
-        model.register_parameter(
-            "w", torch.nn.Parameter(torch.zeros(numel), requires_grad=False)
-        )
-
-        def get_parameter(name):
-            assert name == "w"
-            return model.w
-
-        model.get_parameter = get_parameter
-        return model
-
-    def test_apply_patch_updates_only_selected_entries(self):
-        model = self._make_model(8)
-        engine = self._make_engine(model)
-        engine._apply_patch(
-            SparseWeightPatch(
-                name="w",
-                indices=torch.tensor([1, 3], dtype=torch.int32),
-                values=torch.tensor([5.0, 7.0], dtype=torch.float32),
-            )
-        )
-        expected = torch.zeros(8)
-        expected[1] = 5.0
-        expected[3] = 7.0
-        assert torch.equal(model.w.data, expected)
-
-    def test_apply_patch_rejects_mismatched_lengths(self):
-        model = self._make_model(8)
-        engine = self._make_engine(model)
-        with pytest.raises(ValueError, match="matching lengths"):
-            engine._apply_patch(
-                SparseWeightPatch(
-                    name="w",
-                    indices=torch.tensor([1, 3], dtype=torch.int32),
-                    values=torch.tensor([5.0], dtype=torch.float32),
-                )
-            )
-
-    def test_apply_patch_rejects_non_int32_indices(self):
-        model = self._make_model(8)
-        engine = self._make_engine(model)
-        with pytest.raises(ValueError, match="int32 indices"):
-            engine._apply_patch(
-                SparseWeightPatch(
-                    name="w",
-                    indices=torch.tensor([1], dtype=torch.int64),
-                    values=torch.tensor([5.0], dtype=torch.float32),
-                )
-            )
-
-    def test_apply_patch_rejects_dtype_mismatch(self):
-        model = self._make_model(8)
-        engine = self._make_engine(model)
-        with pytest.raises(ValueError, match="does not match"):
-            engine._apply_patch(
-                SparseWeightPatch(
-                    name="w",
-                    indices=torch.tensor([1], dtype=torch.int32),
-                    values=torch.tensor([5.0], dtype=torch.bfloat16),
-                )
-            )
-
-    def test_apply_patch_rejects_non_contiguous_param(self):
-        model = torch.nn.Module()
-        model.register_parameter(
-            "w",
-            torch.nn.Parameter(
-                torch.arange(12, dtype=torch.float32).view(3, 4).t(),
-                requires_grad=False,
-            ),
-        )
-        model.get_parameter = lambda name: model.w
-        engine = self._make_engine(model)
-        with pytest.raises(NotImplementedError, match="contiguous params"):
-            engine._apply_patch(
-                SparseWeightPatch(
-                    name="w",
-                    indices=torch.tensor([1], dtype=torch.int32),
-                    values=torch.tensor([1.0], dtype=torch.float32),
-                )
-            )
-
-
 # --- Test receive_weights without init raises ---
 
 
@@ -628,162 +532,135 @@ def test_nccl_weight_transfer_between_processes():
     )
 
 
-@ray.remote(num_gpus=1)
-def trainer_broadcast_sparse_tensor(
-    master_address: str,
-    master_port: int,
-    world_size: int,
-) -> bool:
-    """Trainer task that broadcasts sparse patches via the trainer engine.
+def test_sparse_nccl_checkpoint_chunks_to_ep_local_experts_cpu(monkeypatch):
+    """Replay global expert patches through two EP-local loaders on CPU."""
 
-    The worker task drives its own init/receive directly (it is not an RPC
-    endpoint), so the engine gets a no-op control-plane client; the NCCL
-    rendezvous and the patch broadcasts are the real thing.
-    """
-    import torch
-
-    device = _set_ray_assigned_device()
-
-    from vllm.distributed.weight_transfer import WeightTransferTrainerFactory
-    from vllm.distributed.weight_transfer.sparse_nccl_engine import (
-        SparseNCCLTrainerInitInfo,
-        SparseWeightPatch,
-    )
-
-    class NoopClient:
-        def init_weight_transfer_engine(self, init_info):
-            pass
+    class CaptureClient:
+        def __init__(self):
+            self.order: list[str] = []
+            self.update_infos: list[dict] = []
 
         def start_weight_update(self):
-            pass
+            self.order.append("start")
 
         def update_weights(self, update_info):
-            pass
+            self.order.append("update")
+            self.update_infos.append(update_info)
 
         def finish_weight_update(self):
-            pass
+            self.order.append("finish")
 
-    patch = SparseWeightPatch(
-        name="test.weight",
-        indices=torch.tensor([1, 7, 25], dtype=torch.int32, device=device),
-        values=torch.tensor([10.0, 20.0, 30.0], dtype=torch.float32, device=device),
-        full_shape=(10, 10),
-    )
-    engine = WeightTransferTrainerFactory.trainer_init(
-        init_info=SparseNCCLTrainerInitInfo(
-            master_address=master_address,
-            master_port=master_port,
-            world_size=world_size,
-            rank=0,
-        ),
-        client=NoopClient(),
-    )
-    engine.send_weights([patch])
-    torch.accelerator.synchronize()
-    engine.shutdown()
-    return True
+    monkeypatch.setattr(torch.cuda, "current_stream", lambda *_, **__: None)
+    monkeypatch.setattr(torch.cuda, "is_available", lambda: False)
+    monkeypatch.setattr(torch.accelerator, "device_index", lambda *_: nullcontext())
 
-
-@ray.remote(num_gpus=1)
-def inference_receive_sparse_tensor(
-    master_address: str,
-    master_port: int,
-    world_size: int,
-) -> dict:
-    """Inference task that receives sparse patches via the sparse engine."""
-    from unittest.mock import MagicMock
-
-    import torch
-
-    device = _set_ray_assigned_device()
-
-    from vllm.config.parallel import ParallelConfig
-    from vllm.config.weight_transfer import WeightTransferConfig
-    from vllm.distributed.weight_transfer.sparse_nccl_engine import (
-        SparseNCCLWeightTransferEngine,
-        SparseNCCLWeightTransferUpdateInfo,
+    client = CaptureClient()
+    sender = SparseNCCLTrainerWeightTransferEngine(client=client)
+    sender.model_update_group = MagicMock()
+    sender.model_update_group.device = torch.device("cpu")
+    wire_payloads = []
+    sender.model_update_group.broadcast.side_effect = lambda tensor, **_: (
+        wire_payloads.append(tensor.clone())
     )
 
-    config = WeightTransferConfig(backend="sparse_nccl")
-    vllm_config = MagicMock()
-    parallel_config = MagicMock(spec=ParallelConfig)
-    parallel_config.rank = 0
-    parallel_config.world_size = 1
-    parallel_config.data_parallel_rank = 0
-    parallel_config.data_parallel_index = 0
-    vllm_config.parallel_config = parallel_config
-    vllm_config.model_config = MagicMock()
-
-    # Real module holding the target parameter the patch will modify.
-    model = torch.nn.Module()
-    model.register_parameter(
-        "w", torch.nn.Parameter(torch.zeros(30, device="cuda"), requires_grad=False)
-    )
-    model.get_parameter = lambda name: model.w
-
-    update_info = SparseNCCLWeightTransferUpdateInfo(
-        names=["w"],
-        dtype_names=["float32"],
-        shapes=[[30]],
-        num_updates_list=[3],
-    )
-
-    engine = SparseNCCLWeightTransferEngine(
-        config, vllm_config, torch.device("cuda"), model
-    )
-    from vllm.distributed.weight_transfer.nccl_common import (
-        NCCLWeightTransferInitInfo,
-    )
-
-    engine.init_transfer_engine(
-        NCCLWeightTransferInitInfo(
-            master_address=master_address,
-            master_port=master_port,
-            rank_offset=1,
-            world_size=world_size,
+    expert_names = [
+        "model.layers.0.mlp.experts.0.gate_proj.weight",
+        "model.layers.0.mlp.experts.1.gate_proj.weight",
+    ]
+    client.start_weight_update()
+    for name, index, value in zip(
+        expert_names,
+        (0, 3),
+        (5.0, 7.0),
+        strict=True,
+    ):
+        sender.send_weight_chunk(
+            [
+                SparseWeightPatch(
+                    name=name,
+                    indices=torch.tensor([index], dtype=torch.int32),
+                    values=torch.tensor([value]),
+                    full_shape=(2, 2),
+                )
+            ]
         )
+    client.finish_weight_update()
+
+    assert client.order == ["start", "update", "update", "finish"]
+    assert [info["names"] for info in client.update_infos] == [
+        [expert_names[0]],
+        [expert_names[1]],
+    ]
+    assert [info["shapes"] for info in client.update_infos] == [
+        [[2, 2]],
+        [[2, 2]],
+    ]
+    assert [info["num_updates_list"] for info in client.update_infos] == [
+        [1],
+        [1],
+    ]
+
+    expected_gates = (
+        [[5.0, -1.0], [-1.0, -1.0]],
+        [[-1.0, -1.0], [-1.0, 7.0]],
     )
-    engine.receive_weights(update_info)
-    torch.accelerator.synchronize()
+    for ep_rank, expected_gate in enumerate(expected_gates):
+        model = torch.nn.Module()
+        model.register_parameter(
+            "w13_weight",
+            torch.nn.Parameter(torch.full((1, 4, 2), -1.0), requires_grad=False),
+        )
+        local_name = expert_names[ep_rank]
+        seen_names: list[str] = []
+        loaded_names: set[str] = set()
+        copy_count = [0]
 
-    expected = torch.zeros(30, dtype=torch.float32, device=device)
-    expected[[1, 7, 25]] = torch.tensor(
-        [10.0, 20.0, 30.0], dtype=torch.float32, device=device
-    )
-    success = torch.equal(model.w.data, expected)
-    engine.shutdown()
-    return {
-        "success": success,
-        "selected_values": model.w.data[[1, 7, 25]].cpu().tolist(),
-    }
+        def load_weights(
+            weights,
+            local_name=local_name,
+            target=model,
+            seen=seen_names,
+            loaded=loaded_names,
+            count=copy_count,
+        ):
+            for name, checkpoint_weight in weights:
+                seen.append(name)
+                if name != local_name:
+                    continue
+                target.w13_weight.data[0, :2].copy_(checkpoint_weight)
+                loaded.add(name)
+                count[0] += 1
+            return loaded
 
+        model.load_weights = load_weights
+        receiver = SparseNCCLWeightTransferEngine(
+            WeightTransferConfig(backend="sparse_nccl"),
+            create_mock_vllm_config(rank=ep_rank, world_size=2),
+            torch.device("cpu"),
+            model,
+        )
+        receiver.model_update_group = MagicMock()
+        receiver.model_update_group.device = torch.device("cpu")
+        payloads = iter(wire_payloads)
+        receiver.model_update_group.broadcast.side_effect = (
+            lambda tensor, payloads=payloads, **_: tensor.copy_(next(payloads))
+        )
 
-@pytest.mark.skipif(
-    torch.accelerator.device_count() < 2,
-    reason="Need at least 2 GPUs to run NCCL sparse weight transfer test.",
-)
-def test_nccl_sparse_weight_transfer_between_processes():
-    """Test NCCL sparse weight transfer from trainer to inference process."""
-    _init_ray_for_weight_transfer()
+        receiver.start_weight_update()
+        for update_info in client.update_infos:
+            receiver.receive_weights(SparseNCCLWeightTransferUpdateInfo(**update_info))
+        receiver.finish_weight_update()
 
-    master_address = "127.0.0.1"
-    master_port = get_open_port()
-    world_size = 2
+        assert seen_names == expert_names
+        assert loaded_names == {local_name}
+        assert copy_count == [1]
+        assert model.w13_weight[0, :2].tolist() == expected_gate
+        assert model.w13_weight[0, 2:].eq(-1).all()
+        foreign_value = 7.0 if ep_rank == 0 else 5.0
+        assert not model.w13_weight.eq(foreign_value).any()
+        receiver.shutdown()
 
-    inference_future = inference_receive_sparse_tensor.remote(
-        master_address, master_port, world_size
-    )
-    trainer_future = trainer_broadcast_sparse_tensor.remote(
-        master_address, master_port, world_size
-    )
-
-    trainer_result, result = ray.get([trainer_future, inference_future])
-
-    assert trainer_result, "Trainer should complete successfully"
-    assert result["success"], (
-        "Sparse weight transfer failed. "
-        f"Received selected values: {result['selected_values']}"
-    )
+    sender.shutdown()
 
 
 # --- Unit Tests: IPCWeightTransferUpdateInfo Validation ---
@@ -1995,11 +1872,11 @@ def test_sparse_nccl_trainer_init_ships_worker_init_info(monkeypatch):
 
 
 def test_sparse_nccl_trainer_send_weights_drives_client_in_order(monkeypatch):
-    """send_weights takes the round's patches and ships per-patch metadata
-    (names / shapes / num_updates_list) + broadcasts indices + values each."""
+    """One-shot send drives start, update, and finish in order."""
     client = RecordingClient()
     engine = SparseNCCLTrainerWeightTransferEngine(client=client)
     engine.model_update_group = MagicMock()
+    engine.model_update_group.device = torch.device("cpu")
     # The group is a mock, so the stream is just a handle it is handed (and a
     # handle _post_send_sync can synchronize).
     monkeypatch.setattr(torch.cuda, "current_stream", MagicMock())
@@ -2013,6 +1890,7 @@ def test_sparse_nccl_trainer_send_weights_drives_client_in_order(monkeypatch):
     assert client.last_update_info["num_updates_list"] == [2]
     # One broadcast for indices + one for values per patch.
     assert engine.model_update_group.broadcast.call_count == 2
+    engine.shutdown()
 
 
 def test_sparse_nccl_trainer_send_weights_empty_round_is_noop():
@@ -2024,6 +1902,8 @@ def test_sparse_nccl_trainer_send_weights_empty_round_is_noop():
 
     engine.send_weights([])
     engine.send_weights()  # no argument is also a no-op round
+    engine.send_weight_chunk([])
+    engine.send_weight_chunk()
 
     assert client.order == []
 
@@ -2033,6 +1913,7 @@ def test_sparse_nccl_trainer_send_weights_requires_full_shape():
     patch.full_shape = None
     engine = SparseNCCLTrainerWeightTransferEngine(client=RecordingClient())
     engine.model_update_group = MagicMock()
+    engine.model_update_group.device = torch.device("cpu")
 
     with pytest.raises(ValueError, match="full_shape"):
         engine.send_weights([patch])
@@ -2054,6 +1935,7 @@ def test_sparse_nccl_trainer_validates_patch_before_any_rpc():
     client = RecordingClient()
     engine = SparseNCCLTrainerWeightTransferEngine(client=client)
     engine.model_update_group = MagicMock()
+    engine.model_update_group.device = torch.device("cpu")
 
     mismatched = SparseWeightPatch(
         name="w",
@@ -2092,4 +1974,5 @@ def test_sparse_nccl_trainer_non_sender_skips_client():
     assert engine.model_update_group is None
     assert isinstance(engine, SparseNCCLTrainerWeightTransferEngine)
     engine.send_weights([_sparse_patch()])
+    engine.send_weight_chunk([_sparse_patch()])
     assert client.order == []

vllm/distributed/weight_transfer/sparse_nccl_engine.py

@@ -2,20 +2,14 @@
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 """Sparse NCCL weight transfer engine.
 
-A standalone engine (not a subclass of `NCCLWeightTransferEngine`) for applying
-sparse, flat-index weight patches in place. It shares only NCCL process-group
-initialization with the dense engine (via `nccl_common`); the update path
-applies index/value patches directly to existing model parameters and never runs
-layerwise reload.
-
-MVP limitations:
-* TP=1 and PP=1 only
-* uses runtime/kernel-format parameter names
-* not composable with checkpoint-format or packed updates
+Sparse patches use checkpoint names, shapes, and flat indices. The model's native
+weight loader maps them to rank-local runtime parameters, including TP shards and
+packed parameters.
 """
 
 from collections.abc import Iterable
 from concurrent.futures import ThreadPoolExecutor
+from contextlib import nullcontext
 from dataclasses import asdict, dataclass
 from typing import TYPE_CHECKING, ClassVar
 
@@ -42,6 +36,10 @@ from vllm.distributed.weight_transfer.nccl_common import (
 from vllm.distributed.weight_transfer.nccl_common import (
     trainer_init as open_trainer_endpoint,
 )
+from vllm.model_executor.model_loader.checkpoint_weight_patch import (
+    CheckpointWeightPatch,
+    load_checkpoint_weight_patches,
+)
 
 __all__ = [
     "SparseWeightPatch",
@@ -54,15 +52,13 @@ __all__ = [
 
 @dataclass
 class SparseWeightPatch:
-    """A sparse in-place patch for one existing parameter."""
+    """A sparse patch in checkpoint coordinates."""
 
     name: str
     indices: torch.Tensor
     values: torch.Tensor
-    full_shape: tuple[int, ...] | None = None
-    """Full shape of the patched parameter. Required when the patch is sent
-    via `SparseNCCLTrainerWeightTransferEngine` (it ships in the per-round
-    update info); the worker-side apply path does not read it."""
+    full_shape: tuple[int, ...]
+    """Full checkpoint shape."""
 
 
 @dataclass
@@ -119,10 +115,9 @@ class SparseNCCLWeightTransferEngine(
     """
     Sparse weight transfer engine using NCCL.
 
-    Receives flat-index (indices, values) patches broadcast from the trainer
-    (rank 0) and applies them in place to existing model parameters. Weights are
-    applied directly without layerwise reload, so `start_weight_update` and
-    `finish_weight_update` are no-ops.
+    Receives checkpoint-coordinate patches broadcast from the trainer and applies
+    them through the model's native weight loader. Sparse updates modify initialized
+    model tensors in place, so the layerwise reload lifecycle is not used.
     """
 
     init_info_cls = NCCLWeightTransferInitInfo
@@ -147,10 +142,7 @@ class SparseNCCLWeightTransferEngine(
 
     def start_weight_update(self) -> None:
         """No-op: sparse patches are applied in place, no layerwise reload."""
-        if self.parallel_config.world_size != 1:
-            raise NotImplementedError(
-                "Sparse weight updates currently require TP=1 and PP=1"
-            )
+        pass
 
     def finish_weight_update(self) -> None:
         """No-op: sparse patches are applied in place, no layerwise reload."""
@@ -164,64 +156,34 @@ class SparseNCCLWeightTransferEngine(
                 "Call init_transfer_engine() first."
             )
 
-        # Use the worker's assigned device rather than the ambient current
-        # device: the receive path is no longer wrapped in
-        # `with torch.device(self.device)` by the caller, so the current device
-        # is not guaranteed to match self.device. The recv buffers must live on
-        # the same device as the NCCL communicator (created on self.device).
-        device = self.device
-        for name, dtype_name, num_updates in zip(
-            update_info.names,
-            update_info.dtype_names,
-            update_info.num_updates_list,
-        ):
-            dtype = getattr(torch, dtype_name)
-            indices = torch.empty(num_updates, dtype=torch.int32, device=device)
-            values = torch.empty(num_updates, dtype=dtype, device=device)
-            self.model_update_group.broadcast(
-                indices, src=0, stream=torch.cuda.current_stream()
-            )
-            self.model_update_group.broadcast(
-                values, src=0, stream=torch.cuda.current_stream()
-            )
-            self._apply_patch(
-                SparseWeightPatch(name=name, indices=indices, values=values)
-            )
-            del indices
-            del values
-
-    def _apply_patch(self, patch: SparseWeightPatch) -> None:
-        """Apply a single sparse flat-index patch to an existing model param."""
-        param = self.model.get_parameter(patch.name)
-        if not param.data.is_contiguous():
-            raise NotImplementedError(
-                "Sparse weight updates currently require contiguous params: "
-                f"{patch.name}"
-            )
-        if patch.indices.dtype != torch.int32:
-            raise ValueError(
-                f"Sparse weight updates currently require int32 indices: {patch.name}"
-            )
-        if patch.indices.ndim != 1 or patch.values.ndim != 1:
-            raise ValueError(
-                f"Sparse weight patches must be 1D flattened updates: {patch.name}"
-            )
-        if patch.indices.numel() != patch.values.numel():
-            raise ValueError(
-                f"`indices` and `values` must have matching lengths for {patch.name}"
-            )
-        if patch.values.dtype != param.dtype:
-            raise ValueError(
-                f"Sparse values dtype {patch.values.dtype} does not match "
-                f"parameter dtype {param.dtype} for {patch.name}"
-            )
-
-        flat_param = param.data.view(-1)
-        flat_param.index_copy_(
-            0,
-            patch.indices.to(device=flat_param.device, dtype=torch.long),
-            patch.values.to(device=flat_param.device),
-        )
+        # Receive on the communicator's device rather than relying on the
+        # ambient current device in the RPC thread.
+        device = self.model_update_group.device
+        with torch.accelerator.device_index(device.index):
+            stream = torch.cuda.current_stream(device=device)
+            patches = []
+            for name, dtype_name, shape, num_updates in zip(
+                update_info.names,
+                update_info.dtype_names,
+                update_info.shapes,
+                update_info.num_updates_list,
+                strict=True,
+            ):
+                dtype = getattr(torch, dtype_name)
+                indices = torch.empty(num_updates, dtype=torch.int32, device=device)
+                values = torch.empty(num_updates, dtype=dtype, device=device)
+                self.model_update_group.broadcast(indices, src=0, stream=stream)
+                self.model_update_group.broadcast(values, src=0, stream=stream)
+                patches.append(
+                    CheckpointWeightPatch(
+                        name=name,
+                        shape=tuple(shape),
+                        dtype=dtype,
+                        values=values,
+                        indices=indices,
+                    )
+                )
+            load_checkpoint_weight_patches(self.model, patches)
 
     def shutdown(self) -> None:
         if self.model_update_group is not None:
@@ -235,15 +197,16 @@ class SparseNCCLTrainerWeightTransferEngine(
 
     Broadcasts flat-index (indices, values) patches from NCCL rank 0 while the
     inference-side `update_weights` runs concurrently on a side thread (the
-    worker's recvs rendezvous inside the same NCCL broadcasts), then finishes
-    the update. Unlike the full-resync backends, sparse patches differ every
-    round (a fresh set of deltas from each optimizer step), so they are not a
-    stable `WeightSource`: the engine takes no `source`, and each round's
-    patches are passed straight to `send_weights(patches)`. A round with no
-    patches is a no-op.
+    worker's recvs rendezvous inside the same NCCL broadcasts). `send_weights`
+    owns a complete one-shot lifecycle. RL infrastructure that owns the generic
+    client lifecycle can call `send_weight_chunk` between one start and finish.
 
-    The sparse backend assumes a single-rank trainer (matching its TP=1 / PP=1
-    MVP scope), so non-sender ranks skip `send_weights` entirely.
+    Sparse patches differ every round, so they are not a stable `WeightSource`:
+    the engine takes no `source`, and patches are passed directly to the send
+    methods. An empty patch list is a no-op.
+
+    Only the designated trainer sender joins the transfer group; other trainer
+    ranks skip sparse sends.
     """
 
     init_info_cls = SparseNCCLTrainerInitInfo
@@ -303,70 +266,88 @@ class SparseNCCLTrainerWeightTransferEngine(
         return engine
 
     def send_weights(self, patches: Iterable[SparseWeightPatch] | None = None) -> None:
-        """Broadcast this round's sparse patches. `patches` is the per-round
-        payload (sparse deltas differ every round), so it is passed here rather
-        than fixed at init. Every patch must set `full_shape`."""
-        if not self.is_sender:
-            return
-
-        patches = list(patches) if patches is not None else []
+        """Broadcast one sparse update through a one-shot lifecycle."""
+        patches = self._prepare_patches(patches)
         if not patches:
             return
 
-        shapes = []
-        for patch in patches:
-            if patch.full_shape is None:
-                raise ValueError(
-                    "SparseWeightPatch.full_shape must be set to send via the "
-                    f"trainer engine: {patch.name}"
-                )
-            self._validate_patch(patch)
-            shapes.append(list(patch.full_shape))
-
-        update_info = SparseNCCLWeightTransferUpdateInfo(
-            names=[patch.name for patch in patches],
-            dtype_names=[str(patch.values.dtype).split(".")[-1] for patch in patches],
-            shapes=shapes,
-            num_updates_list=[patch.indices.numel() for patch in patches],
-        )
-
-        assert self.model_update_group is not None, (
-            "trainer_init() must be called before send_weights()."
-        )
         self.client.start_weight_update()
-        # update_weights (workers receive) must run concurrently with the
-        # trainer-side broadcasts — both rendezvous inside the same NCCL calls.
-        exe = ThreadPoolExecutor(max_workers=1)
-        try:
-            future = exe.submit(self.client.update_weights, asdict(update_info))
-            # Cheap best-effort: if update_weights already failed (e.g. a bad
-            # request rejected before any NCCL call), surface it now instead of
-            # hanging in broadcast waiting for a peer that will never arrive.
-            if future.done():
-                future.result()
-            stream = torch.cuda.current_stream()
-            for patch in patches:
-                self.model_update_group.broadcast(patch.indices, src=0, stream=stream)
-                self.model_update_group.broadcast(patch.values, src=0, stream=stream)
-            future.result()  # surface inference-side errors
-        finally:
-            # Never wait for the RPC thread here: if a broadcast raised, the
-            # worker is still blocked in the matching NCCL call and will never
-            # return, so joining would turn the error into a permanent hang.
-            # See NCCLTrainerWeightTransferEngine.send_weights.
-            exe.shutdown(wait=False)
+        self._broadcast_chunk(patches)
         self.client.finish_weight_update()
-        self._post_send_sync()
+
+    def send_weight_chunk(
+        self, patches: Iterable[SparseWeightPatch] | None = None
+    ) -> None:
+        """Broadcast one chunk inside a caller-owned weight update lifecycle."""
+        patches = self._prepare_patches(patches)
+        if not patches:
+            return
+        self._broadcast_chunk(patches)
+
+    def _prepare_patches(
+        self, patches: Iterable[SparseWeightPatch] | None
+    ) -> list[SparseWeightPatch]:
+        if not self.is_sender:
+            return []
+
+        prepared = list(patches) if patches is not None else []
+        if not prepared:
+            return []
+        if self.model_update_group is None:
+            raise RuntimeError("trainer_init() must be called before sending weights.")
+        for patch in prepared:
+            self._validate_patch(patch)
+        return prepared
+
+    def _broadcast_chunk(self, patches: list[SparseWeightPatch]) -> None:
+        assert self.model_update_group is not None
+        device = self.model_update_group.device
+        device_context = (
+            torch.accelerator.device_index(device.index)
+            if device.type != "cpu"
+            else nullcontext()
+        )
+        with device_context:
+            update_info = SparseNCCLWeightTransferUpdateInfo(
+                names=[patch.name for patch in patches],
+                dtype_names=[
+                    str(patch.values.dtype).split(".")[-1] for patch in patches
+                ],
+                shapes=[list(patch.full_shape) for patch in patches],
+                num_updates_list=[patch.indices.numel() for patch in patches],
+            )
+
+            # update_weights (workers receive) must run concurrently with the
+            # trainer-side broadcasts — both rendezvous inside the same NCCL calls.
+            executor = ThreadPoolExecutor(max_workers=1)
+            try:
+                future = executor.submit(
+                    self.client.update_weights, asdict(update_info)
+                )
+                # Surface an RPC that failed before the worker entered NCCL rather
+                # than broadcasting to a peer that will never receive.
+                if future.done():
+                    future.result()
+                stream = torch.cuda.current_stream()
+                for patch in patches:
+                    self.model_update_group.broadcast(
+                        patch.indices, src=0, stream=stream
+                    )
+                    self.model_update_group.broadcast(
+                        patch.values, src=0, stream=stream
+                    )
+                future.result()  # surface inference-side errors
+            finally:
+                # A failed broadcast can leave the RPC blocked in the matching
+                # receive. Waiting for that thread would hide the original error.
+                executor.shutdown(wait=False, cancel_futures=True)
+            self._post_send_sync()
 
     @staticmethod
     def _validate_patch(patch: SparseWeightPatch) -> None:
-        """Reject a malformed patch before any NCCL call.
-
-        The worker checks the same invariants in `_apply_patch`, but by then the
-        broadcasts are already under way: a mismatch surfaces as a size mismatch
-        mid-transfer, which wedges both sides instead of raising. Checking here
-        keeps the failure on the trainer, before `start_weight_update`.
-        """
+        """Reject a malformed patch before starting the NCCL transfer."""
+        if patch.full_shape is None:
+            raise ValueError(f"Sparse patch requires full_shape: {patch.name}")
         if patch.indices.dtype != torch.int32:
             raise ValueError(
                 f"Sparse weight updates require int32 indices: {patch.name}"
@@ -382,10 +363,10 @@ class SparseNCCLTrainerWeightTransferEngine(
 
     def _post_send_sync(self) -> None:
         """Wait for the broadcasts to land before returning, so a caller may
-        rebuild or free the patch tensors as soon as `send_weights` returns
-        rather than relying on same-stream ordering. See
+        rebuild or free the patch tensors as soon as a send method returns rather
+        than relying on same-stream ordering. See
         `NCCLTrainerWeightTransferEngine._post_send_sync` for why there is no
-        cross-rank barrier (and sparse is single-rank on the trainer anyway)."""
+        cross-rank barrier."""
         if torch.cuda.is_available():
             torch.cuda.current_stream().synchronize()