From 189bbfb9f432e283e21dcdf2a78357ae5db05ace Mon Sep 17 00:00:00 2001 From: Ilya Boytsov Date: Tue, 13 Oct 2026 00:18:43 +0100 Subject: [PATCH] Add MoNe experts pruning support (#50415) Signed-off-by: Ilya Boytsov Signed-off-by: Misha Goin Signed-off-by: mgoin Co-authored-by: Ilya Boytsov Co-authored-by: Codex Co-authored-by: Misha Goin Co-authored-by: Claude Opus 5.5 --- docs/design/expert_substitution.md | 121 +++++ pyproject.toml | 2 + tests/kernels/moe/test_expert_substitution.py | 422 ++++++++++++++++++ .../moe/test_unquantized_backend_selection.py | 52 ++- .../model_loader/test_expert_substitution.py | 337 ++++++++++++++ tests/models/test_deepseek_v4_fi_moe_ep.py | 1 + .../model_executor/layers/fused_moe/config.py | 2 + .../layers/fused_moe/expert_substitution.py | 380 ++++++++++++++++ vllm/model_executor/layers/fused_moe/layer.py | 26 +- .../layers/fused_moe/modular_kernel.py | 2 + .../model_loader/tensorizer_loader.py | 2 + vllm/model_executor/model_loader/utils.py | 6 + .../models/expert_substitution.py | 109 +++++ 13 files changed, 1458 insertions(+), 4 deletions(-) create mode 100644 docs/design/expert_substitution.md create mode 100644 tests/kernels/moe/test_expert_substitution.py create mode 100644 tests/model_executor/model_loader/test_expert_substitution.py create mode 100644 vllm/model_executor/layers/fused_moe/expert_substitution.py create mode 100644 vllm/model_executor/models/expert_substitution.py diff --git a/docs/design/expert_substitution.md b/docs/design/expert_substitution.md new file mode 100644 index 0000000000..05e69bab62 --- /dev/null +++ b/docs/design/expert_substitution.md @@ -0,0 +1,121 @@ +# Expert substitution (MoNE) + +[Mixture of Novices and Experts (MoNE)](https://arxiv.org/abs/2507.00390) +prunes MoE layers by replacing redundant experts with "novices" that return a +constant vector. vLLM loads the retained experts into compact weights and adds +the novices' output separately, so pruned experts use no MLP weight memory and +no MoE backend needs to know about them. + +## Checkpoint format + +The checkpoint keeps the original architecture. `config.json` lists the +substituted experts of each decoder layer: + +```json +{ + "architectures": ["DeepseekV2ForCausalLM"], + "n_routed_experts": 64, + "approximate_experts": { + "1": [1, 6, 8, 12], + "2": [0, 3, 17, 40] + } +} +``` + +For VLMs, `approximate_experts` may live on the text config. Retained experts +keep their logical checkpoint names. Each substituted expert stores one +`[hidden_size]` tensor instead of its MLP: + +```text +model.layers.1.mlp.experts.0.gate_proj.weight +model.layers.1.mlp.experts.1.approx_value +``` + +Alternatively, `compression_config.transform_config.expert_substitution` +declares the same information with explicitly named constant tensors: + +```json +{ + "compression_config": { + "transform_config": { + "expert_substitution": { + "version": 1, + "router_semantics": { + "preserve_logical_expert_ids": true, + "preserve_router_weights": true, + "renormalize_after_substitution": false + }, + "targets": { + "model.layers.1.mlp.experts": { + "num_logical_experts": 64, + "weight_layout": "compact_retained_experts", + "replacements": { + "1": { + "format": "constant-v1", + "tensors": {"value": "model.layers.1.mlp.expert_replacements.1.value"} + } + } + } + } + } + } + } +} +``` + +Only the values shown for `version`, `router_semantics`, `weight_layout`, and +`format` are supported. Each target path must contain exactly one layer index. +This form is not yet accepted alongside compressed-tensors quantization, whose +`transform_config` schema does not include `expert_substitution`. + +## Semantics + +For a substituted expert `j` with constant `v_j` and router weight `w_j(x)`, +the expert contributes `w_j(x) v_j`. Routing is unchanged: router weights are +neither renormalized nor redistributed to retained experts. + +## Execution + +`FusedMoEFactory` matches each MoE layer to `approximate_experts` by the layer +index in its prefix and, for a match, builds `SubstitutedRoutedExperts`, which: + +1. allocates weights only for retained experts, in ascending logical order; +2. after routing, sums `w_j(x) v_j` over substituted routes in FP32; +3. rewrites substituted routes as zero-weight routes to physical expert 0 and + retained routes as compact physical IDs, preserving the `[num_tokens, top_k]` + contract of every decomposed backend; +4. adds the constant output on one tensor-parallel rank, before the partial + outputs are all-reduced. + +Zero-weight routes may schedule some unnecessary GEMM work; the memory savings +are unaffected. Monolithic backends route internally and are not selected. + +## Weight loading + +Retained experts are loaded through the standard expert mapping by logical ID +and placed in their compact row. Substituted IDs map to no local row, the same +way expert parallelism skips non-local experts. + +Constant tensors are consumed at the model's `load_weights` boundary, so +initial loading and direct weight updates share one path. `approx_value` tensors of +layers that the model does not build (other pipeline stages, MTP layers) are +passed through to the model's loader, which skips them like any other weight +it does not own. Explicitly named constants are always consumed, and dropped +when their layer is built elsewhere. Loading fails if a local substituted +expert has no constant. + +Direct `load_weights` updates may change individual constants. Layerwise +reloads (`reload_weights` in checkpoint format) are rejected, and the model +keeps its previous weights. + +## Supported configurations + +- Models whose MoE layers use `FusedMoEFactory` with the default + `RoutedExperts`, and whose MoE prefix contains exactly one layer index. +- Unquantized experts. +- Tensor and pipeline parallelism, including MTP draft models. +- Any decomposed MoE backend. + +Expert and data parallelism, EPLB, MoE LoRA, fused shared experts, quantized +experts, and routed input/output transforms are rejected at initialization. +Deferred MoE finalize and layerwise weight reloads are rejected when used. diff --git a/pyproject.toml b/pyproject.toml index bc924440fa..db5a2a3b9c 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -226,6 +226,8 @@ view_seperator = "view_seperator" inverse_std_variences = "inverse_std_variences" [tool.typos.default.extend-words] +MoNE = "MoNE" +mone = "mone" Hel = "Hel" wether = "wether" iy = "iy" diff --git a/tests/kernels/moe/test_expert_substitution.py b/tests/kernels/moe/test_expert_substitution.py new file mode 100644 index 0000000000..70720af69f --- /dev/null +++ b/tests/kernels/moe/test_expert_substitution.py @@ -0,0 +1,422 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from types import SimpleNamespace + +import pytest +import torch + +from tests.kernels.moe.modular_kernel_tools.parallel_utils import ( + ProcessGroupInfo, + parallel_launch_with_config, +) +from vllm.config import ParallelConfig, VllmConfig, set_current_vllm_config +from vllm.forward_context import get_forward_context, set_forward_context +from vllm.model_executor.layers.fused_moe import FusedMoEFactory +from vllm.model_executor.layers.fused_moe.expert_substitution import ( + ConstantExpertSubstitution, + SubstitutedRoutedExperts, + get_expert_substitution_spec, + make_expert_substitution, +) +from vllm.model_executor.layers.fused_moe.experts.flashinfer_cutlass_moe import ( + FlashInferExperts, +) +from vllm.model_executor.models.mixtral import MixtralMoE +from vllm.platforms import current_platform +from vllm.v1.worker.workspace import init_workspace_manager + +MOE_PREFIX = "model.layers.1.mlp.experts" + + +def _model_config(approximate_experts, **kwargs) -> SimpleNamespace: + return SimpleNamespace( + hf_config=SimpleNamespace(approximate_experts=approximate_experts), + **kwargs, + ) + + +def _expert_substitution(targets: dict[str, dict[int, str]]) -> dict: + """The ``compression_config.transform_config.expert_substitution`` format.""" + return { + "version": 1, + "router_semantics": { + "preserve_logical_expert_ids": True, + "preserve_router_weights": True, + "renormalize_after_substitution": False, + }, + "targets": { + path: { + "num_logical_experts": 5, + "weight_layout": "compact_retained_experts", + "replacements": { + str(expert_id): { + "format": "constant-v1", + "tensors": {"value": name}, + } + for expert_id, name in replacements.items() + }, + } + for path, replacements in targets.items() + }, + } + + +def _compressed_model_config(expert_substitution: dict, **kwargs) -> SimpleNamespace: + return SimpleNamespace( + hf_config=SimpleNamespace( + compression_config={ + "transform_config": {"expert_substitution": expert_substitution} + } + ), + **kwargs, + ) + + +def _substitution(hidden_size: int, dtype=torch.float32): + # Experts 1 and 3 of 5 are substituted; 0, 2, 4 map to physical 0, 1, 2. + return ConstantExpertSubstitution(1, 5, (1, 3), hidden_size, dtype) + + +def test_parse_approximate_experts(): + assert get_expert_substitution_spec(None) is None + assert get_expert_substitution_spec(_model_config(None)) is None + spec = get_expert_substitution_spec(_model_config({"1": [3, 1], "2": [], 5: (0,)})) + assert spec.experts == {1: (1, 3), 5: (0,)} + assert spec.value_names == {} + text_config = SimpleNamespace( + hf_text_config=SimpleNamespace(approximate_experts={"0": [2]}), + hf_config=SimpleNamespace(), + ) + assert get_expert_substitution_spec(text_config).experts == {0: (2,)} + + +@pytest.mark.parametrize( + "approximate_experts", + [[1, 2], {"layer": [1]}, {"1": [1, 1]}, {"1": [-1]}, {"1": ["x"]}], +) +def test_parse_approximate_experts_rejects_invalid_entries(approximate_experts): + with pytest.raises(ValueError, match="approximate_experts"): + get_expert_substitution_spec(_model_config(approximate_experts)) + + +def test_parse_expert_substitution_with_explicit_value_names(): + spec = get_expert_substitution_spec( + _compressed_model_config( + _expert_substitution( + { + MOE_PREFIX: {3: "shared", 1: "layer_1.expert_1"}, + "model.layers.2.mlp.experts": {0: "shared"}, + } + ) + ) + ) + assert spec.experts == {1: (1, 3), 2: (0,)} + assert spec.value_names == { + "shared": ((1, 3), (2, 0)), + "layer_1.expert_1": ((1, 1),), + } + + +@pytest.mark.parametrize( + "mutate", + [ + lambda raw: raw.update(version=2), + lambda raw: raw["router_semantics"].update(renormalize_after_substitution=True), + lambda raw: raw["targets"].update({"mlp.experts": raw["targets"][MOE_PREFIX]}), + lambda raw: raw["targets"][MOE_PREFIX].update(weight_layout="dense"), + lambda raw: raw["targets"][MOE_PREFIX]["replacements"]["1"].update( + format="linear-v1" + ), + lambda raw: raw["targets"][MOE_PREFIX]["replacements"]["1"].pop("tensors"), + ], +) +def test_parse_expert_substitution_rejects_unsupported_metadata(mutate): + raw = _expert_substitution({MOE_PREFIX: {1: "value"}}) + mutate(raw) + with pytest.raises(ValueError, match="invalid expert_substitution metadata"): + get_expert_substitution_spec(_compressed_model_config(raw)) + + +def test_parse_rejects_both_formats(): + config = _compressed_model_config(_expert_substitution({MOE_PREFIX: {1: "v"}})) + config.hf_config.approximate_experts = {"1": [1]} + with pytest.raises(ValueError, match="both"): + get_expert_substitution_spec(config) + + +def test_make_expert_substitution_matches_layer_index(): + config = _model_config({"1": [3, 1]}) + substitution = make_expert_substitution(config, MOE_PREFIX, 5, 16, torch.bfloat16) + + assert substitution is not None + assert substitution.layer_idx == 1 + assert substitution.substituted_expert_ids == (1, 3) + assert substitution.num_compute_experts == 3 + assert substitution.values.dtype == torch.bfloat16 + assert [substitution.physical_expert_id(i) for i in range(5)] == [0, -1, 1, -1, 2] + assert ( + make_expert_substitution(config, "model.layers.2.mlp.experts", 5, 16, None) + is None + ) + assert make_expert_substitution(None, MOE_PREFIX, 5, 16, None) is None + explicit = _compressed_model_config(_expert_substitution({MOE_PREFIX: {1: "v"}})) + substitution = make_expert_substitution(explicit, MOE_PREFIX, 5, 16, None) + assert substitution is not None + assert substitution.substituted_expert_ids == (1,) + with pytest.raises(ValueError, match="exactly one layer index"): + make_expert_substitution(config, "mlp.experts", 5, 16, None) + with pytest.raises(ValueError, match="within"): + make_expert_substitution(config, MOE_PREFIX, 3, 16, None) + with pytest.raises(ValueError, match="retain at least one expert"): + make_expert_substitution(_model_config({"1": [0, 1]}), MOE_PREFIX, 2, 16, None) + + +def test_constant_substitution_loads_values(): + substitution = _substitution(2) + substitution.weight_loader(substitution.values, torch.tensor([1.0, 2.0]), 3) + with pytest.raises(ValueError, match="missing approx_value .* \\[1\\]"): + substitution.validate_loaded_values("layer") + substitution.weight_loader(substitution.values, torch.tensor([3.0, 4.0]), 1) + substitution.validate_loaded_values("layer") + torch.testing.assert_close( + substitution.values, torch.tensor([[3.0, 4.0], [1.0, 2.0]]) + ) + with pytest.raises(ValueError, match="not listed"): + substitution.weight_loader(substitution.values, torch.ones(2), 2) + with pytest.raises(ValueError, match="has shape"): + substitution.weight_loader(substitution.values, torch.ones(3), 1) + + +@pytest.mark.parametrize("num_tokens", [0, 1, 3, 7]) +@pytest.mark.parametrize("hidden_size", [2, 4]) +def test_constant_substitution_transforms_routes_and_computes_side_output( + num_tokens, hidden_size +): + substitution = _substitution(2) + substitution.values.data.copy_(torch.tensor([[10.0, 20.0], [30.0, 40.0]])) + + rows = torch.arange(num_tokens) % 3 + topk_ids = torch.tensor([[1, 2], [4, 3], [99, -1]])[rows] + topk_weights = torch.tensor([[0.25, 0.75], [0.6, 0.4], [0.5, 0.5]])[rows] + constant_output = substitution.transform_routes( + torch.zeros(num_tokens, hidden_size), topk_weights, topk_ids + ) + + torch.testing.assert_close( + topk_weights, torch.tensor([[0.0, 0.75], [0.6, 0.0], [0.0, 0.0]])[rows] + ) + torch.testing.assert_close(topk_ids, torch.tensor([[0, 1], [2, 0], [0, 0]])[rows]) + expected = torch.zeros(num_tokens, hidden_size) + expected[:, :2] = torch.tensor([[2.5, 5.0], [12.0, 16.0], [0.0, 0.0]])[rows] + torch.testing.assert_close(constant_output, expected) + + +@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16]) +@pytest.mark.parametrize( + "device", + [ + "cpu", + pytest.param( + "cuda", + marks=pytest.mark.skipif( + not current_platform.is_cuda_alike(), + reason="requires a CUDA-like platform", + ), + ), + ], +) +def test_constant_substitution_preserves_router_weight_precision(dtype, device): + substitution = _substitution(1, dtype).to(device) + substitution.values.data.copy_( + torch.tensor([[60000.0], [-60000.0]], dtype=dtype, device=device) + ) + topk_weights = torch.tensor([[0.5001, 0.4999]], dtype=torch.float32, device=device) + topk_ids = torch.tensor([[1, 3]], device=device) + # Nearly cancelling constants must retain the difference between FP32 weights. + expected = (topk_weights.double() @ substitution.values.double()).to(dtype) + + actual = substitution.transform_routes( + torch.zeros(1, 1, dtype=dtype, device=device), topk_weights, topk_ids + ) + + torch.testing.assert_close(actual, expected, atol=0, rtol=0) + + +def _bare_substituted_routed_experts(**attrs) -> SubstitutedRoutedExperts: + layer = SubstitutedRoutedExperts.__new__(SubstitutedRoutedExperts) + torch.nn.Module.__init__(layer) + layer.expert_substitution = _substitution(2) + for name, value in attrs.items(): + setattr(layer, name, value) + return layer + + +@pytest.mark.parametrize("tp_rank", [0, 1]) +@pytest.mark.parametrize("padded_input", [False, True]) +def test_substituted_routed_experts_adds_constant_on_one_tp_rank(tp_rank, padded_input): + def apply(*, layer, x, topk_weights, topk_ids, **kwargs): + torch.testing.assert_close(topk_weights, torch.tensor([[0.0, 0.75]])) + torch.testing.assert_close(topk_ids, torch.tensor([[0, 1]])) + return torch.ones(x.shape[0], 2) + + layer = _bare_substituted_routed_experts( + moe_config=SimpleNamespace(tp_rank=tp_rank), + quant_method=SimpleNamespace(is_monolithic=False, apply=apply), + ) + layer.expert_substitution.values.data.copy_( + torch.tensor([[10.0, 20.0], [30.0, 40.0]]) + ) + actual = layer.forward_modular( + torch.zeros(1, 4 if padded_input else 2), + torch.tensor([[0.25, 0.75]]), + torch.tensor([[1, 2]]), + ) + expected = torch.ones(1, 2) + if tp_rank == 0: + expected += torch.tensor([[2.5, 5.0]]) + torch.testing.assert_close(actual, expected) + + +@pytest.mark.skipif( + not current_platform.is_cuda_alike(), reason="requires a CUDA-like platform" +) +def test_standard_fused_moe_model_discovers_substitution(dist_init): + vllm_config = VllmConfig() + vllm_config.model_config = _model_config( + {"0": [1]}, dtype=torch.bfloat16, is_moe=True + ) + vllm_config.kernel_config.moe_backend = "triton" + + with set_current_vllm_config(vllm_config): + moe = MixtralMoE( + num_experts=4, + top_k=2, + hidden_size=256, + intermediate_size=512, + params_dtype=torch.bfloat16, + prefix="model.layers.0.block_sparse_moe", + ) + + routed_experts = moe.experts.routed_experts + assert isinstance(routed_experts, SubstitutedRoutedExperts) + assert routed_experts.expert_substitution.substituted_expert_ids == (1,) + assert routed_experts.w13_weight.shape[0] == 3 + assert moe.experts.moe_config.require_decomposed_backend + + +@pytest.mark.skipif( + not current_platform.is_cuda_alike(), reason="requires a CUDA-like platform" +) +@pytest.mark.parametrize( + "moe_backend", + [ + pytest.param("triton", id="triton"), + pytest.param( + "flashinfer_cutlass", + id="flashinfer-cutlass", + marks=pytest.mark.skipif( + not FlashInferExperts._supports_current_device(), + reason="FlashInfer CUTLASS is unavailable on this platform", + ), + ), + ], +) +def test_fused_moe_combines_compact_experts_and_substitution(dist_init, moe_backend): + hidden_size = 256 + vllm_config = VllmConfig() + vllm_config.model_config = _model_config( + {"0": [1, 3]}, dtype=torch.bfloat16, is_moe=True + ) + vllm_config.compilation_config.static_forward_context = {} + vllm_config.kernel_config.moe_backend = moe_backend + + with set_current_vllm_config(vllm_config), set_forward_context(None, vllm_config): + init_workspace_manager(torch.accelerator.current_device_index()) + layer = FusedMoEFactory( + num_experts=4, + top_k=2, + hidden_size=hidden_size, + intermediate_size=512, + params_dtype=torch.bfloat16, + prefix="model.layers.0.mlp.experts", + renormalize=False, + ).cuda() + substitution = layer.routed_experts.expert_substitution + assert layer.routed_experts.w13_weight.shape[0] == 2 + with torch.no_grad(): + layer.routed_experts.w13_weight.normal_(0, 0.01) + layer.routed_experts.w2_weight.normal_(0, 0.01) + substitution.values.normal_(0, 0.01) + layer._quant_method.process_weights_after_loading(layer.routed_experts) + + hidden_states = torch.randn(8, hidden_size, dtype=torch.bfloat16, device="cuda") + router_logits = torch.randn(8, 4, dtype=torch.float32, device="cuda") + topk_weights, topk_ids = layer.router.select_experts( + hidden_states.clone(), router_logits + ) + expected = substitution.transform_routes(hidden_states, topk_weights, topk_ids) + expected += layer._quant_method.apply( + layer=layer.routed_experts, + x=hidden_states.clone(), + topk_weights=topk_weights, + topk_ids=topk_ids, + shared_experts=None, + shared_experts_input=None, + ) + + get_forward_context().all_moe_layers = None + actual = layer(hidden_states.clone(), router_logits) + torch.testing.assert_close(actual, expected, atol=0, rtol=0) + + +def _tp_substitution_worker( + process_group: ProcessGroupInfo, + vllm_config: VllmConfig, + _cpu_group, +) -> None: + with set_forward_context(None, vllm_config): + init_workspace_manager(process_group.local_rank) + layer = FusedMoEFactory( + num_experts=4, + top_k=1, + hidden_size=256, + intermediate_size=512, + params_dtype=torch.bfloat16, + prefix="model.layers.0.mlp.experts", + renormalize=True, + ) + with torch.no_grad(): + layer.routed_experts.w13_weight.zero_() + layer.routed_experts.w2_weight.zero_() + layer.routed_experts.expert_substitution.values.fill_(0.25) + layer._quant_method.process_weights_after_loading(layer.routed_experts) + + hidden_states = torch.zeros(2, 256, dtype=torch.bfloat16) + router_logits = torch.tensor( + [[-100.0, 100.0, -100.0, -100.0]] * 2, + dtype=torch.float32, + ) + get_forward_context().all_moe_layers = None + actual = layer(hidden_states, router_logits) + torch.testing.assert_close( + actual, torch.full_like(actual, 0.25), atol=0, rtol=0 + ) + + +@pytest.mark.skipif( + not current_platform.is_cuda_alike() or current_platform.device_count() < 2, + reason="requires two CUDA-like devices", +) +def test_constant_substitution_is_added_once_with_tp2(): + vllm_config = VllmConfig(parallel_config=ParallelConfig(tensor_parallel_size=2)) + vllm_config.model_config = _model_config( + {"0": [1]}, + dtype=torch.bfloat16, + is_moe=True, + sleep_mode_offload_cudagraph=False, + ) + vllm_config.compilation_config.static_forward_context = {} + vllm_config.kernel_config.moe_backend = "triton" + + parallel_launch_with_config(2, _tp_substitution_worker, vllm_config, None) diff --git a/tests/kernels/moe/test_unquantized_backend_selection.py b/tests/kernels/moe/test_unquantized_backend_selection.py index 25e6b7dec6..b64e909fbd 100644 --- a/tests/kernels/moe/test_unquantized_backend_selection.py +++ b/tests/kernels/moe/test_unquantized_backend_selection.py @@ -136,7 +136,13 @@ def test_arm_cpu_unquantized_kernel_selection( ], ) @patch( - "vllm.utils.flashinfer.has_flashinfer", + "vllm.model_executor.layers.fused_moe.experts.flashinfer_cutlass_moe." + "has_flashinfer_cutlass_fused_moe", + return_value=False, +) +@patch( + "vllm.model_executor.layers.fused_moe.experts.trtllm_bf16_moe." + "has_flashinfer_trtllm_fused_moe", return_value=False, ) @patch( @@ -145,13 +151,18 @@ def test_arm_cpu_unquantized_kernel_selection( ) def test_select_default_backend_by_platform( mock_aiter_enabled, - mock_has_flashinfer, + mock_has_flashinfer_trtllm, + mock_has_flashinfer_cutlass, monkeypatch, platform_method, expected_backend, ): """Test default backend selection per platform with all optional - accelerators (FlashInfer, AITER) disabled.""" + accelerators (FlashInfer, AITER) disabled. + + Patch the feature probes directly because their cached results can outlive + changes to has_flashinfer(). + """ with patch( "vllm.model_executor.layers.fused_moe.oracle.unquantized.current_platform" ) as mock_platform: @@ -407,6 +418,41 @@ def test_select_cuda_flashinfer_trtllm_modular_for_standard_all2all( assert experts_cls.__name__ == "TrtLlmBf16ExpertsModular" +@patch( + "vllm.model_executor.layers.fused_moe.experts.trtllm_bf16_moe.TrtLlmBf16ExpertsBase._supports_current_device", + return_value=True, +) +@pytest.mark.parametrize("require_decomposed_backend", [False, True]) +def test_require_decomposed_backend_rejects_monolithic_kernels( + mock_supports_current_device, + require_decomposed_backend, +): + """Expert substitution adjusts routes, so the kernel must not route itself.""" + from vllm.model_executor.layers.fused_moe.experts.trtllm_bf16_moe import ( + TrtLlmBf16ExpertsMonolithic, + ) + from vllm.model_executor.layers.fused_moe.modular_kernel import ( + FusedMoEActivationFormat, + ) + + moe_config = make_dummy_moe_config() + moe_config.routing_method = RoutingMethodType.Renormalize + moe_config.require_decomposed_backend = require_decomposed_backend + + supported, reason = TrtLlmBf16ExpertsMonolithic.is_supported_config( + TrtLlmBf16ExpertsMonolithic, + moe_config, + None, + None, + FusedMoEActivationFormat.Standard, + ) + + assert supported != require_decomposed_backend + assert (reason is not None and "routes adjusted" in reason) == ( + require_decomposed_backend + ) + + @patch( "vllm.model_executor.layers.fused_moe.experts.trtllm_bf16_moe.TrtLlmBf16ExpertsBase._supports_current_device", return_value=True, diff --git a/tests/model_executor/model_loader/test_expert_substitution.py b/tests/model_executor/model_loader/test_expert_substitution.py new file mode 100644 index 0000000000..926720c686 --- /dev/null +++ b/tests/model_executor/model_loader/test_expert_substitution.py @@ -0,0 +1,337 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from pathlib import Path +from types import SimpleNamespace +from unittest.mock import Mock + +import pytest +import torch +from safetensors.torch import save_file +from transformers import DeepseekV2Config, DeepseekV2ForCausalLM + +from vllm.config import ( + AttentionConfig, + DeviceConfig, + KernelConfig, + LoadConfig, + ModelConfig, + VllmConfig, +) +from vllm.model_executor.layers.fused_moe.expert_substitution import ( + ConstantExpertSubstitution, +) +from vllm.model_executor.model_loader import get_model_loader +from vllm.model_executor.models.expert_substitution import as_expert_substitution_model +from vllm.model_executor.models.utils import AutoWeightsLoader, PPMissingLayer +from vllm.platforms import current_platform +from vllm.v1.attention.backends.registry import AttentionBackendEnum + +MODEL_CONFIG = SimpleNamespace( + hf_config=SimpleNamespace(approximate_experts={"0": [1, 2], "1": [1, 2]}) +) + + +def _compression_config(value_names: dict[int, dict[int, str]]) -> dict: + """Explicitly named constants given as ``{layer: {expert: tensor name}}``.""" + targets = { + f"model.layers.{layer}.mlp.experts": { + "num_logical_experts": 3, + "weight_layout": "compact_retained_experts", + "replacements": { + str(expert): {"format": "constant-v1", "tensors": {"value": name}} + for expert, name in names.items() + }, + } + for layer, names in value_names.items() + } + router_semantics = { + "preserve_logical_expert_ids": True, + "preserve_router_weights": True, + "renormalize_after_substitution": False, + } + return { + "transform_config": { + "expert_substitution": { + "version": 1, + "router_semantics": router_semantics, + "targets": targets, + } + } + } + + +# Expert 1 of both layers shares one explicitly named constant. +EXPLICIT_MODEL_CONFIG = SimpleNamespace( + hf_config=SimpleNamespace( + compression_config=_compression_config( + { + layer: {1: "constants.shared", 2: f"constants.layer_{layer}"} + for layer in (0, 1) + } + ) + ) +) + + +def _approx_value(layer: int, expert: int) -> str: + return f"model.layers.{layer}.mlp.experts.{expert}.approx_value" + + +def _cpu_model(local_layers=(0, 1), tracked=True, model_config=MODEL_CONFIG): + """Two decoder layers; the others are placeholders as on another PP stage.""" + + class Model(torch.nn.Module): + def __init__(self, prefix=""): + super().__init__() + self.weight = torch.nn.Parameter(torch.zeros(2), requires_grad=False) + self.model = torch.nn.Module() + self.model.layers = torch.nn.ModuleList() + for index in range(2): + if index not in local_layers: + self.model.layers.append(PPMissingLayer()) + continue + layer = torch.nn.Module() + layer.mlp = torch.nn.Module() + layer.mlp.experts = ConstantExpertSubstitution( + index, 3, (1, 2), 2, torch.float32 + ) + self.model.layers.append(layer) + + def load_weights(self, weights): + loaded = AutoWeightsLoader(self).load_weights(weights) + return loaded if tracked else None + + return as_expert_substitution_model(Model, model_config)() + + +@pytest.mark.parametrize("local_layers", [(0, 1), (0,), (1,), ()]) +def test_loads_local_approx_values_and_passes_through_others(local_layers): + """Other PP stages and MTP drafts own none or some of the configured layers.""" + model = _cpu_model(local_layers) + loaded = model.load_weights( + [ + ("weight", torch.tensor([1.0, 2.0])), + *( + (_approx_value(layer, expert), torch.full((2,), 10.0 * layer + expert)) + for layer in range(2) + for expert in (1, 2) + ), + ] + ) + model.process_weights_after_loading() + + assert loaded == {"weight"} | { + f"model.layers.{layer}.mlp.experts.values" for layer in local_layers + } + for layer in local_layers: + torch.testing.assert_close( + model.model.layers[layer].mlp.experts.values, + torch.tensor([[1.0, 1.0], [2.0, 2.0]]) + 10.0 * layer, + ) + + +@pytest.mark.parametrize("local_layers", [(0, 1), (1,), ()]) +def test_loads_explicitly_named_values(local_layers): + """Named values are consumed even when their layer is built elsewhere.""" + model = _cpu_model(local_layers, model_config=EXPLICIT_MODEL_CONFIG) + loaded = model.load_weights( + [ + ("weight", torch.tensor([1.0, 2.0])), + ("constants.shared", torch.full((2,), 1.0)), + ("constants.layer_0", torch.full((2,), 2.0)), + ("constants.layer_1", torch.full((2,), 3.0)), + ] + ) + model.process_weights_after_loading() + + assert loaded == {"weight"} | { + f"model.layers.{layer}.mlp.experts.values" for layer in local_layers + } + for layer in local_layers: + torch.testing.assert_close( + model.model.layers[layer].mlp.experts.values, + torch.tensor([[1.0, 1.0], [2.0 + layer, 2.0 + layer]]), + ) + + +def test_direct_reload_updates_values_incrementally(): + model = _cpu_model(local_layers=(0,), tracked=False) + values = model.model.layers[0].mlp.experts.values + model.load_weights([(_approx_value(0, 2), torch.ones(2))]) + with pytest.raises(ValueError, match="missing approx_value .* \\[1\\]"): + model.process_weights_after_loading() + + assert model.load_weights([(_approx_value(0, 1), torch.zeros(2))]) is None + model.process_weights_after_loading() + model.load_weights([(_approx_value(0, 2), torch.full((2,), 3.0))]) + torch.testing.assert_close(values, torch.tensor([[0.0, 0.0], [3.0, 3.0]])) + + with pytest.raises(ValueError, match="not listed"): + model.load_weights([(_approx_value(0, 0), torch.ones(2))]) + + +def test_layerwise_reload_is_rejected_and_keeps_weights(): + from vllm.model_executor.model_loader.reload import ( + initialize_layerwise_reload, + record_metadata_for_reloading, + ) + + model = _cpu_model(local_layers=(0,)) + record_metadata_for_reloading(model) + loaded = [ + ("weight", torch.tensor([3.0, 4.0])), + (_approx_value(0, 1), torch.zeros(2)), + (_approx_value(0, 2), torch.ones(2)), + ] + model.load_weights(loaded) + model.process_weights_after_loading() + substitution = model.model.layers[0].mlp.experts + values = substitution.values + + initialize_layerwise_reload(model) + with pytest.raises(NotImplementedError, match="layerwise weight reload"): + model.load_weights([(name, weight + 1) for name, weight in loaded]) + + assert substitution.values is values + torch.testing.assert_close(values, torch.tensor([[0.0, 0.0], [1.0, 1.0]])) + torch.testing.assert_close(model.weight, torch.tensor([3.0, 4.0])) + assert substitution.expert_substitution_routes.tolist() == [0, -1, -2] + + +@pytest.mark.parametrize("missing_value", [False, True]) +def test_tensorizer_finalizes_local_substitution_values(monkeypatch, missing_value): + """Raw Tensorizer loading must reject an unloaded local constant row.""" + from vllm.model_executor.model_loader import tensorizer_loader + + model = _cpu_model(local_layers=(0,)) + weights = [(_approx_value(0, 1), torch.zeros(2))] + if not missing_value: + weights.append((_approx_value(0, 2), torch.ones(2))) + finalize = Mock(wraps=model.process_weights_after_loading) + monkeypatch.setattr(model, "process_weights_after_loading", finalize) + monkeypatch.setattr(tensorizer_loader, "initialize_model", lambda **kwargs: model) + loader = object.__new__(tensorizer_loader.TensorizerLoader) + monkeypatch.setattr(loader, "_get_weights_iterator", lambda: iter(weights)) + config = SimpleNamespace( + model_config=SimpleNamespace(dtype=torch.float32), + device_config=SimpleNamespace(device="cpu"), + ) + + if missing_value: + with pytest.raises(ValueError, match=r"missing approx_value .* \[2\]"): + loader._load_model_serialized_cpu(config) + else: + assert loader._load_model_serialized_cpu(config) is model + torch.testing.assert_close( + model.model.layers[0].mlp.experts.values, + torch.tensor([[0.0, 0.0], [1.0, 1.0]]), + ) + finalize.assert_called_once_with() + + +def _write_tiny_deepseek_mone_checkpoint( + model_dir: Path, explicit_names: bool +) -> tuple[dict[str, torch.Tensor], str]: + config = DeepseekV2Config( + architectures=["DeepseekV2ForCausalLM"], + vocab_size=32, + hidden_size=16, + intermediate_size=32, + moe_intermediate_size=8, + num_hidden_layers=1, + num_attention_heads=2, + num_key_value_heads=2, + max_position_embeddings=32, + first_k_dense_replace=0, + n_routed_experts=3, + n_shared_experts=1, + num_experts_per_tok=1, + n_group=1, + topk_group=1, + norm_topk_prob=False, + q_lora_rank=None, + kv_lora_rank=4, + qk_nope_head_dim=128, + qk_rope_head_dim=64, + v_head_dim=128, + tie_word_embeddings=False, + dtype="float16", + ) + experts = "model.layers.0.mlp.experts" + if explicit_names: + value_name = "model.layers.0.mlp.expert_replacements.1.value" + config.compression_config = _compression_config({0: {1: value_name}}) + else: + value_name = f"{experts}.1.approx_value" + config.approximate_experts = {"0": [1]} + + hf_weights = DeepseekV2ForCausalLM(config).to(torch.float16).state_dict() + gate_up = hf_weights.pop(f"{experts}.gate_up_proj") + down = hf_weights.pop(f"{experts}.down_proj") + weights = {name: value.contiguous() for name, value in hf_weights.items()} + for expert in (0, 2): + gate, up = gate_up[expert].chunk(2, dim=0) + weights[f"{experts}.{expert}.gate_proj.weight"] = gate.contiguous() + weights[f"{experts}.{expert}.up_proj.weight"] = up.contiguous() + weights[f"{experts}.{expert}.down_proj.weight"] = down[expert].contiguous() + weights[value_name] = torch.arange(16, dtype=torch.float16) + config.save_pretrained(model_dir) + save_file(weights, model_dir / "model.safetensors") + return weights, value_name + + +@pytest.mark.skipif( + not current_platform.is_cuda_alike(), reason="requires a CUDA-like device" +) +@pytest.mark.usefixtures("dist_init", "workspace_init") +@pytest.mark.parametrize( + "explicit_names", [False, True], ids=["approximate_experts", "expert_substitution"] +) +def test_loads_mone_checkpoint_into_compact_experts( + tmp_path: Path, explicit_names: bool +): + weights, value_name = _write_tiny_deepseek_mone_checkpoint(tmp_path, explicit_names) + model_config = ModelConfig( + model=str(tmp_path), + tokenizer=str(tmp_path), + skip_tokenizer_init=True, + dtype="float16", + max_model_len=32, + enforce_eager=True, + ) + load_config = LoadConfig(load_format="safetensors", use_tqdm_on_load=False) + vllm_config = VllmConfig( + model_config=model_config, + device_config=DeviceConfig(device="cuda"), + load_config=load_config, + attention_config=AttentionConfig(backend=AttentionBackendEnum.TRITON_MLA), + kernel_config=KernelConfig(moe_backend="triton"), + ) + model = get_model_loader(load_config).load_model(vllm_config, model_config) + + experts = "model.layers.0.mlp.experts" + routed_experts = model.model.layers[0].mlp.experts.routed_experts + # Logical expert 2 is stored in physical row 1. + torch.testing.assert_close( + routed_experts.w13_weight[1].cpu(), + torch.cat( + [ + weights[f"{experts}.2.gate_proj.weight"], + weights[f"{experts}.2.up_proj.weight"], + ] + ), + ) + torch.testing.assert_close( + routed_experts.w2_weight[1].cpu(), weights[f"{experts}.2.down_proj.weight"] + ) + substitution = routed_experts.expert_substitution + torch.testing.assert_close( + substitution.values.cpu(), weights[value_name].unsqueeze(0) + ) + + updated = weights[value_name] + 1 + assert model.load_weights([(value_name, updated)]) == { + f"{experts}.routed_experts.expert_substitution.values" + } + torch.testing.assert_close(substitution.values.cpu(), updated.unsqueeze(0)) diff --git a/tests/models/test_deepseek_v4_fi_moe_ep.py b/tests/models/test_deepseek_v4_fi_moe_ep.py index 7687a6786f..d98bb26a2a 100644 --- a/tests/models/test_deepseek_v4_fi_moe_ep.py +++ b/tests/models/test_deepseek_v4_fi_moe_ep.py @@ -142,6 +142,7 @@ def _megakernel_moe(moe_backend: str, **overrides) -> SimpleNamespace: is_act_and_mul=True, activation=MoEActivation.SILU, has_hash_routing=False, + require_decomposed_backend=False, routing_method=RoutingMethodType.DeepseekV4, router_logits_dtype=torch.float32, hidden_dim=256, diff --git a/vllm/model_executor/layers/fused_moe/config.py b/vllm/model_executor/layers/fused_moe/config.py index f20af8a6f8..e79700ae44 100644 --- a/vllm/model_executor/layers/fused_moe/config.py +++ b/vllm/model_executor/layers/fused_moe/config.py @@ -1293,6 +1293,8 @@ class FusedMoEConfig: has_bias: bool = False is_lora_enabled: bool = False has_hash_routing: bool = False + # Routes are adjusted between routing and expert execution. + require_decomposed_backend: bool = False shared_expert_prefix: str | None = None # When True, the MoE skips its final cross-rank all-reduce (and the separate diff --git a/vllm/model_executor/layers/fused_moe/expert_substitution.py b/vllm/model_executor/layers/fused_moe/expert_substitution.py new file mode 100644 index 0000000000..2483200b01 --- /dev/null +++ b/vllm/model_executor/layers/fused_moe/expert_substitution.py @@ -0,0 +1,380 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Constant expert substitution for MoNE-pruned checkpoints. + +MoNE (https://arxiv.org/abs/2507.00390) replaces redundant experts with +"novices" that return a constant vector. Checkpoints list the replaced experts +per decoder layer either in ``config.approximate_experts``, with constants named +``..approx_value``, or in +``compression_config.transform_config.expert_substitution``, with explicitly +named constants. Retained experts keep their logical checkpoint names and are +loaded into compact physical rows. +""" + +from collections.abc import Mapping +from dataclasses import dataclass, field +from typing import TYPE_CHECKING, Any + +import torch +from torch import nn + +from vllm.model_executor.layers.fused_moe.moe_output import UnfinalizedMoEOutput +from vllm.model_executor.layers.fused_moe.routed_experts import RoutedExperts +from vllm.model_executor.layers.fused_moe.unquantized_fused_moe_method import ( + UnquantizedFusedMoEMethod, +) + +if TYPE_CHECKING: + from vllm.model_executor.layers.fused_moe.runner.shared_experts import SharedExperts + +APPROX_VALUE_SUFFIX = ".approx_value" + + +def decoder_layer_index(path: str) -> int | None: + """Return the only integer component of a module path, if unique.""" + indices = [int(part) for part in path.split(".") if part.isdigit()] + return indices[0] if len(indices) == 1 else None + + +@dataclass(frozen=True) +class ExpertSubstitutionSpec: + # Sorted substituted expert IDs per decoder layer. + experts: dict[int, tuple[int, ...]] + # Explicitly named constant tensors and the (layer, expert) pairs they + # supply. Without an entry, constants use the ``approx_value`` naming. + value_names: dict[str, tuple[tuple[int, int], ...]] = field(default_factory=dict) + + +# The only router semantics that the constant-v1 format supports. +_ROUTER_SEMANTICS = { + "preserve_logical_expert_ids": True, + "preserve_router_weights": True, + "renormalize_after_substitution": False, +} + + +def _expert_ids(layer_idx: int, expert_ids: Any) -> tuple[int, ...]: + ids = tuple(sorted(int(expert_id) for expert_id in expert_ids or ())) + if len(set(ids)) != len(ids) or any(expert_id < 0 for expert_id in ids): + raise ValueError( + f"substituted experts of layer {layer_idx} must be unique, " + f"non-negative expert IDs, got {list(ids)}" + ) + return ids + + +def _parse_approximate_experts(raw: Any) -> ExpertSubstitutionSpec: + if not isinstance(raw, Mapping): + raise ValueError("approximate_experts must map layer indices to expert IDs") + experts: dict[int, tuple[int, ...]] = {} + for layer, expert_ids in raw.items(): + try: + layer_idx = int(layer) + experts[layer_idx] = _expert_ids(layer_idx, expert_ids) + except (TypeError, ValueError) as e: + raise ValueError( + f"approximate_experts has an invalid entry for layer {layer!r}" + ) from e + return ExpertSubstitutionSpec({k: v for k, v in experts.items() if v}) + + +def _parse_expert_substitution(raw: Any) -> ExpertSubstitutionSpec: + experts: dict[int, tuple[int, ...]] = {} + value_names: dict[str, list[tuple[int, int]]] = {} + try: + if raw["version"] != 1 or raw["router_semantics"] != _ROUTER_SEMANTICS: + raise ValueError("only version 1 with unchanged routing is supported") + for path, target in raw["targets"].items(): + layer_idx = decoder_layer_index(path) + if layer_idx is None or layer_idx in experts: + raise ValueError(f"target {path!r} must name one decoder layer") + if target["weight_layout"] != "compact_retained_experts": + raise ValueError(f"unsupported weight layout for {path!r}") + for expert_id, replacement in target["replacements"].items(): + if replacement["format"] != "constant-v1": + raise ValueError(f"unsupported replacement format for {path!r}") + value_names.setdefault(replacement["tensors"]["value"], []).append( + (layer_idx, int(expert_id)) + ) + experts[layer_idx] = _expert_ids(layer_idx, target["replacements"]) + except (AttributeError, KeyError, TypeError, ValueError) as e: + raise ValueError(f"invalid expert_substitution metadata: {e}") from e + return ExpertSubstitutionSpec( + experts, {name: tuple(pairs) for name, pairs in value_names.items()} + ) + + +def get_expert_substitution_spec(model_config: Any) -> ExpertSubstitutionSpec | None: + """Parse the substituted experts declared in a model's HF config.""" + hf_configs = [ + getattr(model_config, name, None) for name in ("hf_text_config", "hf_config") + ] + approximate = None + for config in hf_configs: + approximate = getattr(config, "approximate_experts", None) + if approximate is not None: + break + compression_config = getattr(hf_configs[1], "compression_config", None) + transform_config = ( + compression_config.get("transform_config") + if isinstance(compression_config, Mapping) + else None + ) + declared = ( + transform_config.get("expert_substitution") + if isinstance(transform_config, Mapping) + else None + ) + if approximate is not None and declared is not None: + raise ValueError( + "config declares both approximate_experts and " + "compression_config.transform_config.expert_substitution" + ) + if approximate is not None: + spec = _parse_approximate_experts(approximate) + elif declared is not None: + spec = _parse_expert_substitution(declared) + else: + return None + return spec if spec.experts else None + + +class ConstantExpertSubstitution(nn.Module): + """Constant outputs for the substituted experts of one MoE layer. + + Router outputs keep logical expert IDs. ``transform_routes`` maps retained + experts to compact physical rows and turns substituted routes into + zero-weight routes to physical expert 0, so any decomposed MoE backend can + run them unchanged. + """ + + def __init__( + self, + layer_idx: int, + num_logical_experts: int, + substituted_expert_ids: tuple[int, ...], + hidden_size: int, + params_dtype: torch.dtype, + ) -> None: + super().__init__() + substituted = set(substituted_expert_ids) + if not substituted <= set(range(num_logical_experts)): + raise ValueError( + f"approximate_experts for layer {layer_idx} must be within " + f"[0, {num_logical_experts}), got {sorted(substituted)}" + ) + if len(substituted) == num_logical_experts: + raise ValueError( + f"approximate_experts for layer {layer_idx} must retain at least " + "one expert" + ) + self.layer_idx = layer_idx + self.substituted_expert_ids = tuple(sorted(substituted)) + retained = [i for i in range(num_logical_experts) if i not in substituted] + self.num_compute_experts = len(retained) + + # Retained experts map to their physical row r >= 0, substituted + # experts to -1 - (their row in `values`). + routes = [0] * num_logical_experts + for row, expert_id in enumerate(retained): + routes[expert_id] = row + for row, expert_id in enumerate(self.substituted_expert_ids): + routes[expert_id] = -1 - row + self._routes = tuple(routes) + self.register_buffer( + "expert_substitution_routes", + torch.tensor(routes, dtype=torch.int32), + persistent=False, + ) + + # NaN marks values that the checkpoint has not provided yet. + self.values = nn.Parameter( + torch.full((len(substituted), hidden_size), torch.nan, dtype=params_dtype), + requires_grad=False, + ) + self.values.weight_loader = self.weight_loader + + @property + def num_logical_experts(self) -> int: + return len(self._routes) + + def physical_expert_id(self, expert_id: int) -> int: + """Physical row of a retained expert, or -1 for a substituted one.""" + return max(self._routes[expert_id], -1) + + def weight_loader( + self, param: nn.Parameter, loaded_weight: torch.Tensor, expert_id: int + ) -> None: + if not 0 <= expert_id < self.num_logical_experts or ( + self._routes[expert_id] >= 0 + ): + raise ValueError( + f"expert {expert_id} of layer {self.layer_idx} is not listed in " + "approximate_experts" + ) + row = param.data[-1 - self._routes[expert_id]] + if loaded_weight.shape != row.shape: + raise ValueError( + f"approx_value for expert {expert_id} of layer {self.layer_idx} has " + f"shape {tuple(loaded_weight.shape)}, expected {tuple(row.shape)}" + ) + row.copy_(loaded_weight) + + def validate_loaded_values(self, prefix: str) -> None: + missing_rows = torch.isnan(self.values).any(dim=1).nonzero().flatten().tolist() + if missing_rows: + missing = [self.substituted_expert_ids[row] for row in missing_rows] + raise ValueError( + f"{prefix} is missing approx_value for {len(missing)} " + f"expert(s), first missing expert IDs: {missing[:8]}" + ) + + def transform_routes( + self, + hidden_states: torch.Tensor, + topk_weights: torch.Tensor, + topk_ids: torch.Tensor, + ) -> torch.Tensor: + """Remap routes in place and return the constant experts' output.""" + valid = (topk_ids >= 0) & (topk_ids < self.num_logical_experts) + routes = self.expert_substitution_routes[ + torch.where(valid, topk_ids, torch.zeros_like(topk_ids)).long() + ].long() + substituted = valid & (routes < 0) + + # Sum each token's router weights per substituted expert in FP32; other + # routes land in a trailing column that is dropped. + num_values, value_size = self.values.shape + scales = topk_weights.new_zeros( + (topk_ids.shape[0], num_values + 1), dtype=torch.float32 + ) + scales.scatter_add_( + 1, + torch.where(substituted, -1 - routes, num_values), + topk_weights.float(), + ) + output = torch.zeros_like(hidden_states) + output[..., :value_size] = scales[:, :num_values] @ self.values.float() + + topk_weights.masked_fill_(~valid | substituted, 0.0) + topk_ids.copy_(routes.clamp_min(0).to(topk_ids.dtype)) + return output + + +def make_expert_substitution( + model_config: Any, + prefix: str, + num_experts: int, + hidden_size: int, + params_dtype: torch.dtype | None, +) -> ConstantExpertSubstitution | None: + spec = get_expert_substitution_spec(model_config) + if spec is None: + return None + layer_idx = decoder_layer_index(prefix) + if layer_idx is None: + raise ValueError( + f"cannot match MoE layer {prefix!r} to substituted experts: expected " + "exactly one layer index in its prefix" + ) + expert_ids = spec.experts.get(layer_idx) + if expert_ids is None: + return None + return ConstantExpertSubstitution( + layer_idx, + num_experts, + expert_ids, + hidden_size, + params_dtype or torch.get_default_dtype(), + ) + + +class SubstitutedRoutedExperts(RoutedExperts): + """Routed experts that hold weights only for retained experts.""" + + def __init__( + self, *args: Any, expert_substitution: ConstantExpertSubstitution, **kwargs: Any + ): + super().__init__(*args, **kwargs) + self.expert_substitution = expert_substitution + moe_config = self.moe_config + parallel_config = moe_config.moe_parallel_config + unsupported = [ + name + for name, enabled in ( + ( + "quantized MoE weights", + not isinstance(self.quant_method, UnquantizedFusedMoEMethod), + ), + ("monolithic MoE backends", self.quant_method.is_monolithic), + ("MoE LoRA", moe_config.is_lora_enabled), + ( + "EPLB", + parallel_config.enable_eplb + or moe_config.num_experts + != expert_substitution.num_compute_experts, + ), + ( + "fused shared experts", + self.expert_map_manager.num_fused_shared_experts > 0, + ), + ("expert parallelism", parallel_config.use_ep), + ("data parallelism", parallel_config.dp_size > 1), + ("prefill context parallelism", parallel_config.pcp_size > 1), + ("sequence parallelism", parallel_config.is_sequence_parallel), + ("router weights on expert inputs", self.apply_router_weight_on_input), + ) + if enabled + ] + if unsupported: + raise NotImplementedError( + "expert substitution does not support: " + ", ".join(unsupported) + ) + + def _map_global_expert_id_to_local_expert_id(self, expert_id: int) -> int: + # Checkpoints name retained experts by their logical ID. + physical_id = self.expert_substitution.physical_expert_id(expert_id) + if physical_id < 0: + return -1 + return super()._map_global_expert_id_to_local_expert_id(physical_id) + + def get_expert_mapping( + self, + ckpt_gate_proj_name: str | None = None, + ckpt_down_proj_name: str | None = None, + ckpt_up_proj_name: str | None = None, + include_fused: bool = False, + ) -> list[tuple[str, str, int, str]]: + # Fused checkpoint tensors would cover every logical expert at once. + return self.build_expert_params_mapping( + ckpt_gate_proj_name or self.ckpt_gate_proj_name, + ckpt_down_proj_name or self.ckpt_down_proj_name, + ckpt_up_proj_name or self.ckpt_up_proj_name, + num_experts=self.expert_substitution.num_logical_experts, + routed_experts_prefix="", + lora_base_layer_prefix=self.lora_base_layer_prefix, + ) + + def forward_modular( + self, + x: torch.Tensor, + topk_weights: torch.Tensor, + topk_ids: torch.Tensor, + shared_experts: "SharedExperts | None" = None, + shared_experts_input: torch.Tensor | None = None, + ) -> torch.Tensor: + constant_output = self.expert_substitution.transform_routes( + x, topk_weights, topk_ids + ) + output = super().forward_modular( + x, topk_weights, topk_ids, shared_experts, shared_experts_input + ) + if isinstance(output, UnfinalizedMoEOutput): + raise NotImplementedError( + "expert substitution does not support deferred MoE finalize" + ) + # The constant output is replicated across TP ranks, whose partial + # outputs are all-reduced later, so add it on one rank only. + if self.moe_config.tp_rank == 0: + output.add_(constant_output[..., : output.shape[-1]]) + return output diff --git a/vllm/model_executor/layers/fused_moe/layer.py b/vllm/model_executor/layers/fused_moe/layer.py index 6b7c64bdc4..839eba5d78 100644 --- a/vllm/model_executor/layers/fused_moe/layer.py +++ b/vllm/model_executor/layers/fused_moe/layer.py @@ -23,6 +23,10 @@ from vllm.model_executor.layers.fused_moe.config import ( from vllm.model_executor.layers.fused_moe.expert_map_manager import ( ExpertMapManager, ) +from vllm.model_executor.layers.fused_moe.expert_substitution import ( + SubstitutedRoutedExperts, + make_expert_substitution, +) from vllm.model_executor.layers.fused_moe.routed_experts import RoutedExperts from vllm.model_executor.layers.fused_moe.router.fused_moe_router import ( FusedMoERouter, @@ -253,6 +257,12 @@ def FusedMoEFactory( fuse_shared_experts, ) ) + expert_substitution = make_expert_substitution( + vllm_config.model_config, layer_name, num_experts, hidden_size, params_dtype + ) + if expert_substitution is not None: + # Route over all logical experts but allocate only the retained ones. + global_num_experts -= len(expert_substitution.substituted_expert_ids) # Initialize EPLB manager (or None?) eplb_state: EplbLayerState | None = None @@ -367,6 +377,7 @@ def FusedMoEFactory( device=vllm_config.device_config.device, routing_method=router.routing_method_type, # Not ideal has_hash_routing=hash_indices_table is not None, + require_decomposed_backend=expert_substitution is not None, swiglu_limit=swiglu_limit, swiglu_alpha=swiglu_alpha, swiglu_beta=swiglu_beta, @@ -380,7 +391,20 @@ def FusedMoEFactory( # Create RoutedExperts instance BEFORE create_weights() # This will hold all expert weight parameters - if routed_experts_cls is None: + if expert_substitution is not None: + if routed_experts_cls not in (None, RoutedExperts) or ( + routed_input_transform is not None or routed_output_transform is not None + ): + raise NotImplementedError( + "expert substitution requires the default RoutedExperts and no " + "routed input or output transforms" + ) + routed_experts_cls = SubstitutedRoutedExperts + routed_experts_args = { + **(routed_experts_args or {}), + "expert_substitution": expert_substitution, + } + elif routed_experts_cls is None: routed_experts_cls = RoutedExperts assert params_dtype is not None diff --git a/vllm/model_executor/layers/fused_moe/modular_kernel.py b/vllm/model_executor/layers/fused_moe/modular_kernel.py index 7a4cbd4237..d9a75c55ec 100644 --- a/vllm/model_executor/layers/fused_moe/modular_kernel.py +++ b/vllm/model_executor/layers/fused_moe/modular_kernel.py @@ -569,6 +569,8 @@ class FusedMoEExperts(ABC): ) elif moe_config.has_hash_routing and cls.is_monolithic(): return False, _make_reason("hash routing") + elif moe_config.require_decomposed_backend and cls.is_monolithic(): + return False, _make_reason("routes adjusted after routing") elif not cls._supports_routing_method( moe_config.routing_method, weight_key, activation_key ): diff --git a/vllm/model_executor/model_loader/tensorizer_loader.py b/vllm/model_executor/model_loader/tensorizer_loader.py index 338f9eac07..ae74beb7ed 100644 --- a/vllm/model_executor/model_loader/tensorizer_loader.py +++ b/vllm/model_executor/model_loader/tensorizer_loader.py @@ -84,6 +84,8 @@ class TensorizerLoader(BaseModelLoader): model = initialize_model(vllm_config=vllm_config, prefix=prefix) model.load_weights(self._get_weights_iterator()) + if hasattr(model, "process_weights_after_loading"): + model.process_weights_after_loading() return model.eval() def download_model(self, model_config: ModelConfig) -> None: diff --git a/vllm/model_executor/model_loader/utils.py b/vllm/model_executor/model_loader/utils.py index 65d1ff2808..8a23694100 100644 --- a/vllm/model_executor/model_loader/utils.py +++ b/vllm/model_executor/model_loader/utils.py @@ -84,6 +84,12 @@ def initialize_model( if model_class is None: model_class, _ = get_model_architecture(model_config) + from vllm.model_executor.models.expert_substitution import ( + as_expert_substitution_model, + ) + + model_class = as_expert_substitution_model(model_class, model_config) + if vllm_config.quant_config is not None: configure_quant_config(vllm_config.quant_config, model_class) diff --git a/vllm/model_executor/models/expert_substitution.py b/vllm/model_executor/models/expert_substitution.py new file mode 100644 index 0000000000..0799bca8b7 --- /dev/null +++ b/vllm/model_executor/models/expert_substitution.py @@ -0,0 +1,109 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from collections.abc import Iterable +from functools import wraps +from typing import Any, TypeVar + +import torch +from torch import nn + +from vllm.model_executor.layers.fused_moe.expert_substitution import ( + APPROX_VALUE_SUFFIX, + ConstantExpertSubstitution, + decoder_layer_index, + get_expert_substitution_spec, +) +from vllm.model_executor.model_loader.reload import finalize_layerwise_reload + +_T = TypeVar("_T", bound=nn.Module) + + +def _collect_expert_substitutions( + model: nn.Module, +) -> dict[int, tuple[str, ConstantExpertSubstitution]]: + substitutions: dict[int, tuple[str, ConstantExpertSubstitution]] = {} + for name, module in model.named_modules(): + if not isinstance(module, ConstantExpertSubstitution): + continue + if module.layer_idx in substitutions: + raise ValueError( + f"layer {module.layer_idx} has more than one MoE with substituted " + "experts" + ) + substitutions[module.layer_idx] = (f"{name}.values", module) + return substitutions + + +def _parse_approx_value_name(name: str) -> tuple[int, int] | None: + """Return (layer index, expert ID) of an ``approx_value`` tensor name.""" + if not name.endswith(APPROX_VALUE_SUFFIX): + return None + experts_path, _, expert_id = name.removesuffix(APPROX_VALUE_SUFFIX).rpartition(".") + layer_idx = decoder_layer_index(experts_path) + if layer_idx is None or not expert_id.isdigit(): + return None + return layer_idx, int(expert_id) + + +def as_expert_substitution_model(model_cls: type[_T], model_config: Any) -> type[_T]: + """Load constant expert values at the ``load_weights`` boundary. + + Initial loading and direct weight updates both go through ``load_weights``. + Explicitly named values are always consumed and loaded if their layer is + built by this model. ``approx_value`` tensors of layers built elsewhere + (other pipeline stages, MTP layers) are passed through, so the model's own + loader skips them like any other weight of a layer it does not own. + """ + spec = get_expert_substitution_spec(model_config) + if spec is None: + return model_cls + value_names = spec.value_names + + class ExpertSubstitutionModel(model_cls): # type: ignore[valid-type, misc] + @wraps(model_cls.__init__) # type: ignore[misc] + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + self._expert_substitutions = _collect_expert_substitutions(self) + + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]): + if any(s.values.is_meta for _, s in self._expert_substitutions.values()): + # Put the previous weights back so the model stays usable. + finalize_layerwise_reload(self, model_config) + raise NotImplementedError( + "expert substitution does not support layerwise weight reload" + ) + loaded_substitutions: set[str] = set() + + def remaining_weights(): + for name, loaded_weight in weights: + targets = value_names.get(name) + if targets is None: + key = _parse_approx_value_name(name) + if key is None or key[0] not in self._expert_substitutions: + yield name, loaded_weight + continue + targets = (key,) + for layer_idx, expert_id in targets: + if layer_idx not in self._expert_substitutions: + continue + param_name, substitution = self._expert_substitutions[layer_idx] + param = substitution.values + param.weight_loader(param, loaded_weight, expert_id) + loaded_substitutions.add(param_name) + + loaded = super().load_weights(remaining_weights()) + if loaded is not None: + loaded.update(loaded_substitutions) + return loaded + + def process_weights_after_loading(self) -> None: + for param_name, substitution in self._expert_substitutions.values(): + substitution.validate_loaded_values(param_name) + finalize = getattr(super(), "process_weights_after_loading", None) + if finalize is not None: + finalize() + + ExpertSubstitutionModel.__name__ = model_cls.__name__ + ExpertSubstitutionModel.__qualname__ = model_cls.__qualname__ + return ExpertSubstitutionModel