Spaces:
Sleeping
Sleeping
| from typing import Any | |
| import torch | |
| from cut_cross_entropy import linear_cross_entropy | |
| from kernels import LayerRepository | |
| from transformers.cache_utils import Cache | |
| from transformers.modeling_outputs import CausalLMOutputWithPast | |
| from transformers.models.lfm2 import modeling_lfm2 | |
| from transformers.models.lfm2.configuration_lfm2 import Lfm2Config | |
| from transformers.models.lfm2.modeling_lfm2 import Lfm2DecoderLayer, Lfm2ForCausalLM | |
| from src.modeling.cross_attention import GatedCrossAttention | |
| def enable_causal_convolution_kernels() -> None: | |
| # --------------------------------------------------------- | |
| # Replace LFM reference convolutions with CUDA 13 Hub kernels. | |
| # --------------------------------------------------------- | |
| for kernel_name in ("causal_conv1d_fn", "causal_conv1d_update"): | |
| kernel_class = LayerRepository( | |
| repo_id="kernels-community/mamba-ssm", | |
| layer_name=kernel_name, | |
| version=3, | |
| ).load() | |
| setattr(modeling_lfm2, kernel_name, kernel_class()) | |
| class AudioLfm2DecoderLayer(Lfm2DecoderLayer): | |
| def __init__( | |
| self, | |
| config: Lfm2Config, | |
| layer_idx: int, | |
| audio_cache_index: int, | |
| ) -> None: | |
| super().__init__(config, layer_idx) | |
| self.audio_cache_index = audio_cache_index | |
| self.cross_attention = GatedCrossAttention( | |
| hidden_size=config.hidden_size, | |
| num_heads=config.num_attention_heads, | |
| norm_eps=config.norm_eps, | |
| ) | |
| def forward( | |
| self, | |
| hidden_states: torch.Tensor, | |
| position_embeddings: tuple[torch.Tensor, torch.Tensor] | None = None, | |
| attention_mask: torch.Tensor | None = None, | |
| position_ids: torch.LongTensor | None = None, | |
| past_key_values: Cache | None = None, | |
| *, | |
| audio_key_values: torch.Tensor, | |
| audio_padding_mask: torch.BoolTensor, | |
| **kwargs: Any, | |
| ) -> torch.Tensor: | |
| # --------------------------------------------------------- | |
| # Keep the original LFM operator and residual connection. | |
| # Based on Transformers 5.8.1 Lfm2DecoderLayer. | |
| # --------------------------------------------------------- | |
| residual = hidden_states | |
| hidden_states, _ = self.self_attn( | |
| hidden_states=self.operator_norm(hidden_states), | |
| position_embeddings=position_embeddings, | |
| attention_mask=attention_mask, | |
| position_ids=position_ids, | |
| past_key_values=past_key_values, | |
| **kwargs, | |
| ) | |
| hidden_states = hidden_states + residual | |
| # --------------------------------------------------------- | |
| # Add audio information before the original feed-forward block. | |
| # --------------------------------------------------------- | |
| hidden_states = self.cross_attention( | |
| hidden_states, | |
| audio_key_values[:, self.audio_cache_index], | |
| audio_padding_mask, | |
| ) | |
| hidden_states = hidden_states + self.feed_forward(self.ffn_norm(hidden_states)) | |
| return hidden_states | |
| class AudioLfm2ForCausalLM(Lfm2ForCausalLM): | |
| def project_audio_key_values( | |
| self, | |
| audio_features: torch.Tensor, | |
| ) -> torch.Tensor: | |
| # --------------------------------------------------------- | |
| # Build one reusable audio K/V cache for every cross-attention layer. | |
| # --------------------------------------------------------- | |
| key_values = [layer.cross_attention.project_key_value(audio_features) for layer in self.model.layers if isinstance(layer, AudioLfm2DecoderLayer)] | |
| return torch.stack(key_values, dim=1) | |
| def forward( | |
| self, | |
| input_ids: torch.LongTensor | None = None, | |
| attention_mask: torch.Tensor | None = None, | |
| position_ids: torch.LongTensor | None = None, | |
| past_key_values: Cache | None = None, | |
| inputs_embeds: torch.FloatTensor | None = None, | |
| labels: torch.LongTensor | None = None, | |
| use_cache: bool | None = None, | |
| logits_to_keep: int | torch.Tensor = 0, | |
| *, | |
| audio_key_values: torch.Tensor, | |
| audio_padding_mask: torch.BoolTensor, | |
| **kwargs: Any, | |
| ) -> CausalLMOutputWithPast: | |
| outputs = self.model( | |
| input_ids=input_ids, | |
| attention_mask=attention_mask, | |
| position_ids=position_ids, | |
| past_key_values=past_key_values, | |
| inputs_embeds=inputs_embeds, | |
| use_cache=use_cache, | |
| audio_key_values=audio_key_values, | |
| audio_padding_mask=audio_padding_mask, | |
| **kwargs, | |
| ) | |
| hidden_states = outputs.last_hidden_state | |
| # --------------------------------------------------------- | |
| # Fuse the classifier and training loss without materializing logits. | |
| # --------------------------------------------------------- | |
| if labels is not None: | |
| loss = linear_cross_entropy( | |
| hidden_states.to(torch.bfloat16), | |
| self.lm_head.weight, | |
| labels, | |
| reduction="sum", | |
| shift=True, | |
| filter_eps="high", | |
| impl="cce", | |
| ) | |
| logits = None | |
| else: | |
| slice_indices = slice(-logits_to_keep, None) if isinstance(logits_to_keep, int) else logits_to_keep | |
| logits = self.lm_head(hidden_states[:, slice_indices, :]) | |
| loss = None | |
| return CausalLMOutputWithPast( | |
| loss=loss, | |
| logits=logits, | |
| past_key_values=outputs.past_key_values, | |
| hidden_states=outputs.hidden_states, | |
| attentions=outputs.attentions, | |
| ) | |
| def add_cross_attention(model: Lfm2ForCausalLM) -> None: | |
| # --------------------------------------------------------- | |
| # Replace only full-attention layers and preserve LFM weights. | |
| # --------------------------------------------------------- | |
| audio_cache_index = 0 | |
| for layer_idx, layer_type in enumerate(model.config.layer_types): | |
| if layer_type != "full_attention": | |
| continue | |
| original_layer = model.model.layers[layer_idx] | |
| new_layer = AudioLfm2DecoderLayer( | |
| model.config, | |
| layer_idx, | |
| audio_cache_index, | |
| ) | |
| new_layer.load_state_dict(original_layer.state_dict(), strict=False) | |
| parameter = next(original_layer.parameters()) | |
| new_layer.to(device=parameter.device, dtype=parameter.dtype) | |
| model.model.layers[layer_idx] = new_layer | |
| audio_cache_index += 1 | |