| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | """Trainer-side engine for the sharded-RDT (pull-based NIXL) backend. |
| 4 | |
| 5 | RDT is pull-based, so unlike NCCL this engine broadcasts nothing. It owns a |
| 6 | per-rank producer server (an internal Ray actor exposing the NIXL serve surface |
| 7 | the worker engine dials by name), and on each ``send_weights`` gathers this |
| 8 | rank's weights group-by-group from the ``WeightSource``, shares each group into |
| 9 | the server over CUDA IPC, and — on the sender — drives the inference-side |
| 10 | start/update/finish handshake, whose single empty ``update_weights`` unblocks the |
| 11 | workers to pull. |
| 12 | |
| 13 | All serve-side state lives on the server actor, so trainer processes need no |
| 14 | mixin, no named actors and no special actor options. |
| 15 | |
| 16 | See docs/training/weight_transfer/sharded_rdt.md for the publish -> serve -> |
| 17 | free_group -> release lifecycle and the ownership model. |
| 18 | """ |
| 19 | |
| 20 | import contextlib |
| 21 | import threading |
| 22 | import time |
| 23 | import uuid |
| 24 | from collections.abc import Callable, Collection |
| 25 | from concurrent.futures import ThreadPoolExecutor |
| 26 | from dataclasses import asdict, dataclass, field |
| 27 | from typing import Any, ClassVar |
| 28 | |
| 29 | import ray |
| 30 | import torch |
| 31 | from torch.multiprocessing.reductions import ( |
| 32 | StorageWeakRef, |
| 33 | rebuild_cuda_tensor, |
| 34 | reduce_tensor, |
| 35 | ) |
| 36 | from typing_extensions import Self |
| 37 | |
| 38 | from vllm.distributed.weight_transfer.base import ( |
| 39 | ParamMeta, |
| 40 | TrainerInitInfo, |
| 41 | TrainerWeightTransferEngine, |
| 42 | VLLMWeightSyncClient, |
| 43 | WeightSource, |
| 44 | layerwise_groups, |
| 45 | ) |
| 46 | from vllm.distributed.weight_transfer.sharded_rdt_common import ( |
| 47 | ALLOWED_OPS, |
| 48 | buffer_alloc_bytes, |
| 49 | check_ray_rdt_version, |
| 50 | register_nixl_memory, |
| 51 | ) |
| 52 | from vllm.logger import init_logger |
| 53 | |
| 54 | logger = init_logger(__name__) |
| 55 | |
| 56 | # Gathered-but-unfreed groups the gather loop may run ahead by before it stops |
| 57 | # gathering. A gathered group is published immediately (publish_group never |
| 58 | # blocks), so at most ``lookahead + 1`` groups are resident on the trainer — the |
| 59 | # memory bound larger models size against, pinned by the tests. |
| 60 | # |
| 61 | # At 1, group N+1 is gathered and serveable while the consumers pull N, so the |
| 62 | # boundary's free-barrier latency hides behind live pulls. Raise it only if one |
| 63 | # group's gather is slower than its pulls. |
| 64 | DEFAULT_GATHER_LOOKAHEAD = 1 |
| 65 | |
| 66 | # Seconds with no publish, produce or free before the |
| 67 | # producer declares the sync dead. A consumer that dies mid-sync never signals |
| 68 | # ``free_group``, so the group is never released and the waits below block |
| 69 | # forever — which stops this rank iterating its WeightSource, a collective, and |
| 70 | # wedges every other trainer rank with no exception anywhere. This converts that |
| 71 | # into one real error. |
| 72 | # |
| 73 | # A liveness backstop, not a latency target: nominal gaps are sub-second, so this |
| 74 | # is ~100x margin. It rides the init info because the producer is a Ray actor and |
| 75 | # does not inherit the trainer's environment. |
| 76 | DEFAULT_STALL_TIMEOUT_S = 300.0 |
| 77 | # How often a blocked waiter re-checks the progress stamp. |
| 78 | _STALL_POLL_S = 5.0 |
| 79 | |
| 80 | # Tensors whose STORAGE is at most this are packed into a ring slot and exported |
| 81 | # as one CUDA-IPC handle; larger storages are exported directly. Batching the |
| 82 | # small ones is what keeps sync time steady -- a per-tensor handle for each of |
| 83 | # ~37k names dominated the variance -- and `_pack_group_for_export` is also the |
| 84 | # only publish path that orders the pack against the sidecar's IPC read, so the |
| 85 | # packing is not optional. |
| 86 | _EXPORT_RING_MAX_BYTES = 64 << 20 |
| 87 | |
| 88 | # The actor method the worker engine dials for the NIXL pull; fixed by contract. |
| 89 | PRODUCE_METHOD_NAME = "rdt_produce_weights_batched" |
| 90 | |
| 91 | |
| 92 | # Cross-deployment serve-slot sharing. |
| 93 | # |
| 94 | # Consumers whose ids differ by a multiple of ``workers_per_replica`` are the |
| 95 | # same worker of different inference deployments, so they share a baked plan and |
| 96 | # pull byte-identical chunks. A serve ring each would cost one full ring per |
| 97 | # deployment of the producer's GPU and repack the same bytes once per deployment. |
| 98 | # |
| 99 | # NIXL reads are one-sided, so R readers can read ONE registered slot; what the |
| 100 | # producer cannot see is when a reader has FINISHED. The release edge is |
| 101 | # therefore the consumer's ISSUE order, sent as ``seq``: the pipeline drains pull |
| 102 | # i before issuing i+K, so slot ``seq % K`` was last packed for ``seq - K``, |
| 103 | # whose read is over. A counter on THIS side would follow execution order |
| 104 | # instead, which is wrong -- Ray may start a consumer's K concurrent produce |
| 105 | # calls in any order, so a call that executes K-before another can be a pull |
| 106 | # still being read, and its slot gets repacked underneath the reader. Silent, |
| 107 | # showing up only as a logprob drift. |
| 108 | # |
| 109 | # Hence one rendezvous per generation, keyed by ``seq``: the group's live sharers |
| 110 | # meet there, the LAST to arrive packs, and all return that one blob. It carries |
| 111 | # the release edge too, since ``seq`` is packed only once every live sharer has |
| 112 | # arrived, so each has drained ``seq - K``. Mismatched plans would rendezvous on |
| 113 | # differing bytes, so they are rejected at init by ``reserve_serve_buffer``'s |
| 114 | # ``plan_digest``. A sharer that dies mid-sync stalls its group, exactly as a |
| 115 | # dead consumer already stalls the per-group free barrier; the next |
| 116 | # ``begin_sync`` narrows the live set. |
| 117 | # |
| 118 | # LIMITATION: the rendezvous synchronizes deployments that a pull-based transfer |
| 119 | # would otherwise leave independent -- per chunk every sharer waits for the |
| 120 | # slowest, so skew between their DP pause/resume brackets lands in every |
| 121 | # deployment's sync time. Nothing about P2P requires this. It is the price of |
| 122 | # serving out of GPU memory, which is too scarce for a ring per consumer, so |
| 123 | # slots must be reused, and reuse needs a release edge. |
| 124 | # |
| 125 | # TODO: explore staging packed chunks in HOST memory. Host buffers are cheap |
| 126 | # enough for a ring per consumer, which would drop the rendezvous, the |
| 127 | # cross-replica coupling and the slot-reuse hazard together. The cost is a D2H |
| 128 | # bounce on the serve path, which has been the bottleneck before, so it needs |
| 129 | # measuring rather than assuming. |
| 130 | |
| 131 | |
| 132 | @dataclass |
| 133 | class _SharedPack: |
| 134 | """One generation of one slot-sharing group: the rendezvous state for a |
| 135 | single chunk, identified by the consumers' issue index.""" |
| 136 | |
| 137 | group: int |
| 138 | slot: int |
| 139 | """``seq % ring_depth``. Fixed at creation, so slot reuse follows the |
| 140 | consumers' issue order rather than this side's execution order.""" |
| 141 | arrived: set[int] = field(default_factory=set) |
| 142 | blob: Any = None |
| 143 | packing: bool = False |
| 144 | done: bool = False |
| 145 | error: BaseException | None = None |
| 146 | |
| 147 | |
| 148 | @dataclass |
| 149 | class ShardedRDTTrainerInitInfo(TrainerInitInfo): |
| 150 | """Trainer init info for the sharded-RDT backend. |
| 151 | |
| 152 | Identical on every rank except ``rank`` (rank 0 is the sender). Carries only |
| 153 | the must-agree wire params; the sender forwards them verbatim onto the |
| 154 | worker-side init info so the two cannot drift. Server-actor names are |
| 155 | generated per rank and all-gathered by the engine, not supplied here. |
| 156 | """ |
| 157 | |
| 158 | backend: ClassVar[str] = "sharded_rdt" |
| 159 | |
| 160 | num_consumers: int |
| 161 | """Total inference-worker (consumer) count across the whole fleet |
| 162 | (DP*TP*PP*PCP), for the M:N block assignment / free ref-count.""" |
| 163 | workers_per_replica: int = 0 |
| 164 | """Consumers per inference DEPLOYMENT (``num_consumers // num_replicas``). |
| 165 | |
| 166 | Fixes the slot-sharing groups: consumers whose ids differ by a multiple of |
| 167 | this are the same worker of different deployments, so they bake identical |
| 168 | plans and pull byte-identical chunks, and ONE registered serve slot can |
| 169 | serve all of them. 0 disables sharing, and so does the single-deployment |
| 170 | value ``num_consumers``, which makes every group a singleton and the serve |
| 171 | path identical to the unshared one. |
| 172 | """ |
| 173 | trainer_actor_namespace: str | None = None |
| 174 | """Ray namespace the engine spawns its serve actors in. The inference |
| 175 | workers (which run in their own EngineCore subprocess with its own |
| 176 | ``ray.init``) resolve those actors by name, so this must be the namespace |
| 177 | they can see. Forwarded to the worker-side init info.""" |
| 178 | num_rdt_buffers: int = 2 |
| 179 | """Serve/receive ring depth K (must match the worker).""" |
| 180 | buffer_presize_gb: float = 0.0 |
| 181 | """Serve-buffer pre-size floor in GiB (avoids NIXL desc-cache churn).""" |
| 182 | gather_lookahead: int = DEFAULT_GATHER_LOOKAHEAD |
| 183 | """Gathered-but-unfreed groups the gather loop may run ahead by. Bounds |
| 184 | trainer-resident memory at ``gather_lookahead + 1`` groups.""" |
| 185 | stall_timeout_s: float = DEFAULT_STALL_TIMEOUT_S |
| 186 | """Seconds of no publish/serve/free progress before the producer fails the |
| 187 | sync (see ``DEFAULT_STALL_TIMEOUT_S``).""" |
| 188 | |
| 189 | |
| 190 | class _RDTProducerServer: |
| 191 | """Per-rank NIXL serve surface: an internal Ray actor sharing the trainer |
| 192 | rank's GPU over CUDA IPC. Holds the gather cache, per-consumer serve rings, |
| 193 | the per-group free barrier and the packed serve. The engine feeds it with |
| 194 | ``publish_group``; workers pull with ``rdt_produce_weights_batched`` and |
| 195 | signal with ``free_group``. |
| 196 | |
| 197 | A plain class — the engine wraps it with ``ray.remote(...)`` at spawn so the |
| 198 | actor options live in one place. |
| 199 | """ |
| 200 | |
| 201 | def __init__( |
| 202 | self, |
| 203 | *, |
| 204 | num_rdt_buffers: int, |
| 205 | buffer_presize_gb: float, |
| 206 | gather_lookahead: int, |
| 207 | served_names: list[str] | None = None, |
| 208 | stall_timeout_s: float = DEFAULT_STALL_TIMEOUT_S, |
| 209 | workers_per_replica: int = 0, |
| 210 | ) -> None: |
| 211 | import gc |
| 212 | |
| 213 | self._device_index = torch.accelerator.current_device_index() |
| 214 | # name -> rebuilt CUDA-IPC tensor (or view); guarded by _cache_cond. |
| 215 | self._cache: dict[str, torch.Tensor] = {} |
| 216 | self._cache_cond = threading.Condition() |
| 217 | self._gather_error: BaseException | None = None |
| 218 | |
| 219 | # Monotonic stamp of the last forward step of the |
| 220 | # publish -> serve -> free loop. The waits below poll it instead of |
| 221 | # blocking forever, so a dead consumer fails the sync with a real error. |
| 222 | self._stall_timeout = float(stall_timeout_s) |
| 223 | self._last_progress = time.monotonic() |
| 224 | |
| 225 | # Names this producer publishes, so a misrouted pull fails loudly |
| 226 | # instead of blocking forever in the cache wait. None = serve anything. |
| 227 | self._served_names = set(served_names) if served_names is not None else None |
| 228 | |
| 229 | # Every live consumer signals free_group(gi) at every |
| 230 | # owner, once per sync; the group frees when the count reaches |
| 231 | # begin_sync's live total — one uniform target, no routed per-producer |
| 232 | # ones. Signals may precede the publish, which completes them. |
| 233 | self._live_count = 1 |
| 234 | self._free_counts: dict[int, int] = {} |
| 235 | # gi -> the names this producer published for it (what release drops). |
| 236 | self._group_names: dict[int, list[str]] = {} |
| 237 | |
| 238 | # Published-but-unfreed groups, plus freed ones not |
| 239 | # yet handed back. The memory gate lives in the ENGINE's gather loop; this |
| 240 | # side only accounts — free_group moves a group from _inflight_groups to |
| 241 | # _freed_pending, which the engine collects via wait_freed / end_sync to |
| 242 | # drop its storage refs. _lookahead is engine-enforced, unused here. |
| 243 | self._lookahead = max(0, gather_lookahead) |
| 244 | self._inflight_groups: list[int] = [] |
| 245 | self._freed_pending: list[int] = [] |
| 246 | |
| 247 | # Ring of packed serve buffers, one ring per SLOT-SHARING GROUP (one |
| 248 | # consumer per group unless several deployments share); the slot within |
| 249 | # it is the consumer's ``seq % nring``. |
| 250 | self._nring = max(1, num_rdt_buffers) |
| 251 | self._serve_rings: dict[int, list[torch.Tensor | None]] = {} |
| 252 | self._serve_lock = threading.Lock() |
| 253 | # registerMem on a shared NIXL agent is not concurrency-safe; serialize. |
| 254 | self._reg_lock = threading.Lock() |
| 255 | self._buffer_presize = int(buffer_presize_gb * (1 << 30)) |
| 256 | self._serve_device = torch.device("cuda", self._device_index) |
| 257 | |
| 258 | # Slot-sharing state, all guarded by _cache_cond, so the rendezvous |
| 259 | # waits ride the stall watchdog like every other wait here. |
| 260 | # ``_share_width`` = consumers per deployment; 0 means every group is a |
| 261 | # singleton and this degenerates to the unshared serve path. |
| 262 | self._share_width = max(0, int(workers_per_replica)) |
| 263 | self._live_ids: set[int] | None = None |
| 264 | self._sharing_active = False |
| 265 | self._sharers: dict[int, frozenset] = {} # group -> live members |
| 266 | self._gens: dict[int, dict] = {} # group -> {seq: _SharedPack} |
| 267 | # group -> {cid: plan digest}. Init-time state, not cleared per sync: it |
| 268 | # is what proves the sharers of a group pull the same chunks. |
| 269 | self._plan_digests: dict[int, dict[int, str]] = {} |
| 270 | |
| 271 | # (sharing group, ring idx, packed layout) -> |
| 272 | # (buffer data_ptr, destination views). Keyed on the buffer pointer too, so |
| 273 | # a ring regrow invalidates rather than writing into a freed buffer. See |
| 274 | # the serve path for why the layout, not the spec names, is the key. |
| 275 | self._pack_dsts: dict[tuple, tuple[int, list[torch.Tensor]]] = {} |
| 276 | # [RDT-EXPORT-RING] handle -> rebuilt base. The trainer reuses buffers |
| 277 | # across groups, so the same handle arrives on many publishes; without |
| 278 | # this we reopen it every group. Cleared per sync so no IPC mapping into |
| 279 | # trainer memory outlives the sync that made it. |
| 280 | self._base_cache: dict[tuple, torch.Tensor] = {} |
| 281 | |
| 282 | # Freeze the static post-init object graph so gen-2 GC never stops the |
| 283 | # world mid-serve, which shows up as serve stragglers. |
| 284 | gc.collect() |
| 285 | gc.freeze() |
| 286 | |
| 287 | def ping(self) -> int: |
| 288 | return self._device_index |
| 289 | |
| 290 | # ---------------- stall watchdog ---------------- |
| 291 | |
| 292 | def _note_progress_locked(self) -> None: |
| 293 | """Record that the credit loop moved. Caller holds ``_cache_cond``.""" |
| 294 | self._last_progress = time.monotonic() |
| 295 | |
| 296 | def _wait_for(self, blocked: Callable[[], bool], what: str) -> None: |
| 297 | """``_cache_cond.wait()`` with a liveness bound. |
| 298 | |
| 299 | Waits while ``blocked()`` holds, returning early on a gather error. If |
| 300 | nothing on this producer progresses for ``_stall_timeout``, self-fires |
| 301 | ``set_gather_error``, which every waiter here already checks — so the rank |
| 302 | unwinds through one path and the driver gets a real exception. |
| 303 | |
| 304 | The progress stamp is global to the producer, not per-waiter: a merely |
| 305 | slow waiter is kept alive by its peers' progress. Caller holds |
| 306 | ``_cache_cond``. |
| 307 | """ |
| 308 | while blocked(): |
| 309 | if self._gather_error is not None: |
| 310 | return |
| 311 | self._cache_cond.wait(_STALL_POLL_S) |
| 312 | if not blocked() or self._gather_error is not None: |
| 313 | return |
| 314 | stalled = time.monotonic() - self._last_progress |
| 315 | if stalled >= self._stall_timeout: |
| 316 | msg = ( |
| 317 | f"RDT stall: no progress for {stalled:.0f}s while waiting for " |
| 318 | f"{what} (timeout {self._stall_timeout:.0f}s). A consumer most " |
| 319 | f"likely died mid-sync: {len(self._inflight_groups)} group(s) " |
| 320 | f"published and unfreed." |
| 321 | ) |
| 322 | logger.error("[rdt-stall] %s", msg) |
| 323 | # Same channel a gather failure uses: wakes every waiter here, and |
| 324 | # each raises rather than returning a half-served result. |
| 325 | self._gather_error = RuntimeError(msg) |
| 326 | self._cache_cond.notify_all() |
| 327 | return |
| 328 | |
| 329 | def warmup_nixl(self) -> None: |
| 330 | """Create this server's NIXL agent now, while the rank's GPU is quiet. |
| 331 | |
| 332 | Called at spawn, before the server-name all-gather, so no rank can be |
| 333 | spinning in a collective. Creating the agent lazily instead deadlocks on |
| 334 | EFA-class fabrics (see "warmup_nixl breaks a startup deadlock" in the |
| 335 | doc). The warmup buffer stays registered so the agent's CUDA-HMEM path |
| 336 | stays initialized. |
| 337 | """ |
| 338 | self._nixl_warmup_buf = torch.zeros(1 << 20, dtype=torch.uint8, device="cuda") |
| 339 | with self._reg_lock: |
| 340 | register_nixl_memory(self._nixl_warmup_buf) |
| 341 | |
| 342 | # ---------------- engine-facing (per sync) ---------------- |
| 343 | |
| 344 | def begin_sync( |
| 345 | self, live_count: int, live_consumer_ids: list | None = None |
| 346 | ) -> None: |
| 347 | """Reset per-sync free/backpressure state and set this sync's barrier |
| 348 | target. |
| 349 | |
| 350 | ``live_count`` is how many consumers take part in THIS sync. Required, no |
| 351 | default: a forgotten argument silently targeting 1 would free groups after |
| 352 | the FIRST signal while others still pull — use-after-free, not an error. |
| 353 | The driver awaits the previous sync's finish before the next begins, so |
| 354 | nothing is in flight; a straggler signal would otherwise credit the wrong |
| 355 | sync, which is why the consumer drains its signals before finishing. |
| 356 | |
| 357 | ``live_consumer_ids`` is that same live set enumerated, which is what |
| 358 | sizes each slot-sharing group's rendezvous; a count cannot, since a group |
| 359 | is a specific set of ids. ``None`` leaves every group a singleton. |
| 360 | |
| 361 | The packed-destination cache deliberately survives: the layout repeats |
| 362 | every sync. So does ``_plan_digests``, which is init-time state. |
| 363 | """ |
| 364 | with self._cache_cond: |
| 365 | self._gather_error = None |
| 366 | self._live_count = max(1, int(live_count)) |
| 367 | self._live_ids = ( |
| 368 | set(int(c) for c in live_consumer_ids) |
| 369 | if live_consumer_ids is not None |
| 370 | else None |
| 371 | ) |
| 372 | # Sharing is on only when the geometry shows several deployments; |
| 373 | # one deployment takes the singleton path. |
| 374 | self._sharing_active = ( |
| 375 | self._share_width > 0 |
| 376 | and self._live_ids is not None |
| 377 | and len(self._live_ids) > self._share_width |
| 378 | ) |
| 379 | self._sharers.clear() |
| 380 | self._gens.clear() |
| 381 | self._inflight_groups.clear() |
| 382 | self._freed_pending.clear() |
| 383 | self._free_counts.clear() |
| 384 | self._group_names.clear() |
| 385 | self._base_cache.clear() |
| 386 | self._note_progress_locked() |
| 387 | |
| 388 | def _release_group_locked(self, group_idx: int) -> None: |
| 389 | """Drop a freed group's cache entries and queue it for the engine |
| 390 | (whose gather-credit gate blocks in ``wait_freed`` on exactly this). |
| 391 | |
| 392 | Shared by the last ``free_group`` and by ``publish_group`` completing an |
| 393 | early-signaled group. Caller must hold ``_cache_cond``. |
| 394 | """ |
| 395 | for name in self._group_names.pop(group_idx, ()): |
| 396 | self._cache.pop(name, None) |
| 397 | self._free_counts.pop(group_idx, None) |
| 398 | if group_idx in self._inflight_groups: |
| 399 | self._inflight_groups.remove(group_idx) |
| 400 | self._freed_pending.append(group_idx) |
| 401 | |
| 402 | def publish_group(self, group_idx: int, entries: tuple) -> None: |
| 403 | """Rebuild one gather group's CUDA-IPC tensors into the serve cache. |
| 404 | |
| 405 | NEVER blocks: a gathered group is serveable immediately. The memory bound |
| 406 | lives in the engine's gather loop, which stops GATHERING (not publishing) |
| 407 | past ``gather_lookahead`` unfreed groups. |
| 408 | |
| 409 | ``entries`` is ``(storages, views)``: one CUDA-IPC export per storage, |
| 410 | plus per-name ``(sid, dtype_name, shape, stride, storage_offset)`` rebuilt |
| 411 | here as ``as_strided`` views. |
| 412 | |
| 413 | Signals can arrive BEFORE their publish — a consumer pulling nothing of a |
| 414 | group signals it as its plan starts — so a group whose barrier is already |
| 415 | satisfied is released here. Freed groups reach the engine only through |
| 416 | ``wait_freed`` / ``end_sync``: a freed notice riding an unharvested async |
| 417 | publish result would wedge the loop. |
| 418 | """ |
| 419 | storages, views = entries |
| 420 | bases: dict[int, torch.Tensor] = {} |
| 421 | for sid, reduce_args in storages.items(): |
| 422 | cached = self._base_cache.get(reduce_args) |
| 423 | if cached is not None: |
| 424 | bases[sid] = cached |
| 425 | continue |
| 426 | list_args = list(reduce_args) |
| 427 | # Index 6 of reduce_tensor's args is the exporter's device index; |
| 428 | # rebuild on this server's device (same physical GPU as the rank). |
| 429 | list_args[6] = self._device_index |
| 430 | bases[sid] = self._base_cache[reduce_args] = rebuild_cuda_tensor(*list_args) |
| 431 | rebuilt: dict[str, torch.Tensor] = {} |
| 432 | for name, (sid, dtype_name, shape, stride, storage_offset) in views.items(): |
| 433 | typed = bases[sid].view(getattr(torch, dtype_name)) |
| 434 | rebuilt[name] = torch.as_strided(typed, shape, stride, storage_offset) |
| 435 | del bases |
| 436 | |
| 437 | with self._cache_cond: |
| 438 | self._cache.update(rebuilt) |
| 439 | self._inflight_groups.append(group_idx) |
| 440 | self._group_names[group_idx] = list(rebuilt) |
| 441 | if self._free_counts.get(group_idx, 0) >= self._live_count: |
| 442 | self._release_group_locked(group_idx) |
| 443 | self._note_progress_locked() |
| 444 | self._cache_cond.notify_all() |
| 445 | |
| 446 | def wait_freed(self) -> list[int]: |
| 447 | """Block until at least one published group has been freed; return and |
| 448 | clear the freed backlog. The engine's gather-credit gate calls this once |
| 449 | its loop is ``gather_lookahead`` groups ahead, so this wait is where the |
| 450 | trainer paces to the consumers' pull rate. |
| 451 | |
| 452 | Raises rather than returning empty when the sync errors or stalls: the |
| 453 | engine is blocked here inside its gather loop, and an empty return would |
| 454 | spin it straight back in. |
| 455 | """ |
| 456 | with self._cache_cond: |
| 457 | self._wait_for(lambda: not self._freed_pending, "a freed-group credit") |
| 458 | if not self._freed_pending: |
| 459 | raise RuntimeError( |
| 460 | "gather errored while waiting for a freed-group credit: " |
| 461 | f"{self._gather_error!r}" |
| 462 | ) |
| 463 | freed = self._freed_pending |
| 464 | self._freed_pending = [] |
| 465 | return freed |
| 466 | |
| 467 | def end_sync(self) -> list[int]: |
| 468 | """Block until every published group has been freed by its consumers; |
| 469 | return the remaining freed keys so the engine drops its last refs.""" |
| 470 | with self._cache_cond: |
| 471 | self._wait_for( |
| 472 | lambda: bool(self._inflight_groups), |
| 473 | "every published group to be freed", |
| 474 | ) |
| 475 | freed = self._freed_pending |
| 476 | self._freed_pending = [] |
| 477 | self._base_cache.clear() |
| 478 | return freed |
| 479 | |
| 480 | def set_gather_error(self, message: str) -> None: |
| 481 | """Record a trainer-side gather failure so blocked serves / publishes |
| 482 | stop waiting and surface it.""" |
| 483 | with self._cache_cond: |
| 484 | self._gather_error = RuntimeError(message) |
| 485 | self._cache_cond.notify_all() |
| 486 | |
| 487 | # ---------------- consumer-facing (called by name over Ray) ---------------- |
| 488 | |
| 489 | def free_group(self, group_idx: int) -> None: |
| 490 | """Consumer back-edge: one consumer is done with group ``group_idx``, |
| 491 | either because its last chunk landed or because it had nothing to pull. |
| 492 | |
| 493 | The per-group barrier counts one signal per live consumer against the |
| 494 | ``begin_sync`` count; every consumer signals every owner, so the target is |
| 495 | the same integer everywhere. The last signal drops the cache entries and |
| 496 | queues the group as gather credit for ``wait_freed``. A signal arriving |
| 497 | before its publish is only counted — ``publish_group`` completes it. |
| 498 | """ |
| 499 | gi = int(group_idx) |
| 500 | with self._cache_cond: |
| 501 | count = self._free_counts.get(gi, 0) + 1 |
| 502 | self._free_counts[gi] = count |
| 503 | if gi in self._group_names and count >= self._live_count: |
| 504 | self._release_group_locked(gi) |
| 505 | self._note_progress_locked() |
| 506 | self._cache_cond.notify_all() |
| 507 | |
| 508 | def _new_serve_buffer(self, nbytes: int) -> torch.Tensor: |
| 509 | """Allocate + NIXL-register one serve slot: the single allocation seam, |
| 510 | so registration cannot be skipped on either of the two paths that make |
| 511 | buffers (the init-time reservation and the serve-path backstop).""" |
| 512 | t = torch.empty(nbytes, dtype=torch.uint8, device=self._serve_device) |
| 513 | with self._reg_lock: |
| 514 | register_nixl_memory(t) |
| 515 | return t |
| 516 | |
| 517 | def reserve_serve_buffer( |
| 518 | self, consumer_id: int, nbytes: int, plan_digest: str | None = None |
| 519 | ) -> None: |
| 520 | """Pre-allocate + NIXL-register this consumer's serve ring before any |
| 521 | pull, while the fabric is idle (avoids registration races during the |
| 522 | sync-0 RDMA churn under M:N fan-in). Idempotent; grows only if needed. |
| 523 | |
| 524 | The ring is keyed by SLOT-SHARING GROUP, so the sharers of one group |
| 525 | reserve ONE ring between them. They all call this, same group and same |
| 526 | size, so the whole body runs under ``_serve_lock``, or two concurrent |
| 527 | first-callers would each allocate a ring and one would be dropped while |
| 528 | still registered. |
| 529 | |
| 530 | ``plan_digest`` is the caller's ordered chunk plan for THIS producer. |
| 531 | Sharers must agree on it, and this is where a disagreement is caught |
| 532 | rather than at serve time, where it is a rendezvous nobody completes. |
| 533 | |
| 534 | Raises: |
| 535 | RuntimeError: two consumers of one sharing group disagree on the |
| 536 | chunks they pull from this producer. |
| 537 | |
| 538 | """ |
| 539 | sg = self._share_group(consumer_id) |
| 540 | if plan_digest is not None: |
| 541 | with self._cache_cond: |
| 542 | seen = self._plan_digests.setdefault(sg, {}) |
| 543 | clash = next( |
| 544 | ((c, d) for c, d in seen.items() if d != plan_digest), None |
| 545 | ) |
| 546 | if clash is None: |
| 547 | seen[consumer_id] = plan_digest |
| 548 | if clash is not None: |
| 549 | raise RuntimeError( |
| 550 | f"consumers {clash[0]} and {consumer_id} are in slot-sharing " |
| 551 | f"group {sg} but do not pull the same chunks from this " |
| 552 | f"producer ({clash[1]} != {plan_digest}). Sharing a serve " |
| 553 | f"slot requires the deployments to be identical; check that " |
| 554 | f"workers_per_replica matches the inference fleet's " |
| 555 | f"geometry." |
| 556 | ) |
| 557 | # This is where the producer's NIXL serve memory is committed: one ring |
| 558 | # of K slots per SHARING GROUP that pulls from this rank -- with a single |
| 559 | # deployment, one ring per consumer, since every group is a singleton. |
| 560 | # Routing only sends a consumer here for names this rank owns, so the |
| 561 | # count is the fan-in, not the fleet. |
| 562 | # |
| 563 | # Every slot is sized by the LARGEST chunk that consumer pulls from us, |
| 564 | # rounded up to a 256 MB multiple, so the whole ring pays for the worst |
| 565 | # chunk in the plan while most chunks are far smaller. |
| 566 | # |
| 567 | # TODO: that over-allocates. Sizing slots per chunk instead of per plan, |
| 568 | # or capping chunk bytes so no chunk is an outlier, would shrink it; so |
| 569 | # would dropping the 256 MB round-up, which is slack once a buffer never |
| 570 | # regrows. Row-splitting the vocab matrix shrinks the outlier itself. |
| 571 | # The host-staging TODO above would remove these rings altogether. |
| 572 | alloc = buffer_alloc_bytes(nbytes, self._buffer_presize) |
| 573 | with self._serve_lock: |
| 574 | rings = self._serve_rings.setdefault(sg, [None] * self._nring) |
| 575 | for i in range(self._nring): |
| 576 | slot = rings[i] |
| 577 | if slot is None or slot.numel() < alloc: |
| 578 | rings[i] = self._new_serve_buffer(alloc) |
| 579 | |
| 580 | def _share_group(self, consumer_id: int) -> int: |
| 581 | """The slot-sharing group ``consumer_id`` belongs to: its index within |
| 582 | its own deployment, since that is what fixes the plan. Without a width |
| 583 | the group is the consumer itself and nothing is shared.""" |
| 584 | return consumer_id % self._share_width if self._share_width > 0 else consumer_id |
| 585 | |
| 586 | def _sharers_of(self, sg: int) -> frozenset: |
| 587 | """The LIVE consumers of group ``sg``: the rendezvous width. Derived from |
| 588 | ``begin_sync``'s live set, so a degraded sync narrows the barrier instead |
| 589 | of waiting forever on a dead deployment. Caller holds ``_cache_cond``.""" |
| 590 | cached = self._sharers.get(sg) |
| 591 | if cached is None: |
| 592 | if not self._sharing_active or self._live_ids is None: |
| 593 | cached = frozenset((sg,)) |
| 594 | else: |
| 595 | cached = frozenset( |
| 596 | c for c in self._live_ids if c % self._share_width == sg |
| 597 | ) or frozenset((sg,)) |
| 598 | self._sharers[sg] = cached |
| 599 | return cached |
| 600 | |
| 601 | def _join_share_group(self, consumer_id: int, seq: int) -> tuple: |
| 602 | """Register this call's arrival in its group's rendezvous for ``seq`` and |
| 603 | return ``(generation, is_packer)``. |
| 604 | |
| 605 | Exactly one caller per generation gets ``is_packer=True``, whichever |
| 606 | completes the arrival set, and it is the only one that touches the GPU. |
| 607 | The slot is ``seq % nring``: fixed by the consumers' issue order, so it |
| 608 | needs no release accounting on this side. |
| 609 | """ |
| 610 | sg = self._share_group(consumer_id) |
| 611 | with self._cache_cond: |
| 612 | gens = self._gens.setdefault(sg, {}) |
| 613 | gen = gens.get(seq) |
| 614 | if gen is None: |
| 615 | gen = gens[seq] = _SharedPack(group=sg, slot=seq % self._nring) |
| 616 | gen.arrived.add(consumer_id) |
| 617 | if gen.packing or gen.done or not gen.arrived >= self._sharers_of(sg): |
| 618 | return gen, False |
| 619 | gen.packing = True |
| 620 | return gen, True |
| 621 | |
| 622 | def _await_shared_pack(self, gen: _SharedPack) -> torch.Tensor: |
| 623 | """Block until this generation's packer published its blob, and return |
| 624 | it. A failed pack is re-raised here, so every sharer of a bad pack fails |
| 625 | instead of reading a half-written slot.""" |
| 626 | with self._cache_cond: |
| 627 | self._wait_for( |
| 628 | lambda: not gen.done and gen.error is None, |
| 629 | "a co-replica's shared serve pack", |
| 630 | ) |
| 631 | if gen.error is not None: |
| 632 | raise RuntimeError( |
| 633 | f"the sharer packing this chunk failed: {gen.error!r}" |
| 634 | ) from gen.error |
| 635 | if not gen.done: |
| 636 | raise RuntimeError( |
| 637 | f"gather errored while awaiting a shared pack: " |
| 638 | f"{self._gather_error!r}" |
| 639 | ) |
| 640 | return gen.blob |
| 641 | |
| 642 | def _serve_slot(self, sg: int, idx: int, need: int) -> torch.Tensor: |
| 643 | """Group ``sg``'s ring slot ``idx``, grown if this chunk outgrew the |
| 644 | reservation. Growing registers memory while the fabric is busy, the |
| 645 | hazard ``reserve_serve_buffer`` exists to avoid, so it is a backstop: the |
| 646 | reservation is sized from the same static plan as this pack.""" |
| 647 | with self._serve_lock: |
| 648 | buffer = self._serve_rings.setdefault(sg, [None] * self._nring)[idx] |
| 649 | if buffer is not None and buffer.numel() >= need: |
| 650 | return buffer |
| 651 | buffer = self._new_serve_buffer(buffer_alloc_bytes(need, self._buffer_presize)) |
| 652 | with self._serve_lock: |
| 653 | self._serve_rings.setdefault(sg, [None] * self._nring)[idx] = buffer |
| 654 | return buffer |
| 655 | |
| 656 | def _fail_shared_pack(self, gen: _SharedPack, exc: BaseException) -> None: |
| 657 | """Publish a pack failure so waiting sharers raise instead of hanging on |
| 658 | a generation that will never complete.""" |
| 659 | with self._cache_cond: |
| 660 | gen.error = exc |
| 661 | self._cache_cond.notify_all() |
| 662 | |
| 663 | @ray.method(tensor_transport="nixl") |
| 664 | def rdt_produce_weights_batched( |
| 665 | self, specs: list, consumer_id: int = 0, seq: int = -1 |
| 666 | ): |
| 667 | """Serve one batched slice request over NIXL. |
| 668 | |
| 669 | Waits until the specs' names are cached, then rendezvouses with the |
| 670 | other live sharers of this consumer's slot-sharing group. The last to |
| 671 | arrive replays each spec's op chain (pure views into cached tensors, |
| 672 | guarded by ALLOWED_OPS) and byte-packs the slices 16B-aligned into the |
| 673 | group's ring slot ``seq % nring``, mirroring the consumer's identical |
| 674 | layout; every sharer then returns that one packed blob, so R deployments |
| 675 | cost one pack and one slot instead of R. With a single deployment every |
| 676 | group is a singleton, so the arriving call is always its own packer and |
| 677 | this is the unshared serve path exactly. |
| 678 | |
| 679 | ``seq`` is the caller's index in its pull stream to this producer, which |
| 680 | is what fixes the slot. Callers that want a single slice pass one spec |
| 681 | and read the blob back with that slice's dtype/shape. |
| 682 | |
| 683 | Raises: |
| 684 | ValueError: ``seq`` was not supplied. |
| 685 | |
| 686 | """ |
| 687 | needed = sorted({n for n, _ in specs}) |
| 688 | if self._served_names is not None: |
| 689 | unserved = [n for n in needed if n not in self._served_names] |
| 690 | if unserved: |
| 691 | # Without this the cache wait below would block forever: this |
| 692 | # producer never gathers these names, so they can never arrive. |
| 693 | raise RuntimeError( |
| 694 | f"pull routed to the wrong producer: {unserved[:3]} " |
| 695 | f"({len(unserved)} names) are not served here" |
| 696 | ) |
| 697 | with self._cache_cond: |
| 698 | self._wait_for( |
| 699 | lambda: not all(n in self._cache for n in needed), |
| 700 | f"{len(needed)} name(s) to be published", |
| 701 | ) |
| 702 | if not all(n in self._cache for n in needed): |
| 703 | # _wait_for only gives up on a gather error (its own stall |
| 704 | # included), so the names are never coming. |
| 705 | raise RuntimeError( |
| 706 | f"gather errored before {needed}: {self._gather_error!r}" |
| 707 | ) |
| 708 | |
| 709 | if seq < 0: |
| 710 | # The slot is derived from it, so a caller that does not send one |
| 711 | # cannot be served safely. |
| 712 | raise ValueError( |
| 713 | "rdt_produce_weights_batched requires the consumer's issue index `seq`" |
| 714 | ) |
| 715 | gen, is_packer = self._join_share_group(consumer_id, seq) |
| 716 | if not is_packer: |
| 717 | blob = self._await_shared_pack(gen) |
| 718 | else: |
| 719 | try: |
| 720 | blob = self._pack_shared(gen, specs) |
| 721 | except BaseException as e: |
| 722 | self._fail_shared_pack(gen, e) |
| 723 | raise |
| 724 | # A served pull is the third of the four progress signals the stall |
| 725 | # watchdog reads (publish / produce / free / begin_sync): a long sync whose |
| 726 | # consumers pull steadily but slowly must never trip it. |
| 727 | with self._cache_cond: |
| 728 | self._note_progress_locked() |
| 729 | return [blob] |
| 730 | |
| 731 | def _pack_shared(self, gen: _SharedPack, specs: list) -> torch.Tensor: |
| 732 | """Replay the op chains into this generation's slot and publish the blob |
| 733 | to the sharers waiting on it. |
| 734 | |
| 735 | Runs on exactly one call per generation and OUTSIDE ``_cache_cond``: the |
| 736 | pack is GPU work and must not block publishes, frees or other groups. |
| 737 | """ |
| 738 | sliced: list = [] # (byte_off, tensor) |
| 739 | pack_cur = 0 |
| 740 | for name, chain in specs: |
| 741 | t = self._cache[name] |
| 742 | for op, args, kw in chain: |
| 743 | if op not in ALLOWED_OPS: |
| 744 | raise ValueError(f"{name!r}: disallowed op {op!r}") |
| 745 | t = getattr(t, op)(*args, **dict(kw)) |
| 746 | off = (pack_cur + 15) & ~15 |
| 747 | pack_cur = off + t.numel() * t.element_size() |
| 748 | sliced.append((off, t)) |
| 749 | |
| 750 | buffer = self._serve_slot(gen.group, gen.slot, pack_cur) |
| 751 | |
| 752 | # The destination views are a pure function of the packed |
| 753 | # layout, which is byte-identical every sync, so build them once per |
| 754 | # (sharing group, ring slot, layout) and reuse — rebuilding per call cost |
| 755 | # 5.2ms of a 7.5ms 384-spec 235B group. |
| 756 | # |
| 757 | # The key is the LAYOUT, not the spec names: a name can appear in two |
| 758 | # requests with different op chains (one name's copies can split across |
| 759 | # owner-class chunks), and serving the second through the first's views |
| 760 | # would write the wrong bytes with nothing downstream to catch it. |
| 761 | dst_key = ( |
| 762 | gen.group, |
| 763 | gen.slot, |
| 764 | tuple((off, t.dtype, t.shape) for off, t in sliced), |
| 765 | ) |
| 766 | cached = self._pack_dsts.get(dst_key) |
| 767 | if cached is None or cached[0] != buffer.data_ptr(): |
| 768 | dsts = [] |
| 769 | for off, t in sliced: |
| 770 | nb = t.numel() * t.element_size() |
| 771 | dsts.append(buffer[off : off + nb].view(t.dtype).reshape(t.shape)) |
| 772 | self._pack_dsts[dst_key] = (buffer.data_ptr(), dsts) |
| 773 | else: |
| 774 | dsts = cached[1] |
| 775 | torch._foreach_copy_(dsts, [t for _off, t in sliced]) |
| 776 | |
| 777 | blob = buffer[:pack_cur] |
| 778 | with self._cache_cond: |
| 779 | gen.blob = blob |
| 780 | gen.done = True |
| 781 | self._note_progress_locked() |
| 782 | self._cache_cond.notify_all() |
| 783 | return blob |
| 784 | |
| 785 | def shutdown(self) -> None: |
| 786 | with self._cache_cond: |
| 787 | self._cache.clear() |
| 788 | with self._cache_cond: |
| 789 | # Generations hold blob views into the rings, so they go first. |
| 790 | self._gens.clear() |
| 791 | self._sharers.clear() |
| 792 | with self._serve_lock: |
| 793 | self._serve_rings.clear() |
| 794 | # Must go with the rings: these are views INTO them, and the |
| 795 | # data_ptr guard that normally invalidates them cannot tell a freed |
| 796 | # buffer from a new one recycled at the same address. |
| 797 | self._pack_dsts.clear() |
| 798 | |
| 799 | |
| 800 | class ShardedRDTTrainerWeightTransferEngine( |
| 801 | TrainerWeightTransferEngine[ShardedRDTTrainerInitInfo] |
| 802 | ): |
| 803 | """Trainer-side engine for the pull-based sharded-RDT backend. |
| 804 | |
| 805 | Lives on every trainer rank. Owns a per-rank `_RDTProducerServer` actor |
| 806 | (the NIXL serve surface). `send_weights` gathers this rank's weights |
| 807 | group-by-group from the `WeightSource`, shares each group into the server |
| 808 | over CUDA IPC, and — on the sender — drives the inference-side handshake so |
| 809 | the workers pull. Non-sender ranks only gather (staying in the collective). |
| 810 | """ |
| 811 | |
| 812 | init_info_cls = ShardedRDTTrainerInitInfo |
| 813 | |
| 814 | def __init__( |
| 815 | self, |
| 816 | *, |
| 817 | client: VLLMWeightSyncClient, |
| 818 | source: WeightSource, |
| 819 | is_sender: bool = True, |
| 820 | init_info: ShardedRDTTrainerInitInfo, |
| 821 | ) -> None: |
| 822 | super().__init__(client=client, source=source, is_sender=is_sender) |
| 823 | self._init_info = init_info |
| 824 | self._server: Any = None # the _RDTProducerServer actor handle |
| 825 | self._server_name: str | None = None |
| 826 | # Group-major metadata / partition, computed at trainer_init. |
| 827 | self._meta: list[ParamMeta] = [] |
| 828 | self._groups: list[list[str]] = [] |
| 829 | # Ownership resolved at trainer_init from the fleet's held names: the |
| 830 | # distinct owner sets, the per-name index into them, this rank's groups |
| 831 | # and its held set. Routing itself is consumer-side; the trainer only |
| 832 | # ships the table. |
| 833 | self._owner_sets: list[list[int]] = [] |
| 834 | self._name_owner_class: list[int] = [] |
| 835 | self._owned_idx: list[int] = [] |
| 836 | self._held_names: set[str] | None = None |
| 837 | # Strong refs to gathered tensors we've shared into the server, keyed by |
| 838 | # group index. CUDA-IPC exports must outlive the importer, so we hold |
| 839 | # them until the server reports the group freed. See send_weights. |
| 840 | self._inflight: dict[int, dict[str, torch.Tensor]] = {} |
| 841 | # [RDT-EXPORT-RING] Reusable packed slots, one per resident group. See |
| 842 | # _pack_group_for_export. Persistent across syncs: freeing and |
| 843 | # reallocating them each sync fragmented the caching allocator badly |
| 844 | # enough to grow reserved memory ~1.5 GB per sync. |
| 845 | self._export_ring: list[torch.Tensor | None] = [] |
| 846 | self._export_ring_args: list[tuple | None] = [] |
| 847 | # Which live group holds which slot, and the slots free to hand out. A |
| 848 | # modular counter is unsafe here: the credit gate bounds the NUMBER of |
| 849 | # unfreed groups, not their order, so it can hand a later group the slot |
| 850 | # a still-live group is being served from. |
| 851 | self._slot_of_group: dict[int, int] = {} |
| 852 | self._free_slots: list[int] = [] |
| 853 | # Set by trainer_init from the server-name all-gather; retained so a |
| 854 | # restarted consumer can be re-initialized without another collective. |
| 855 | self._server_names: list[str] | None = None |
| 856 | |
| 857 | def _rpc(self, method: str, *args: Any) -> Any: |
| 858 | """Call one of the server actor's methods and block for the result. |
| 859 | The single seam through which the engine talks to its server, so tests |
| 860 | can inject a local (non-Ray) fake server.""" |
| 861 | import ray |
| 862 | |
| 863 | return ray.get(getattr(self._server, method).remote(*args)) |
| 864 | |
| 865 | @staticmethod |
| 866 | def _contig_stride(shape) -> tuple: |
| 867 | stride: list[int] = [] |
| 868 | acc = 1 |
| 869 | for s in reversed(list(shape)): |
| 870 | stride.append(acc) |
| 871 | acc *= int(s) |
| 872 | return tuple(reversed(stride)) |
| 873 | |
| 874 | def _pack_group_for_export(self, held: list, slot_idx: int) -> tuple: |
| 875 | """Copy ``held`` into ring slot ``slot_idx``; return ``(storages, views, |
| 876 | refs)`` in the same shape the per-storage path returns, but with ONE |
| 877 | storage. Views are built exactly as ``publish_group`` rebuilds them, so |
| 878 | the two sides cannot disagree about layout. |
| 879 | |
| 880 | Slot safety is not the credit gate alone: that bounds how MANY groups |
| 881 | are unfreed, not the order they free in, and barriers do complete out of |
| 882 | order. The caller takes ``slot_idx`` from a pool keyed by live group |
| 883 | (``_slot_of_group`` / ``_free_slots``); ``_drop_inflight`` returns a slot |
| 884 | only once its group is freed everywhere. |
| 885 | """ |
| 886 | offsets: list[int] = [] |
| 887 | cur = 0 |
| 888 | for _name, t in held: |
| 889 | cur = (cur + 15) & ~15 # 16B keeps every element size we ship aligned |
| 890 | offsets.append(cur) |
| 891 | cur += t.numel() * t.element_size() |
| 892 | need = max(16, (cur + 15) & ~15) |
| 893 | |
| 894 | slot = self._export_ring[slot_idx] |
| 895 | device = held[0][1].device |
| 896 | if slot is None or slot.numel() < need or slot.device != device: |
| 897 | self._export_ring[slot_idx] = self._export_ring_args[slot_idx] = None |
| 898 | slot = torch.empty(need, dtype=torch.uint8, device=device) |
| 899 | self._export_ring[slot_idx] = slot |
| 900 | |
| 901 | ust = slot.untyped_storage() |
| 902 | sid = ust.data_ptr() |
| 903 | reduce_args = self._export_ring_args[slot_idx] |
| 904 | if reduce_args is None: |
| 905 | base = torch.empty(0, dtype=torch.uint8, device=device) |
| 906 | base.set_(ust, 0, (ust.nbytes(),)) |
| 907 | _rebuild, reduce_args = reduce_tensor(base) |
| 908 | self._export_ring_args[slot_idx] = reduce_args |
| 909 | |
| 910 | views: dict[str, tuple] = {} |
| 911 | refs: dict[str, torch.Tensor] = {} |
| 912 | for (name, t), off in zip(held, offsets): |
| 913 | shape, esz = tuple(t.shape), t.element_size() |
| 914 | stride = self._contig_stride(shape) |
| 915 | dst = torch.as_strided(slot.view(t.dtype), shape, stride, off // esz) |
| 916 | dst.copy_(t) |
| 917 | refs[name] = dst |
| 918 | views[name] = ( |
| 919 | sid, |
| 920 | str(t.dtype).split(".")[-1], |
| 921 | list(shape), |
| 922 | list(stride), |
| 923 | off // esz, |
| 924 | ) |
| 925 | # The producer sidecar reads this slot from another process over CUDA |
| 926 | # IPC, while these copies are merely enqueued on this rank's stream; |
| 927 | # nothing orders the two. Without this sync it packs bytes the copy has |
| 928 | # not written yet -- silent corruption, and only for ring-packed tensors. |
| 929 | if slot.is_cuda: |
| 930 | torch.cuda.current_stream(slot.device).synchronize() |
| 931 | return {sid: reduce_args}, views, refs |
| 932 | |
| 933 | def _publish_async(self, group_idx: int, entries): |
| 934 | """Fire publish_group WITHOUT waiting on the RPC (the gather loop |
| 935 | overlaps the publish's server-side rebuild with the next group's |
| 936 | gather) and return a handle that ``_await_publish`` resolves. Ray actor |
| 937 | handle in production; a plain (non-Ray) fake server runs inline.""" |
| 938 | method = self._server.publish_group |
| 939 | remote = getattr(method, "remote", None) |
| 940 | if remote is not None: |
| 941 | return remote(group_idx, entries) |
| 942 | return method(group_idx, entries) |
| 943 | |
| 944 | def _await_publish(self, ref) -> None: |
| 945 | """Resolve one async publish so a server-side rebuild error surfaces at |
| 946 | window depth instead of at end_sync. Publishes carry nothing back — |
| 947 | freed groups flow only through wait_freed/end_sync (one channel; see |
| 948 | publish_group).""" |
| 949 | import ray |
| 950 | |
| 951 | if isinstance(ref, ray.ObjectRef): |
| 952 | ray.get(ref) |
| 953 | |
| 954 | # ---------------- construction ---------------- |
| 955 | |
| 956 | @classmethod |
| 957 | def trainer_init( |
| 958 | cls, |
| 959 | init_info: ShardedRDTTrainerInitInfo, |
| 960 | *, |
| 961 | client: VLLMWeightSyncClient, |
| 962 | source: WeightSource | None = None, |
| 963 | ) -> Self: |
| 964 | if source is None: |
| 965 | raise ValueError( |
| 966 | "Sharded RDT trainer weight transfer requires a WeightSource." |
| 967 | ) |
| 968 | engine = cls( |
| 969 | client=client, |
| 970 | source=source, |
| 971 | is_sender=init_info.is_sender, |
| 972 | init_info=init_info, |
| 973 | ) |
| 974 | |
| 975 | # The VMM-based expandable_segments allocator makes IPC storage opens ~9x |
| 976 | # slower on both export and rebuild, and CUDA-IPC publish is this engine's |
| 977 | # hot path. Frameworks enabling it should disable it around send_weights. |
| 978 | # Only the env var is visible; runtime changes are not introspectable. |
| 979 | import os as _os |
| 980 | |
| 981 | if "expandable_segments:True" in _os.environ.get("PYTORCH_CUDA_ALLOC_CONF", ""): |
| 982 | logger.warning( |
| 983 | "Sharded-RDT trainer: PYTORCH_CUDA_ALLOC_CONF enables " |
| 984 | "expandable_segments; CUDA-IPC weight publishing will be " |
| 985 | "several times slower. Disable expandable segments around " |
| 986 | "weight sync (allocate gather buffers in classic segments)." |
| 987 | ) |
| 988 | |
| 989 | engine._meta = list(source.metadata()) |
| 990 | names = [m.name for m in engine._meta] |
| 991 | engine._groups = layerwise_groups(names) |
| 992 | flat = [n for g in engine._groups for n in g] |
| 993 | if flat != names: |
| 994 | raise ValueError( |
| 995 | "Sharded RDT requires a WeightSource whose metadata order is " |
| 996 | "group-contiguous. Reorder the source so all names sharing a " |
| 997 | "layer index are adjacent." |
| 998 | ) |
| 999 | |
| 1000 | world, rank = engine._world_and_rank() |
| 1001 | engine._resolve_ownership(world, rank) |
| 1002 | engine._spawn_server(sorted(engine._held_names or [])) |
| 1003 | |
| 1004 | # Every rank's server must exist before the sender's init RPC (the worker |
| 1005 | # init calls reserve_serve_buffer back on ALL producer servers). The |
| 1006 | # all-gather of server names doubles as that barrier. |
| 1007 | server_names = engine._all_gather_server_names(world, rank) |
| 1008 | # Retained so a RESTARTED consumer can be re-initialized without another |
| 1009 | # all-gather (see get_worker_init_payload). The uuid actor names stay valid |
| 1010 | # for the run, because producers never restart. |
| 1011 | engine._server_names = server_names |
| 1012 | |
| 1013 | if engine.is_sender: |
| 1014 | worker_init = engine._build_worker_init_info(server_names) |
| 1015 | engine.client.init_weight_transfer_engine(asdict(worker_init)) |
| 1016 | return engine |
| 1017 | |
| 1018 | def get_worker_init_payload(self) -> dict: |
| 1019 | """The consumer-side init payload, rebuilt on demand. Pure — no collective. |
| 1020 | |
| 1021 | Reads only retained state, so a restarted inference engine can rejoin at a |
| 1022 | sync boundary: the driver can ask any time, including mid-run with every |
| 1023 | rank in its own training step. A collective here would deadlock, since the |
| 1024 | ranks are not at a matching one. |
| 1025 | |
| 1026 | Raises: |
| 1027 | RuntimeError: called before ``trainer_init`` cached the server names. |
| 1028 | |
| 1029 | """ |
| 1030 | if self._server_names is None: |
| 1031 | raise RuntimeError( |
| 1032 | "get_worker_init_payload requires trainer_init to have run (the " |
| 1033 | "producer server names are gathered there)." |
| 1034 | ) |
| 1035 | return asdict(self._build_worker_init_info(self._server_names)) |
| 1036 | |
| 1037 | def _world_and_rank(self) -> tuple[int, int]: |
| 1038 | if torch.distributed.is_available() and torch.distributed.is_initialized(): |
| 1039 | return torch.distributed.get_world_size(), torch.distributed.get_rank() |
| 1040 | return 1, self._init_info.rank |
| 1041 | |
| 1042 | def _resolve_ownership(self, world: int, rank: int) -> None: |
| 1043 | """Resolve which rank holds which name, and this rank's publish plan. |
| 1044 | |
| 1045 | A source may hold only part of the model — pipeline stages, expert |
| 1046 | parallelism, or any mix — so each rank declares its held names and the |
| 1047 | fleet all-gathers them. The consumers route per name, so the wire carries |
| 1048 | the transposed result: the distinct owner sets, and a per-name index into |
| 1049 | them. |
| 1050 | |
| 1051 | The masks are positional over metadata order, which is why the metadata |
| 1052 | digest is checked first: a rank whose names disagree would transpose into |
| 1053 | the wrong owners entirely. |
| 1054 | """ |
| 1055 | assert self.source is not None # guaranteed by trainer_init |
| 1056 | num_groups = len(self._groups) |
| 1057 | names = [m.name for m in self._meta] |
| 1058 | held = self.source.held_names() |
| 1059 | if held is None: |
| 1060 | held_set = set(names) |
| 1061 | else: |
| 1062 | held_set = {str(n) for n in held} |
| 1063 | unknown = sorted(held_set - set(names)) |
| 1064 | if unknown: |
| 1065 | raise ValueError( |
| 1066 | f"WeightSource.held_names() lists {len(unknown)} name(s) not " |
| 1067 | f"in metadata(), e.g. {unknown[:3]}." |
| 1068 | ) |
| 1069 | if not held_set: |
| 1070 | raise ValueError( |
| 1071 | "WeightSource.held_names() is empty; a rank with nothing to " |
| 1072 | "serve cannot take part in the gather" |
| 1073 | ) |
| 1074 | |
| 1075 | # One collective carries this rank's holdings as a bitmask over metadata |
| 1076 | # order, plus the digest of the metadata it indexes into. The digest is |
| 1077 | # what makes partial ownership safe: only the sender's metadata reaches |
| 1078 | # the consumers, so a rank describing just its own share would silently |
| 1079 | # serve the wrong weights. |
| 1080 | digest = self._meta_digest() |
| 1081 | mask = bytearray((len(names) + 7) // 8) |
| 1082 | for i, n in enumerate(names): |
| 1083 | if n in held_set: |
| 1084 | mask[i >> 3] |= 1 << (i & 7) |
| 1085 | per_rank = self._all_gather_owned(world, (digest, bytes(mask))) |
| 1086 | mismatched = [r for r, (d, *_rest) in enumerate(per_rank) if d != digest] |
| 1087 | if mismatched: |
| 1088 | raise ValueError( |
| 1089 | f"WeightSource.metadata() disagrees across trainer ranks " |
| 1090 | f"(rank {rank} digest {digest}, differing ranks {mismatched[:4]}). " |
| 1091 | "Every rank must describe the WHOLE model, even when it holds " |
| 1092 | "only some of it." |
| 1093 | ) |
| 1094 | masks = [m for _d, m in per_rank] |
| 1095 | |
| 1096 | # Transpose to per-name owners, then dedup into classes. Numbering by |
| 1097 | # FIRST APPEARANCE in metadata order keeps it a pure function of |
| 1098 | # rank-identical inputs, which is what lets a rejoining consumer rebuild |
| 1099 | # the identical table from get_worker_init_payload. |
| 1100 | owner_sets: list[list[int]] = [] |
| 1101 | class_of_owners: dict[tuple[int, ...], int] = {} |
| 1102 | name_owner_class: list[int] = [] |
| 1103 | for i, n in enumerate(names): |
| 1104 | owners = tuple(r for r, m in enumerate(masks) if m[i >> 3] & (1 << (i & 7))) |
| 1105 | if not owners: |
| 1106 | raise ValueError( |
| 1107 | f"no trainer rank holds {n!r}; every name in metadata() must " |
| 1108 | "be held by at least one rank or it can never be served." |
| 1109 | ) |
| 1110 | ci = class_of_owners.get(owners) |
| 1111 | if ci is None: |
| 1112 | ci = len(owner_sets) |
| 1113 | class_of_owners[owners] = ci |
| 1114 | owner_sets.append(list(owners)) |
| 1115 | name_owner_class.append(ci) |
| 1116 | |
| 1117 | self._owner_sets = owner_sets |
| 1118 | self._name_owner_class = name_owner_class |
| 1119 | # Groups holding anything here: exactly what iteration must cover. |
| 1120 | self._owned_idx = [ |
| 1121 | gi |
| 1122 | for gi in range(num_groups) |
| 1123 | if any(n in held_set for n in self._groups[gi]) |
| 1124 | ] |
| 1125 | self._held_names = held_set |
| 1126 | |
| 1127 | def _validate_held_yields(self, gi: int, names, tensors) -> None: |
| 1128 | """Check ``held_names()`` against what the source actually yields — the |
| 1129 | one place both sit side by side. |
| 1130 | |
| 1131 | Without it, a source that claims a name but yields ``None`` for it dies |
| 1132 | by stall watchdog 300s later: consumers route pulls here, the pull |
| 1133 | passes the served-names guard, and the cache wait never completes. With |
| 1134 | it, that is an immediate error naming the weight. One set lookup per |
| 1135 | name per sync.""" |
| 1136 | if self._held_names is None: |
| 1137 | return |
| 1138 | for name, tensor in zip(names, tensors): |
| 1139 | if (tensor is None) == (name in self._held_names): |
| 1140 | claim = "does not hold" if tensor is not None else "holds" |
| 1141 | have = "a real tensor" if tensor is not None else "None" |
| 1142 | raise RuntimeError( |
| 1143 | f"held_names() disagrees with the yielded tensors: group {gi} " |
| 1144 | f"name {name!r} yielded {have} but this rank {claim} it. A " |
| 1145 | "WeightSource must yield a real tensor for every held name " |
| 1146 | "and None for the rest (see WeightSource.held_names)." |
| 1147 | ) |
| 1148 | |
| 1149 | def _meta_digest(self) -> str: |
| 1150 | """Stable digest of this rank's metadata (name order + count).""" |
| 1151 | import hashlib |
| 1152 | |
| 1153 | h = hashlib.sha256() |
| 1154 | h.update(f"{len(self._meta)}\n".encode()) |
| 1155 | for m in self._meta: |
| 1156 | h.update(m.name.encode()) |
| 1157 | h.update(b"\n") |
| 1158 | return h.hexdigest()[:16] |
| 1159 | |
| 1160 | def _all_gather_owned(self, world: int, mine: tuple) -> list[tuple]: |
| 1161 | """All-gather each rank's (metadata digest, held-name bitmask).""" |
| 1162 | if world <= 1 or not ( |
| 1163 | torch.distributed.is_available() and torch.distributed.is_initialized() |
| 1164 | ): |
| 1165 | return [mine] |
| 1166 | gathered: list[Any] = [None] * world |
| 1167 | torch.distributed.all_gather_object(gathered, mine) |
| 1168 | return gathered |
| 1169 | |
| 1170 | def _all_gather_server_names(self, world: int, rank: int) -> list[str]: |
| 1171 | assert self._server_name is not None |
| 1172 | if world <= 1 or not ( |
| 1173 | torch.distributed.is_available() and torch.distributed.is_initialized() |
| 1174 | ): |
| 1175 | return [self._server_name] |
| 1176 | gathered: list[str | None] = [None] * world |
| 1177 | torch.distributed.all_gather_object(gathered, self._server_name) |
| 1178 | return [n for n in gathered if n is not None] |
| 1179 | |
| 1180 | def _spawn_server(self, served_names: list[str]) -> None: |
| 1181 | import ray |
| 1182 | from ray.util.scheduling_strategies import NodeAffinitySchedulingStrategy |
| 1183 | |
| 1184 | # The trainer is usually a separate install from the workers, so it gets |
| 1185 | # its own check rather than trusting the consumer side's. |
| 1186 | check_ray_rdt_version() |
| 1187 | |
| 1188 | ii = self._init_info |
| 1189 | self._server_name = f"vllm_rdt_producer_{uuid.uuid4().hex[:12]}_rk{ii.rank}" |
| 1190 | node_id = ray.get_runtime_context().get_node_id() |
| 1191 | # Pin the server to this rank's physical GPU: num_gpus=0 so Ray does not |
| 1192 | # allocate a second, CUDA_VISIBLE_DEVICES so CUDA IPC to the rank's |
| 1193 | # gathered tensors works. max_concurrency > 1 serves pulls while control |
| 1194 | # calls are in flight; enable_tensor_transport gives the NIXL serve. |
| 1195 | gpu_ids = ray.get_gpu_ids() |
| 1196 | # The server is the trainer rank's process twin: forward the env the |
| 1197 | # rank runs under (library paths etc.) so it imports torch/vllm the |
| 1198 | # same way, then pin it to the rank's physical GPU for CUDA IPC. |
| 1199 | import os |
| 1200 | |
| 1201 | env_vars = { |
| 1202 | k: os.environ[k] |
| 1203 | for k in ( |
| 1204 | "LD_LIBRARY_PATH", |
| 1205 | "LD_PRELOAD", |
| 1206 | "NCCL_CUMEM_ENABLE", |
| 1207 | "VLLM_NCCL_SO_PATH", |
| 1208 | "PATH", |
| 1209 | ) |
| 1210 | if k in os.environ |
| 1211 | } |
| 1212 | if gpu_ids: |
| 1213 | # A num_gpus=0 actor sharing the rank's GPU: Ray would set |
| 1214 | # CUDA_VISIBLE_DEVICES="" and hide every GPU, so tell it not to touch |
| 1215 | # the var and pin the device ourselves. |
| 1216 | env_vars["RAY_EXPERIMENTAL_NOSET_CUDA_VISIBLE_DEVICES"] = "1" |
| 1217 | env_vars["CUDA_VISIBLE_DEVICES"] = str(gpu_ids[0]) |
| 1218 | runtime_env = {"env_vars": env_vars} if env_vars else {} |
| 1219 | server_cls = ray.remote(_RDTProducerServer).options( |
| 1220 | name=self._server_name, |
| 1221 | namespace=ii.trainer_actor_namespace, |
| 1222 | num_cpus=0, |
| 1223 | num_gpus=0, |
| 1224 | # Thread budget: under EP-local routing every consumer pulls from every |
| 1225 | # producer, and with issue-ahead each can park K produce calls in the |
| 1226 | # cache-wait — C*K blocked calls. A smaller pool queues the control |
| 1227 | # plane behind them and deadlocks: the parked produces wait for a |
| 1228 | # publish that can never get a thread (the sync-2 wedge at 235B tp8, |
| 1229 | # 16 parked vs 8 threads). Size for the worst case plus control-plane |
| 1230 | # slack. |
| 1231 | max_concurrency=ii.num_consumers * ii.num_rdt_buffers + 4, |
| 1232 | enable_tensor_transport=True, |
| 1233 | scheduling_strategy=NodeAffinitySchedulingStrategy( |
| 1234 | node_id=node_id, soft=False |
| 1235 | ), |
| 1236 | runtime_env=runtime_env, |
| 1237 | ) |
| 1238 | self._server = server_cls.remote( |
| 1239 | served_names=served_names, |
| 1240 | num_rdt_buffers=ii.num_rdt_buffers, |
| 1241 | buffer_presize_gb=ii.buffer_presize_gb, |
| 1242 | gather_lookahead=ii.gather_lookahead, |
| 1243 | stall_timeout_s=ii.stall_timeout_s, |
| 1244 | workers_per_replica=ii.workers_per_replica, |
| 1245 | ) |
| 1246 | ray.get(self._server.ping.remote()) |
| 1247 | # Pre-barrier NIXL warmup: must complete before the server-name |
| 1248 | # all-gather so no rank can be inside a collective yet (see warmup_nixl). |
| 1249 | ray.get(self._server.warmup_nixl.remote()) |
| 1250 | |
| 1251 | def _build_worker_init_info(self, server_names: list[str]): |
| 1252 | from vllm.distributed.weight_transfer.sharded_rdt_engine import ( |
| 1253 | ShardedRDTWeightTransferInitInfo, |
| 1254 | ) |
| 1255 | |
| 1256 | group_lens = [len(g) for g in self._groups] |
| 1257 | names = [m.name for m in self._meta] |
| 1258 | dtype_names = [str(m.dtype).split(".")[-1] for m in self._meta] |
| 1259 | shapes = [list(m.shape) for m in self._meta] |
| 1260 | return ShardedRDTWeightTransferInitInfo( |
| 1261 | trainer_actor_names=server_names, |
| 1262 | trainer_actor_namespace=self._init_info.trainer_actor_namespace, |
| 1263 | produce_method_name=PRODUCE_METHOD_NAME, |
| 1264 | names=names, |
| 1265 | dtype_names=dtype_names, |
| 1266 | shapes=shapes, |
| 1267 | group_lens=group_lens, |
| 1268 | owner_sets=self._owner_sets, |
| 1269 | name_owner_class=self._name_owner_class, |
| 1270 | num_consumers=self._init_info.num_consumers, |
| 1271 | num_rdt_buffers=self._init_info.num_rdt_buffers, |
| 1272 | buffer_presize_gb=self._init_info.buffer_presize_gb, |
| 1273 | ) |
| 1274 | |
| 1275 | # ---------------- per-round ---------------- |
| 1276 | |
| 1277 | def send_weights(self, live_consumer_ids: Collection[int] | None = None) -> None: |
| 1278 | """Gather this rank's weights and publish them for the consumers to pull. |
| 1279 | |
| 1280 | ``live_consumer_ids`` restricts the sync to the consumers still alive; |
| 1281 | ``None`` serves the whole provisioned set. The provisioned geometry is |
| 1282 | FROZEN for the run — a degraded sync only lowers the live count handed to |
| 1283 | ``begin_sync``. Every rank must get the SAME live set, since they share |
| 1284 | the gather collectives, so the caller computes it once for all of them. |
| 1285 | """ |
| 1286 | assert self.source is not None |
| 1287 | if live_consumer_ids is None: |
| 1288 | live_ids = list(range(self._init_info.num_consumers)) |
| 1289 | live_count = len(live_ids) |
| 1290 | else: |
| 1291 | live_ids = sorted(set(int(c) for c in live_consumer_ids)) |
| 1292 | live_count = len(live_ids) |
| 1293 | logger.warning( |
| 1294 | "[rdt-degraded] serving %d/%d live consumers; every group's " |
| 1295 | "free barrier counts to the live total", |
| 1296 | live_count, |
| 1297 | self._init_info.num_consumers, |
| 1298 | ) |
| 1299 | self._send_weights_inner(live_count, live_ids) |
| 1300 | |
| 1301 | def _send_weights_inner(self, live_count: int, live_ids: list[int]) -> None: |
| 1302 | if not self.is_sender: |
| 1303 | self._run_gather_loop( |
| 1304 | update_future=None, live_count=live_count, live_ids=live_ids |
| 1305 | ) |
| 1306 | return |
| 1307 | |
| 1308 | self.client.start_weight_update() |
| 1309 | |
| 1310 | from vllm.distributed.weight_transfer.sharded_rdt_engine import ( |
| 1311 | ShardedRDTWeightTransferUpdateInfo, |
| 1312 | ) |
| 1313 | |
| 1314 | empty_update = asdict(ShardedRDTWeightTransferUpdateInfo()) |
| 1315 | with ThreadPoolExecutor(max_workers=1) as exe: |
| 1316 | # The workers block inside update_weights until they've pulled every |
| 1317 | # group, so it runs concurrently with the gather/publish loop. |
| 1318 | future = exe.submit(self.client.update_weights, empty_update) |
| 1319 | self._run_gather_loop( |
| 1320 | update_future=future, live_count=live_count, live_ids=live_ids |
| 1321 | ) |
| 1322 | future.result() # surface inference-side errors |
| 1323 | |
| 1324 | self.client.finish_weight_update() |
| 1325 | |
| 1326 | def _run_gather_loop( |
| 1327 | self, update_future, live_count: int, live_ids: list[int] | None = None |
| 1328 | ) -> None: |
| 1329 | """Gather this rank's weights group-by-group and publish each into the |
| 1330 | server over CUDA IPC. A gathered group is published — serveable — |
| 1331 | immediately; the loop gates BEFORE the next gather while more than |
| 1332 | `gather_lookahead` groups are unfreed (the per-group free barrier: a |
| 1333 | credit releases when every live consumer has signaled the group). So |
| 1334 | the loop self-paces to the consumers' pull rate with at most |
| 1335 | `gather_lookahead + 1` groups resident. Runs on every rank; only the |
| 1336 | sender has an `update_future` to fail fast on.""" |
| 1337 | assert self.source is not None # guaranteed by trainer_init |
| 1338 | # The live IDS ride along with the count: the free barrier needs only |
| 1339 | # the count, but the slot-sharing rendezvous needs to know WHICH |
| 1340 | # consumers are live. |
| 1341 | self._rpc("begin_sync", live_count, live_ids) |
| 1342 | # One generator resume per GROUP: `iter_groups` yields (names, tensors) |
| 1343 | # per owned group in metadata order. Every owner must reach a group in the |
| 1344 | # same order or their shared gather collective mismatches. |
| 1345 | groups = self.source.iter_groups() |
| 1346 | # Publishes fire without an inline ray.get and are harvested this window |
| 1347 | # deep, so the RPC and server-side rebuild overlap the NEXT group's |
| 1348 | # gather/export. The window also surfaces a server-side error at depth 2 |
| 1349 | # rather than at end_sync. |
| 1350 | _PUBLISH_WINDOW = 2 |
| 1351 | pending_publish: list = [] |
| 1352 | # The memory bound: gate BEFORE gathering while more |
| 1353 | # than `bound` groups are unfreed, so at most bound + 1 are resident. The |
| 1354 | # count is exact — `_inflight` gains its entry right after each gather and |
| 1355 | # only wait_freed/end_sync shrinks it, so a group freed server-side but not |
| 1356 | # yet collected still holds its refs and still counts. The publish window |
| 1357 | # drains first, so every resident group is pullable while we wait. |
| 1358 | # |
| 1359 | # 0 is legal and means NO pipelining: the loop waits for group N to be |
| 1360 | # freed by every live consumer before gathering N+1, so exactly one group |
| 1361 | # is resident and the sync costs the full serialized |
| 1362 | # gather + publish + pull + free RTT per group. 1 (the default) overlaps |
| 1363 | # N+1's gather with N's pulls. |
| 1364 | bound = max(0, self._init_info.gather_lookahead) |
| 1365 | |
| 1366 | # [RDT-EXPORT-RING] `bound + 1` slots is exactly the residency the credit |
| 1367 | # gate enforces. Small storages are packed into a slot and exported once |
| 1368 | # per slot instead of once per storage (518 -> 7 cudaIpcGetMemHandle |
| 1369 | # calls per rank per sync at 235B, each 0.5-1.5ms). |
| 1370 | nslots = bound + 1 |
| 1371 | # data_ptr -> (storage weakref, reduce_args), this sync only. `storages` |
| 1372 | # is per-group, so a source that reuses a buffer across groups would |
| 1373 | # otherwise re-export it every group. The weakref is load-bearing: a |
| 1374 | # freed allocation can be reissued at the same address, and serving that |
| 1375 | # from a stale handle is silently wrong bytes. |
| 1376 | handle_cache: dict[int, tuple] = {} |
| 1377 | if len(getattr(self, "_export_ring", [])) != nslots: |
| 1378 | self._export_ring = [None] * nslots |
| 1379 | self._export_ring_args = [None] * nslots |
| 1380 | # Slots are handed out per group and returned by _drop_inflight, so the |
| 1381 | # pool starts full every sync. |
| 1382 | self._slot_of_group = {} |
| 1383 | self._free_slots = list(range(nslots)) |
| 1384 | |
| 1385 | try: |
| 1386 | for gi in self._owned_idx: |
| 1387 | group = self._groups[gi] |
| 1388 | while len(self._inflight) > bound: |
| 1389 | while pending_publish: |
| 1390 | self._await_publish(pending_publish.pop(0)) |
| 1391 | self._drop_inflight(self._rpc("wait_freed")) |
| 1392 | names, tensors = next(groups) |
| 1393 | if list(names) != list(group): |
| 1394 | raise RuntimeError( |
| 1395 | f"WeightSource group yielded {len(names)} names starting " |
| 1396 | f"{names[:2]!r} but expected {len(group)} starting " |
| 1397 | f"{group[:2]!r}; iteration order must match metadata." |
| 1398 | ) |
| 1399 | self._validate_held_yields(gi, names, tensors) |
| 1400 | # Share each unique STORAGE once and describe every name as an |
| 1401 | # as_strided view onto it. ``None`` means a name this rank does not |
| 1402 | # hold (a foreign expert); the source keeps it in the list so the |
| 1403 | # order check stays rank-uniform, and it is dropped here, before |
| 1404 | # the IPC export, matching the sidecar's served_names. |
| 1405 | storages: dict[int, tuple] = {} |
| 1406 | views: dict[str, tuple] = {} |
| 1407 | refs: dict[str, torch.Tensor] = {} |
| 1408 | small: list = [] |
| 1409 | for name, tensor in zip(names, tensors): |
| 1410 | if tensor is None: |
| 1411 | continue |
| 1412 | tensor = tensor.detach() |
| 1413 | if not tensor.is_cuda: |
| 1414 | tensor = tensor.cuda() |
| 1415 | # Threshold on the STORAGE, not the tensor: a source may hand |
| 1416 | # out many small views of one large stacked buffer, and a |
| 1417 | # per-tensor test would pack the whole thing, costing +1 |
| 1418 | # group of peak memory for a storage already coalesced. |
| 1419 | if tensor.untyped_storage().nbytes() <= _EXPORT_RING_MAX_BYTES: |
| 1420 | small.append((name, tensor)) |
| 1421 | continue |
| 1422 | tensor = tensor.contiguous() |
| 1423 | refs[name] = tensor # keep the export alive |
| 1424 | ust = tensor.untyped_storage() |
| 1425 | sid = ust.data_ptr() |
| 1426 | if sid not in storages: |
| 1427 | cached = handle_cache.get(sid) |
| 1428 | if cached is not None and not cached[0].expired(): |
| 1429 | reduce_args = cached[1] |
| 1430 | else: |
| 1431 | base = torch.empty( |
| 1432 | 0, dtype=torch.uint8, device=tensor.device |
| 1433 | ) |
| 1434 | base.set_(ust, 0, (ust.nbytes(),)) |
| 1435 | _rebuild, reduce_args = reduce_tensor(base) |
| 1436 | handle_cache[sid] = (StorageWeakRef(ust), reduce_args) |
| 1437 | storages[sid] = reduce_args |
| 1438 | views[name] = ( |
| 1439 | sid, |
| 1440 | str(tensor.dtype).split(".")[-1], |
| 1441 | list(tensor.shape), |
| 1442 | list(tensor.stride()), |
| 1443 | tensor.storage_offset(), |
| 1444 | ) |
| 1445 | if small: |
| 1446 | # Take a slot no live group holds. The gate above already |
| 1447 | # guarantees one exists, so this is a backstop -- but it |
| 1448 | # must drain publishes first, like the credit gate. |
| 1449 | while not self._free_slots: |
| 1450 | while pending_publish: |
| 1451 | self._await_publish(pending_publish.pop(0)) |
| 1452 | self._drop_inflight(self._rpc("wait_freed")) |
| 1453 | slot_idx = self._free_slots.pop(0) |
| 1454 | self._slot_of_group[gi] = slot_idx |
| 1455 | st, vw, rf = self._pack_group_for_export(small, slot_idx) |
| 1456 | storages.update(st) |
| 1457 | views.update(vw) |
| 1458 | refs.update(rf) |
| 1459 | del small |
| 1460 | del tensors |
| 1461 | # The loop leaves its last `tensor` / `ust` / `base` bound |
| 1462 | # through the credit gate and into the next gather, and `base` |
| 1463 | # is a whole-storage view -- so one tensor and its entire |
| 1464 | # allocation stay alive past `_drop_inflight`, unseen by the |
| 1465 | # `_inflight` accounting the residency bound rests on. |
| 1466 | # Assignment, not `del`: unbound if no direct-path tensor. |
| 1467 | tensor = ust = base = None |
| 1468 | if not refs: |
| 1469 | # A group with nothing held here cannot occur (every group |
| 1470 | # carries replicated names), but publishing an empty group |
| 1471 | # would park a credit nobody's pull is waiting on; skip. |
| 1472 | continue |
| 1473 | # Hold our refs before publishing; drop them only when the |
| 1474 | # server reports the group freed (IPC export must outlive import). |
| 1475 | self._inflight[gi] = refs |
| 1476 | pending_publish.append(self._publish_async(gi, (storages, views))) |
| 1477 | while len(pending_publish) >= _PUBLISH_WINDOW: |
| 1478 | self._await_publish(pending_publish.pop(0)) |
| 1479 | if update_future is not None and update_future.done(): |
| 1480 | # update_weights returned/failed early — surface now instead |
| 1481 | # of blocking further gathers. |
| 1482 | update_future.result() |
| 1483 | while pending_publish: |
| 1484 | self._await_publish(pending_publish.pop(0)) |
| 1485 | freed = self._rpc("end_sync") |
| 1486 | self._drop_inflight(freed) |
| 1487 | except BaseException as e: |
| 1488 | with contextlib.suppress(Exception): |
| 1489 | self._rpc("set_gather_error", repr(e)) |
| 1490 | self._inflight.clear() |
| 1491 | raise |
| 1492 | finally: |
| 1493 | # Keep the slots, drop the handles: the sidecar's base cache is |
| 1494 | # cleared per sync, so a handle must not outlive the sync either. |
| 1495 | self._export_ring_args = [None] * len(getattr(self, "_export_ring", [])) |
| 1496 | |
| 1497 | def _drop_inflight(self, freed_keys: list) -> None: |
| 1498 | """Release a freed group's refs and return its export-ring slot.""" |
| 1499 | free_slots = getattr(self, "_free_slots", None) |
| 1500 | slot_of_group = getattr(self, "_slot_of_group", None) |
| 1501 | for k in freed_keys: |
| 1502 | gi = int(k) |
| 1503 | self._inflight.pop(gi, None) |
| 1504 | if slot_of_group is None: |
| 1505 | continue |
| 1506 | slot = slot_of_group.pop(gi, None) |
| 1507 | if slot is not None and free_slots is not None and slot not in free_slots: |
| 1508 | free_slots.append(slot) |
| 1509 | |
| 1510 | # ---------------- misc ---------------- |
| 1511 | |
| 1512 | def shutdown(self) -> None: |
| 1513 | if self._server is None: |
| 1514 | return |
| 1515 | import ray |
| 1516 | |
| 1517 | with contextlib.suppress(Exception): |
| 1518 | ray.get(self._server.shutdown.remote()) |
| 1519 | ray.kill(self._server) |
| 1520 | self._server = None |
| 1521 | self._inflight.clear() |