Spaces:
Sleeping
Sleeping
File size: 6,610 Bytes
82e11f0 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 | 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
|