CofeAI/FLM-Audio
132.9k
1# coding=utf-82"""PyTorch FLM-Audio model, based on LLAMA implementation."""3 4import math5import warnings6from typing import List, Optional, Tuple, Union7from dataclasses import dataclass8 9import torch10import torch.nn as nn11import torch.nn.functional as F12 13from transformers.activations import ACT2FN14from transformers.modeling_rope_utils import ROPE_INIT_FUNCTIONS, dynamic_rope_update15from transformers.cache_utils import Cache, DynamicCache, StaticCache16from transformers.modeling_attn_mask_utils import AttentionMaskConverter17from transformers.modeling_outputs import (18 ModelOutput,19 BaseModelOutputWithPast,20 CausalLMOutputWithPast,21)22from transformers.modeling_utils import PreTrainedModel23from transformers.pytorch_utils import ALL_LAYERNORM_LAYERS24from transformers.utils import (25 add_start_docstrings,26 add_start_docstrings_to_model_forward,27 is_flash_attn_2_available,28 is_flash_attn_greater_or_equal_2_10,29 logging,30 replace_return_docstrings,31)32from .configuration_flmaudio import FLMAudioConfig33from .depth_gpt import DepthGPT, DepthGPTConfig34 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 = "FLMAudioConfig"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 FLMAudioRMSNorm(nn.Module):58 def __init__(self, hidden_size, eps=1e-6):59 """60 FLMAudioRMSNorm 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(FLMAudioRMSNorm)75 76class FLMAudioRotaryEmbedding(nn.Module):77 def __init__(self, config, device=None):78 super().__init__()79 # BC: "rope_type" was originally "type"80 if hasattr(config, "rope_scaling") and config.rope_scaling is not None:81 self.rope_type = config.rope_scaling.get("rope_type", config.rope_scaling.get("type"))82 else:83 self.rope_type = "default"84 self.max_seq_len_cached = config.max_position_embeddings85 self.original_max_seq_len = config.max_position_embeddings86 87 self.config = config88 self.rope_init_fn = ROPE_INIT_FUNCTIONS[self.rope_type]89 90 inv_freq, self.attention_scaling = self.rope_init_fn(self.config, device)91 self.register_buffer("inv_freq", inv_freq, persistent=False)92 self.original_inv_freq = self.inv_freq93 94 @torch.no_grad()95 @dynamic_rope_update # power user: used with advanced RoPE types (e.g. dynamic rope)96 def forward(self, x, position_ids):97 inv_freq_expanded = self.inv_freq[None, None, :, None].float().expand(3, position_ids.shape[1], -1, 1)98 position_ids_expanded = position_ids[:, :, None, :].float() # shape (3, bs, 1, positions)99 100 device_type = x.device.type if isinstance(x.device.type, str) and x.device.type != "mps" else "cpu"101 with torch.autocast(device_type=device_type, enabled=False): # Force float32102 freqs = (inv_freq_expanded.float() @ position_ids_expanded.float()).transpose(2, 3)103 emb = torch.cat((freqs, freqs), dim=-1)104 cos = emb.cos() * self.attention_scaling105 sin = emb.sin() * self.attention_scaling106 107 return cos.to(dtype=x.dtype), sin.to(dtype=x.dtype)108 109 110def rotate_half(x):111 """Rotates half the hidden dims of the input."""112 x1 = x[..., : x.shape[-1] // 2]113 x2 = x[..., x.shape[-1] // 2 :]114 return torch.cat((-x2, x1), dim=-1)115 116 117def apply_multimodal_rotary_pos_emb(q, k, cos, sin, mrope_section, unsqueeze_dim=1):118 mrope_section = mrope_section * 2119 cos = torch.cat([m[i % 3] for i, m in enumerate(cos.split(mrope_section, dim=-1))], dim=-1).unsqueeze(120 unsqueeze_dim121 )122 sin = torch.cat([m[i % 3] for i, m in enumerate(sin.split(mrope_section, dim=-1))], dim=-1).unsqueeze(123 unsqueeze_dim124 )125 126 q_embed = (q * cos) + (rotate_half(q) * sin)127 k_embed = (k * cos) + (rotate_half(k) * sin)128 return q_embed, k_embed129 130 131class FLMAudioMLP(nn.Module):132 def __init__(self, config):133 super().__init__()134 self.config = config135 self.hidden_size = config.hidden_size136 self.intermediate_size = config.intermediate_size137 self.gate_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=False)138 self.up_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=False)139 self.down_proj = nn.Linear(self.intermediate_size, self.hidden_size, bias=False)140 self.act_fn = ACT2FN[config.hidden_act]141 142 def forward(self, x):143 if self.config.pretraining_tp > 1:144 slice = self.intermediate_size // self.config.pretraining_tp145 gate_proj_slices = self.gate_proj.weight.split(slice, dim=0)146 up_proj_slices = self.up_proj.weight.split(slice, dim=0)147 down_proj_slices = self.down_proj.weight.split(slice, dim=1)148 149 gate_proj = torch.cat(150 [151 F.linear(x, gate_proj_slices[i])152 for i in range(self.config.pretraining_tp)153 ],154 dim=-1,155 )156 up_proj = torch.cat(157 [158 F.linear(x, up_proj_slices[i])159 for i in range(self.config.pretraining_tp)160 ],161 dim=-1,162 )163 164 intermediate_states = (self.act_fn(gate_proj) * up_proj).split(slice, dim=2)165 down_proj = [166 F.linear(intermediate_states[i], down_proj_slices[i])167 for i in range(self.config.pretraining_tp)168 ]169 down_proj = sum(down_proj)170 else:171 down_proj = self.down_proj(self.act_fn(self.gate_proj(x)) * self.up_proj(x))172 173 return down_proj174 175 176def repeat_kv(hidden_states: torch.Tensor, n_rep: int) -> torch.Tensor:177 """178 This is the equivalent of torch.repeat_interleave(x, dim=1, repeats=n_rep). The hidden states go from (batch,179 num_key_value_heads, seqlen, head_dim) to (batch, num_attention_heads, seqlen, head_dim)180 """181 batch, num_key_value_heads, slen, head_dim = hidden_states.shape182 if n_rep == 1:183 return hidden_states184 hidden_states = hidden_states[:, :, None, :, :].expand(185 batch, num_key_value_heads, n_rep, slen, head_dim186 )187 return hidden_states.reshape(batch, num_key_value_heads * n_rep, slen, head_dim)188 189 190class FLMAudioAttention(nn.Module):191 """Multi-headed attention from 'Attention Is All You Need' paper"""192 193 def __init__(self, config: FLMAudioConfig, layer_idx: Optional[int] = None):194 super().__init__()195 self.config = config196 self.layer_idx = layer_idx197 if layer_idx is None:198 logger.warning_once(199 f"Instantiating {self.__class__.__name__} without passing a `layer_idx` is not recommended and will "200 "lead to errors during the forward call if caching is used. Please make sure to provide a `layer_idx` "201 "when creating this class."202 )203 204 self.attention_dropout = config.attention_dropout205 self.hidden_size = config.hidden_size206 self.num_heads = config.num_attention_heads207 self.head_dim = self.hidden_size // self.num_heads208 self.num_key_value_heads = config.num_key_value_heads209 self.num_key_value_groups = self.num_heads // self.num_key_value_heads210 self.is_causal = True211 self.rope_scaling = config.rope_scaling212 213 if (self.head_dim * self.num_heads) != self.hidden_size:214 raise ValueError(215 f"hidden_size must be divisible by num_heads (got `hidden_size`: {self.hidden_size}"216 f" and `num_heads`: {self.num_heads})."217 )218 219 self.q_proj = nn.Linear(220 self.hidden_size, self.num_heads * self.head_dim, bias=config.attention_bias221 )222 self.k_proj = nn.Linear(223 self.hidden_size,224 self.num_key_value_heads * self.head_dim,225 bias=config.attention_bias,226 )227 self.v_proj = nn.Linear(228 self.hidden_size,229 self.num_key_value_heads * self.head_dim,230 bias=config.attention_bias,231 )232 self.o_proj = nn.Linear(233 self.hidden_size, self.hidden_size, bias=config.attention_bias and not config.disable_att_o_bias234 )235 236 237 def forward(238 self,239 hidden_states: torch.Tensor,240 attention_mask: Optional[torch.Tensor] = None,241 position_ids: Optional[torch.LongTensor] = None,242 past_key_value: Optional[Cache] = None,243 output_attentions: bool = False,244 use_cache: bool = False,245 cache_position: Optional[torch.LongTensor] = None,246 position_embeddings: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, # necessary, but kept here for BC247 **kwargs,248 ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]:249 bsz, q_len, _ = hidden_states.size()250 251 if self.config.pretraining_tp > 1:252 key_value_slicing = (253 self.num_key_value_heads * self.head_dim254 ) // self.config.pretraining_tp255 query_slices = self.q_proj.weight.split(256 (self.num_heads * self.head_dim) // self.config.pretraining_tp, dim=0257 )258 key_slices = self.k_proj.weight.split(key_value_slicing, dim=0)259 value_slices = self.v_proj.weight.split(key_value_slicing, dim=0)260 261 query_states = [262 F.linear(hidden_states, query_slices[i])263 for i in range(self.config.pretraining_tp)264 ]265 query_states = torch.cat(query_states, dim=-1)266 267 key_states = [268 F.linear(hidden_states, key_slices[i])269 for i in range(self.config.pretraining_tp)270 ]271 key_states = torch.cat(key_states, dim=-1)272 273 value_states = [274 F.linear(hidden_states, value_slices[i])275 for i in range(self.config.pretraining_tp)276 ]277 value_states = torch.cat(value_states, dim=-1)278 279 else:280 query_states = self.q_proj(hidden_states)281 key_states = self.k_proj(hidden_states)282 value_states = self.v_proj(hidden_states)283 284 query_states = query_states.view(285 bsz, q_len, self.num_heads, self.head_dim286 ).transpose(1, 2)287 key_states = key_states.view(288 bsz, q_len, self.num_key_value_heads, self.head_dim289 ).transpose(1, 2)290 value_states = value_states.view(291 bsz, q_len, self.num_key_value_heads, self.head_dim292 ).transpose(1, 2)293 294 cos, sin = position_embeddings295 296 query_states, key_states = apply_multimodal_rotary_pos_emb(297 query_states, key_states, cos, sin, self.rope_scaling["mrope_section"]298 )299 300 if past_key_value is not None:301 # sin and cos are specific to RoPE models; cache_position needed for the static cache302 cache_kwargs = {"sin": sin, "cos": cos, "cache_position": cache_position}303 key_states, value_states = past_key_value.update(304 key_states, value_states, self.layer_idx, cache_kwargs305 )306 307 key_states = repeat_kv(key_states, self.num_key_value_groups)308 value_states = repeat_kv(value_states, self.num_key_value_groups)309 310 attn_weights = torch.matmul(311 query_states, key_states.transpose(2, 3)312 ) / math.sqrt(self.head_dim)313 314 if attention_mask is not None: # no matter the length, we just slice it315 causal_mask = attention_mask[:, :, :, : key_states.shape[-2]]316 attn_weights = attn_weights + causal_mask317 318 # upcast attention to fp32319 attn_weights = nn.functional.softmax(320 attn_weights, dim=-1, dtype=torch.float32321 ).to(query_states.dtype)322 attn_weights = nn.functional.dropout(323 attn_weights, p=self.attention_dropout, training=self.training324 )325 attn_output = torch.matmul(attn_weights, value_states)326 327 if attn_output.size() != (bsz, self.num_heads, q_len, self.head_dim):328 raise ValueError(329 f"`attn_output` should be of size {(bsz, self.num_heads, q_len, self.head_dim)}, but is"330 f" {attn_output.size()}"331 )332 333 attn_output = attn_output.transpose(1, 2).contiguous()334 335 attn_output = attn_output.reshape(bsz, q_len, self.hidden_size)336 337 if self.config.pretraining_tp > 1:338 attn_output = attn_output.split(339 self.hidden_size // self.config.pretraining_tp, dim=2340 )341 o_proj_slices = self.o_proj.weight.split(342 self.hidden_size // self.config.pretraining_tp, dim=1343 )344 attn_output = sum(345 [346 F.linear(attn_output[i], o_proj_slices[i])347 for i in range(self.config.pretraining_tp)348 ]349 )350 else:351 attn_output = self.o_proj(attn_output)352 353 if not output_attentions:354 attn_weights = None355 356 return attn_output, attn_weights, past_key_value357 358 359class FLMAudioFlashAttention2(FLMAudioAttention):360 """361 FLM-Audio flash attention module. This module inherits from `FLMAudioAttention` as the weights of the module stays362 untouched. The only required change would be on the forward pass where it needs to correctly call the public API of363 flash attention and deal with padding tokens in case the input contains any of them.364 """365 366 def __init__(self, *args, **kwargs):367 super().__init__(*args, **kwargs)368 369 # TODO: Should be removed once Flash Attention for RoCm is bumped to 2.1.370 # 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.371 # 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).372 self._flash_attn_uses_top_left_mask = not is_flash_attn_greater_or_equal_2_10()373 374 def forward(375 self,376 hidden_states: torch.Tensor,377 attention_mask: Optional[torch.LongTensor] = None,378 position_ids: Optional[torch.LongTensor] = None,379 past_key_value: Optional[Cache] = None,380 output_attentions: bool = False,381 use_cache: bool = False,382 cache_position: Optional[torch.LongTensor] = None,383 position_embeddings: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, # necessary, but kept here for BC384 **kwargs,385 ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]:386 output_attentions = False387 388 bsz, q_len, _ = hidden_states.size()389 390 query_states = self.q_proj(hidden_states)391 key_states = self.k_proj(hidden_states)392 value_states = self.v_proj(hidden_states)393 394 # Flash attention requires the input to have the shape395 # batch_size x seq_length x head_dim x hidden_dim396 # therefore we just need to keep the original shape397 query_states = query_states.view(398 bsz, q_len, self.num_heads, self.head_dim399 ).transpose(1, 2)400 key_states = key_states.view(401 bsz, q_len, self.num_key_value_heads, self.head_dim402 ).transpose(1, 2)403 value_states = value_states.view(404 bsz, q_len, self.num_key_value_heads, self.head_dim405 ).transpose(1, 2)406 407 # cos, sin = self.rotary_emb(value_states, position_ids)408 cos, sin = position_embeddings409 # query_states, key_states = apply_rotary_pos_emb(410 # query_states, key_states, cos, sin411 # )412 query_states, key_states = apply_multimodal_rotary_pos_emb(413 query_states, key_states, cos, sin, self.rope_scaling["mrope_section"]414 )415 416 past_key_value = getattr(self, "past_key_value", past_key_value)417 418 if past_key_value is not None:419 # sin and cos are specific to RoPE models; cache_position needed for the static cache420 cache_kwargs = {"sin": sin, "cos": cos, "cache_position": cache_position}421 key_states, value_states = past_key_value.update(422 key_states, value_states, self.layer_idx, cache_kwargs423 )424 425 # 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 cache426 # to be able to avoid many of these transpose/reshape/view.427 query_states = query_states.transpose(1, 2)428 key_states = key_states.transpose(1, 2)429 value_states = value_states.transpose(1, 2)430 431 dropout_rate = self.attention_dropout if self.training else 0.0432 433 # In PEFT, usually we cast the layer norms in float32 for training stability reasons434 # therefore the input hidden states gets silently casted in float32. Hence, we need435 # cast them back in the correct dtype just to be sure everything works as expected.436 # This might slowdown training & inference so it is recommended to not cast the LayerNorms437 # in fp32. (FLMAudioRMSNorm handles it correctly)438 439 input_dtype = query_states.dtype440 if input_dtype == torch.float32:441 if torch.is_autocast_enabled():442 target_dtype = torch.get_autocast_gpu_dtype()443 # Handle the case where the model is quantized444 elif hasattr(self.config, "_pre_quantization_dtype"):445 target_dtype = self.config._pre_quantization_dtype446 else:447 target_dtype = self.q_proj.weight.dtype448 449 logger.warning_once(450 f"The input hidden states seems to be silently casted in float32, this might be related to"451 f" the fact you have upcasted embedding or layer norm layers in float32. We will cast back the input in"452 f" {target_dtype}."453 )454 455 query_states = query_states.to(target_dtype)456 key_states = key_states.to(target_dtype)457 value_states = value_states.to(target_dtype)458 459 attn_output = self._flash_attention_forward(460 query_states,461 key_states,462 value_states,463 attention_mask,464 q_len,465 dropout=dropout_rate,466 )467 468 attn_output = attn_output.reshape(bsz, q_len, self.hidden_size).contiguous()469 attn_output = self.o_proj(attn_output)470 471 if not output_attentions:472 attn_weights = None473 474 return attn_output, attn_weights, past_key_value475 476 def _flash_attention_forward(477 self,478 query_states,479 key_states,480 value_states,481 attention_mask,482 query_length,483 dropout=0.0,484 softmax_scale=None,485 ):486 """487 Calls the forward method of Flash Attention - if the input hidden states contain at least one padding token488 first unpad the input, then computes the attention scores and pad the final attention scores.489 490 Args:491 query_states (`torch.Tensor`):492 Input query states to be passed to Flash Attention API493 key_states (`torch.Tensor`):494 Input key states to be passed to Flash Attention API495 value_states (`torch.Tensor`):496 Input value states to be passed to Flash Attention API497 attention_mask (`torch.Tensor`):498 The padding mask - corresponds to a tensor of size `(batch_size, seq_len)` where 0 stands for the499 position of padding tokens and 1 for the position of non-padding tokens.500 dropout (`float`):501 Attention dropout502 softmax_scale (`float`, *optional*):503 The scaling of QK^T before applying softmax. Default to 1 / sqrt(head_dim)504 """505 if not self._flash_attn_uses_top_left_mask:506 causal = self.is_causal507 else:508 # TODO: Remove the `query_length != 1` check once Flash Attention for RoCm is bumped to 2.1. For details, please see the comment in FLMAudioFlashAttention2 __init__.509 causal = self.is_causal and query_length != 1510 511 # Contains at least one padding token in the sequence512 if attention_mask is not None:513 batch_size = query_states.shape[0]514 (515 query_states,516 key_states,517 value_states,518 indices_q,519 cu_seq_lens,520 max_seq_lens,521 ) = self._upad_input(522 query_states, key_states, value_states, attention_mask, query_length523 )524 525 cu_seqlens_q, cu_seqlens_k = cu_seq_lens526 max_seqlen_in_batch_q, max_seqlen_in_batch_k = max_seq_lens527 528 attn_output_unpad = flash_attn_varlen_func(529 query_states,530 key_states,531 value_states,532 cu_seqlens_q=cu_seqlens_q,533 cu_seqlens_k=cu_seqlens_k,534 max_seqlen_q=max_seqlen_in_batch_q,535 max_seqlen_k=max_seqlen_in_batch_k,536 dropout_p=dropout,537 softmax_scale=softmax_scale,538 causal=causal,539 )540 541 attn_output = pad_input(542 attn_output_unpad, indices_q, batch_size, query_length543 )544 else:545 attn_output = flash_attn_func(546 query_states,547 key_states,548 value_states,549 dropout,550 softmax_scale=softmax_scale,551 causal=causal,552 )553 554 return attn_output555 556 def _upad_input(557 self, query_layer, key_layer, value_layer, attention_mask, query_length558 ):559 indices_k, cu_seqlens_k, max_seqlen_in_batch_k = _get_unpad_data(attention_mask)560 batch_size, kv_seq_len, num_key_value_heads, head_dim = key_layer.shape561 562 key_layer = index_first_axis(563 key_layer.reshape(batch_size * kv_seq_len, num_key_value_heads, head_dim),564 indices_k,565 )566 value_layer = index_first_axis(567 value_layer.reshape(batch_size * kv_seq_len, num_key_value_heads, head_dim),568 indices_k,569 )570 if query_length == kv_seq_len:571 query_layer = index_first_axis(572 query_layer.reshape(batch_size * kv_seq_len, self.num_heads, head_dim),573 indices_k,574 )575 cu_seqlens_q = cu_seqlens_k576 max_seqlen_in_batch_q = max_seqlen_in_batch_k577 indices_q = indices_k578 elif query_length == 1:579 max_seqlen_in_batch_q = 1580 cu_seqlens_q = torch.arange(581 batch_size + 1, dtype=torch.int32, device=query_layer.device582 ) # There is a memcpy here, that is very bad.583 indices_q = cu_seqlens_q[:-1]584 query_layer = query_layer.squeeze(1)585 else:586 # The -q_len: slice assumes left padding.587 attention_mask = attention_mask[:, -query_length:]588 query_layer, indices_q, cu_seqlens_q, max_seqlen_in_batch_q = unpad_input(589 query_layer, attention_mask590 )591 592 return (593 query_layer,594 key_layer,595 value_layer,596 indices_q,597 (cu_seqlens_q, cu_seqlens_k),598 (max_seqlen_in_batch_q, max_seqlen_in_batch_k),599 )600 601 602class FLMAudioSdpaAttention(FLMAudioAttention):603 """604 FLM-Audio attention module using torch.nn.functional.scaled_dot_product_attention. This module inherits from605 `FLMAudioAttention` as the weights of the module stays untouched. The only changes are on the forward pass to adapt to606 SDPA API.607 """608 609 # Adapted from FLMAudioAttention.forward610 def forward(611 self,612 hidden_states: torch.Tensor,613 attention_mask: Optional[torch.Tensor] = None,614 position_ids: Optional[torch.LongTensor] = None,615 past_key_value: Optional[Cache] = None,616 output_attentions: bool = False,617 use_cache: bool = False,618 cache_position: Optional[torch.LongTensor] = None,619 position_embeddings: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, # necessary, but kept here for BC620 ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]:621 if output_attentions:622 # TODO: Improve this warning with e.g. `model.config.attn_implementation = "manual"` once this is implemented.623 logger.warning_once(624 "FLMAudioModel is using FLMAudioSdpaAttention, but `torch.nn.functional.scaled_dot_product_attention` does not support `output_attentions=True`. Falling back to the manual attention implementation, "625 '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.'626 )627 return super().forward(628 hidden_states=hidden_states,629 attention_mask=attention_mask,630 position_ids=position_ids,631 past_key_value=past_key_value,632 output_attentions=output_attentions,633 use_cache=use_cache,634 cache_position=cache_position,635 )636 637 bsz, q_len, _ = hidden_states.size()638 639 query_states = self.q_proj(hidden_states)640 key_states = self.k_proj(hidden_states)641 value_states = self.v_proj(hidden_states)642 643 query_states = query_states.view(644 bsz, q_len, self.num_heads, self.head_dim645 ).transpose(1, 2)646 key_states = key_states.view(647 bsz, q_len, self.num_key_value_heads, self.head_dim648 ).transpose(1, 2)649 value_states = value_states.view(650 bsz, q_len, self.num_key_value_heads, self.head_dim651 ).transpose(1, 2)652 653 cos, sin = position_embeddings654 655 query_states, key_states = apply_multimodal_rotary_pos_emb(656 query_states, key_states, cos, sin, self.rope_scaling["mrope_section"]657 )658 659 if past_key_value is not None:660 # sin and cos are specific to RoPE models; cache_position needed for the static cache661 cache_kwargs = {"sin": sin, "cos": cos, "cache_position": cache_position}662 key_states, value_states = past_key_value.update(663 key_states, value_states, self.layer_idx, cache_kwargs664 )665 666 key_states = repeat_kv(key_states, self.num_key_value_groups)667 value_states = repeat_kv(value_states, self.num_key_value_groups)668 669 causal_mask = attention_mask670 # if attention_mask is not None and cache_position is not None:671 if attention_mask is not None:672 causal_mask = causal_mask[:, :, :, : key_states.shape[-2]]673 674 # SDPA with memory-efficient backend is currently (torch==2.1.2) bugged with non-contiguous inputs with custom attn_mask,675 # Reference: https://github.com/pytorch/pytorch/issues/112577.676 if query_states.device.type == "cuda" and causal_mask is not None:677 query_states = query_states.contiguous()678 key_states = key_states.contiguous()679 value_states = value_states.contiguous()680 681 attn_output = F.scaled_dot_product_attention(682 query_states,683 key_states,684 value_states,685 attn_mask=causal_mask,686 dropout_p=self.attention_dropout if self.training else 0.0,687 )688 689 attn_output = attn_output.transpose(1, 2).contiguous()690 attn_output = attn_output.view(bsz, q_len, self.hidden_size)691 692 attn_output = self.o_proj(attn_output)693 694 return attn_output, None, past_key_value695 696 697FLMAUDIO_ATTENTION_CLASSES = {698 "eager": FLMAudioAttention,699 "flash_attention_2": FLMAudioFlashAttention2,700 "sdpa": FLMAudioSdpaAttention,701}702 703 704class FLMAudioDecoderLayer(nn.Module):705 def __init__(self, config: FLMAudioConfig, layer_idx: int):706 super().__init__()707 self.hidden_size = config.hidden_size708 self.self_attn = FLMAUDIO_ATTENTION_CLASSES.get(709 config._attn_implementation, FLMAudioAttention710 )(config=config, layer_idx=layer_idx)711 self.mlp = FLMAudioMLP(config)712 self.input_layernorm = FLMAudioRMSNorm(713 config.hidden_size, eps=config.rms_norm_eps714 )715 self.post_attention_layernorm = FLMAudioRMSNorm(716 config.hidden_size, eps=config.rms_norm_eps717 )718 719 def forward(720 self,721 hidden_states: torch.Tensor,722 attention_mask: Optional[torch.Tensor] = None,723 position_ids: Optional[torch.LongTensor] = None,724 past_key_value: Optional[Tuple[torch.Tensor]] = None,725 output_attentions: Optional[bool] = False,726 use_cache: Optional[bool] = False,727 cache_position: Optional[torch.LongTensor] = None,728 position_embeddings: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, # necessary, but kept here for BC729 **kwargs,730 ) -> Tuple[731 torch.FloatTensor, Optional[Tuple[torch.FloatTensor, torch.FloatTensor]]732 ]:733 """734 Args:735 hidden_states (`torch.FloatTensor`): input to the layer of shape `(batch, seq_len, embed_dim)`736 attention_mask (`torch.FloatTensor`, *optional*):737 attention mask of size `(batch_size, sequence_length)` if flash attention is used or `(batch_size, 1,738 query_sequence_length, key_sequence_length)` if default attention is used.739 output_attentions (`bool`, *optional*):740 Whether or not to return the attentions tensors of all attention layers. See `attentions` under741 returned tensors for more detail.742 use_cache (`bool`, *optional*):743 If set to `True`, `past_key_values` key value states are returned and can be used to speed up decoding744 (see `past_key_values`).745 past_key_value (`Tuple(torch.FloatTensor)`, *optional*): cached past key and value projection states746 """747 if "padding_mask" in kwargs:748 warnings.warn(749 "Passing `padding_mask` is deprecated and will be removed in v4.37. Please make sure use `attention_mask` instead.`"750 )751 752 residual = hidden_states753 754 hidden_states = self.input_layernorm(hidden_states)755 756 # Self Attention757 hidden_states, self_attn_weights, present_key_value = self.self_attn(758 hidden_states=hidden_states,759 attention_mask=attention_mask,760 position_ids=position_ids,761 past_key_value=past_key_value,762 output_attentions=output_attentions,763 use_cache=use_cache,764 cache_position=cache_position,765 position_embeddings=position_embeddings,766 **kwargs,767 )768 hidden_states = residual + hidden_states769 770 # Fully Connected771 residual = hidden_states772 hidden_states = self.post_attention_layernorm(hidden_states)773 hidden_states = self.mlp(hidden_states)774 hidden_states = residual + hidden_states775 776 outputs = (hidden_states,)777 778 if output_attentions:779 outputs += (self_attn_weights,)780 781 if use_cache:782 outputs += (present_key_value,)783 784 return outputs785 786 787FLMAUDIO_START_DOCSTRING = r"""788 This model inherits from [`PreTrainedModel`]. Check the superclass documentation for the generic methods the789 library implements for all its model (such as downloading or saving, resizing the input embeddings, pruning heads790 etc.)791 792 This model is also a PyTorch [torch.nn.Module](https://pytorch.org/docs/stable/nn.html#torch.nn.Module) subclass.793 Use it as a regular PyTorch Module and refer to the PyTorch documentation for all matter related to general usage794 and behavior.795 796 Parameters:797 config ([`FLMAudioConfig`]):798 Model configuration class with all the parameters of the model. Initializing with a config file does not799 load the weights associated with the model, only the configuration. Check out the800 [`~PreTrainedModel.from_pretrained`] method to load the model weights.801"""802 803 804@add_start_docstrings(805 "The bare FLM-Audio Model outputting raw hidden-states without any specific head on top.",806 FLMAUDIO_START_DOCSTRING,807)808class FLMAudioPreTrainedModel(PreTrainedModel):809 config_class = FLMAudioConfig810 base_model_prefix = "model"811 supports_gradient_checkpointing = True812 _no_split_modules = ["FLMAudioDecoderLayer"]813 _skip_keys_device_placement = ["past_key_values"]814 _supports_flash_attn_2 = True815 _supports_sdpa = True816 _supports_cache_class = True817 818 def _init_weights(self, module):819 std = self.config.initializer_range820 if isinstance(module, nn.Linear):821 module.weight.data.normal_(mean=0.0, std=std)822 if module.bias is not None:823 module.bias.data.zero_()824 elif isinstance(module, nn.Embedding):825 module.weight.data.normal_(mean=0.0, std=std)826 if module.padding_idx is not None:827 module.weight.data[module.padding_idx].zero_()828 829 def _setup_cache(830 self, cache_cls, max_batch_size, max_cache_len: Optional[int] = None831 ):832 if (833 self.config._attn_implementation == "flash_attention_2"834 and cache_cls == StaticCache835 ):836 raise ValueError(837 "`static` cache implementation is not compatible with `attn_implementation==flash_attention_2` "838 "make sure to use `sdpa` in the mean time, and open an issue at https://github.com/huggingface/transformers"839 )840 841 for layer in self.model.layers:842 device = layer.input_layernorm.weight.device843 if hasattr(self.config, "_pre_quantization_dtype"):844 dtype = self.config._pre_quantization_dtype845 else:846 dtype = layer.self_attn.o_proj.weight.dtype847 layer.self_attn.past_key_value = cache_cls(848 self.config, max_batch_size, max_cache_len, device=device, dtype=dtype849 )850 851 def _reset_cache(self):852 for layer in self.model.layers:853 layer.self_attn.past_key_value = None854 855 856class MultiModalEmbedding(nn.Module):857 def __init__(self, config):858 super().__init__()859 self.config = config860 self.use_mup = config.use_mup861 self.input_mult = config.input_mult862 self.hidden_size = config.hidden_size863 864 self.vocab_size = config.vocab_size865 self.aud_vocab_size = config.aud_vocab_size866 867 self.aud_channel = config.aud_channel868 869 self.aud_emp_token_id = config.mm_token_info.aud_emp_token_id870 871 self.text_embeddings = nn.Embedding(self.vocab_size, self.hidden_size)872 873 self.aud_listen_embeddings = nn.ModuleList(874 [875 nn.Embedding(self.aud_vocab_size, self.hidden_size)876 for _ in range(self.aud_channel)877 ]878 )879 self.aud_speak_embeddings = nn.ModuleList(880 [881 nn.Embedding(self.aud_vocab_size, self.hidden_size)882 for _ in range(self.aud_channel)883 ]884 )885 886 @staticmethod887 def merge_multichannel_embeddings(888 token_ids, embedding_layer, emp_token_id, embeddings889 ):890 if token_ids is not None and embedding_layer is not None:891 assert token_ids.shape[2] == len(embedding_layer)892 for c in range(token_ids.shape[2]):893 _emb_state = embedding_layer[c](token_ids[:, :, c])894 _emb_state[token_ids[:, :, c] == emp_token_id] = 0.0895 embeddings += _emb_state896 _emb_state = None897 del _emb_state898 return embeddings899 900 def forward(901 self,902 text_ids,903 speak_ids,904 listen_ids,905 ):906 assert text_ids is not None907 embeddings = self.text_embeddings(text_ids)908 mask = ~(text_ids == self.config.pad_token_id)909 910 for aud_chn_idx in range(self.aud_channel):911 aud_speak_embed = self.aud_speak_embeddings[aud_chn_idx](912 speak_ids[..., aud_chn_idx]913 ).squeeze(0)914 aud_listen_embed = self.aud_listen_embeddings[aud_chn_idx](915 listen_ids[..., aud_chn_idx]916 ).squeeze(0)917 embeddings[mask] += aud_speak_embed + aud_listen_embed918 919 if self.use_mup:920 embeddings = embeddings * self.input_mult921 922 return embeddings923 924 925FLMAUDIO_INPUTS_DOCSTRING = r"""926 Args:927 input_ids (`torch.LongTensor` of shape `(batch_size, sequence_length)`):928 Indices of input sequence tokens in the vocabulary. Padding will be ignored by default should you provide929 it.930 931 Indices can be obtained using [`AutoTokenizer`]. See [`PreTrainedTokenizer.encode`] and932 [`PreTrainedTokenizer.__call__`] for details.933 934 [What are input IDs?](../glossary#input-ids)935 attention_mask (`torch.Tensor` of shape `(batch_size, sequence_length)`, *optional*):936 Mask to avoid performing attention on padding token indices. Mask values selected in `[0, 1]`:937 938 - 1 for tokens that are **not masked**,939 - 0 for tokens that are **masked**.940 941 [What are attention masks?](../glossary#attention-mask)942 943 Indices can be obtained using [`AutoTokenizer`]. See [`PreTrainedTokenizer.encode`] and944 [`PreTrainedTokenizer.__call__`] for details.945 946 If `past_key_values` is used, optionally only the last `input_ids` have to be input (see947 `past_key_values`).948 949 If you want to change padding behavior, you should read [`modeling_opt._prepare_decoder_attention_mask`]950 and modify to your needs. See diagram 1 in [the paper](https://arxiv.org/abs/1910.13461) for more951 information on the default strategy.952 953 - 1 indicates the head is **not masked**,954 - 0 indicates the head is **masked**.955 position_ids (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*):956 Indices of positions of each input sequence tokens in the position embeddings. Selected in the range `[0,957 config.n_positions - 1]`.958 959 [What are position IDs?](../glossary#position-ids)960 past_key_values (`Cache` or `tuple(tuple(torch.FloatTensor))`, *optional*):961 Pre-computed hidden-states (key and values in the self-attention blocks and in the cross-attention962 blocks) that can be used to speed up sequential decoding. This typically consists in the `past_key_values`963 returned by the model at a previous stage of decoding, when `use_cache=True` or `config.use_cache=True`.964 965 Two formats are allowed:966 - a [`~cache_utils.Cache`] instance;967 - Tuple of `tuple(torch.FloatTensor)` of length `config.n_layers`, with each tuple having 2 tensors of968 shape `(batch_size, num_heads, sequence_length, embed_size_per_head)`). This is also known as the legacy969 cache format.970 971 The model will output the same cache format that is fed as input. If no `past_key_values` are passed, the972 legacy cache format will be returned.973 974 If `past_key_values` are used, the user can optionally input only the last `input_ids` (those that don't975 have their past key value states given to this model) of shape `(batch_size, 1)` instead of all `input_ids`976 of shape `(batch_size, sequence_length)`.977 inputs_embeds (`torch.FloatTensor` of shape `(batch_size, sequence_length, hidden_size)`, *optional*):978 Optionally, instead of passing `input_ids` you can choose to directly pass an embedded representation. This979 is useful if you want more control over how to convert `input_ids` indices into associated vectors than the980 model's internal embedding lookup matrix.981 use_cache (`bool`, *optional*):982 If set to `True`, `past_key_values` key value states are returned and can be used to speed up decoding (see983 `past_key_values`).984 output_attentions (`bool`, *optional*):985 Whether or not to return the attentions tensors of all attention layers. See `attentions` under returned986 tensors for more detail.987 output_hidden_states (`bool`, *optional*):988 Whether or not to return the hidden states of all layers. See `hidden_states` under returned tensors for989 more detail.990 return_dict (`bool`, *optional*):991 Whether or not to return a [`~utils.ModelOutput`] instead of a plain tuple.992 cache_position (`torch.LongTensor` of shape `(sequence_length)`, *optional*):993 Indices depicting the position of the input sequence tokens in the sequence. Contrarily to `position_ids`,994 this tensor is not affected by padding. It is used to update the cache in the correct position and to infer995 the complete sequence length.996"""997 998 999@add_start_docstrings(1000 "The bare FLM-Audio Model outputting raw hidden-states without any specific head on top.",1001 FLMAUDIO_START_DOCSTRING,1002)1003class FLMAudioModel(FLMAudioPreTrainedModel):1004 """1005 Transformer decoder consisting of *config.num_hidden_layers* layers. Each layer is a [`FLMAudioDecoderLayer`]1006 1007 Args:1008 config: FLMAudioConfig1009 """1010 1011 def __init__(self, config: FLMAudioConfig):1012 super().__init__(config)1013 self.padding_idx = config.pad_token_id1014 self.vocab_size = config.vocab_size1015 1016 self.embed_tokens = MultiModalEmbedding(config)1017 self.layers = nn.ModuleList(1018 [1019 FLMAudioDecoderLayer(config, layer_idx)1020 for layer_idx in range(config.num_hidden_layers)1021 ]1022 )1023 self.norm = FLMAudioRMSNorm(config.hidden_size, eps=config.rms_norm_eps)1024 self.rotary_emb = FLMAudioRotaryEmbedding(config=config)1025 self.gradient_checkpointing = False1026 self.rope_deltas = None # cache rope_deltas here1027 1028 # Initialize weights and apply final processing1029 self.post_init()1030 1031 def get_input_embeddings(self) -> MultiModalEmbedding:1032 return self.embed_tokens1033 1034 def set_input_embeddings(self, value: MultiModalEmbedding):1035 self.embed_tokens = value1036 1037 def get_rope_index(1038 self,1039 input_ids: Optional[torch.LongTensor] = None,1040 second_per_grid_ts: Optional[torch.Tensor] = None,1041 attention_mask: Optional[torch.Tensor] = None,1042 ) -> Tuple[torch.Tensor, torch.Tensor]:1043 1044 mrope_position_deltas = []1045 1046 if attention_mask is not None:1047 position_ids = attention_mask.long().cumsum(-1) - 11048 position_ids.masked_fill_(attention_mask == 0, 1)1049 position_ids = position_ids.unsqueeze(0).expand(3, -1, -1).to(attention_mask.device)1050 max_position_ids = position_ids.max(0, keepdim=False)[0].max(-1, keepdim=True)[0]1051 mrope_position_deltas = max_position_ids + 1 - attention_mask.shape[-1]1052 else:1053 position_ids = (1054 torch.arange(input_ids.shape[1], device=input_ids.device)1055 .view(1, 1, -1)1056 .expand(3, input_ids.shape[0], -1)1057 )1058 mrope_position_deltas = torch.zeros(1059 [input_ids.shape[0], 1],1060 device=input_ids.device,1061 dtype=input_ids.dtype,1062 )1063 1064 return position_ids, mrope_position_deltas1065 1066 1067 @add_start_docstrings_to_model_forward(FLMAUDIO_INPUTS_DOCSTRING)1068 def forward(1069 self,1070 text_ids: torch.LongTensor = None,1071 listen_ids: torch.LongTensor = None,1072 speak_ids: torch.LongTensor = None,1073 attention_mask: Optional[torch.Tensor] = None,1074 position_ids: Optional[torch.LongTensor] = None,1075 past_key_values: Optional[List[torch.FloatTensor]] = None,1076 inputs_embeds: Optional[torch.FloatTensor] = None,1077 use_cache: Optional[bool] = None,1078 output_attentions: Optional[bool] = None,1079 output_hidden_states: Optional[bool] = None,1080 return_dict: Optional[bool] = None,1081 rope_deltas: Optional[torch.LongTensor] = None,1082 cache_position: Optional[torch.LongTensor] = None,1083 second_per_grid_ts: Optional[torch.Tensor] = None,1084 **kwargs,1085 ) -> Union[Tuple, BaseModelOutputWithPast]:1086 output_attentions = (1087 output_attentions1088 if output_attentions is not None1089 else self.config.output_attentions1090 )1091 output_hidden_states = (1092 output_hidden_states1093 if output_hidden_states is not None1094 else self.config.output_hidden_states1095 )1096 use_cache = use_cache if use_cache is not None else self.config.use_cache1097 return_dict = (1098 return_dict if return_dict is not None else self.config.use_return_dict1099 )1100 1101 if (text_ids is None) ^ (inputs_embeds is not None):1102 raise ValueError(1103 "You cannot specify both input_ids and inputs_embeds at the same time, and must specify either one"1104 )1105 1106 if self.gradient_checkpointing and self.training and use_cache:1107 logger.warning_once(1108 "`use_cache=True` is incompatible with gradient checkpointing. Setting `use_cache=False`."1109 )1110 use_cache = False1111 1112 if inputs_embeds is None:1113 inputs_embeds = self.embed_tokens(1114 text_ids,1115 speak_ids,1116 listen_ids,1117 )1118 1119 past_seen_tokens = 01120 if use_cache: # kept for BC (cache positions)1121 if not isinstance(past_key_values, StaticCache):1122 past_key_values = DynamicCache.from_legacy_cache(past_key_values)1123 past_seen_tokens = past_key_values.get_seq_length()1124 1125 if cache_position is None:1126 if isinstance(past_key_values, StaticCache):1127 raise ValueError(1128 "cache_position is a required argument when using StaticCache."1129 )1130 cache_position = torch.arange(1131 past_seen_tokens,1132 past_seen_tokens + inputs_embeds.shape[1],1133 device=inputs_embeds.device,1134 )1135 1136 # if we get 4D attention mask we cannot calculate rope deltas anymore. TODO @raushan fixme1137 if position_ids is None:1138 # calculate RoPE index once per generation in the pre-fill stage only1139 if (1140 (cache_position is not None and cache_position[0] == 0)1141 or self.rope_deltas is None1142 or (past_key_values is None or past_key_values.get_seq_length() == 0)1143 ):1144 position_ids, rope_deltas = self.get_rope_index(1145 text_ids,1146 second_per_grid_ts,1147 attention_mask,1148 )1149 self.rope_deltas = rope_deltas1150 # then use the prev pre-calculated rope-deltas to get the correct position ids1151 else:1152 batch_size, seq_length, _ = inputs_embeds.shape1153 delta = (1154 (cache_position[0] + self.rope_deltas).to(inputs_embeds.device)1155 if cache_position is not None1156 else 01157 )1158 position_ids = torch.arange(seq_length, device=inputs_embeds.device)1159 position_ids = position_ids.view(1, -1).expand(batch_size, -1)1160 if cache_position is not None: # otherwise `deltas` is an int `0`1161 delta = delta.repeat_interleave(batch_size // delta.shape[0], dim=0)1162 position_ids = position_ids.add(delta)1163 position_ids = position_ids.unsqueeze(0).expand(3, -1, -1)1164 1165 causal_mask = self._update_causal_mask(1166 attention_mask, inputs_embeds, cache_position1167 )1168 1169 # embed positions1170 hidden_states = inputs_embeds1171 1172 position_embeddings = self.rotary_emb(hidden_states, position_ids)1173 1174 # decoder layers1175 all_hidden_states = () if output_hidden_states else None1176 all_self_attns = () if output_attentions else None1177 next_decoder_cache = None1178 1179 for decoder_layer in self.layers:1180 if output_hidden_states:1181 all_hidden_states += (hidden_states,)1182 1183 if self.gradient_checkpointing and self.training:1184 layer_outputs = self._gradient_checkpointing_func(1185 decoder_layer.__call__,1186 hidden_states,1187 causal_mask,1188 position_ids,1189 past_key_values,1190 output_attentions,1191 use_cache,1192 cache_position,1193 position_embeddings,1194 )1195 else:1196 layer_outputs = decoder_layer(1197 hidden_states,1198 attention_mask=causal_mask,1199 position_ids=position_ids,1200 past_key_value=past_key_values,