[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(
|
self.deepstack_embed_to_decoder_layer = range(
|
||||||
len(config.vision_config.deepstack_visual_indexes)
|
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(
|
def get_deepstack_embeds(
|
||||||
self, layer_idx: int, input_deepstack_embeds: Optional[torch.Tensor]
|
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
|
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.
|
if self.use_hf_deepstack_order:
|
||||||
# See: https://github.com/huggingface/transformers/blob/v5.0.0rc0/src/transformers/models/qwen3_vl/modeling_qwen3_vl.py#L549
|
# HF-order path (RL on-policy / FSDP). SGLang applies residual at the START of the
|
||||||
# To match HF behavior, deepstack must be added AFTER residual: (hidden_states + residual) + deepstack
|
# next layer, so to match HF's (hidden_states + residual) + deepstack, deepstack for
|
||||||
# The order matters because addition with different tensors is not associative in practice.
|
# the previous layer is added after residual via post_residual_addition.
|
||||||
# Deepstack for prev_layer is applied at the start of current layer via post_residual_addition.
|
deepstack_embeds = self.get_deepstack_embeds(
|
||||||
deepstack_embeds = self.get_deepstack_embeds(
|
layer_idx - 1, input_deepstack_embeds
|
||||||
layer_idx - 1, input_deepstack_embeds
|
)
|
||||||
)
|
hidden_states, residual = layer(
|
||||||
hidden_states, residual = layer(
|
positions,
|
||||||
positions,
|
hidden_states,
|
||||||
hidden_states,
|
forward_batch,
|
||||||
forward_batch,
|
residual,
|
||||||
residual,
|
post_residual_addition=deepstack_embeds,
|
||||||
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.
|
# Handle deepstack for the last processed layer (HF-order path only).
|
||||||
last_deepstack = self.get_deepstack_embeds(
|
last_deepstack = (
|
||||||
self.end_layer - 1, input_deepstack_embeds
|
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:
|
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.model_loader.weight_utils import default_weight_loader
|
||||||
from sglang.srt.models.qwen3_moe import Qwen3MoeDecoderLayer, Qwen3MoeModel
|
from sglang.srt.models.qwen3_moe import Qwen3MoeDecoderLayer, Qwen3MoeModel
|
||||||
from sglang.srt.models.qwen3_vl import Qwen3VLForConditionalGeneration
|
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
|
from sglang.srt.utils.hf_transformers_utils import get_processor
|
||||||
|
|
||||||
logger = logging.getLogger(__name__)
|
logger = logging.getLogger(__name__)
|
||||||
@@ -58,6 +59,11 @@ class Qwen3MoeLLMModel(Qwen3MoeModel):
|
|||||||
# This approach follows the original implementation.
|
# This approach follows the original implementation.
|
||||||
# TODO: make config of type Qwen3VLMoeConfig, so that we can directly obtain deepstack_visual_indexes.
|
# TODO: make config of type Qwen3VLMoeConfig, so that we can directly obtain deepstack_visual_indexes.
|
||||||
self.deepstack_embed_to_decoder_layer = range(3)
|
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:
|
def get_input_embeddings(self) -> nn.Embedding:
|
||||||
return self.embed_tokens
|
return self.embed_tokens
|
||||||
@@ -104,25 +110,43 @@ class Qwen3MoeLLMModel(Qwen3MoeModel):
|
|||||||
hidden_states + residual if residual is not None else hidden_states
|
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.
|
if self.use_hf_deepstack_order:
|
||||||
# See: https://github.com/huggingface/transformers/blob/v5.0.0rc0/src/transformers/models/qwen3_vl/modeling_qwen3_vl.py#L549
|
# HF-order path (RL on-policy / FSDP). SGLang applies residual at the START of the
|
||||||
# To match HF behavior, deepstack must be added AFTER residual: (hidden_states + residual) + deepstack
|
# next layer, so to match HF's (hidden_states + residual) + deepstack, deepstack for
|
||||||
# The order matters because addition with different tensors is not associative in practice.
|
# the previous layer is added after residual via post_residual_addition.
|
||||||
# Deepstack for prev_layer is applied at the start of current layer via post_residual_addition.
|
deepstack_embeds = self.get_deepstack_embeds(
|
||||||
deepstack_embeds = self.get_deepstack_embeds(
|
layer_idx - 1, input_deepstack_embeds
|
||||||
layer_idx - 1, input_deepstack_embeds
|
)
|
||||||
)
|
hidden_states, residual = layer(
|
||||||
hidden_states, residual = layer(
|
positions,
|
||||||
positions,
|
hidden_states,
|
||||||
hidden_states,
|
forward_batch,
|
||||||
forward_batch,
|
residual,
|
||||||
residual,
|
post_residual_addition=deepstack_embeds,
|
||||||
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.
|
# Handle deepstack for the last processed layer (HF-order path only).
|
||||||
last_deepstack = self.get_deepstack_embeds(
|
last_deepstack = (
|
||||||
self.end_layer - 1, input_deepstack_embeds
|
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:
|
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
|
python3 -m unittest test_vision_openai_server.TestOpenAIVisionServer.test_multi_images_chat_completion
|
||||||
"""
|
"""
|
||||||
|
|
||||||
|
import base64
|
||||||
|
import io
|
||||||
|
import re
|
||||||
import unittest
|
import unittest
|
||||||
|
|
||||||
import openai
|
import openai
|
||||||
|
from PIL import Image, ImageDraw
|
||||||
|
|
||||||
from sglang.srt.environ import envs
|
from sglang.srt.environ import envs
|
||||||
from sglang.test.ci.ci_register import register_cuda_ci
|
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")
|
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):
|
class TestLlavaServer(ImageOpenAITestMixin):
|
||||||
model = "lmms-lab/llava-onevision-qwen2-0.5b-ov"
|
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):
|
with envs.SGLANG_MM_FEATURE_CACHE_MB.override(512):
|
||||||
super().setUpClass()
|
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):
|
class TestQwen2VLContextLengthServer(CustomTestCase):
|
||||||
# --context-length 300 is calibrated to this model's mm-token expansion:
|
# --context-length 300 is calibrated to this model's mm-token expansion:
|
||||||
|
|||||||
Reference in New Issue
Block a user