# 🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨 # This file was automatically generated from src/transformers/models/zamba2/modular_zamba2.py. # Do NOT edit this file manually as any edits will be overwritten by the generation of # the file from the modular. If any change should be done, please apply the change to the # modular_zamba2.py file directly. One of our CI enforces this. # 🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨 # Copyright 2024 Zyphra Technologies and the HuggingFace Inc. team. All rights reserved. # # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. # You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. import math from collections.abc import Callable from itertools import cycle import torch import torch.nn.functional as F from torch import nn from ... import initialization as init from ...activations import ACT2FN from ...cache_utils import Cache, DynamicCache from ...generation import GenerationMixin from ...integrations import use_kernel_forward_from_hub, use_kernel_func_from_hub_with_fallback, use_kernelized_func from ...integrations.accelerate import force_accelerate_hooks from ...masking_utils import create_causal_mask, create_recurrent_attention_mask from ...modeling_layers import GradientCheckpointingLayer from ...modeling_outputs import BaseModelOutputWithPast, CausalLMOutputWithPast, SequenceClassifierOutputWithPast from ...modeling_rope_utils import ROPE_INIT_FUNCTIONS, dynamic_rope_update from ...modeling_utils import ALL_ATTENTION_FUNCTIONS, PreTrainedModel from ...processing_utils import Unpack from ...utils import TransformersKwargs, auto_docstring, can_return_tuple, logging from ...utils.deprecation import deprecate_kwarg from ...utils.generic import maybe_autocast, merge_with_config_defaults from ...utils.output_capturing import capture_outputs from .configuration_zamba2 import Zamba2Config logger = logging.get_logger(__name__) class Zamba2RMSNormGated(torch.nn.Module): def __init__(self, hidden_size, group_size, eps=1e-6): super().__init__() self.weight = nn.Parameter(torch.ones(hidden_size)) self.variance_epsilon = eps self.group_size = group_size def forward(self, hidden_states, gate=None): input_dtype = hidden_states.dtype hidden_states = hidden_states.to(torch.float32) if gate is not None: hidden_states = hidden_states * nn.functional.silu(gate.to(torch.float32)) *prefix_dims, last_dim = hidden_states.shape group_count = last_dim // self.group_size hidden_states_group = hidden_states.view(*prefix_dims, group_count, self.group_size) variance = hidden_states_group.pow(2).mean(-1, keepdim=True) hidden_states_group = hidden_states_group * torch.rsqrt(variance + self.variance_epsilon) hidden_states = hidden_states_group.view(*prefix_dims, group_count * self.group_size) return self.weight * hidden_states.to(input_dtype) class Zamba2RMSNorm(nn.Module): def __init__(self, hidden_size, eps: float = 1e-6) -> None: """ Zamba2RMSNorm is equivalent to T5LayerNorm """ super().__init__() self.weight = nn.Parameter(torch.ones(hidden_size)) self.variance_epsilon = eps def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: input_dtype = hidden_states.dtype hidden_states = hidden_states.to(torch.float32) variance = hidden_states.pow(2).mean(-1, keepdim=True) hidden_states = hidden_states * torch.rsqrt(variance + self.variance_epsilon) return self.weight * hidden_states.to(input_dtype) def extra_repr(self): return f"{tuple(self.weight.shape)}, eps={self.variance_epsilon}" class Zamba2RotaryEmbedding(nn.Module): @deprecate_kwarg("device", version="5.18") def __init__(self, config: Zamba2Config, device=None): super().__init__() self.max_seq_len_cached = config.max_position_embeddings self.original_max_seq_len = config.max_position_embeddings self.config = config self.rope_type = self.config.rope_parameters["rope_type"] rope_init_fn: Callable = self.compute_default_rope_parameters if self.rope_type != "default": rope_init_fn = ROPE_INIT_FUNCTIONS[self.rope_type] inv_freq, self.attention_scaling = rope_init_fn(self.config, device) self.inv_freq = nn.Buffer(inv_freq, persistent=False) self.original_inv_freq = nn.Buffer(inv_freq.clone(), persistent=False) @staticmethod @deprecate_kwarg("device", version="5.18") def compute_default_rope_parameters(config: Zamba2Config, device=None, **kwargs) -> tuple[torch.Tensor, float]: """ Computes the inverse frequencies according to the original RoPE implementation Args: config ([`~transformers.PreTrainedConfig`]): The model configuration. Returns: Tuple of (`torch.Tensor`, `float`), containing the inverse frequencies for the RoPE embeddings and the post-processing scaling factor applied to the computed cos/sin (unused in this type of RoPE). """ base = config.rope_parameters["rope_theta"] dim = getattr(config, "head_dim", None) or config.hidden_size // config.num_attention_heads attention_factor = 1.0 # Unused in this type of RoPE # Compute the inverse frequencies inv_freq = 1.0 / (base ** (torch.arange(0, dim, 2, dtype=torch.float) / dim)) return inv_freq.to(device), attention_factor @torch.no_grad() @dynamic_rope_update # power user: used with advanced RoPE types (e.g. dynamic rope) def forward(self, x, position_ids): inv_freq_expanded = ( self.inv_freq[None, :, None].expand(position_ids.shape[0], -1, 1).to(dtype=torch.float, device=x.device) ) position_ids_expanded = position_ids[:, None, :].float() device_type = x.device.type if isinstance(x.device.type, str) and x.device.type != "mps" else "cpu" # Disable any outside autocast context if any, to really force fp32 with maybe_autocast(device_type=device_type, enabled=False): freqs = (inv_freq_expanded @ position_ids_expanded).transpose(1, 2) emb = torch.cat((freqs, freqs), dim=-1) cos = emb.cos() * self.attention_scaling sin = emb.sin() * self.attention_scaling return cos.to(dtype=x.dtype), sin.to(dtype=x.dtype) def repeat_kv(hidden_states: torch.Tensor, n_rep: int) -> torch.Tensor: """ This is the equivalent of torch.repeat_interleave(x, dim=1, repeats=n_rep). The hidden states go from (batch, num_key_value_heads, seqlen, head_dim) to (batch, num_attention_heads, seqlen, head_dim) """ batch, num_key_value_heads, slen, head_dim = hidden_states.shape if n_rep == 1: return hidden_states hidden_states = hidden_states[:, :, None, :, :].expand(batch, num_key_value_heads, n_rep, slen, head_dim) return hidden_states.reshape(batch, num_key_value_heads * n_rep, slen, head_dim) def eager_attention_forward( module: nn.Module, query: torch.Tensor, key: torch.Tensor, value: torch.Tensor, attention_mask: torch.Tensor | None, scaling: float, dropout: float = 0.0, **kwargs, ): key_states = repeat_kv(key, module.num_key_value_groups) value_states = repeat_kv(value, module.num_key_value_groups) attn_weights = torch.matmul(query, key_states.transpose(2, 3)) * scaling if attention_mask is not None: attn_weights = attn_weights + attention_mask attn_weights = nn.functional.softmax(attn_weights, dim=-1, dtype=torch.float32).to(query.dtype) attn_weights = nn.functional.dropout(attn_weights, p=dropout, training=module.training) attn_output = torch.matmul(attn_weights, value_states) attn_output = attn_output.transpose(1, 2).contiguous() return attn_output, attn_weights def rotate_half(x): """Rotates half the hidden dims of the input.""" x1 = x[..., : x.shape[-1] // 2] x2 = x[..., x.shape[-1] // 2 :] return torch.cat((-x2, x1), dim=-1) @use_kernel_forward_from_hub("rotary_pos_emb") def apply_rotary_pos_emb(q, k, cos, sin, unsqueeze_dim=1): """Applies Rotary Position Embedding to the query and key tensors. Args: q (`torch.Tensor`): The query tensor. k (`torch.Tensor`): The key tensor. cos (`torch.Tensor`): The cosine part of the rotary embedding. sin (`torch.Tensor`): The sine part of the rotary embedding. unsqueeze_dim (`int`, *optional*, defaults to 1): The 'unsqueeze_dim' argument specifies the dimension along which to unsqueeze cos[position_ids] and sin[position_ids] so that they can be properly broadcasted to the dimensions of q and k. For example, note that cos[position_ids] and sin[position_ids] have the shape [batch_size, seq_len, head_dim]. Then, if q and k have the shape [batch_size, heads, seq_len, head_dim], then setting unsqueeze_dim=1 makes cos[position_ids] and sin[position_ids] broadcastable to the shapes of q and k. Similarly, if q and k have the shape [batch_size, seq_len, heads, head_dim], then set unsqueeze_dim=2. Returns: `tuple(torch.Tensor)` comprising of the query and key tensors rotated using the Rotary Position Embedding. """ cos = cos.unsqueeze(unsqueeze_dim) sin = sin.unsqueeze(unsqueeze_dim) q_embed = (q * cos) + (rotate_half(q) * sin) k_embed = (k * cos) + (rotate_half(k) * sin) return q_embed, k_embed class Zamba2Attention(nn.Module): """ Multi-headed attention from 'Attention Is All You Need' paper. Adapted from transformers.models.mistral.modeling_mistral.MistralAttention: The input dimension here is attention_hidden_size = 2 * hidden_size, and head_dim = attention_hidden_size // num_heads. The extra factor of 2 comes from the input being the concatenation of original_hidden_states with the output of the previous (mamba) layer (see fig. 2 in https://huggingface.co/papers/2405.16712). Additionally, replaced attn_weights = torch.matmul(query_states, key_states.transpose(2, 3)) / math.sqrt(self.head_dim) with attn_weights = torch.matmul(query_states, key_states.transpose(2, 3)) / math.sqrt(self.head_dim/2) Finally, this attention layer contributes to tied transformer blocks aimed to increasing compute without increasing model size. Because this 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 increase expressivity with a small memory overhead (see Fig. 2 of https://huggingface.co/papers/2411.15242). """ def __init__( self, config: Zamba2Config, layer_idx: int | None = None, num_fwd_mem_blocks: int | None = None, block_id: int | None = None, ): super().__init__() self.config = config self.layer_idx = layer_idx self.attention_hidden_size = config.attention_hidden_size self.head_dim = config.attention_head_dim self.num_key_value_groups = config.num_attention_heads // config.num_key_value_heads self.max_position_embeddings = config.max_position_embeddings self.scaling = (self.head_dim / 2) ** -0.5 self.is_causal = True self.attention_dropout = config.attention_dropout self.q_proj = nn.Linear(config.attention_hidden_size, config.num_attention_heads * self.head_dim, bias=False) self.k_proj = nn.Linear(config.attention_hidden_size, config.num_key_value_heads * self.head_dim, bias=False) self.v_proj = nn.Linear(config.attention_hidden_size, config.num_key_value_heads * self.head_dim, bias=False) self.o_proj = nn.Linear(config.num_attention_heads * self.head_dim, config.hidden_size, bias=False) self.num_fwd_mem_blocks = num_fwd_mem_blocks self.layer_block_map = config.hybrid_layer_ids self.block_id = block_id if config.use_shared_attention_adapter: self.linear_q_adapter_list = nn.ModuleList([]) self.linear_k_adapter_list = nn.ModuleList([]) self.linear_v_adapter_list = nn.ModuleList([]) for i in range(self.num_fwd_mem_blocks): if i % config.num_mem_blocks == block_id: linear_q_adapter = nn.Sequential( nn.Linear(self.attention_hidden_size, self.config.adapter_rank, bias=False), nn.Linear(self.config.adapter_rank, self.attention_hidden_size, bias=False), ) linear_k_adapter = nn.Sequential( nn.Linear(self.attention_hidden_size, self.config.adapter_rank, bias=False), nn.Linear(self.config.adapter_rank, self.attention_hidden_size, bias=False), ) linear_v_adapter = nn.Sequential( nn.Linear(self.attention_hidden_size, self.config.adapter_rank, bias=False), nn.Linear(self.config.adapter_rank, self.attention_hidden_size, bias=False), ) else: linear_q_adapter = nn.Identity() linear_k_adapter = nn.Identity() linear_v_adapter = nn.Identity() self.linear_q_adapter_list.append(linear_q_adapter) self.linear_k_adapter_list.append(linear_k_adapter) self.linear_v_adapter_list.append(linear_v_adapter) self.layer_dic = {value: index for index, value in enumerate(self.layer_block_map)} def forward( self, hidden_states: torch.Tensor, layer_idx: int, attention_mask: torch.Tensor | None = None, past_key_values: Cache | None = None, position_embeddings: tuple[torch.Tensor, torch.Tensor] | None = None, **kwargs: Unpack[TransformersKwargs], ) -> tuple[torch.Tensor, torch.Tensor | None, tuple[torch.Tensor] | None]: input_shape = hidden_states.shape[:-1] hidden_shape = (*input_shape, -1, self.head_dim) query_states = self.q_proj(hidden_states) key_states = self.k_proj(hidden_states) value_states = self.v_proj(hidden_states) if self.config.use_shared_attention_adapter: adapter_layer_idx = self.layer_dic[layer_idx] query_states = query_states + self.linear_q_adapter_list[adapter_layer_idx](hidden_states) key_states = key_states + self.linear_k_adapter_list[adapter_layer_idx](hidden_states) value_states = value_states + self.linear_v_adapter_list[adapter_layer_idx](hidden_states) query_states = query_states.view(hidden_shape).transpose(1, 2) key_states = key_states.view(hidden_shape).transpose(1, 2) value_states = value_states.view(hidden_shape).transpose(1, 2) if self.config.use_mem_rope: cos, sin = position_embeddings query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin) if past_key_values is not None: key_states, value_states = past_key_values.update(key_states, value_states, layer_idx) attention_interface: Callable = ALL_ATTENTION_FUNCTIONS.get_interface( self.config._attn_implementation, eager_attention_forward ) attn_output, attn_weights = attention_interface( self, query_states, key_states, value_states, attention_mask, dropout=0.0 if not self.training else self.attention_dropout, scaling=self.scaling, **kwargs, ) attn_output = attn_output.reshape(*input_shape, -1).contiguous() attn_output = self.o_proj(attn_output) return attn_output, attn_weights # Helper methods for segment sum computation def pad_tensor_by_size(input_tensor: torch.Tensor, pad_size: int): """ Padding x tensor with `pad_size` on the seq_len dim (dim=1) Assumes that we only have tensors of either size 4 or 3 """ pad_shape = (0, 0, 0, 0, 0, pad_size, 0, 0) if len(input_tensor.shape) == 4 else (0, 0, 0, pad_size, 0, 0) return torch.nn.functional.pad(input_tensor, pad_shape, mode="constant", value=0) def reshape_into_chunks(input_tensor, pad_size, chunk_size): """ Padding input_tensor with `pad_size` on the seq_len dim (dim=1) and simultaneously splitting it into chunk sequences. Assumes that we only have tensors of either size 4 or 3 """ # [bsz, seq_len, ...] -> [bsz, seq_len multiple of chunk_size, ...] input_tensor = pad_tensor_by_size(input_tensor, pad_size) if len(input_tensor.shape) == 3: # [bsz, seq_len multiple of chunk_size, num_heads] -> [bsz, -1, chunk_size, num_heads] return input_tensor.reshape(input_tensor.shape[0], -1, chunk_size, input_tensor.shape[2]) else: # [bsz, seq_len multiple of chunk_size, num_heads, head_dim or state_size] -> [bsz, -1, chunk_size, num_heads, head_dim or state_size] return input_tensor.reshape( input_tensor.shape[0], -1, chunk_size, input_tensor.shape[2], input_tensor.shape[3] ) def segment_sum(input_tensor): """ More stable segment sum calculation. Uses cumulative sums and masking instead of direct subtractions. """ chunk_size = input_tensor.size(-1) # 1. expand input tensor to have an additional dimension and repeat along that dimension # [..., chunk_size] -> [..., chunk_size, chunk_size] input_tensor = input_tensor[..., None].expand(*input_tensor.size(), chunk_size) # 2. create a lower triangular mask with the diagonal set to 0 to 0 out elements above diag mask = torch.tril(torch.ones(chunk_size, chunk_size, device=input_tensor.device, dtype=torch.bool), diagonal=-1) input_tensor = input_tensor.masked_fill(~mask, 0) # 3. compute actual cumsum tensor_segsum = torch.cumsum(input_tensor, dim=-2) # 4. apply mask to keep only the lower triangular part of the cumulative sum result (incl diagonal this time) mask = torch.tril(torch.ones(chunk_size, chunk_size, device=input_tensor.device, dtype=torch.bool), diagonal=0) tensor_segsum = tensor_segsum.masked_fill(~mask, -torch.inf) return tensor_segsum def apply_mask_to_padding_states(hidden_states, attention_mask): """ Tunes out the hidden states for padding tokens, see https://github.com/state-spaces/mamba/issues/66 """ # NOTE: attention mask is a 2D boolean tensor if attention_mask is not None: dtype = hidden_states.dtype hidden_states = (hidden_states * attention_mask[:, :, None]).to(dtype) return hidden_states @use_kernel_func_from_hub_with_fallback("causal_conv1d_update", "causal_conv1d") def causal_conv1d_update( hidden_states: torch.Tensor, conv_state: torch.Tensor, weight: nn.Parameter, bias: nn.Parameter | None = None, activation: str | None = None, ): _, hidden_size, seq_len = hidden_states.shape state_len = conv_state.shape[-1] hidden_states_new = torch.cat([conv_state, hidden_states], dim=-1).to(weight.dtype) conv_state.copy_(hidden_states_new[:, :, -state_len:]) out = F.conv1d(hidden_states_new, weight.unsqueeze(1), bias, padding=0, groups=hidden_size) out = out[:, :, -seq_len:] if activation is not None: out = ACT2FN[activation](out) return out.to(hidden_states.dtype) @use_kernel_func_from_hub_with_fallback("causal_conv1d_fn", "causal_conv1d") def causal_conv1d_fn( hidden_states: torch.Tensor, weight: nn.Parameter, bias: nn.Parameter | None = None, activation: str | None = None, **kwargs, ): _, hidden_size, seq_len = hidden_states.shape padding = weight.shape[-1] - 1 out = F.conv1d( hidden_states.to(weight.dtype), weight=weight.unsqueeze(1), bias=bias, padding=padding, groups=hidden_size, )[:, :, :seq_len] if activation is not None: out = ACT2FN[activation](out) return out.to(hidden_states.dtype) @use_kernel_func_from_hub_with_fallback("mamba_split_conv1d_scan_combined", "mamba_ssm") def mamba2_split_conv1d_scan_combined( zxbcdt: torch.Tensor, conv1d_weight: torch.Tensor, conv1d_bias: torch.Tensor | None, dt_bias: torch.Tensor, A: torch.Tensor, D: torch.Tensor, chunk_size: int, initial_states: torch.Tensor | None = None, dt_limit: tuple[float, float] = (0.0, float("inf")), return_final_states: bool = False, activation: str = "silu", rmsnorm_weight: torch.Tensor | None = None, rmsnorm_eps: float = 1e-6, outproj_weight: torch.Tensor | None = None, outproj_bias: torch.Tensor | None = None, headdim: int | None = None, ngroups: int = 1, norm_before_gate: bool = True, **kwargs, ): return None @use_kernel_func_from_hub_with_fallback("selective_state_update", "mamba_ssm") def mamba2_selective_state_update( state: torch.Tensor, hidden_states: torch.Tensor, dt: torch.Tensor, A: torch.Tensor, B: torch.Tensor, C: torch.Tensor, D: torch.Tensor | None = None, dt_bias: torch.Tensor | None = None, dt_softplus: bool = False, z: torch.Tensor | None = None, **kwargs, ): batch_size, num_heads, head_dim = hidden_states.shape num_groups = B.shape[1] state_size = B.shape[-1] if dt_bias is not None: dt = dt + dt_bias.to(dt.dtype) if dt_softplus: dt = F.softplus(dt) dt = dt[..., None] # Discretize A dA = torch.exp(dt.float() * A.float()).to(device=state.device) # Discretize B B = B.reshape(batch_size, num_groups, 1, state_size) B = B.expand(batch_size, num_groups, num_heads // num_groups, state_size).contiguous() B = B.reshape(batch_size, num_heads, 1, state_size) dB = dt * B # Discretize x into dB dBx = (dB * hidden_states[..., None]).to(device=state.device) # State calculation ssm_states = state * dA + dBx state.copy_(ssm_states.to(state.dtype)) # Subsequent output C = C.reshape(batch_size, num_groups, 1, state_size) C = C.expand(batch_size, num_groups, num_heads // num_groups, state_size).contiguous() C = C.reshape(batch_size, num_heads, state_size) # Reshape ssm_states to merge the first two dimensions ssm_states = ssm_states.to(device=C.device, dtype=C.dtype) ssm_states_reshaped = ssm_states.view(batch_size * num_heads, head_dim, state_size) C_reshaped = C.view(batch_size * num_heads, state_size, 1) out = torch.bmm(ssm_states_reshaped, C_reshaped) out = out.view(batch_size, num_heads, head_dim) # D skip connection if D is not None: out = (out + hidden_states * D).to(out.dtype) if z is not None: out = out * F.silu(z) return out.to(hidden_states.dtype) @use_kernel_func_from_hub_with_fallback("mamba_chunk_scan_combined", "mamba_ssm") def mamba2_chunk_scan( hidden_states: torch.Tensor, dt: torch.Tensor, A: torch.Tensor, B: torch.Tensor, C: torch.Tensor, chunk_size: int, D: torch.Tensor | None = None, dt_bias: torch.Tensor | None = None, initial_states: torch.Tensor | None = None, dt_softplus: bool = False, dt_limit: tuple[float, float] = (0.0, float("inf")), return_final_states: bool = False, **kwargs, ): batch_size, sequence_length, num_heads, head_dim = hidden_states.shape num_groups = B.shape[2] if dt_bias is not None: dt = dt + dt_bias.to(dt.dtype) if dt_softplus: dt = F.softplus(dt) dt = torch.clamp(dt, min=dt_limit[0], max=dt_limit[1]) hidden_states = hidden_states.float() B = B.float().repeat_interleave(num_heads // num_groups, dim=2, output_size=num_heads) C = C.float().repeat_interleave(num_heads // num_groups, dim=2, output_size=num_heads) pad_size = (chunk_size - sequence_length % chunk_size) % chunk_size D_residual = None if D is not None: D_residual = D[..., None] * pad_tensor_by_size(hidden_states, pad_size) # Discretize x and A hidden_states = hidden_states * dt[..., None].float() A = A.to(hidden_states.dtype) * dt.float() # Rearrange into blocks/chunks hidden_states, A, B, C = [reshape_into_chunks(tensor, pad_size, chunk_size) for tensor in (hidden_states, A, B, C)] A = A.permute(0, 3, 1, 2) A_cumsum = torch.cumsum(A, dim=-1) # 1. Compute the output for each intra-chunk (diagonal blocks) # This is the analog of a causal mask L = torch.exp(segment_sum(A)) # Contraction of C and B to get G (attention-weights like) G = (C[:, :, :, None, :, :] * B[:, :, None, :, :, :]).sum(dim=-1) # Compute M, equivalent to applying attention mask to weights M = (G[..., None] * L.permute(0, 2, 3, 4, 1)[..., None]).sum(dim=-1) # Compute Y_diag (apply to values) Y_diag = (M[..., None] * hidden_states[:, :, None]).sum(dim=3) # 2. Compute the state for each intra-chunk # (right term of low-rank factorization of off-diagonal blocks; B terms) decay_states = torch.exp(A_cumsum[:, :, :, -1:] - A_cumsum) B_decay = B * decay_states.permute(0, -2, -1, 1)[..., None] states = (B_decay[..., None, :] * hidden_states[..., None]).sum(dim=2) # 3. Compute the inter-chunk SSM recurrence; produces correct SSM states at chunk boundaries # (middle term of factorization of off-diag blocks; A terms) previous_states = ( initial_states[:, None].to(dtype=states.dtype, device=states.device) if initial_states is not None else torch.zeros_like(states[:, :1]) ) states = torch.cat([previous_states, states], dim=1) decay_chunk = torch.exp(segment_sum(F.pad(A_cumsum[:, :, :, -1], (1, 0)))).transpose(1, 3) new_states = (decay_chunk[..., None, None] * states[:, :, None, ...]).sum(dim=1) states, final_state = new_states[:, :-1], new_states[:, -1] # 4. Compute state -> output conversion per chunk # (left term of low-rank factorization of off-diagonal blocks; C terms) state_decay_out = torch.exp(A_cumsum) C_times_states = C[..., None, :] * states[:, :, None, ...] Y_off = C_times_states.sum(-1) * state_decay_out.permute(0, 2, 3, 1)[..., None] # Add output of intra-chunk and inter-chunk terms (diagonal and off-diagonal blocks) output = Y_diag + Y_off output = output.reshape(batch_size, -1, num_heads, head_dim) if D_residual is not None: output = output + D_residual # Cutting off padded chunks if pad_size > 0: output = output[:, :sequence_length] if return_final_states: return output, final_state return output @use_kernelized_func( [ causal_conv1d_fn, causal_conv1d_update, mamba2_split_conv1d_scan_combined, mamba2_selective_state_update, mamba2_chunk_scan, ] ) class Zamba2MambaMixer(nn.Module): """ Compute ∆, A, B, C, and D the state space parameters and compute the `contextualized_states`. A, D are input independent (see Mamba paper [1] Section 3.5.2 "Interpretation of A" for why A isn't selective) ∆, B, C are input-dependent (this is a key difference between Mamba and the linear time invariant S4, and is why Mamba is called **selective** state spaces) """ def __init__(self, config: Zamba2Config, layer_idx: int | None = None, initialize_mixer_weights: bool = True): super().__init__() self.config = config self.num_heads = self.config.n_mamba_heads self.hidden_size = config.hidden_size self.ssm_state_size = config.mamba_d_state self.conv_kernel_size = config.mamba_d_conv self.intermediate_size = int(config.mamba_expand * self.hidden_size) self.layer_idx = layer_idx self.use_conv_bias = config.use_conv_bias self.activation = "silu" self.act = nn.SiLU() self.n_groups = config.mamba_ngroups self.head_dim = config.mamba_headdim self.chunk_size = config.chunk_size # No upper limit self.time_step_limit = (config.time_step_min, float("inf")) self.conv_dim = self.intermediate_size + 2 * self.n_groups * self.ssm_state_size self.conv1d = nn.Conv1d( in_channels=self.conv_dim, out_channels=self.conv_dim, bias=True, kernel_size=config.mamba_d_conv, groups=self.conv_dim, padding=config.mamba_d_conv - 1, ) # projection of the input hidden states projection_size = self.intermediate_size + self.conv_dim + self.num_heads self.in_proj = nn.Linear( self.hidden_size, projection_size, bias=config.add_bias_linear, ) # selective projection used to make dt, B and C input dependent # time step projection (discretization) self.dt_bias = nn.Parameter(torch.empty(self.num_heads)) # S4D real initialization. These are not discretized! # The core is to load them, compute the discrete states, then write the updated state. Keeps the memory bounded self.A_log = nn.Parameter(torch.empty(self.num_heads)) self.norm = Zamba2RMSNormGated( self.intermediate_size, group_size=self.intermediate_size // self.n_groups, eps=1e-5 ) self.D = nn.Parameter(torch.empty(self.num_heads)) if initialize_mixer_weights and self.dt_bias.device.type != "meta": self.init_zamba2_weights() self.out_proj = nn.Linear(self.intermediate_size, self.hidden_size, bias=config.add_bias_linear) self.layer_type = config.layer_types[layer_idx] @torch.no_grad() def init_zamba2_weights(self): A = torch.arange(1, self.num_heads + 1, device=self.A_log.device, dtype=torch.float32) init.copy_(self.A_log, torch.log(A)) init.ones_(self.D) init.ones_(self.dt_bias) @force_accelerate_hooks("conv1d") def forward( self, hidden_states: torch.Tensor, cache_params: Cache | None = None, attention_mask: torch.Tensor | None = None, **kwargs, ): batch_size, seq_len, _ = hidden_states.shape dtype = hidden_states.dtype use_precomputed_states = cache_params is not None and cache_params.has_previous_state(self.layer_idx) # 1. Gated MLP's linear projection hidden_states = apply_mask_to_padding_states(hidden_states, attention_mask) projected_states = self.in_proj(hidden_states) A = -torch.exp(self.A_log.float()) fused_kwargs = kwargs | {"dt_limit": self.time_step_limit} if self.training and cache_params is None: fused_output = mamba2_split_conv1d_scan_combined( projected_states, self.conv1d.weight.squeeze(1), self.conv1d.bias, self.dt_bias, A, D=self.D, chunk_size=self.chunk_size, activation=self.activation, rmsnorm_weight=self.norm.weight, rmsnorm_eps=self.norm.variance_epsilon, outproj_weight=self.out_proj.weight, outproj_bias=self.out_proj.bias, headdim=self.head_dim, ngroups=self.n_groups, norm_before_gate=False, return_final_states=False, **fused_kwargs, ) # Only kernels can use this shortcircuit, fallback to normal torch otherwise if fused_output is not None: return fused_output gate, hidden_states_B_C, dt = projected_states.split( [self.intermediate_size, self.conv_dim, self.num_heads], dim=-1 ) if use_precomputed_states: conv_state = cache_params.layers[self.layer_idx].conv_states[0] recurrent_state = cache_params.layers[self.layer_idx].recurrent_states[0] # 2. Convolution sequence transformation hidden_states_B_C = hidden_states_B_C.transpose(1, 2) if use_precomputed_states and seq_len == 1 and not cache_params.layers[self.layer_idx].record_past: hidden_states_B_C = causal_conv1d_update( hidden_states_B_C, conv_state, self.conv1d.weight.squeeze(1), self.conv1d.bias, activation=self.activation, ) else: if cache_params is not None: hidden_states_B_C = cache_params.update_conv_state( hidden_states_B_C, self.layer_idx, conv_kernel_size=self.conv_kernel_size, ) hidden_states_B_C = causal_conv1d_fn( hidden_states_B_C, self.conv1d.weight.squeeze(1), self.conv1d.bias, activation=self.activation, **kwargs, ) if cache_params is not None: hidden_states_B_C = hidden_states_B_C[:, :, -seq_len:] # 3. SSM transformation hidden_states_B_C = apply_mask_to_padding_states(hidden_states_B_C.transpose(1, 2), attention_mask) hidden_states, B, C = torch.split( hidden_states_B_C, [self.intermediate_size, self.n_groups * self.ssm_state_size, self.n_groups * self.ssm_state_size], dim=-1, ) # Recurrent form if use_precomputed_states and seq_len == 1: hidden_states = hidden_states.view(batch_size, self.num_heads, self.head_dim) dt = dt.transpose(1, 2).expand(-1, -1, self.head_dim) A = A[:, None, None].expand(-1, self.head_dim, self.ssm_state_size) B = B.view(batch_size, self.n_groups, self.ssm_state_size) C = C.view(batch_size, self.n_groups, self.ssm_state_size) D = self.D[:, None].expand(-1, self.head_dim) dt_bias = self.dt_bias[:, None].expand(-1, self.head_dim) scan_output = mamba2_selective_state_update( recurrent_state, hidden_states, dt, A, B, C, D, z=None, dt_bias=dt_bias, dt_softplus=True, **kwargs, ) scan_output = scan_output.view(batch_size, 1, -1) # Chunk form else: output_final_state = cache_params is not None scan_result = mamba2_chunk_scan( hidden_states.view(batch_size, seq_len, self.num_heads, self.head_dim), dt, A, B.view(batch_size, seq_len, self.n_groups, self.ssm_state_size), C.view(batch_size, seq_len, self.n_groups, self.ssm_state_size), chunk_size=self.chunk_size, D=self.D, z=None, return_final_states=output_final_state, dt_bias=self.dt_bias, dt_softplus=True, initial_states=recurrent_state if use_precomputed_states else None, dt_limit=self.time_step_limit, **kwargs, ) if output_final_state: scan_output, final_state = scan_result cache_params.update_recurrent_state(final_state, layer_idx=self.layer_idx) else: scan_output = scan_result scan_output = scan_output.reshape(batch_size, seq_len, -1) scan_output = self.norm(scan_output, gate) # 4. Final linear projection contextualized_states = self.out_proj(scan_output.to(dtype)) return contextualized_states class Zamba2MLP(nn.Module): def __init__(self, config: Zamba2Config, num_fwd_mem_blocks=None, block_id: int | None = None): """ This MLP layer contributes to tied transformer blocks aimed to increasing compute without increasing model size. Because this layer 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. """ super().__init__() self.config = config self.hidden_size = config.hidden_size self.intermediate_size = config.intermediate_size self.num_fwd_mem_blocks = num_fwd_mem_blocks self.block_id = block_id self.gate_up_proj = nn.Linear(self.hidden_size, 2 * self.intermediate_size, bias=config.add_bias_linear) self.down_proj = nn.Linear(self.intermediate_size, self.hidden_size, bias=config.add_bias_linear) self.act_fn = ACT2FN[config.hidden_act] self.gate_up_proj_adapter_list = nn.ModuleList([]) for i in range(self.num_fwd_mem_blocks): if i % config.num_mem_blocks == block_id: gate_up_proj_adapter = nn.Sequential( nn.Linear(self.config.hidden_size, self.config.adapter_rank, bias=False), nn.Linear(self.config.adapter_rank, 2 * self.intermediate_size, bias=False), ) else: gate_up_proj_adapter = nn.Identity() self.gate_up_proj_adapter_list.append(gate_up_proj_adapter) layer_block_map = config.hybrid_layer_ids self.layer_dic = {value: index for index, value in enumerate(layer_block_map)} def forward(self, hidden_state, layer_idx=None): gate_up_state = self.gate_up_proj(hidden_state) layer_idx = self.layer_dic[layer_idx] gate_up_state = gate_up_state + self.gate_up_proj_adapter_list[layer_idx](hidden_state) gate_up_state = torch.chunk(gate_up_state, 2, dim=-1) hidden_state = self.act_fn(gate_up_state[0]) * gate_up_state[1] output = self.down_proj(hidden_state) return output class Zamba2AttentionDecoderLayer(nn.Module): def __init__(self, config: Zamba2Config, block_id: int | None = None, layer_idx: int | None = None): super().__init__() self.block_id = block_id num_gs = len(config.hybrid_layer_ids) self.self_attn = Zamba2Attention(config, layer_idx=-1, num_fwd_mem_blocks=num_gs, block_id=block_id) self.feed_forward = Zamba2MLP(config, num_fwd_mem_blocks=num_gs, block_id=block_id) self.input_layernorm = Zamba2RMSNorm(config.attention_hidden_size, eps=config.rms_norm_eps) self.pre_ff_layernorm = Zamba2RMSNorm(config.hidden_size, eps=config.rms_norm_eps) def forward( self, hidden_states: torch.Tensor, original_hidden_states: torch.Tensor, layer_idx: int, attention_mask: torch.Tensor | None = None, past_key_values: Cache | None = None, position_embeddings: torch.LongTensor | None = None, **kwargs: Unpack[TransformersKwargs], ) -> tuple[torch.FloatTensor]: """ Args: hidden_states (`torch.FloatTensor`): output of previous Mamba layer of shape `(batch, seq_len, embed_dim)` original_hidden_states (`torch.FloatTensor`): word embedding output of shape `(batch, seq_len, embed_dim)`. This is concatenated with `hidden_states` (which is the output of the previous (mamba) layer). The concatenated tensor is then used as input of the pre-attention RMSNorm (see fig. 2 in https://huggingface.co/papers/2405.16712). attention_mask (`torch.FloatTensor`, *optional*): attention mask of size `(batch, sequence_length)` where padding elements are indicated by 0. past_key_values (`Cache`, *optional*): cached past key and value projection states use_cache (`bool`, *optional*): If set to `True`, `past_key_values` key value states are returned and can be used to speed up decoding (see `past_key_values`). position_embeddings (`tuple[torch.FloatTensor, torch.FloatTensor]`, *optional*): Tuple containing the cosine and sine positional embeddings of shape `(batch_size, seq_len, head_dim)`, with `head_dim` being the embedding dimension of each attention head. """ hidden_states = torch.concatenate([hidden_states, original_hidden_states], dim=-1) hidden_states = self.input_layernorm(hidden_states) hidden_states, _ = self.self_attn( hidden_states=hidden_states, layer_idx=layer_idx, attention_mask=attention_mask, past_key_values=past_key_values, position_embeddings=position_embeddings, **kwargs, ) hidden_states = self.pre_ff_layernorm(hidden_states) hidden_states = self.feed_forward(hidden_states, layer_idx) return hidden_states class Zamba2MambaDecoderLayer(GradientCheckpointingLayer): def __init__(self, config: Zamba2Config, layer_idx: int): super().__init__() self.mamba = Zamba2MambaMixer(config=config, layer_idx=layer_idx) self.input_layernorm = Zamba2RMSNorm(config.hidden_size, eps=config.rms_norm_eps) self.layer_idx = layer_idx def forward( self, hidden_states: torch.Tensor, original_hidden_states: torch.Tensor | None = None, layer_idx: int | None = None, attention_mask: torch.Tensor | None = None, causal_mask: torch.Tensor | None = None, past_key_values: Cache | None = None, use_cache: bool | None = False, position_ids: torch.LongTensor | None = None, transformer_hidden_states: torch.Tensor | None = None, **kwargs: Unpack[TransformersKwargs], ) -> tuple[torch.FloatTensor, tuple[torch.FloatTensor, torch.FloatTensor] | None]: """ Args: hidden_states (`torch.FloatTensor`): input to the layer of shape `(batch, seq_len, embed_dim)` attention_mask (`torch.FloatTensor`, *optional*): attention mask of size `(batch, sequence_length)` where padding elements are indicated by 0. past_key_values (`Cache`, *optional*): cached past key and value projection states use_cache (`bool`, *optional*): If set to `True`, `past_key_values` key value states are returned and can be used to speed up decoding (see `past_key_values`). """ residual = hidden_states # `transformer_hidden_states` is the output from shared transformer + linear layer (see fig. 2 in https://huggingface.co/papers/2405.16712). # `transformer_hidden_states` is then added to the input to the mamba layer below (as described in eq. (6) of https://huggingface.co/papers/2405.16712). hidden_states = ( hidden_states + transformer_hidden_states if transformer_hidden_states is not None else hidden_states ) hidden_states = self.input_layernorm(hidden_states) hidden_states = self.mamba( hidden_states=hidden_states, cache_params=past_key_values, attention_mask=attention_mask, **kwargs, ) # residual connection after mamba hidden_states = residual + hidden_states return hidden_states class Zamba2HybridLayer(GradientCheckpointingLayer): def __init__( self, shared_transformer: Zamba2AttentionDecoderLayer, linear: nn.Linear, mamba: Zamba2MambaDecoderLayer ): super().__init__() self.linear = linear self.mamba_decoder = mamba self.shared_transformer = shared_transformer def forward( self, hidden_states: torch.Tensor, original_hidden_states: torch.Tensor | None = None, layer_idx: int | None = None, attention_mask: torch.Tensor | None = None, causal_mask: torch.Tensor | None = None, past_key_values: Cache | None = None, use_cache: bool | None = False, position_embeddings: torch.LongTensor | None = None, position_ids: torch.LongTensor | None = None, **kwargs: Unpack[TransformersKwargs], ) -> tuple[torch.FloatTensor, tuple[torch.FloatTensor, torch.FloatTensor] | None]: """ Args: hidden_states (`torch.FloatTensor`): input to the layer of shape `(batch, seq_len, embed_dim)` original_hidden_states (`torch.FloatTensor`): word embedding output that will be concatenated with hidden activations to form the input of the shared transformer layer. layer_idx (`int`): layer number. attention_mask (`torch.FloatTensor`, *optional*): attention mask of size `(batch, sequence_length)` where padding elements are indicated by 0. past_key_values (`Cache`, *optional*): cached past key and value projection states use_cache (`bool`, *optional*): If set to `True`, `past_key_values` key value states are returned and can be used to speed up decoding (see `past_key_values`). position_embeddings (`tuple[torch.FloatTensor, torch.FloatTensor]`, *optional*): Tuple containing the cosine and sine positional embeddings of shape `(batch_size, seq_len, head_dim)`, with `head_dim` being the embedding dimension of each attention head. """ transformer_hidden_states = self.shared_transformer( hidden_states, original_hidden_states=original_hidden_states, layer_idx=layer_idx, attention_mask=causal_mask, past_key_values=past_key_values, position_embeddings=position_embeddings, position_ids=position_ids, **kwargs, ) transformer_hidden_states = self.linear(transformer_hidden_states) hidden_states = self.mamba_decoder( hidden_states, transformer_hidden_states=transformer_hidden_states, attention_mask=attention_mask, past_key_values=past_key_values, use_cache=use_cache, position_embeddings=position_embeddings, **kwargs, ) return hidden_states @auto_docstring class Zamba2PreTrainedModel(PreTrainedModel): config: Zamba2Config base_model_prefix = "model" supports_gradient_checkpointing = True _no_split_modules = ["Zamba2HybridLayer", "Zamba2MambaDecoderLayer"] _skip_keys_device_placement = ["past_key_values"] _supports_flash_attn = True _supports_flex_attn = True _supports_sdpa = True _is_stateful = True _can_record_outputs = { "hidden_states": Zamba2MambaDecoderLayer, "attentions": Zamba2Attention, } @torch.no_grad() def _init_weights(self, module): super()._init_weights(module) if isinstance(module, Zamba2MambaMixer): dt = torch.exp( torch.rand(self.config.n_mamba_heads) * (math.log(self.config.time_step_max) - math.log(self.config.time_step_min)) + math.log(self.config.time_step_min) ).clamp(min=self.config.time_step_floor) # # Inverse of softplus: https://github.com/pytorch/pytorch/issues/72759 inv_dt = dt + torch.log(-torch.expm1(-dt)) init.copy_(module.dt_bias, inv_dt) A = torch.arange(1, module.num_heads + 1) init.copy_(module.A_log, torch.log(A)) init.ones_(module.D) @auto_docstring class Zamba2Model(Zamba2PreTrainedModel): """ Model consisting of *config.num_hidden_layers* layers. Args: config: Zamba2Config """ def __init__(self, config: Zamba2Config): super().__init__(config) self.config = config self.padding_idx = config.pad_token_id self.vocab_size = config.vocab_size self.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size, self.padding_idx) self.layers_block_type = config.layers_block_type self.layers = self.get_layers() self._attn_implementation = config._attn_implementation self.final_layernorm = Zamba2RMSNorm(config.hidden_size, eps=config.rms_norm_eps) if config.use_mem_rope: if config.use_long_context: logger.warning_once( "`use_long_context` set to `True`: using rescaled `rope_theta` and extended `max_position_embeddings`." ) self.rotary_emb = Zamba2RotaryEmbedding(config) self.gradient_checkpointing = False # Initialize weights and apply final processing self.post_init() @merge_with_config_defaults @capture_outputs @auto_docstring def forward( self, input_ids: torch.LongTensor | None = None, attention_mask: torch.Tensor | None = None, position_ids: torch.LongTensor | None = None, past_key_values: Cache | None = None, inputs_embeds: torch.FloatTensor | None = None, use_cache: bool | None = None, **kwargs: Unpack[TransformersKwargs], ) -> tuple | BaseModelOutputWithPast: if (input_ids is None) ^ (inputs_embeds is not None): raise ValueError( "You cannot specify both input_ids and inputs_embeds at the same time, and must specify either one" ) if inputs_embeds is None: inputs_embeds = self.embed_tokens(input_ids) hidden_states = inputs_embeds original_hidden_states = torch.clone(inputs_embeds) # original_hidden_states: word embedding output that will be concatenated with hidden activations to form the input of the shared transformer layer if use_cache and past_key_values is None: past_key_values = DynamicCache(config=self.config) if position_ids is None: past_seen_tokens = past_key_values.get_seq_length() if past_key_values is not None else 0 position_ids = torch.arange(inputs_embeds.shape[1], device=inputs_embeds.device) + past_seen_tokens position_ids = position_ids.unsqueeze(0) if not isinstance(causal_mask_mapping := attention_mask, dict): mask_kwargs = { "config": self.config, "inputs_embeds": inputs_embeds, "attention_mask": attention_mask, "past_key_values": past_key_values, "position_ids": position_ids, } causal_mask_mapping = { "full_attention": create_causal_mask(**mask_kwargs), "linear_attention": create_recurrent_attention_mask(**mask_kwargs), } # create position embeddings to be shared across the decoder layers if self.config.use_mem_rope: position_embeddings = self.rotary_emb(hidden_states, position_ids=position_ids) else: position_embeddings = None for layer_idx, layer in enumerate(self.layers): hidden_states = layer( hidden_states, original_hidden_states, layer_idx, causal_mask_mapping["linear_attention"], causal_mask_mapping["full_attention"], past_key_values=past_key_values, use_cache=use_cache, position_embeddings=position_embeddings, position_ids=position_ids, **kwargs, ) hidden_states = self.final_layernorm(hidden_states) return BaseModelOutputWithPast( last_hidden_state=hidden_states, past_key_values=past_key_values if use_cache else None, ) def get_layers(self): layers = [] self._tied_weights_keys = {} self.first_transformer_layer_id = 0 unique_hybrid_blocks = [] for layer_id, layer_type in enumerate(self.layers_block_type): mamba_layer = Zamba2MambaDecoderLayer(self.config, layer_idx=layer_id) if layer_type == "hybrid": prefix_pattern = f"layers.{layer_id}.shared_transformer" # Zamba ties Hybrid module weights by repeating blocks after every # `num_mem_blocks`. So if `num_mem_blocks=2`, the blocks looks like # [1, 2, 1, 2, 1, 2] where all "ones" share the same set of weights. if ( not isinstance(unique_hybrid_blocks, list) or len(unique_hybrid_blocks) >= self.config.num_mem_blocks ): if isinstance(unique_hybrid_blocks, list): unique_hybrid_blocks = cycle(unique_hybrid_blocks) target_pattern = next(unique_hybrid_blocks) self._tied_weights_keys.update({prefix_pattern: target_pattern}) else: # Store source patterns to which the subsequent modules will be tied unique_hybrid_blocks.append(prefix_pattern) block_id = layer_id % self.config.num_mem_blocks attn_block = Zamba2AttentionDecoderLayer(self.config, block_id=block_id) linear_layer = nn.Linear(self.config.hidden_size, self.config.hidden_size, bias=False) layers.append(Zamba2HybridLayer(attn_block, linear_layer, mamba_layer)) else: layers.append(mamba_layer) return nn.ModuleList(layers) # Adapted from transformers.models.jamba.modeling_jamba.JambaForCausalLM with Jamba->Zamba2, JAMBA->ZAMBA2 class Zamba2ForCausalLM(Zamba2PreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} def __init__(self, config: Zamba2Config): super().__init__(config) self.model = Zamba2Model(config) self.vocab_size = config.vocab_size self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False) # Initialize weights and apply final processing self.post_init() @can_return_tuple @auto_docstring def forward( self, input_ids: torch.LongTensor | None = None, attention_mask: torch.Tensor | None = None, position_ids: torch.LongTensor | None = None, past_key_values: Cache | None = None, inputs_embeds: torch.FloatTensor | None = None, labels: torch.LongTensor | None = None, use_cache: bool | None = None, logits_to_keep: int | torch.Tensor = 0, **kwargs: Unpack[TransformersKwargs], ) -> tuple | CausalLMOutputWithPast: r""" labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*): Labels for computing the masked language modeling loss. Indices should either be in `[0, ..., config.vocab_size]` or -100 (see `input_ids` docstring). Tokens with indices set to `-100` are ignored (masked), the loss is only computed for the tokens with labels in `[0, ..., config.vocab_size]`. Example: ```python >>> from transformers import AutoTokenizer, Zamba2ForCausalLM >>> model = Zamba2ForCausalLM.from_pretrained("Zyphra/Zamba2-7B-v1") >>> tokenizer = AutoTokenizer.from_pretrained("Zyphra/Zamba2-7B-v1") >>> prompt = "Hey, are you conscious? Can you talk to me?" >>> inputs = tokenizer(prompt, return_tensors="pt") >>> # Generate >>> generate_ids = model.generate(inputs.input_ids, max_length=30) >>> tokenizer.batch_decode(generate_ids, skip_special_tokens=True, clean_up_tokenization_spaces=False)[0] "Hey, are you conscious? Can you talk to me?\nI'm not conscious, but I can talk to you." ```""" outputs: BaseModelOutputWithPast = self.model( input_ids=input_ids, attention_mask=attention_mask, position_ids=position_ids, past_key_values=past_key_values, inputs_embeds=inputs_embeds, use_cache=use_cache, **kwargs, ) hidden_states = outputs.last_hidden_state # Only compute necessary logits, and do not upcast them to float if we are not computing the loss slice_indices = slice(-logits_to_keep, None) if isinstance(logits_to_keep, int) else logits_to_keep logits = self.lm_head(hidden_states[:, slice_indices, :]) loss = None if labels is not None: loss = self.loss_function( logits, labels, self.vocab_size, **kwargs, ) return CausalLMOutputWithPast( loss=loss, logits=logits, past_key_values=outputs.past_key_values, hidden_states=outputs.hidden_states, attentions=outputs.attentions, ) def prepare_inputs_for_generation(self, input_ids, **kwargs): kwargs["logits_to_keep"] = self.config.num_logits_to_keep model_inputs = super().prepare_inputs_for_generation(input_ids, **kwargs) return model_inputs @auto_docstring( custom_intro=""" The Zamba2 Model with a sequence classification head on top (linear layer). [`Zamba2ForSequenceClassification`] uses the last token in order to do the classification, as other causal models (e.g. GPT-2) do. Since it does classification on the last token, it requires to know the position of the last token. If a `pad_token_id` is defined in the configuration, it finds the last token that is not a padding token in each row. If no `pad_token_id` is defined, it simply takes the last value in each row of the batch. Since it cannot guess the padding tokens when `inputs_embeds` are passed instead of `input_ids`, it does the same (take the last value in each row of the batch). """ ) class Zamba2ForSequenceClassification(Zamba2PreTrainedModel): def __init__(self, config: Zamba2Config): super().__init__(config) self.num_labels = config.num_labels self.model = Zamba2Model(config) self.score = nn.Linear(config.hidden_size, self.num_labels, bias=False) # Initialize weights and apply final processing self.post_init() @can_return_tuple @auto_docstring def forward( self, input_ids: torch.LongTensor | None = None, attention_mask: torch.Tensor | None = None, position_ids: torch.LongTensor | None = None, past_key_values: Cache | None = None, inputs_embeds: torch.FloatTensor | None = None, labels: torch.LongTensor | None = None, use_cache: bool | None = None, logits_to_keep: int | torch.Tensor = 0, **kwargs: Unpack[TransformersKwargs], ) -> tuple | SequenceClassifierOutputWithPast: r""" labels (`torch.LongTensor` of shape `(batch_size,)`, *optional*): Labels for computing the sequence classification/regression loss. Indices should be in `[0, ..., config.num_labels - 1]`. If `config.num_labels == 1` a regression loss is computed (Mean-Square loss), If `config.num_labels > 1` a classification loss is computed (Cross-Entropy). """ transformer_outputs: BaseModelOutputWithPast = self.model( input_ids, attention_mask=attention_mask, position_ids=position_ids, past_key_values=past_key_values, inputs_embeds=inputs_embeds, use_cache=use_cache, **kwargs, ) hidden_states = transformer_outputs[0] logits = self.score(hidden_states) if input_ids is not None: batch_size = input_ids.shape[0] else: batch_size = inputs_embeds.shape[0] if self.config.pad_token_id is None and batch_size != 1: raise ValueError("Cannot handle batch sizes > 1 if no padding token is defined.") if self.config.pad_token_id is None: last_non_pad_token = -1 elif input_ids is not None: non_pad_mask = (input_ids != self.config.pad_token_id).to(logits.device, torch.int32) token_indices = torch.arange(input_ids.shape[-1], device=logits.device, dtype=torch.int32) last_non_pad_token = (token_indices * non_pad_mask).argmax(-1) else: last_non_pad_token = -1 logger.warning_once( f"{self.__class__.__name__} will not detect padding tokens in `inputs_embeds`. Results may be " "unexpected if using padding tokens in conjunction with `inputs_embeds.`" ) pooled_logits = logits[torch.arange(batch_size, device=logits.device), last_non_pad_token] loss = None if labels is not None: loss = self.loss_function( logits=pooled_logits, labels=labels, pooled_logits=pooled_logits, config=self.config, **kwargs ) return SequenceClassifierOutputWithPast( loss=loss, logits=pooled_logits, past_key_values=transformer_outputs.past_key_values, hidden_states=transformer_outputs.hidden_states, attentions=transformer_outputs.attentions, ) __all__ = ["Zamba2ForCausalLM", "Zamba2ForSequenceClassification", "Zamba2Model", "Zamba2PreTrainedModel"]