minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

[diffusion] chore: reuse SRT SigLIP in Pi0.5 (#34992)

Browse Source
This commit is contained in:
Mick
2026-08-17 19:33:36 +08:00
committed by GitHub
parent f33b83b4cc
commit e9ad8102a2
5 changed files with 155 additions and 148 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/multimodal_gen/runtime/models/vlas/pi05_core.py
+46 -114
View File
@@ -5,6 +5,7 @@ from __future__ import annotations
import math
from dataclasses import dataclass
from functools import partial
from typing import Any, Literal
import torch
@@ -12,8 +13,7 @@ import torch.nn.functional as F
from torch import Tensor, nn
from transformers.modeling_outputs import BaseModelOutputWithPooling
from transformers.models.auto import CONFIG_MAPPING
from transformers.models.gemma.modeling_gemma import GemmaConfig
from transformers.models.paligemma.modeling_paligemma import PaliGemmaModel
from transformers.models.gemma.configuration_gemma import GemmaConfig
from sglang.multimodal_gen.configs.pipeline_configs.pi05 import Pi05PipelineConfig
from sglang.multimodal_gen.runtime.distributed.parallel_state import (
@@ -22,6 +22,7 @@ from sglang.multimodal_gen.runtime.distributed.parallel_state import (
get_ulysses_parallel_world_size,
model_parallel_is_initialized,
)
from sglang.multimodal_gen.runtime.layers.activation import get_act_fn
from sglang.multimodal_gen.runtime.layers.attention import LocalAttention, USPAttention
from sglang.multimodal_gen.runtime.layers.linear import (
MergedColumnParallelLinear,
@@ -30,12 +31,16 @@ from sglang.multimodal_gen.runtime.layers.linear import (
)
from sglang.multimodal_gen.runtime.layers.rotary_embedding import RotaryEmbedding
from sglang.multimodal_gen.runtime.managers.forward_context import set_forward_context
from sglang.multimodal_gen.runtime.managers.memory_managers.layerwise_offload import (
LayerwiseOffloadableModuleMixin,
)
from sglang.multimodal_gen.runtime.platforms import AttentionBackendEnum
from sglang.multimodal_gen.runtime.vla.prefix_cache import VLADensePrefixCache
from sglang.srt.layers.activation import GeluAndMul
from sglang.srt.layers.rotary_embedding import (
apply_rotary_pos_emb as native_apply_rotary_pos_emb,
)
from sglang.srt.models.siglip import SiglipVisionModel
def config_compute_dtype(config: GemmaConfig) -> torch.dtype | None:
@@ -106,73 +111,6 @@ def _use_ulysses_action_attention(num_heads: int) -> bool:
)
class Pi05SiglipAttention(nn.Module):
def __init__(self, attention: nn.Module):
super().__init__()
self.embed_dim = attention.embed_dim
self.num_heads = attention.num_heads
self.head_dim = attention.head_dim
self.scale = getattr(attention, "scale", self.head_dim**-0.5)
self.dropout = getattr(attention, "dropout", 0.0)
self.q_proj = attention.q_proj
self.k_proj = attention.k_proj
self.v_proj = attention.v_proj
self.out_proj = attention.out_proj
self.attn = LocalAttention(
num_heads=self.num_heads,
head_size=self.head_dim,
num_kv_heads=self.num_heads,
softmax_scale=self.scale,
causal=False,
supported_attention_backends={
AttentionBackendEnum.FA,
AttentionBackendEnum.FA2,
AttentionBackendEnum.TORCH_SDPA,
},
compute_dtype=self.q_proj.weight.dtype,
allow_cudnn_sdp=True,
)
def forward(
self,
hidden_states: torch.Tensor,
attention_mask: torch.Tensor | None = None,
output_attentions: bool = False,
**kwargs,
) -> tuple[torch.Tensor, None]:
input_shape = hidden_states.shape[:-1]
query_states = self.q_proj(hidden_states).view(
*input_shape,
self.num_heads,
self.head_dim,
)
key_states = self.k_proj(hidden_states).view(
*input_shape,
self.num_heads,
self.head_dim,
)
value_states = self.v_proj(hidden_states).view(
*input_shape,
self.num_heads,
self.head_dim,
)
attn_output = self.attn(
query_states,
key_states,
value_states,
attn_mask=attention_mask,
)
attn_output = attn_output.reshape(*input_shape, self.embed_dim).contiguous()
return self.out_proj(attn_output), None
def patch_siglip_vision_attention_to_native(vision_model: nn.Module) -> None:
for layer in vision_model.encoder.layers:
if isinstance(layer.self_attn, Pi05SiglipAttention):
continue
layer.self_attn = Pi05SiglipAttention(layer.self_attn)
class PiGemmaRMSNorm(nn.Module):
def __init__(self, dim: int, eps: float = 1e-6, cond_dim: int | None = None):
super().__init__()
@@ -788,15 +726,51 @@ class PiGemmaForCausalLM(nn.Module):
self.lm_head = None
class PaliGemmaModelWithPiGemma(PaliGemmaModel):
class PaliGemmaMultiModalProjector(nn.Module):
def __init__(self, config):
super().__init__()
self.linear = nn.Linear(
config.vision_config.hidden_size,
config.vision_config.projection_dim,
bias=True,
)
def forward(self, image_features: torch.Tensor) -> torch.Tensor:
return self.linear(image_features)
class Pi05SiglipVisionModel(SiglipVisionModel, LayerwiseOffloadableModuleMixin):
layerwise_offload_dit_group_enabled = False
layer_names = ["vision_model.encoder.layers"]
class PaliGemmaModelWithPiGemma(nn.Module):
def __init__(self, config, *, tensor_parallel: bool = False):
super().__init__(config)
del self.language_model
super().__init__()
self.config = config
self.vision_tower = Pi05SiglipVisionModel(
config.vision_config,
act_layer=partial(get_act_fn, config.vision_config.hidden_act),
qkv_backend="sdpa",
flatten_batch=False,
use_data_parallel=True,
)
self.multi_modal_projector = PaliGemmaMultiModalProjector(config)
self.language_model = PiGemmaModel(
config.text_config,
tensor_parallel=tensor_parallel,
)
def get_image_features(
self, pixel_values: torch.Tensor
) -> BaseModelOutputWithPooling:
vision_features = self.vision_tower(pixel_values)
image_features = self.multi_modal_projector(vision_features)
return BaseModelOutputWithPooling(
last_hidden_state=vision_features,
pooler_output=image_features,
)
class PaliGemmaForConditionalGenerationWithPiGemma(nn.Module):
def __init__(self, config, *, tensor_parallel: bool = False):
@@ -907,34 +881,6 @@ def prepare_optional_full_attention_mask(
return torch.where(masks_4d, 0.0, OPENPI_ATTENTION_MASK_VALUE)
def siglip_vision_forward_with_openpi_dtype(
self,
pixel_values,
interpolate_pos_encoding: bool | None = False,
**kwargs,
) -> BaseModelOutputWithPooling:
hidden_states = self.embeddings(
pixel_values,
interpolate_pos_encoding=interpolate_pos_encoding,
)
if (
len(self.encoder.layers) > 0
and self.encoder.layers[0].self_attn.q_proj.weight.dtype == torch.bfloat16
):
hidden_states = hidden_states.to(torch.bfloat16)
encoder_outputs = self.encoder(inputs_embeds=hidden_states, **kwargs)
last_hidden_state = encoder_outputs.last_hidden_state
last_hidden_state = self.post_layernorm(last_hidden_state)
pooler_output = self.head(last_hidden_state) if self.use_head else None
return BaseModelOutputWithPooling(
last_hidden_state=last_hidden_state,
pooler_output=pooler_output,
hidden_states=encoder_outputs.hidden_states,
attentions=encoder_outputs.attentions,
)
def compute_layer_complete(
inputs_embeds,
attention_mask,
@@ -1059,12 +1005,6 @@ class PaliGemmaWithExpertModel(nn.Module):
config=vlm_config_hf,
tensor_parallel=prefix_tensor_parallel,
)
vision_tower = self.paligemma.model.vision_tower
vision_model = getattr(vision_tower, "vision_model", vision_tower)
vision_model.forward = siglip_vision_forward_with_openpi_dtype.__get__(
vision_model,
type(vision_model),
)
self.paligemma.lm_head = None
if runtime_role in ("all", "action"):
@@ -1090,14 +1030,6 @@ class PaliGemmaWithExpertModel(nn.Module):
self.gemma_expert.lm_head = None
self.gemma_expert.model.embed_tokens = None
self.to_selected_dtype(precision)
self.patch_native_attention_after_dtype_finalize()
def patch_native_attention_after_dtype_finalize(self) -> None:
if self.paligemma is None:
return
vision_tower = self.paligemma.model.vision_tower
vision_model = getattr(vision_tower, "vision_model", vision_tower)
patch_siglip_vision_attention_to_native(vision_model)
def to_selected_dtype(
self, precision: Literal["bfloat16", "float32"] = "bfloat16"
python/sglang/multimodal_gen/runtime/models/vlas/pi05_policy.py
+4 -2
View File
@@ -569,6 +569,7 @@ class Pi05PolicyModel(nn.Module):
candidates = [key]
replacements = {
".vision_tower.vision_model.": ".vision_tower.",
".self_attn.out_proj.": ".self_attn.proj.",
".paligemma.language_model.": ".paligemma.model.language_model.",
".paligemma.vision_tower.": ".paligemma.model.vision_tower.",
".paligemma.multi_modal_projector.": (
@@ -576,8 +577,9 @@ class Pi05PolicyModel(nn.Module):
),
}
for old, new in replacements.items():
if old in key:
candidates.append(key.replace(old, new))
for candidate in list(candidates):
if old in candidate:
candidates.append(candidate.replace(old, new))
if key in {
"paligemma_with_expert.paligemma.lm_head.weight",
python/sglang/multimodal_gen/test/unit/test_pi05_runtime_helpers.py
+58 -23
View File
@@ -10,8 +10,7 @@ import sglang.multimodal_gen.runtime.models.vlas.pi05_policy as pi05_policy_modu
from sglang.multimodal_gen.configs.pipeline_configs.pi05 import Pi05PipelineConfig
from sglang.multimodal_gen.runtime.models.vlas.pi05_core import (
Pi05CoreModel,
Pi05SiglipAttention,
patch_siglip_vision_attention_to_native,
Pi05SiglipVisionModel,
)
from sglang.multimodal_gen.runtime.models.vlas.pi05_policy import (
Pi05CheckpointManifest,
@@ -30,6 +29,8 @@ from sglang.multimodal_gen.runtime.vla.prefix_cache import (
PrefixContext,
VLADensePrefixCache,
)
from sglang.srt.models.siglip import SiglipVisionModel
from sglang.srt.runtime_context import get_context
def _prefix_context(value: float, digest: str | None) -> PrefixContext:
@@ -183,30 +184,64 @@ def test_action_parallel_info_reports_single_rank_without_process_group():
}
class _FakeSiglipAttention(nn.Module):
def __init__(self):
super().__init__()
self.embed_dim = 8
self.num_heads = 2
self.head_dim = 4
self.scale = self.head_dim**-0.5
self.dropout = 0.0
self.q_proj = nn.Linear(self.embed_dim, self.embed_dim)
self.k_proj = nn.Linear(self.embed_dim, self.embed_dim)
self.v_proj = nn.Linear(self.embed_dim, self.embed_dim)
self.out_proj = nn.Linear(self.embed_dim, self.embed_dim)
def test_pi05_siglip_reuses_srt_model_with_layerwise_groups():
config = SimpleNamespace(
hidden_size=8,
intermediate_size=16,
num_hidden_layers=1,
num_attention_heads=2,
layer_norm_eps=1e-6,
image_size=4,
patch_size=2,
num_channels=3,
hidden_act="gelu_pytorch_tanh",
)
with get_context().override_server_args():
model = Pi05SiglipVisionModel(
config,
act_layer=lambda: nn.GELU(approximate="tanh"),
qkv_backend="sdpa",
flatten_batch=False,
use_data_parallel=True,
)
state_keys = set(model.state_dict())
prefix = "vision_model.encoder.layers.0.self_attn"
vision_model = model.vision_model
layer = vision_model.encoder.layers[0]
assert isinstance(model, SiglipVisionModel)
assert f"{prefix}.qkv_proj.weight" in state_keys
assert f"{prefix}.proj.weight" in state_keys
assert vision_model.embeddings.position_embedding.tp_size == 1
assert layer.self_attn.tp_size == 1
assert layer.self_attn.qkv_backend.flatten_batch is False
assert layer.mlp.fc1.tp_size == 1
assert layer.mlp.fc2.tp_size == 1
assert isinstance(layer.mlp.act, nn.GELU)
assert layer.mlp.act.approximate == "tanh"
assert model.device == vision_model.embeddings.patch_embedding.weight.device
assert model.layer_names == ["vision_model.encoder.layers"]
def test_siglip_attention_patch_uses_native_wrapper_once():
layer = SimpleNamespace(self_attn=_FakeSiglipAttention())
vision_model = SimpleNamespace(encoder=SimpleNamespace(layers=[layer]))
def test_pi05_siglip_checkpoint_names_map_to_srt_layers():
source_prefix = (
"paligemma_with_expert.paligemma.vision_tower.vision_model."
"encoder.layers.0.self_attn"
)
target_prefix = (
"paligemma_with_expert.paligemma.model.vision_tower.vision_model."
"encoder.layers.0.self_attn"
)
patch_siglip_vision_attention_to_native(vision_model)
first = layer.self_attn
patch_siglip_vision_attention_to_native(vision_model)
assert isinstance(first, Pi05SiglipAttention)
assert layer.self_attn is first
assert (
f"{target_prefix}.qkv_proj.weight",
"q",
) in Pi05PolicyModel._candidate_target_weights(f"{source_prefix}.q_proj.weight")
assert (
f"{target_prefix}.proj.weight",
None,
) in Pi05PolicyModel._candidate_target_weights(f"{source_prefix}.out_proj.weight")
def test_prefix_language_embedding_matches_openpi_scale():
python/sglang/srt/models/siglip.py
+45 -9
View File
@@ -2,7 +2,7 @@
# https://github.com/huggingface/transformers/blob/af9b2eaa54c150741f298d6db939af6328e1dc38/src/transformers/models/siglip/modeling_siglip.py
from functools import partial
from typing import Optional, Type, Union
from typing import Callable, Optional, Type, Union
import torch
import torch.nn as nn
@@ -14,13 +14,14 @@ from sglang.srt.layers.conv import Conv2dLayer
from sglang.srt.layers.linear import ColumnParallelLinear, RowParallelLinear
from sglang.srt.layers.quantization.base_config import QuantizationConfig
from sglang.srt.layers.vocab_parallel_embedding import VocabParallelEmbedding
from sglang.srt.runtime_context import get_parallel
from sglang.srt.utils import add_prefix
# Adapted from transformers.models.siglip.modeling_siglip.SiglipVisionTransformer
class SiglipVisionEmbeddings(nn.Module):
def __init__(self, config: SiglipVisionConfig):
def __init__(self, config: SiglipVisionConfig, use_data_parallel: bool = False):
super().__init__()
self.config = config
self.embed_dim = config.hidden_size
@@ -38,7 +39,9 @@ class SiglipVisionEmbeddings(nn.Module):
self.num_patches = (self.image_size // self.patch_size) ** 2
self.num_positions = self.num_patches
self.position_embedding = VocabParallelEmbedding(
self.num_positions, self.embed_dim
self.num_positions,
self.embed_dim,
enable_tp=not use_data_parallel,
)
self.register_buffer(
"position_ids",
@@ -64,16 +67,21 @@ class SiglipMLP(nn.Module):
def __init__(
self,
config,
act_layer: Type[nn.Module] = QuickGELU,
act_layer: Callable[[], nn.Module] = QuickGELU,
quant_config: Optional[QuantizationConfig] = None,
prefix: str = "",
use_data_parallel: bool = False,
):
super().__init__()
tp_size = 1 if use_data_parallel else get_parallel().tp_size
tp_rank = 0 if use_data_parallel else get_parallel().tp_rank
self.fc1 = ColumnParallelLinear(
config.hidden_size,
config.intermediate_size,
quant_config=quant_config,
prefix=add_prefix("fc1", prefix),
tp_size=tp_size,
tp_rank=tp_rank,
)
self.act = act_layer()
self.fc2 = RowParallelLinear(
@@ -81,6 +89,8 @@ class SiglipMLP(nn.Module):
config.hidden_size,
quant_config=quant_config,
prefix=add_prefix("fc2", prefix),
tp_size=tp_size,
tp_rank=tp_rank,
)
def forward(self, x: torch.Tensor) -> torch.Tensor:
@@ -96,11 +106,13 @@ class SiglipEncoderLayer(nn.Module):
def __init__(
self,
config: SiglipVisionConfig,
act_layer: Type[nn.Module] = QuickGELU,
act_layer: Callable[[], nn.Module] = QuickGELU,
norm_layer: Type[nn.Module] = None,
quant_config: Optional[QuantizationConfig] = None,
prefix: str = "",
qkv_backend: Optional[str] = None,
flatten_batch: bool = True,
use_data_parallel: bool = False,
) -> None:
super().__init__()
if norm_layer is None:
@@ -112,14 +124,16 @@ class SiglipEncoderLayer(nn.Module):
num_heads=config.num_attention_heads,
projection_size=config.hidden_size,
use_qkv_parallel=True,
flatten_batch=True,
flatten_batch=flatten_batch,
qkv_backend=qkv_backend,
use_data_parallel=use_data_parallel,
quant_config=quant_config,
prefix=add_prefix("self_attn", prefix),
)
self.mlp = SiglipMLP(
config,
act_layer=act_layer,
use_data_parallel=use_data_parallel,
quant_config=quant_config,
prefix=add_prefix("mlp", prefix),
)
@@ -170,6 +184,9 @@ class SiglipEncoder(nn.Module):
quant_config: Optional[QuantizationConfig] = None,
prefix: str = "",
qkv_backend: Optional[str] = None,
act_layer: Callable[[], nn.Module] = QuickGELU,
flatten_batch: bool = True,
use_data_parallel: bool = False,
) -> None:
super().__init__()
@@ -183,6 +200,9 @@ class SiglipEncoder(nn.Module):
config=config,
norm_layer=norm_layer,
qkv_backend=qkv_backend,
act_layer=act_layer,
flatten_batch=flatten_batch,
use_data_parallel=use_data_parallel,
quant_config=quant_config,
prefix=add_prefix(f"layers.{layer_idx}", prefix),
)
@@ -220,17 +240,25 @@ class SiglipVisionTransformer(nn.Module):
quant_config: Optional[QuantizationConfig] = None,
prefix: str = "",
qkv_backend: Optional[str] = None,
act_layer: Callable[[], nn.Module] = QuickGELU,
flatten_batch: bool = True,
use_data_parallel: bool = False,
) -> None:
super().__init__()
self.config = config
embed_dim = config.hidden_size
self.embeddings = SiglipVisionEmbeddings(config)
self.embeddings = SiglipVisionEmbeddings(
config, use_data_parallel=use_data_parallel
)
self.encoder = SiglipEncoder(
config=config,
qkv_backend=qkv_backend,
act_layer=act_layer,
flatten_batch=flatten_batch,
use_data_parallel=use_data_parallel,
quant_config=quant_config,
prefix=add_prefix("encoder", prefix),
)
@@ -247,13 +275,15 @@ class SiglipVisionTransformer(nn.Module):
@property
def device(self) -> torch.device:
return self.encoder.layers[0].layer_norm1.weight.device
return self.embeddings.patch_embedding.weight.device
def forward(
self,
pixel_values: torch.Tensor,
) -> torch.Tensor:
hidden_states = self.embeddings(pixel_values.to(self.device))
hidden_states = self.embeddings(pixel_values.to(self.device)).to(
self.post_layernorm.weight.dtype
)
return_all_hidden_states = False
@@ -275,11 +305,17 @@ class SiglipVisionModel(nn.Module):
quant_config: Optional[QuantizationConfig] = None,
prefix: str = "",
qkv_backend: Optional[str] = None,
act_layer: Callable[[], nn.Module] = QuickGELU,
flatten_batch: bool = True,
use_data_parallel: bool = False,
):
super().__init__()
self.vision_model = SiglipVisionTransformer(
config,
qkv_backend=qkv_backend,
act_layer=act_layer,
flatten_batch=flatten_batch,
use_data_parallel=use_data_parallel,
quant_config=quant_config,
prefix=add_prefix("vision_model", prefix),
)