diff --git a/docs/cookbook/vla/OpenPI/Pi0.5.mdx b/docs/cookbook/vla/OpenPI/Pi0.5.mdx index 1e6f57f9f..166eb2f2d 100644 --- a/docs/cookbook/vla/OpenPI/Pi0.5.mdx +++ b/docs/cookbook/vla/OpenPI/Pi0.5.mdx @@ -18,6 +18,8 @@ Pi0.5 is an OpenPI / LeRobot diffusion Vision-Language-Action (dVLA) policy. It SGLang serves Pi0.5 through the native `multimodal_gen` runtime. The implementation uses a SigLIP/PaliGemma prefix encoder and a Gemma action expert: the prefix is encoded once, then the action expert runs the flow-matching denoising loop. This is not a token decode workload, so the Pi0.5 path does not use the LLM sampler, logits processor, token streaming, paged decode KV cache, or a separate SRT serving engine. +The SigLIP vision tower, PaliGemma language stack, and action expert are all SGLang-native modules. Transformers is used for checkpoint configuration and tokenization, not for the runtime neural network. + The prefix encoder covers both stages of observation encoding: SigLIP turns resized camera pixels into continuous patch embeddings, then the PaliGemma transformer jointly encodes those patches with tokenized task/state inputs and produces per-layer prefix K/V. At flow timestep `t`, the action expert projects the noisy continuous action chunk `x_t` into action embeddings. Its queries attend to both the fixed prefix K/V and the current action K/V, while the timestep follows a separate sinusoidal-MLP path and conditions every action-expert layer through AdaRMSNorm gates. Supported public checkpoints: diff --git a/python/sglang/multimodal_gen/runtime/models/vlas/pi05_core.py b/python/sglang/multimodal_gen/runtime/models/vlas/pi05_core.py index 3146dd65d..01b57c309 100644 --- a/python/sglang/multimodal_gen/runtime/models/vlas/pi05_core.py +++ b/python/sglang/multimodal_gen/runtime/models/vlas/pi05_core.py @@ -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" diff --git a/python/sglang/multimodal_gen/runtime/models/vlas/pi05_policy.py b/python/sglang/multimodal_gen/runtime/models/vlas/pi05_policy.py index 2b189eb17..7a7d18ee5 100644 --- a/python/sglang/multimodal_gen/runtime/models/vlas/pi05_policy.py +++ b/python/sglang/multimodal_gen/runtime/models/vlas/pi05_policy.py @@ -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", diff --git a/python/sglang/multimodal_gen/test/unit/test_pi05_runtime_helpers.py b/python/sglang/multimodal_gen/test/unit/test_pi05_runtime_helpers.py index 6271654fb..ea81e89a7 100644 --- a/python/sglang/multimodal_gen/test/unit/test_pi05_runtime_helpers.py +++ b/python/sglang/multimodal_gen/test/unit/test_pi05_runtime_helpers.py @@ -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(): diff --git a/python/sglang/srt/models/siglip.py b/python/sglang/srt/models/siglip.py index 60443b76c..6f8dec4d6 100644 --- a/python/sglang/srt/models/siglip.py +++ b/python/sglang/srt/models/siglip.py @@ -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), )