chenchenghao/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 11 Packages Projects Releases Wiki Activity

[BUGFIX] Fix dp size > 1 for qwen3 vl model (#17624)

Browse Source 
Co-authored-by: yizhang2077 <1109276519@qq.com>
This commit is contained in:
Zheng Li
2026-01-30 20:44:25 +08:00
committed by GitHub
parent c04efe030a
commit 0c5a81acb8
5 changed files with 48 additions and 19 deletions
Download Patch File Download Diff File Expand all files Collapse all files

2
python/sglang/srt/layers/attention/vision.py
View File

@@ -596,6 +596,7 @@ class VisionAttention(nn.Module):
[torch.Tensor, torch.Tensor, Any, Any], Tuple[torch.Tensor, torch.Tensor]
] = None,
use_data_parallel: bool = False,
use_dp_attention_reduce: bool = False,
aux_stream: Optional[torch.cuda.Stream] = None,
**kwargs,
):
@@ -688,6 +689,7 @@ class VisionAttention(nn.Module):
tp_size=self.tp_size,
reduce_results=False,
prefix=add_prefix("proj", prefix),
use_dp_attention_reduce=use_dp_attention_reduce,
)
self.aux_stream = aux_stream
self.ln_events = [torch.cuda.Event(), torch.cuda.Event()] if aux_stream else []

12
python/sglang/srt/layers/linear.py
View File

