Add MoNe experts pruning support (#50415)
Signed-off-by: Ilya Boytsov <ieboytsov@nebius.com> Signed-off-by: Misha Goin <mgoin64@gmail.com> Signed-off-by: mgoin <mgoin64@gmail.com> Co-authored-by: Ilya Boytsov <ieboytsov@nebius.com> Co-authored-by: Codex <noreply@openai.com> Co-authored-by: Misha Goin <mgoin64@gmail.com> Co-authored-by: Claude Opus 5.5 <noreply@anthropic.com>
This commit is contained in:
Ilya Boytsov
co-authored by
Ilya Boytsov
Codex
Misha Goin
Claude Opus 5.5
Ilya Boytsov
Codex
Misha Goin
Claude Opus 5.5
parent
1484266938
commit
189bbfb9f4
@@ -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.
|
||||
@@ -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"
|
||||
|
||||
@@ -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)
|
||||
@@ -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,
|
||||
|
||||
@@ -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))
|
||||
@@ -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,
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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
|
||||
``<experts prefix>.<expert id>.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
|
||||
@@ -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
|
||||
|
||||
@@ -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
|
||||
):
|
||||
|
||||
@@ -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:
|
||||
|
||||
@@ -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)
|
||||
|
||||
|
||||
@@ -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
|
||||
Reference in New Issue
Block a user