[Doc][RL] Document the weight-transfer HTTP protocol for trainers without vLLM (#62267)

Signed-off-by: aoshen02 <aoshen@inferact.ai>
Co-authored-by: Claude Opus 5.5 <noreply@anthropic.com>
This commit is contained in:
aoshen02
2026-10-18 09:26:56 +00:00
committed by GitHub
co-authored by Claude Opus 5.5
parent 1b2a666391
commit 2365177ddd
6 changed files with 435 additions and 1 deletions
+8 -1
View File
@@ -145,7 +145,7 @@ engine.send_weights() # called on every rank
## API Endpoints ## API Endpoints
When running vLLM as an HTTP server, the following endpoints are available for weight transfer. `HTTPVLLMWeightSyncClient` speaks the first four for you. When running vLLM as an HTTP server, the following endpoints are available for weight transfer. `HTTPVLLMWeightSyncClient` speaks the first five for you. A trainer may also implement the protocol directly without importing vLLM: see the wire contracts in [NCCL](nccl.md#trainer-without-vllm-over-http) and [IPC](ipc.md#trainer-without-vllm-over-http). This keeps the trainer decoupled from vLLM: it needs no vLLM install or version pin, only this documented HTTP and wire contract, so trainer and server can upgrade independently (non-PyTorch trainers such as JAX can use it too).
| Endpoint | Method | Description | | Endpoint | Method | Description |
| -------- | ------ | ----------- | | -------- | ------ | ----------- |
@@ -153,6 +153,7 @@ When running vLLM as an HTTP server, the following endpoints are available for w
| `/start_weight_update` | POST | Start a weight update | | `/start_weight_update` | POST | Start a weight update |
| `/update_weights` | POST | Transfer a batch of weights with backend-specific metadata | | `/update_weights` | POST | Transfer a batch of weights with backend-specific metadata |
| `/finish_weight_update` | POST | Finish the update and optionally commit its `weight_version` | | `/finish_weight_update` | POST | Finish the update and optionally commit its `weight_version` |
| `/shutdown_weight_transfer_engine` | POST | Release the transfer channel without stopping inference |
| `/update_weight_version` | POST | Update `weight_version` without changing model weights | | `/update_weight_version` | POST | Update `weight_version` without changing model weights |
| `/weight_info` | GET | Get the latest committed weight version | | `/weight_info` | GET | Get the latest committed weight version |
| `/pause` | POST | Pause generation before weight sync to handle inflight requests | | `/pause` | POST | Pause generation before weight sync to handle inflight requests |
@@ -162,6 +163,12 @@ When running vLLM as an HTTP server, the following endpoints are available for w
!!! note !!! note
The HTTP weight transfer endpoints require `VLLM_SERVER_DEV_MODE=1` to be set. The HTTP weight transfer endpoints require `VLLM_SERVER_DEV_MODE=1` to be set.
!!! warning
If `/update_weights` fails, the model is left mid-reload while `/health`
stays 200, and `/finish_weight_update` is rejected. Before generating again,
resend a full `/start_weight_update` → `/update_weights` →
`/finish_weight_update` (see [#57854](https://github.com/vllm-project/vllm/pull/57854)).
The Rust frontend's optional gRPC `Control` service exposes the same pause, sleep, weight-transfer, and weight-version lifecycle for trusted sidecars. The `ServerInfo.rl_capabilities` response reports whether weight transfer and sleep mode were configured. Backend-specific `init_info` and `update_info` remain JSON metadata; model tensors continue to move over the configured NCCL, IPC, sparse-NCCL, or sharded-RDT transport. The Rust frontend's optional gRPC `Control` service exposes the same pause, sleep, weight-transfer, and weight-version lifecycle for trusted sidecars. The `ServerInfo.rl_capabilities` response reports whether weight transfer and sleep mode were configured. Backend-specific `init_info` and `update_info` remain JSON metadata; model tensors continue to move over the configured NCCL, IPC, sparse-NCCL, or sharded-RDT transport.
## Extending the System ## Extending the System
+38
View File
@@ -110,6 +110,43 @@ mode carries a per-chunk barrier across ranks: without it, a rank could overwrit
its buffer while its colocated worker is still reading the current chunk. This is its buffer while its colocated worker is still reading the current chunk. This is
handled inside `send_weights()`. handled inside `send_weights()`.
## Trainer Without vLLM over HTTP
A colocated trainer can speak the [HTTP endpoints](README.md#api-endpoints)
directly, with only PyTorch and an HTTP client. The server needs
`VLLM_ALLOW_INSECURE_SERIALIZATION=1`; without it, init and start return 200 and
the first `/update_weights` returns 500. Reference:
[`rlhf_http_ipc_no_vllm.py`](../../../examples/rl/rlhf_http_ipc_no_vllm.py).
`POST /init_weight_transfer_engine` body: `{"init_info": {"packed": false}}`.
`POST /update_weights` body `{"update_info": {...}}`:
| Field | Meaning |
| ----- | ------- |
| `names`, `dtype_names`, `shapes` | Checkpoint names, `torch` dtype attribute names (`"bfloat16"`), full shapes |
| `ipc_handles_pickled` | `base64(pickle.dumps(handles))`. Unpacked: a list with one dict per name. Packed: one dict for the shared buffer |
| `tensor_sizes` | Packed only: byte size of each tensor in the chunk, in `names` order |
- A handle dict is `{gpu_uuid: reduce_tensor(t)[1]}` with one entry per GPU the
server uses (`torch.multiprocessing.reductions.reduce_tensor`).
`gpu_uuid = str(torch.cuda.get_device_properties(i).uuid)`, **without** the
`GPU-` prefix that `nvidia-smi` prints. The worker overwrites `args[6]`
(device index) with its own.
- `ipc_handles` (unpickled) is not JSON-serializable; over HTTP use
`ipc_handles_pickled`.
- Packed: export one `uint8` buffer once and reuse it for every chunk. Copy
`t.contiguous().view(-1).view(torch.uint8)` in `names` order and close a chunk
**before** the next tensor would overflow it (unlike NCCL, which closes after).
A tensor larger than the buffer is an error. Synchronize before each
`/update_weights`; do not overwrite the buffer until the call returns.
- Keep every exported tensor alive until `/finish_weight_update` returns, then
drop the references and call `torch.cuda.ipc_collect()`.
- Multi-rank trainers: each rank exports handles for its own GPU, the ranks
gather the `{uuid: args}` dicts to rank 0 over their own `torch.distributed`
group and merge them per index, and rank 0 sends all HTTP. In packed mode,
barrier before and after each `/update_weights`.
## Rank-Local Updates ## Rank-Local Updates
For different parameter subsets per worker, `update_info` may be a list indexed For different parameter subsets per worker, `update_info` may be a list indexed
@@ -119,4 +156,5 @@ worker to participate.
## Examples ## Examples
- [RLHF with IPC weight syncing (`vllm serve`, HTTP)](../../../examples/rl/rlhf_http_ipc.py) - **Start here.** Server and training model share a single GPU; HTTP control plane, CUDA IPC data plane. Launches and tears down its own server - [RLHF with IPC weight syncing (`vllm serve`, HTTP)](../../../examples/rl/rlhf_http_ipc.py) - **Start here.** Server and training model share a single GPU; HTTP control plane, CUDA IPC data plane. Launches and tears down its own server
- [IPC weight syncing from a trainer that does not import vLLM](../../../examples/rl/rlhf_http_ipc_no_vllm.py) - `torch` + `requests` only, `ipc_handles_pickled` over HTTP; see [Trainer Without vLLM over HTTP](#trainer-without-vllm-over-http)
- [RLHF with IPC + FSDP2 and expert parallelism](../../../examples/rl/rlhf_ipc_fsdp_ep.py) - Multi-rank trainer colocated with a `--data-parallel-size 4` server on the same 4 GPUs: every FSDP rank builds an engine and joins the handle all-gather, with packed chunking and sleep/wake around the transfer - [RLHF with IPC + FSDP2 and expert parallelism](../../../examples/rl/rlhf_ipc_fsdp_ep.py) - Multi-rank trainer colocated with a `--data-parallel-size 4` server on the same 4 GPUs: every FSDP rank builds an engine and joins the handle all-gather, with packed chunking and sleep/wake around the transfer
+46
View File
@@ -130,6 +130,51 @@ pass — this is the invariant to test first.
IPC does not read `metadata()` at all: it derives the update info from IPC does not read `metadata()` at all: it derives the update info from
iteration as it goes, so it cannot observe a divergence. iteration as it goes, so it cannot observe a divergence.
## Trainer Without vLLM over HTTP
A trainer can speak the [HTTP endpoints](README.md#api-endpoints) directly, with
only PyTorch and an HTTP client. It mints the `ncclUniqueId` itself and is NCCL
rank 0. Reference: [`rlhf_http_nccl_no_vllm.py`](../../../examples/rl/rlhf_http_nccl_no_vllm.py).
`POST /init_weight_transfer_engine` body `{"init_info": {...}}`:
| Field | Default | Meaning |
| ----- | ------- | ------- |
| `nccl_unique_id_b64` | — | Standard (not URL-safe) base64 of the 128-byte `ncclUniqueId`. Replaces `master_address` / `master_port` |
| `rank_offset` | — | First worker rank; `1` for one trainer rank |
| `world_size` | — | Trainer ranks + `GET /get_world_size` |
| `packed` | `false` | Packed broadcasting |
| `packed_buffer_size_bytes` | 1 GiB | Must match the trainer's chunking |
| `packed_num_buffers` | 2 | Worker-side receive buffers |
- Worker rank = `rank_offset + dp_index * (TP * PP) + rank`.
- `/init_weight_transfer_engine` returns only after all ranks join: send it in
the background while the trainer calls `ncclCommInitRank(world_size, uid, 0)`.
There is no timeout; if the trainer never joins, the workers wait forever.
- Right after init every worker issues a warm-up `all_reduce` (sum) of
`torch.zeros(1)` (one float32 element). The trainer must issue the same one
before any other collective.
Each sync: `POST /start_weight_update`, then `POST /update_weights` with
`{"update_info": {"names", "dtype_names", "shapes"}}` **in the background**
while rank 0 broadcasts in `names` order, then `POST /finish_weight_update`.
`dtype_names` are `torch` attribute names (`"bfloat16"`); `shapes` are full
checkpoint shapes.
- Unpacked: one broadcast per tensor.
- Packed: concatenate `t.contiguous().view(-1).view(torch.uint8)` in `names`
order; broadcast a chunk as soon as its size **exceeds**
`packed_buffer_size_bytes` (the crossing tensor stays in the chunk), then the
remainder.
Trainer-side communicators without vLLM:
| Communicator | Status |
| ------------ | ------ |
| `torch.distributed.ProcessGroupNCCL` on a private `TCPStore` + `eager_connect_single_device` | Works; the unique id is read from store key `"0"`, a torch implementation detail (used by the example) |
| `ctypes` against torch's bundled `libnccl.so.2` | Works |
| `torch.cuda.nccl.init_rank` | Fails on Python 3.10+ (`PY_SSIZE_T_CLEAN` `SystemError`) |
## Sparse NCCL ## Sparse NCCL
Sparse, flat-index weight patches use Sparse, flat-index weight patches use
@@ -196,6 +241,7 @@ Sparse NCCL sends only `O(nnz)` indices and values. Checkpoint application still
## Examples ## Examples
- [RLHF with NCCL weight syncing (`vllm serve`, HTTP)](../../../examples/rl/rlhf_http_nccl.py) - **Start here.** Trainer on one GPU, 2x tensor-parallel fp8 server on two others; HTTP control plane, NCCL data plane. Launches and tears down its own server - [RLHF with NCCL weight syncing (`vllm serve`, HTTP)](../../../examples/rl/rlhf_http_nccl.py) - **Start here.** Trainer on one GPU, 2x tensor-parallel fp8 server on two others; HTTP control plane, NCCL data plane. Launches and tears down its own server
- [NCCL weight syncing from a trainer that does not import vLLM](../../../examples/rl/rlhf_http_nccl_no_vllm.py) - `torch` + `requests` only, `nccl_unique_id_b64` rendezvous; see [Trainer Without vLLM over HTTP](#trainer-without-vllm-over-http)
- [RLHF with NCCL + FSDP2 and expert parallelism](../../../examples/rl/rlhf_nccl_fsdp_ep.py) - Multi-rank trainer: every FSDP rank builds an engine and joins the `full_tensor()` gather while only rank 0 touches the wire - [RLHF with NCCL + FSDP2 and expert parallelism](../../../examples/rl/rlhf_nccl_fsdp_ep.py) - Multi-rank trainer: every FSDP rank builds an engine and joins the `full_tensor()` gather while only rank 0 touches the wire
- [RLHF with sparse NCCL weight syncing (offline, Ray)](../../../examples/rl/rlhf_sparse_nccl.py) - Qwen3 MoE per-expert checkpoint updates with one trainer GPU and two TP2/EP2 inference GPUs - [RLHF with sparse NCCL weight syncing (offline, Ray)](../../../examples/rl/rlhf_sparse_nccl.py) - Qwen3 MoE per-expert checkpoint updates with one trainer GPU and two TP2/EP2 inference GPUs
- [RLHF with async weight syncing (offline, Ray)](../../../examples/rl/rlhf_async_new_apis.py) - Async generation with mid-flight pause, weight sync, resume, and validation against a fresh model; uses an in-process `AsyncLLMEngine` with `RayVLLMWeightSyncClient` - [RLHF with async weight syncing (offline, Ray)](../../../examples/rl/rlhf_async_new_apis.py) - Async generation with mid-flight pause, weight sync, resume, and validation against a fresh model; uses an in-process `AsyncLLMEngine` with `RayVLLMWeightSyncClient`
+168
View File
@@ -0,0 +1,168 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""CUDA IPC weight sync over HTTP from a trainer that never imports vLLM.
The trainer uses only `torch`, `requests`, `safetensors`, `pybase64` and the
stdlib, and speaks the wire protocol in
`docs/training/weight_transfer/ipc.md` ("Trainer Without vLLM over HTTP"):
it exports its tensors with `torch.multiprocessing.reductions.reduce_tensor`
and sends the handles, pickled and base64-encoded, in `ipc_handles_pickled`.
IPC requires colocation: this process holds a full copy of the weights on every
GPU the server uses, and keys each handle by that GPU's UUID. A multi-process
trainer would instead gather the per-rank `{uuid: args}` dicts to one rank.
Start a server; IPC handles are pickled, so it needs
`VLLM_ALLOW_INSECURE_SERIALIZATION=1` (the weights it starts with are discarded):
VLLM_SERVER_DEV_MODE=1 VLLM_ALLOW_INSECURE_SERIALIZATION=1 \
CUDA_VISIBLE_DEVICES=0 vllm serve Qwen/Qwen3-0.6B \
--load-format dummy --enforce-eager --gpu-memory-utilization 0.4 \
--weight-transfer-config '{"backend": "ipc"}'
Then run the trainer on the same GPU(s) (`--packed` for chunked transfer):
CUDA_VISIBLE_DEVICES=0 python examples/rl/rlhf_http_ipc_no_vllm.py \
--checkpoint /path/to/Qwen3-0.6B
For a TP=2 server on GPUs 0-1, use `CUDA_VISIBLE_DEVICES=0,1` on both sides and
pass `--num-gpus 2` to the trainer.
"""
import sys
class _BlockVLLM:
def find_spec(self, name, path=None, target=None):
if name.split(".")[0] == "vllm":
raise ImportError(f"this trainer must not import {name!r}")
return None
sys.meta_path.insert(0, _BlockVLLM())
import argparse # noqa: E402
import glob # noqa: E402
import pickle # noqa: E402
import pybase64 as base64 # noqa: E402
import requests # noqa: E402
import torch # noqa: E402
from safetensors.torch import load_file # noqa: E402
from torch.multiprocessing.reductions import reduce_tensor # noqa: E402
parser = argparse.ArgumentParser()
parser.add_argument("--url", default="http://localhost:8000")
parser.add_argument("--checkpoint", required=True, help="dir with *.safetensors")
parser.add_argument("--num-gpus", type=int, default=1)
parser.add_argument("--packed", action="store_true")
parser.add_argument("--packed-buffer-size-bytes", type=int, default=512 << 20)
args = parser.parse_args()
def post(path: str, body: dict | None = None) -> dict:
r = requests.post(args.url + path, json=body or {}, timeout=600)
r.raise_for_status()
return r.json()
def generate(prompts: list[str]) -> list[str]:
model = requests.get(args.url + "/v1/models", timeout=60).json()["data"][0]["id"]
body = {"model": model, "prompt": prompts, "max_tokens": 16, "temperature": 0}
choices = post("/v1/completions", body)["choices"]
return [c["text"] for c in sorted(choices, key=lambda c: c["index"])]
# Weights are {gpu_uuid: {name: tensor}}: one full copy per colocated GPU. The
# UUID is str(torch.cuda.get_device_properties(i).uuid), with no "GPU-" prefix.
Weights = dict[str, dict[str, torch.Tensor]]
def update_weights(
weights: Weights,
names: list[str],
ipc_handles: list[dict[str, tuple]] | dict[str, tuple],
tensor_sizes: list[int] | None = None,
) -> None:
w = next(iter(weights.values()))
update_info = {
"names": names,
"dtype_names": [str(w[n].dtype).removeprefix("torch.") for n in names],
"shapes": [list(w[n].shape) for n in names],
"ipc_handles_pickled": base64.b64encode(pickle.dumps(ipc_handles)).decode(),
}
if tensor_sizes is not None:
update_info["tensor_sizes"] = tensor_sizes
post("/update_weights", {"update_info": update_info})
def send_unpacked(weights: Weights, names: list[str]) -> list[torch.Tensor]:
"""One /update_weights call; one {uuid: args} dict per parameter."""
copies = {u: {n: w[n].contiguous() for n in names} for u, w in weights.items()}
handles = [{u: reduce_tensor(copies[u][n])[1] for u in copies} for n in names]
update_weights(weights, names, handles)
return [t for c in copies.values() for t in c.values()]
def send_packed(weights: Weights, names: list[str], cap: int) -> list[torch.Tensor]:
"""One /update_weights call per chunk; one {uuid: args} dict per chunk."""
buffers = {
u: torch.empty(cap, dtype=torch.uint8, device=w[names[0]].device)
for u, w in weights.items()
}
handle = {u: reduce_tensor(b)[1] for u, b in buffers.items()} # export once
chunk: list[str] = []
sizes: list[int] = []
def flush() -> None:
for b in buffers.values():
torch.accelerator.synchronize(b.device) # copies must land before the read
update_weights(weights, chunk, handle, sizes)
chunk.clear()
sizes.clear()
for n in names:
nbytes = next(iter(weights.values()))[n].nbytes
assert nbytes <= cap, f"{n} ({nbytes} B) exceeds the buffer"
if chunk and sum(sizes) + nbytes > cap: # close *before* overflowing
flush()
offset = sum(sizes)
for u, w in weights.items():
flat = w[n].contiguous().view(-1).view(torch.uint8)
buffers[u][offset : offset + nbytes].copy_(flat)
chunk.append(n)
sizes.append(nbytes)
if chunk:
flush()
return list(buffers.values())
def main() -> None:
weights: Weights = {}
for i in range(args.num_gpus):
uuid = str(torch.cuda.get_device_properties(i).uuid)
weights[uuid] = {}
for path in sorted(glob.glob(f"{args.checkpoint}/*.safetensors")):
weights[uuid].update(load_file(path, device=f"cuda:{i}"))
names = sorted(next(iter(weights.values())))
prompts = ["The capital of France is", "def fibonacci(n):"]
print("before:", generate(prompts))
post("/init_weight_transfer_engine", {"init_info": {"packed": args.packed}})
post("/start_weight_update")
if args.packed:
exported = send_packed(weights, names, args.packed_buffer_size_bytes)
else:
exported = send_unpacked(weights, names)
post("/finish_weight_update", {"weight_version": None})
# The worker reads the exported memory until finish; only then release it.
del exported
torch.cuda.ipc_collect()
print("after: ", generate(prompts))
assert not [m for m in sys.modules if m.split(".")[0] == "vllm"]
if __name__ == "__main__":
main()
+174
View File
@@ -0,0 +1,174 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""NCCL weight sync over HTTP from a trainer that never imports vLLM.
The trainer uses only `torch`, `requests`, `safetensors`, `pybase64` and the
stdlib, and speaks the wire protocol in
`docs/training/weight_transfer/nccl.md` ("Trainer Without vLLM over HTTP"):
`nccl_unique_id_b64` rendezvous, a one-element warm-up `all_reduce`, then
`/start_weight_update` -> `/update_weights` (concurrent with the rank-0
broadcast) -> `/finish_weight_update`.
The NCCL communicator is a plain `torch.distributed.ProcessGroupNCCL` over a
private `TCPStore` that only the trainer joins. Its `ncclUniqueId` is read back
from store key "0", which is a torch implementation detail.
Start a server on GPUs 0-1 (the weights it starts with are discarded):
VLLM_SERVER_DEV_MODE=1 CUDA_VISIBLE_DEVICES=0,1 vllm serve Qwen/Qwen3-0.6B \
--tensor-parallel-size 2 --load-format dummy --enforce-eager \
--weight-transfer-config '{"backend": "nccl"}'
Then run the trainer on GPU 2 (`--packed` for packed broadcasting):
CUDA_VISIBLE_DEVICES=2 python examples/rl/rlhf_http_nccl_no_vllm.py \
--checkpoint /path/to/Qwen3-0.6B
"""
import sys
class _BlockVLLM:
def find_spec(self, name, path=None, target=None):
if name.split(".")[0] == "vllm":
raise ImportError(f"this trainer must not import {name!r}")
return None
sys.meta_path.insert(0, _BlockVLLM())
import argparse # noqa: E402
import datetime # noqa: E402
import glob # noqa: E402
import threading # noqa: E402
import pybase64 as base64 # noqa: E402
import requests # noqa: E402
import torch # noqa: E402
import torch.distributed as dist # noqa: E402
from safetensors.torch import load_file # noqa: E402
parser = argparse.ArgumentParser()
parser.add_argument("--url", default="http://localhost:8000")
parser.add_argument("--checkpoint", required=True, help="dir with *.safetensors")
parser.add_argument("--packed", action="store_true")
parser.add_argument("--packed-buffer-size-bytes", type=int, default=64 << 20)
parser.add_argument("--store-port", type=int, default=29612)
args = parser.parse_args()
def post(path: str, body: dict | None = None) -> dict:
r = requests.post(args.url + path, json=body or {}, timeout=600)
r.raise_for_status()
return r.json()
class PostInBackground(threading.Thread):
"""An HTTP call that only returns once this process joins the collective."""
def __init__(self, path: str, body: dict) -> None:
super().__init__()
self.path, self.body = path, body
self.error: BaseException | None = None
self.start()
def run(self) -> None:
try:
post(self.path, self.body)
except BaseException as e:
self.error = e
def result(self) -> None:
self.join()
if self.error is not None:
raise self.error
def generate(prompts: list[str]) -> list[str]:
model = requests.get(args.url + "/v1/models", timeout=60).json()["data"][0]["id"]
body = {"model": model, "prompt": prompts, "max_tokens": 16, "temperature": 0}
choices = post("/v1/completions", body)["choices"]
return [c["text"] for c in sorted(choices, key=lambda c: c["index"])]
def main() -> None:
device = torch.device("cuda:0")
torch.accelerator.set_device_index(device)
weights: dict[str, torch.Tensor] = {}
for path in sorted(glob.glob(f"{args.checkpoint}/*.safetensors")):
weights.update(load_file(path, device="cuda:0"))
names = sorted(weights)
prompts = ["The capital of France is", "def fibonacci(n):"]
print("before:", generate(prompts))
# Trainer is NCCL rank 0; vLLM worker i is rank 1 + i.
world_size = 1 + requests.get(args.url + "/get_world_size").json()["world_size"]
store = dist.TCPStore(
"127.0.0.1", args.store_port, is_master=True, wait_for_workers=False
)
pg = dist.ProcessGroupNCCL(store, 0, world_size, dist.ProcessGroupNCCL.Options())
connect = threading.Thread(target=pg.eager_connect_single_device, args=(device,))
connect.start() # publishes the unique id, then waits in ncclCommInitRank
reader = dist.TCPStore("127.0.0.1", args.store_port, is_master=False)
reader.wait(["0"], datetime.timedelta(seconds=60))
unique_id = reader.get("0")
assert len(unique_id) == 128
init = PostInBackground(
"/init_weight_transfer_engine",
{
"init_info": {
"nccl_unique_id_b64": base64.b64encode(unique_id).decode(),
"rank_offset": 1,
"world_size": world_size,
"packed": args.packed,
"packed_buffer_size_bytes": args.packed_buffer_size_bytes,
}
},
)
connect.join()
pg.allreduce([torch.zeros(1, device=device)]).wait() # vLLM's warm-up
init.result()
def broadcast(t: torch.Tensor) -> None:
opts = dist.BroadcastOptions()
opts.rootRank = 0
pg.broadcast([t], opts).wait()
post("/start_weight_update")
update_info = {
"names": names,
"dtype_names": [str(weights[n].dtype).removeprefix("torch.") for n in names],
"shapes": [list(weights[n].shape) for n in names],
}
update = PostInBackground("/update_weights", {"update_info": update_info})
if not args.packed:
for n in names:
broadcast(weights[n].contiguous())
else:
# Same chunking as the worker: a chunk closes right after the tensor
# that pushes it past the buffer size.
chunk: list[torch.Tensor] = []
size = 0
for n in names:
chunk.append(weights[n].contiguous().view(-1).view(torch.uint8))
size += chunk[-1].numel()
if size > args.packed_buffer_size_bytes:
broadcast(torch.cat(chunk))
chunk, size = [], 0
if chunk:
broadcast(torch.cat(chunk))
torch.accelerator.synchronize()
update.result()
post("/finish_weight_update", {"weight_version": None})
print("after: ", generate(prompts))
post("/shutdown_weight_transfer_engine")
pg.shutdown()
assert not [m for m in sys.modules if m.split(".")[0] == "vllm"]
if __name__ == "__main__":
main()
@@ -59,6 +59,7 @@ CHECK_IMPORTS = {
"vllm/distributed/device_communicators/shm_tensor_arena.py", "vllm/distributed/device_communicators/shm_tensor_arena.py",
"vllm/distributed/weight_transfer/ipc_engine.py", "vllm/distributed/weight_transfer/ipc_engine.py",
"vllm/distributed/weight_transfer/clients.py", "vllm/distributed/weight_transfer/clients.py",
"examples/rl/rlhf_http_ipc_no_vllm.py",
"vllm/model_executor/model_loader/weight_cache/protocol.py", "vllm/model_executor/model_loader/weight_cache/protocol.py",
"tests/distributed/test_shm_broadcast.py", "tests/distributed/test_shm_broadcast.py",
"tests/distributed/test_weight_transfer.py", "tests/distributed/test_weight_transfer.py",