| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | """Sharded Ray Direct Transport (RDT) weight transfer engine (consumer side). |
| 4 | |
| 5 | Pulls only the slice each vLLM worker consumes under tensor/expert parallelism, |
| 6 | not the full HF-format tensor. |
| 7 | |
| 8 | Two phases. BAKE, once at ``init_transfer_engine``: drive ``model.load_weights`` |
| 9 | against ``FakeRDTTensor`` placeholders and record, per leaf module, how each |
| 10 | destination slice is fetched (an op chain) and where it lands (an ``as_strided`` |
| 11 | descriptor). REPLAY, every sync: no ``load_weights``, no ``FakeRDTTensor`` |
| 12 | dispatch, no discovery — pull the recorded slices in packed chunks over a ring |
| 13 | of receive buffers, scatter them into freshly materialized params, then quant |
| 14 | and kernel-copy. A live name with no recorded plan fails the plan build: there |
| 15 | is no fallback load. |
| 16 | |
| 17 | Weights arrive in checkpoint format: the engine drives layerwise reload itself, |
| 18 | in ``start_weight_update`` / ``finish_weight_update``. |
| 19 | |
| 20 | Data flow |
| 21 | --------- |
| 22 | One thing at four resolutions, over three lifetimes. ``FetchKey`` -- |
| 23 | ``(name, op_chain)``, "which slice of which trainer tensor" -- is the atom; |
| 24 | everything else is bookkeeping around it. |
| 25 | |
| 26 | BAKE once at init, kept for the engine's life |
| 27 | FakeRDTTensor intercepts the model's loaders; each copy_ records a |
| 28 | _Scatter (sharded_rdt_fake): src FetchKey, owning layer, the destination |
| 29 | as_strided region, and the produced dtype/nbytes. |
| 30 | -> _name_to_plan: name -> that module's scatter list |
| 31 | |
| 32 | PLAN once, cached, _build_call_plan |
| 33 | _Chunk one packed pull: its scatters, deduped keys, the byte-exact |
| 34 | pack_layout, which producer serves it, and what to run after |
| 35 | (materialize / quant / free). One per (group, owner class). |
| 36 | _CallPlan all chunks + pre_free. |
| 37 | |
| 38 | RUN per chunk, per sync, _run_chunk_pipeline |
| 39 | _Chunk -> _PendingPull issued, not yet landed: Ray ref, buffer views, |
| 40 | ring slot. [RPC thread] |
| 41 | -> _ProcItem chunk + results + slot. [hand-off to the |
| 42 | background scatter thread] |
| 43 | |
| 44 | The RUN pair stays split on purpose: it is the thread boundary, and only |
| 45 | ``targets`` should outlive the get -- carrying the Ray ref and the whole-buffer |
| 46 | blob into the queue would keep both alive for the scatter's lifetime. |
| 47 | |
| 48 | See docs/training/weight_transfer/sharded_rdt.md for the design and the measured |
| 49 | results behind the choices here. |
| 50 | """ |
| 51 | |
| 52 | import time |
| 53 | from dataclasses import dataclass, field |
| 54 | from math import prod |
| 55 | from typing import TYPE_CHECKING, Any, cast |
| 56 | |
| 57 | import torch |
| 58 | |
| 59 | from vllm.config.weight_transfer import WeightTransferConfig |
| 60 | from vllm.distributed.weight_transfer.base import ( |
| 61 | WeightTransferEngine, |
| 62 | WeightTransferInitInfo, |
| 63 | WeightTransferUpdateInfo, |
| 64 | ) |
| 65 | from vllm.distributed.weight_transfer.sharded_rdt_common import ( |
| 66 | RdtRouter, |
| 67 | buffer_alloc_bytes, |
| 68 | check_ray_rdt_version, |
| 69 | register_nixl_memory, |
| 70 | ) |
| 71 | from vllm.distributed.weight_transfer.sharded_rdt_fake import ( |
| 72 | BakeSink, |
| 73 | FakeRDTTensor, |
| 74 | FetchKey, |
| 75 | _Scatter, |
| 76 | ) |
| 77 | from vllm.logger import init_logger |
| 78 | |
| 79 | if TYPE_CHECKING: |
| 80 | from vllm.config import VllmConfig |
| 81 | |
| 82 | logger = init_logger(__name__) |
| 83 | |
| 84 | |
| 85 | @dataclass |
| 86 | class _Chunk: |
| 87 | """One packed pull plus its post-processing (see Data flow). |
| 88 | |
| 89 | A module's copies span chunks when its experts sit in several owner classes, |
| 90 | so ``materialize`` fires on its FIRST chunk and ``quant`` on its LAST -- |
| 91 | materialize-once by construction, not by a runtime counter. ``pack_layout`` |
| 92 | mirrors the producer's rule byte-exactly. |
| 93 | """ |
| 94 | |
| 95 | scatters: "list[_Scatter]" |
| 96 | keys: "list[FetchKey]" |
| 97 | pack_layout: "list[tuple[int, torch.dtype, int, tuple[int, ...]]]" # off,dt,n,shape |
| 98 | pack_bytes: int |
| 99 | materialize: "list[Any]" |
| 100 | quant: "list[Any]" |
| 101 | free: "list[int]" |
| 102 | # Trainer rank serving this chunk, resolved at plan time. One per chunk: |
| 103 | # every name in it shares an owner class, so one producer holds them all. |
| 104 | owner: int |
| 105 | # This chunk's index among the chunks this worker pulls from ``owner``, in |
| 106 | # issue order. It picks the producer's serve slot (``seq % ring_depth``), |
| 107 | # which is what makes slot reuse safe: the pipeline drains pull i before |
| 108 | # issuing i+K, so the slot seq+K takes was last used by seq, whose read is |
| 109 | # finished. Choosing the slot in EXECUTION order instead -- as a per-call |
| 110 | # counter on the producer would -- can hand a slot to a pull whose |
| 111 | # predecessor is still being read, because Ray may start a consumer's |
| 112 | # concurrent produce calls in any order. Static, since the plan is. |
| 113 | seq: int = 0 |
| 114 | |
| 115 | |
| 116 | @dataclass |
| 117 | class _CallPlan: |
| 118 | """The static plan for one sync (see Data flow). Pure, so it is built once |
| 119 | and reused; runtime is then execution only. |
| 120 | |
| 121 | ``pre_free`` = groups with NO chunk on this worker, signaled at sync start |
| 122 | (owners tolerate a signal preceding its publish). With the last-chunk signals |
| 123 | this keeps the completeness invariant consumer-local: every group is signaled |
| 124 | exactly once. |
| 125 | """ |
| 126 | |
| 127 | chunks: "list[_Chunk]" |
| 128 | pre_free: "list[int]" |
| 129 | |
| 130 | |
| 131 | @dataclass |
| 132 | class _PendingPull: |
| 133 | """A dispatched pull whose blocking ``ray.get`` has not run (see Data flow). |
| 134 | |
| 135 | ``targets``/``blob`` must stay strongly referenced until it completes: |
| 136 | ``set_target_for_ref`` stores WEAKREFS, so dropping them silently reroutes the |
| 137 | transfer into a fallback buffer.""" |
| 138 | |
| 139 | ref: "Any" |
| 140 | keys: "list[FetchKey]" |
| 141 | targets: "list[torch.Tensor]" |
| 142 | blob: "torch.Tensor" |
| 143 | slot: int |
| 144 | |
| 145 | |
| 146 | @dataclass |
| 147 | class _ProcItem: |
| 148 | """A landed pull handed to the background scatter thread (see Data flow). |
| 149 | |
| 150 | ``results`` alias the ring buffer ``slot``, held as strong refs so they |
| 151 | outlive the RPC-thread frame until the scatter consumes them. |
| 152 | """ |
| 153 | |
| 154 | chunk: "_Chunk" |
| 155 | results: "dict[FetchKey, torch.Tensor]" |
| 156 | slot: int |
| 157 | |
| 158 | |
| 159 | @dataclass |
| 160 | class ShardedRDTWeightTransferInitInfo(WeightTransferInitInfo): |
| 161 | """Initialization info for the sharded RDT backend.""" |
| 162 | |
| 163 | trainer_actor_names: list[str] = field(default_factory=list) |
| 164 | """Names of all trainer Ray actors exposing the producer method (set via |
| 165 | ``.options(name=...)``), ordered by trainer rank. ``RdtRouter`` picks one of |
| 166 | them per pull, out of the name's owner set. Every actor is bound regardless, |
| 167 | because ``free_group`` fans out to every owner. Must be non-empty; a |
| 168 | single-producer trainer passes a one-element list.""" |
| 169 | |
| 170 | trainer_actor_namespace: str | None = None |
| 171 | """Optional Ray namespace the trainer actor(s) live in.""" |
| 172 | |
| 173 | produce_method_name: str = "rdt_produce_weights_batched" |
| 174 | """Name of the trainer-side producer method. It and the rest of the serve |
| 175 | surface (``free_group``, ``reserve_serve_buffer``) are documented where they |
| 176 | are implemented, on ``_RDTProducerServer`` in ``sharded_rdt_trainer.py``.""" |
| 177 | |
| 178 | names: list[str] = field(default_factory=list) |
| 179 | """The trainer's complete, flat param name list. The bake drives |
| 180 | ``model.load_weights`` over all of them once and keys the plan by source |
| 181 | name.""" |
| 182 | |
| 183 | dtype_names: list[str] = field(default_factory=list) |
| 184 | """Dtype name (e.g. 'bfloat16') for each entry of ``names``.""" |
| 185 | |
| 186 | shapes: list[list[int]] = field(default_factory=list) |
| 187 | """Full HF shape for each entry of ``names``.""" |
| 188 | |
| 189 | group_lens: list[int] = field(default_factory=list) |
| 190 | """Partition of ``names`` into gather groups, in the SAME order the trainers |
| 191 | gather and publish them (group-major; ``sum(group_lens) == len(names)``, and |
| 192 | ``names`` must be ordered to match). Required: the engine pre-builds the whole |
| 193 | static chunk/signal plan from it at init, and fires ``free_group`` at every |
| 194 | owner as each group's last chunk completes.""" |
| 195 | |
| 196 | owner_sets: list[list[int]] = field(default_factory=list) |
| 197 | """The distinct producer sets that occur, each a sorted list of trainer ranks |
| 198 | (indices into ``trainer_actor_names``). Indexed by ``name_owner_class``. |
| 199 | Empty means every producer holds every name.""" |
| 200 | |
| 201 | name_owner_class: list[int] = field(default_factory=list) |
| 202 | """Per-name index into ``owner_sets``, parallel to ``names``: which producers |
| 203 | hold each name. Empty means every producer holds every name. |
| 204 | |
| 205 | This one table expresses every layout the trainer can have — pipeline stages |
| 206 | (the names of these groups have this owner set), expert parallelism (this |
| 207 | expert name has this one-rank owner set), and combinations of the two — so |
| 208 | the engine cuts each group's baked copies into one chunk per distinct class |
| 209 | present and routes each chunk to that class's owner. Derived by the trainer |
| 210 | from ``WeightSource.held_names()``.""" |
| 211 | |
| 212 | num_consumers: int = 0 |
| 213 | """Total consumer count across the fleet, for M:N routing. Authoritative when |
| 214 | > 0; at 0 the engine infers it from ``parallel_config``, which is correct for |
| 215 | the supported serving modes but worth setting explicitly under M:N. Each |
| 216 | worker's distinct index comes from ``_global_worker_index``.""" |
| 217 | |
| 218 | num_rdt_buffers: int = 2 |
| 219 | """Depth of the consumer receive-buffer ring. Must match the |
| 220 | producer's — ``_run_chunk_pipeline``'s slot-safety argument rests on it. 2 = |
| 221 | double buffer: chunk i+1's serve overlaps chunk i's RDMA, and scatter(i-1) |
| 222 | overlaps RDMA(i) in the other slot. Keep depth x chunk_bytes under the |
| 223 | fabric's address-translation reach (~2-3 GB/flow on the reference 8xB200 RoCE |
| 224 | cluster, where K=3 measurably hurt).""" |
| 225 | |
| 226 | buffer_presize_gb: float = 0.0 |
| 227 | """Pre-size each packed receive-buffer slot to this many GiB |
| 228 | (0 = size to the first chunk + coarse 256MB round-up). Set it to cover the |
| 229 | model's largest atomic chunk (e.g. an untied lm_head). Sizing buffers ONCE |
| 230 | matters beyond perf -- see the doc's "Sizing buffers once matters beyond |
| 231 | throughput".""" |
| 232 | |
| 233 | replica_rank: int = 0 |
| 234 | """This inference engine's ordinal in the fleet (0..``num_replicas``-1). |
| 235 | |
| 236 | Multi-engine deployments run several INDEPENDENT inference engines, each with |
| 237 | its own self-contained parallel config, so every engine's |
| 238 | ``_global_worker_index`` restarts at 0 and would collide across engines. The |
| 239 | driver gives each engine a distinct ``replica_rank`` (with identical |
| 240 | ``num_replicas``) so the engine offsets its consumers into a globally distinct |
| 241 | range for the M:N block assignment.Default 0/1 (single |
| 242 | engine) needs no override.""" |
| 243 | |
| 244 | num_replicas: int = 1 |
| 245 | """Number of independent inference engines in the fleet. Default 1 => the |
| 246 | per-replica offset is 0 and consumer identity is exactly |
| 247 | ``_global_worker_index`` (preserves single-engine and single-DP-deployment |
| 248 | behavior). When > 1, ``workers_per_replica = num_consumers // num_replicas`` |
| 249 | and this engine's consumers occupy |
| 250 | ``replica_rank * workers_per_replica + _global_worker_index()``. Assumes a |
| 251 | uniform fleet (every replica has the same worker count). Set alongside |
| 252 | ``replica_rank``, by the driver.""" |
| 253 | |
| 254 | |
| 255 | @dataclass |
| 256 | class ShardedRDTWeightTransferUpdateInfo(WeightTransferUpdateInfo): |
| 257 | """Update info for the sharded RDT backend: intentionally EMPTY. |
| 258 | |
| 259 | The chunk/free plan is a pure function of the baked plan and the driver's |
| 260 | gather-group partition, both fixed for the engine's lifetime, so it is built |
| 261 | once at ``init_transfer_engine`` from ``ShardedRDTWeightTransferInitInfo``'s |
| 262 | ``names`` + ``group_lens``. ONE ``update_weights`` per sync then just re-runs |
| 263 | that plan; there is nothing per-sync to carry. |
| 264 | """ |
| 265 | |
| 266 | |
| 267 | class ShardedRDTWeightTransferEngine( |
| 268 | WeightTransferEngine[ |
| 269 | ShardedRDTWeightTransferInitInfo, |
| 270 | ShardedRDTWeightTransferUpdateInfo, |
| 271 | ] |
| 272 | ): |
| 273 | """Pull-based RDT/NIXL backend that transports only the slice each worker |
| 274 | consumes. |
| 275 | |
| 276 | Requires ``distributed_executor_backend="ray"``, ``nixl`` in the shared env, |
| 277 | a named trainer actor exposing a ``@ray.method(tensor_transport="nixl")`` |
| 278 | producer, and weight loaders that stay inside ``SUPPORTED_OPS`` — anything |
| 279 | needing real data (``.to``, ``.item``, arithmetic, bool-mask indexing) |
| 280 | raises during the bake. |
| 281 | |
| 282 | The plan is baked once at ``init_transfer_engine`` into one scatter list |
| 283 | per fully-loaded leaf module, indexed by source name; every |
| 284 | ``update_weights`` replays the modules its gathered names cover. |
| 285 | """ |
| 286 | |
| 287 | init_info_cls = ShardedRDTWeightTransferInitInfo |
| 288 | update_info_cls = ShardedRDTWeightTransferUpdateInfo |
| 289 | # receive_weights pulls synchronously but defers GPU post-processing to a |
| 290 | # background thread so it overlaps the next chunk's pull, so ``update_weights`` |
| 291 | # skips the base's device sync and ``finish_weight_update`` drains the |
| 292 | # deferred work before finalize. |
| 293 | defers_processing = True |
| 294 | # The baked replay plan is a function of one concrete model's parameter |
| 295 | # layout, so a separate draft model cannot reuse it. |
| 296 | supports_draft_weight_update = False |
| 297 | |
| 298 | def __init__( |
| 299 | self, |
| 300 | config: WeightTransferConfig, |
| 301 | vllm_config: "VllmConfig", |
| 302 | device: torch.device, |
| 303 | model: torch.nn.Module, |
| 304 | ) -> None: |
| 305 | super().__init__(config, vllm_config, device, model) |
| 306 | # This engine's data plane is Ray's: the consumer resolves each trainer's |
| 307 | # serve actor with ``ray.get_actor`` and pulls over NIXL through the |
| 308 | # workers' tensor transport, neither of which exists unless the workers |
| 309 | # ARE Ray actors. Caught here, at construction, because the alternative |
| 310 | # is an opaque failure during the first handshake. A single-worker |
| 311 | # deployment defaults to "uni", so "ray" has to be asked for explicitly. |
| 312 | backend = vllm_config.parallel_config.distributed_executor_backend |
| 313 | # A custom Executor subclass is a deliberate override we cannot judge; |
| 314 | # only reject the named backends that are definitely not Ray. |
| 315 | if isinstance(backend, str) and backend != "ray": |
| 316 | raise ValueError( |
| 317 | f"The 'sharded_rdt' weight transfer backend requires " |
| 318 | f"distributed_executor_backend='ray', got {backend!r}. Pass " |
| 319 | f"--distributed-executor-backend ray (and keep " |
| 320 | f"VLLM_USE_RAY_V2_EXECUTOR_BACKEND=1, the default, so the " |
| 321 | f"workers are RayExecutorV2 actors with tensor transport)." |
| 322 | ) |
| 323 | # Ownership table, routing rules, and (once bound) every producer's |
| 324 | # handle — pulls go to one owner per chunk, but free_group fans out to |
| 325 | # every owner of the group, so all handles are bound. |
| 326 | self._router: RdtRouter | None = None |
| 327 | # Driver-supplied total consumer count (init_info.num_consumers); 0 => infer |
| 328 | # from parallel_config. See _num_consumers. |
| 329 | self._num_consumers_override: int = 0 |
| 330 | # Source name -> the plan for the module consuming it. Several names of a |
| 331 | # fused module share one entry; replay dedups. A live name absent here |
| 332 | # fails the plan build. |
| 333 | self._name_to_plan: dict[str, list[_Scatter]] = {} |
| 334 | # name -> (dtype_name, shape) for every init name; the bake builds its |
| 335 | # lazies from this. |
| 336 | self._name_meta: dict[str, tuple[str, list[int]]] = {} |
| 337 | # Names whose copy_ fired during the bake. Unbaked names not here never |
| 338 | # move data (e.g. non-local EP experts), so receive_weights skips them. |
| 339 | self._live_names: set[str] = set() |
| 340 | |
| 341 | # ---- Consumer-side pre-registered receive buffers ---------------------- |
| 342 | # One persistent 1-D buffer per dtype per ring slot, STORAGE-registered with |
| 343 | # NIXL once, so every per-slice view hits Ray's registration cache and the |
| 344 | # recv path does zero registration in steady state. See the doc's "Sizing |
| 345 | # buffers once matters beyond throughput" for why they rarely grow. |
| 346 | # |
| 347 | # Ring depth is both the slot count and the in-flight pull count, one |
| 348 | # quantity from ``num_rdt_buffers``. Decoupling them with a spare slot |
| 349 | # measurably slows the sync. |
| 350 | self._ring_depth = 1 |
| 351 | self._dest_buffers: list[dict[torch.dtype, torch.Tensor]] = [{}] |
| 352 | self._slot_read_done: list[Any] = [] # one torch.cuda.Event per slot |
| 353 | # Generation counters guarding the events (see _ensure_proc_worker). |
| 354 | self._slot_queued: list[int] = [] |
| 355 | self._slot_done: list[int] = [] |
| 356 | self._slot_cv: Any = None |
| 357 | self._pull_slot = 0 |
| 358 | # (id(chunk), slot) -> (buffer data_ptr, per-key views). |
| 359 | # Lives on the engine, not on the _Chunk: the plan is static and shared |
| 360 | # across syncs, so a per-slot runtime cache has no business inside it. |
| 361 | self._targets_cache: dict[tuple[int, int], tuple[int, list[torch.Tensor]]] = {} |
| 362 | self._buffer_presize = 0 # bytes; set from init_info |
| 363 | self._pending_frees: list[Any] = [] # free_group signal refs, drained per sync |
| 364 | # The STATIC plan: built once — at init from init_info.group_lens, else |
| 365 | # lazily on the first update_weights — and reused. Non-None means |
| 366 | # update_weights ignores its per-sync names. |
| 367 | self._cached_plan: _CallPlan | None = None |
| 368 | # Completed syncs. The chunk pipeline runs SERIAL during sync 0, when both |
| 369 | # sides still register buffers, and pipelines from sync 1. See the doc. |
| 370 | self._completed_syncs = 0 |
| 371 | |
| 372 | # ---- Background post-processing thread (pull/process pipelining) ------- |
| 373 | # The RPC thread hands pulled buffer views to this worker, which runs |
| 374 | # materialize/scatter/quant/kernel_copy on its own stream while the next |
| 375 | # pull proceeds. drain_pending() joins it before finalize. |
| 376 | self._proc_queue: Any | None = None |
| 377 | self._proc_thread: Any | None = None |
| 378 | self._proc_stream: Any | None = None |
| 379 | # Dedicated quant thread (pass 2): own queue + stream so quant never |
| 380 | # sits between two items' scatters on the scatter thread (measured to |
| 381 | # stall the RPC thread's slot handshake ~0.5-0.75s/iter). |
| 382 | self._quant_queue: Any | None = None |
| 383 | self._quant_thread: Any | None = None |
| 384 | self._quant_stream: Any | None = None |
| 385 | self._proc_error: BaseException | None = None |
| 386 | |
| 387 | def init_transfer_engine(self, init_info: ShardedRDTWeightTransferInitInfo) -> None: |
| 388 | """Configure the ring, bind the producers, bake the replay plan, and |
| 389 | pre-register every NIXL buffer -- in that order, because each step depends |
| 390 | on the previous one. |
| 391 | |
| 392 | The bake drives ``model.load_weights`` and the pre-registration blocks on |
| 393 | RPCs to the producers, so this is a heavyweight one-off; every later |
| 394 | ``update_weights`` is pure replay. |
| 395 | """ |
| 396 | # Read the attribute directly: if a vLLM bump renames it, this guard |
| 397 | # must fail loudly at init rather than silently disappear. |
| 398 | if self.parallel_config.enable_eplb: |
| 399 | raise RuntimeError( |
| 400 | "sharded_rdt does not support EPLB (enable_eplb=true): dynamic " |
| 401 | "expert rearrangement invalidates the baked replay plan, which " |
| 402 | "records each expert's destination slot once at init — after a " |
| 403 | "rearrangement the replay would silently load weights into the " |
| 404 | "wrong expert slots. Disable EPLB or use another weight-transfer " |
| 405 | "backend." |
| 406 | ) |
| 407 | # The worker is often a different install from whatever spawned it, so |
| 408 | # check here too and not only where the actor options are set. |
| 409 | check_ray_rdt_version() |
| 410 | self._configure_ring(init_info) |
| 411 | self._resolve_producers(init_info) |
| 412 | |
| 413 | # A pure dry run: the trainer's gather cache is empty at init, so nothing |
| 414 | # can (or does) get pulled -- we only record how each slice is fetched and |
| 415 | # where it lands, then restore the model. |
| 416 | self._bake(init_info) |
| 417 | self._build_static_plan(init_info) |
| 418 | # Register ALL NIXL memory now, while the fabric is idle: dma-buf GPUDirect |
| 419 | # registration concurrent with in-flight RDMA intermittently fails |
| 420 | # (ibv_reg_mr 'Bad address'), which bites under M:N fan-in. Sizes come from |
| 421 | # the static plan, so this is exact. |
| 422 | self._preregister_at_init() |
| 423 | # Start the background post-processing worker (pull/process pipelining). |
| 424 | self._ensure_proc_worker() |
| 425 | |
| 426 | def _configure_ring(self, init_info: ShardedRDTWeightTransferInitInfo) -> None: |
| 427 | """Ring depth K. |
| 428 | |
| 429 | Must run before ``_ensure_proc_worker`` creates the per-slot events and |
| 430 | counters, and before any buffer is grown (both happen on the first pull). |
| 431 | """ |
| 432 | self._num_consumers_override = int(init_info.num_consumers or 0) |
| 433 | k = max(1, int(init_info.num_rdt_buffers)) |
| 434 | self._ring_depth = k |
| 435 | self._dest_buffers = [{} for _ in range(k)] |
| 436 | self._targets_cache = {} |
| 437 | self._buffer_presize = int(float(init_info.buffer_presize_gb) * (1 << 30)) |
| 438 | logger.info( |
| 439 | "[RDT-RING] active_pulls=%d slots=%d presize=%.2fGiB", |
| 440 | k, |
| 441 | k, |
| 442 | self._buffer_presize / (1 << 30), |
| 443 | ) |
| 444 | |
| 445 | def _resolve_producers(self, init_info: ShardedRDTWeightTransferInitInfo) -> None: |
| 446 | """Work out this worker's consumer identity, build the router, and bind |
| 447 | EVERY producer actor. |
| 448 | |
| 449 | Pull routing is M:N: each chunk goes to ONE producer holding every name |
| 450 | in it (see ``RdtRouter``), so one producer serves a whole pull. All |
| 451 | producers are bound regardless, because the per-group ``free_group`` |
| 452 | signal fans out to every owner of a group, including producers this |
| 453 | worker never pulls from. |
| 454 | """ |
| 455 | try: |
| 456 | import ray |
| 457 | except ImportError as e: |
| 458 | raise RuntimeError( |
| 459 | "Ray is required for the 'sharded_rdt' weight transfer " |
| 460 | "backend. Install Ray and run workers as Ray actors " |
| 461 | "(distributed_executor_backend='ray')." |
| 462 | ) from e |
| 463 | |
| 464 | producer_names = list(init_info.trainer_actor_names) |
| 465 | if not producer_names: |
| 466 | raise RuntimeError( |
| 467 | "Sharded RDT engine requires a trainer producer: set " |
| 468 | "init_info.trainer_actor_names." |
| 469 | ) |
| 470 | |
| 471 | consumer_id = self._resolve_consumer_id(init_info) |
| 472 | name_owner_class = list(init_info.name_owner_class or []) |
| 473 | if name_owner_class and len(name_owner_class) != len(init_info.names): |
| 474 | raise RuntimeError( |
| 475 | f"Sharded RDT engine: {len(name_owner_class)} owner-class entries " |
| 476 | f"for {len(init_info.names)} names." |
| 477 | ) |
| 478 | router = RdtRouter( |
| 479 | len(producer_names), |
| 480 | self._num_consumers(), |
| 481 | list(init_info.owner_sets) or None, |
| 482 | name_owner_class or None, |
| 483 | list(init_info.names), |
| 484 | list(init_info.group_lens), |
| 485 | # Carve the block per deployment so every replica's worker w pulls |
| 486 | # each name from the same producer, which is what lets that producer |
| 487 | # serve all of them from one slot. |
| 488 | workers_per_replica=self._workers_per_replica(init_info), |
| 489 | ) |
| 490 | router.validate() |
| 491 | |
| 492 | # A fresh router per init, never appended to: every owner index is a |
| 493 | # position in these lists, so a rejoining engine's second |
| 494 | # init_transfer_engine must not shift them. |
| 495 | actors: list[Any] = [] |
| 496 | methods: list[Any] = [] |
| 497 | |
| 498 | for chosen_name in producer_names: |
| 499 | try: |
| 500 | actor = ray.get_actor( |
| 501 | chosen_name, |
| 502 | namespace=init_info.trainer_actor_namespace, |
| 503 | ) |
| 504 | except ValueError as e: |
| 505 | raise RuntimeError( |
| 506 | f"Sharded RDT engine could not find trainer actor " |
| 507 | f"{chosen_name!r} (namespace=" |
| 508 | f"{init_info.trainer_actor_namespace!r})." |
| 509 | ) from e |
| 510 | # ``ray.get_actor`` resolves by NAME, in a process that never |
| 511 | # imported the producer's class, and ``enable_tensor_transport`` |
| 512 | # lives in that class's metadata: the creating process passes |
| 513 | # ``meta.enable_tensor_transport``, and Ray infers it from the |
| 514 | # class's ``@ray.method(tensor_transport=...)`` decorators. A |
| 515 | # name-resolved handle therefore always reports False and the |
| 516 | # dispatch guard rejects the pull, even though the trainer set the |
| 517 | # option. Inherent to resolving by name across processes, not a |
| 518 | # version bug; it fails on Ray 2.56.0. |
| 519 | # |
| 520 | # Forcing it skips no validation. Ray's next guard, |
| 521 | # ``actor_has_tensor_transport``, asks the LIVE actor whether it can |
| 522 | # build a NIXL agent and still runs, and ``_spawn_server`` hardcodes |
| 523 | # the actor option, so a misconfigured producer cannot reach here. |
| 524 | actor._ray_enable_tensor_transport = True |
| 525 | actors.append(actor) |
| 526 | methods.append(getattr(actor, init_info.produce_method_name)) |
| 527 | router.bind(actors, methods, consumer_id) |
| 528 | self._router = router |
| 529 | logger.info( |
| 530 | "Sharded RDT engine (consumer %d) bound to all %d producers " |
| 531 | "(batched method %r, %d owner class(es))", |
| 532 | consumer_id, |
| 533 | len(producer_names), |
| 534 | init_info.produce_method_name, |
| 535 | len(init_info.owner_sets) or 1, |
| 536 | ) |
| 537 | |
| 538 | def _workers_per_replica(self, init_info: ShardedRDTWeightTransferInitInfo) -> int: |
| 539 | """Consumers per inference deployment, assuming a uniform fleet. |
| 540 | |
| 541 | Read by this worker's consumer id and by the router's block carve, which |
| 542 | must not disagree, so it is derived once here. |
| 543 | """ |
| 544 | return max(1, self._num_consumers() // max(1, int(init_info.num_replicas or 1))) |
| 545 | |
| 546 | def _resolve_consumer_id(self, init_info: ShardedRDTWeightTransferInitInfo) -> int: |
| 547 | """This worker's DISTINCT index in 0..C-1 across the whole fleet. |
| 548 | |
| 549 | Within one engine that is ``_global_worker_index()``. But a fleet of |
| 550 | INDEPENDENT engines (each with its own parallel config) restarts that |
| 551 | index at 0 per engine, so each engine offsets into its own range using |
| 552 | ``replica_rank``: with a uniform fleet, |
| 553 | ``workers_per_replica = C // num_replicas``. ``num_replicas`` defaults to |
| 554 | 1 (offset 0), preserving single-engine and single-DP-deployment behaviour. |
| 555 | """ |
| 556 | replica_rank = max(0, int(init_info.replica_rank or 0)) |
| 557 | return ( |
| 558 | replica_rank * self._workers_per_replica(init_info) |
| 559 | + self._global_worker_index() |
| 560 | ) |
| 561 | |
| 562 | def _build_static_plan(self, init_info: ShardedRDTWeightTransferInitInfo) -> None: |
| 563 | """Build the chunk/free plan once. It never changes across syncs, so |
| 564 | ``update_weights`` needs no per-sync names.""" |
| 565 | if not init_info.group_lens: |
| 566 | raise ValueError( |
| 567 | "Sharded RDT engine requires init_info.group_lens (the gather-" |
| 568 | "group partition of init_info.names)." |
| 569 | ) |
| 570 | if sum(init_info.group_lens) != len(init_info.names): |
| 571 | raise ValueError( |
| 572 | f"init_info.group_lens sums to {sum(init_info.group_lens)} " |
| 573 | f"but {len(init_info.names)} names were given." |
| 574 | ) |
| 575 | self._cached_plan = self._build_call_plan(init_info.names, init_info.group_lens) |
| 576 | logger.info( |
| 577 | "[RDT-PLAN] pre-built static call plan at init: %d chunks", |
| 578 | len(self._cached_plan.chunks), |
| 579 | ) |
| 580 | |
| 581 | def _preregister_at_init(self) -> None: |
| 582 | """Register every NIXL buffer this worker will use at init, before any |
| 583 | transfer runs, so nothing registers during the sync-0 RDMA churn. |
| 584 | |
| 585 | Both sides are sized from the static plan: receive buffers are |
| 586 | ``ring_depth`` slots at the largest chunk's ``pack_bytes``, and each bound |
| 587 | producer is asked to pre-register a serve ring at the max bytes this |
| 588 | consumer will pull from it. A no-op when this worker has no chunks.""" |
| 589 | plan = self._cached_plan |
| 590 | if plan is None or not plan.chunks: |
| 591 | return |
| 592 | # (a) consumer receive buffers — one per ring slot, at the largest chunk. |
| 593 | max_pack = max(c.pack_bytes for c in plan.chunks) |
| 594 | alloc = buffer_alloc_bytes(max_pack, self._buffer_presize) |
| 595 | for slot in range(self._ring_depth): |
| 596 | buffer = self._dest_buffers[slot].get(torch.uint8) |
| 597 | if buffer is None or buffer.numel() < alloc: |
| 598 | buffer = torch.empty(alloc, dtype=torch.uint8, device=self.device) |
| 599 | register_nixl_memory(buffer) |
| 600 | self._dest_buffers[slot][torch.uint8] = buffer |
| 601 | |
| 602 | # (b) producer serve rings — max bytes this consumer pulls from each |
| 603 | # bound producer, plus the per-producer plan digest a producer sharing |
| 604 | # one ring across deployments checks them against. |
| 605 | assert self._router is not None |
| 606 | serve_bytes = [0] * self._router.num_producers |
| 607 | per_owner_keys: list[list] = [[] for _ in range(self._router.num_producers)] |
| 608 | for c in plan.chunks: |
| 609 | serve_bytes[c.owner] = max(serve_bytes[c.owner], c.pack_bytes) |
| 610 | per_owner_keys[c.owner].append(c.keys) |
| 611 | digests = [_plan_digest(keys) for keys in per_owner_keys] |
| 612 | import ray |
| 613 | |
| 614 | refs = self._router.reserve_serve_buffers(serve_bytes, digests) |
| 615 | if refs: |
| 616 | ray.get(refs) # block until every serve ring is registered |
| 617 | logger.info( |
| 618 | "[RDT-PLAN] pre-registered %d receive slots (%.0f MiB each) + serve " |
| 619 | "rings on %d producer(s) %s", |
| 620 | self._ring_depth, |
| 621 | alloc / (1 << 20), |
| 622 | len(refs), |
| 623 | [nb // (1 << 20) for nb in serve_bytes], |
| 624 | ) |
| 625 | |
| 626 | def _global_worker_index(self) -> int: |
| 627 | """This worker's stable, distinct global index across the inference fleet: |
| 628 | ``data_parallel_index * world_size + rank`` over the TP*PP world. |
| 629 | |
| 630 | ``data_parallel_index``, not ``data_parallel_rank``: vLLM resets the |
| 631 | latter to 0 in a dense worker but keeps the former as the distinct global |
| 632 | DP rank. Same formula as the sibling ``nccl_engine``, so dense-via-TP and |
| 633 | MoE-via-DP+EP both yield distinct 0..C-1.""" |
| 634 | pc = self.parallel_config |
| 635 | return pc.data_parallel_index * pc.world_size + pc.rank # world_size = TP*PP |
| 636 | |
| 637 | def _num_consumers(self) -> int: |
| 638 | """Total inference-worker count. Prefers the driver-supplied |
| 639 | ``init_info.num_consumers`` (authoritative -- the driver knows the whole |
| 640 | fleet); else ``world_size_across_dp``, the same stride |
| 641 | ``_global_worker_index`` indexes with, so the two agree at any pp.""" |
| 642 | if self._num_consumers_override > 0: |
| 643 | return self._num_consumers_override |
| 644 | return self.parallel_config.world_size_across_dp |
| 645 | |
| 646 | def start_weight_update(self) -> None: |
| 647 | """Put the model's params on meta so layerwise reload streams them in |
| 648 | as each layer's slices land. Baked replay uses checkpoint format.""" |
| 649 | from vllm.model_executor.model_loader.reload import ( |
| 650 | initialize_layerwise_reload, |
| 651 | ) |
| 652 | |
| 653 | initialize_layerwise_reload(self.model) |
| 654 | |
| 655 | def finish_weight_update(self) -> None: |
| 656 | """Drain the deferred pull/process pipeline (so every layer is fully |
| 657 | loaded) before finalizing the layerwise reload.""" |
| 658 | from vllm.model_executor.model_loader.reload import ( |
| 659 | finalize_layerwise_reload, |
| 660 | ) |
| 661 | |
| 662 | self.drain_pending() |
| 663 | finalize_layerwise_reload(self.model, self.model_config) |
| 664 | |
| 665 | def update_weights(self, update_info: dict[str, Any]) -> None: |
| 666 | """Receive one update. Unlike the base, does NOT issue a per-update |
| 667 | device sync: post-processing is deferred to background threads and a |
| 668 | sync here would block on them and serialize the pull/process pipeline. |
| 669 | Completion is guaranteed by ``drain_pending`` in |
| 670 | ``finish_weight_update``.""" |
| 671 | self.receive_weights(self.parse_update_info(update_info)) |
| 672 | |
| 673 | def receive_weights( |
| 674 | self, |
| 675 | update_info: ShardedRDTWeightTransferUpdateInfo, |
| 676 | ) -> None: |
| 677 | """Pull + replay the baked leaf modules the sync covers. |
| 678 | |
| 679 | The chunk/free plan is STATIC across syncs — a pure function of the baked |
| 680 | plan and the driver's group partition — so it was built once at init and |
| 681 | every sync just re-runs the pipeline over its self-describing chunks, with |
| 682 | no per-sync bookkeeping and an empty ``update_info``. |
| 683 | |
| 684 | Assumes each baked module's source names fall within one gather group, |
| 685 | which the per-layer / pre / post partition guarantees (a leaf module's |
| 686 | sources all live in one decoder layer). A module that did span groups |
| 687 | would be planned once per group and could pull a name whose group the |
| 688 | pipeline already freed, which parks the pull until the producer's stall |
| 689 | watchdog fires. |
| 690 | """ |
| 691 | del update_info # the plan is static; nothing arrives per sync |
| 692 | if self._router is None: |
| 693 | raise RuntimeError( |
| 694 | "Sharded RDT engine not initialized. Call init_transfer_engine() first." |
| 695 | ) |
| 696 | # Surface any error the background thread hit on a prior item promptly. |
| 697 | self._raise_proc_error() |
| 698 | if self._cached_plan is None: |
| 699 | raise RuntimeError( |
| 700 | "Sharded RDT engine has no call plan: init_info.group_lens must " |
| 701 | "be supplied at init_transfer_engine()." |
| 702 | ) |
| 703 | self._run_chunk_pipeline(self._cached_plan) |
| 704 | |
| 705 | def _build_fake_weights( |
| 706 | self, |
| 707 | names: list[str], |
| 708 | sink: "BakeSink", |
| 709 | device: torch.device, |
| 710 | ) -> list[tuple[str, torch.Tensor]]: |
| 711 | """Zero-storage lazies for ``names``, dtype/shape from the init metadata, |
| 712 | all feeding the bake's recording sink.""" |
| 713 | return [ |
| 714 | ( |
| 715 | name, |
| 716 | FakeRDTTensor( |
| 717 | name=name, |
| 718 | shape=torch.Size(self._name_meta[name][1]), |
| 719 | dtype=_dtype_from_name(self._name_meta[name][0]), |
| 720 | device=device, |
| 721 | sink=sink, |
| 722 | ), |
| 723 | ) |
| 724 | for name in names |
| 725 | ] |
| 726 | |
| 727 | def _issue_pull(self, chunk: "_Chunk", slot: int) -> "_PendingPull": |
| 728 | """Reserve ``slot``, lay the targets out in its buffer, dispatch the |
| 729 | produce RPC and point the transfer at the buffer — WITHOUT the blocking |
| 730 | ``ray.get`` (that is ``_complete_pull``). The chunked pipeline issues |
| 731 | chunk i+1 before completing chunk i, so the producer serves the next |
| 732 | chunk while the in-flight RDMA streams. |
| 733 | |
| 734 | Slot-reuse guard, both stages required: a generation wait (the CUDA event |
| 735 | binds only to its LAST record, so synchronizing before the background |
| 736 | thread recorded this item's event passes silently — observed as |
| 737 | nondeterministic weight corruption), then the event synchronize. It must |
| 738 | precede ``set_target_for_ref``, not just the get: the transfer may start |
| 739 | any time after the metadata push. See the doc's "slot generation |
| 740 | handshake". |
| 741 | """ |
| 742 | from ray.experimental import set_target_for_ref |
| 743 | |
| 744 | if self._slot_read_done: |
| 745 | with self._slot_cv: |
| 746 | while self._slot_done[slot] < self._slot_queued[slot]: |
| 747 | self._slot_cv.wait(timeout=1.0) |
| 748 | self._slot_read_done[slot].synchronize() |
| 749 | |
| 750 | # Every slice is byte-packed into ONE uint8 buffer, 16B-aligned in |
| 751 | # keys order, under one NIXL descriptor. ``chunk.pack_layout`` mirrors the |
| 752 | # producer's rule byte-exactly; ``targets`` carves the dtype views back out. |
| 753 | keys = chunk.keys |
| 754 | cur = chunk.pack_bytes |
| 755 | buffers = self._dest_buffers[slot] |
| 756 | buffer = buffers.get(torch.uint8) |
| 757 | if buffer is None or buffer.numel() < cur: |
| 758 | # Size ONCE with headroom: Ray's desc cache is keyed by data_ptr and |
| 759 | # outlives its tensor, so a regrowth can false-hit a recycled pointer |
| 760 | # and skip registering the new extent -> NIXL_ERR_NOT_FOUND. |
| 761 | alloc = buffer_alloc_bytes(cur, self._buffer_presize) |
| 762 | buffer = torch.empty(alloc, dtype=torch.uint8, device=self.device) |
| 763 | register_nixl_memory(buffer) |
| 764 | buffers[torch.uint8] = buffer |
| 765 | targets = self._pull_targets(chunk, slot, buffer) |
| 766 | # One produce RPC serves the whole packed chunk, since the group's owner |
| 767 | # holds every slice. The buffer view must stay strongly referenced through |
| 768 | # the get: set_target_for_ref stores WEAKREFS, and a dropped target |
| 769 | # reroutes the transfer into a fallback buffer. |
| 770 | blob = buffer[:cur] |
| 771 | assert self._router is not None |
| 772 | ref = self._router.pull(chunk.owner, keys, chunk.seq) |
| 773 | set_target_for_ref(ref, [blob]) |
| 774 | return _PendingPull( |
| 775 | ref=ref, |
| 776 | keys=keys, |
| 777 | targets=targets, |
| 778 | blob=blob, |
| 779 | slot=slot, |
| 780 | ) |
| 781 | |
| 782 | def _pull_targets( |
| 783 | self, chunk: "_Chunk", slot: int, buffer: torch.Tensor |
| 784 | ) -> "list[torch.Tensor]": |
| 785 | """Per-key dtype views into ``slot``'s buffer. |
| 786 | |
| 787 | The packed layout is static, so the views are built once per (chunk, slot) |
| 788 | rather than once per pull -- rebuilding them cost ~1150 Python ops per pull |
| 789 | at 235B. Keyed on the buffer pointer as well, so a regrow invalidates |
| 790 | instead of handing back views into a freed buffer. |
| 791 | """ |
| 792 | cached = self._targets_cache.get((id(chunk), slot)) |
| 793 | if cached is not None and cached[0] == buffer.data_ptr(): |
| 794 | return cached[1] |
| 795 | targets = [ |
| 796 | buffer[off : off + n * dt.itemsize].view(dt).reshape(shape) |
| 797 | for off, dt, n, shape in chunk.pack_layout |
| 798 | ] |
| 799 | self._targets_cache[(id(chunk), slot)] = (buffer.data_ptr(), targets) |
| 800 | return targets |
| 801 | |
| 802 | def _complete_pull(self, pending: "_PendingPull") -> "dict[FetchKey, torch.Tensor]": |
| 803 | """Blocking half of a pull: the NIXL read lands during this ``ray.get``.""" |
| 804 | import ray |
| 805 | |
| 806 | ray.get(pending.ref) |
| 807 | return dict(zip(pending.keys, pending.targets)) |
| 808 | |
| 809 | # ---------------- Bake (dry run, at init) / replay ---------------- |
| 810 | |
| 811 | def _bake(self, init_info: ShardedRDTWeightTransferInitInfo) -> None: |
| 812 | """Bake the replay plan once, as a self-driven meta dry run. |
| 813 | |
| 814 | Puts the params on meta, then drives ``model.load_weights`` over |
| 815 | ``init_info.names`` through the model's ORIGINAL loaders (the stamps |
| 816 | bypass ``online_process_loader``, so ``_layerwise_process`` is never in |
| 817 | the path). Nothing materializes or pulls; the fake's ``copy_`` records |
| 818 | the source op chain and the meta destination's geometry. Afterwards one |
| 819 | scatter list per fully-loaded leaf module (copied numel == loadable |
| 820 | size) is indexed by source name; a partial or unrecordable module fails |
| 821 | the plan build. The model is restored. |
| 822 | |
| 823 | This leans on layerwise internals a public API should expose first-class: |
| 824 | a currently-loading hook instead of monkeypatched stamps, a dry-run mode |
| 825 | instead of bypassing ``online_process_loader``, and an |
| 826 | ``abort_layerwise_reload`` instead of ``_restore_after_dry_run``. |
| 827 | """ |
| 828 | from vllm.model_executor.model_loader.reload.layerwise import ( |
| 829 | initialize_layerwise_reload, |
| 830 | ) |
| 831 | from vllm.model_executor.model_loader.reload.utils import get_layer_size |
| 832 | |
| 833 | names, dtype_names, shapes = ( |
| 834 | init_info.names, |
| 835 | init_info.dtype_names, |
| 836 | init_info.shapes, |
| 837 | ) |
| 838 | self._name_meta = {n: (d, s) for n, d, s in zip(names, dtype_names, shapes)} |
| 839 | if not names: |
| 840 | return |
| 841 | |
| 842 | model = self.model |
| 843 | recorder = BakeSink() |
| 844 | |
| 845 | _t0 = time.perf_counter() |
| 846 | with torch.device(self.device): |
| 847 | # Meta-restore params + save kernel tensors (we bypass the loader |
| 848 | # wrapping it installs, below). |
| 849 | initialize_layerwise_reload(model) |
| 850 | # Stamp the *original* loaders (bypassing online_process_loader), so |
| 851 | # the single load pass runs the loaders on meta and records via the |
| 852 | # fake's copy_ — with no inline _layerwise_process, no deferral. |
| 853 | self._install_recording_stamps(model, recorder) |
| 854 | model.load_weights(self._build_fake_weights(names, recorder, self.device)) |
| 855 | # Keep only fully-loaded modules (copied numel >= loadable size, the |
| 856 | # test online_process_loader uses): a partial module leaves unwritten |
| 857 | # regions that finalize would init, so baking it scatters garbage. |
| 858 | for module, recorded in recorder.copies_by_layer.items(): |
| 859 | if not recorded or any(c is None for c in recorded): |
| 860 | continue # unrecordable copy_ -> slow path |
| 861 | # Guard above guarantees every entry is a real _Scatter. |
| 862 | copies = cast("list[_Scatter]", recorded) |
| 863 | copied = sum(prod(c.shape) for c in copies) |
| 864 | if copied < get_layer_size(module): |
| 865 | continue # partial -> slow path |
| 866 | for c in copies: |
| 867 | # Every name of the module shares ONE list, so identity |
| 868 | # dedups them back to one module at plan time. |
| 869 | self._name_to_plan[c.src[0]] = copies |
| 870 | self._restore_after_dry_run(model) |
| 871 | |
| 872 | # Names whose copy_ fired during the bake. Names not in here no-op for |
| 873 | # this worker (e.g. foreign-EP experts) and are skipped; live names that |
| 874 | # did not bake fail the plan build. |
| 875 | self._live_names = set(recorder.copied_names) |
| 876 | |
| 877 | n_groups = len({id(g) for g in self._name_to_plan.values()}) |
| 878 | logger.info( |
| 879 | "Sharded RDT dry-run baked %d/%d names into %d leaf modules " |
| 880 | "(%d live) in %.3fs", |
| 881 | len(self._name_to_plan), |
| 882 | len(names), |
| 883 | n_groups, |
| 884 | len(self._live_names), |
| 885 | time.perf_counter() - _t0, |
| 886 | ) |
| 887 | |
| 888 | def _install_recording_stamps( |
| 889 | self, model: torch.nn.Module, recorder: "BakeSink" |
| 890 | ) -> None: |
| 891 | """Wrap each loadable param's ``weight_loader`` to stamp |
| 892 | ``recorder.current = (leaf_module, param_name)`` before delegating to the |
| 893 | original loader, so the fake's ``copy_`` can attribute each recorded copy. |
| 894 | ``functools.wraps`` keeps the loader's real signature (so vLLM's |
| 895 | ``_layerwise_process`` ``param`` redirect still works if a stamp leaks), |
| 896 | and ``_rdt_stamp_inner`` tags it so ``_restore_after_dry_run`` can unwrap it. |
| 897 | """ |
| 898 | import functools |
| 899 | |
| 900 | from vllm.model_executor.model_loader.reload.layerwise import ( |
| 901 | _get_original_loader, |
| 902 | ) |
| 903 | from vllm.model_executor.model_loader.reload.utils import get_layer_tensors |
| 904 | from vllm.model_executor.model_loader.weight_utils import default_weight_loader |
| 905 | |
| 906 | def _make_stamp(layer, name, inner, added=False): |
| 907 | @functools.wraps(inner) # keep ``inner``'s signature (incl. ``param``) |
| 908 | def stamp(*args, **kwargs): |
| 909 | recorder.current = (layer, name) |
| 910 | try: |
| 911 | return inner(*args, **kwargs) |
| 912 | finally: |
| 913 | recorder.current = None |
| 914 | |
| 915 | # Tag so _restore_after_dry_run can detect and unwrap leaked stamps, |
| 916 | # and so a second bake doesn't double-wrap. |
| 917 | stamp._rdt_stamp_inner = inner # type: ignore[attr-defined] |
| 918 | stamp._rdt_stamp_added = added # type: ignore[attr-defined] |
| 919 | return stamp |
| 920 | |
| 921 | for module in model.modules(): |
| 922 | for name, tensor in get_layer_tensors(module).items(): |
| 923 | if getattr(tensor, "weight_loader", None) is None: |
| 924 | # A param with NO loader (e.g. GLM's router bias, a plain |
| 925 | # nn.Parameter) is still loaded by the model's load_weights |
| 926 | # through the getattr(param, "weight_loader", |
| 927 | # default_weight_loader) fallback — unstamped, its bake copy |
| 928 | # would be unattributable, failing the module's coverage |
| 929 | # gate and the plan build. Stamp the same default loader |
| 930 | # the fallback would pick; the restore deletes it again. |
| 931 | tensor.weight_loader = _make_stamp( |
| 932 | module, name, default_weight_loader, added=True |
| 933 | ) |
| 934 | continue |
| 935 | # Bypass online_process_loader: stamp the *original* loader. |
| 936 | original = _get_original_loader(tensor) |
| 937 | tensor.weight_loader = _make_stamp(module, name, original) |
| 938 | |
| 939 | def _restore_after_dry_run(self, model: torch.nn.Module) -> None: |
| 940 | """Restore each layerwise layer's saved kernel tensors without pulling |
| 941 | (a real ``finalize_layerwise_reload`` would materialize/load) and reset |
| 942 | its info. Also unwrap any recording ``stamp`` left on the params, since a |
| 943 | leaked stamp would sit under the next sync's ``online_process_loader`` and |
| 944 | silently break ``_layerwise_process``'s ``param`` redirect. |
| 945 | """ |
| 946 | from vllm.model_executor.model_loader.reload.layerwise import ( |
| 947 | LAYERWISE_INFO, |
| 948 | _place_kernel_tensors, |
| 949 | ) |
| 950 | from vllm.model_executor.model_loader.reload.utils import get_layer_tensors |
| 951 | |
| 952 | for layer in model.modules(): |
| 953 | info = LAYERWISE_INFO.get(layer) |
| 954 | if info is not None and info.can_load(): |
| 955 | if info.kernel_tensors is not None: |
| 956 | _place_kernel_tensors(layer, info) |
| 957 | info.reset() |
| 958 | # Unwrap any recording stamps left on the (now-restored) params so they |
| 959 | # never leak into a later update_weights. ``_rdt_stamp_inner`` is set by |
| 960 | # ``_install_recording_stamps``; unwrap repeatedly in case of nesting. |
| 961 | for module in model.modules(): |
| 962 | for _name, tensor in get_layer_tensors(module).items(): |
| 963 | loader = getattr(tensor, "weight_loader", None) |
| 964 | added = False |
| 965 | while loader is not None and hasattr(loader, "_rdt_stamp_inner"): |
| 966 | added = added or getattr(loader, "_rdt_stamp_added", False) |
| 967 | loader = loader._rdt_stamp_inner |
| 968 | tensor.weight_loader = loader |
| 969 | if added: |
| 970 | # The stamp was ATTACHED to a param that had no loader |
| 971 | # (see _install_recording_stamps); leave none behind. |
| 972 | del tensor.weight_loader |
| 973 | if hasattr(model, "_original_do_torchao_reload"): |
| 974 | model._do_torchao_reload = model._original_do_torchao_reload |
| 975 | |
| 976 | # ---------------- Background post-processing (pull/process pipeline) ------- |
| 977 | |
| 978 | def _ensure_proc_worker(self) -> None: |
| 979 | """Lazily create the per-slot events, the background CUDA stream, the work |
| 980 | queue, and the single processing thread. Idempotent.""" |
| 981 | if self._proc_thread is not None: |
| 982 | return |
| 983 | import queue |
| 984 | import threading |
| 985 | |
| 986 | self._slot_read_done = [torch.cuda.Event() for _ in range(self._ring_depth)] |
| 987 | # Generation handshake for the events above: the RPC thread counts items |
| 988 | # queued per slot, this thread counts records, and a pull may synchronize() |
| 989 | # only once done[slot] has caught up with queued[slot]. Without it the |
| 990 | # synchronize can bind to a stale record and pass, letting the next RDMA |
| 991 | # overwrite a slot under a pending scatter. See the doc's "The slot |
| 992 | # generation handshake". |
| 993 | self._slot_queued = [0] * self._ring_depth |
| 994 | self._slot_done = [0] * self._ring_depth |
| 995 | self._slot_cv = threading.Condition() |
| 996 | self._proc_stream = torch.cuda.Stream(device=self.device) |
| 997 | self._proc_queue = queue.Queue() |
| 998 | self._proc_error = None |
| 999 | t = threading.Thread( |
| 1000 | target=self._proc_worker_loop, name="rdt-postprocess", daemon=True |
| 1001 | ) |
| 1002 | self._proc_thread = t |
| 1003 | t.start() |
| 1004 | self._quant_stream = torch.cuda.Stream(device=self.device) |
| 1005 | self._quant_queue = queue.Queue() |
| 1006 | qt = threading.Thread( |
| 1007 | target=self._quant_worker_loop, name="rdt-quant", daemon=True |
| 1008 | ) |
| 1009 | self._quant_thread = qt |
| 1010 | qt.start() |
| 1011 | |
| 1012 | def _proc_worker_loop(self) -> None: |
| 1013 | """Single persistent thread: run each queued item's process phase on the |
| 1014 | background stream. Exits on the ``None`` sentinel (shutdown). An item that |
| 1015 | raises is recorded in ``_proc_error`` and re-raised on the RPC thread / |
| 1016 | at drain, so a failed sync fails loudly rather than corrupting silently. |
| 1017 | """ |
| 1018 | torch.cuda.set_device(self.device) |
| 1019 | q = self._proc_queue |
| 1020 | assert q is not None |
| 1021 | while True: |
| 1022 | item = q.get() |
| 1023 | try: |
| 1024 | if item is None: |
| 1025 | return |
| 1026 | # _process_item publishes the slot generation internally (in a |
| 1027 | # finally around its scatter pass), so an error here still |
| 1028 | # unblocks the RPC thread's generation wait. |
| 1029 | self._process_item(item) |
| 1030 | except BaseException as e: # noqa: BLE001 - surfaced on the RPC thread |
| 1031 | self._proc_error = e |
| 1032 | logger.exception("RDT background post-processing failed") |
| 1033 | finally: |
| 1034 | q.task_done() |
| 1035 | |
| 1036 | def _mark_slot_done(self, slot: int) -> None: |
| 1037 | """Publish that a queued item's read-done event has been recorded (or the |
| 1038 | item failed) so a pull waiting to reuse ``slot`` can proceed to its |
| 1039 | CUDA-event synchronize.""" |
| 1040 | with self._slot_cv: |
| 1041 | self._slot_done[slot] += 1 |
| 1042 | self._slot_cv.notify_all() |
| 1043 | |
| 1044 | def _raise_proc_error(self) -> None: |
| 1045 | """Re-raise (once) any error captured by the background thread.""" |
| 1046 | if self._proc_error is not None: |
| 1047 | err = self._proc_error |
| 1048 | self._proc_error = None |
| 1049 | raise RuntimeError("RDT background post-processing failed") from err |
| 1050 | |
| 1051 | def drain_pending(self) -> None: |
| 1052 | """Block until the background thread has processed every queued item and |
| 1053 | its stream work is complete, then re-raise any error it hit. Called from |
| 1054 | the worker's ``finish_weight_update`` before ``finalize_layerwise_reload`` |
| 1055 | so every baked layer is fully loaded (and ``info.reset()``-ed) first.""" |
| 1056 | if self._proc_queue is not None: |
| 1057 | self._proc_queue.join() # every put() item task_done()'d |
| 1058 | # The scatter thread feeds the quant thread, so join it SECOND (all |
| 1059 | # completed-group batches have been put by now), then sync both streams |
| 1060 | # so finalize sees fully-materialized, quanted, reset layers. |
| 1061 | if self._quant_queue is not None: |
| 1062 | self._quant_queue.join() |
| 1063 | if self._proc_stream is not None: |
| 1064 | self._proc_stream.synchronize() |
| 1065 | if self._quant_stream is not None: |
| 1066 | self._quant_stream.synchronize() |
| 1067 | self._raise_proc_error() |
| 1068 | # Ensure every fired free_group signal has EXECUTED on |
| 1069 | # the producer before the sync ends: ``begin_sync`` resets the producer's |
| 1070 | # per-group signal counts, so a signal landing after the next sync |
| 1071 | # started would credit a group it does not belong to, over-crediting the |
| 1072 | # gather-lookahead backpressure (extra gather resident -> trainer OOM |
| 1073 | # risk). |
| 1074 | if self._pending_frees: |
| 1075 | import ray |
| 1076 | |
| 1077 | try: |
| 1078 | ray.get(self._pending_frees) |
| 1079 | finally: |
| 1080 | self._pending_frees.clear() |
| 1081 | # One sync iteration fully drained; buffers/registrations are at (or |
| 1082 | # nearer) high-water — the chunk pipeline may issue ahead from now on. |
| 1083 | self._completed_syncs += 1 |
| 1084 | |
| 1085 | def _dispatch_item(self, item: "_ProcItem") -> None: |
| 1086 | """Hand one chunk item to the background scatter thread. |
| 1087 | |
| 1088 | Counts the item against its slot BEFORE dispatch: the next pull into |
| 1089 | that slot must wait until the background thread has processed (and |
| 1090 | RECORDED the read-done event for) every item ever queued on it. |
| 1091 | """ |
| 1092 | with self._slot_cv: |
| 1093 | self._slot_queued[item.slot] += 1 |
| 1094 | assert self._proc_queue is not None |
| 1095 | self._proc_queue.put(item) |
| 1096 | |
| 1097 | def _chunk_module_scatters( |
| 1098 | self, modules: "list[list[_Scatter]]" |
| 1099 | ) -> "list[tuple[int, list[_Scatter]]]": |
| 1100 | """Cut the modules' copies into one chunk per distinct owner class |
| 1101 | present, ascending by class index, as ``(class_idx, scatters)`` pairs. |
| 1102 | |
| 1103 | A chunk is one packed pull, so every name in it must share a producer — |
| 1104 | which is exactly what an owner class is. The cut is a pure function of |
| 1105 | the bake and the ownership table. Copy order within a chunk is bake |
| 1106 | order, and a module's copies may span chunks (materialize/quant fire on |
| 1107 | its first/last chunk; see ``_build_call_plan``).""" |
| 1108 | assert self._router is not None |
| 1109 | by_class: dict[int, list[_Scatter]] = {} |
| 1110 | for copies in modules: |
| 1111 | for c in copies: |
| 1112 | by_class.setdefault(self._router.class_of(c.src[0]), []).append(c) |
| 1113 | return [(ci, by_class[ci]) for ci in sorted(by_class)] |
| 1114 | |
| 1115 | def _build_call_plan(self, names: list[str], group_lens: list[int]) -> "_CallPlan": |
| 1116 | """Build the STATIC plan for one whole-sync call. |
| 1117 | |
| 1118 | Pure — no pulls, no engine state touched — so the result is cached and |
| 1119 | reused every sync. Three passes: |
| 1120 | 1. Split ``names`` into gather groups, one chunk per owner class |
| 1121 | present in this worker's baked copies, recording each group's last |
| 1122 | chunk for its ``free_group`` signal (or ``pre_free`` when this worker |
| 1123 | pulls nothing from it). Every group is signaled exactly once by |
| 1124 | construction. The stream has no per-group call boundaries, so group |
| 1125 | L+1's first chunk issues while L's still stream. |
| 1126 | 2. Per leaf module, find its FIRST and LAST chunk — materialize on the |
| 1127 | first, quant/kernel/reset on the last, correct by construction |
| 1128 | instead of by runtime counters. |
| 1129 | 3. Assemble ``_Chunk``s: dedup keys and precompute the packed layout. |
| 1130 | """ |
| 1131 | router = self._router |
| 1132 | assert router is not None, "init_transfer_engine() must run before planning" |
| 1133 | |
| 1134 | # --- pass 1: gather groups -> per-owner-class scatter chunks ----------- |
| 1135 | raw_chunks: list[list[_Scatter]] = [] |
| 1136 | free_at: dict[int, list[int]] = {} # chunk idx -> groups to signal after it |
| 1137 | pre_free: list[int] = [] # groups with no chunk here (signal at start) |
| 1138 | unbaked: list[str] = [] |
| 1139 | pos = 0 |
| 1140 | for gi, glen in enumerate(group_lens): |
| 1141 | gnames = names[pos : pos + glen] |
| 1142 | pos += glen |
| 1143 | modules: list[list[_Scatter]] = [] |
| 1144 | seen: set[int] = set() |
| 1145 | for n in gnames: |
| 1146 | mod = self._name_to_plan.get(n) |
| 1147 | if mod is None: |
| 1148 | if n in self._live_names: |
| 1149 | unbaked.append(n) |
| 1150 | elif id(mod) not in seen: |
| 1151 | seen.add(id(mod)) |
| 1152 | modules.append(mod) |
| 1153 | if not modules: |
| 1154 | # Nothing to pull for this group on this worker; its owners |
| 1155 | # still published it, so signal it done at sync start (they |
| 1156 | # tolerate a signal that arrives before the publish). |
| 1157 | pre_free.append(gi) |
| 1158 | continue |
| 1159 | for _ci, scatters in self._chunk_module_scatters(modules): |
| 1160 | raw_chunks.append(scatters) |
| 1161 | free_at.setdefault(len(raw_chunks) - 1, []).append(gi) |
| 1162 | if unbaked: |
| 1163 | # A live name with no baked plan cannot be loaded: there is no |
| 1164 | # fallback, and its pull would target groups the pipeline has |
| 1165 | # already freed. Fail here, where the names are known. |
| 1166 | raise RuntimeError( |
| 1167 | f"Sharded RDT: {len(unbaked)} live name(s) did not bake, " |
| 1168 | f"first: {unbaked[:3]}. Fix the loader so these names bake " |
| 1169 | "(fully load their module in one pass), or exclude them from " |
| 1170 | "the synced set." |
| 1171 | ) |
| 1172 | |
| 1173 | # --- pass 2: per-module first/last chunk -> materialize/quant ---------- |
| 1174 | first_at: dict[int, int] = {} |
| 1175 | last_at: dict[int, int] = {} |
| 1176 | layer_by_id: dict[int, Any] = {} |
| 1177 | for ci, scatters in enumerate(raw_chunks): |
| 1178 | for sc in scatters: |
| 1179 | lid = id(sc.layer) |
| 1180 | layer_by_id[lid] = sc.layer |
| 1181 | first_at.setdefault(lid, ci) |
| 1182 | last_at[lid] = ci |
| 1183 | materialize_at: dict[int, list[Any]] = {} |
| 1184 | quant_at: dict[int, list[Any]] = {} |
| 1185 | for lid, ci in first_at.items(): |
| 1186 | materialize_at.setdefault(ci, []).append(layer_by_id[lid]) |
| 1187 | for lid, ci in last_at.items(): |
| 1188 | quant_at.setdefault(ci, []).append(layer_by_id[lid]) |
| 1189 | |
| 1190 | # --- pass 3: assemble _Chunks (dedup keys + precompute pack layout) ---- |
| 1191 | from collections import defaultdict |
| 1192 | |
| 1193 | chunks: list[_Chunk] = [] |
| 1194 | per_owner_seq: dict[int, int] = defaultdict(int) |
| 1195 | for ci, scatters in enumerate(raw_chunks): |
| 1196 | keys: list[FetchKey] = [] |
| 1197 | kmeta: dict[FetchKey, tuple[torch.dtype, tuple[int, ...]]] = {} |
| 1198 | for sc in scatters: |
| 1199 | if sc.src not in kmeta: |
| 1200 | kmeta[sc.src] = (sc.dtype, sc.shape) |
| 1201 | keys.append(sc.src) |
| 1202 | pack_layout: list[tuple[int, torch.dtype, int, tuple[int, ...]]] = [] |
| 1203 | cur = 0 |
| 1204 | for k in keys: |
| 1205 | dt, shape = kmeta[k] |
| 1206 | numel = prod(shape) or 1 # type: ignore[arg-type] |
| 1207 | off = (cur + 15) & ~15 |
| 1208 | pack_layout.append((off, dt, numel, shape)) |
| 1209 | cur = off + numel * dt.itemsize |
| 1210 | owner = router.producer_for(router.consumer_id, scatters[0].src[0]) |
| 1211 | chunks.append( |
| 1212 | _Chunk( |
| 1213 | seq=per_owner_seq[owner], |
| 1214 | scatters=scatters, |
| 1215 | keys=keys, |
| 1216 | pack_layout=pack_layout, |
| 1217 | pack_bytes=cur, |
| 1218 | materialize=materialize_at.get(ci, []), |
| 1219 | quant=quant_at.get(ci, []), |
| 1220 | free=free_at.get(ci, []), |
| 1221 | # Every name of a chunk shares an owner class and a group, |
| 1222 | # so any of them resolves the same producer. |
| 1223 | owner=owner, |
| 1224 | ) |
| 1225 | ) |
| 1226 | per_owner_seq[owner] += 1 |
| 1227 | return _CallPlan(chunks=chunks, pre_free=pre_free) |
| 1228 | |
| 1229 | def _signal_group_done(self, group_idx: int) -> None: |
| 1230 | """Fire-and-forget ``free_group`` signal at EVERY owner of the group. |
| 1231 | |
| 1232 | The per-group barrier: each owner counts one signal per live consumer |
| 1233 | and frees the group (releasing its lookahead credit) on the last one, so |
| 1234 | every owner must hear from every consumer — including owners this worker |
| 1235 | pulls nothing from. Refs are held and drained in ``drain_pending`` so |
| 1236 | every signal has EXECUTED before the sync ends: ``begin_sync`` resets |
| 1237 | the counters, and a straggler landing in the next sync would credit a |
| 1238 | group it does not belong to.""" |
| 1239 | assert self._router is not None |
| 1240 | self._pending_frees.extend(self._router.free_group(group_idx)) |
| 1241 | |
| 1242 | def _run_chunk_pipeline(self, plan: "_CallPlan") -> None: |
| 1243 | """Pipelined chunk pulls over the ring of receive slots. |
| 1244 | |
| 1245 | Issues up to ``ring_depth`` produce RPCs ahead of the blocking gets, so |
| 1246 | while chunk i's RDMA streams the producer serves i+1 into its own ring |
| 1247 | slot and the background thread scatters i-1 out of another. Reads stay |
| 1248 | serialized on the shared NIC — the bandwidth floor, not a loss. |
| 1249 | |
| 1250 | Slot safety rests on two arguments spelled out in the doc: the producer's |
| 1251 | ring is no shallower than this one and drain-before-issue orders its |
| 1252 | reuse; the consumer's slots are held by ``_issue_pull``'s generation |
| 1253 | handshake. |
| 1254 | """ |
| 1255 | from collections import deque |
| 1256 | |
| 1257 | inflight: deque[tuple[_PendingPull, _Chunk]] = deque() |
| 1258 | |
| 1259 | # Gather groups with no chunk on this worker (plan.pre_free): signal |
| 1260 | # them done before the pipeline — their owners still published them, and |
| 1261 | # a signal that arrives before its publish is counted (fire-and-forget). |
| 1262 | for gi in plan.pre_free: |
| 1263 | self._signal_group_done(gi) |
| 1264 | |
| 1265 | def drain_one() -> None: |
| 1266 | pending, chunk = inflight.popleft() |
| 1267 | results = self._complete_pull(pending) |
| 1268 | self._dispatch_item( |
| 1269 | _ProcItem(chunk=chunk, results=results, slot=pending.slot) |
| 1270 | ) |
| 1271 | # Each gather group whose LAST chunk this is: its read is done, so |
| 1272 | # this consumer is finished with the group -> signal every owner |
| 1273 | # (fire-and-forget, off the critical path). |
| 1274 | for gi in chunk.free: |
| 1275 | self._signal_group_done(gi) |
| 1276 | |
| 1277 | for chunk in plan.chunks: |
| 1278 | if not chunk.keys: |
| 1279 | # _chunk_module_scatters never emits empty chunks; keep the |
| 1280 | # signals safe anyway if this ever changes. |
| 1281 | for gi in chunk.free: |
| 1282 | self._signal_group_done(gi) |
| 1283 | continue |
| 1284 | # Drain BEFORE issue once the ring is full: frees this chunk's slot |
| 1285 | # (generation-wise) and guarantees the producer-slot invariant above. |
| 1286 | # Sync 0 runs SERIAL (max 1 in flight): both sides still grow and |
| 1287 | # register buffers, and a producer registration mid-flight churns its |
| 1288 | # agent-metadata version under an in-flight pull (see __init__). |
| 1289 | depth = 1 if self._completed_syncs == 0 else self._ring_depth |
| 1290 | if len(inflight) >= depth: |
| 1291 | drain_one() |
| 1292 | slot = self._pull_slot |
| 1293 | self._pull_slot = (slot + 1) % self._ring_depth |
| 1294 | inflight.append((self._issue_pull(chunk, slot), chunk)) |
| 1295 | while inflight: |
| 1296 | drain_one() |
| 1297 | |
| 1298 | def _process_item(self, item: "_ProcItem") -> None: |
| 1299 | """Scatter-thread half: materialize this chunk's first-seen modules, |
| 1300 | scatter its slices on the process stream, publish the slot, then hand the |
| 1301 | modules it COMPLETES to the quant thread. |
| 1302 | |
| 1303 | Mirrors ``_layerwise_process`` minus the loader replay. Once every scatter |
| 1304 | reading ``item.slot`` is enqueued, records the slot's read-done event so |
| 1305 | the RPC thread can block on it before overwriting the slot. |
| 1306 | """ |
| 1307 | from vllm.model_executor.model_loader.reload.layerwise import ( |
| 1308 | LAYERWISE_INFO, |
| 1309 | ) |
| 1310 | from vllm.model_executor.model_loader.reload.meta import materialize_layer |
| 1311 | |
| 1312 | results = item.results |
| 1313 | chunk = item.chunk |
| 1314 | with ( |
| 1315 | torch.cuda.device(self.device), |
| 1316 | torch.cuda.stream(self._proc_stream), |
| 1317 | torch.device(self.device), |
| 1318 | ): |
| 1319 | # PASS 1 — slot readers: materialize the modules whose FIRST scatter |
| 1320 | # is in this chunk (empty HF params, once per module by construction) |
| 1321 | # then scatter this chunk's copies. The scatter copies are the ONLY |
| 1322 | # reads of the receive buffer; quant and kernel-copy operate on the |
| 1323 | # scattered params. Releasing the slot right after the scatters lets |
| 1324 | # the NEXT chunk's RDMA overwrite the buffer while quant still runs. |
| 1325 | try: |
| 1326 | for layer in chunk.materialize: |
| 1327 | info = LAYERWISE_INFO.get(layer) |
| 1328 | if info is None or not info.can_load(): |
| 1329 | raise RuntimeError( |
| 1330 | f"Baked replay: layer {type(layer).__name__} " |
| 1331 | "was not set up for reload this sync " |
| 1332 | "(start_weight_update must run before " |
| 1333 | "update_weights)." |
| 1334 | ) |
| 1335 | materialize_layer(layer, info) |
| 1336 | for sc in chunk.scatters: |
| 1337 | param = getattr(sc.layer, sc.param_name) |
| 1338 | # Through .data, exactly like default_weight_loader: params |
| 1339 | # that keep requires_grad=True (plain nn.Parameters like |
| 1340 | # GLM's router bias) reject an in-place copy through an |
| 1341 | # autograd view of the leaf. |
| 1342 | dst = param.data.as_strided(sc.shape, sc.stride, sc.offset) |
| 1343 | with torch._C.DisableTorchFunctionSubclass(): |
| 1344 | dst.copy_(results[sc.src]) |
| 1345 | # All reads of this slot's buffer are now enqueued on the process |
| 1346 | # stream; record + publish so the RPC thread can reuse the slot. |
| 1347 | self._slot_read_done[item.slot].record(self._proc_stream) |
| 1348 | finally: |
| 1349 | # Publish even on error so the RPC thread's generation wait |
| 1350 | # unblocks (the failure itself surfaces via _raise_proc_error). |
| 1351 | self._mark_slot_done(item.slot) |
| 1352 | |
| 1353 | # PASS 2 — param readers only: quant / kernel-copy / reset for the |
| 1354 | # modules whose LAST scatter is in this chunk. Handed to the |
| 1355 | # DEDICATED quant thread (own CUDA stream, event-chained after this |
| 1356 | # chunk's scatters) so it never delays this thread's next pass-1. |
| 1357 | # Running it in order here stalls the RPC thread's slot handshake by |
| 1358 | # ~0.5-0.75s/iter, since every group's quant delays the next publish. |
| 1359 | if chunk.quant: |
| 1360 | ready = torch.cuda.Event() |
| 1361 | ready.record(self._proc_stream) |
| 1362 | if self._quant_queue is None: |
| 1363 | self._run_quant(chunk.quant, ready) |
| 1364 | else: |
| 1365 | self._quant_queue.put((chunk.quant, ready)) |
| 1366 | |
| 1367 | def _run_quant(self, layers: "list[Any]", ready: "torch.cuda.Event") -> None: |
| 1368 | """Quant/kernel-copy/reset the given COMPLETED leaf modules, exactly as |
| 1369 | _layerwise_process. Runs on the quant thread's own stream, ordered after |
| 1370 | the modules' scatters via ``ready``; touches only the scattered params |
| 1371 | (never a receive slot), so it can overlap subsequent chunks' RDMA and |
| 1372 | scatters. ``info.reset()`` is what makes finalize skip the layer — |
| 1373 | drain_pending joins the quant queue before finalize runs.""" |
| 1374 | from vllm.model_executor.layers.quantization.base_config import ( |
| 1375 | QuantizeMethodBase, |
| 1376 | ) |
| 1377 | from vllm.model_executor.model_loader.reload.layerwise import ( |
| 1378 | LAYERWISE_INFO, |
| 1379 | _copy_and_restore_kernel_tensors, |
| 1380 | ) |
| 1381 | |
| 1382 | stream = self._quant_stream or self._proc_stream |
| 1383 | assert stream is not None # created in _ensure_proc_worker before use |
| 1384 | with ( |
| 1385 | torch.cuda.device(self.device), |
| 1386 | torch.cuda.stream(stream), |
| 1387 | torch.device(self.device), |
| 1388 | ): |
| 1389 | stream.wait_event(ready) |
| 1390 | for layer in layers: |
| 1391 | info = LAYERWISE_INFO.get(layer) |
| 1392 | assert info is not None # completed leaf module is set up for reload |
| 1393 | quant_method = getattr(layer, "quant_method", None) |
| 1394 | if isinstance(quant_method, QuantizeMethodBase): |
| 1395 | if hasattr(layer, "_already_called_process_weights_after_loading"): |
| 1396 | delattr(layer, "_already_called_process_weights_after_loading") |
| 1397 | quant_method.process_weights_after_loading(layer) |
| 1398 | # Copy into persistent kernel storage (preserves cudagraph refs). |
| 1399 | if info.kernel_tensors is not None: |
| 1400 | _copy_and_restore_kernel_tensors(layer, info) |
| 1401 | # Reset so finalize_layerwise_reload skips this (loaded) layer. |
| 1402 | info.reset() |
| 1403 | |
| 1404 | def _quant_worker_loop(self) -> None: |
| 1405 | """Dedicated quant thread: drains (completed_modules, scatter-done event) |
| 1406 | batches. Errors surface via _proc_error like the scatter thread's.""" |
| 1407 | torch.cuda.set_device(self.device) |
| 1408 | q = self._quant_queue |
| 1409 | assert q is not None |
| 1410 | while True: |
| 1411 | batch = q.get() |
| 1412 | try: |
| 1413 | if batch is None: |
| 1414 | return |
| 1415 | self._run_quant(*batch) |
| 1416 | except BaseException as e: # noqa: BLE001 - surfaced on the RPC thread |
| 1417 | self._proc_error = e |
| 1418 | logger.exception("RDT quant thread failed") |
| 1419 | finally: |
| 1420 | q.task_done() |
| 1421 | |
| 1422 | def shutdown(self) -> None: |
| 1423 | # Stop the background post-processing thread (drain, then sentinel + join) |
| 1424 | # before dropping the state it touches. |
| 1425 | if self._proc_thread is not None: |
| 1426 | try: |
| 1427 | self.drain_pending() |
| 1428 | except Exception: |
| 1429 | logger.exception("RDT drain during shutdown failed") |
| 1430 | assert self._proc_queue is not None |
| 1431 | self._proc_queue.put(None) # sentinel |
| 1432 | self._proc_thread.join(timeout=30) |
| 1433 | self._proc_thread = None |
| 1434 | self._proc_queue = None |
| 1435 | self._proc_stream = None |
| 1436 | if self._quant_thread is not None: |
| 1437 | assert self._quant_queue is not None |
| 1438 | self._quant_queue.put(None) # sentinel |
| 1439 | self._quant_thread.join(timeout=30) |
| 1440 | self._quant_thread = None |
| 1441 | self._quant_queue = None |
| 1442 | self._quant_stream = None |
| 1443 | self._slot_read_done = [] |
| 1444 | self._router = None |
| 1445 | # Drop strong references to baked modules so the model can be freed. |
| 1446 | self._name_to_plan.clear() |
| 1447 | self._name_meta.clear() |
| 1448 | self._live_names.clear() |
| 1449 | # Drop the cached plan (holds _Scatter refs to the baked layers). |
| 1450 | self._cached_plan = None |
| 1451 | # Release the receive buffers (their NIXL registration is pinned for the |
| 1452 | # process lifetime; freeing the tensors just drops our strong refs). |
| 1453 | self._dest_buffers = [{} for _ in range(self._ring_depth)] |
| 1454 | |
| 1455 | |
| 1456 | def _plan_digest(keys_per_chunk: list) -> str: |
| 1457 | """Digest of the chunks one consumer pulls from one producer, in pull order. |
| 1458 | |
| 1459 | Two consumers a producer serves out of ONE shared serve ring must agree on |
| 1460 | this, since sharing rests on their plans being identical; the producer |
| 1461 | compares it at init. Over the whole (name, op-chain) list rather than the |
| 1462 | names, because the chains decide the bytes each pull returns. |
| 1463 | """ |
| 1464 | import hashlib |
| 1465 | |
| 1466 | h = hashlib.sha256() |
| 1467 | for keys in keys_per_chunk: |
| 1468 | h.update(repr(keys).encode()) |
| 1469 | h.update(b"\n") |
| 1470 | return h.hexdigest()[:32] |
| 1471 | |
| 1472 | |
| 1473 | def _dtype_from_name(name: str) -> torch.dtype: |
| 1474 | """Resolve a string like 'bfloat16' to torch.bfloat16.""" |
| 1475 | dtype = getattr(torch, name, None) |
| 1476 | if not isinstance(dtype, torch.dtype): |
| 1477 | raise ValueError(f"Unknown torch dtype name: {name!r}") |
| 1478 | return dtype |