From 21ea5b4fa1062a379dd7e6795497ad6becd5a856 Mon Sep 17 00:00:00 2001 From: Aaron Hao Date: Fri, 7 Aug 2026 00:35:59 -0700 Subject: [PATCH] [rl] Stateful Trainer Send: NCCL + Sparse NCCL [3/N] (#50902) Signed-off-by: haoaaron Co-authored-by: Claude Opus 5 (1M context) --- examples/rl/rlhf_async_new_apis.py | 95 +--- examples/rl/rlhf_http_nccl.py | 132 +---- examples/rl/rlhf_nccl.py | 105 ++-- examples/rl/rlhf_nccl_fsdp_ep.py | 466 +++++++++-------- examples/rl/rlhf_sparse_nccl.py | 186 +++---- tests/distributed/test_packed_tensor.py | 27 - tests/distributed/test_weight_transfer.py | 493 +++++++++++++++++- vllm/distributed/weight_transfer/base.py | 20 +- vllm/distributed/weight_transfer/factory.py | 21 +- .../distributed/weight_transfer/ipc_engine.py | 5 +- .../weight_transfer/nccl_common.py | 29 +- .../weight_transfer/nccl_engine.py | 389 ++++++++++---- .../weight_transfer/packed_tensor.py | 8 + .../weight_transfer/sparse_nccl_engine.py | 229 +++++++- 14 files changed, 1422 insertions(+), 783 deletions(-) diff --git a/examples/rl/rlhf_async_new_apis.py b/examples/rl/rlhf_async_new_apis.py index 7043182ab1..acd1a42128 100644 --- a/examples/rl/rlhf_async_new_apis.py +++ b/examples/rl/rlhf_async_new_apis.py @@ -41,7 +41,6 @@ causes unexpected behavior. import asyncio import uuid -from dataclasses import asdict import ray import torch @@ -50,16 +49,12 @@ from transformers import AutoModelForCausalLM, AutoTokenizer import vllm from vllm import SamplingParams from vllm.config import WeightTransferConfig -from vllm.distributed.weight_transfer.base import ( - WeightTransferInitRequest, - WeightTransferUpdateRequest, -) -from vllm.distributed.weight_transfer.nccl_engine import ( - NCCLTrainerSendWeightsArgs, - NCCLWeightTransferEngine, - NCCLWeightTransferInitInfo, - NCCLWeightTransferUpdateInfo, +from vllm.distributed.weight_transfer import ( + ModuleSource, + RayVLLMWeightSyncClient, + WeightTransferTrainerFactory, ) +from vllm.distributed.weight_transfer.nccl_engine import NCCLTrainerInitInfo from vllm.platforms import current_platform from vllm.utils.network_utils import get_ip, get_open_port from vllm.v1.executor import Executor @@ -141,40 +136,23 @@ class TrainModel: self.port = get_open_port() self.master_address = get_ip() - def get_master_address_and_port(self): - return self.master_address, self.port - - def get_weight_metadata(self): - """Return weight names, dtypes, and shapes for weight transfer.""" - names = [] - dtype_names = [] - shapes = [] - for name, p in self.model.named_parameters(): - names.append(name) - dtype_names.append(str(p.dtype).split(".")[-1]) - shapes.append(list(p.shape)) - return names, dtype_names, shapes - - def init_weight_transfer_group(self, world_size): - """Initialize the NCCL process group for weight transfer.""" - self.model_update_group = NCCLWeightTransferEngine.trainer_init( - dict( + def init_weight_transfer(self, world_size, llm_handle): + """Build the trainer-side weight-transfer engine and rendezvous.""" + self.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), ) - def broadcast_weights(self, packed: bool = True): - """Broadcast weights to the inference engine.""" - trainer_args = NCCLTrainerSendWeightsArgs( - group=self.model_update_group, - packed=packed, - ) - NCCLWeightTransferEngine.trainer_send_weights( - iterator=self.model.named_parameters(), - trainer_args=trainer_args, - ) + def broadcast_weights(self): + """Push weights to the inference engine (drives start/update/finish).""" + self.engine.send_weights() @torch.inference_mode() def generate(self, token_ids: list[int], max_new_tokens: int) -> list[int]: @@ -263,32 +241,12 @@ batch_prompt_token_ids = [ # Set up the communication channel between the training process and the # inference engine. -master_address, master_port = ray.get(train_model.get_master_address_and_port.remote()) - world_size = 2 # 1 trainer + 1 inference worker -inference_handle = llm.init_weight_transfer_engine.remote( - WeightTransferInitRequest( - init_info=asdict( - NCCLWeightTransferInitInfo( - master_address=master_address, - master_port=master_port, - rank_offset=1, - world_size=world_size, - ) - ) - ) -) - -# Initialize weight transfer group on both the training actor and inference engine -train_handle = train_model.init_weight_transfer_group.remote(world_size) -ray.get([train_handle, inference_handle]) +ray.get(train_model.init_weight_transfer.remote(world_size, llm)) N_NEW_TOKENS = 100 -# Collect weight metadata once -names, dtype_names, shapes = ray.get(train_model.get_weight_metadata.remote()) - # ── Phase 1: concurrent requests with weight sync ─────────────────── print(f"\n{'=' * 50}") print(f"Prompts ({len(PROMPTS)}):") @@ -306,24 +264,7 @@ gen_futures = [ ray.get(llm.pause_after_n_tokens.remote()) -ray.get(llm.start_weight_update.remote()) - -inference_handle = llm.update_weights.remote( - WeightTransferUpdateRequest( - update_info=asdict( - NCCLWeightTransferUpdateInfo( - names=names, - dtype_names=dtype_names, - shapes=shapes, - packed=True, - ) - ) - ) -) -train_handle = train_model.broadcast_weights.remote(packed=True) -ray.get([train_handle, inference_handle]) - -ray.get(llm.finish_weight_update.remote()) +ray.get(train_model.broadcast_weights.remote()) ray.get(llm.resume_generation.remote()) results = ray.get(gen_futures) diff --git a/examples/rl/rlhf_http_nccl.py b/examples/rl/rlhf_http_nccl.py index b40b8de32f..911dc788ae 100644 --- a/examples/rl/rlhf_http_nccl.py +++ b/examples/rl/rlhf_http_nccl.py @@ -38,10 +38,12 @@ import torch from openai import OpenAI from transformers import AutoModelForCausalLM -from vllm.distributed.weight_transfer.nccl_engine import ( - NCCLTrainerSendWeightsArgs, - NCCLWeightTransferEngine, +from vllm.distributed.weight_transfer import ( + HTTPVLLMWeightSyncClient, + ModuleSource, + WeightTransferTrainerFactory, ) +from vllm.distributed.weight_transfer.nccl_engine import NCCLTrainerInitInfo from vllm.utils.network_utils import get_ip, get_open_port BASE_URL = "http://localhost:8000" @@ -62,62 +64,6 @@ def generate_completions(client: OpenAI, model: str, prompts: list[str]) -> list return results -def init_weight_transfer_engine( - base_url: str, - master_address: str, - master_port: int, - rank_offset: int, - world_size: int, -) -> None: - """Initialize weight transfer via HTTP endpoint.""" - url = f"{base_url}/init_weight_transfer_engine" - payload = { - "init_info": dict( - master_address=master_address, - master_port=master_port, - rank_offset=rank_offset, - world_size=world_size, - ) - } - response = requests.post(url, json=payload, timeout=60) - response.raise_for_status() - - -def start_weight_update(base_url: str) -> None: - """Start a weight update via HTTP endpoint.""" - url = f"{base_url}/start_weight_update" - response = requests.post(url, json={}, timeout=60) - response.raise_for_status() - - -def update_weights( - base_url: str, - names: list[str], - dtype_names: list[str], - shapes: list[list[int]], - packed: bool = False, -) -> None: - """Update weights via HTTP endpoint.""" - url = f"{base_url}/update_weights" - payload = { - "update_info": dict( - names=names, - dtype_names=dtype_names, - shapes=shapes, - packed=packed, - ) - } - response = requests.post(url, json=payload, timeout=300) - response.raise_for_status() - - -def finish_weight_update(base_url: str) -> None: - """Finish a weight update via HTTP endpoint.""" - url = f"{base_url}/finish_weight_update" - response = requests.post(url, json={}, timeout=60) - response.raise_for_status() - - def pause_generation(base_url: str) -> None: """Pause generation via HTTP endpoint.""" url = f"{base_url}/pause" @@ -177,75 +123,35 @@ def main(): print("-" * 50) # Set up the communication channel between the training process and the - # vLLM server. The trainer is rank 0, vLLM worker(s) start at rank_offset. + # vLLM server. The trainer is rank 0, vLLM worker(s) start after it. master_address = get_ip() master_port = get_open_port() - rank_offset = 1 print(f"Initializing weight transfer: master={master_address}:{master_port}") - # Initialize weight transfer on vLLM server (this is async, server will - # wait for NCCL connection) - import threading - - init_thread = threading.Thread( - target=init_weight_transfer_engine, - args=(BASE_URL, master_address, master_port, rank_offset, world_size), - ) - init_thread.start() - - # Initialize NCCL process group on trainer side - model_update_group = NCCLWeightTransferEngine.trainer_init( - dict( + # The trainer engine owns the full handshake. `trainer_init` kicks off the + # server's init_weight_transfer_engine (via the HTTP client) on a worker + # thread while opening the trainer's NCCL endpoint, so both ends rendezvous + # together — no manual threading needed here. + engine = WeightTransferTrainerFactory.trainer_init( + init_info=NCCLTrainerInitInfo( master_address=master_address, master_port=master_port, world_size=world_size, + rank=0, # single-GPU trainer is the sole (sender) rank + packed=True, ), + client=HTTPVLLMWeightSyncClient(BASE_URL), + source=ModuleSource(train_model), ) - # Wait for init_weight_transfer_engine to complete - init_thread.join() - # Pause generation before weight sync pause_generation(BASE_URL) - # Collect weight metadata for the update request - names = [] - dtype_names = [] - shapes = [] - for name, p in train_model.named_parameters(): - names.append(name) - dtype_names.append(str(p.dtype).split(".")[-1]) - shapes.append(list(p.shape)) - - # Start weight update - start_weight_update(BASE_URL) - - # Start the update_weights call in a separate thread since it will block - # waiting for NCCL broadcasts - # packed=True enables efficient batched tensor broadcasting - update_thread = threading.Thread( - target=update_weights, - args=(BASE_URL, names, dtype_names, shapes, True), # packed=True - ) - update_thread.start() - - # Broadcast all weights from trainer to vLLM workers + # One call drives start_weight_update / update_weights / finish_weight_update + # over HTTP, concurrent with the NCCL broadcast. print("Broadcasting weights via NCCL...") - trainer_args = NCCLTrainerSendWeightsArgs( - group=model_update_group, - packed=True, - ) - NCCLWeightTransferEngine.trainer_send_weights( - iterator=train_model.named_parameters(), - trainer_args=trainer_args, - ) - - # Wait for update_weights to complete - update_thread.join() - - # Finish weight update - finish_weight_update(BASE_URL) + engine.send_weights() # Resume generation after weight sync resume_generation(BASE_URL) diff --git a/examples/rl/rlhf_nccl.py b/examples/rl/rlhf_nccl.py index bebd6bc70d..e06d9ec677 100644 --- a/examples/rl/rlhf_nccl.py +++ b/examples/rl/rlhf_nccl.py @@ -36,10 +36,12 @@ from transformers import AutoModelForCausalLM from vllm import LLM, SamplingParams from vllm.config import WeightTransferConfig -from vllm.distributed.weight_transfer.nccl_engine import ( - NCCLTrainerSendWeightsArgs, - NCCLWeightTransferEngine, +from vllm.distributed.weight_transfer import ( + ModuleSource, + RayVLLMWeightSyncClient, + WeightTransferTrainerFactory, ) +from vllm.distributed.weight_transfer.nccl_engine import NCCLTrainerInitInfo from vllm.platforms import current_platform from vllm.utils.network_utils import get_ip, get_open_port @@ -80,40 +82,33 @@ class TrainModel: self.port = get_open_port() self.master_address = get_ip() - def get_master_address_and_port(self): - return self.master_address, self.port + def init_weight_transfer(self, world_size, llm_handle): + """Build the trainer-side weight-transfer engine. - def get_weight_metadata(self): - """Return weight names, dtypes, and shapes for weight transfer.""" - names = [] - dtype_names = [] - shapes = [] - for name, p in self.model.named_parameters(): - names.append(name) - dtype_names.append(str(p.dtype).split(".")[-1]) - shapes.append(list(p.shape)) - return names, dtype_names, shapes - - def init_weight_transfer_group(self, world_size): - """Initialize the NCCL process group for weight transfer.""" - self.model_update_group = NCCLWeightTransferEngine.trainer_init( - dict( + `trainer_init` drives the full handshake: it kicks off the inference + side's `init_weight_transfer_engine` (via the Ray client) on a worker + thread while opening the trainer's NCCL endpoint, so both ends + rendezvous together. After this returns, `send_weights()` is callable. + """ + self.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), ) - def broadcast_weights(self, packed: bool = True): - """Broadcast weights to the inference engine.""" - trainer_args = NCCLTrainerSendWeightsArgs( - group=self.model_update_group, - packed=packed, - ) - NCCLWeightTransferEngine.trainer_send_weights( - iterator=self.model.named_parameters(), - trainer_args=trainer_args, - ) + def broadcast_weights(self): + """Push the current weights to the inference engine. + + Drives start/update/finish on the inference side and the NCCL + broadcast internally — one call. + """ + self.engine.send_weights() # Initialize Ray and set the visible devices. The vLLM engine will @@ -178,51 +173,15 @@ for output in outputs: ray.get(llm.sleep.remote(level=0)) # Set up the communication channel between the training process and the -# inference engine. -master_address, master_port = ray.get(train_model.get_master_address_and_port.remote()) - +# inference engine. The trainer engine now owns the full handshake — the +# driver only has to hand it the vLLM actor handle and the world size. world_size = ray.get(llm.get_world_size.remote()) + 1 # +1 for the trainer -inference_handle = llm.init_weight_transfer_engine.remote( - dict( - init_info=dict( - master_address=master_address, - master_port=master_port, - rank_offset=1, - world_size=world_size, - ) - ) -) +ray.get(train_model.init_weight_transfer.remote(world_size, llm)) -# Initialize weight transfer group on both the training actor and inference engine -train_handle = train_model.init_weight_transfer_group.remote(world_size) -ray.get([train_handle, inference_handle]) - -# Synchronize the updated weights to the inference engine using batched API. -# Collect all weight metadata from the training actor -names, dtype_names, shapes = ray.get(train_model.get_weight_metadata.remote()) - -# Start weight update -ray.get(llm.start_weight_update.remote()) - -# Issue update_weights call with NCCL-specific update info -# packed=True enables efficient batched tensor broadcasting -inference_handle = llm.update_weights.remote( - dict( - update_info=dict( - names=names, - dtype_names=dtype_names, - shapes=shapes, - packed=True, - ) - ) -) - -# Broadcast all weights from trainer using the weight transfer API -train_handle = train_model.broadcast_weights.remote(packed=True) -ray.get([train_handle, inference_handle]) - -# Finish weight update -ray.get(llm.finish_weight_update.remote()) +# Synchronize the updated weights to the inference engine. The engine drives +# start_weight_update / update_weights / finish_weight_update on the inference +# side internally, concurrent with the NCCL broadcast. +ray.get(train_model.broadcast_weights.remote()) ray.get(llm.wake_up.remote(tags=["scheduling"])) diff --git a/examples/rl/rlhf_nccl_fsdp_ep.py b/examples/rl/rlhf_nccl_fsdp_ep.py index 860db058ca..e4380bea94 100644 --- a/examples/rl/rlhf_nccl_fsdp_ep.py +++ b/examples/rl/rlhf_nccl_fsdp_ep.py @@ -4,69 +4,79 @@ RLHF with FSDP2 training (4 GPUs) and vLLM expert-parallel inference (4 GPUs). 8-GPU layout: - Training — 4 GPUs, PyTorch FSDP2 (fully_shard) - Inference — 4 GPUs, vLLM AsyncLLMEngine with expert parallelism + + Training — 4 GPUs, PyTorch FSDP2 (fully_shard), as Ray actors + Inference — 4 GPUs, a `vllm serve` HTTP server with expert parallelism + data parallelism (TP=1, DP=4, enable_expert_parallel → EP_SIZE = TP×DP = 4) -FSDP workers are Ray actors that form a single FSDP2 process group. -Rank 0 gathers full parameters via DTensor.full_tensor() and broadcasts -them to the vLLM inference engine through the NCCL weight-transfer API. +The inference side is a standalone HTTP server (spawned by this script with +`vllm serve`), so both the weight-sync control plane (HTTP) and the NCCL data +plane run inside the rank-0 FSDP Ray actor. That lets the trainer use the +unified `TrainerWeightTransferEngine.send_weights()` with an +`HTTPVLLMWeightSyncClient` — one call drives start/update/finish on the server +concurrently with the NCCL broadcast. Every FSDP rank builds an engine and calls +`send_weights()`, so all 4 participate in the incremental `full_tensor()` +all-gather; only rank 0 holds a communicator and broadcasts (it is the only +trainer rank in the NCCL group). -The inference engine uses AsyncLLMEngine which automatically spawns -DP worker processes (no manual placement group needed). Weight sync -uses pause_generation / resume_generation. +GPU split (single node): the server takes GPUs 0-3 (CUDA_VISIBLE_DEVICES), and +Ray (training) is restricted to GPUs 4-7. Steps: - 1. Launch 4 FSDP training workers. - 2. Launch AsyncLLMEngine with EP+DP (dummy weights). - 3. Generate from prompts → gibberish (random weights). - 4. Pause generation, transfer weights from FSDP, resume. + 1. Launch the vLLM HTTP server (EP+DP, dummy weights) on GPUs 0-3. + 2. Launch 4 FSDP training workers (Ray) on GPUs 4-7. + 3. Generate from prompts over HTTP → gibberish (random weights). + 4. Pause generation, transfer weights FSDP → server over NCCL, resume. 5. Generate from prompts → sensible output (synced weights). Assumes a single-node cluster with 8 GPUs. """ -import asyncio +import json import os -import uuid -from dataclasses import asdict +import subprocess +import sys +import time import ray +import requests import torch import torch.distributed as dist from huggingface_hub import snapshot_download +from openai import OpenAI from torch.distributed.fsdp import fully_shard from transformers import AutoModelForCausalLM -import vllm -from vllm import SamplingParams -from vllm.config import WeightTransferConfig -from vllm.distributed.weight_transfer.base import ( - WeightTransferInitRequest, - WeightTransferUpdateRequest, -) -from vllm.distributed.weight_transfer.nccl_engine import ( - NCCLTrainerSendWeightsArgs, - NCCLWeightTransferEngine, - NCCLWeightTransferInitInfo, - NCCLWeightTransferUpdateInfo, +from vllm.distributed.weight_transfer import ( + HTTPVLLMWeightSyncClient, + ModuleSource, + WeightTransferTrainerFactory, ) +from vllm.distributed.weight_transfer.nccl_engine import NCCLTrainerInitInfo from vllm.utils.network_utils import get_ip, get_open_port -from vllm.v1.executor import Executor MODEL_NAME = "Qwen/Qwen3-30B-A3B" +SERVED_MODEL_NAME = "policy" FSDP_WORLD_SIZE = 4 INFERENCE_TP_SIZE = 1 INFERENCE_DP_SIZE = 4 +# Training (FSDP) GPUs are reserved through Ray; the inference server then runs +# on the complementary GPUs (see main()). We do NOT hard-code the split via +# CUDA_VISIBLE_DEVICES before ray.init(): that only restricts Ray when ray.init() +# *starts* a local cluster, and is silently ignored when it connects to an +# existing one (e.g. a shared/managed Ray cluster), causing training and the +# server to collide on the same physical GPUs. +SERVER_PORT = 8000 +BASE_URL = f"http://localhost:{SERVER_PORT}" + @ray.remote(num_gpus=1) class FSDPTrainWorker: """ One FSDP2 training worker per GPU. Four of these form the FSDP group. - Rank 0 additionally handles weight transfer to the vLLM engine. + Rank 0 additionally drives weight transfer to the vLLM server. """ def __init__( @@ -78,6 +88,7 @@ class FSDPTrainWorker: fsdp_master_port: int, ): self.rank = rank + self.engine = None os.environ["MASTER_ADDR"] = fsdp_master_addr os.environ["MASTER_PORT"] = str(fsdp_master_port) @@ -89,12 +100,6 @@ class FSDPTrainWorker: model_name, torch_dtype=torch.bfloat16 ) - self.weight_names = [n for n, _ in model.named_parameters()] - self.weight_dtype_names = [ - str(p.dtype).split(".")[-1] for _, p in model.named_parameters() - ] - self.weight_shapes = [list(p.shape) for _, p in model.named_parameters()] - for layer in model.model.layers: fully_shard(layer) fully_shard(model) @@ -103,11 +108,14 @@ class FSDPTrainWorker: self.transfer_port = None self.transfer_master_address = None - self.model_update_group = None def get_rank(self): return self.rank + def get_gpu_ids(self): + """Physical GPU id(s) Ray assigned to this worker (for server/train split).""" + return ray.get_gpu_ids() + # ---- weight-transfer setup (rank 0 only) ---- def setup_transfer_endpoint(self): @@ -117,95 +125,130 @@ class FSDPTrainWorker: self.transfer_master_address = get_ip() return self.transfer_master_address, self.transfer_port - def init_weight_transfer_group(self, transfer_world_size: int): - """Join the weight-transfer NCCL group as rank 0 (the source).""" - assert self.rank == 0 - self.model_update_group = NCCLWeightTransferEngine.trainer_init( - dict( - master_address=self.transfer_master_address, - master_port=self.transfer_port, - world_size=transfer_world_size, - ), - ) + def setup_engine( + self, + base_url: str, + transfer_master_address: str, + transfer_port: int, + transfer_world_size: int, + ): + """Build the trainer engine on every FSDP rank. - def get_weight_metadata(self): - """Return weight names, dtypes, and shapes captured before FSDP wrapping.""" - return self.weight_names, self.weight_dtype_names, self.weight_shapes + Called on all ranks with the shared rendezvous endpoint. Rank 0 is the + sender: `trainer_init` opens its rank-0 NCCL endpoint and, on a worker + thread, calls the server's `init_weight_transfer_engine` over HTTP so + both ends rendezvous together. The other ranks skip the rendezvous and + only join the FSDP all-gather during send_weights. + """ + self.engine = WeightTransferTrainerFactory.trainer_init( + init_info=NCCLTrainerInitInfo( + master_address=transfer_master_address, + master_port=transfer_port, + world_size=transfer_world_size, + rank=self.rank, # FSDP rank; sender is rank 0 + packed=True, + ), + client=HTTPVLLMWeightSyncClient(base_url), + # Yields sharded DTensors; the engine reads global shape/dtype for + # metadata (no gather) and calls full_tensor() at broadcast time. + source=ModuleSource(self.model), + ) # ---- collective ops (ALL FSDP ranks must call concurrently) ---- - def gather_and_broadcast_weights(self, packed: bool = True): + def gather_and_broadcast_weights(self): + """All-gather full parameters and broadcast them to the vLLM server. + + Called on all FSDP ranks. `send_weights` gathers each param via + `full_tensor()` (a collective every rank must enter in the same order); + only rank 0 (the sender) drives the server-side update_weights + concurrently with the NCCL broadcast — the other ranks only gather. """ - All-gather full parameters and broadcast them to vLLM. - Only rank 0 performs the actual NCCL broadcast; others just - participate in the FSDP all-gather. - - full_tensor() is a collective — all FSDP ranks must call it - for each parameter in the same order. Rank 0 additionally - feeds each gathered tensor to the weight-transfer engine. - """ - if self.rank == 0: - - def _full_param_iter(): - for name, param in self.model.named_parameters(): - yield name, param.full_tensor() - - trainer_args = NCCLTrainerSendWeightsArgs( - group=self.model_update_group, - packed=packed, - ) - NCCLWeightTransferEngine.trainer_send_weights( - iterator=_full_param_iter(), - trainer_args=trainer_args, - ) - else: - for _, param in self.model.named_parameters(): - param.full_tensor() + self.engine.send_weights() -def create_async_engine(**kwargs): - """Create an AsyncLLMEngine directly (no subclass needed).""" - engine_args = vllm.AsyncEngineArgs(**kwargs) - vllm_config = engine_args.create_engine_config() - executor_class = Executor.get_class(vllm_config) - return vllm.AsyncLLMEngine( - vllm_config=vllm_config, - executor_class=executor_class, - log_requests=engine_args.enable_log_requests, - log_stats=not engine_args.disable_log_stats, +def start_vllm_server(server_gpus: str) -> subprocess.Popen: + """Spawn a `vllm serve` HTTP server (EP+DP) on `server_gpus` and wait for it.""" + serve_args = [ + "vllm", + "serve", + MODEL_NAME, + "--served-model-name", + SERVED_MODEL_NAME, + "--tensor-parallel-size", + str(INFERENCE_TP_SIZE), + "--data-parallel-size", + str(INFERENCE_DP_SIZE), + "--enable-expert-parallel", + "--enforce-eager", + "--load-format", + "dummy", + "--gpu-memory-utilization", + "0.7", + "--port", + str(SERVER_PORT), + "--weight-transfer-config", + json.dumps({"backend": "nccl"}), + ] + env = os.environ.copy() + env["CUDA_VISIBLE_DEVICES"] = server_gpus + env["VLLM_SERVER_DEV_MODE"] = "1" # exposes the weight-transfer endpoints + env["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn" + print(f"[server] Launching: {' '.join(serve_args)} (GPUs {server_gpus})") + proc = subprocess.Popen( + serve_args, + env=env, + stdout=sys.stdout, + stderr=sys.stderr, + start_new_session=True, ) - -async def generate_batch(engine, prompts, sampling_params): - """Generate completions for a batch of prompts.""" - - async def gen_one(prompt): - output = None - async for request_output in engine.generate( - {"prompt": prompt}, - sampling_params, - request_id=str(uuid.uuid4()), - ): - output = request_output - return output - - return await asyncio.gather(*[gen_one(p) for p in prompts]) + # Wait for the server to come up (model load can take a while). + deadline = time.monotonic() + 1800 + while True: + if proc.poll() is not None: + raise RuntimeError("vLLM server exited before becoming ready.") + try: + if requests.get(f"{BASE_URL}/health", timeout=5).status_code == 200: + break + except requests.RequestException: + pass + if time.monotonic() > deadline: + raise RuntimeError("vLLM server failed to start in time.") + time.sleep(2) + print("[server] Ready.") + return proc -async def main(): - ray.init() +def generate_completions(client: OpenAI, prompts: list[str]) -> list[str]: + """Generate completions for a batch of prompts via the OpenAI HTTP API.""" + results = [] + for prompt in prompts: + response = client.completions.create( + model=SERVED_MODEL_NAME, + prompt=prompt, + max_tokens=32, + temperature=0, + ) + results.append(response.choices[0].text) + return results + +def main(): # Download model weights to local/shared disk once. local_model_path = snapshot_download(MODEL_NAME) print(f"[init] Model downloaded to {local_model_path}") - # FSDP rendezvous address (single-node) + ray.init() + + # FSDP rendezvous address (single-node). fsdp_master_addr = get_ip() fsdp_master_port = get_open_port() - # Launch 4 FSDP training workers. - # Ray allocates 1 GPU per worker; AsyncLLMEngine's internal DP - # placement groups will land on the remaining 4 GPUs. + # Launch the FSDP training workers first so Ray reserves their GPUs, then + # place the inference server on the GPUs Ray did NOT use. This keeps the two + # on disjoint physical GPUs whether ray.init() started a fresh cluster or + # connected to an existing one. fsdp_workers = [ FSDPTrainWorker.remote( local_model_path, @@ -219,126 +262,109 @@ async def main(): ray.get([w.get_rank.remote() for w in fsdp_workers]) print(f"[init] {FSDP_WORLD_SIZE} FSDP training workers ready.") - # Launch vLLM with expert parallelism + data parallelism. - # AsyncLLMEngine with data_parallel_backend="ray" creates its own - # placement groups internally — no manual placement group needed. - print("[engine] Creating AsyncLLMEngine...") - engine = create_async_engine( - model=local_model_path, - enforce_eager=True, - tensor_parallel_size=INFERENCE_TP_SIZE, - data_parallel_size=INFERENCE_DP_SIZE, - enable_expert_parallel=True, - distributed_executor_backend="ray", - data_parallel_backend="ray", - weight_transfer_config=WeightTransferConfig(backend="nccl"), - load_format="dummy", - gpu_memory_utilization=0.7, - ) - print("[engine] AsyncLLMEngine created.") - - prompts = [ - "Hello, my name is", - "The president of the United States is", - "The capital of France is", - "The future of AI is", + # Discover the physical GPUs Ray assigned to training; run the server on the + # complementary GPUs. + training_gpus = { + int(g) + for ids in ray.get([w.get_gpu_ids.remote() for w in fsdp_workers]) + for g in ids + } + num_gpus = int(ray.cluster_resources().get("GPU", 0)) + num_server_gpus = INFERENCE_TP_SIZE * INFERENCE_DP_SIZE + server_gpu_ids = [g for g in range(num_gpus) if g not in training_gpus][ + :num_server_gpus ] - sampling_params = SamplingParams(temperature=0) - - # Generate with dummy weights — expect gibberish. - print("[generate] Starting generation with dummy weights...") - outputs = await generate_batch(engine, prompts, sampling_params) - print("[generate] Generation complete.") - - print("-" * 60) - print("BEFORE weight sync (dummy weights):") - print("-" * 60) - for output in outputs: - print(f"Prompt: {output.prompt!r}") - print(f"Generated: {output.outputs[0].text!r}") - print("-" * 60) - - # --- Weight-transfer setup --- - print("[transfer] Setting up weight-transfer endpoint...") - transfer_addr, transfer_port = ray.get( - fsdp_workers[0].setup_transfer_endpoint.remote() - ) - print(f"[transfer] Endpoint ready at {transfer_addr}:{transfer_port}") - - transfer_world_size = INFERENCE_TP_SIZE * INFERENCE_DP_SIZE + 1 - print( - f"[transfer] World size: {transfer_world_size} " - f"(1 trainer + {INFERENCE_TP_SIZE * INFERENCE_DP_SIZE} vLLM workers)" - ) - - print("[transfer] Initializing NCCL groups...") - train_handle = fsdp_workers[0].init_weight_transfer_group.remote( - transfer_world_size - ) - await engine.init_weight_transfer_engine( - WeightTransferInitRequest( - init_info=asdict( - NCCLWeightTransferInitInfo( - master_address=transfer_addr, - master_port=transfer_port, - rank_offset=1, - world_size=transfer_world_size, - ) - ) + if len(server_gpu_ids) < num_server_gpus: + raise RuntimeError( + f"Need {num_server_gpus} free GPUs for the inference server but only " + f"found {server_gpu_ids} (training uses {sorted(training_gpus)} of " + f"{num_gpus} cluster GPUs)." ) - ) - ray.get(train_handle) - print("[transfer] NCCL groups initialized.") + server_gpus = ",".join(str(g) for g in server_gpu_ids) + print(f"[init] Training GPUs {sorted(training_gpus)}; server GPUs [{server_gpus}].") - # --- Pause, transfer weights, resume --- - print("[sync] Pausing generation...") - await engine.pause_generation(mode="abort") - print("[sync] Generation paused.") + # Start the inference server on the complementary GPUs. + server_proc = start_vllm_server(server_gpus) + try: + client = OpenAI(base_url=f"{BASE_URL}/v1", api_key="EMPTY") - names, dtype_names, shapes = ray.get(fsdp_workers[0].get_weight_metadata.remote()) - print(f"[sync] Got metadata for {len(names)} parameters.") + prompts = [ + "Hello, my name is", + "The president of the United States is", + "The capital of France is", + "The future of AI is", + ] - print("[sync] Starting weight update...") - await engine.start_weight_update() - - print("[sync] Broadcasting weights from FSDP → vLLM...") - broadcast_handles = [ - w.gather_and_broadcast_weights.remote(packed=True) for w in fsdp_workers - ] - await engine.update_weights( - WeightTransferUpdateRequest( - update_info=asdict( - NCCLWeightTransferUpdateInfo( - names=names, - dtype_names=dtype_names, - shapes=shapes, - packed=True, - ) - ) - ) - ) - ray.get(broadcast_handles) - - await engine.finish_weight_update() - print("[sync] Weight broadcast complete.") - - print("[sync] Resuming generation...") - await engine.resume_generation() - print("[sync] Generation resumed.") - - # Generate with synced weights — expect sensible output. - print("[generate] Starting generation with synced weights...") - outputs_updated = await generate_batch(engine, prompts, sampling_params) - print("[generate] Generation complete.") - - print("-" * 60) - print("AFTER weight sync (real weights):") - print("-" * 60) - for output in outputs_updated: - print(f"Prompt: {output.prompt!r}") - print(f"Generated: {output.outputs[0].text!r}") + # Generate with dummy weights — expect gibberish. + print("[generate] Generating with dummy weights...") + outputs = generate_completions(client, prompts) print("-" * 60) + print("BEFORE weight sync (dummy weights):") + print("-" * 60) + for prompt, text in zip(prompts, outputs): + print(f"Prompt: {prompt!r}") + print(f"Generated: {text!r}") + print("-" * 60) + + # --- Weight-transfer setup --- + print("[transfer] Setting up weight-transfer endpoint...") + transfer_addr, transfer_port = ray.get( + fsdp_workers[0].setup_transfer_endpoint.remote() + ) + print(f"[transfer] Endpoint ready at {transfer_addr}:{transfer_port}") + + transfer_world_size = INFERENCE_TP_SIZE * INFERENCE_DP_SIZE + 1 + print( + f"[transfer] World size: {transfer_world_size} " + f"(1 trainer + {INFERENCE_TP_SIZE * INFERENCE_DP_SIZE} vLLM workers)" + ) + + # Build the trainer engine on all FSDP ranks (rank 0 is the sender). The + # sender drives the server's init_weight_transfer_engine (HTTP) while + # opening the trainer NCCL endpoint, so both ends rendezvous together; + # the other ranks build a null-client engine that only gathers. + print("[transfer] Initializing NCCL groups (all FSDP ranks)...") + ray.get( + [ + w.setup_engine.remote( + BASE_URL, transfer_addr, transfer_port, transfer_world_size + ) + for w in fsdp_workers + ] + ) + print("[transfer] NCCL groups initialized.") + + # --- Pause, transfer weights, resume --- + print("[sync] Pausing generation...") + requests.post(f"{BASE_URL}/pause", timeout=60).raise_for_status() + + # All ranks participate in the FSDP all-gather; rank 0 additionally + # drives start/update/finish on the server and the NCCL broadcast. + print("[sync] Broadcasting weights from FSDP → vLLM...") + ray.get([w.gather_and_broadcast_weights.remote() for w in fsdp_workers]) + print("[sync] Weight broadcast complete.") + + print("[sync] Resuming generation...") + requests.post(f"{BASE_URL}/resume", timeout=60).raise_for_status() + + # Generate with synced weights — expect sensible output. + print("[generate] Generating with synced weights...") + outputs_updated = generate_completions(client, prompts) + print("-" * 60) + print("AFTER weight sync (real weights):") + print("-" * 60) + for prompt, text in zip(prompts, outputs_updated): + print(f"Prompt: {prompt!r}") + print(f"Generated: {text!r}") + print("-" * 60) + finally: + print("[server] Shutting down...") + server_proc.terminate() + try: + server_proc.wait(timeout=30) + except subprocess.TimeoutExpired: + server_proc.kill() if __name__ == "__main__": - asyncio.run(main()) + main() diff --git a/examples/rl/rlhf_sparse_nccl.py b/examples/rl/rlhf_sparse_nccl.py index 09cf5bfbaa..f28edf3e7e 100644 --- a/examples/rl/rlhf_sparse_nccl.py +++ b/examples/rl/rlhf_sparse_nccl.py @@ -44,12 +44,14 @@ from transformers import AutoModelForCausalLM, AutoTokenizer from vllm import LLM, SamplingParams from vllm.config import WeightTransferConfig -from vllm.distributed.weight_transfer.nccl_engine import ( - NCCLTrainerSendWeightsArgs, - NCCLWeightTransferEngine, +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 ( - SparseNCCLWeightTransferEngine, + SparseNCCLTrainerInitInfo, SparseWeightPatch, ) from vllm.utils.network_utils import get_ip, get_open_port @@ -87,7 +89,8 @@ class TrainModel: self.model = None self.patched_param = None self.pending_sparse_patches: list[SparseWeightPatch] | None = None - self.model_update_group = None + self.dense_engine = None + self.sparse_engine = None self.master_address = get_ip() self.port = get_open_port() self.reset_model() @@ -112,34 +115,32 @@ class TrainModel: self.port = get_open_port() return self.master_address, self.port - def init_weight_transfer_group(self, world_size: int) -> None: - self.model_update_group = NCCLWeightTransferEngine.trainer_init( - dict( + 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), ) - def get_dense_update_info(self, packed: bool = False) -> tuple[dict, int]: - names = [] - dtype_names = [] - shapes = [] - payload_bytes = 0 - for name, param in self.model.named_parameters(): - names.append(name) - dtype_names.append(str(param.dtype).split(".")[-1]) - shapes.append(list(param.shape)) - payload_bytes += param.numel() * param.element_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) - return ( - dict( - names=names, - dtype_names=dtype_names, - shapes=shapes, - packed=packed, - ), - payload_bytes, + def dense_payload_bytes(self) -> int: + return sum( + param.numel() * param.element_size() for param in self.model.parameters() ) @torch.inference_mode() @@ -173,7 +174,7 @@ class TrainModel: self, prompts: Sequence[str], max_patch_rows: int = MAX_PATCH_ROWS, - ) -> tuple[dict[str, object], list[int], str, int]: + ) -> tuple[list[int], str, int]: selected_token_ids: list[int] = [] special_ids = set(self.tokenizer.all_special_ids) for prompt in prompts: @@ -220,63 +221,60 @@ class TrainModel: flat_indices = ( row_ids.unsqueeze(1).mul(hidden_size).add(column_offsets).reshape(-1) ) - flat_values = self.patched_param[row_ids].reshape(-1).contiguous() + 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.detach() - .float() - .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() ) - update_info = dict( - names=[PATCHED_PARAM_NAME], - dtype_names=[str(self.patched_param.dtype).split(".")[-1]], - shapes=[list(self.patched_param.shape)], - num_updates_list=[flat_indices.numel()], - ) - return update_info, selected_token_ids, patch_digest, sparse_payload_bytes + return selected_token_ids, patch_digest, sparse_payload_bytes - def broadcast_weights(self, packed: bool = False) -> float: - if self.model_update_group is None: - raise RuntimeError("Weight transfer group is not initialized") + 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") - trainer_args = NCCLTrainerSendWeightsArgs( - group=self.model_update_group, - packed=packed, - ) start = time.perf_counter() - NCCLWeightTransferEngine.trainer_send_weights( - iterator=self.model.named_parameters(), - trainer_args=trainer_args, - ) + self.dense_engine.send_weights() torch.accelerator.synchronize() return (time.perf_counter() - start) * 1000.0 - def broadcast_pending_sparse_patch(self) -> float: - if self.model_update_group is None: - raise RuntimeError("Weight transfer group is not initialized") + def init_sparse_engine(self, world_size: int, llm_handle) -> None: + """Build the sparse trainer engine""" + self.sparse_engine = WeightTransferTrainerFactory.trainer_init( + init_info=SparseNCCLTrainerInitInfo( + master_address=self.master_address, + master_port=self.port, + world_size=world_size, + rank=0, # single-GPU trainer is the sole (sender) rank + ), + 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") start = time.perf_counter() - SparseNCCLWeightTransferEngine.trainer_send_weights( - iter(self.pending_sparse_patches), - NCCLTrainerSendWeightsArgs(group=self.model_update_group), - ) + self.sparse_engine.send_weights(self.pending_sparse_patches) torch.accelerator.synchronize() self.pending_sparse_patches = None return (time.perf_counter() - start) * 1000.0 @@ -340,39 +338,17 @@ def run_dense_phase( dense_before = collect_vllm_generations(llm) ray.get(llm.sleep.remote(level=0)) - master_address, master_port = ray.get(train_model.create_rendezvous.remote()) + ray.get(train_model.create_rendezvous.remote()) world_size = ray.get(llm.get_world_size.remote()) + 1 - inference_init = llm.init_weight_transfer_engine.remote( - dict( - init_info=dict( - master_address=master_address, - master_port=master_port, - rank_offset=1, - world_size=world_size, - ) - ) - ) - trainer_init = train_model.init_weight_transfer_group.remote(world_size) - ray.get([trainer_init, inference_init]) - ray.get(llm.start_weight_update.remote()) + ray.get(train_model.init_dense_engine.remote(world_size, llm)) - dense_update_info, dense_payload_bytes = ray.get( - train_model.get_dense_update_info.remote() - ) - _, selected_token_ids, patch_digest, _ = ray.get( + 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) ) - inference_update = llm.update_weights.remote( - dict(update_info=dense_update_info) - ) - dense_send_ms, _ = ray.get( - [ - train_model.broadcast_weights.remote(packed=False), - inference_update, - ] - ) - ray.get(llm.finish_weight_update.remote()) + # 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) @@ -386,6 +362,7 @@ def run_dense_phase( "dense_send_ms": dense_send_ms, } finally: + ray.get(train_model.shutdown_engines.remote()) ray.kill(llm) @@ -399,36 +376,14 @@ def run_sparse_phase( sparse_before = collect_vllm_generations(llm) ray.get(llm.sleep.remote(level=0)) - master_address, master_port = ray.get(train_model.create_rendezvous.remote()) + ray.get(train_model.create_rendezvous.remote()) world_size = ray.get(llm.get_world_size.remote()) + 1 - inference_init = llm.init_weight_transfer_engine.remote( - dict( - init_info=dict( - master_address=master_address, - master_port=master_port, - rank_offset=1, - world_size=world_size, - ) - ) - ) - trainer_init = train_model.init_weight_transfer_group.remote(world_size) - ray.get([trainer_init, inference_init]) - ray.get(llm.start_weight_update.remote()) + ray.get(train_model.init_sparse_engine.remote(world_size, llm)) - sparse_update_info, selected_token_ids, patch_digest, sparse_payload_bytes = ( - ray.get(train_model.prepare_sparse_patch.remote(PROMPTS)) + selected_token_ids, patch_digest, sparse_payload_bytes = ray.get( + train_model.prepare_sparse_patch.remote(PROMPTS) ) - - inference_update = llm.update_weights.remote( - dict(update_info=sparse_update_info) - ) - sparse_send_ms, _ = ray.get( - [ - train_model.broadcast_pending_sparse_patch.remote(), - inference_update, - ] - ) - ray.get(llm.finish_weight_update.remote()) + sparse_send_ms = ray.get(train_model.send_pending_sparse_patch.remote()) ray.get(llm.wake_up.remote(tags=["scheduling"])) sparse_after = collect_vllm_generations(llm) @@ -442,6 +397,7 @@ def run_sparse_phase( "sparse_send_ms": sparse_send_ms, } finally: + ray.get(train_model.shutdown_engines.remote()) ray.kill(llm) diff --git a/tests/distributed/test_packed_tensor.py b/tests/distributed/test_packed_tensor.py index 6a2ab78446..5b6846af76 100644 --- a/tests/distributed/test_packed_tensor.py +++ b/tests/distributed/test_packed_tensor.py @@ -9,7 +9,6 @@ These utilities enable efficient batched tensor transfer over NCCL. import pytest import torch -from vllm.distributed.weight_transfer.nccl_engine import NCCLWeightTransferUpdateInfo from vllm.distributed.weight_transfer.packed_tensor import ( pack_tensors, packed_ipc_consumer, @@ -68,32 +67,6 @@ def create_state_dict_info( return {name: (tuple(tensor.shape), tensor.dtype) for name, tensor in params} -# --- Unit Tests: NCCLWeightTransferUpdateInfo packed field --- - - -class TestNCCLWeightTransferUpdateInfoPacked: - """Test NCCLWeightTransferUpdateInfo dataclass packed field.""" - - def test_packed_default_false(self): - """Test that packed defaults to False.""" - info = NCCLWeightTransferUpdateInfo( - names=["layer.weight"], - dtype_names=["float32"], - shapes=[[10, 10]], - ) - assert info.packed is False - - def test_packed_can_be_set_true(self): - """Test that packed can be set to True.""" - info = NCCLWeightTransferUpdateInfo( - names=["layer.weight"], - dtype_names=["float32"], - shapes=[[10, 10]], - packed=True, - ) - assert info.packed is True - - # --- Unit Tests: packed_nccl_broadcast_producer --- diff --git a/tests/distributed/test_weight_transfer.py b/tests/distributed/test_weight_transfer.py index ddc6dec04e..05a57dbea5 100644 --- a/tests/distributed/test_weight_transfer.py +++ b/tests/distributed/test_weight_transfer.py @@ -7,6 +7,8 @@ Integration tests for NCCL and IPC weight transfer between processes using Ray. """ import pickle +import threading +import time from unittest.mock import MagicMock import pybase64 as base64 @@ -20,9 +22,11 @@ from vllm.config.weight_transfer import WeightTransferConfig from vllm.distributed.weight_transfer import ( HTTPVLLMWeightSyncClient, ModuleSource, + ParamMeta, RayVLLMWeightSyncClient, TrainerWeightTransferEngine, VLLMWeightSyncClient, + WeightSource, WeightTransferEngineFactory, WeightTransferTrainerFactory, ) @@ -39,11 +43,19 @@ from vllm.distributed.weight_transfer.ipc_engine import ( IPCWeightTransferUpdateInfo, ) from vllm.distributed.weight_transfer.nccl_engine import ( + NCCLTrainerInitInfo, + NCCLTrainerWeightTransferEngine, NCCLWeightTransferEngine, NCCLWeightTransferInitInfo, NCCLWeightTransferUpdateInfo, ) +from vllm.distributed.weight_transfer.packed_tensor import ( + DEFAULT_PACKED_BUFFER_SIZE_BYTES, + DEFAULT_PACKED_NUM_BUFFERS, +) from vllm.distributed.weight_transfer.sparse_nccl_engine import ( + SparseNCCLTrainerInitInfo, + SparseNCCLTrainerWeightTransferEngine, SparseNCCLWeightTransferEngine, SparseNCCLWeightTransferUpdateInfo, SparseWeightPatch, @@ -535,11 +547,15 @@ def inference_receive_tensor( _vllm_config_mod.set_current_vllm_config = lambda cfg: contextlib.nullcontext() # Initialize the engine (joins as rank 1) + # Trainer broadcasts a single tensor unpacked, so the worker must not + # expect the packed wire format (packed is a must-agree wire param shipped + # on the init info). init_info = NCCLWeightTransferInitInfo( master_address=master_address, master_port=master_port, rank_offset=1, # Trainer is rank 0, we become rank 1 world_size=world_size, + packed=False, ) engine.init_transfer_engine(init_info) @@ -616,39 +632,53 @@ def trainer_broadcast_sparse_tensor( master_port: int, world_size: int, ) -> bool: - """Trainer task that broadcasts sparse patches via NCCL.""" + """Trainer task that broadcasts sparse patches via the trainer engine. + + 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.device_communicators.pynccl import PyNcclCommunicator - from vllm.distributed.utils import StatelessProcessGroup - from vllm.distributed.weight_transfer.nccl_engine import ( - NCCLTrainerSendWeightsArgs, - ) + from vllm.distributed.weight_transfer import WeightTransferTrainerFactory from vllm.distributed.weight_transfer.sparse_nccl_engine import ( - SparseNCCLWeightTransferEngine, + SparseNCCLTrainerInitInfo, SparseWeightPatch, ) - pg = StatelessProcessGroup.create( - host=master_address, - port=master_port, - rank=0, - world_size=world_size, - ) - comm = PyNcclCommunicator(pg, device=device.index) + class NoopClient: + def init_weight_transfer_engine(self, init_info): + pass + + def start_weight_update(self): + pass + + def update_weights(self, update_info): + pass + + def finish_weight_update(self): + pass 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), ) - SparseNCCLWeightTransferEngine.trainer_send_weights( - iter([patch]), - NCCLTrainerSendWeightsArgs(group=comm), + 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 @@ -1369,14 +1399,26 @@ class TestModuleSource: source = ModuleSource(_module_with(("w", torch.zeros(2)))) assert [n for n, _ in source] == [n for n, _ in source] == ["w"] + def test_metadata_agrees_with_iteration(self): + """The two channels must line up element-for-element: engines declare + the round from `metadata()` and then send what iteration yields.""" + source = ModuleSource( + _module_with(("w", torch.zeros(2, 3)), ("b", torch.zeros(3))) + ) + meta = source.metadata() + pairs = list(source) + assert [m.name for m in meta] == [name for name, _ in pairs] + assert [m.dtype for m in meta] == [t.dtype for _, t in pairs] + assert [m.shape for m in meta] == [tuple(t.shape) for _, t in pairs] + class TestTrainerFactory: """WeightTransferTrainerFactory registry mechanics.""" - def test_registry_has_ipc(self): - # IPC is the first backend migrated to the trainer engine; NCCL / - # sparse NCCL register in later PRs. + def test_registry_has_all_backends(self): + assert "nccl" in WeightTransferTrainerFactory._registry assert "ipc" in WeightTransferTrainerFactory._registry + assert "sparse_nccl" in WeightTransferTrainerFactory._registry def test_register_and_dispatch(self): saved = dict(WeightTransferTrainerFactory._registry) @@ -1406,6 +1448,12 @@ class TestTrainerFactory: def test_ipc_init_info_declares_backend(self): assert IPCTrainerInitInfo.backend == "ipc" + def test_nccl_init_info_declares_backend(self): + assert NCCLTrainerInitInfo.backend == "nccl" + + def test_sparse_nccl_init_info_declares_backend(self): + assert SparseNCCLTrainerInitInfo.backend == "sparse_nccl" + def test_trainer_init_info_subclass_must_set_backend(self): with pytest.raises(TypeError, match="class-level `backend`"): @@ -1475,3 +1523,408 @@ def test_ipc_trainer_init_ships_packed_to_worker(): assert engine.packed is True assert client.order == ["init"] assert client.last_init_info == {"packed": True} + + +def test_nccl_trainer_init_ships_worker_init_info(monkeypatch): + """The sender's trainer_init drives the inference-side init handshake with + the worker-shaped init info (rank_offset=1) while opening its own endpoint, + and propagates the must-agree wire params to the worker.""" + import vllm.distributed.weight_transfer.nccl_engine as nccl_engine_mod + + # Bypass the real NCCL rendezvous. + monkeypatch.setattr( + nccl_engine_mod, "open_trainer_endpoint", lambda info: MagicMock() + ) + + client = RecordingClient() + engine = WeightTransferTrainerFactory.trainer_init( + init_info=NCCLTrainerInitInfo( + master_address="127.0.0.1", + master_port=29500, + world_size=3, + rank=0, + packed=True, + packed_buffer_size_bytes=1024, + packed_num_buffers=3, + ), + client=client, + source=ModuleSource(_module_with(("w", torch.zeros(4)))), + ) + + assert isinstance(engine, NCCLTrainerWeightTransferEngine) + assert engine.is_sender is True + assert engine.packed is True + assert client.order == ["init"] + assert client.last_init_info == { + "master_address": "127.0.0.1", + "master_port": 29500, + "rank_offset": 1, + "world_size": 3, + "packed": True, + "packed_buffer_size_bytes": 1024, + "packed_num_buffers": 3, + } + + +def test_nccl_worker_learns_wire_params_from_init_handshake(monkeypatch): + """The worker engine reads packed + buffer geometry from the + trainer-supplied init info at the handshake, not from the config or the + per-round update info.""" + import vllm.distributed.weight_transfer.nccl_engine as nccl_engine_mod + + monkeypatch.setattr( + nccl_engine_mod, "worker_init_process_group", lambda info, pc: MagicMock() + ) + + engine = NCCLWeightTransferEngine( + WeightTransferConfig(backend="nccl"), + create_mock_vllm_config(), + torch.device("cuda:0"), + MagicMock(spec=torch.nn.Module), + ) + assert engine.packed is False # pre-handshake default (legacy unpacked) + engine.init_transfer_engine( + NCCLWeightTransferInitInfo( + master_address="127.0.0.1", + master_port=29500, + rank_offset=1, + world_size=2, + packed=True, + packed_buffer_size_bytes=2048, + packed_num_buffers=4, + ) + ) + + assert engine.packed is True + assert engine.packed_buffer_size_bytes == 2048 + assert engine.packed_num_buffers == 4 + + +def test_nccl_trainer_init_non_sender_skips_rendezvous_and_client(): + """Non-sender trainer ranks build an engine without opening an endpoint or + touching the client; they only join the collectives in send_weights.""" + client = RecordingClient() + engine = WeightTransferTrainerFactory.trainer_init( + init_info=NCCLTrainerInitInfo( + master_address="127.0.0.1", + master_port=29500, + world_size=3, + rank=1, + ), + client=client, + source=ModuleSource(_module_with(("w", torch.zeros(4)))), + ) + + assert engine.is_sender is False + assert engine.model_update_group is None + assert client.order == [] + + # send_weights on a non-sender only iterates the source (packed mode needs + # no CUDA stream on non-senders), never the client. + engine.send_weights() + assert client.order == [] + + +@pytest.mark.skipif( + torch.accelerator.device_count() < 1, + reason="Need at least 1 GPU (NCCL broadcast / CUDA stream).", +) +def test_nccl_trainer_send_weights_drives_client_in_order(): + """send_weights issues start -> update -> finish and ships per-round + metadata; the packed wire params ride the init handshake, not the + per-round update_info.""" + client = RecordingClient() + engine = NCCLTrainerWeightTransferEngine( + client=client, + source=ModuleSource(_module_with(("w", torch.zeros(4, device="cuda")))), + packed=False, + ) + # Bypass the real NCCL rendezvous; broadcast is a no-op. + engine.model_update_group = MagicMock() + + engine.send_weights() + + assert client.order == ["start", "update", "finish"] + assert client.last_update_info is not None + assert client.last_update_info["names"] == ["w"] + assert client.last_update_info["shapes"] == [[4]] + assert "packed" not in client.last_update_info + + +class _ScriptedSource(WeightSource): + """Declares `meta` but yields whatever `pairs` says — used to drive the + metadata/iteration agreement checks.""" + + def __init__(self, meta, pairs): + self._meta = meta + self._pairs = pairs + + def metadata(self): + return list(self._meta) + + def __iter__(self): + yield from self._pairs + + +def _mock_group_engine(source, monkeypatch, **kwargs): + """Unpacked trainer engine with a mocked group and stream (no GPU needed).""" + engine = NCCLTrainerWeightTransferEngine( + client=RecordingClient(), source=source, packed=False, **kwargs + ) + engine.model_update_group = MagicMock() + monkeypatch.setattr(torch.cuda, "current_stream", MagicMock()) + return engine + + +def test_nccl_trainer_init_requires_source(): + """NCCL is a full-resync backend: it cannot run without a WeightSource.""" + with pytest.raises(ValueError, match="requires a WeightSource"): + NCCLTrainerWeightTransferEngine.trainer_init( + NCCLTrainerInitInfo( + master_address="127.0.0.1", master_port=29500, world_size=2, rank=0 + ), + client=RecordingClient(), + ) + + +def test_nccl_trainer_send_weights_rejects_reordered_source(monkeypatch): + """The worker sizes its buffers (and cuts packed chunks) from metadata(), so + iteration disagreeing with it must raise rather than corrupt the stream.""" + meta = [ + ParamMeta("w", torch.float32, (4,)), + ParamMeta("b", torch.float32, (2,)), + ] + reordered = [("b", torch.zeros(2)), ("w", torch.zeros(4))] + engine = _mock_group_engine(_ScriptedSource(meta, reordered), monkeypatch) + + with pytest.raises(ValueError, match="disagrees with iteration at index 0"): + engine.send_weights() + + +def test_nccl_trainer_send_weights_rejects_dtype_disagreement(monkeypatch): + """A source that declares one wire dtype and materializes another would make + the two sides disagree on every byte offset.""" + meta = [ParamMeta("w", torch.float32, (4,))] + engine = _mock_group_engine( + _ScriptedSource(meta, [("w", torch.zeros(4, dtype=torch.bfloat16))]), + monkeypatch, + ) + + with pytest.raises(ValueError, match="disagrees with iteration"): + engine.send_weights() + + +def test_nccl_trainer_send_weights_rejects_truncated_source(monkeypatch): + """Yielding fewer parameters than declared leaves the worker waiting.""" + meta = [ + ParamMeta("w", torch.float32, (4,)), + ParamMeta("b", torch.float32, (2,)), + ] + engine = _mock_group_engine( + _ScriptedSource(meta, [("w", torch.zeros(4))]), monkeypatch + ) + + with pytest.raises(ValueError, match="yielded 1 parameters"): + engine.send_weights() + + +def test_nccl_trainer_send_weights_broadcasts_contiguous(monkeypatch): + """NCCL sends numel elements from data_ptr(), so a non-contiguous view must + be linearized first or the worker receives unrelated memory.""" + base = torch.arange(6, dtype=torch.float32).reshape(2, 3) + view = base.t() # non-contiguous + meta = [ParamMeta("w", torch.float32, tuple(view.shape))] + engine = _mock_group_engine(_ScriptedSource(meta, [("w", view)]), monkeypatch) + + engine.send_weights() + + sent = engine.model_update_group.broadcast.call_args.args[0] + assert sent.is_contiguous() + assert torch.equal(sent, view) + + +def test_nccl_trainer_send_weights_raises_instead_of_hanging(monkeypatch): + """A failed broadcast must surface even while the inference-side + update_weights is still blocked in its matching NCCL call. + + Joining the RPC thread there would deadlock: the worker only returns once + the broadcast it is waiting for arrives, which never happens. + """ + rpc_entered = threading.Event() + release_rpc = threading.Event() + + class _BlockingClient(RecordingClient): + def update_weights(self, update_info): + rpc_entered.set() + release_rpc.wait(timeout=30) # stands in for a wedged NCCL recv + super().update_weights(update_info) + + class _FailingSource(WeightSource): + def metadata(self): + return [ParamMeta("w", torch.float32, (4,))] + + def __iter__(self): + # Fail only once the RPC is provably in flight. + rpc_entered.wait(timeout=60) + raise RuntimeError("broadcast blew up") + + engine = NCCLTrainerWeightTransferEngine( + client=_BlockingClient(), source=_FailingSource(), packed=False + ) + engine.model_update_group = MagicMock() + monkeypatch.setattr(torch.cuda, "current_stream", MagicMock()) + + started = time.perf_counter() + try: + with pytest.raises(RuntimeError, match="broadcast blew up"): + engine.send_weights() + elapsed = time.perf_counter() - started + assert rpc_entered.is_set(), "the RPC was never in flight" + assert not release_rpc.is_set(), "send_weights waited for the wedged RPC" + # Regression guard: joining the RPC thread would park here until the + # client's own timeout expires instead of raising immediately. + assert elapsed < 5.0, f"send_weights blocked {elapsed:.1f}s on the RPC" + finally: + release_rpc.set() + + +def _sparse_patch(device: str = "cpu") -> SparseWeightPatch: + return SparseWeightPatch( + name="w", + indices=torch.tensor([1, 3], dtype=torch.int32, device=device), + values=torch.tensor([1.0, 2.0], dtype=torch.float32, device=device), + full_shape=(4, 4), + ) + + +def test_sparse_nccl_trainer_init_ships_worker_init_info(monkeypatch): + """The sender's trainer_init drives the init handshake with the + worker-shaped init info; sparse ships no packed wire params, so the worker + keeps its unpacked defaults. Sparse takes no `source`.""" + import vllm.distributed.weight_transfer.sparse_nccl_engine as sparse_mod + + monkeypatch.setattr(sparse_mod, "open_trainer_endpoint", lambda info: MagicMock()) + + client = RecordingClient() + engine = WeightTransferTrainerFactory.trainer_init( + init_info=SparseNCCLTrainerInitInfo( + master_address="127.0.0.1", + master_port=29500, + world_size=2, + rank=0, + ), + client=client, + ) + + assert isinstance(engine, SparseNCCLTrainerWeightTransferEngine) + assert client.order == ["init"] + assert client.last_init_info == { + "master_address": "127.0.0.1", + "master_port": 29500, + "rank_offset": 1, + "world_size": 2, + "packed": False, + "packed_buffer_size_bytes": DEFAULT_PACKED_BUFFER_SIZE_BYTES, + "packed_num_buffers": DEFAULT_PACKED_NUM_BUFFERS, + } + + +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.""" + client = RecordingClient() + engine = SparseNCCLTrainerWeightTransferEngine(client=client) + engine.model_update_group = MagicMock() + # 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()) + + engine.send_weights([_sparse_patch()]) + + assert client.order == ["start", "update", "finish"] + assert client.last_update_info is not None + assert client.last_update_info["names"] == ["w"] + assert client.last_update_info["shapes"] == [[4, 4]] + 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 + + +def test_sparse_nccl_trainer_send_weights_empty_round_is_noop(): + """A round with no patches must not touch the client (an empty sparse + update info is invalid by construction).""" + client = RecordingClient() + engine = SparseNCCLTrainerWeightTransferEngine(client=client) + engine.model_update_group = MagicMock() + + engine.send_weights([]) + engine.send_weights() # no argument is also a no-op round + + assert client.order == [] + + +def test_sparse_nccl_trainer_send_weights_requires_full_shape(): + patch = _sparse_patch() + patch.full_shape = None + engine = SparseNCCLTrainerWeightTransferEngine(client=RecordingClient()) + engine.model_update_group = MagicMock() + + with pytest.raises(ValueError, match="full_shape"): + engine.send_weights([patch]) + + +def test_sparse_nccl_trainer_rejects_source(): + """Sparse is a delta backend; a WeightSource would silently never be sent.""" + with pytest.raises(ValueError, match="takes no WeightSource"): + SparseNCCLTrainerWeightTransferEngine( + client=RecordingClient(), + source=ModuleSource(_module_with(("w", torch.zeros(2)))), + ) + + +def test_sparse_nccl_trainer_validates_patch_before_any_rpc(): + """Malformed patches must fail on the trainer, before start_weight_update: + the worker's own checks only run once the broadcasts are already in flight, + where a size mismatch wedges both sides instead of raising.""" + client = RecordingClient() + engine = SparseNCCLTrainerWeightTransferEngine(client=client) + engine.model_update_group = MagicMock() + + mismatched = SparseWeightPatch( + name="w", + indices=torch.tensor([1, 3], dtype=torch.int32), + values=torch.tensor([1.0], dtype=torch.float32), + full_shape=(4, 4), + ) + with pytest.raises(ValueError, match="matching lengths"): + engine.send_weights([mismatched]) + + wrong_index_dtype = SparseWeightPatch( + name="w", + indices=torch.tensor([1, 3], dtype=torch.int64), + values=torch.tensor([1.0, 2.0], dtype=torch.float32), + full_shape=(4, 4), + ) + with pytest.raises(ValueError, match="int32 indices"): + engine.send_weights([wrong_index_dtype]) + + assert client.order == [] + + +def test_sparse_nccl_trainer_non_sender_skips_client(): + client = RecordingClient() + engine = WeightTransferTrainerFactory.trainer_init( + init_info=SparseNCCLTrainerInitInfo( + master_address="127.0.0.1", + master_port=29500, + world_size=2, + rank=1, + ), + client=client, + ) + + assert engine.is_sender is False + assert engine.model_update_group is None + assert isinstance(engine, SparseNCCLTrainerWeightTransferEngine) + engine.send_weights([_sparse_patch()]) + assert client.order == [] diff --git a/vllm/distributed/weight_transfer/base.py b/vllm/distributed/weight_transfer/base.py index 3fd21101e7..ac36f90378 100644 --- a/vllm/distributed/weight_transfer/base.py +++ b/vllm/distributed/weight_transfer/base.py @@ -77,6 +77,7 @@ class WeightSource(ABC): @abstractmethod def metadata(self) -> list[ParamMeta]: + """Declare what iteration will yield, without transferring anything.""" raise NotImplementedError @abstractmethod @@ -402,8 +403,7 @@ class TrainerWeightTransferEngine(ABC, Generic[TTrainerInitInfo]): Symmetric to `WeightTransferEngine` but lives in the training process. Constructed via the `trainer_init` factory classmethod; carries any backend-specific state (NCCL communicators, IPC device info, transfer - plans) on `self`. The `WeightSource` is required at `trainer_init`, - then replayed each round by the no-argument `send_weights()`. + plans) on `self`. Unlike the worker engine, the trainer side does not take a `WeightTransferConfig`: the backend is selected from the init info's @@ -431,7 +431,7 @@ class TrainerWeightTransferEngine(ABC, Generic[TTrainerInitInfo]): self, *, client: "VLLMWeightSyncClient", - source: "WeightSource", + source: "WeightSource | None" = None, is_sender: bool = True, ) -> None: self.is_sender = is_sender @@ -447,24 +447,24 @@ class TrainerWeightTransferEngine(ABC, Generic[TTrainerInitInfo]): init_info: TTrainerInitInfo, *, client: "VLLMWeightSyncClient", - source: "WeightSource", + source: "WeightSource | None" = None, ) -> Self: """Rendezvous with the inference side and return a ready instance. Called on every trainer rank. The sender drives the full handshake via `client` (build the worker-side init info, call `client.init_weight_transfer_engine`, open the trainer-side endpoint); - non-sender ranks skip the rendezvous and the RPC. `source` is stored on - `self.source`; after return, `send_weights()` is callable. + non-sender ranks skip the rendezvous and the RPC. """ raise NotImplementedError @abstractmethod def send_weights(self) -> None: - """Push `self.source`'s weights to inference workers and drive the full - update round trip: `start_weight_update`, `update_weights` (run - concurrently with the trainer-side broadcast when the backend requires - it), then `finish_weight_update`. Called on every trainer rank.""" + """Push weights to inference workers and drive the full update round + trip: `start_weight_update`, `update_weights` (run concurrently with the + trainer-side broadcast when the backend requires it), then + `finish_weight_update`. Called on every trainer rank. + """ raise NotImplementedError def shutdown(self) -> None: diff --git a/vllm/distributed/weight_transfer/factory.py b/vllm/distributed/weight_transfer/factory.py index 1172b46480..fe2ab1bc9e 100644 --- a/vllm/distributed/weight_transfer/factory.py +++ b/vllm/distributed/weight_transfer/factory.py @@ -167,7 +167,7 @@ class WeightTransferTrainerFactory: init_info: "TrainerInitInfo", *, client: "VLLMWeightSyncClient", - source: "WeightSource", + source: "WeightSource | None" = None, ) -> TrainerWeightTransferEngine: """Build and rendezvous a ready-to-send trainer engine. @@ -184,7 +184,9 @@ class WeightTransferTrainerFactory: selects the engine; it also carries the wire params (e.g. `packed`). client: Inference-side control-plane client. - source: `WeightSource` of `(name, tensor)` pairs to send each round. + source: `WeightSource` of `(name, tensor)` pairs to send each round, + for full-resync backends (NCCL, IPC). Sparse backend + omits it and passes its per-round payload to `send_weights`. Raises: ValueError: If `init_info.backend` is not registered. @@ -233,10 +235,21 @@ WeightTransferEngineFactory.register_engine( ) -# Trainer-side engines. Backends register here as they migrate to the stateful -# trainer engine; NCCL / sparse NCCL keep their static trainer path until then. +# Trainer-side engines, parallel to the worker registry above. +WeightTransferTrainerFactory.register_engine( + "nccl", + "vllm.distributed.weight_transfer.nccl_engine", + "NCCLTrainerWeightTransferEngine", +) + WeightTransferTrainerFactory.register_engine( "ipc", "vllm.distributed.weight_transfer.ipc_engine", "IPCTrainerWeightTransferEngine", ) + +WeightTransferTrainerFactory.register_engine( + "sparse_nccl", + "vllm.distributed.weight_transfer.sparse_nccl_engine", + "SparseNCCLTrainerWeightTransferEngine", +) diff --git a/vllm/distributed/weight_transfer/ipc_engine.py b/vllm/distributed/weight_transfer/ipc_engine.py index d512c9898c..8a7d8744c8 100644 --- a/vllm/distributed/weight_transfer/ipc_engine.py +++ b/vllm/distributed/weight_transfer/ipc_engine.py @@ -291,8 +291,10 @@ class IPCTrainerWeightTransferEngine(TrainerWeightTransferEngine[IPCTrainerInitI init_info: IPCTrainerInitInfo, *, client: VLLMWeightSyncClient, - source: WeightSource, + source: WeightSource | None = None, ) -> Self: + if source is None: + raise ValueError("IPC trainer weight transfer requires a WeightSource.") engine = cls( client=client, source=source, @@ -307,6 +309,7 @@ class IPCTrainerWeightTransferEngine(TrainerWeightTransferEngine[IPCTrainerInitI return engine def send_weights(self) -> None: + assert self.source is not None source = self.source if self.is_sender: self.client.start_weight_update() diff --git a/vllm/distributed/weight_transfer/nccl_common.py b/vllm/distributed/weight_transfer/nccl_common.py index fd2f2a3514..de72063c32 100644 --- a/vllm/distributed/weight_transfer/nccl_common.py +++ b/vllm/distributed/weight_transfer/nccl_common.py @@ -9,7 +9,7 @@ sparse engine does not have to subclass the dense one. """ from dataclasses import dataclass -from typing import TYPE_CHECKING +from typing import TYPE_CHECKING, Protocol import torch @@ -18,16 +18,36 @@ if TYPE_CHECKING: from vllm.distributed.device_communicators.pynccl import PyNcclCommunicator from vllm.distributed.weight_transfer.base import WeightTransferInitInfo +from vllm.distributed.weight_transfer.packed_tensor import ( + DEFAULT_PACKED_BUFFER_SIZE_BYTES, + DEFAULT_PACKED_NUM_BUFFERS, +) @dataclass class NCCLWeightTransferInitInfo(WeightTransferInitInfo): - """Initialization info for NCCL-based weight transfer backends.""" + """Worker-side initialization info for NCCL-based weight transfer backends.""" master_address: str master_port: int rank_offset: int world_size: int + packed: bool = False + packed_buffer_size_bytes: int = DEFAULT_PACKED_BUFFER_SIZE_BYTES + packed_num_buffers: int = DEFAULT_PACKED_NUM_BUFFERS + + +class NCCLRendezvous(Protocol): + """The three fields `trainer_init` needs to open the trainer endpoint. + + Structural so each backend can keep its own trainer init info next to its + engine (dense `NCCLTrainerInitInfo`, `SparseNCCLTrainerInitInfo`) without + this shared module importing either. + """ + + master_address: str + master_port: int + world_size: int def stateless_init_process_group( @@ -83,7 +103,7 @@ def worker_init_process_group( def trainer_init( - init_info: NCCLWeightTransferInitInfo | dict, + init_info: NCCLRendezvous | dict, ) -> "PyNcclCommunicator": """ Initialize NCCL process group for trainer-side weight transfer. @@ -92,7 +112,8 @@ def trainer_init( CUDA device (torch.accelerator.current_device_index()). Args: - init_info: Either an NCCLWeightTransferInitInfo object or a dict with keys: + init_info: Any object carrying the `NCCLRendezvous` fields (a trainer or + worker NCCL init info), or a dict with keys: - master_address: str - master_port: int - world_size: int diff --git a/vllm/distributed/weight_transfer/nccl_engine.py b/vllm/distributed/weight_transfer/nccl_engine.py index 0ba44f3486..7fcca51651 100644 --- a/vllm/distributed/weight_transfer/nccl_engine.py +++ b/vllm/distributed/weight_transfer/nccl_engine.py @@ -2,11 +2,13 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project """NCCL-based (dense) weight transfer engine.""" -from collections.abc import Callable, Iterator -from dataclasses import dataclass -from typing import TYPE_CHECKING, Any +from collections.abc import Iterator +from concurrent.futures import ThreadPoolExecutor +from dataclasses import asdict, dataclass +from typing import TYPE_CHECKING, Any, ClassVar import torch +from typing_extensions import Self if TYPE_CHECKING: from vllm.config import VllmConfig @@ -14,72 +16,76 @@ if TYPE_CHECKING: from vllm.config.weight_transfer import WeightTransferConfig from vllm.distributed.weight_transfer.base import ( + ParamMeta, + TrainerInitInfo, + TrainerWeightTransferEngine, + VLLMWeightSyncClient, + WeightSource, WeightTransferEngine, WeightTransferUpdateInfo, ) from vllm.distributed.weight_transfer.nccl_common import ( NCCLWeightTransferInitInfo, - trainer_init, worker_init_process_group, ) +from vllm.distributed.weight_transfer.nccl_common import ( + trainer_init as open_trainer_endpoint, +) from vllm.distributed.weight_transfer.packed_tensor import ( DEFAULT_PACKED_BUFFER_SIZE_BYTES, DEFAULT_PACKED_NUM_BUFFERS, packed_nccl_broadcast_consumer, + packed_nccl_broadcast_producer, ) -# Re-exported for backward compatibility; canonical home is nccl_common. +# NCCLWeightTransferInitInfo is re-exported here for convenience; its canonical +# home is nccl_common, shared with the sparse backend. __all__ = [ "NCCLWeightTransferInitInfo", - "NCCLTrainerSendWeightsArgs", + "NCCLTrainerInitInfo", "NCCLWeightTransferUpdateInfo", "NCCLWeightTransferEngine", + "NCCLTrainerWeightTransferEngine", ] @dataclass -class NCCLTrainerSendWeightsArgs: - """Arguments for NCCL trainer_send_weights method.""" +class NCCLTrainerInitInfo(TrainerInitInfo): + """Trainer-side init info for the dense NCCL weight transfer backend. - group: Any - """Process group (PyNcclCommunicator) for NCCL communication.""" - src: int = 0 - """Source rank (default 0, trainer is typically rank 0).""" - post_iter_func: Callable[[tuple[str, torch.Tensor]], torch.Tensor] | None = None - """Optional function to apply to each (name, tensor) pair before broadcasting. - If None, extracts just the tensor.""" - packed: bool = False - """Whether to use packed tensor broadcasting for efficiency. - When True, multiple tensors are batched together before broadcasting - to reduce NCCL communication overhead.""" - stream: torch.cuda.Stream | None = None - """CUDA stream to use for broadcasting if packed is False. - If packed is True, new streams will be created for each buffer.""" + The sender opens its endpoint as NCCL rank 0, so it needs no `rank_offset`. + `world_size` is the full trainer+worker NCCL group size. `rank` (from + `TrainerInitInfo`) identifies this trainer process; rank 0 is the sender. + + `packed` / buffer sizes are the transfer's wire params. The trainer + propagates them to the worker at `trainer_init` so the two sides cannot + disagree. Note this defaults to packed, unlike the worker-side + `NCCLWeightTransferInitInfo`, whose default only applies when no trainer + ships a value. `backend` is the factory dispatch key.""" + + backend: ClassVar[str] = "nccl" + + master_address: str + master_port: int + world_size: int + packed: bool = True packed_buffer_size_bytes: int = DEFAULT_PACKED_BUFFER_SIZE_BYTES - """Size in bytes for each packed tensor buffer. - Must match the value used in NCCLWeightTransferUpdateInfo.""" packed_num_buffers: int = DEFAULT_PACKED_NUM_BUFFERS - """Number of buffers for double/triple buffering during packed transfer. - Must match the value used in NCCLWeightTransferUpdateInfo.""" @dataclass class NCCLWeightTransferUpdateInfo(WeightTransferUpdateInfo): - """Update info for NCCL weight transfer backend.""" + """Per-round update info for the dense NCCL weight transfer backend. + + Whether the transfer is packed (and the buffer geometry) is a must-agree + wire param carried on the init info (`NCCLTrainerInitInfo` / + `NCCLWeightTransferInitInfo`), not here; this carries only the per-round + parameter metadata. + """ names: list[str] dtype_names: list[str] shapes: list[list[int]] - packed: bool = False - """Whether to use packed tensor broadcasting for efficiency. - When True, multiple tensors are batched together before broadcasting - to reduce NCCL communication overhead.""" - packed_buffer_size_bytes: int = DEFAULT_PACKED_BUFFER_SIZE_BYTES - """Size in bytes for each packed tensor buffer. - Both producer and consumer must use the same value.""" - packed_num_buffers: int = DEFAULT_PACKED_NUM_BUFFERS - """Number of buffers for double/triple buffering during packed transfer. - Both producer and consumer must use the same value.""" def __post_init__(self): """Validate that all lists have the same length.""" @@ -121,15 +127,25 @@ class NCCLWeightTransferEngine( ) -> None: super().__init__(config, vllm_config, device, model) self.model_update_group: PyNcclCommunicator | None = None + # Set from the trainer-supplied init info at the handshake; defaults are + # only for the (unreachable) receive-before-init case. + self.packed = False + self.packed_buffer_size_bytes = DEFAULT_PACKED_BUFFER_SIZE_BYTES + self.packed_num_buffers = DEFAULT_PACKED_NUM_BUFFERS def init_transfer_engine(self, init_info: NCCLWeightTransferInitInfo) -> None: """ - Initialize NCCL process group with the trainer. + Initialize NCCL process group with the trainer and record the + trainer-supplied wire params so the worker decodes exactly as the + trainer encodes. Args: init_info: NCCL initialization info containing master address, port, - rank offset, and world size + rank offset, world size, and the packed wire params """ + self.packed = init_info.packed + self.packed_buffer_size_bytes = init_info.packed_buffer_size_bytes + self.packed_num_buffers = init_info.packed_num_buffers self.model_update_group = worker_init_process_group( init_info, self.parallel_config ) @@ -154,13 +170,14 @@ class NCCLWeightTransferEngine( """ Receive weights from trainer via NCCL broadcast. - If update_info.packed is True, uses packed tensor broadcasting for - efficient transfer of multiple weights in batches. Otherwise, uses simple - one-by-one broadcasting. + Whether to use packed broadcasting (and the buffer geometry) is read + from `self.packed` / `self.packed_*`, set at the init handshake from the + trainer's init info, so it is guaranteed to match how the trainer + encoded. Args: - update_info: NCCL update info containing parameter names, dtypes, shapes, - and packed flag + update_info: NCCL update info containing parameter names, dtypes, + and shapes """ if self.model_update_group is None: raise RuntimeError( @@ -173,7 +190,7 @@ class NCCLWeightTransferEngine( ) with disable_mtp_completeness_check(): - if update_info.packed: + if self.packed: # Build iterator of (name, (shape, dtype)) from update_info def state_dict_info_iterator(): for name, dtype_name, shape in zip( @@ -187,8 +204,8 @@ class NCCLWeightTransferEngine( group=self.model_update_group, src=0, post_unpack_func=self.model.load_weights, - buffer_size_bytes=update_info.packed_buffer_size_bytes, - num_buffers=update_info.packed_num_buffers, + buffer_size_bytes=self.packed_buffer_size_bytes, + num_buffers=self.packed_num_buffers, device=self.device, ) else: @@ -210,63 +227,241 @@ class NCCLWeightTransferEngine( self.model_update_group = None @staticmethod - def trainer_send_weights( - iterator: Iterator[tuple[str, torch.Tensor]], - trainer_args: dict[str, Any] | NCCLTrainerSendWeightsArgs, - ) -> None: - """Broadcast dense weights from trainer to vLLM workers. + def trainer_send_weights(*args: Any, **kwargs: Any) -> None: + """Removed. Use the stateful `NCCLTrainerWeightTransferEngine` instead. - Args: - iterator: Iterator of model parameters. Returns (name, tensor) tuples - trainer_args: Dictionary or NCCLTrainerSendWeightsArgs instance containing - NCCL-specific arguments. If a dict, should contain keys from - NCCLTrainerSendWeightsArgs. - - Example: - >>> from vllm.distributed.weight_transfer.nccl_engine import ( - ... NCCLWeightTransferEngine, - ... NCCLTrainerSendWeightsArgs, - ... ) - >>> param_iter = ((n, p) for n, p in model.named_parameters()) - >>> args = NCCLTrainerSendWeightsArgs(group=group, packed=True) - >>> NCCLWeightTransferEngine.trainer_send_weights(param_iter, args) + Transitional stub kept only to satisfy the (still abstract) + `WeightTransferEngine.trainer_send_weights`; that member is dropped from + the worker ABC once every backend has migrated to the trainer engine. """ - # Parse trainer args - accept either dict or dataclass instance - if isinstance(trainer_args, dict): - args = NCCLTrainerSendWeightsArgs(**trainer_args) - else: - args = trainer_args + raise NotImplementedError( + "The static NCCL trainer path has been replaced by " + "NCCLTrainerWeightTransferEngine. Build it via " + "WeightTransferTrainerFactory.trainer_init(NCCLTrainerInitInfo(...), " + "client=..., source=...) and drive it with send_weights()." + ) - if args.post_iter_func is None: - # Default: extract just the tensor from (name, tensor) tuple - post_iter_func = lambda x: x[1] - else: - post_iter_func = args.post_iter_func - if args.packed: - # Use packed tensor broadcasting for efficiency - from vllm.distributed.weight_transfer.packed_tensor import ( - packed_nccl_broadcast_producer, +class NCCLTrainerWeightTransferEngine(TrainerWeightTransferEngine[NCCLTrainerInitInfo]): + """Trainer-side NCCL weight transfer engine. + + On the sender (rank 0) holds the NCCL communicator and drives the full + update round trip: it runs the inference-side `update_weights` concurrently + with the trainer-side broadcast (both rendezvous inside the same NCCL + calls), then finishes the update. Non-sender trainer ranks hold no + communicator; they only iterate the source to stay in the trainer-side + collective (e.g. FSDP `full_tensor()`) and skip the client RPCs and the + broadcast (all guarded on `is_sender`). + + `packed` / buffer sizes come from `NCCLTrainerInitInfo`; the sender + propagates them to the worker at `trainer_init` (on the worker-side init + info), so per-round payloads carry only parameter metadata. + """ + + init_info_cls = NCCLTrainerInitInfo + + def __init__( + self, + *, + client: VLLMWeightSyncClient, + source: WeightSource, + is_sender: bool = True, + packed: bool = True, + packed_buffer_size_bytes: int = DEFAULT_PACKED_BUFFER_SIZE_BYTES, + packed_num_buffers: int = DEFAULT_PACKED_NUM_BUFFERS, + ) -> None: + super().__init__(client=client, source=source, is_sender=is_sender) + self.packed = packed + self.packed_buffer_size_bytes = packed_buffer_size_bytes + self.packed_num_buffers = packed_num_buffers + self.model_update_group: PyNcclCommunicator | None = None + + @classmethod + def trainer_init( + cls, + init_info: NCCLTrainerInitInfo, + *, + client: VLLMWeightSyncClient, + source: WeightSource | None = None, + ) -> Self: + if source is None: + raise ValueError("NCCL trainer weight transfer requires a WeightSource.") + engine = cls( + client=client, + source=source, + is_sender=init_info.is_sender, + packed=init_info.packed, + packed_buffer_size_bytes=init_info.packed_buffer_size_bytes, + packed_num_buffers=init_info.packed_num_buffers, + ) + if not engine.is_sender: + # Non-sender trainer ranks aren't part of the transfer NCCL group and + # don't drive the inference side; they only participate in the + # trainer-side gather during send_weights. + return engine + + # Workers sit at rank_offset 1, after the single trainer sender rank 0. + worker_init_info = NCCLWeightTransferInitInfo( + master_address=init_info.master_address, + master_port=init_info.master_port, + rank_offset=1, + world_size=init_info.world_size, + packed=init_info.packed, + packed_buffer_size_bytes=init_info.packed_buffer_size_bytes, + packed_num_buffers=init_info.packed_num_buffers, + ) + + # The inference workers block inside init_weight_transfer_engine waiting + # for the NCCL rendezvous, so we kick that off on a side thread while we + # open the trainer endpoint (rank 0); both sides must rendezvous together. + with ThreadPoolExecutor(max_workers=1) as exe: + future = exe.submit( + engine.client.init_weight_transfer_engine, asdict(worker_init_info) ) + engine.model_update_group = open_trainer_endpoint(init_info) + future.result() # surface any inference-side init error + return engine + + def send_weights(self) -> None: + assert self.source is not None # guaranteed by trainer_init / __init__ + source = self.source + + # Metadata is declared without gathering. For Megatron it is itself a + # collective, so every rank runs it; only the sender ships it. + meta = source.metadata() + + if not self.is_sender: + # Non-sender ranks only join the trainer-side gather collective. + self._broadcast(source, meta) + self._post_send_sync() + return + + update_info = NCCLWeightTransferUpdateInfo( + names=[m.name for m in meta], + dtype_names=[str(m.dtype).split(".")[-1] for m in meta], + shapes=[list(m.shape) for m in meta], + ) + + self.client.start_weight_update() + # update_weights (workers receive) must run concurrently with the + # trainer-side broadcast — 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() + self._broadcast(source, meta) + future.result() # surface inference-side errors + finally: + # Never wait for the RPC thread here. If the 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. + # Let the exception out instead; the transfer group is unusable + # either way and the caller has to tear it down. (The orphaned + # thread is still joined at interpreter exit, so a caller that wants + # to exit cleanly after such a failure must not wait on it.) + exe.shutdown(wait=False) + self.client.finish_weight_update() + self._post_send_sync() + + def _broadcast(self, source: WeightSource, meta: list[ParamMeta]) -> None: + """Iterate the source (materializing each tensor — a collective on all + ranks) and, on the sender, broadcast from rank 0, packed or one-by-one. + Non-sender ranks only replay the iteration to stay in the collective.""" + if not self.is_sender: + for _ in source: + pass + return + + assert self.model_update_group is not None, ( + "trainer_init() must be called before _broadcast()." + ) + pairs = self._checked_iter(source, meta) + if self.packed: packed_nccl_broadcast_producer( - iterator=iterator, - group=args.group, - src=args.src, - post_iter_func=post_iter_func, - buffer_size_bytes=args.packed_buffer_size_bytes, - num_buffers=args.packed_num_buffers, + iterator=pairs, + group=self.model_update_group, + src=0, + post_iter_func=lambda item: item[1], + buffer_size_bytes=self.packed_buffer_size_bytes, + num_buffers=self.packed_num_buffers, ) else: - # Use simple one-by-one broadcasting - for item in iterator: - tensor = post_iter_func(item) - args.group.broadcast( - tensor, - src=args.src, - stream=args.stream or torch.cuda.current_stream(), - ) + stream = torch.cuda.current_stream() + for _name, tensor in pairs: + # NCCL sends `numel` elements straight from `data_ptr()`, so a + # non-contiguous view would ship whatever follows its base + # pointer. Keep the copy referenced until the broadcast is + # enqueued. (The packed path linearizes in `pack_tensors`.) + send = tensor if tensor.is_contiguous() else tensor.contiguous() + self.model_update_group.broadcast(send, src=0, stream=stream) - # Trainer-side process-group setup. Delegates to the shared helper so the - # sparse engine can reuse the exact same rendezvous without subclassing. - trainer_init = staticmethod(trainer_init) + @staticmethod + def _checked_iter( + source: WeightSource, meta: list[ParamMeta] + ) -> Iterator[tuple[str, torch.Tensor]]: + """Yield the source's pairs, checking each against what the worker was + told to expect. + + The worker sizes its receive buffers — and in packed mode cuts its chunk + boundaries — from the update info, which is built from `metadata()`. If + iteration disagrees with it, the two sides split the stream differently + and the transfer hangs in NCCL or loads garbage. Checking here costs one + comparison per parameter and turns that into an error naming the first + divergent parameter. Sender-only: under pipeline parallelism a + non-sender's yielded tensor is not meaningful. + """ + sent = 0 + for name, tensor in source: + if sent >= len(meta): + raise ValueError( + f"WeightSource yielded more parameters than metadata() " + f"declared ({len(meta)}); first extra is {name!r}." + ) + expected = meta[sent] + if ( + name != expected.name + or tensor.dtype != expected.dtype + or tuple(tensor.shape) != expected.shape + ): + raise ValueError( + "WeightSource metadata() disagrees with iteration at index " + f"{sent}: declared {expected.name!r} " + f"{expected.dtype} {tuple(expected.shape)}, got {name!r} " + f"{tensor.dtype} {tuple(tensor.shape)}. Both channels must " + "enumerate the same parameters in the same order." + ) + sent += 1 + yield name, tensor + if sent != len(meta): + raise ValueError( + f"WeightSource yielded {sent} parameters but metadata() " + f"declared {len(meta)}; the worker is waiting for the rest." + ) + + def _post_send_sync(self) -> None: + """Wait for this rank's transfer work to land before returning. + + Broadcasts are only *enqueued* by `send_weights`: the unpacked path on + the current stream, the packed path on the producer's own streams (which + it drains itself). Waiting here lets a caller mutate parameters, or + start the next step on another stream, as soon as `send_weights` + returns, instead of silently depending on same-stream ordering. Every + rank waits: a non-sender's `full_tensor()` gathers feed the sender's + broadcast, so they must have landed before it may touch its shards. + + Unlike IPC there is no cross-rank barrier here. Nothing in this backend + outlives the collective it travelled in (IPC's barrier keeps shared + buffers alive until every consumer has opened them), so a barrier would + only add a dependency on the default process group that this backend + otherwise does not have. + """ + if torch.cuda.is_available(): + torch.cuda.current_stream().synchronize() + + def shutdown(self) -> None: + self.model_update_group = None diff --git a/vllm/distributed/weight_transfer/packed_tensor.py b/vllm/distributed/weight_transfer/packed_tensor.py index b34c58d53d..59bdbafd40 100644 --- a/vllm/distributed/weight_transfer/packed_tensor.py +++ b/vllm/distributed/weight_transfer/packed_tensor.py @@ -177,6 +177,10 @@ def packed_nccl_broadcast_producer( # Move to the next buffer buffer_idx = (buffer_idx + 1) % num_buffers + # Drain every slot before returning. + for stream in streams: + stream.synchronize() + def packed_nccl_broadcast_consumer( iterator: Iterator[tuple[str, tuple[list[int], torch.dtype]]], @@ -279,6 +283,10 @@ def packed_nccl_broadcast_consumer( ) break + # Drain every slot before returning. + for stream in streams: + stream.synchronize() + # ── IPC packed transfer ──────────────────────────────────────────────── diff --git a/vllm/distributed/weight_transfer/sparse_nccl_engine.py b/vllm/distributed/weight_transfer/sparse_nccl_engine.py index 2666afcf4e..0bd499aad4 100644 --- a/vllm/distributed/weight_transfer/sparse_nccl_engine.py +++ b/vllm/distributed/weight_transfer/sparse_nccl_engine.py @@ -14,11 +14,13 @@ MVP limitations: * not composable with checkpoint-format or packed updates """ -from collections.abc import Iterator -from dataclasses import dataclass -from typing import TYPE_CHECKING, Any +from collections.abc import Iterable +from concurrent.futures import ThreadPoolExecutor +from dataclasses import asdict, dataclass +from typing import TYPE_CHECKING, Any, ClassVar import torch +from typing_extensions import Self if TYPE_CHECKING: from vllm.config import VllmConfig @@ -26,20 +28,27 @@ if TYPE_CHECKING: from vllm.config.weight_transfer import WeightTransferConfig from vllm.distributed.weight_transfer.base import ( + TrainerInitInfo, + TrainerWeightTransferEngine, + VLLMWeightSyncClient, + WeightSource, WeightTransferEngine, WeightTransferUpdateInfo, ) from vllm.distributed.weight_transfer.nccl_common import ( NCCLWeightTransferInitInfo, - trainer_init, worker_init_process_group, ) -from vllm.distributed.weight_transfer.nccl_engine import NCCLTrainerSendWeightsArgs +from vllm.distributed.weight_transfer.nccl_common import ( + trainer_init as open_trainer_endpoint, +) __all__ = [ "SparseWeightPatch", + "SparseNCCLTrainerInitInfo", "SparseNCCLWeightTransferUpdateInfo", "SparseNCCLWeightTransferEngine", + "SparseNCCLTrainerWeightTransferEngine", ] @@ -50,6 +59,25 @@ class SparseWeightPatch: 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.""" + + +@dataclass +class SparseNCCLTrainerInitInfo(TrainerInitInfo): + """Trainer-side init info for the sparse NCCL weight transfer backend. + + Same rendezvous shape as the dense NCCL backend (the sender opens its + endpoint as NCCL rank 0), but with no packed wire params: sparse transfers + are never packed. `backend` is the factory dispatch key.""" + + backend: ClassVar[str] = "sparse_nccl" + + master_address: str + master_port: int + world_size: int @dataclass @@ -200,25 +228,182 @@ class SparseNCCLWeightTransferEngine( self.model_update_group = None @staticmethod - def trainer_send_weights( - iterator: Iterator[SparseWeightPatch], - trainer_args: dict[str, Any] | NCCLTrainerSendWeightsArgs, - ) -> None: - """Broadcast sparse flat-index patches from trainer to vLLM workers.""" - if isinstance(trainer_args, dict): - args = NCCLTrainerSendWeightsArgs(**trainer_args) - else: - args = trainer_args + def trainer_send_weights(*args: Any, **kwargs: Any) -> None: + """Removed. Use the stateful `SparseNCCLTrainerWeightTransferEngine`. - if args.packed: + Transitional stub kept only to satisfy the (still abstract) + `WeightTransferEngine.trainer_send_weights`; that member is dropped from + the worker ABC once every backend has migrated to the trainer engine. + """ + raise NotImplementedError( + "The static sparse NCCL trainer path has been replaced by " + "SparseNCCLTrainerWeightTransferEngine. Build it via " + "WeightTransferTrainerFactory.trainer_init(" + "SparseNCCLTrainerInitInfo(...), client=..., source=...) and drive " + "it with send_weights()." + ) + + +class SparseNCCLTrainerWeightTransferEngine( + TrainerWeightTransferEngine[SparseNCCLTrainerInitInfo] +): + """Trainer-side sparse NCCL weight transfer engine. + + 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. + + The sparse backend assumes a single-rank trainer (matching its TP=1 / PP=1 + MVP scope), so non-sender ranks skip `send_weights` entirely. + """ + + init_info_cls = SparseNCCLTrainerInitInfo + + def __init__( + self, + *, + client: VLLMWeightSyncClient, + source: WeightSource | None = None, + is_sender: bool = True, + ) -> None: + # Sparse is a delta backend: it takes per-round patches via + # send_weights, so a `source` would silently never be sent. The + # parameter exists only to match the base/factory signature. + if source is not None: raise ValueError( - "Sparse NCCL updates cannot be combined with `packed=True`" + "Sparse NCCL weight transfer takes no WeightSource; pass each " + "round's patches to send_weights(patches) instead." + ) + super().__init__(client=client, source=source, is_sender=is_sender) + self.model_update_group: PyNcclCommunicator | None = None + + @classmethod + def trainer_init( + cls, + init_info: SparseNCCLTrainerInitInfo, + *, + client: VLLMWeightSyncClient, + source: WeightSource | None = None, + ) -> Self: + engine = cls(client=client, source=source, is_sender=init_info.is_sender) + if not engine.is_sender: + return engine + + # Workers sit at rank_offset 1, after the single trainer sender rank 0. + # Sparse transfers are never packed, so the worker keeps the unpacked + # defaults on its init info. + worker_init_info = NCCLWeightTransferInitInfo( + master_address=init_info.master_address, + master_port=init_info.master_port, + rank_offset=1, + world_size=init_info.world_size, + ) + + # The inference workers block inside init_weight_transfer_engine waiting + # for the NCCL rendezvous, so we kick that off on a side thread while we + # open the trainer endpoint (rank 0); both sides must rendezvous together. + with ThreadPoolExecutor(max_workers=1) as exe: + future = exe.submit( + engine.client.init_weight_transfer_engine, asdict(worker_init_info) + ) + # Open the trainer endpoint as NCCL rank 0 on the current device + # (the init info satisfies the helper's rendezvous protocol). + engine.model_update_group = open_trainer_endpoint(init_info) + future.result() # surface any inference-side init error + + 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 [] + 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.client.finish_weight_update() + 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`. + """ + if patch.indices.dtype != torch.int32: + raise ValueError( + f"Sparse weight updates 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}" ) - stream = args.stream or torch.cuda.current_stream() - for patch in iterator: - args.group.broadcast(patch.indices, src=args.src, stream=stream) - args.group.broadcast(patch.values, src=args.src, stream=stream) + 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 + `NCCLTrainerWeightTransferEngine._post_send_sync` for why there is no + cross-rank barrier (and sparse is single-rank on the trainer anyway).""" + if torch.cuda.is_available(): + torch.cuda.current_stream().synchronize() - # Trainer-side process-group setup (shared with the dense engine). - trainer_init = staticmethod(trainer_init) + def shutdown(self) -> None: + self.model_update_group = None