CofeAI/Tele-FLM
3018k
1# coding=utf-82""" PyTorch Tele-FLM model, based on LLAMA implementation. """3 4import math5import warnings6from typing import List, Optional, Tuple, Union7 8import torch9import torch.nn.functional as F10import torch.utils.checkpoint11from torch import nn12from torch.nn import BCEWithLogitsLoss, CrossEntropyLoss, MSELoss13 14from transformers.activations import ACT2FN15from transformers.cache_utils import Cache, DynamicCache, StaticCache16from transformers.modeling_attn_mask_utils import AttentionMaskConverter17from transformers.modeling_outputs import (18 BaseModelOutputWithPast,19 CausalLMOutputWithPast,20 QuestionAnsweringModelOutput,21 SequenceClassifierOutputWithPast,22)23from transformers.modeling_utils import PreTrainedModel24from transformers.pytorch_utils import ALL_LAYERNORM_LAYERS25from transformers.utils import (26 add_start_docstrings,27 add_start_docstrings_to_model_forward,28 is_flash_attn_2_available,29 is_flash_attn_greater_or_equal_2_10,30 logging,31 replace_return_docstrings,32)33from .configuration_teleflm import TeleFLMConfig34 35if is_flash_attn_2_available():36 from flash_attn import flash_attn_func, flash_attn_varlen_func37 from flash_attn.bert_padding import index_first_axis, pad_input, unpad_input # noqa38 39 40logger = logging.get_logger(__name__)41 42_CONFIG_FOR_DOC = "TeleFLMConfig"43 44 45def _get_unpad_data(attention_mask):46 seqlens_in_batch = attention_mask.sum(dim=-1, dtype=torch.int32)47 indices = torch.nonzero(attention_mask.flatten(), as_tuple=False).flatten()48 max_seqlen_in_batch = seqlens_in_batch.max().item()49 cu_seqlens = F.pad(torch.cumsum(seqlens_in_batch, dim=0, dtype=torch.int32), (1, 0))50 return (51 indices,52 cu_seqlens,53 max_seqlen_in_batch,54 )55 56 57class TeleFLMRMSNorm(nn.Module):58 def __init__(self, hidden_size, eps=1e-6):59 """60 TeleFLMRMSNorm is equivalent to T5LayerNorm61 """62 super().__init__()63 self.weight = nn.Parameter(torch.ones(hidden_size))64 self.variance_epsilon = eps65 66 def forward(self, hidden_states):67 input_dtype = hidden_states.dtype68 hidden_states = hidden_states.to(torch.float32)69 variance = hidden_states.pow(2).mean(-1, keepdim=True)70 hidden_states = hidden_states * torch.rsqrt(variance + self.variance_epsilon)71 return self.weight * hidden_states.to(input_dtype)72 73 74ALL_LAYERNORM_LAYERS.append(TeleFLMRMSNorm)75 76 77class TeleFLMRotaryEmbedding(nn.Module):78 def __init__(self, dim, max_position_embeddings=4096, base=10000, device=None, scaling_factor=1.0):79 super().__init__()80 self.scaling_factor = scaling_factor81 self.dim = dim82 self.max_position_embeddings = max_position_embeddings83 self.base = base84 inv_freq = 1.0 / (self.base ** (torch.arange(0, self.dim, 2, dtype=torch.int64).float().to(device) / self.dim))85 self.register_buffer("inv_freq", inv_freq, persistent=False)86 # For BC we register cos and sin cached87 self.max_seq_len_cached = max_position_embeddings88 t = torch.arange(self.max_seq_len_cached, device=device, dtype=torch.int64).type_as(self.inv_freq)89 t = t / self.scaling_factor90 freqs = torch.outer(t, self.inv_freq)91 # Different from paper, but it uses a different permutation in order to obtain the same calculation92 emb = torch.cat((freqs, freqs), dim=-1)93 self.register_buffer("_cos_cached", emb.cos().to(torch.get_default_dtype()), persistent=False)94 self.register_buffer("_sin_cached", emb.sin().to(torch.get_default_dtype()), persistent=False)95 96 97 @torch.no_grad()98 def forward(self, x, position_ids):99 # x: [bs, num_attention_heads, seq_len, head_size]100 inv_freq_expanded = self.inv_freq[None, :, None].float().expand(position_ids.shape[0], -1, 1)101 position_ids_expanded = position_ids[:, None, :].float()102 # Force float32 since bfloat16 loses precision on long contexts103 # See https://github.com/huggingface/transformers/pull/29285104 device_type = x.device.type105 device_type = device_type if isinstance(device_type, str) and device_type != "mps" else "cpu"106 with torch.autocast(device_type=device_type, enabled=False):107 freqs = (inv_freq_expanded.float() @ position_ids_expanded.float()).transpose(1, 2)108 emb = torch.cat((freqs, freqs), dim=-1)109 cos = emb.cos()110 sin = emb.sin()111 return cos.to(dtype=x.dtype), sin.to(dtype=x.dtype)112 113 114class TeleFLMLinearScalingRotaryEmbedding(TeleFLMRotaryEmbedding):115 """TeleFLMRotaryEmbedding extended with linear scaling. Credits to the Reddit user /u/kaiokendev"""116 117 def forward(self, x, position_ids):118 # difference to the original RoPE: a scaling factor is aplied to the position ids119 position_ids = position_ids.float() / self.scaling_factor120 cos, sin = super().forward(x, position_ids)121 return cos, sin122 123 124class TeleFLMDynamicNTKScalingRotaryEmbedding(TeleFLMRotaryEmbedding):125 """TeleFLMRotaryEmbedding extended with Dynamic NTK scaling. Credits to the Reddit users /u/bloc97 and /u/emozilla"""126 127 def forward(self, x, position_ids):128 # difference to the original RoPE: inv_freq is recomputed when the sequence length > original length129 seq_len = torch.max(position_ids) + 1130 if seq_len > self.max_position_embeddings:131 base = self.base * (132 (self.scaling_factor * seq_len / self.max_position_embeddings) - (self.scaling_factor - 1)133 ) ** (self.dim / (self.dim - 2))134 inv_freq = 1.0 / (135 base ** (torch.arange(0, self.dim, 2, dtype=torch.int64).float().to(x.device) / self.dim)136 )137 self.register_buffer("inv_freq", inv_freq, persistent=False) # TODO joao: this may break with compilation138 139 cos, sin = super().forward(x, position_ids)140 return cos, sin141 142 143def rotate_half(x):144 """Rotates half the hidden dims of the input."""145 x1 = x[..., : x.shape[-1] // 2]146 x2 = x[..., x.shape[-1] // 2 :]147 return torch.cat((-x2, x1), dim=-1)148 149 150def apply_rotary_pos_emb(q, k, cos, sin, position_ids=None, unsqueeze_dim=1):151 """Applies Rotary Position Embedding to the query and key tensors.152 153 Args:154 q (`torch.Tensor`): The query tensor.155 k (`torch.Tensor`): The key tensor.156 cos (`torch.Tensor`): The cosine part of the rotary embedding.157 sin (`torch.Tensor`): The sine part of the rotary embedding.158 position_ids (`torch.Tensor`, *optional*):159 Deprecated and unused.160 unsqueeze_dim (`int`, *optional*, defaults to 1):161 The 'unsqueeze_dim' argument specifies the dimension along which to unsqueeze cos[position_ids] and162 sin[position_ids] so that they can be properly broadcasted to the dimensions of q and k. For example, note163 that cos[position_ids] and sin[position_ids] have the shape [batch_size, seq_len, head_dim]. Then, if q and164 k have the shape [batch_size, heads, seq_len, head_dim], then setting unsqueeze_dim=1 makes165 cos[position_ids] and sin[position_ids] broadcastable to the shapes of q and k. Similarly, if q and k have166 the shape [batch_size, seq_len, heads, head_dim], then set unsqueeze_dim=2.167 Returns:168 `tuple(torch.Tensor)` comprising of the query and key tensors rotated using the Rotary Position Embedding.169 """170 cos = cos.unsqueeze(unsqueeze_dim)171 sin = sin.unsqueeze(unsqueeze_dim)172 q_embed = (q * cos) + (rotate_half(q) * sin)173 k_embed = (k * cos) + (rotate_half(k) * sin)174 return q_embed, k_embed175 176 177class TeleFLMMLP(nn.Module):178 def __init__(self, config):179 super().__init__()180 self.config = config181 self.hidden_size = config.hidden_size182 self.intermediate_size = config.intermediate_size183 self.gate_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=False)184 self.up_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=False)185 self.down_proj = nn.Linear(self.intermediate_size, self.hidden_size, bias=False)186 self.act_fn = ACT2FN[config.hidden_act]187 188 def forward(self, x):189 if self.config.pretraining_tp > 1:190 slice = self.intermediate_size // self.config.pretraining_tp191 gate_proj_slices = self.gate_proj.weight.split(slice, dim=0)192 up_proj_slices = self.up_proj.weight.split(slice, dim=0)193 down_proj_slices = self.down_proj.weight.split(slice, dim=1)194 195 gate_proj = torch.cat(196 [F.linear(x, gate_proj_slices[i]) for i in range(self.config.pretraining_tp)], dim=-1197 )198 up_proj = torch.cat([F.linear(x, up_proj_slices[i]) for i in range(self.config.pretraining_tp)], dim=-1)199 200 intermediate_states = (self.act_fn(gate_proj) * up_proj).split(slice, dim=2)201 down_proj = [202 F.linear(intermediate_states[i], down_proj_slices[i]) for i in range(self.config.pretraining_tp)203 ]204 down_proj = sum(down_proj)205 else:206 down_proj = self.down_proj(self.act_fn(self.gate_proj(x)) * self.up_proj(x))207 208 return down_proj209 210 211def repeat_kv(hidden_states: torch.Tensor, n_rep: int) -> torch.Tensor:212 """213 This is the equivalent of torch.repeat_interleave(x, dim=1, repeats=n_rep). The hidden states go from (batch,214 num_key_value_heads, seqlen, head_dim) to (batch, num_attention_heads, seqlen, head_dim)215 """216 batch, num_key_value_heads, slen, head_dim = hidden_states.shape217 if n_rep == 1:218 return hidden_states219 hidden_states = hidden_states[:, :, None, :, :].expand(batch, num_key_value_heads, n_rep, slen, head_dim)220 return hidden_states.reshape(batch, num_key_value_heads * n_rep, slen, head_dim)221 222 223class TeleFLMAttention(nn.Module):224 """Multi-headed attention from 'Attention Is All You Need' paper"""225 226 def __init__(self, config: TeleFLMConfig, layer_idx: Optional[int] = None):227 super().__init__()228 self.config = config229 self.layer_idx = layer_idx230 if layer_idx is None:231 logger.warning_once(232 f"Instantiating {self.__class__.__name__} without passing a `layer_idx` is not recommended and will "233 "lead to errors during the forward call if caching is used. Please make sure to provide a `layer_idx` "234 "when creating this class."235 )236 237 self.attention_dropout = config.attention_dropout238 self.hidden_size = config.hidden_size239 self.num_heads = config.num_attention_heads240 self.head_dim = self.hidden_size // self.num_heads241 self.num_key_value_heads = config.num_key_value_heads242 self.num_key_value_groups = self.num_heads // self.num_key_value_heads243 self.max_position_embeddings = config.max_position_embeddings244 self.rope_theta = config.rope_theta245 self.is_causal = True246 247 if (self.head_dim * self.num_heads) != self.hidden_size:248 raise ValueError(249 f"hidden_size must be divisible by num_heads (got `hidden_size`: {self.hidden_size}"250 f" and `num_heads`: {self.num_heads})."251 )252 253 self.q_proj = nn.Linear(self.hidden_size, self.num_heads * self.head_dim, bias=config.attention_bias)254 self.k_proj = nn.Linear(self.hidden_size, self.num_key_value_heads * self.head_dim, bias=config.attention_bias)255 self.v_proj = nn.Linear(self.hidden_size, self.num_key_value_heads * self.head_dim, bias=config.attention_bias)256 self.o_proj = nn.Linear(self.hidden_size, self.hidden_size, bias=config.attention_bias)257 self._init_rope()258 259 def _init_rope(self):260 if self.config.rope_scaling is None:261 self.rotary_emb = TeleFLMRotaryEmbedding(262 self.head_dim,263 max_position_embeddings=self.max_position_embeddings,264 base=self.rope_theta,265 )266 else:267 scaling_type = self.config.rope_scaling["type"]268 scaling_factor = self.config.rope_scaling["factor"]269 if scaling_type == "linear":270 self.rotary_emb = TeleFLMLinearScalingRotaryEmbedding(271 self.head_dim,272 max_position_embeddings=self.max_position_embeddings,273 scaling_factor=scaling_factor,274 base=self.rope_theta,275 )276 elif scaling_type == "dynamic":277 self.rotary_emb = TeleFLMDynamicNTKScalingRotaryEmbedding(278 self.head_dim,279 max_position_embeddings=self.max_position_embeddings,280 scaling_factor=scaling_factor,281 base=self.rope_theta,282 )283 else:284 raise ValueError(f"Unknown RoPE scaling type {scaling_type}")285 286 def forward(287 self,288 hidden_states: torch.Tensor,289 attention_mask: Optional[torch.Tensor] = None,290 position_ids: Optional[torch.LongTensor] = None,291 past_key_value: Optional[Cache] = None,292 output_attentions: bool = False,293 use_cache: bool = False,294 cache_position: Optional[torch.LongTensor] = None,295 **kwargs,296 ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]:297 bsz, q_len, _ = hidden_states.size()298 299 if self.config.pretraining_tp > 1:300 key_value_slicing = (self.num_key_value_heads * self.head_dim) // self.config.pretraining_tp301 query_slices = self.q_proj.weight.split(302 (self.num_heads * self.head_dim) // self.config.pretraining_tp, dim=0303 )304 key_slices = self.k_proj.weight.split(key_value_slicing, dim=0)305 value_slices = self.v_proj.weight.split(key_value_slicing, dim=0)306 307 query_states = [F.linear(hidden_states, query_slices[i]) for i in range(self.config.pretraining_tp)]308 query_states = torch.cat(query_states, dim=-1)309 310 key_states = [F.linear(hidden_states, key_slices[i]) for i in range(self.config.pretraining_tp)]311 key_states = torch.cat(key_states, dim=-1)312 313 value_states = [F.linear(hidden_states, value_slices[i]) for i in range(self.config.pretraining_tp)]314 value_states = torch.cat(value_states, dim=-1)315 316 else:317 query_states = self.q_proj(hidden_states)318 key_states = self.k_proj(hidden_states)319 value_states = self.v_proj(hidden_states)320 321 query_states = query_states.view(bsz, q_len, self.num_heads, self.head_dim).transpose(1, 2)322 key_states = key_states.view(bsz, q_len, self.num_key_value_heads, self.head_dim).transpose(1, 2)323 value_states = value_states.view(bsz, q_len, self.num_key_value_heads, self.head_dim).transpose(1, 2)324 325 past_key_value = getattr(self, "past_key_value", past_key_value)326 cos, sin = self.rotary_emb(value_states, position_ids)327 query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin)328 329 if past_key_value is not None:330 # sin and cos are specific to RoPE models; cache_position needed for the static cache331 cache_kwargs = {"sin": sin, "cos": cos, "cache_position": cache_position}332 key_states, value_states = past_key_value.update(key_states, value_states, self.layer_idx, cache_kwargs)333 334 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 356 attn_output = attn_output.reshape(bsz, q_len, self.hidden_size)357 358 if self.config.pretraining_tp > 1:359 attn_output = attn_output.split(self.hidden_size // self.config.pretraining_tp, dim=2)360 o_proj_slices = self.o_proj.weight.split(self.hidden_size // self.config.pretraining_tp, dim=1)361 attn_output = sum([F.linear(attn_output[i], o_proj_slices[i]) for i in range(self.config.pretraining_tp)])362 else:363 attn_output = self.o_proj(attn_output)364 365 if not output_attentions:366 attn_weights = None367 368 return attn_output, attn_weights, past_key_value369 370 371class TeleFLMFlashAttention2(TeleFLMAttention):372 """373 Tele-FLM flash attention module. This module inherits from `TeleFLMAttention` as the weights of the module stays374 untouched. The only required change would be on the forward pass where it needs to correctly call the public API of375 flash attention and deal with padding tokens in case the input contains any of them.376 """377 378 def __init__(self, *args, **kwargs):379 super().__init__(*args, **kwargs)380 381 # TODO: Should be removed once Flash Attention for RoCm is bumped to 2.1.382 # 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.383 # 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).384 self._flash_attn_uses_top_left_mask = not is_flash_attn_greater_or_equal_2_10()385 386 def forward(387 self,388 hidden_states: torch.Tensor,389 attention_mask: Optional[torch.LongTensor] = None,390 position_ids: Optional[torch.LongTensor] = None,391 past_key_value: Optional[Cache] = None,392 output_attentions: bool = False,393 use_cache: bool = False,394 cache_position: Optional[torch.LongTensor] = None,395 **kwargs,396 ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]:397 output_attentions = False398 399 bsz, q_len, _ = hidden_states.size()400 401 query_states = self.q_proj(hidden_states)402 key_states = self.k_proj(hidden_states)403 value_states = self.v_proj(hidden_states)404 405 # Flash attention requires the input to have the shape406 # batch_size x seq_length x head_dim x hidden_dim407 # therefore we just need to keep the original shape408 query_states = query_states.view(bsz, q_len, self.num_heads, self.head_dim).transpose(1, 2)409 key_states = key_states.view(bsz, q_len, self.num_key_value_heads, self.head_dim).transpose(1, 2)410 value_states = value_states.view(bsz, q_len, self.num_key_value_heads, self.head_dim).transpose(1, 2)411 412 cos, sin = self.rotary_emb(value_states, position_ids)413 query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin)414 415 past_key_value = getattr(self, "past_key_value", past_key_value)416 417 if past_key_value is not None:418 # sin and cos are specific to RoPE models; cache_position needed for the static cache419 cache_kwargs = {"sin": sin, "cos": cos, "cache_position": cache_position}420 key_states, value_states = past_key_value.update(key_states, value_states, self.layer_idx, cache_kwargs)421 422 # 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 cache423 # to be able to avoid many of these transpose/reshape/view.424 query_states = query_states.transpose(1, 2)425 key_states = key_states.transpose(1, 2)426 value_states = value_states.transpose(1, 2)427 428 dropout_rate = self.attention_dropout if self.training else 0.0429 430 # In PEFT, usually we cast the layer norms in float32 for training stability reasons431 # therefore the input hidden states gets silently casted in float32. Hence, we need432 # cast them back in the correct dtype just to be sure everything works as expected.433 # This might slowdown training & inference so it is recommended to not cast the LayerNorms434 # in fp32. (TeleFLMRMSNorm handles it correctly)435 436 input_dtype = query_states.dtype437 if input_dtype == torch.float32:438 if torch.is_autocast_enabled():439 target_dtype = torch.get_autocast_gpu_dtype()440 # Handle the case where the model is quantized441 elif hasattr(self.config, "_pre_quantization_dtype"):442 target_dtype = self.config._pre_quantization_dtype443 else:444 target_dtype = self.q_proj.weight.dtype445 446 logger.warning_once(447 f"The input hidden states seems to be silently casted in float32, this might be related to"448 f" the fact you have upcasted embedding or layer norm layers in float32. We will cast back the input in"449 f" {target_dtype}."450 )451 452 query_states = query_states.to(target_dtype)453 key_states = key_states.to(target_dtype)454 value_states = value_states.to(target_dtype)455 456 attn_output = self._flash_attention_forward(457 query_states, key_states, value_states, attention_mask, q_len, dropout=dropout_rate458 )459 460 attn_output = attn_output.reshape(bsz, q_len, self.hidden_size).contiguous()461 attn_output = self.o_proj(attn_output)462 463 if not output_attentions:464 attn_weights = None465 466 return attn_output, attn_weights, past_key_value467 468 def _flash_attention_forward(469 self, query_states, key_states, value_states, attention_mask, query_length, dropout=0.0, softmax_scale=None470 ):471 """472 Calls the forward method of Flash Attention - if the input hidden states contain at least one padding token473 first unpad the input, then computes the attention scores and pad the final attention scores.474 475 Args:476 query_states (`torch.Tensor`):477 Input query states to be passed to Flash Attention API478 key_states (`torch.Tensor`):479 Input key states to be passed to Flash Attention API480 value_states (`torch.Tensor`):481 Input value states to be passed to Flash Attention API482 attention_mask (`torch.Tensor`):483 The padding mask - corresponds to a tensor of size `(batch_size, seq_len)` where 0 stands for the484 position of padding tokens and 1 for the position of non-padding tokens.485 dropout (`float`):486 Attention dropout487 softmax_scale (`float`, *optional*):488 The scaling of QK^T before applying softmax. Default to 1 / sqrt(head_dim)489 """490 if not self._flash_attn_uses_top_left_mask:491 causal = self.is_causal492 else:493 # TODO: Remove the `query_length != 1` check once Flash Attention for RoCm is bumped to 2.1. For details, please see the comment in TeleFLMFlashAttention2 __init__.494 causal = self.is_causal and query_length != 1495 496 # Contains at least one padding token in the sequence497 if attention_mask is not None:498 batch_size = query_states.shape[0]499 query_states, key_states, value_states, indices_q, cu_seq_lens, max_seq_lens = self._upad_input(500 query_states, key_states, value_states, attention_mask, query_length501 )502 503 cu_seqlens_q, cu_seqlens_k = cu_seq_lens504 max_seqlen_in_batch_q, max_seqlen_in_batch_k = max_seq_lens505 506 attn_output_unpad = flash_attn_varlen_func(507 query_states,508 key_states,509 value_states,510 cu_seqlens_q=cu_seqlens_q,511 cu_seqlens_k=cu_seqlens_k,512 max_seqlen_q=max_seqlen_in_batch_q,513 max_seqlen_k=max_seqlen_in_batch_k,514 dropout_p=dropout,515 softmax_scale=softmax_scale,516 causal=causal,517 )518 519 attn_output = pad_input(attn_output_unpad, indices_q, batch_size, query_length)520 else:521 attn_output = flash_attn_func(522 query_states, key_states, value_states, dropout, softmax_scale=softmax_scale, causal=causal523 )524 525 return attn_output526 527 def _upad_input(self, query_layer, key_layer, value_layer, attention_mask, query_length):528 indices_k, cu_seqlens_k, max_seqlen_in_batch_k = _get_unpad_data(attention_mask)529 batch_size, kv_seq_len, num_key_value_heads, head_dim = key_layer.shape530 531 key_layer = index_first_axis(532 key_layer.reshape(batch_size * kv_seq_len, num_key_value_heads, head_dim), indices_k533 )534 value_layer = index_first_axis(535 value_layer.reshape(batch_size * kv_seq_len, num_key_value_heads, head_dim), indices_k536 )537 if query_length == kv_seq_len:538 query_layer = index_first_axis(539 query_layer.reshape(batch_size * kv_seq_len, self.num_heads, head_dim), indices_k540 )541 cu_seqlens_q = cu_seqlens_k542 max_seqlen_in_batch_q = max_seqlen_in_batch_k543 indices_q = indices_k544 elif query_length == 1:545 max_seqlen_in_batch_q = 1546 cu_seqlens_q = torch.arange(547 batch_size + 1, dtype=torch.int32, device=query_layer.device548 ) # There is a memcpy here, that is very bad.549 indices_q = cu_seqlens_q[:-1]550 query_layer = query_layer.squeeze(1)551 else:552 # The -q_len: slice assumes left padding.553 attention_mask = attention_mask[:, -query_length:]554 query_layer, indices_q, cu_seqlens_q, max_seqlen_in_batch_q = unpad_input(query_layer, attention_mask)555 556 return (557 query_layer,558 key_layer,559 value_layer,560 indices_q,561 (cu_seqlens_q, cu_seqlens_k),562 (max_seqlen_in_batch_q, max_seqlen_in_batch_k),563 )564 565 566class TeleFLMSdpaAttention(TeleFLMAttention):567 """568 Tele-FLM attention module using torch.nn.functional.scaled_dot_product_attention. This module inherits from569 `TeleFLMAttention` as the weights of the module stays untouched. The only changes are on the forward pass to adapt to570 SDPA API.571 """572 573 # Adapted from TeleFLMAttention.forward574 def forward(575 self,576 hidden_states: torch.Tensor,577 attention_mask: Optional[torch.Tensor] = None,578 position_ids: Optional[torch.LongTensor] = None,579 past_key_value: Optional[Cache] = None,580 output_attentions: bool = False,581 use_cache: bool = False,582 cache_position: Optional[torch.LongTensor] = None,583 ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]:584 if output_attentions:585 # TODO: Improve this warning with e.g. `model.config.attn_implementation = "manual"` once this is implemented.586 logger.warning_once(587 "TeleFLMModel is using TeleFLMSdpaAttention, but `torch.nn.functional.scaled_dot_product_attention` does not support `output_attentions=True`. Falling back to the manual attention implementation, "588 '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.'589 )590 return super().forward(591 hidden_states=hidden_states,592 attention_mask=attention_mask,593 position_ids=position_ids,594 past_key_value=past_key_value,595 output_attentions=output_attentions,596 use_cache=use_cache,597 cache_position=cache_position,598 )599 600 bsz, q_len, _ = hidden_states.size()601 602 query_states = self.q_proj(hidden_states)603 key_states = self.k_proj(hidden_states)604 value_states = self.v_proj(hidden_states)605 606 query_states = query_states.view(bsz, q_len, self.num_heads, self.head_dim).transpose(1, 2)607 key_states = key_states.view(bsz, q_len, self.num_key_value_heads, self.head_dim).transpose(1, 2)608 value_states = value_states.view(bsz, q_len, self.num_key_value_heads, self.head_dim).transpose(1, 2)609 610 cos, sin = self.rotary_emb(value_states, position_ids)611 query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin)612 613 # In case static cache is used, it is an instance attribute.614 past_key_value = getattr(self, "past_key_value", past_key_value)615 616 if past_key_value is not None:617 # sin and cos are specific to RoPE models; cache_position needed for the static cache618 cache_kwargs = {"sin": sin, "cos": cos, "cache_position": cache_position}619 key_states, value_states = past_key_value.update(key_states, value_states, self.layer_idx, cache_kwargs)620 621 key_states = repeat_kv(key_states, self.num_key_value_groups)622 value_states = repeat_kv(value_states, self.num_key_value_groups)623 624 causal_mask = attention_mask625 # if attention_mask is not None and cache_position is not None:626 if attention_mask is not None:627 causal_mask = causal_mask[:, :, :, : key_states.shape[-2]]628 629 # SDPA with memory-efficient backend is currently (torch==2.1.2) bugged with non-contiguous inputs with custom attn_mask,630 # Reference: https://github.com/pytorch/pytorch/issues/112577.631 if query_states.device.type == "cuda" and causal_mask is not None:632 query_states = query_states.contiguous()633 key_states = key_states.contiguous()634 value_states = value_states.contiguous()635 636 attn_output = torch.nn.functional.scaled_dot_product_attention(637 query_states,638 key_states,639 value_states,640 attn_mask=causal_mask,641 dropout_p=self.attention_dropout if self.training else 0.0,642 )643 644 attn_output = attn_output.transpose(1, 2).contiguous()645 attn_output = attn_output.view(bsz, q_len, self.hidden_size)646 647 attn_output = self.o_proj(attn_output)648 649 return attn_output, None, past_key_value650 651 652TELEFLM_ATTENTION_CLASSES = {653 "eager": TeleFLMAttention,654 "flash_attention_2": TeleFLMFlashAttention2,655 "sdpa": TeleFLMSdpaAttention,656}657 658 659class TeleFLMDecoderLayer(nn.Module):660 def __init__(self, config: TeleFLMConfig, layer_idx: int):661 super().__init__()662 self.hidden_size = config.hidden_size663 self.self_attn = TELEFLM_ATTENTION_CLASSES.get(config._attn_implementation, TeleFLMAttention)(config=config, layer_idx=layer_idx)664 self.mlp = TeleFLMMLP(config)665 self.input_layernorm = TeleFLMRMSNorm(config.hidden_size, eps=config.rms_norm_eps)666 self.post_attention_layernorm = TeleFLMRMSNorm(config.hidden_size, eps=config.rms_norm_eps)667 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[Tuple[torch.Tensor]] = 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 """693 if "padding_mask" in kwargs:694 warnings.warn(695 "Passing `padding_mask` is deprecated and will be removed in v4.37. Please make sure use `attention_mask` instead.`"696 )697 698 residual = hidden_states699 700 hidden_states = self.input_layernorm(hidden_states)701 702 # Self Attention703 hidden_states, self_attn_weights, present_key_value = self.self_attn(704 hidden_states=hidden_states,705 attention_mask=attention_mask,706 position_ids=position_ids,707 past_key_value=past_key_value,708 output_attentions=output_attentions,709 use_cache=use_cache,710 cache_position=cache_position,711 **kwargs,712 )713 hidden_states = residual + hidden_states714 715 # Fully Connected716 residual = hidden_states717 hidden_states = self.post_attention_layernorm(hidden_states)718 hidden_states = self.mlp(hidden_states)719 hidden_states = residual + hidden_states720 721 outputs = (hidden_states,)722 723 if output_attentions:724 outputs += (self_attn_weights,)725 726 if use_cache:727 outputs += (present_key_value,)728 729 return outputs730 731 732TELEFLM_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 ([`TeleFLMConfig`]):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 Tele-FLM Model outputting raw hidden-states without any specific head on top.",751 TELEFLM_START_DOCSTRING,752)753class TeleFLMPreTrainedModel(PreTrainedModel):754 config_class = TeleFLMConfig755 base_model_prefix = "model"756 supports_gradient_checkpointing = True757 _no_split_modules = ["TeleFLMDecoderLayer"]758 _skip_keys_device_placement = ["past_key_values"]759 _supports_flash_attn_2 = True760 _supports_sdpa = True761 _supports_cache_class = True762 763 def _init_weights(self, module):764 std = self.config.initializer_range765 if isinstance(module, nn.Linear):766 module.weight.data.normal_(mean=0.0, std=std)767 if module.bias is not None:768 module.bias.data.zero_()769 elif isinstance(module, nn.Embedding):770 module.weight.data.normal_(mean=0.0, std=std)771 if module.padding_idx is not None:772 module.weight.data[module.padding_idx].zero_()773 774 def _setup_cache(self, cache_cls, max_batch_size, max_cache_len: Optional[int] = None):775 if self.config._attn_implementation == "flash_attention_2" and cache_cls == StaticCache:776 raise ValueError(777 "`static` cache implementation is not compatible with `attn_implementation==flash_attention_2` "778 "make sure to use `sdpa` in the mean time, and open an issue at https://github.com/huggingface/transformers"779 )780 781 for layer in self.model.layers:782 device = layer.input_layernorm.weight.device783 if hasattr(self.config, "_pre_quantization_dtype"):784 dtype = self.config._pre_quantization_dtype785 else:786 dtype = layer.self_attn.o_proj.weight.dtype787 layer.self_attn.past_key_value = cache_cls(788 self.config, max_batch_size, max_cache_len, device=device, dtype=dtype789 )790 791 def _reset_cache(self):792 for layer in self.model.layers:793 layer.self_attn.past_key_value = None794 795 796TELEFLM_INPUTS_DOCSTRING = r"""797 Args:798 input_ids (`torch.LongTensor` of shape `(batch_size, sequence_length)`):799 Indices of input sequence tokens in the vocabulary. Padding will be ignored by default should you provide800 it.801 802 Indices can be obtained using [`AutoTokenizer`]. See [`PreTrainedTokenizer.encode`] and803 [`PreTrainedTokenizer.__call__`] for details.804 805 [What are input IDs?](../glossary#input-ids)806 attention_mask (`torch.Tensor` of shape `(batch_size, sequence_length)`, *optional*):807 Mask to avoid performing attention on padding token indices. Mask values selected in `[0, 1]`:808 809 - 1 for tokens that are **not masked**,810 - 0 for tokens that are **masked**.811 812 [What are attention masks?](../glossary#attention-mask)813 814 Indices can be obtained using [`AutoTokenizer`]. See [`PreTrainedTokenizer.encode`] and815 [`PreTrainedTokenizer.__call__`] for details.816 817 If `past_key_values` is used, optionally only the last `input_ids` have to be input (see818 `past_key_values`).819 820 If you want to change padding behavior, you should read [`modeling_opt._prepare_decoder_attention_mask`]821 and modify to your needs. See diagram 1 in [the paper](https://arxiv.org/abs/1910.13461) for more822 information on the default strategy.823 824 - 1 indicates the head is **not masked**,825 - 0 indicates the head is **masked**.826 position_ids (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*):827 Indices of positions of each input sequence tokens in the position embeddings. Selected in the range `[0,828 config.n_positions - 1]`.829 830 [What are position IDs?](../glossary#position-ids)831 past_key_values (`Cache` or `tuple(tuple(torch.FloatTensor))`, *optional*):832 Pre-computed hidden-states (key and values in the self-attention blocks and in the cross-attention833 blocks) that can be used to speed up sequential decoding. This typically consists in the `past_key_values`834 returned by the model at a previous stage of decoding, when `use_cache=True` or `config.use_cache=True`.835 836 Two formats are allowed:837 - a [`~cache_utils.Cache`] instance;838 - Tuple of `tuple(torch.FloatTensor)` of length `config.n_layers`, with each tuple having 2 tensors of839 shape `(batch_size, num_heads, sequence_length, embed_size_per_head)`). This is also known as the legacy840 cache format.841 842 The model will output the same cache format that is fed as input. If no `past_key_values` are passed, the843 legacy cache format will be returned.844 845 If `past_key_values` are used, the user can optionally input only the last `input_ids` (those that don't846 have their past key value states given to this model) of shape `(batch_size, 1)` instead of all `input_ids`847 of shape `(batch_size, sequence_length)`.848 inputs_embeds (`torch.FloatTensor` of shape `(batch_size, sequence_length, hidden_size)`, *optional*):849 Optionally, instead of passing `input_ids` you can choose to directly pass an embedded representation. This850 is useful if you want more control over how to convert `input_ids` indices into associated vectors than the851 model's internal embedding lookup matrix.852 use_cache (`bool`, *optional*):853 If set to `True`, `past_key_values` key value states are returned and can be used to speed up decoding (see854 `past_key_values`).855 output_attentions (`bool`, *optional*):856 Whether or not to return the attentions tensors of all attention layers. See `attentions` under returned857 tensors for more detail.858 output_hidden_states (`bool`, *optional*):859 Whether or not to return the hidden states of all layers. See `hidden_states` under returned tensors for860 more detail.861 return_dict (`bool`, *optional*):862 Whether or not to return a [`~utils.ModelOutput`] instead of a plain tuple.863 cache_position (`torch.LongTensor` of shape `(sequence_length)`, *optional*):864 Indices depicting the position of the input sequence tokens in the sequence. Contrarily to `position_ids`,865 this tensor is not affected by padding. It is used to update the cache in the correct position and to infer866 the complete sequence length.867"""868 869 870@add_start_docstrings(871 "The bare Tele-FLM Model outputting raw hidden-states without any specific head on top.",872 TELEFLM_START_DOCSTRING,873)874class TeleFLMModel(TeleFLMPreTrainedModel):875 """876 Transformer decoder consisting of *config.num_hidden_layers* layers. Each layer is a [`TeleFLMDecoderLayer`]877 878 Args:879 config: TeleFLMConfig880 """881 882 def __init__(self, config: TeleFLMConfig):883 super().__init__(config)884 self.padding_idx = config.pad_token_id885 self.vocab_size = config.vocab_size886 # Mup887 self.use_mup = config.use_mup888 if self.use_mup:889 self.input_mult = config.input_mult890 self.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size, self.padding_idx)891 self.layers = nn.ModuleList(892 [TeleFLMDecoderLayer(config, layer_idx) for layer_idx in range(config.num_hidden_layers)]893 )894 self.norm = TeleFLMRMSNorm(config.hidden_size, eps=config.rms_norm_eps)895 self.gradient_checkpointing = False896 897 # Initialize weights and apply final processing898 self.post_init()899 900 def get_input_embeddings(self):901 return self.embed_tokens902 903 def set_input_embeddings(self, value):904 self.embed_tokens = value905 906 @add_start_docstrings_to_model_forward(TELEFLM_INPUTS_DOCSTRING)907 def forward(908 self,909 input_ids: torch.LongTensor = None,910 attention_mask: Optional[torch.Tensor] = None,911 position_ids: Optional[torch.LongTensor] = None,912 past_key_values: Optional[List[torch.FloatTensor]] = None,913 inputs_embeds: Optional[torch.FloatTensor] = None,914 use_cache: Optional[bool] = None,915 output_attentions: Optional[bool] = None,916 output_hidden_states: Optional[bool] = None,917 return_dict: Optional[bool] = None,918 cache_position: Optional[torch.LongTensor] = None,919 ) -> Union[Tuple, BaseModelOutputWithPast]:920 output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions921 output_hidden_states = (922 output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states923 )924 use_cache = use_cache if use_cache is not None else self.config.use_cache925 return_dict = return_dict if return_dict is not None else self.config.use_return_dict926 927 if (input_ids is None) ^ (inputs_embeds is not None):928 raise ValueError(929 "You cannot specify both input_ids and inputs_embeds at the same time, and must specify either one"930 )931 932 if self.gradient_checkpointing and self.training and use_cache:933 logger.warning_once(934 "`use_cache=True` is incompatible with gradient checkpointing. Setting `use_cache=False`."935 )936 use_cache = False937 938 if inputs_embeds is None:939 inputs_embeds = self.embed_tokens(input_ids)940 941 # Mup942 if self.use_mup:943 inputs_embeds = inputs_embeds * self.input_mult944 945 past_seen_tokens = 0946 if use_cache: # kept for BC (cache positions)947 if not isinstance(past_key_values, StaticCache):948 past_key_values = DynamicCache.from_legacy_cache(past_key_values)949 past_seen_tokens = past_key_values.get_seq_length()950 951 if cache_position is None:952 if isinstance(past_key_values, StaticCache):953 raise ValueError("cache_position is a required argument when using StaticCache.")954 cache_position = torch.arange(955 past_seen_tokens, past_seen_tokens + inputs_embeds.shape[1], device=inputs_embeds.device956 )957 958 if position_ids is None:959 position_ids = cache_position.unsqueeze(0)960 961 causal_mask = self._update_causal_mask(attention_mask, inputs_embeds, cache_position)962 963 # embed positions964 hidden_states = inputs_embeds965 966 # decoder layers967 all_hidden_states = () if output_hidden_states else None968 all_self_attns = () if output_attentions else None969 next_decoder_cache = None970 971 for decoder_layer in self.layers:972 if output_hidden_states:973 all_hidden_states += (hidden_states,)974 975 if self.gradient_checkpointing and self.training:976 layer_outputs = self._gradient_checkpointing_func(977 decoder_layer.__call__,978 hidden_states,979 causal_mask,980 position_ids,981 past_key_values,982 output_attentions,983 use_cache,984 cache_position,985 )986 else:987 layer_outputs = decoder_layer(988 hidden_states,989 attention_mask=causal_mask,990 position_ids=position_ids,991 past_key_value=past_key_values,992 output_attentions=output_attentions,993 use_cache=use_cache,994 cache_position=cache_position,995 )996 997 hidden_states = layer_outputs[0]998 999 if use_cache:1000 next_decoder_cache = layer_outputs[2 if output_attentions else 1]1001 1002 if output_attentions:1003 all_self_attns += (layer_outputs[1],)1004 1005 hidden_states = self.norm(hidden_states)1006 1007 # add hidden states from the last decoder layer1008 if output_hidden_states:1009 all_hidden_states += (hidden_states,)1010 1011 next_cache = None1012 if use_cache:1013 next_cache = (1014 next_decoder_cache.to_legacy_cache() if isinstance(next_decoder_cache, Cache) else next_decoder_cache1015 )1016 if not return_dict:1017 return tuple(v for v in [hidden_states, next_cache, all_hidden_states, all_self_attns] if v is not None)1018 return BaseModelOutputWithPast(1019 last_hidden_state=hidden_states,1020 past_key_values=next_cache,1021 hidden_states=all_hidden_states,1022 attentions=all_self_attns,1023 )1024 1025 # TODO: As of torch==2.2.0, the `attention_mask` passed to the model in `generate` is 2D and of dynamic length even when the static1026 # KV cache is used. This is an issue for torch.compile which then recaptures cudagraphs at each decode steps due to the dynamic shapes.1027 # (`recording cudagraph tree for symint key 13`, etc.), which is VERY slow. A workaround is `@torch.compiler.disable`, but this prevents using1028 # `fullgraph=True`. See more context in https://github.com/huggingface/transformers/pull/291141029 def _update_causal_mask(self, attention_mask, input_tensor, cache_position):1030 if self.config._attn_implementation == "flash_attention_2":1031 if attention_mask is not None and 0.0 in attention_mask:1032 return attention_mask1033 return None1034 1035 dtype, device = input_tensor.dtype, input_tensor.device1036 min_dtype = torch.finfo(dtype).min1037 sequence_length = input_tensor.shape[1]1038 if hasattr(getattr(self.layers[0], "self_attn", {}), "past_key_value"): # static cache1039 target_length = self.config.max_position_embeddings1040 else: # dynamic cache1041 target_length = (1042 attention_mask.shape[-1] if isinstance(attention_mask, torch.Tensor) else cache_position[-1] + 11043 )1044 1045 causal_mask = torch.full((sequence_length, target_length), fill_value=min_dtype, dtype=dtype, device=device)1046 if sequence_length != 1:1047 causal_mask = torch.triu(causal_mask, diagonal=1)1048 causal_mask *= torch.arange(target_length, device=device) > cache_position.reshape(-1, 1)1049 causal_mask = causal_mask[None, None, :, :].expand(input_tensor.shape[0], 1, -1, -1)1050 if attention_mask is not None:1051 causal_mask = causal_mask.clone() # copy to contiguous memory for in-place edit1052 if attention_mask.dim() == 2:1053 mask_length = attention_mask.shape[-1]1054 padding_mask = causal_mask[..., :mask_length].eq(0.0) * attention_mask[:, None, None, :].eq(0.0)1055 causal_mask[..., :mask_length] = causal_mask[..., :mask_length].masked_fill(padding_mask, min_dtype)1056 elif attention_mask.dim() == 4:1057 # backwards compatibility: we allow passing a 4D attention mask shorter than the input length with1058 # cache. In that case, the 4D attention mask attends to the newest tokens only.1059 if attention_mask.shape[-2] < cache_position[0] + sequence_length:1060 offset = cache_position[0]1061 else:1062 offset = 01063 mask_shape = attention_mask.shape1064 mask_slice = (attention_mask.eq(0.0)).to(dtype=dtype) * min_dtype1065 causal_mask[1066 : mask_shape[0], : mask_shape[1], offset : mask_shape[2] + offset, : mask_shape[3]1067 ] = mask_slice1068 1069 if (1070 self.config._attn_implementation == "sdpa"1071 and attention_mask is not None1072 and attention_mask.device.type == "cuda"1073 ):1074 # Attend to all tokens in fully masked rows in the causal_mask, for example the relevant first rows when1075 # using left padding. This is required by F.scaled_dot_product_attention memory-efficient attention path.1076 # Details: https://github.com/pytorch/pytorch/issues/1102131077 causal_mask = AttentionMaskConverter._unmask_unattended(causal_mask, min_dtype)1078 1079 return causal_mask1080 1081 1082class TeleFLMForCausalLM(TeleFLMPreTrainedModel):1083 _tied_weights_keys = ["lm_head.weight"]1084 1085 def __init__(self, config):1086 super().__init__(config)1087 self.model = TeleFLMModel(config)1088 self.vocab_size = config.vocab_size1089 self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)1090 self.use_mup = config.use_mup1091 if self.use_mup:1092 self.mup_scale_factor = config.mup_scale_factor1093 self.output_mult = config.output_mult / self.mup_scale_factor1094 # Initialize weights and apply final processing1095 self.post_init()1096 1097 def get_input_embeddings(self):1098 return self.model.embed_tokens1099 1100 def set_input_embeddings(self, value):1101 self.model.embed_tokens = value1102 1103 def get_output_embeddings(self):1104 return self.lm_head1105 1106 def set_output_embeddings(self, new_embeddings):1107 self.lm_head = new_embeddings1108 1109 def set_decoder(self, decoder):1110 self.model = decoder1111 1112 def get_decoder(self):1113 return self.model1114 1115 @add_start_docstrings_to_model_forward(TELEFLM_INPUTS_DOCSTRING)1116 @replace_return_docstrings(output_type=CausalLMOutputWithPast, config_class=_CONFIG_FOR_DOC)1117 def forward(1118 self,1119 input_ids: torch.LongTensor = None,1120 attention_mask: Optional[torch.Tensor] = None,1121 position_ids: Optional[torch.LongTensor] = None,1122 past_key_values: Optional[List[torch.FloatTensor]] = None,1123 inputs_embeds: Optional[torch.FloatTensor] = None,1124 labels: Optional[torch.LongTensor] = None,1125 use_cache: Optional[bool] = None,1126 output_attentions: Optional[bool] = None,1127 output_hidden_states: Optional[bool] = None,1128 return_dict: Optional[bool] = None,1129 cache_position: Optional[torch.LongTensor] = None,1130 ) -> Union[Tuple, CausalLMOutputWithPast]:1131 r"""1132 Args:1133 labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*):1134 Labels for computing the masked language modeling loss. Indices should either be in `[0, ...,1135 config.vocab_size]` or -100 (see `input_ids` docstring). Tokens with indices set to `-100` are ignored1136 (masked), the loss is only computed for the tokens with labels in `[0, ..., config.vocab_size]`.1137 1138 Returns:1139 1140 Example:1141 1142 ```python1143 >>> from transformers import AutoTokenizer, TeleFLMForCausalLM1144 1145 >>> model = TeleFLMForCausalLM.from_pretrained("CofeAI/Tele-FLM")1146 >>> tokenizer = AutoTokenizer.from_pretrained("CofeAI/Tele-FLM")1147 1148 >>> prompt = "Hey, are you conscious? Can you talk to me?"1149 >>> inputs = tokenizer(prompt, return_tensors="pt")1150 1151 >>> # Generate1152 >>> generate_ids = model.generate(inputs.input_ids, max_length=30)1153 >>> tokenizer.batch_decode(generate_ids, skip_special_tokens=True, clean_up_tokenization_spaces=False)[0]1154 "Hey, are you conscious? Can you talk to me?\nI'm not conscious, but I can talk to you."1155 ```"""1156 output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions1157 output_hidden_states = (1158 output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states1159 )1160 return_dict = return_dict if return_dict is not None else self.config.use_return_dict1161 1162 # decoder outputs consists of (dec_features, layer_state, dec_hidden, dec_attn)1163 outputs = self.model(1164 input_ids=input_ids,1165 attention_mask=attention_mask,1166 position_ids=position_ids,1167 past_key_values=past_key_values,1168 inputs_embeds=inputs_embeds,1169 use_cache=use_cache,1170 output_attentions=output_attentions,1171 output_hidden_states=output_hidden_states,1172 return_dict=return_dict,1173 cache_position=cache_position,1174 )1175 1176 hidden_states = outputs[0]1177 if self.config.pretraining_tp > 1:1178 lm_head_slices = self.lm_head.weight.split(self.vocab_size // self.config.pretraining_tp, dim=0)1179 logits = [F.linear(hidden_states, lm_head_slices[i]) for i in range(self.config.pretraining_tp)]1180 logits = torch.cat(logits, dim=-1)1181 else:1182 logits = self.lm_head(hidden_states)1183 logits = logits.float()1184 # Mup1185 if self.use_mup:1186 logits = logits * self.output_mult1187 1188 loss = None1189 if labels is not None:1190 # Shift so that tokens < n predict n1191 shift_logits = logits[..., :-1, :].contiguous()1192 shift_labels = labels[..., 1:].contiguous()1193 # Flatten the tokens1194 loss_fct = CrossEntropyLoss()1195 shift_logits = shift_logits.view(-1, self.config.vocab_size)1196 shift_labels = shift_labels.view(-1)1197 # Enable model parallelism1198 shift_labels = shift_labels.to(shift_logits.device)1199 loss = loss_fct(shift_logits, shift_labels)1200