Aluode/PerceptionLabPortable
0
1# coding=utf-82# Copyright 2024 Zyphra Technologies and the HuggingFace Inc. team. All rights reserved.3#4#5# Licensed under the Apache License, Version 2.0 (the "License");6# you may not use this file except in compliance with the License.7# You may obtain a copy of the License at8#9# http://www.apache.org/licenses/LICENSE-2.010#11# Unless required by applicable law or agreed to in writing, software12# distributed under the License is distributed on an "AS IS" BASIS,13# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.14# See the License for the specific language governing permissions and15# limitations under the License.16import math17import re18from itertools import cycle19from typing import Callable, Optional, Union20 21import torch22from torch import nn23 24from ...activations import ACT2FN25from ...modeling_flash_attention_utils import FlashAttentionKwargs26from ...modeling_outputs import BaseModelOutputWithPast27from ...modeling_utils import ALL_ATTENTION_FUNCTIONS, PreTrainedModel28from ...processing_utils import Unpack29from ...utils import (30 logging,31)32from ...utils.deprecation import deprecate_kwarg33from ...utils.import_utils import (34 is_causal_conv1d_available,35 is_mamba_ssm_available,36)37from ..llama.modeling_llama import LlamaRotaryEmbedding, apply_rotary_pos_emb38from ..mamba2.modeling_mamba2 import pad_tensor_by_size, reshape_into_chunks, segment_sum39from ..zamba.modeling_zamba import (40 ZambaAttention,41 ZambaAttentionDecoderLayer,42 ZambaForCausalLM,43 ZambaForSequenceClassification,44 ZambaHybridDynamicCache,45 ZambaHybridLayer,46 ZambaMambaDecoderLayer,47 ZambaModel,48 ZambaRMSNorm,49 eager_attention_forward,50)51from .configuration_zamba2 import Zamba2Config52 53 54if is_mamba_ssm_available():55 from mamba_ssm.ops.triton.selective_state_update import selective_state_update56 from mamba_ssm.ops.triton.ssd_combined import mamba_chunk_scan_combined, mamba_split_conv1d_scan_combined57else:58 selective_state_update, mamba_chunk_scan_combined, mamba_split_conv1d_scan_combined = None, None, None59 60if is_causal_conv1d_available():61 from causal_conv1d import causal_conv1d_fn, causal_conv1d_update62else:63 causal_conv1d_update, causal_conv1d_fn = None, None64 65is_fast_path_available = all((selective_state_update, causal_conv1d_fn, causal_conv1d_update))66 67 68_CONFIG_FOR_DOC = "Zyphra/Zamba2-2.7B"69 70logger = logging.get_logger(__name__)71 72 73class Zamba2RMSNormGated(torch.nn.Module):74 def __init__(self, hidden_size, group_size, eps=1e-6):75 super().__init__()76 self.weight = nn.Parameter(torch.ones(hidden_size))77 self.variance_epsilon = eps78 self.group_size = group_size79 80 def forward(self, hidden_states, gate=None):81 input_dtype = hidden_states.dtype82 hidden_states = hidden_states.to(torch.float32)83 if gate is not None:84 hidden_states = hidden_states * nn.functional.silu(gate.to(torch.float32))85 *prefix_dims, last_dim = hidden_states.shape86 group_count = last_dim // self.group_size87 hidden_states_group = hidden_states.view(*prefix_dims, group_count, self.group_size)88 variance = hidden_states_group.pow(2).mean(-1, keepdim=True)89 hidden_states_group = hidden_states_group * torch.rsqrt(variance + self.variance_epsilon)90 hidden_states = hidden_states_group.view(*prefix_dims, group_count * self.group_size)91 return self.weight * hidden_states.to(input_dtype)92 93 94class Zamba2RMSNorm(ZambaRMSNorm):95 pass96 97 98class Zamba2HybridDynamicCache(ZambaHybridDynamicCache):99 """100 A dynamic cache that can handle both the attention cache (which has a seq_len dimension) and the mamba cache101 (which has a constant shape regardless of seq_len).102 103 This cache has two sets of lists of tensors: `key_cache` and `value_cache` for attention cache and `conv_states`104 and `ssm_states` for mamba cache. Each of these lists has `num_layers` tensors. The expected shape for each tensor105 For attention layers, `key_cache` and `value_cache` have a shape of `(batch_size, num_heads, seq_len, head_dim)`,106 while `conv_states` and `ssm_states` have a shape of `(batch_size, 0)` (empty tensors).107 For mamba layers, `key_cache` and `value_cache` have a shape of `(batch_size, 0)` (empty tensors),108 while `conv_states` represents the convolution state and has a shape of `(batch_size, d_inner, d_conv)`,109 and `ssm_states` represents the ssm state and has a shape of `(batch_size, d_inner, d_state)`.110 """111 112 def __init__(113 self, config: Zamba2Config, batch_size: int, dtype: torch.dtype = torch.float16, device: Optional[str] = None114 ):115 self.dtype = dtype116 self.layers_block_type = config.layers_block_type117 self.has_previous_state = False118 self.intermediate_size = int(config.mamba_expand * config.hidden_size)119 self.ssm_state_size = config.mamba_d_state120 self.conv_kernel_size = config.mamba_d_conv121 self.n_mamba_heads = config.n_mamba_heads122 self.transformer_layers = []123 self._modules = {}124 self._parameters = {}125 self._buffers = {}126 self.conv_states = {}127 self.ssm_states = {}128 for i in range(config.num_hidden_layers):129 self.conv_states[i] = torch.zeros(130 batch_size,131 self.intermediate_size + 2 * config.mamba_ngroups * config.mamba_d_state,132 self.conv_kernel_size,133 device=device,134 dtype=dtype,135 )136 self.ssm_states[i] = torch.zeros(137 batch_size, self.n_mamba_heads, config.mamba_headdim, self.ssm_state_size, device=device, dtype=dtype138 )139 if self.layers_block_type[i] == "hybrid":140 self.transformer_layers.append(i)141 self.key_cache = [torch.tensor([[]] * batch_size, device=device) for _ in range(config.num_hidden_layers)]142 self.value_cache = [torch.tensor([[]] * batch_size, device=device) for _ in range(config.num_hidden_layers)]143 144 def update_conv_state(145 self, layer_idx: int, new_conv_state: torch.Tensor, cache_position: torch.LongTensor146 ) -> torch.Tensor:147 conv_state = self.conv_states[layer_idx]148 cache_position = cache_position.clamp(0, self.conv_kernel_size - 1)149 150 conv_state = conv_state.roll(shifts=-1, dims=-1)151 conv_state[:, :, cache_position] = new_conv_state.to(conv_state.device)152 self.conv_states[layer_idx].zero_()153 self.conv_states[layer_idx] += conv_state154 return self.conv_states[layer_idx]155 156 def reset(self):157 self.conv_states.zero_()158 self.ssm_states.zero_()159 160 def get_seq_length(self, layer_idx: Optional[int] = 0) -> int:161 """Returns the sequence length of the cached states. A layer index can be optionally passed."""162 # take any layer that contains cache and not empty tensor163 layer_idx = self.transformer_layers[0] if layer_idx not in self.transformer_layers else layer_idx164 if len(self.key_cache) <= layer_idx or self.key_cache[layer_idx].numel() == 0:165 return 0166 return self.key_cache[layer_idx].shape[-2]167 168 169class Zamba2RotaryEmbedding(LlamaRotaryEmbedding):170 pass171 172 173class Zamba2Attention(ZambaAttention):174 """175 Multi-headed attention from 'Attention Is All You Need' paper.176 177 Adapted from transformers.models.mistral.modeling_mistral.MistralAttention:178 The input dimension here is attention_hidden_size = 2 * hidden_size, and head_dim = attention_hidden_size // num_heads.179 The extra factor of 2 comes from the input being the concatenation of original_hidden_states with the output of the previous (mamba) layer180 (see fig. 2 in https://huggingface.co/papers/2405.16712).181 Additionally, replaced182 attn_weights = torch.matmul(query_states, key_states.transpose(2, 3)) / math.sqrt(self.head_dim) with183 attn_weights = torch.matmul(query_states, key_states.transpose(2, 3)) / math.sqrt(self.head_dim/2)184 Finally, this attention layer contributes to tied transformer blocks aimed to increasing compute without increasing model size. Because this185 layer is tied, un-tied adapters (formally the same as LoRA but used in the base model) modules are added to the q, k, v projectors to increase186 expressivity with a small memory overhead (see Fig. 2 of https://huggingface.co/papers/2411.15242).187 """188 189 def __init__(190 self,191 config: Zamba2Config,192 layer_idx: Optional[int] = None,193 num_fwd_mem_blocks: Optional[int] = None,194 block_id: Optional[int] = None,195 ):196 super().__init__(config, layer_idx)197 self.num_fwd_mem_blocks = num_fwd_mem_blocks198 self.layer_block_map = config.hybrid_layer_ids199 self.block_id = block_id200 201 if config.use_shared_attention_adapter:202 self.linear_q_adapter_list = nn.ModuleList([])203 self.linear_k_adapter_list = nn.ModuleList([])204 self.linear_v_adapter_list = nn.ModuleList([])205 206 for i in range(self.num_fwd_mem_blocks):207 if i % config.num_mem_blocks == block_id:208 linear_q_adapter = nn.Sequential(209 nn.Linear(self.attention_hidden_size, self.config.adapter_rank, bias=False),210 nn.Linear(self.config.adapter_rank, self.attention_hidden_size, bias=False),211 )212 linear_k_adapter = nn.Sequential(213 nn.Linear(self.attention_hidden_size, self.config.adapter_rank, bias=False),214 nn.Linear(self.config.adapter_rank, self.attention_hidden_size, bias=False),215 )216 linear_v_adapter = nn.Sequential(217 nn.Linear(self.attention_hidden_size, self.config.adapter_rank, bias=False),218 nn.Linear(self.config.adapter_rank, self.attention_hidden_size, bias=False),219 )220 else:221 linear_q_adapter = nn.Identity()222 linear_k_adapter = nn.Identity()223 linear_v_adapter = nn.Identity()224 self.linear_q_adapter_list.append(linear_q_adapter)225 self.linear_k_adapter_list.append(linear_k_adapter)226 self.linear_v_adapter_list.append(linear_v_adapter)227 228 self.layer_dic = {value: index for index, value in enumerate(self.layer_block_map)}229 230 @deprecate_kwarg("past_key_value", new_name="past_key_values", version="4.58")231 def forward(232 self,233 hidden_states: torch.Tensor,234 layer_idx: int,235 attention_mask: Optional[torch.Tensor] = None,236 past_key_values: Optional[Zamba2HybridDynamicCache] = None,237 position_embeddings: Optional[tuple[torch.Tensor, torch.Tensor]] = None,238 **kwargs: Unpack[FlashAttentionKwargs],239 ) -> tuple[torch.Tensor, Optional[torch.Tensor], Optional[tuple[torch.Tensor]]]:240 input_shape = hidden_states.shape[:-1]241 hidden_shape = (*input_shape, -1, self.head_dim)242 243 query_states = self.q_proj(hidden_states)244 key_states = self.k_proj(hidden_states)245 value_states = self.v_proj(hidden_states)246 if self.config.use_shared_attention_adapter:247 adapter_layer_idx = self.layer_dic[layer_idx]248 query_states = query_states + self.linear_q_adapter_list[adapter_layer_idx](hidden_states)249 key_states = key_states + self.linear_k_adapter_list[adapter_layer_idx](hidden_states)250 value_states = value_states + self.linear_v_adapter_list[adapter_layer_idx](hidden_states)251 252 query_states = query_states.view(hidden_shape).transpose(1, 2)253 key_states = key_states.view(hidden_shape).transpose(1, 2)254 value_states = value_states.view(hidden_shape).transpose(1, 2)255 256 if self.config.use_mem_rope:257 cos, sin = position_embeddings258 query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin)259 260 if past_key_values is not None:261 key_states, value_states = past_key_values.update(key_states, value_states, layer_idx)262 263 attention_interface: Callable = eager_attention_forward264 if self.config._attn_implementation != "eager":265 attention_interface = ALL_ATTENTION_FUNCTIONS[self.config._attn_implementation]266 267 attn_output, attn_weights = attention_interface(268 self,269 query_states,270 key_states,271 value_states,272 attention_mask,273 dropout=0.0 if not self.training else self.attention_dropout,274 scaling=self.scaling,275 **kwargs,276 )277 278 attn_output = attn_output.reshape(*input_shape, -1).contiguous()279 attn_output = self.o_proj(attn_output)280 return attn_output, attn_weights281 282 283class Zamba2MambaMixer(nn.Module):284 """285 Compute ∆, A, B, C, and D the state space parameters and compute the `contextualized_states`.286 A, D are input independent (see Mamba paper [1] Section 3.5.2 "Interpretation of A" for why A isn't selective)287 ∆, B, C are input-dependent (this is a key difference between Mamba and the linear time invariant S4,288 and is why Mamba is called **selective** state spaces)289 """290 291 def __init__(self, config: Zamba2Config, layer_idx: Optional[int] = None):292 super().__init__()293 self.config = config294 self.hidden_size = config.hidden_size295 self.ssm_state_size = config.mamba_d_state296 self.conv_kernel_size = config.mamba_d_conv297 self.intermediate_size = int(config.mamba_expand * self.hidden_size)298 self.layer_idx = layer_idx299 self.use_conv_bias = config.use_conv_bias300 self.activation = "silu"301 self.act = nn.SiLU()302 self.use_mem_eff_path = config.use_mem_eff_path303 304 self.n_groups = config.mamba_ngroups305 self.head_dim = config.mamba_headdim306 self.num_heads = self.config.n_mamba_heads307 self.chunk_size = config.chunk_size308 309 self.time_step_limit = config.time_step_limit310 self.time_step_min = config.time_step_min311 self.time_step_max = config.time_step_max312 313 self.conv_dim = self.intermediate_size + 2 * self.n_groups * self.ssm_state_size314 self.conv1d = nn.Conv1d(315 in_channels=self.conv_dim,316 out_channels=self.conv_dim,317 bias=True,318 kernel_size=config.mamba_d_conv,319 groups=self.conv_dim,320 padding=config.mamba_d_conv - 1,321 )322 323 # projection of the input hidden states324 projection_size = self.intermediate_size + self.conv_dim + self.num_heads325 self.in_proj = nn.Linear(326 self.hidden_size,327 projection_size,328 bias=config.add_bias_linear,329 )330 # selective projection used to make dt, B and C input dependent331 332 # time step projection (discretization)333 # instantiate once and copy inv_dt in init_weights of PretrainedModel334 self.dt_bias = nn.Parameter(torch.ones(self.num_heads))335 336 # S4D real initialization. These are not discretized!337 # The core is to load them, compute the discrete states, then write the updated state. Keeps the memory bounded338 A = torch.arange(1, self.num_heads + 1)339 self.A_log = nn.Parameter(torch.log(A))340 self.norm = Zamba2RMSNormGated(341 self.intermediate_size, group_size=self.intermediate_size // self.n_groups, eps=1e-5342 )343 self.D = nn.Parameter(torch.ones(self.num_heads))344 345 self.out_proj = nn.Linear(self.intermediate_size, self.hidden_size, bias=config.add_bias_linear)346 347 if not is_fast_path_available:348 logger.warning_once(349 "The fast path is not available because one of `(selective_state_update, causal_conv1d_fn, causal_conv1d_update)`"350 " is None. Falling back to the naive implementation. To install follow https://github.com/state-spaces/mamba/#installation and"351 " https://github.com/Dao-AILab/causal-conv1d"352 )353 354 def cuda_kernels_forward(355 self,356 hidden_states: torch.Tensor,357 cache_params: Optional[Zamba2HybridDynamicCache] = None,358 attention_mask: Optional[torch.Tensor] = None,359 ):360 # set up dimensions for reshapes later361 362 batch_size, seq_len, _ = hidden_states.shape363 groups_time_state_size = self.n_groups * self.ssm_state_size364 d_to_remove = 2 * self.intermediate_size + 2 * self.n_groups * self.ssm_state_size + self.num_heads365 366 # getting projected states from cache if it exists367 if cache_params is not None and cache_params.has_previous_state:368 in_projected_states = self.in_proj(hidden_states.squeeze(1)) # (B 2D)369 d_mlp = (in_projected_states.shape[-1] - d_to_remove) // 2370 split_projection_dim = [d_mlp, d_mlp, self.intermediate_size, self.conv_dim, self.num_heads]371 _, _, gate, hidden_states_B_C, dt = torch.split(in_projected_states, split_projection_dim, dim=-1)372 373 hidden_states_B_C = causal_conv1d_update(374 hidden_states_B_C,375 cache_params.conv_states[self.layer_idx],376 self.conv1d.weight.squeeze(1),377 self.conv1d.bias,378 self.activation,379 )380 381 hidden_states, B, C = torch.split(382 hidden_states_B_C,383 [self.intermediate_size, groups_time_state_size, groups_time_state_size],384 dim=-1,385 )386 A = -torch.exp(self.A_log.float()) # (nheads,)387 388 A = A[:, None, ...][:, :, None].expand(-1, self.head_dim, self.ssm_state_size).to(dtype=torch.float32)389 dt = dt[:, :, None].expand(-1, -1, self.head_dim)390 dt_bias = self.dt_bias[:, None, ...].expand(-1, self.head_dim)391 D = self.D[:, None, ...].expand(-1, self.head_dim)392 B = B.view(batch_size, self.n_groups, B.shape[1] // self.n_groups)393 C = C.view(batch_size, self.n_groups, C.shape[1] // self.n_groups)394 hidden_states_reshaped = hidden_states.view(batch_size, self.num_heads, self.head_dim)395 hidden_states = selective_state_update(396 cache_params.ssm_states[self.layer_idx],397 hidden_states_reshaped,398 dt,399 A,400 B,401 C,402 D,403 z=None,404 dt_bias=dt_bias,405 dt_softplus=True,406 )407 hidden_states = hidden_states.view(batch_size, self.num_heads * self.head_dim)408 hidden_states = self.norm(hidden_states, gate)409 out = self.out_proj(hidden_states)[:, None, ...]410 # if no cache is found, calling the kernel411 else:412 if attention_mask is not None and not torch.all(attention_mask == 1):413 # tune out hidden states for pad tokens, see https://github.com/state-spaces/mamba/issues/66414 dtype = hidden_states.dtype415 hidden_states = (hidden_states * attention_mask[:, :, None]).to(dtype)416 # 1. Gated MLP's linear projection417 projected_states = self.in_proj(hidden_states)418 A = -torch.exp(self.A_log.float()) # (num_heads) or (intermediate_size, state_size)419 dt_limit_kwargs = {} if self.time_step_limit is None else {"dt_limit": self.time_step_limit}420 if attention_mask is not None:421 input_not_masked = torch.all(attention_mask == 1)422 else:423 input_not_masked = True424 425 if self.use_mem_eff_path and self.training and cache_params is None and input_not_masked:426 out, ssm_state = mamba_split_conv1d_scan_combined(427 projected_states,428 self.conv1d.weight.squeeze(1),429 self.conv1d.bias,430 self.dt_bias,431 A,432 D=self.D,433 chunk_size=self.chunk_size,434 seq_idx=None,435 activation=self.activation,436 rmsnorm_weight=self.norm.weight,437 rmsnorm_eps=self.norm.variance_epsilon,438 outproj_weight=self.out_proj.weight,439 outproj_bias=self.out_proj.bias,440 headdim=self.head_dim,441 ngroups=self.n_groups,442 norm_before_gate=False,443 return_final_states=True,444 **dt_limit_kwargs,445 )446 447 else:448 gate, hidden_states_B_C, time_step = torch.split(449 projected_states,450 [self.intermediate_size, self.conv_dim, self.num_heads],451 dim=-1,452 )453 454 # 1D Convolution455 if cache_params is not None:456 hidden_states_B_C_t = hidden_states_B_C.transpose(1, 2)457 conv_state = nn.functional.pad(458 hidden_states_B_C_t, (self.conv_kernel_size - hidden_states_B_C_t.shape[-1], 0)459 )460 cache_params.conv_states[self.layer_idx].copy_(conv_state)461 if causal_conv1d_fn is None or self.activation not in ["silu", "swish"]:462 hidden_states_B_C = self.act(463 self.conv1d(hidden_states_B_C.transpose(1, 2)).transpose(1, 2)[:, :seq_len]464 ) # (B, L, self.d_inner + 2 * ngroups * d_state)465 else:466 hidden_states_B_C = causal_conv1d_fn(467 x=hidden_states_B_C.transpose(1, 2),468 weight=self.conv1d.weight.squeeze(1),469 bias=self.conv1d.bias,470 activation=self.activation,471 ).transpose(1, 2)[:, :seq_len]472 hidden_states, B, C = torch.split(473 hidden_states_B_C,474 [self.intermediate_size, groups_time_state_size, groups_time_state_size],475 dim=-1,476 )477 if attention_mask is not None and not torch.all(attention_mask == 1):478 # tune out hidden states for pad tokens, see https://github.com/state-spaces/mamba/issues/66479 dtype = hidden_states.dtype480 hidden_states = (hidden_states * attention_mask[:, :, None]).to(dtype)481 scan_output, ssm_state = mamba_chunk_scan_combined(482 hidden_states.view(batch_size, seq_len, -1, self.head_dim),483 time_step,484 A,485 B.view(batch_size, seq_len, self.n_groups, -1),486 C.view(batch_size, seq_len, self.n_groups, -1),487 chunk_size=self.chunk_size,488 D=self.D,489 z=None,490 seq_idx=None,491 return_final_states=True,492 dt_bias=self.dt_bias,493 dt_softplus=True,494 **dt_limit_kwargs,495 )496 if ssm_state is not None and cache_params is not None:497 cache_params.ssm_states[self.layer_idx].copy_(ssm_state)498 scan_output = scan_output.view(batch_size, seq_len, -1)499 # Multiply "gate" branch and apply extra normalization layer500 scan_output = self.norm(scan_output, gate)501 out = self.out_proj(scan_output)502 return out503 504 # fmt: off505 def torch_forward(self, input_states, cache_params: Optional[Zamba2HybridDynamicCache]=None, attention_mask: Optional[torch.Tensor]=None):506 batch_size, seq_len, _ = input_states.shape507 dtype = input_states.dtype508 # Gated MLP's linear projection509 if cache_params is not None and cache_params.has_previous_state:510 projected_states = self.in_proj(input_states.squeeze(1))511 else:512 if attention_mask is not None and not torch.all(attention_mask==1):513 # tune out hidden states for pad tokens, see https://github.com/state-spaces/mamba/issues/66514 input_states = (input_states * attention_mask[:, :, None]).to(dtype)515 projected_states = self.in_proj(input_states)516 d_mlp = (projected_states.shape[-1] - 2 * self.intermediate_size - 2 * self.n_groups * self.ssm_state_size- self.num_heads) // 2517 _, _, gate, hidden_states, dt = projected_states.split(518 [d_mlp, d_mlp, self.intermediate_size, self.conv_dim, self.num_heads], dim=-1519 )520 521 # Convolution sequence transformation522 if cache_params is not None:523 ssm_state = cache_params.ssm_states[self.layer_idx].clone()524 ssm_state = ssm_state.to(hidden_states.device)525 if cache_params.has_previous_state:526 gate = gate.unsqueeze(1)527 conv_state = cache_params.conv_states[self.layer_idx] # [batch, intermediate_size, conv_kernel_size]528 conv_state = torch.roll(conv_state, shifts=-1, dims=-1)529 # handle batched generation - states are copied through530 conv_state[:, :, -1] = hidden_states[:, 0, :] if hidden_states.ndim == 3 else hidden_states531 cache_params.conv_states[self.layer_idx].copy_(conv_state)532 hidden_states = torch.sum(conv_state.to(projected_states.device) * self.conv1d.weight[:, 0, :], dim=-1)533 if self.use_conv_bias:534 hidden_states += self.conv1d.bias535 hidden_states = self.act(hidden_states).to(dtype)[:, None, ...] # [batch, 1, intermediate_size] : decoding536 else:537 hidden_states = hidden_states.transpose(1,2)538 conv_state = nn.functional.pad(539 hidden_states,540 (self.conv_kernel_size - hidden_states.shape[-1], 0)541 )542 cache_params.conv_states[self.layer_idx].copy_(conv_state)543 hidden_states = self.act(self.conv1d(hidden_states).transpose(1,2))[:, :seq_len, :] # [batch, intermediate_size, seq_len]544 if attention_mask is not None and not torch.all(attention_mask==1):545 dtype = hidden_states.dtype546 # tune out hidden states for pad tokens, see https://github.com/state-spaces/mamba/issues/66547 hidden_states = (hidden_states * attention_mask[:, :, None]).to(dtype)548 else:549 ssm_state = torch.zeros(550 (batch_size, self.num_heads, self.head_dim, self.ssm_state_size),551 device=hidden_states.device, dtype=dtype552 )553 hidden_states = self.act(self.conv1d(hidden_states.transpose(1, 2))[..., :seq_len].transpose(1, 2))554 hidden_states, B, C = torch.split(hidden_states, [self.intermediate_size, self.n_groups * self.ssm_state_size, self.n_groups * self.ssm_state_size], dim=-1)555 A = -torch.exp(self.A_log.float()) # [num_heads]556 if cache_params is not None and cache_params.has_previous_state:557 # Note: there is no need to pad parameter matrices here, as there is just one new token558 # for batched generation559 dt = dt[:, None, ...] if dt.ndim == 2 else dt[:, 0, :][:, None, ...]560 dt = dt.transpose(1, 2).expand(batch_size, dt.shape[-1], self.head_dim)561 # [num_heads] -> [num_heads, head_dim]562 dt_bias = self.dt_bias[..., None].expand(self.dt_bias.shape[0], self.head_dim)563 564 dt = torch.nn.functional.softplus(dt + dt_bias.to(dt.dtype))565 dt = torch.clamp(dt, self.time_step_min) #, self.time_step_max)566 A = A[..., None, None].expand(self.num_heads, self.head_dim, self.ssm_state_size).to(dtype=torch.float32)567 # [bsz, num_heads, head_dim, state_size]568 dA = torch.exp(dt[..., None] * A)569 570 # Discretize B571 # [bsz, n_groups * state_size] -> [bsz, n_groups, 1, state_size] ->572 # -> [bsz, n_groups, group to head repetition factor, state_size] -> [bsz, num_heads, state_size]573 B = B.reshape(batch_size, self.n_groups, -1)[..., None, :]574 B = B.expand(batch_size, self.n_groups, self.num_heads // self.n_groups, B.shape[-1]).contiguous()575 B = B.reshape(batch_size, -1, B.shape[-1])576 # [bsz, num_heads, head_dim, state_size]577 dB = dt[..., None] * B[..., None, :]578 579 # Discretize x into dB580 # [bsz, intermediate_size] -> [bsz, num_heads, head_dim]581 hidden_states = hidden_states.reshape(batch_size, -1, self.head_dim)582 dBx = dB * hidden_states[..., None]583 584 # State calculation585 cache_params.ssm_states[self.layer_idx].copy_(586 cache_params.ssm_states[self.layer_idx] * dA + dBx587 )588 589 # Subsequent output590 # [bsz, n_groups * state_size] -> [bsz, num_heads, state_size]591 C = C.reshape(batch_size, self.n_groups, -1)[..., None, :]592 C = C.expand(batch_size, self.n_groups, self.num_heads // self.n_groups, C.shape[-1]).contiguous()593 C = C.reshape(batch_size, -1, C.shape[-1])594 # [bsz, num_heads, head_dim]595 596 ssm_states = cache_params.ssm_states[self.layer_idx].to(C.dtype) # Shape: [b, h, d, n]597 # Reshape ssm_states to merge the first two dimensions598 ssm_states_reshaped = ssm_states.view(batch_size * self.num_heads, self.head_dim, self.ssm_state_size) # Shape: [b*h, d, n]599 C_reshaped = C.view(batch_size * self.num_heads, self.ssm_state_size, 1) # Shape: [b*h, n, 1]600 y = torch.bmm(ssm_states_reshaped, C_reshaped)601 y = y.view(batch_size, self.num_heads, self.head_dim)602 603 # D skip connection604 # [num_heads] -> [num_heads, head_dim]605 D = self.D[..., None].expand(self.D.shape[0], self.head_dim)606 y = (y + hidden_states * D).to(y.dtype)607 608 # [bsz, num_heads, head_dim] -> [bsz, 1, intermediate_size]609 y = y.reshape(batch_size, -1)[:, None, ...]610 else:611 # begin ssd naive implementation without einsums612 dt = nn.functional.softplus(dt + self.dt_bias)613 dt = torch.clamp(dt, self.time_step_min)614 hidden_states = hidden_states.reshape(batch_size, seq_len, -1, self.head_dim).float()615 B = B.reshape(batch_size, seq_len, -1, self.ssm_state_size).float()616 C = C.reshape(batch_size, seq_len, -1, self.ssm_state_size).float()617 B = B.repeat_interleave(self.num_heads // self.n_groups, dim=2, output_size=self.num_heads)618 C = C.repeat_interleave(self.num_heads // self.n_groups, dim=2, output_size=self.num_heads)619 pad_size = (self.chunk_size - seq_len % self.chunk_size) % self.chunk_size620 621 D_residual = self.D[..., None] * pad_tensor_by_size(hidden_states, pad_size)622 623 # Discretize x and A624 hidden_states = hidden_states * dt[..., None]625 A = A.to(hidden_states.dtype) * dt626 627 # Rearrange into blocks/chunks628 hidden_states, A, B, C = [reshape_into_chunks(t, pad_size, self.chunk_size) for t in (hidden_states, A, B, C)]629 630 631 # [bsz, -1, chunk_size, num_heads] -> [bsz, num_heads, -1, chunk_size]632 A = A.permute(0, 3, 1, 2)633 A_cumsum = torch.cumsum(A, dim=-1)634 635 # 1. Compute the output for each intra-chunk (diagonal blocks)636 # This is the analog of a causal mask637 L = torch.exp(segment_sum(A))638 639 # First, contraction of C and B to get G (attention-weights like)640 G_intermediate = C[:, :, :, None, :, :] * B[:, :, None, :, : ,:] # shape: (b, c, l, s, h, n)641 G = G_intermediate.sum(dim=-1) # shape: (b, c, l, s, h)642 643 644 # Step 2: Compute M, equivalent to applying attention mask to weights645 M_intermediate = G[..., None] * L.permute(0, 2, 3, 4, 1)[..., None]646 M = M_intermediate.sum(dim=-1)647 648 # Step 3: Compute Y_diag (apply to values)649 Y_diag = (M[..., None] * hidden_states[:, :, None]).sum(3)650 651 # (right term of low-rank factorization of off-diagonal blocks; B terms)652 653 decay_states = torch.exp(A_cumsum[:, :, :, -1:] - A_cumsum)654 B_decay_contraction = B * decay_states.permute(0, 2, 3, 1)[..., None]655 # permute back B * decay states656 states = (B_decay_contraction.permute(0, 1, 3, 2, 4)[..., None] * hidden_states.permute(0, 1, 3, 2, 4)[..., None, :]).sum(dim=3).permute(0, 1, 2, 4, 3)657 if cache_params is not None and cache_params.has_previous_state:658 previous_states = cache_params.ssm_states[self.layer_idx][:, None, ...]659 else:660 previous_states = torch.zeros_like(states[:, :1])661 states = torch.cat([previous_states, states], dim=1)662 decay_chunk = torch.exp(segment_sum(nn.functional.pad(A_cumsum[:, :, :, -1], (1, 0))))663 664 states_permuted = states.permute(0, 2, 1, 3, 4)665 result = (decay_chunk[..., None, None] * states_permuted[:, :, None, ...]).sum(dim=2)666 new_states = result.permute(0, 2, 1, 3, 4)667 states, ssm_state = new_states[:, :-1], new_states[:, -1]668 669 # Compute state -> output conversion per chunk670 # (left term of low-rank factorization of off-diagonal blocks; C terms)671 state_decay_out = torch.exp(A_cumsum)672 # compute Yoff673 C_times_states = (C[..., None, :] * states[:, :, None, ...])674 state_decay_out_permuted = state_decay_out.permute(0, 2, 3, 1)675 Y_off = (C_times_states.sum(-1) * state_decay_out_permuted[..., None])676 # Add output of intra-chunk and inter-chunk terms (diagonal and off-diagonal blocks)677 678 y = Y_diag + Y_off679 # [bsz, -1, self.chunk_size, num_heads, head_dim] -> [bsz, (padded) seq_len, num_heads, head_dim]680 y = y.reshape(batch_size, -1, self.num_heads, self.head_dim)681 682 y = y + D_residual683 # Cutting off padded chunks684 if pad_size > 0:685 y = y[:, :seq_len, :, :]686 y = y.reshape(batch_size, seq_len, -1)687 if ssm_state is not None and cache_params is not None:688 cache_params.ssm_states[self.layer_idx].copy_(ssm_state)689 690 scan_output = self.norm(y, gate)691 692 # end ssd naive693 694 # 4. Final linear projection695 contextualized_states = self.out_proj(scan_output.to(dtype)) # [batch, seq_len, hidden_size]696 return contextualized_states697 # fmt: on698 699 def forward(700 self,701 hidden_states,702 cache_params: Optional[Zamba2HybridDynamicCache] = None,703 attention_mask: Optional[torch.Tensor] = None,704 ):705 if is_fast_path_available and "cuda" in self.in_proj.weight.device.type:706 return self.cuda_kernels_forward(hidden_states, cache_params, attention_mask)707 708 return self.torch_forward(hidden_states, cache_params, attention_mask)709 710 711class Zamba2MLP(nn.Module):712 def __init__(self, config: Zamba2Config, num_fwd_mem_blocks=None, block_id: Optional[int] = None):713 """714 This MLP layer contributes to tied transformer blocks aimed to increasing compute without increasing model size. Because this layer715 is tied, un-tied adapter modules (formally same as LoRA, but used in the base model) are added to the up and gate projectors to increase expressivity with a small memory overhead.716 """717 super().__init__()718 self.config = config719 self.hidden_size = config.hidden_size720 self.intermediate_size = config.intermediate_size721 self.num_fwd_mem_blocks = num_fwd_mem_blocks722 self.block_id = block_id723 724 self.gate_up_proj = nn.Linear(self.hidden_size, 2 * self.intermediate_size, bias=config.add_bias_linear)725 self.down_proj = nn.Linear(self.intermediate_size, self.hidden_size, bias=config.add_bias_linear)726 self.act_fn = ACT2FN[config.hidden_act]727 728 self.gate_up_proj_adapter_list = nn.ModuleList([])729 for i in range(self.num_fwd_mem_blocks):730 if i % config.num_mem_blocks == block_id:731 gate_up_proj_adapter = nn.Sequential(732 nn.Linear(self.config.hidden_size, self.config.adapter_rank, bias=False),733 nn.Linear(self.config.adapter_rank, 2 * self.intermediate_size, bias=False),734 )735 else:736 gate_up_proj_adapter = nn.Identity()737 self.gate_up_proj_adapter_list.append(gate_up_proj_adapter)738 739 layer_block_map = config.hybrid_layer_ids740 self.layer_dic = {value: index for index, value in enumerate(layer_block_map)}741 742 def forward(self, hidden_state, layer_idx=None):743 gate_up_state = self.gate_up_proj(hidden_state)744 layer_idx = self.layer_dic[layer_idx]745 gate_up_state = gate_up_state + self.gate_up_proj_adapter_list[layer_idx](hidden_state)746 747 gate_up_state = torch.chunk(gate_up_state, 2, dim=-1)748 hidden_state = self.act_fn(gate_up_state[0]) * gate_up_state[1]749 output = self.down_proj(hidden_state)750 return output751 752 753class Zamba2AttentionDecoderLayer(ZambaAttentionDecoderLayer):754 def __init__(self, config: Zamba2Config, block_id: Optional[int] = None, layer_idx: Optional[int] = None):755 self.block_id = block_id756 num_gs = len(config.hybrid_layer_ids)757 super().__init__(config, layer_idx)758 self.self_attn = Zamba2Attention(config, layer_idx=-1, num_fwd_mem_blocks=num_gs, block_id=block_id)759 self.feed_forward = Zamba2MLP(config, num_fwd_mem_blocks=num_gs, block_id=block_id)760 761 @deprecate_kwarg("past_key_value", new_name="past_key_values", version="4.58")762 def forward(763 self,764 hidden_states: torch.Tensor,765 original_hidden_states: torch.Tensor,766 layer_idx: int,767 attention_mask: Optional[torch.Tensor] = None,768 past_key_values: Optional[Zamba2HybridDynamicCache] = None,769 output_attentions: Optional[bool] = False,770 position_embeddings: Optional[torch.LongTensor] = None,771 **kwargs: Unpack[FlashAttentionKwargs],772 ) -> tuple[torch.FloatTensor, Optional[tuple[torch.FloatTensor, torch.FloatTensor]]]:773 """774 Args:775 hidden_states (`torch.FloatTensor`): output of previous Mamba layer of shape `(batch, seq_len, embed_dim)`776 original_hidden_states (`torch.FloatTensor`): word embedding output of shape `(batch, seq_len, embed_dim)`.777 This is concatenated with `hidden_states` (which is the output of the previous (mamba) layer). The778 concatenated tensor is then used as input of the pre-attention RMSNorm779 (see fig. 2 in https://huggingface.co/papers/2405.16712).780 attention_mask (`torch.FloatTensor`, *optional*): attention mask of size781 `(batch, sequence_length)` where padding elements are indicated by 0.782 past_key_values (`Zamba2HybridDynamicCache`, *optional*): cached past key and value projection states783 output_attentions (`bool`, *optional*):784 Whether or not to return the attentions tensors of all attention layers. See `attentions` under785 returned tensors for more detail.786 use_cache (`bool`, *optional*):787 If set to `True`, `past_key_values` key value states are returned and can be used to speed up decoding788 (see `past_key_values`).789 position_embeddings (`tuple[torch.FloatTensor, torch.FloatTensor]`, *optional*):790 Tuple containing the cosine and sine positional embeddings of shape `(batch_size, seq_len, head_dim)`,791 with `head_dim` being the embedding dimension of each attention head.792 """793 hidden_states = torch.concatenate([hidden_states, original_hidden_states], dim=-1)794 hidden_states = self.input_layernorm(hidden_states)795 hidden_states, self_attn_weights = self.self_attn(796 hidden_states=hidden_states,797 layer_idx=layer_idx,798 attention_mask=attention_mask,799 past_key_values=past_key_values,800 output_attentions=output_attentions,801 position_embeddings=position_embeddings,802 **kwargs,803 )804 805 hidden_states = self.pre_ff_layernorm(hidden_states)806 hidden_states = self.feed_forward(hidden_states, layer_idx)807 808 outputs = (hidden_states,)809 810 if output_attentions:811 outputs += (self_attn_weights,)812 813 return outputs814 815 816class Zamba2MambaDecoderLayer(ZambaMambaDecoderLayer):817 def __init__(self, config: Zamba2Config, layer_idx: int):818 super().__init__(config, layer_idx)819 self.mamba = Zamba2MambaMixer(config=config, layer_idx=layer_idx)820 self.input_layernorm = Zamba2RMSNorm(config.hidden_size, eps=config.rms_norm_eps)821 822 823class Zamba2HybridLayer(ZambaHybridLayer):824 def __init__(825 self, shared_transformer: Zamba2AttentionDecoderLayer, linear: nn.Linear, mamba: Zamba2MambaDecoderLayer826 ):827 super().__init__(shared_transformer, linear, mamba)828 del self.shared_transf829 self.shared_transformer = shared_transformer830 831 @deprecate_kwarg("past_key_value", new_name="past_key_values", version="4.58")832 def forward(833 self,834 hidden_states: torch.Tensor,835 original_hidden_states: Optional[torch.Tensor] = None,836 layer_idx: Optional[int] = None,837 attention_mask: Optional[torch.Tensor] = None,838 causal_mask: Optional[torch.Tensor] = None,839 past_key_values: Optional[Zamba2HybridDynamicCache] = None,840 output_attentions: Optional[bool] = False,841 use_cache: Optional[bool] = False,842 position_embeddings: Optional[torch.LongTensor] = None,843 ) -> tuple[torch.FloatTensor, Optional[tuple[torch.FloatTensor, torch.FloatTensor]]]:844 """845 Args:846 hidden_states (`torch.FloatTensor`): input to the layer of shape `(batch, seq_len, embed_dim)`847 original_hidden_states (`torch.FloatTensor`): word embedding output that will be concatenated with848 hidden activations to form the input of the shared transformer layer.849 layer_idx (`int`): layer number.850 attention_mask (`torch.FloatTensor`, *optional*): attention mask of size851 `(batch, sequence_length)` where padding elements are indicated by 0.852 past_key_values (`Zamba2HybridDynamicCache`, *optional*): cached past key and value projection states853 output_attentions (`bool`, *optional*):854 Whether or not to return the attentions tensors of all attention layers. See `attentions` under855 returned tensors for more detail.856 use_cache (`bool`, *optional*):857 If set to `True`, `past_key_values` key value states are returned and can be used to speed up decoding858 (see `past_key_values`).859 position_embeddings (`tuple[torch.FloatTensor, torch.FloatTensor]`, *optional*):860 Tuple containing the cosine and sine positional embeddings of shape `(batch_size, seq_len, head_dim)`,861 with `head_dim` being the embedding dimension of each attention head.862 """863 864 layer_outputs = self.shared_transformer(865 hidden_states,866 original_hidden_states=original_hidden_states,867 layer_idx=layer_idx,868 attention_mask=causal_mask,869 past_key_values=past_key_values,870 output_attentions=output_attentions,871 position_embeddings=position_embeddings,872 )873 874 transformer_hidden_states = layer_outputs[0]875 876 if output_attentions:877 self_attn_weights = layer_outputs[1]878 879 transformer_hidden_states = self.linear(transformer_hidden_states)880 881 layer_outputs = self.mamba_decoder(882 hidden_states,883 transformer_hidden_states=transformer_hidden_states,884 attention_mask=attention_mask,885 past_key_values=past_key_values,886 output_attentions=output_attentions,887 use_cache=use_cache,888 position_embeddings=position_embeddings,889 )890 891 if output_attentions:892 layer_outputs = (layer_outputs[0], self_attn_weights) + layer_outputs[2:]893 894 return layer_outputs895 896 897class Zamba2PreTrainedModel(PreTrainedModel):898 config: Zamba2Config899 base_model_prefix = "model"900 supports_gradient_checkpointing = True901 _no_split_modules = ["Zamba2AttentionDecoderLayer", "Zamba2MambaDecoderLayer"]902 _skip_keys_device_placement = "past_key_values"903 _supports_flash_attn = True904 _supports_flex_attn = True905 _supports_sdpa = True906 # Note: only supports Zamba2HybridDynamicCache907 _is_stateful = True908 909 def _init_weights(self, module):910 super()._init_weights(module)911 if isinstance(module, Zamba2MambaMixer):912 dt = torch.exp(913 torch.rand(self.config.n_mamba_heads)914 * (math.log(self.config.time_step_max) - math.log(self.config.time_step_min))915 + math.log(self.config.time_step_min)916 ).clamp(min=self.config.time_step_floor)917 # # Inverse of softplus: https://github.com/pytorch/pytorch/issues/72759918 inv_dt = dt + torch.log(-torch.expm1(-dt))919 module.dt_bias.data.copy_(inv_dt)920 921 A = torch.arange(1, module.num_heads + 1)922 module.A_log.data.copy_(torch.log(A))923 module.D.data.fill_(1.0)924 925 926class Zamba2Model(ZambaModel, Zamba2PreTrainedModel):927 """928 Model consisting of *config.num_hidden_layers* layers.929 930 Args:931 config: Zamba2Config932 """933 934 def __init__(self, config: Zamba2Config):935 Zamba2PreTrainedModel.__init__(self, config)936 self.config = config937 self.padding_idx = config.pad_token_id938 self.vocab_size = config.vocab_size939 940 self.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size, self.padding_idx)941 blocks = [Zamba2AttentionDecoderLayer(config, block_id=k) for k in range(config.num_mem_blocks)]942 mamba_layers = []943 linear_layers = []944 self.layers_block_type = config.layers_block_type945 for i in range(config.num_hidden_layers):946 if config.layers_block_type[i] == "mamba":947 mamba_layers.append(Zamba2MambaDecoderLayer(config, layer_idx=i))948 elif config.layers_block_type[i] == "hybrid":949 linear_layers.append(nn.Linear(self.config.hidden_size, self.config.hidden_size, bias=False))950 mamba_layers.append(Zamba2MambaDecoderLayer(config, layer_idx=i))951 mamba_layers = iter(mamba_layers)952 linear_layers = iter(linear_layers)953 blocks = cycle(blocks)954 layers = self.get_layers(blocks, linear_layers, mamba_layers)955 self.layers = nn.ModuleList(layers)956 957 self._attn_implementation = config._attn_implementation958 self.final_layernorm = Zamba2RMSNorm(config.hidden_size, eps=config.rms_norm_eps)959 if config.use_mem_rope:960 if config.use_long_context:961 logger.warning_once(962 "`use_long_context` set to `True`: using rescaled `rope_theta` and extended `max_position_embeddings`."963 )964 self.rotary_emb = Zamba2RotaryEmbedding(config)965 self.gradient_checkpointing = False966 967 # Initialize weights and apply final processing968 self.post_init()969 970 def get_layers(self, blocks, linear_layers, mamba_layers):971 layers = []972 self._tied_weights_keys = []973 self.first_transformer_layer_id = 0974 for layer_id, layer_type in enumerate(self.layers_block_type):975 if layer_type == "hybrid":976 if self.first_transformer_layer_id == 0:977 self.first_transformer_layer_id = layer_id978 block = next(blocks)979 if self.config.num_mem_blocks * len(self.config.hybrid_layer_ids) > 1:980 prefix_pattern = rf"^layers\.{layer_id}\.shared_transformer\."981 main_keys_pattern = re.compile(982 prefix_pattern983 + r"(?:"984 + r"self_attn\.(?:q_proj|k_proj|v_proj|o_proj)\.weight|"985 + r"feed_forward\.(?:gate_up_proj|down_proj)\.weight|"986 + r"(?:input_layernorm|pre_ff_layernorm)\.weight"987 + r")$"988 )989 self._tied_weights_keys.append(main_keys_pattern)990 991 adapter_id = 0992 for _layer_type in self.layers_block_type:993 if _layer_type == "hybrid" and adapter_id % self.config.num_mem_blocks == block.block_id:994 adapter_pattern = re.compile(995 r"^shared_transformer\.feed_forward\.gate_up_proj_adapter_list\."996 + str(adapter_id)997 + r"\.(?:0|1)\.weight$"998 )999 self._tied_weights_keys.append(adapter_pattern)1000 adapter_id += 11001 if self.config.use_shared_attention_adapter:1002 adapter_id = 01003 for _layer_type in self.layers_block_type:1004 if _layer_type == "hybrid" and adapter_id % self.config.num_mem_blocks == block.block_id:1005 attn_adapter_pattern = re.compile(1006 r"^shared_transformer\.self_attn\."1007 + r"(?:linear_q_adapter_list|linear_k_adapter_list|linear_v_adapter_list)\."1008 + str(adapter_id)1009 + r"\.(?:0|1)\.weight$"1010 )1011 self._tied_weights_keys.append(attn_adapter_pattern)1012 adapter_id += 11013 layers.append(Zamba2HybridLayer(block, next(linear_layers), next(mamba_layers)))1014 else:1015 layers.append(next(mamba_layers))1016 return layers1017 1018 def forward(1019 self,1020 input_ids: Optional[torch.LongTensor] = None,1021 attention_mask: Optional[torch.Tensor] = None,1022 position_ids: Optional[torch.LongTensor] = None,1023 past_key_values: Optional[Zamba2HybridDynamicCache] = None,1024 inputs_embeds: Optional[torch.FloatTensor] = None,1025 use_cache: Optional[bool] = None,1026 output_attentions: Optional[bool] = None,1027 output_hidden_states: Optional[bool] = None,1028 return_dict: Optional[bool] = None,1029 cache_position: Optional[torch.LongTensor] = None,1030 ) -> Union[tuple, BaseModelOutputWithPast]:1031 output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions1032 output_hidden_states = (1033 output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states1034 )1035 use_cache = use_cache if use_cache is not None else self.config.use_cache1036 1037 return_dict = return_dict if return_dict is not None else self.config.use_return_dict1038 1039 if (input_ids is None) ^ (inputs_embeds is not None):1040 raise ValueError(1041 "You cannot specify both input_ids and inputs_embeds at the same time, and must specify either one"1042 )1043 1044 if self.gradient_checkpointing and self.training and use_cache:1045 logger.warning_once(1046 "`use_cache=True` is incompatible with gradient checkpointing. Setting `use_cache=False`."1047 )1048 use_cache = False1049 1050 if inputs_embeds is None:1051 inputs_embeds = self.embed_tokens(input_ids)1052 1053 hidden_states = inputs_embeds1054 1055 original_hidden_states = torch.clone(inputs_embeds)1056 # original_hidden_states: word embedding output that will be concatenated with hidden activations to form the input of the shared transformer layer1057 1058 if use_cache and past_key_values is None:1059 batch_size = input_ids.shape[0] if input_ids is not None else inputs_embeds.shape[0]1060 past_key_values = Zamba2HybridDynamicCache(self.config, batch_size, dtype=self.dtype, device=self.device)1061 1062 if cache_position is None:1063 past_seen_tokens = (1064 past_key_values.get_seq_length(layer_idx=self.first_transformer_layer_id)1065 if past_key_values is not None1066 else 01067 )1068 cache_position = torch.arange(1069 past_seen_tokens, past_seen_tokens + inputs_embeds.shape[1], device=inputs_embeds.device1070 )1071 if position_ids is None:1072 position_ids = cache_position.unsqueeze(0)1073 1074 causal_mask = self._update_causal_mask(attention_mask, inputs_embeds, cache_position)1075 1076 # create position embeddings to be shared across the decoder layers1077 if self.config.use_mem_rope:1078 position_embeddings = self.rotary_emb(hidden_states, position_ids)1079 else:1080 position_embeddings = None1081 1082 all_hidden_states = () if output_hidden_states else None1083 all_self_attns = () if output_attentions else None1084 1085 for layer_idx, layer in enumerate(self.layers):1086 if output_hidden_states:1087 all_hidden_states += (hidden_states,)1088 1089 if self.gradient_checkpointing and self.training:1090 layer_outputs = self._gradient_checkpointing_func(1091 layer.__call__,1092 hidden_states,1093 original_hidden_states,1094 layer_idx,1095 attention_mask,1096 causal_mask,1097 past_key_values,1098 output_attentions,1099 use_cache,1100 position_embeddings,1101 )1102 else:1103 layer_outputs = layer(1104 hidden_states,1105 original_hidden_states=original_hidden_states,1106 layer_idx=layer_idx,1107 attention_mask=attention_mask,1108 causal_mask=causal_mask,1109 past_key_values=past_key_values,1110 output_attentions=output_attentions,1111 use_cache=use_cache,1112 position_embeddings=position_embeddings,1113 )1114 hidden_states = layer_outputs[0]1115 1116 if output_attentions:1117 if layer_outputs[1] is not None:1118 # append attentions only of attention layers. Mamba layers return `None` as the attention weights1119 all_self_attns += (layer_outputs[1],)1120 1121 hidden_states = self.final_layernorm(hidden_states)1122 1123 # add hidden states from the last decoder layer1124 if output_hidden_states:1125 all_hidden_states += (hidden_states,)1126 1127 if past_key_values is not None and not past_key_values.has_previous_state:1128 past_key_values.has_previous_state = True1129 1130 output = BaseModelOutputWithPast(1131 last_hidden_state=hidden_states,1132 past_key_values=past_key_values if use_cache else None,1133 hidden_states=all_hidden_states,1134 attentions=all_self_attns,1135 )1136 return output if return_dict else output.to_tuple()1137 1138 1139class Zamba2ForCausalLM(ZambaForCausalLM):1140 pass1141 1142 1143class Zamba2ForSequenceClassification(ZambaForSequenceClassification):1144 pass1145 1146 1147__all__ = [1148 "Zamba2ForCausalLM",1149 "Zamba2ForSequenceClassification",1150 "Zamba2Model",1151 "Zamba2PreTrainedModel",1152]1153