@@ -21,7 +21,10 @@ from sglang.srt.distributed import (
from sglang.srt.distributed.device_communicators.pynccl_allocator import (
use_symmetric_memory,
)
from sglang.srt.layers.dp_attention import is_allocation_symmetric
from sglang.srt.layers.dp_attention import (
get_attention_tp_group,
is_allocation_symmetric,
)
from sglang.srt.layers.parameter import (
BasevLLMParameter,
BlockQuantScaleParameter,
@@ -1263,6 +1266,7 @@ class RowParallelLinear(LinearBase):
tp_rank: Optional[int] = None,
tp_size: Optional[int] = None,
use_presharded_weights: bool = False,
use_dp_attention_reduce: bool = False,
):
quant_config = None if _disable_hip_linear_quant else quant_config
super().__init__(
@@ -1271,6 +1275,7 @@ class RowParallelLinear(LinearBase):
self.input_is_parallel = input_is_parallel
self.reduce_results = reduce_results
self.use_dp_attention_reduce = use_dp_attention_reduce
# Divide the weight matrix along the last dimension.
if tp_rank is None:
@@ -1419,7 +1424,10 @@ class RowParallelLinear(LinearBase):
output_parallel = self.quant_method.apply(self, input_parallel, bias=bias_)
if self.reduce_results and self.tp_size > 1 and not skip_all_reduce:
output = tensor_model_parallel_all_reduce(output_parallel)
if self.use_dp_attention_reduce:
output = get_attention_tp_group().all_reduce(output_parallel)
else:
output = tensor_model_parallel_all_reduce(output_parallel)
else:
output = output_parallel

10
python/sglang/srt/model_executor/forward_batch_info.py
View File

@@ -860,7 +860,15 @@ class ForwardBatch(ForwardBatchDeepSeekMHAMixin):
)
if self.mrope_positions is not None:
self.mrope_positions = self._pad_tensor_to_size(self.mrope_positions, bs)
self.mrope_positions = torch.cat(
[
self.mrope_positions,
self.mrope_positions.new_zeros(
3, num_tokens - self.mrope_positions.shape[1]
),
],
dim=1,
)
# TODO: check if we need to pad other tensors
if self.extend_seq_lens is not None:

27
python/sglang/srt/models/qwen3_vl.py
View File

@@ -25,14 +25,15 @@ from einops import rearrange
from transformers.activations import ACT2FN
from sglang.srt.configs.qwen3_vl import Qwen3VLConfig, Qwen3VLVisionConfig
from sglang.srt.distributed import (
get_tensor_model_parallel_rank,
get_tensor_model_parallel_world_size,
)
from sglang.srt.distributed import get_tensor_model_parallel_world_size
from sglang.srt.distributed.parallel_state import get_pp_group
from sglang.srt.environ import envs
from sglang.srt.layers.attention.vision import VisionAttention
from sglang.srt.layers.dp_attention import is_dp_attention_enabled
from sglang.srt.layers.dp_attention import (
get_attention_tp_rank,
get_attention_tp_size,
is_dp_attention_enabled,
)
from sglang.srt.layers.linear import ColumnParallelLinear, RowParallelLinear
from sglang.srt.layers.logits_processor import LogitsProcessor
from sglang.srt.layers.pooler import Pooler, PoolingType
@@ -85,10 +86,8 @@ class Qwen3_VisionMLP(nn.Module):
use_data_parallel: bool = False,
):
super().__init__()
self.tp_size = (
1 if use_data_parallel else get_tensor_model_parallel_world_size()
)
self.tp_rank = 0 if use_data_parallel else get_tensor_model_parallel_rank()
self.tp_size = 1 if use_data_parallel else get_attention_tp_size()
self.tp_rank = 0 if use_data_parallel else get_attention_tp_rank()
self.linear_fc1 = ColumnParallelLinear(
in_features,
hidden_features,
@@ -106,6 +105,7 @@ class Qwen3_VisionMLP(nn.Module):
prefix=add_prefix("linear_fc2", prefix),
tp_size=self.tp_size,
tp_rank=self.tp_rank,
use_dp_attention_reduce=is_dp_attention_enabled(),
)
self.act = ACT2FN[hidden_act]
@@ -176,6 +176,7 @@ class Qwen3_VisionBlock(nn.Module):
quant_config=quant_config,
prefix=add_prefix("attn", prefix),
use_data_parallel=use_data_parallel,
use_dp_attention_reduce=is_dp_attention_enabled(),
)
self.mlp = Qwen3_VisionMLP(
dim,
@@ -235,10 +236,8 @@ class Qwen3VLMoeVisionPatchMerger(nn.Module):
self.norm = norm_layer(
self.hidden_size if use_postshuffle_norm else context_dim
)
self.tp_size = (
1 if use_data_parallel else get_tensor_model_parallel_world_size()
)
self.tp_rank = 0 if use_data_parallel else get_tensor_model_parallel_rank()
self.tp_size = 1 if use_data_parallel else get_attention_tp_size()
self.tp_rank = 0 if use_data_parallel else get_attention_tp_rank()
self.linear_fc1 = ColumnParallelLinear(
self.hidden_size,
self.hidden_size,
@@ -257,6 +256,7 @@ class Qwen3VLMoeVisionPatchMerger(nn.Module):
prefix=add_prefix("linear_fc2", prefix),
tp_size=self.tp_size,
tp_rank=self.tp_rank,
use_dp_attention_reduce=is_dp_attention_enabled(),
)
def forward(self, x: torch.Tensor) -> torch.Tensor:
@@ -713,6 +713,7 @@ class Qwen3VLForConditionalGeneration(nn.Module):
self.config.vocab_size,
self.config.hidden_size,
quant_config=quant_config,
use_attn_tp_group=get_global_server_args().enable_dp_lm_head,
prefix=add_prefix("lm_head", prefix),
)
else:

16
python/sglang/srt/multimodal/mm_utils.py
View File

@@ -495,11 +495,19 @@ def run_dp_sharded_mrope_vision_model(
```
"""
tp_size = get_tensor_model_parallel_world_size()
from sglang.srt.layers.dp_attention import (
get_attention_tp_group,
get_attention_tp_rank,
get_attention_tp_size,
)
tp_size = get_attention_tp_size()
if tp_size == 1:
return vision_model(pixel_values, grid_thw=torch.tensor(grid_thw_list))
# GPU_0 tp_rank_local = 0
# GPU_1 tp_rank_local = 1
tp_rank_local = get_tensor_model_parallel_rank()
tp_rank_local = get_attention_tp_rank()
# patches_per_image = [1000, 100, 200, 50]
patches_per_image = [math.prod(grid_thw) for grid_thw in grid_thw_list]
@@ -611,7 +619,9 @@ def run_dp_sharded_mrope_vision_model(
image_embeds_local_padded = image_embeds_local
# Do all_gather to collect embeddings from all ranks
gathered_embeds = tensor_model_parallel_all_gather(image_embeds_local_padded, dim=0)
gathered_embeds = get_attention_tp_group().all_gather(
image_embeds_local_padded, dim=0
)
# Remove padding and reconstruct per-rank embeddings
rank_embeddings = list[torch.Tensor]()