[CI Bug] Fully solve accuracy issue for DSv3.2 + MTP + Sequence Parallel (#48036)
Signed-off-by: yewentao256 <zhyanwentao@126.com>
Wentao Ye committed
Jul 14, 2026 at 10:00 UTC
1ff9429655f01c804a7b17412144323e804f1324
3 files changed
+34
-3
vllm/config/parallel.py
-1
@@ -652,7 +652,6 @@ class ParallelConfig:
652
)
653
and self.enable_expert_parallel
654
and self.tensor_parallel_size > 1
655
- and self.data_parallel_size > 1
655
)
656
657
@property
vllm/model_executor/models/deepseek_mtp.py
+24
@@ -10,6 +10,7 @@ from transformers import PretrainedConfig
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.distributed import tensor_model_parallel_all_gather
14
from vllm.logger import init_logger
15
from vllm.model_executor.layers.fused_moe import (
16
fused_moe_make_expert_params_mapping,
@@ -40,6 +41,24 @@ from .utils import get_pp_missing_layer_names, maybe_prefix
41
logger = init_logger(__name__)
42
43
44
+def _restore_full_token_layout_if_needed(
45
+ hidden_states: torch.Tensor,
46
+ residual: torch.Tensor,
47
+ num_tokens: int,
48
+) -> tuple[torch.Tensor, torch.Tensor]:
49
+ """Restore full token rows for the MTP proposer after SP MoE layers."""
50
+ if hidden_states.shape[0] == num_tokens:
51
+ return hidden_states, residual
52
+
53
+ combined_states = torch.cat([hidden_states, residual], dim=-1)
54
+ combined_states = tensor_model_parallel_all_gather(combined_states, 0)
55
+ combined_states = combined_states[:num_tokens]
56
+ hidden_states, residual = combined_states.split(
57
+ [hidden_states.shape[-1], residual.shape[-1]], dim=-1
58
+ )
59
+ return hidden_states, residual
60
+
61
+
62
class SharedHead(nn.Module):
63
def __init__(
64
self,
@@ -119,6 +138,11 @@ class DeepSeekMultiTokenPredictorLayer(nn.Module):
138
hidden_states=hidden_states,
139
residual=None,
140
)
141
+ hidden_states, residual = _restore_full_token_layout_if_needed(
142
+ hidden_states,
143
+ residual,
144
+ positions.shape[0],
145
+ )
146
hidden_states = residual + hidden_states # pre-final-norm (logits hidden)
147
# Recycle the post-final-norm hidden into the next draft step.
148
# compute_logits applies shared_head (== final norm) to the pre-norm
vllm/models/deepseek_v32/nvidia/mtp.py
+10
-2
@@ -21,7 +21,10 @@ from vllm.model_executor.model_loader.weight_utils import (
21
default_weight_loader,
22
maybe_remap_kv_scale_name,
23
)
24
-from vllm.model_executor.models.deepseek_mtp import SharedHead
24
+from vllm.model_executor.models.deepseek_mtp import (
25
+ SharedHead,
26
+ _restore_full_token_layout_if_needed,
27
+)
28
from vllm.model_executor.models.deepseek_v2 import (
29
DeepseekV2MixtureOfExperts,
30
DeepseekV2MoE,
@@ -89,7 +92,12 @@ class DeepseekV32MultiTokenPredictorLayer(nn.Module):
92
hidden_states, residual = self.mtp_block(
93
positions=positions, hidden_states=hidden_states, residual=None
94
)
92
- # mtp_block's MoE output is left un-reduced (reduce_results=False); the
95
+ hidden_states, residual = _restore_full_token_layout_if_needed(
96
+ hidden_states,
97
+ residual,
98
+ positions.shape[0],
99
+ )
100
+ # mtp_block's MoE output is left un-reduced (skip_final_all_reduce); the
101
# main model fuses that all-reduce into the next norm, but here the
102
# recycle hidden is consumed directly, so reduce it now.
103
hidden_states = tensor_model_parallel_all_reduce(hidden_states)