| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | """Op-chain recording for the sharded-RDT backend. |
| 4 | |
| 5 | The consumer asks the trainer for the exact slice a worker consumes, described as |
| 6 | an op chain replayed on the trainer's live tensor. ``FakeRDTTensor`` builds the |
| 7 | chain by intercepting the model's own weight loaders; ``copy_`` is its data sink, |
| 8 | ``BakeSink`` — the dry run records how each slice would be fetched and where it |
| 9 | lands. |
| 10 | |
| 11 | Chains are built against meta tensors, so nothing here touches Ray or a GPU. |
| 12 | """ |
| 13 | |
| 14 | from collections import defaultdict |
| 15 | from collections.abc import Callable |
| 16 | from dataclasses import dataclass, field |
| 17 | from typing import Any |
| 18 | |
| 19 | import torch |
| 20 | |
| 21 | from vllm.distributed.weight_transfer.sharded_rdt_common import SUPPORTED_OPS |
| 22 | |
| 23 | # ("op_name", positional_args, sorted_kwargs_items). Entries must be hashable so |
| 24 | # the chain can serve as a FetchKey. |
| 25 | OpSpec = tuple[str, tuple[Any, ...], tuple[tuple[str, Any], ...]] |
| 26 | OpChain = tuple[OpSpec, ...] |
| 27 | FetchKey = tuple[str, OpChain] |
| 28 | |
| 29 | |
| 30 | def _freeze_kwargs(kwargs: dict[str, Any]) -> tuple[tuple[str, Any], ...]: |
| 31 | """Sort kwargs into a tuple of items for hashable storage in OpSpec.""" |
| 32 | return tuple(sorted(kwargs.items())) |
| 33 | |
| 34 | |
| 35 | class _UnsupportedFakeOp(NotImplementedError): |
| 36 | """Raised when a weight loader does something to a FakeRDTTensor that cannot |
| 37 | be expressed as a slice request. |
| 38 | |
| 39 | Surfaced as NotImplementedError so callers can distinguish "this backend |
| 40 | can't handle this loader" from genuine bugs. |
| 41 | """ |
| 42 | |
| 43 | |
| 44 | @dataclass |
| 45 | class _Scatter: |
| 46 | """One recorded scatter: pull ``src``, copy it into ``layer``'s |
| 47 | ``param_name`` at the recorded strided region. The bake's output unit; see |
| 48 | the engine module's Data flow. |
| 49 | |
| 50 | Self-contained so replay needs no lookups. Destination geometry is read off |
| 51 | the meta view and rebuilt each sync as ``as_strided(shape, stride, offset)`` |
| 52 | -- ``layer`` is resolved to a param at replay time, never baked, since every |
| 53 | sync re-materializes fresh tensors. |
| 54 | |
| 55 | ``shape`` and ``dtype`` also size the slice ON THE WIRE, so ``dtype`` is the |
| 56 | PRODUCED dtype (the fake's after its op chain), not the source name's: |
| 57 | ``view(dtype)`` is allowlisted, and taking the source's would size the slice |
| 58 | with the wrong itemsize and carve the packed blob differently on the two |
| 59 | sides. |
| 60 | """ |
| 61 | |
| 62 | layer: Any |
| 63 | param_name: str |
| 64 | src: FetchKey |
| 65 | offset: int |
| 66 | shape: tuple[int, ...] |
| 67 | stride: tuple[int, ...] |
| 68 | dtype: torch.dtype |
| 69 | |
| 70 | |
| 71 | def _meta_copy_(dest: torch.Tensor, src: "FakeRDTTensor") -> torch.Tensor: |
| 72 | """Fire ``dest.copy_`` from a zero-storage meta source of ``src``'s geometry. |
| 73 | |
| 74 | Moves no data but still counts against the layer's loaded numel, which drives |
| 75 | ``_layerwise_process``; skipping it would leave the layer looking unloaded |
| 76 | forever. |
| 77 | |
| 78 | A non-meta ``dest`` is one the layerwise reload never moved there (the |
| 79 | ``SKIP_LOAD_TENSORS`` set, e.g. GLM's router bias). Torch forbids |
| 80 | ``real.copy_(meta)``, and skipping the count is consistent -- ``get_layer_size`` |
| 81 | excludes the same set. The caller still RECORDS the copy, so the replay writes |
| 82 | the real param. |
| 83 | """ |
| 84 | if not dest.is_meta: |
| 85 | return dest |
| 86 | meta_src = torch.empty(src.shape, dtype=src.dtype, device="meta") |
| 87 | with torch._C.DisableTorchFunctionSubclass(): |
| 88 | return dest.copy_(meta_src) |
| 89 | |
| 90 | |
| 91 | @dataclass |
| 92 | class BakeSink: |
| 93 | """``copy_`` sink for the dry-run bake: record how each slice would be fetched |
| 94 | and where it would land, and move nothing. |
| 95 | |
| 96 | ``_install_recording_stamps`` sets ``current = (leaf_module, param_name)`` |
| 97 | around each loader so ``accept_copy`` can attribute the copy. An unstamped |
| 98 | ``copy_`` cannot be attributed and stays unrecorded, so its module fails the |
| 99 | coverage gate and takes the plain load. ``copies_by_layer`` is keyed by module |
| 100 | object, so iterating it yields each leaf module once. |
| 101 | """ |
| 102 | |
| 103 | copies_by_layer: "dict[Any, list[_Scatter | None]]" = field( |
| 104 | default_factory=lambda: defaultdict(list) |
| 105 | ) |
| 106 | current: "tuple[Any, str] | None" = None |
| 107 | # Source names whose ``copy_`` fired during the bake. Names not here never |
| 108 | # moved data (e.g. experts owned by another EP rank) and are skipped; names |
| 109 | # here but unbaked fail the plan build. |
| 110 | copied_names: "set[str]" = field(default_factory=set) |
| 111 | |
| 112 | def accept_copy(self, dest: torch.Tensor, src: "FakeRDTTensor") -> torch.Tensor: |
| 113 | self.copied_names.add(src._name) |
| 114 | if self.current is not None: |
| 115 | layer, param_name = self.current |
| 116 | if dest.numel() != src.numel(): |
| 117 | # The packed layout sizes each slice from the DESTINATION shape, |
| 118 | # while the producer packs what the replayed chain produces. A |
| 119 | # broadcasting copy_ makes those two lengths differ, and neither |
| 120 | # side can notice: the consumer's buffer view is exactly |
| 121 | # prod(dest.shape) elements, so it reshapes cleanly over bytes |
| 122 | # the producer laid out at different offsets. Refuse at bake. |
| 123 | raise _UnsupportedFakeOp( |
| 124 | f"FakeRDTTensor: broadcasting copy_ into {param_name!r} of " |
| 125 | f"{type(layer).__name__} ({src.numel()} elements from " |
| 126 | f"{src._name!r} into {dest.numel()}). The sharded RDT " |
| 127 | "backend sends exactly the slice the loader asks for, so " |
| 128 | "the source and destination must have equal numel." |
| 129 | ) |
| 130 | self.copies_by_layer[layer].append( |
| 131 | _Scatter( |
| 132 | layer=layer, |
| 133 | param_name=param_name, |
| 134 | src=src._key(), |
| 135 | offset=dest.storage_offset(), |
| 136 | shape=tuple(dest.shape), |
| 137 | stride=tuple(dest.stride()), |
| 138 | dtype=src.dtype, |
| 139 | ) |
| 140 | ) |
| 141 | return _meta_copy_(dest, src) |
| 142 | |
| 143 | |
| 144 | class FakeRDTTensor(torch.Tensor): |
| 145 | """Zero-storage tensor that records how to fetch a weight slice. |
| 146 | |
| 147 | ``_make_wrapper_subclass`` gives it shape/dtype/device without storage. Every |
| 148 | op in ``SUPPORTED_OPS`` returns a child with the spec appended; ``copy_`` |
| 149 | delegates to the installed sink. Anything else reaches ``__torch_dispatch__`` |
| 150 | and raises ``_UnsupportedFakeOp``, so failures are loud rather than silently |
| 151 | fetching the wrong bytes. |
| 152 | """ |
| 153 | |
| 154 | # Declared at class scope so mypy can infer attribute types on |
| 155 | # instances built via ``_make_wrapper_subclass`` (where ``__new__`` |
| 156 | # returns a tensor it can't annotate as ``self``). |
| 157 | _name: str |
| 158 | _ops: OpChain |
| 159 | # Handles this fake's copy_. ``None`` only on bare construction. |
| 160 | _sink: "BakeSink | None" |
| 161 | |
| 162 | @staticmethod |
| 163 | def __new__( |
| 164 | cls, |
| 165 | name: str, |
| 166 | shape: torch.Size, |
| 167 | dtype: torch.dtype, |
| 168 | device: torch.device, |
| 169 | ops: OpChain = (), |
| 170 | sink: "BakeSink | None" = None, |
| 171 | ) -> "FakeRDTTensor": |
| 172 | t = torch.Tensor._make_wrapper_subclass( |
| 173 | cls, |
| 174 | shape, |
| 175 | dtype=dtype, |
| 176 | device=device, |
| 177 | requires_grad=False, |
| 178 | ) |
| 179 | t._name = name |
| 180 | t._ops = tuple(ops) |
| 181 | t._sink = sink |
| 182 | return t |
| 183 | |
| 184 | def _key(self) -> FetchKey: |
| 185 | return (self._name, self._ops) |
| 186 | |
| 187 | def _make_child( |
| 188 | self, |
| 189 | new_shape: torch.Size, |
| 190 | new_dtype: torch.dtype, |
| 191 | *new_ops: OpSpec, |
| 192 | ) -> "FakeRDTTensor": |
| 193 | """Append one or more ops to the chain and return a fresh child. |
| 194 | |
| 195 | Variadic so multi-return ops (e.g. chunk) can append both the base |
| 196 | op and an indexing op in a single call. |
| 197 | """ |
| 198 | return FakeRDTTensor( |
| 199 | name=self._name, |
| 200 | shape=new_shape, |
| 201 | dtype=new_dtype, |
| 202 | device=self.device, |
| 203 | ops=self._ops + new_ops, |
| 204 | sink=self._sink, |
| 205 | ) |
| 206 | |
| 207 | def _meta(self) -> torch.Tensor: |
| 208 | """A zero-storage meta tensor of this fake's shape/dtype, so PyTorch |
| 209 | itself computes post-op geometry rather than us reimplementing it.""" |
| 210 | return torch.empty(self.shape, dtype=self.dtype, device="meta") |
| 211 | |
| 212 | @classmethod |
| 213 | def __torch_function__( |
| 214 | cls, |
| 215 | func, |
| 216 | types, |
| 217 | args=(), |
| 218 | kwargs=None, |
| 219 | ): |
| 220 | kwargs = kwargs or {} |
| 221 | |
| 222 | # Checked before the allowlist: copy_ must not be recorded into the chain. |
| 223 | if func is torch.Tensor.copy_: |
| 224 | dest = args[0] |
| 225 | src = args[1] if len(args) > 1 else kwargs.get("src") |
| 226 | if isinstance(src, cls) and src._sink is not None: |
| 227 | return src._sink.accept_copy(dest, src) |
| 228 | |
| 229 | # Allowlisted slice/view/shape ops: append to chain and return child. |
| 230 | op_name = SUPPORTED_OPS.get(func) |
| 231 | if op_name is not None: |
| 232 | self_ = args[0] |
| 233 | if isinstance(self_, cls): |
| 234 | rest = tuple(args[1:]) |
| 235 | return cls._intercept(self_, func, op_name, rest, kwargs) |
| 236 | |
| 237 | # Metadata reads are answered by the wrapper subclass and stop here; ops |
| 238 | # that need data fall through to __torch_dispatch__, which raises. |
| 239 | with torch._C.DisableTorchFunctionSubclass(): |
| 240 | return func(*args, **kwargs) |
| 241 | |
| 242 | @classmethod |
| 243 | def _intercept( |
| 244 | cls, |
| 245 | self_: "FakeRDTTensor", |
| 246 | func: Callable, |
| 247 | op_name: str, |
| 248 | args: tuple, |
| 249 | kwargs: dict, |
| 250 | ): |
| 251 | """Append the op and return a child, or a tuple of children for |
| 252 | multi-return ops. Each child's geometry comes from running the op on a |
| 253 | meta tensor.""" |
| 254 | meta = self_._meta() |
| 255 | with torch._C.DisableTorchFunctionSubclass(): |
| 256 | meta_result = func(meta, *args, **kwargs) |
| 257 | |
| 258 | base_op: OpSpec = (op_name, tuple(args), _freeze_kwargs(kwargs)) |
| 259 | |
| 260 | # Single-tensor result: one child carrying base_op alone. |
| 261 | if isinstance(meta_result, torch.Tensor): |
| 262 | return self_._make_child(meta_result.shape, meta_result.dtype, base_op) |
| 263 | |
| 264 | # Multi-return: one child per output, each carrying base_op plus a |
| 265 | # trailing __getitem__(i) so the replay can index the result. |
| 266 | if isinstance(meta_result, tuple | list): |
| 267 | return tuple( |
| 268 | self_._make_child( |
| 269 | m.shape, |
| 270 | m.dtype, |
| 271 | base_op, |
| 272 | ("__getitem__", (i,), ()), |
| 273 | ) |
| 274 | for i, m in enumerate(meta_result) |
| 275 | ) |
| 276 | |
| 277 | # Non-tensor result (e.g. .item() snuck into the allowlist). |
| 278 | raise _UnsupportedFakeOp( |
| 279 | f"FakeRDTTensor: {op_name!r} returned a non-tensor " |
| 280 | f"({type(meta_result).__name__}); cannot defer." |
| 281 | ) |
| 282 | |
| 283 | @classmethod |
| 284 | def __torch_dispatch__(cls, func, types, args=(), kwargs=None): |
| 285 | kwargs = kwargs or {} |
| 286 | |
| 287 | # A fake at the aten level means an unsupported loader op. Name the op and |
| 288 | # the chain; materializing instead would mask a correctness bug. |
| 289 | for arg in (*args, *kwargs.values()): |
| 290 | if isinstance(arg, cls): |
| 291 | raise _UnsupportedFakeOp( |
| 292 | f"FakeRDTTensor: unsupported op {func} reached " |
| 293 | f"__torch_dispatch__ on fake {arg._name!r} " |
| 294 | f"(chain={arg._ops}). Supported ops are: " |
| 295 | f"{sorted(SUPPORTED_OPS.values())}, plus copy_. " |
| 296 | "Loaders that need .to(), .float(), .item(), arithmetic, " |
| 297 | "bool-mask indexing, or .data access are not supported by " |
| 298 | "the sharded RDT backend." |
| 299 | ) |
| 300 | return func(*args, **kwargs) |