[BugFix][VLM] keep Qwen3-VL MoE inference deepstack order (#34690)
This commit is contained in:
@@ -1152,6 +1152,11 @@ class Qwen3LLMModel(Qwen3Model):
|
||||
self.deepstack_embed_to_decoder_layer = range(
|
||||
len(config.vision_config.deepstack_visual_indexes)
|
||||
)
|
||||
# Use HF deepstack order only if rl_on_policy_target is set;
|
||||
# otherwise, retain original order for inference accuracy.
|
||||
self.use_hf_deepstack_order = (
|
||||
get_exec().deterministic.rl_on_policy_target is not None
|
||||
)
|
||||
|
||||
def get_deepstack_embeds(
|
||||
self, layer_idx: int, input_deepstack_embeds: Optional[torch.Tensor]
|
||||
@@ -1196,25 +1201,43 @@ class Qwen3LLMModel(Qwen3Model):
|
||||
hidden_states + residual if residual is not None else hidden_states
|
||||
)
|
||||
|
||||
# SGLang applies residual at the START of the next layer, not at the END like HuggingFace.
|
||||
# See: https://github.com/huggingface/transformers/blob/v5.0.0rc0/src/transformers/models/qwen3_vl/modeling_qwen3_vl.py#L549
|
||||
# To match HF behavior, deepstack must be added AFTER residual: (hidden_states + residual) + deepstack
|
||||
# The order matters because addition with different tensors is not associative in practice.
|
||||
# Deepstack for prev_layer is applied at the start of current layer via post_residual_addition.
|
||||
deepstack_embeds = self.get_deepstack_embeds(
|
||||
layer_idx - 1, input_deepstack_embeds
|
||||
)
|
||||
hidden_states, residual = layer(
|
||||
positions,
|
||||
hidden_states,
|
||||
forward_batch,
|
||||
residual,
|
||||
post_residual_addition=deepstack_embeds,
|
||||
)
|
||||
if self.use_hf_deepstack_order:
|
||||
# HF-order path (RL on-policy / FSDP). SGLang applies residual at the START of the
|
||||
# next layer, so to match HF's (hidden_states + residual) + deepstack, deepstack for
|
||||
# the previous layer is added after residual via post_residual_addition.
|
||||
deepstack_embeds = self.get_deepstack_embeds(
|
||||
layer_idx - 1, input_deepstack_embeds
|
||||
)
|
||||
hidden_states, residual = layer(
|
||||
positions,
|
||||
hidden_states,
|
||||
forward_batch,
|
||||
residual,
|
||||
post_residual_addition=deepstack_embeds,
|
||||
)
|
||||
else:
|
||||
# Inference path: add deepstack directly to hidden_states at the end of the layer
|
||||
# (original, grounding-correct order).
|
||||
hidden_states, residual = layer(
|
||||
positions,
|
||||
hidden_states,
|
||||
forward_batch,
|
||||
residual,
|
||||
)
|
||||
if (
|
||||
input_deepstack_embeds is not None
|
||||
and layer_idx in self.deepstack_embed_to_decoder_layer
|
||||
):
|
||||
sep = self.hidden_size * layer_idx
|
||||
hidden_states.add_(
|
||||
input_deepstack_embeds[:, sep : sep + self.hidden_size]
|
||||
)
|
||||
|
||||
# Handle deepstack for the last processed layer if it exists.
|
||||
last_deepstack = self.get_deepstack_embeds(
|
||||
self.end_layer - 1, input_deepstack_embeds
|
||||
# Handle deepstack for the last processed layer (HF-order path only).
|
||||
last_deepstack = (
|
||||
self.get_deepstack_embeds(self.end_layer - 1, input_deepstack_embeds)
|
||||
if self.use_hf_deepstack_order
|
||||
else None
|
||||
)
|
||||
|
||||
if not self.pp_group.is_last_rank:
|
||||
|
||||
@@ -31,6 +31,7 @@ from sglang.srt.model_executor.forward_batch_info import ForwardBatch, PPProxyTe
|
||||
from sglang.srt.model_loader.weight_utils import default_weight_loader
|
||||
from sglang.srt.models.qwen3_moe import Qwen3MoeDecoderLayer, Qwen3MoeModel
|
||||
from sglang.srt.models.qwen3_vl import Qwen3VLForConditionalGeneration
|
||||
from sglang.srt.runtime_context import get_exec
|
||||
from sglang.srt.utils.hf_transformers_utils import get_processor
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
@@ -58,6 +59,11 @@ class Qwen3MoeLLMModel(Qwen3MoeModel):
|
||||
# This approach follows the original implementation.
|
||||
# TODO: make config of type Qwen3VLMoeConfig, so that we can directly obtain deepstack_visual_indexes.
|
||||
self.deepstack_embed_to_decoder_layer = range(3)
|
||||
# Use HF deepstack order only if rl_on_policy_target is set;
|
||||
# otherwise, retain original order for inference accuracy.
|
||||
self.use_hf_deepstack_order = (
|
||||
get_exec().deterministic.rl_on_policy_target is not None
|
||||
)
|
||||
|
||||
def get_input_embeddings(self) -> nn.Embedding:
|
||||
return self.embed_tokens
|
||||
@@ -104,25 +110,43 @@ class Qwen3MoeLLMModel(Qwen3MoeModel):
|
||||
hidden_states + residual if residual is not None else hidden_states
|
||||
)
|
||||
|
||||
# SGLang applies residual at the START of the next layer, not at the END like HuggingFace.
|
||||
# See: https://github.com/huggingface/transformers/blob/v5.0.0rc0/src/transformers/models/qwen3_vl/modeling_qwen3_vl.py#L549
|
||||
# To match HF behavior, deepstack must be added AFTER residual: (hidden_states + residual) + deepstack
|
||||
# The order matters because addition with different tensors is not associative in practice.
|
||||
# Deepstack for prev_layer is applied at the start of current layer via post_residual_addition.
|
||||
deepstack_embeds = self.get_deepstack_embeds(
|
||||
layer_idx - 1, input_deepstack_embeds
|
||||
)
|
||||
hidden_states, residual = layer(
|
||||
positions,
|
||||
hidden_states,
|
||||
forward_batch,
|
||||
residual,
|
||||
post_residual_addition=deepstack_embeds,
|
||||
)
|
||||
if self.use_hf_deepstack_order:
|
||||
# HF-order path (RL on-policy / FSDP). SGLang applies residual at the START of the
|
||||
# next layer, so to match HF's (hidden_states + residual) + deepstack, deepstack for
|
||||
# the previous layer is added after residual via post_residual_addition.
|
||||
deepstack_embeds = self.get_deepstack_embeds(
|
||||
layer_idx - 1, input_deepstack_embeds
|
||||
)
|
||||
hidden_states, residual = layer(
|
||||
positions,
|
||||
hidden_states,
|
||||
forward_batch,
|
||||
residual,
|
||||
post_residual_addition=deepstack_embeds,
|
||||
)
|
||||
else:
|
||||
# Inference path: add deepstack directly to hidden_states at the end of the layer
|
||||
# (original, grounding-correct order).
|
||||
hidden_states, residual = layer(
|
||||
positions,
|
||||
hidden_states,
|
||||
forward_batch,
|
||||
residual,
|
||||
)
|
||||
if (
|
||||
input_deepstack_embeds is not None
|
||||
and layer_idx in self.deepstack_embed_to_decoder_layer
|
||||
):
|
||||
sep = self.hidden_size * layer_idx
|
||||
hidden_states.add_(
|
||||
input_deepstack_embeds[:, sep : sep + self.hidden_size]
|
||||
)
|
||||
|
||||
# Handle deepstack for the last processed layer if it exists.
|
||||
last_deepstack = self.get_deepstack_embeds(
|
||||
self.end_layer - 1, input_deepstack_embeds
|
||||
# Handle deepstack for the last processed layer (HF-order path only).
|
||||
last_deepstack = (
|
||||
self.get_deepstack_embeds(self.end_layer - 1, input_deepstack_embeds)
|
||||
if self.use_hf_deepstack_order
|
||||
else None
|
||||
)
|
||||
|
||||
if not self.pp_group.is_last_rank:
|
||||
|
||||
@@ -4,9 +4,13 @@ python3 -m unittest test_vision_openai_server.TestOpenAIVisionServer.test_mixed_
|
||||
python3 -m unittest test_vision_openai_server.TestOpenAIVisionServer.test_multi_images_chat_completion
|
||||
"""
|
||||
|
||||
import base64
|
||||
import io
|
||||
import re
|
||||
import unittest
|
||||
|
||||
import openai
|
||||
from PIL import Image, ImageDraw
|
||||
|
||||
from sglang.srt.environ import envs
|
||||
from sglang.test.ci.ci_register import register_cuda_ci
|
||||
@@ -28,6 +32,34 @@ from sglang.test.vlm_utils import (
|
||||
register_cuda_ci(est_time=560, stage="base-b", runner_config="1-gpu-large")
|
||||
|
||||
|
||||
# --- Qwen3-VL grounding regression (deepstack fusion) --------------------------
|
||||
# Guards Qwen3MoeLLMModel.forward: deepstack (multi-scale ViT features) injection
|
||||
# must keep its original inference order. PR #14636 rerouted it through
|
||||
# post_residual_addition (for RL on-policy / FSDP), which is FP-order-sensitive
|
||||
# and regresses FP8 visual grounding (the predicted point drifts by ~150+ px).
|
||||
_GROUNDING_IMG_SIZE = 1000
|
||||
# Target box in pixels; on a 1000x1000 canvas this equals the 0-1000 normalized
|
||||
# coordinate, so the check is robust to normalized-vs-pixel conventions.
|
||||
_GROUNDING_BOX = (620, 180, 880, 360) # (x0, y0, x1, y1), center (750, 270)
|
||||
_GROUNDING_MARGIN = 60
|
||||
_GROUNDING_SYSTEM = (
|
||||
"You are a UI grounding model. Treat the image as a 1000x1000 normalized "
|
||||
"coordinate system with the top-left at (0,0) and the bottom-right at "
|
||||
"(1000,1000). Return the geometric center of the requested element. "
|
||||
"Output ONLY one coordinate in the form (x, y) and nothing else."
|
||||
)
|
||||
_GROUNDING_COORD_RE = re.compile(r"\(?\s*(\d{1,4})\s*,\s*(\d{1,4})\s*\)?")
|
||||
|
||||
|
||||
def _make_grounding_image() -> str:
|
||||
"""White canvas with a single red box at _GROUNDING_BOX; base64 data URI."""
|
||||
img = Image.new("RGB", (_GROUNDING_IMG_SIZE, _GROUNDING_IMG_SIZE), (255, 255, 255))
|
||||
ImageDraw.Draw(img).rectangle(_GROUNDING_BOX, fill=(220, 30, 30))
|
||||
buf = io.BytesIO()
|
||||
img.save(buf, format="PNG")
|
||||
return "data:image/png;base64," + base64.b64encode(buf.getvalue()).decode("utf-8")
|
||||
|
||||
|
||||
class TestLlavaServer(ImageOpenAITestMixin):
|
||||
model = "lmms-lab/llava-onevision-qwen2-0.5b-ov"
|
||||
|
||||
@@ -52,6 +84,47 @@ class TestQwen3VLServer(ImageOpenAITestMixin, VideoOpenAITestMixin):
|
||||
with envs.SGLANG_MM_FEATURE_CACHE_MB.override(512):
|
||||
super().setUpClass()
|
||||
|
||||
def test_deepstack_grounding_hits_target_box(self):
|
||||
# Regression guard for the Qwen3-VL MoE deepstack fusion order: the
|
||||
# predicted point must land inside the target box; a deepstack corruption
|
||||
# drifts it out (see PR #14636).
|
||||
client = openai.Client(api_key=self.api_key, base_url=self.base_url)
|
||||
response = client.chat.completions.create(
|
||||
model="default",
|
||||
messages=[
|
||||
{"role": "system", "content": _GROUNDING_SYSTEM},
|
||||
{
|
||||
"role": "user",
|
||||
"content": [
|
||||
{
|
||||
"type": "image_url",
|
||||
"image_url": {"url": _make_grounding_image()},
|
||||
},
|
||||
{
|
||||
"type": "text",
|
||||
"text": "Point at the center of the red rectangle.",
|
||||
},
|
||||
],
|
||||
},
|
||||
],
|
||||
temperature=0,
|
||||
**(self.get_vision_request_kwargs()),
|
||||
)
|
||||
out = response.choices[0].message.content
|
||||
match = _GROUNDING_COORD_RE.search(out or "")
|
||||
self.assertIsNotNone(match, f"could not parse a coordinate from: {out!r}")
|
||||
x, y = int(match.group(1)), int(match.group(2))
|
||||
x0, y0, x1, y1 = _GROUNDING_BOX
|
||||
inside = (x0 - _GROUNDING_MARGIN <= x <= x1 + _GROUNDING_MARGIN) and (
|
||||
y0 - _GROUNDING_MARGIN <= y <= y1 + _GROUNDING_MARGIN
|
||||
)
|
||||
self.assertTrue(
|
||||
inside,
|
||||
f"grounding output {out!r} -> ({x}, {y}) fell outside target box "
|
||||
f"{_GROUNDING_BOX} (margin {_GROUNDING_MARGIN}); deepstack fusion "
|
||||
f"likely regressed grounding.",
|
||||
)
|
||||
|
||||
|
||||
class TestQwen2VLContextLengthServer(CustomTestCase):
|
||||
# --context-length 300 is calibrated to this model's mm-token expansion:
|
||||
|
||||
Reference in New Issue
Block a user