A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 5, 18:30 UTC
[Bugfix] Fix FusedMoE weight loading with padded hidden dimensions (#37010)
Signed-off-by: SandishKumarHN <sandish@fb.com>
SandishKumarHN committed
Mar 31, 2026 at 09:22 UTC
3896e021a02fb84e55c998f9fc3a4f0e6a9a9e5b
2 files changed
+374
-14
tests/kernels/moe/test_moe_weight_loading_padded.py
new
+292
@@ -0,0 +1,292 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+"""Tests for FusedMoE weight loading with padded hidden dimensions.
4
+
5
+When using DeepEP backends or NIXL EP with models like nemotron_h,
6
+hidden_size may be rounded up (e.g., 2688 -> 3072) for backend requirements.
7
+Weight parameters are created with the padded size, but checkpoint weights
8
+have the original unpadded size. These tests verify that weight loading
9
+correctly handles this mismatch.
10
+"""
11
+
12
+import pytest
13
+import torch
14
+
15
+from vllm.model_executor.layers.fused_moe.layer import FusedMoE
16
+
17
+
18
+class TestGetHiddenDim:
19
+ """Unit tests for _get_hidden_dim."""
20
+
21
+ def test_2d_non_transposed_w2(self):
22
+ # w2: shard_dim=1 (intermediate), hidden=0
23
+ assert FusedMoE._get_hidden_dim(shard_dim=1, ndim=2) == 0
24
+
25
+ def test_2d_non_transposed_w13(self):
26
+ # w1/w3: shard_dim=0 (intermediate), hidden=1
27
+ assert FusedMoE._get_hidden_dim(shard_dim=0, ndim=2) == 1
28
+
29
+ def test_2d_transposed_w2(self):
30
+ # transposed w2: shard_dim=0, hidden=1
31
+ assert FusedMoE._get_hidden_dim(shard_dim=0, ndim=2) == 1
32
+
33
+ def test_2d_transposed_w13(self):
34
+ # transposed w1/w3: shard_dim=1, hidden=0
35
+ assert FusedMoE._get_hidden_dim(shard_dim=1, ndim=2) == 0
36
+
37
+ def test_3d_non_transposed_w2(self):
38
+ # 3D w2: shard_dim=2, hidden=1
39
+ assert FusedMoE._get_hidden_dim(shard_dim=2, ndim=3) == 1
40
+
41
+ def test_3d_non_transposed_w13(self):
42
+ # 3D w1/w3: shard_dim=1, hidden=2
43
+ assert FusedMoE._get_hidden_dim(shard_dim=1, ndim=3) == 2
44
+
45
+ def test_3d_transposed_w2(self):
46
+ # transposed 3D w2: shard_dim=1, hidden=2
47
+ assert FusedMoE._get_hidden_dim(shard_dim=1, ndim=3) == 2
48
+
49
+ def test_3d_transposed_w13(self):
50
+ # transposed 3D w1/w3: shard_dim=2, hidden=1
51
+ assert FusedMoE._get_hidden_dim(shard_dim=2, ndim=3) == 1
52
+
53
+ def test_1d_returns_zero(self):
54
+ # 1D per-channel scales: always returns 0
55
+ assert FusedMoE._get_hidden_dim(shard_dim=0, ndim=1) == 0
56
+ assert FusedMoE._get_hidden_dim(shard_dim=1, ndim=1) == 0
57
+
58
+ def test_invalid_shard_dim_raises(self):
59
+ # shard_dim outside the data dimensions should raise
60
+ with pytest.raises(ValueError, match="not a valid data dimension"):
61
+ FusedMoE._get_hidden_dim(shard_dim=0, ndim=3)
62
+
63
+
64
+class TestNarrowExpertDataForPadding:
65
+ """Unit tests for _narrow_expert_data_for_padding."""
66
+
67
+ def test_no_narrowing_when_shapes_match(self):
68
+ expert_data = torch.zeros(1024, 1024)
69
+ loaded_weight = torch.randn(1024, 1024)
70
+ result = FusedMoE._narrow_expert_data_for_padding(
71
+ expert_data, loaded_weight, hidden_dim=0
72
+ )
73
+ assert result.shape == loaded_weight.shape
74
+ assert result.data_ptr() == expert_data.data_ptr()
75
+
76
+ def test_narrow_w2_hidden_dim(self):
77
+ # w2: (hidden_size, intermediate_size) - hidden_size padded at dim 0
78
+ expert_data = torch.zeros(3072, 1024)
79
+ loaded_weight = torch.randn(2688, 1024)
80
+ result = FusedMoE._narrow_expert_data_for_padding(
81
+ expert_data, loaded_weight, hidden_dim=0
82
+ )
83
+ assert result.shape == (2688, 1024)
84
+
85
+ def test_narrow_w13_hidden_dim(self):
86
+ # w1/w3: (intermediate_size, hidden_size) - hidden_size padded at dim 1
87
+ expert_data = torch.zeros(2048, 3072)
88
+ loaded_weight = torch.randn(2048, 2688)
89
+ result = FusedMoE._narrow_expert_data_for_padding(
90
+ expert_data, loaded_weight, hidden_dim=1
91
+ )
92
+ assert result.shape == (2048, 2688)
93
+
94
+ def test_narrow_transposed_w2(self):
95
+ # transposed w2: (intermediate_size, hidden_size) - hidden at dim 1
96
+ expert_data = torch.zeros(1024, 3072)
97
+ loaded_weight = torch.randn(1024, 2688)
98
+ hidden_dim = FusedMoE._get_hidden_dim(shard_dim=0, ndim=2)
99
+ result = FusedMoE._narrow_expert_data_for_padding(
100
+ expert_data, loaded_weight, hidden_dim=hidden_dim
101
+ )
102
+ assert result.shape == (1024, 2688)
103
+
104
+ def test_narrow_3d_full_load(self):
105
+ # 3D tensor for full_load path: w2 (num_experts, hidden_size, intermediate)
106
+ expert_data = torch.zeros(8, 3072, 1024)
107
+ loaded_weight = torch.randn(8, 2688, 1024)
108
+ result = FusedMoE._narrow_expert_data_for_padding(
109
+ expert_data, loaded_weight, hidden_dim=1
110
+ )
111
+ assert result.shape == (8, 2688, 1024)
112
+
113
+ def test_narrow_1d_scale(self):
114
+ # 1D scale tensor: per-channel w2 scale (hidden_size,)
115
+ expert_data = torch.zeros(3072)
116
+ loaded_weight = torch.randn(2688)
117
+ result = FusedMoE._narrow_expert_data_for_padding(
118
+ expert_data, loaded_weight, hidden_dim=0
119
+ )
120
+ assert result.shape == (2688,)
121
+
122
+ def test_scalar_weight_no_op(self):
123
+ # 0-dim tensor should be a no-op
124
+ expert_data = torch.zeros(3072)
125
+ loaded_weight = torch.tensor(1.0)
126
+ result = FusedMoE._narrow_expert_data_for_padding(
127
+ expert_data, loaded_weight, hidden_dim=0
128
+ )
129
+ # ndim == 0, so no narrowing
130
+ assert result.shape == (3072,)
131
+
132
+ def test_no_narrowing_when_loaded_weight_larger(self):
133
+ # Guard: don't narrow if loaded_weight is larger than expert_data
134
+ expert_data = torch.zeros(2688, 1024)
135
+ loaded_weight = torch.randn(3072, 1024)
136
+ result = FusedMoE._narrow_expert_data_for_padding(
137
+ expert_data, loaded_weight, hidden_dim=0
138
+ )
139
+ assert result.shape == (2688, 1024)
140
+ assert result.data_ptr() == expert_data.data_ptr()
141
+
142
+ def test_negative_hidden_dim_is_noop(self):
143
+ # Negative hidden_dim should be a safe no-op (0 <= check)
144
+ expert_data = torch.zeros(3072, 1024)
145
+ loaded_weight = torch.randn(2688, 1024)
146
+ result = FusedMoE._narrow_expert_data_for_padding(
147
+ expert_data, loaded_weight, hidden_dim=-1
148
+ )
149
+ # -1 fails the 0 <= check, so no narrowing
150
+ assert result.shape == (3072, 1024)
151
+ assert result.data_ptr() == expert_data.data_ptr()
152
+
153
+ def test_only_narrows_hidden_dim(self):
154
+ # Verify that only the specified hidden_dim is narrowed,
155
+ # even when other dimensions also differ
156
+ expert_data = torch.zeros(3072, 2048)
157
+ loaded_weight = torch.randn(2688, 1024)
158
+ result = FusedMoE._narrow_expert_data_for_padding(
159
+ expert_data, loaded_weight, hidden_dim=0
160
+ )
161
+ # Only dim 0 (hidden) should be narrowed; dim 1 stays at 2048
162
+ assert result.shape == (2688, 2048)
163
+
164
+ def test_narrowed_data_shares_storage(self):
165
+ # Verify narrowing returns a view (writes go to original tensor)
166
+ expert_data = torch.zeros(3072, 1024)
167
+ loaded_weight = torch.randn(2688, 1024)
168
+ result = FusedMoE._narrow_expert_data_for_padding(
169
+ expert_data, loaded_weight, hidden_dim=0
170
+ )
171
+ result.copy_(loaded_weight)
172
+ # The first 2688 rows of expert_data should now have loaded_weight
173
+ assert torch.equal(expert_data[:2688, :], loaded_weight)
174
+ # Padded region should remain zero
175
+ assert torch.equal(expert_data[2688:, :], torch.zeros(3072 - 2688, 1024))
176
+
177
+
178
+class TestWeightLoadingWithPaddedHiddenSize:
179
+ """Integration-style tests that simulate padded weight loading."""
180
+
181
+ def test_load_w2_with_padding(self):
182
+ """Simulate loading w2 weights when hidden_size is padded."""
183
+ padded_hidden = 3072
184
+ original_hidden = 2688
185
+ intermediate = 1024
186
+
187
+ expert_data_full = torch.zeros(padded_hidden, intermediate)
188
+ loaded_weight = torch.randn(original_hidden, intermediate)
189
+
190
+ # w2 non-transposed: shard_dim=1, hidden_dim=0
191
+ hidden_dim = FusedMoE._get_hidden_dim(shard_dim=1, ndim=2)
192
+ expert_data = FusedMoE._narrow_expert_data_for_padding(
193
+ expert_data_full, loaded_weight, hidden_dim=hidden_dim
194
+ )
195
+ expert_data.copy_(loaded_weight)
196
+
197
+ assert torch.equal(expert_data_full[:original_hidden, :], loaded_weight)
198
+ assert torch.equal(
199
+ expert_data_full[original_hidden:, :],
200
+ torch.zeros(padded_hidden - original_hidden, intermediate),
201
+ )
202
+
203
+ def test_load_w13_with_padding(self):
204
+ """Simulate loading w1/w3 weights when hidden_size is padded."""
205
+ padded_hidden = 3072
206
+ original_hidden = 2688
207
+ intermediate = 1024
208
+
209
+ # w1/w3: (intermediate_size, hidden_size)
210
+ expert_data_full = torch.zeros(intermediate, padded_hidden)
211
+ loaded_weight = torch.randn(intermediate, original_hidden)
212
+
213
+ # w1 non-transposed: shard_dim=0, hidden_dim=1
214
+ hidden_dim = FusedMoE._get_hidden_dim(shard_dim=0, ndim=2)
215
+ expert_data = FusedMoE._narrow_expert_data_for_padding(
216
+ expert_data_full, loaded_weight, hidden_dim=hidden_dim
217
+ )
218
+ expert_data.copy_(loaded_weight)
219
+
220
+ assert torch.equal(expert_data_full[:, :original_hidden], loaded_weight)
221
+ assert torch.equal(
222
+ expert_data_full[:, original_hidden:],
223
+ torch.zeros(intermediate, padded_hidden - original_hidden),
224
+ )
225
+
226
+ def test_load_transposed_w2_with_padding(self):
227
+ """Simulate loading transposed w2 (GPTQ) with padded hidden_size."""
228
+ padded_hidden = 3072
229
+ original_hidden = 2688
230
+ intermediate = 1024
231
+
232
+ # transposed w2: (intermediate_size, hidden_size), shard_dim=0
233
+ expert_data_full = torch.zeros(intermediate, padded_hidden)
234
+ loaded_weight = torch.randn(intermediate, original_hidden)
235
+
236
+ hidden_dim = FusedMoE._get_hidden_dim(shard_dim=0, ndim=2)
237
+ expert_data = FusedMoE._narrow_expert_data_for_padding(
238
+ expert_data_full, loaded_weight, hidden_dim=hidden_dim
239
+ )
240
+ expert_data.copy_(loaded_weight)
241
+
242
+ assert torch.equal(expert_data_full[:, :original_hidden], loaded_weight)
243
+
244
+ def test_no_padding_is_noop(self):
245
+ """Verify that when sizes match, behavior is unchanged."""
246
+ hidden = 2048
247
+ intermediate = 1024
248
+
249
+ expert_data_full = torch.zeros(hidden, intermediate)
250
+ loaded_weight = torch.randn(hidden, intermediate)
251
+
252
+ hidden_dim = FusedMoE._get_hidden_dim(shard_dim=1, ndim=2)
253
+ expert_data = FusedMoE._narrow_expert_data_for_padding(
254
+ expert_data_full, loaded_weight, hidden_dim=hidden_dim
255
+ )
256
+ expert_data.copy_(loaded_weight)
257
+
258
+ assert torch.equal(expert_data_full, loaded_weight)
259
+
260
+ def test_bnb_shape_mismatch_raises(self):
261
+ """BnB + padded hidden_size should raise via weight_loader."""
262
+ from unittest.mock import MagicMock
263
+
264
+ num_experts = 1
265
+ padded_packed = 3072 # padded packed size
266
+ original_packed = 2688 # original packed size
267
+
268
+ # Build a param that looks like a BnB 4-bit MoE weight.
269
+ param_data = torch.zeros(num_experts, padded_packed, 1, dtype=torch.uint8)
270
+ param = torch.nn.Parameter(param_data, requires_grad=False)
271
+ param.use_bitsandbytes_4bit = True
272
+
273
+ loaded_weight = torch.randint(0, 255, (original_packed, 1), dtype=torch.uint8)
274
+
275
+ # Minimal FusedMoE mock so weight_loader reaches the BnB path.
276
+ moe = MagicMock(spec=FusedMoE)
277
+ moe.quant_config = None
278
+ moe.quant_method = MagicMock()
279
+ moe.quant_method.__class__.__name__ = "BitsAndBytesMethod"
280
+ moe._expert_map = None
281
+ moe.tp_rank = 0
282
+
283
+ # Call the real weight_loader (unbound) with our mock as self.
284
+ with pytest.raises(ValueError, match="BitsAndBytes"):
285
+ FusedMoE.weight_loader(
286
+ moe,
287
+ param,
288
+ loaded_weight,
289
+ weight_name="w2",
290
+ shard_id="w2",
291
+ expert_id=0,
292
+ )
vllm/model_executor/layers/fused_moe/layer.py
+82
-14
@@ -860,6 +860,10 @@ class FusedMoE(CustomOp):
860
):
861
# for per channel weight quantization
862
if shard_id == "w2":
863
+ hidden_dim = self._get_hidden_dim(shard_dim, expert_data.ndim)
864
+ expert_data = self._narrow_expert_data_for_padding(
865
+ expert_data, loaded_weight, hidden_dim=hidden_dim
866
+ )
867
expert_data.copy_(loaded_weight)
868
elif shard_id in ("w1", "w3"):
869
self._load_w13(
@@ -870,6 +874,59 @@ class FusedMoE(CustomOp):
874
tp_rank=tp_rank,
875
)
876
877
+ @staticmethod
878
+ def _get_hidden_dim(shard_dim: int, ndim: int) -> int:
879
+ """Compute the hidden dimension index from the shard (intermediate)
880
+ dimension and tensor rank.
881
+
882
+ For 2D weight tensors the two data dims are (0, 1). For 3D tensors
883
+ with an expert dimension at dim 0, they are (1, 2). ``shard_dim``
884
+ occupies one of these; the hidden dimension is the other.
885
+ For 1D tensors (e.g. per-channel scales) returns 0.
886
+ """
887
+ if ndim < 2:
888
+ return 0
889
+ dim_a = ndim - 2
890
+ dim_b = ndim - 1
891
+ if shard_dim == dim_a:
892
+ return dim_b
893
+ if shard_dim == dim_b:
894
+ return dim_a
895
+ raise ValueError(
896
+ f"shard_dim={shard_dim} is not a valid data dimension "
897
+ f"for a {ndim}D tensor (expected {dim_a} or {dim_b})"
898
+ )
899
+
900
+ @staticmethod
901
+ def _narrow_expert_data_for_padding(
902
+ expert_data: torch.Tensor,
903
+ loaded_weight: torch.Tensor,
904
+ hidden_dim: int,
905
+ ) -> torch.Tensor:
906
+ """Narrow expert_data hidden dim to match loaded_weight for padded
907
+ hidden_size.
908
+
909
+ When backends (e.g., DeepEP) round up hidden_size, weight parameters
910
+ are larger than checkpoint weights. Narrow the padded hidden dimension
911
+ before copying.
912
+
913
+ Args:
914
+ expert_data: The (possibly padded) parameter tensor to narrow.
915
+ loaded_weight: The checkpoint weight tensor with original size.
916
+ hidden_dim: The dimension index corresponding to hidden_size.
917
+ Must be non-negative.
918
+ """
919
+ if (
920
+ loaded_weight.ndim > 0
921
+ and 0 <= hidden_dim < expert_data.ndim
922
+ and hidden_dim < loaded_weight.ndim
923
+ and expert_data.shape[hidden_dim] > loaded_weight.shape[hidden_dim]
924
+ ):
925
+ expert_data = expert_data.narrow(
926
+ hidden_dim, 0, loaded_weight.shape[hidden_dim]
927
+ )
928
+ return expert_data
929
+
930
def _load_w13(
931
self,
932
expert_data: torch.Tensor,
@@ -907,13 +964,10 @@ class FusedMoE(CustomOp):
964
else:
965
assert shard_id == "w3"
966
expert_data = expert_data.narrow(shard_dim, shard_size, shard_size)
910
-
911
- # Handle padding: if loaded_weight is smaller than expert_data (can happen
912
- # on last TP shard with padding), copy to top-left corner
913
- if expert_data.shape != loaded_weight.shape:
914
- expert_data = expert_data[
915
- : loaded_weight.shape[0], : loaded_weight.shape[1]
916
- ]
967
+ hidden_dim = self._get_hidden_dim(shard_dim, expert_data.ndim)
968
+ expert_data = self._narrow_expert_data_for_padding(
969
+ expert_data, loaded_weight, hidden_dim=hidden_dim
970
+ )
971
expert_data.copy_(loaded_weight)
972
973
def _load_w2(
@@ -943,12 +997,10 @@ class FusedMoE(CustomOp):
997
narrow_size = min(shard_size, available)
998
loaded_weight = loaded_weight.narrow(shard_dim, start_offset, narrow_size)
999
# w2, down_proj: Load into only logical weight of w2.
946
- # Handle padding: if loaded_weight is smaller than expert_data (can happen
947
- # on last TP shard with padding), copy to top-left corner
948
- if expert_data.shape != loaded_weight.shape:
949
- expert_data = expert_data[
950
- : loaded_weight.shape[0], : loaded_weight.shape[1]
951
- ]
1000
+ hidden_dim = self._get_hidden_dim(shard_dim, expert_data.ndim)
1001
+ expert_data = self._narrow_expert_data_for_padding(
1002
+ expert_data, loaded_weight, hidden_dim=hidden_dim
1003
+ )
1004
expert_data.copy_(loaded_weight)
1005
1006
def _load_single_value(
@@ -1095,9 +1147,25 @@ class FusedMoE(CustomOp):
1147
1148
expert_data = param.data[expert_id]
1149
if shard_id == "w2":
1150
+ # BnB params are stored as flat packed tensors (e.g.
1151
+ # (packed_size, 1)), not in the logical weight layout.
1152
+ # Narrowing packed data for hidden-dim padding is not
1153
+ # meaningful, so require an exact shape match.
1154
+ if expert_data.shape != loaded_weight.shape:
1155
+ raise ValueError(
1156
+ "BitsAndBytes quantization with padded hidden_size "
1157
+ "(e.g., from DeepEP) is not supported. "
1158
+ f"Parameter shape {tuple(expert_data.shape)} != "
1159
+ f"checkpoint shape {tuple(loaded_weight.shape)}"
1160
+ )
1161
expert_data.copy_(loaded_weight)
1162
elif shard_id in ("w1", "w3"):
1100
- # BNB inflight quantization has already sharded the weights
1163
+ # BnB stores weights as flat packed tensors. _load_w13 is
1164
+ # still used to split the w1/w3 portions along shard_dim.
1165
+ # _narrow_expert_data_for_padding will be a no-op since
1166
+ # packed sizes should already match; if DeepEP padding
1167
+ # causes a mismatch the copy_() will fail with a clear
1168
+ # shape error.
1169
full_load = True
1170
self._load_w13(
1171
shard_id=shard_id,