diff --git "a/higgs_audio/model/modeling_higgs_audio.py" "b/higgs_audio/model/modeling_higgs_audio.py" deleted file mode 100644--- "a/higgs_audio/model/modeling_higgs_audio.py" +++ /dev/null @@ -1,2388 +0,0 @@ -"""Higgs-Audio is an end-to-end multimodal model with the capability to understand and generate text / audio.""" - -import torch -import torch.nn as nn -import math -import glob -import functools -import os -from collections import defaultdict, OrderedDict -from dataclasses import dataclass -from enum import Enum -from safetensors.torch import load_file -from typing import Optional, Tuple, Union, List, Dict, Any - -from transformers import AutoTokenizer -from transformers.modeling_outputs import BaseModelOutput -from transformers.models.whisper.modeling_whisper import WhisperEncoderLayer -from transformers.models.llama.modeling_llama import ( - LlamaDecoderLayer, - LlamaRMSNorm, - LlamaRotaryEmbedding, - LLAMA_ATTENTION_CLASSES, - LlamaMLP, - LlamaRMSNorm, -) -from transformers.modeling_attn_mask_utils import AttentionMaskConverter -from transformers.cache_utils import Cache, DynamicCache, StaticCache -from transformers.generation import ( - GenerationMixin, - GenerationConfig, - LogitsProcessorList, - StoppingCriteriaList, -) -from transformers.generation.utils import GenerateNonBeamOutput -from transformers.utils import logging, ModelOutput - -from .common import HiggsAudioPreTrainedModel -from .utils import ( - merge_input_ids_with_audio_features, - count_parameters, -) -from .configuration_higgs_audio import HiggsAudioConfig, HiggsAudioEncoderConfig -from .custom_modules import PartiallyFrozenLinear, PartiallyFrozenEmbedding -from .cuda_graph_runner import CUDAGraphRunner -from .audio_head import HiggsAudioDecoderProjector - -logger = logging.get_logger(__name__) - - -class GenerationMode(Enum): - """Enum for different generation modes in HiggsAudio model.""" - - TEXT = 0 # Text generation mode - AUDIO_INIT = 1 # Audio generation mode initialization - AUDIO_IN_PROGRESS = 2 # Audio generation mode in progress - - -def _whisper_encoder_zero_shape_forward(whisper_encoder, *args, **kwargs): - """The whisper encoder does not support zero-shape tensor by default due to the following implementations - - key_states = self._shape(self.k_proj(current_states), -1, bsz) - - If `bsz` is 0, the "-1" dimension will be ambiguous and triggers error in the shape inference pass. - - See also: https://github.com/huggingface/transformers/blob/30335093276212ce74938bdfd85bfd5df31a668a/src/transformers/models/whisper/modeling_whisper.py#L306-L307 - - This function monkey-patches all `_shape` functions in the whisper encoder's self-attention layers to ensure function supports zero-shape tensor. - - #FIXME!!!! This is a temporary workaround and should be removed once the upstream issue is resolved. - - """ - - global _higgs_flash_attention_forward - - def _patched_shape(tensor: torch.Tensor, seq_len: int, bsz: int, num_heads: int, head_dim: int): - if seq_len == -1: - return tensor.view(bsz, tensor.shape[1], num_heads, head_dim).transpose(1, 2).contiguous() - else: - return tensor.view(bsz, seq_len, num_heads, head_dim).transpose(1, 2).contiguous() - - def _patched_scaled_dot_product_attention( - query, - key, - value, - attn_mask=None, - dropout_p=0.0, - is_causal=False, - scale=None, - enable_gqa=False, - ) -> torch.Tensor: - # IMPORTANT! Implementation here is wrong and is only for the purpose of obtaining the correct attn_weight shape - if enable_gqa: - key = key.repeat_interleave(query.size(-3) // key.size(-3), -3) - value = value.repeat_interleave(query.size(-3) // value.size(-3), -3) - - attn_weight = query @ key.transpose(-2, -1) - return attn_weight @ value - - # Apply monkey-patch - if whisper_encoder.config._attn_implementation != "flash_attention_2": - old_shape_functions = [] - for layer in whisper_encoder.layers: - old_shape_functions.append(getattr(layer.self_attn, "_shape")) - layer.self_attn._shape = functools.partial( - _patched_shape, - num_heads=layer.self_attn.num_heads, - head_dim=layer.self_attn.head_dim, - ) - - original_scaled_dot_product_attention = torch.nn.functional.scaled_dot_product_attention - torch.nn.functional.scaled_dot_product_attention = _patched_scaled_dot_product_attention - - out = whisper_encoder(*args, **kwargs) - torch.nn.functional.scaled_dot_product_attention = original_scaled_dot_product_attention - - # Restore the original shape functions - if whisper_encoder.config._attn_implementation != "flash_attention_2": - for layer, old_shape_function in zip(whisper_encoder.layers, old_shape_functions): - layer.self_attn._shape = old_shape_function - - return out - - -def _prepare_4d_causal_attention_mask_with_cache_position( - attention_mask: torch.Tensor, - sequence_length: int, - target_length: int, - dtype: torch.dtype, - device: torch.device, - min_dtype: float, - cache_position: torch.Tensor, - batch_size: int, -): - """ - Creates a causal 4D mask of shape `(batch_size, 1, query_length, key_value_length)` from a 2D mask of shape - `(batch_size, key_value_length)`, or if the input `attention_mask` is already 4D, do nothing. - - Args: - attention_mask (`torch.Tensor`): - A 2D attention mask of shape `(batch_size, key_value_length)` or a 4D attention mask of shape `(batch_size, 1, query_length, key_value_length)`. - sequence_length (`int`): - The sequence length being processed. - target_length (`int`): - The target length: when generating with static cache, the mask should be as long as the static cache, to account for the 0 padding, the part of the cache that is not filled yet. - dtype (`torch.dtype`): - The dtype to use for the 4D attention mask. - device (`torch.device`): - The device to plcae the 4D attention mask on. - min_dtype (`float`): - The minimum value representable with the dtype `dtype`. - cache_position (`torch.Tensor`): - Indices depicting the position of the input sequence tokens in the sequence. - batch_size (`torch.Tensor`): - Batch size. - """ - if attention_mask is not None and attention_mask.dim() == 4: - # In this case we assume that the mask comes already in inverted form and requires no inversion or slicing. - causal_mask = attention_mask - else: - causal_mask = torch.full( - (sequence_length, target_length), - fill_value=min_dtype, - dtype=dtype, - device=device, - ) - if sequence_length != 1: - causal_mask = torch.triu(causal_mask, diagonal=1) - causal_mask *= torch.arange(target_length, device=device) > cache_position.reshape(-1, 1) - causal_mask = causal_mask[None, None, :, :].expand(batch_size, 1, -1, -1) - if attention_mask is not None: - causal_mask = causal_mask.clone() # copy to contiguous memory for in-place edit - mask_length = attention_mask.shape[-1] - padding_mask = causal_mask[:, :, :, :mask_length] + attention_mask[:, None, None, :] - padding_mask = padding_mask == 0 - causal_mask[:, :, :, :mask_length] = causal_mask[:, :, :, :mask_length].masked_fill( - padding_mask, min_dtype - ) - - return causal_mask - - -class HiggsAudioFeatureProjector(nn.Module): - """Projector that maps audio features extracted by Whisper to hidden state of the text model.""" - - def __init__(self, config: HiggsAudioConfig): - super().__init__() - self.linear = nn.Linear( - config.audio_encoder_config.d_model, - config.text_config.hidden_size, - bias=True, - ) - - def forward(self, audio_features): - hidden_states = self.linear(audio_features) - return hidden_states - - -# Revised on top of transformers.models.qwen2_audio.modeling_qwen2_audio with Qwen2AudioEncoder --> HiggsAudioEncoder -# The code was originally borrowed from WhisperEncoder -class HiggsAudioEncoder(HiggsAudioPreTrainedModel): - """ - Transformer encoder consisting of *config.encoder_layers* self attention layers. Each layer is a - [`WhisperEncoderLayer`]. - - Args: - config: HiggsAudioEncoderConfig - """ - - # Ignore copy - config_class = HiggsAudioEncoderConfig - main_input_name = "input_features" - _no_split_modules = ["WhisperEncoderLayer"] - - def __init__(self, config: HiggsAudioEncoderConfig): - super().__init__(config) - self.dropout = config.dropout - self.layerdrop = config.encoder_layerdrop - - embed_dim = config.d_model - self.num_mel_bins = config.num_mel_bins - self.padding_idx = config.pad_token_id - self.max_source_positions = config.max_source_positions - self.embed_scale = math.sqrt(embed_dim) if config.scale_embedding else 1.0 - - self.conv1 = nn.Conv1d(self.num_mel_bins, embed_dim, kernel_size=3, padding=1) - self.conv2 = nn.Conv1d(embed_dim, embed_dim, kernel_size=3, stride=2, padding=1) - - self.embed_positions = nn.Embedding(self.max_source_positions, embed_dim) - self.embed_positions.requires_grad_(False) - - # Flash Attention 2 does not support zero shape tensor, so we have to use sdpa implementation for the Whisper component. - self.layers = nn.ModuleList([WhisperEncoderLayer(config) for _ in range(config.encoder_layers)]) - self.layer_norm = nn.LayerNorm(config.d_model) - # Ignore copy - self.avg_pooler = nn.AvgPool1d(2, stride=2) - - self.gradient_checkpointing = False - # Initialize weights and apply final processing - self.post_init() - - def _freeze_parameters(self): - for param in self.parameters(): - param.requires_grad = False - self._requires_grad = False - - def get_input_embeddings(self) -> nn.Module: - return self.conv1 - - def set_input_embeddings(self, value: nn.Module): - self.conv1 = value - - def forward( - self, - input_features, - attention_mask=None, - head_mask=None, - output_attentions=None, - output_hidden_states=None, - return_dict=None, - check_seq_length=True, - ): - r""" - Args: - input_features (`torch.LongTensor` of shape `(batch_size, feature_size, sequence_length)`): - Float values of mel features extracted from the raw speech waveform. Raw speech waveform can be - obtained by loading a `.flac` or `.wav` audio file into an array of type `List[float]` or a - `numpy.ndarray`, *e.g.* via the soundfile library (`pip install soundfile`). To prepare the array into - `input_features`, the [`AutoFeatureExtractor`] should be used for extracting the mel features, padding - and conversion into a tensor of type `torch.FloatTensor`. See [`~WhisperFeatureExtractor.__call__`] - attention_mask (`torch.Tensor`)`, *optional*): - HiggsAudio does not support masking of the `input_features`, this argument is preserved for compatibility, - but it is not used. By default the silence in the input log mel spectrogram are ignored. - head_mask (`torch.Tensor` of shape `(encoder_layers, encoder_attention_heads)`, *optional*): - Mask to nullify selected heads of the attention modules. Mask values selected in `[0, 1]`: - - - 1 indicates the head is **not masked**, - - 0 indicates the head is **masked**. - output_attentions (`bool`, *optional*): - Whether or not to return the attentions tensors of all attention layers. See `attentions` under - returned tensors for more detail. - output_hidden_states (`bool`, *optional*): - Whether or not to return the hidden states of all layers. See `hidden_states` under returned tensors - for more detail. - return_dict (`bool`, *optional*): - Whether or not to return a [`~utils.ModelOutput`] instead of a plain tuple. - """ - - expected_seq_length = self.config.max_source_positions * self.conv1.stride[0] * self.conv2.stride[0] - if check_seq_length and (input_features.shape[-1] != expected_seq_length): - raise ValueError( - f"HiggsAudio expects the mel input features to be of length {expected_seq_length}, but found {input_features.shape[-1]}. Make sure to pad the input mel features to {expected_seq_length}." - ) - - output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions - output_hidden_states = ( - output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states - ) - return_dict = return_dict if return_dict is not None else self.config.use_return_dict - - # Ignore copy - input_features = input_features.to(dtype=self.conv1.weight.dtype, device=self.conv1.weight.device) - - inputs_embeds = nn.functional.gelu(self.conv1(input_features)) - inputs_embeds = nn.functional.gelu(self.conv2(inputs_embeds)) - - inputs_embeds = inputs_embeds.permute(0, 2, 1) - embed_pos = self.embed_positions.weight - - hidden_states = inputs_embeds + embed_pos - hidden_states = nn.functional.dropout(hidden_states, p=self.dropout, training=self.training) - - encoder_states = () if output_hidden_states else None - all_attentions = () if output_attentions else None - - # check if head_mask has a correct number of layers specified if desired - if head_mask is not None: - assert head_mask.size()[0] == (len(self.layers)), ( - f"The head_mask should be specified for {len(self.layers)} layers, but it is for {head_mask.size()[0]}." - ) - - for idx, encoder_layer in enumerate(self.layers): - if output_hidden_states: - encoder_states = encoder_states + (hidden_states,) - # add LayerDrop (see https://arxiv.org/abs/1909.11556 for description) - to_drop = False - if self.training: - dropout_probability = torch.rand([]) - if dropout_probability < self.layerdrop: # skip the layer - to_drop = True - - # Ignore copy - if to_drop: - layer_outputs = (None, None) - else: - if self.gradient_checkpointing and self.training: - layer_outputs = self._gradient_checkpointing_func( - encoder_layer.__call__, - hidden_states, - attention_mask, - (head_mask[idx] if head_mask is not None else None), - output_attentions, - ) - else: - layer_outputs = encoder_layer( - hidden_states, - attention_mask, - layer_head_mask=(head_mask[idx] if head_mask is not None else None), - output_attentions=output_attentions, - ) - - hidden_states = layer_outputs[0] - - if output_attentions: - all_attentions = all_attentions + (layer_outputs[1],) - - # Ignore copy - hidden_states = hidden_states.permute(0, 2, 1) - # If the sequence length after average pooling is not divisible by the sequence parallel size, we would duplicate it across the sequence parallel ranks. - # In this case, gradients need to be scaled up because the subsequent scaling up in the function _apply_audio_tower is skipped. - hidden_states = self.avg_pooler(hidden_states) - - hidden_states = hidden_states.permute(0, 2, 1) - - hidden_states = self.layer_norm(hidden_states) - - if output_hidden_states: - encoder_states = encoder_states + (hidden_states,) - - if not return_dict: - return tuple(v for v in [hidden_states, encoder_states, all_attentions] if v is not None) - return BaseModelOutput( - last_hidden_state=hidden_states, - hidden_states=encoder_states, - attentions=all_attentions, - ) - - # Ignore copy - def _get_feat_extract_output_lengths(self, input_lengths: torch.LongTensor): - """ - Computes the output length of the convolutional layers and the output length of the audio encoder - """ - # TODO(sxjscience) Double confirm the formula - input_lengths = (input_lengths - 1) // 2 + 1 - output_lengths = (input_lengths - 2) // 2 + 1 - return input_lengths, output_lengths - - -class HiggsAudioDualFFNDecoderLayer(nn.Module): - """We implement a dual-path FFN decoder layer where the audio tokens and text tokens go through separate FFN layers. - - The audio and text tokens share the text-attention layer, but will be encoded with separate feedforward layers. - In addition, the audio tokens can be configured to go through separate attention layer. - - Following is an illustration: - - t t t a a a t t t - | - | (audio self-attention layer) - v - t t t h'_a h'_a h'_a t t t - | - | (shared attention layer) - v - h_t h_t h_t h_a h_a h_a h_t h_t h_t - | - | (separate text/audio hidden states) - v - [h_t h_t h_t h_t h_t h_t], [h_a, h_a, h_a] - | | - | (separate FFNs) | - v v - [o_t o_t o_t o_t o_t o_t], [o_a, o_a, o_a] - | - | (reorder) - v - o_t o_t o_t o_a o_a o_a o_t o_t o_t - - This has a few advantages: - 1) We are able to use a smaller FFN, or even bypass the FFN for audio tokens. This accelerates the inference speed. - 2) The Audio-FFN introduces more trainable parameters to the model. - This should have the same effect as the mixture-of-expert layer and we may expect better performance due to the scaling law. - 3) We can replace the original FFN in LLMs with the dual-path FFN without changing the model architecture. - - - """ - - def __init__( - self, - config: HiggsAudioConfig, - layer_idx: int, - fast_forward: bool = False, - use_audio_attention: bool = False, - ): - super().__init__() - text_config = config.text_config - self.hidden_size = text_config.hidden_size - self.layer_idx = layer_idx - self.self_attn = LLAMA_ATTENTION_CLASSES[config._attn_implementation](config=text_config, layer_idx=layer_idx) - - self.mlp = LlamaMLP(text_config) - - if not fast_forward: - if use_audio_attention: - self.audio_attn = LLAMA_ATTENTION_CLASSES[config._attn_implementation]( - config=text_config, layer_idx=layer_idx + 1 - ) - self.audio_post_audio_attn_layer_norm = LlamaRMSNorm( - text_config.hidden_size, eps=text_config.rms_norm_eps - ) - - self.audio_mlp = LlamaMLP(text_config) - self.audio_input_layernorm = LlamaRMSNorm(text_config.hidden_size, eps=text_config.rms_norm_eps) - self.audio_post_attention_layernorm = LlamaRMSNorm(text_config.hidden_size, eps=text_config.rms_norm_eps) - - self.use_audio_attention = use_audio_attention - self.fast_forward = fast_forward - if self.fast_forward: - assert not self.use_audio_attention, ( - "We cannot use audio_attention if the layer is marked as fast-forward." - ) - self.input_layernorm = LlamaRMSNorm(text_config.hidden_size, eps=text_config.rms_norm_eps) - self.post_attention_layernorm = LlamaRMSNorm(text_config.hidden_size, eps=text_config.rms_norm_eps) - - def forward( - self, - hidden_states: torch.Tensor, - attention_mask: Optional[torch.Tensor] = None, - audio_attention_mask: Optional[torch.Tensor] = None, - fast_forward_attention_mask: Optional[torch.Tensor] = None, - position_ids: Optional[torch.LongTensor] = None, - audio_out_mask: Optional[torch.BoolTensor] = None, - is_decoding_audio_token: Optional[bool] = None, - past_key_value: Optional[Cache] = None, - output_attentions: Optional[bool] = False, - use_cache: Optional[bool] = False, - cache_position: Optional[torch.LongTensor] = None, - position_embeddings: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, # will become mandatory in v4.46 - is_using_cuda_graph: Optional[bool] = False, - **kwargs, - ): - """ - Args: - hidden_states (`torch.FloatTensor`): input to the layer of shape `(batch, seq_len, embed_dim)` - attention_mask (`torch.FloatTensor`, *optional*): - attention mask of size `(batch_size, sequence_length)` if flash attention is used or `(batch_size, 1, - query_sequence_length, key_sequence_length)` if default attention is used. - position_ids - IDs of positions in the input sequence - audio_out_mask - Mask for identifying the audio tokens. Size (batch_size, sequence_length) - 1 --> location contains audio_out - 0 --> location does not contain audio_out - - When use_cache is True and not in torch compile mode, the audio_out_mask contains audio_out masks for - all tokens up to the current token. That means, it has size (batch_size, sequence_length) while - hidden_states will have size (batch_size, 1). In the torch compile mode, the audio_out_mask will have - size (batch_size, 1). - is_decoding_audio_token - Used in the torch compile mode to determine if the current token is an audio token or not. - past_key_value (`Cache`, *optional*): cached past key and value projection states. We fetch the corresponding cached key/value via the layer_idx. - output_attentions (`bool`, *optional*): - Whether or not to return the attentions tensors of all attention layers. See `attentions` under - returned tensors for more detail. - use_cache (`bool`, *optional*): - If set to `True`, `past_key_values` key value states are returned and can be used to speed up decoding - (see `past_key_values`). - cache_position (`torch.LongTensor` of shape `(sequence_length)`, *optional*): - Indices depicting the position of the input sequence tokens in the sequence - position_embeddings (`Tuple[torch.FloatTensor, torch.FloatTensor]`, *optional*): - Tuple containing the cosine and sine positional embeddings of shape `(batch_size, seq_len, head_dim)`, - with `head_dim` being the embedding dimension of each attention head. - is_using_cuda_graph (`bool`, *optional*): - Indicates whether the model is running by cuda graph. - kwargs (`dict`, *optional*): - Arbitrary kwargs to be ignored, used for FSDP and other methods that injects code - into the model - """ - residual = hidden_states - target_length = hidden_states.shape[1] - use_static_cache = isinstance(past_key_value, StaticCache) - decode_stage = hidden_states.shape[1] == 1 - if is_using_cuda_graph: - assert decode_stage and use_static_cache, ( - "The CUDA graph mode should only be used in the decoding stage with static cache." - ) - - # If we are decoding an audio token and the layer is marked as fast-forward, - # we can skip it. - if is_decoding_audio_token and self.fast_forward: - return (hidden_states,) - - has_audio_out = audio_out_mask is not None and audio_out_mask.shape[0] > 0 - - audio_out_mask_sq = audio_out_mask - - if self.fast_forward and has_audio_out: - original_hidden_states = hidden_states.clone() - min_dtype = torch.finfo(hidden_states.dtype).min - if attention_mask is None: - attention_mask = ~audio_out_mask - - if self.self_attn.config._attn_implementation != "flash_attention_2": - sequence_length = audio_out_mask.shape[1] - attention_mask = _prepare_4d_causal_attention_mask_with_cache_position( - attention_mask=attention_mask, - sequence_length=sequence_length, - target_length=sequence_length, - dtype=hidden_states.dtype, - min_dtype=min_dtype, - device=hidden_states.device, - cache_position=cache_position, - batch_size=hidden_states.shape[0], - ) - if use_cache: - attention_mask = attention_mask[:, :, -target_length:, :] - elif len(attention_mask.shape) == 2: - # Attention mask has shape (batch_size, sequence_length) - # We should be using flash attention 2 - attention_mask = attention_mask * ~audio_out_mask - elif len(attention_mask.shape) == 4: - # When using static cache, the attention mask was already preprocessed in the previous layer - if use_static_cache: - attention_mask = fast_forward_attention_mask - else: - if use_cache: - # Attention mask has shape (batch_size, 1, query_length, key_length) - # In addition, the attention mask should be inverted, that means "1" (attend_to) --> "0", and "0" --> minimal dtype value. - attention_mask = attention_mask.masked_fill( - audio_out_mask[:, -target_length:].reshape(audio_out_mask.shape[0], 1, target_length, 1) - | audio_out_mask.reshape(audio_out_mask.shape[0], 1, 1, audio_out_mask.shape[1]), - min_dtype, - ) - else: - attention_mask = attention_mask.masked_fill( - audio_out_mask.reshape(audio_out_mask.shape[0], 1, audio_out_mask.shape[1], 1) - | audio_out_mask.reshape(audio_out_mask.shape[0], 1, 1, audio_out_mask.shape[1]), - min_dtype, - ) - else: - raise NotImplementedError(f"Unsupported attention_mask format, attention_mask={attention_mask}") - - if ( - self.self_attn.config._attn_implementation == "sdpa" - and attention_mask is not None - and attention_mask.device.type == "cuda" - and not output_attentions - ): - # Attend to all tokens in fully masked rows in the causal_mask, for example the relevant first rows when - # using left padding. This is required by F.scaled_dot_product_attention memory-efficient attention path. - # Details: https://github.com/pytorch/pytorch/issues/110213 - attention_mask = AttentionMaskConverter._unmask_unattended(attention_mask, min_dtype) - - if has_audio_out and not self.fast_forward: - # Apply separate layernorm layers for audio tokens and text tokens - if use_cache: - hidden_states = torch.where( - audio_out_mask_sq[:, -target_length:].unsqueeze(-1), - self.audio_input_layernorm(hidden_states), - self.input_layernorm(hidden_states), - ) - else: - hidden_states = torch.where( - audio_out_mask_sq.unsqueeze(-1), - self.audio_input_layernorm(hidden_states), - self.input_layernorm(hidden_states), - ) - else: - hidden_states = self.input_layernorm(hidden_states) - - # Audio Attention - if self.use_audio_attention and has_audio_out: - if use_static_cache: - assert audio_attention_mask is not None, ( - "audio_attention_mask should not be None when using static cache." - ) - - if audio_attention_mask is None: - no_audio_out_mask = (~audio_out_mask)[:, -target_length:].reshape( - audio_out_mask.shape[0], 1, target_length, 1 - ) | (~audio_out_mask).reshape(audio_out_mask.shape[0], 1, 1, audio_out_mask.shape[1]) - min_dtype = torch.finfo(hidden_states.dtype).min - - if attention_mask is None: - audio_attention_mask = audio_out_mask - - if self.audio_attn.config._attn_implementation != "flash_attention_2": - sequence_length = audio_out_mask.shape[1] - audio_attention_mask = _prepare_4d_causal_attention_mask_with_cache_position( - attention_mask=audio_attention_mask, - sequence_length=sequence_length, - target_length=sequence_length, - dtype=hidden_states.dtype, - min_dtype=min_dtype, - device=hidden_states.device, - cache_position=cache_position, - batch_size=hidden_states.shape[0], - ) - if use_cache: - audio_attention_mask = audio_attention_mask[:, :, -target_length:, :] - audio_attention_mask = audio_attention_mask.masked_fill(no_audio_out_mask, min_dtype) - elif len(attention_mask.shape) == 2: - # Attention mask has shape (batch_size, sequence_length) - audio_attention_mask = attention_mask * audio_out_mask - elif len(attention_mask.shape) == 4: - # Attention mask has shape (batch_size, 1, query_length, key_length) - # In addition, the attention mask should be inverted. This means "1" (attend_to) --> "0", and "0" --> minimal dtype value. - audio_attention_mask = attention_mask.masked_fill(no_audio_out_mask, min_dtype) - else: - raise NotImplementedError(f"Unsupported attention_mask format, attention_mask={attention_mask}") - - if ( - self.audio_attn.config._attn_implementation == "sdpa" - and audio_attention_mask is not None - and audio_attention_mask.device.type == "cuda" - and not output_attentions - ): - # Attend to all tokens in fully masked rows in the causal_mask, for example the relevant first rows when - # using left padding. This is required by F.scaled_dot_product_attention memory-efficient attention path. - # Details: https://github.com/pytorch/pytorch/issues/110213 - audio_attention_mask = AttentionMaskConverter._unmask_unattended(audio_attention_mask, min_dtype) - - audio_attention_mask = audio_attention_mask.contiguous() - - audio_hidden_states, audio_self_attn_weights, audio_present_key_value = self.audio_attn( - hidden_states=hidden_states, - attention_mask=audio_attention_mask, - position_ids=position_ids, - past_key_value=past_key_value, - output_attentions=output_attentions, - use_cache=use_cache, - cache_position=cache_position, - position_embeddings=position_embeddings, - **kwargs, - ) - audio_hidden_states = residual + audio_hidden_states - if use_cache: - residual = torch.where( - audio_out_mask_sq[:, -target_length:].unsqueeze(-1), - audio_hidden_states, - residual, - ) - else: - residual = torch.where(audio_out_mask_sq.unsqueeze(-1), audio_hidden_states, residual) - audio_hidden_states = self.audio_post_audio_attn_layer_norm(audio_hidden_states) - if use_cache: - hidden_states = torch.where( - audio_out_mask_sq[:, -target_length:].unsqueeze(-1), - audio_hidden_states, - hidden_states, - ) - else: - hidden_states = torch.where(audio_out_mask_sq.unsqueeze(-1), audio_hidden_states, hidden_states) - - # Text Attention - hidden_states, self_attn_weights, present_key_value = self.self_attn( - hidden_states=hidden_states, - attention_mask=attention_mask, - position_ids=position_ids, - past_key_value=past_key_value, - output_attentions=output_attentions, - use_cache=use_cache, - cache_position=cache_position, - position_embeddings=position_embeddings, - **kwargs, - ) - hidden_states = residual + hidden_states - - # Apply Dual-path FFN - residual = hidden_states - - if has_audio_out and not self.fast_forward: - if use_cache: - real_audio_out_mask = audio_out_mask_sq[:, -target_length:] - else: - real_audio_out_mask = audio_out_mask_sq - - # Make whole graph in decode stage - if decode_stage and is_using_cuda_graph: - assert is_decoding_audio_token is not None, ( - "is_decoding_audio_token should be present in the decoding stage." - ) - if is_decoding_audio_token: - hidden_states = self.audio_post_attention_layernorm(hidden_states) - hidden_states = self.audio_mlp(hidden_states) - else: - hidden_states = self.post_attention_layernorm(hidden_states) - hidden_states = self.mlp(hidden_states) - residual = residual + hidden_states - else: - text_hidden_states = self.post_attention_layernorm(hidden_states[~real_audio_out_mask]) - audio_hidden_states = self.audio_post_attention_layernorm(hidden_states[real_audio_out_mask]) - - text_hidden_states = self.mlp(text_hidden_states) - residual[~real_audio_out_mask] += text_hidden_states - - audio_hidden_states = self.audio_mlp(audio_hidden_states) - residual[real_audio_out_mask] += audio_hidden_states - - hidden_states = residual - else: - hidden_states = self.post_attention_layernorm(hidden_states) - hidden_states = self.mlp(hidden_states) - hidden_states = residual + hidden_states - - if self.fast_forward and has_audio_out: - if use_cache: - hidden_states = torch.where( - audio_out_mask_sq[:, -target_length:].unsqueeze(-1), - original_hidden_states, - hidden_states, - ) - else: - hidden_states = torch.where( - audio_out_mask_sq.unsqueeze(-1), - original_hidden_states, - hidden_states, - ) - - outputs = (hidden_states,) - - if output_attentions: - if self.use_audio_attention: - # The returned attn weights have shape (batch_size, num_heads + num_audio_attn_heads, seq_length, seq_length) - outputs += (torch.concat([self_attn_weights, audio_self_attn_weights], dim=1),) - else: - # The returned attn weights have shape (batch_size, num_heads, seq_length, seq_length) - outputs += (self_attn_weights,) - - if use_cache: - outputs += (present_key_value,) - - return outputs - - -@dataclass -class HiggsAudioModelOutputWithPast(ModelOutput): - loss: Optional[torch.FloatTensor] = None - llm_loss: Optional[torch.FloatTensor] = None - audio_loss: Optional[torch.FloatTensor] = None - codebook_losses: Optional[torch.FloatTensor] = None - logits: Optional[torch.FloatTensor] = None - expanded_input_ids: Optional[torch.LongTensor] = None - expanded_labels: Optional[torch.LongTensor] = None - audio_in_mask: Optional[torch.BoolTensor] = None - audio_in_discrete_codes_mask: Optional[torch.BoolTensor] = None - audio_out_mask: Optional[torch.BoolTensor] = None - attention_mask: Optional[torch.BoolTensor] = None - audio_logits: Optional[torch.FloatTensor] = None - past_key_values: Optional[Cache] = None - hidden_states: Optional[Tuple[torch.FloatTensor, ...]] = None - audio_hidden_states: Optional[Tuple[torch.FloatTensor, ...]] = None - attentions: Optional[Tuple[torch.FloatTensor, ...]] = None - - -@dataclass -class HiggsAudioGenerationOutput(ModelOutput): - """ - Outputs of HiggsAudio generation models, when using non-beam methods. - - Args: - sequences (`torch.LongTensor` of shape `(batch_size, sequence_length)`): - The generated sequences. The second dimension (sequence_length) is either equal to `max_length` or shorter - if all batches finished early due to the `eos_token_id`. - audio_sequences (`tuple(torch.LongTensor)` *optional*): - The generated discrete audio codes. These codes can be used to fill-in related locations of <|AUDIO_OUT|> at input sequences. - scores (`tuple(torch.FloatTensor)` *optional*, returned when `output_scores=True`): - Processed prediction scores of the language modeling head (scores for each vocabulary token before SoftMax) - at each generation step. Tuple of `torch.FloatTensor` with up to `max_new_tokens` elements (one element for - each generated token). - If the generated token is a text token, the tensor will have shape `(batch_size, config.vocab_size)`. - If the generated token is an audio token, the tensor will have shape `(config.audio_num_codebooks, self.audio_codebook_size)` - logits (`tuple(torch.FloatTensor)` *optional*, returned when `output_logits=True`): - Unprocessed prediction scores of the language modeling head or the audio head (scores for each vocabulary token before SoftMax) - at each generation step. Tuple of `torch.FloatTensor` with up to `max_new_tokens` elements (one element for - each generated token). - If the generated token is a text token, the tensor will have shape `(batch_size, config.vocab_size)`. - If the generated token is an audio token, the tensor will have shape `(config.audio_num_codebooks, self.audio_codebook_size)` - attentions (`tuple(tuple(torch.FloatTensor))`, *optional*, returned when `output_attentions=True`): - Tuple (one element for each generated token) of tuples (one element for each layer of the decoder) of - `torch.FloatTensor` of shape `(batch_size, num_heads, generated_length, sequence_length)`. - hidden_states (`tuple(tuple(torch.FloatTensor))`, *optional*, returned when `output_hidden_states=True`): - Tuple (one element for each generated token) of tuples (one element for each layer of the decoder) of - `torch.FloatTensor` of shape `(batch_size, generated_length, hidden_size)`. - past_key_values (`tuple(tuple(torch.FloatTensor)))`, *optional*, returned when `use_cache=True`): - Returns the model cache, used to speed up decoding. Different models have a different cache format, check - the model's documentation. Usually, a [`~cache_utils.Cache`] instance. - """ - - sequences: torch.LongTensor = None - audio_sequences: Optional[List[torch.LongTensor]] = None - scores: Optional[Tuple[torch.FloatTensor]] = None - logits: Optional[Tuple[torch.FloatTensor]] = None - attentions: Optional[Tuple[Tuple[torch.FloatTensor]]] = None - hidden_states: Optional[Tuple[Tuple[torch.FloatTensor]]] = None - past_key_values: Optional[Tuple[Tuple[Tuple[torch.FloatTensor]]]] = None - - -class HiggsAudioModel(HiggsAudioPreTrainedModel, GenerationMixin): - """Higgs-Audio is an end-to-end multimodal model with the capability to understand and generate text / audio. - - Consider the following example for mixed text/audio understanding / generation: - - - input_tokens: <|audio_bos|>[AUDIO]<|audio_eos|><|audio_bos|>[AUDIO]<|audio_eos|> - - input_tokens: <|audio_bos|>[AUDIO]<|audio_eos|><|audio_out_bos|>[AUDIO_OUT]<|audio_eos|> - - We will fill [AUDIO] with the audio features extracted by Whisper and fill [AUDIO_OUT] with the audio tokens. - - Consider the following example for mixed text/audio generation: - - text: <|audio_out_bos|> MASK MASK MASK MASK MASK <|audio_eos|> [text_token1] - audio: MASK <|audio_stream_bos|> [audio_token1] [audio_token2] [audio_token3] <|audio_stream_eos|> MASK MASK - token_type: 0 1 1 1 1 1 0 0 - - """ - - _supports_cache_class = True - _supports_static_cache = True - - def __init__(self, config: HiggsAudioConfig): - super().__init__(config) - self.padding_idx = config.pad_token_id - self.audio_in_token_idx = config.audio_in_token_idx - self.audio_out_token_idx = config.audio_out_token_idx - self.audio_out_bos_token_id = config.audio_out_bos_token_id if "audio_out_bos_token_id" in config else None - self.audio_eos_token_id = config.audio_eos_token_id if "audio_eos_token_id" in config else None - self.vocab_size = config.text_config.vocab_size - self.audio_num_codebooks = config.audio_num_codebooks - self.use_delay_pattern = config.use_delay_pattern - self.use_audio_out_embed_projector = config.use_audio_out_embed_projector - self.use_audio_out_self_attention = config.use_audio_out_self_attention - - self.embed_tokens = nn.Embedding(self.vocab_size, config.text_config.hidden_size, self.padding_idx) - - if config.audio_adapter_type == "dual_ffn": - layer_idx = 0 - layers = [] - for j in range(config.text_config.num_hidden_layers): - if j in config.audio_dual_ffn_layers: - layers.append( - HiggsAudioDualFFNDecoderLayer( - config, - layer_idx, - use_audio_attention=self.use_audio_out_self_attention, - ) - ) - layer_idx += 2 if self.use_audio_out_self_attention else 1 - else: - layers.append(LlamaDecoderLayer(config.text_config, layer_idx)) - layer_idx += 1 - self.layers = nn.ModuleList(layers) - elif config.audio_adapter_type == "dual_ffn_fast_forward": - layer_idx = 0 - layers = [] - for j in range(config.text_config.num_hidden_layers): - if j in config.audio_dual_ffn_layers: - layers.append( - HiggsAudioDualFFNDecoderLayer( - config, - layer_idx, - fast_forward=False, - use_audio_attention=self.use_audio_out_self_attention, - ) - ) - layer_idx += 2 if self.use_audio_out_self_attention else 1 - else: - layers.append( - HiggsAudioDualFFNDecoderLayer( - config, - layer_idx, - fast_forward=True, - use_audio_attention=False, - ) - ) - layer_idx += 1 - self.layers = nn.ModuleList(layers) - elif config.audio_adapter_type == "stack": - self.layers = nn.ModuleList( - [ - LlamaDecoderLayer(config.text_config, layer_idx) - for layer_idx in range(config.text_config.num_hidden_layers) - ] - ) - layer_idx = config.text_config.num_hidden_layers - else: - raise NotImplementedError(f"Audio adapter type {config.audio_adapter_type} not implemented.") - - self.num_activation_checkpointing_layers = len(self.layers) - - self.decode_graph_runners = defaultdict(dict[bool, CUDAGraphRunner]) - self.norm = LlamaRMSNorm(config.text_config.hidden_size, eps=config.text_config.rms_norm_eps) - self.rotary_emb = LlamaRotaryEmbedding(config=config.text_config) - - if not config.skip_audio_tower: - self.audio_tower = HiggsAudioEncoder(config.audio_encoder_config) - self.audio_encoder_proj = HiggsAudioFeatureProjector(config) - else: - self.audio_tower = None - self.audio_encoder_proj = None - self.audio_decoder_proj = HiggsAudioDecoderProjector(config, layer_idx=layer_idx) - self.audio_codebook_size = ( - config.audio_codebook_size + 2 - ) # We add 1 for the audio_stream_bos token and 1 for the audio_stream_eos token - - if config.use_audio_out_embed_projector: - self.audio_out_embed_projector = nn.Linear( - config.text_config.hidden_size, - config.text_config.hidden_size, - bias=False, - ) - - self.audio_codebook_embeddings = nn.Embedding( - config.audio_num_codebooks * self.audio_codebook_size, - config.text_config.hidden_size, - ) - - self.audio_codebook_weights = ( - torch.ones(config.audio_num_codebooks) / config.audio_num_codebooks - ) # default to equal weights - self.post_init() - - def set_num_activation_checkpointing_layers(self, num_layers): - self.num_activation_checkpointing_layers = num_layers - - def set_delay_pattern(self): - self.config.use_delay_pattern = True - self.use_delay_pattern = True - - def set_audio_special_tokens(self, tokenizer: AutoTokenizer): - self.audio_out_bos_token_id = tokenizer.convert_tokens_to_ids("<|audio_out_bos|>") - self.audio_eos_token_id = tokenizer.convert_tokens_to_ids("<|audio_eos|>") - - def _embed_audio_ids(self, audio_ids): - """Embed the audio ids - - Args: - audio_ids: torch.LongTensor of shape (num_codebooks, audio_in_total_length) - - Returns: - audio_embed: torch.LongTensor of shape (audio_in_total_length, hidden_size) - """ - codebook_shift = ( - torch.arange(self.config.audio_num_codebooks, device=audio_ids.device) * self.audio_codebook_size - ) - audio_embed = self.audio_codebook_embeddings(audio_ids + codebook_shift.unsqueeze(-1)) - if self.config.audio_embed_avg: - audio_embed = torch.mean(audio_embed, dim=0) - else: - audio_embed = torch.sum(audio_embed, dim=0) - if self.use_audio_out_embed_projector: - audio_embed = self.audio_out_embed_projector(audio_embed) - return audio_embed - - def _apply_audio_tower(self, audio_features, audio_feature_attention_mask): - """Apply the audio tower to the audio features""" - - if audio_features.shape[0] == 0: - if torch.is_grad_enabled(): - # FIXME!!!!!!!! - # This is a hack to ensure that the forward+backward pass of audio_tower and audio_encoder_proj get triggered. - # The monkey patch won't overwrite the backward pass of nn.Module. - audio_outputs = _whisper_encoder_zero_shape_forward( - self.audio_tower, - audio_features, - attention_mask=None, - check_seq_length=False, - ) - selected_audio_feature = audio_outputs.last_hidden_state - audio_features_embed = self.audio_encoder_proj(selected_audio_feature) - audio_feat_out_lengths = None - return audio_features_embed, audio_feat_out_lengths - else: - return None, None - - audio_feat_lengths, audio_feat_out_lengths = self.audio_tower._get_feat_extract_output_lengths( - audio_feature_attention_mask.sum(-1) - ) - batch_size, _, max_mel_seq_len = audio_features.shape - max_seq_len = (max_mel_seq_len - 1) // 2 + 1 - # Create a sequence tensor of shape (batch_size, max_seq_len) - seq_range = ( - torch.arange( - 0, - max_seq_len, - dtype=audio_feat_lengths.dtype, - device=audio_feat_lengths.device, - ) - .unsqueeze(0) - .expand(batch_size, max_seq_len) - ) - lengths_expand = audio_feat_lengths.unsqueeze(1).expand(batch_size, max_seq_len) - # Create mask - padding_mask = seq_range < lengths_expand - - if self.config._attn_implementation != "flash_attention_2": - audio_attention_mask = padding_mask.view(batch_size, 1, 1, max_seq_len).expand( - batch_size, 1, max_seq_len, max_seq_len - ) - else: - audio_attention_mask = padding_mask - - audio_outputs = self.audio_tower(audio_features, attention_mask=audio_attention_mask) - selected_audio_feature = audio_outputs.last_hidden_state - audio_features_embed = self.audio_encoder_proj(selected_audio_feature) - - return audio_features_embed, audio_feat_out_lengths - - def _update_causal_mask( - self, - attention_mask: torch.Tensor, - input_tensor: torch.Tensor, - cache_position: torch.Tensor, - past_key_values: Cache, - output_attentions: bool, - ): - if self.config._attn_implementation == "flash_attention_2": - if attention_mask is not None and 0.0 in attention_mask: - return attention_mask - return None - - # For SDPA, when possible, we will rely on its `is_causal` argument instead of its `attn_mask` argument, in - # order to dispatch on Flash Attention 2. This feature is not compatible with static cache, as SDPA will fail - # to infer the attention mask. - past_seen_tokens = past_key_values.get_seq_length() if past_key_values is not None else 0 - using_static_cache = isinstance(past_key_values, StaticCache) - - # When output attentions is True, sdpa implementation's forward method calls the eager implementation's forward - if self.config._attn_implementation == "sdpa" and not using_static_cache and not output_attentions: - if AttentionMaskConverter._ignore_causal_mask_sdpa( - attention_mask, - inputs_embeds=input_tensor, - past_key_values_length=past_seen_tokens, - is_training=self.training, - ): - return None - - dtype, device = input_tensor.dtype, input_tensor.device - min_dtype = torch.finfo(dtype).min - sequence_length = input_tensor.shape[1] - if using_static_cache: - target_length = past_key_values.get_max_length() - else: - target_length = ( - attention_mask.shape[-1] - if isinstance(attention_mask, torch.Tensor) - else past_seen_tokens + sequence_length + 1 - ) - - # In case the provided `attention` mask is 2D, we generate a causal mask here (4D). - causal_mask = _prepare_4d_causal_attention_mask_with_cache_position( - attention_mask, - sequence_length=sequence_length, - target_length=target_length, - dtype=dtype, - device=device, - min_dtype=min_dtype, - cache_position=cache_position, - batch_size=input_tensor.shape[0], - ) - - if ( - self.config._attn_implementation == "sdpa" - and attention_mask is not None - and attention_mask.device.type == "cuda" - and not output_attentions - ): - # Attend to all tokens in fully masked rows in the causal_mask, for example the relevant first rows when - # using left padding. This is required by F.scaled_dot_product_attention memory-efficient attention path. - # Details: https://github.com/pytorch/pytorch/issues/110213 - causal_mask = AttentionMaskConverter._unmask_unattended(causal_mask, min_dtype) - - return causal_mask - - def _prepare_all_static_kv_cache_masks(self, hidden_states, attention_mask, audio_out_mask, past_key_values): - target_length = hidden_states.shape[1] - cur_pos = audio_out_mask.shape[1] - min_dtype = torch.finfo(hidden_states.dtype).min - assert len(attention_mask.shape) == 4, "Only support SDPA for now" - kv_cache_len = past_key_values.get_max_cache_shape() - audio_out_mask_padded = torch.nn.functional.pad(audio_out_mask, (0, kv_cache_len - cur_pos), value=True) - fast_forward_attention_mask = attention_mask.masked_fill( - audio_out_mask_padded[:, audio_out_mask.shape[1] - target_length : audio_out_mask.shape[1]].reshape( - audio_out_mask_padded.shape[0], 1, target_length, 1 - ) - | audio_out_mask_padded.reshape(audio_out_mask_padded.shape[0], 1, 1, audio_out_mask_padded.shape[1]), - min_dtype, - ) - - no_audio_out_mask = ~audio_out_mask - no_audio_out_mask = torch.nn.functional.pad( - no_audio_out_mask, (0, kv_cache_len - audio_out_mask.shape[1]), value=False - ) - no_audio_out_mask = no_audio_out_mask[ - :, audio_out_mask.shape[1] - target_length : audio_out_mask.shape[1] - ].reshape(audio_out_mask.shape[0], 1, target_length, 1) | no_audio_out_mask.reshape( - audio_out_mask.shape[0], 1, 1, kv_cache_len - ) - audio_attention_mask = attention_mask.masked_fill(no_audio_out_mask, min_dtype) - return fast_forward_attention_mask, audio_attention_mask - - def _forward_core( - self, - hidden_states: torch.Tensor, - causal_mask: torch.Tensor, - position_ids: torch.Tensor, - audio_discrete_codes_mask: torch.Tensor, - cache_position: torch.Tensor, - past_key_values: Optional[Union[Cache, List[torch.FloatTensor]]], - use_cache: bool, - audio_attention_mask: torch.Tensor, - fast_forward_attention_mask: torch.Tensor, - output_attentions: bool, - output_hidden_states: bool, - is_decoding_audio_token: Optional[bool] = None, - is_using_cuda_graph: Optional[bool] = False, - ): - # create position embeddings to be shared across the decoder layers - # When past_key_values is passed in, we need to offset the position ids when calculating the position embeddings. - # Therefore, cache_position is used. - position_id_offset = cache_position[0] if use_cache else 0 - position_embeddings = self.rotary_emb(hidden_states, position_ids + position_id_offset) - - # decoder layers - all_hidden_states = () if output_hidden_states else None - all_self_attns = () if output_attentions else None - - for decoder_layer in self.layers: - if output_hidden_states: - all_hidden_states += (hidden_states,) - if isinstance(decoder_layer, HiggsAudioDualFFNDecoderLayer): - layer_outputs = decoder_layer( - hidden_states, - attention_mask=causal_mask, - audio_attention_mask=audio_attention_mask, - fast_forward_attention_mask=fast_forward_attention_mask, - position_ids=position_ids, - audio_out_mask=audio_discrete_codes_mask, - is_decoding_audio_token=is_decoding_audio_token, - past_key_value=past_key_values, - output_attentions=output_attentions, - use_cache=use_cache, - cache_position=cache_position, - position_embeddings=position_embeddings, - is_using_cuda_graph=is_using_cuda_graph, - ) - else: - layer_outputs = decoder_layer( - hidden_states, - attention_mask=causal_mask, - position_ids=position_ids, - past_key_value=past_key_values, - output_attentions=output_attentions, - use_cache=use_cache, - cache_position=cache_position, - position_embeddings=position_embeddings, - ) - - hidden_states = layer_outputs[0] - - if output_attentions: - all_self_attns += (layer_outputs[1],) - - return hidden_states, all_hidden_states, all_self_attns - - def forward( - self, - input_ids: Optional[torch.LongTensor] = None, - inputs_embeds: Optional[torch.FloatTensor] = None, - attention_mask: Optional[torch.BoolTensor] = None, - audio_features: Optional[torch.FloatTensor] = None, - audio_feature_attention_mask: Optional[torch.BoolTensor] = None, - audio_in_ids: Optional[torch.LongTensor] = None, - audio_in_ids_start: Optional[torch.LongTensor] = None, - audio_out_ids: Optional[torch.LongTensor] = None, - audio_out_ids_start: Optional[torch.LongTensor] = None, - audio_out_ids_start_group_loc: Optional[torch.LongTensor] = None, - label_ids: Optional[torch.LongTensor] = None, - label_audio_ids: Optional[torch.LongTensor] = None, - past_key_values: Optional[Union[Cache, List[torch.FloatTensor]]] = None, - use_cache: Optional[bool] = None, - output_attentions: Optional[bool] = None, - output_hidden_states: Optional[bool] = None, - output_audio_hidden_states: Optional[bool] = False, - return_dict: Optional[bool] = None, - cache_position: Optional[torch.LongTensor] = None, - cache_audio_discrete_codes_mask: Optional[torch.LongTensor] = None, - past_key_values_buckets: Optional[OrderedDict[int, Cache]] = None, - reward: Optional[torch.FloatTensor] = None, - ): - """Forward pass for the Higgs-Audio model. - - Args: - input_ids (:obj:`torch.LongTensor`): - The input ids of the prompt. It will have shape (bsz, seq_len). - When use_cache is enabled, the input_ids will have - shape (bsz, 1) for incremental decode or None - inputs_embeds: - Input embeddings. This flag won't be used. - attention_mask (:obj:`torch.LongTensor`): - The attention mask of the prompt. It will have shape (bsz, seq_len). - audio_features (:obj:`torch.FloatTensor`): - The audio features extracted by Whisper. It will have shape (num_audio_in, feature_dim, max_mel_seq_len). - audio_feature_attention_mask (:obj:`torch.LongTensor`): - The attention mask of the audio features. It will have shape (num_audio_in, max_mel_seq_len). - audio_in_ids (:obj:`torch.LongTensor`): - The discretized audio tokens. It will have shape (num_codebooks, audio_in_total_length). - audio_in_ids_start (:obj:`torch.LongTensor`): - The start indices for each audio in audio_in_ids. It will have shape (num_audio_in,) - audio_out_ids (:obj:`torch.LongTensor`): - The discretized audio tokens. It will have shape (num_codebooks, audio_out_total_length). - audio_out_ids_start (:obj:`torch.LongTensor`): - The start indices for each audio in audio_out_ids. It will have shape (num_audio_out,) - audio_out_ids_start_group_loc (:obj:`torch.LongTensor`): - The sample indices in a batch that map to each element in the audio_out_ids_start. It will have shape (num_audio_out,) - label_text_ids (:obj:`torch.LongTensor`): - The labels of the prompt. It will have shape (bsz, seq_len). - label_audio_ids (:obj:`torch.LongTensor`): - The labels of the audio tokens. It will have the same shape as audio_out_ids, i.e., (num_codebooks, audio_out_total_length) - past_key_values (:obj:`Tuple`): - Tuple of past key values. - use_cache (:obj:`bool`): - Whether to use cache. - output_attentions (:obj:`bool`): - Whether to output attentions. - output_hidden_states (:obj:`bool`): - Whether to output hidden states. - output_audio_hidden_states (:obj:`bool`): - Whether to output audio hidden states. - return_dict (:obj:`bool`): - Whether to return a dictionary. - cache_position (:obj:`torch.LongTensor`): - The position of the cache. - cache_audio_discrete_codes_mask (:obj:`torch.LongTensor`): - The cached audio discrete codes mask. It will only be used when use_cache is turned on. - past_key_values_buckets (:obj:`OrderedDict`): - The buckets of past key values. - """ - target_device = input_ids.device - - # not used - del inputs_embeds - - if audio_features is not None: - audio_features = audio_features.to(target_device) - audio_feature_attention_mask = audio_feature_attention_mask.to(target_device) - - # 1. Extract the input embeddings - inputs_embeds = self.embed_tokens(input_ids) - - # 2. Extract audio embeddings - if self.config.skip_audio_tower: - audio_features_embed = audio_features_length = None - else: - audio_features_embed, audio_features_length = self._apply_audio_tower( - audio_features, audio_feature_attention_mask - ) - - if self.config.encode_audio_in_tokens: - if audio_in_ids is not None and audio_in_ids.shape[-1] > 0: - audio_in_ids = audio_in_ids.to(target_device) - else: - audio_in_ids = torch.zeros( - (self.audio_num_codebooks, 0), - device=target_device, - dtype=torch.long, - ) - audio_in_embed = self._embed_audio_ids(audio_in_ids) - else: - audio_in_embed = None - - if audio_out_ids is not None and audio_out_ids.shape[-1] > 0: - audio_out_ids = audio_out_ids.to(target_device) - else: - audio_out_ids = torch.zeros((self.audio_num_codebooks, 0), device=target_device, dtype=torch.long) - audio_out_embed = self._embed_audio_ids(audio_out_ids) - - # 3. Merge text, audio-in embeddings, and audio-out embeddings - - # use_cache is turned on during inference time, we should set round_to to 1 to avoid extra padding in the end. - round_to = 1 if use_cache else 8 - left_padding = True if use_cache or input_ids.shape[0] == 1 else False - ( - inputs_embeds, - attention_mask, - labels, - position_ids, - input_ids, - audio_in_mask, - audio_in_discrete_codes_mask, - audio_out_mask, - ) = merge_input_ids_with_audio_features( - audio_features_embed, - audio_features_length, - audio_in_embed, - audio_in_ids_start, - audio_out_embed, - audio_out_ids_start, - self.audio_in_token_idx, - self.audio_out_token_idx, - inputs_embeds, - input_ids, - attention_mask, - label_ids, - pad_token_id=self.padding_idx, - round_to=round_to, - left_padding=left_padding, - ) - - # re-check if we use the correct kv cache bucket after - # the input_embeds has been merged with audio features - if past_key_values_buckets is not None and inputs_embeds.shape[1] > past_key_values.get_max_cache_shape(): - past_key_values, self.current_past_key_values_bucket = self._prepare_kv_cache( - inputs_embeds.shape[1], None, past_key_values_buckets - ) - - if use_cache and past_key_values is None: - past_key_values = DynamicCache() - - if cache_position is None: - past_seen_tokens = past_key_values.get_seq_length() if past_key_values is not None else 0 - cache_position = torch.arange( - past_seen_tokens, - past_seen_tokens + inputs_embeds.shape[1], - device=inputs_embeds.device, - ) - if isinstance(past_key_values, StaticCache) and past_seen_tokens >= past_key_values.get_max_cache_shape(): - raise ValueError( - f"The current sequence length ({past_seen_tokens}) exceeds " - f"the maximum cache shape. " - f"Please consider increasing the cache size." - ) - - # Use torch compile - use_static_cache = isinstance(past_key_values, StaticCache) - - # Apply the LLM component - causal_mask = self._update_causal_mask( - attention_mask, - inputs_embeds, - cache_position, - past_key_values, - output_attentions, - ) - - hidden_states = inputs_embeds - - audio_discrete_codes_mask = audio_in_discrete_codes_mask | audio_out_mask - if cache_audio_discrete_codes_mask is not None and use_cache: - audio_discrete_codes_mask = torch.concat( - [cache_audio_discrete_codes_mask, audio_discrete_codes_mask], dim=1 - ) - - # Generate the audio attention mask outside the layer to avoid recompilation - if use_static_cache: - fast_forward_attention_mask, audio_attention_mask = self._prepare_all_static_kv_cache_masks( - hidden_states, - causal_mask, - audio_discrete_codes_mask, - past_key_values, - ) - # Set the audio out mask to the last token - if hidden_states.shape[1] == 1: - audio_discrete_codes_mask = audio_discrete_codes_mask[:, -1:] - audio_discrete_codes_mask = audio_discrete_codes_mask.reshape((-1, 1)).contiguous() - is_decoding_audio_token = audio_discrete_codes_mask.item() - else: - is_decoding_audio_token = False - - # Use the captured cuda graph runner for decoding - # if it exists, otherwise use the normal forward pass - if ( - past_key_values is not None - and past_key_values.get_max_cache_shape() in self.decode_graph_runners - and (input_ids.shape[-1] == 1) - ): - _forward_core = self.decode_graph_runners[past_key_values.get_max_cache_shape()][is_decoding_audio_token] - is_using_cuda_graph = True - else: - _forward_core = self._forward_core - is_using_cuda_graph = False - - hidden_states, all_hidden_states, all_self_attns = _forward_core( - hidden_states=hidden_states, - causal_mask=causal_mask, - position_ids=position_ids, - audio_discrete_codes_mask=audio_discrete_codes_mask, - is_decoding_audio_token=is_decoding_audio_token if use_static_cache else None, - cache_position=cache_position, - past_key_values=past_key_values, - use_cache=use_cache, - audio_attention_mask=audio_attention_mask if use_static_cache else None, - fast_forward_attention_mask=fast_forward_attention_mask if use_static_cache else None, - output_attentions=output_attentions, - output_hidden_states=output_hidden_states, - is_using_cuda_graph=is_using_cuda_graph, - ) - hidden_states = self.norm(hidden_states) - - # add hidden states from the last decoder layer - if output_hidden_states: - all_hidden_states += (hidden_states,) - - # Apply the audio decoder projector - ( - logits, - audio_logits, - decoder_all_self_attns, - decoder_all_hidden_states, - audio_hidden_states, - _, - ) = self.audio_decoder_proj( - hidden_states, - audio_out_mask, - label_audio_ids=label_audio_ids, - attention_mask=causal_mask, - position_ids=position_ids, - past_key_values=past_key_values, - use_cache=use_cache, - output_attentions=output_attentions, - output_audio_hidden_states=output_audio_hidden_states, - cache_position=cache_position, - ) - - if audio_logits is not None: - audio_logits = audio_logits.view( - audio_logits.shape[0], - self.audio_num_codebooks, - self.audio_codebook_size, - ).float() - - if output_hidden_states: - if decoder_all_hidden_states is not None and len(decoder_all_hidden_states) > 1: - all_hidden_states += decoder_all_hidden_states[1:] - - if output_attentions: - all_self_attns += decoder_all_self_attns - - next_cache = past_key_values if use_cache else None - - ret = HiggsAudioModelOutputWithPast( - logits=logits, - audio_logits=audio_logits, - expanded_input_ids=input_ids, - expanded_labels=labels, - audio_in_mask=audio_in_mask, - audio_in_discrete_codes_mask=audio_in_discrete_codes_mask, - audio_out_mask=audio_out_mask, - attention_mask=attention_mask, - past_key_values=next_cache, - hidden_states=all_hidden_states, - audio_hidden_states=audio_hidden_states, - attentions=all_self_attns, - ) - - return_dict = return_dict if return_dict is not None else self.config.use_return_dict - if not return_dict: - outputs = ret.to_tuple() - return outputs - - return ret - - # Overwrite GenerationMixin._update_model_kwargs_for_generation - def _update_model_kwargs_for_generation( - self, - outputs: ModelOutput, - model_kwargs: Dict[str, Any], - is_encoder_decoder: bool = False, - num_new_tokens: int = 1, - extend_attention_mask: bool = True, - ) -> Dict[str, Any]: - """Update the model kwargs for each step.""" - model_kwargs["past_key_values"] = outputs.past_key_values - - # update attention mask - if "attention_mask" in model_kwargs: - attention_mask = model_kwargs["attention_mask"] - if extend_attention_mask: - model_kwargs["attention_mask"] = torch.cat( - [ - attention_mask, - attention_mask.new_ones((attention_mask.shape[0], 1)), - ], - dim=-1, - ) - if "cache_audio_discrete_codes_mask" in model_kwargs: - if model_kwargs["cache_audio_discrete_codes_mask"] is None: - model_kwargs["cache_audio_discrete_codes_mask"] = ( - outputs.audio_in_discrete_codes_mask | outputs.audio_out_mask - ) - else: - model_kwargs["cache_audio_discrete_codes_mask"] = torch.concat( - [ - model_kwargs["cache_audio_discrete_codes_mask"], - outputs.audio_in_discrete_codes_mask | outputs.audio_out_mask, - ], - 1, - ) - - return model_kwargs - - def _copy_kv_cache(self, from_cache: Cache, to_cache: Cache): - num_layers = self.config.text_config.num_hidden_layers - if self.config.audio_dual_ffn_layers is not None: - num_layers += len(self.config.audio_dual_ffn_layers) - """ Copy the key-value pairs from one cache to another. """ - for layer_idx in range(num_layers): - from_cache_size = from_cache.get_max_cache_shape() - assert to_cache.get_max_cache_shape() >= from_cache_size, ( - f"The target cache size {to_cache.get_max_cache_shape()} is smaller than the source cache size {from_cache_size}." - ) - to_cache.key_cache[layer_idx][:, :, :from_cache_size, :] = from_cache.key_cache[layer_idx] - to_cache.value_cache[layer_idx][:, :, :from_cache_size, :] = from_cache.value_cache[layer_idx] - - def _prepare_kv_cache( - self, - current_sequence_length: int, - current_past_key_values_bucket: Optional[int], - past_key_values_buckets: OrderedDict[int, Cache], - ) -> Tuple[Optional[Cache], Optional[int]]: - """Prepare the KV cache for the current sequence length.""" - for cache_length in past_key_values_buckets.keys(): - if cache_length >= current_sequence_length: - # Promote to the next KV cache bucket, copy the current KV cache bucket - # to the new one. - if current_past_key_values_bucket is not None and cache_length != current_past_key_values_bucket: - self._copy_kv_cache( - past_key_values_buckets[current_past_key_values_bucket], - past_key_values_buckets[cache_length], - ) - - return past_key_values_buckets[cache_length], cache_length - - raise ValueError( - f"The current sequence length {current_sequence_length} is larger than " - f"all past key values buckets {past_key_values_buckets.keys()}." - ) - - def _sample_audio_tokens( - self, - hidden_states: torch.Tensor, - audio_logits: torch.Tensor, - audio_out_ids: torch.Tensor, - do_sample: bool, - logits_processor: LogitsProcessorList, - device: torch.device, - torch_generator: Optional[torch.Generator], - generation_config: GenerationConfig, - num_delay: int, - num_remaining_delays: Optional[int], - ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, int, Optional[int]]: - """Sample audio tokens and its corresponding text tokens from the logits""" - - # parameters related to repetition aware sampling - ras_win_len = generation_config.generation_kwargs.get("ras_win_len", None) - ras_win_max_num_repeat = generation_config.generation_kwargs.get("ras_win_max_num_repeat", 2) - audio_eos_token_id = generation_config.generation_kwargs.get("audio_eos_token_id", None) - # In the audio generation mode, we sample from audio_logits and keep updating audio_out_ids. - next_audio_token_logits = audio_logits.clone()[-1, :, :].float().to(device) - # TopP, TopK logits processor supports empty input_ids - next_audio_token_scores = logits_processor(None, next_audio_token_logits) - - # token selection - if do_sample: - # next_audio_token_scores has been applied top_p, top_k, and temperature. - probs = nn.functional.softmax(next_audio_token_scores, dim=-1) - # TODO (joao): this OP throws "skipping cudagraphs due to ['incompatible ops']", find solution - next_audio_tokens = torch.multinomial(probs, num_samples=1, generator=torch_generator).squeeze(1) - else: - next_audio_tokens = torch.argmax(next_audio_token_scores, dim=-1) - - # next_tokens: (num_codebooks, ) - if ras_win_len is not None: - # check if there are repetitions over a window of tokens. - rep_num = (audio_out_ids[:, -ras_win_len:] == next_audio_tokens.unsqueeze(1)).sum(dim=1) - - # if we saw repeated tokens in the most recent window of tokens, resample without temperature. - row_indices = torch.nonzero(rep_num >= ras_win_max_num_repeat).squeeze(1) - resampled_next_tokens = ( - next_audio_token_logits[row_indices] - .softmax(dim=-1) - .multinomial(1, replacement=True, generator=torch_generator) - .squeeze(1) - ) - next_audio_tokens[row_indices] = resampled_next_tokens - - # Force the next text tokens to be <|AUDIO_OUT|> in audio generation mode - next_tokens = torch.full( - (audio_logits.shape[0],), - self.config.audio_out_token_idx, - dtype=torch.long, - device=device, - ) - - # Handle delay_pattern - if self.use_delay_pattern: - if num_delay + 1 < next_audio_tokens.shape[0]: - next_audio_tokens[(num_delay + 1) :] = self.config.audio_stream_bos_id - num_delay += 1 - if num_remaining_delays is not None: - next_audio_tokens[: (self.audio_num_codebooks - num_remaining_delays)] = ( - self.config.audio_stream_eos_id - ) - num_remaining_delays -= 1 - else: - all_eos_indices = (next_audio_tokens == self.config.audio_stream_eos_id).nonzero() - if torch.numel(all_eos_indices) > 0: - all_eos_indices = all_eos_indices[0] - last_eos_idx = all_eos_indices[-1] - next_audio_tokens[:last_eos_idx] = self.config.audio_stream_eos_id - num_remaining_delays = self.audio_num_codebooks - last_eos_idx - 1 - if num_remaining_delays is not None and num_remaining_delays <= 0: - next_tokens[...] = audio_eos_token_id - num_delay = 0 - num_remaining_delays = None - - return ( - next_tokens, - next_audio_tokens, - next_audio_token_logits, - next_audio_token_scores, - num_delay, - num_remaining_delays, - ) - - def _sample_text_tokens( - self, - logits: torch.Tensor, - input_ids: torch.Tensor, - do_sample: bool, - logits_processor: LogitsProcessorList, - device: torch.device, - generation_mode: GenerationMode, - torch_generator: Optional[torch.Generator], - ) -> torch.Tensor: - """Sample text tokens from the logits""" - # Clone is needed to avoid keeping a hanging ref to outputs.logits which may be very large for first iteration - # (the clone itself is always small) - next_token_logits = logits.clone()[:, -1, :].float() - next_token_logits = next_token_logits.to(input_ids.device) - - # pre-process distribution - next_token_scores = logits_processor(input_ids, next_token_logits) - - if generation_mode == GenerationMode.AUDIO_INIT: - # See the audio bos token, we should start generating audio tokens - next_tokens = torch.full( - (input_ids.shape[0],), - self.audio_out_token_idx, - dtype=torch.long, - device=device, - ) - next_audio_tokens = torch.full( - (self.config.audio_num_codebooks,), - self.config.audio_stream_bos_id, - dtype=torch.long, - device=device, - ) - else: - if do_sample: - probs = nn.functional.softmax(next_token_scores, dim=-1) - # TODO (joao): this OP throws "skipping cudagraphs due to ['incompatible ops']", find solution - next_tokens = torch.multinomial(probs, num_samples=1, generator=torch_generator).squeeze(1) - else: - next_tokens = torch.argmax(next_token_scores, dim=-1) - - next_audio_tokens = None - - return next_tokens, next_audio_tokens, next_token_logits, next_token_scores - - # Built on top of GenerationMixin._sample. - # We revise the implementation to support generating both audio / text. - def _sample( - self, - input_ids: torch.LongTensor, - logits_processor: LogitsProcessorList, - stopping_criteria: StoppingCriteriaList, - generation_config: GenerationConfig, - synced_gpus: bool, - streamer: Optional["BaseStreamer"], - past_key_values_buckets: Optional[OrderedDict[int, Cache]], - **model_kwargs, - ) -> Union[GenerateNonBeamOutput, torch.LongTensor]: - r""" - Generates sequences of token ids for joint text/audio models using **multinomial sampling**. - - This function may also be revised to support generating samples from HiggsAudio-like end-to-end text/audio models built on top of LLMs. - If the input_ids ends with <|audio_out_bos|>, we will switch to the audio-generation mode. - - ``` - ...<|start_header_id|>assistant<|end_header_id|>\n\n<|audio_out_bos|> - ``` - - Otherwise, we will keep generating the text tokens. - - Parameters: - input_ids (`torch.LongTensor` of shape `(batch_size, sequence_length)`): - The sequence used as a prompt for the generation. - logits_processor (`LogitsProcessorList`): - An instance of [`LogitsProcessorList`]. List of instances of class derived from [`LogitsProcessor`] - used to modify the prediction scores of the language modeling head applied at each generation step. - stopping_criteria (`StoppingCriteriaList`): - An instance of [`StoppingCriteriaList`]. List of instances of class derived from [`StoppingCriteria`] - used to tell if the generation loop should stop. - generation_config ([`~generation.GenerationConfig`]): - The generation configuration to be used as parametrization of the decoding method. - synced_gpus (`bool`): - Whether to continue running the while loop until max_length (needed to avoid deadlocking with - `FullyShardedDataParallel` and DeepSpeed ZeRO Stage 3). - streamer (`BaseStreamer`, *optional*): - Streamer object that will be used to stream the generated sequences. Generated tokens are passed - through `streamer.put(token_ids)` and the streamer is responsible for any further processing. - model_kwargs: - Additional model specific kwargs will be forwarded to the `forward` function of the model. If model is - an encoder-decoder model the kwargs should include `encoder_outputs`. - - Return: - [`~generation.GenerateDecoderOnlyOutput`], [`~generation.GenerateEncoderDecoderOutput`] or `torch.LongTensor`: - A `torch.LongTensor` containing the generated tokens (default behaviour) or a - [`~generation.GenerateDecoderOnlyOutput`] if `model.config.is_encoder_decoder=False` and - `return_dict_in_generate=True` or a [`~generation.GenerateEncoderDecoderOutput`] if - `model.config.is_encoder_decoder=True`. - """ - assert input_ids.shape[0] == 1, "Only support batch_size=1 in _sample()" - audio_out_bos_token_id = generation_config.generation_kwargs.get("audio_out_bos_token_id", None) - - # torch generator for sampling - seed = generation_config.generation_kwargs.get("seed", None) - if seed is not None: - torch_generator = torch.Generator(device=input_ids.device).manual_seed(seed) - else: - torch_generator = None - - # init values - pad_token_id = generation_config._pad_token_tensor - output_attentions = generation_config.output_attentions - output_hidden_states = generation_config.output_hidden_states - output_scores = generation_config.output_scores - output_logits = generation_config.output_logits - return_dict_in_generate = generation_config.return_dict_in_generate - max_length = generation_config.max_length - has_eos_stopping_criteria = any(hasattr(criteria, "eos_token_id") for criteria in stopping_criteria) - do_sample = generation_config.do_sample - # Used to track which past_key_va - self.current_past_key_values_bucket = None - - # init attention / hidden states / scores tuples - scores = () if (return_dict_in_generate and output_scores) else None - raw_logits = () if (return_dict_in_generate and output_logits) else None - - decoder_attentions = () if (return_dict_in_generate and output_attentions) else None - decoder_hidden_states = () if (return_dict_in_generate and output_hidden_states) else None - - # keep track of which sequences are already finished - batch_size, cur_len = input_ids.shape - this_peer_finished = False - unfinished_sequences = torch.ones(batch_size, dtype=torch.long, device=input_ids.device) - if generation_config.use_cache: - model_kwargs["cache_audio_discrete_codes_mask"] = None - - init_model_input = True - num_delay = 0 - num_remaining_delays = None - audio_sequences = [] - # A tensor to keep track of all the audio placeholder tokens. - input_ids_full = input_ids.clone() - - # Initialize the audio variables based on the input prompt. - if input_ids[0][-1] == self.config.audio_out_token_idx: - audio_sequences = [model_kwargs["audio_out_ids"][:, model_kwargs["audio_out_ids_start"][-1] :]] - if self.use_delay_pattern: - num_delay = ( - self.audio_num_codebooks - - (model_kwargs["audio_out_ids"][:, -1] == self.config.audio_stream_bos_id).sum() - ) - all_eos_indices = (model_kwargs["audio_out_ids"][:, -1] == self.config.audio_stream_eos_id).nonzero() - if torch.numel(all_eos_indices) > 0: - all_eos_indices = all_eos_indices[0] - last_eos_idx = all_eos_indices[-1] - num_remaining_delays = self.audio_num_codebooks - last_eos_idx - 1 - - while self._has_unfinished_sequences( - this_peer_finished, - synced_gpus, - device=input_ids.device, - cur_len=cur_len, - max_length=max_length, - ): - # Check which multimodal stage we are in - # FIXME: Assume single input generation - if input_ids[0][-1] == audio_out_bos_token_id: - generation_mode = GenerationMode.AUDIO_INIT - elif input_ids[0][-1] == self.audio_out_token_idx: - generation_mode = GenerationMode.AUDIO_IN_PROGRESS - else: - generation_mode = GenerationMode.TEXT - - is_audio_generation_mode = generation_mode == GenerationMode.AUDIO_IN_PROGRESS - - if init_model_input or not generation_config.use_cache: - model_inputs = {"input_ids": input_ids, **model_kwargs} - else: - model_inputs = {"input_ids": input_ids[:, -1:], **model_kwargs} - - if is_audio_generation_mode and generation_config.use_cache: - model_inputs["audio_out_ids"] = model_kwargs["audio_out_ids"][:, -1:] - model_inputs["audio_out_ids_start"] = torch.tensor([0], dtype=torch.long, device=input_ids.device) - elif not is_audio_generation_mode: - del model_inputs["audio_out_ids"] - del model_inputs["audio_out_ids_start"] - - if generation_config.use_cache: - if "audio_features" in model_inputs and model_inputs["audio_features"] is not None: - model_inputs["audio_features"] = model_inputs["audio_features"][:0, ...] - model_inputs["audio_feature_attention_mask"] = model_inputs["audio_feature_attention_mask"][ - :0, ... - ] - - if "audio_in_ids" in model_inputs and model_inputs["audio_in_ids"] is not None: - model_inputs["audio_in_ids"] = None - model_inputs["audio_in_ids_start"] = None - - # prepare variable output controls (note: some models won't accept all output controls) - model_inputs.update({"output_attentions": output_attentions} if output_attentions else {}) - model_inputs.update({"output_hidden_states": output_hidden_states} if output_hidden_states else {}) - - if past_key_values_buckets is not None: - past_key_values, self.current_past_key_values_bucket = self._prepare_kv_cache( - cur_len, - self.current_past_key_values_bucket, - past_key_values_buckets, - ) - if past_key_values is not None: - model_inputs.update({"past_key_values": past_key_values}) - model_inputs["past_key_values_buckets"] = past_key_values_buckets - - # forward pass to get next token - outputs = self(**model_inputs, return_dict=True) - - # Update the actual sequence length after the first forward pass - if init_model_input and past_key_values_buckets is not None: - cur_len = past_key_values_buckets[self.current_past_key_values_bucket].get_seq_length().item() - - # synced_gpus: don't waste resources running the code we don't need; kwargs must be updated before skipping - model_kwargs = self._update_model_kwargs_for_generation( - outputs, - model_kwargs, - is_encoder_decoder=self.config.is_encoder_decoder, - extend_attention_mask=True, - ) - - # After the first forward pass, we can set init_model_input to False. - init_model_input = False - - if synced_gpus and this_peer_finished: - continue - - if is_audio_generation_mode: - # In audio generation mode, we sample the audio tokens from audio logits. - # It might also generate the audio eos token to end the audio generation. - ( - next_tokens, - next_audio_tokens, - next_audio_token_logits, - next_audio_token_scores, - num_delay, - num_remaining_delays, - ) = self._sample_audio_tokens( - hidden_states=outputs.audio_hidden_states, - audio_logits=outputs.audio_logits, - audio_out_ids=model_kwargs["audio_out_ids"], - do_sample=do_sample, - logits_processor=logits_processor, - device=input_ids.device, - torch_generator=torch_generator, - generation_config=generation_config, - num_delay=num_delay, - num_remaining_delays=num_remaining_delays, - ) - - # update generated ids, model inputs, and length for next step - model_kwargs["audio_out_ids"] = torch.cat( - [model_kwargs["audio_out_ids"], next_audio_tokens[:, None]], dim=-1 - ) - audio_sequences[-1] = torch.cat([audio_sequences[-1], next_audio_tokens[:, None]], dim=-1) - - if streamer is not None: - streamer.put(next_audio_tokens.cpu()) - else: - # In text generation mode, we sample the text tokens from text logits. - # It might also generate the audio placeholder token to start the audio generation. - next_tokens, next_audio_tokens, next_token_logits, next_token_scores = self._sample_text_tokens( - input_ids=input_ids, - logits=outputs.logits, - do_sample=do_sample, - logits_processor=logits_processor, - device=input_ids.device, - generation_mode=generation_mode, - torch_generator=torch_generator, - ) - - if streamer is not None: - streamer.put(next_tokens.cpu()) - - if next_audio_tokens is not None: - # If the token is audio bos token, we will generate the audio placeholder token - # and the corrensponding audio stream bos token to start the audio generation. - audio_sequences.append(next_audio_tokens[:, None]) - if streamer is not None: - streamer.put(next_audio_tokens.cpu()) - if model_kwargs["audio_out_ids"] is None or model_kwargs["audio_out_ids"].shape[0] == 0: - # Initialize audio_out_ids - model_kwargs["audio_out_ids"] = next_audio_tokens[:, None] - model_kwargs["audio_out_ids_start"] = torch.tensor( - [0], dtype=torch.long, device=input_ids.device - ) - else: - model_kwargs["audio_out_ids_start"] = torch.concat( - [ - model_kwargs["audio_out_ids_start"], - torch.tensor( - [model_kwargs["audio_out_ids"].shape[1]], - dtype=torch.long, - device=input_ids.device, - ), - ], - dim=0, - ) - model_kwargs["audio_out_ids"] = torch.concat( - [model_kwargs["audio_out_ids"], next_audio_tokens[:, None]], - dim=1, - ) - - if return_dict_in_generate: - if output_scores: - if is_audio_generation_mode: - scores += (next_audio_token_scores,) - else: - scores += (next_token_scores,) - if output_logits: - if is_audio_generation_mode: - raw_logits += (next_audio_token_logits,) - else: - raw_logits += (next_token_logits,) - if output_attentions: - decoder_attentions += (outputs.attentions,) - if output_hidden_states: - decoder_hidden_states += (outputs.hidden_states,) - - # finished sentences should have their next token be a padding token - if has_eos_stopping_criteria: - next_tokens = next_tokens * unfinished_sequences + pad_token_id * (1 - unfinished_sequences) - - if "tokenizer_length" in generation_config.generation_kwargs: - tokenizer_length = generation_config.generation_kwargs["tokenizer_length"] - if torch.max(next_tokens) >= tokenizer_length: - raise ValueError( - f"Next generated token has max value {torch.max(next_tokens)} which is greater than the tokenizer's vocabulary size {tokenizer_length}, this is undesired behavior." - ) - - # update generated ids, model inputs, and length for next step - if not is_audio_generation_mode or next_tokens[0] != self.audio_out_token_idx: - # We only add one <|AUDIO_OUT|> token to the input_ids for simplicity. - input_ids = torch.cat([input_ids, next_tokens[:, None]], dim=-1) - input_ids_full = torch.cat([input_ids_full, next_tokens[:, None]], dim=-1) - unfinished_sequences = unfinished_sequences & ~stopping_criteria(input_ids_full, scores) - this_peer_finished = unfinished_sequences.max() == 0 - cur_len += 1 - - # This is needed to properly delete outputs.logits which may be very large for first iteration - # Otherwise a reference to outputs is kept which keeps the logits alive in the next iteration - del outputs - - if streamer is not None: - streamer.end() - - if return_dict_in_generate: - return HiggsAudioGenerationOutput( - sequences=input_ids, - audio_sequences=audio_sequences, - scores=scores, - logits=raw_logits, - attentions=decoder_attentions, - hidden_states=decoder_hidden_states, - past_key_values=model_kwargs.get("past_key_values"), - ) - else: - return input_ids, audio_sequences - - @torch.inference_mode() - def generate( - self, - input_ids: Optional[torch.LongTensor] = None, - audio_features: Optional[torch.FloatTensor] = None, - audio_feature_attention_mask: Optional[torch.BoolTensor] = None, - audio_in_ids: Optional[torch.LongTensor] = None, - audio_in_ids_start: Optional[torch.LongTensor] = None, - audio_out_ids: Optional[torch.LongTensor] = None, - audio_out_ids_start: Optional[torch.LongTensor] = None, - past_key_values: Optional[Union[Cache, List[torch.FloatTensor]]] = None, - audio_out_bos_token_id: int = None, - audio_eos_token_id: int = None, - past_key_values_buckets: Optional[OrderedDict[int, Cache]] = None, - seed: Optional[int] = None, - **kwargs, - ): - """ - The generate function in huggingface generally follows these steps: - - for sample_step in 1, 2, 3, 4, 5, ... - ... - - """ - # Right now, it's a very simplified version of generate, we should revisit this after our model architecture stabilizes. - assert input_ids.shape[0] == 1, ( - "Currently HiggsAudioModel.generate() only supports batch_size=1. See the implementation of " - ) - generation_config, kwargs = self._prepare_generation_config(kwargs.pop("generation_config", None), **kwargs) - if audio_out_bos_token_id is not None: - generation_config.generation_kwargs["audio_out_bos_token_id"] = audio_out_bos_token_id - else: - try: - generation_config.generation_kwargs["audio_out_bos_token_id"] = self.audio_out_bos_token_id - except: - generation_config.generation_kwargs["audio_out_bos_token_id"] = None - - if audio_eos_token_id is not None: - generation_config.generation_kwargs["audio_eos_token_id"] = audio_eos_token_id - else: - try: - generation_config.generation_kwargs["audio_eos_token_id"] = self.audio_eos_token_id - except: - generation_config.generation_kwargs["audio_eos_token_id"] = None - - has_default_max_length = kwargs.get("max_length") is None and generation_config.max_length is not None - has_default_min_length = kwargs.get("min_length") is None and generation_config.min_length is not None - - generation_config.generation_kwargs["ras_win_len"] = kwargs.pop("ras_win_len", None) - generation_config.generation_kwargs["ras_win_max_num_repeat"] = kwargs.pop("ras_win_max_num_repeat", 2) - # Set generation seed if determinstic generation is required - if seed is not None: - generation_config.generation_kwargs["seed"] = seed - - # Store tokenizer in generation config if it is in kwargs without popping it - if "tokenizer" in kwargs: - generation_config.generation_kwargs["tokenizer_length"] = len(kwargs["tokenizer"]) - - # input_ids: [bsz, seq_len] - # The merging of audio features happens inside the forward path. The input_ids does not need to change. - # TODO: prepare the final input embeddings to improve generation performance - input_ids_length = input_ids.shape[-1] - generation_config = self._prepare_generated_length( - generation_config=generation_config, - has_default_max_length=has_default_max_length, - has_default_min_length=has_default_min_length, - model_input_name=None, - inputs_tensor=None, - input_ids_length=input_ids_length, - ) - assert generation_config.num_beams == 1, "Currently, we only support beam search with num_beams=1" - return_dict_in_generate = generation_config.return_dict_in_generate - output_scores = generation_config.output_scores - - # When attn_implement is spda or flash-attention, it will create causal mask automatically. - attention_mask = kwargs.pop("attention_mask", None) - return super().generate( - input_ids=input_ids, - attention_mask=attention_mask, - audio_features=audio_features, - audio_feature_attention_mask=audio_feature_attention_mask, - audio_in_ids=audio_in_ids, - audio_in_ids_start=audio_in_ids_start, - audio_out_ids=audio_out_ids, - audio_out_ids_start=audio_out_ids_start, - past_key_values=past_key_values, - generation_config=generation_config, - output_scores=output_scores, - return_dict_in_generate=return_dict_in_generate, - past_key_values_buckets=past_key_values_buckets, - **kwargs, - ) - - def parameter_count_per_component(self): - """Count the number of parameters per component in the model. - - HiggsAudio has the following main components: - audio_tower: For mapping audio features to hidden states), - llm_embed: The size of embedding layer of the LLM - llm_non_embed: The size of non-embedding layer of the LLM - audio_adapter: The overall size of additional layers for audio generation - - """ - trainable_stats = { - "audio_tower": 0, - "llm_embed": 0, - "llm_non_embed": 0, - "audio_embed": 0, - "audio_adapter": 0, - "overall": 0, - } - total_stats = { - "audio_tower": 0, - "llm_embed": 0, - "llm_non_embed": 0, - "audio_embed": 0, - "audio_adapter": 0, - "overall": 0, - } - - total_stats["overall"] = count_parameters(self, trainable_only=False) - trainable_stats["overall"] = count_parameters(self, trainable_only=True) - - for mod in [self.audio_tower]: - if mod is not None: - total_stats["audio_tower"] += count_parameters(mod, trainable_only=False) - trainable_stats["audio_tower"] += count_parameters(mod, trainable_only=True) - - total_stats["llm_embed"] = count_parameters(self.embed_tokens, trainable_only=False) - trainable_stats["llm_embed"] = count_parameters(self.embed_tokens, trainable_only=True) - - total_stats["audio_embed"] = count_parameters(self.audio_codebook_embeddings, trainable_only=False) - trainable_stats["audio_embed"] = count_parameters(self.audio_codebook_embeddings, trainable_only=True) - - # Calculate number of parameters for LLM - for layer in self.layers: - if isinstance(layer, HiggsAudioDualFFNDecoderLayer): - total_param_count = count_parameters(layer, trainable_only=False) - total_trainable_param_count = count_parameters(layer, trainable_only=True) - total_stats["llm_non_embed"] += total_param_count - trainable_stats["llm_non_embed"] += total_trainable_param_count - if not layer.fast_forward: - audio_mlp_param_count = count_parameters(layer.audio_mlp, trainable_only=False) - audio_mlp_trainable_param_count = count_parameters(layer.audio_mlp, trainable_only=True) - - audio_norm_param_count = count_parameters( - layer.audio_post_attention_layernorm, trainable_only=False - ) + count_parameters(layer.audio_input_layernorm, trainable_only=False) - audio_norm_trainable_param_count = count_parameters( - layer.audio_post_attention_layernorm, trainable_only=True - ) + count_parameters(layer.audio_input_layernorm, trainable_only=True) - total_stats["llm_non_embed"] -= audio_mlp_param_count + audio_norm_param_count - trainable_stats["llm_non_embed"] -= ( - audio_mlp_trainable_param_count + audio_norm_trainable_param_count - ) - total_stats["audio_adapter"] += audio_mlp_param_count + audio_norm_param_count - trainable_stats["audio_adapter"] += ( - audio_mlp_trainable_param_count + audio_norm_trainable_param_count - ) - - if layer.use_audio_attention: - audio_attn_param_count = count_parameters( - layer.audio_attn, trainable_only=False - ) + count_parameters(layer.audio_post_audio_attn_layer_norm, trainable_only=False) - audio_attn_trainable_param_count = count_parameters( - layer.audio_attn, trainable_only=True - ) + count_parameters(layer.audio_post_audio_attn_layer_norm, trainable_only=True) - total_stats["llm_non_embed"] -= audio_attn_param_count - trainable_stats["llm_non_embed"] -= audio_attn_trainable_param_count - total_stats["audio_adapter"] += audio_attn_param_count - trainable_stats["audio_adapter"] += audio_attn_trainable_param_count - else: - total_stats["llm_non_embed"] += count_parameters(layer, trainable_only=False) - trainable_stats["llm_non_embed"] += count_parameters(layer, trainable_only=True) - total_stats["llm_non_embed"] += count_parameters(self.norm, trainable_only=False) - trainable_stats["llm_non_embed"] += count_parameters(self.norm, trainable_only=True) - - total_stats["audio_adapter"] += count_parameters(self.audio_decoder_proj.audio_lm_head, trainable_only=False) - trainable_stats["audio_adapter"] += count_parameters( - self.audio_decoder_proj.audio_lm_head, trainable_only=True - ) - total_stats["llm_embed"] += count_parameters(self.audio_decoder_proj.text_lm_head, trainable_only=False) - trainable_stats["llm_embed"] += count_parameters(self.audio_decoder_proj.text_lm_head, trainable_only=True) - - other_audio_modules = [self.audio_encoder_proj] - if self.use_audio_out_embed_projector: - other_audio_modules.append(self.audio_out_embed_projector) - - for mod in other_audio_modules: - if mod is not None: - total_stats["audio_adapter"] += count_parameters(mod, trainable_only=False) - trainable_stats["audio_adapter"] += count_parameters(mod, trainable_only=True) - return {"trainable": trainable_stats, "total": total_stats} - - def set_skip_audio_tower(self): - self.config.skip_audio_tower = True - self.config.encode_whisper_embed = False - - def set_encode_audio_in_tokens(self): - self.config.encode_audio_in_tokens = True - - def freeze_audio_tower(self): - if self.audio_tower is not None: - for param in self.audio_tower.parameters(): - param.requires_grad = False - - def freeze_audio_encoder_proj(self): - if self.audio_encoder_proj is not None: - for param in self.audio_encoder_proj.parameters(): - param.requires_grad = False - - def freeze_llm(self, freeze_embed=True, freeze_embed_until_idx: Optional[int] = None): - for layer in self.layers: - if isinstance(layer, HiggsAudioDualFFNDecoderLayer): - for param in layer.self_attn.parameters(): - param.requires_grad = False - for param in layer.mlp.parameters(): - param.requires_grad = False - - for param in layer.post_attention_layernorm.parameters(): - param.requires_grad = False - - for param in layer.input_layernorm.parameters(): - param.requires_grad = False - else: - for param in layer.parameters(): - param.requires_grad = False - - for param in self.norm.parameters(): - param.requires_grad = False - - if freeze_embed: - if freeze_embed_until_idx is None: - for param in self.embed_tokens.parameters(): - param.requires_grad = False - else: - assert isinstance(self.embed_tokens, nn.Embedding) - self.embed_tokens = PartiallyFrozenEmbedding( - original_embedding=self.embed_tokens, - freeze_until_idx=freeze_embed_until_idx, - ) - - def freeze_text_head(self, freeze_text_head_until_idx: Optional[int] = None): - """Freeze the final text head""" - if freeze_text_head_until_idx is None: - for param in self.audio_decoder_proj.text_lm_head.parameters(): - param.requires_grad = False - - else: - assert isinstance(self.audio_decoder_proj.text_lm_head, nn.Linear) - self.audio_decoder_proj.text_lm_head = PartiallyFrozenLinear( - original_linear=self.audio_decoder_proj.text_lm_head, - freeze_until_idx=freeze_text_head_until_idx, - ) - - @classmethod - def merge_weights_from_checkpoint(cls, checkpoint_dir: str, merged_output_dir: str, *model_args, **kwargs): - # For users' convenience, we merge back embedding and text_lm_head if they are splitted - splitted_model = super().from_pretrained( - checkpoint_dir, - *model_args, - torch_dtype=torch.bfloat16, - device_map="cpu", - **{**kwargs, "state_dict": None}, # Prevent auto-loading state_dict - ) - - # Load all safetensor shards - state_dict = {} - shard_paths = sorted(glob.glob(os.path.join(checkpoint_dir, "*.safetensors"))) - - for shard_path in shard_paths: - shard_dict = load_file(shard_path) # Load each shard - state_dict.update(shard_dict) # Merge into a single dict - - # Merge weights - if ( - "audio_decoder_proj.text_lm_head.linear_frozen.weight" in state_dict - and "audio_decoder_proj.text_lm_head.linear_trainable.weight" in state_dict - ): - state_dict["audio_decoder_proj.text_lm_head.weight"] = torch.cat( - [ - state_dict["audio_decoder_proj.text_lm_head.linear_frozen.weight"], - state_dict["audio_decoder_proj.text_lm_head.linear_trainable.weight"], - ], - dim=0, - ) - - del state_dict["audio_decoder_proj.text_lm_head.linear_frozen.weight"] - del state_dict["audio_decoder_proj.text_lm_head.linear_trainable.weight"] - - if ( - "embed_tokens.embedding_frozen.weight" in state_dict - and "embed_tokens.embedding_trainable.weight" in state_dict - ): - state_dict["embed_tokens.weight"] = torch.cat( - [ - state_dict["embed_tokens.embedding_frozen.weight"], - state_dict["embed_tokens.embedding_trainable.weight"], - ], - dim=0, - ) - - del state_dict["embed_tokens.embedding_frozen.weight"] - del state_dict["embed_tokens.embedding_trainable.weight"] - - # Load the final state_dict - splitted_model.load_state_dict(state_dict, strict=True) - - if merged_output_dir: - splitted_model.save_pretrained(merged_output_dir, is_main_process=True, state_dict=state_dict) - - @torch.inference_mode() - def capture_model(self, past_key_values: list[Union[Cache, List[torch.FloatTensor]]]) -> None: - """Capture CUDA graphs for the model's forward pass with different KV cache lengths. - - Args: - past_key_values: List of KV caches to capture graphs for - """ - for past_key_value in past_key_values: - kv_cache_length = past_key_value.get_max_cache_shape() - # We capture two graphs, one for decoding audio tokens and one for decoding text tokens - for is_decoding_audio_token in [True, False]: - runner = CUDAGraphRunner(self._forward_core) - - # Create dummy inputs for graph capture - batch_size = 1 - hidden_dim = self.config.hidden_size - - hidden_states = torch.zeros( - (batch_size, 1, hidden_dim), - dtype=self.config.torch_dtype, - device="cuda", - ) - causal_mask = torch.ones( - (batch_size, 1, 1, kv_cache_length), - dtype=self.config.torch_dtype, - device="cuda", - ) - position_ids = torch.zeros((batch_size, 1), dtype=torch.long, device="cuda") - audio_discrete_codes_mask = torch.tensor([[is_decoding_audio_token]], dtype=torch.bool, device="cuda") - cache_position = torch.tensor([kv_cache_length - 1], dtype=torch.long, device="cuda") - audio_attention_mask = torch.ones_like(causal_mask) - fast_forward_attention_mask = torch.ones_like(causal_mask) - - runner.capture( - hidden_states=hidden_states, - causal_mask=causal_mask, - position_ids=position_ids, - audio_discrete_codes_mask=audio_discrete_codes_mask, - cache_position=cache_position, - past_key_values=past_key_value, - use_cache=True, - audio_attention_mask=audio_attention_mask, - fast_forward_attention_mask=fast_forward_attention_mask, - output_attentions=False, - output_hidden_states=False, - is_decoding_audio_token=is_decoding_audio_token, - is_using_cuda_graph=True, - ) - - self.decode_graph_runners[kv_cache_length][is_decoding_audio_token] = runner