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 12, 16:00 UTC
[Bugfix][Rocm] Fix shared expert weight loading failure in DeepSeek-MTP (#27563)
Signed-off-by: zhyajie <yajizhan@amd.com> Co-authored-by: zhyajie <yajizhan@amd.com>
杰兮 committed
Nov 24, 2025 at 18:16 UTC
8005e606bf280b7b6002f57e95ae3210ddc6f041
2 files changed
+124
-49
vllm/model_executor/models/deepseek_mtp.py
+117
-42
@@ -1,15 +1,17 @@
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
-from collections.abc import Iterable
3
+import typing
4
+from collections.abc import Callable, Iterable
5
6
import torch
7
import torch.nn as nn
8
from transformers import PretrainedConfig
9
10
+from vllm._aiter_ops import rocm_aiter_ops
11
from vllm.compilation.decorators import support_torch_compile
12
from vllm.config import VllmConfig
13
from vllm.logger import init_logger
12
-from vllm.model_executor.layers.fused_moe import FusedMoE
14
+from vllm.model_executor.layers.fused_moe import SharedFusedMoE
15
from vllm.model_executor.layers.layernorm import RMSNorm
16
from vllm.model_executor.layers.logits_processor import LogitsProcessor
17
from vllm.model_executor.layers.quantization import QuantizationConfig
@@ -231,6 +233,9 @@ class DeepSeekMTP(nn.Module, SupportsPP, DeepseekV2MixtureOfExperts):
233
return self.model.compute_logits(hidden_states, spec_step_idx)
234
235
def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
236
+ rocm_aiter_moe_shared_expert_enabled = (
237
+ rocm_aiter_ops.is_fusion_moe_shared_experts_enabled()
238
+ )
239
stacked_params_mapping = [
240
("gate_up_proj", "gate_proj", 0),
241
("gate_up_proj", "up_proj", 1),
@@ -238,11 +243,16 @@ class DeepSeekMTP(nn.Module, SupportsPP, DeepseekV2MixtureOfExperts):
243
("fused_qkv_a_proj", "kv_a_proj_with_mqa", 1),
244
]
245
241
- expert_params_mapping = FusedMoE.make_expert_params_mapping(
246
+ expert_params_mapping = SharedFusedMoE.make_expert_params_mapping(
247
ckpt_gate_proj_name="gate_proj",
248
ckpt_down_proj_name="down_proj",
249
ckpt_up_proj_name="up_proj",
245
- num_experts=self.config.n_routed_experts,
250
+ num_experts=self.config.n_routed_experts
251
+ + (
252
+ self.config.n_shared_experts
253
+ if rocm_aiter_moe_shared_expert_enabled
254
+ else 0
255
+ ),
256
)
257
258
params_dict = dict(self.named_parameters())
@@ -253,6 +263,9 @@ class DeepSeekMTP(nn.Module, SupportsPP, DeepseekV2MixtureOfExperts):
263
spec_layer = get_spec_layer_idx_from_weight_name(self.config, name)
264
if spec_layer is None:
265
continue
266
+ is_fusion_moe_shared_experts_layer = (
267
+ rocm_aiter_moe_shared_expert_enabled and ("mlp.shared_experts" in name)
268
+ )
269
name = self._rewrite_spec_layer_name(spec_layer, name)
270
for param_name, weight_name, shard_id in stacked_params_mapping:
271
# Skip non-stacked layers and experts (experts handled below).
@@ -266,6 +279,8 @@ class DeepSeekMTP(nn.Module, SupportsPP, DeepseekV2MixtureOfExperts):
279
# for mlp.experts[0].gate_gate_up_proj, which breaks load.
280
if ("mlp.experts." in name) and name not in params_dict:
281
continue
282
+ if is_fusion_moe_shared_experts_layer:
283
+ continue
284
name_mapped = name.replace(weight_name, param_name)
285
286
# QKV fusion is optional, fall back to normal
@@ -286,45 +301,105 @@ class DeepSeekMTP(nn.Module, SupportsPP, DeepseekV2MixtureOfExperts):
301
weight_loader(param, loaded_weight, shard_id)
302
break
303
else:
289
- for mapping in expert_params_mapping:
290
- param_name, weight_name, expert_id, shard_id = mapping
291
- if weight_name not in name:
292
- continue
293
- name = name.replace(weight_name, param_name)
294
-
295
- param = params_dict[name]
296
- weight_loader = param.weight_loader
297
- weight_loader(
298
- param,
299
- loaded_weight,
300
- name,
301
- shard_id=shard_id,
302
- expert_id=expert_id,
303
- )
304
- break
305
- else:
306
- # Skip loading extra bias for GPTQ models.
307
- if name.endswith(".bias") and name not in params_dict:
308
- continue
309
-
310
- name = maybe_remap_kv_scale_name(name, params_dict)
311
- if name is None:
312
- continue
313
-
314
- # According to DeepSeek-V3 Technical Report, MTP modules
315
- # shares embedding layer. We only load the first weights.
316
- if (
317
- spec_layer != self.model.mtp_start_layer_idx
318
- and ".layers" not in name
319
- ):
320
- continue
321
-
322
- param = params_dict[name]
323
- weight_loader = getattr(
324
- param, "weight_loader", default_weight_loader
304
+ # Special handling: when AITER fusion_shared_experts is enabled,
305
+ # checkpoints may provide a single widened shared_experts tensor
306
+ # without explicit expert indices
307
+ # (e.g. ...mlp.shared_experts.gate_proj.weight).
308
+ # For models with multiple shared experts, split that tensor
309
+ # evenly into per-shared-expert slices and load them into
310
+ # appended expert slots mlp.experts.{n_routed_experts + j}.*
311
+ # accordingly.
312
+ num_chunks = 1
313
+ if is_fusion_moe_shared_experts_layer:
314
+ num_chunks = getattr(self.config, "n_shared_experts", 1) or 1
315
+ # Determine split axis based on op type
316
+ # gate/up: ColumnParallel → split along dim 0
317
+ # down: RowParallel → split along dim 1
318
+ split_dim = 1 if "down_proj.weight" in name else 0
319
+ total = loaded_weight.shape[split_dim]
320
+ assert total % num_chunks == 0, (
321
+ f"Shared expert weight dim {total} "
322
+ f"not divisible by num_chunks {num_chunks}"
323
)
326
- weight_loader(param, loaded_weight)
327
- loaded_params.add(name)
324
+ chunk_size = total // num_chunks
325
+
326
+ for j in range(num_chunks):
327
+ chunk_name = name
328
+ weight_to_load = loaded_weight
329
+
330
+ if is_fusion_moe_shared_experts_layer:
331
+ if split_dim == 0:
332
+ weight_to_load = loaded_weight[
333
+ j * chunk_size : (j + 1) * chunk_size, :
334
+ ]
335
+ else:
336
+ weight_to_load = loaded_weight[
337
+ :, j * chunk_size : (j + 1) * chunk_size
338
+ ]
339
+ # Synthesize an expert-style name so expert mapping
340
+ # can route it
341
+ chunk_name = name.replace(
342
+ "mlp.shared_experts",
343
+ f"mlp.experts.{self.config.n_routed_experts + j}",
344
+ )
345
+
346
+ # Use expert_params_mapping to locate the destination
347
+ # param and delegate to its expert-aware weight_loader
348
+ # with expert_id.
349
+ for mapping in expert_params_mapping:
350
+ param_name, weight_name, expert_id, shard_id = mapping
351
+ if weight_name not in chunk_name:
352
+ continue
353
+
354
+ # Do not modify `name` since the loop may continue here
355
+ # Instead, create a new variable
356
+ name_mapped = chunk_name.replace(weight_name, param_name)
357
+
358
+ param = params_dict[name_mapped]
359
+ # We should ask the weight loader to return success or
360
+ # not here since otherwise we may skip experts with
361
+ # other available replicas.
362
+ weight_loader = typing.cast(
363
+ Callable[..., bool], param.weight_loader
364
+ )
365
+ success = weight_loader(
366
+ param,
367
+ weight_to_load,
368
+ name_mapped,
369
+ shard_id=shard_id,
370
+ expert_id=expert_id,
371
+ return_success=True,
372
+ )
373
+ if success:
374
+ if not is_fusion_moe_shared_experts_layer:
375
+ name = name_mapped
376
+ else:
377
+ loaded_params.add(name_mapped)
378
+ break
379
+ else:
380
+ # Skip loading extra bias for GPTQ models.
381
+ if name.endswith(".bias") and name not in params_dict:
382
+ continue
383
+
384
+ name = maybe_remap_kv_scale_name(name, params_dict)
385
+ if name is None:
386
+ continue
387
+
388
+ # According to DeepSeek-V3 Technical Report, MTP modules
389
+ # shares embedding layer. We only load the first weights.
390
+ if (
391
+ spec_layer != self.model.mtp_start_layer_idx
392
+ and ".layers" not in name
393
+ ):
394
+ continue
395
+
396
+ param = params_dict[name]
397
+ weight_loader = getattr(
398
+ param, "weight_loader", default_weight_loader
399
+ )
400
+ weight_loader(param, loaded_weight)
401
+ if not is_fusion_moe_shared_experts_layer:
402
+ loaded_params.add(name)
403
return loaded_params
404
405
def _rewrite_spec_layer_name(self, spec_layer: int, name: str) -> str:
vllm/model_executor/models/deepseek_v2.py
+7
-7
@@ -1479,8 +1479,8 @@ class DeepseekV2ForCausalLM(
1479
if spec_layer is not None:
1480
continue # skip spec decode layers for main model
1481
1482
- is_fuse_shared_experts_layer = rocm_aiter_moe_shared_expert_enabled and (
1483
- "mlp.shared_experts" in name
1482
+ is_fusion_moe_shared_experts_layer = (
1483
+ rocm_aiter_moe_shared_expert_enabled and ("mlp.shared_experts" in name)
1484
)
1485
1486
for param_name, weight_name, shard_id in stacked_params_mapping:
@@ -1495,7 +1495,7 @@ class DeepseekV2ForCausalLM(
1495
# for mlp.experts[0].gate_gate_up_proj, which breaks load.
1496
if ("mlp.experts." in name) and name not in params_dict:
1497
continue
1498
- if is_fuse_shared_experts_layer:
1498
+ if is_fusion_moe_shared_experts_layer:
1499
continue
1500
name_mapped = name.replace(weight_name, param_name)
1501
@@ -1531,7 +1531,7 @@ class DeepseekV2ForCausalLM(
1531
# appended expert slots mlp.experts.{n_routed_experts + j}.*
1532
# accordingly.
1533
num_chunks = 1
1534
- if is_fuse_shared_experts_layer:
1534
+ if is_fusion_moe_shared_experts_layer:
1535
num_chunks = getattr(self.config, "n_shared_experts", 1) or 1
1536
# Determine split axis based on op type
1537
# gate/up: ColumnParallel → split along dim 0
@@ -1548,7 +1548,7 @@ class DeepseekV2ForCausalLM(
1548
chunk_name = name
1549
weight_to_load = loaded_weight
1550
1551
- if is_fuse_shared_experts_layer:
1551
+ if is_fusion_moe_shared_experts_layer:
1552
if split_dim == 0:
1553
weight_to_load = loaded_weight[
1554
j * chunk_size : (j + 1) * chunk_size, :
@@ -1599,7 +1599,7 @@ class DeepseekV2ForCausalLM(
1599
return_success=True,
1600
)
1601
if success:
1602
- if not is_fuse_shared_experts_layer:
1602
+ if not is_fusion_moe_shared_experts_layer:
1603
name = name_mapped
1604
else:
1605
loaded_params.add(name_mapped)
@@ -1628,7 +1628,7 @@ class DeepseekV2ForCausalLM(
1628
param, "weight_loader", default_weight_loader
1629
)
1630
weight_loader(param, loaded_weight)
1631
- if not is_fuse_shared_experts_layer:
1631
+ if not is_fusion_moe_shared_experts_layer:
1632
loaded_params.add(name)
1633
1634
return loaded_params