9069a57139 · tduncan/Nvidia-Optimization-Stack-VLLM

docs/serving/offline_inference.md

+2

@@ -65,6 +65,8 @@ For further details on Weight Transfer, please refer to [this page](../training/

65 - `LLM.start_weight_update` - Starts a new weight update cycle.

66 - `LLM.update_weights` - Updates the model weights.

67 - `LLM.finish_weight_update` - Finishes the current weight update cycle.

68 +- `LLM.update_weight_version` - Sets the weight version without updating model weights.

69 +- `LLM.get_weight_version` - Returns the latest committed weight version.

70

71 ## Additional APIs

72

docs/serving/online_serving/README.md

+2

@@ -179,6 +179,8 @@ For further details on Weight Transfer, please refer to [this page](../../traini

179 - `/start_weight_update` - Prepares the inference engine for a weight update.

180 - `/update_weights` - Update model weights (can alter model behavior)

181 - `/finish_weight_update` - Finalizes the weight update

182 +- `/update_weight_version` - Set the weight version without updating model weights

183 +- `/weight_info` - Get the latest committed weight version

184 - `/get_world_size` - Get distributed world size

185

186 ### Collective RPC

docs/training/async_rl.md

+2 -1

@@ -38,11 +38,12 @@ Resumes the scheduler after a pause. Any requests frozen with `mode="keep"` will

38

39 ### HTTP Endpoints

40

41 -When using the vLLM HTTP server, the same functionality is available via:

41 +With `VLLM_SERVER_DEV_MODE=1`, the vLLM HTTP server exposes the same functionality via:

42

43 - `POST /pause?mode=keep` - Pause generation

44 - `POST /resume` - Resume generation

45 - `POST /abort_requests` - Abort in-flight requests without pausing the scheduler (send `{}` to abort all, or `{"request_ids": [...]}`)

46 +- `GET /weight_info` - Return the latest committed `weight_version`

47

48 !!! note "Data Parallelism"

49 When using data parallelism with vLLM's **internal load balancer** (i.e. `data_parallel_backend="ray"`), pause and resume are handled automatically across all DP ranks -- a single call is sufficient. When using an **external load balancer** (i.e. multiple independent vLLM instances behind a proxy), you must send pause and resume requests to **every** engine instance individually before and after the weight update.

docs/training/weight_transfer/README.md

+4 -2

@@ -53,7 +53,9 @@ When running vLLM as an HTTP server, the following endpoints are available for w

53 | `/init_weight_transfer_engine` | POST | Initialize the weight transfer engine with backend-specific info |

54 | `/start_weight_update` | POST | Start a weight update |

55 | `/update_weights` | POST | Transfer a batch of weights with backend-specific metadata |

56 -| `/finish_weight_update` | POST | Finish the weight update and run post-processing |

56 +| `/finish_weight_update` | POST | Finish the update and optionally commit its `weight_version` |

57 +| `/update_weight_version` | POST | Update `weight_version` without changing model weights |

58 +| `/weight_info` | GET | Get the latest committed weight version |

59 | `/pause` | POST | Pause generation before weight sync to handle inflight requests |

60 | `/resume` | POST | Resume generation after weight sync |

61 | `/get_world_size` | GET | Get the number of inference workers (useful for NCCL world size calculation) |

@@ -79,7 +81,7 @@ EngineClass.trainer_send_weights(

81 )

82

83 # 4. Finish weight update on inference side

82 -llm.finish_weight_update()

84 +llm.finish_weight_update(weight_version="step-42")

85 ```

86

87 See the [NCCL](nccl.md) and [IPC](ipc.md) pages for backend-specific trainer APIs and full examples.

tests/distributed/test_weight_transfer.py

+8 -1

@@ -1247,7 +1247,7 @@ class RecordingClient:

1247 self.order.append("update")

1248 self.last_update_info = update_info

1249

1250 - def finish_weight_update(self) -> None:

1250 + def finish_weight_update(self, weight_version: str | None = None) -> None:

1251 self.order.append("finish")

1252

1253

@@ -1303,6 +1303,10 @@ class TestTrainerClients:

1303 assert isinstance(update_req, WeightTransferUpdateRequest)

1304 assert update_req.update_info == {"names": ["w"]}

1305

1306 + client.finish_weight_update("step-42")

1307 + handle.finish_weight_update.remote.assert_called_once_with()

1308 + handle.update_weight_version.remote.assert_called_once_with("step-42")

1309 +

1310 def test_http_client_pickles_ipc_handles_for_json(self, monkeypatch):

1311 """HTTP update_weights must encode raw ipc_handles as a base64 pickle."""

1312 captured = {}

@@ -1334,6 +1338,9 @@ class TestTrainerClients:

1338 client.update_weights(update_info)

1339 assert captured["json"]["update_info"] == update_info

1340

1341 + client.finish_weight_update("step-42")

1342 + assert captured["json"] == {"weight_version": "step-42"}

1343 +

1344

1345 class TestModuleSource:

1346 """`ModuleSource` metadata vs. materialized iteration (dense, no GPU)."""

tests/entrypoints/openai/test_openai_schema.py

+1

@@ -148,6 +148,7 @@ def test_openapi_stateless(case: schemathesis.Case):

148 "/start_draft_weight_update",

149 "/update_weights",

150 "/finish_weight_update",

151 + "/update_weight_version",

152 ):

153 return

154

tests/entrypoints/weight_transfer/test_weight_transfer_llm.py

+11 -1

@@ -234,6 +234,7 @@ def test_update_weights_calls_engine():

234 assert shapes == test_shapes

235

236 llm.finish_weight_update()

237 + assert llm.get_weight_version() == "default"

238

239

240 @create_new_process_for_each_test()

@@ -259,6 +260,8 @@ def test_full_weight_transfer_flow():

260 weight_transfer_config=WeightTransferConfig(backend="nccl"),

261 )

262

263 + assert llm.get_weight_version() == "default"

264 +

265 # Step 1: Initialize weight transfer engine

266 llm.init_weight_transfer_engine(

267 WeightTransferInitRequest(init_info={"test_param": "flow_test"})

@@ -278,8 +281,15 @@ def test_full_weight_transfer_flow():

281 )

282 )

283

284 + assert llm.get_weight_version() == "default"

285 +

286 # Step 4: Finish weight update

282 - llm.finish_weight_update()

287 + llm.finish_weight_update("step-42")

288 +

289 + assert llm.get_weight_version() == "step-42"

290 +

291 + llm.update_weight_version("manual-version")

292 + assert llm.get_weight_version() == "manual-version"

293

294 # Verify the full flow completed

295 def check_flow(self):

vllm/distributed/weight_transfer/base.py

+1 -1

@@ -370,7 +370,7 @@ class VLLMWeightSyncClient(Protocol):

370

371 def update_weights(self, update_info: dict[str, Any]) -> None: ...

372

373 - def finish_weight_update(self) -> None: ...

373 + def finish_weight_update(self, weight_version: str | None = None) -> None: ...

374

375

376 class TrainerWeightTransferEngine(ABC, Generic[TConfig, TInitInfo]):

vllm/distributed/weight_transfer/clients.py

+10 -3

@@ -77,8 +77,11 @@ class HTTPVLLMWeightSyncClient:

77 "update_weights", {"update_info": _json_safe_update_info(update_info)}

78 )

79

80 - def finish_weight_update(self) -> None:

81 - self._post("finish_weight_update")

80 + def finish_weight_update(self, weight_version: str | None = None) -> None:

81 + json = (

82 + {"weight_version": weight_version} if weight_version is not None else None

83 + )

84 + self._post("finish_weight_update", json)

85

86

87 class RayVLLMWeightSyncClient:

@@ -108,7 +111,11 @@ class RayVLLMWeightSyncClient:

111 request = WeightTransferUpdateRequest(update_info=update_info)

112 ray.get([h.update_weights.remote(request) for h in self.handles])

113

111 - def finish_weight_update(self) -> None:

114 + def finish_weight_update(self, weight_version: str | None = None) -> None:

115 import ray

116

117 ray.get([h.finish_weight_update.remote() for h in self.handles])

118 + if weight_version is not None:

119 + ray.get(

120 + [h.update_weight_version.remote(weight_version) for h in self.handles]

121 + )

vllm/engine/protocol.py

+10 -2

@@ -267,6 +267,14 @@ class EngineClient(ABC):

267 """Batched weight update for RL training."""

268 raise NotImplementedError

269

270 - async def finish_weight_update(self) -> None:

271 - """Finish the current weight update."""

270 + async def finish_weight_update(self, weight_version: str | None = None) -> None:

271 + """Finish the weight update and set its version if provided."""

272 + raise NotImplementedError

273 +

274 + async def update_weight_version(self, new_version: str) -> None:

275 + """Set the weight version without updating weights."""

276 + raise NotImplementedError

277 +

278 + async def get_weight_version(self) -> str:

279 + """Return the latest committed weight version."""

280 raise NotImplementedError

vllm/entrypoints/llm.py

+12 -2

@@ -885,9 +885,19 @@ class LLM(BeamSearchOfflineMixin, PoolingOfflineMixin, OfflineInferenceMixin):

885 "update_weights", kwargs={"update_info": update_info_dict}

886 )

887

888 - def finish_weight_update(self) -> None:

889 - """Finish the current weight update."""

888 + def finish_weight_update(self, weight_version: str | None = None) -> None:

889 + """Finish the weight update and set its version if provided."""

890 self.llm_engine.collective_rpc("finish_weight_update")

891 + if weight_version is not None:

892 + self.llm_engine.set_weight_version(weight_version)

893 +

894 + def update_weight_version(self, new_version: str) -> None:

895 + """Set the weight version without updating weights."""

896 + self.llm_engine.set_weight_version(new_version)

897 +

898 + def get_weight_version(self) -> str:

899 + """Return the latest committed weight version."""

900 + return self.llm_engine.get_weight_version()

901

902 def __repr__(self) -> str:

903 """Return a transformers-style hierarchical view of the model."""

vllm/entrypoints/serve/dev/rlhf/api_router.py

+21 -3

@@ -5,7 +5,7 @@ import json

5 from http import HTTPStatus

6 from typing import Annotated

7

8 -from fastapi import APIRouter, FastAPI, HTTPException, Query, Request

8 +from fastapi import APIRouter, Body, FastAPI, HTTPException, Query, Request

9 from fastapi.responses import JSONResponse

10

11 from vllm.distributed.weight_transfer.base import (

@@ -203,11 +203,29 @@ async def update_weights(raw_request: Request):

203

204

205 @router.post("/finish_weight_update")

206 -async def finish_weight_update(raw_request: Request):

207 - await engine_client(raw_request).finish_weight_update()

206 +async def finish_weight_update(

207 + raw_request: Request,

208 + weight_version: Annotated[str | None, Body(embed=True)] = None,

209 +):

210 + await engine_client(raw_request).finish_weight_update(weight_version)

211 return JSONResponse(content={"message": "Weight update finished"})

212

213

214 +@router.post("/update_weight_version")

215 +async def update_weight_version(

216 + raw_request: Request,

217 + new_version: Annotated[str, Body(embed=True)],

218 +):

219 + await engine_client(raw_request).update_weight_version(new_version)

220 + return JSONResponse(content={"success": True, "new_version": new_version})

221 +

222 +

223 +@router.get("/weight_info")

224 +async def weight_info(raw_request: Request):

225 + weight_version = await engine_client(raw_request).get_weight_version()

226 + return JSONResponse(content={"weight_version": weight_version})

227 +

228 +

229 @router.get("/get_world_size")

230 async def get_world_size(

231 raw_request: Request,

vllm/v1/engine/async_llm.py

+12 -2

@@ -1106,6 +1106,16 @@ class AsyncLLM(EngineClient):

1106 "update_weights", kwargs={"update_info": request.update_info}

1107 )

1108

1109 - async def finish_weight_update(self) -> None:

1110 - """Finish the current weight update."""

1109 + async def finish_weight_update(self, weight_version: str | None = None) -> None:

1110 + """Finish the weight update and set its version if provided."""

1111 await self.collective_rpc("finish_weight_update")

1112 + if weight_version is not None:

1113 + await self.update_weight_version(weight_version)

1114 +

1115 + async def update_weight_version(self, new_version: str) -> None:

1116 + """Set the weight version without updating weights."""

1117 + await self.engine_core.set_weight_version_async(new_version)

1118 +

1119 + async def get_weight_version(self) -> str:

1120 + """Return the latest committed weight version."""

1121 + return await self.engine_core.get_weight_version_async()

vllm/v1/engine/core.py

+9

@@ -125,6 +125,8 @@ class EngineCore:

125 )

126

127 self.log_stats = log_stats

128 + # Opaque weight version supplied by the caller.

129 + self._weight_version = "default"

130

131 # Setup Model.

132 self.model_executor = executor_class(vllm_config)

@@ -956,6 +958,13 @@ class EngineCore:

958 ) -> list[_R]:

959 return self.model_executor.collective_rpc(method, timeout, args, kwargs)

960

961 + def set_weight_version(self, weight_version: str) -> None:

962 + self._weight_version = weight_version

963 +

964 + def get_weight_version(self) -> str:

965 + """Return the latest committed weight version."""

966 + return self._weight_version

967 +

968 def preprocess_add_request(self, request: EngineCoreRequest) -> tuple[Request, int]:

969 """Preprocess the request.

970

vllm/v1/engine/core_client.py

+30

@@ -176,9 +176,21 @@ class EngineCoreClient(ABC):

176 def execute_dummy_batch(self) -> None:

177 raise NotImplementedError

178

179 + def set_weight_version(self, weight_version: str) -> None:

180 + raise NotImplementedError

181 +

182 + def get_weight_version(self) -> str:

183 + raise NotImplementedError

184 +

185 async def execute_dummy_batch_async(self) -> None:

186 raise NotImplementedError

187

188 + async def set_weight_version_async(self, weight_version: str) -> None:

189 + raise NotImplementedError

190 +

191 + async def get_weight_version_async(self) -> str:

192 + raise NotImplementedError

193 +

194 def abort_requests(self, request_ids: list[str]) -> None:

195 raise NotImplementedError

196

@@ -351,6 +363,12 @@ class InprocClient(EngineCoreClient):

363 def execute_dummy_batch(self) -> None:

364 self.engine_core.execute_dummy_batch()

365

366 + def set_weight_version(self, weight_version: str) -> None:

367 + self.engine_core.set_weight_version(weight_version)

368 +

369 + def get_weight_version(self) -> str:

370 + return self.engine_core.get_weight_version()

371 +

372 def add_lora(self, lora_request: LoRARequest) -> bool:

373 return self.engine_core.add_lora(lora_request)

374

@@ -947,6 +965,12 @@ class SyncMPClient(MPClient):

965 def execute_dummy_batch(self) -> None:

966 self.call_utility("execute_dummy_batch")

967

968 + def set_weight_version(self, weight_version: str) -> None:

969 + self.call_utility("set_weight_version", weight_version)

970 +

971 + def get_weight_version(self) -> str:

972 + return self.call_utility("get_weight_version")

973 +

974 def collective_rpc(

975 self,

976 method: str | Callable[..., _R],

@@ -1199,6 +1223,12 @@ class AsyncMPClient(MPClient):

1223 async def execute_dummy_batch_async(self) -> None:

1224 await self.call_utility_async("execute_dummy_batch")

1225

1226 + async def set_weight_version_async(self, weight_version: str) -> None:

1227 + await self.call_utility_async("set_weight_version", weight_version)

1228 +

1229 + async def get_weight_version_async(self) -> str:

1230 + return await self.call_utility_async("get_weight_version")

1231 +

1232 async def add_lora_async(self, lora_request: LoRARequest) -> bool:

1233 return await self.call_utility_async("add_lora", lora_request)

1234

vllm/v1/engine/llm_engine.py

+7

@@ -425,6 +425,13 @@ class LLMEngine:

425 ) -> list[_R]:

426 return self.engine_core.collective_rpc(method, timeout, args, kwargs)

427

428 + def set_weight_version(self, weight_version: str) -> None:

429 + self.engine_core.set_weight_version(weight_version)

430 +

431 + def get_weight_version(self) -> str:

432 + """Return the latest committed weight version."""

433 + return self.engine_core.get_weight_version()

434 +

435 def apply_model(self, func: Callable[[nn.Module], _R]) -> list[_R]:

436 return self.collective_rpc("apply_model", args=(func,))

437

[Core][Frontend] Add weight version tagging for RL rollouts (#49040)