@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / e3b65a5ba0

Fork 0

A high-throughput and memory-efficient inference and serving engine for LLMs

Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git

Synced Nov 4, 02:30 UTC

[feat] Add explicit `/start_weight_update` and `/finish_weight_update` APIs for weight transfer (#39212)

Sumanth R Hegde committed May 8, 2026 at 18:03 UTC e3b65a5ba069b350120ca7a614a010787d2de867

19 files changed +288 -54

docs/training/weight_transfer/README.md

+14 -4

@@ -4,10 +4,12 @@ vLLM provides a pluggable weight transfer system for synchronizing model weights

5 ## Architecture

7 -The weight transfer system follows a **two-phase protocol** with a pluggable backend design:

7 +The weight transfer system follows a **four-phase protocol** with a pluggable backend design:

9 1. **Initialization** (`init_weight_transfer_engine`): Establishes the communication channel between the trainer and inference workers. Called once before the training loop begins.

10 -2. **Weight Update** (`update_weights`): Transfers updated weights from the trainer to the inference engine. Called after each training step (or batch of steps).

10 +2. **Start** (`start_weight_update`): Prepares the inference engine for a weight update.

11 +3. **Weight Update** (`update_weights`): Transfers updated weights from the trainer to the inference engine. May be called one or more times (e.g., for chunked transfers).

12 +4. **Finish** (`finish_weight_update`): Finalizes the weight update (e.g., runs post-processing for checkpoint-format weights). Called once after all weights have been transferred.

14 ## Available Backends

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

50 | Endpoint | Method | Description |

51 | -------- | ------ | ----------- |

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

51 -| `/update_weights` | POST | Trigger a weight update with backend-specific metadata |

53 +| `/start_weight_update` | POST | Start a weight update |

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

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

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

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

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

@@ -64,11 +68,17 @@ Both backends provide static methods that the trainer calls to send weights. The

68 # 1. Initialize the transfer engine (backend-specific)

69 EngineClass.trainer_init(init_info)

67 -# 2. Send weights to inference workers

71 +# 2. Start weight update on inference side

72 +llm.start_weight_update(is_checkpoint_format=True)

73 +

74 +# 3. Send weights to inference workers

75 EngineClass.trainer_send_weights(

76 iterator=model.named_parameters(),

77 trainer_args=backend_specific_args,

78 )

79 +

80 +# 4. Finish weight update on inference side

81 +llm.finish_weight_update()

82 ```

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

docs/training/weight_transfer/base.md

+2 -4

@@ -43,16 +43,14 @@ update_request = WeightTransferUpdateRequest(

44 ### WeightTransferUpdateInfo

46 -The base `WeightTransferUpdateInfo` includes an `is_checkpoint_format` flag:

46 +The base `WeightTransferUpdateInfo` is a marker class for backend-specific update info:

48 ```python

49 @dataclass

50 class WeightTransferUpdateInfo(ABC):

51 - is_checkpoint_format: bool = True

51 + pass

52 ```

54 -When `is_checkpoint_format=True` (the default), vLLM applies layerwise weight processing (repacking, renaming, etc.) on the received weights before loading them. Set to `False` if the trainer has already converted weights to the kernel format expected by the model.

55 -

54 ## Implementing a Custom Engine

56 To create a custom weight transfer backend:

docs/training/weight_transfer/ipc.md

+16 -2

@@ -38,11 +38,15 @@ trainer_args = IPCTrainerSendWeightsArgs(

38 mode="ray",

39 llm_handle=llm_actor_handle,

40 )

41 -

41 +# start

42 +ray.get(llm_actor_handle.start_weight_update.remote(is_checkpoint_format=True))

43 +# send weights

44 IPCWeightTransferEngine.trainer_send_weights(

45 iterator=model.named_parameters(),

46 trainer_args=trainer_args,

47 )

48 +# finish

49 +ray.get(llm_actor_handle.finish_weight_update.remote())

50 ```

52 In Ray mode, the engine calls `llm_handle.update_weights.remote(...)` directly, passing the IPC handles via Ray's serialization.

@@ -57,13 +61,23 @@ trainer_args = IPCTrainerSendWeightsArgs(

61 url="http://localhost:8000",

62 )

64 +# start

65 +base_url = "http://localhost:8000"

66 +url = f"{base_url}/start_weight_update"

67 +response = requests.post(url, json={"is_checkpoint_format": True}, timeout=60)

68 +response.raise_for_status()

69 +# send weights

70 IPCWeightTransferEngine.trainer_send_weights(

71 iterator=model.named_parameters(),

72 trainer_args=trainer_args,

73 )

74 +# finish

75 +url = f"{base_url}/finish_weight_update"

76 +response = requests.post(url, json={}, timeout=60)

77 +response.raise_for_status()

78 ```

66 -In HTTP mode, IPC handles are pickled, base64-encoded, and sent as JSON to the `/update_weights` endpoint.

80 +In HTTP mode, IPC handles are pickled, base64-encoded, and sent as JSON to the `/update_weights` endpoint. As with Ray mode, you must call `start_weight_update` before and `finish_weight_update` after.

82 See [`IPCTrainerSendWeightsArgs`](https://github.com/vllm-project/vllm/blob/main/vllm/distributed/weight_transfer/ipc_engine.py) for the full list of configurable fields.

docs/training/weight_transfer/nccl.md

+10 -1

@@ -84,11 +84,15 @@ Both the trainer (`NCCLTrainerSendWeightsArgs`) and inference side (`NCCLWeightT

85 ## Receiving Weights (Inference Side)

87 -The inference side triggers weight reception by calling `update_weights`:

87 +The inference side triggers weight reception using the four-phase protocol — `init_weight_transfer_engine`, `start_weight_update`, `update_weights`, `finish_weight_update`. The init phase is shown [above](#initialization). The remaining three steps are:

89 ```python

90 from vllm.distributed.weight_transfer.base import WeightTransferUpdateRequest

92 +# 1. Start the weight update

93 +llm.start_weight_update(is_checkpoint_format=True)

94 +

95 +# 2. Receive weights (can be called multiple times for chunked transfers)

96 llm.update_weights(

97 WeightTransferUpdateRequest(

98 update_info=dict(

@@ -99,10 +103,15 @@ llm.update_weights(

103 )

104 )

105 )

106 +

107 +# 3. Finish the weight update

108 +llm.finish_weight_update()

109 ```

110

111 The `names`, `dtype_names`, and `shapes` lists describe each parameter. These must match the order in which the trainer iterates over its parameters.

112

113 +`start_weight_update` must be called before `update_weights`, and `finish_weight_update` must be called after all weight chunks have been transferred. The `is_checkpoint_format` flag controls whether layerwise reload processing is applied (`True` for checkpoint-format weights, `False` for pre-processed kernel-format weights).

114 +

115 ## Examples

116

117 - [RLHF with NCCL weight syncing (offline, Ray)](../../examples/rl/rlhf_nccl.md) - Trainer on one GPU, 2x tensor-parallel vLLM engine on two others, with packed NCCL weight broadcast

examples/rl/rlhf_async_new_apis.py

@@ -310,6 +310,8 @@ gen_futures = [

310

311 ray.get(llm.pause_after_n_tokens.remote())

312

313 +ray.get(llm.start_weight_update.remote(is_checkpoint_format=True))

314 +

315 inference_handle = llm.update_weights.remote(

316 WeightTransferUpdateRequest(

317 update_info=asdict(

@@ -325,6 +327,8 @@ inference_handle = llm.update_weights.remote(

327 train_handle = train_model.broadcast_weights.remote(packed=True)

328 ray.get([train_handle, inference_handle])

329

330 +ray.get(llm.finish_weight_update.remote())

331 +

332 ray.get(llm.resume_generation.remote())

333 results = ray.get(gen_futures)

334

examples/rl/rlhf_http_ipc.py

+23 -1

@@ -80,6 +80,24 @@ def init_weight_transfer_engine(base_url: str) -> None:

80 response.raise_for_status()

83 +def start_weight_update(

84 + base_url: str,

85 + is_checkpoint_format: bool = True,

86 +) -> None:

87 + """Start a weight update via HTTP endpoint."""

88 + url = f"{base_url}/start_weight_update"

89 + payload = {"is_checkpoint_format": is_checkpoint_format}

90 + response = requests.post(url, json=payload, timeout=60)

91 + response.raise_for_status()

92 +

93 +

94 +def finish_weight_update(base_url: str) -> None:

95 + """Finish a weight update via HTTP endpoint."""

96 + url = f"{base_url}/finish_weight_update"

97 + response = requests.post(url, json={}, timeout=60)

98 + response.raise_for_status()

99 +

100 +

101 def pause_generation(base_url: str) -> None:

102 """Pause generation via HTTP endpoint."""

103 url = f"{base_url}/pause"

@@ -151,7 +169,9 @@ def main():

169 # Pause generation before weight sync

170 pause_generation(BASE_URL)

171

154 - # Broadcast weights via IPC handles using HTTP mode

172 + # Start weight update, broadcast via IPC, then finish

173 + start_weight_update(BASE_URL, is_checkpoint_format=False)

174 +

175 print("Broadcasting weights via CUDA IPC (HTTP)...")

176 trainer_args = IPCTrainerSendWeightsArgs(mode="http", url=BASE_URL)

177 IPCWeightTransferEngine.trainer_send_weights(

@@ -159,6 +179,8 @@ def main():

179 trainer_args=trainer_args,

180 )

181

182 + finish_weight_update(BASE_URL)

183 +

184 # Resume generation after weight sync

185 resume_generation(BASE_URL)

186

examples/rl/rlhf_http_nccl.py

+24

@@ -83,6 +83,17 @@ def init_weight_transfer_engine(

83 response.raise_for_status()

86 +def start_weight_update(

87 + base_url: str,

88 + is_checkpoint_format: bool = True,

89 +) -> None:

90 + """Start a weight update via HTTP endpoint."""

91 + url = f"{base_url}/start_weight_update"

92 + payload = {"is_checkpoint_format": is_checkpoint_format}

93 + response = requests.post(url, json=payload, timeout=60)

94 + response.raise_for_status()

95 +

96 +

97 def update_weights(

98 base_url: str,

99 names: list[str],

@@ -104,6 +115,13 @@ def update_weights(

115 response.raise_for_status()

116

117

118 +def finish_weight_update(base_url: str) -> None:

119 + """Finish a weight update via HTTP endpoint."""

120 + url = f"{base_url}/finish_weight_update"

121 + response = requests.post(url, json={}, timeout=60)

122 + response.raise_for_status()

123 +

124 +

125 def pause_generation(base_url: str) -> None:

126 """Pause generation via HTTP endpoint."""

127 url = f"{base_url}/pause"

@@ -204,6 +222,9 @@ def main():

222 dtype_names.append(str(p.dtype).split(".")[-1])

223 shapes.append(list(p.shape))

224

225 + # Start weight update

226 + start_weight_update(BASE_URL, is_checkpoint_format=True)

227 +

228 # Start the update_weights call in a separate thread since it will block

229 # waiting for NCCL broadcasts

230 # packed=True enables efficient batched tensor broadcasting

@@ -227,6 +248,9 @@ def main():

248 # Wait for update_weights to complete

249 update_thread.join()

250

251 + # Finish weight update

252 + finish_weight_update(BASE_URL)

253 +

254 # Resume generation after weight sync

255 resume_generation(BASE_URL)

256

examples/rl/rlhf_ipc.py

+3 -1

@@ -134,8 +134,10 @@ for output in outputs:

134 ray.get(llm.sleep.remote(level=0))

135

136 ray.get(train_model.init_weight_transfer.remote())

137 -# Synchronize the updated weights to the inference engine using batched API.

137 +# Start weight update, sync weights, then finish

138 +ray.get(llm.start_weight_update.remote(is_checkpoint_format=True))

139 ray.get(train_model.broadcast_weights.remote(llm))

140 +ray.get(llm.finish_weight_update.remote())

141

142 ray.get(llm.wake_up.remote(tags=["scheduling"]))

143

examples/rl/rlhf_nccl.py

@@ -186,6 +186,9 @@ ray.get([train_handle, inference_handle])

186 # Collect all weight metadata from the training actor

187 names, dtype_names, shapes = ray.get(train_model.get_weight_metadata.remote())

188

189 +# Start weight update

190 +ray.get(llm.start_weight_update.remote(is_checkpoint_format=True))

191 +

192 # Issue update_weights call with NCCL-specific update info

193 # packed=True enables efficient batched tensor broadcasting

194 inference_handle = llm.update_weights.remote(

@@ -203,6 +206,9 @@ inference_handle = llm.update_weights.remote(

206 train_handle = train_model.broadcast_weights.remote(packed=True)

207 ray.get([train_handle, inference_handle])

208

209 +# Finish weight update

210 +ray.get(llm.finish_weight_update.remote())

211 +

212 ray.get(llm.wake_up.remote(tags=["scheduling"]))

213

214 # Generate text with the updated model. The output is expected to be normal

examples/rl/rlhf_nccl_fsdp_ep.py

@@ -298,6 +298,9 @@ async def main():

298 names, dtype_names, shapes = ray.get(fsdp_workers[0].get_weight_metadata.remote())

299 print(f"[sync] Got metadata for {len(names)} parameters.")

300

301 + print("[sync] Starting weight update...")

302 + await engine.start_weight_update(is_checkpoint_format=True)

303 +

304 print("[sync] Broadcasting weights from FSDP → vLLM...")

305 broadcast_handles = [

306 w.gather_and_broadcast_weights.remote(packed=True) for w in fsdp_workers

@@ -315,6 +318,8 @@ async def main():

318 )

319 )

320 ray.get(broadcast_handles)

321 +

322 + await engine.finish_weight_update()

323 print("[sync] Weight broadcast complete.")

324

325 print("[sync] Resuming generation...")

tests/entrypoints/openai/test_openai_schema.py

@@ -133,7 +133,9 @@ def test_openapi_stateless(case: Case):

133 # (weight_transfer_config) and are meant to be stateful.

134 if case.operation.path in (

135 "/init_weight_transfer_engine",

136 + "/start_weight_update",

137 "/update_weights",

138 + "/finish_weight_update",

139 ):

140 return

141

tests/entrypoints/weight_transfer/test_weight_transfer_llm.py

+16 -3

@@ -199,6 +199,9 @@ def test_update_weights_calls_engine():

199 WeightTransferInitRequest(init_info={"test_param": "init"})

200 )

201

202 + # Start weight update (required before update_weights)

203 + llm.start_weight_update(is_checkpoint_format=True)

204 +

205 # Call update_weights

206 test_names = ["layer.weight", "layer.bias"]

207 test_dtypes = ["float32", "float32"]

@@ -229,10 +232,14 @@ def test_update_weights_calls_engine():

232 assert dtypes == test_dtypes

233 assert shapes == test_shapes

234

235 + # Finish weight update

236 + llm.finish_weight_update()

237 +

238

239 @create_new_process_for_each_test()

240 def test_full_weight_transfer_flow():

235 - """Test the complete weight transfer flow: init -> update."""

241 + """Test the complete weight transfer flow:

242 + init -> start -> update -> finish."""

243 if torch.accelerator.device_count() < 1:

244 pytest.skip("Need at least 1 GPU for this test")

245

@@ -253,12 +260,15 @@ def test_full_weight_transfer_flow():

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

261 )

262

256 - # Step 1: Initialize

263 + # Step 1: Initialize weight transfer engine

264 llm.init_weight_transfer_engine(

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

266 )

267

261 - # Step 2: Update weights

268 + # Step 2: Start weight update

269 + llm.start_weight_update(is_checkpoint_format=True)

270 +

271 + # Step 3: Update weights

272 llm.update_weights(

273 WeightTransferUpdateRequest(

274 update_info={

@@ -269,6 +279,9 @@ def test_full_weight_transfer_flow():

279 )

280 )

281

282 + # Step 4: Finish weight update

283 + llm.finish_weight_update()

284 +

285 # Verify the full flow completed

286 def check_flow(self):

287 engine = self.weight_transfer_engine

vllm/distributed/weight_transfer/base.py

+2 -6

@@ -4,7 +4,7 @@

5 from abc import ABC, abstractmethod

6 from collections.abc import Callable, Iterator

7 -from dataclasses import KW_ONLY, dataclass, field

7 +from dataclasses import dataclass, field

8 from typing import Any, Generic, TypeVar

10 import torch

@@ -28,11 +28,7 @@ class WeightTransferInitInfo(ABC): # noqa: B024

28 class WeightTransferUpdateInfo(ABC): # noqa: B024

29 """Base class for backend-specific weight update info."""

31 - _: KW_ONLY

32 - is_checkpoint_format: bool = True

33 - """Set to True if weights are in checkpoint/original model format and need

34 - layerwise processing. Set to False if weights have already been processed

35 - into kernel format (repacking, renaming, etc.)."""

31 + pass

34 # API-level request classes (accept dicts for backend-agnostic serialization)

vllm/distributed/weight_transfer/ipc_engine.py

@@ -211,6 +211,11 @@ class IPCWeightTransferEngine(

211 - 'ray': Sends weights via Ray RPC to a Ray-based LLM handle

212 - 'http': Sends weights via HTTP POST to a vLLM HTTP server

213

214 + .. note::

215 + This method calls ``update_weights`` internally. The caller must

216 + call ``start_weight_update`` before and ``finish_weight_update``

217 + after this method.

218 +

219 Args:

220 iterator: Iterator of model parameters. Returns (name, tensor) tuples.

221 Tensors should be on the same GPU as the inference workers.

vllm/engine/protocol.py

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

230 """Initialize weight transfer for RL training."""

231 raise NotImplementedError

232

233 + async def start_weight_update(self, is_checkpoint_format: bool = True) -> None:

234 + """Start a new weight update."""

235 + raise NotImplementedError

236 +

237 async def update_weights(self, request: WeightTransferUpdateRequest) -> None:

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

239 raise NotImplementedError

240 +

241 + async def finish_weight_update(self) -> None:

242 + """Finish the current weight update."""

243 + raise NotImplementedError

vllm/entrypoints/llm.py

+20

@@ -1908,6 +1908,20 @@ class LLM:

1908 "init_weight_transfer_engine", kwargs={"init_info": init_info_dict}

1909 )

1910

1911 + def start_weight_update(self, is_checkpoint_format: bool = True) -> None:

1912 + """

1913 + Start a new weight update.

1914 +

1915 + Args:

1916 + is_checkpoint_format: Whether incoming weights are in checkpoint

1917 + format (need layerwise processing) or kernel format (direct

1918 + copy).

1919 + """

1920 + self.llm_engine.collective_rpc(

1921 + "start_weight_update",

1922 + kwargs={"is_checkpoint_format": is_checkpoint_format},

1923 + )

1924 +

1925 def update_weights(self, request: WeightTransferUpdateRequest | dict) -> None:

1926 """

1927 Update the weights of the model.

@@ -1923,6 +1937,12 @@ class LLM:

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

1938 )

1939

1940 + def finish_weight_update(self) -> None:

1941 + """

1942 + Finish the current weight update.

1943 + """

1944 + self.llm_engine.collective_rpc("finish_weight_update")

1945 +

1946 def __repr__(self) -> str:

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

1948 # Cache the result to avoid repeated collective_rpc calls

vllm/entrypoints/serve/rlhf/api_router.py

+19

@@ -128,6 +128,19 @@ async def init_weight_transfer_engine(raw_request: Request):

128 return JSONResponse(content={"message": "Weight transfer initialized"})

129

130

131 +@router.post("/start_weight_update")

132 +async def start_weight_update(raw_request: Request):

133 + try:

134 + body = await raw_request.json()

135 + except json.JSONDecodeError as e:

136 + raise HTTPException(status_code=400, detail="Invalid JSON format") from e # noqa: B904

137 + is_checkpoint_format = body.get("is_checkpoint_format", True)

138 + await engine_client(raw_request).start_weight_update(

139 + is_checkpoint_format=is_checkpoint_format

140 + )

141 + return JSONResponse(content={"message": "Weight update started"})

142 +

143 +

144 @router.post("/update_weights")

145 async def update_weights(raw_request: Request):

146 try:

@@ -146,6 +159,12 @@ async def update_weights(raw_request: Request):

159 return JSONResponse(content={"message": "Weights updated"})

160

161

162 +@router.post("/finish_weight_update")

163 +async def finish_weight_update(raw_request: Request):

164 + await engine_client(raw_request).finish_weight_update()

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

166 +

167 +

168 @router.get("/get_world_size")

169 async def get_world_size(

170 raw_request: Request,

vllm/v1/engine/async_llm.py

+11

@@ -1049,6 +1049,13 @@ class AsyncLLM(EngineClient):

1049 "init_weight_transfer_engine", kwargs={"init_info": init_info_dict}

1050 )

1051

1052 + async def start_weight_update(self, is_checkpoint_format: bool = True) -> None:

1053 + """Start a new weight update."""

1054 + await self.collective_rpc(

1055 + "start_weight_update",

1056 + kwargs={"is_checkpoint_format": is_checkpoint_format},

1057 + )

1058 +

1059 async def update_weights(self, request: WeightTransferUpdateRequest) -> None:

1060 """

1061 Batched weight update for RL training.

@@ -1067,3 +1074,7 @@ class AsyncLLM(EngineClient):

1074 await self.collective_rpc(

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

1076 )

1077 +

1078 + async def finish_weight_update(self) -> None:

1079 + """Finish the current weight update."""

1080 + await self.collective_rpc("finish_weight_update")

vllm/v1/worker/gpu_worker.py

+98 -32

@@ -140,6 +140,8 @@ class Worker(WorkerBase):

140 if self.vllm_config.weight_transfer_config is not None

141 else None

142 )

143 + self._weight_update_active = False

144 + self._is_checkpoint_format = True

145

146 # Torch/CUDA profiler. Enabled and configured through profiler_config.

147 # Profiler wrapper is created lazily in profile() when start is called,

@@ -966,6 +968,13 @@ class Worker(WorkerBase):

968 model_config=self.model_config,

969 )

970

971 + def _check_weight_transfer_engine(self) -> None:

972 + if self.weight_transfer_engine is None:

973 + raise RuntimeError(

974 + "Weight transfer not configured. "

975 + "Please set weight_transfer_config to enable weight transfer."

976 + )

977 +

978 def init_weight_transfer_engine(self, init_info: dict) -> None:

979 """

980 Initialize weight transfer mechanism.

@@ -974,26 +983,62 @@ class Worker(WorkerBase):

983 Args:

984 init_info: Dictionary containing backend-specific initialization info

985 """

977 - if self.weight_transfer_engine is None:

978 - raise RuntimeError(

979 - "Weight transfer not configured. "

980 - "Please set weight_transfer_config to enable weight transfer."

981 - )

986 + self._check_weight_transfer_engine()

987 + assert self.weight_transfer_engine is not None

988 # Parse dict into backend-specific typed dataclass

989 typed_init_info = self.weight_transfer_engine.parse_init_info(init_info)

990 self.weight_transfer_engine.init_transfer_engine(typed_init_info)

991

992 + def start_weight_update(self, is_checkpoint_format: bool = True) -> None:

993 + """

994 + Start a new weight update.

995 +

996 + Prepares the model for receiving weights. For checkpoint format,

997 + this initializes state for layerwise processing. For kernel format, this is

998 + a no-op but must still be called for consistency.

999 +

1000 + Args:

1001 + is_checkpoint_format: Whether incoming weights are in checkpoint

1002 + format (need layerwise processing) or kernel format (direct

1003 + copy). Stored as state for finish_weight_update.

1004 + """

1005 + self._check_weight_transfer_engine()

1006 +

1007 + if self._weight_update_active:

1008 + raise RuntimeError(

1009 + "start_weight_update called while a weight update is "

1010 + "already active. Call finish_weight_update first."

1011 + )

1012 +

1013 + if is_checkpoint_format:

1014 + from vllm.model_executor.model_loader.reload import (

1015 + initialize_layerwise_reload,

1016 + )

1017 +

1018 + model = self.model_runner.model

1019 + with torch.device(self.device):

1020 + initialize_layerwise_reload(model)

1021 +

1022 + # Store state so update_weights/finish_weight_update can check

1023 + self._is_checkpoint_format = is_checkpoint_format

1024 + self._weight_update_active = True

1025 +

1026 def update_weights(self, update_info: dict) -> None:

1027 """

988 - Batched weight update from the trainer.

1028 + Receive weights from the trainer (one or more chunks).

1029 +

1030 + start_weight_update must be called before update_weights and

1031 + finish_weight_update must be called after.

1032

1033 Args:

1034 update_info: Dictionary containing backend-specific update info

1035 """

993 - if self.weight_transfer_engine is None:

1036 + self._check_weight_transfer_engine()

1037 + assert self.weight_transfer_engine is not None

1038 +

1039 + if not self._weight_update_active:

1040 raise RuntimeError(

995 - "Weight transfer not configured. "

996 - "Please set weight_transfer_config to enable weight transfer."

1041 + "start_weight_update must be called before update_weights."

1042 )

1043

1044 # Parse dict into backend-specific typed dataclass

@@ -1001,38 +1046,59 @@ class Worker(WorkerBase):

1046

1047 model = self.model_runner.model

1048

1004 - if typed_update_info.is_checkpoint_format:

1005 - from vllm.model_executor.model_loader.reload import (

1006 - finalize_layerwise_reload,

1007 - initialize_layerwise_reload,

1008 - )

1009 -

1010 - # Use layerwise reload pattern for checkpoint format weights

1011 - with torch.device(self.device):

1012 - initialize_layerwise_reload(model)

1049 + with torch.device(self.device):

1050 + if self._is_checkpoint_format:

1051 self.weight_transfer_engine.receive_weights(

1052 typed_update_info,

1053 load_weights=model.load_weights,

1054 )

1017 - finalize_layerwise_reload(model, self.model_config)

1018 - else:

1019 - # Weights are already in kernel format, copy directly

1020 - def load_weights_direct(

1021 - weights: list[tuple[str, torch.Tensor]],

1022 - ) -> None:

1023 - for name, weight in weights:

1024 - param = model.get_parameter(name)

1025 - param.copy_(weight)

1026 -

1027 - self.weight_transfer_engine.receive_weights(

1028 - typed_update_info,

1029 - load_weights=load_weights_direct,

1030 - )

1055 + else:

1056 + # Weights are already in kernel format, copy directly

1057 + def load_weights_direct(

1058 + weights: list[tuple[str, torch.Tensor]],

1059 + ) -> None:

1060 + for name, weight in weights:

1061 + param = model.get_parameter(name)

1062 + param.copy_(weight)

1063 +

1064 + self.weight_transfer_engine.receive_weights(

1065 + typed_update_info,

1066 + load_weights=load_weights_direct,

1067 + )

1068

1069 # NCCL broadcast/packed path are asynchronous.

1070 # Sync here so the next step uses the new weights.

1071 torch.accelerator.synchronize()

1072

1073 + def finish_weight_update(self) -> None:

1074 + """

1075 + Finish the current weight update.

1076 +

1077 + For checkpoint format, this runs layerwise postprocessing.

1078 + Uses the is_checkpoint_format state stored by start_weight_update.

1079 + """

1080 + self._check_weight_transfer_engine()

1081 +

1082 + if not self._weight_update_active:

1083 + raise RuntimeError(

1084 + "start_weight_update must be called before finish_weight_update."

1085 + )

1086 +

1087 + is_checkpoint_format = self._is_checkpoint_format

1088 +

1089 + if is_checkpoint_format:

1090 + from vllm.model_executor.model_loader.reload import (

1091 + finalize_layerwise_reload,

1092 + )

1093 +

1094 + model = self.model_runner.model

1095 + with torch.device(self.device):

1096 + finalize_layerwise_reload(model, self.model_config)

1097 +

1098 + # Reset state

1099 + self._weight_update_active = False

1100 + self._is_checkpoint_format = True

1101 +

1102 def shutdown(self) -> None:

1103 # has_kv_transfer_group can be None during interpreter shutdown.

1104 if ensure_kv_transfer_shutdown is not None: