amd/Instella-3B-Instruct
59233
1# This code has been adapter from the Olmo2 codebase and updated to match the Instella model details. 2# https://github.com/huggingface/transformers/tree/v4.47.1/src/transformers/models/olmo23 4import math5from typing import List, Optional, Tuple, Union6 7import torch8from torch import nn9 10from transformers.activations import ACT2FN11from transformers.cache_utils import Cache, DynamicCache, StaticCache12from transformers.generation import GenerationMixin13from transformers.modeling_attn_mask_utils import AttentionMaskConverter14from transformers.modeling_outputs import BaseModelOutputWithPast, CausalLMOutputWithPast15from transformers.modeling_utils import PreTrainedModel16from transformers.utils import (17 add_start_docstrings,18 add_start_docstrings_to_model_forward,19 is_flash_attn_2_available,20 is_flash_attn_greater_or_equal_2_10,21 logging,22 replace_return_docstrings,23)24 25"""26Instella configuration27"""28 29from transformers import AutoConfig, PretrainedConfig30 31class InstellaConfig(PretrainedConfig):32 r"""33 This is the configuration class to store the configuration of a [`Instella2Model`]. It is used to instantiate an Instella234 model according to the specified arguments, defining the model architecture. 35 36 Configuration objects inherit from [`PretrainedConfig`] and can be used to control the model outputs. Read the37 documentation from [`PretrainedConfig`] for more information.38 39 40 Args:41 vocab_size (`int`, *optional*, defaults to 50304):42 Vocabulary size of the Instella2 model. Defines the number of different tokens that can be represented by the43 `inputs_ids` passed when calling [`Instella2Model`]44 hidden_size (`int`, *optional*, defaults to 4096):45 Dimension of the hidden representations.46 intermediate_size (`int`, *optional*, defaults to 11008):47 Dimension of the MLP representations.48 num_hidden_layers (`int`, *optional*, defaults to 32):49 Number of hidden layers in the Transformer decoder.50 num_attention_heads (`int`, *optional*, defaults to 32):51 Number of attention heads for each attention layer in the Transformer decoder.52 num_key_value_heads (`int`, *optional*):53 This is the number of key_value heads that should be used to implement Grouped Query Attention. If54 `num_key_value_heads=num_attention_heads`, the model will use Multi Head Attention (MHA), if55 `num_key_value_heads=1` the model will use Multi Query Attention (MQA) otherwise GQA is used. When56 converting a multi-head checkpoint to a GQA checkpoint, each group key and value head should be constructed57 by meanpooling all the original heads within that group. For more details checkout [this58 paper](https://arxiv.org/pdf/2305.13245.pdf). If it is not specified, will default to59 `num_attention_heads`.60 hidden_act (`str` or `function`, *optional*, defaults to `"silu"`):61 The non-linear activation function (function or string) in the decoder.62 max_position_embeddings (`int`, *optional*, defaults to 2048):63 The maximum sequence length that this model might ever be used with.64 initializer_range (`float`, *optional*, defaults to 0.02):65 The standard deviation of the truncated_normal_initializer for initializing all weight matrices.66 use_cache (`bool`, *optional*, defaults to `True`):67 Whether or not the model should return the last key/values attentions (not used by all models). Only68 relevant if `config.is_decoder=True`.69 pad_token_id (`int`, *optional*, defaults to 1):70 Padding token id.71 bos_token_id (`int`, *optional*):72 Beginning of stream token id.73 eos_token_id (`int`, *optional*, defaults to 50279):74 End of stream token id.75 tie_word_embeddings (`bool`, *optional*, defaults to `False`):76 Whether to tie weight embeddings77 rope_theta (`float`, *optional*, defaults to 10000.0):78 The base period of the RoPE embeddings.79 rope_scaling (`Dict`, *optional*):80 Dictionary containing the scaling configuration for the RoPE embeddings. Currently supports two scaling81 strategies: linear and dynamic. Their scaling factor must be a float greater than 1. The expected format is82 `{"type": strategy name, "factor": scaling factor}`. When using this flag, don't update83 `max_position_embeddings` to the expected new maximum. See the following thread for more information on how84 these scaling strategies behave:85 https://www.reddit.com/r/LocalLLaMA/comments/14mrgpr/dynamically_scaled_rope_further_increases/. This is an86 experimental feature, subject to breaking API changes in future versions.87 attention_bias (`bool`, defaults to `False`, *optional*, defaults to `False`):88 Whether to use a bias in the query, key, value and output projection layers during self-attention.89 attention_dropout (`float`, *optional*, defaults to 0.0):90 The dropout ratio for the attention probabilities.91 rms_norm_eps (`float`, *optional*, defaults to 1e-05):92 The epsilon used by the rms normalization layers.93 94 ```python95 >>> from transformers import Instella2Model, Instella2Config96 97 >>> configuration = Instella2Config()98 >>> model = Instella2Model(configuration)99 100 >>> # Accessing the model configuration101 >>> configuration = model.config102 ```103 """104 105 model_type = "instella"106 keys_to_ignore_at_inference = ["past_key_values"]107 108 def __init__(109 self,110 vocab_size=50304,111 hidden_size=4096,112 intermediate_size=11008,113 num_hidden_layers=32,114 num_attention_heads=32,115 num_key_value_heads=None,116 hidden_act="silu",117 max_position_embeddings=2048,118 initializer_range=0.02,119 use_cache=True,120 pad_token_id=1,121 bos_token_id=None,122 eos_token_id=50279,123 tie_word_embeddings=False,124 rope_theta=10000.0,125 rope_scaling=None,126 attention_bias=False,127 attention_dropout=0.0,128 rms_norm_eps=1e-5,129 **kwargs,130 ):131 super().__init__(132 pad_token_id=pad_token_id,133 bos_token_id=bos_token_id,134 eos_token_id=eos_token_id,135 tie_word_embeddings=tie_word_embeddings,136 **kwargs,137 )138 self.vocab_size = vocab_size139 self.max_position_embeddings = max_position_embeddings140 self.hidden_size = hidden_size141 self.intermediate_size = intermediate_size142 self.num_hidden_layers = num_hidden_layers143 self.num_attention_heads = num_attention_heads144 145 # for backward compatibility146 if num_key_value_heads is None:147 num_key_value_heads = num_attention_heads148 149 self.num_key_value_heads = num_key_value_heads150 self.hidden_act = hidden_act151 self.initializer_range = initializer_range152 self.use_cache = use_cache153 self.rope_theta = rope_theta154 self.rope_scaling = rope_scaling155 self._rope_scaling_validation()156 self.attention_bias = attention_bias157 self.attention_dropout = attention_dropout158 159 self.rms_norm_eps = rms_norm_eps160 161 def _rope_scaling_validation(self):162 """163 Validate the `rope_scaling` configuration.164 """165 if self.rope_scaling is None:166 return167 168 if not isinstance(self.rope_scaling, dict) or len(self.rope_scaling) != 2:169 raise ValueError(170 "`rope_scaling` must be a dictionary with two fields, `type` and `factor`, " f"got {self.rope_scaling}"171 )172 rope_scaling_type = self.rope_scaling.get("type", None)173 rope_scaling_factor = self.rope_scaling.get("factor", None)174 if rope_scaling_type is None or rope_scaling_type not in ["linear", "dynamic"]:175 raise ValueError(176 f"`rope_scaling`'s type field must be one of ['linear', 'dynamic'], got {rope_scaling_type}"177 )178 if rope_scaling_factor is None or not isinstance(rope_scaling_factor, float) or rope_scaling_factor <= 1.0:179 raise ValueError(f"`rope_scaling`'s factor field must be a float > 1, got {rope_scaling_factor}")180 181 182if is_flash_attn_2_available():183 from transformers.modeling_flash_attention_utils import _flash_attention_forward184 185 186logger = logging.get_logger(__name__)187 188_CONFIG_FOR_DOC = "InstellaConfig"189 190 191class InstellaRMSNorm(nn.Module):192 def __init__(self, hidden_size, eps=1e-6):193 """194 InstellaRMSNorm is equivalent to T5LayerNorm195 """196 super().__init__()197 self.weight = nn.Parameter(torch.ones(hidden_size))198 self.variance_epsilon = eps199 200 def forward(self, hidden_states):201 input_dtype = hidden_states.dtype202 hidden_states = hidden_states.to(torch.float32)203 variance = hidden_states.pow(2).mean(-1, keepdim=True)204 hidden_states = hidden_states * torch.rsqrt(variance + self.variance_epsilon)205 return self.weight * hidden_states.to(input_dtype)206 207 def extra_repr(self):208 return f"{tuple(self.weight.shape)}, eps={self.variance_epsilon}"209 210 211# copied from transformers.models.llama.modeling_llama.LlamaRotaryEmbedding with Llama->Instella212# TODO(joao): add me back asap :)213class InstellaRotaryEmbedding(nn.Module):214 def __init__(self, dim, max_position_embeddings=2048, base=10000, device=None, scaling_factor=1.0):215 super().__init__()216 self.scaling_factor = scaling_factor217 self.dim = dim218 self.max_position_embeddings = max_position_embeddings219 self.base = base220 inv_freq = 1.0 / (self.base ** (torch.arange(0, self.dim, 2, dtype=torch.int64).float().to(device) / self.dim))221 self.register_buffer("inv_freq", inv_freq, persistent=False)222 # For BC we register cos and sin cached223 self.max_seq_len_cached = max_position_embeddings224 225 @torch.no_grad()226 def forward(self, x, position_ids):227 # x: [bs, num_attention_heads, seq_len, head_size]228 inv_freq_expanded = self.inv_freq[None, :, None].float().expand(position_ids.shape[0], -1, 1)229 position_ids_expanded = position_ids[:, None, :].float()230 # Force float32 since bfloat16 loses precision on long contexts231 # See https://github.com/huggingface/transformers/pull/29285232 device_type = x.device.type233 device_type = device_type if isinstance(device_type, str) and device_type != "mps" else "cpu"234 with torch.autocast(device_type=device_type, enabled=False):235 freqs = (inv_freq_expanded.float() @ position_ids_expanded.float()).transpose(1, 2)236 emb = torch.cat((freqs, freqs), dim=-1)237 cos = emb.cos()238 sin = emb.sin()239 return cos.to(dtype=x.dtype), sin.to(dtype=x.dtype)240 241 242# copied from transformers.models.llama.modeling_llama.LlamaLinearScalingRotaryEmbedding with Llama->Instella243# TODO(joao): add me back asap :)244class InstellaLinearScalingRotaryEmbedding(InstellaRotaryEmbedding):245 """InstellaRotaryEmbedding extended with linear scaling. Credits to the Reddit user /u/kaiokendev"""246 247 def forward(self, x, position_ids):248 # difference to the original RoPE: a scaling factor is aplied to the position ids249 position_ids = position_ids.float() / self.scaling_factor250 cos, sin = super().forward(x, position_ids)251 return cos, sin252 253 254# copied from transformers.models.llama.modeling_llama.LlamaDynamicNTKScalingRotaryEmbedding with Llama->Instella255# TODO(joao): add me back asap :)256class InstellaDynamicNTKScalingRotaryEmbedding(InstellaRotaryEmbedding):257 """InstellaRotaryEmbedding extended with Dynamic NTK scaling. Credits to the Reddit users /u/bloc97 and /u/emozilla"""258 259 def forward(self, x, position_ids):260 # difference to the original RoPE: inv_freq is recomputed when the sequence length > original length261 seq_len = torch.max(position_ids) + 1262 if seq_len > self.max_position_embeddings:263 base = self.base * (264 (self.scaling_factor * seq_len / self.max_position_embeddings) - (self.scaling_factor - 1)265 ) ** (self.dim / (self.dim - 2))266 inv_freq = 1.0 / (267 base ** (torch.arange(0, self.dim, 2, dtype=torch.int64).float().to(x.device) / self.dim)268 )269 self.register_buffer("inv_freq", inv_freq, persistent=False) # TODO joao: this may break with compilation270 271 cos, sin = super().forward(x, position_ids)272 return cos, sin273 274 275def rotate_half(x):276 """Rotates half the hidden dims of the input."""277 x1 = x[..., : x.shape[-1] // 2]278 x2 = x[..., x.shape[-1] // 2 :]279 return torch.cat((-x2, x1), dim=-1)280 281 282def apply_rotary_pos_emb(q, k, cos, sin, position_ids=None, unsqueeze_dim=1):283 """Applies Rotary Position Embedding to the query and key tensors.284 285 Args:286 q (`torch.Tensor`): The query tensor.287 k (`torch.Tensor`): The key tensor.288 cos (`torch.Tensor`): The cosine part of the rotary embedding.289 sin (`torch.Tensor`): The sine part of the rotary embedding.290 position_ids (`torch.Tensor`, *optional*):291 Deprecated and unused.292 unsqueeze_dim (`int`, *optional*, defaults to 1):293 The 'unsqueeze_dim' argument specifies the dimension along which to unsqueeze cos[position_ids] and294 sin[position_ids] so that they can be properly broadcasted to the dimensions of q and k. For example, note295 that cos[position_ids] and sin[position_ids] have the shape [batch_size, seq_len, head_dim]. Then, if q and296 k have the shape [batch_size, heads, seq_len, head_dim], then setting unsqueeze_dim=1 makes297 cos[position_ids] and sin[position_ids] broadcastable to the shapes of q and k. Similarly, if q and k have298 the shape [batch_size, seq_len, heads, head_dim], then set unsqueeze_dim=2.299 Returns:300 `tuple(torch.Tensor)` comprising of the query and key tensors rotated using the Rotary Position Embedding.301 """302 cos = cos.unsqueeze(unsqueeze_dim)303 sin = sin.unsqueeze(unsqueeze_dim)304 q_embed = (q * cos) + (rotate_half(q) * sin)305 k_embed = (k * cos) + (rotate_half(k) * sin)306 return q_embed, k_embed307 308 309def repeat_kv(hidden_states: torch.Tensor, n_rep: int) -> torch.Tensor:310 """311 This is the equivalent of torch.repeat_interleave(x, dim=1, repeats=n_rep). The hidden states go from (batch,312 num_key_value_heads, seqlen, head_dim) to (batch, num_attention_heads, seqlen, head_dim)313 """314 batch, num_key_value_heads, slen, head_dim = hidden_states.shape315 if n_rep == 1:316 return hidden_states317 hidden_states = hidden_states[:, :, None, :, :].expand(batch, num_key_value_heads, n_rep, slen, head_dim)318 return hidden_states.reshape(batch, num_key_value_heads * n_rep, slen, head_dim)319 320 321class InstellaAttention(nn.Module):322 """Multi-headed attention from 'Attention Is All You Need' paper"""323 324 # copied from transformers.models.llama.modeling_llama.LlamaAttention.__init__ with Llama->Instella325 # TODO(joao): add me back asap :)326 def __init__(self, config: InstellaConfig, layer_idx: Optional[int] = None):327 super().__init__()328 self.config = config329 self.layer_idx = layer_idx330 if layer_idx is None:331 logger.warning_once(332 f"Instantiating {self.__class__.__name__} without passing a `layer_idx` is not recommended and will "333 "lead to errors during the forward call if caching is used. Please make sure to provide a `layer_idx` "334 "when creating this class."335 )336 337 self.attention_dropout = config.attention_dropout338 self.hidden_size = config.hidden_size339 self.num_heads = config.num_attention_heads340 self.head_dim = self.hidden_size // self.num_heads341 self.num_key_value_heads = config.num_key_value_heads342 self.num_key_value_groups = self.num_heads // self.num_key_value_heads343 self.max_position_embeddings = config.max_position_embeddings344 self.rope_theta = config.rope_theta345 self.is_causal = True346 347 if (self.head_dim * self.num_heads) != self.hidden_size:348 raise ValueError(349 f"hidden_size must be divisible by num_heads (got `hidden_size`: {self.hidden_size}"350 f" and `num_heads`: {self.num_heads})."351 )352 353 self.q_proj = nn.Linear(self.hidden_size, self.num_heads * self.head_dim, bias=config.attention_bias)354 self.k_proj = nn.Linear(self.hidden_size, self.num_key_value_heads * self.head_dim, bias=config.attention_bias)355 self.v_proj = nn.Linear(self.hidden_size, self.num_key_value_heads * self.head_dim, bias=config.attention_bias)356 self.o_proj = nn.Linear(self.hidden_size, self.hidden_size, bias=config.attention_bias)357 self._init_rope()358 self.q_norm = InstellaRMSNorm(self.num_heads * self.head_dim, config.rms_norm_eps)359 self.k_norm = InstellaRMSNorm(self.num_key_value_heads * self.head_dim, config.rms_norm_eps)360 361 def _init_rope(self):362 if self.config.rope_scaling is None:363 self.rotary_emb = InstellaRotaryEmbedding(364 self.head_dim,365 max_position_embeddings=self.max_position_embeddings,366 base=self.rope_theta,367 )368 else:369 scaling_type = self.config.rope_scaling["type"]370 scaling_factor = self.config.rope_scaling["factor"]371 if scaling_type == "linear":372 self.rotary_emb = InstellaLinearScalingRotaryEmbedding(373 self.head_dim,374 max_position_embeddings=self.max_position_embeddings,375 scaling_factor=scaling_factor,376 base=self.rope_theta,377 )378 elif scaling_type == "dynamic":379 self.rotary_emb = InstellaDynamicNTKScalingRotaryEmbedding(380 self.head_dim,381 max_position_embeddings=self.max_position_embeddings,382 scaling_factor=scaling_factor,383 base=self.rope_theta,384 )385 else:386 raise ValueError(f"Unknown RoPE scaling type {scaling_type}")387 388 def forward(389 self,390 hidden_states: torch.Tensor,391 attention_mask: Optional[torch.Tensor] = None,392 position_ids: Optional[torch.LongTensor] = None,393 past_key_value: Optional[Cache] = None,394 output_attentions: bool = False,395 use_cache: bool = False,396 cache_position: Optional[torch.LongTensor] = None,397 **kwargs,398 ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]:399 bsz, q_len, _ = hidden_states.size()400 401 query_states = self.q_norm(self.q_proj(hidden_states))402 key_states = self.k_norm(self.k_proj(hidden_states))403 value_states = self.v_proj(hidden_states)404 405 query_states = query_states.view(bsz, q_len, self.num_heads, self.head_dim).transpose(1, 2)406 key_states = key_states.view(bsz, q_len, self.num_key_value_heads, self.head_dim).transpose(1, 2)407 value_states = value_states.view(bsz, q_len, self.num_key_value_heads, self.head_dim).transpose(1, 2)408 409 cos, sin = self.rotary_emb(value_states, position_ids)410 query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin)411 412 if past_key_value is not None:413 # sin and cos are specific to RoPE models; cache_position needed for the static cache414 cache_kwargs = {"sin": sin, "cos": cos, "cache_position": cache_position}415 key_states, value_states = past_key_value.update(key_states, value_states, self.layer_idx, cache_kwargs)416 417 key_states = repeat_kv(key_states, self.num_key_value_groups)418 value_states = repeat_kv(value_states, self.num_key_value_groups)419 420 attn_weights = torch.matmul(query_states, key_states.transpose(2, 3)) / math.sqrt(self.head_dim)421 422 if attention_mask is not None: # no matter the length, we just slice it423 causal_mask = attention_mask[:, :, :, : key_states.shape[-2]]424 attn_weights = attn_weights + causal_mask425 426 # upcast attention to fp32427 attn_weights = nn.functional.softmax(attn_weights, dim=-1, dtype=torch.float32).to(query_states.dtype)428 attn_weights = nn.functional.dropout(attn_weights, p=self.attention_dropout, training=self.training)429 attn_output = torch.matmul(attn_weights, value_states)430 431 if attn_output.size() != (bsz, self.num_heads, q_len, self.head_dim):432 raise ValueError(433 f"`attn_output` should be of size {(bsz, self.num_heads, q_len, self.head_dim)}, but is"434 f" {attn_output.size()}"435 )436 437 attn_output = attn_output.transpose(1, 2).contiguous()438 439 attn_output = attn_output.reshape(bsz, q_len, self.hidden_size)440 441 attn_output = self.o_proj(attn_output)442 443 if not output_attentions:444 attn_weights = None445 446 return attn_output, attn_weights, past_key_value447 448 449class InstellaFlashAttention2(InstellaAttention):450 """451 Instella flash attention module. This module inherits from `InstellaAttention` as the weights of the module stays452 untouched. The only required change would be on the forward pass where it needs to correctly call the public API of453 flash attention and deal with padding tokens in case the input contains any of them.454 455 Instella flash attention module. This module inherits from `InstellaAttention` as the weights of the module stays456 untouched. The only required change would be on the forward pass where it needs to correctly call the public API of457 flash attention and deal with padding tokens in case the input contains any of them.458 """459 460 def __init__(self, *args, **kwargs):461 super().__init__(*args, **kwargs)462 463 # TODO: Should be removed once Flash Attention for RoCm is bumped to 2.1.464 # flash_attn<2.1 generates top-left aligned causal mask, while what is needed here is bottom-right alignement, 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.465 # 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).466 self._flash_attn_uses_top_left_mask = not is_flash_attn_greater_or_equal_2_10()467 468 def forward(469 self,470 hidden_states: torch.Tensor,471 attention_mask: Optional[torch.LongTensor] = None,472 position_ids: Optional[torch.LongTensor] = None,473 past_key_value: Optional[Cache] = None,474 output_attentions: bool = False,475 use_cache: bool = False,476 cache_position: Optional[torch.LongTensor] = None,477 **kwargs,478 ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]:479 output_attentions = False480 481 bsz, q_len, _ = hidden_states.size()482 483 query_states = self.q_norm(self.q_proj(hidden_states))484 key_states = self.k_norm(self.k_proj(hidden_states))485 value_states = self.v_proj(hidden_states)486 487 # Flash attention requires the input to have the shape488 # batch_size x seq_length x head_dim x hidden_dim489 # therefore we just need to keep the original shape490 query_states = query_states.view(bsz, q_len, self.num_heads, self.head_dim).transpose(1, 2)491 key_states = key_states.view(bsz, q_len, self.num_key_value_heads, self.head_dim).transpose(1, 2)492 value_states = value_states.view(bsz, q_len, self.num_key_value_heads, self.head_dim).transpose(1, 2)493 494 cos, sin = self.rotary_emb(value_states, position_ids)495 query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin)496 497 if past_key_value is not None:498 # sin and cos are specific to RoPE models; cache_position needed for the static cache499 cache_kwargs = {"sin": sin, "cos": cos, "cache_position": cache_position}500 key_states, value_states = past_key_value.update(key_states, value_states, self.layer_idx, cache_kwargs)501 502 # TODO: These transpose are quite inefficient but Flash Attention requires the layout [batch_size, sequence_length, num_heads, head_dim]. We would need to refactor the KV cache503 # to be able to avoid many of these transpose/reshape/view.504 query_states = query_states.transpose(1, 2)505 key_states = key_states.transpose(1, 2)506 value_states = value_states.transpose(1, 2)507 508 dropout_rate = self.attention_dropout if self.training else 0.0509 510 # In PEFT, usually we cast the layer norms in float32 for training stability reasons511 # therefore the input hidden states gets silently casted in float32. Hence, we need512 # cast them back in the correct dtype just to be sure everything works as expected.513 # This might slowdown training & inference so it is recommended to not cast the LayerNorms514 # in fp32. (InstellaRMSNorm handles it correctly)515 516 input_dtype = query_states.dtype517 if input_dtype == torch.float32:518 if torch.is_autocast_enabled():519 target_dtype = torch.get_autocast_gpu_dtype()520 # Handle the case where the model is quantized521 elif hasattr(self.config, "_pre_quantization_dtype"):522 target_dtype = self.config._pre_quantization_dtype523 else:524 target_dtype = self.q_proj.weight.dtype525 526 logger.warning_once(527 f"The input hidden states seems to be silently casted in float32, this might be related to"528 f" the fact you have upcasted embedding or layer norm layers in float32. We will cast back the input in"529 f" {target_dtype}."530 )531 532 query_states = query_states.to(target_dtype)533 key_states = key_states.to(target_dtype)534 value_states = value_states.to(target_dtype)535 536 attn_output = _flash_attention_forward(537 query_states,538 key_states,539 value_states,540 attention_mask,541 q_len,542 position_ids=position_ids,543 dropout=dropout_rate,544 use_top_left_mask=self._flash_attn_uses_top_left_mask,545 is_causal=self.is_causal,546 )547 548 attn_output = attn_output.reshape(bsz, q_len, self.hidden_size).contiguous()549 attn_output = self.o_proj(attn_output)550 551 if not output_attentions:552 attn_weights = None553 554 return attn_output, attn_weights, past_key_value555 556 557class InstellaSdpaAttention(InstellaAttention):558 """559 Instella attention module using torch.nn.functional.scaled_dot_product_attention. This module inherits from560 `InstellaAttention` as the weights of the module stays untouched. The only changes are on the forward pass to adapt to561 SDPA API.562 """563 564 # Adapted from InstellaAttention.forward565 def forward(566 self,567 hidden_states: torch.Tensor,568 attention_mask: Optional[torch.Tensor] = None,569 position_ids: Optional[torch.LongTensor] = None,570 past_key_value: Optional[Cache] = None,571 output_attentions: bool = False,572 use_cache: bool = False,573 cache_position: Optional[torch.LongTensor] = None,574 ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]:575 if output_attentions:576 # TODO: Improve this warning with e.g. `model.config.attn_implementation = "manual"` once this is implemented.577 logger.warning_once(578 "InstellaModel is using InstellaSdpaAttention, but `torch.nn.functional.scaled_dot_product_attention` does not support `output_attentions=True`. Falling back to the manual attention implementation, "579 '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.'580 )581 return super().forward(582 hidden_states=hidden_states,583 attention_mask=attention_mask,584 position_ids=position_ids,585 past_key_value=past_key_value,586 output_attentions=output_attentions,587 use_cache=use_cache,588 cache_position=cache_position,589 )590 bsz, q_len, _ = hidden_states.size()591 query_states = self.q_norm(self.q_proj(hidden_states))592 key_states = self.k_norm(self.k_proj(hidden_states))593 value_states = self.v_proj(hidden_states)594 595 query_states = query_states.view(bsz, q_len, self.num_heads, self.head_dim).transpose(1, 2)596 key_states = key_states.view(bsz, q_len, self.num_key_value_heads, self.head_dim).transpose(1, 2)597 value_states = value_states.view(bsz, q_len, self.num_key_value_heads, self.head_dim).transpose(1, 2)598 cos, sin = self.rotary_emb(value_states, position_ids)599 query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin)600 if past_key_value is not None:601 # sin and cos are specific to RoPE models; cache_position needed for the static cache602 cache_kwargs = {"sin": sin, "cos": cos, "cache_position": cache_position}603 key_states, value_states = past_key_value.update(key_states, value_states, self.layer_idx, cache_kwargs)604 key_states = repeat_kv(key_states, self.num_key_value_groups)605 value_states = repeat_kv(value_states, self.num_key_value_groups)606 causal_mask = attention_mask607 # if attention_mask is not None and cache_position is not None:608 if attention_mask is not None:609 causal_mask = causal_mask[:, :, :, : key_states.shape[-2]]610 # SDPA with memory-efficient backend is currently (torch==2.1.2) bugged with non-contiguous inputs with custom attn_mask,611 # Reference: https://github.com/pytorch/pytorch/issues/112577.612 if query_states.device.type == "cuda" and causal_mask is not None:613 query_states = query_states.contiguous()614 key_states = key_states.contiguous()615 value_states = value_states.contiguous()616 # We dispatch to SDPA's Flash Attention or Efficient kernels via this `is_causal` if statement instead of an inline conditional assignment617 # in SDPA to support both torch.compile's dynamic shapes and full graph options. An inline conditional prevents dynamic shapes from compiling.618 is_causal = True if causal_mask is None and q_len > 1 else False619 attn_output = torch.nn.functional.scaled_dot_product_attention(620 query_states,621 key_states,622 value_states,623 attn_mask=causal_mask,624 dropout_p=self.attention_dropout if self.training else 0.0,625 is_causal=is_causal,626 )627 attn_output = attn_output.transpose(1, 2).contiguous()628 attn_output = attn_output.view(bsz, q_len, self.hidden_size)629 attn_output = self.o_proj(attn_output)630 return attn_output, None, past_key_value631 632 633class InstellaMLP(nn.Module):634 def __init__(self, config):635 super().__init__()636 self.config = config637 self.hidden_size = config.hidden_size638 self.intermediate_size = config.intermediate_size639 self.gate_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=False)640 self.up_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=False)641 self.down_proj = nn.Linear(self.intermediate_size, self.hidden_size, bias=False)642 self.act_fn = ACT2FN[config.hidden_act]643 644 def forward(self, x):645 return self.down_proj(self.act_fn(self.gate_proj(x)) * self.up_proj(x))646 647 648Instella_ATTENTION_CLASSES = {649 "eager": InstellaAttention,650 "flash_attention_2": InstellaFlashAttention2,651 "sdpa": InstellaSdpaAttention,652}653 654 655class InstellaDecoderLayer(nn.Module):656 def __init__(self, config: InstellaConfig, layer_idx: int):657 super().__init__()658 self.hidden_size = config.hidden_size659 660 self.self_attn = Instella_ATTENTION_CLASSES[config._attn_implementation](config=config, layer_idx=layer_idx)661 662 self.mlp = InstellaMLP(config)663 self.pre_attention_layernorm = InstellaRMSNorm(config.hidden_size, eps=config.rms_norm_eps)664 self.pre_feedforward_layernorm = InstellaRMSNorm(config.hidden_size, eps=config.rms_norm_eps)665 666 # copied from transformers.models.llama.modeling_llama.LlamaDecoderLayer.forward667 # TODO(joao): add me back asap :)668 def forward(669 self,670 hidden_states: torch.Tensor,671 attention_mask: Optional[torch.Tensor] = None,672 position_ids: Optional[torch.LongTensor] = None,673 past_key_value: Optional[Cache] = None,674 output_attentions: Optional[bool] = False,675 use_cache: Optional[bool] = False,676 cache_position: Optional[torch.LongTensor] = None,677 **kwargs,678 ) -> Tuple[torch.FloatTensor, Optional[Tuple[torch.FloatTensor, torch.FloatTensor]]]:679 """680 Args:681 hidden_states (`torch.FloatTensor`): input to the layer of shape `(batch, seq_len, embed_dim)`682 attention_mask (`torch.FloatTensor`, *optional*):683 attention mask of size `(batch_size, sequence_length)` if flash attention is used or `(batch_size, 1,684 query_sequence_length, key_sequence_length)` if default attention is used.685 output_attentions (`bool`, *optional*):686 Whether or not to return the attentions tensors of all attention layers. See `attentions` under687 returned tensors for more detail.688 use_cache (`bool`, *optional*):689 If set to `True`, `past_key_values` key value states are returned and can be used to speed up decoding690 (see `past_key_values`).691 past_key_value (`Tuple(torch.FloatTensor)`, *optional*): cached past key and value projection states692 cache_position (`torch.LongTensor` of shape `(sequence_length)`, *optional*):693 Indices depicting the position of the input sequence tokens in the sequence694 kwargs (`dict`, *optional*):695 Arbitrary kwargs to be ignored, used for FSDP and other methods that injects code696 into the model697 """698 residual = hidden_states699 700 # Self Attention701 hidden_states = self.pre_attention_layernorm(hidden_states)702 hidden_states, self_attn_weights, present_key_value = self.self_attn(703 hidden_states=hidden_states,704 attention_mask=attention_mask,705 position_ids=position_ids,706 past_key_value=past_key_value,707 output_attentions=output_attentions,708 use_cache=use_cache,709 cache_position=cache_position,710 **kwargs,711 )712 # hidden_states = self.post_attention_layernorm(hidden_states)713 hidden_states = residual + hidden_states714 # print(hidden_states)715 716 # Fully Connected717 residual = hidden_states718 hidden_states = self.pre_feedforward_layernorm(hidden_states)719 hidden_states = self.mlp(hidden_states)720 # hidden_states = self.post_feedforward_layernorm(hidden_states)721 hidden_states = residual + hidden_states722 # print(hidden_states)723 724 outputs = (hidden_states,)725 if output_attentions:726 outputs += (self_attn_weights,)727 if use_cache:728 outputs += (present_key_value,)729 return outputs730 731 732Instella_START_DOCSTRING = r"""733 This model inherits from [`PreTrainedModel`]. Check the superclass documentation for the generic methods the734 library implements for all its model (such as downloading or saving, resizing the input embeddings, pruning heads735 etc.)736 737 This model is also a PyTorch [torch.nn.Module](https://pytorch.org/docs/stable/nn.html#torch.nn.Module) subclass.738 Use it as a regular PyTorch Module and refer to the PyTorch documentation for all matter related to general usage739 and behavior.740 741 Parameters:742 config ([`InstellaConfig`]):743 Model configuration class with all the parameters of the model. Initializing with a config file does not744 load the weights associated with the model, only the configuration. Check out the745 [`~PreTrainedModel.from_pretrained`] method to load the model weights.746"""747 748 749@add_start_docstrings(750 "The bare Instella Model outputting raw hidden-states without any specific head on top.",751 Instella_START_DOCSTRING,752)753class InstellaPreTrainedModel(PreTrainedModel):754 config_class = InstellaConfig755 base_model_prefix = "model"756 supports_gradient_checkpointing = True757 _no_split_modules = ["InstellaDecoderLayer"]758 _skip_keys_device_placement = ["past_key_values"]759 _supports_flash_attn_2 = True760 _supports_sdpa = True761 _supports_cache_class = True762 _supports_quantized_cache = True763 _supports_static_cache = True764 765 def _init_weights(self, module):766 std = self.config.initializer_range767 if isinstance(module, nn.Linear):768 module.weight.data.normal_(mean=0.0, std=std)769 if module.bias is not None:770 module.bias.data.zero_()771 elif isinstance(module, nn.Embedding):772 module.weight.data.normal_(mean=0.0, std=std)773 if module.padding_idx is not None:774 module.weight.data[module.padding_idx].zero_()775 776 777Instella_INPUTS_DOCSTRING = r"""778 Args:779 input_ids (`torch.LongTensor` of shape `(batch_size, sequence_length)`):780 Indices of input sequence tokens in the vocabulary. Padding will be ignored by default should you provide781 it.782 783 Indices can be obtained using [`AutoTokenizer`]. See [`PreTrainedTokenizer.encode`] and784 [`PreTrainedTokenizer.__call__`] for details.785 786 [What are input IDs?](../glossary#input-ids)787 attention_mask (`torch.Tensor` of shape `(batch_size, sequence_length)`, *optional*):788 Mask to avoid performing attention on padding token indices. Mask values selected in `[0, 1]`:789 790 - 1 for tokens that are **not masked**,791 - 0 for tokens that are **masked**.792 793 [What are attention masks?](../glossary#attention-mask)794 795 Indices can be obtained using [`AutoTokenizer`]. See [`PreTrainedTokenizer.encode`] and796 [`PreTrainedTokenizer.__call__`] for details.797 798 If `past_key_values` is used, optionally only the last `input_ids` have to be input (see799 `past_key_values`).800 801 If you want to change padding behavior, you should read [`modeling_opt._prepare_decoder_attention_mask`]802 and modify to your needs. See diagram 1 in [the paper](https://arxiv.org/abs/1910.13461) for more803 information on the default strategy.804 805 - 1 indicates the head is **not masked**,806 - 0 indicates the head is **masked**.807 position_ids (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*):808 Indices of positions of each input sequence tokens in the position embeddings. Selected in the range `[0,809 config.n_positions - 1]`.810 811 [What are position IDs?](../glossary#position-ids)812 past_key_values (`Cache` or `tuple(tuple(torch.FloatTensor))`, *optional*):813 Pre-computed hidden-states (key and values in the self-attention blocks and in the cross-attention814 blocks) that can be used to speed up sequential decoding. This typically consists in the `past_key_values`815 returned by the model at a previous stage of decoding, when `use_cache=True` or `config.use_cache=True`.816 817 Two formats are allowed:818 - a [`~cache_utils.Cache`] instance, see our819 [kv cache guide](https://huggingface.co/docs/transformers/en/kv_cache);820 - Tuple of `tuple(torch.FloatTensor)` of length `config.n_layers`, with each tuple having 2 tensors of821 shape `(batch_size, num_heads, sequence_length, embed_size_per_head)`). This is also known as the legacy822 cache format.823 824 The model will output the same cache format that is fed as input. If no `past_key_values` are passed, the825 legacy cache format will be returned.826 827 If `past_key_values` are used, the user can optionally input only the last `input_ids` (those that don't828 have their past key value states given to this model) of shape `(batch_size, 1)` instead of all `input_ids`829 of shape `(batch_size, sequence_length)`.830 inputs_embeds (`torch.FloatTensor` of shape `(batch_size, sequence_length, hidden_size)`, *optional*):831 Optionally, instead of passing `input_ids` you can choose to directly pass an embedded representation. This832 is useful if you want more control over how to convert `input_ids` indices into associated vectors than the833 model's internal embedding lookup matrix.834 use_cache (`bool`, *optional*):835 If set to `True`, `past_key_values` key value states are returned and can be used to speed up decoding (see836 `past_key_values`).837 output_attentions (`bool`, *optional*):838 Whether or not to return the attentions tensors of all attention layers. See `attentions` under returned839 tensors for more detail.840 output_hidden_states (`bool`, *optional*):841 Whether or not to return the hidden states of all layers. See `hidden_states` under returned tensors for842 more detail.843 return_dict (`bool`, *optional*):844 Whether or not to return a [`~utils.ModelOutput`] instead of a plain tuple.845 cache_position (`torch.LongTensor` of shape `(sequence_length)`, *optional*):846 Indices depicting the position of the input sequence tokens in the sequence. Contrarily to `position_ids`,847 this tensor is not affected by padding. It is used to update the cache in the correct position and to infer848 the complete sequence length.849"""850 851 852@add_start_docstrings(853 "The bare Instella Model outputting raw hidden-states without any specific head on top.",854 Instella_START_DOCSTRING,855)856class InstellaModel(InstellaPreTrainedModel):857 """858 Transformer decoder consisting of *config.num_hidden_layers* layers. Each layer is a [`InstellaDecoderLayer`]859 860 Args:861 config: InstellaConfig862 """863 864 def __init__(self, config: InstellaConfig):865 super().__init__(config)866 self.padding_idx = config.pad_token_id867 self.vocab_size = config.vocab_size868 869 self.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size, self.padding_idx)870 self.layers = nn.ModuleList(871 [InstellaDecoderLayer(config, layer_idx) for layer_idx in range(config.num_hidden_layers)]872 )873 # self.layers = self.layers[:5]874 self.norm = InstellaRMSNorm(config.hidden_size, eps=config.rms_norm_eps)875 self.gradient_checkpointing = False876 877 # Initialize weights and apply final processing878 self.post_init()879 880 def get_input_embeddings(self):881 return self.embed_tokens882 883 def set_input_embeddings(self, value):884 self.embed_tokens = value885 886 @add_start_docstrings_to_model_forward(Instella_INPUTS_DOCSTRING)887 # copied from transformers.models.llama.modeling_llama.LlamaModel.forward888 # TODO(joao): add me back asap :)889 def forward(890 self,891 input_ids: torch.LongTensor = None,892 attention_mask: Optional[torch.Tensor] = None,893 position_ids: Optional[torch.LongTensor] = None,894 past_key_values: Optional[Union[Cache, List[torch.FloatTensor]]] = None,895 inputs_embeds: Optional[torch.FloatTensor] = None,896 use_cache: Optional[bool] = None,897 output_attentions: Optional[bool] = None,898 output_hidden_states: Optional[bool] = None,899 return_dict: Optional[bool] = None,900 cache_position: Optional[torch.LongTensor] = None,901 ) -> Union[Tuple, BaseModelOutputWithPast]:902 output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions903 output_hidden_states = (904 output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states905 )906 use_cache = use_cache if use_cache is not None else self.config.use_cache907 return_dict = return_dict if return_dict is not None else self.config.use_return_dict908 909 if (input_ids is None) ^ (inputs_embeds is not None):910 raise ValueError("You must specify exactly one of input_ids or inputs_embeds")911 912 if self.gradient_checkpointing and self.training and use_cache:913 logger.warning_once(914 "`use_cache=True` is incompatible with gradient checkpointing. Setting `use_cache=False`."915 )916 use_cache = False917 918 if inputs_embeds is None:919 inputs_embeds = self.embed_tokens(input_ids)920 # print(inputs_embeds)921 922 # kept for BC (non `Cache` `past_key_values` inputs)923 return_legacy_cache = False924 if use_cache and not isinstance(past_key_values, Cache):925 return_legacy_cache = True926 if past_key_values is None:927 past_key_values = DynamicCache()928 else:929 past_key_values = DynamicCache.from_legacy_cache(past_key_values)930 logger.warning_once(931 "We detected that you are passing `past_key_values` as a tuple of tuples. This is deprecated and "932 "will be removed in v4.47. Please convert your cache or use an appropriate `Cache` class "933 "(https://huggingface.co/docs/transformers/kv_cache#legacy-cache-format)"934 )935 936 if cache_position is None:937 past_seen_tokens = past_key_values.get_seq_length() if past_key_values is not None else 0938 cache_position = torch.arange(939 past_seen_tokens, past_seen_tokens + inputs_embeds.shape[1], device=inputs_embeds.device940 )941 if position_ids is None:942 position_ids = cache_position.unsqueeze(0)943 944 causal_mask = self._update_causal_mask(945 attention_mask, inputs_embeds, cache_position, past_key_values, output_attentions946 )947 948 # embed positions949 hidden_states = inputs_embeds950 951 # decoder layers952 all_hidden_states = () if output_hidden_states else None953 all_self_attns = () if output_attentions else None954 next_decoder_cache = None955 956 for decoder_layer in self.layers:957 if output_hidden_states:958 all_hidden_states += (hidden_states,)959 960 if self.gradient_checkpointing and self.training:961 layer_outputs = self._gradient_checkpointing_func(962 decoder_layer.__call__,963 hidden_states,964 causal_mask,965 position_ids,966 past_key_values,967 output_attentions,968 use_cache,969 cache_position,970 )971 else:972 layer_outputs = decoder_layer(973 hidden_states,974 attention_mask=causal_mask,975 position_ids=position_ids,976 past_key_value=past_key_values,977 output_attentions=output_attentions,978 use_cache=use_cache,979 cache_position=cache_position,980 )981 982 hidden_states = layer_outputs[0]983 984 if use_cache:985 next_decoder_cache = layer_outputs[2 if output_attentions else 1]986 987 if output_attentions:988 all_self_attns += (layer_outputs[1],)989 990 hidden_states = self.norm(hidden_states)991 # print(hidden_states)992 993 # add hidden states from the last decoder layer994 if output_hidden_states:995 all_hidden_states += (hidden_states,)996 997 next_cache = next_decoder_cache if use_cache else None998 if return_legacy_cache:999 next_cache = next_cache.to_legacy_cache()1000 1001 if not return_dict:1002 return tuple(v for v in [hidden_states, next_cache, all_hidden_states, all_self_attns] if v is not None)1003 return BaseModelOutputWithPast(1004 last_hidden_state=hidden_states,1005 past_key_values=next_cache,1006 hidden_states=all_hidden_states,1007 attentions=all_self_attns,1008 )1009 1010 def _update_causal_mask(1011 self,1012 attention_mask: torch.Tensor,1013 input_tensor: torch.Tensor,1014 cache_position: torch.Tensor,1015 past_key_values: Cache,1016 output_attentions: bool,1017 ):1018 if self.config._attn_implementation == "flash_attention_2":1019 if attention_mask is not None and 0.0 in attention_mask:1020 return attention_mask1021 return None1022 1023 # For SDPA, when possible, we will rely on its `is_causal` argument instead of its `attn_mask` argument, in1024 # order to dispatch on Flash Attention 2. This feature is not compatible with static cache, as SDPA will fail1025 # to infer the attention mask.1026 past_seen_tokens = past_key_values.get_seq_length() if past_key_values is not None else 01027 using_static_cache = isinstance(past_key_values, StaticCache)1028 1029 # When output attentions is True, sdpa implementation's forward method calls the eager implementation's forward1030 if self.config._attn_implementation == "sdpa" and not using_static_cache and not output_attentions:1031 if AttentionMaskConverter._ignore_causal_mask_sdpa(1032 attention_mask,1033 inputs_embeds=input_tensor,1034 past_key_values_length=past_seen_tokens,1035 is_training=self.training,1036 ):1037 return None1038 1039 dtype, device = input_tensor.dtype, input_tensor.device1040 sequence_length = input_tensor.shape[1]1041 if using_static_cache:1042 target_length = past_key_values.get_max_cache_shape()1043 else:1044 target_length = (1045 attention_mask.shape[-1]1046 if isinstance(attention_mask, torch.Tensor)1047 else past_seen_tokens + sequence_length + 11048 )1049 1050 # In case the provided `attention` mask is 2D, we generate a causal mask here (4D).1051 causal_mask = self._prepare_4d_causal_attention_mask_with_cache_position(1052 attention_mask,1053 sequence_length=sequence_length,1054 target_length=target_length,1055 dtype=dtype,1056 device=device,1057 cache_position=cache_position,1058 batch_size=input_tensor.shape[0],1059 )1060 1061 if (1062 self.config._attn_implementation == "sdpa"1063 and attention_mask is not None1064 and attention_mask.device.type == "cuda"1065 and not output_attentions1066 ):1067 # Attend to all tokens in fully masked rows in the causal_mask, for example the relevant first rows when1068 # using left padding. This is required by F.scaled_dot_product_attention memory-efficient attention path.1069 # Details: https://github.com/pytorch/pytorch/issues/1102131070 min_dtype = torch.finfo(dtype).min1071 causal_mask = AttentionMaskConverter._unmask_unattended(causal_mask, min_dtype)1072 1073 return causal_mask1074 1075 @staticmethod1076 def _prepare_4d_causal_attention_mask_with_cache_position(1077 attention_mask: torch.Tensor,1078 sequence_length: int,1079 target_length: int,1080 dtype: torch.dtype,1081 device: torch.device,1082 cache_position: torch.Tensor,1083 batch_size: int,1084 **kwargs,1085 ):1086 """1087 Creates a causal 4D mask of shape `(batch_size, 1, query_length, key_value_length)` from a 2D mask of shape1088 `(batch_size, key_value_length)`, or if the input `attention_mask` is already 4D, do nothing.1089 1090 Args:1091 attention_mask (`torch.Tensor`):1092 A 2D attention mask of shape `(batch_size, key_value_length)` or a 4D attention mask of shape1093 `(batch_size, 1, query_length, key_value_length)`.1094 sequence_length (`int`):1095 The sequence length being processed.1096 target_length (`int`):1097 The target length: when generating with static cache, the mask should be as long as the static cache,1098 to account for the 0 padding, the part of the cache that is not filled yet.1099 dtype (`torch.dtype`):1100 The dtype to use for the 4D attention mask.1101 device (`torch.device`):1102 The device to plcae the 4D attention mask on.1103 cache_position (`torch.Tensor`):1104 Indices depicting the position of the input sequence tokens in the sequence.1105 batch_size (`torch.Tensor`):1106 Batch size.1107 """1108 if attention_mask is not None and attention_mask.dim() == 4:1109 # In this case we assume that the mask comes already in inverted form and requires no inversion or slicing.1110 causal_mask = attention_mask1111 else:1112 min_dtype = torch.finfo(dtype).min1113 causal_mask = torch.full(1114 (sequence_length, target_length), fill_value=min_dtype, dtype=dtype, device=device1115 )1116 if sequence_length != 1:1117 causal_mask = torch.triu(causal_mask, diagonal=1)1118 causal_mask *= torch.arange(target_length, device=device) > cache_position.reshape(-1, 1)1119 causal_mask = causal_mask[None, None, :, :].expand(batch_size, 1, -1, -1)1120 if attention_mask is not None:1121 causal_mask = causal_mask.clone() # copy to contiguous memory for in-place edit1122 mask_length = attention_mask.shape[-1]1123 padding_mask = causal_mask[:, :, :, :mask_length] + attention_mask[:, None, None, :]1124 padding_mask = padding_mask == 01125 causal_mask[:, :, :, :mask_length] = causal_mask[:, :, :, :mask_length].masked_fill(1126 padding_mask, min_dtype1127 )1128 1129 return causal_mask1130 1131# TODO: re-enable check: Copied from transformers.models.llama.modeling_llama.LlamaForCausalLM with LLAMA->Instella,Llama->Instella1132class InstellaForCausalLM(InstellaPreTrainedModel, GenerationMixin):1133 _tied_weights_keys = ["lm_head.weight"]1134 1135 def __init__(self, config: InstellaConfig):1136 super().__init__(config)1137 self.model = InstellaModel(config)1138 self.vocab_size = config.vocab_size1139 self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)1140 1141 # Initialize weights and apply final processing1142 self.post_init()1143 1144 def get_input_embeddings(self):1145 return self.model.embed_tokens1146 1147 def set_input_embeddings(self, value):1148 self.model.embed_tokens = value1149 1150 def get_output_embeddings(self):1151 return self.lm_head1152 1153 def set_output_embeddings(self, new_embeddings):1154 self.lm_head = new_embeddings1155 1156 def set_decoder(self, decoder):1157 self.model = decoder1158 1159 def get_decoder(self):1160 return self.model1161 1162 @add_start_docstrings_to_model_forward(Instella_INPUTS_DOCSTRING)1163 @replace_return_docstrings(output_type=CausalLMOutputWithPast, config_class=_CONFIG_FOR_DOC)1164 # Ignore copy1165 def forward(1166 self,1167 input_ids: torch.LongTensor = None,1168 attention_mask: Optional[torch.Tensor] = None,1169 position_ids: Optional[torch.LongTensor] = None,1170 past_key_values: Optional[List[torch.FloatTensor]] = None,1171 inputs_embeds: Optional[torch.FloatTensor] = None,1172 labels: Optional[torch.LongTensor] = None,1173 use_cache: Optional[bool] = None,1174 output_attentions: Optional[bool] = None,1175 output_hidden_states: Optional[bool] = None,1176 return_dict: Optional[bool] = None,1177 cache_position: Optional[torch.LongTensor] = None,1178 num_logits_to_keep: int = 0,1179 **loss_kwargs,1180 ) -> Union[Tuple, CausalLMOutputWithPast]:1181 r"""1182 Args:1183 labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*):1184 Labels for computing the masked language modeling loss. Indices should either be in `[0, ...,1185 config.vocab_size]` or -100 (see `input_ids` docstring). Tokens with indices set to `-100` are ignored1186 (masked), the loss is only computed for the tokens with labels in `[0, ..., config.vocab_size]`.1187 1188 num_logits_to_keep (`int`, *optional*):1189 Calculate logits for the last `num_logits_to_keep` tokens. If `0`, calculate logits for all1190 `input_ids` (special case). Only last token logits are needed for generation, and calculating them only for that1191 token can save memory, which becomes pretty significant for long sequences or large vocabulary size.1192 1193 Returns:1194 1195 Example:1196 1197 ```python1198 >>> from transformers import AutoTokenizer, InstellaForCausalLM1199 1200 >>> model = InstellaForCausalLM.from_pretrained("allenai/Instella2-1B-hf")