Skip to content

GitLab

Unverified Commit be6931ee authored by Pleaplusone's avatar Pleaplusone Committed by GitHub
Browse files

[ROCm][Bugfix] Fix ptpc scale load issue for fused shared expert path in deepseek mtp (#33018) 


Signed-off-by: default avatarganyi <ygan@amd.com>
parent 9ef3b718
Hide whitespace changes
Inline Side-by-side
Showing with 11 additions and 8 deletions
vllm/model_executor/models/deepseek_mtp.py
View file @ be6931ee
...@@ -316,7 +316,11 @@ class DeepSeekMTP(nn.Module, DeepseekV2MixtureOfExperts): ...@@ -316,7 +316,11 @@ class DeepSeekMTP(nn.Module, DeepseekV2MixtureOfExperts):
# Determine split axis based on op type # Determine split axis based on op type
# gate/up: ColumnParallel → split along dim 0 # gate/up: ColumnParallel → split along dim 0
# down: RowParallel → split along dim 1 # down: RowParallel → split along dim 1
split_dim = 1 if "down_proj.weight" in name else 0 split_dim = (
1
if ("down_proj.weight" in name and loaded_weight.ndim > 1)
else 0
)
total = loaded_weight.shape[split_dim] total = loaded_weight.shape[split_dim]
assert total % num_chunks == 0, ( assert total % num_chunks == 0, (
f"Shared expert weight dim {total} " f"Shared expert weight dim {total} "
...@@ -329,14 +333,13 @@ class DeepSeekMTP(nn.Module, DeepseekV2MixtureOfExperts): ...@@ -329,14 +333,13 @@ class DeepSeekMTP(nn.Module, DeepseekV2MixtureOfExperts):
weight_to_load = loaded_weight weight_to_load = loaded_weight
if is_fusion_moe_shared_experts_layer: if is_fusion_moe_shared_experts_layer:
if split_dim == 0: chunk_slice = slice(j * chunk_size, (j + 1) * chunk_size)
weight_to_load = loaded_weight[ if loaded_weight.ndim == 1:
j * chunk_size : (j + 1) * chunk_size, : weight_to_load = loaded_weight[chunk_slice]
] elif split_dim == 0:
weight_to_load = loaded_weight[chunk_slice, :]
else: else:
weight_to_load = loaded_weight[ weight_to_load = loaded_weight[:, chunk_slice]
:, j * chunk_size : (j + 1) * chunk_size
]
# Synthesize an expert-style name so expert mapping # Synthesize an expert-style name so expert mapping
# can route it # can route it
chunk_name = name.replace( chunk_name = name.replace(
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment