1754 lines
73 KiB
Python
1754 lines
73 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
|
"""Consumer-side unit tests for the sharded-RDT backend.
|
|
|
|
The consumer engine's planning core is pure — it turns the baked plan plus the
|
|
driver's gather-group partition into a static `_CallPlan` with no pulls, no
|
|
model and no Ray — so most of it is exercised here on CPU/meta tensors only.
|
|
These tests pin the values the engine produces today (chunk boundaries, packed
|
|
byte offsets, recorded op chains) so a refactor that changes them fails loudly
|
|
instead of silently shipping different bytes.
|
|
|
|
`TestBakeOnARealModel` at the end is the exception: it needs a GPU, because the
|
|
bake only exists to be driven through a real `model.load_weights`, and the parts
|
|
that matter there — that the monkeypatched loaders are put back, and that the
|
|
recorded op chains reproduce what the loaders did — cannot be faked.
|
|
|
|
`tests/distributed/test_sharded_rdt_trainer.py` covers the trainer (producer)
|
|
side; `test_sharded_rdt_producer.py` covers its serve-actor protocol.
|
|
"""
|
|
|
|
from types import SimpleNamespace
|
|
|
|
import pytest
|
|
import torch
|
|
|
|
from vllm.distributed.weight_transfer.base import layerwise_groups
|
|
from vllm.distributed.weight_transfer.sharded_rdt_common import (
|
|
ALLOWED_OPS,
|
|
SUPPORTED_OPS,
|
|
RdtRouter,
|
|
assign_producer_indices,
|
|
buffer_alloc_bytes,
|
|
)
|
|
from vllm.distributed.weight_transfer.sharded_rdt_engine import (
|
|
ShardedRDTWeightTransferEngine,
|
|
ShardedRDTWeightTransferInitInfo,
|
|
_dtype_from_name,
|
|
)
|
|
from vllm.distributed.weight_transfer.sharded_rdt_fake import (
|
|
BakeSink,
|
|
FakeRDTTensor,
|
|
_Scatter,
|
|
_UnsupportedFakeOp,
|
|
)
|
|
|
|
META = torch.device("meta")
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Harness
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
class _FakeLayer:
|
|
"""Stand-in for a leaf module. The planner only stores it and keys on
|
|
``id()``; nothing in the plan path touches its attributes."""
|
|
|
|
def __init__(self, tag: str) -> None:
|
|
self.tag = tag
|
|
|
|
def __repr__(self) -> str:
|
|
return f"_FakeLayer({self.tag!r})"
|
|
|
|
|
|
def _module(layer, copies):
|
|
"""What the bake records for one leaf module: a scatter list whose entries
|
|
all carry that module as their ``layer``."""
|
|
for c in copies:
|
|
c.layer = layer
|
|
return copies
|
|
|
|
|
|
def _copy(
|
|
name,
|
|
layer_param="weight",
|
|
*,
|
|
offset=0,
|
|
shape=(4,),
|
|
ops=(),
|
|
layer=None,
|
|
dtype=torch.bfloat16,
|
|
):
|
|
"""A `_Scatter` as the bake would record it: source key, owning layer, and
|
|
the meta destination region."""
|
|
stride = (
|
|
(1,)
|
|
if len(shape) == 1
|
|
else tuple(
|
|
int(torch.empty(shape, device=META).stride()[i]) for i in range(len(shape))
|
|
)
|
|
)
|
|
return _Scatter(
|
|
layer=layer,
|
|
param_name=layer_param,
|
|
src=(name, tuple(ops)),
|
|
offset=offset,
|
|
shape=tuple(shape),
|
|
stride=stride,
|
|
dtype=dtype,
|
|
)
|
|
|
|
|
|
def _planner(baked, *, name_meta, live=None, held_by=None, group_lens=None):
|
|
"""A planning-only engine.
|
|
|
|
`__init__` needs a config/model/device the planner never reads, so build the
|
|
instance without it and set exactly the state `_build_call_plan` consumes:
|
|
the baked plan, the name metadata, the live-name set, and a bound router.
|
|
|
|
``held_by``: name -> the ranks holding it. Names left out are held by every
|
|
producer (the replicated class 0). Default: one producer holding everything,
|
|
the degenerate router. Stand-in ranks 0-3 let plan tests assert chunk
|
|
structure; owner VALUES are pinned by TestCallPlanRouting.
|
|
"""
|
|
names = list(name_meta)
|
|
if held_by:
|
|
n_prod = max(r for owners in held_by.values() for r in owners) + 1
|
|
owner_sets = [list(range(n_prod))]
|
|
seen = {tuple(range(n_prod)): 0}
|
|
classes = []
|
|
for n in names:
|
|
key = tuple(sorted(held_by.get(n, range(n_prod))))
|
|
if key not in seen:
|
|
seen[key] = len(owner_sets)
|
|
owner_sets.append(list(key))
|
|
classes.append(seen[key])
|
|
else:
|
|
n_prod, owner_sets, classes = 1, [[0]], [0] * len(names)
|
|
eng = object.__new__(ShardedRDTWeightTransferEngine)
|
|
eng._name_to_plan = dict(baked)
|
|
eng._name_meta = dict(name_meta)
|
|
eng._live_names = set(live or name_meta)
|
|
router = RdtRouter(
|
|
n_prod, 1, owner_sets, classes, names, group_lens or [1] * len(names)
|
|
)
|
|
router.bind([object()] * n_prod, [object()] * n_prod, 0)
|
|
eng._router = router
|
|
return eng
|
|
|
|
|
|
def _one_module_per_layer(n_layers, *, dtype="bfloat16", numel=4):
|
|
"""`(baked, name_meta, names, group_lens)` for a pre / N-layer / post model
|
|
where every gather group is one fully-baked leaf module."""
|
|
names = ["embed.weight"]
|
|
names += [f"model.layers.{i}.w" for i in range(n_layers)]
|
|
names += ["norm.weight"]
|
|
baked = {}
|
|
for n in names:
|
|
layer = _FakeLayer(n)
|
|
baked[n] = _module(layer, [_copy(n, shape=(numel,))])
|
|
name_meta = {n: (dtype, [numel]) for n in names}
|
|
group_lens = [1] * len(names)
|
|
return baked, name_meta, names, group_lens
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# FakeRDTTensor: op-chain recording
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
class TestFakeOpChains:
|
|
"""Every allowlisted op must append itself to the chain and hand back a
|
|
child whose shape/dtype PyTorch itself computed. The chain is the wire
|
|
format the producer replays, so its exact contents are load-bearing."""
|
|
|
|
def _fake(self, shape=(4, 6), dtype=torch.bfloat16, sink=None):
|
|
return FakeRDTTensor(
|
|
name="w", shape=torch.Size(shape), dtype=dtype, device=META, sink=sink
|
|
)
|
|
|
|
def test_bare_fake_has_metadata_but_no_chain(self):
|
|
t = self._fake()
|
|
assert t.shape == (4, 6)
|
|
assert t.dtype is torch.bfloat16
|
|
assert t._key() == ("w", ())
|
|
|
|
@pytest.mark.parametrize(
|
|
"call,expect_op,expect_shape",
|
|
[
|
|
(lambda t: t.narrow(0, 1, 2), ("narrow", (0, 1, 2), ()), (2, 6)),
|
|
(lambda t: t.view(6, 4), ("view", (6, 4), ()), (6, 4)),
|
|
(lambda t: t.reshape(24), ("reshape", (24,), ()), (24,)),
|
|
(lambda t: t[1], ("__getitem__", (1,), ()), (6,)),
|
|
(lambda t: t.unsqueeze(0), ("unsqueeze", (0,), ()), (1, 4, 6)),
|
|
(lambda t: t.transpose(0, 1), ("transpose", (0, 1), ()), (6, 4)),
|
|
(lambda t: t.t(), ("t", (), ()), (6, 4)),
|
|
(lambda t: t.permute(1, 0), ("permute", (1, 0), ()), (6, 4)),
|
|
(lambda t: t.flatten(), ("flatten", (), ()), (24,)),
|
|
(lambda t: t.contiguous(), ("contiguous", (), ()), (4, 6)),
|
|
],
|
|
)
|
|
def test_single_return_op_appends_one_spec(self, call, expect_op, expect_shape):
|
|
child = call(self._fake())
|
|
assert isinstance(child, FakeRDTTensor)
|
|
assert child._key() == ("w", (expect_op,))
|
|
assert tuple(child.shape) == expect_shape
|
|
|
|
def test_squeeze_records_its_argument(self):
|
|
child = self._fake(shape=(1, 4)).squeeze(0)
|
|
assert child._key() == ("w", (("squeeze", (0,), ()),))
|
|
assert tuple(child.shape) == (4,)
|
|
|
|
def test_chains_compose_in_call_order(self):
|
|
child = self._fake().t().narrow(0, 0, 3).flatten()
|
|
assert child._key() == (
|
|
"w",
|
|
(
|
|
("t", (), ()),
|
|
("narrow", (0, 0, 3), ()),
|
|
("flatten", (), ()),
|
|
),
|
|
)
|
|
assert tuple(child.shape) == (12,)
|
|
|
|
def test_kwargs_are_frozen_sorted_for_hashability(self):
|
|
child = self._fake().narrow(dim=0, start=1, length=2)
|
|
(op, args, kwargs) = child._key()[1][0]
|
|
assert (op, args) == ("narrow", ())
|
|
assert kwargs == (("dim", 0), ("length", 2), ("start", 1))
|
|
hash(child._key()) # the chain is used as a dict key on both sides
|
|
|
|
@pytest.mark.parametrize("op,n", [("chunk", 2), ("unbind", 4)])
|
|
def test_multi_return_op_emits_one_child_per_output(self, op, n):
|
|
"""chunk/unbind hand back a tuple; each child carries the base op plus a
|
|
trailing __getitem__(i) so the producer can index the replayed result."""
|
|
parts = getattr(self._fake(), op)(*((n,) if op == "chunk" else ()), 0)
|
|
assert isinstance(parts, tuple)
|
|
assert len(parts) == n
|
|
for i, part in enumerate(parts):
|
|
base, index = part._key()[1]
|
|
assert base[0] == op
|
|
assert index == ("__getitem__", (i,), ())
|
|
|
|
def test_op_chain_is_hashable_as_a_fetch_key(self):
|
|
keys = {self._fake().t()._key(), self._fake().t()._key()}
|
|
assert len(keys) == 1, "equal chains must collapse — they dedup pull keys"
|
|
|
|
|
|
class TestFakeUnsupportedOps:
|
|
"""Anything that needs real data must fail loudly at bake time rather than
|
|
silently transferring the wrong bytes."""
|
|
|
|
def _fake(self):
|
|
return FakeRDTTensor(
|
|
name="w",
|
|
shape=torch.Size((4,)),
|
|
dtype=torch.bfloat16,
|
|
device=META,
|
|
sink=None,
|
|
)
|
|
|
|
@pytest.mark.parametrize(
|
|
"call",
|
|
[
|
|
lambda t: t.to(torch.float32),
|
|
lambda t: t.float(),
|
|
lambda t: t + 1,
|
|
lambda t: t * 2,
|
|
lambda t: t.sum(),
|
|
],
|
|
ids=["to", "float", "add", "mul", "sum"],
|
|
)
|
|
def test_data_dependent_ops_raise(self, call):
|
|
with pytest.raises(_UnsupportedFakeOp):
|
|
call(self._fake())
|
|
|
|
def test_error_names_the_weight_and_the_chain(self):
|
|
with pytest.raises(_UnsupportedFakeOp) as exc:
|
|
self._fake().narrow(0, 0, 2).float()
|
|
msg = str(exc.value)
|
|
assert "'w'" in msg
|
|
assert "narrow" in msg
|
|
|
|
def test_unsupported_is_a_notimplementederror(self):
|
|
"""Callers distinguish "this backend can't handle the loader" from bugs."""
|
|
assert issubclass(_UnsupportedFakeOp, NotImplementedError)
|
|
|
|
|
|
class TestBakeRecording:
|
|
"""During the dry run the fake's ``copy_`` is the data sink: it records the
|
|
source chain plus the meta destination's strided region and moves nothing."""
|
|
|
|
def _recorder_and_fake(self, shape=(4, 6)):
|
|
rec = BakeSink()
|
|
fake = FakeRDTTensor(
|
|
name="w",
|
|
shape=torch.Size(shape),
|
|
dtype=torch.bfloat16,
|
|
device=META,
|
|
sink=rec,
|
|
)
|
|
return rec, fake
|
|
|
|
def test_copy_records_the_destination_region(self):
|
|
rec, fake = self._recorder_and_fake()
|
|
layer = _FakeLayer("q_proj")
|
|
param = torch.empty((8, 6), dtype=torch.bfloat16, device=META)
|
|
dest = param.narrow(0, 4, 4) # the second half of a fused param
|
|
rec.current = (layer, "weight")
|
|
dest.copy_(fake)
|
|
|
|
(recorded,) = rec.copies_by_layer[layer]
|
|
assert recorded.src == ("w", ())
|
|
assert recorded.param_name == "weight"
|
|
assert recorded.offset == dest.storage_offset() == 24
|
|
assert recorded.shape == (4, 6)
|
|
assert recorded.stride == (6, 1)
|
|
|
|
def test_copy_marks_the_source_name_live(self):
|
|
rec, fake = self._recorder_and_fake()
|
|
rec.current = (_FakeLayer("l"), "weight")
|
|
torch.empty((4, 6), dtype=torch.bfloat16, device=META).copy_(fake)
|
|
assert rec.copied_names == {"w"}
|
|
|
|
def test_unattributed_copy_is_live_but_unrecorded(self):
|
|
"""A copy_ with no loader stamp cannot be attributed to a param, so its
|
|
module must fall back to the plain load — but the name still moved data."""
|
|
rec, fake = self._recorder_and_fake()
|
|
rec.current = None
|
|
torch.empty((4, 6), dtype=torch.bfloat16, device=META).copy_(fake)
|
|
assert rec.copied_names == {"w"}
|
|
assert dict(rec.copies_by_layer) == {}
|
|
|
|
def test_copies_are_grouped_by_module_in_call_order(self):
|
|
rec, _ = self._recorder_and_fake()
|
|
layer = _FakeLayer("gate_up")
|
|
for i, name in enumerate(("gate", "up")):
|
|
fake = FakeRDTTensor(
|
|
name=name,
|
|
shape=torch.Size((4,)),
|
|
dtype=torch.bfloat16,
|
|
device=META,
|
|
sink=rec,
|
|
)
|
|
param = torch.empty((8,), dtype=torch.bfloat16, device=META)
|
|
rec.current = (layer, "weight")
|
|
param.narrow(0, 4 * i, 4).copy_(fake)
|
|
assert [c.src[0] for c in rec.copies_by_layer[layer]] == ["gate", "up"]
|
|
assert [c.offset for c in rec.copies_by_layer[layer]] == [0, 4]
|
|
|
|
def test_recording_a_sliced_source(self):
|
|
"""The chain on the source and the region on the dest are independent."""
|
|
rec, fake = self._recorder_and_fake(shape=(8, 6))
|
|
layer = _FakeLayer("k_proj")
|
|
rec.current = (layer, "weight")
|
|
param = torch.empty((4, 6), dtype=torch.bfloat16, device=META)
|
|
param.copy_(fake.narrow(0, 2, 4))
|
|
(recorded,) = rec.copies_by_layer[layer]
|
|
assert recorded.src == ("w", (("narrow", (0, 2, 4), ()),))
|
|
assert recorded.shape == (4, 6)
|
|
|
|
def test_the_recorded_dtype_is_the_chains_output_not_the_sources(self):
|
|
"""``view(dtype)`` is allowlisted, so a chain can reinterpret dtype. The
|
|
producer packs what the replay yields, so the record must carry the
|
|
POST-chain dtype: taking it from the source name's metadata instead sizes
|
|
the slice with the wrong itemsize and shifts every later slice in the
|
|
chunk, carving the packed blob differently on the two sides."""
|
|
rec, fake = self._recorder_and_fake(shape=(4,))
|
|
layer = _FakeLayer("reinterpreted")
|
|
rec.current = (layer, "weight")
|
|
viewed = fake.view(torch.float32) # 4 x bf16 -> 2 x f32
|
|
param = torch.empty((2,), dtype=torch.float32, device=META)
|
|
param.copy_(viewed)
|
|
|
|
(recorded,) = rec.copies_by_layer[layer]
|
|
assert recorded.dtype is torch.float32, "recorded the source dtype"
|
|
assert recorded.shape == (2,)
|
|
# What the source name alone would have said, and what the producer
|
|
# would NOT have sent.
|
|
assert 2 * torch.float32.itemsize == 8 != 2 * torch.bfloat16.itemsize
|
|
|
|
def test_a_broadcasting_copy_is_refused_at_bake(self):
|
|
"""The packed layout sizes each slice from the DESTINATION shape while
|
|
the producer packs what the chain yields, and nothing downstream can
|
|
catch a mismatch — the consumer's buffer view is exactly prod(dest.shape)
|
|
elements, so it reshapes cleanly over bytes laid out at other offsets.
|
|
So the bake refuses rather than recording a slice it cannot carve."""
|
|
rec, fake = self._recorder_and_fake(shape=(1, 6))
|
|
rec.current = (_FakeLayer("broadcast"), "weight")
|
|
param = torch.empty((4, 6), dtype=torch.bfloat16, device=META)
|
|
with pytest.raises(_UnsupportedFakeOp, match="broadcasting"):
|
|
param.copy_(fake)
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Consumer identity
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
class TestConsumerIdentity:
|
|
"""Every worker in the fleet needs a DISTINCT id in 0..C-1: it selects the
|
|
worker's producer block and keys the producer's per-consumer serve ring, so
|
|
a collision silently serves two workers out of one ring."""
|
|
|
|
def _engine(self, *, dp_index, rank, world_size):
|
|
eng = object.__new__(ShardedRDTWeightTransferEngine)
|
|
eng.parallel_config = SimpleNamespace(
|
|
data_parallel_index=dp_index, rank=rank, world_size=world_size
|
|
)
|
|
return eng
|
|
|
|
def _ids(self, *, num_consumers, workers, replica_rank=0, num_replicas=1):
|
|
"""Consumer ids for a DP-only engine whose workers are ``dp_index``es."""
|
|
info = ShardedRDTWeightTransferInitInfo(
|
|
num_consumers=num_consumers,
|
|
replica_rank=replica_rank,
|
|
num_replicas=num_replicas,
|
|
)
|
|
out = []
|
|
for dp_index in workers:
|
|
eng = self._engine(dp_index=dp_index, rank=0, world_size=1)
|
|
eng._num_consumers_override = num_consumers
|
|
out.append(eng._resolve_consumer_id(info))
|
|
return out
|
|
|
|
def test_one_engine_indexes_by_dp_and_tp(self):
|
|
"""dense-via-TP and MoE-via-DP+EP both flatten to 0..C-1."""
|
|
eng = self._engine(dp_index=0, rank=3, world_size=8)
|
|
eng._num_consumers_override = 8
|
|
info = ShardedRDTWeightTransferInitInfo(num_consumers=8)
|
|
assert eng._resolve_consumer_id(info) == 3
|
|
|
|
eng = self._engine(dp_index=5, rank=0, world_size=1)
|
|
eng._num_consumers_override = 8
|
|
assert eng._resolve_consumer_id(info) == 5
|
|
|
|
def test_independent_engines_offset_into_distinct_ranges(self):
|
|
"""A fleet of separate engines restarts _global_worker_index at 0 in
|
|
each, so without the replica offset every engine would claim 0..w-1.
|
|
The driver sets replica_rank on the payload; these are the ids it buys."""
|
|
ids = [
|
|
self._ids(num_consumers=8, workers=range(4), replica_rank=r, num_replicas=2)
|
|
for r in (0, 1)
|
|
]
|
|
assert ids == [[0, 1, 2, 3], [4, 5, 6, 7]]
|
|
assert len(set(ids[0] + ids[1])) == 8, "every worker distinct fleet-wide"
|
|
|
|
def test_the_default_is_a_single_engine_with_no_offset(self):
|
|
assert self._ids(num_consumers=4, workers=range(4)) == [0, 1, 2, 3]
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# RdtRouter: per-name ownership -> one producer per pull
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
def _router(
|
|
num_producers,
|
|
num_consumers,
|
|
owner_sets=None,
|
|
name_owner_class=None,
|
|
names=None,
|
|
group_lens=None,
|
|
):
|
|
"""A router over ``n`` single-name groups unless told otherwise, so a test
|
|
can state expectations in group indices."""
|
|
if names is None:
|
|
n = len(name_owner_class or []) or (len(group_lens or []) or 1)
|
|
names = [f"g{i}" for i in range(n)]
|
|
if group_lens is None:
|
|
group_lens = [1] * len(names)
|
|
return RdtRouter(
|
|
num_producers, num_consumers, owner_sets, name_owner_class, names, group_lens
|
|
)
|
|
|
|
|
|
class TestRdtRouter:
|
|
"""Who serves each name.
|
|
|
|
A wrong answer is not a wrong number but a hang or a loud misroute: a
|
|
consumer pulling from a producer that never gathered the name trips its
|
|
served-names guard. Freeing does NOT route this way — every consumer signals
|
|
every owner of a group and the producers count to the live total — so these
|
|
pin PULL routing, plus ``group_owners`` for the free fan-out.
|
|
"""
|
|
|
|
def test_identity_when_fleets_match(self):
|
|
r = _router(8, 8, group_lens=[1] * 6)
|
|
assert all(r.producer_for(3, f"g{g}") == 3 for g in range(6))
|
|
assert all(r.producer_for(c, "g0") == c for c in range(8))
|
|
|
|
def test_every_producer_nic_carries_traffic(self):
|
|
"""16 producers / 8 consumers: the block rule spreads each consumer's
|
|
pulls over its block, alternating by group, so no producer sits idle.
|
|
Plain `consumer_id % len(owners)` would idle half of them."""
|
|
r = _router(16, 8, group_lens=[1] * 95)
|
|
names = [f"g{i}" for i in range(95)]
|
|
for c in range(8):
|
|
block = assign_producer_indices(16, 8, c)
|
|
assert {r.producer_for(c, n) for n in names} == set(block)
|
|
assert [r.producer_for(0, f"g{g}") for g in range(4)] == [0, 1, 0, 1]
|
|
assert {r.producer_for(c, n) for c in range(8) for n in names} == set(range(16))
|
|
|
|
def test_fan_in_shares_one_producer(self):
|
|
r = _router(2, 8, group_lens=[1] * 5)
|
|
names = [f"g{i}" for i in range(5)]
|
|
assert [{r.producer_for(c, n) for n in names} for c in range(8)] == [
|
|
{c // 4} for c in range(8)
|
|
]
|
|
|
|
def test_pipeline_stages_route_to_the_owning_stage(self):
|
|
"""2 stages x 8 ranks: groups 0-2 on stage 0, 3-5 on stage 1. Two owner
|
|
sets, one per stage; each consumer reaches both."""
|
|
r = _router(
|
|
16,
|
|
8,
|
|
owner_sets=[list(range(8)), list(range(8, 16))],
|
|
name_owner_class=[0, 0, 0, 1, 1, 1],
|
|
)
|
|
for c in range(8):
|
|
picks = [r.producer_for(c, f"g{g}") for g in range(6)]
|
|
assert picks == [c] * 3 + [c + 8] * 3
|
|
assert r.group_owners(0) == list(range(8))
|
|
assert r.group_owners(3) == list(range(8, 16))
|
|
r.validate()
|
|
|
|
def test_expert_names_route_to_the_holding_ranks(self):
|
|
"""A name held by only some ranks routes to those; a replicated name
|
|
keeps the full set."""
|
|
r = RdtRouter(
|
|
8,
|
|
4,
|
|
[[0, 1, 2, 3], [2, 3], [4, 5, 6, 7], [6, 7]],
|
|
[0, 1, 2, 3],
|
|
["repl0", "expert0", "repl1", "expert1"],
|
|
[2, 2],
|
|
)
|
|
r.validate()
|
|
assert r.owners("expert0") == [2, 3]
|
|
assert r.owners("expert1") == [6, 7]
|
|
assert r.owners("repl1") == [4, 5, 6, 7]
|
|
for c in range(4):
|
|
assert r.producer_for(c, "expert0") in (2, 3)
|
|
assert r.producer_for(c, "repl1") in (4, 5, 6, 7)
|
|
|
|
def test_consumers_spread_over_an_owner_set(self):
|
|
"""Several ranks hold a name (its TP peers); the block rule must spread
|
|
consumers across them, not funnel through one."""
|
|
r = _router(
|
|
4, 4, owner_sets=[[0, 1]], name_owner_class=[0] * 6, group_lens=[1] * 6
|
|
)
|
|
served = {r.producer_for(c, f"g{g}") for c in range(4) for g in range(6)}
|
|
assert served == {0, 1}
|
|
|
|
def test_a_group_frees_at_every_owner_of_any_of_its_names(self):
|
|
"""The free barrier is per group, so its fan-out is the UNION over the
|
|
group's names: a rank holding none of them is not signalled."""
|
|
r = RdtRouter(4, 2, [[0, 1], [2, 3]], [0, 1, 0], ["a", "b", "c"], [2, 1])
|
|
assert r.group_owners(0) == [0, 1, 2, 3]
|
|
assert r.group_owners(1) == [0, 1]
|
|
|
|
def test_an_unowned_name_raises(self):
|
|
r = RdtRouter(2, 2, [[]], [0], ["w"], [1])
|
|
with pytest.raises(ValueError, match="has no owner"):
|
|
r.producer_for(0, "w")
|
|
|
|
def test_validate_rejects_an_empty_owner_set(self):
|
|
with pytest.raises(ValueError, match="empty"):
|
|
RdtRouter(
|
|
4, 2, [[0, 1], [], [2, 3]], [0, 1, 2], ["a", "b", "c"], [3]
|
|
).validate()
|
|
|
|
def test_validate_rejects_an_out_of_range_owner(self):
|
|
with pytest.raises(ValueError, match="out of range"):
|
|
RdtRouter(2, 2, [[0, 5]], [0], ["a"], [1]).validate()
|
|
|
|
def test_validate_rejects_an_unresolvable_class(self):
|
|
with pytest.raises(ValueError, match="owner class"):
|
|
RdtRouter(4, 2, [[0, 1]], [0, 7], ["a", "b"], [2]).validate()
|
|
|
|
|
|
class TestReplicaOverlay:
|
|
"""[RDT-SHARE-SLOTS] Several inference deployments carve the same blocks.
|
|
|
|
The producer can serve R deployments out of ONE slot only if the R copies of
|
|
a worker meet on the same producer asking for the same bytes. Identical
|
|
plans give the second half for free; this is the first half. Carving over
|
|
the whole fleet instead spreads the multi-owner (replicated) names of
|
|
different deployments onto different producers, so a producer ends up
|
|
holding a ring per distinct worker index and nothing can merge.
|
|
"""
|
|
|
|
def _fleet(
|
|
self,
|
|
num_producers,
|
|
workers,
|
|
replicas,
|
|
owner_sets=None,
|
|
name_owner_class=None,
|
|
n_names=6,
|
|
):
|
|
names = [f"g{i}" for i in range(n_names)]
|
|
return (
|
|
RdtRouter(
|
|
num_producers,
|
|
workers * replicas,
|
|
owner_sets,
|
|
name_owner_class,
|
|
names,
|
|
[1] * n_names,
|
|
workers_per_replica=workers,
|
|
),
|
|
names,
|
|
)
|
|
|
|
def test_every_replica_of_a_worker_resolves_the_same_producer(self):
|
|
r, names = self._fleet(16, 8, 4)
|
|
for worker in range(8):
|
|
for name in names:
|
|
picks = {r.producer_for(worker + rep * 8, name) for rep in range(4)}
|
|
assert len(picks) == 1, (
|
|
f"worker {worker} split across producers {picks}"
|
|
)
|
|
|
|
def test_adding_a_deployment_does_not_move_the_first_one(self):
|
|
"""The strongest form: an existing consumer's routes are byte-identical
|
|
before and after another deployment is provisioned."""
|
|
one, names = self._fleet(16, 8, 1)
|
|
two, _ = self._fleet(16, 8, 2)
|
|
for consumer_id in range(8):
|
|
assert [one.producer_for(consumer_id, n) for n in names] == [
|
|
two.producer_for(consumer_id, n) for n in names
|
|
]
|
|
|
|
def test_one_deployment_is_the_historical_carve(self):
|
|
"""Default (no width) and width == the whole fleet must agree with the
|
|
plain block rule, so a single-deployment fleet is untouched."""
|
|
plain = _router(16, 8, group_lens=[1] * 6)
|
|
explicit, names = self._fleet(16, 8, 1)
|
|
for consumer_id in range(8):
|
|
block = assign_producer_indices(16, 8, consumer_id)
|
|
for name in names:
|
|
assert plain.producer_for(consumer_id, name) == explicit.producer_for(
|
|
consumer_id, name
|
|
)
|
|
assert explicit.producer_for(consumer_id, name) in block
|
|
|
|
def test_the_overlay_holds_for_partially_owned_names(self):
|
|
"""Pipeline stages and experts too: two owner sets, one narrow. The
|
|
narrow one already forced agreement (one owner, no block to carve); the
|
|
wide one is what the carve decides."""
|
|
r, names = self._fleet(
|
|
8,
|
|
4,
|
|
2,
|
|
owner_sets=[[0, 1, 2, 3], [5]],
|
|
name_owner_class=[0, 1, 0, 1, 0, 1],
|
|
)
|
|
for worker in range(4):
|
|
for name in names:
|
|
assert r.producer_for(worker, name) == r.producer_for(worker + 4, name)
|
|
assert {r.producer_for(c, "g1") for c in range(8)} == {5}
|
|
|
|
def test_a_width_that_does_not_divide_the_fleet_raises(self):
|
|
"""A uniform fleet is asserted upstream in ``get_world_size``; if it ever
|
|
is not, the overlay would map two workers of one deployment onto the same
|
|
block index, so refuse rather than serve the wrong bytes."""
|
|
with pytest.raises(ValueError, match="does not divide"):
|
|
RdtRouter(4, 6, None, None, ["a"], [1], workers_per_replica=4)
|
|
|
|
def test_pulls_still_carry_the_global_consumer_id(self):
|
|
"""The carve uses the index WITHIN a deployment, but the wire must keep
|
|
the fleet-global id: it is what lets the producer tell the sharers of a
|
|
slot apart, and count their arrivals separately."""
|
|
|
|
class _Method:
|
|
def __init__(self):
|
|
self.calls = []
|
|
|
|
def remote(self, keys, consumer_id, seq):
|
|
self.calls.append((keys, consumer_id, seq))
|
|
|
|
r, _ = self._fleet(4, 2, 2)
|
|
methods = [_Method() for _ in range(4)]
|
|
r.bind([object()] * 4, methods, consumer_id=3) # replica 1, worker 1
|
|
r.pull(2, ["k"], 7)
|
|
|
|
assert methods[2].calls == [(["k"], 3, 7)]
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Packed layout: the cross-process invariant
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
def _producer_pack_offsets(slices):
|
|
"""The producer's rule, transcribed from ``rdt_produce_weights_batched``
|
|
(sharded_rdt_trainer.py): 16B-aligned offsets in specs order.
|
|
|
|
Kept as an independent implementation on purpose: the consumer computing the
|
|
same offsets is the invariant that makes the packed blob readable, and this
|
|
is the only thing that catches the two rules drifting apart. A divergence is
|
|
silent at runtime — the bytes still arrive, they are just carved up wrong.
|
|
"""
|
|
pack_cur = 0
|
|
offsets = []
|
|
for numel, dtype in slices:
|
|
off = (pack_cur + 15) & ~15
|
|
pack_cur = off + numel * dtype.itemsize
|
|
offsets.append(off)
|
|
return offsets, pack_cur
|
|
|
|
|
|
class TestPackedLayout:
|
|
def test_consumer_layout_matches_the_producer_rule(self):
|
|
"""Mixed dtypes, sizes that do not land on 16B — i.e. the Kimi-style
|
|
group (fp8 weights + fp32 scales + bf16 norms) that exposes alignment."""
|
|
specs = [
|
|
("w.fp8", "float8_e4m3fn", [17]),
|
|
("w.scale", "float32", [3]),
|
|
("w.norm", "bfloat16", [5]),
|
|
("w.big", "bfloat16", [1024]),
|
|
]
|
|
baked, name_meta = {}, {}
|
|
layer = _FakeLayer("mixed")
|
|
copies = []
|
|
for name, dtype_name, shape in specs:
|
|
name_meta[name] = (dtype_name, shape)
|
|
copies.append(
|
|
_copy(name, shape=tuple(shape), dtype=_dtype_from_name(dtype_name))
|
|
)
|
|
mod_plan = _module(layer, copies)
|
|
for name, _d, _s in specs:
|
|
baked[name] = mod_plan
|
|
|
|
eng = _planner(baked, name_meta=name_meta)
|
|
plan = eng._build_call_plan([n for n, _, _ in specs], [len(specs)])
|
|
(chunk,) = plan.chunks
|
|
|
|
want_offsets, want_bytes = _producer_pack_offsets(
|
|
[(s[0], _dtype_from_name(d)) for _n, d, s in specs]
|
|
)
|
|
assert [off for off, _dt, _n, _shape in chunk.pack_layout] == want_offsets
|
|
assert chunk.pack_bytes == want_bytes
|
|
|
|
def test_every_offset_is_16b_aligned(self):
|
|
specs = [(f"w{i}", "bfloat16", [i * 3 + 1]) for i in range(6)]
|
|
layer = _FakeLayer("l")
|
|
copies = [_copy(n, shape=tuple(s)) for n, _d, s in specs]
|
|
mod_plan = _module(layer, copies)
|
|
eng = _planner(
|
|
{n: mod_plan for n, _d, _s in specs},
|
|
name_meta={n: (d, s) for n, d, s in specs},
|
|
)
|
|
(chunk,) = eng._build_call_plan([n for n, _, _ in specs], [len(specs)]).chunks
|
|
assert all(off % 16 == 0 for off, _dt, _n, _s in chunk.pack_layout)
|
|
|
|
def test_layout_carries_dtype_numel_and_shape_for_the_view_carve(self):
|
|
baked, name_meta, names, group_lens = _one_module_per_layer(1, numel=6)
|
|
eng = _planner(baked, name_meta=name_meta)
|
|
chunk = eng._build_call_plan(names, group_lens).chunks[0]
|
|
(off, dt, numel, shape) = chunk.pack_layout[0]
|
|
assert (off, dt, numel, shape) == (0, torch.bfloat16, 6, (6,))
|
|
|
|
def test_duplicate_source_keys_are_packed_once(self):
|
|
"""Two destinations fed by the SAME (name, chain) pull one slice."""
|
|
layer = _FakeLayer("l")
|
|
shared = ("w", ())
|
|
copies = [
|
|
_copy("w", "a", shape=(4,)),
|
|
_copy("w", "b", shape=(4,)),
|
|
]
|
|
assert copies[0].src == shared and copies[1].src == shared
|
|
eng = _planner(
|
|
{"w": _module(layer, copies)},
|
|
name_meta={"w": ("bfloat16", [4])},
|
|
)
|
|
(chunk,) = eng._build_call_plan(["w"], [1]).chunks
|
|
assert len(chunk.scatters) == 2
|
|
assert chunk.keys == [shared]
|
|
assert len(chunk.pack_layout) == 1
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Chunk splitting
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
class TestChunkModuleScatters:
|
|
"""The chunk cut: one chunk per distinct owner class present in the copies,
|
|
ascending by class index. Derived purely from the bake plus the ownership
|
|
table, so vLLM's expert placement needs no cases."""
|
|
|
|
def test_unstamped_copies_form_one_chunk(self):
|
|
layer = _FakeLayer("l")
|
|
copies = [_copy(f"w{i}", shape=(4,)) for i in range(5)]
|
|
eng = _planner({}, name_meta={f"w{i}": ("bfloat16", [4]) for i in range(5)})
|
|
chunks = eng._chunk_module_scatters([_module(layer, copies)])
|
|
assert [ci for ci, _s in chunks] == [0]
|
|
assert len(chunks[0][1]) == 5
|
|
|
|
def test_one_chunk_per_owner_class_ascending(self):
|
|
"""Classes are numbered by first appearance in metadata order, so the
|
|
replicated class (0) leads and the rest follow — the ordering the old
|
|
"-1 first, then ascending stamp" rule produced, now falling out of the
|
|
table instead of a special case."""
|
|
layer = _FakeLayer("moe")
|
|
names = ["norm", "e5", "e0", "e2"]
|
|
copies = [_copy(n, shape=(4,)) for n in names]
|
|
eng = _planner(
|
|
{},
|
|
name_meta={n: ("bfloat16", [4]) for n in names},
|
|
held_by={"e0": [0], "e2": [1], "e5": [2]},
|
|
)
|
|
chunks = eng._chunk_module_scatters([_module(layer, copies)])
|
|
assert [ci for ci, _s in chunks] == [0, 1, 2, 3]
|
|
assert [[s.src[0] for s in sc] for _ci, sc in chunks] == [
|
|
["norm"],
|
|
["e5"],
|
|
["e0"],
|
|
["e2"],
|
|
]
|
|
|
|
def test_copy_order_within_a_chunk_is_bake_order(self):
|
|
"""Scatter order within a chunk must follow the bake, not the stamp
|
|
sort: the packed layout and the producer's replay agree on keys order."""
|
|
layer = _FakeLayer("moe")
|
|
names = ["e3", "e1", "e2"] # one stamp, deliberately unsorted names
|
|
copies = [_copy(n, shape=(4,)) for n in names]
|
|
eng = _planner(
|
|
{},
|
|
name_meta={n: ("bfloat16", [4]) for n in names},
|
|
held_by={n: [0] for n in names},
|
|
)
|
|
((_ci, scatters),) = eng._chunk_module_scatters([_module(layer, copies)])
|
|
assert [s.src[0] for s in scatters] == names
|
|
|
|
def test_stamps_interleave_across_modules_without_reordering_within(self):
|
|
"""Two modules' copies bucket by stamp; within each bucket the order is
|
|
module-then-bake order."""
|
|
a, b = _FakeLayer("a"), _FakeLayer("b")
|
|
mods = [
|
|
_module(a, [_copy("a0", shape=(4,)), _copy("a1", shape=(4,))]),
|
|
_module(b, [_copy("b0", shape=(4,)), _copy("b1", shape=(4,))]),
|
|
]
|
|
eng = _planner(
|
|
{},
|
|
name_meta={n: ("bfloat16", [4]) for n in ("a0", "a1", "b0", "b1")},
|
|
held_by={"a0": [0], "a1": [1], "b0": [0], "b1": [1]},
|
|
)
|
|
chunks = eng._chunk_module_scatters(mods)
|
|
assert [(ci, [s.src[0] for s in sc]) for ci, sc in chunks] == [
|
|
(1, ["a0", "b0"]),
|
|
(2, ["a1", "b1"]),
|
|
]
|
|
|
|
def test_scatters_carry_their_own_dtype(self):
|
|
"""dtype rides the record from the bake, where it is the fake's dtype
|
|
AFTER its op chain — not a plan-time lookup of the source name, which
|
|
would be wrong for any chain that reinterprets dtype."""
|
|
layer = _FakeLayer("l")
|
|
eng = _planner({}, name_meta={"w": ("float32", [10])})
|
|
((_er, scatters),) = eng._chunk_module_scatters(
|
|
[_module(layer, [_copy("w", shape=(10,), dtype=torch.float32)])]
|
|
)
|
|
(sc,) = scatters
|
|
assert sc.dtype is torch.float32
|
|
assert sc.layer is layer
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# _build_call_plan
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
class TestBuildCallPlan:
|
|
def test_one_chunk_per_group_when_unstamped(self):
|
|
baked, name_meta, names, group_lens = _one_module_per_layer(3)
|
|
eng = _planner(baked, name_meta=name_meta)
|
|
plan = eng._build_call_plan(names, group_lens)
|
|
assert len(plan.chunks) == len(group_lens) == 5
|
|
assert plan.pre_free == []
|
|
|
|
def test_stamps_multiply_chunks_within_a_group(self):
|
|
layer = _FakeLayer("l")
|
|
names = [f"w{i}" for i in range(4)]
|
|
copies = [_copy(n, shape=(4,)) for n in names]
|
|
mod_plan = _module(layer, copies)
|
|
eng = _planner(
|
|
{n: mod_plan for n in names},
|
|
name_meta={n: ("bfloat16", [4]) for n in names},
|
|
held_by={"w2": [0], "w3": [1]},
|
|
)
|
|
plan = eng._build_call_plan(names, [4])
|
|
assert len(plan.chunks) == 3 # replicated, w2's class, w3's class
|
|
|
|
def test_materialize_fires_on_a_modules_first_chunk_only(self):
|
|
"""Empty HF params are allocated once per module, by construction —
|
|
including a FusedMoE-like module whose copies span owner-class chunks."""
|
|
layer = _FakeLayer("spanning")
|
|
names = [f"w{i}" for i in range(4)]
|
|
copies = [_copy(n, shape=(4,)) for n in names]
|
|
mod_plan = _module(layer, copies)
|
|
eng = _planner(
|
|
{n: mod_plan for n in names},
|
|
name_meta={n: ("bfloat16", [4]) for n in names},
|
|
held_by={"w0": [0], "w1": [0], "w2": [1], "w3": [1]},
|
|
)
|
|
plan = eng._build_call_plan(names, [4])
|
|
assert [c.materialize for c in plan.chunks] == [[layer], []]
|
|
|
|
def test_quant_defers_to_a_modules_last_chunk(self):
|
|
layer = _FakeLayer("spanning")
|
|
names = [f"w{i}" for i in range(4)]
|
|
copies = [_copy(n, shape=(4,)) for n in names]
|
|
mod_plan = _module(layer, copies)
|
|
eng = _planner(
|
|
{n: mod_plan for n in names},
|
|
name_meta={n: ("bfloat16", [4]) for n in names},
|
|
held_by={"w0": [0], "w1": [0], "w2": [1], "w3": [1]},
|
|
)
|
|
plan = eng._build_call_plan(names, [4])
|
|
assert [c.quant for c in plan.chunks] == [[], [layer]]
|
|
|
|
def test_free_signal_fires_on_each_groups_last_chunk(self):
|
|
baked, name_meta, names, group_lens = _one_module_per_layer(2)
|
|
eng = _planner(baked, name_meta=name_meta)
|
|
plan = eng._build_call_plan(names, group_lens)
|
|
assert [c.free for c in plan.chunks] == [[0], [1], [2], [3]]
|
|
|
|
def test_free_signal_waits_for_a_groups_last_chunk_when_it_is_split(self):
|
|
"""The load-bearing case: with the group cut across owner-class chunks the
|
|
signal must hang off the LAST one. Signaling earlier lets the producers
|
|
drop the gather buffers while a later chunk's RDMA is still reading."""
|
|
layer = _FakeLayer("l")
|
|
names = [f"w{i}" for i in range(4)]
|
|
copies = [_copy(n, shape=(4,)) for n in names]
|
|
mod_plan = _module(layer, copies)
|
|
eng = _planner(
|
|
{n: mod_plan for n in names},
|
|
name_meta={n: ("bfloat16", [4]) for n in names},
|
|
held_by={"w2": [0], "w3": [1]},
|
|
)
|
|
plan = eng._build_call_plan(names, [4])
|
|
assert len(plan.chunks) == 3
|
|
assert [c.free for c in plan.chunks] == [[], [], [0]]
|
|
|
|
def test_free_signal_timing_across_two_split_groups(self):
|
|
"""Two groups x 2 owner-class chunks each: signals land on each group's
|
|
last chunk only."""
|
|
names_a = [f"model.layers.0.w{i}" for i in range(4)]
|
|
names_b = [f"model.layers.1.w{i}" for i in range(4)]
|
|
baked, name_meta, stamps = {}, {}, {}
|
|
for names in (names_a, names_b):
|
|
layer = _FakeLayer(names[0])
|
|
mod_plan = _module(layer, [_copy(n, shape=(4,)) for n in names])
|
|
for i, n in enumerate(names):
|
|
baked[n] = mod_plan
|
|
name_meta[n] = ("bfloat16", [4])
|
|
stamps[n] = [i // 2] # two owner classes per group
|
|
eng = _planner(baked, name_meta=name_meta, held_by=stamps)
|
|
plan = eng._build_call_plan(names_a + names_b, [4, 4])
|
|
assert [c.free for c in plan.chunks] == [[], [0], [], [1]]
|
|
|
|
def test_a_group_with_nothing_local_is_signaled_at_sync_start(self):
|
|
"""The owners still published it, so this consumer must signal it —
|
|
immediately, before the pipeline (signal-before-publish is tolerated),
|
|
never hung off another group's chunk."""
|
|
baked, name_meta, names, group_lens = _one_module_per_layer(2)
|
|
del baked["model.layers.1.w"] # nothing baked, and not live => no pull
|
|
eng = _planner(
|
|
baked, name_meta=name_meta, live=set(names) - {"model.layers.1.w"}
|
|
)
|
|
plan = eng._build_call_plan(names, group_lens)
|
|
assert len(plan.chunks) == 3
|
|
assert plan.pre_free == [2]
|
|
assert [c.free for c in plan.chunks] == [[0], [1], [3]]
|
|
|
|
def test_leading_groups_with_nothing_local_go_to_pre_free(self):
|
|
baked, name_meta, names, group_lens = _one_module_per_layer(2)
|
|
del baked["embed.weight"]
|
|
eng = _planner(baked, name_meta=name_meta, live=set(names) - {"embed.weight"})
|
|
plan = eng._build_call_plan(names, group_lens)
|
|
assert plan.pre_free == [0]
|
|
assert len(plan.chunks) == 3
|
|
|
|
def test_modules_are_deduped_within_a_group(self):
|
|
"""Several names of one fused module map to the same scatter list; the
|
|
plan must not chunk it twice."""
|
|
layer = _FakeLayer("fused")
|
|
names = ["qkv.q", "qkv.k", "qkv.v"]
|
|
copies = [_copy(n, shape=(4,)) for n in names]
|
|
mod_plan = _module(layer, copies)
|
|
eng = _planner(
|
|
{n: mod_plan for n in names},
|
|
name_meta={n: ("bfloat16", [4]) for n in names},
|
|
)
|
|
plan = eng._build_call_plan(names, [3])
|
|
(chunk,) = plan.chunks
|
|
assert chunk.materialize == [layer]
|
|
assert len(chunk.scatters) == 3
|
|
|
|
def test_plan_is_a_pure_function_of_its_inputs(self):
|
|
"""Called twice with the same arguments it must produce the same plan —
|
|
that is what lets the engine build it once and reuse it every sync."""
|
|
baked, name_meta, names, group_lens = _one_module_per_layer(2)
|
|
eng = _planner(baked, name_meta=name_meta)
|
|
a = eng._build_call_plan(names, group_lens)
|
|
b = eng._build_call_plan(names, group_lens)
|
|
assert [c.pack_bytes for c in a.chunks] == [c.pack_bytes for c in b.chunks]
|
|
assert [c.keys for c in a.chunks] == [c.keys for c in b.chunks]
|
|
assert [c.free for c in a.chunks] == [c.free for c in b.chunks]
|
|
assert a.pre_free == b.pre_free
|
|
|
|
def test_single_producer_routes_every_chunk_to_local_owner_zero(self):
|
|
baked, name_meta, names, group_lens = _one_module_per_layer(2)
|
|
eng = _planner(baked, name_meta=name_meta)
|
|
plan = eng._build_call_plan(names, group_lens)
|
|
assert [c.owner for c in plan.chunks] == [0] * len(plan.chunks)
|
|
|
|
|
|
class TestUnbakedNamesGuard:
|
|
"""A live name with no baked plan has no way to load — there is no fallback,
|
|
and its pull would target groups the pipeline already freed — so the plan
|
|
build fails at init, naming the names."""
|
|
|
|
def test_a_live_unbaked_name_fails_the_plan_build(self):
|
|
baked, name_meta, names, group_lens = _one_module_per_layer(2)
|
|
del baked["norm.weight"] # unbaked but live
|
|
eng = _planner(baked, name_meta=name_meta, live=set(names))
|
|
with pytest.raises(RuntimeError, match="norm.weight"):
|
|
eng._build_call_plan(names, group_lens)
|
|
|
|
def test_never_copied_names_are_dropped_entirely(self):
|
|
"""Experts owned by another EP rank no-op in their loader: not an error,
|
|
and not a chunk — their group is pre-freed if nothing else fills it."""
|
|
baked, name_meta, names, group_lens = _one_module_per_layer(2)
|
|
del baked["norm.weight"]
|
|
eng = _planner(baked, name_meta=name_meta, live=set(names) - {"norm.weight"})
|
|
plan = eng._build_call_plan(names, group_lens)
|
|
assert len(names) - 1 == len(plan.chunks)
|
|
assert plan.pre_free == [len(group_lens) - 1]
|
|
|
|
|
|
class TestCallPlanRouting:
|
|
"""Under partial ownership each chunk must be pulled from a producer that
|
|
actually holds every name in it. Consumers bind EVERY producer (signals fan
|
|
out to all owners), so local index == trainer rank."""
|
|
|
|
def _routed_planner(
|
|
self, held_by_group, num_consumers, consumer_id, *, held_by=None
|
|
):
|
|
"""``held_by_group[g]`` = the ranks holding group g's names, so a test
|
|
can still speak in groups; ``held_by`` overrides individual names."""
|
|
baked, name_meta, names, group_lens = _one_module_per_layer(
|
|
len(held_by_group) - 2
|
|
)
|
|
per_name = {n: list(held_by_group[gi]) for gi, n in enumerate(names)}
|
|
per_name.update(held_by or {})
|
|
n_prod = max(max(o) for o in per_name.values()) + 1
|
|
owner_sets: list[list[int]] = []
|
|
seen: dict[tuple, int] = {}
|
|
classes = []
|
|
for n in names:
|
|
key = tuple(sorted(per_name[n]))
|
|
if key not in seen:
|
|
seen[key] = len(owner_sets)
|
|
owner_sets.append(list(key))
|
|
classes.append(seen[key])
|
|
eng = _planner(baked, name_meta=name_meta)
|
|
router = RdtRouter(
|
|
n_prod, num_consumers, owner_sets, classes, names, group_lens
|
|
)
|
|
router.validate()
|
|
router.bind([object()] * n_prod, [object()] * n_prod, consumer_id)
|
|
eng._router = router
|
|
return eng, names, group_lens
|
|
|
|
def test_two_stage_ownership_splits_chunks_between_owners(self):
|
|
"""Two PP stages, 4 groups: stage 0 holds the first half, stage 1 the
|
|
second. Every chunk goes to a holder of its names."""
|
|
eng, names, group_lens = self._routed_planner([[0], [0], [1], [1]], 1, 0)
|
|
plan = eng._build_call_plan(names, group_lens)
|
|
assert [c.owner for c in plan.chunks] == [0, 0, 1, 1]
|
|
|
|
def test_a_multi_owner_group_still_resolves_to_one_producer(self):
|
|
"""Every chunk is served by exactly ONE producer: splitting a pull only
|
|
multiplies produce calls, since the consumer's NIC bounds it."""
|
|
eng, names, group_lens = self._routed_planner([[0, 1]] * 4, 1, 0)
|
|
plan = eng._build_call_plan(names, group_lens)
|
|
assert all(0 <= c.owner < 2 for c in plan.chunks)
|
|
|
|
def test_an_expert_name_routes_to_the_rank_holding_it(self):
|
|
"""PP ∩ EP: stage 0 = ranks 0-1, stage 1 = ranks 2-3, and one expert
|
|
name inside a stage-1 group is held by rank 3 alone. The name's own owner
|
|
set is what routes it."""
|
|
eng, names, group_lens = self._routed_planner(
|
|
[[0, 1], [0, 1], [2, 3], [2, 3]],
|
|
1,
|
|
0,
|
|
held_by={"model.layers.1.w": [3]},
|
|
)
|
|
plan = eng._build_call_plan(names, group_lens)
|
|
assert plan.chunks[2].owner == 3
|
|
assert plan.chunks[0].owner in (0, 1)
|
|
assert plan.chunks[3].owner in (2, 3)
|
|
|
|
def test_an_unheld_name_is_rejected_at_init(self):
|
|
"""A name no rank holds can never be served. The trainer raises when it
|
|
transposes the holdings; the router's validate() is the consumer-side
|
|
backstop, and either way it fails at init rather than at first pull."""
|
|
with pytest.raises(ValueError, match="empty"):
|
|
RdtRouter(2, 1, [[]], [0], ["w"], [1]).validate()
|
|
|
|
|
|
class TestSignalCompleteness:
|
|
"""Every gather group is signaled exactly once — after its last chunk when
|
|
the worker pulls from it, at sync start otherwise — for ANY expert
|
|
placement. A group signaled twice over-credits the barrier; a group never
|
|
signaled parks a producer credit and hangs end_sync."""
|
|
|
|
def _moe_planner(self, worker_experts, *, ep_size=4, n_experts=8, n_layers=2):
|
|
"""pre / n_layers MoE layers / post. Each layer: one norm (its own
|
|
module, stamp -1) + n_experts expert names stamped ``e // n_local``.
|
|
The worker's bake covers only ``worker_experts`` (its placement);
|
|
foreign expert names never copied => dropped entirely, as in the real
|
|
bake."""
|
|
n_local = n_experts // ep_size
|
|
names = ["embed.weight"]
|
|
group_lens = [1]
|
|
baked = {
|
|
"embed.weight": _module(
|
|
_FakeLayer("embed"), [_copy("embed.weight", shape=(4,))]
|
|
)
|
|
}
|
|
name_meta = {"embed.weight": ("bfloat16", [4])}
|
|
stamps: dict = {}
|
|
live = {"embed.weight"}
|
|
for li in range(n_layers):
|
|
norm = f"model.layers.{li}.norm"
|
|
enames = [f"model.layers.{li}.experts.{e}.w" for e in range(n_experts)]
|
|
names += [norm] + enames
|
|
group_lens.append(1 + n_experts)
|
|
name_meta[norm] = ("bfloat16", [4])
|
|
live.add(norm)
|
|
baked[norm] = _module(_FakeLayer(f"norm{li}"), [_copy(norm, shape=(4,))])
|
|
fused = _module(
|
|
_FakeLayer(f"moe{li}"),
|
|
[
|
|
_copy(f"model.layers.{li}.experts.{e}.w", shape=(4,))
|
|
for e in worker_experts
|
|
],
|
|
)
|
|
for e in range(n_experts):
|
|
n = f"model.layers.{li}.experts.{e}.w"
|
|
name_meta[n] = ("bfloat16", [4])
|
|
stamps[n] = [e // n_local] # held by that EP rank alone
|
|
if e in worker_experts:
|
|
baked[n] = fused
|
|
live.add(n)
|
|
names.append("lm_head.weight")
|
|
group_lens.append(1)
|
|
baked["lm_head.weight"] = _module(
|
|
_FakeLayer("head"), [_copy("lm_head.weight", shape=(4,))]
|
|
)
|
|
name_meta["lm_head.weight"] = ("bfloat16", [4])
|
|
live.add("lm_head.weight")
|
|
eng = _planner(baked, name_meta=name_meta, live=live, held_by=stamps)
|
|
return eng, names, group_lens
|
|
|
|
@staticmethod
|
|
def _signals(plan):
|
|
return sorted(list(plan.pre_free) + [gi for c in plan.chunks for gi in c.free])
|
|
|
|
def test_linear_placement_signals_every_group_exactly_once(self):
|
|
eng, names, group_lens = self._moe_planner(worker_experts=[0, 1, 2, 3])
|
|
plan = eng._build_call_plan(names, group_lens)
|
|
assert self._signals(plan) == list(range(len(group_lens)))
|
|
|
|
def test_round_robin_placement_signals_every_group_exactly_once(self):
|
|
eng, names, group_lens = self._moe_planner(worker_experts=[0, 2, 4, 6])
|
|
plan = eng._build_call_plan(names, group_lens)
|
|
assert self._signals(plan) == list(range(len(group_lens)))
|
|
|
|
def test_placement_changes_only_the_chunk_count_never_the_signals(self):
|
|
"""linear experts 0-3 hit 2 owner classes; round_robin 0,2,4,6
|
|
hits all 4. More chunks per group, identical signal set."""
|
|
plans = {}
|
|
for label, experts in (("linear", [0, 1, 2, 3]), ("round_robin", [0, 2, 4, 6])):
|
|
eng, names, group_lens = self._moe_planner(worker_experts=experts)
|
|
plans[label] = eng._build_call_plan(names, group_lens)
|
|
# per MoE group: the replicated chunk + one per owner class present
|
|
assert (
|
|
len(plans["linear"].chunks) == 2 + 2 * 3
|
|
) # embed+head + 2 layers x (1 + 2)
|
|
assert len(plans["round_robin"].chunks) == 2 + 2 * 5 # 2 layers x (1 + 4)
|
|
assert self._signals(plans["linear"]) == self._signals(plans["round_robin"])
|
|
|
|
def test_the_fused_module_materializes_first_and_quants_last_across_chunks(self):
|
|
"""The FusedMoE shape: one module's copies span every owner-class chunk of
|
|
its group; materialize on its first, quant on its last."""
|
|
eng, names, group_lens = self._moe_planner(
|
|
worker_experts=[0, 2, 4, 6], n_layers=1
|
|
)
|
|
plan = eng._build_call_plan(names, group_lens)
|
|
# group 1 = [norm chunk, coord 0..3 chunks] at chunk indices 1..5
|
|
fused_chunks = [
|
|
c
|
|
for c in plan.chunks[1:6]
|
|
if any("experts" in sc.src[0] for sc in c.scatters)
|
|
]
|
|
assert len(fused_chunks) == 4
|
|
(fused_layer,) = fused_chunks[0].materialize
|
|
assert fused_layer.tag == "moe0"
|
|
assert fused_chunks[0].quant == []
|
|
assert fused_chunks[-1].quant == [fused_layer]
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Shared helpers
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
class TestBufferAllocBytes:
|
|
def test_rounds_up_to_a_coarse_256mb_boundary(self):
|
|
assert buffer_alloc_bytes(1) == 256 << 20
|
|
assert buffer_alloc_bytes(256 << 20) == 256 << 20
|
|
assert buffer_alloc_bytes((256 << 20) + 1) == 512 << 20
|
|
|
|
def test_presize_is_a_floor_not_a_cap(self):
|
|
assert buffer_alloc_bytes(1, presize=3 << 30) == 3 << 30
|
|
big = 4 << 30
|
|
assert buffer_alloc_bytes(big, presize=1 << 30) == big
|
|
|
|
def test_never_returns_less_than_requested(self):
|
|
for nbytes in (1, 1 << 20, (700 << 20) + 3):
|
|
assert buffer_alloc_bytes(nbytes) >= nbytes
|
|
|
|
|
|
class TestLayerwiseGroups:
|
|
def test_partition_is_one_group_per_layer_between_unindexed_blocks(self):
|
|
names = [
|
|
"embed.weight",
|
|
"model.layers.0.a",
|
|
"model.layers.0.b",
|
|
"model.layers.1.a",
|
|
"norm.weight",
|
|
]
|
|
assert layerwise_groups(names) == [
|
|
["embed.weight"],
|
|
["model.layers.0.a", "model.layers.0.b"],
|
|
["model.layers.1.a"],
|
|
["norm.weight"],
|
|
]
|
|
|
|
def test_flattening_a_partition_returns_the_input_order(self):
|
|
"""Un-indexed names group by POSITION, not by name class — so one that
|
|
looks like a pre-block name still lands after the layers it follows.
|
|
Group index therefore means the same thing on every rank."""
|
|
names = [
|
|
"embed.weight",
|
|
"model.layers.0.a",
|
|
"embed.weight.tied", # after the layers => its own trailing block
|
|
]
|
|
groups = layerwise_groups(names)
|
|
assert [n for g in groups for n in g] == names
|
|
|
|
def test_layers_are_ordered_by_index_not_appearance(self):
|
|
names = ["model.layers.10.a", "model.layers.2.a"]
|
|
assert layerwise_groups(names) == [
|
|
["model.layers.2.a"],
|
|
["model.layers.10.a"],
|
|
]
|
|
|
|
def test_stacks_keep_prefix_appearance_order_and_sort_within(self):
|
|
"""Two stacks stay in the order they appear — the vision tower before the
|
|
text stack — while each sorts by index internally."""
|
|
names = [
|
|
"visual.blocks.1.w",
|
|
"visual.blocks.0.w",
|
|
"model.language_model.layers.1.w",
|
|
"model.language_model.layers.0.w",
|
|
]
|
|
assert layerwise_groups(names) == [
|
|
["visual.blocks.0.w"],
|
|
["visual.blocks.1.w"],
|
|
["model.language_model.layers.0.w"],
|
|
["model.language_model.layers.1.w"],
|
|
]
|
|
|
|
def test_the_post_block_lands_last_however_early_it_arrives(self):
|
|
"""Megatron-Bridge streams the last pipeline stage's output block before
|
|
its layers. Sweeping un-indexed names after the first layer into a
|
|
trailing group is what lets the gather loop walk groups ascending."""
|
|
names = ["model.layers.0.w", "model.norm.weight", "model.layers.1.w"]
|
|
assert layerwise_groups(names) == [
|
|
["model.layers.0.w"],
|
|
["model.layers.1.w"],
|
|
["model.norm.weight"],
|
|
]
|
|
|
|
def test_names_sharing_a_layer_coalesce_when_the_source_interleaves(self):
|
|
"""A raw-checkpoint source may yield in shard-packing order. Coalescing on
|
|
the key keeps that from shattering one layer into many groups."""
|
|
names = [
|
|
"model.layers.0.a",
|
|
"model.layers.1.a",
|
|
"model.layers.0.b",
|
|
]
|
|
assert layerwise_groups(names) == [
|
|
["model.layers.0.a", "model.layers.0.b"],
|
|
["model.layers.1.a"],
|
|
]
|
|
|
|
def test_no_empty_groups(self):
|
|
for names in ([], ["only.pre"], ["model.layers.0.a"]):
|
|
assert all(g for g in layerwise_groups(names))
|
|
|
|
def test_a_model_with_no_layers_is_one_pre_group(self):
|
|
assert layerwise_groups(["a", "b"]) == [["a", "b"]]
|
|
|
|
@pytest.mark.parametrize(
|
|
"prefix",
|
|
[
|
|
"model.layers.", # Llama, Qwen3, GLM-4.5, DeepSeek, Mixtral
|
|
"model.language_model.layers.", # Qwen3.5/3.6, Qwen2.5-VL, Gemma3
|
|
"language_model.layers.", # Kimi-K2.5, Nemotron-VL
|
|
"thinker.model.layers.", # Qwen2.5/3-Omni
|
|
"transformer.h.", # GPT-2, GPT-J, Bloom, Falcon
|
|
"transformer.blocks.", # MPT, DBRX
|
|
"transformer.encoder.layers.", # ChatGLM
|
|
"backbone.layers.", # Mamba, Mamba2
|
|
"gpt_neox.layers.", # GPT-NeoX
|
|
"model.decoder.layers.", # OPT
|
|
"model.h.", # Exaone
|
|
"model.layers.layers.", # Plamo3
|
|
"bert.encoder.layer.", # BERT, RoBERTa
|
|
"visual.blocks.", # Qwen-VL vision tower
|
|
"vision_tower.vision_model.encoder.layers.", # SigLIP/CLIP towers
|
|
"audio_tower.layers.", # Whisper-style audio towers
|
|
],
|
|
)
|
|
def test_every_supported_naming_convention_partitions_per_layer(self, prefix):
|
|
"""One decoder layer per group under any prefix, with no per-architecture
|
|
table. A literal `model.layers.` match silently yields one whole-model
|
|
group for most of these."""
|
|
names = [f"{prefix}{i}.self_attn.q_proj.weight" for i in range(3)]
|
|
assert layerwise_groups(names) == [[n] for n in names]
|
|
|
|
def test_a_moe_layer_stays_whole(self):
|
|
"""The OUTERMOST index wins, so per-expert names group by layer. Splitting
|
|
on the expert would cut FusedMoE — one leaf module — across groups."""
|
|
names = [
|
|
"model.layers.0.mlp.experts.0.w1.weight",
|
|
"model.layers.0.mlp.experts.1.w1.weight",
|
|
"model.layers.1.mlp.experts.0.w1.weight",
|
|
]
|
|
assert layerwise_groups(names) == [
|
|
[
|
|
"model.layers.0.mlp.experts.0.w1.weight",
|
|
"model.layers.0.mlp.experts.1.w1.weight",
|
|
],
|
|
["model.layers.1.mlp.experts.0.w1.weight"],
|
|
]
|
|
|
|
def test_a_vlm_partitions_the_vision_tower_and_the_text_stack(self):
|
|
"""Qwen-VL layout. Today's literal `model.layers.` match puts every one of
|
|
these in ONE group — the whole model — because the text stack ships as
|
|
`model.language_model.layers.`."""
|
|
names = [
|
|
"visual.patch_embed.proj.weight",
|
|
"visual.blocks.0.attn.qkv.weight",
|
|
"visual.blocks.1.attn.qkv.weight",
|
|
"visual.merger.mlp.weight",
|
|
"model.language_model.layers.0.self_attn.q_proj.weight",
|
|
"lm_head.weight",
|
|
]
|
|
assert layerwise_groups(names) == [
|
|
["visual.patch_embed.proj.weight"],
|
|
["visual.blocks.0.attn.qkv.weight"],
|
|
["visual.blocks.1.attn.qkv.weight"],
|
|
["model.language_model.layers.0.self_attn.q_proj.weight"],
|
|
["visual.merger.mlp.weight", "lm_head.weight"],
|
|
]
|
|
|
|
def test_two_stacks_sharing_an_index_do_not_merge(self):
|
|
"""The key carries the prefix, so vision block 0 and text layer 0 are
|
|
different groups even though both are index 0."""
|
|
names = [
|
|
"visual.blocks.0.attn.qkv.weight",
|
|
"model.language_model.layers.0.self_attn.q_proj.weight",
|
|
]
|
|
assert layerwise_groups(names) == [[n] for n in names]
|
|
|
|
|
|
class TestDtypeFromName:
|
|
def test_resolves_wire_names(self):
|
|
assert _dtype_from_name("bfloat16") is torch.bfloat16
|
|
assert _dtype_from_name("float8_e4m3fn") is torch.float8_e4m3fn
|
|
|
|
def test_rejects_a_non_dtype_attribute(self):
|
|
with pytest.raises(ValueError):
|
|
_dtype_from_name("nn") # torch.nn exists but is not a dtype
|
|
|
|
def test_rejects_an_unknown_name(self):
|
|
with pytest.raises(ValueError):
|
|
_dtype_from_name("float9")
|
|
|
|
|
|
class TestOpAllowlistAgreement:
|
|
"""The consumer records op chains from ``SUPPORTED_OPS``; the producer
|
|
replays a chain only if every op is in ``ALLOWED_OPS``. The two must describe
|
|
the same contract, so they are derived from one table.
|
|
|
|
A divergence is silent until first pull: an op the consumer can emit but the
|
|
producer rejects bakes fine at init, then fails a whole sync later with
|
|
"disallowed op".
|
|
"""
|
|
|
|
def test_the_two_allowlists_describe_the_same_contract(self):
|
|
assert set(SUPPORTED_OPS.values()) == set(ALLOWED_OPS)
|
|
|
|
def test_every_bake_recordable_op_is_replayable(self):
|
|
"""Anything the bake can record must be serveable, or the failure lands a
|
|
whole sync later than the mistake."""
|
|
assert not set(SUPPORTED_OPS.values()) - set(ALLOWED_OPS)
|
|
|
|
def test_the_producer_allows_nothing_the_consumer_cannot_emit(self):
|
|
"""The allowlist is also a guard against a misbehaving or spoofed
|
|
consumer invoking arbitrary methods on trainer tensors, so it must not be
|
|
wider than the ops the bake can actually produce. ``to`` in particular
|
|
would let a replay change dtype or device."""
|
|
assert not set(ALLOWED_OPS) - set(SUPPORTED_OPS.values())
|
|
assert "to" not in ALLOWED_OPS
|
|
|
|
def test_transpose_via_t_is_serveable(self):
|
|
"""The specific regression: a chain recorded from ``.t()`` must pass the
|
|
producer's guard."""
|
|
fake = FakeRDTTensor(
|
|
name="w",
|
|
shape=torch.Size((4, 6)),
|
|
dtype=torch.bfloat16,
|
|
device=META,
|
|
sink=None,
|
|
)
|
|
(op, _args, _kwargs) = fake.t()._key()[1][0]
|
|
assert op in ALLOWED_OPS
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Bake against a real model (GPU)
|
|
# ---------------------------------------------------------------------------
|
|
|
|
TINY_QWEN = "Qwen/Qwen3-0.6B"
|
|
|
|
# One small model per attention/FFN shape the bake has to survive. The point is
|
|
# the checkpoint->module mapping, which differs structurally between them: GQA
|
|
# fuses Q/K/V of unequal head counts into one param, MLA replaces them with the
|
|
# compressed kv_a/kv_b pair, and MoE turns per-expert checkpoint entries into
|
|
# stacked expert params.
|
|
#
|
|
# Keep them small: `load_format="dummy"` skips the download but not the
|
|
# allocation, so every parameter is still built on the device at full size.
|
|
BAKE_MODELS = [
|
|
pytest.param(TINY_QWEN, id="gqa-dense"),
|
|
pytest.param("TitanML/tiny-mixtral", id="moe"),
|
|
# MLA (kv_lora_rank set, so `ModelConfig.use_mla`) plus routed experts. Only
|
|
# the STRUCTURE is under test -- `_bake_probe` synthesizes the checkpoint
|
|
# from the safetensors headers and never reads a weight value -- so a random
|
|
# 0.06B model does the job, as long as it has the MLA head shape the prefill
|
|
# backends accept: (qk_nope=128, qk_rope=64, v=128).
|
|
#
|
|
# Not an fp8 checkpoint: `process_weights_after_loading` swaps the vLLM
|
|
# parameter subclasses out on those, and the plain second `load_weights`
|
|
# this test compares against then cannot run.
|
|
pytest.param("hmellor/tiny-random-DeepseekV2ForCausalLM", id="mla-moe"),
|
|
]
|
|
|
|
|
|
class TestRequiresTheRayExecutor:
|
|
"""The engine's data plane is Ray's, so a non-Ray executor cannot work. The
|
|
check is at construction because the alternative is an opaque failure during
|
|
the first handshake, long after the misconfiguration."""
|
|
|
|
def _construct(self, backend):
|
|
cfg = SimpleNamespace(
|
|
parallel_config=SimpleNamespace(distributed_executor_backend=backend)
|
|
)
|
|
eng = object.__new__(ShardedRDTWeightTransferEngine)
|
|
# Only the guard is under test, so stand in for the base __init__.
|
|
with pytest.MonkeyPatch.context() as m:
|
|
m.setattr(
|
|
ShardedRDTWeightTransferEngine.__bases__[0],
|
|
"__init__",
|
|
lambda self, *a, **k: None,
|
|
)
|
|
ShardedRDTWeightTransferEngine.__init__(eng, None, cfg, META, None)
|
|
return eng
|
|
|
|
@pytest.mark.parametrize("backend", ["uni", "mp", "external_launcher"])
|
|
def test_a_non_ray_executor_is_refused(self, backend):
|
|
with pytest.raises(ValueError, match="requires distributed_executor"):
|
|
self._construct(backend)
|
|
|
|
def test_ray_is_accepted(self):
|
|
self._construct("ray")
|
|
|
|
def test_a_custom_executor_class_is_left_alone(self):
|
|
"""A `type[Executor]` override is deliberate and unjudgeable here, so it
|
|
must not be rejected for merely not being the string "ray"."""
|
|
|
|
class _CustomExecutor:
|
|
pass
|
|
|
|
self._construct(_CustomExecutor)
|
|
|
|
|
|
def _bake_probe(self, model_name): # noqa: PLR0915 (runs in the vLLM worker)
|
|
"""Drive `_bake` against the worker's real model and check the three
|
|
properties. Returns a small summary; every assertion fires here, in the
|
|
worker, so a failure surfaces as the RPC's exception.
|
|
|
|
Self-contained on purpose: this is pickled to the worker, so it imports what
|
|
it needs rather than closing over this module's globals.
|
|
"""
|
|
import torch
|
|
from huggingface_hub import get_safetensors_metadata
|
|
|
|
from vllm.distributed.weight_transfer.sharded_rdt_common import ALLOWED_OPS
|
|
from vllm.distributed.weight_transfer.sharded_rdt_engine import (
|
|
ShardedRDTWeightTransferEngine,
|
|
ShardedRDTWeightTransferInitInfo,
|
|
)
|
|
from vllm.model_executor.model_loader.reload.utils import get_layer_tensors
|
|
|
|
model = self.model_runner.model
|
|
device = next(model.parameters()).device
|
|
|
|
# ---- the checkpoint, exactly as a trainer would describe and serve it ----
|
|
# Only the safetensors HEADERS are fetched (a few KB); the values are
|
|
# synthesized. Nothing here depends on real weights -- both the reference
|
|
# load and the replay read this same dict, so the comparison is unaffected,
|
|
# and it keeps a multi-GB download out of CI.
|
|
st_dtype = {
|
|
"F64": torch.float64,
|
|
"F32": torch.float32,
|
|
"F16": torch.float16,
|
|
"BF16": torch.bfloat16,
|
|
"F8_E4M3": torch.float8_e4m3fn,
|
|
"F8_E5M2": torch.float8_e5m2,
|
|
"I64": torch.int64,
|
|
"I32": torch.int32,
|
|
"I8": torch.int8,
|
|
"U8": torch.uint8,
|
|
"BOOL": torch.bool,
|
|
}
|
|
meta = get_safetensors_metadata(model_name)
|
|
specs = {}
|
|
for file_meta in meta.files_metadata.values():
|
|
for name, t in file_meta.tensors.items():
|
|
specs[name] = (st_dtype[t.dtype], tuple(t.shape))
|
|
|
|
gen = torch.Generator(device="cpu").manual_seed(0)
|
|
ckpt: dict[str, torch.Tensor] = {}
|
|
for name, (dt, shape) in specs.items():
|
|
if dt.is_floating_point:
|
|
# Small magnitudes so any narrowing cast stays finite.
|
|
v = torch.randn(shape, generator=gen, dtype=torch.float32) * 0.1
|
|
ckpt[name] = v.to(dtype=dt)
|
|
else:
|
|
ckpt[name] = torch.zeros(shape, dtype=dt)
|
|
names = sorted(ckpt)
|
|
dtype_names = [str(ckpt[n].dtype).split(".")[-1] for n in names]
|
|
shapes = [list(ckpt[n].shape) for n in names]
|
|
|
|
def _feed():
|
|
return iter([(n, ckpt[n]) for n in names])
|
|
|
|
# ---- reference: what the model's OWN loader produces from that checkpoint --
|
|
# Poison first, so "which params did load_weights actually write?" is
|
|
# answerable. Some params have no checkpoint source at all -- derived ones,
|
|
# or ones process_weights_after_loading builds -- and those keep their NaNs
|
|
# and are legitimately outside the plan. That set is what the plan must
|
|
# cover, not every parameter of the model.
|
|
def _poison():
|
|
with torch.no_grad():
|
|
for prm in model.parameters():
|
|
if prm.is_floating_point():
|
|
prm.fill_(float("nan"))
|
|
else:
|
|
prm.zero_()
|
|
|
|
_poison()
|
|
model.load_weights(_feed())
|
|
loaded = {
|
|
n
|
|
for n, prm in model.named_parameters()
|
|
if not (prm.is_floating_point() and torch.isnan(prm).any())
|
|
}
|
|
# On the host: the weights themselves already fill the GPU.
|
|
reference = {
|
|
n: prm.detach().to("cpu", copy=True)
|
|
for n, prm in model.named_parameters()
|
|
if n in loaded
|
|
}
|
|
|
|
# ---- pre-bake state, for the restoration check --------------------------
|
|
param_name_of = {}
|
|
for mod_name, mod in model.named_modules():
|
|
for pname in get_layer_tensors(mod):
|
|
param_name_of[(id(mod), pname)] = f"{mod_name}.{pname}".lstrip(".")
|
|
loaders_before = {}
|
|
for mod in model.modules():
|
|
for pname, tensor in get_layer_tensors(mod).items():
|
|
loaders_before[(id(mod), pname)] = getattr(tensor, "weight_loader", None)
|
|
|
|
# ---- bake ---------------------------------------------------------------
|
|
eng = object.__new__(ShardedRDTWeightTransferEngine)
|
|
eng.model = model
|
|
eng.device = device
|
|
eng._name_to_plan = {}
|
|
eng._name_meta = {}
|
|
eng._live_names = set()
|
|
eng._bake(
|
|
ShardedRDTWeightTransferInitInfo(
|
|
names=names, dtype_names=dtype_names, shapes=shapes
|
|
)
|
|
)
|
|
|
|
# (1) the bake ran and produced a plan
|
|
assert eng._name_to_plan, "bake produced no plan"
|
|
assert eng._live_names, "bake recorded no live names"
|
|
n_modules = len({id(v) for v in eng._name_to_plan.values()})
|
|
|
|
# (2) every intermediate change to model state was undone
|
|
for mod in model.modules():
|
|
for pname, tensor in get_layer_tensors(mod).items():
|
|
key = (id(mod), pname)
|
|
loader = getattr(tensor, "weight_loader", None)
|
|
assert not hasattr(loader, "_rdt_stamp_inner"), (
|
|
f"recording stamp left on {param_name_of.get(key, key)}"
|
|
)
|
|
assert loader is loaders_before[key], (
|
|
f"weight_loader not restored on {param_name_of.get(key, key)}"
|
|
)
|
|
assert not tensor.is_meta, (
|
|
f"{param_name_of.get(key, key)} left on meta after the dry run"
|
|
)
|
|
for pname, p in model.named_parameters():
|
|
assert torch.equal(p.detach().cpu(), reference[pname]), f"bake mutated {pname}"
|
|
|
|
# (3) replaying the recorded op chains reproduces that same load
|
|
scatters = {id(v): v for v in eng._name_to_plan.values()}
|
|
|
|
# No deferred-attention layer may be a plan target. The engine's quant
|
|
# thread calls `info.reset()` on every module it processes, which makes
|
|
# finalize_layerwise_reload SKIP that module -- and finalize is the only
|
|
# thing that calls the LAYER's process_weights_after_loading, which is where
|
|
# MLA derives W_UK_T/W_UV from kv_b_proj. An attention layer in the plan
|
|
# would therefore leave those derived weights stale after a sync, silently.
|
|
from vllm.model_executor.layers.attention import is_deferred_attention_layer
|
|
|
|
attn_in_plan = sorted(
|
|
{
|
|
type(c.layer).__name__
|
|
for plan in scatters.values()
|
|
for c in plan
|
|
if is_deferred_attention_layer(c.layer)
|
|
}
|
|
)
|
|
assert not attn_in_plan, (
|
|
f"plan targets deferred-attention layer(s) {attn_in_plan}; the quant "
|
|
f"thread would reset them and finalize would skip their "
|
|
f"process_weights_after_loading, staling any derived weights"
|
|
)
|
|
|
|
touched = set()
|
|
_poison()
|
|
with torch.no_grad():
|
|
for plan in scatters.values():
|
|
for c in plan:
|
|
src_name, chain = c.src
|
|
t = ckpt[src_name]
|
|
for op, args, kw in chain:
|
|
assert op in ALLOWED_OPS, f"disallowed op {op!r}"
|
|
t = getattr(t, op)(*args, **dict(kw))
|
|
dest = getattr(c.layer, c.param_name)
|
|
dest.as_strided(c.shape, c.stride, c.offset).copy_(t)
|
|
touched.add(param_name_of[(id(c.layer), c.param_name)])
|
|
|
|
assert loaded, "the reference load wrote nothing"
|
|
# The plan owns everything with a checkpoint source. Anything else the model
|
|
# ends up holding is DERIVED -- under MLA, `process_weights_after_loading`
|
|
# splits kv_b_proj into W_UK_T/W_UV via `replace_parameter` -- and belongs to
|
|
# finalize_layerwise_reload, which `finish_weight_update` runs after the
|
|
# replay. So the uncovered set must contain only params owned by a module
|
|
# that post-processes; a plain loadable weight going missing still fails.
|
|
post_processed = {
|
|
id(mod)
|
|
for mod in model.modules()
|
|
if hasattr(mod, "process_weights_after_loading")
|
|
}
|
|
modules_by_name = dict(model.named_modules())
|
|
uncovered = sorted(loaded - touched)
|
|
derived: list[str] = []
|
|
unexplained: list[str] = []
|
|
for n in uncovered:
|
|
mod_name = n.rsplit(".", 1)[0]
|
|
mod = modules_by_name.get(mod_name)
|
|
is_derived = mod is not None and id(mod) in post_processed
|
|
(derived if is_derived else unexplained).append(n)
|
|
assert not unexplained, (
|
|
f"the plan misses {len(unexplained)} param(s) load_weights wrote that "
|
|
f"nothing derives: {unexplained[:5]}"
|
|
)
|
|
mismatched = [
|
|
n
|
|
for n, prm in model.named_parameters()
|
|
if n in loaded and not torch.equal(prm.detach().cpu(), reference[n])
|
|
]
|
|
assert not mismatched, f"replay differs from load_weights for: {mismatched[:5]}"
|
|
|
|
return {
|
|
"names": len(names),
|
|
"baked": len(eng._name_to_plan),
|
|
"live": len(eng._live_names),
|
|
"modules": n_modules,
|
|
"loaded_params": len(loaded),
|
|
"derived_not_in_plan": len(derived),
|
|
"replayed_params": len(touched),
|
|
"total_params": sum(1 for _ in model.named_parameters()),
|
|
}
|
|
|
|
|
|
@pytest.mark.skipif(
|
|
not torch.cuda.is_available(), reason="the bake drives a real model on GPU"
|
|
)
|
|
class TestBakeOnARealModel:
|
|
"""`_bake` against a real vLLM model, which the CPU tests above cannot reach:
|
|
they drive the planner with hand-built scatters, so nothing there exercises
|
|
the dry run through a real `load_weights` — the part that monkeypatches the
|
|
model and has to put it back.
|
|
"""
|
|
|
|
@pytest.mark.parametrize("model", BAKE_MODELS)
|
|
def test_bake_records_a_replayable_plan_and_restores_the_model(
|
|
self, vllm_runner, monkeypatch, model
|
|
):
|
|
monkeypatch.setenv("VLLM_ALLOW_INSECURE_SERIALIZATION", "1")
|
|
with vllm_runner(
|
|
model, load_format="dummy", enforce_eager=True, max_model_len=1024
|
|
) as llm:
|
|
(summary,) = llm.llm.collective_rpc(_bake_probe, args=(model,))
|
|
print(f"bake summary [{model}]: {summary}")
|
|
# Non-vacuity: the checks above are only worth something if the plan
|
|
# covers the model. Fusion means far fewer leaf modules than names --
|
|
# QKV and gate/up for GQA, the stacked experts for MoE.
|
|
assert summary["baked"] >= 0.95 * summary["names"], summary
|
|
assert summary["live"] >= 0.95 * summary["names"], summary
|
|
assert 0 < summary["modules"] < summary["baked"], summary
|
|
# The probe already asserts the plan covers every param load_weights
|
|
# wrote; this just keeps that from being vacuously few.
|
|
assert summary["loaded_params"] > 0.5 * summary["total_params"], summary
|