| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | """Tests for the sharded-RDT trainer-side engine. |
| 4 | |
| 5 | The trainer engine's own surface: init info and factory registration, the |
| 6 | send_weights round trip against an in-process fake producer server, the |
| 7 | ownership table it resolves and ships to the workers, and the producer's packed |
| 8 | destination-view cache. |
| 9 | |
| 10 | The plan/replay half (including `RdtRouter`, which is consumer-side) lives in |
| 11 | test_sharded_rdt_plan.py; the producer server's protocol in |
| 12 | test_sharded_rdt_producer.py. |
| 13 | """ |
| 14 | |
| 15 | import itertools |
| 16 | from dataclasses import asdict |
| 17 | |
| 18 | import pytest |
| 19 | import torch |
| 20 | |
| 21 | from tests.distributed.test_weight_transfer import RecordingClient |
| 22 | from vllm.distributed.weight_transfer import ( |
| 23 | ModuleSource, |
| 24 | ParamMeta, |
| 25 | WeightSource, |
| 26 | WeightTransferTrainerFactory, |
| 27 | ) |
| 28 | from vllm.distributed.weight_transfer.base import layerwise_groups |
| 29 | from vllm.distributed.weight_transfer.sharded_rdt_trainer import ( |
| 30 | ShardedRDTTrainerInitInfo, |
| 31 | ShardedRDTTrainerWeightTransferEngine, |
| 32 | ) |
| 33 | |
| 34 | |
| 35 | class _ListSource(WeightSource): |
| 36 | """A WeightSource over an explicit ordered (name, cpu-tensor) list, so the |
| 37 | sharded-RDT group/order logic can be tested without a real model.""" |
| 38 | |
| 39 | def __init__(self, pairs): |
| 40 | self._pairs = list(pairs) |
| 41 | |
| 42 | def metadata(self): |
| 43 | return [ParamMeta(n, t.dtype, tuple(t.shape)) for n, t in self._pairs] |
| 44 | |
| 45 | def __iter__(self): |
| 46 | return iter(self._pairs) |
| 47 | |
| 48 | |
| 49 | class _FakeProducerServer: |
| 50 | """In-process stand-in for the _RDTProducerServer Ray actor. Records the |
| 51 | engine->server call sequence and, by default, frees each group as soon as |
| 52 | it is published (simulating the consumers' free_group barrier) so the |
| 53 | gather loop's credit gate never parks. Mirrors the real server's |
| 54 | per-group barrier: publish/free are keyed by GROUP INDEX, signals count |
| 55 | to the ``begin_sync`` live total, publish_group returns nothing, and freed |
| 56 | groups flow back only through wait_freed / end_sync (see |
| 57 | test_sharded_rdt_producer.TestFakeServerAgreesWithTheRealOne).""" |
| 58 | |
| 59 | def __init__(self, auto_free=True): |
| 60 | self.order: list[str] = [] |
| 61 | self.published: list[int] = [] |
| 62 | self.live_count = 1 |
| 63 | self._inflight_groups: list[int] = [] |
| 64 | self.auto_free = auto_free |
| 65 | self.free_counts: dict[int, int] = {} |
| 66 | self._pending_freed: list[int] = [] |
| 67 | |
| 68 | # Alias for tests that inspect in-flight groups. |
| 69 | @property |
| 70 | def inflight(self): |
| 71 | return self._inflight_groups |
| 72 | |
| 73 | def begin_sync(self, live_count, live_consumer_ids=None): |
| 74 | self.order.append("begin") |
| 75 | self.live_count = max(1, int(live_count)) |
| 76 | self.live_ids = live_consumer_ids |
| 77 | |
| 78 | def publish_group(self, group_idx, entries): |
| 79 | self.order.append("publish") |
| 80 | self.published.append(group_idx) |
| 81 | self._inflight_groups.append(group_idx) |
| 82 | if self.auto_free or self.free_counts.get(group_idx, 0) >= self.live_count: |
| 83 | self._inflight_groups.remove(group_idx) |
| 84 | self._pending_freed.append(group_idx) |
| 85 | |
| 86 | def free_group(self, group_idx): |
| 87 | """Consumer back-edge; may arrive before the group's publish.""" |
| 88 | self.free_counts[group_idx] = self.free_counts.get(group_idx, 0) + 1 |
| 89 | if ( |
| 90 | group_idx in self._inflight_groups |
| 91 | and self.free_counts[group_idx] >= self.live_count |
| 92 | ): |
| 93 | self._inflight_groups.remove(group_idx) |
| 94 | self._pending_freed.append(group_idx) |
| 95 | |
| 96 | def wait_freed(self): |
| 97 | """The engine's credit gate. The real server blocks here; the fake |
| 98 | must already have a banked credit when the gate asks (auto_free, or a |
| 99 | test's own free_group) — anything else is the deadlock the real |
| 100 | watchdog would kill, so fail loudly. Not appended to ``order``: it is |
| 101 | pacing, not a lifecycle milestone.""" |
| 102 | assert self._pending_freed, ( |
| 103 | "wait_freed with nothing freed: the gather loop would deadlock" |
| 104 | ) |
| 105 | freed = self._pending_freed |
| 106 | self._pending_freed = [] |
| 107 | return freed |
| 108 | |
| 109 | def end_sync(self): |
| 110 | self.order.append("end") |
| 111 | freed = self._pending_freed |
| 112 | self._pending_freed = [] |
| 113 | return freed |
| 114 | |
| 115 | def set_gather_error(self, message): |
| 116 | self.order.append("error") |
| 117 | |
| 118 | |
| 119 | def _rdt_engine_with_fake_server( |
| 120 | source, *, is_sender, client, server, monkeypatch, fleet_owned=None |
| 121 | ): |
| 122 | """Build a ShardedRDTTrainerWeightTransferEngine wired to an in-process fake |
| 123 | server (no Ray, no CUDA IPC): bypass trainer_init's spawn, set the |
| 124 | group-major metadata, and route _rpc to the fake.""" |
| 125 | import vllm.distributed.weight_transfer.sharded_rdt_trainer as mod |
| 126 | |
| 127 | # reduce_tensor needs CUDA; the fake server never rebuilds, so stub it. |
| 128 | monkeypatch.setattr(mod, "reduce_tensor", lambda t: (None, ("fake",))) |
| 129 | |
| 130 | init_info = ShardedRDTTrainerInitInfo(num_consumers=1, rank=0 if is_sender else 1) |
| 131 | engine = ShardedRDTTrainerWeightTransferEngine( |
| 132 | client=client, source=source, is_sender=is_sender, init_info=init_info |
| 133 | ) |
| 134 | engine._meta = list(source.metadata()) |
| 135 | names = [m.name for m in engine._meta] |
| 136 | engine._groups = layerwise_groups(names) |
| 137 | engine._server = server |
| 138 | engine._rpc = lambda method, *args: getattr(server, method)(*args) |
| 139 | # What trainer_init resolves from the source's held names + the fleet's |
| 140 | # all-gather. ``fleet_owned`` stands in for that all-gather so a partial- |
| 141 | # ownership rank can be tested without a real process group; the fleet must |
| 142 | # cover every name or the transpose rejects it (nothing would serve the rest). |
| 143 | if fleet_owned is None: |
| 144 | engine._resolve_ownership(1, 0) |
| 145 | else: |
| 146 | # Stands in for the (metadata digest, held-name bitmask) all-gather: |
| 147 | # each rank's group list becomes the mask its held names would set. |
| 148 | names = [m.name for m in engine._meta] |
| 149 | groups = engine._groups |
| 150 | |
| 151 | def _fake_gather(_world, mine): |
| 152 | out = [] |
| 153 | for owned in fleet_owned: |
| 154 | held = {n for gi in owned for n in groups[gi]} |
| 155 | mask = bytearray((len(names) + 7) // 8) |
| 156 | for i, n in enumerate(names): |
| 157 | if n in held: |
| 158 | mask[i >> 3] |= 1 << (i & 7) |
| 159 | out.append((mine[0], bytes(mask))) |
| 160 | return out |
| 161 | |
| 162 | monkeypatch.setattr(engine, "_all_gather_owned", _fake_gather) |
| 163 | engine._resolve_ownership(len(fleet_owned), 0) |
| 164 | return engine |
| 165 | |
| 166 | |
| 167 | def _rdt_source_two_layers(): |
| 168 | return _ListSource( |
| 169 | [ |
| 170 | ("embed.weight", torch.zeros(2)), |
| 171 | ("model.layers.0.w", torch.zeros(2)), |
| 172 | ("model.layers.1.w", torch.zeros(2)), |
| 173 | ("norm.weight", torch.zeros(2)), |
| 174 | ] |
| 175 | ) |
| 176 | |
| 177 | |
| 178 | class TestShardedRDTTrainerInitInfo: |
| 179 | def test_declares_backend(self): |
| 180 | assert ShardedRDTTrainerInitInfo.backend == "sharded_rdt" |
| 181 | |
| 182 | def test_rank_is_keyword_only_and_drives_is_sender(self): |
| 183 | assert ShardedRDTTrainerInitInfo(num_consumers=4, rank=0).is_sender is True |
| 184 | assert ShardedRDTTrainerInitInfo(num_consumers=4, rank=1).is_sender is False |
| 185 | with pytest.raises(TypeError): |
| 186 | # rank is keyword-only. |
| 187 | ShardedRDTTrainerInitInfo(4, 0) # type: ignore[misc] |
| 188 | |
| 189 | def test_registered_in_trainer_factory(self): |
| 190 | cls = WeightTransferTrainerFactory._registry["sharded_rdt"]() |
| 191 | assert cls is ShardedRDTTrainerWeightTransferEngine |
| 192 | |
| 193 | |
| 194 | def test_sharded_rdt_trainer_init_requires_source(): |
| 195 | with pytest.raises(ValueError, match="requires a WeightSource"): |
| 196 | ShardedRDTTrainerWeightTransferEngine.trainer_init( |
| 197 | ShardedRDTTrainerInitInfo(num_consumers=1, rank=0), |
| 198 | client=RecordingClient(), |
| 199 | source=None, |
| 200 | ) |
| 201 | |
| 202 | |
| 203 | def test_sharded_rdt_worker_init_info_is_group_major(monkeypatch): |
| 204 | source = _rdt_source_two_layers() |
| 205 | engine = _rdt_engine_with_fake_server( |
| 206 | source, |
| 207 | is_sender=True, |
| 208 | client=RecordingClient(), |
| 209 | server=_FakeProducerServer(), |
| 210 | monkeypatch=monkeypatch, |
| 211 | ) |
| 212 | worker_init = engine._build_worker_init_info(["srv_rk0"]) |
| 213 | # 4 params -> pre / layer0 / layer1 / post = 4 groups of length 1. |
| 214 | assert worker_init.names == [ |
| 215 | "embed.weight", |
| 216 | "model.layers.0.w", |
| 217 | "model.layers.1.w", |
| 218 | "norm.weight", |
| 219 | ] |
| 220 | assert worker_init.group_lens == [1, 1, 1, 1] |
| 221 | assert worker_init.trainer_actor_names == ["srv_rk0"] |
| 222 | assert worker_init.produce_method_name == "rdt_produce_weights_batched" |
| 223 | assert sum(worker_init.group_lens) == len(worker_init.names) |
| 224 | |
| 225 | |
| 226 | @pytest.mark.skipif( |
| 227 | not torch.cuda.is_available(), |
| 228 | reason="the gather loop's CUDA-IPC export needs a device", |
| 229 | ) |
| 230 | def test_sharded_rdt_send_weights_drives_client_in_order(monkeypatch): |
| 231 | server = _FakeProducerServer(auto_free=True) |
| 232 | client = RecordingClient() |
| 233 | engine = _rdt_engine_with_fake_server( |
| 234 | _rdt_source_two_layers(), |
| 235 | is_sender=True, |
| 236 | client=client, |
| 237 | server=server, |
| 238 | monkeypatch=monkeypatch, |
| 239 | ) |
| 240 | engine.send_weights() |
| 241 | |
| 242 | assert client.order == ["start", "update", "finish"] |
| 243 | # begin, one publish per group (4), end. |
| 244 | assert server.order == ["begin", "publish", "publish", "publish", "publish", "end"] |
| 245 | assert len(server.published) == 4 |
| 246 | # every group freed -> no engine-held refs remain. |
| 247 | assert engine._inflight == {} |
| 248 | |
| 249 | |
| 250 | def test_sharded_rdt_send_weights_group_order_mismatch_raises(monkeypatch): |
| 251 | # Source whose iteration order disagrees with its metadata order. |
| 252 | class _BadSource(_ListSource): |
| 253 | def __iter__(self): |
| 254 | reordered = list(self._pairs) |
| 255 | reordered[0], reordered[1] = reordered[1], reordered[0] |
| 256 | return iter(reordered) |
| 257 | |
| 258 | server = _FakeProducerServer() |
| 259 | engine = _rdt_engine_with_fake_server( |
| 260 | _BadSource(_rdt_source_two_layers()._pairs), |
| 261 | is_sender=True, |
| 262 | client=RecordingClient(), |
| 263 | server=server, |
| 264 | monkeypatch=monkeypatch, |
| 265 | ) |
| 266 | with pytest.raises(RuntimeError, match="iteration order must match"): |
| 267 | engine.send_weights() |
| 268 | assert "error" in server.order # gather error propagated to the server |
| 269 | |
| 270 | |
| 271 | @pytest.mark.skipif( |
| 272 | not torch.cuda.is_available(), |
| 273 | reason="the gather loop's CUDA-IPC export needs a device", |
| 274 | ) |
| 275 | def test_sharded_rdt_non_sender_skips_client(monkeypatch): |
| 276 | class _RaisingClient(RecordingClient): |
| 277 | def start_weight_update(self): |
| 278 | raise AssertionError("non-sender must not touch the client") |
| 279 | |
| 280 | def update_weights(self, update_info): |
| 281 | raise AssertionError("non-sender must not touch the client") |
| 282 | |
| 283 | def finish_weight_update(self, weight_version: str | None = None): |
| 284 | raise AssertionError("non-sender must not touch the client") |
| 285 | |
| 286 | server = _FakeProducerServer(auto_free=True) |
| 287 | client = _RaisingClient() |
| 288 | engine = _rdt_engine_with_fake_server( |
| 289 | _rdt_source_two_layers(), |
| 290 | is_sender=False, |
| 291 | client=client, |
| 292 | server=server, |
| 293 | monkeypatch=monkeypatch, |
| 294 | ) |
| 295 | engine.send_weights() # gathers only; must not raise |
| 296 | assert client.order == [] |
| 297 | assert server.order == ["begin", "publish", "publish", "publish", "publish", "end"] |
| 298 | |
| 299 | |
| 300 | @pytest.mark.skipif( |
| 301 | not torch.cuda.is_available(), |
| 302 | reason="the gather loop's CUDA-IPC export needs a device", |
| 303 | ) |
| 304 | def test_sharded_rdt_send_weights_surfaces_update_error(monkeypatch): |
| 305 | class _FailingUpdateClient(RecordingClient): |
| 306 | def update_weights(self, update_info): |
| 307 | self.order.append("update") |
| 308 | raise RuntimeError("inference side rejected update") |
| 309 | |
| 310 | server = _FakeProducerServer(auto_free=True) |
| 311 | engine = _rdt_engine_with_fake_server( |
| 312 | _rdt_source_two_layers(), |
| 313 | is_sender=True, |
| 314 | client=_FailingUpdateClient(), |
| 315 | server=server, |
| 316 | monkeypatch=monkeypatch, |
| 317 | ) |
| 318 | with pytest.raises(RuntimeError, match="inference side rejected update"): |
| 319 | engine.send_weights() |
| 320 | |
| 321 | |
| 322 | class _OwnedSource(_ListSource): |
| 323 | """A source holding only some groups' names, like a pipeline-parallel rank. |
| 324 | Takes group indices for convenience and declares the names inside them.""" |
| 325 | |
| 326 | def __init__(self, pairs, owned_group_idx): |
| 327 | super().__init__(pairs) |
| 328 | self._owned = list(owned_group_idx) |
| 329 | groups = layerwise_groups([n for n, _ in pairs]) |
| 330 | self._owned_names = [n for gi in self._owned for n in groups[gi]] |
| 331 | |
| 332 | def held_names(self): |
| 333 | return list(self._owned_names) |
| 334 | |
| 335 | def __iter__(self): |
| 336 | by_name = dict(self._pairs) |
| 337 | return iter([(n, by_name[n]) for n in self._owned_names]) |
| 338 | |
| 339 | |
| 340 | @pytest.mark.skipif( |
| 341 | not torch.cuda.is_available(), |
| 342 | reason="the gather loop's CUDA-IPC export needs a device", |
| 343 | ) |
| 344 | def test_sharded_rdt_publishes_only_the_groups_it_holds(monkeypatch): |
| 345 | server = _FakeProducerServer(auto_free=True) |
| 346 | engine = _rdt_engine_with_fake_server( |
| 347 | _OwnedSource(_rdt_source_two_layers()._pairs, [1, 2]), |
| 348 | is_sender=False, |
| 349 | client=RecordingClient(), |
| 350 | server=server, |
| 351 | monkeypatch=monkeypatch, |
| 352 | fleet_owned=[[1, 2], [0, 3]], # this rank holds the layers, rank 1 the rest |
| 353 | ) |
| 354 | engine.send_weights() |
| 355 | |
| 356 | assert server.published == [1, 2] |
| 357 | assert server.order == ["begin", "publish", "publish", "end"] |
| 358 | assert engine._inflight == {} |
| 359 | # embed + norm on rank 1, the two layers here: two distinct owner sets, |
| 360 | # numbered by first appearance in metadata order. |
| 361 | assert engine._owner_sets == [[1], [0]] |
| 362 | assert engine._name_owner_class == [0, 1, 1, 0] |
| 363 | |
| 364 | |
| 365 | def test_sharded_rdt_owned_group_order_mismatch_raises(monkeypatch): |
| 366 | class _MisorderedOwned(_OwnedSource): |
| 367 | def __iter__(self): |
| 368 | return iter(list(super().__iter__())[::-1]) |
| 369 | |
| 370 | server = _FakeProducerServer(auto_free=True) |
| 371 | engine = _rdt_engine_with_fake_server( |
| 372 | _MisorderedOwned(_rdt_source_two_layers()._pairs, [1, 2]), |
| 373 | is_sender=False, |
| 374 | client=RecordingClient(), |
| 375 | server=server, |
| 376 | monkeypatch=monkeypatch, |
| 377 | fleet_owned=[[1, 2], [0, 3]], |
| 378 | ) |
| 379 | with pytest.raises(RuntimeError, match="iteration order must match"): |
| 380 | engine.send_weights() |
| 381 | assert "error" in server.order |
| 382 | |
| 383 | |
| 384 | @pytest.mark.skipif( |
| 385 | not torch.cuda.is_available(), |
| 386 | reason="the gather loop's CUDA-IPC export needs a device", |
| 387 | ) |
| 388 | def test_sharded_rdt_begin_sync_carries_the_live_count(monkeypatch): |
| 389 | """The free barrier's target is one integer per sync: the live consumer |
| 390 | count, defaulting to the whole provisioned fleet.""" |
| 391 | server = _FakeProducerServer(auto_free=True) |
| 392 | engine = _rdt_engine_with_fake_server( |
| 393 | _rdt_source_two_layers(), |
| 394 | is_sender=False, |
| 395 | client=RecordingClient(), |
| 396 | server=server, |
| 397 | monkeypatch=monkeypatch, |
| 398 | ) |
| 399 | engine.send_weights() |
| 400 | assert server.live_count == 1 # num_consumers=1 in the harness |
| 401 | engine.send_weights(live_consumer_ids=[0]) |
| 402 | assert server.live_count == 1 |
| 403 | |
| 404 | |
| 405 | class TestLiveCountPlumbing: |
| 406 | """``send_weights(live_consumer_ids)`` -> the barrier target. The |
| 407 | provisioned geometry is frozen; a degraded sync only lowers the target and |
| 408 | narrows each slot-sharing group's rendezvous.""" |
| 409 | |
| 410 | @staticmethod |
| 411 | def _engine(num_consumers, world=2, rank=0): |
| 412 | engine = ShardedRDTTrainerWeightTransferEngine.__new__( |
| 413 | ShardedRDTTrainerWeightTransferEngine |
| 414 | ) |
| 415 | engine._init_info = ShardedRDTTrainerInitInfo( |
| 416 | rank=rank, num_consumers=num_consumers |
| 417 | ) |
| 418 | engine.source = object() # send_weights asserts a source is present |
| 419 | received: list = [] |
| 420 | engine._send_weights_inner = lambda count, ids: received.append((count, ids)) |
| 421 | return engine, received |
| 422 | |
| 423 | def test_none_counts_the_whole_provisioned_fleet(self): |
| 424 | engine, got = self._engine(num_consumers=8) |
| 425 | engine.send_weights(None) |
| 426 | assert got == [(8, list(range(8)))] |
| 427 | |
| 428 | def test_a_live_set_counts_its_distinct_members(self): |
| 429 | engine, got = self._engine(num_consumers=8) |
| 430 | engine.send_weights([0, 1, 4, 5, 5]) |
| 431 | assert got == [(4, [0, 1, 4, 5])] |
| 432 | |
| 433 | def test_a_full_live_set_matches_the_provisioned_count(self): |
| 434 | engine, got = self._engine(num_consumers=8) |
| 435 | engine.send_weights(list(range(8))) |
| 436 | engine.send_weights(None) |
| 437 | assert got == [(8, list(range(8)))] * 2 |
| 438 | |
| 439 | def test_which_consumers_died_does_not_matter_only_how_many(self): |
| 440 | """The whole point of the barrier: no routed per-producer targets, so |
| 441 | the identity of the dead consumer is irrelevant to the producers.""" |
| 442 | counts = [] |
| 443 | for live in ([0, 1, 2, 3], [4, 5, 6, 7], [0, 2, 4, 6]): |
| 444 | engine, got = self._engine(num_consumers=8) |
| 445 | engine.send_weights(live) |
| 446 | counts += [c for c, _ids in got] |
| 447 | assert counts == [4, 4, 4] |
| 448 | |
| 449 | def test_the_live_ids_travel_with_the_count(self): |
| 450 | """The count sizes the free barrier, the ids size the slot-sharing |
| 451 | rendezvous, so a producer that shares slots can tell WHICH consumers it |
| 452 | is still waiting for. They must describe the same set.""" |
| 453 | engine, got = self._engine(num_consumers=8) |
| 454 | engine.send_weights([6, 0, 2, 2]) |
| 455 | ((count, ids),) = got |
| 456 | assert ids == [0, 2, 6] and count == len(ids) |
| 457 | |
| 458 | |
| 459 | def test_sharded_rdt_worker_init_info_carries_the_ownership_table(monkeypatch): |
| 460 | import json |
| 461 | |
| 462 | server = _FakeProducerServer(auto_free=True) |
| 463 | engine = _rdt_engine_with_fake_server( |
| 464 | _OwnedSource(_rdt_source_two_layers()._pairs, [0, 1]), |
| 465 | is_sender=True, |
| 466 | client=RecordingClient(), |
| 467 | server=server, |
| 468 | monkeypatch=monkeypatch, |
| 469 | fleet_owned=[[0, 1], [2, 3]], |
| 470 | ) |
| 471 | worker_init = engine._build_worker_init_info(["srv_rk0", "srv_rk1"]) |
| 472 | assert worker_init.owner_sets == [[0], [1]] |
| 473 | assert worker_init.name_owner_class == [0, 0, 1, 1] |
| 474 | assert len(worker_init.name_owner_class) == len(worker_init.names) |
| 475 | # The payload crosses the control plane as JSON: nested lists must survive. |
| 476 | round_tripped = json.loads(json.dumps(asdict(worker_init))) |
| 477 | assert round_tripped["owner_sets"] == [[0], [1]] |
| 478 | assert round_tripped["name_owner_class"] == [0, 0, 1, 1] |
| 479 | |
| 480 | |
| 481 | def test_weight_source_holds_everything_by_default(): |
| 482 | """The contract's default: a source that says nothing holds the whole model.""" |
| 483 | src = _rdt_source_two_layers() |
| 484 | assert src.held_names() is None |
| 485 | assert ModuleSource(torch.nn.Linear(2, 2)).held_names() is None |
| 486 | |
| 487 | |
| 488 | def test_sharded_rdt_rejects_a_held_name_outside_metadata(monkeypatch): |
| 489 | class _BadHeld(_OwnedSource): |
| 490 | def held_names(self): |
| 491 | return ["embed.weight", "not.a.real.name"] |
| 492 | |
| 493 | with pytest.raises(ValueError, match="not"): |
| 494 | _rdt_engine_with_fake_server( |
| 495 | _BadHeld(_rdt_source_two_layers()._pairs, [0]), |
| 496 | is_sender=False, |
| 497 | client=RecordingClient(), |
| 498 | server=_FakeProducerServer(), |
| 499 | monkeypatch=monkeypatch, |
| 500 | ) |
| 501 | |
| 502 | |
| 503 | def test_sharded_rdt_rejects_a_rank_holding_nothing(monkeypatch): |
| 504 | class _HoldsNothing(_ListSource): |
| 505 | def held_names(self): |
| 506 | return [] |
| 507 | |
| 508 | with pytest.raises(ValueError, match="empty"): |
| 509 | _rdt_engine_with_fake_server( |
| 510 | _HoldsNothing(_rdt_source_two_layers()._pairs), |
| 511 | is_sender=False, |
| 512 | client=RecordingClient(), |
| 513 | server=_FakeProducerServer(), |
| 514 | monkeypatch=monkeypatch, |
| 515 | ) |
| 516 | |
| 517 | |
| 518 | def test_sharded_rdt_rejects_a_name_no_rank_holds(monkeypatch): |
| 519 | """Every name must be held somewhere or it can never be served — caught |
| 520 | when the holdings are transposed, naming the orphan.""" |
| 521 | with pytest.raises(ValueError, match="no trainer rank holds"): |
| 522 | _rdt_engine_with_fake_server( |
| 523 | _OwnedSource(_rdt_source_two_layers()._pairs, [0, 1]), |
| 524 | is_sender=False, |
| 525 | client=RecordingClient(), |
| 526 | server=_FakeProducerServer(), |
| 527 | monkeypatch=monkeypatch, |
| 528 | fleet_owned=[[0, 1], [2]], # nobody holds group 3 |
| 529 | ) |
| 530 | |
| 531 | |
| 532 | def test_sharded_rdt_rejects_metadata_disagreement_across_ranks(monkeypatch): |
| 533 | """Only the sender's metadata reaches the consumers, so a rank describing |
| 534 | just its own share must fail loudly rather than silently drop the rest.""" |
| 535 | engine = _rdt_engine_with_fake_server( |
| 536 | _OwnedSource(_rdt_source_two_layers()._pairs, [1, 2]), |
| 537 | is_sender=False, |
| 538 | client=RecordingClient(), |
| 539 | server=_FakeProducerServer(auto_free=True), |
| 540 | monkeypatch=monkeypatch, |
| 541 | fleet_owned=[[1, 2], [0, 3]], # a covering fleet, so construction succeeds |
| 542 | ) |
| 543 | # Now rank 1 reports a DIFFERENT metadata digest for the same model. |
| 544 | monkeypatch.setattr( |
| 545 | engine, |
| 546 | "_all_gather_owned", |
| 547 | lambda w, mine: [mine, ("deadbeefdeadbeef", mine[1])], |
| 548 | ) |
| 549 | with pytest.raises(ValueError, match="disagrees across trainer ranks"): |
| 550 | engine._resolve_ownership(2, 0) |
| 551 | |
| 552 | |
| 553 | def _serve_ring_server(src_name, src): |
| 554 | """A producer server with one cached tensor and a pre-seeded serve ring, so a |
| 555 | pull needs no Ray and no NIXL registration. Returns (server, serve) where |
| 556 | ``serve(chain)`` packs one spec into the SAME ring slot every time — which is |
| 557 | what puts the destination-view cache, and only it, under test.""" |
| 558 | from vllm.distributed.weight_transfer.sharded_rdt_trainer import ( |
| 559 | _RDTProducerServer, |
| 560 | ) |
| 561 | |
| 562 | srv = _RDTProducerServer( |
| 563 | num_rdt_buffers=2, |
| 564 | buffer_presize_gb=0.0, |
| 565 | gather_lookahead=2, |
| 566 | ) |
| 567 | srv._cache[src_name] = src |
| 568 | srv._serve_rings[0] = [ |
| 569 | torch.empty(1 << 16, dtype=torch.uint8, device="cuda") for _ in range(2) |
| 570 | ] |
| 571 | |
| 572 | # The slot is ``seq % nring``, so stepping seq by nring keeps every pull in |
| 573 | # slot 0 while still making each one its own generation -- reusing a seq |
| 574 | # would rendezvous with the finished pack and return it unchanged. |
| 575 | seqs = itertools.count(0, srv._nring) |
| 576 | |
| 577 | def serve(chain): |
| 578 | return srv.rdt_produce_weights_batched( |
| 579 | [(src_name, chain)], consumer_id=0, seq=next(seqs) |
| 580 | )[0] |
| 581 | |
| 582 | return srv, serve |
| 583 | |
| 584 | |
| 585 | @pytest.mark.skipif( |
| 586 | not torch.cuda.is_available(), reason="the producer server needs a CUDA device" |
| 587 | ) |
| 588 | def test_serve_does_not_reuse_packed_views_of_another_shape(): |
| 589 | """Two requests can share a name yet pack different slices of it. |
| 590 | |
| 591 | The producer caches the destination views it carves into a serve ring slot. |
| 592 | Keyed by name alone, the second request is packed through the first's views. |
| 593 | Reachable when one name's copies split across owner-class chunks. |
| 594 | """ |
| 595 | name = "model.layers.0.w" |
| 596 | src = torch.arange(64, dtype=torch.bfloat16, device="cuda").reshape(8, 8) |
| 597 | _srv, serve = _serve_ring_server(name, src) |
| 598 | |
| 599 | serve((("narrow", (0, 0, 2), ()),)) # 2 rows |
| 600 | wide = src.narrow(0, 0, 6) # 6 rows, same name, same slot |
| 601 | blob = serve((("narrow", (0, 0, 6), ()),)) |
| 602 | |
| 603 | got = blob[: wide.numel() * wide.element_size()].view(wide.dtype).reshape(6, 8) |
| 604 | assert torch.equal(got, wide) |
| 605 | |
| 606 | |
| 607 | @pytest.mark.skipif( |
| 608 | not torch.cuda.is_available(), reason="the producer server needs a CUDA device" |
| 609 | ) |
| 610 | def test_serve_does_not_reuse_packed_views_of_another_dtype(): |
| 611 | """The SILENT case of the same cache: two requests whose slices have the same |
| 612 | name and the same shape but different dtypes pack at identical offsets, so |
| 613 | reusing the stale views raises nothing — ``copy_`` just casts, and the blob |
| 614 | carries the wrong bytes with no check downstream. |
| 615 | |
| 616 | ``view(dtype)`` rather than ``to(dtype)`` because the chain is replayed under |
| 617 | ``ALLOWED_OPS``, which rejects ``to``: same-itemsize dtypes keep the byte |
| 618 | length identical, which is exactly what makes the stale view fit. |
| 619 | """ |
| 620 | name = "model.layers.0.w" |
| 621 | src = torch.arange(64, dtype=torch.bfloat16, device="cuda").reshape(8, 8) |
| 622 | _srv, serve = _serve_ring_server(name, src) |
| 623 | |
| 624 | serve((("view", (torch.float16,), ()),)) # same shape and bytes, fp16 |
| 625 | blob = serve(()) # same name and shape, bf16 |
| 626 | |
| 627 | got = blob[: src.numel() * src.element_size()].view(src.dtype).reshape(8, 8) |
| 628 | assert torch.equal(got, src), "packed through a cached view of the wrong dtype" |