mlx-community/MolmoPoint-8B-5bit
016
1import math2from copy import deepcopy3from dataclasses import dataclass4from typing import Optional, Union, Callable5 6import torch7from torch import nn8from torch.nn import functional as F9 10from transformers.models.auto import AutoModelForImageTextToText11from transformers.activations import ACT2FN12from transformers.configuration_utils import PretrainedConfig13from transformers.cache_utils import Cache, DynamicCache14from transformers.generation import GenerationMixin15from transformers.masking_utils import create_causal_mask, create_masks_for_generate16from transformers.modeling_flash_attention_utils import (17 _flash_attention_forward,18 FlashAttentionKwargs,19 flash_attn_supports_top_left_mask,20)21from transformers.modeling_layers import GradientCheckpointingLayer22from transformers.modeling_outputs import (23 BaseModelOutputWithPast,24)25from transformers.modeling_rope_utils import ROPE_INIT_FUNCTIONS, dynamic_rope_update26from transformers.modeling_utils import ALL_ATTENTION_FUNCTIONS, PreTrainedModel27from transformers.processing_utils import Unpack28from transformers.utils import (29 ModelOutput,30 TransformersKwargs,31 can_return_tuple,32 logging,33)34 35from .configuration_molmo2 import Molmo2Config, Molmo2VitConfig, Molmo2AdapterConfig, Molmo2TextConfig36 37 38logger = logging.get_logger(__name__)39 40 41@dataclass42class Molmo2CausalLMOutputWithPast(ModelOutput):43 """44 Base class for Molmo2 causal language model (or autoregressive) outputs.45 46 Args:47 loss (`torch.FloatTensor` of shape `(1,)`, *optional*, returned when `labels` is provided):48 Language modeling loss (for next-token prediction).49 logits (`torch.FloatTensor` of shape `(batch_size, sequence_length, config.vocab_size)`):50 Prediction scores of the language modeling head (scores for each vocabulary token before SoftMax).51 past_key_values (`Cache`, *optional*, returned when `use_cache=True` is passed or when `config.use_cache=True`):52 It is a [`~cache_utils.Cache`] instance. For more details, see our [kv cache guide](https://huggingface.co/docs/transformers/en/kv_cache).53 54 Contains pre-computed hidden-states (key and values in the self-attention blocks) that can be used (see55 `past_key_values` input) to speed up sequential decoding.56 image_hidden_states (`torch.FloatTensor`, *optional*):57 A `torch.FloatTensor` of size `(batch_size, num_images, sequence_length, hidden_size)`.58 image_hidden_states of the model produced by the vision encoder and after projecting the last hidden state.59 """60 61 loss: Optional[torch.FloatTensor] = None62 logits: Optional[torch.FloatTensor] = None63 past_key_values: Optional[Cache] = None64 hidden_states: Optional[tuple[torch.FloatTensor]] = None65 attentions: Optional[tuple[torch.FloatTensor]] = None66 image_hidden_states: Optional[torch.FloatTensor] = None67 68 69@dataclass70class Molmo2ModelOutputWithPast(BaseModelOutputWithPast):71 """72 Base class for Molmo2 outputs, with hidden states and attentions.73 74 Args:75 image_hidden_states (`torch.FloatTensor`, *optional*):76 A `torch.FloatTensor` of size `(batch_num_patches, hidden_size)`.77 image_hidden_states of the model produced by the vision backbone78 """79 last_hidden_state: Optional[torch.FloatTensor] = None80 past_key_values: Optional[Cache] = None81 hidden_states: Optional[tuple[torch.FloatTensor]] = None82 attentions: Optional[tuple[torch.FloatTensor]] = None83 image_hidden_states: Optional[torch.FloatTensor] = None84 85 86class ViTMLP(nn.Module):87 def __init__(self, dim: int, hidden_dim: int, hidden_act: str, device: Union[str, torch.device] = None):88 super().__init__()89 self.w1 = nn.Linear(dim, hidden_dim, bias=True, device=device)90 self.act = ACT2FN[hidden_act]91 self.w2 = nn.Linear(hidden_dim, dim, bias=True, device=device)92 93 def forward(self, x: torch.Tensor) -> torch.Tensor:94 return self.w2(self.act(self.w1(x)))95 96 97class ViTMultiHeadDotProductAttention(nn.Module):98 def __init__(99 self,100 hidden_size: int,101 num_heads: int,102 num_key_value_heads: int,103 head_dim: int,104 use_bias: bool = True,105 input_dim: Optional[int] = None,106 float32_attention: bool = True,107 attention_dropout: float = 0.0,108 residual_dropout: float = 0.0,109 device: Union[str, torch.device] = None,110 attn_implementation: str = "eager",111 ):112 super().__init__()113 114 self.hidden_size = hidden_size115 self.num_heads = num_heads116 self.head_dim = head_dim117 self.num_key_value_heads = num_key_value_heads118 self.num_key_value_groups = self.num_heads // self.num_key_value_heads119 self.attn_implementation = attn_implementation120 self.is_causal = False121 122 input_dim = input_dim or hidden_size123 124 self.wq = nn.Linear(125 input_dim,126 self.num_heads * self.head_dim,127 bias=use_bias,128 device=device,129 )130 self.wk = nn.Linear(131 input_dim,132 self.num_key_value_heads * self.head_dim,133 bias=use_bias,134 device=device,135 )136 self.wv = nn.Linear(137 input_dim,138 self.num_key_value_heads * self.head_dim,139 bias=use_bias,140 device=device,141 )142 self.wo = nn.Linear(143 self.num_heads * self.head_dim,144 self.hidden_size,145 )146 self.float32_attention = float32_attention147 self.attention_dropout = attention_dropout148 self.residual_dropout = nn.Dropout(residual_dropout)149 150 def _split_heads(self, hidden_states, num_heads) -> torch.Tensor:151 return hidden_states.reshape(hidden_states.shape[:2] + (num_heads, self.head_dim))152 153 def _merge_heads(self, hidden_states) -> torch.Tensor:154 return hidden_states.reshape(hidden_states.shape[:2] + (self.hidden_size,))155 156 def forward(157 self,158 inputs_q: torch.Tensor,159 inputs_kv: Optional[torch.Tensor] = None,160 attn_mask: Optional[torch.Tensor] = None,161 ) -> torch.Tensor:162 163 if inputs_kv is not None:164 inputs_k = inputs_kv165 inputs_v = inputs_kv166 else:167 inputs_k = inputs_q168 inputs_v = inputs_q169 170 xq, xk, xv = self.wq(inputs_q), self.wk(inputs_k), self.wv(inputs_v)171 172 xq = self._split_heads(xq, self.num_heads)173 xk = self._split_heads(xk, self.num_key_value_heads)174 xv = self._split_heads(xv, self.num_key_value_heads)175 176 if self.num_heads != self.num_key_value_heads:177 xk = xk.repeat_interleave(self.num_key_value_groups, dim=2, output_size=self.num_heads)178 xv = xv.repeat_interleave(self.num_key_value_groups, dim=2, output_size=self.num_heads)179 180 og_dtype = xq.dtype181 182 if self.float32_attention:183 xq = xq.to(torch.float)184 xk = xk.to(torch.float)185 186 dropout_p = 0.0 if not self.training else self.attention_dropout187 188 if self.attn_implementation == "eager":189 attn_weights = torch.einsum("...qhd,...khd->...hqk", xq / math.sqrt(xq.size(-1)), xk)190 attn_weights = F.softmax(attn_weights, dim=-1, dtype=torch.float32).to(xq.dtype)191 attn_weights = F.dropout(192 attn_weights,193 p=dropout_p,194 training=self.training195 )196 attn_output = torch.einsum("...hqk,...khd->...qhd", attn_weights.to(xv.dtype), xv)197 198 elif self.attn_implementation == "sdpa":199 if not torch.is_autocast_enabled():200 xv = xv.to(torch.float)201 202 attn_output = F.scaled_dot_product_attention(203 xq.transpose(1, 2).contiguous(),204 xk.transpose(1, 2).contiguous(),205 xv.transpose(1, 2).contiguous(),206 attn_mask=attn_mask,207 is_causal=False,208 dropout_p=dropout_p,209 ).transpose(1, 2)210 211 elif self.attn_implementation == "flash_attention_2":212 if xq.dtype == torch.float32:213 if torch.is_autocast_enabled():214 target_dtype = torch.get_autocast_gpu_dtype()215 else:216 target_dtype = self.wq.weight.dtype217 attn_output = _flash_attention_forward(218 xq,219 xk,220 xv,221 attention_mask=attn_mask,222 query_length=inputs_q.shape[1],223 is_causal=False,224 dropout=dropout_p,225 softmax_scale=xq.shape[-1] ** -0.5,226 use_top_left_mask=flash_attn_supports_top_left_mask(),227 target_dtype=target_dtype,228 implementation=self.attn_implementation,229 )230 else:231 raise ValueError(f"Attention implementation {self.attn_implementation} not supported")232 233 attn_output = attn_output.to(og_dtype)234 attn_output = self._merge_heads(attn_output)235 attn_output = self.wo(attn_output)236 attn_output = self.residual_dropout(attn_output)237 238 return attn_output239 240 241class Molmo2VisionBlock(nn.Module):242 243 def __init__(self, config: Molmo2VitConfig, device: Union[str, torch.device] = None):244 super().__init__()245 self.attention = ViTMultiHeadDotProductAttention(246 hidden_size=config.hidden_size,247 num_heads=config.num_attention_heads,248 num_key_value_heads=config.num_key_value_heads,249 head_dim=config.head_dim,250 float32_attention=config.float32_attention,251 attention_dropout=config.attention_dropout,252 residual_dropout=config.residual_dropout,253 device=device,254 attn_implementation=config._attn_implementation,255 )256 self.feed_forward = ViTMLP(config.hidden_size, config.intermediate_size, config.hidden_act, device=device)257 self.attention_norm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps, device=device)258 self.ffn_norm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps, device=device)259 260 def forward(self, x: torch.Tensor) -> torch.Tensor:261 x = x + self.attention(self.attention_norm(x))262 x = x + self.feed_forward(self.ffn_norm(x))263 return x264 265 266class Molmo2VisionBlockCollection(nn.Module):267 268 def __init__(self, config: Molmo2VitConfig, device: Union[str, torch.device] = None):269 super().__init__()270 self.conifg = config271 self.resblocks = nn.ModuleList([272 Molmo2VisionBlock(config, device) for _ in range(config.num_hidden_layers)273 ])274 275 def forward(self, x: torch.Tensor) -> list[torch.Tensor]:276 hidden_states = []277 for r in self.resblocks:278 x = r(x)279 hidden_states.append(x)280 return hidden_states281 282 283class Molmo2VisionTransformer(nn.Module):284 285 def __init__(self, config: Molmo2VitConfig, device: Union[str, torch.device] = None):286 super().__init__()287 self.config = config288 289 # positional embeddings290 self.scale = config.hidden_size ** -0.5291 self.num_prefix_tokens: int = 0 # no class embeddings292 self.positional_embedding = nn.Parameter(293 torch.zeros(config.image_num_pos, config.hidden_size, device=device),294 )295 296 image_patch_size = config.image_patch_size297 self.patch_embedding = nn.Linear(298 image_patch_size * image_patch_size * 3,299 config.hidden_size,300 bias=True,301 device=device,302 )303 304 self.transformer = Molmo2VisionBlockCollection(config, device)305 306 def add_pos_emb(self, x: torch.Tensor, patch_num: int) -> torch.Tensor:307 pos_emb = self.positional_embedding308 309 pos_emb = pos_emb.reshape(310 (int(math.sqrt(pos_emb.shape[0])), int(math.sqrt(pos_emb.shape[0])), pos_emb.shape[1])311 )312 313 (patch_num_0, patch_num_1) = patch_num314 315 if pos_emb.shape[0] != patch_num_0 or pos_emb.shape[1] != patch_num_1:316 # Dervied from https://github.com/facebookresearch/mae/blob/main/util/pos_embed.py317 # antialias: default True in jax.image.resize318 pos_emb = pos_emb.unsqueeze(0).permute(0, 3, 1, 2)319 pos_emb = F.interpolate(320 pos_emb, size=(patch_num_0, patch_num_1), mode="bicubic", align_corners=False, antialias=True,321 )322 pos_emb = pos_emb.permute(0, 2, 3, 1).squeeze(0)323 324 pos_emb = pos_emb.reshape(-1, pos_emb.shape[-1])325 x = x + pos_emb[None, :, :].to(x.dtype)326 return x327 328 def forward(self, x: torch.Tensor, patch_num: int = None) -> list[torch.Tensor]:329 """330 : param x: (batch_size, num_patch, n_pixels)331 """332 if patch_num is None:333 patch_num = self.config.image_num_patch334 335 B, N, D = x.shape336 337 x = self.patch_embedding(x)338 339 # class embeddings and positional embeddings340 x = self.add_pos_emb(x, patch_num)341 342 hidden_states = self.transformer(x)343 return hidden_states344 345 346class ImageProjectorMLP(nn.Module):347 348 def __init__(349 self,350 input_dim: int,351 hidden_dim: int,352 output_dim: int,353 hidden_act: str,354 device: Union[str, torch.device] = None,355 ):356 super().__init__()357 self.w1 = nn.Linear(input_dim, hidden_dim, bias=False, device=device)358 self.w2 = nn.Linear(hidden_dim, output_dim, bias=False, device=device)359 self.w3 = nn.Linear(input_dim, hidden_dim, bias=False, device=device)360 self.act = ACT2FN[hidden_act]361 362 def forward(self, x: torch.Tensor) -> torch.Tensor:363 return self.w2(self.act(self.w1(x)) * self.w3(x))364 365 366class Molmo2VisionBackbone(nn.Module):367 def __init__(self, vit_config: Molmo2VitConfig, adapter_config: Molmo2AdapterConfig):368 super().__init__()369 self.vit_config = vit_config370 self.adapter_config = adapter_config371 372 self.vit_layers = []373 for layer in adapter_config.vit_layers:374 if layer >= 0:375 self.vit_layers.append(layer)376 else:377 self.vit_layers.append(layer + vit_config.num_hidden_layers)378 379 last_layer_needed = max(self.vit_layers) + 1380 if last_layer_needed < vit_config.num_hidden_layers:381 new_vit_config = deepcopy(vit_config)382 new_vit_config.num_hidden_layers = last_layer_needed383 self.image_vit = Molmo2VisionTransformer(new_vit_config)384 else:385 self.image_vit = Molmo2VisionTransformer(vit_config)386 387 self.num_prefix_tokens: int = self.image_vit.num_prefix_tokens388 389 pool_dim = vit_config.hidden_size * len(adapter_config.vit_layers)390 self.image_pooling_2d = ViTMultiHeadDotProductAttention(391 hidden_size=adapter_config.hidden_size,392 num_heads=adapter_config.num_attention_heads,393 num_key_value_heads=adapter_config.num_key_value_heads,394 head_dim=adapter_config.head_dim,395 input_dim=pool_dim,396 float32_attention=adapter_config.float32_attention,397 attention_dropout=adapter_config.attention_dropout,398 residual_dropout=adapter_config.residual_dropout,399 attn_implementation=adapter_config._attn_implementation,400 )401 self.image_projector = ImageProjectorMLP(402 adapter_config.hidden_size,403 adapter_config.intermediate_size,404 adapter_config.text_hidden_size,405 adapter_config.hidden_act,406 )407 self.image_feature_dropout = nn.Dropout(adapter_config.image_feature_dropout)408 409 def encode_image(self, images: torch.Tensor) -> torch.Tensor:410 """411 : param images: (batch_size, num_crops, num_patch, n_pixels)412 """413 B, T, N, D = images.shape414 images = images.view(B * T, N, D)415 image_features = self.image_vit(images)416 417 features = []418 for layer in self.vit_layers:419 features.append(image_features[layer])420 image_features = torch.cat(features, dim=-1)421 422 if self.num_prefix_tokens > 0:423 image_features = image_features[:, 1:]424 image_features = image_features.view(B, T, N, -1)425 return image_features426 427 @property428 def dtype(self) -> torch.dtype:429 return self.image_vit.patch_embedding.weight.dtype430 431 @property432 def device(self) -> torch.device:433 return self.image_vit.patch_embedding.weight.device434 435 def forward(436 self,437 images: torch.Tensor,438 pooled_patches_idx: torch.Tensor,439 ) -> tuple[torch.Tensor, Optional[torch.Tensor]]:440 441 # image_features: (batch_size, num_crops(=num_image), num_patch, nximage_emb_dim)442 batch_size, num_image = images.shape[:2]443 images = images.to(device=self.device, dtype=self.dtype)444 image_features = self.encode_image(images)445 446 image_features = self.image_feature_dropout(image_features)447 dim = image_features.shape[-1]448 valid = pooled_patches_idx >= 0449 valid_token = torch.any(valid, -1)450 451 # Use `pooled_patches_idx` to arange the features for image pooling452 batch_idx = torch.arange(pooled_patches_idx.shape[0], dtype=torch.long, device=pooled_patches_idx.device)453 batch_idx = torch.tile(batch_idx.view(batch_size, 1, 1), [1, pooled_patches_idx.shape[1], pooled_patches_idx.shape[2]])454 455 # Now [batch, num_high_res_features, pool_dim, dim]456 to_pool = image_features.reshape(batch_size, -1, dim)[batch_idx, torch.clip(pooled_patches_idx, 0)]457 to_pool = to_pool * valid.to(self.dtype)[:, :, :, None]458 to_pool = to_pool.reshape([-1, pooled_patches_idx.shape[-1], dim])459 if self.adapter_config.pooling_attention_mask:460 attn_mask = valid.reshape([-1, 1, 1, valid.shape[-1]])461 denom = valid.view(-1, to_pool.shape[-2]).float().sum(-1)462 denom = torch.where(denom == 0, 1, denom)463 query = to_pool.sum(-2, keepdim=True) / denom[:, None, None].to(to_pool.dtype)464 else:465 attn_mask = None466 query = to_pool.mean(-2, keepdim=True)467 pooled_features = self.image_pooling_2d(query, to_pool, attn_mask=attn_mask)468 pooled_features = pooled_features.reshape([batch_size, -1, pooled_features.shape[-1]])469 470 # MLP layer to map the feature.471 pooled_features = self.image_projector(pooled_features)472 return pooled_features.view(-1, pooled_features.shape[-1])[valid_token.flatten()]473 474 475# Copied from transformers.models.llama.modeling_llama.rotate_half476def rotate_half(x):477 """Rotates half the hidden dims of the input."""478 x1 = x[..., : x.shape[-1] // 2]479 x2 = x[..., x.shape[-1] // 2 :]480 return torch.cat((-x2, x1), dim=-1)481 482 483# Copied from transformers.models.llama.modeling_llama.apply_rotary_pos_emb484def apply_rotary_pos_emb(q, k, cos, sin, position_ids=None, unsqueeze_dim=1):485 """Applies Rotary Position Embedding to the query and key tensors.486 487 Args:488 q (`torch.Tensor`): The query tensor.489 k (`torch.Tensor`): The key tensor.490 cos (`torch.Tensor`): The cosine part of the rotary embedding.491 sin (`torch.Tensor`): The sine part of the rotary embedding.492 position_ids (`torch.Tensor`, *optional*):493 Deprecated and unused.494 unsqueeze_dim (`int`, *optional*, defaults to 1):495 The 'unsqueeze_dim' argument specifies the dimension along which to unsqueeze cos[position_ids] and496 sin[position_ids] so that they can be properly broadcasted to the dimensions of q and k. For example, note497 that cos[position_ids] and sin[position_ids] have the shape [batch_size, seq_len, head_dim]. Then, if q and498 k have the shape [batch_size, heads, seq_len, head_dim], then setting unsqueeze_dim=1 makes499 cos[position_ids] and sin[position_ids] broadcastable to the shapes of q and k. Similarly, if q and k have500 the shape [batch_size, seq_len, heads, head_dim], then set unsqueeze_dim=2.501 Returns:502 `tuple(torch.Tensor)` comprising of the query and key tensors rotated using the Rotary Position Embedding.503 """504 cos = cos.unsqueeze(unsqueeze_dim)505 sin = sin.unsqueeze(unsqueeze_dim)506 q_embed = (q * cos) + (rotate_half(q) * sin)507 k_embed = (k * cos) + (rotate_half(k) * sin)508 return q_embed, k_embed509 510 511class Molmo2RotaryEmbedding(nn.Module):512 inv_freq: torch.Tensor # fix linting for `register_buffer`513 514 def __init__(515 self,516 config: Molmo2TextConfig,517 device: Union[str, torch.device] = None,518 rope_type: Optional[str] = None,519 ):520 super().__init__()521 if rope_type is not None:522 self.rope_type = rope_type523 elif hasattr(config, "rope_scaling") and isinstance(config.rope_scaling, dict):524 # BC: "rope_type" was originally "type"525 self.rope_type = config.rope_scaling.get("rope_type", config.rope_scaling.get("type"))526 else:527 self.rope_type = "default"528 self.max_seq_len_cached = config.max_position_embeddings529 self.original_max_seq_len = config.max_position_embeddings530 531 self.config = config532 self.rope_init_fn = ROPE_INIT_FUNCTIONS[self.rope_type]533 534 inv_freq, self.attention_scaling = self.rope_init_fn(self.config, device)535 self.register_buffer("inv_freq", inv_freq, persistent=False)536 self.original_inv_freq = self.inv_freq537 538 @torch.no_grad()539 @dynamic_rope_update # power user: used with advanced RoPE types (e.g. dynamic rope)540 def forward(self, x, position_ids: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:541 inv_freq_expanded = self.inv_freq[None, :, None].float().expand(position_ids.shape[0], -1, 1).to(x.device)542 position_ids_expanded = position_ids[:, None, :].float()543 544 device_type = x.device.type if isinstance(x.device.type, str) and x.device.type != "mps" else "cpu"545 with torch.autocast(device_type=device_type, enabled=False): # Force float32546 freqs = (inv_freq_expanded.float() @ position_ids_expanded.float()).transpose(1, 2)547 emb = torch.cat((freqs, freqs), dim=-1)548 cos = emb.cos() * self.attention_scaling549 sin = emb.sin() * self.attention_scaling550 551 return cos.to(dtype=x.dtype), sin.to(dtype=x.dtype)552 553 554class Molmo2RMSNorm(nn.Module):555 556 def __init__(557 self,558 size: int,559 eps: float = 1e-6,560 device: Union[str, torch.device] = None,561 ):562 super().__init__()563 self.weight = nn.Parameter(torch.ones(size, device=device))564 self.eps = eps565 566 def forward(self, x: torch.Tensor) -> torch.Tensor:567 with torch.autocast(enabled=False, device_type=x.device.type):568 og_dtype = x.dtype569 x = x.to(torch.float32)570 variance = x.pow(2).mean(-1, keepdim=True)571 x = x * torch.rsqrt(variance + self.eps)572 x = x.to(og_dtype)573 574 return self.weight * x575 576 def extra_repr(self):577 return f"{tuple(self.weight.shape)}, eps={self.eps}"578 579 580# Copied from transformers.models.llama.modeling_llama.repeat_kv581def repeat_kv(hidden_states: torch.Tensor, n_rep: int) -> torch.Tensor:582 """583 This is the equivalent of torch.repeat_interleave(x, dim=1, repeats=n_rep). The hidden states go from (batch,584 num_key_value_heads, seqlen, head_dim) to (batch, num_attention_heads, seqlen, head_dim)585 """586 batch, num_key_value_heads, slen, head_dim = hidden_states.shape587 if n_rep == 1:588 return hidden_states589 hidden_states = hidden_states[:, :, None, :, :].expand(batch, num_key_value_heads, n_rep, slen, head_dim)590 return hidden_states.reshape(batch, num_key_value_heads * n_rep, slen, head_dim)591 592 593def eager_attention_forward(594 module: nn.Module,595 query: torch.Tensor,596 key: torch.Tensor,597 value: torch.Tensor,598 attention_mask: Optional[torch.Tensor],599 scaling: float,600 dropout: float = 0.0,601 **kwargs,602) -> tuple[torch.Tensor, Optional[torch.Tensor]]:603 key_states = repeat_kv(key, module.num_key_value_groups)604 value_states = repeat_kv(value, module.num_key_value_groups)605 606 attn_weights = torch.matmul(query, key_states.transpose(2, 3)) * scaling607 if attention_mask is not None:608 causal_mask = attention_mask[:, :, :, : key_states.shape[-2]]609 attn_weights = attn_weights + causal_mask610 611 attn_weights = nn.functional.softmax(attn_weights, dim=-1, dtype=torch.float32).to(query.dtype)612 attn_weights = nn.functional.dropout(attn_weights, p=dropout, training=module.training)613 attn_output = torch.matmul(attn_weights, value_states)614 attn_output = attn_output.transpose(1, 2).contiguous()615 616 return attn_output, attn_weights617 618 619class Molmo2Attention(nn.Module):620 """Multi-headed attention from 'Attention Is All You Need' paper"""621 622 def __init__(self, config: Molmo2TextConfig, layer_idx: int) -> None:623 super().__init__()624 self.config = config625 self.layer_idx = layer_idx626 self.num_heads = config.num_attention_heads627 self.num_key_value_heads = config.num_key_value_heads628 self.num_key_value_groups = config.num_attention_heads // config.num_key_value_heads629 self.head_dim = config.head_dim630 self.scaling = self.head_dim**-0.5631 self.is_causal = True632 633 self.fused_dims = (634 config.num_attention_heads * config.head_dim,635 config.head_dim * config.num_key_value_heads,636 config.head_dim * config.num_key_value_heads,637 )638 self.att_proj = nn.Linear(639 config.hidden_size,640 sum(self.fused_dims),641 bias=config.qkv_bias,642 )643 644 # Layer norms.645 self.k_norm: Optional[Molmo2RMSNorm] = None646 self.q_norm: Optional[Molmo2RMSNorm] = None647 self.qk_norm_type: Optional[str] = None648 if config.use_qk_norm:649 k_norm_size = (650 config.head_dim651 if config.qk_norm_type == "qwen3" else652 config.num_key_value_heads * config.head_dim653 )654 self.k_norm = Molmo2RMSNorm(k_norm_size, eps=config.layer_norm_eps)655 q_norm_size = (656 config.head_dim657 if config.qk_norm_type == "qwen3" else658 config.num_attention_heads * config.head_dim659 )660 self.q_norm = Molmo2RMSNorm(q_norm_size, eps=config.layer_norm_eps)661 self.qk_norm_type = config.qk_norm_type662 663 self.attention_dropout = config.attention_dropout664 665 self.attn_out = nn.Linear(666 config.head_dim * config.num_attention_heads,667 config.hidden_size,668 bias=False,669 )670 671 def forward(672 self,673 hidden_states: torch.Tensor,674 position_embeddings: tuple[torch.Tensor, torch.Tensor],675 attention_mask: Optional[torch.Tensor],676 past_key_values: Optional[Cache] = None,677 cache_position: Optional[torch.LongTensor] = None,678 **kwargs: Unpack[FlashAttentionKwargs],679 ) -> tuple[torch.Tensor, Optional[torch.Tensor], Optional[tuple[torch.Tensor]]]:680 input_shape = hidden_states.shape[:-1]681 hidden_shape = (*input_shape, -1, self.head_dim)682 683 qkv = self.att_proj(hidden_states)684 query_states, key_states, value_states = qkv.split(self.fused_dims, dim=-1)685 value_states = value_states.view(hidden_shape)686 687 # Optionally apply layer norm to keys and queries.688 if self.q_norm is not None and self.k_norm is not None and self.qk_norm_type != "qwen3":689 query_states = self.q_norm(query_states)690 key_states = self.k_norm(key_states)691 692 query_states = query_states.view(hidden_shape)693 key_states = key_states.view(hidden_shape)694 if self.q_norm is not None and self.k_norm is not None and self.qk_norm_type == "qwen3":695 query_states = self.q_norm(query_states)696 key_states = self.k_norm(key_states)697 query_states = query_states.transpose(1, 2)698 key_states = key_states.transpose(1, 2)699 value_states = value_states.transpose(1, 2)700 701 cos, sin = position_embeddings702 query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin)703 704 if past_key_values is not None: 705 # sin and cos are specific to RoPE models; cache_position needed for the static cache706 cache_kwargs = {"sin": sin, "cos": cos, "cache_position": cache_position}707 key_states, value_states = past_key_values.update(key_states, value_states, self.layer_idx, cache_kwargs)708 709 attention_interface: Callable = eager_attention_forward710 if self.config._attn_implementation != "eager":711 attention_interface = ALL_ATTENTION_FUNCTIONS[self.config._attn_implementation]712 713 attn_output, attn_weights = attention_interface(714 self,715 query_states,716 key_states,717 value_states,718 attention_mask,719 dropout=0.0 if not self.training else self.attention_dropout,720 scaling=self.scaling,721 **kwargs,722 )723 724 attn_output = attn_output.reshape(*input_shape, -1).contiguous()725 attn_output = self.attn_out(attn_output)726 return attn_output, attn_weights727 728 729class LanguageModelMLP(nn.Module):730 731 def __init__(732 self,733 input_dim: int,734 intermediate_size: int,735 hidden_act: str,736 device: Union[str, torch.device] = None,737 ):738 super().__init__()739 self.ff_proj = nn.Linear(input_dim, intermediate_size * 2, bias=False, device=device)740 self.ff_out = nn.Linear(intermediate_size, input_dim, bias=False, device=device)741 self.act = ACT2FN[hidden_act]742 743 def forward(self, x: torch.Tensor) -> torch.Tensor:744 x = self.ff_proj(x)745 x, gate = x.chunk(2, dim=-1)746 x = self.act(gate) * x747 x = self.ff_out(x)748 return x749 750 751class Molmo2DecoderLayer(GradientCheckpointingLayer):752 753 def __init__(754 self,755 config: Molmo2TextConfig,756 layer_idx: Optional[int] = None,757 device: Union[str, torch.device] = None758 ):759 super().__init__()760 self.config = config761 762 self.self_attn = Molmo2Attention(config, layer_idx)763 self.attn_norm = Molmo2RMSNorm(764 config.hidden_size, eps=config.layer_norm_eps, device=device)765 self.dropout = nn.Dropout(config.residual_dropout)766 self.mlp = LanguageModelMLP(767 config.hidden_size, config.intermediate_size, config.hidden_act, device=device)768 self.ff_norm = Molmo2RMSNorm(769 config.hidden_size, eps=config.layer_norm_eps, device=device)770 771 def forward(772 self,773 hidden_states: torch.Tensor,774 position_embeddings: tuple[torch.Tensor, torch.Tensor],775 attention_mask: Optional[torch.Tensor] = None,776 position_ids: Optional[torch.LongTensor] = None,777 past_key_values: Optional[Cache] = None,778 output_attentions: Optional[bool] = False,779 use_cache: Optional[bool] = False,780 cache_position: Optional[torch.LongTensor] = None,781 **kwargs: Unpack[TransformersKwargs],782 ) -> tuple[torch.FloatTensor, Optional[tuple[torch.FloatTensor, torch.FloatTensor]]]:783 784 residual = hidden_states785 hidden_states = self.attn_norm(hidden_states)786 787 # Self Attention788 hidden_states, self_attn_weights = self.self_attn(789 hidden_states=hidden_states,790 position_embeddings=position_embeddings,791 attention_mask=attention_mask,792 position_ids=position_ids,793 past_key_values=past_key_values,794 output_attentions=output_attentions,795 use_cache=use_cache,796 cache_position=cache_position,797 **kwargs,798 )799 800 hidden_states = residual + self.dropout(hidden_states)801 802 # Fully Connected803 residual = hidden_states804 hidden_states = self.ff_norm(hidden_states)805 hidden_states = self.mlp(hidden_states)806 807 hidden_states = residual + self.dropout(hidden_states)808 809 outputs = (hidden_states,)810 811 if output_attentions:812 outputs += (self_attn_weights,)813 814 return outputs815 816 817class Molmo2PostNormDecoderLayer(Molmo2DecoderLayer):818 def forward(819 self,820 hidden_states: torch.Tensor,821 position_embeddings: tuple[torch.Tensor, torch.Tensor],822 attention_mask: Optional[torch.Tensor] = None,823 position_ids: Optional[torch.LongTensor] = None,824 past_key_values: Optional[Cache] = None,825 output_attentions: Optional[bool] = False,826 use_cache: Optional[bool] = False,827 cache_position: Optional[torch.LongTensor] = None,828 **kwargs,829 ) -> tuple[torch.FloatTensor, Optional[tuple[torch.FloatTensor, torch.FloatTensor]]]:830 831 residual = hidden_states832 833 # Self Attention834 hidden_states, self_attn_weights = self.self_attn(835 hidden_states=hidden_states,836 position_embeddings=position_embeddings,837 attention_mask=attention_mask,838 position_ids=position_ids,839 past_key_values=past_key_values,840 output_attentions=output_attentions,841 use_cache=use_cache,842 cache_position=cache_position,843 )844 hidden_states = self.attn_norm(hidden_states)845 846 hidden_states = residual + self.dropout(hidden_states)847 848 # Fully Connected849 residual = hidden_states850 hidden_states = self.mlp(hidden_states)851 hidden_states = self.ff_norm(hidden_states)852 853 hidden_states = residual + self.dropout(hidden_states)854 855 outputs = (hidden_states,)856 857 if output_attentions:858 outputs += (self_attn_weights,)859 860 return outputs861 862 863class Molmo2Embedding(nn.Module):864 def __init__(865 self,866 num_embeddings: int,867 num_new_embeddings: int,868 features: int,869 device: Union[str, torch.device] = None,870 ):871 super().__init__()872 self.embedding = nn.Parameter(873 torch.zeros(num_embeddings, features, device=device),874 )875 self.new_embedding = nn.Parameter(876 torch.zeros(num_new_embeddings, features, device=device),877 )878 879 def forward(self, x: torch.Tensor) -> torch.Tensor:880 return F.embedding(x, torch.cat([self.embedding, self.new_embedding], dim=0))881 882 883class Molmo2PreTrainedModel(PreTrainedModel):884 config: Molmo2Config885 base_model_prefix = "model"886 supports_gradient_checkpointing = True887 _no_split_modules = [888 "Molmo2DecoderLayer",889 "Molmo2PostNormDecoderLayer",890 "Molmo2VisionBlock",891 "ViTMultiHeadDotProductAttention",892 ]893 _skip_keys_device_placement = "past_key_values"894 _supports_flash_attn = True895 _supports_sdpa = True896 897 _can_compile_fullgraph = True898 _supports_attention_backend = True899 _can_record_outputs = {900 "hidden_states": Molmo2DecoderLayer,901 "attentions": Molmo2Attention,902 }903 904 def _init_weights(self, module):905 std = self.config.initializer_range906 if isinstance(module, (nn.Linear,)):907 module.weight.data.normal_(mean=0.0, std=std)908 if module.bias is not None:909 module.bias.data.zero_()910 elif isinstance(module, Molmo2Embedding):911 module.embedding.data.normal_(mean=0.0, std=std)912 module.new_embedding.data.normal_(mean=0.0, std=std)913 elif isinstance(module, nn.Embedding):914 module.weight.data.normal_(mean=0.0, std=std)915 if module.padding_idx is not None:916 module.weight.data[module.padding_idx].zero_()917 elif isinstance(module, Molmo2RMSNorm):918 module.weight.data.fill_(1.0)919 elif isinstance(module, nn.LayerNorm):920 module.weight.data.fill_(1.0)921 if module.bias is not None:922 module.bias.data.zero_()923 924 925class Molmo2TextModel(Molmo2PreTrainedModel):926 config: Molmo2TextConfig927 _no_split_modules = ["Molmo2DecoderLayer", "Molmo2PostNormDecoderLayer"]928 929 def __init__(self, config: Molmo2TextConfig):930 super().__init__(config)931 if config.additional_vocab_size is not None:932 self.wte = Molmo2Embedding(933 config.vocab_size,934 config.additional_vocab_size,935 config.hidden_size,936 )937 else:938 self.wte = nn.Embedding(config.vocab_size, config.hidden_size)939 self.emb_drop = nn.Dropout(config.embedding_dropout)940 decoder_layer = Molmo2PostNormDecoderLayer if config.norm_after else Molmo2DecoderLayer941 self.blocks = nn.ModuleList(942 [decoder_layer(config, layer_idx) for layer_idx in range(config.num_hidden_layers)]943 )944 self.ln_f = Molmo2RMSNorm(config.hidden_size, eps=config.layer_norm_eps)945 if config.rope_scaling_layers is not None:946 self.rotary_embs = nn.ModuleDict(947 {948 "default": Molmo2RotaryEmbedding(config, rope_type="default"),949 "scaling": Molmo2RotaryEmbedding(config),950 }951 )952 else:953 self.rotary_emb = Molmo2RotaryEmbedding(config)954 self.gradient_checkpointing = False955 956 # Initialize weights and apply final processing957 self.post_init()958 959 def get_input_embeddings(self) -> torch.nn.Module:960 return self.wte961 962 def set_input_embeddings(self, value: torch.nn.Module) -> None:963 self.wte = value964 965 @can_return_tuple966 def forward(967 self,968 input_ids: Optional[torch.LongTensor] = None,969 attention_mask: Optional[torch.Tensor] = None,970 position_ids: Optional[torch.LongTensor] = None,971 past_key_values: Optional[Cache] = None,972 inputs_embeds: Optional[torch.FloatTensor] = None,973 use_cache: Optional[bool] = None,974 output_attentions: Optional[bool] = None,975 output_hidden_states: Optional[bool] = None,976 cache_position: Optional[torch.LongTensor] = None,977 **kwargs: Unpack[TransformersKwargs],978 ) -> BaseModelOutputWithPast:979 output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions980 output_hidden_states = (981 output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states982 )983 use_cache = use_cache if use_cache is not None else self.config.use_cache984 985 if (input_ids is None) ^ (inputs_embeds is not None):986 raise ValueError("You must specify exactly one of input_ids or inputs_embeds")987 988 if self.gradient_checkpointing and self.training and use_cache:989 logger.warning_once(990 "`use_cache=True` is incompatible with gradient checkpointing. Setting `use_cache=False`."991 )992 use_cache = False993 994 if inputs_embeds is None:995 input_ids = input_ids * (input_ids != -1).to(input_ids.dtype)996 inputs_embeds = self.wte(input_ids)997 998 # torch.jit.trace() doesn't support cache objects in the output999 if use_cache and past_key_values is None and not torch.jit.is_tracing():1000 past_key_values = DynamicCache(config=self.config)1001 1002 if cache_position is None:1003 past_seen_tokens = past_key_values.get_seq_length() if past_key_values is not None else 01004 cache_position = torch.arange(1005 past_seen_tokens,1006 past_seen_tokens + inputs_embeds.shape[1],1007 device=inputs_embeds.device,1008 )1009 1010 if position_ids is None:1011 position_ids = cache_position.unsqueeze(0)1012 1013 # It may already have been prepared by e.g. `generate`1014 if not isinstance(causal_mask_mapping := attention_mask, dict):1015 # Prepare mask arguments1016 mask_kwargs = {1017 "config": self.config,1018 "input_embeds": inputs_embeds,1019 "attention_mask": attention_mask,1020 "cache_position": cache_position,1021 "past_key_values": past_key_values,1022 "position_ids": position_ids,1023 }1024 1025 # Create the mask1026 causal_mask_mapping = create_causal_mask(**mask_kwargs)1027 1028 hidden_states = inputs_embeds1029 1030 # create position embeddings to be shared across the decoder layers1031 if self.config.rope_scaling_layers is not None:1032 position_embeddings_mapping = {1033 "default": self.rotary_embs["default"](hidden_states, position_ids),1034 "scaling": self.rotary_embs["scaling"](hidden_states, position_ids),1035 }1036 else:1037 position_embeddings = self.rotary_emb(hidden_states, position_ids)1038 1039 # decoder layers1040 all_hidden_states = () if output_hidden_states else None1041 all_self_attns = () if output_attentions else None1042 1043 for layer_idx, decoder_block in enumerate(self.blocks[: self.config.num_hidden_layers]):1044 if output_hidden_states:1045 all_hidden_states += (hidden_states,)1046 1047 if self.config.rope_scaling_layers is not None:1048 position_embeddings_i = (1049 position_embeddings_mapping["scaling"]1050 if layer_idx in self.config.rope_scaling_layers1051 else position_embeddings_mapping["default"]1052 )1053 else:1054 position_embeddings_i = position_embeddings1055 1056 layer_outputs = decoder_block(1057 hidden_states,1058 attention_mask=causal_mask_mapping,1059 position_ids=position_ids,1060 past_key_values=past_key_values,1061 output_attentions=output_attentions,1062 use_cache=use_cache,1063 cache_position=cache_position,1064 position_embeddings=position_embeddings_i,1065 **kwargs,1066 )1067 1068 hidden_states = layer_outputs[0]1069 1070 if output_attentions:1071 all_self_attns += (layer_outputs[1],)1072 1073 hidden_states = self.ln_f(hidden_states)1074 1075 # add hidden states from the last decoder layer1076 if output_hidden_states:1077 all_hidden_states += (hidden_states,)1078 1079 return BaseModelOutputWithPast(1080 last_hidden_state=hidden_states,1081 past_key_values=past_key_values,1082 hidden_states=all_hidden_states,1083 attentions=all_self_attns,1084 )1085 1086# Adapted from transformers.models.gemma3.modeling_gemma31087def token_type_ids_mask_function(1088 token_type_ids: Optional[torch.Tensor] = None,1089) -> Optional[Callable]:1090 """1091 This function adds the correct offsets to the `q_idx` and `kv_idx` as the torch API can only accept lengths,1092 not start and end indices.1093 """1094 # Do not return an additional mask in this case1095 if token_type_ids is None:1096 return None1097 1098 def inner_mask(batch_idx: int, head_idx: int, q_idx: int, kv_idx: int) -> bool:1099 # If it's 1 for both query and key/value, we are in an image block1100 # NOTE: static cache shape goes beyond input seq length, while token_type_ids.shape[1] == input seq length1101 # Since vmap doesn't support `if statement` we workaround it with `torch.where`1102 safe_idx = torch.where(kv_idx < token_type_ids.shape[1], kv_idx, 0)1103 token_type_ids_at_kv_idx = token_type_ids[batch_idx, safe_idx]1104 token_type_ids_at_kv_idx = torch.where(kv_idx < token_type_ids.shape[1], token_type_ids_at_kv_idx, 0)1105 1106 is_image_block = (token_type_ids[batch_idx, q_idx] == 1) & (token_type_ids_at_kv_idx == 1)1107 1108 # This is bidirectional attention whenever we are dealing with image tokens1109 return is_image_block & is_image_block1110 1111 return inner_mask1112 1113 1114class Molmo2Model(Molmo2PreTrainedModel):1115 base_model_prefix = ""1116 _checkpoint_conversion_mapping = {}1117 # Reference: fix gemma3 grad acc #372081118 accepts_loss_kwargs = False1119 config: Molmo2Config1120 1121 1122 def __init__(self, config: Molmo2Config):1123 super().__init__(config)1124 self.transformer: Molmo2TextModel = Molmo2TextModel(config.text_config)1125 self.vision_backbone: Optional[Molmo2VisionBackbone] = None1126 if config.vit_config is not None and config.adapter_config is not None:1127 self.vision_backbone = Molmo2VisionBackbone(config.vit_config, config.adapter_config)1128 1129 # Initialize weights and apply final processing1130 self.post_init()1131 1132 def get_input_embeddings(self) -> torch.nn.Module:1133 return self.transformer.wte1134 1135 def set_input_embeddings(self, value: torch.nn.Module) -> None:1136 self.transformer.wte = value1137 1138 def set_decoder(self, decoder):1139 self.transformer = decoder1140 1141 def get_decoder(self):1142 return self.transformer1143 1144 @property1145 def device(self) -> torch.device:1146 return self.transformer.ln_f.weight.device1147 1148 def build_batched_images(1149 self,1150 input_ids: torch.LongTensor,1151 pixel_values: torch.Tensor,1152 image_token_pooling: torch.Tensor,1153 image_grids: torch.Tensor,1154 image_num_crops: torch.Tensor,1155 ) -> tuple[torch.Tensor, torch.Tensor]:1156 # 1) Count the number of images in each example1157 raw_counts = (input_ids == self.config.image_end_token_id).sum(1) # [N]1158 # Each image is represented by global view and high-res view1159 # so we divide by 2 to get the number of images1160 counts = raw_counts // 21161 N = counts.size(0)1162 device = input_ids.device1163 1164 # Total number of images in the batch1165 num_images = int(counts.sum().item())1166 1167 # Sanity check1168 assert image_grids.size(0) == num_images, \1169 f"Expected {num_images} image grids, but got {image_grids.size(0)}"1170 assert image_num_crops.size(0) == num_images, \1171 f"Expected {num_images} image num crops, but got {image_num_crops.size(0)}"1172 1173 # 1-1) Compute per-image pooled patch count from image grids1174 with torch.no_grad():1175 first_prod = image_grids[:, :2].prod(dim=1) # [num_images]1176 second_prod = image_grids[:, 2:].prod(dim=1) # [num_images]1177 num_pooled_patches_per_image = (first_prod + second_prod).to(image_num_crops.dtype) # [num_images]1178 1179 # pixel_values: [n_crops, n_patches, pixels_per_patch]1180 n_crops, n_patches, pixels_per_patch = pixel_values.shape1181 1182 # 2) Map each image index → example index1183 # Example: if counts = [2, 1, 3], then this becomes [0,0,1,2,2,2]1184 example_ids_for_image = torch.arange(N, device=device).repeat_interleave(counts) # [num_images]1185 assert example_ids_for_image.numel() == num_images1186 1187 # 2-1) Compute crops_per_example by summing per-image crop counts1188 crops_per_example = torch.zeros(1189 N, dtype=image_num_crops.dtype, device=image_num_crops.device1190 )1191 crops_per_example.index_add_(0, example_ids_for_image, image_num_crops) # [N]1192 1193 # 2-2) Per-image number of patches = (crops per image) * n_patches1194 patches_per_image = image_num_crops * n_patches # [num_images]1195 1196 # 2-3) Compute per-example per-image patch offsets1197 counts_list = counts.tolist()1198 index_offset_per_example_list = []1199 offset_img = 01200 for c in counts_list: