Aluode/PerceptionLabPortable
0
1# coding=utf-82# Copyright 2024 AI21 Labs Ltd. and the HuggingFace Inc. team. All rights reserved.3#4# This code is based on EleutherAI's GPT-NeoX library and the GPT-NeoX5# and OPT implementations in this library. It has been modified from its6# original forms to accommodate minor architectural differences compared7# to GPT-NeoX and OPT used by the Meta AI team that trained the model.8#9# Licensed under the Apache License, Version 2.0 (the "License");10# you may not use this file except in compliance with the License.11# You may obtain a copy of the License at12#13# http://www.apache.org/licenses/LICENSE-2.014#15# Unless required by applicable law or agreed to in writing, software16# distributed under the License is distributed on an "AS IS" BASIS,17# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.18# See the License for the specific language governing permissions and19# limitations under the License.20"""PyTorch Jamba model."""21 22import math23from typing import Any, Optional, Union24 25import torch26import torch.nn.functional as F27from torch import nn28 29from ...activations import ACT2FN30from ...generation import GenerationMixin31from ...modeling_attn_mask_utils import AttentionMaskConverter32from ...modeling_flash_attention_utils import flash_attn_supports_top_left_mask, is_flash_attn_available33from ...modeling_layers import (34 GenericForSequenceClassification,35 GradientCheckpointingLayer,36)37from ...modeling_outputs import MoeCausalLMOutputWithPast, MoeModelOutputWithPast38from ...modeling_utils import PreTrainedModel39from ...processing_utils import Unpack40from ...utils import TransformersKwargs, auto_docstring, can_return_tuple, logging41from ...utils.deprecation import deprecate_kwarg42from ...utils.import_utils import is_causal_conv1d_available, is_mamba_ssm_available43from .configuration_jamba import JambaConfig44 45 46if is_flash_attn_available():47 from ...modeling_flash_attention_utils import _flash_attention_forward48 49 50if is_mamba_ssm_available():51 from mamba_ssm.ops.selective_scan_interface import mamba_inner_fn, selective_scan_fn52 from mamba_ssm.ops.triton.selective_state_update import selective_state_update53else:54 selective_state_update, selective_scan_fn, mamba_inner_fn = None, None, None55 56if is_causal_conv1d_available():57 from causal_conv1d import causal_conv1d_fn, causal_conv1d_update58else:59 causal_conv1d_update, causal_conv1d_fn = None, None60 61is_fast_path_available = all(62 (selective_state_update, selective_scan_fn, causal_conv1d_fn, causal_conv1d_update, mamba_inner_fn)63)64 65 66logger = logging.get_logger(__name__)67 68 69# Copied from transformers.models.qwen2_moe.modeling_qwen2_moe.load_balancing_loss_func with gate->router70def load_balancing_loss_func(71 router_logits: Union[torch.Tensor, tuple[torch.Tensor], None],72 num_experts: Optional[int] = None,73 top_k=2,74 attention_mask: Optional[torch.Tensor] = None,75) -> Union[torch.Tensor, int]:76 r"""77 Computes auxiliary load balancing loss as in Switch Transformer - implemented in Pytorch.78 79 See Switch Transformer (https://huggingface.co/papers/2101.03961) for more details. This function implements the loss80 function presented in equations (4) - (6) of the paper. It aims at penalizing cases where the routing between81 experts is too unbalanced.82 83 Args:84 router_logits:85 Logits from the `router`, should be a tuple of model.config.num_hidden_layers tensors of86 shape [batch_size X sequence_length, num_experts].87 num_experts:88 Number of experts89 top_k:90 The number of experts to route per-token, can be also interpreted as the `top-k` routing91 parameter.92 attention_mask (`torch.Tensor`, *optional*):93 The attention_mask used in forward function94 shape [batch_size X sequence_length] if not None.95 96 Returns:97 The auxiliary loss.98 """99 if router_logits is None or not isinstance(router_logits, tuple):100 return 0101 102 if isinstance(router_logits, tuple):103 compute_device = router_logits[0].device104 concatenated_router_logits = torch.cat(105 [layer_router.to(compute_device) for layer_router in router_logits], dim=0106 )107 108 routing_weights = torch.nn.functional.softmax(concatenated_router_logits, dim=-1)109 110 _, selected_experts = torch.topk(routing_weights, top_k, dim=-1)111 112 expert_mask = torch.nn.functional.one_hot(selected_experts, num_experts)113 114 if attention_mask is None:115 # Compute the percentage of tokens routed to each experts116 tokens_per_expert = torch.mean(expert_mask.float(), dim=0)117 118 # Compute the average probability of routing to these experts119 router_prob_per_expert = torch.mean(routing_weights, dim=0)120 else:121 batch_size, sequence_length = attention_mask.shape122 num_hidden_layers = concatenated_router_logits.shape[0] // (batch_size * sequence_length)123 124 # Compute the mask that masks all padding tokens as 0 with the same shape of expert_mask125 expert_attention_mask = (126 attention_mask[None, :, :, None, None]127 .expand((num_hidden_layers, batch_size, sequence_length, top_k, num_experts))128 .reshape(-1, top_k, num_experts)129 .to(compute_device)130 )131 132 # Compute the percentage of tokens routed to each experts133 tokens_per_expert = torch.sum(expert_mask.float() * expert_attention_mask, dim=0) / torch.sum(134 expert_attention_mask, dim=0135 )136 137 # Compute the mask that masks all padding tokens as 0 with the same shape of tokens_per_expert138 router_per_expert_attention_mask = (139 attention_mask[None, :, :, None]140 .expand((num_hidden_layers, batch_size, sequence_length, routing_weights.shape[1]))141 .reshape(-1, routing_weights.shape[1])142 .to(compute_device)143 )144 145 # Compute the average probability of routing to these experts146 router_prob_per_expert = torch.sum(routing_weights * router_per_expert_attention_mask, dim=0) / torch.sum(147 router_per_expert_attention_mask, dim=0148 )149 150 device_index = routing_weights.device.index if routing_weights.device.index is not None else 0151 rank = routing_weights.shape[1] * int(device_index)152 overall_loss = torch.sum(153 tokens_per_expert[:, rank : rank + routing_weights.shape[1]] * router_prob_per_expert.unsqueeze(0)154 )155 return overall_loss * num_experts156 157 158# Copied from transformers.models.llama.modeling_llama.LlamaRMSNorm with Llama->Jamba159class JambaRMSNorm(nn.Module):160 def __init__(self, hidden_size, eps=1e-6):161 """162 JambaRMSNorm is equivalent to T5LayerNorm163 """164 super().__init__()165 self.weight = nn.Parameter(torch.ones(hidden_size))166 self.variance_epsilon = eps167 168 def forward(self, hidden_states):169 input_dtype = hidden_states.dtype170 hidden_states = hidden_states.to(torch.float32)171 variance = hidden_states.pow(2).mean(-1, keepdim=True)172 hidden_states = hidden_states * torch.rsqrt(variance + self.variance_epsilon)173 return self.weight * hidden_states.to(input_dtype)174 175 def extra_repr(self):176 return f"{tuple(self.weight.shape)}, eps={self.variance_epsilon}"177 178 179# Copied from transformers.models.llama.modeling_llama.repeat_kv180def repeat_kv(hidden_states: torch.Tensor, n_rep: int) -> torch.Tensor:181 """182 This is the equivalent of torch.repeat_interleave(x, dim=1, repeats=n_rep). The hidden states go from (batch,183 num_key_value_heads, seqlen, head_dim) to (batch, num_attention_heads, seqlen, head_dim)184 """185 batch, num_key_value_heads, slen, head_dim = hidden_states.shape186 if n_rep == 1:187 return hidden_states188 hidden_states = hidden_states[:, :, None, :, :].expand(batch, num_key_value_heads, n_rep, slen, head_dim)189 return hidden_states.reshape(batch, num_key_value_heads * n_rep, slen, head_dim)190 191 192class HybridMambaAttentionDynamicCache:193 """194 A dynamic cache that can handle both the attention cache (which has a seq_len dimension) and the mamba cache195 (which has a constant shape regardless of seq_len).196 197 This cache has two sets of lists of tensors: `key_cache` and `value_cache` for attention cache and `conv_states`198 and `ssm_states` for mamba cache. Each of these lists has `num_layers` tensors. The expected shape for each tensor199 For attention layers, `key_cache` and `value_cache` have a shape of `(batch_size, num_heads, seq_len, head_dim)`,200 while `conv_states` and `ssm_states` have a shape of `(batch_size, 0)` (empty tensors).201 For mamba layers, `key_cache` and `value_cache` have a shape of `(batch_size, 0)` (empty tensors),202 while `conv_states` represents the convolution state and has a shape of `(batch_size, d_inner, d_conv)`,203 and `ssm_states` represents the ssm state and has a shape of `(batch_size, d_inner, d_state)`.204 """205 206 is_compileable = False207 208 def __init__(self, config, batch_size, dtype=torch.float16, device=None):209 self.dtype = dtype210 self.layers_block_type = config.layers_block_type211 self.has_previous_state = False # only used by mamba212 intermediate_size = config.mamba_expand * config.hidden_size213 ssm_state_size = config.mamba_d_state214 conv_kernel_size = config.mamba_d_conv215 self.conv_states = []216 self.ssm_states = []217 self.transformer_layers = []218 for i in range(config.num_hidden_layers):219 if self.layers_block_type[i] == "mamba":220 self.conv_states += [221 torch.zeros(batch_size, intermediate_size, conv_kernel_size, device=device, dtype=dtype)222 ]223 self.ssm_states += [224 torch.zeros(batch_size, intermediate_size, ssm_state_size, device=device, dtype=dtype)225 ]226 else:227 self.conv_states += [torch.tensor([[]] * batch_size, device=device)]228 self.ssm_states += [torch.tensor([[]] * batch_size, device=device)]229 self.transformer_layers.append(i)230 231 self.key_cache = [torch.tensor([[]] * batch_size, device=device) for _ in range(config.num_hidden_layers)]232 self.value_cache = [torch.tensor([[]] * batch_size, device=device) for _ in range(config.num_hidden_layers)]233 234 def update(235 self,236 key_states: torch.Tensor,237 value_states: torch.Tensor,238 layer_idx: int,239 cache_kwargs: Optional[dict[str, Any]] = None,240 ) -> tuple[torch.Tensor, torch.Tensor]:241 # Update the cache242 if self.key_cache[layer_idx].shape[-1] == 0:243 self.key_cache[layer_idx] = key_states244 self.value_cache[layer_idx] = value_states245 else:246 self.key_cache[layer_idx] = torch.cat([self.key_cache[layer_idx], key_states], dim=2)247 self.value_cache[layer_idx] = torch.cat([self.value_cache[layer_idx], value_states], dim=2)248 249 return self.key_cache[layer_idx], self.value_cache[layer_idx]250 251 def reorder_cache(self, beam_idx: torch.LongTensor):252 """Reorders the cache for beam search, given the selected beam indices."""253 for layer_idx in range(len(self.key_cache)):254 device = self.key_cache[layer_idx].device255 self.key_cache[layer_idx] = self.key_cache[layer_idx].index_select(0, beam_idx.to(device))256 device = self.value_cache[layer_idx].device257 self.value_cache[layer_idx] = self.value_cache[layer_idx].index_select(0, beam_idx.to(device))258 259 device = self.conv_states[layer_idx].device260 self.conv_states[layer_idx] = self.conv_states[layer_idx].index_select(0, beam_idx.to(device))261 device = self.ssm_states[layer_idx].device262 self.ssm_states[layer_idx] = self.ssm_states[layer_idx].index_select(0, beam_idx.to(device))263 264 def get_seq_length(self, layer_idx: Optional[int] = 0) -> int:265 """Returns the sequence length of the cached states. A layer index can be optionally passed."""266 # take any layer that contains cache and not empty tensor267 layer_idx = self.transformer_layers[0] if layer_idx not in self.transformer_layers else layer_idx268 if len(self.key_cache) <= layer_idx:269 return 0270 return self.key_cache[layer_idx].shape[-2]271 272 273# Adapted from transformers.models.mistral.modeling_mistral.MistralAttention with Mistral->Jamba274class JambaAttention(nn.Module):275 """276 Multi-headed attention from 'Attention Is All You Need' paper. Modified to use sliding window attention: Longformer277 and "Generating Long Sequences with Sparse Transformers".278 """279 280 def __init__(self, config: JambaConfig, layer_idx: Optional[int] = None):281 super().__init__()282 self.config = config283 self.layer_idx = layer_idx284 if layer_idx is None:285 logger.warning_once(286 f"Instantiating {self.__class__.__name__} without passing a `layer_idx` is not recommended and will "287 "lead to errors during the forward call if caching is used. Please make sure to provide a `layer_idx` "288 "when creating this class."289 )290 291 self.hidden_size = config.hidden_size292 self.num_heads = config.num_attention_heads293 self.head_dim = self.hidden_size // self.num_heads294 self.num_key_value_heads = config.num_key_value_heads295 self.num_key_value_groups = self.num_heads // self.num_key_value_heads296 self.is_causal = True297 self.attention_dropout = config.attention_dropout298 299 if (self.head_dim * self.num_heads) != self.hidden_size:300 raise ValueError(301 f"hidden_size must be divisible by num_heads (got `hidden_size`: {self.hidden_size}"302 f" and `num_heads`: {self.num_heads})."303 )304 self.q_proj = nn.Linear(self.hidden_size, self.num_heads * self.head_dim, bias=False)305 self.k_proj = nn.Linear(self.hidden_size, self.num_key_value_heads * self.head_dim, bias=False)306 self.v_proj = nn.Linear(self.hidden_size, self.num_key_value_heads * self.head_dim, bias=False)307 self.o_proj = nn.Linear(self.num_heads * self.head_dim, self.hidden_size, bias=False)308 309 @deprecate_kwarg("past_key_value", new_name="past_key_values", version="4.58")310 def forward(311 self,312 hidden_states: torch.Tensor,313 attention_mask: Optional[torch.Tensor] = None,314 position_ids: Optional[torch.LongTensor] = None,315 past_key_values: Optional[HybridMambaAttentionDynamicCache] = None,316 output_attentions: bool = False,317 use_cache: bool = False,318 cache_position: Optional[torch.LongTensor] = None,319 ) -> tuple[torch.Tensor, Optional[torch.Tensor], Optional[tuple[torch.Tensor]]]:320 bsz, q_len, _ = hidden_states.size()321 322 query_states = self.q_proj(hidden_states)323 key_states = self.k_proj(hidden_states)324 value_states = self.v_proj(hidden_states)325 326 query_states = query_states.view(bsz, q_len, self.num_heads, self.head_dim).transpose(1, 2)327 key_states = key_states.view(bsz, q_len, self.num_key_value_heads, self.head_dim).transpose(1, 2)328 value_states = value_states.view(bsz, q_len, self.num_key_value_heads, self.head_dim).transpose(1, 2)329 330 if past_key_values is not None:331 key_states, value_states = past_key_values.update(key_states, value_states, self.layer_idx)332 333 # repeat k/v heads if n_kv_heads < n_heads334 key_states = repeat_kv(key_states, self.num_key_value_groups)335 value_states = repeat_kv(value_states, self.num_key_value_groups)336 337 attn_weights = torch.matmul(query_states, key_states.transpose(2, 3)) / math.sqrt(self.head_dim)338 339 if attention_mask is not None: # no matter the length, we just slice it340 causal_mask = attention_mask[:, :, :, : key_states.shape[-2]]341 attn_weights = attn_weights + causal_mask342 343 # upcast attention to fp32344 attn_weights = nn.functional.softmax(attn_weights, dim=-1, dtype=torch.float32).to(query_states.dtype)345 attn_weights = nn.functional.dropout(attn_weights, p=self.attention_dropout, training=self.training)346 attn_output = torch.matmul(attn_weights, value_states)347 348 if attn_output.size() != (bsz, self.num_heads, q_len, self.head_dim):349 raise ValueError(350 f"`attn_output` should be of size {(bsz, self.num_heads, q_len, self.head_dim)}, but is"351 f" {attn_output.size()}"352 )353 354 attn_output = attn_output.transpose(1, 2).contiguous()355 attn_output = attn_output.reshape(bsz, q_len, self.hidden_size)356 357 attn_output = self.o_proj(attn_output)358 359 if not output_attentions:360 attn_weights = None361 362 return attn_output, attn_weights, past_key_values363 364 365# Adapted from transformers.models.mistral.modeling_mistral.MistralFlashAttention2 with Mistral->Jamba366class JambaFlashAttention2(JambaAttention):367 """368 Jamba flash attention module. This module inherits from `JambaAttention` as the weights of the module stays369 untouched. The only required change would be on the forward pass where it needs to correctly call the public API of370 flash attention and deal with padding tokens in case the input contains any of them.371 """372 373 def __init__(self, *args, **kwargs):374 super().__init__(*args, **kwargs)375 376 # TODO: Should be removed once Flash Attention for RoCm is bumped to 2.1.377 # flash_attn<2.1 generates top-left aligned causal mask, while what is needed here is bottom-right alignment, that was made default for flash_attn>=2.1. This attribute is used to handle this difference. Reference: https://github.com/Dao-AILab/flash-attention/releases/tag/v2.1.0.378 # Beware that with flash_attn<2.1, using q_seqlen != k_seqlen (except for the case q_seqlen == 1) produces a wrong mask (top-left).379 self._flash_attn_uses_top_left_mask = flash_attn_supports_top_left_mask()380 381 def forward(382 self,383 hidden_states: torch.Tensor,384 attention_mask: Optional[torch.Tensor] = None,385 position_ids: Optional[torch.LongTensor] = None,386 past_key_values: Optional[HybridMambaAttentionDynamicCache] = None,387 output_attentions: bool = False,388 use_cache: bool = False,389 cache_position: Optional[torch.LongTensor] = None,390 **kwargs,391 ):392 bsz, q_len, _ = hidden_states.size()393 394 query_states = self.q_proj(hidden_states)395 key_states = self.k_proj(hidden_states)396 value_states = self.v_proj(hidden_states)397 398 # Flash attention requires the input to have the shape399 # batch_size x seq_length x head_dim x hidden_dim400 # therefore we just need to keep the original shape401 query_states = query_states.view(bsz, q_len, self.num_heads, self.head_dim)402 key_states = key_states.view(bsz, q_len, self.num_key_value_heads, self.head_dim).transpose(1, 2)403 value_states = value_states.view(bsz, q_len, self.num_key_value_heads, self.head_dim).transpose(1, 2)404 405 if past_key_values is not None:406 key_states, value_states = past_key_values.update(key_states, value_states, self.layer_idx)407 408 # repeat k/v heads if n_kv_heads < n_heads409 key_states = repeat_kv(key_states, self.num_key_value_groups)410 value_states = repeat_kv(value_states, self.num_key_value_groups)411 dropout_rate = 0.0 if not self.training else self.attention_dropout412 413 # In PEFT, usually we cast the layer norms in float32 for training stability reasons414 # therefore the input hidden states gets silently casted in float32. Hence, we need415 # cast them back in float16 just to be sure everything works as expected.416 input_dtype = query_states.dtype417 device_type = query_states.device.type if query_states.device.type != "mps" else "cpu"418 if input_dtype == torch.float32:419 if torch.is_autocast_enabled():420 target_dtype = (421 torch.get_autocast_dtype(device_type)422 if hasattr(torch, "get_autocast_dtype")423 else torch.get_autocast_gpu_dtype()424 )425 # Handle the case where the model is quantized426 elif hasattr(self.config, "_pre_quantization_dtype"):427 target_dtype = self.config._pre_quantization_dtype428 else:429 target_dtype = self.q_proj.weight.dtype430 431 logger.warning_once(432 f"The input hidden states seems to be silently casted in float32, this might be related to"433 f" the fact you have upcasted embedding or layer norm layers in float32. We will cast back the input in"434 f" {target_dtype}."435 )436 437 query_states = query_states.to(target_dtype)438 key_states = key_states.to(target_dtype)439 value_states = value_states.to(target_dtype)440 441 # Reashape to the expected shape for Flash Attention442 key_states = key_states.transpose(1, 2)443 value_states = value_states.transpose(1, 2)444 445 attn_output = _flash_attention_forward(446 query_states,447 key_states,448 value_states,449 attention_mask,450 q_len,451 dropout=dropout_rate,452 sliding_window=getattr(self.config, "sliding_window", None),453 is_causal=self.is_causal,454 use_top_left_mask=self._flash_attn_uses_top_left_mask,455 )456 457 attn_output = attn_output.reshape(bsz, q_len, self.hidden_size).contiguous()458 attn_output = self.o_proj(attn_output)459 460 if not output_attentions:461 attn_weights = None462 463 return attn_output, attn_weights, past_key_values464 465 466# Adapted from transformers.models.mistral.modeling_mistral.MistralSdpaAttention with Mistral->Jamba467class JambaSdpaAttention(JambaAttention):468 """469 Jamba attention module using torch.nn.functional.scaled_dot_product_attention. This module inherits from470 `JambaAttention` as the weights of the module stays untouched. The only changes are on the forward pass to adapt to471 SDPA API.472 """473 474 # Adapted from JambaAttention.forward475 @deprecate_kwarg("past_key_value", new_name="past_key_values", version="4.58")476 def forward(477 self,478 hidden_states: torch.Tensor,479 attention_mask: Optional[torch.Tensor] = None,480 position_ids: Optional[torch.LongTensor] = None,481 past_key_values: Optional[HybridMambaAttentionDynamicCache] = None,482 output_attentions: bool = False,483 use_cache: bool = False,484 cache_position: Optional[torch.LongTensor] = None,485 ) -> tuple[torch.Tensor, Optional[torch.Tensor], Optional[tuple[torch.Tensor]]]:486 if output_attentions:487 # TODO: Improve this warning with e.g. `model.config.attn_implementation = "manual"` once this is implemented.488 logger.warning_once(489 "JambaModel is using JambaSdpaAttention, but `torch.nn.functional.scaled_dot_product_attention` does not support `output_attentions=True`. Falling back to the manual attention implementation, "490 'but specifying the manual implementation will be required from Transformers version v5.0.0 onwards. This warning can be removed using the argument `attn_implementation="eager"` when loading the model.'491 )492 return super().forward(493 hidden_states=hidden_states,494 attention_mask=attention_mask,495 position_ids=position_ids,496 past_key_values=past_key_values,497 output_attentions=output_attentions,498 use_cache=use_cache,499 )500 501 bsz, q_len, _ = hidden_states.size()502 503 query_states = self.q_proj(hidden_states)504 key_states = self.k_proj(hidden_states)505 value_states = self.v_proj(hidden_states)506 507 query_states = query_states.view(bsz, q_len, self.num_heads, self.head_dim).transpose(1, 2)508 key_states = key_states.view(bsz, q_len, self.num_key_value_heads, self.head_dim).transpose(1, 2)509 value_states = value_states.view(bsz, q_len, self.num_key_value_heads, self.head_dim).transpose(1, 2)510 511 if past_key_values is not None:512 key_states, value_states = past_key_values.update(key_states, value_states, self.layer_idx)513 514 key_states = repeat_kv(key_states, self.num_key_value_groups)515 value_states = repeat_kv(value_states, self.num_key_value_groups)516 517 causal_mask = attention_mask518 if attention_mask is not None:519 causal_mask = causal_mask[:, :, :, : key_states.shape[-2]]520 521 # SDPA with memory-efficient backend is currently (torch==2.1.2) bugged with non-contiguous inputs with custom attn_mask,522 # Reference: https://github.com/pytorch/pytorch/issues/112577.523 if query_states.device.type == "cuda" and attention_mask is not None:524 query_states = query_states.contiguous()525 key_states = key_states.contiguous()526 value_states = value_states.contiguous()527 528 # We dispatch to SDPA's Flash Attention or Efficient kernels via this `is_causal` if statement instead of an inline conditional assignment529 # in SDPA to support both torch.compile's dynamic shapes and full graph options. An inline conditional prevents dynamic shapes from compiling.530 # The q_len > 1 is necessary to match with AttentionMaskConverter.to_causal_4d that does not create a causal mask in case q_len == 1.531 is_causal = self.is_causal and causal_mask is None and q_len > 1532 533 attn_output = torch.nn.functional.scaled_dot_product_attention(534 query_states,535 key_states,536 value_states,537 attn_mask=causal_mask,538 dropout_p=self.attention_dropout if self.training else 0.0,539 is_causal=is_causal,540 )541 542 attn_output = attn_output.transpose(1, 2).contiguous()543 attn_output = attn_output.view(bsz, q_len, self.hidden_size)544 545 attn_output = self.o_proj(attn_output)546 547 return attn_output, None, past_key_values548 549 550JAMBA_ATTENTION_CLASSES = {551 "eager": JambaAttention,552 "flash_attention_2": JambaFlashAttention2,553 "sdpa": JambaSdpaAttention,554}555 556 557# Adapted from transformers.models.mamba.modeling_mamba.MambaMixer558class JambaMambaMixer(nn.Module):559 """560 Compute ∆, A, B, C, and D the state space parameters and compute the `contextualized_states`.561 A, D are input independent (see Mamba paper [1] Section 3.5.2 "Interpretation of A" for why A isn't selective)562 ∆, B, C are input-dependent (this is a key difference between Mamba and the linear time invariant S4,563 and is why Mamba is called **selective** state spaces)564 """565 566 def __init__(self, config: JambaConfig, layer_idx):567 super().__init__()568 self.config = config569 self.layer_idx = layer_idx570 self.hidden_size = config.hidden_size571 self.ssm_state_size = config.mamba_d_state572 self.conv_kernel_size = config.mamba_d_conv573 self.intermediate_size = config.mamba_expand * config.hidden_size574 self.time_step_rank = config.mamba_dt_rank575 self.use_conv_bias = config.mamba_conv_bias576 self.use_bias = config.mamba_proj_bias577 self.conv1d = nn.Conv1d(578 in_channels=self.intermediate_size,579 out_channels=self.intermediate_size,580 bias=self.use_conv_bias,581 kernel_size=self.conv_kernel_size,582 groups=self.intermediate_size,583 padding=self.conv_kernel_size - 1,584 )585 586 self.activation = config.hidden_act587 self.act = ACT2FN[config.hidden_act]588 589 self.use_fast_kernels = config.use_mamba_kernels590 591 # projection of the input hidden states592 self.in_proj = nn.Linear(self.hidden_size, self.intermediate_size * 2, bias=self.use_bias)593 # selective projection used to make dt, B and C input dependent594 self.x_proj = nn.Linear(self.intermediate_size, self.time_step_rank + self.ssm_state_size * 2, bias=False)595 # time step projection (discretization)596 self.dt_proj = nn.Linear(self.time_step_rank, self.intermediate_size, bias=True)597 598 # S4D real initialization. These are not discretized!599 # The core is to load them, compute the discrete states, then write the updated state. Keeps the memory bounded600 A = torch.arange(1, self.ssm_state_size + 1)[None, :]601 A = A.expand(self.intermediate_size, -1).contiguous()602 603 self.A_log = nn.Parameter(torch.log(A))604 self.D = nn.Parameter(torch.ones(self.intermediate_size))605 self.out_proj = nn.Linear(self.intermediate_size, self.hidden_size, bias=self.use_bias)606 607 self.dt_layernorm = JambaRMSNorm(self.time_step_rank, eps=config.rms_norm_eps)608 self.b_layernorm = JambaRMSNorm(self.ssm_state_size, eps=config.rms_norm_eps)609 self.c_layernorm = JambaRMSNorm(self.ssm_state_size, eps=config.rms_norm_eps)610 611 if not is_fast_path_available:612 logger.warning_once(613 "The fast path is not available because one of `(selective_state_update, selective_scan_fn, causal_conv1d_fn, causal_conv1d_update, mamba_inner_fn)`"614 " is None. To install follow https://github.com/state-spaces/mamba/#installation and"615 " https://github.com/Dao-AILab/causal-conv1d. If you want to use the naive implementation, set `use_mamba_kernels=False` in the model config"616 )617 618 def cuda_kernels_forward(619 self,620 hidden_states: torch.Tensor,621 cache_params: Optional[HybridMambaAttentionDynamicCache] = None,622 attention_mask: Optional[torch.LongTensor] = None,623 ):624 batch_size, seq_len, _ = hidden_states.shape625 use_precomputed_states = (626 cache_params is not None627 and cache_params.has_previous_state628 and seq_len == 1629 and cache_params.conv_states[self.layer_idx].shape[0]630 == cache_params.ssm_states[self.layer_idx].shape[0]631 == batch_size632 )633 # 1. Gated MLP's linear projection634 projected_states = self.in_proj(hidden_states).transpose(1, 2)635 636 # We can't use `mamba_inner_fn` even if in training and without cache params because we have the637 # inner layernorms which isn't supported by this fused kernel638 hidden_states, gate = projected_states.chunk(2, dim=1)639 640 if attention_mask is not None:641 hidden_states = hidden_states * attention_mask.unsqueeze(1)642 643 # 2. Convolution sequence transformation644 conv_weights = self.conv1d.weight.view(self.conv1d.weight.size(0), self.conv1d.weight.size(2))645 if use_precomputed_states:646 hidden_states = causal_conv1d_update(647 hidden_states.squeeze(-1),648 cache_params.conv_states[self.layer_idx],649 conv_weights,650 self.conv1d.bias,651 self.activation,652 )653 hidden_states = hidden_states.unsqueeze(-1)654 else:655 if cache_params is not None:656 conv_states = nn.functional.pad(hidden_states, (self.conv_kernel_size - hidden_states.shape[-1], 0))657 cache_params.conv_states[self.layer_idx].copy_(conv_states)658 hidden_states = causal_conv1d_fn(hidden_states, conv_weights, self.conv1d.bias, activation=self.activation)659 660 if attention_mask is not None:661 hidden_states = hidden_states * attention_mask.unsqueeze(1)662 663 # 3. State Space Model sequence transformation664 # 3.a. input varying initialization of time_step, B and C665 ssm_parameters = self.x_proj(hidden_states.transpose(1, 2))666 time_step, B, C = torch.split(667 ssm_parameters, [self.time_step_rank, self.ssm_state_size, self.ssm_state_size], dim=-1668 )669 670 time_step = self.dt_layernorm(time_step)671 B = self.b_layernorm(B)672 C = self.c_layernorm(C)673 674 # Here we need to apply dt_proj without the bias, as the bias is added in the selective scan kernel.675 # This is a hack to apply dt_proj while still using the forward pass of `torch.nn.Linear`, which is needed676 # in order to make quantization work. Quantization code replaces `torch.nn.Linear` layers with quantized677 # linear layers, and requires to call the forward pass directly.678 # Quantized model can't work with the original code:679 # ```discrete_time_step = self.dt_proj.weight @ time_step.transpose(1, 2)```680 time_proj_bias = self.dt_proj.bias.data681 with torch.no_grad():682 self.dt_proj.bias.data = torch.zeros_like(self.dt_proj.bias.data)683 discrete_time_step = self.dt_proj(time_step).transpose(1, 2)684 with torch.no_grad():685 self.dt_proj.bias.data = time_proj_bias686 687 A = -torch.exp(self.A_log.float())688 # 3.c perform the recurrence y ← SSM(A, B, C)(x)689 time_proj_bias = time_proj_bias.float() if time_proj_bias is not None else None690 if use_precomputed_states:691 scan_outputs = selective_state_update(692 cache_params.ssm_states[self.layer_idx],693 hidden_states[..., 0],694 discrete_time_step[..., 0],695 A,696 B[:, 0],697 C[:, 0],698 self.D,699 gate[..., 0],700 time_proj_bias,701 dt_softplus=True,702 ).unsqueeze(-1)703 else:704 scan_outputs, ssm_state = selective_scan_fn(705 hidden_states,706 discrete_time_step,707 A,708 B.transpose(1, 2),709 C.transpose(1, 2),710 self.D.float(),711 gate,712 time_proj_bias,713 delta_softplus=True,714 return_last_state=True,715 )716 if ssm_state is not None and cache_params is not None:717 cache_params.ssm_states[self.layer_idx].copy_(ssm_state)718 719 # 4. Final linear projection720 contextualized_states = self.out_proj(scan_outputs.transpose(1, 2))721 722 return contextualized_states723 724 # fmt: off725 def slow_forward(self, input_states, cache_params: Optional[HybridMambaAttentionDynamicCache] = None, attention_mask: Optional[torch.LongTensor] = None):726 batch_size, seq_len, _ = input_states.shape727 dtype = input_states.dtype728 # 1. Gated MLP's linear projection729 projected_states = self.in_proj(input_states).transpose(1, 2) # [batch, 2 * intermediate_size, seq_len]730 hidden_states, gate = projected_states.chunk(2, dim=1)731 732 if attention_mask is not None:733 hidden_states = hidden_states * attention_mask.unsqueeze(1)734 735 use_cache = isinstance(cache_params, HybridMambaAttentionDynamicCache)736 # 2. Convolution sequence transformation737 if use_cache and cache_params.ssm_states[self.layer_idx].shape[0] == batch_size:738 if self.training:739 # In training mode, we don't want to perform in-place operations on ssm_state so we can compute the backwards pass740 ssm_state = cache_params.ssm_states[self.layer_idx].clone()741 else:742 ssm_state = cache_params.ssm_states[self.layer_idx]743 744 ssm_state = ssm_state.to(hidden_states.device)745 746 if cache_params.has_previous_state and seq_len == 1 and \747 cache_params.conv_states[self.layer_idx].shape[0] == batch_size:748 conv_state = cache_params.conv_states[self.layer_idx] # [batch, intermediate_size, conv_kernel_size]749 conv_state = torch.roll(conv_state, shifts=-1, dims=-1)750 conv_state[:, :, -1] = hidden_states[:, :, 0]751 cache_params.conv_states[self.layer_idx] = conv_state752 hidden_states = torch.sum(conv_state * self.conv1d.weight[:, 0, :], dim=-1)753 if self.use_conv_bias:754 hidden_states += self.conv1d.bias755 hidden_states = self.act(hidden_states).to(dtype).unsqueeze(-1) # [batch, intermediate_size, 1] : decoding756 else:757 conv_state = nn.functional.pad(758 hidden_states,759 (self.conv_kernel_size - hidden_states.shape[-1], 0)760 )761 cache_params.conv_states[self.layer_idx] = conv_state762 hidden_states = self.act(self.conv1d(hidden_states)[..., :seq_len]) # [batch, intermediate_size, seq_len]763 else:764 ssm_state = torch.zeros(765 (batch_size, self.intermediate_size, self.ssm_state_size),766 device=hidden_states.device, dtype=dtype767 )768 hidden_states = self.act(self.conv1d(hidden_states)[..., :seq_len]) # [batch, intermediate_size, seq_len]769 770 if attention_mask is not None:771 hidden_states = hidden_states * attention_mask.unsqueeze(1)772 773 # 3. State Space Model sequence transformation774 # 3.a. Selection: [batch, seq_len, self.time_step_rank + self.ssm_state_size * 2]775 ssm_parameters = self.x_proj(hidden_states.transpose(1, 2))776 time_step, B, C = torch.split(777 ssm_parameters, [self.time_step_rank, self.ssm_state_size, self.ssm_state_size], dim=-1778 )779 780 time_step = self.dt_layernorm(time_step)781 B = self.b_layernorm(B)782 C = self.c_layernorm(C)783 784 discrete_time_step = self.dt_proj(time_step) # [batch, seq_len, intermediate_size]785 discrete_time_step = nn.functional.softplus(discrete_time_step).transpose(1, 2) # [batch, intermediate_size, seq_len]786 787 # 3.b. Discretization: B and C to [batch, seq_len, intermediate_size, ssm_state_size] (SRAM)788 A = -torch.exp(self.A_log.float()) # [intermediate_size, ssm_state_size]789 discrete_A = torch.exp(A[None, :, None, :] * discrete_time_step[:, :, :, None]) # [batch, intermediate_size, seq_len, ssm_state_size]790 discrete_B = discrete_time_step[:, :, :, None] * B[:, None, :, :].float() # [batch, intermediate_size, seq_len, ssm_state_size]791 deltaB_u = discrete_B * hidden_states[:, :, :, None].float()792 # 3.c perform the recurrence y ← SSM(A, B, C)(x)793 scan_outputs = []794 for i in range(seq_len):795 ssm_state = discrete_A[:, :, i, :] * ssm_state + deltaB_u[:, :, i, :] # [batch, intermediate_size, ssm_state]796 scan_output = torch.matmul(ssm_state.to(dtype), C[:, i, :].unsqueeze(-1)) # [batch, intermediate_size, 1]797 scan_outputs.append(scan_output[:, :, 0])798 scan_output = torch.stack(scan_outputs, dim=-1) # [batch, intermediate_size, seq_len]799 scan_output = scan_output + (hidden_states * self.D[None, :, None])800 scan_output = (scan_output * self.act(gate))801 802 if use_cache:803 cache_params.ssm_states[self.layer_idx] = ssm_state804 805 # 4. Final linear projection806 contextualized_states = self.out_proj(scan_output.transpose(1, 2)) # [batch, seq_len, hidden_size]807 return contextualized_states808 # fmt: on809 810 def forward(811 self,812 hidden_states,813 cache_params: Optional[HybridMambaAttentionDynamicCache] = None,814 attention_mask: Optional[torch.LongTensor] = None,815 ):816 if self.use_fast_kernels:817 if not is_fast_path_available or "cuda" not in self.x_proj.weight.device.type:818 raise ValueError(819 "Fast Mamba kernels are not available. Make sure to they are installed and that the mamba module is on a CUDA device"820 )821 return self.cuda_kernels_forward(hidden_states, cache_params, attention_mask)822 return self.slow_forward(hidden_states, cache_params, attention_mask)823 824 825# Copied from transformers.models.mistral.modeling_mistral.MistralMLP with Mistral->Jamba826class JambaMLP(nn.Module):827 def __init__(self, config):828 super().__init__()829 self.config = config830 self.hidden_size = config.hidden_size831 self.intermediate_size = config.intermediate_size832 self.gate_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=False)833 self.up_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=False)834 self.down_proj = nn.Linear(self.intermediate_size, self.hidden_size, bias=False)835 self.act_fn = ACT2FN[config.hidden_act]836 837 def forward(self, x):838 down_proj = self.down_proj(self.act_fn(self.gate_proj(x)) * self.up_proj(x))839 return down_proj840 841 842# Adapted from transformers.models.mixtral.modeling_mixtral.MixtralSparseMoeBlock with Mistral->Jamba843class JambaSparseMoeBlock(nn.Module):844 """845 This implementation is846 strictly equivalent to standard MoE with full capacity (no847 dropped tokens). It's faster since it formulates MoE operations848 in terms of block-sparse operations to accommodate imbalanced849 assignments of tokens to experts, whereas standard MoE either850 (1) drop tokens at the cost of reduced performance or (2) set851 capacity factor to number of experts and thus waste computation852 and memory on padding.853 """854 855 def __init__(self, config: JambaConfig):856 super().__init__()857 self.hidden_dim = config.hidden_size858 self.ffn_dim = config.intermediate_size859 self.num_experts = config.num_experts860 self.top_k = config.num_experts_per_tok861 862 self.router = nn.Linear(self.hidden_dim, self.num_experts, bias=False)863 self.experts = nn.ModuleList([JambaMLP(config) for _ in range(self.num_experts)])864 865 def forward(self, hidden_states: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:866 """ """867 batch_size, sequence_length, hidden_dim = hidden_states.shape868 869 hidden_states = hidden_states.view(-1, hidden_dim)870 # router_logits: (batch * sequence_length, n_experts)871 router_logits = self.router(hidden_states)872 routing_weights = F.softmax(router_logits, dim=1, dtype=torch.float)873 routing_weights, selected_experts = torch.topk(routing_weights, self.top_k, dim=-1)874 # we cast back to the input dtype875 routing_weights = routing_weights.to(hidden_states.dtype)876 877 final_hidden_states = torch.zeros(878 (batch_size * sequence_length, hidden_dim), dtype=hidden_states.dtype, device=hidden_states.device879 )880 881 # One hot encode the selected experts to create an expert mask882 # this will be used to easily index which expert is going to be sollicitated883 expert_mask = torch.nn.functional.one_hot(selected_experts, num_classes=self.num_experts).permute(2, 1, 0)884 885 # Loop over all available experts in the model and perform the computation on each expert886 for expert_idx in range(self.num_experts):887 expert_layer = self.experts[expert_idx]888 idx, top_x = torch.where(expert_mask[expert_idx])889 890 if top_x.shape[0] == 0:891 continue892 893 # Index the correct hidden states and compute the expert hidden state for894 # the current expert. We need to make sure to multiply the output hidden895 # states by `routing_weights` on the corresponding tokens (top-1 and top-2)896 current_state = hidden_states[None, top_x].reshape(-1, hidden_dim)897 current_hidden_states = expert_layer(current_state) * routing_weights[top_x, idx, None]898 899 # However `index_add_` only support torch tensors for indexing so we'll use900 # the `top_x` tensor here.901 final_hidden_states.index_add_(0, top_x, current_hidden_states.to(hidden_states.dtype))902 final_hidden_states = final_hidden_states.reshape(batch_size, sequence_length, hidden_dim)903 return final_hidden_states, router_logits904 905 906class JambaAttentionDecoderLayer(GradientCheckpointingLayer):907 def __init__(self, config: JambaConfig, layer_idx: int):908 super().__init__()909 num_experts = config.layers_num_experts[layer_idx]910 self.self_attn = JAMBA_ATTENTION_CLASSES[config._attn_implementation](config, layer_idx)911 912 ffn_layer_class = JambaSparseMoeBlock if num_experts > 1 else JambaMLP913 self.feed_forward = ffn_layer_class(config)914 self.input_layernorm = JambaRMSNorm(config.hidden_size, eps=config.rms_norm_eps)915 self.pre_ff_layernorm = JambaRMSNorm(config.hidden_size, eps=config.rms_norm_eps)916 917 @deprecate_kwarg("past_key_value", new_name="past_key_values", version="4.58")918 def forward(919 self,920 hidden_states: torch.Tensor,921 attention_mask: Optional[torch.Tensor] = None,922 position_ids: Optional[torch.LongTensor] = None,923 past_key_values: Optional[HybridMambaAttentionDynamicCache] = None,924 output_attentions: Optional[bool] = False,925 output_router_logits: Optional[bool] = False,926 use_cache: Optional[bool] = False,927 cache_position: Optional[torch.LongTensor] = None,928 ) -> tuple[torch.FloatTensor, Optional[tuple[torch.FloatTensor, torch.FloatTensor]]]:929 """930 Args:931 hidden_states (`torch.FloatTensor`): input to the layer of shape `(batch, seq_len, embed_dim)`932 attention_mask (`torch.FloatTensor`, *optional*): attention mask of size933 `(batch, sequence_length)` where padding elements are indicated by 0.934 past_key_values (`HybridMambaAttentionDynamicCache`, *optional*): cached past key and value projection states935 output_attentions (`bool`, *optional*):936 Whether or not to return the attentions tensors of all attention layers. See `attentions` under937 returned tensors for more detail.938 output_router_logits (`bool`, *optional*):939 Whether or not to return the logits of all the routers. They are useful for computing the router loss, and940 should not be returned during inference.941 use_cache (`bool`, *optional*):942 If set to `True`, `past_key_values` key value states are returned and can be used to speed up decoding943 (see `past_key_values`).944 cache_position (`torch.LongTensor` of shape `(sequence_length)`, *optional*):945 Indices depicting the position of the input sequence tokens in the sequence.946 """947 948 residual = hidden_states949 950 hidden_states = self.input_layernorm(hidden_states)951 952 hidden_states, self_attn_weights, present_key_value = self.self_attn(953 hidden_states=hidden_states,954 attention_mask=attention_mask,955 position_ids=position_ids,956 past_key_values=past_key_values,957 output_attentions=output_attentions,958 use_cache=use_cache,959 cache_position=cache_position,960 )961 962 # residual connection after attention963 hidden_states = residual + hidden_states964 965 # feed-forward (experts/MLP)966 residual = hidden_states967 hidden_states = self.pre_ff_layernorm(hidden_states)968 ff_outputs = self.feed_forward(hidden_states)969 if isinstance(ff_outputs, tuple):970 hidden_states, router_logits = ff_outputs971 else:972 hidden_states, router_logits = ff_outputs, None973 hidden_states = residual + hidden_states974 975 outputs = (hidden_states,)976 977 if output_attentions:978 outputs += (self_attn_weights,)979 980 if use_cache:981 outputs += (present_key_value,)982 983 if output_router_logits:984 outputs += (router_logits,)985 986 return outputs987 988 989class JambaMambaDecoderLayer(GradientCheckpointingLayer):990 def __init__(self, config: JambaConfig, layer_idx: int):991 super().__init__()992 num_experts = config.layers_num_experts[layer_idx]993 self.mamba = JambaMambaMixer(config=config, layer_idx=layer_idx)994 995 ffn_layer_class = JambaSparseMoeBlock if num_experts > 1 else JambaMLP996 self.feed_forward = ffn_layer_class(config)997 self.input_layernorm = JambaRMSNorm(config.hidden_size, eps=config.rms_norm_eps)998 self.pre_ff_layernorm = JambaRMSNorm(config.hidden_size, eps=config.rms_norm_eps)999 1000 @deprecate_kwarg("past_key_value", new_name="past_key_values", version="4.58")1001 def forward(1002 self,1003 hidden_states: torch.Tensor,1004 attention_mask: Optional[torch.Tensor] = None,1005 position_ids: Optional[torch.LongTensor] = None,1006 past_key_values: Optional[HybridMambaAttentionDynamicCache] = None,1007 output_attentions: Optional[bool] = False,1008 output_router_logits: Optional[bool] = False,1009 use_cache: Optional[bool] = False,1010 cache_position: Optional[torch.LongTensor] = None,1011 ) -> tuple[torch.FloatTensor, Optional[tuple[torch.FloatTensor, torch.FloatTensor]]]:1012 """1013 Args:1014 hidden_states (`torch.FloatTensor`): input to the layer of shape `(batch, seq_len, embed_dim)`1015 attention_mask (`torch.FloatTensor`, *optional*): attention mask of size1016 `(batch, sequence_length)` where padding elements are indicated by 0.1017 past_key_values (`HybridMambaAttentionDynamicCache`, *optional*): cached past key and value projection states1018 output_attentions (`bool`, *optional*):1019 Whether or not to return the attentions tensors of all attention layers. See `attentions` under1020 returned tensors for more detail.1021 output_router_logits (`bool`, *optional*):1022 Whether or not to return the logits of all the routers. They are useful for computing the router loss, and1023 should not be returned during inference.1024 use_cache (`bool`, *optional*):1025 If set to `True`, `past_key_values` key value states are returned and can be used to speed up decoding1026 (see `past_key_values`).1027 cache_position (`torch.LongTensor` of shape `(sequence_length)`, *optional*):1028 Indices depicting the position of the input sequence tokens in the sequence.1029 """1030 1031 residual = hidden_states1032 1033 hidden_states = self.input_layernorm(hidden_states)1034 1035 hidden_states = self.mamba(1036 hidden_states=hidden_states,1037 cache_params=past_key_values,1038 attention_mask=attention_mask,1039 )1040 self_attn_weights = None1041 1042 # residual connection after mamba1043 hidden_states = residual + hidden_states1044 1045 # feed-forward (experts/MLP)1046 residual = hidden_states1047 hidden_states = self.pre_ff_layernorm(hidden_states)1048 ff_outputs = self.feed_forward(hidden_states)1049 if isinstance(ff_outputs, tuple):1050 hidden_states, router_logits = ff_outputs1051 else:1052 hidden_states, router_logits = ff_outputs, None1053 hidden_states = residual + hidden_states1054 1055 outputs = (hidden_states,)1056 1057 if output_attentions:1058 outputs += (self_attn_weights,)1059 1060 if use_cache:1061 outputs += (past_key_values,)1062 1063 if output_router_logits:1064 outputs += (router_logits,)1065 1066 return outputs1067 1068 1069@auto_docstring1070class JambaPreTrainedModel(PreTrainedModel):1071 config: JambaConfig1072 base_model_prefix = "model"1073 supports_gradient_checkpointing = True1074 _no_split_modules = ["JambaAttentionDecoderLayer", "JambaMambaDecoderLayer"]1075 _skip_keys_device_placement = "past_key_values"1076 _supports_flash_attn = True1077 _supports_sdpa = True1078 # Note: only supports HybridMambaAttentionDynamicCache1079 _is_stateful = True1080 1081 def _init_weights(self, module):1082 std = self.config.initializer_range1083 if isinstance(module, (nn.Linear, nn.Conv1d)):1084 module.weight.data.normal_(mean=0.0, std=std)1085 if module.bias is not None:1086 module.bias.data.zero_()1087 elif isinstance(module, nn.Embedding):1088 module.weight.data.normal_(mean=0.0, std=std)1089 if module.padding_idx is not None:1090 module.weight.data[module.padding_idx].zero_()1091 elif isinstance(module, JambaRMSNorm):1092 module.weight.data.fill_(1.0)1093 elif isinstance(module, JambaMambaMixer):1094 A = torch.arange(1, module.ssm_state_size + 1)[None, :]1095 A = A.expand(module.intermediate_size, -1).contiguous()1096 module.A_log.data.copy_(torch.log(A))1097 module.D.data.fill_(1.0)1098 1099 1100ALL_DECODER_LAYER_TYPES = {"attention": JambaAttentionDecoderLayer, "mamba": JambaMambaDecoderLayer}1101 1102 1103# Adapted from transformers.models.mistral.modeling_mistral.MistralModel with MISTRAL->JAMBA, Mistral->Jamba1104@auto_docstring1105class JambaModel(JambaPreTrainedModel):1106 """1107 Transformer decoder consisting of *config.num_hidden_layers* layers. Each layer is a [`JambaDecoderLayer`]1108 1109 Args:1110 config: JambaConfig1111 """1112 1113 def __init__(self, config: JambaConfig):1114 super().__init__(config)1115 self.padding_idx = config.pad_token_id1116 self.vocab_size = config.vocab_size1117 1118 self.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size, self.padding_idx)1119 decoder_layers = []1120 for i in range(config.num_hidden_layers):1121 layer_class = ALL_DECODER_LAYER_TYPES[config.layers_block_type[i]]1122 decoder_layers.append(layer_class(config, layer_idx=i))1123 self.layers = nn.ModuleList(decoder_layers)1124 1125 self._attn_implementation = config._attn_implementation1126 self.final_layernorm = JambaRMSNorm(config.hidden_size, eps=config.rms_norm_eps)1127 1128 self.gradient_checkpointing = False1129 # Initialize weights and apply final processing1130 self.post_init()1131 1132 @can_return_tuple1133 @auto_docstring1134 def forward(1135 self,1136 input_ids: Optional[torch.LongTensor] = None,1137 attention_mask: Optional[torch.Tensor] = None,1138 position_ids: Optional[torch.LongTensor] = None,1139 past_key_values: Optional[HybridMambaAttentionDynamicCache] = None,1140 inputs_embeds: Optional[torch.FloatTensor] = None,1141 use_cache: Optional[bool] = None,1142 output_attentions: Optional[bool] = None,1143 output_hidden_states: Optional[bool] = None,1144 output_router_logits: Optional[bool] = None,1145 cache_position: Optional[torch.LongTensor] = None,1146 **kwargs: Unpack[TransformersKwargs],1147 ) -> MoeModelOutputWithPast:1148 output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions1149 output_router_logits = (1150 output_router_logits if output_router_logits is not None else self.config.output_router_logits1151 )1152 output_hidden_states = (1153 output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states1154 )1155 use_cache = use_cache if use_cache is not None else self.config.use_cache1156 1157 if (input_ids is None) ^ (inputs_embeds is not None):1158 raise ValueError("You must specify exactly one of input_ids or inputs_embeds")1159 1160 if self.gradient_checkpointing and self.training and use_cache:1161 logger.warning_once(1162 "`use_cache=True` is incompatible with gradient checkpointing. Setting `use_cache=False`."1163 )1164 use_cache = False1165 1166 if inputs_embeds is None:1167 inputs_embeds = self.embed_tokens(input_ids)1168 hidden_states = inputs_embeds1169 1170 if use_cache and past_key_values is None:1171 logger.warning_once(1172 "Jamba requires an initialized `HybridMambaAttentionDynamicCache` to return a cache. None was "1173 "provided, so no cache will be returned."1174 )1175 1176 if cache_position is None:1177 cache_position = torch.arange(hidden_states.shape[1], device=hidden_states.device)1178 if position_ids is None:1179 position_ids = cache_position.unsqueeze(0)1180 1181 causal_mask = self._update_causal_mask(attention_mask, inputs_embeds, cache_position)1182 mamba_mask = self._update_mamba_mask(attention_mask, cache_position)1183 1184 all_hidden_states = () if output_hidden_states else None1185 all_self_attns = () if output_attentions else None1186 all_router_logits = () if output_router_logits else None1187 1188 for decoder_layer in self.layers:1189 # Depending on the layer type we opt for 2D base attention mask (Mamba) or 4D causal mask (Attention)1190 layer_mask = mamba_mask if isinstance(decoder_layer, JambaMambaDecoderLayer) else causal_mask1191 1192 if output_hidden_states:1193 all_hidden_states += (hidden_states,)1194 1195 layer_outputs = decoder_layer(1196 hidden_states,1197 attention_mask=layer_mask,1198 position_ids=position_ids,1199 past_key_values=past_key_values,1200 output_attentions=output_attentions,