wehub/vllm-project--vllm
1
0
Fork 0
Code Issues Pull Requests Actions Deployments Models Agent Notes Packages Projects Releases Wiki Activity

chore: import upstream snapshot with attribution
pre-commit / pre-run-check (push) Has been cancelled
Details
pre-commit / pre-commit (push) Has been cancelled
Details

Browse Source
This commit is contained in:
wehub-resource-sync
2026-07-13 12:55:37 +08:00
commit 7ce4c8e27e
5900 changed files with 1668062 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files
examples/rl/rlhf_async_new_apis.py
+424
View File
@@ -0,0 +1,424 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""
Demonstrates async reinforcement learning using vLLM and Ray,
with native weight syncing APIs and batch-invariant generation.
The script separates training and inference workloads onto distinct GPUs
so that Ray can manage process placement and inter-process communication.
A Hugging Face Transformer model occupies one GPU for training, and a
vLLM AsyncLLMEngine occupies another GPU for inference.
Batch invariance is enabled so that generation output is deterministic
regardless of how many requests are batched together. This is required
for the validation phase to succeed. Batch invariance currently requires
NVIDIA GPUs with compute capability 9.0 or higher:
- H-series: H100, H200
- B-series: B100, B200
The example performs the following steps:
* Load the training model (Qwen3-1.7B) on one GPU via a Ray actor.
* Initialize the inference engine with a base model (Qwen3-1.7B-Base)
on a separate GPU using vLLM's AsyncLLMEngine with Ray as the
distributed executor backend.
* Set up an NCCL-based weight transfer channel between the trainer
and the inference engine.
* Submit generation requests for a batch of prompts.
* Pause generation once any request reaches a token threshold.
* Broadcast the training model's weights to the inference engine
via the NCCL weight transfer engine, replacing the base weights.
* Resume generation and collect results, noting which tokens were
generated before vs. after the weight swap.
* Validate correctness by launching a fresh vLLM instance loaded
directly with the training model and comparing its output to the
post-swap tokens from the weight-synced engine.
This example assumes a single-node cluster with two GPUs, but Ray
supports multi-node clusters. vLLM expects the GPUs are only used for vLLM
workloads. Residual GPU activity interferes with vLLM memory profiling and
causes unexpected behavior.
"""
import asyncio
import uuid
from dataclasses import asdict
import ray
import torch
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.platforms import current_platform
from vllm.utils.network_utils import get_ip, get_open_port
from vllm.v1.executor import Executor
MODEL_NAME_V1 = "Qwen/Qwen3-1.7B-Base"
MODEL_NAME_V2 = "Qwen/Qwen3-1.7B"
PAUSE_TOKEN_THRESHOLD = 10
ATTN_BACKEND = "TRITON_ATTN" if current_platform.is_rocm() else "FLASH_ATTN"
class MyLLM(vllm.AsyncLLMEngine):
"""Configure the vLLM worker for Ray placement group execution."""
def __init__(self, **kwargs):
engine_args = vllm.AsyncEngineArgs(**kwargs)
vllm_config = engine_args.create_engine_config()
executor_class = Executor.get_class(vllm_config)
super().__init__(
vllm_config=vllm_config,
executor_class=executor_class,
log_requests=engine_args.enable_log_requests,
log_stats=not engine_args.disable_log_stats,
)
self._generation_paused = False
self._request_pause_flag = False
async def do_generate(
self, prompt_token_ids: list[int], sampling_params: vllm.SamplingParams
) -> tuple[vllm.RequestOutput, int]:
"""Generate a single request, setting the request pause flag once the
token count reaches the threshold.
Returns (output, pause_token_index). pause_token_index is the number
of tokens generated before the weight change, or -1 if no pause.
"""
pause_token_index = -1
prev_token_count = 0
async for request_output in self.generate(
{"prompt_token_ids": prompt_token_ids},
sampling_params,
request_id=str(uuid.uuid4()),
):
output = request_output
cur_token_count = len(output.outputs[0].token_ids)
if (
cur_token_count >= PAUSE_TOKEN_THRESHOLD
and not self._request_pause_flag
):
self._request_pause_flag = True
if self._generation_paused and pause_token_index == -1:
pause_token_index = prev_token_count
prev_token_count = cur_token_count
return output, pause_token_index
async def pause_after_n_tokens(self):
"""Wait for any request to set the pause flag, then pause."""
while not self._request_pause_flag:
await asyncio.sleep(0)
await super().pause_generation(mode="keep")
await asyncio.sleep(5)
self._generation_paused = True
@ray.remote(num_gpus=1)
class TrainModel:
"""Ray actor that wraps the training model on a dedicated GPU."""
def __init__(self, model_name: str):
from vllm.model_executor.layers.batch_invariant import (
init_batch_invariance,
)
# need to init all env vars for batch invariance which affect nccl ops
init_batch_invariance()
self.model = AutoModelForCausalLM.from_pretrained(
model_name, dtype=torch.bfloat16
).to("cuda:0")
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(
master_address=self.master_address,
master_port=self.port,
world_size=world_size,
),
)
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,
)
@torch.inference_mode()
def generate(self, token_ids: list[int], max_new_tokens: int) -> list[int]:
"""Greedy-decode max_new_tokens from the given context."""
input_ids = torch.tensor([token_ids], device="cuda:0")
output = self.model.generate(
input_ids,
max_new_tokens=max_new_tokens,
do_sample=False,
)
new_token_ids = output[0, len(token_ids) :].tolist()
return new_token_ids
# Build platform-specific env vars for Ray
ray_env_vars = {}
if current_platform.is_rocm():
# Workaround for RCCL bug. See https://github.com/ROCm/rocm-systems/issues/5756
ray_env_vars["RAY_EXPERIMENTAL_NOSET_HIP_VISIBLE_DEVICES"] = "1"
# For ROCm, BATCH_INVARIANT vllm is not supported
ray_env_vars["VLLM_ROCM_USE_SKINNY_GEMM"] = "0"
else:
# Enable batch invariance for deterministic outputs on NVIDIA
ray_env_vars["VLLM_BATCH_INVARIANT"] = "1"
ray.init(runtime_env={"env_vars": ray_env_vars})
# Launch the training model actor. Ray's resource scheduler will allocate
# 1 GPU (via num_gpus=1 in the decorator), ensuring pg_inference gets different GPUs.
train_model = TrainModel.remote(MODEL_NAME_V2)
rocm_determinism_kwargs = {}
if current_platform.is_rocm():
# ROCm: To minimize non-determinism, we set fixed seed, no prefix caching, and
# sequential request processing (max_num_seqs=1).
rocm_determinism_kwargs = {
"seed": 0,
"enable_prefix_caching": False,
"max_num_seqs": 1,
}
# Build platform-specific LLM kwargs
llm_kwargs = dict(
model=MODEL_NAME_V1,
enforce_eager=True,
max_model_len=8192,
distributed_executor_backend="ray",
attention_backend=ATTN_BACKEND,
gpu_memory_utilization=0.75,
weight_transfer_config=WeightTransferConfig(backend="nccl"),
)
llm_kwargs.update(rocm_determinism_kwargs)
# Launch the vLLM inference engine.
# With data_parallel_backend="ray", vLLM's CoreEngineActorManager creates
# its own placement groups internally for each DP rank, so we must NOT
# create an outer placement group (it would reserve GPUs and hide them
# from the internal DP resource check).
llm = ray.remote(
num_cpus=0,
num_gpus=0,
)(MyLLM).remote(**llm_kwargs)
PROMPTS = [
"The president of the United States is",
"The capital of France is",
"The largest ocean on Earth is",
"The speed of light in a vacuum is",
"The chemical formula for water is",
"The tallest mountain in the world is",
"The first person to walk on the moon was",
"The Great Wall of China was built to",
"Photosynthesis is the process by which",
"The theory of general relativity was proposed by",
"The boiling point of water at sea level is",
"The largest planet in our solar system is",
"DNA stands for deoxyribonucleic acid and it",
]
tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME_V1)
batch_prompt_token_ids = [
tokenizer.encode(prompt, add_special_tokens=False) for prompt in PROMPTS
]
# 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])
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)}):")
for p in PROMPTS:
print(f" - {p!r}")
print(f"{'=' * 50}")
sampling_params = SamplingParams(
temperature=0, max_tokens=PAUSE_TOKEN_THRESHOLD + N_NEW_TOKENS
)
gen_futures = [
llm.do_generate.remote(ptids, sampling_params) for ptids in batch_prompt_token_ids
]
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(llm.resume_generation.remote())
results = ray.get(gen_futures)
for i, (output, pause_idx) in enumerate(results):
all_token_ids = list(output.outputs[0].token_ids)
before_text = tokenizer.decode(all_token_ids[:pause_idx])
after_text = tokenizer.decode(all_token_ids[pause_idx:])
print(f"\n Request {i} ({PROMPTS[i]!r}):")
print(f" Old weights ({pause_idx} tokens): {before_text!r}")
n_after = len(all_token_ids) - pause_idx
print(f" New weights ({n_after} tokens): {after_text!r}")
# ── Phase 2: validate with a fresh V2 vLLM instance ────────────────
# This validation relies on batch-invariant (deterministic) generation to
# compare outputs from the weight-synced engine against a fresh V2 instance.
# On NVIDIA, batch invariance is fully supported, so we require 100% exact
# token match. On ROCm, batch invariance is not yet fully implemented
# (see https://github.com/vllm-project/vllm/issues/27433 and
# https://github.com/vllm-project/vllm/issues/33123), so residual
# non-determinism (e.g. GEMM accumulation order, missing kernel overrides)
# can cause single-token divergences that don't indicate a weight-sync
# failure. We relax the pass rate to 90% on ROCm to accommodate this; a
# real regression (broken weight transfer) would cause ~0% pass rate, not 90%+.
MIN_PASS_RATE = 1.0 if not current_platform.is_rocm() else 0.9
print(f"\n{'=' * 50}")
print("VALIDATION: comparing weight-synced vLLM with fresh V2 instance")
if current_platform.is_rocm():
print(f" (ROCm mode: requiring >= {MIN_PASS_RATE:.0%} exact match rate)")
print(f"{'=' * 50}")
ray.get(llm.shutdown.remote())
ray.kill(llm)
ray.kill(train_model)
llm_v2_kwargs = dict(
model=MODEL_NAME_V2,
enforce_eager=True,
max_model_len=8192,
gpu_memory_utilization=0.75,
distributed_executor_backend="ray",
attention_backend=ATTN_BACKEND,
)
llm_v2_kwargs.update(rocm_determinism_kwargs)
llm_v2 = ray.remote(
num_cpus=0,
num_gpus=0,
)(MyLLM).remote(**llm_v2_kwargs)
val_futures = [
llm_v2.do_generate.remote(
list(output.prompt_token_ids) + list(output.outputs[0].token_ids)[:pause_idx],
SamplingParams(
temperature=0, max_tokens=len(output.outputs[0].token_ids) - pause_idx
),
)
for output, pause_idx in results
]
val_results = ray.get(val_futures)
num_pass = 0
num_total = len(results)
for i, ((output, pause_idx), (val_output, _)) in enumerate(zip(results, val_results)):
expected = list(output.outputs[0].token_ids)[pause_idx:]
actual = list(val_output.outputs[0].token_ids)
match = actual == expected
if match:
num_pass += 1
print(f" [PASS] {PROMPTS[i]!r}")
else:
print(f" [FAIL] {PROMPTS[i]!r}")
print(f" weight-synced vLLM: {tokenizer.decode(expected)!r}")
print(f" V2 vLLM: {tokenizer.decode(actual)!r}")
for j, (e, a) in enumerate(zip(expected, actual)):
if e != a:
print(
f" first divergence at output token {j}: "
f"expected {e} ({tokenizer.decode([e])!r}) vs "
f"actual {a} ({tokenizer.decode([a])!r})"
)
break
ray.get(llm_v2.shutdown.remote())
ray.kill(llm_v2)
pass_rate = num_pass / num_total
print(f"\n Result: {num_pass}/{num_total} prompts passed ({pass_rate:.0%})")
print(f" Required: >= {MIN_PASS_RATE:.0%}")
assert pass_rate >= MIN_PASS_RATE, (
f"Validation pass rate {pass_rate:.0%} ({num_pass}/{num_total}) "
f"is below the required {MIN_PASS_RATE:.0%} threshold. "
f"See failures above for details."
)
print("=" * 50)
examples/rl/rlhf_http_ipc.py
+199
View File
@@ -0,0 +1,199 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""
Demonstrates reinforcement learning from human feedback (RLHF) using vLLM
via HTTP API, with IPC-based weight syncing APIs.
Unlike rlhf_nccl.py which uses NCCL and can use separate GPUs, this script
uses CUDA IPC which requires the training model and vLLM server to be on the
same GPU. Memory must be carefully managed to fit both models.
Unlike rlhf.py which creates a vLLM instance programmatically, this script
assumes you have already started a vLLM server using `vllm serve`. It uses:
- OpenAI-compatible API for inference requests
- HTTP endpoints for weight transfer control plane
- CUDA IPC for actual weight data transfer
Prerequisites:
Start a vLLM server with weight transfer enabled and reduced GPU memory
utilization to leave room for the training model:
$ VLLM_SERVER_DEV_MODE=1 VLLM_ALLOW_INSECURE_SERIALIZATION=1 \
vllm serve facebook/opt-125m --enforce-eager \
--weight-transfer-config '{"backend": "ipc"}' \
--load-format dummy \
--gpu-memory-utilization 0.5
Then run this script:
$ python rlhf_http_ipc.py
The example performs the following steps:
* Load the training model on GPU 0 (same GPU as the vLLM server).
* Generate text using the vLLM server via OpenAI-compatible API. The output
is expected to be nonsense because the server is initialized with dummy weights.
* Initialize weight transfer via HTTP endpoint (no-op for IPC).
* Broadcast the real weights from the training model to the vLLM server
using CUDA IPC handles.
* Generate text again to show normal output after the weight update.
"""
import os
import requests
import torch
from openai import OpenAI
from transformers import AutoModelForCausalLM
from vllm.distributed.weight_transfer.ipc_engine import (
IPCTrainerSendWeightsArgs,
IPCWeightTransferEngine,
)
BASE_URL = "http://localhost:8000"
MODEL_NAME = "facebook/opt-125m"
# Enable insecure serialization for IPC handle serialization
os.environ["VLLM_ALLOW_INSECURE_SERIALIZATION"] = "1"
def generate_completions(client: OpenAI, model: str, prompts: list[str]) -> list[str]:
"""Generate completions using the OpenAI-compatible API."""
results = []
for prompt in prompts:
response = client.completions.create(
model=model,
prompt=prompt,
max_tokens=32,
temperature=0,
)
results.append(response.choices[0].text)
return results
def init_weight_transfer_engine(base_url: str) -> None:
"""Initialize weight transfer via HTTP endpoint (no-op for IPC)."""
url = f"{base_url}/init_weight_transfer_engine"
payload = {"init_info": dict()}
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 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"
response = requests.post(url, timeout=60)
response.raise_for_status()
def resume_generation(base_url: str) -> None:
"""Resume generation via HTTP endpoint."""
url = f"{base_url}/resume"
response = requests.post(url, timeout=60)
response.raise_for_status()
def get_world_size(base_url: str) -> int:
"""Get world size from the vLLM server."""
url = f"{base_url}/get_world_size"
response = requests.get(url, timeout=10)
response.raise_for_status()
return response.json()["world_size"]
def main():
# IPC requires the training model to be on the same GPU as the vLLM server
# The server should be started on GPU 0 with reduced memory utilization
device = "cuda:0"
torch.accelerator.set_device_index(device)
# Load the training model on the same GPU as the server
# Use bfloat16 to reduce memory footprint
print(f"Loading training model: {MODEL_NAME} on {device}")
print(
"Note: Ensure the vLLM server was started with --gpu-memory-utilization 0.5 "
"or lower to leave room for the training model."
)
train_model = AutoModelForCausalLM.from_pretrained(MODEL_NAME, dtype=torch.bfloat16)
train_model.to(device)
train_model.eval() # Set to eval mode to save memory
# Create OpenAI client pointing to the vLLM server
client = OpenAI(
base_url=f"{BASE_URL}/v1",
api_key="EMPTY", # vLLM doesn't require an API key by default
)
# Test prompts
prompts = [
"Hello, my name is",
"The president of the United States is",
"The capital of France is",
"The future of AI is",
]
# Generate text before weight update. The output is expected to be nonsense
# because the server is initialized with dummy weights.
print("-" * 50)
print("Generating text BEFORE weight update (expect nonsense):")
print("-" * 50)
outputs = generate_completions(client, MODEL_NAME, prompts)
for prompt, generated_text in zip(prompts, outputs):
print(f"Prompt: {prompt!r}\nGenerated text: {generated_text!r}")
print("-" * 50)
print("Initializing weight transfer (IPC backend)...")
# Initialize weight transfer on vLLM server (no-op for IPC, but still required)
init_weight_transfer_engine(BASE_URL)
# Pause generation before weight sync
pause_generation(BASE_URL)
# Start weight update, broadcast via IPC, then finish
start_weight_update(BASE_URL)
print("Broadcasting weights via CUDA IPC (HTTP)...")
trainer_args = IPCTrainerSendWeightsArgs(send_mode="http", url=BASE_URL)
IPCWeightTransferEngine.trainer_send_weights(
iterator=train_model.named_parameters(),
trainer_args=trainer_args,
)
finish_weight_update(BASE_URL)
# Resume generation after weight sync
resume_generation(BASE_URL)
# Generate text after weight update. The output is expected to be normal
# because the real weights are now loaded.
print("-" * 50)
print("Generating text AFTER weight update:")
print("-" * 50)
outputs_updated = generate_completions(client, MODEL_NAME, prompts)
for prompt, generated_text in zip(prompts, outputs_updated):
print(f"Prompt: {prompt!r}\nGenerated text: {generated_text!r}")
print("-" * 50)
# Note: The training model and IPC handles remain in memory.
# In a real RLHF training loop, you would update the training model
# and create new IPC handles for each weight update.
if __name__ == "__main__":
main()
examples/rl/rlhf_http_nccl.py
+265
View File
@@ -0,0 +1,265 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""
Demonstrates reinforcement learning from human feedback (RLHF) using vLLM
via HTTP API, with native weight syncing APIs.
Unlike rlhf.py which creates a vLLM instance programmatically, this script
assumes you have already started a vLLM server using `vllm serve`. It uses:
- OpenAI-compatible API for inference requests
- HTTP endpoints for weight transfer control plane
- NCCL for actual weight data transfer
Prerequisites:
Start a vLLM server with weight transfer enabled:
$ VLLM_SERVER_DEV_MODE=1 vllm serve facebook/opt-125m \
--enforce-eager \
--weight-transfer-config '{"backend": "nccl"}' \
--load-format dummy
Then run this script:
$ python rlhf_http.py
The example performs the following steps:
* Load the training model on GPU 0.
* Generate text using the vLLM server via OpenAI-compatible API. The output
is expected to be nonsense because the server is initialized with dummy weights.
* Initialize weight transfer via HTTP endpoint.
* Broadcast the real weights from the training model to the vLLM server
using NCCL.
* Generate text again to show normal output after the weight update.
"""
import requests
import torch
from openai import OpenAI
from transformers import AutoModelForCausalLM
from vllm.distributed.weight_transfer.nccl_engine import (
NCCLTrainerSendWeightsArgs,
NCCLWeightTransferEngine,
)
from vllm.utils.network_utils import get_ip, get_open_port
BASE_URL = "http://localhost:8000"
MODEL_NAME = "facebook/opt-125m"
def generate_completions(client: OpenAI, model: str, prompts: list[str]) -> list[str]:
"""Generate completions using the OpenAI-compatible API."""
results = []
for prompt in prompts:
response = client.completions.create(
model=model,
prompt=prompt,
max_tokens=32,
temperature=0,
)
results.append(response.choices[0].text)
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"
response = requests.post(url, timeout=60)
response.raise_for_status()
def resume_generation(base_url: str) -> None:
"""Resume generation via HTTP endpoint."""
url = f"{base_url}/resume"
response = requests.post(url, timeout=60)
response.raise_for_status()
def get_world_size(base_url: str) -> int:
"""Get world size from the vLLM server."""
url = f"{base_url}/get_world_size"
response = requests.get(url, timeout=10)
response.raise_for_status()
return response.json()["world_size"]
def main():
# Get the inference world size from the vLLM server
inference_world_size = get_world_size(BASE_URL)
world_size = inference_world_size + 1 # +1 for the trainer
device = f"cuda:{inference_world_size}"
torch.accelerator.set_device_index(device)
# Load the training model
print(f"Loading training model: {MODEL_NAME}")
train_model = AutoModelForCausalLM.from_pretrained(MODEL_NAME, dtype=torch.bfloat16)
train_model.to(device)
# Create OpenAI client pointing to the vLLM server
client = OpenAI(
base_url=f"{BASE_URL}/v1",
api_key="EMPTY", # vLLM doesn't require an API key by default
)
# Test prompts
prompts = [
"Hello, my name is",
"The president of the United States is",
"The capital of France is",
"The future of AI is",
]
# Generate text before weight update. The output is expected to be nonsense
# because the server is initialized with dummy weights.
print("-" * 50)
print("Generating text BEFORE weight update (expect nonsense):")
print("-" * 50)
outputs = generate_completions(client, MODEL_NAME, prompts)
for prompt, generated_text in zip(prompts, outputs):
print(f"Prompt: {prompt!r}\nGenerated text: {generated_text!r}")
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.
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(
master_address=master_address,
master_port=master_port,
world_size=world_size,
),
)
# 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
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)
# Resume generation after weight sync
resume_generation(BASE_URL)
# Generate text after weight update. The output is expected to be normal
# because the real weights are now loaded.
print("-" * 50)
print("Generating text AFTER weight update:")
print("-" * 50)
outputs_updated = generate_completions(client, MODEL_NAME, prompts)
for prompt, generated_text in zip(prompts, outputs_updated):
print(f"Prompt: {prompt!r}\nGenerated text: {generated_text!r}")
print("-" * 50)
if __name__ == "__main__":
main()
examples/rl/rlhf_ipc.py
+155
View File
@@ -0,0 +1,155 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""
Demonstrates reinforcement learning from human feedback (RLHF) using vLLM and Ray,
with IPC-based weight syncing APIs
The script colocates the training and inference workloads onto the same GPU using Ray.
The example performs the following steps:
* Request a placement group of 1 GPU.
* Place the inference model on the above GPU using the placement group.
* Place and load the training model on the same GPU using the placement group.
* Generate text from a list of prompts using the inference engine.
* Update the weights of the training model and broadcast the updated weights
to the inference engine by using CUDA IPC handles. Note that
for demonstration purposes we simply zero out the weights.
This example assumes a single-node cluster with a single GPU,
but can be extended to multiple GPUs.
"""
import os
import ray
from ray.util.placement_group import placement_group
from ray.util.scheduling_strategies import PlacementGroupSchedulingStrategy
from transformers import AutoModelForCausalLM
from vllm import LLM, SamplingParams
from vllm.config import WeightTransferConfig
from vllm.distributed.weight_transfer.ipc_engine import (
IPCTrainerSendWeightsArgs,
IPCWeightTransferEngine,
)
class MyLLM(LLM):
"""Configure the vLLM worker for Ray placement group execution."""
def __init__(self, *args, **kwargs):
# Remove the top-level CUDA_VISIBLE_DEVICES variable set by Ray
# so that vLLM can manage its own device placement within the worker.
os.environ.pop("CUDA_VISIBLE_DEVICES", None)
# Each worker uses 0.4 GPU so that two instances fit on the same GPU.
os.environ["VLLM_RAY_PER_WORKER_GPUS"] = "0.4"
os.environ["VLLM_RAY_BUNDLE_INDICES"] = "0"
# needed for ipc handle serialization
os.environ["VLLM_ALLOW_INSECURE_SERIALIZATION"] = "1"
super().__init__(*args, **kwargs)
# Load the OPT-125M model onto GPU 0 for the training workload.
MODEL_NAME = "facebook/opt-125m"
@ray.remote
class TrainModel:
def __init__(self, llm_handle: ray.actor.ActorHandle):
self.train_model = AutoModelForCausalLM.from_pretrained(
MODEL_NAME,
)
self.train_model.to("cuda:0")
self.llm_handle = llm_handle
def init_weight_transfer(self):
# IPC backend doesn't need initialization info
ray.get(
self.llm_handle.init_weight_transfer_engine.remote(dict(init_info=dict()))
)
def broadcast_weights(
self, llm_handle: ray.actor.ActorHandle, packed: bool = False
):
"""Broadcast weights to the inference engine using IPC."""
self.llm_handle = llm_handle
trainer_args = IPCTrainerSendWeightsArgs(
send_mode="ray", llm_handle=llm_handle, packed=packed
)
IPCWeightTransferEngine.trainer_send_weights(
iterator=self.train_model.named_parameters(),
trainer_args=trainer_args,
)
ray.init()
pg_colocate = placement_group([{"GPU": 1, "CPU": 0}])
ray.get(pg_colocate.ready())
llm = ray.remote(
num_cpus=0,
num_gpus=0,
scheduling_strategy=PlacementGroupSchedulingStrategy(
placement_group=pg_colocate,
placement_group_capture_child_tasks=True,
),
)(MyLLM).remote(
model=MODEL_NAME,
enforce_eager=True,
tensor_parallel_size=1,
distributed_executor_backend="ray",
gpu_memory_utilization=0.7,
weight_transfer_config=WeightTransferConfig(backend="ipc"),
load_format="dummy",
)
train_model = TrainModel.options(
num_gpus=0.1,
num_cpus=0,
scheduling_strategy=PlacementGroupSchedulingStrategy(
placement_group=pg_colocate, placement_group_capture_child_tasks=True
),
).remote(llm)
# Generate text from the prompts.
prompts = [
"Hello, my name is",
"The president of the United States is",
"The capital of France is",
"The future of AI is",
]
sampling_params = SamplingParams(temperature=0)
outputs = ray.get(llm.generate.remote(prompts, sampling_params))
print("-" * 50)
for output in outputs:
prompt = output.prompt
generated_text = output.outputs[0].text
print(f"Prompt: {prompt!r}\nGenerated text: {generated_text!r}")
print("-" * 50)
ray.get(llm.sleep.remote(level=0))
ray.get(train_model.init_weight_transfer.remote())
# Start weight update, sync weights, then finish
ray.get(llm.start_weight_update.remote())
ray.get(train_model.broadcast_weights.remote(llm))
ray.get(llm.finish_weight_update.remote())
ray.get(llm.wake_up.remote(tags=["scheduling"]))
outputs_packed = ray.get(llm.generate.remote(prompts, sampling_params))
print("-" * 50)
print("Results after packed/chunked IPC weight sync:")
for output in outputs_packed:
prompt = output.prompt
generated_text = output.outputs[0].text
print(f"Prompt: {prompt!r}\nGenerated text: {generated_text!r}")
print("-" * 50)
examples/rl/rlhf_ipc_fsdp_ep.py
+418
View File
@@ -0,0 +1,418 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""
RLHF with FSDP2 training and vLLM expert-parallel inference using **CUDA IPC**
weight transfer and **packed** tensors.
Layout (4 GPUs, TP=1, DP=4, EP):
* One Ray placement group per GPU.
* Each PG holds one FSDP training worker and one vLLM ``LLM`` instance
(sync API) using fractional GPUs so both fit on the same device.
* The 4 ``LLM`` instances form a DP group via env-var-based SPMD
coordination (``VLLM_DP_RANK``, ``VLLM_DP_SIZE``, etc.), the same
mechanism used by ``examples/offline_inference/data_parallel.py``.
* A ``DataParallelInferenceEngine`` actor spawns all 4 LLM actors,
waits for initialization, and orchestrates generation / weight-sync.
Uses the built-in ``ray`` send_mode: each FSDP worker calls
``trainer_send_weights`` targeting its colocated LLM actor.
This example was run on 4xH100.
"""
from __future__ import annotations
import os
from dataclasses import asdict
import ray
import torch
import torch.distributed as dist
from huggingface_hub import snapshot_download
from ray.util.placement_group import placement_group
from ray.util.scheduling_strategies import PlacementGroupSchedulingStrategy
from torch.distributed._tensor import DTensor
from torch.distributed.fsdp import fully_shard
from transformers import AutoModelForCausalLM
from vllm import LLM, SamplingParams
from vllm.config import WeightTransferConfig
from vllm.distributed.weight_transfer.ipc_engine import (
IPCTrainerSendWeightsArgs,
IPCWeightTransferEngine,
IPCWeightTransferInitInfo,
)
from vllm.utils.network_utils import get_ip, get_open_port
TRAIN_GPU_FRACTION = float(os.environ.get("RLHF_IPC_TRAIN_GPU_FRACTION", "0.42"))
VLLM_GPU_FRACTION = float(os.environ.get("RLHF_IPC_VLLM_GPU_FRACTION", "0.42"))
MODEL_NAME = "Qwen/Qwen3-30B-A3B"
FSDP_WORLD_SIZE = 4
INFERENCE_TP_SIZE = 1
INFERENCE_DP_SIZE = 4
class MyLLM(LLM):
"""LLM subclass that configures DP env vars for SPMD coordination."""
def __init__(
self,
*args,
dp_rank: int = 0,
dp_size: int = 1,
dp_master_ip: str = "127.0.0.1",
dp_master_port: int = 0,
**kwargs,
):
os.environ.pop("CUDA_VISIBLE_DEVICES", None)
os.environ["VLLM_RAY_PER_WORKER_GPUS"] = str(VLLM_GPU_FRACTION)
os.environ["VLLM_RAY_BUNDLE_INDICES"] = "0"
os.environ["VLLM_ALLOW_INSECURE_SERIALIZATION"] = "1"
os.environ["VLLM_DP_RANK"] = str(dp_rank)
os.environ["VLLM_DP_RANK_LOCAL"] = str(dp_rank)
os.environ["VLLM_DP_SIZE"] = str(dp_size)
os.environ["VLLM_DP_MASTER_IP"] = dp_master_ip
os.environ["VLLM_DP_MASTER_PORT"] = str(dp_master_port)
super().__init__(*args, **kwargs)
def ready(self):
return True
@ray.remote(num_cpus=0, num_gpus=TRAIN_GPU_FRACTION)
class FSDPTrainWorker:
"""One FSDP2 worker per GPU; colocated with vLLM DP rank via placement group."""
def __init__(
self,
model_name: str,
rank: int,
fsdp_world_size: int,
fsdp_master_addr: str,
fsdp_master_port: int,
):
self.rank = rank
os.environ["MASTER_ADDR"] = fsdp_master_addr
os.environ["MASTER_PORT"] = str(fsdp_master_port)
dist.init_process_group(backend="nccl", rank=rank, world_size=fsdp_world_size)
torch.accelerator.set_device_index(0)
model = AutoModelForCausalLM.from_pretrained(
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)
self.model = model
def get_rank(self):
return self.rank
def get_weight_metadata(self):
return self.weight_names, self.weight_dtype_names, self.weight_shapes
def gather_and_broadcast_weights_ipc(self, llm_handle, packed: bool = True):
"""All-gather full params; all ranks create IPC handles, rank 0 sends.
All ranks must call trainer_send_weights so they participate in the
all_gather_object collective inside _all_gather_and_merge_handles.
Only rank 0 actually sends the payload to vLLM (gated by _is_rank_zero).
"""
def _full_param_iter():
# HF's Qwen3MoeExperts (and other recent HF MoE impls) packs
# all experts into two fused 3-D tensors per layer:
# experts.gate_up_proj shape (E, 2*I, H)
# experts.down_proj shape (E, H, I)
# vLLM's Qwen3MoE load_weights still expects the older
# per-expert HF layout (experts.<i>.gate_proj.weight,
# experts.<i>.up_proj.weight, experts.<i>.down_proj.weight),
# so we un-fuse on the fly. Split order matches HF's forward:
# gate, up = linear(x, gate_up_proj[i]).chunk(2, dim=-1)
# → rows [:I] of gate_up_proj[i] are gate, rows [I:] are up.
params = self.model.state_dict()
for name in list(params.keys()):
param = params.pop(name)
if isinstance(param, DTensor):
tensor = param.full_tensor().detach().contiguous()
else:
tensor = param.detach().contiguous()
del param
if name.endswith(".experts.gate_up_proj") and tensor.dim() == 3:
prefix = name[: -len(".gate_up_proj")]
num_experts, two_inter, _ = tensor.shape
inter = two_inter // 2
for i in range(num_experts):
expert = tensor[i]
yield (
f"{prefix}.{i}.gate_proj.weight",
expert[:inter].contiguous(),
)
yield (
f"{prefix}.{i}.up_proj.weight",
expert[inter:].contiguous(),
)
del tensor
elif name.endswith(".experts.down_proj") and tensor.dim() == 3:
prefix = name[: -len(".down_proj")]
num_experts = tensor.shape[0]
for i in range(num_experts):
yield (
f"{prefix}.{i}.down_proj.weight",
tensor[i].contiguous(),
)
del tensor
else:
yield name, tensor
trainer_args = IPCTrainerSendWeightsArgs(
send_mode="ray",
llm_handle=llm_handle,
packed=packed,
packed_buffer_size_bytes=1024 * 1024 * 1024, # 1 GB
)
IPCWeightTransferEngine.trainer_send_weights(
iterator=_full_param_iter(),
trainer_args=trainer_args,
)
@ray.remote(num_cpus=1)
class DataParallelInferenceEngine:
"""Manages a pool of DP-sharded vLLM LLM actors.
Spawns one MyLLM actor per placement group, waits for all engines to
finish initializing, and exposes generation / weight-sync helpers.
"""
def __init__(
self,
model: str,
pgs: list,
dp_master_ip: str,
dp_master_port: int,
):
dp_size = len(pgs)
self.llm_actors = []
for r in range(dp_size):
sched = PlacementGroupSchedulingStrategy(
placement_group=pgs[r],
placement_group_capture_child_tasks=True,
)
actor = (
ray.remote(num_cpus=0, num_gpus=0)(MyLLM)
.options(scheduling_strategy=sched)
.remote(
model=model,
enforce_eager=True,
tensor_parallel_size=INFERENCE_TP_SIZE,
distributed_executor_backend="ray",
enable_expert_parallel=True,
gpu_memory_utilization=0.35,
weight_transfer_config=WeightTransferConfig(backend="ipc"),
enable_sleep_mode=True,
load_format="dummy",
dp_rank=r,
dp_size=dp_size,
dp_master_ip=dp_master_ip,
dp_master_port=dp_master_port,
)
)
self.llm_actors.append(actor)
ray.get([actor.ready.remote() for actor in self.llm_actors])
def get_llm_actors(self):
return self.llm_actors
def generate(self, prompts: list[str], sampling_params):
"""Distribute prompts round-robin across DP ranks and collect results."""
dp_size = len(self.llm_actors)
per_rank: list[list[str]] = [[] for _ in range(dp_size)]
indices: list[list[int]] = [[] for _ in range(dp_size)]
for i, prompt in enumerate(prompts):
rank = i % dp_size
per_rank[rank].append(prompt)
indices[rank].append(i)
refs = [
actor.generate.remote(per_rank[r], sampling_params)
for r, actor in enumerate(self.llm_actors)
if per_rank[r]
]
all_outputs = ray.get(refs)
ordered = [None] * len(prompts)
rank_idx = 0
for r in range(dp_size):
if per_rank[r]:
for local_i, orig_i in enumerate(indices[r]):
ordered[orig_i] = all_outputs[rank_idx][local_i]
rank_idx += 1
return ordered
def init_weight_transfer(self):
ray.get(
[
actor.init_weight_transfer_engine.remote(
dict(init_info=asdict(IPCWeightTransferInitInfo()))
)
for actor in self.llm_actors
]
)
def start_weight_update(self):
ray.get([actor.start_weight_update.remote() for actor in self.llm_actors])
def finish_weight_update(self):
ray.get([actor.finish_weight_update.remote() for actor in self.llm_actors])
def sleep(self, level: int = 0):
ray.get([actor.sleep.remote(level=level) for actor in self.llm_actors])
def wake_up(self, tags: list[str] | None = None):
ray.get([actor.wake_up.remote(tags=tags) for actor in self.llm_actors])
def main():
ray.init(
runtime_env={
"env_vars": {
"VLLM_ALLOW_INSECURE_SERIALIZATION": "1",
}
}
)
assert TRAIN_GPU_FRACTION + VLLM_GPU_FRACTION <= 1.0, (
"Train + vLLM GPU fractions must sum to at most 1.0 per bundle."
)
local_model_path = snapshot_download(MODEL_NAME)
print(f"[init] Model downloaded to {local_model_path}")
fsdp_master_addr = get_ip()
fsdp_master_port = get_open_port()
dp_master_port = get_open_port()
dp_master_ip = get_ip()
# Create one placement group per DP rank (one GPU each).
pgs = []
for _ in range(INFERENCE_DP_SIZE):
pg = placement_group([{"GPU": 1, "CPU": 1}])
pgs.append(pg)
ray.get([pg.ready() for pg in pgs])
print(f"[init] {len(pgs)} placement groups ready.")
# Launch FSDP training workers, one per PG.
scheduling = [
PlacementGroupSchedulingStrategy(
placement_group=pgs[r],
placement_group_capture_child_tasks=True,
)
for r in range(FSDP_WORLD_SIZE)
]
fsdp_workers = [
FSDPTrainWorker.options(scheduling_strategy=scheduling[r]).remote(
local_model_path,
r,
FSDP_WORLD_SIZE,
fsdp_master_addr,
fsdp_master_port,
)
for r in range(FSDP_WORLD_SIZE)
]
ray.get([w.get_rank.remote() for w in fsdp_workers])
print(f"[init] {FSDP_WORLD_SIZE} FSDP workers ready.")
# Launch DP inference engine (spawns and initializes all LLM actors).
inference_engine = DataParallelInferenceEngine.remote(
model=local_model_path,
pgs=pgs,
dp_master_ip=dp_master_ip,
dp_master_port=dp_master_port,
)
llm_actors = ray.get(inference_engine.get_llm_actors.remote())
print(f"[init] {INFERENCE_DP_SIZE} LLM actors ready.")
prompts = [
"Hello, my name is",
"The president of the United States is",
"The capital of France is",
"The future of AI is",
]
sampling_params = SamplingParams(temperature=0)
print("[generate] Generating with dummy weights...")
outputs = ray.get(inference_engine.generate.remote(prompts, sampling_params))
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 ---
print("[transfer] Initializing IPC weight transfer...")
ray.get(inference_engine.init_weight_transfer.remote())
# Two-phase sleep/wake pattern:
# 1. sleep(level=1) — offload weights to CPU, discard KV cache
# 2. wake_up(tags=["weights"]) — bring weights back to GPU (KV cache still free)
# 3. IPC weight transfer — overwrite weights, plenty of room without KV cache
# 4. wake_up(tags=["kv_cache"]) — re-allocate KV cache for inference
print("[sync] Sleeping engines (offload weights + free KV cache)...")
ray.get(inference_engine.sleep.remote(level=1))
print("[sync] Waking weights (KV cache stays free)...")
ray.get(inference_engine.wake_up.remote(tags=["weights"]))
print("[sync] Starting weight update...")
ray.get(inference_engine.start_weight_update.remote())
print("[sync] Packed IPC transfer FSDP → vLLM...")
ray.get(
[
w.gather_and_broadcast_weights_ipc.remote(llm_actors, packed=True)
for w in fsdp_workers
]
)
ray.get(inference_engine.finish_weight_update.remote())
print("[sync] Weight transfer complete.")
print("[sync] Waking KV cache + scheduling...")
ray.get(inference_engine.wake_up.remote(tags=["kv_cache", "scheduling"]))
print("[generate] Generating with synced weights...")
outputs_updated = ray.get(
inference_engine.generate.remote(prompts, sampling_params)
)
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}")
print("-" * 60)
if __name__ == "__main__":
main()
examples/rl/rlhf_nccl.py
+237
View File
@@ -0,0 +1,237 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""
Demonstrates reinforcement learning using vLLM and Ray,
with native weight syncing APIs at engine instance.
The script separates training and inference workloads onto distinct GPUs
so that Ray can manage process placement and inter-process communication.
A Hugging Face Transformer model occupies one GPU for training, whereas a
2x tensor-parallel vLLM inference engine occupies two GPUs.
The example performs the following steps:
* Load the training model on one gpu (scheduled via ray)
* Initialize the inference model with dummy weights across
two gpus using vLLM's tensor parallelism and Ray placement groups.
* Generate gibberish from a list of prompts using the randomly initialized
inference engine.
* Update the weights of the training model and broadcast the updated weights
to the inference engine by using a Ray collective RPC group.
* Generating from the list of prompts after weight sync should result
in sensible outputs.
This example assumes a single-node cluster with three GPUs, but Ray
supports multi-node clusters. vLLM expects the GPUs are only used for vLLM
workloads. Residual GPU activity interferes with vLLM memory profiling and
causes unexpected behavior.
"""
import os
import ray
import torch
from ray.util.placement_group import placement_group
from ray.util.scheduling_strategies import PlacementGroupSchedulingStrategy
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.platforms import current_platform
from vllm.utils.network_utils import get_ip, get_open_port
MODEL_NAME = "facebook/opt-125m"
# MODEL_NAME = "inference-optimization/Qwen3-0.6B-W4A16-G128"
def get_assigned_gpu():
"""This is a temporary workaround for a runtime bug in RCCL on ROCm."""
if not current_platform.is_rocm():
return 0
assigned_gpu = int(ray.get_gpu_ids()[0])
os.environ.pop("CUDA_VISIBLE_DEVICES", None)
os.environ.pop("HIP_VISIBLE_DEVICES", None)
torch.accelerator.set_device_idx(assigned_gpu)
return assigned_gpu
class MyLLM(LLM):
"""Configure the vLLM worker for Ray placement group execution."""
def __init__(self, *args, **kwargs):
os.environ["VLLM_RAY_BUNDLE_INDICES"] = "0,1"
super().__init__(*args, **kwargs)
@ray.remote(num_gpus=1)
class TrainModel:
"""Ray actor that wraps the training model on a dedicated GPU."""
def __init__(self, model_name: str):
assigned_gpu = get_assigned_gpu()
self.model = AutoModelForCausalLM.from_pretrained(
model_name,
).to(f"cuda:{assigned_gpu}")
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(
master_address=self.master_address,
master_port=self.port,
world_size=world_size,
),
)
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,
)
# Initialize Ray and set the visible devices. The vLLM engine will
# be placed on GPUs 1 and 2.
ray.init()
# Create a placement group that reserves GPU 1–2 for the vLLM inference engine.
# Learn more about Ray placement groups:
# https://docs.ray.io/en/latest/placement-groups.html
# Launch the training model actor. Ray's resource scheduler will allocate
# 1 GPU (via num_gpus=1 in the decorator), ensuring pg_inference gets different GPUs.
train_model = TrainModel.remote(MODEL_NAME)
pg_inference = placement_group([{"GPU": 1, "CPU": 0}] * 2)
ray.get(pg_inference.ready())
scheduling_inference = PlacementGroupSchedulingStrategy(
placement_group=pg_inference,
placement_group_capture_child_tasks=True,
placement_group_bundle_index=0,
)
# Launch the vLLM inference engine. The `enforce_eager` flag reduces
# start-up latency.
# Note: Weight transfer APIs (init_weight_transfer_engine, update_weights)
# are now native to vLLM workers.
llm = ray.remote(
num_cpus=0,
num_gpus=0,
scheduling_strategy=scheduling_inference,
)(MyLLM).remote(
model=MODEL_NAME,
enforce_eager=True,
tensor_parallel_size=2,
data_parallel_size=1,
distributed_executor_backend="ray",
weight_transfer_config=WeightTransferConfig(backend="nccl"),
load_format="dummy",
quantization="fp8",
)
# Generate text from the prompts.
prompts = [
"Hello, my name is",
"The president of the United States is",
"The capital of France is",
"The future of AI is",
]
sampling_params = SamplingParams(temperature=0)
outputs = ray.get(llm.generate.remote(prompts, sampling_params))
# Generate text with the initial model. The output is expected to be nonsense
# because the weights are randomly initialized.
print("-" * 50)
for output in outputs:
prompt = output.prompt
generated_text = output.outputs[0].text
print(f"Prompt: {prompt!r}\nGenerated text: {generated_text!r}")
print("-" * 50)
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())
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,
)
)
)
# 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())
ray.get(llm.wake_up.remote(tags=["scheduling"]))
# Generate text with the updated model. The output is expected to be normal
# because the weights are updated.
outputs_updated = ray.get(llm.generate.remote(prompts, sampling_params))
print("-" * 50)
for output in outputs_updated:
prompt = output.prompt
generated_text = output.outputs[0].text
print(f"Prompt: {prompt!r}\nGenerated text: {generated_text!r}")
print("-" * 50)
examples/rl/rlhf_nccl_fsdp_ep.py
+344
View File
@@ -0,0 +1,344 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""
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 +
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 engine uses AsyncLLMEngine which automatically spawns
DP worker processes (no manual placement group needed). Weight sync
uses pause_generation / resume_generation.
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.
5. Generate from prompts → sensible output (synced weights).
Assumes a single-node cluster with 8 GPUs.
"""
import asyncio
import os
import uuid
from dataclasses import asdict
import ray
import torch
import torch.distributed as dist
from huggingface_hub import snapshot_download
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.utils.network_utils import get_ip, get_open_port
from vllm.v1.executor import Executor
MODEL_NAME = "Qwen/Qwen3-30B-A3B"
FSDP_WORLD_SIZE = 4
INFERENCE_TP_SIZE = 1
INFERENCE_DP_SIZE = 4
@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.
"""
def __init__(
self,
model_name: str,
rank: int,
fsdp_world_size: int,
fsdp_master_addr: str,
fsdp_master_port: int,
):
self.rank = rank
os.environ["MASTER_ADDR"] = fsdp_master_addr
os.environ["MASTER_PORT"] = str(fsdp_master_port)
dist.init_process_group(backend="nccl", rank=rank, world_size=fsdp_world_size)
torch.accelerator.set_device_index(0)
model = AutoModelForCausalLM.from_pretrained(
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)
self.model = model
self.transfer_port = None
self.transfer_master_address = None
self.model_update_group = None
def get_rank(self):
return self.rank
# ---- weight-transfer setup (rank 0 only) ----
def setup_transfer_endpoint(self):
"""Create the NCCL rendezvous endpoint for weight transfer."""
assert self.rank == 0
self.transfer_port = get_open_port()
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 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
# ---- collective ops (ALL FSDP ranks must call concurrently) ----
def gather_and_broadcast_weights(self, packed: bool = True):
"""
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()
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,
)
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])
async def main():
ray.init()
# 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)
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.
fsdp_workers = [
FSDPTrainWorker.remote(
local_model_path,
rank,
FSDP_WORLD_SIZE,
fsdp_master_addr,
fsdp_master_port,
)
for rank in range(FSDP_WORLD_SIZE)
]
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",
]
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,
)
)
)
)
ray.get(train_handle)
print("[transfer] NCCL groups initialized.")
# --- Pause, transfer weights, resume ---
print("[sync] Pausing generation...")
await engine.pause_generation(mode="abort")
print("[sync] Generation paused.")
names, dtype_names, shapes = ray.get(fsdp_workers[0].get_weight_metadata.remote())
print(f"[sync] Got metadata for {len(names)} parameters.")
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}")
print("-" * 60)
if __name__ == "__main__":
asyncio.run(main())
examples/rl/rlhf_sparse_nccl.py
+529
View File
@@ -0,0 +1,529 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""
Demonstrates dense-vs-sparse NCCL weight syncing with a real model.
This example mirrors the validation story used for the sparse NCCL MVP:
both the dense update path and the sparse patch path start from the same real
checkpoint and apply the same deterministic trainer-side patch. The script then
checks that greedy 1-token outputs match between the dense and sparse vLLM
engines after the update.
The example performs the following steps:
* Load a training model on one GPU via a Ray actor.
* Launch a vLLM engine with the same real model on a second GPU.
* Verify trainer vs vLLM baseline agreement before any update.
* Apply a deterministic patch to ``model.embed_tokens.weight`` on the trainer.
* Run a dense NCCL update into a fresh vLLM engine and collect post-update
outputs.
* Reset the trainer back to the baseline checkpoint.
* Apply the same deterministic patch again.
* Run a sparse NCCL update into another fresh vLLM engine and collect
post-update outputs.
* Compare dense vs sparse baseline outputs, dense vs sparse post-update
outputs, estimated payload sizes, and trainer-side send times.
Current sparse weight transfer MVP limitations:
* ``TP=1`` and ``PP=1`` only
* sparse updates use runtime/kernel-format parameter names
* sparse updates are not composable with checkpoint-format or packed updates
This example assumes a single-node cluster with two GPUs.
"""
import hashlib
import os
import time
from collections.abc import Sequence
import ray
import torch
from ray.util.placement_group import placement_group
from ray.util.scheduling_strategies import PlacementGroupSchedulingStrategy
from transformers import AutoModelForCausalLM, AutoTokenizer
from vllm import LLM, SamplingParams
from vllm.config import WeightTransferConfig
from vllm.distributed.weight_transfer.nccl_engine import (
NCCLTrainerSendWeightsArgs,
NCCLWeightTransferEngine,
)
from vllm.distributed.weight_transfer.sparse_nccl_engine import (
SparseNCCLWeightTransferEngine,
SparseWeightPatch,
)
from vllm.utils.network_utils import get_ip, get_open_port
MODEL_NAME = "Qwen/Qwen2.5-0.5B-Instruct"
PATCHED_PARAM_NAME = "model.embed_tokens.weight"
MAX_PATCH_ROWS = 32
PROMPTS = [
"Hello, my name is",
"The president of the United States is",
"The capital of France is",
"The future of AI is",
]
SAMPLING_PARAMS = SamplingParams(temperature=0.0, max_tokens=1)
class MyLLM(LLM):
"""Configure the vLLM worker for Ray placement group execution."""
def __init__(self, *args, **kwargs):
os.environ["VLLM_RAY_BUNDLE_INDICES"] = "0"
super().__init__(*args, **kwargs)
@ray.remote(num_gpus=1)
class TrainModel:
"""Ray actor that owns the trainer-side model and deterministic patch state."""
def __init__(self, model_name: str):
self.model_name = model_name
self.tokenizer = AutoTokenizer.from_pretrained(model_name)
if self.tokenizer.pad_token_id is None:
self.tokenizer.pad_token = self.tokenizer.eos_token
self.model = None
self.patched_param = None
self.pending_sparse_patches: list[SparseWeightPatch] | None = None
self.model_update_group = None
self.master_address = get_ip()
self.port = get_open_port()
self.reset_model()
def reset_model(self) -> None:
self.model = AutoModelForCausalLM.from_pretrained(
self.model_name,
torch_dtype=torch.bfloat16,
).to("cuda:0")
self.model.eval()
try:
self.patched_param = self.model.get_parameter(PATCHED_PARAM_NAME)
except AttributeError as exc:
raise RuntimeError(
f"Expected trainer model to expose `{PATCHED_PARAM_NAME}`"
) from exc
self.pending_sparse_patches = None
def create_rendezvous(self) -> tuple[str, int]:
self.port = get_open_port()
return self.master_address, self.port
def init_weight_transfer_group(self, world_size: int) -> None:
self.model_update_group = NCCLWeightTransferEngine.trainer_init(
dict(
master_address=self.master_address,
master_port=self.port,
world_size=world_size,
)
)
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()
return (
dict(
names=names,
dtype_names=dtype_names,
shapes=shapes,
packed=packed,
),
payload_bytes,
)
@torch.inference_mode()
def generate(
self,
prompts: Sequence[str],
max_new_tokens: int = 1,
) -> list[dict[str, object]]:
generations = []
for prompt in prompts:
model_inputs = self.tokenizer(prompt, return_tensors="pt").to("cuda:0")
output = self.model.generate(
**model_inputs,
max_new_tokens=max_new_tokens,
do_sample=False,
pad_token_id=self.tokenizer.pad_token_id,
)
new_token_ids = output[0, model_inputs["input_ids"].shape[1] :].tolist()
generations.append(
{
"token_ids": new_token_ids,
"text": self.tokenizer.decode(
new_token_ids,
skip_special_tokens=False,
),
}
)
return generations
def prepare_sparse_patch(
self,
prompts: Sequence[str],
max_patch_rows: int = MAX_PATCH_ROWS,
) -> tuple[dict[str, object], list[int], str, int]:
selected_token_ids: list[int] = []
special_ids = set(self.tokenizer.all_special_ids)
for prompt in prompts:
token_ids = self.tokenizer(prompt, add_special_tokens=False)["input_ids"]
for token_id in token_ids:
if token_id in special_ids or token_id in selected_token_ids:
continue
selected_token_ids.append(token_id)
if len(selected_token_ids) == max_patch_rows:
break
if len(selected_token_ids) == max_patch_rows:
break
if not selected_token_ids:
raise ValueError("Could not derive any non-special token IDs to patch")
vocab_size = self.patched_param.shape[0]
next_token_id = selected_token_ids[-1]
while len(selected_token_ids) < max_patch_rows:
next_token_id = (next_token_id + 1) % vocab_size
if next_token_id in special_ids or next_token_id in selected_token_ids:
continue
selected_token_ids.append(next_token_id)
row_ids = torch.tensor(
selected_token_ids,
device=self.patched_param.device,
dtype=torch.long,
)
hidden_size = self.patched_param.shape[1]
column_offsets = torch.arange(
hidden_size,
device=self.patched_param.device,
dtype=torch.long,
)
with torch.no_grad():
# Rotate the selected embedding rows instead of zeroing them so the
# patch remains deterministic while avoiding a degenerate collapse
# to the same special token after the update.
replacement_rows = self.patched_param[row_ids].roll(shifts=1, dims=0)
self.patched_param[row_ids] = replacement_rows
flat_indices = (
row_ids.unsqueeze(1).mul(hidden_size).add(column_offsets).reshape(-1)
)
flat_values = self.patched_param[row_ids].reshape(-1).contiguous()
self.pending_sparse_patches = [
SparseWeightPatch(
name=PATCHED_PARAM_NAME,
indices=flat_indices.to(torch.int32),
values=flat_values,
)
]
patch_digest = hashlib.sha256(
self.pending_sparse_patches[0].indices.cpu().numpy().tobytes()
+ self.pending_sparse_patches[0]
.values.detach()
.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
def broadcast_weights(self, packed: bool = False) -> float:
if self.model_update_group is None:
raise RuntimeError("Weight transfer group 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,
)
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")
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),
)
torch.accelerator.synchronize()
self.pending_sparse_patches = None
return (time.perf_counter() - start) * 1000.0
def launch_llm(
scheduling_inference: PlacementGroupSchedulingStrategy,
backend: str = "nccl",
):
return ray.remote(
num_cpus=0,
num_gpus=0,
scheduling_strategy=scheduling_inference,
)(MyLLM).remote(
model=MODEL_NAME,
enforce_eager=True,
tensor_parallel_size=1,
distributed_executor_backend="ray",
gpu_memory_utilization=0.7,
weight_transfer_config=WeightTransferConfig(backend=backend),
)
def collect_vllm_generations(llm_handle) -> list[dict[str, object]]:
outputs = ray.get(llm_handle.generate.remote(PROMPTS, SAMPLING_PARAMS))
generations = []
for output in outputs:
generations.append(
{
"token_ids": output.outputs[0].token_ids,
"text": output.outputs[0].text,
}
)
return generations
def token_sequences_match(
left: Sequence[dict[str, object]],
right: Sequence[dict[str, object]],
) -> bool:
return [item["token_ids"] for item in left] == [item["token_ids"] for item in right]
def print_generations(label: str, prompts: Sequence[str], generations) -> None:
print(f"\n{label}")
print("-" * 50)
for prompt, generation in zip(prompts, generations):
print(f"Prompt: {prompt!r}")
print(f"Token IDs: {generation['token_ids']}")
print(f"Text: {generation['text']!r}")
print("-" * 50)
def run_dense_phase(
train_model,
scheduling_inference: PlacementGroupSchedulingStrategy,
) -> dict[str, object]:
ray.get(train_model.reset_model.remote())
llm = launch_llm(scheduling_inference, backend="nccl")
try:
dense_before = collect_vllm_generations(llm)
ray.get(llm.sleep.remote(level=0))
master_address, master_port = 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())
dense_update_info, dense_payload_bytes = ray.get(
train_model.get_dense_update_info.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())
ray.get(llm.wake_up.remote(tags=["scheduling"]))
dense_after = collect_vllm_generations(llm)
return {
"dense_before": dense_before,
"dense_after": dense_after,
"selected_token_ids": selected_token_ids,
"patch_digest": patch_digest,
"dense_payload_bytes": dense_payload_bytes,
"dense_send_ms": dense_send_ms,
}
finally:
ray.kill(llm)
def run_sparse_phase(
train_model,
scheduling_inference: PlacementGroupSchedulingStrategy,
) -> dict[str, object]:
ray.get(train_model.reset_model.remote())
llm = launch_llm(scheduling_inference, backend="sparse_nccl")
try:
sparse_before = collect_vllm_generations(llm)
ray.get(llm.sleep.remote(level=0))
master_address, master_port = 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())
sparse_update_info, 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())
ray.get(llm.wake_up.remote(tags=["scheduling"]))
sparse_after = collect_vllm_generations(llm)
return {
"sparse_before": sparse_before,
"sparse_after": sparse_after,
"selected_token_ids": selected_token_ids,
"patch_digest": patch_digest,
"sparse_payload_bytes": sparse_payload_bytes,
"sparse_send_ms": sparse_send_ms,
}
finally:
ray.kill(llm)
ray.init()
try:
train_model = TrainModel.remote(MODEL_NAME)
pg_inference = placement_group([{"GPU": 1, "CPU": 0}])
ray.get(pg_inference.ready())
scheduling_inference = PlacementGroupSchedulingStrategy(
placement_group=pg_inference,
placement_group_capture_child_tasks=True,
placement_group_bundle_index=0,
)
dense_results = run_dense_phase(train_model, scheduling_inference)
sparse_results = run_sparse_phase(train_model, scheduling_inference)
baseline_equal = token_sequences_match(
dense_results["dense_before"],
sparse_results["sparse_before"],
)
patch_selection_equal = (
dense_results["selected_token_ids"] == sparse_results["selected_token_ids"]
)
patch_digest_equal = dense_results["patch_digest"] == sparse_results["patch_digest"]
after_equal = token_sequences_match(
dense_results["dense_after"],
sparse_results["sparse_after"],
)
any_output_changed = any(
before["token_ids"] != after["token_ids"]
for before, after in zip(
dense_results["dense_before"],
dense_results["dense_after"],
)
)
dense_payload_mb = dense_results["dense_payload_bytes"] / (1024 * 1024)
sparse_payload_mb = sparse_results["sparse_payload_bytes"] / (1024 * 1024)
print_generations(
"Dense baseline outputs",
PROMPTS,
dense_results["dense_before"],
)
print_generations(
"Sparse baseline outputs", PROMPTS, sparse_results["sparse_before"]
)
print_generations(
"Dense outputs after update", PROMPTS, dense_results["dense_after"]
)
print_generations(
"Sparse outputs after update",
PROMPTS,
sparse_results["sparse_after"],
)
print(f"patched_token_ids = {dense_results['selected_token_ids']}")
print(f"patch_selection_equal = {patch_selection_equal}")
print(f"dense_patch_digest = {dense_results['patch_digest']}")
print(f"sparse_patch_digest = {sparse_results['patch_digest']}")
print(f"patch_digest_equal = {patch_digest_equal}")
print(f"baseline_equal = {baseline_equal}")
print(f"after_equal = {after_equal}")
print(f"any_output_changed = {any_output_changed}")
print(f"dense_payload_mb = {dense_payload_mb:.2f}")
print(f"sparse_payload_mb = {sparse_payload_mb:.2f}")
print(f"dense_send_ms = {dense_results['dense_send_ms']:.2f}")
print(f"sparse_send_ms = {sparse_results['sparse_send_ms']:.2f}")
if not baseline_equal:
raise RuntimeError(
"Dense and sparse phases did not start from the same baseline"
)
if not patch_selection_equal:
raise RuntimeError("Dense and sparse phases used different sparse patches")
if not patch_digest_equal:
raise RuntimeError("Dense and sparse phases produced different patch values")
if not after_equal:
raise RuntimeError("Dense and sparse updates produced different outputs")
if not any_output_changed:
raise RuntimeError("Patch did not change the observed outputs")
finally:
ray.shutdown()
examples/rl/routed_experts_e2e.py
+384
View File
@@ -0,0 +1,384 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""
End-to-end example for routed experts capture with hybrid models.
Validates that:
1. routed_experts is returned in CompletionOutput for MoE models.
2. Expert IDs are within valid range.
3. Results are deterministic across runs (baseline vs reference).
Usage:
python examples/rl/routed_experts_e2e.py \
--model Qwen/Qwen3-30B-A3B \
--tp 4 \
--max-model-len 4096 \
--num-prompts 20 \
--max-new-tokens 50
"""
from __future__ import annotations
import argparse
import asyncio
import logging
import os
import uuid
from dataclasses import dataclass, field
import numpy as np
from vllm.engine.arg_utils import AsyncEngineArgs
logger = logging.getLogger(__name__)
DEFAULT_MODEL = "Qwen/Qwen3-30B-A3B"
TEST_PROMPTS = [
"Hello, my name is",
"The capital of France is",
"Explain quantum computing in simple terms:",
"Write a Python function that sorts a list:",
"The meaning of life is",
"In a distant galaxy, there was a",
"The best way to learn programming is",
"Once upon a time in a land far away,",
"The theory of relativity states that",
"How does photosynthesis work?",
"Describe the process of machine learning:",
"What are the benefits of exercise?",
"The history of artificial intelligence began",
"Translate the following to French: Hello world",
"Summarize the plot of Romeo and Juliet:",
"What is the difference between TCP and UDP?",
"The water cycle consists of",
"Explain how a neural network learns:",
"The periodic table organizes elements by",
"Write a haiku about the ocean:",
]
@dataclass
class InferenceResult:
"""Result from a single inference run."""
experts_list: list[np.ndarray] = field(default_factory=list)
token_ids_list: list[list[int]] = field(default_factory=list)
num_experts: int = 0
# ---------------------------------------------------------------------------
# Inference helpers
# ---------------------------------------------------------------------------
async def _run_async_inference(
engine_args: AsyncEngineArgs,
prompts: list[str],
max_new_tokens: int,
) -> InferenceResult:
"""Run inference using AsyncLLM."""
from vllm.sampling_params import SamplingParams
from vllm.v1.engine.async_llm import AsyncLLM
engine = AsyncLLM.from_engine_args(engine_args)
hf_config = engine.model_config.hf_text_config
num_experts: int = getattr(hf_config, "num_experts", 0) or getattr(
hf_config, "num_local_experts", 0
)
assert num_experts > 0, "Could not determine num_experts from model config"
sampling_params = SamplingParams(
temperature=0,
max_tokens=max_new_tokens,
)
async def _generate_one(prompt: str, idx: int):
request_id = str(uuid.uuid4())
final_output = None
async for output in engine.generate(prompt, sampling_params, request_id):
final_output = output
assert final_output is not None
completion = final_output.outputs[0]
routed = completion.routed_experts
num_prompt_tokens = len(final_output.prompt_token_ids)
num_generated_tokens = len(completion.token_ids)
expected_len = num_prompt_tokens + num_generated_tokens - 1
assert routed is not None, f"Prompt {idx}: routed_experts is None"
assert routed.shape[0] == expected_len, (
f"Prompt {idx}: routed_experts length {routed.shape[0]} != "
f"prompt ({num_prompt_tokens}) + generated ({num_generated_tokens})"
f" - 1 = {expected_len}"
)
return idx, routed, list(completion.token_ids)
tasks = [_generate_one(p, i) for i, p in enumerate(prompts)]
outputs = await asyncio.gather(*tasks)
# Sort by original index to maintain prompt order
outputs.sort(key=lambda x: x[0])
result = InferenceResult(num_experts=num_experts)
for _, routed, token_ids in outputs:
result.experts_list.append(routed)
result.token_ids_list.append(token_ids)
engine.shutdown()
return result
def run_inference(
model: str,
prompts: list[str],
max_new_tokens: int = 50,
tp: int = 1,
max_model_len: int = 4096,
) -> InferenceResult:
"""Run inference with routed experts capture enabled via AsyncLLM."""
engine_args = AsyncEngineArgs(
model=model,
enable_return_routed_experts=True,
tensor_parallel_size=tp,
max_model_len=max_model_len,
disable_log_stats=True,
attention_backend="FLASH_ATTN",
)
result = asyncio.run(_run_async_inference(engine_args, prompts, max_new_tokens))
from vllm.platforms import current_platform
if current_platform.is_cuda_alike():
current_platform.empty_cache()
return result
# ---------------------------------------------------------------------------
# Validation helpers
# ---------------------------------------------------------------------------
def validate_expert_ids(
experts_list: list[np.ndarray],
num_experts: int,
) -> None:
"""Check that all expert IDs are within valid range [0, num_experts)."""
for i, experts in enumerate(experts_list):
assert np.all(experts >= 0), (
f"Prompt {i}: negative expert IDs found, min={experts.min()}"
)
assert np.all(experts < num_experts), (
f"Prompt {i}: expert ID out of range [0, {num_experts}), "
f"max={experts.max()}"
)
def validate_shapes(experts_list: list[np.ndarray]) -> None:
"""Check that all routed_experts arrays have at least 2 dimensions."""
for i, experts in enumerate(experts_list):
assert experts.ndim >= 2, (
f"Prompt {i}: expected at least 2D array, got shape {experts.shape}"
)
logger.info("Prompt %d: routed_experts shape = %s", i, experts.shape)
# ---------------------------------------------------------------------------
# Comparison helpers
# ---------------------------------------------------------------------------
def compare_token_ids(
baseline: list[list[int]],
reference: list[list[int]],
) -> float:
"""Compare token IDs from two runs. Returns mismatch ratio."""
assert len(baseline) == len(reference), (
f"Length mismatch: {len(baseline)} vs {len(reference)}"
)
total_tokens = 0
total_mismatches = 0
for i, (base, ref) in enumerate(zip(baseline, reference)):
min_len = min(len(base), len(ref))
max_len = max(len(base), len(ref))
matches = 0
for a, b in zip(base[:min_len], ref[:min_len]):
if a != b:
break
matches += 1
total_mismatches += max_len - matches
total_tokens += max_len
if matches < min_len or len(base) != len(ref):
print(
f" Prompt {i}: token_ids len={len(base)} vs {len(ref)}, "
f"mismatches={max_len - matches}/{max_len}"
)
if total_tokens == 0:
raise ValueError("No tokens to compare")
mismatch_ratio = total_mismatches / total_tokens
print(
f"Token ID mismatches: {total_mismatches}/{total_tokens} ({mismatch_ratio:.4%})"
)
return mismatch_ratio
def compare_routed_experts(
baseline: list[np.ndarray],
reference: list[np.ndarray],
threshold: float = 0.05,
) -> float:
"""Compare two runs of routed experts. Returns mismatch ratio.
Raises AssertionError if ratio exceeds threshold.
"""
assert len(baseline) == len(reference), (
f"Length mismatch: {len(baseline)} vs {len(reference)}"
)
total_elements = 0
total_mismatches = 0
for i, (base, ref) in enumerate(zip(baseline, reference)):
min_len = min(len(base), len(ref))
max_len = max(len(base), len(ref))
if min_len == 0:
continue
base_trimmed = base[:min_len]
ref_trimmed = ref[:min_len]
matches = 0
for a, b in zip(base_trimmed, ref_trimmed):
if a.sum() != b.sum():
break
matches += 1
total_mismatches += max_len - matches
total_elements += max_len
if matches < min_len or len(base) != len(ref):
print(
f" Prompt {i}: routed_experts len={len(base)} vs {len(ref)}, "
f"mismatches={max_len - matches}/{max_len}"
)
if total_elements == 0:
raise ValueError("No elements to compare")
mismatch_ratio = total_mismatches / total_elements
print(
f"Routed experts mismatches: {total_mismatches}/{total_elements} "
f"({mismatch_ratio:.4%})"
)
assert mismatch_ratio < threshold, (
f"Too many mismatches: {total_mismatches}/{total_elements} "
f"({mismatch_ratio:.4%}) exceeds threshold {threshold:.4%}"
)
return mismatch_ratio
# ---------------------------------------------------------------------------
# CLI entry point
# ---------------------------------------------------------------------------
def main():
os.environ.setdefault("VLLM_BATCH_INVARIANT", "1")
parser = argparse.ArgumentParser(
description="Test routed experts capture for MoE models"
)
parser.add_argument("--model", type=str, default=DEFAULT_MODEL)
parser.add_argument("--tp", type=int, default=1)
parser.add_argument("--max-model-len", type=int, default=4096)
parser.add_argument("--num-prompts", type=int, default=20)
parser.add_argument("--max-new-tokens", type=int, default=50)
parser.add_argument(
"--deterministic",
action="store_true",
help="Run twice and compare results for determinism check",
)
parser.add_argument(
"--threshold",
type=float,
default=0.05,
help="Maximum allowed mismatch ratio for determinism check",
)
args = parser.parse_args()
logging.basicConfig(level=logging.INFO)
prompts = TEST_PROMPTS[: args.num_prompts]
print(f"Model: {args.model}")
print(f"TP: {args.tp}")
print(f"Prompts: {len(prompts)}")
print(f"Max new tokens: {args.max_new_tokens}")
print()
print("=== Run 1 (baseline) ===")
baseline = run_inference(
model=args.model,
prompts=prompts,
max_new_tokens=args.max_new_tokens,
tp=args.tp,
max_model_len=args.max_model_len,
)
print(f"num_experts (from model config): {baseline.num_experts}")
print("\n=== Validation ===")
validate_shapes(baseline.experts_list)
validate_expert_ids(baseline.experts_list, num_experts=baseline.num_experts)
print(f"All {len(baseline.experts_list)} results passed validation.")
for i, experts in enumerate(baseline.experts_list):
print(
f" Prompt {i}: shape={experts.shape}, "
f"min={experts.min()}, max={experts.max()}"
)
if args.deterministic:
print("\n=== Run 2 (reference) ===")
reference = run_inference(
model=args.model,
prompts=prompts,
max_new_tokens=args.max_new_tokens,
tp=args.tp,
max_model_len=args.max_model_len,
)
print("\n=== Determinism Check ===")
validate_expert_ids(reference.experts_list, num_experts=baseline.num_experts)
print("\n--- Token IDs ---")
token_mismatch = compare_token_ids(
baseline.token_ids_list, reference.token_ids_list
)
print("\n--- Routed Experts ---")
expert_mismatch = compare_routed_experts(
baseline.experts_list,
reference.experts_list,
threshold=args.threshold,
)
print(
f"\nDeterminism check passed. "
f"Token mismatch: {token_mismatch:.4%}, "
f"Expert mismatch: {expert_mismatch:.4%}"
)
print("\nAll tests passed!")
if __name__ == "__main__":
main()
examples/rl/skip_loading_weights_in_engine_init.py
+53
View File
@@ -0,0 +1,53 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
from vllm import LLM, RequestOutput, SamplingParams
# Sample prompts.
prompts = [
"Hello, my name is",
"The president of the United States is",
"The capital of France is",
"The future of AI is",
]
# Create a sampling params object.
sampling_params = SamplingParams(temperature=0.8, top_p=0.95)
def print_prompts_and_outputs(outputs: list[RequestOutput]) -> None:
print("-" * 60)
for output in outputs:
prompt = output.prompt
generated_text = output.outputs[0].text
print(f"Prompt: {prompt!r}")
print(f"Output: {generated_text!r}")
print("-" * 60)
def main():
# Create an LLM without loading real weights
llm = LLM(
model="Qwen/Qwen3-0.6B",
load_format="dummy",
enforce_eager=True,
tensor_parallel_size=4,
)
outputs = llm.generate(prompts, sampling_params)
print("\nOutputs do not make sense:")
print_prompts_and_outputs(outputs)
# Update load format from `dummy` to `auto`
llm.collective_rpc(
"update_config", args=({"load_config": {"load_format": "auto"}},)
)
# Now reload real weights inplace
llm.collective_rpc("reload_weights")
# Check outputs make sense
outputs = llm.generate(prompts, sampling_params)
print("\nOutputs make sense after loading real weights:")
print_prompts_and_outputs(outputs)
if __name__ == "__main__":
main()