@tduncan / Nvidia-Optimization-Stack-VLLM / commits / 42d9f8ccdb

support rl feature : weight checker (#51350)

Signed-off-by: shiyuan680 <yangcheng104@huawei.com> Signed-off-by: Ronald1995 <ronaldautomobile@163.com> Signed-off-by: aoshen02 <aoshen02@users.noreply.github.com> Co-authored-by: shiyuan680 <yangcheng104@huawei.com> Co-authored-by: Ronald1995 <ronaldautomobile@163.com> Co-authored-by: DeepSeek Harness <noreply@deepseek.com> Co-authored-by: aoshen02 <aoshen02@users.noreply.github.com> Co-authored-by: Claude Opus 5.5 <noreply@anthropic.com> Co-authored-by: aoshen02 <aoshen@inferact.ai>

shiyuan680 committed Sep 30, 2026 at 11:13 UTC 42d9f8ccdbbb7e24c173272809c66ad627a0de9d
15 files changed +611 -2
docs/training/rlhf.md
+1 -1
@@ -16,7 +16,7 @@ The following open-source RL libraries use vLLM for fast rollouts (sorted alphab
16 - [Unsloth](https://github.com/unslothai/unsloth)
17 - [verl](https://github.com/volcengine/verl)
18
19 -For weight synchronization between training and inference, see the [Weight Transfer](weight_transfer/README.md) documentation, which covers the pluggable backend system with [NCCL](weight_transfer/nccl.md) (multi-GPU) and [IPC](weight_transfer/ipc.md) (same-GPU) engines.
19 +For weight synchronization between training and inference, see the [Weight Transfer](weight_transfer/README.md) documentation, which covers the pluggable backend system with [NCCL](weight_transfer/nccl.md) (multi-GPU) and [IPC](weight_transfer/ipc.md) (same-GPU) engines. To verify that an update landed on every rank, see the [Weight Checker](weight_checker.md).
20
21 For pipelining generation and training to improve GPU utilization and throughput, see the [Async Reinforcement Learning](async_rl.md) guide, which covers the pause/resume API for safely updating weights mid-flight.
22
docs/training/weight_checker.md new
+60
@@ -0,0 +1,60 @@
1 +# Weight Checker
2 +
3 +The Weight Checker verifies that an RL weight update landed on every rank. It
4 +hashes each model parameter with SHA-256, so a caller can save the original
5 +digests, reset the weights, transfer them back, and confirm they match.
6 +
7 +Enable it with `VLLM_SERVER_DEV_MODE=1`. All operations use
8 +`POST /weight_checker`:
9 +
10 +| Request | Response |
11 +| --- | --- |
12 +| `{"action": "checksum"}` | `{"checksums": {key: sha256_hex}}` |
13 +| `{"action": "reset"}` | `{"status": "reset"}` |
14 +| `{"action": "compare", "baseline": {key: sha256_hex}}` | `{"match": bool, "mismatches": [key]}` |
15 +
16 +Keys have the form `dp{dp}:pp{pp}:pcp{pcp}:tp{tp}:{tensor_name}`, so each
17 +shard is checked separately. All actions cover the target model's parameters,
18 +which is what weight loading writes; `reset` zeroes them. The endpoint keeps no
19 +state, so the caller holds the baseline. Invalid requests return HTTP 400.
20 +
21 +## RL weight-update check
22 +
23 +```bash
24 +URL=http://localhost:8000
25 +
26 +# 1. Stop serving, then save the original digests.
27 +curl -X POST "$URL/pause?mode=abort"
28 +curl -X POST $URL/weight_checker -H 'Content-Type: application/json' \
29 + -d '{"action":"checksum"}' > baseline.json
30 +
31 +# 2. Zero the weights.
32 +curl -X POST $URL/weight_checker -H 'Content-Type: application/json' \
33 + -d '{"action":"reset"}'
34 +
35 +# 3. Transfer the original weights with start/update/finish_weight_update.
36 +
37 +# 4. Compare against the baseline; expect {"match": true, "mismatches": []}.
38 +curl -X POST $URL/weight_checker -H 'Content-Type: application/json' \
39 + -d "{\"action\":\"compare\",\"baseline\":$(jq .checksums baseline.json)}"
40 +
41 +curl -X POST $URL/resume
42 +```
43 +
44 +## Limitations
45 +
46 +- Stay paused from `checksum` to `compare`: the weights are invalid between
47 + `reset` and the transfer, and EPLB moves experts while serving.
48 +- A request covers the engines managed by the API server it reaches. With
49 + several API servers (for example `--data-parallel-external-lb`), send it to
50 + each one.
51 +- Buffers, draft models and LoRA adapters are not checked.
52 +- CPU backends that repack linear weights hide them from the checker.
53 +- compressed-tensors online transforms are not supported: their Hadamard
54 + weights are normalized once per storage, so restored copies differ.
55 +- MoE weights padded beyond the checkpoint shape (DeepEP hidden size, MXFP4)
56 + keep uninitialized padding, so `compare` reports them as mismatches.
57 +- The engine must be awake: sleep levels 1 and 2 release the weight memory.
58 +- `--offload-backend prefetch` is not supported: parameters point at staging
59 + buffers, not the offloaded weights.
60 +- Hashing copies every weight to CPU, so keep it off latency-sensitive paths.
rust/src/server/src/routes.rs
+2
@@ -18,6 +18,7 @@ mod server_info;
18 mod sleep;
19 mod tokenize;
20 mod version;
21 +mod weight_checker;
22 mod weight_transfer;
23 mod world_size;
24
@@ -142,6 +143,7 @@ fn build_router_with_options(
143 post(weight_transfer::update_weight_version),
144 )
145 .route("/weight_info", get(weight_transfer::weight_info))
146 + .route("/weight_checker", post(weight_checker::weight_checker))
147 .route("/abort_requests", post(abort_requests::abort_requests))
148 .route("/sleep", post(sleep::sleep))
149 .route(
rust/src/server/src/routes/tests.rs
+76
@@ -5813,6 +5813,82 @@ async fn collective_rpc_route_sends_expected_utility_call_and_returns_results()
5813 );
5814 }
5815
5816 +#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
5817 +#[serial]
5818 +async fn weight_checker_route_merges_worker_checksums_and_compares() {
5819 + let (mut app, engine_task) = test_admin_app_with_engine_script(|dealer, push| {
5820 + boxed_test_future(async move {
5821 + let workers = Value::Array(vec![
5822 + Value::Map(vec![(Value::from("tp0:a"), Value::from("1"))]),
5823 + Value::Map(vec![(Value::from("tp1:a"), Value::from("2"))]),
5824 + ]);
5825 + for (method, result) in [
5826 + ("compute_weight_checksums", workers.clone()),
5827 + ("compute_weight_checksums", workers),
5828 + ("reset_weights", Value::Array(vec![Value::Nil, Value::Nil])),
5829 + ] {
5830 + let utility = recv_engine_message(dealer).await;
5831 + let payload = decode_value(&utility[1]).expect("decode utility payload");
5832 + let array = payload.as_array().expect("utility payload array");
5833 + assert_eq!(array[2], Value::from("collective_rpc"));
5834 + assert_eq!(
5835 + array[3].as_array().expect("rpc args")[0],
5836 + Value::from(method)
5837 + );
5838 + let call_id = array[1].as_u64().expect("call id");
5839 + let envelope = UtilityResultEnvelope::without_type_info(result);
5840 + send_outputs(push, utility_outputs(call_id, envelope)).await;
5841 + }
5842 + })
5843 + })
5844 + .await;
5845 +
5846 + let (status, body) =
5847 + post_json(&mut app, "/weight_checker", json!({"action": "checksum"})).await;
5848 + assert_eq!(status, StatusCode::OK);
5849 + assert_eq!(body, json!({"checksums": {"tp0:a": "1", "tp1:a": "2"}}));
5850 +
5851 + let baseline = json!({"tp0:a": "1", "tp1:a": "9", "tp2:a": "3"});
5852 + let (status, body) = post_json(
5853 + &mut app,
5854 + "/weight_checker",
5855 + json!({"action": "compare", "baseline": baseline}),
5856 + )
5857 + .await;
5858 + assert_eq!(status, StatusCode::OK);
5859 + assert_eq!(
5860 + body,
5861 + json!({"match": false, "mismatches": ["tp1:a", "tp2:a"]})
5862 + );
5863 +
5864 + let (status, body) = post_json(&mut app, "/weight_checker", json!({"action": "reset"})).await;
5865 + assert_eq!(status, StatusCode::OK);
5866 + assert_eq!(body, json!({"status": "reset"}));
5867 + engine_task.abort_and_join().await;
5868 +}
5869 +
5870 +#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
5871 +#[serial]
5872 +async fn weight_checker_route_rejects_invalid_requests_before_engine_calls() {
5873 + let (mut app, engine_task) = test_admin_app_with_engine_script(|dealer, _push| {
5874 + boxed_test_future(async move {
5875 + let message = recv_engine_message(dealer).await;
5876 + panic!("invalid request reached engine: {message:?}");
5877 + })
5878 + })
5879 + .await;
5880 + for body in [
5881 + json!({}),
5882 + json!({"action": "frobnicate"}),
5883 + json!({"action": "compare"}),
5884 + json!({"action": "compare", "baseline": {"k": 1}}),
5885 + ] {
5886 + let (status, _) = post_json(&mut app, "/weight_checker", body.clone()).await;
5887 + assert_eq!(status, StatusCode::BAD_REQUEST, "{body}");
5888 + }
5889 + engine_task.abort_and_join().await;
5890 +}
5891 +
5892 #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
5893 #[serial]
5894 async fn sleep_route_uses_python_compatible_default_query_values() {
rust/src/server/src/routes/weight_checker.rs new
+113
@@ -0,0 +1,113 @@
1 +// SPDX-License-Identifier: Apache-2.0
2 +// SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +
4 +use std::collections::{BTreeMap, BTreeSet};
5 +use std::sync::Arc;
6 +
7 +use axum::Json;
8 +use axum::extract::State;
9 +use axum::extract::rejection::JsonRejection;
10 +use serde::{Deserialize, Serialize};
11 +use serde_json::Value as JsonValue;
12 +
13 +use crate::error::{ApiError, invalid_request};
14 +use crate::state::AppState;
15 +use crate::utils::utility_call_error;
16 +
17 +#[derive(Debug, Deserialize, PartialEq)]
18 +#[serde(rename_all = "snake_case")]
19 +pub(crate) enum WeightCheckerAction {
20 + Checksum,
21 + Reset,
22 + Compare,
23 +}
24 +
25 +#[derive(Debug, Deserialize)]
26 +pub(crate) struct WeightCheckerRequest {
27 + action: WeightCheckerAction,
28 + #[serde(default)]
29 + baseline: Option<BTreeMap<String, String>>,
30 +}
31 +
32 +#[derive(Serialize)]
33 +#[serde(untagged)]
34 +pub(crate) enum WeightCheckerResponse {
35 + Checksums {
36 + checksums: BTreeMap<String, String>,
37 + },
38 + Reset {
39 + status: &'static str,
40 + },
41 + Compare {
42 + #[serde(rename = "match")]
43 + matched: bool,
44 + mismatches: Vec<String>,
45 + },
46 +}
47 +
48 +/// Checksum, reset, or compare model weights against a baseline.
49 +pub async fn weight_checker(
50 + State(state): State<Arc<AppState>>,
51 + body: Result<Json<WeightCheckerRequest>, JsonRejection>,
52 +) -> Result<Json<WeightCheckerResponse>, ApiError> {
53 + let Json(body) = body?;
54 + let client = state.engine_core_client();
55 + let no_args = || {
56 + (
57 + Vec::<JsonValue>::new(),
58 + BTreeMap::<String, JsonValue>::new(),
59 + )
60 + };
61 +
62 + if body.action == WeightCheckerAction::Reset {
63 + let (args, kwargs) = no_args();
64 + client
65 + .collective_rpc("reset_weights", None, args, kwargs)
66 + .await
67 + .map_err(|error| utility_call_error("reset_weights", error))?;
68 + return Ok(Json(WeightCheckerResponse::Reset { status: "reset" }));
69 + }
70 + if body.action == WeightCheckerAction::Compare && body.baseline.is_none() {
71 + return Err(invalid_request!(
72 + param = Some("baseline"),
73 + "action='compare' requires a 'baseline' object"
74 + ));
75 + }
76 +
77 + let (args, kwargs) = no_args();
78 + let workers = client
79 + .collective_rpc("compute_weight_checksums", None, args, kwargs)
80 + .await
81 + .map_err(|error| utility_call_error("compute_weight_checksums", error))?;
82 + let mut checksums = BTreeMap::new();
83 + for worker in workers {
84 + let worker: BTreeMap<String, String> = rmpv::ext::from_value(worker).map_err(|error| {
85 + ApiError::server_error(format!("invalid weight checksums: {error}"))
86 + })?;
87 + for (key, digest) in worker {
88 + if checksums.contains_key(&key) {
89 + return Err(ApiError::server_error(format!(
90 + "Duplicate weight checksum keys: {key}"
91 + )));
92 + }
93 + checksums.insert(key, digest);
94 + }
95 + }
96 +
97 + let Some(baseline) = body.baseline.filter(|_| body.action == WeightCheckerAction::Compare)
98 + else {
99 + return Ok(Json(WeightCheckerResponse::Checksums { checksums }));
100 + };
101 + let mismatches: Vec<String> = checksums
102 + .keys()
103 + .chain(baseline.keys())
104 + .collect::<BTreeSet<_>>()
105 + .into_iter()
106 + .filter(|key| checksums.get(*key) != baseline.get(*key))
107 + .cloned()
108 + .collect();
109 + Ok(Json(WeightCheckerResponse::Compare {
110 + matched: mismatches.is_empty(),
111 + mismatches,
112 + }))
113 +}
tests/entrypoints/serve/dev/rlhf/state_transitions/test_weight_checker.py new
+133
@@ -0,0 +1,133 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +import os
4 +from unittest.mock import patch
5 +
6 +import pytest
7 +import regex as re
8 +import requests
9 +import torch
10 +from transformers import AutoModelForCausalLM
11 +
12 +from tests.entrypoints.serve.dev.rlhf.conftest import (
13 + MODEL_NAME,
14 + gen,
15 + health,
16 + ok,
17 + server,
18 +)
19 +from tests.utils import multi_gpu_test
20 +from vllm.distributed.weight_transfer import (
21 + HTTPVLLMWeightSyncClient,
22 + ModuleSource,
23 + WeightTransferTrainerFactory,
24 +)
25 +from vllm.distributed.weight_transfer.ipc_engine import IPCTrainerInitInfo
26 +
27 +
28 +def weight_checker(url: str, action: str, baseline=None) -> requests.Response:
29 + body = {"action": action}
30 + if baseline is not None:
31 + body["baseline"] = baseline
32 + return requests.post(f"{url}/weight_checker", json=body, timeout=300)
33 +
34 +
35 +def checksums(url: str) -> dict[str, str]:
36 + response = weight_checker(url, "checksum")
37 + response.raise_for_status()
38 + return response.json()["checksums"]
39 +
40 +
41 +def ranks(keys) -> set[str]:
42 + return {key.rsplit(":", 1)[0] for key in keys}
43 +
44 +
45 +def reset_reload_and_compare(url: str) -> None:
46 + baseline = checksums(url)
47 + try:
48 + weight_checker(url, "reset").raise_for_status()
49 + comparison = weight_checker(url, "compare", baseline).json()
50 + assert ranks(comparison["mismatches"]) == ranks(baseline)
51 + finally:
52 + requests.post(
53 + f"{url}/collective_rpc", json={"method": "reload_weights"}, timeout=300
54 + ).raise_for_status()
55 + assert weight_checker(url, "compare", baseline).json() == {
56 + "match": True,
57 + "mismatches": [],
58 + }
59 +
60 +
61 +@pytest.fixture(scope="class")
62 +def server_url():
63 + with server() as url:
64 + yield url
65 +
66 +
67 +class TestSingleGPU:
68 + @pytest.mark.parametrize(
69 + "payload",
70 + [{}, {"action": "frobnicate"}, {"action": "compare"}, [], None],
71 + )
72 + def test_invalid_request_returns_400(self, server_url, payload):
73 + response = requests.post(
74 + f"{server_url}/weight_checker", json=payload, timeout=10
75 + )
76 + assert response.status_code == 400, response.text
77 + assert health(server_url) == 200
78 +
79 + def test_checksum_is_stable_and_compare_reports_differences(self, server_url):
80 + baseline = checksums(server_url)
81 + assert all(
82 + re.fullmatch(r"[0-9a-f]{64}", digest) for digest in baseline.values()
83 + )
84 + assert checksums(server_url) == baseline
85 +
86 + changed, missing, *_ = baseline
87 + altered = {**baseline, changed: "0" * 64}
88 + del altered[missing]
89 + assert weight_checker(server_url, "compare", altered).json() == {
90 + "match": False,
91 + "mismatches": sorted([changed, missing]),
92 + }
93 +
94 + def test_reset_changes_weights_and_reload_restores_them(self, server_url):
95 + reset_reload_and_compare(server_url)
96 + assert ok(gen(server_url))
97 +
98 +
99 +def test_ipc_weight_transfer_restores_reset_weights():
100 + """The documented RL flow: reset, then transfer the checkpoint back."""
101 + args = ["--weight-transfer-config", '{"backend": "ipc"}']
102 + with (
103 + patch.dict(os.environ, {"VLLM_ALLOW_INSECURE_SERIALIZATION": "1"}),
104 + server(extra_args=args, port=8772) as url,
105 + ):
106 + baseline = checksums(url)
107 + weight_checker(url, "reset").raise_for_status()
108 + trainer_model = AutoModelForCausalLM.from_pretrained(
109 + MODEL_NAME, torch_dtype=torch.bfloat16
110 + ).cuda()
111 + WeightTransferTrainerFactory.trainer_init(
112 + IPCTrainerInitInfo(rank=0),
113 + client=HTTPVLLMWeightSyncClient(url),
114 + source=ModuleSource(trainer_model),
115 + ).send_weights()
116 + assert weight_checker(url, "compare", baseline).json() == {
117 + "match": True,
118 + "mismatches": [],
119 + }
120 +
121 +
122 +@multi_gpu_test(num_gpus=2)
123 +@pytest.mark.parametrize(
124 + "args, expected",
125 + [
126 + (["--data-parallel-size", "2"], {"dp0:pp0:pcp0:tp0", "dp1:pp0:pcp0:tp0"}),
127 + (["--tensor-parallel-size", "2"], {"dp0:pp0:pcp0:tp0", "dp0:pp0:pcp0:tp1"}),
128 + ],
129 +)
130 +def test_checksum_and_reset_cover_every_worker(args, expected):
131 + with server(extra_args=args, port=8771, timeout=600) as url:
132 + assert ranks(checksums(url)) == expected
133 + reset_reload_and_compare(url)
tests/model_executor/model_loader/test_weight_checksum.py new
+93
@@ -0,0 +1,93 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +from types import SimpleNamespace
4 +
5 +import pytest
6 +import torch
7 +import torch.nn as nn
8 +from torch.testing._internal.two_tensor import TwoTensor
9 +
10 +from vllm.config import ParallelConfig
11 +from vllm.model_executor import parameter
12 +from vllm.model_executor.model_loader.weight_checksum import (
13 + compute_tensor_digests,
14 + zero_weights,
15 +)
16 +from vllm.v1.worker import gpu_worker
17 +from vllm.v1.worker.gpu_worker import Worker
18 +
19 +pytestmark = pytest.mark.cpu_test
20 +
21 +
22 +class _Model(nn.Module):
23 + def __init__(self):
24 + super().__init__()
25 + self.linear = nn.Linear(2, 2)
26 + nn.init.ones_(self.linear.weight)
27 + nn.init.ones_(self.linear.bias)
28 + # A strided view, which cannot be reinterpreted as bytes in place.
29 + self.strided = nn.Parameter(torch.arange(1.0, 7.0)[::2])
30 + self.scalar = nn.Parameter(torch.tensor(2.0))
31 + # Loading never restores buffers, so reset must leave them alone.
32 + self.register_buffer("k_scale", torch.tensor(2.0))
33 +
34 +
35 +def test_digests_cover_every_parameter():
36 + assert set(compute_tensor_digests(_Model())) == {
37 + "linear.weight",
38 + "linear.bias",
39 + "strided",
40 + "scalar",
41 + }
42 +
43 +
44 +def test_zero_weights_changes_every_parameter_and_no_buffer():
45 + model = _Model()
46 + before = compute_tensor_digests(model)
47 + zero_weights(model)
48 + after = compute_tensor_digests(model)
49 + assert all(before[name] != after[name] for name in before)
50 + assert model.k_scale.item() == 2.0
51 +
52 +
53 +def test_shared_weight_partitions_are_hashed_and_reset(monkeypatch):
54 + monkeypatch.setattr(parameter, "get_tensor_model_parallel_rank", lambda: 0)
55 + monkeypatch.setattr(parameter, "get_tensor_model_parallel_world_size", lambda: 1)
56 + model = nn.Module()
57 + model.transform = parameter.SharedWeightParameter(weight_loader=None)
58 + model.transform.add_partition(0, object(), 2, 2)
59 + model.transform.partitions[0].data.fill_(3.0)
60 +
61 + before = compute_tensor_digests(model)
62 + zero_weights(model)
63 + assert set(before) == {"transform.0"}
64 + assert compute_tensor_digests(model) != before
65 +
66 +
67 +def test_tensor_subclass_inner_tensors_are_hashed_and_reset():
68 + model = nn.Module()
69 + model.weight = nn.Parameter(
70 + TwoTensor(torch.ones(2, 2), torch.full((2, 2), 2.0)), requires_grad=False
71 + )
72 +
73 + before = compute_tensor_digests(model)
74 + zero_weights(model)
75 + assert set(before) == {"weight.a", "weight.b"}
76 + assert all(
77 + before[name] != digest for name, digest in compute_tensor_digests(model).items()
78 + )
79 +
80 +
81 +def test_dense_dp_replicas_get_distinct_key_prefixes(monkeypatch):
82 + for name in ("get_tp_group", "get_pp_group", "get_pcp_group"):
83 + monkeypatch.setattr(gpu_worker, name, lambda: SimpleNamespace(rank_in_group=0))
84 + prefixes = []
85 + for dp_rank in (0, 1):
86 + worker = object.__new__(Worker)
87 + worker.parallel_config = ParallelConfig(
88 + data_parallel_size=2, data_parallel_rank=dp_rank
89 + )
90 + # Dense DP resets data_parallel_rank to 0 in each engine.
91 + worker.parallel_config.reconfigure_for_independent_dp_rank()
92 + prefixes.append(worker._weight_checksum_prefix())
93 + assert prefixes == ["dp0:pp0:pcp0:tp0:", "dp1:pp0:pcp0:tp0:"]
tests/v1/engine/test_engine_core_client.py
+18
@@ -222,6 +222,24 @@ def _make_dplb_client(num_engines: int = 3, client_count: int = 1) -> DPLBAsyncM
222 return client
223
224
225 +def test_dplb_weight_checksums_gather_every_engine():
226 + """call_utility_async on a DPLB client would keep only the first engine."""
227 + client = _make_dplb_client(num_engines=2)
228 +
229 + async def call_utility(method, *args, engine):
230 + assert (method, args) == ("compute_weight_checksums", ())
231 + return [{f"dp{engine[0]}:tp{tp}:w": "x"} for tp in (0, 1)]
232 +
233 + client._call_utility_async = call_utility
234 + workers = asyncio.run(client.compute_weight_checksums_async())
235 + assert [key for worker in workers for key in worker] == [
236 + "dp0:tp0:w",
237 + "dp0:tp1:w",
238 + "dp1:tp0:w",
239 + "dp1:tp1:w",
240 + ]
241 +
242 +
243 def test_dplb_late_interaction_sticky_routing():
244 client = _make_dplb_client()
245
vllm/engine/protocol.py
+4
@@ -268,6 +268,10 @@ class EngineClient(ABC):
268 """Perform a collective RPC call to the given path."""
269 raise NotImplementedError
270
271 + async def compute_weight_checksums(self) -> list[dict[str, str]]:
272 + """Return rank-qualified weight digests from every worker of every engine."""
273 + raise NotImplementedError
274 +
275 async def handle_fault(
276 self, fault_tolerance_request: FaultToleranceRequest
277 ) -> FaultToleranceResult:
vllm/entrypoints/serve/dev/rlhf/api_router.py
+38 -1
@@ -3,7 +3,7 @@
3
4 import json
5 from http import HTTPStatus
6 -from typing import Annotated, Any
6 +from typing import Annotated, Any, Literal
7
8 from fastapi import APIRouter, Body, FastAPI, HTTPException, Query, Request
9 from fastapi.responses import JSONResponse
@@ -216,6 +216,43 @@ async def weight_info(raw_request: Request):
216 return JSONResponse(content={"weight_version": weight_version})
217
218
219 +@router.post("/weight_checker")
220 +async def weight_checker(
221 + raw_request: Request,
222 + action: Annotated[Literal["checksum", "reset", "compare"], Body(embed=True)],
223 + baseline: Annotated[dict[str, str] | None, Body(embed=True)] = None,
224 +) -> JSONResponse:
225 + """Checksum, reset, or compare model weights against a baseline."""
226 + client = engine_client(raw_request)
227 + if action == "reset":
228 + await client.collective_rpc("reset_weights")
229 + return JSONResponse(content={"status": "reset"})
230 + if action == "compare" and baseline is None:
231 + raise HTTPException(
232 + status_code=HTTPStatus.BAD_REQUEST.value,
233 + detail="action='compare' requires a 'baseline' object",
234 + )
235 +
236 + checksums: dict[str, str] = {}
237 + for worker_checksums in await client.compute_weight_checksums():
238 + if duplicates := checksums.keys() & worker_checksums.keys():
239 + raise HTTPException(
240 + status_code=HTTPStatus.INTERNAL_SERVER_ERROR.value,
241 + detail=f"Duplicate weight checksum keys: {sorted(duplicates)}",
242 + )
243 + checksums.update(worker_checksums)
244 + if action == "checksum":
245 + return JSONResponse(content={"checksums": checksums})
246 +
247 + assert baseline is not None
248 + mismatches = sorted(
249 + key
250 + for key in checksums.keys() | baseline.keys()
251 + if checksums.get(key) != baseline.get(key)
252 + )
253 + return JSONResponse(content={"match": not mismatches, "mismatches": mismatches})
254 +
255 +
256 @router.get("/get_world_size")
257 async def get_world_size(
258 raw_request: Request,
vllm/model_executor/model_loader/weight_checksum.py new
+43
@@ -0,0 +1,43 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""Per-parameter weight digests and reset, for verifying weight updates."""
4 +
5 +import hashlib
6 +from collections.abc import Iterator
7 +
8 +import torch
9 +import torch.nn as nn
10 +from torch.utils._python_dispatch import is_traceable_wrapper_subclass
11 +
12 +from vllm.model_executor.parameter import SharedWeightParameter
13 +
14 +
15 +def _tensors(name: str, tensor: torch.Tensor) -> Iterator[tuple[str, torch.Tensor]]:
16 + # Both keep their bytes in inner tensors (e.g. TorchAO's qdata and scale).
17 + if isinstance(tensor, SharedWeightParameter):
18 + for index, partition in tensor.partitions.items():
19 + yield from _tensors(f"{name}.{index}", partition)
20 + elif is_traceable_wrapper_subclass(tensor):
21 + for attr in tensor.__tensor_flatten__()[0]:
22 + yield from _tensors(f"{name}.{attr}", getattr(tensor, attr))
23 + else:
24 + yield name, tensor
25 +
26 +
27 +def _weights(model: nn.Module) -> Iterator[tuple[str, torch.Tensor]]:
28 + for name, param in model.named_parameters():
29 + yield from _tensors(name, param)
30 +
31 +
32 +def compute_tensor_digests(model: nn.Module) -> dict[str, str]:
33 + return {
34 + name: hashlib.sha256(
35 + weight.detach().cpu().contiguous().view(-1).view(torch.uint8).numpy()
36 + ).hexdigest()
37 + for name, weight in _weights(model)
38 + }
39 +
40 +
41 +def zero_weights(model: nn.Module) -> None:
42 + for _, weight in _weights(model):
43 + weight.data.zero_()
vllm/v1/engine/async_llm.py
+3
@@ -1125,6 +1125,9 @@ class AsyncLLM(EngineClient):
1125 async def is_sleeping(self) -> bool:
1126 return await self.engine_core.is_sleeping_async()
1127
1128 + async def compute_weight_checksums(self) -> list[dict[str, str]]:
1129 + return await self.engine_core.compute_weight_checksums_async()
1130 +
1131 async def add_lora(self, lora_request: LoRARequest) -> bool:
1132 """Load a new LoRA adapter into the engine for future requests."""
1133 return await self.engine_core.add_lora_async(lora_request)
vllm/v1/engine/core.py
+3
@@ -1005,6 +1005,9 @@ class EngineCore:
1005 def execute_dummy_batch(self):
1006 self.model_executor.execute_dummy_batch()
1007
1008 + def compute_weight_checksums(self) -> list[dict[str, str]]:
1009 + return self.collective_rpc("compute_weight_checksums")
1010 +
1011 def add_lora(self, lora_request: LoRARequest) -> bool:
1012 return self.model_executor.add_lora(lora_request)
1013
vllm/v1/engine/core_client.py
+4
@@ -1362,6 +1362,10 @@ class AsyncMPClient(MPClient):
1362 "collective_rpc", method, timeout, args, kwargs
1363 )
1364
1365 + async def compute_weight_checksums_async(self) -> list[dict[str, str]]:
1366 + per_engine = await self.call_utility_all_async("compute_weight_checksums")
1367 + return [worker for workers in per_engine for worker in workers]
1368 +
1369 async def handle_fault(
1370 self, ft_request: FaultToleranceRequest
1371 ) -> FaultToleranceResult:
vllm/v1/worker/gpu_worker.py
+20
@@ -58,6 +58,10 @@ from vllm.distributed.weight_transfer import (
58 from vllm.logger import init_logger
59 from vllm.lora.request import LoRARequest
60 from vllm.model_executor.model_loader import get_model_loader
61 +from vllm.model_executor.model_loader.weight_checksum import (
62 + compute_tensor_digests,
63 + zero_weights,
64 +)
65 from vllm.model_executor.warmup.kernel_warmup import kernel_warmup
66 from vllm.multimodal.gpu_ipc_memory import reserve_mm_ipc_gpu_memory
67 from vllm.platforms import current_platform
@@ -342,6 +346,22 @@ class Worker(WorkerBase):
346 def checkpoint_restore(self) -> None:
347 checkpoint_restore_distributed_state()
348
349 + def _weight_checksum_prefix(self) -> str:
350 + # data_parallel_rank is 0 in dense workers; only the index is unique.
351 + dp = self.parallel_config.data_parallel_index
352 + pp = get_pp_group().rank_in_group
353 + pcp = get_pcp_group().rank_in_group
354 + tp = get_tp_group().rank_in_group
355 + return f"dp{dp}:pp{pp}:pcp{pcp}:tp{tp}:"
356 +
357 + def compute_weight_checksums(self) -> dict[str, str]:
358 + prefix = self._weight_checksum_prefix()
359 + digests = compute_tensor_digests(self.model_runner.get_model())
360 + return {prefix + name: digest for name, digest in digests.items()}
361 +
362 + def reset_weights(self) -> None:
363 + zero_weights(self.model_runner.get_model())
364 +
365 def _maybe_get_memory_pool_context(self, tag: str) -> AbstractContextManager:
366 if (
367 current_platform.is_cuda_alike()