internlm/Intern-S2-Mobius
63820
1# Copyright 2025 InternS2Mobius Team. All rights reserved.2#3# Licensed under the Apache License, Version 2.0 (the "License");4# you may not use this file except in compliance with the License.5# You may obtain a copy of the License at6#7# http://www.apache.org/licenses/LICENSE-2.08#9# Unless required by applicable law or agreed to in writing, software10# distributed under the License is distributed on an "AS IS" BASIS,11# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.12# See the License for the specific language governing permissions and13# limitations under the License.14 15import math16from collections.abc import Callable17from dataclasses import dataclass18from typing import Any, Optional19 20import torch21import torch.nn as nn22import torch.nn.functional as F23 24from transformers import initialization as init25from transformers.activations import ACT2FN26from transformers.cache_utils import Cache27from transformers.generation import GenerationMixin28from transformers.integrations import use_experts_implementation, use_kernelized_func29from transformers.masking_utils import create_causal_mask30from transformers.modeling_flash_attention_utils import FlashAttentionKwargs31from transformers.modeling_layers import GradientCheckpointingLayer32from transformers.modeling_outputs import BaseModelOutputWithPast, BaseModelOutputWithPooling, ModelOutput33from transformers.modeling_rope_utils import ROPE_INIT_FUNCTIONS, dynamic_rope_update34from transformers.modeling_utils import ALL_ATTENTION_FUNCTIONS, PreTrainedModel35from transformers.processing_utils import Unpack36from transformers.utils import TransformersKwargs, auto_docstring, can_return_tuple, logging, torch_compilable_check37from transformers.utils.generic import is_flash_attention_requested, maybe_autocast, merge_with_config_defaults38from transformers.utils.import_utils import is_causal_conv1d_available, is_flash_linear_attention_available39from transformers.utils.output_capturing import OutputRecorder, capture_outputs40 41from .configuration_interns2_mobius import (42 InternS2MobiusConfig,43 InternS2MobiusTextConfig,44 InternS2MobiusVisionConfig,45)46 47 48if is_causal_conv1d_available():49 from causal_conv1d import causal_conv1d_fn, causal_conv1d_update50else:51 causal_conv1d_update, causal_conv1d_fn = None, None52 53if is_flash_linear_attention_available():54 from fla.modules import FusedRMSNormGated55 from fla.ops.gated_delta_rule import chunk_gated_delta_rule, fused_recurrent_gated_delta_rule56else:57 chunk_gated_delta_rule, fused_recurrent_gated_delta_rule = None, None58 FusedRMSNormGated = None59 60logger = logging.get_logger(__name__)61 62 63class InternS2MobiusVisionRotaryEmbedding(nn.Module):64 inv_freq: torch.Tensor # fix linting for `register_buffer`65 66 def __init__(self, dim: int, theta: float = 10000.0) -> None:67 super().__init__()68 self.dim = dim69 self.theta = theta70 inv_freq = 1.0 / (theta ** (torch.arange(0, dim, 2, dtype=torch.float) / dim))71 self.register_buffer("inv_freq", inv_freq, persistent=False)72 73 def forward(self, seqlen: int) -> torch.Tensor:74 seq = torch.arange(seqlen, device=self.inv_freq.device, dtype=self.inv_freq.dtype)75 freqs = torch.outer(seq, self.inv_freq)76 return freqs77 78 79class InternS2MobiusDynamicCache:80 """81 A dynamic cache that can handle both the attention cache (which has a seq_len dimension) and the linear attention82 cache (which has a constant shape regardless of seq_len).83 84 This cache has two sets of lists of tensors: `key_cache` and `value_cache` for attention cache and `conv_states`85 and `ssm_states` for gated deltanet cache. Each of these lists has `num_layers` tensors. The expected shape for each tensor86 For attention layers, `key_cache` and `value_cache` have a shape of `(batch_size, num_heads, seq_len, head_dim)`,87 while `conv_states` and `ssm_states` have a shape of `(batch_size, 0)` (empty tensors).88 For linear attention layers, `key_cache` and `value_cache` have a shape of `(batch_size, 0)` (empty tensors),89 while `conv_states` represents the convolution state and has a shape of `(batch_size, d_inner, d_conv)`,90 and `recurrent_states` represents the recurrent state and has a shape of `(batch_size, d_inner, d_state)`.91 """92 93 is_compileable = False94 95 def __init__(self, config: InternS2MobiusConfig):96 super().__init__()97 self.layer_types = config.layer_types98 self.transformer_layers = [99 i for i in range(config.num_hidden_layers) if self.layer_types[i] == "full_attention"100 ]101 self.last_linear_layer = len(self.layer_types) - 1 - self.layer_types[::-1].index("linear_attention")102 103 # Initialize everything to None -> will be lazy initialized to allow multi-gpu (device_map) inference104 self.conv_states = [None for _ in range(config.num_hidden_layers)]105 self.recurrent_states = [None for _ in range(config.num_hidden_layers)]106 self.key_cache = [None for _ in range(config.num_hidden_layers)]107 self.value_cache = [None for _ in range(config.num_hidden_layers)]108 109 def __len__(self):110 return len(self.layer_types)111 112 def update(113 self,114 key_states: torch.Tensor,115 value_states: torch.Tensor,116 layer_idx: int,117 cache_kwargs: dict[str, Any] | None = None,118 ) -> tuple[torch.Tensor, torch.Tensor]:119 if self.key_cache[layer_idx] is None:120 self.key_cache[layer_idx] = key_states121 self.value_cache[layer_idx] = value_states122 else:123 self.key_cache[layer_idx] = torch.cat([self.key_cache[layer_idx], key_states], dim=2)124 self.value_cache[layer_idx] = torch.cat([self.value_cache[layer_idx], value_states], dim=2)125 126 return self.key_cache[layer_idx], self.value_cache[layer_idx]127 128 def reorder_cache(self, beam_idx: torch.LongTensor):129 """Reorders the cache for beam search, given the selected beam indices."""130 for layer_idx in range(len(self.key_cache)):131 if self.key_cache[layer_idx] is not None:132 device = self.key_cache[layer_idx].device133 beam_idx = beam_idx.to(device)134 self.key_cache[layer_idx] = self.key_cache[layer_idx].index_select(0, beam_idx)135 self.value_cache[layer_idx] = self.value_cache[layer_idx].index_select(0, beam_idx)136 137 if self.conv_states[layer_idx] is not None:138 device = self.conv_states[layer_idx].device139 beam_idx = beam_idx.to(device)140 self.conv_states[layer_idx] = self.conv_states[layer_idx].index_select(0, beam_idx)141 self.recurrent_states[layer_idx] = self.recurrent_states[layer_idx].index_select(0, beam_idx)142 143 def get_seq_length(self, layer_idx: int | None = 0) -> int:144 """Returns the sequence length of the cached states. A layer index can be optionally passed."""145 # take any layer that contains cache and not empty tensor146 layer_idx = self.transformer_layers[0] if layer_idx not in self.transformer_layers else layer_idx147 if len(self.key_cache) <= layer_idx or self.key_cache[layer_idx] is None:148 return 0149 return self.key_cache[layer_idx].shape[-2]150 151 def get_mask_sizes(self, cache_position: torch.Tensor, layer_idx: int) -> tuple[int, int]:152 """153 Return a tuple (kv_length, kv_offset) corresponding to the length and offset that will be returned for154 the given layer at `layer_idx`.155 The masks are then prepared according to the given lengths (kv_length, kv_offset) and patterns for each layer.156 """157 kv_offset = 0158 query_length = cache_position.shape[0]159 past_seen_tokens = self.get_seq_length(layer_idx)160 kv_length = query_length + past_seen_tokens161 return kv_length, kv_offset162 163 @property164 def has_previous_state(self):165 """We have a previous state if the last linear (conv) layer was already updated."""166 return self.conv_states[self.last_linear_layer] is not None167 168 169class InternS2MobiusRMSNormGated(nn.Module):170 def __init__(self, hidden_size, eps=1e-6, **kwargs):171 super().__init__()172 self.weight = nn.Parameter(torch.ones(hidden_size))173 self.variance_epsilon = eps174 175 def forward(self, hidden_states, gate=None):176 input_dtype = hidden_states.dtype177 hidden_states = hidden_states.to(torch.float32)178 variance = hidden_states.pow(2).mean(-1, keepdim=True)179 # Norm before gate180 hidden_states = hidden_states * torch.rsqrt(variance + self.variance_epsilon)181 hidden_states = self.weight * hidden_states.to(input_dtype)182 hidden_states = hidden_states * F.silu(gate.to(torch.float32))183 184 return hidden_states.to(input_dtype)185 186 187def apply_mask_to_padding_states(hidden_states, attention_mask):188 """189 Tunes out the hidden states for padding tokens, see https://github.com/state-spaces/mamba/issues/66190 """191 # NOTE: attention mask is a 2D boolean tensor192 if attention_mask is not None and attention_mask.shape[1] > 1 and attention_mask.shape[0] > 1:193 dtype = hidden_states.dtype194 hidden_states = (hidden_states * attention_mask[:, :, None]).to(dtype)195 196 return hidden_states197 198 199is_fast_path_available = all(200 (causal_conv1d_fn, causal_conv1d_update, chunk_gated_delta_rule, fused_recurrent_gated_delta_rule)201)202 203 204def torch_causal_conv1d_update(205 hidden_states,206 conv_state,207 weight,208 bias=None,209 activation=None,210):211 _, hidden_size, seq_len = hidden_states.shape212 state_len = conv_state.shape[-1]213 214 hidden_states_new = torch.cat([conv_state, hidden_states], dim=-1).to(weight.dtype)215 conv_state.copy_(hidden_states_new[:, :, -state_len:])216 out = F.conv1d(hidden_states_new, weight.unsqueeze(1), bias, padding=0, groups=hidden_size)217 out = F.silu(out[:, :, -seq_len:])218 out = out.to(hidden_states.dtype)219 return out220 221 222def l2norm(x: torch.FloatTensor, dim: int = -1, eps: float = 1e-6):223 """This function is intended to align with the l2norm implementation in the FLA library."""224 inv_norm = torch.rsqrt((x * x).sum(dim=dim, keepdim=True) + eps)225 return x * inv_norm226 227 228def torch_chunk_gated_delta_rule(229 query,230 key,231 value,232 g,233 beta,234 chunk_size=64,235 initial_state=None,236 output_final_state=False,237 use_qk_l2norm_in_kernel=False,238):239 initial_dtype = query.dtype240 if use_qk_l2norm_in_kernel:241 query = l2norm(query, dim=-1, eps=1e-6)242 key = l2norm(key, dim=-1, eps=1e-6)243 query, key, value, beta, g = [244 x.transpose(1, 2).contiguous().to(torch.float32) for x in (query, key, value, beta, g)245 ]246 247 batch_size, num_heads, sequence_length, k_head_dim = key.shape248 v_head_dim = value.shape[-1]249 pad_size = (chunk_size - sequence_length % chunk_size) % chunk_size250 query = F.pad(query, (0, 0, 0, pad_size))251 key = F.pad(key, (0, 0, 0, pad_size))252 value = F.pad(value, (0, 0, 0, pad_size))253 beta = F.pad(beta, (0, pad_size))254 g = F.pad(g, (0, pad_size))255 total_sequence_length = sequence_length + pad_size256 scale = 1 / (query.shape[-1] ** 0.5)257 query = query * scale258 259 v_beta = value * beta.unsqueeze(-1)260 k_beta = key * beta.unsqueeze(-1)261 # reshape to chunks262 query, key, value, k_beta, v_beta = [263 x.reshape(x.shape[0], x.shape[1], -1, chunk_size, x.shape[-1]) for x in (query, key, value, k_beta, v_beta)264 ]265 g = g.reshape(g.shape[0], g.shape[1], -1, chunk_size)266 mask = torch.triu(torch.ones(chunk_size, chunk_size, dtype=torch.bool, device=query.device), diagonal=0)267 268 # chunk decay269 g = g.cumsum(dim=-1)270 decay_mask = ((g.unsqueeze(-1) - g.unsqueeze(-2)).tril().exp().float()).tril()271 attn = -((k_beta @ key.transpose(-1, -2)) * decay_mask).masked_fill(mask, 0)272 for i in range(1, chunk_size):273 row = attn[..., i, :i].clone()274 sub = attn[..., :i, :i].clone()275 attn[..., i, :i] = row + (row.unsqueeze(-1) * sub).sum(-2)276 attn = attn + torch.eye(chunk_size, dtype=attn.dtype, device=attn.device)277 value = attn @ v_beta278 k_cumdecay = attn @ (k_beta * g.exp().unsqueeze(-1))279 last_recurrent_state = (280 torch.zeros(batch_size, num_heads, k_head_dim, v_head_dim).to(value)281 if initial_state is None282 else initial_state.to(value)283 )284 core_attn_out = torch.zeros_like(value)285 mask = torch.triu(torch.ones(chunk_size, chunk_size, dtype=torch.bool, device=query.device), diagonal=1)286 287 # for each chunk288 for i in range(0, total_sequence_length // chunk_size):289 q_i, k_i, v_i = query[:, :, i], key[:, :, i], value[:, :, i]290 attn = (q_i @ k_i.transpose(-1, -2) * decay_mask[:, :, i]).masked_fill_(mask, 0)291 v_prime = (k_cumdecay[:, :, i]) @ last_recurrent_state292 v_new = v_i - v_prime293 attn_inter = (q_i * g[:, :, i, :, None].exp()) @ last_recurrent_state294 core_attn_out[:, :, i] = attn_inter + attn @ v_new295 last_recurrent_state = (296 last_recurrent_state * g[:, :, i, -1, None, None].exp()297 + (k_i * (g[:, :, i, -1, None] - g[:, :, i]).exp()[..., None]).transpose(-1, -2) @ v_new298 )299 300 if not output_final_state:301 last_recurrent_state = None302 core_attn_out = core_attn_out.reshape(core_attn_out.shape[0], core_attn_out.shape[1], -1, core_attn_out.shape[-1])303 core_attn_out = core_attn_out[:, :, :sequence_length]304 core_attn_out = core_attn_out.transpose(1, 2).contiguous().to(initial_dtype)305 return core_attn_out, last_recurrent_state306 307 308def torch_recurrent_gated_delta_rule(309 query, key, value, g, beta, initial_state, output_final_state, use_qk_l2norm_in_kernel=False310):311 initial_dtype = query.dtype312 if use_qk_l2norm_in_kernel:313 query = l2norm(query, dim=-1, eps=1e-6)314 key = l2norm(key, dim=-1, eps=1e-6)315 query, key, value, beta, g = [316 x.transpose(1, 2).contiguous().to(torch.float32) for x in (query, key, value, beta, g)317 ]318 319 batch_size, num_heads, sequence_length, k_head_dim = key.shape320 v_head_dim = value.shape[-1]321 scale = 1 / (query.shape[-1] ** 0.5)322 query = query * scale323 324 core_attn_out = torch.zeros(batch_size, num_heads, sequence_length, v_head_dim).to(value)325 last_recurrent_state = (326 torch.zeros(batch_size, num_heads, k_head_dim, v_head_dim).to(value)327 if initial_state is None328 else initial_state.to(value)329 )330 331 for i in range(sequence_length):332 q_t = query[:, :, i]333 k_t = key[:, :, i]334 v_t = value[:, :, i]335 g_t = g[:, :, i].exp().unsqueeze(-1).unsqueeze(-1)336 beta_t = beta[:, :, i].unsqueeze(-1)337 338 last_recurrent_state = last_recurrent_state * g_t339 kv_mem = (last_recurrent_state * k_t.unsqueeze(-1)).sum(dim=-2)340 delta = (v_t - kv_mem) * beta_t341 last_recurrent_state = last_recurrent_state + k_t.unsqueeze(-1) * delta.unsqueeze(-2)342 core_attn_out[:, :, i] = (last_recurrent_state * q_t.unsqueeze(-1)).sum(dim=-2)343 344 if not output_final_state:345 last_recurrent_state = None346 core_attn_out = core_attn_out.transpose(1, 2).contiguous().to(initial_dtype)347 return core_attn_out, last_recurrent_state348 349 350class InternS2MobiusGatedDeltaNet(nn.Module):351 def __init__(self, config: InternS2MobiusConfig, layer_idx: int):352 super().__init__()353 self.hidden_size = config.hidden_size354 self.num_v_heads = config.linear_num_value_heads355 self.num_k_heads = config.linear_num_key_heads356 self.head_k_dim = config.linear_key_head_dim357 self.head_v_dim = config.linear_value_head_dim358 self.key_dim = self.head_k_dim * self.num_k_heads359 self.value_dim = self.head_v_dim * self.num_v_heads360 361 self.conv_kernel_size = config.linear_conv_kernel_dim362 self.layer_idx = layer_idx363 self.activation = config.hidden_act364 self.act = ACT2FN[config.hidden_act]365 self.layer_norm_epsilon = config.rms_norm_eps366 367 # QKV368 self.conv_dim = self.key_dim * 2 + self.value_dim369 self.conv1d = nn.Conv1d(370 in_channels=self.conv_dim,371 out_channels=self.conv_dim,372 bias=False,373 kernel_size=self.conv_kernel_size,374 groups=self.conv_dim,375 padding=self.conv_kernel_size - 1,376 )377 378 # time step projection (discretization)379 # instantiate once and copy inv_dt in init_weights of PretrainedModel380 self.dt_bias = nn.Parameter(torch.ones(self.num_v_heads))381 382 A = torch.empty(self.num_v_heads).uniform_(0, 16)383 self.A_log = nn.Parameter(torch.log(A))384 385 self.norm = (386 InternS2MobiusRMSNormGated(self.head_v_dim, eps=self.layer_norm_epsilon)387 if FusedRMSNormGated is None388 else FusedRMSNormGated(389 self.head_v_dim,390 eps=self.layer_norm_epsilon,391 activation=self.activation,392 device=torch.cuda.current_device(),393 dtype=config.dtype if config.dtype is not None else torch.get_default_dtype(),394 )395 )396 397 self.out_proj = nn.Linear(self.value_dim, self.hidden_size, bias=False)398 399 self.causal_conv1d_fn = causal_conv1d_fn400 self.causal_conv1d_update = causal_conv1d_update or torch_causal_conv1d_update401 self.chunk_gated_delta_rule = chunk_gated_delta_rule or torch_chunk_gated_delta_rule402 self.recurrent_gated_delta_rule = fused_recurrent_gated_delta_rule or torch_recurrent_gated_delta_rule403 404 if not is_fast_path_available:405 logger.warning_once(406 "The fast path is not available because one of the required library is not installed. Falling back to "407 "torch implementation. To install follow https://github.com/fla-org/flash-linear-attention#installation and"408 " https://github.com/Dao-AILab/causal-conv1d"409 )410 411 self.in_proj_qkv = nn.Linear(self.hidden_size, self.key_dim * 2 + self.value_dim, bias=False)412 self.in_proj_z = nn.Linear(self.hidden_size, self.value_dim, bias=False)413 self.in_proj_b = nn.Linear(self.hidden_size, self.num_v_heads, bias=False)414 self.in_proj_a = nn.Linear(self.hidden_size, self.num_v_heads, bias=False)415 416 def forward(417 self,418 hidden_states: torch.Tensor,419 cache_params: InternS2MobiusDynamicCache | None = None,420 cache_position: torch.LongTensor | None = None,421 attention_mask: torch.Tensor | None = None,422 ):423 hidden_states = apply_mask_to_padding_states(hidden_states, attention_mask)424 425 # Set up dimensions for reshapes later426 batch_size, seq_len, _ = hidden_states.shape427 428 use_precomputed_states = (429 cache_params is not None430 and cache_params.has_previous_state431 and seq_len == 1432 and cache_position is not None433 )434 435 # getting projected states from cache if it exists436 if cache_params is not None:437 conv_state = cache_params.conv_states[self.layer_idx]438 recurrent_state = cache_params.recurrent_states[self.layer_idx]439 440 mixed_qkv = self.in_proj_qkv(hidden_states)441 mixed_qkv = mixed_qkv.transpose(1, 2)442 443 z = self.in_proj_z(hidden_states)444 z = z.reshape(batch_size, seq_len, -1, self.head_v_dim)445 446 b = self.in_proj_b(hidden_states)447 a = self.in_proj_a(hidden_states)448 449 if use_precomputed_states:450 # 2. Convolution sequence transformation451 # NOTE: the conv state is updated in `causal_conv1d_update`452 mixed_qkv = self.causal_conv1d_update(453 mixed_qkv,454 conv_state,455 self.conv1d.weight.squeeze(1),456 self.conv1d.bias,457 self.activation,458 )459 else:460 if cache_params is not None:461 conv_state = F.pad(mixed_qkv, (self.conv_kernel_size - mixed_qkv.shape[-1], 0))462 cache_params.conv_states[self.layer_idx] = conv_state463 if self.causal_conv1d_fn is not None:464 mixed_qkv = self.causal_conv1d_fn(465 x=mixed_qkv,466 weight=self.conv1d.weight.squeeze(1),467 bias=self.conv1d.bias,468 activation=self.activation,469 seq_idx=None,470 )471 else:472 mixed_qkv = F.silu(self.conv1d(mixed_qkv)[:, :, :seq_len])473 474 mixed_qkv = mixed_qkv.transpose(1, 2)475 query, key, value = torch.split(476 mixed_qkv,477 [478 self.key_dim,479 self.key_dim,480 self.value_dim,481 ],482 dim=-1,483 )484 485 query = query.reshape(batch_size, seq_len, -1, self.head_k_dim)486 key = key.reshape(batch_size, seq_len, -1, self.head_k_dim)487 value = value.reshape(batch_size, seq_len, -1, self.head_v_dim)488 489 beta = b.sigmoid()490 # If the model is loaded in fp16, without the .float() here, A might be -inf491 g = -self.A_log.float().exp() * F.softplus(a.float() + self.dt_bias)492 if self.num_v_heads // self.num_k_heads > 1:493 query = query.repeat_interleave(self.num_v_heads // self.num_k_heads, dim=2)494 key = key.repeat_interleave(self.num_v_heads // self.num_k_heads, dim=2)495 496 if not use_precomputed_states:497 core_attn_out, last_recurrent_state = self.chunk_gated_delta_rule(498 query,499 key,500 value,501 g=g,502 beta=beta,503 initial_state=None,504 output_final_state=cache_params is not None,505 use_qk_l2norm_in_kernel=True,506 )507 508 else:509 core_attn_out, last_recurrent_state = self.recurrent_gated_delta_rule(510 query,511 key,512 value,513 g=g,514 beta=beta,515 initial_state=recurrent_state,516 output_final_state=cache_params is not None,517 use_qk_l2norm_in_kernel=True,518 )519 520 # Update cache521 if cache_params is not None:522 cache_params.recurrent_states[self.layer_idx] = last_recurrent_state523 524 # reshape input data into 2D tensor525 core_attn_out = core_attn_out.reshape(-1, self.head_v_dim)526 z = z.reshape(-1, self.head_v_dim)527 core_attn_out = self.norm(core_attn_out, z)528 core_attn_out = core_attn_out.reshape(batch_size, seq_len, -1)529 530 output = self.out_proj(core_attn_out)531 return output532 533 534def rotate_half(x):535 """Rotates half the hidden dims of the input."""536 x1 = x[..., : x.shape[-1] // 2]537 x2 = x[..., x.shape[-1] // 2 :]538 return torch.cat((-x2, x1), dim=-1)539 540 541# Adapted from transformers.models.glm.modular_glm.apply_rotary_pos_emb542def apply_rotary_pos_emb(q, k, cos, sin, unsqueeze_dim=1):543 """Applies Rotary Position Embedding to the query and key tensors.544 545 Removes the interleaving of cos and sin from GLM546 547 Args:548 q (`torch.Tensor`): The query tensor.549 k (`torch.Tensor`): The key tensor.550 cos (`torch.Tensor`): The cosine part of the rotary embedding.551 sin (`torch.Tensor`): The sine part of the rotary embedding.552 unsqueeze_dim (`int`, *optional*, defaults to 1):553 The 'unsqueeze_dim' argument specifies the dimension along which to unsqueeze cos[position_ids] and554 sin[position_ids] so that they can be properly broadcasted to the dimensions of q and k. For example, note555 that cos[position_ids] and sin[position_ids] have the shape [batch_size, seq_len, head_dim]. Then, if q and556 k have the shape [batch_size, heads, seq_len, head_dim], then setting unsqueeze_dim=1 makes557 cos[position_ids] and sin[position_ids] broadcastable to the shapes of q and k. Similarly, if q and k have558 the shape [batch_size, seq_len, heads, head_dim], then set unsqueeze_dim=2.559 Returns:560 `tuple(torch.Tensor)` comprising of the query and key tensors rotated using the Rotary Position Embedding.561 """562 cos = cos.unsqueeze(unsqueeze_dim)563 sin = sin.unsqueeze(unsqueeze_dim)564 565 # Keep half or full tensor for later concatenation566 rotary_dim = cos.shape[-1]567 q_rot, q_pass = q[..., :rotary_dim], q[..., rotary_dim:]568 k_rot, k_pass = k[..., :rotary_dim], k[..., rotary_dim:]569 570 # Apply rotary embeddings on the first half or full tensor571 q_embed = (q_rot * cos) + (rotate_half(q_rot) * sin)572 k_embed = (k_rot * cos) + (rotate_half(k_rot) * sin)573 574 # Concatenate back to full shape575 q_embed = torch.cat([q_embed, q_pass], dim=-1)576 k_embed = torch.cat([k_embed, k_pass], dim=-1)577 return q_embed, k_embed578 579 580def repeat_kv(hidden_states: torch.Tensor, n_rep: int) -> torch.Tensor:581 """582 This is the equivalent of torch.repeat_interleave(x, dim=1, repeats=n_rep). The hidden states go from (batch,583 num_key_value_heads, seqlen, head_dim) to (batch, num_attention_heads, seqlen, head_dim)584 """585 batch, num_key_value_heads, slen, head_dim = hidden_states.shape586 if n_rep == 1:587 return hidden_states588 hidden_states = hidden_states[:, :, None, :, :].expand(batch, num_key_value_heads, n_rep, slen, head_dim)589 return hidden_states.reshape(batch, num_key_value_heads * n_rep, slen, head_dim)590 591 592def eager_attention_forward(593 module: nn.Module,594 query: torch.Tensor,595 key: torch.Tensor,596 value: torch.Tensor,597 attention_mask: torch.Tensor | None,598 scaling: float,599 dropout: float = 0.0,600 **kwargs: Unpack[TransformersKwargs],601):602 key_states = repeat_kv(key, module.num_key_value_groups)603 value_states = repeat_kv(value, module.num_key_value_groups)604 605 attn_weights = torch.matmul(query, key_states.transpose(2, 3)) * scaling606 if attention_mask is not None:607 attn_weights = attn_weights + attention_mask608 609 attn_weights = nn.functional.softmax(attn_weights, dim=-1, dtype=torch.float32).to(query.dtype)610 attn_weights = nn.functional.dropout(attn_weights, p=dropout, training=module.training)611 attn_output = torch.matmul(attn_weights, value_states)612 attn_output = attn_output.transpose(1, 2).contiguous()613 614 return attn_output, attn_weights615 616 617class InternS2MobiusRMSNorm(nn.Module):618 def __init__(self, dim: int, eps: float = 1e-6):619 super().__init__()620 self.eps = eps621 self.weight = nn.Parameter(torch.zeros(dim))622 623 def _norm(self, x):624 return x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + self.eps)625 626 def forward(self, x):627 output = self._norm(x.float())628 # Llama does x.to(float16) * w whilst InternS2Mobius is (x * w).to(float16)629 # See https://github.com/huggingface/transformers/pull/29402630 output = output * (1.0 + self.weight.float())631 return output.type_as(x)632 633 def extra_repr(self):634 return f"{tuple(self.weight.shape)}, eps={self.eps}"635 636 637@use_kernelized_func(apply_rotary_pos_emb)638class InternS2MobiusAttention(nn.Module):639 """Multi-headed attention from 'Attention Is All You Need' paper"""640 641 def __init__(self, config: InternS2MobiusConfig, layer_idx: int):642 super().__init__()643 self.config = config644 self.layer_idx = layer_idx645 self.head_dim = getattr(config, "head_dim", config.hidden_size // config.num_attention_heads)646 self.num_key_value_groups = config.num_attention_heads // config.num_key_value_heads647 self.scaling = self.head_dim**-0.5648 self.attention_dropout = config.attention_dropout649 self.is_causal = True650 self.q_proj = nn.Linear(651 config.hidden_size, config.num_attention_heads * self.head_dim * 2, bias=config.attention_bias652 )653 self.k_proj = nn.Linear(654 config.hidden_size, config.num_key_value_heads * self.head_dim, bias=config.attention_bias655 )656 self.v_proj = nn.Linear(657 config.hidden_size, config.num_key_value_heads * self.head_dim, bias=config.attention_bias658 )659 self.o_proj = nn.Linear(660 config.num_attention_heads * self.head_dim, config.hidden_size, bias=config.attention_bias661 )662 self.q_norm = InternS2MobiusRMSNorm(663 self.head_dim, eps=config.rms_norm_eps664 ) # unlike olmo, only on the head dim!665 self.k_norm = InternS2MobiusRMSNorm(666 self.head_dim, eps=config.rms_norm_eps667 ) # thus post q_norm does not need reshape668 669 def forward(670 self,671 hidden_states: torch.Tensor,672 position_embeddings: tuple[torch.Tensor, torch.Tensor],673 attention_mask: torch.Tensor | None,674 past_key_values: Cache | None = None,675 cache_position: torch.LongTensor | None = None,676 **kwargs: Unpack[FlashAttentionKwargs],677 ) -> tuple[torch.Tensor, torch.Tensor | None]:678 input_shape = hidden_states.shape[:-1]679 hidden_shape = (*input_shape, -1, self.head_dim)680 681 query_states, gate = torch.chunk(682 self.q_proj(hidden_states).view(*input_shape, -1, self.head_dim * 2), 2, dim=-1683 )684 gate = gate.reshape(*input_shape, -1)685 686 query_states = self.q_norm(query_states.view(hidden_shape)).transpose(1, 2)687 key_states = self.k_norm(self.k_proj(hidden_states).view(hidden_shape)).transpose(1, 2)688 value_states = self.v_proj(hidden_states).view(hidden_shape).transpose(1, 2)689 690 cos, sin = position_embeddings691 query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin)692 693 if past_key_values is not None:694 # sin and cos are specific to RoPE models; cache_position needed for the static cache695 cache_kwargs = {"sin": sin, "cos": cos, "cache_position": cache_position}696 key_states, value_states = past_key_values.update(key_states, value_states, self.layer_idx, cache_kwargs)697 698 attention_interface: Callable = ALL_ATTENTION_FUNCTIONS.get_interface(699 self.config._attn_implementation, eager_attention_forward700 )701 702 attn_output, attn_weights = attention_interface(703 self,704 query_states,705 key_states,706 value_states,707 attention_mask,708 dropout=0.0 if not self.training else self.attention_dropout,709 scaling=self.scaling,710 **kwargs,711 )712 713 attn_output = attn_output.reshape(*input_shape, -1).contiguous()714 attn_output = attn_output * torch.sigmoid(gate)715 716 attn_output = self.o_proj(attn_output)717 return attn_output, attn_weights718 719 720class InternS2MobiusMLP(nn.Module):721 def __init__(self, config: InternS2MobiusConfig, intermediate_size: int):722 super().__init__()723 self.config = config724 self.hidden_size = config.hidden_size725 self.intermediate_size = intermediate_size726 self.gate_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=False)727 self.up_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=False)728 self.down_proj = nn.Linear(self.intermediate_size, self.hidden_size, bias=False)729 self.act_fn = ACT2FN[config.hidden_act]730 731 def forward(self, x):732 down_proj = self.down_proj(self.act_fn(self.gate_proj(x)) * self.up_proj(x))733 return down_proj734 735 736@use_experts_implementation737class InternS2MobiusExperts(nn.Module):738 """Collection of expert weights stored as 3D tensors."""739 740 def __init__(self, config):741 super().__init__()742 self.num_experts = config.num_experts743 self.hidden_dim = config.hidden_size744 self.intermediate_dim = config.moe_intermediate_size745 self.gate_up_proj = nn.Parameter(torch.empty(self.num_experts, 2 * self.intermediate_dim, self.hidden_dim))746 self.down_proj = nn.Parameter(torch.empty(self.num_experts, self.hidden_dim, self.intermediate_dim))747 self.act_fn = ACT2FN[config.hidden_act]748 749 def forward(750 self,751 hidden_states: torch.Tensor,752 top_k_index: torch.Tensor,753 top_k_weights: torch.Tensor,754 ) -> torch.Tensor:755 final_hidden_states = torch.zeros_like(hidden_states)756 with torch.no_grad():757 expert_mask = torch.nn.functional.one_hot(top_k_index, num_classes=self.num_experts)758 expert_mask = expert_mask.permute(2, 1, 0)759 expert_hit = torch.greater(expert_mask.sum(dim=(-1, -2)), 0).nonzero()760 761 for expert_idx in expert_hit:762 expert_idx = expert_idx[0]763 if expert_idx == self.num_experts:764 continue765 top_k_pos, token_idx = torch.where(expert_mask[expert_idx])766 current_state = hidden_states[token_idx]767 gate, up = nn.functional.linear(current_state, self.gate_up_proj[expert_idx]).chunk(2, dim=-1)768 current_hidden_states = self.act_fn(gate) * up769 current_hidden_states = nn.functional.linear(current_hidden_states, self.down_proj[expert_idx])770 current_hidden_states = current_hidden_states * top_k_weights[token_idx, top_k_pos, None]771 final_hidden_states.index_add_(0, token_idx, current_hidden_states.to(final_hidden_states.dtype))772 773 return final_hidden_states774 775 776class InternS2MobiusTopKRouter(nn.Module):777 def __init__(self, config):778 super().__init__()779 self.top_k = config.num_experts_per_tok780 self.num_experts = config.num_experts781 self.hidden_dim = config.hidden_size782 self.weight = nn.Parameter(torch.zeros(self.num_experts, self.hidden_dim))783 784 def forward(self, hidden_states):785 hidden_states = hidden_states.reshape(-1, self.hidden_dim)786 router_logits = F.linear(hidden_states, self.weight) # (seq_len, num_experts)787 router_logits = torch.nn.functional.softmax(router_logits, dtype=torch.float, dim=-1)788 router_top_value, router_indices = torch.topk(router_logits, self.top_k, dim=-1) # (seq_len, top_k)789 router_top_value /= router_top_value.sum(dim=-1, keepdim=True)790 router_top_value = router_top_value.to(router_logits.dtype)791 router_scores = router_top_value792 return router_logits, router_scores, router_indices793 794 795class InternS2MobiusMetaMoeBlock(nn.Module):796 """A shared MoE block containing both the router (gate) and experts."""797 798 def __init__(self, config):799 super().__init__()800 self.gate = InternS2MobiusTopKRouter(config)801 self.experts = InternS2MobiusExperts(config)802 803 def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:804 _, routing_weights, selected_experts = self.gate(hidden_states)805 return self.experts(hidden_states, selected_experts, routing_weights)806 807 808class InternS2MobiusSharedExpertBlock(nn.Module):809 """Per-layer shared expert. Routed experts are shared across layers via meta blocks."""810 811 def __init__(self, config):812 super().__init__()813 self.shared_expert = InternS2MobiusMLP(config, intermediate_size=config.shared_expert_intermediate_size)814 self.shared_expert_gate = torch.nn.Linear(config.hidden_size, 1, bias=False)815 816 def forward(self, hidden_states: torch.Tensor, expert_output: torch.Tensor) -> torch.Tensor:817 shared_expert_output = self.shared_expert(hidden_states)818 shared_expert_output = F.sigmoid(self.shared_expert_gate(hidden_states)) * shared_expert_output819 return expert_output + shared_expert_output820 821 822class InternS2MobiusDecoderLayer(GradientCheckpointingLayer):823 def __init__(self, config: InternS2MobiusTextConfig, layer_idx: int):824 super().__init__()825 self.layer_idx = layer_idx826 self.hidden_size = config.hidden_size827 self.num_blocks = config.num_blocks828 self.layer_type = config.layer_types[layer_idx]829 if self.layer_type == "linear_attention":830 self.linear_attn = InternS2MobiusGatedDeltaNet(config, layer_idx)831 elif self.layer_type == "full_attention":832 self.self_attn = InternS2MobiusAttention(config, layer_idx)833 self.mlp = InternS2MobiusSharedExpertBlock(config)834 self.input_layernorm = InternS2MobiusRMSNorm(config.hidden_size, eps=config.rms_norm_eps)835 self.post_attention_layernorm = InternS2MobiusRMSNorm(config.hidden_size, eps=config.rms_norm_eps)836 837 def forward(838 self,839 hidden_states: torch.Tensor,840 position_embeddings: tuple[torch.Tensor, torch.Tensor],841 attention_mask: torch.Tensor | None = None,842 position_ids: torch.LongTensor | None = None,843 past_key_values: Cache | None = None,844 cache_position: torch.LongTensor | None = None,845 meta_mlp: nn.ModuleList | None = None,846 **kwargs: Unpack[FlashAttentionKwargs],847 ) -> torch.FloatTensor:848 residual = hidden_states849 850 hidden_states = self.input_layernorm(hidden_states)851 852 # Token Mixer853 if self.layer_type == "linear_attention":854 hidden_states = self.linear_attn(855 hidden_states=hidden_states,856 cache_params=past_key_values,857 cache_position=cache_position,858 attention_mask=attention_mask,859 )860 elif self.layer_type == "full_attention":861 # Self Attention862 hidden_states, _ = self.self_attn(863 hidden_states=hidden_states,864 attention_mask=attention_mask,865 position_ids=position_ids,866 past_key_values=past_key_values,867 cache_position=cache_position,868 position_embeddings=position_embeddings,869 **kwargs,870 )871 872 hidden_states = residual + hidden_states873 874 # Fully Connected875 residual = hidden_states876 hidden_states = self.post_attention_layernorm(hidden_states)877 878 batch_size, sequence_length, hidden_dim = hidden_states.shape879 hidden_states_2d = hidden_states.view(-1, hidden_dim)880 881 # Routed experts from shared meta block882 block_idx = self.layer_idx % self.num_blocks883 expert_output = meta_mlp[block_idx](hidden_states_2d)884 885 # Per-layer shared expert + combine886 hidden_states = self.mlp(hidden_states_2d, expert_output)887 hidden_states = hidden_states.view(batch_size, sequence_length, hidden_dim)888 889 hidden_states = residual + hidden_states890 891 return hidden_states892 893 894class InternS2MobiusPreTrainedModel(PreTrainedModel):895 config: InternS2MobiusConfig896 base_model_prefix = "model"897 supports_gradient_checkpointing = True898 _no_split_modules = ["InternS2MobiusDecoderLayer", "InternS2MobiusVisionBlock"]899 _skip_keys_device_placement = "past_key_values"900 _supports_flash_attn = True901 _supports_sdpa = True902 _keys_to_ignore_on_load_unexpected = [r"^mtp.*"]903 _can_record_outputs = {904 "router_logits": OutputRecorder(InternS2MobiusTopKRouter, index=0),905 "hidden_states": InternS2MobiusDecoderLayer,906 "attentions": InternS2MobiusAttention,907 }908 _is_stateful = True909 910 @torch.no_grad()911 def _init_weights(self, module):912 super()._init_weights(module)913 if isinstance(module, InternS2MobiusGatedDeltaNet):914 init.ones_(module.dt_bias)915 init.copy_(module.A_log, torch.empty_like(module.A_log).uniform_(0, 16).log_())916 # We initialize with 0s to be 1 centered as the RMSNorm here does (1 + weight)917 elif isinstance(module, InternS2MobiusRMSNorm):918 init.zeros_(module.weight)919 elif isinstance(module, InternS2MobiusExperts):920 init.normal_(module.gate_up_proj, mean=0.0, std=self.config.initializer_range)921 init.normal_(module.down_proj, mean=0.0, std=self.config.initializer_range)922 elif isinstance(module, InternS2MobiusTopKRouter):923 init.normal_(module.weight, mean=0.0, std=self.config.initializer_range)924 elif isinstance(module, InternS2MobiusVisionRotaryEmbedding):925 inv_freq = 1.0 / (module.theta ** (torch.arange(0, module.dim, 2, dtype=torch.float) / module.dim))926 init.copy_(module.inv_freq, inv_freq)927 928 929class InternS2MobiusVisionMLP(nn.Module):930 def __init__(self, config):931 super().__init__()932 self.hidden_size = config.hidden_size933 self.intermediate_size = config.intermediate_size934 self.linear_fc1 = nn.Linear(self.hidden_size, self.intermediate_size, bias=True)935 self.linear_fc2 = nn.Linear(self.intermediate_size, self.hidden_size, bias=True)936 self.act_fn = ACT2FN[config.hidden_act]937 938 def forward(self, hidden_state):939 return self.linear_fc2(self.act_fn(self.linear_fc1(hidden_state)))940 941 942class InternS2MobiusVisionPatchEmbed(nn.Module):943 def __init__(self, config) -> None:944 super().__init__()945 self.patch_size = config.patch_size946 self.temporal_patch_size = config.temporal_patch_size947 self.in_channels = config.in_channels948 self.embed_dim = config.hidden_size949 950 kernel_size = [self.temporal_patch_size, self.patch_size, self.patch_size]951 self.proj = nn.Conv3d(self.in_channels, self.embed_dim, kernel_size=kernel_size, stride=kernel_size, bias=True)952 953 def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:954 target_dtype = self.proj.weight.dtype955 hidden_states = hidden_states.view(956 -1, self.in_channels, self.temporal_patch_size, self.patch_size, self.patch_size957 )958 hidden_states = self.proj(hidden_states.to(dtype=target_dtype)).view(-1, self.embed_dim)959 return hidden_states960 961 962class InternS2MobiusVisionPatchMerger(nn.Module):963 def __init__(self, config: InternS2MobiusVisionConfig, use_postshuffle_norm=False) -> None:964 super().__init__()965 self.hidden_size = config.hidden_size * (config.spatial_merge_size**2)966 self.use_postshuffle_norm = use_postshuffle_norm967 self.norm = nn.LayerNorm(self.hidden_size if use_postshuffle_norm else config.hidden_size, eps=1e-6)968 self.linear_fc1 = nn.Linear(self.hidden_size, self.hidden_size)969 self.act_fn = nn.GELU()970 self.linear_fc2 = nn.Linear(self.hidden_size, config.out_hidden_size)971 972 def forward(self, x: torch.Tensor) -> torch.Tensor:973 x = self.norm(x.view(-1, self.hidden_size) if self.use_postshuffle_norm else x).view(-1, self.hidden_size)974 x = self.linear_fc2(self.act_fn(self.linear_fc1(x)))975 return x976 977 978def apply_rotary_pos_emb_vision(979 q: torch.Tensor, k: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor980) -> tuple[torch.Tensor, torch.Tensor]:981 orig_q_dtype = q.dtype982 orig_k_dtype = k.dtype983 q, k = q.float(), k.float()984 cos, sin = cos.unsqueeze(-2).float(), sin.unsqueeze(-2).float()985 q_embed = (q * cos) + (rotate_half(q) * sin)986 k_embed = (k * cos) + (rotate_half(k) * sin)987 q_embed = q_embed.to(orig_q_dtype)988 k_embed = k_embed.to(orig_k_dtype)989 return q_embed, k_embed990 991 992class InternS2MobiusVisionAttention(nn.Module):993 def __init__(self, config: InternS2MobiusVisionConfig) -> None:994 super().__init__()995 self.dim = config.hidden_size996 self.num_heads = config.num_heads997 self.head_dim = self.dim // self.num_heads998 self.num_key_value_groups = 1 # needed for eager attention999 self.qkv = nn.Linear(self.dim, self.dim * 3, bias=True)1000 self.proj = nn.Linear(self.dim, self.dim)1001 self.scaling = self.head_dim**-0.51002 self.config = config1003 self.attention_dropout = 0.01004 self.is_causal = False1005 1006 def forward(1007 self,1008 hidden_states: torch.Tensor,1009 cu_seqlens: torch.Tensor,1010 rotary_pos_emb: torch.Tensor | None = None,1011 position_embeddings: tuple[torch.Tensor, torch.Tensor] | None = None,1012 **kwargs,1013 ) -> torch.Tensor:1014 seq_length = hidden_states.shape[0]1015 query_states, key_states, value_states = (1016 self.qkv(hidden_states).reshape(seq_length, 3, self.num_heads, -1).permute(1, 0, 2, 3).unbind(0)1017 )1018 cos, sin = position_embeddings1019 query_states, key_states = apply_rotary_pos_emb_vision(query_states, key_states, cos, sin)1020 1021 query_states = query_states.transpose(0, 1).unsqueeze(0)1022 key_states = key_states.transpose(0, 1).unsqueeze(0)1023 value_states = value_states.transpose(0, 1).unsqueeze(0)1024 1025 attention_interface: Callable = ALL_ATTENTION_FUNCTIONS.get_interface(1026 self.config._attn_implementation, eager_attention_forward1027 )1028 1029 if is_flash_attention_requested(self.config):1030 # Flash Attention: Use cu_seqlens for variable length attention1031 max_seqlen = (cu_seqlens[1:] - cu_seqlens[:-1]).max()1032 attn_output, _ = attention_interface(1033 self,1034 query_states,1035 key_states,1036 value_states,1037 attention_mask=None,1038 scaling=self.scaling,1039 dropout=0.0 if not self.training else self.attention_dropout,1040 cu_seq_lens_q=cu_seqlens,1041 cu_seq_lens_k=cu_seqlens,1042 max_length_q=max_seqlen,1043 max_length_k=max_seqlen,1044 is_causal=False,1045 **kwargs,1046 )1047 else:1048 # Other implementations: Process each chunk separately1049 lengths = cu_seqlens[1:] - cu_seqlens[:-1]1050 splits = [1051 torch.split(tensor, lengths.tolist(), dim=2) for tensor in (query_states, key_states, value_states)1052 ]1053 1054 attn_outputs = [1055 attention_interface(1056 self,1057 q,1058 k,1059 v,1060 attention_mask=None,1061 scaling=self.scaling,1062 dropout=0.0 if not self.training else self.attention_dropout,1063 is_causal=False,1064 **kwargs,1065 )[0]1066 for q, k, v in zip(*splits)1067 ]1068 attn_output = torch.cat(attn_outputs, dim=1)1069 1070 attn_output = attn_output.reshape(seq_length, -1).contiguous()1071 attn_output = self.proj(attn_output)1072 return attn_output1073 1074 1075class InternS2MobiusVisionBlock(GradientCheckpointingLayer):1076 def __init__(self, config, attn_implementation: str = "sdpa") -> None:1077 super().__init__()1078 self.norm1 = nn.LayerNorm(config.hidden_size, eps=1e-6)1079 self.norm2 = nn.LayerNorm(config.hidden_size, eps=1e-6)1080 self.attn = InternS2MobiusVisionAttention(config=config)1081 self.mlp = InternS2MobiusVisionMLP(config=config)1082 1083 def forward(1084 self,1085 hidden_states: torch.Tensor,1086 cu_seqlens: torch.Tensor,1087 rotary_pos_emb: torch.Tensor | None = None,1088 position_embeddings: tuple[torch.Tensor, torch.Tensor] | None = None,1089 **kwargs,1090 ) -> torch.Tensor:1091 hidden_states = hidden_states + self.attn(1092 self.norm1(hidden_states),1093 cu_seqlens=cu_seqlens,1094 rotary_pos_emb=rotary_pos_emb,1095 position_embeddings=position_embeddings,1096 **kwargs,1097 )1098 hidden_states = hidden_states + self.mlp(self.norm2(hidden_states))1099 return hidden_states1100 1101 1102class InternS2MobiusVisionModel(InternS2MobiusPreTrainedModel):1103 config: InternS2MobiusVisionConfig1104 _no_split_modules = ["InternS2MobiusVisionBlock"]1105 _can_record_outputs = {1106 "hidden_states": InternS2MobiusVisionBlock,1107 "attentions": InternS2MobiusVisionAttention,1108 }1109 1110 def __init__(self, config, *inputs, **kwargs) -> None:1111 super().__init__(config, *inputs, **kwargs)1112 self.spatial_merge_size = config.spatial_merge_size1113 self.patch_size = config.patch_size1114 self.spatial_merge_unit = self.spatial_merge_size * self.spatial_merge_size1115 1116 self.patch_embed = InternS2MobiusVisionPatchEmbed(1117 config=config,1118 )1119 1120 self.pos_embed = nn.Embedding(config.num_position_embeddings, config.hidden_size)1121 self.num_grid_per_side = int(config.num_position_embeddings**0.5)1122 1123 head_dim = config.hidden_size // config.num_heads1124 self.rotary_pos_emb = InternS2MobiusVisionRotaryEmbedding(head_dim // 2)1125 1126 self.blocks = nn.ModuleList([InternS2MobiusVisionBlock(config) for _ in range(config.depth)])1127 self.merger = InternS2MobiusVisionPatchMerger(1128 config=config,1129 use_postshuffle_norm=False,1130 )1131 1132 self.gradient_checkpointing = False1133 1134 self.post_init()1135 1136 def rot_pos_emb(self, grid_thw: torch.Tensor) -> torch.Tensor:1137 merge_size = self.spatial_merge_size1138 grid_thw_list = grid_thw.tolist()1139 1140 max_hw = max(max(h, w) for _, h, w in grid_thw_list)1141 freq_table = self.rotary_pos_emb(max_hw) # (max_hw, dim // 2)1142 device = freq_table.device1143 1144 total_tokens = sum(t * h * w for t, h, w in grid_thw_list)1145 pos_ids = torch.empty((total_tokens, 2), dtype=torch.long, device=device)1146 1147 offset = 01148 for num_frames, height, width in grid_thw_list:1149 merged_h, merged_w = height // merge_size, width // merge_size1150 1151 block_rows = torch.arange(merged_h, device=device) # block row indices1152 block_cols = torch.arange(merged_w, device=device) # block col indices1153 intra_row = torch.arange(merge_size, device=device) # intra-block row offsets1154 intra_col = torch.arange(merge_size, device=device) # intra-block col offsets1155 1156 # Compute full-resolution positions1157 row_idx = block_rows[:, None, None, None] * merge_size + intra_row[None, None, :, None]1158 col_idx = block_cols[None, :, None, None] * merge_size + intra_col[None, None, None, :]1159 1160 row_idx = row_idx.expand(merged_h, merged_w, merge_size, merge_size).reshape(-1)1161 col_idx = col_idx.expand(merged_h, merged_w, merge_size, merge_size).reshape(-1)1162 1163 coords = torch.stack((row_idx, col_idx), dim=-1)1164 1165 if num_frames > 1:1166 coords = coords.repeat(num_frames, 1)1167 1168 num_tokens = coords.shape[0]1169 pos_ids[offset : offset + num_tokens] = coords1170 offset += num_tokens1171 1172 embeddings = freq_table[pos_ids] # lookup rotary embeddings1173 embeddings = embeddings.flatten(1)1174 return embeddings1175 1176 def fast_pos_embed_interpolate(self, grid_thw):1177 grid_thw_list = grid_thw.tolist()1178 grid_ts = [row[0] for row in grid_thw_list]1179 grid_hs = [row[1] for row in grid_thw_list]1180 grid_ws = [row[2] for row in grid_thw_list]1181 device = self.pos_embed.weight.device1182 1183 idx_list = [[] for _ in range(4)]1184 weight_list = [[] for _ in range(4)]1185 1186 for t, h, w in grid_thw_list:1187 h_idxs = torch.linspace(0, self.num_grid_per_side - 1, h)1188 w_idxs = torch.linspace(0, self.num_grid_per_side - 1, w)1189 1190 h_idxs_floor = h_idxs.int()1191 w_idxs_floor = w_idxs.int()1192 h_idxs_ceil = (h_idxs.int() + 1).clip(max=self.num_grid_per_side - 1)1193 w_idxs_ceil = (w_idxs.int() + 1).clip(max=self.num_grid_per_side - 1)1194 1195 dh = h_idxs - h_idxs_floor1196 dw = w_idxs - w_idxs_floor1197 1198 base_h = h_idxs_floor * self.num_grid_per_side1199 base_h_ceil = h_idxs_ceil * self.num_grid_per_side1200