[Bugfix][Model] Fix Aria expert weight names and layout (#57487)
Signed-off-by: Jiang Wu <245355730+JiangLLM@users.noreply.github.com> Co-authored-by: Jiang Wu <245355730+JiangLLM@users.noreply.github.com>
Jiang Wu committed
Sep 18, 2026 at 21:06 UTC
01c7bf88135458940d89584c05ae47404aac77cb
2 files changed
+109
-4
tests/models/multimodal/test_aria.py
new
+100
@@ -0,0 +1,100 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+
4
+import pytest
5
+import torch
6
+from transformers import AriaTextConfig
7
+from transformers.models.aria.modeling_aria import AriaTextMoELayer as HFMoE
8
+
9
+from vllm.config import VllmConfig
10
+from vllm.model_executor.layers.quantization.compressed_tensors.compressed_tensors import ( # noqa: E501
11
+ CompressedTensorsConfig,
12
+)
13
+from vllm.model_executor.models.aria import AriaTextModel, AriaTextMoELayer
14
+from vllm.model_executor.models.utils import AutoWeightsLoader
15
+
16
+
17
+@pytest.mark.cpu_test
18
+@pytest.mark.usefixtures("dist_init")
19
+@pytest.mark.parametrize(
20
+ "checkpoint_name,loaded_name,param_name",
21
+ [
22
+ (
23
+ "experts.fc1.weight",
24
+ "experts.w13_weight",
25
+ "experts.routed_experts.w13_weight",
26
+ ),
27
+ (
28
+ "experts.fc2.weight",
29
+ "experts.w2_weight",
30
+ "experts.routed_experts.w2_weight",
31
+ ),
32
+ (
33
+ "shared_experts.down_proj.weight",
34
+ "shared_experts.down_proj.weight",
35
+ "shared_experts.down_proj.weight",
36
+ ),
37
+ ],
38
+)
39
+def test_aria_expert_weights_load_with_checkpoint_layout(
40
+ checkpoint_name: str, loaded_name: str, param_name: str
41
+):
42
+ """Real Aria checkpoint weights must reach the right parameter and layout."""
43
+ config = AriaTextConfig(
44
+ hidden_size=32,
45
+ intermediate_size=64,
46
+ moe_num_experts=2,
47
+ moe_topk=1,
48
+ moe_num_shared_experts=1,
49
+ )
50
+ checkpoint_weight = HFMoE(config).state_dict()[checkpoint_name]
51
+ checkpoint_weight.copy_(
52
+ torch.arange(
53
+ checkpoint_weight.numel(), dtype=checkpoint_weight.dtype
54
+ ).reshape_as(checkpoint_weight)
55
+ )
56
+ expected = checkpoint_weight.clone()
57
+ if checkpoint_name.startswith("experts."):
58
+ expected = expected.transpose(-1, -2)
59
+
60
+ layer = AriaTextMoELayer(config, quant_config=None, prefix="mlp")
61
+ loaded = AutoWeightsLoader(layer).load_weights(
62
+ [(checkpoint_name, checkpoint_weight)], mapper=AriaTextModel.hf_to_vllm_mapper
63
+ )
64
+
65
+ assert loaded == {loaded_name}
66
+ torch.testing.assert_close(
67
+ dict(layer.named_parameters())[param_name], expected, rtol=0, atol=0
68
+ )
69
+
70
+
71
+@pytest.mark.cpu_test
72
+def test_aria_quant_config_renames_expert_modules():
73
+ quant_config = CompressedTensorsConfig(
74
+ target_scheme_map={},
75
+ ignore=[
76
+ "model.layers.0.mlp.experts.fc1",
77
+ "model.layers.0.mlp.experts.fc2",
78
+ ],
79
+ quant_format="pack-quantized",
80
+ )
81
+ AriaTextModel.__new__(
82
+ AriaTextModel, vllm_config=VllmConfig(quant_config=quant_config)
83
+ )
84
+
85
+ assert quant_config.ignore == [
86
+ "model.layers.0.mlp.experts.gate_up_proj",
87
+ "model.layers.0.mlp.experts.down_proj",
88
+ ]
89
+
90
+
91
+@pytest.mark.cpu_test
92
+@pytest.mark.parametrize(
93
+ "name,expected", [("fc1", "gate_up_proj"), ("fc2", "down_proj")]
94
+)
95
+def test_aria_expert_scale_names(name: str, expected: str):
96
+ suffixes = ["weight_scale", "input_scale"]
97
+ names = [f"model.layers.0.mlp.experts.{name}.{suffix}" for suffix in suffixes]
98
+ assert AriaTextModel.hf_to_vllm_mapper.apply_list(names) == [
99
+ f"model.layers.0.mlp.experts.{expected}.{suffix}" for suffix in suffixes
100
+ ]
vllm/model_executor/models/aria.py
+9
-4
@@ -3,6 +3,7 @@
3
from collections.abc import Iterable, Mapping, Sequence
4
from typing import Annotated, Literal
5
6
+import regex as re
7
import torch
8
import torch.nn as nn
9
from transformers import AriaConfig, AriaTextConfig, BatchFeature
@@ -222,6 +223,7 @@ class AriaTextMoELayer(nn.Module):
223
intermediate_size=config.intermediate_size,
224
quant_config=quant_config,
225
prefix=f"{prefix}.experts",
226
+ is_fused_checkpoint_transposed=True,
227
)
228
229
def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
@@ -267,14 +269,17 @@ class AriaTextModel(LlamaModel, SupportsQuant):
269
"gate_up_proj": ["gate_proj", "up_proj"],
270
}
271
270
- # Aria packs all experts into single (transposed) fc1/fc2 tensors, which is
271
- # exactly the pre-fused checkpoint layout FusedMoE self-loads once fc1/fc2
272
- # are renamed to the fused gate_up_proj/down_proj names.
272
+ # The fused expert loader expects names without the .weight suffix.
273
hf_to_vllm_mapper = LlamaModel.hf_to_vllm_mapper | WeightsMapper(
274
+ orig_to_new_regex={
275
+ re.compile(r"experts\.fc1\.weight$"): "experts.gate_up_proj",
276
+ re.compile(r"experts\.fc2\.weight$"): "experts.down_proj",
277
+ },
278
+ # Quantization configs also use the expert module names.
279
orig_to_new_substr={
280
"experts.fc1": "experts.gate_up_proj",
281
"experts.fc2": "experts.down_proj",
277
- }
282
+ },
283
)
284
285
def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):