"""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