diff --git a/docs/training/weight_transfer/README.md b/docs/training/weight_transfer/README.md index f26c033281..a389dc7165 100644 --- a/docs/training/weight_transfer/README.md +++ b/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 diff --git a/docs/training/weight_transfer/nccl.md b/docs/training/weight_transfer/nccl.md index 5f1a61dedc..3a3e207fdc 100644 --- a/docs/training/weight_transfer/nccl.md +++ b/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` diff --git a/examples/rl/rlhf_sparse_nccl.py b/examples/rl/rlhf_sparse_nccl.py index f28edf3e7e..ee987c43f5 100644 --- a/examples/rl/rlhf_sparse_nccl.py +++ b/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() diff --git a/tests/distributed/test_weight_transfer.py b/tests/distributed/test_weight_transfer.py index 5f9c6ff531..1353def738 100644 --- a/tests/distributed/test_weight_transfer.py +++ b/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 == [] diff --git a/vllm/distributed/weight_transfer/sparse_nccl_engine.py b/vllm/distributed/weight_transfer/sparse_nccl_engine.py index 5d952684c8..89eb892a67 100644 --- a/vllm/distributed/weight_transfer/sparse_nccl_engine.py +++ b/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()