iHore/vLLM
8
0
Fork 0
Code Issues Pull Requests Actions 1 Packages Projects Releases Wiki Activity

[rl] Stateful Trainer Send: NCCL + Sparse NCCL [3/N] (#50902)

Browse Source 
Signed-off-by: haoaaron <ahao@anyscale.com>
Co-authored-by: Claude Opus 5 (1M context) <noreply@anthropic.com>
This commit is contained in:
Aaron Hao
2026-08-07 00:35:59 -07:00
committed by GitHub
co-authored by Claude Opus 5
parent 43d691ec6b
commit 21ea5b4fa1
14 changed files with 1422 additions and 783 deletions
Download Patch File Download Diff File Expand all files Collapse all files
examples/rl/rlhf_async_new_apis.py
+18 -77
View File
@@ -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)
examples/rl/rlhf_http_nccl.py
+19 -113
View File
@@ -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)
examples/rl/rlhf_nccl.py
+32 -73
View File
@@ -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"]))
examples/rl/rlhf_nccl_fsdp_ep.py
+246 -220
View File
@@ -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()
examples/rl/rlhf_sparse_nccl.py
+71 -115
View File
@@ -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)
tests/distributed/test_packed_tensor.py
-27
View File
@@ -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 ---
tests/distributed/test_weight_transfer.py
+473 -20
View File
@@ -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 == []
vllm/distributed/weight_transfer/base.py
+10 -10
View File
@@ -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:
vllm/distributed/weight_transfer/factory.py
+17 -4
View File
@@ -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",
)
vllm/distributed/weight_transfer/ipc_engine.py
+4 -1
View File
@@ -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()
vllm/distributed/weight_transfer/nccl_common.py
+25 -4
View File
@@ -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
vllm/distributed/weight_transfer/nccl_engine.py
+292 -97
View File
@@ -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
vllm/distributed/weight_transfer/packed_tensor.py
+8
View File
@@ -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 ────────────────────────────────────────────────
vllm/distributed/weight_transfer/sparse_nccl_engine.py
+207 -22
View File
@@ -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