File size: 5,811 Bytes
58785a2 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 | from __future__ import annotations
from types import MethodType
import torch
import torch.nn.functional as F
from torch import nn
def _rotate_half(value: torch.Tensor) -> torch.Tensor:
first, second = value.chunk(2, dim=-1)
return torch.cat((-second, first), dim=-1)
def _apply_penguin_rotary(query, key, cos, sin):
section = [cos.shape[-1] // 2, cos.shape[-1] // 2]
cos = torch.cat([axis[index % 2] for index, axis in enumerate(cos.split(section, dim=-1))], dim=-1)
sin = torch.cat([axis[index % 2] for index, axis in enumerate(sin.split(section, dim=-1))], dim=-1)
cos = cos.unsqueeze(1)
sin = sin.unsqueeze(1)
return query * cos + _rotate_half(query) * sin, key * cos + _rotate_half(key) * sin
def _penguin_sdpa_attention_forward(
module,
hidden_states: torch.Tensor,
position_embeddings,
attention_mask=None,
past_key_value=None,
cache_position=None,
cu_seqlens=None,
**kwargs,
):
"""SDPA equivalent of Penguin's packed, non-causal FlashAttention path."""
input_shape = hidden_states.shape[:-1]
hidden_shape = (*input_shape, -1, module.head_dim)
query = module.q_norm(module.q_proj(hidden_states).view(hidden_shape)).transpose(1, 2)
key = module.k_norm(module.k_proj(hidden_states).view(hidden_shape)).transpose(1, 2)
value = module.v_proj(hidden_states).view(hidden_shape).transpose(1, 2)
cos, sin = position_embeddings
query, key = _apply_penguin_rotary(query, key, cos, sin)
if past_key_value is not None:
cache_kwargs = {"sin": sin, "cos": cos, "cache_position": cache_position}
key, value = past_key_value.update(key, value, module.layer_idx, cache_kwargs)
if cu_seqlens is None:
attended = F.scaled_dot_product_attention(
query,
key,
value,
attn_mask=attention_mask,
dropout_p=0.0,
is_causal=False,
enable_gqa=query.shape[1] != key.shape[1],
)
else:
# Penguin packs image sequences along S inside a singleton batch. Splitting
# here prevents patches from different images attending to one another.
boundaries = cu_seqlens.detach().cpu().tolist()
chunks = []
for start, end in zip(boundaries[:-1], boundaries[1:]):
chunks.append(
F.scaled_dot_product_attention(
query[:, :, start:end],
key[:, :, start:end],
value[:, :, start:end],
dropout_p=0.0,
is_causal=False,
enable_gqa=query.shape[1] != key.shape[1],
)
)
attended = torch.cat(chunks, dim=2)
attended = attended.transpose(1, 2).reshape(*input_shape, -1).contiguous()
return module.o_proj(attended), None
def install_penguin_sdpa(model: nn.Module) -> int:
patched = 0
for module in model.modules():
if module.__class__.__name__ == "PenguinVLAttention":
module.forward = MethodType(_penguin_sdpa_attention_forward, module)
patched += 1
if patched == 0:
raise RuntimeError("No PenguinVLAttention modules found for SDPA patch")
return patched
def install_penguin_transformers_compatibility() -> None:
"""Bridge Penguin remote code written for Transformers 4.51 to 5.x RoPE APIs."""
from transformers.modeling_rope_utils import ROPE_INIT_FUNCTIONS
if "default" not in ROPE_INIT_FUNCTIONS:
ROPE_INIT_FUNCTIONS["default"] = compute_default_rope_parameters
def compute_default_rope_parameters(config, device=None, **kwargs):
parameters = getattr(config, "rope_parameters", None) or getattr(config, "rope_scaling", None) or {}
theta = float(parameters.get("rope_theta", getattr(config, "rope_theta", 10000.0)))
dimension = int(getattr(config, "head_dim", config.hidden_size // config.num_attention_heads))
inverse_frequency = 1.0 / (
theta ** (torch.arange(0, dimension, 2, dtype=torch.float32, device=device) / dimension)
)
return inverse_frequency, 1.0
class PenguinVisionEncoder(nn.Module):
"""Thin checked wrapper around the official Penguin encoder."""
def __init__(self, model: nn.Module, hidden_size: int = 1024):
super().__init__()
self.model = model
self.hidden_size = hidden_size
def forward(self, pixel_values, grid_sizes, merge_sizes) -> torch.Tensor:
features = self.model(
pixel_values=pixel_values,
grid_sizes=grid_sizes,
merge_sizes=merge_sizes,
)
if hasattr(features, "last_hidden_state"):
features = features.last_hidden_state
if features.ndim != 2 or features.shape[-1] != self.hidden_size:
raise RuntimeError(
f"Expected packed Penguin output [sum(N_i), {self.hidden_size}], "
f"got {tuple(features.shape)}"
)
if grid_sizes.ndim != 2 or grid_sizes.shape[1] != 3:
raise RuntimeError(f"Invalid grid_sizes shape: {tuple(grid_sizes.shape)}")
if merge_sizes.ndim != 1 or merge_sizes.shape[0] != grid_sizes.shape[0]:
raise RuntimeError(f"Invalid merge_sizes shape: {tuple(merge_sizes.shape)}")
spatial = grid_sizes[:, 1:] // merge_sizes[:, None]
if bool((spatial * merge_sizes[:, None] != grid_sizes[:, 1:]).any()):
raise RuntimeError("Grid dimensions must be divisible by merge_size")
counts = grid_sizes[:, 0] * spatial[:, 0] * spatial[:, 1]
if features.shape[0] != int(counts.sum()):
raise RuntimeError(
"Packed Penguin features do not match dynamic grids: "
f"features={features.shape[0]}, counts={counts.tolist()}"
)
return features
|