anodyine/vllm-omni
1
0 Fork 0
Code Issues 4 Pull Requests 2 Actions Packages Projects Releases Wiki Activity

[Refactor] Refactor HunyuanImage3 SigLIP2 ViT to vLLM layers (#3297)

Browse Source 
Signed-off-by: gcanlin <canlinguosdu@gmail.com>
This commit is contained in:
Canlin Guo authored and GitHub committed 2026-06-03 17:51:06 +08:00
1 parent 07737e6c61
commit ce69eb35e9
3 files changed
+342 -388
Download Patch File Download Diff File Expand all files Collapse all files

No files matched your search

vllm_omni/diffusion/models/hunyuan_image3/pipeline_hunyuan_image3.py
+3 -1
View File
@@ -573,7 +573,9 @@ class HunyuanImage3Pipeline(
cond_vit_image_embeds = []
for batch_idx, image in enumerate(cond_vit_images):
cur_kwargs = {k: v[batch_idx] for k, v in vit_kwargs.items()}
image_embed = self.vision_model(image, **cur_kwargs).last_hidden_state
# Siglip2VisionTransformer now returns a plain tensor (B, max_patches,
# hidden_size) after the vLLM-layers refactor, not an HF-style output.
image_embed = self.vision_model(image, **cur_kwargs)
image_embed = self.vision_aligner(image_embed)
n, seq_len, dim = image_embed.shape
image_embed = image_embed.reshape(n * seq_len, dim)
vllm_omni/model_executor/models/hunyuan_image3/hunyuan_image3.py
+73 -6
View File
@@ -23,7 +23,9 @@ from vllm.config.multimodal import BaseDummyOptions
from vllm.distributed import (
get_ep_group,
get_pp_group,
get_tensor_model_parallel_rank,
get_tensor_model_parallel_world_size,
tensor_model_parallel_all_gather,
)
from vllm.inputs import MultiModalDataDict
from vllm.logger import init_logger
@@ -1451,6 +1453,12 @@ class HunyuanImage3ForConditionalGeneration(nn.Module, SupportsMultiModal, Suppo
prefer_model_sampler = True
# Siglip2 ViT supports data-parallel encoding (mm_encoder_tp_mode="data"):
# weights are replicated and the image batch is sharded across TP ranks.
# Without this flag, vLLM silently falls back to "weights" (TP). See
# _vit_encode_dp.
supports_encoder_tp_data = True
packed_modules_mapping = {
"qkv_proj": [
"q_proj",
@@ -1519,7 +1527,14 @@ class HunyuanImage3ForConditionalGeneration(nn.Module, SupportsMultiModal, Suppo
self.time_embed = TimestepEmbedder(hidden_size=config.hidden_size)
# vision
self.vision_model = Siglip2VisionTransformer(config.vit)
multimodal_config = vllm_config.model_config.multimodal_config
self.multimodal_config = multimodal_config
self.use_data_parallel = multimodal_config.mm_encoder_tp_mode == "data"
self.vision_model = Siglip2VisionTransformer(
config.vit,
quant_config=quant_config,
prefix="vision_model",
)
self.vision_aligner = LightProjector(config.vit_aligner)
# Used to embed timestep information into the input sequence.
@@ -1825,17 +1840,69 @@ class HunyuanImage3ForConditionalGeneration(nn.Module, SupportsMultiModal, Suppo
"""
Encode pixel_values through ViT encoder (vision_model and vision_aligner).
"""
# Handle empty batch
# Handle empty batch. Multimodal inputs are identical across TP ranks,
# so this returns on all ranks together without collective-op deadlock.
if pixel_values.shape[0] == 0:
return None
vision_output = self.vision_model(
pixel_values, attention_mask=vit_attention_mask, spatial_shapes=vit_spatial_shapes
)
image_embed = vision_output.last_hidden_state
if self.use_data_parallel:
# ViT weights are replicated (disable_tp on every linear); shard the
# image batch across ranks instead of doing redundant full-batch
# work on each one.
image_embed = self._vit_encode_dp(pixel_values, vit_attention_mask, vit_spatial_shapes)
else:
image_embed = self.vision_model(
pixel_values,
attention_mask=vit_attention_mask,
spatial_shapes=vit_spatial_shapes,
)
image_embed = self.vision_aligner(image_embed)
return image_embed
def _vit_encode_dp(
self,
pixel_values: torch.Tensor,
vit_attention_mask: torch.Tensor,
vit_spatial_shapes: torch.Tensor,
) -> torch.Tensor:
"""Data-parallel ViT: each rank runs a slice of the image batch, then
outputs are all-gathered. Mirrors ``run_dp_sharded_vision_model`` but
shards the three batched ViT inputs together (pixel_values /
attention_mask / spatial_shapes) instead of a single tensor.
"""
num_images = pixel_values.shape[0]
world_size = get_tensor_model_parallel_world_size()
rank = get_tensor_model_parallel_rank()
num_per_rank = (num_images + world_size - 1) // world_size
num_padded = num_per_rank * world_size - num_images
if num_padded > 0:
# Pad the batch so every rank gets an equal slice (required for
# all_gather). Dummy images use spatial_shapes (1, 1) with a single
# valid patch: Siglip2 interpolates position embeddings per image
# and rejects size 0, so (0, 0) padding would crash. Their outputs
# are dropped after the gather.
pad_pixels = pixel_values.new_zeros((num_padded, *pixel_values.shape[1:]))
pixel_values = torch.cat([pixel_values, pad_pixels], dim=0)
pad_mask = vit_attention_mask.new_zeros((num_padded, *vit_attention_mask.shape[1:]))
pad_mask[:, 0] = 1
vit_attention_mask = torch.cat([vit_attention_mask, pad_mask], dim=0)
pad_shapes = vit_spatial_shapes.new_ones((num_padded, *vit_spatial_shapes.shape[1:]))
vit_spatial_shapes = torch.cat([vit_spatial_shapes, pad_shapes], dim=0)
start = rank * num_per_rank
end = start + num_per_rank
image_embed = self.vision_model(
pixel_values[start:end],
attention_mask=vit_attention_mask[start:end],
spatial_shapes=vit_spatial_shapes[start:end],
)
# max_patches (dim 1) is identical across ranks, so dim-0 gather is safe.
image_embed = tensor_model_parallel_all_gather(image_embed.contiguous(), dim=0)
return image_embed[:num_images]
def _timestep_encode(
self,
timestep: torch.Tensor,
vllm_omni/model_executor/models/hunyuan_image3/siglip2.py
+266 -381
View File
@@ -25,18 +25,37 @@
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
"""SigLIP2 Vision Transformer for HunyuanImage3, rewritten in vLLM style.
import warnings
Key optimizations over the original HuggingFace-style implementation:
- QKVParallelLinear: fused QKV projection with tensor parallelism support
- MMEncoderAttention: FlashAttention / xFormers backend
- ColumnParallelLinear / RowParallelLinear: TP-aware MLP layers
- Packed sequence processing: eliminates padding waste via cu_seqlens
- Data parallel support for multi-GPU ViT inference
"""
from collections.abc import Iterable
import torch
import torch.nn as nn
import torch.nn.functional as F
from transformers.activations import ACT2FN
from transformers.modeling_attn_mask_utils import _prepare_4d_attention_mask
from transformers.modeling_outputs import BaseModelOutput, BaseModelOutputWithPooling
from vllm.distributed import divide, get_tensor_model_parallel_world_size
from vllm.model_executor.layers.activation import get_act_fn
from vllm.model_executor.layers.attention import MMEncoderAttention
from vllm.model_executor.layers.linear import (
ColumnParallelLinear,
QKVParallelLinear,
RowParallelLinear,
)
from vllm.model_executor.layers.quantization import QuantizationConfig
from vllm.model_executor.model_loader.weight_utils import default_weight_loader
from vllm.model_executor.models.vision import is_vit_use_data_parallel
class Config:
"""Convert dict config to object with attribute access."""
def __init__(self, config):
if config is not None:
for key, value in config.items():
@@ -65,470 +84,336 @@ class Siglip2VisionEmbeddings(nn.Module):
self.position_embedding_size = int(self.num_patches**0.5)
self.position_embedding = nn.Embedding(self.num_patches, self.embed_dim)
@staticmethod
def resize_positional_embeddings(
positional_embeddings: torch.Tensor,
def forward(
self,
pixel_values: torch.FloatTensor,
spatial_shapes: torch.LongTensor,
max_length: int,
) -> torch.Tensor:
"""
Resize positional embeddings to image-specific size and pad to a fixed size.
Process packed pixel values with per-image position embeddings.
Args:
positional_embeddings (`torch.Tensor`):
Position embeddings of shape (height, width, embed_dim)
spatial_shapes (`torch.LongTensor`):
Spatial shapes of shape (batch_size, 2) to resize the positional embeddings to
max_length (`int`):
Maximum length of the positional embeddings to pad resized positional embeddings to
pixel_values: Packed pixel values
(total_real_patches, num_channels * patch_size * patch_size)
spatial_shapes: Per-image spatial shapes (B, 2) as [(h, w), ...]
Returns:
`torch.Tensor`: Embeddings of shape (batch_size, max_length, embed_dim)
Packed embeddings (total_real_patches, embed_dim)
"""
batch_size = spatial_shapes.shape[0]
embed_dim = positional_embeddings.shape[-1]
source_dtype = positional_embeddings.dtype
target_dtype = self.patch_embedding.weight.dtype
patch_embeds = self.patch_embedding(pixel_values.to(dtype=target_dtype))
resulted_positional_embeddings = torch.empty(
(batch_size, max_length, embed_dim),
device=positional_embeddings.device,
dtype=source_dtype,
# Resize position embeddings per image and concatenate (packed)
positional_embeddings = self.position_embedding.weight.reshape(
self.position_embedding_size, self.position_embedding_size, -1
)
# (H_pe, W_pe, embed_dim) → (1, embed_dim, H_pe, W_pe) for interpolation
pe_for_resize = positional_embeddings.permute(2, 0, 1).unsqueeze(0)
# Upcast on CPU since antialias is not supported for bfloat16/float16
if pe_for_resize.device.type == "cpu":
pe_for_resize = pe_for_resize.to(torch.float32)
# (height, width, embed_dim) -> (1, embed_dim, height, width) for interpolation
positional_embeddings = positional_embeddings.permute(2, 0, 1).unsqueeze(0)
# Upcast to float32 on CPU because antialias is not supported for bfloat16/float16 on CPU
if positional_embeddings.device.type == "cpu":
positional_embeddings = positional_embeddings.to(torch.float32)
for i in range(batch_size):
# (1, dim, height, width) -> (1, dim, target_height, target_width)
height, width = spatial_shapes[i]
resized_embeddings = F.interpolate(
positional_embeddings,
position_embs: list[torch.Tensor] = []
for i in range(spatial_shapes.shape[0]):
height, width = int(spatial_shapes[i, 0]), int(spatial_shapes[i, 1])
resized = F.interpolate(
pe_for_resize,
size=(height, width),
mode="bilinear",
align_corners=False,
antialias=True,
)
# (1, embed_dim, h, w) → (h*w, embed_dim)
resized = resized.reshape(self.embed_dim, height * width).transpose(0, 1)
position_embs.append(resized.to(target_dtype))
# (1, dim, target_height, target_width) -> (target_height * target_width, dim)
resized_embeddings = resized_embeddings.reshape(embed_dim, height * width).transpose(0, 1)
# Cast to original dtype
resized_embeddings = resized_embeddings.to(source_dtype)
resulted_positional_embeddings[i, : height * width] = resized_embeddings
resulted_positional_embeddings[i, height * width :] = resized_embeddings[0]
return resulted_positional_embeddings
def forward(self, pixel_values: torch.FloatTensor, spatial_shapes: torch.LongTensor) -> torch.Tensor:
"""
Args:
pixel_values (`torch.FloatTensor`):
Pixel values of shape (batch_size, max_num_patches, num_channels * patch_size * patch_size)
spatial_shapes (`List[Tuple[int, int]]`):
Spatial shapes of shape (batch_size, 2) to resize the positional embeddings to
"""
# Apply patch embeddings to already patchified pixel values
target_dtype = self.patch_embedding.weight.dtype
patch_embeds = self.patch_embedding(pixel_values.to(dtype=target_dtype))
# Get positional resized and padded positional embeddings
positional_embeddings = self.position_embedding.weight.reshape(
self.position_embedding_size, self.position_embedding_size, -1
)
resized_positional_embeddings = self.resize_positional_embeddings(
positional_embeddings, spatial_shapes, max_length=pixel_values.shape[1]
)
# Add positional embeddings to patch embeddings
embeddings = patch_embeds + resized_positional_embeddings
return embeddings
packed_position_embs = torch.cat(position_embs, dim=0)
return patch_embeds + packed_position_embs
class Siglip2Attention(nn.Module):
"""Multi-headed attention from 'Attention Is All You Need' paper"""
"""Multi-headed attention using QKVParallelLinear and MMEncoderAttention."""
def __init__(self, config):
def __init__(
self,
config,
quant_config: QuantizationConfig | None = None,
prefix: str = "",
):
super().__init__()
self.config = config
self.embed_dim = config.hidden_size
self.num_heads = config.num_attention_heads
self.head_dim = self.embed_dim // self.num_heads
if self.head_dim * self.num_heads != self.embed_dim:
raise ValueError(
f"embed_dim must be divisible by num_heads (got `embed_dim`: {self.embed_dim} and `num_heads`:"
f" {self.num_heads})."
)
self.scale = self.head_dim**-0.5
self.dropout = config.attention_dropout
self.k_proj = nn.Linear(self.embed_dim, self.embed_dim)
self.v_proj = nn.Linear(self.embed_dim, self.embed_dim)
self.q_proj = nn.Linear(self.embed_dim, self.embed_dim)
self.out_proj = nn.Linear(self.embed_dim, self.embed_dim)
def forward(
self,
hidden_states: torch.Tensor,
attention_mask: torch.Tensor | None = None,
output_attentions: bool | None = False,
) -> tuple[torch.Tensor, torch.Tensor | None]:
"""Input shape: Batch x Time x Channel"""
batch_size, q_len, _ = hidden_states.size()
query_states = self.q_proj(hidden_states)
key_states = self.k_proj(hidden_states)
value_states = self.v_proj(hidden_states)
query_states = query_states.view(batch_size, q_len, self.num_heads, self.head_dim).transpose(1, 2)
key_states = key_states.view(batch_size, q_len, self.num_heads, self.head_dim).transpose(1, 2)
value_states = value_states.view(batch_size, q_len, self.num_heads, self.head_dim).transpose(1, 2)
k_v_seq_len = key_states.shape[-2]
attn_weights = torch.matmul(query_states, key_states.transpose(2, 3)) * self.scale
if attn_weights.size() != (batch_size, self.num_heads, q_len, k_v_seq_len):
raise ValueError(
f"Attention weights should be of size {(batch_size, self.num_heads, q_len, k_v_seq_len)}, but is"
f" {attn_weights.size()}"
)
if attention_mask is not None:
if attention_mask.size() != (batch_size, 1, q_len, k_v_seq_len):
raise ValueError(
f"Attention mask should be of size {(batch_size, 1, q_len, k_v_seq_len)}, "
f"but is {attention_mask.size()}"
)
attn_weights = attn_weights + attention_mask
# upcast attention to fp32
attn_weights = nn.functional.softmax(attn_weights, dim=-1, dtype=torch.float32).to(query_states.dtype)
attn_weights = nn.functional.dropout(attn_weights, p=self.dropout, training=self.training)
attn_output = torch.matmul(attn_weights, value_states)
if attn_output.size() != (batch_size, self.num_heads, q_len, self.head_dim):
raise ValueError(
f"`attn_output` should be of size {(batch_size, self.num_heads, q_len, self.head_dim)}, but is"
f" {attn_output.size()}"
)
attn_output = attn_output.transpose(1, 2).contiguous()
attn_output = attn_output.reshape(batch_size, q_len, self.embed_dim)
attn_output = self.out_proj(attn_output)
return attn_output, attn_weights
class Siglip2SdpaAttention(Siglip2Attention):
"""
Siglip2 attention module using torch.nn.functional.scaled_dot_product_attention. This module inherits from
`Siglip2Attention` as the weights of the module stays untouched. The only changes are on the forward pass to adapt
to SDPA API.
"""
is_causal = False
# Adapted from Siglip2Attention.forward and transformers.models.llama.modeling_llama.LlamaSdpaAttention.forward
def forward(
self,
hidden_states: torch.Tensor,
attention_mask: torch.Tensor | None = None,
output_attentions: bool | None = False,
) -> tuple[torch.Tensor, torch.Tensor | None]:
if output_attentions:
# TODO: Improve this warning with e.g. `model.config.attn_implementation = "manual"`
# once this is implemented.
warnings.warn(
"Siglip2Model is using Siglip2SdpaAttention, but `torch.nn.functional.scaled_dot_product_attention` "
"does not support `output_attentions=True`. Falling back to the manual attention implementation, "
"but specifying the manual implementation will be required from Transformers version v5.0.0 onwards. "
'This warning can be removed using the argument `attn_implementation="eager"` when loading the model.'
)
return super().forward(
hidden_states=hidden_states,
attention_mask=attention_mask,
output_attentions=output_attentions,
)
batch_size, q_len, _ = hidden_states.size()
query_states = self.q_proj(hidden_states)
key_states = self.k_proj(hidden_states)
value_states = self.v_proj(hidden_states)
query_states = query_states.view(batch_size, q_len, self.num_heads, self.head_dim).transpose(1, 2)
key_states = key_states.view(batch_size, q_len, self.num_heads, self.head_dim).transpose(1, 2)
value_states = value_states.view(batch_size, q_len, self.num_heads, self.head_dim).transpose(1, 2)
# SDPA with memory-efficient backend is currently (torch==2.1.2) bugged with non-contiguous inputs with
# custom attn_mask,
# Reference: https://github.com/pytorch/pytorch/issues/112577.
if query_states.device.type == "cuda" and attention_mask is not None:
query_states = query_states.contiguous()
key_states = key_states.contiguous()
value_states = value_states.contiguous()
# We dispatch to SDPA's Flash Attention or Efficient kernels via this `is_causal` if statement instead of an
# inline conditional assignment in SDPA to support both torch.compile's dynamic shapes and full graph options.
# An inline conditional prevents dynamic shapes from compiling.
is_causal = True if self.is_causal and q_len > 1 else False
attn_output = torch.nn.functional.scaled_dot_product_attention(
query_states,
key_states,
value_states,
attn_mask=attention_mask,
dropout_p=self.dropout if self.training else 0.0,
is_causal=is_causal,
use_data_parallel = is_vit_use_data_parallel()
self.qkv_proj = QKVParallelLinear(
hidden_size=self.embed_dim,
head_size=self.head_dim,
total_num_heads=self.num_heads,
quant_config=quant_config,
prefix=f"{prefix}.qkv_proj",
disable_tp=use_data_parallel,
)
self.out_proj = RowParallelLinear(
input_size=self.embed_dim,
output_size=self.embed_dim,
quant_config=quant_config,
prefix=f"{prefix}.out_proj",
disable_tp=use_data_parallel,
)
attn_output = attn_output.transpose(1, 2).contiguous()
attn_output = attn_output.view(batch_size, q_len, self.embed_dim)
self.tp_size = 1 if use_data_parallel else get_tensor_model_parallel_world_size()
self.num_heads_per_partition = divide(self.num_heads, self.tp_size)
attn_output = self.out_proj(attn_output)
self.attn = MMEncoderAttention(
num_heads=self.num_heads_per_partition,
head_size=self.head_dim,
scale=self.scale,
prefix=f"{prefix}.attn",
)
return attn_output, None
def forward(
self,
hidden_states: torch.Tensor,
cu_seqlens: torch.Tensor,
) -> torch.Tensor:
"""
Args:
hidden_states: Packed input (total_tokens, embed_dim)
cu_seqlens: Cumulative sequence lengths (B+1,)
"""
seq_length = hidden_states.shape[0]
qkv, _ = self.qkv_proj(hidden_states)
q, k, v = qkv.chunk(3, dim=-1)
q = q.view(seq_length, self.num_heads_per_partition, self.head_dim)
k = k.view(seq_length, self.num_heads_per_partition, self.head_dim)
v = v.view(seq_length, self.num_heads_per_partition, self.head_dim)
max_seqlen = (cu_seqlens[1:] - cu_seqlens[:-1]).max()
attn_output = self.attn(
query=q.unsqueeze(0),
key=k.unsqueeze(0),
value=v.unsqueeze(0),
cu_seqlens=cu_seqlens,
max_seqlen=max_seqlen,
)
attn_output = attn_output.reshape(seq_length, self.num_heads_per_partition * self.head_dim)
attn_output, _ = self.out_proj(attn_output)
return attn_output
class Siglip2MLP(nn.Module):
def __init__(self, config):
def __init__(
self,
config,
quant_config: QuantizationConfig | None = None,
prefix: str = "",
):
super().__init__()
self.config = config
self.activation_fn = ACT2FN[config.hidden_act]
self.fc1 = nn.Linear(config.hidden_size, config.intermediate_size)
self.fc2 = nn.Linear(config.intermediate_size, config.hidden_size)
use_data_parallel = is_vit_use_data_parallel()
self.activation_fn = get_act_fn(config.hidden_act)
self.fc1 = ColumnParallelLinear(
config.hidden_size,
config.intermediate_size,
quant_config=quant_config,
prefix=f"{prefix}.fc1",
disable_tp=use_data_parallel,
)
self.fc2 = RowParallelLinear(
config.intermediate_size,
config.hidden_size,
quant_config=quant_config,
prefix=f"{prefix}.fc2",
disable_tp=use_data_parallel,
)
def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
hidden_states = self.fc1(hidden_states)
hidden_states, _ = self.fc1(hidden_states)
hidden_states = self.activation_fn(hidden_states)
hidden_states = self.fc2(hidden_states)
hidden_states, _ = self.fc2(hidden_states)
return hidden_states
class Siglip2EncoderLayer(nn.Module):
def __init__(self, config):
def __init__(
self,
config,
quant_config: QuantizationConfig | None = None,
prefix: str = "",
):
super().__init__()
self.embed_dim = config.hidden_size
self.self_attn = Siglip2Attention(config=config)
self.self_attn = Siglip2Attention(
config,
quant_config=quant_config,
prefix=f"{prefix}.self_attn",
)
self.layer_norm1 = nn.LayerNorm(self.embed_dim, eps=config.layer_norm_eps)
self.mlp = Siglip2MLP(config)
self.mlp = Siglip2MLP(
config,
quant_config=quant_config,
prefix=f"{prefix}.mlp",
)
self.layer_norm2 = nn.LayerNorm(self.embed_dim, eps=config.layer_norm_eps)
# Ignore copy
def forward(
self,
hidden_states: torch.Tensor,
attention_mask: torch.Tensor,
output_attentions: bool | None = False,
) -> tuple[torch.FloatTensor]:
cu_seqlens: torch.Tensor,
) -> torch.Tensor:
"""
Args:
hidden_states (`torch.FloatTensor`):
Input to the layer of shape `(batch, seq_len, embed_dim)`.
attention_mask (`torch.FloatTensor`):
Attention mask of shape `(batch, 1, q_len, k_v_seq_len)` where padding elements are indicated by very
large negative values.
output_attentions (`bool`, *optional*, defaults to `False`):
Whether or not to return the attentions tensors of all attention layers. See `attentions` under
returned tensors for more detail.
hidden_states: Packed input (total_tokens, embed_dim)
cu_seqlens: Cumulative sequence lengths (B+1,)
"""
residual = hidden_states
hidden_states = self.layer_norm1(hidden_states)
hidden_states, attn_weights = self.self_attn(
hidden_states=hidden_states,
attention_mask=attention_mask,
output_attentions=output_attentions,
)
hidden_states = self.self_attn(hidden_states, cu_seqlens=cu_seqlens)
hidden_states = residual + hidden_states
residual = hidden_states
hidden_states = self.layer_norm2(hidden_states)
hidden_states = self.mlp(hidden_states)
hidden_states = residual + hidden_states
outputs = (hidden_states,)
if output_attentions:
outputs += (attn_weights,)
return outputs
return hidden_states
class Siglip2Encoder(nn.Module):
"""
Transformer encoder consisting of `config.num_hidden_layers` self attention layers. Each layer is a
[`Siglip2EncoderLayer`].
Args:
config: Siglip2Config
"""
"""Transformer encoder with packed sequence processing."""
def __init__(self, config):
def __init__(
self,
config,
quant_config: QuantizationConfig | None = None,
prefix: str = "",
):
super().__init__()
self.config = config
self.layers = nn.ModuleList([Siglip2EncoderLayer(config) for _ in range(config.num_hidden_layers)])
self.gradient_checkpointing = True
self.layers = nn.ModuleList(
[
Siglip2EncoderLayer(
config,
quant_config=quant_config,
prefix=f"{prefix}.layers.{idx}",
)
for idx in range(config.num_hidden_layers)
]
)
# Ignore copy
def forward(
self,
inputs_embeds,
attention_mask: torch.Tensor | None = None,
output_attentions: bool | None = None,
output_hidden_states: bool | None = None,
return_dict: bool | None = None,
) -> tuple | BaseModelOutput:
r"""
Args:
inputs_embeds (`torch.FloatTensor` of shape `(batch_size, sequence_length, hidden_size)`):
Optionally, instead of passing `input_ids` you can choose to directly pass an embedded representation.
This is useful if you want more control over how to convert `input_ids` indices into associated vectors
than the model's internal embedding lookup matrix.
attention_mask (`torch.Tensor` of shape `(batch_size, sequence_length)`, *optional*):
Mask to avoid performing attention on padding token indices. Mask values selected in `[0, 1]`:
- 1 for tokens that are **not masked**,
- 0 for tokens that are **masked**.
What are attention masks? See HuggingFace documentation for details.
output_attentions (`bool`, *optional*):
Whether or not to return the attentions tensors of all attention layers. See `attentions` under
returned tensors for more detail.
output_hidden_states (`bool`, *optional*):
Whether or not to return the hidden states of all layers. See `hidden_states` under returned tensors
for more detail.
return_dict (`bool`, *optional*):
Whether or not to return a [`~utils.ModelOutput`] instead of a plain tuple.
hidden_states: torch.Tensor,
cu_seqlens: torch.Tensor,
) -> torch.Tensor:
"""
output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
output_hidden_states = (
output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
)
return_dict = return_dict if return_dict is not None else self.config.use_return_dict
encoder_states = () if output_hidden_states else None
all_attentions = () if output_attentions else None
hidden_states = inputs_embeds
for layer_index, encoder_layer in enumerate(self.layers): # len(self.layers): 27
if output_hidden_states:
encoder_states = encoder_states + (hidden_states,)
layer_outputs = encoder_layer(
hidden_states,
attention_mask,
output_attentions=output_attentions,
)
hidden_states = layer_outputs[0]
if output_attentions:
all_attentions = all_attentions + (layer_outputs[1],)
if output_hidden_states:
encoder_states = encoder_states + (hidden_states,)
if not return_dict:
return tuple(v for v in [hidden_states, encoder_states, all_attentions] if v is not None)
return BaseModelOutput(last_hidden_state=hidden_states, hidden_states=encoder_states, attentions=all_attentions)
class Siglip2MultiheadAttentionPoolingHead(nn.Module):
"""Multihead Attention Pooling."""
def __init__(self, config):
super().__init__()
self.probe = nn.Parameter(torch.randn(1, 1, config.hidden_size))
self.attention = torch.nn.MultiheadAttention(config.hidden_size, config.num_attention_heads, batch_first=True)
self.layernorm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)
self.mlp = Siglip2MLP(config)
self.num_heads = config.num_attention_heads
def forward(self, hidden_state: torch.Tensor, attention_mask: torch.Tensor | None = None):
batch_size = hidden_state.shape[0]
probe = self.probe.repeat(batch_size, 1, 1)
if attention_mask is not None:
target_len, source_len = probe.shape[1], hidden_state.shape[1]
attention_mask = _prepare_4d_attention_mask(attention_mask, hidden_state.dtype, target_len)
attention_mask = attention_mask.repeat(1, self.num_heads, target_len, 1)
attention_mask = attention_mask.reshape(-1, target_len, source_len)
hidden_state = self.attention(probe, hidden_state, hidden_state, attn_mask=attention_mask)[0]
residual = hidden_state
hidden_state = self.layernorm(hidden_state)
hidden_state = residual + self.mlp(hidden_state)
return hidden_state[:, 0]
Args:
hidden_states: Packed input (total_tokens, embed_dim)
cu_seqlens: Cumulative sequence lengths (B+1,)
"""
for layer in self.layers:
hidden_states = layer(hidden_states, cu_seqlens)
return hidden_states
class Siglip2VisionTransformer(nn.Module):
def __init__(self, config):
def __init__(
self,
config,
quant_config: QuantizationConfig | None = None,
prefix: str = "",
):
super().__init__()
config = Config(config)
self.config = config
embed_dim = config.hidden_size
self.embed_dim = config.hidden_size
self.embeddings = Siglip2VisionEmbeddings(config)
self.encoder = Siglip2Encoder(config)
self.post_layernorm = nn.LayerNorm(embed_dim, eps=config.layer_norm_eps)
self.use_head = True if not hasattr(config, "vision_use_head") else config.vision_use_head
if self.use_head:
self.head = Siglip2MultiheadAttentionPoolingHead(config)
self._use_flash_attention_2 = config._attn_implementation == "flash_attention_2"
self.encoder = Siglip2Encoder(
config,
quant_config=quant_config,
prefix=f"{prefix}.encoder" if prefix else "encoder",
)
self.post_layernorm = nn.LayerNorm(self.embed_dim, eps=config.layer_norm_eps)
def forward(
self,
pixel_values: torch.FloatTensor,
attention_mask: torch.Tensor,
spatial_shapes: torch.LongTensor,
output_attentions: bool | None = None,
output_hidden_states: bool | None = None,
return_dict: bool | None = None,
) -> tuple | BaseModelOutputWithPooling:
r"""
Returns:
) -> torch.Tensor:
"""
output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
output_hidden_states = (
output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
Args:
pixel_values: Batched pixel values
(B, max_num_patches, num_channels * patch_size * patch_size)
attention_mask: (B, max_num_patches) with 1 for real, 0 for padding
spatial_shapes: (B, 2) with (height, width) per image
Returns:
(B, max_num_patches, hidden_size) with zeros at padding positions
"""
batch_size, max_patches, _ = pixel_values.shape
# Pack: extract real tokens using attention_mask
mask_bool = attention_mask.bool()
packed_pixels = pixel_values[mask_bool] # (total_real_patches, C*P*P)
# Compute cu_seqlens from spatial_shapes (h * w per image)
seq_lens = (spatial_shapes[:, 0] * spatial_shapes[:, 1]).to(torch.int32)
cu_seqlens = torch.zeros(batch_size + 1, dtype=torch.int32, device=pixel_values.device)
cu_seqlens[1:] = seq_lens.cumsum(0)
# Embeddings (packed)
hidden_states = self.embeddings(packed_pixels, spatial_shapes)
# Encoder (packed)
hidden_states = self.encoder(hidden_states, cu_seqlens)
# Post layernorm
hidden_states = self.post_layernorm(hidden_states)
# Unpack: scatter back to (B, max_patches, hidden_size)
output = torch.zeros(
batch_size,
max_patches,
self.embed_dim,
dtype=hidden_states.dtype,
device=hidden_states.device,
)
return_dict = return_dict if return_dict is not None else self.config.use_return_dict
output[mask_bool] = hidden_states
hidden_states = self.embeddings(pixel_values, spatial_shapes)
return output
if attention_mask is not None and not self._use_flash_attention_2:
# [batch_size, seq_len] -> [batch_size, 1, tgt_seq_len, src_seq_len]
encoder_attention_mask = _prepare_4d_attention_mask(attention_mask, hidden_states.dtype)
else:
encoder_attention_mask = attention_mask
def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
stacked_params_mapping = [
# (param_name, shard_name, shard_id)
("qkv_proj", "q_proj", "q"),
("qkv_proj", "k_proj", "k"),
("qkv_proj", "v_proj", "v"),
]
params_dict = dict(self.named_parameters())
loaded_params: set[str] = set()
encoder_outputs = self.encoder(
inputs_embeds=hidden_states,
attention_mask=encoder_attention_mask,
output_attentions=output_attentions,
output_hidden_states=output_hidden_states,
return_dict=return_dict,
)
last_hidden_state = encoder_outputs[0]
last_hidden_state = self.post_layernorm(last_hidden_state)
pooler_output = self.head(last_hidden_state, attention_mask) if self.use_head else None
if not return_dict:
return (last_hidden_state, pooler_output) + encoder_outputs[1:]
return BaseModelOutputWithPooling(
last_hidden_state=last_hidden_state,
pooler_output=pooler_output,
hidden_states=encoder_outputs.hidden_states,
attentions=encoder_outputs.attentions,
)
for name, loaded_weight in weights:
for param_name, weight_name, shard_id in stacked_params_mapping:
if weight_name not in name:
continue
name = name.replace(weight_name, param_name)
param = params_dict[name]
weight_loader = param.weight_loader
weight_loader(param, loaded_weight, shard_id)
break
else:
if name not in params_dict:
# Skip weights for removed modules (e.g., pooling head)
continue
param = params_dict[name]
weight_loader = getattr(param, "weight_loader", default_weight_loader)
weight_loader(param, loaded_weight)
loaded_params.add(name)
return loaded_params
class LightProjector(nn.Module):