import logging from typing import Any, Dict, Optional, Set, Tuple, Union import peft import torch import torch.nn as nn import torch.nn.functional as F import transformers import transformers.activations import transformers.modeling_outputs import transformers.models # We must use relative import in this directory to allow uploading to HF Hub # Even "from . import X" pattern doesn't work (undocumented and unclear why) from .ultravox_config import LossConfig from .ultravox_config import LossFunction from .ultravox_config import UltravoxConfig from .whisper_model_modified import WhisperEncoder as ModifiedWhisperEncoder class UltravoxModel(transformers.LlamaPreTrainedModel): """ The Ultravox model which consists of an audio encoder and a language model. Audio input is processed by the audio encoder, then every `stack_factor` frames are stacked together and projected to the language model's embedding space using a few linear layers. The text is embedded by the language model as usual and then the audio and text embeddings are merged together. A special token `<|audio|>` is used to indicate the start of the audio embeddings in the merged embeddings. Parameters: config: Model configuration class with all the parameters of the model. """ config_class = UltravoxConfig config: UltravoxConfig # for type hinting _no_split_modules = ["Wav2Vec2Model", "WhisperEncoder", "LlamaDecoderLayer"] # We minimize the weights in state_dict in order to reduce the size of the checkpoint # The issue is that load_pretrained() uses state_dict() keys to know what keys are expected # As such we have to tell is to ignore some keys that are not always in the model _keys_to_ignore_on_load_unexpected = ["audio_tower.*", "language_model.*"] # Usually we load encoder weights from a pretrained model, so we don't want to load the decoder weights # Technically we never hit this issue because these keys are already removed from state_dict() however, # but there's no harm in keeping it here for when we change that behavior. _keys_to_ignore_on_load_missing = ["audio_tower.*"] def __init__(self, config: UltravoxConfig): super().__init__(config) self.keep_params: Set[str] = set() self.vocab_size = config.vocab_size self.audio_tower = self._create_audio_tower(config) self.multi_modal_projector = UltravoxProjector(config) self.language_model = self._create_language_model(config) self.loss_config = LossConfig() self.post_init() def get_input_embeddings(self): return self.language_model.get_input_embeddings() def set_input_embeddings(self, value): self.language_model.set_input_embeddings(value) def get_output_embeddings(self): return self.language_model.get_output_embeddings() def set_output_embeddings(self, new_embeddings): self.language_model.set_output_embeddings(new_embeddings) def set_decoder(self, decoder): self.language_model.set_decoder(decoder) def get_decoder(self): return self.language_model.get_decoder() def tie_weights(self): return self.language_model.tie_weights() def set_loss_config(self, loss_config: LossConfig): self.loss_config = loss_config def _setup_cache( self, cache_cls, max_batch_size: int, max_cache_len: Optional[int] = None ): self.language_model._setup_cache(cache_cls, max_batch_size, max_cache_len) def _reorder_cache(self, past_key_values, beam_idx): return self.language_model._reorder_cache(past_key_values, beam_idx) def resize_token_embeddings( self, new_num_tokens: Optional[int] = None, pad_to_multiple_of: Optional[int] = None, ) -> nn.Embedding: model_embeds = self.language_model.resize_token_embeddings( new_num_tokens, pad_to_multiple_of ) # update vocab size self.config.text_config.vocab_size = model_embeds.num_embeddings self.config.vocab_size = model_embeds.num_embeddings self.vocab_size = model_embeds.num_embeddings return model_embeds def _compute_kl_loss( self, lm_output: transformers.modeling_outputs.CausalLMOutputWithPast, labels: Optional[torch.Tensor] = None, past_key_values: Optional[Union[Tuple, transformers.cache_utils.Cache]] = None, alt_input_ids: Optional[torch.Tensor] = None, alt_attention_mask: Optional[torch.Tensor] = None, alt_labels: Optional[torch.Tensor] = None, **kwargs, ): # disable gradient computation for the teacher model with torch.no_grad(): # compute the teacher (text-only) model's distribution alt_inputs_embeds = self.get_input_embeddings().forward(alt_input_ids) alt_lm_output = self.language_model.forward( inputs_embeds=alt_inputs_embeds, labels=alt_labels, attention_mask=alt_attention_mask, past_key_values=past_key_values, **kwargs, ) # compute the KL divergence loss between the two models kl_loss = F.kl_div( F.log_softmax( lm_output.logits[labels != -100] / self.loss_config.kl_temperature, dim=-1, ), F.softmax( alt_lm_output.logits[alt_labels != -100] / self.loss_config.kl_temperature, dim=-1, ), reduction="batchmean", ) return {"loss": kl_loss} def forward( self, input_ids: torch.Tensor, audio_values: Optional[torch.FloatTensor] = None, inputs_embeds: Optional[torch.FloatTensor] = None, labels: Optional[torch.Tensor] = None, attention_mask: Optional[torch.Tensor] = None, audio_token_start_idx: Optional[torch.Tensor] = None, audio_token_len: Optional[torch.Tensor] = None, past_key_values: Optional[Union[Tuple, transformers.cache_utils.Cache]] = None, # the alt_* fields are needed for KL divergence loss alt_input_ids: Optional[torch.Tensor] = None, alt_attention_mask: Optional[torch.Tensor] = None, alt_labels: Optional[torch.Tensor] = None, **kwargs, ) -> Union[Tuple, transformers.modeling_outputs.CausalLMOutputWithPast]: """ Forward pass for the Ultravox model. `input_ids` are the tokenized text input. They are embedded by the language model as usual. `audio_values` are processed by the audio encoder and then every `stack_factor` frames are stacked together and projected to the language model's embedding space using a few linear layers. The audio and text embeddings are merged together. A special token `<|audio|>` is used to indicate the start of the audio embeddings in the merged embeddings. Args: input_ids: The tokenized text input. audio_values: The processed audio values. inputs_embeds: The embeddings for the input tokens. labels: The tokenized text labels. attention_mask: The attention mask for the input. position_ids: The position ids for the input. past_key_values: The past key value cache for the language model attention layers. **kwargs: Additional keyword arguments. Passed directly to the language model. """ if inputs_embeds is None: # B x T -> B x T x D inputs_embeds = self.get_input_embeddings().forward(input_ids) if audio_values is not None: assert ( audio_token_start_idx is not None and audio_token_len is not None ), "audio_token_start_idx and audio_token_len must be provided if audio_values are provided." assert ( len(audio_token_start_idx) == len(audio_token_len) == len(audio_values) ), "audio_token_start_idx, audio_token_len, and audio_values must have the same batch size." # B x A/3200 x D audio_tower_output = self.audio_tower.forward( audio_values ).last_hidden_state audio_tower_output = audio_tower_output.to(inputs_embeds.dtype) audio_embeds = self.multi_modal_projector.forward(audio_tower_output) # combine audio and text embeddings for i, (audio, start, length) in enumerate( zip(audio_embeds, audio_token_start_idx, audio_token_len) ): length = min(length, audio.shape[0]) inputs_embeds[i, start : start + length] = audio[:length] lm_output = self.language_model.forward( inputs_embeds=inputs_embeds, labels=labels, attention_mask=attention_mask, past_key_values=past_key_values, **kwargs, ) if self.training: if self.loss_config.loss_function == LossFunction.CrossEntropy: return lm_output elif self.loss_config.loss_function == LossFunction.KL_Divergence: return self._compute_kl_loss( lm_output=lm_output, labels=labels, past_key_values=past_key_values, alt_input_ids=alt_input_ids, alt_attention_mask=alt_attention_mask, alt_labels=alt_labels, **kwargs, ) else: raise ValueError( f"Unsupported loss function: {self.loss_config.loss_function}" ) else: return lm_output def prepare_inputs_for_generation( self, input_ids: torch.Tensor, audio_values: Optional[torch.FloatTensor] = None, audio_token_start_idx: Optional[torch.Tensor] = None, audio_token_len: Optional[torch.Tensor] = None, past_key_values: Optional[Union[Tuple, transformers.cache_utils.Cache]] = None, attention_mask: Optional[torch.Tensor] = None, inputs_embeds: Optional[torch.Tensor] = None, **kwargs, ) -> Dict[str, Any]: model_input = self.language_model.prepare_inputs_for_generation( input_ids=input_ids, past_key_values=past_key_values, attention_mask=attention_mask, inputs_embeds=inputs_embeds, **kwargs, ) if is_cache_empty(past_key_values) and audio_values is not None: # We only want to use audio features in the 1st generation step model_input["audio_values"] = audio_values model_input["audio_token_start_idx"] = audio_token_start_idx model_input["audio_token_len"] = audio_token_len return model_input @classmethod def _create_audio_tower( cls, config: UltravoxConfig ) -> Union[transformers.Wav2Vec2Model, ModifiedWhisperEncoder]: if config.audio_model_id is not None: if "whisper" in config.audio_model_id is not None: audio_tower = ModifiedWhisperEncoder.from_pretrained( config.audio_model_id ) else: audio_tower = transformers.AutoModel.from_pretrained( config.audio_model_id ) else: if "whisper" in config.audio_config._name_or_path: audio_tower = ModifiedWhisperEncoder(config.audio_config) else: with transformers.modeling_utils.no_init_weights(): # we only ever use from_config if the weights are retrained, hence initializing is not # required. This makes the model quite creation faster since init on CPU is quite slow. audio_tower = transformers.AutoModel.from_config( config.audio_config ) if isinstance( audio_tower, (transformers.Wav2Vec2BertModel, transformers.WhisperModel), ): # For these models we only need the encoder part # Wav2Vec2BertModel -> Wav2Vec2BertEncoder # WhisperModel -> WhisperEncoder audio_tower = audio_tower.encoder audio_tower = apply_lora(audio_tower, config.audio_model_lora_config) return audio_tower @classmethod def _create_language_model( cls, config: UltravoxConfig ) -> transformers.LlamaForCausalLM: if config.text_model_id is not None: language_model = transformers.AutoModelForCausalLM.from_pretrained( config.text_model_id, attn_implementation=config._attn_implementation ) else: with transformers.modeling_utils.no_init_weights(): # we only ever use from_config if the weights are retrained, hence initializing is not # required. This makes the model quite creation faster since init on CPU is quite slow. language_model = transformers.AutoModelForCausalLM.from_config( config.text_config, attn_implementation=config._attn_implementation ) language_model = apply_lora(language_model, config.text_model_lora_config) return language_model def _add_language_model_weights_to_keep(self): if self.config.text_model_id is not None: self.config.text_model_id = None self.keep_params.update( set( [ f"language_model.{name}" for name, _ in self.language_model.named_parameters() ] ) ) def _add_audio_tower_weights_to_keep(self): if self.config.audio_model_id is not None: self.config.audio_model_id = None self.keep_params.update( set( [ f"audio_tower.{name}" for name, _ in self.audio_tower.named_parameters() ] ) ) def merge_and_unload(self): if isinstance(self.language_model, peft.PeftModel): self.language_model = self.language_model.merge_and_unload() # no need to download base language model weights anymore, so we can remove the id self._add_language_model_weights_to_keep() if isinstance(self.audio_tower, peft.PeftModel): self.audio_tower = self.audio_tower.merge_and_unload() # no need to download base audio model weights anymore, so we can remove the id self._add_audio_tower_weights_to_keep() for param in ["text_model_lora_config", "audio_model_lora_config"]: if hasattr(self.config, param): delattr(self.config, param) def push_to_hub(self, *args, **kwargs): self.merge_and_unload() self.to(self.language_model.dtype) return super().push_to_hub(*args, **kwargs) def state_dict(self, *args, **kwargs): named_params = dict(self.named_parameters()) state_dict = super().state_dict(*args, **kwargs) state_dict = { k: v for k, v in state_dict.items() if k in self.keep_params or (k in named_params and named_params[k].requires_grad) } return state_dict def load_state_dict( self, state_dict: Dict[str, Any], *args, **kwargs, ): self.keep_params.update(set(state_dict.keys())) return super().load_state_dict(state_dict, *args, **kwargs) def print_trainable_parameters(self): """ Prints the number of trainable parameters in the model (reuses Peft model's method) """ count_params = peft.peft_model.PeftModel.get_nb_trainable_parameters trainable_params, all_param = count_params(self) logging.info( f"trainable params: {trainable_params:,d} || all params: {all_param:,d}" f" || trainable%: {100 * trainable_params / all_param:.1f}%" ) lm_trainable_params, lm_all_params = count_params(self.language_model) audio_trainable_params, audio_all_params = count_params(self.audio_tower) projector_trainable_params = ( trainable_params - lm_trainable_params - audio_trainable_params ) projector_all_params = all_param - lm_all_params - audio_all_params logging.info( f"Trainable%: " f" LLM: {100 * lm_trainable_params / lm_all_params:.1f}%" f" || Audio Encoder: {100 * audio_trainable_params / audio_all_params:.1f}%" f" || Projector: {100 * projector_trainable_params / projector_all_params:.1f}%" ) def is_cache_empty( past_key_values: Optional[Union[Tuple, transformers.cache_utils.Cache]] ) -> bool: """ Check if the cache is empty. """ if past_key_values is None: return True if isinstance(past_key_values, tuple): return all(len(c) == 0 for c in past_key_values) return past_key_values.get_seq_length() == 0 def apply_lora(model: torch.nn.Module, lora_config: dict) -> torch.nn.Module: """ Applies LoRA finetuning to the model. If the `r` parameter is set to 0, the model is frozen instead. """ lora_config = peft.LoraConfig(**lora_config or {}) if lora_config.r == 0: # freeze the model entirely for param in model.parameters(): param.requires_grad = False else: model = peft.get_peft_model(model, lora_config) return model class StackAudioFrames(nn.Module): """ Stack the audio embedding frames to reduce the sequence length by a factor of `stack_factor`. The number of output frames will be `ceil(T / stack_factor) + 1` where `T` is the number of input frames. NOTE: the extra +1 is intentional: in case the number of audio tokens are over-estimated by the processor, we want to make sure `processor.audio_token_replacement` (i.e. EOS) doesn't get leaked into the middle of embeddings. In most cases this extra padding will get removed in the model's forward function so it has no effect. """ def __init__(self, stack_factor: int = 8): super().__init__() self.stack_factor = stack_factor def forward(self, audio_embeds: torch.Tensor) -> torch.Tensor: B, T, C = audio_embeds.shape T_pad = (T + self.stack_factor - 1) // self.stack_factor * self.stack_factor audio_embeds = F.pad(audio_embeds, (0, 0, 0, T_pad - T + self.stack_factor)) B, T, C = audio_embeds.shape audio_embeds = audio_embeds.view( B, T // self.stack_factor, C * self.stack_factor ) return audio_embeds class RMSNorm(transformers.models.llama.modeling_llama.LlamaRMSNorm): def __init__(self, hidden_size: int, init: float = 1, eps: float = 1e-6): super().__init__(hidden_size=hidden_size, eps=eps) self.weight.data.fill_(init) class SwiGLU(nn.Module): def forward(self, x): x, gate = x.chunk(2, dim=-1) return F.silu(gate) * x class UltravoxProjector(nn.Sequential): def __init__(self, config: UltravoxConfig): super().__init__() self.hidden_dim = config.hidden_size self._pad_and_stack = StackAudioFrames(config.stack_factor) dim = config.audio_config.hidden_size * config.stack_factor self.ln_pre = RMSNorm(dim, init=config.norm_init) self.linear_1 = nn.Linear(dim, self.hidden_dim, bias=False) dim = self.hidden_dim self.act = transformers.activations.get_activation(config.projector_act) dim = dim // 2 if config.projector_act == "swiglu" else dim self.linear_2 = nn.Linear(dim, config.text_config.hidden_size, bias=False) self.ln_post = RMSNorm(config.text_config.hidden_size, init=config.norm_init) def forward(self, audio_features: torch.Tensor) -> torch.Tensor: audio_features = self._pad_and_stack(audio_features) audio_features = self.ln_pre(audio_features) hidden_states = self.linear_1(audio_features) hidden_states = self.act(hidden_states) hidden_states = self.linear_2(hidden_states) hidden_states = self.ln_post(hidden_states) return hidden_states UltravoxConfig.register_for_auto_class() UltravoxModel.register_for_auto_class() transformers.AutoConfig.register("ultravox", UltravoxConfig) transformers.AutoModel.register(UltravoxConfig, UltravoxModel) # transformers.AutoProcessor.register(UltravoxConfig, UltravoxProcessor) # TODO: make processor work standalone transformers.activations.ACT2FN["swiglu"] = SwiGLU