| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | import importlib |
| 4 | from types import SimpleNamespace |
| 5 | |
| 6 | import pytest |
| 7 | import torch |
| 8 | import torch.nn as nn |
| 9 | |
| 10 | from vllm.model_executor.model_loader.mtp_validation import ( |
| 11 | disable_mtp_completeness_check, |
| 12 | ) |
| 13 | from vllm.model_executor.models.utils import ( |
| 14 | PPMissingLayer, |
| 15 | spec_decode_needs_target_embed, |
| 16 | ) |
| 17 | from vllm.platforms import current_platform |
| 18 | from vllm.v1.worker.gpu.spec_decode.eagle import utils as eagle_utils |
| 19 | |
| 20 | VOCAB, HIDDEN = 32, 8 |
| 21 | |
| 22 | |
| 23 | def _fake_pp(world_size: int, is_last_rank: bool = True): |
| 24 | return lambda: SimpleNamespace( |
| 25 | world_size=world_size, |
| 26 | is_last_rank=is_last_rank, |
| 27 | is_first_rank=world_size == 1, |
| 28 | ) |
| 29 | |
| 30 | |
| 31 | def _inner(embed: nn.Module | None) -> nn.Module: |
| 32 | inner = nn.Module() |
| 33 | if embed is not None: |
| 34 | inner.embed_tokens = embed |
| 35 | return inner |
| 36 | |
| 37 | |
| 38 | def _embed(fill: float | None = None) -> nn.Embedding: |
| 39 | embed = nn.Embedding(VOCAB, HIDDEN) |
| 40 | if fill is not None: |
| 41 | with torch.no_grad(): |
| 42 | embed.weight.fill_(fill) |
| 43 | return embed |
| 44 | |
| 45 | |
| 46 | @pytest.mark.parametrize("draft_embed", ["loaded", "unset"]) |
| 47 | def test_drafter_without_own_embedding_gets_the_targets(monkeypatch, draft_embed): |
| 48 | monkeypatch.setattr(eagle_utils, "get_pp_group", _fake_pp(2)) |
| 49 | target_embed = _embed() |
| 50 | draft_inner = _inner(_embed() if draft_embed == "loaded" else None) |
| 51 | if draft_embed == "unset": |
| 52 | draft_inner.embed_tokens = None |
| 53 | draft = SimpleNamespace(has_own_embed_tokens=False) |
| 54 | |
| 55 | eagle_utils.maybe_share_target_embed(draft, draft_inner, _inner(target_embed)) |
| 56 | |
| 57 | assert draft_inner.embed_tokens is target_embed |
| 58 | |
| 59 | |
| 60 | def test_missing_target_embedding_raises_instead_of_running_on_garbage(monkeypatch): |
| 61 | monkeypatch.setattr(eagle_utils, "get_pp_group", _fake_pp(2)) |
| 62 | draft_inner = _inner(_embed()) |
| 63 | draft = SimpleNamespace(has_own_embed_tokens=False) |
| 64 | |
| 65 | with pytest.raises(RuntimeError, match="needs the target input embedding"): |
| 66 | eagle_utils.maybe_share_target_embed( |
| 67 | draft, draft_inner, _inner(PPMissingLayer()) |
| 68 | ) |
| 69 | |
| 70 | |
| 71 | def test_drafter_with_distinct_weights_keeps_them(monkeypatch): |
| 72 | monkeypatch.setattr(eagle_utils, "get_pp_group", _fake_pp(2)) |
| 73 | draft_embed = _embed(fill=1.0) |
| 74 | draft_inner = _inner(draft_embed) |
| 75 | draft = SimpleNamespace(has_own_embed_tokens=True) |
| 76 | |
| 77 | eagle_utils.maybe_share_target_embed(draft, draft_inner, _inner(_embed(fill=2.0))) |
| 78 | |
| 79 | assert draft_inner.embed_tokens is draft_embed |
| 80 | |
| 81 | |
| 82 | def test_mtp_style_drafter_is_left_alone_under_pp(monkeypatch): |
| 83 | monkeypatch.setattr(eagle_utils, "get_pp_group", _fake_pp(2)) |
| 84 | draft_embed = _embed() |
| 85 | draft_inner = _inner(draft_embed) |
| 86 | |
| 87 | eagle_utils.maybe_share_target_embed(nn.Module(), draft_inner, _inner(_embed())) |
| 88 | |
| 89 | assert draft_inner.embed_tokens is draft_embed |
| 90 | |
| 91 | |
| 92 | @pytest.mark.parametrize( |
| 93 | "method,pp_size,is_last_rank,expected", |
| 94 | [ |
| 95 | ("eagle", 2, True, True), |
| 96 | ("eagle3", 2, True, True), |
| 97 | ("dflash", 2, True, True), |
| 98 | ("dspark", 2, True, True), |
| 99 | ("eagle3", 1, True, False), |
| 100 | ("eagle3", 2, False, False), |
| 101 | ("mtp", 2, True, False), |
| 102 | (None, 2, True, False), |
| 103 | ], |
| 104 | ) |
| 105 | def test_target_embedding_provisioning( |
| 106 | monkeypatch, method, pp_size, is_last_rank, expected |
| 107 | ): |
| 108 | monkeypatch.setattr( |
| 109 | "vllm.distributed.parallel_state.get_pp_group", |
| 110 | _fake_pp(pp_size, is_last_rank), |
| 111 | raising=True, |
| 112 | ) |
| 113 | speculative_config = None if method is None else SimpleNamespace(method=method) |
| 114 | config = SimpleNamespace(speculative_config=speculative_config) |
| 115 | assert spec_decode_needs_target_embed(config) is expected |
| 116 | |
| 117 | |
| 118 | @pytest.mark.parametrize( |
| 119 | "pp_size,is_last_rank,expected", |
| 120 | [(2, True, True), (2, False, False), (1, True, False)], |
| 121 | ) |
| 122 | def test_mtp_target_embedding_requires_model_opt_in( |
| 123 | monkeypatch, pp_size, is_last_rank, expected |
| 124 | ): |
| 125 | monkeypatch.setattr( |
| 126 | "vllm.distributed.parallel_state.get_pp_group", |
| 127 | _fake_pp(pp_size, is_last_rank), |
| 128 | ) |
| 129 | config = SimpleNamespace(speculative_config=SimpleNamespace(method="mtp")) |
| 130 | assert spec_decode_needs_target_embed(config, include_mtp=True) is expected |
| 131 | assert not spec_decode_needs_target_embed(config) |
| 132 | |
| 133 | |
| 134 | @pytest.fixture |
| 135 | def m3_modules(): |
| 136 | if not current_platform.is_cuda_alike(): |
| 137 | pytest.skip("M3 model imports require CUDA or ROCm dependencies") |
| 138 | backend = "amd" if current_platform.is_rocm() else "nvidia" |
| 139 | root = f"vllm.models.minimax_m3.{backend}" |
| 140 | return importlib.import_module(f"{root}.mtp"), importlib.import_module( |
| 141 | f"{root}.model" |
| 142 | ) |
| 143 | |
| 144 | |
| 145 | @pytest.fixture |
| 146 | def m3_config(): |
| 147 | hf_config = SimpleNamespace( |
| 148 | vocab_size=VOCAB, |
| 149 | hidden_size=HIDDEN, |
| 150 | num_local_experts=0, |
| 151 | num_hidden_layers=0, |
| 152 | rms_norm_eps=1e-5, |
| 153 | ) |
| 154 | return SimpleNamespace( |
| 155 | use_v2_model_runner=True, |
| 156 | load_config=SimpleNamespace(load_format="auto"), |
| 157 | model_config=SimpleNamespace(hf_text_config=hf_config), |
| 158 | speculative_config=SimpleNamespace( |
| 159 | method="mtp", draft_model_config=SimpleNamespace(hf_config=hf_config) |
| 160 | ), |
| 161 | quant_config=None, |
| 162 | cache_config=None, |
| 163 | ) |
| 164 | |
| 165 | |
| 166 | @pytest.fixture |
| 167 | def make_m3_draft(monkeypatch, m3_modules, m3_config): |
| 168 | mtp, _ = m3_modules |
| 169 | inner = _inner(_embed()) |
| 170 | inner.layers = nn.ModuleDict({"0": nn.Linear(1, 1, bias=False)}) |
| 171 | inner.num_mtp_layers = 1 |
| 172 | monkeypatch.setattr(mtp, "MiniMaxM3MultiTokenPredictor", lambda **_: inner) |
| 173 | monkeypatch.setattr(mtp, "ParallelLMHead", lambda *_, **kw: _embed()) |
| 174 | monkeypatch.setattr(mtp, "LogitsProcessor", lambda *_: None) |
| 175 | monkeypatch.setattr( |
| 176 | mtp, "fused_moe_make_expert_params_mapping", lambda *_, **kw: [] |
| 177 | ) |
| 178 | monkeypatch.setattr(mtp, "get_pp_group", _fake_pp(2)) |
| 179 | monkeypatch.setattr(eagle_utils, "get_pp_group", _fake_pp(2)) |
| 180 | return lambda: mtp.MiniMaxM3MTP(vllm_config=m3_config) |
| 181 | |
| 182 | |
| 183 | @pytest.mark.parametrize( |
| 184 | "key", |
| 185 | [None, "model.embed_tokens.weight", "language_model.model.embed_tokens.weight"], |
| 186 | ) |
| 187 | def test_m3_checkpoint_embedding_ownership(make_m3_draft, key): |
| 188 | """Actual M3 loading determines sharing and survives an unrelated reload.""" |
| 189 | draft = make_m3_draft() |
| 190 | assert draft.has_own_embed_tokens |
| 191 | weights = [("model.mtp.layers.0.weight", torch.ones(1, 1))] |
| 192 | if key is not None: |
| 193 | weights.append((key, torch.ones(VOCAB, HIDDEN))) |
| 194 | loaded = draft.load_weights(weights) |
| 195 | assert ("model.embed_tokens.weight" in loaded) is (key is not None) |
| 196 | with disable_mtp_completeness_check(): |
| 197 | draft.load_weights([("lm_head.weight", torch.zeros(VOCAB, HIDDEN))]) |
| 198 | assert draft.has_own_embed_tokens is (key is not None) |
| 199 | target = _inner(_embed(fill=2.0)) |
| 200 | original = draft.model.embed_tokens |
| 201 | eagle_utils.maybe_share_target_embed(draft, draft.model, target) |
| 202 | if key is None: |
| 203 | assert draft.model.embed_tokens is target.embed_tokens |
| 204 | else: |
| 205 | assert draft.model.embed_tokens is original |
| 206 | torch.testing.assert_close(original.weight, torch.ones(VOCAB, HIDDEN)) |
| 207 | |
| 208 | |
| 209 | @pytest.mark.parametrize( |
| 210 | "load_format", ["dummy", "sharded_state", "runai_streamer_sharded", "tensorizer"] |
| 211 | ) |
| 212 | def test_m3_direct_loader_embedding_ownership(make_m3_draft, m3_config, load_format): |
| 213 | """Directly loaded embeddings survive transfers and share only when equal.""" |
| 214 | m3_config.load_config.load_format = load_format |
| 215 | draft = make_m3_draft() |
| 216 | draft.load_state_dict( |
| 217 | {key: torch.ones_like(value) for key, value in draft.state_dict().items()} |
| 218 | ) |
| 219 | with disable_mtp_completeness_check(): |
| 220 | draft.load_weights([("model.mtp.layers.0.weight", torch.zeros(1, 1))]) |
| 221 | target = _inner(_embed(fill=2.0)) |
| 222 | original = draft.model.embed_tokens |
| 223 | eagle_utils.maybe_share_target_embed(draft, draft.model, target) |
| 224 | assert draft.model.embed_tokens is ( |
| 225 | target.embed_tokens if load_format == "dummy" else original |
| 226 | ) |
| 227 | if load_format != "dummy": |
| 228 | equal_target = _inner(_embed(fill=1.0)) |
| 229 | eagle_utils.maybe_share_target_embed(draft, draft.model, equal_target) |
| 230 | assert draft.model.embed_tokens is equal_target.embed_tokens |
| 231 | |
| 232 | |
| 233 | @pytest.mark.parametrize("use_v2,pp_size", [(True, 1), (False, 1), (False, 2)]) |
| 234 | def test_m3_preserves_legacy_embedding_sharing( |
| 235 | monkeypatch, m3_modules, m3_config, make_m3_draft, use_v2, pp_size |
| 236 | ): |
| 237 | mtp, _ = m3_modules |
| 238 | m3_config.use_v2_model_runner = use_v2 |
| 239 | monkeypatch.setattr(mtp, "get_pp_group", _fake_pp(pp_size)) |
| 240 | draft = make_m3_draft() |
| 241 | draft.load_weights( |
| 242 | [ |
| 243 | ("model.mtp.layers.0.weight", torch.ones(1, 1)), |
| 244 | ("model.embed_tokens.weight", torch.ones(VOCAB, HIDDEN)), |
| 245 | ] |
| 246 | ) |
| 247 | assert not hasattr(draft, "has_own_embed_tokens") |
| 248 | |
| 249 | |
| 250 | @pytest.mark.parametrize( |
| 251 | "rank,use_v2,method,expected", |
| 252 | [ |
| 253 | (0, True, "mtp", True), |
| 254 | (1, True, "mtp", False), |
| 255 | (2, True, "mtp", True), |
| 256 | (2, False, "mtp", False), |
| 257 | (2, True, None, False), |
| 258 | ], |
| 259 | ) |
| 260 | def test_m3_target_allocates_embedding_on_required_stage( |
| 261 | monkeypatch, m3_modules, m3_config, rank, use_v2, method, expected |
| 262 | ): |
| 263 | _, model = m3_modules |
| 264 | pp = lambda: SimpleNamespace( |
| 265 | world_size=3, is_first_rank=rank == 0, is_last_rank=rank == 2 |
| 266 | ) |
| 267 | monkeypatch.setattr(model, "get_pp_group", pp) |
| 268 | monkeypatch.setattr("vllm.distributed.parallel_state.get_pp_group", pp) |
| 269 | monkeypatch.setattr(model, "VocabParallelEmbedding", lambda *_, **kw: _embed()) |
| 270 | monkeypatch.setattr(model, "make_layers", lambda *_, **kw: (0, 0, nn.ModuleList())) |
| 271 | monkeypatch.setattr(model, "MiniMAXGemmaRMSNorm", lambda *_, **kw: nn.Identity()) |
| 272 | m3_config.use_v2_model_runner = use_v2 |
| 273 | if method is None: |
| 274 | m3_config.speculative_config = None |
| 275 | target = model.MiniMaxM3Model(vllm_config=m3_config) |
| 276 | assert isinstance(target.embed_tokens, nn.Embedding) is expected |
| 277 | |
| 278 | |
| 279 | @pytest.mark.parametrize("initial_embed", [False, True]) |
| 280 | def test_m3_weight_transfer_updates_embedding_ownership(make_m3_draft, initial_embed): |
| 281 | """Layer transfers retain ownership; newly loaded embeddings become owned.""" |
| 282 | draft = make_m3_draft() |
| 283 | weights = [("model.mtp.layers.0.weight", torch.ones(1, 1))] |
| 284 | if initial_embed: |
| 285 | weights.append(("model.embed_tokens.weight", torch.ones(VOCAB, HIDDEN))) |
| 286 | draft.load_weights(weights) |
| 287 | with disable_mtp_completeness_check(): |
| 288 | draft.load_weights([("model.mtp.layers.0.weight", torch.zeros(1, 1))]) |
| 289 | assert draft.has_own_embed_tokens is initial_embed |
| 290 | draft.load_weights( |
| 291 | [("model.embed_tokens.weight", torch.full((VOCAB, HIDDEN), 3.0))] |
| 292 | ) |
| 293 | assert draft.has_own_embed_tokens |
| 294 | target = _inner(_embed(fill=2.0)) |
| 295 | original = draft.model.embed_tokens |
| 296 | eagle_utils.maybe_share_target_embed(draft, draft.model, target) |
| 297 | assert draft.model.embed_tokens is original |
| 298 | torch.testing.assert_close(original.weight, torch.full((VOCAB, HIDDEN), 3.0)) |