File size: 12,701 Bytes
9d45136 | 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 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227 228 229 230 231 232 233 234 235 236 237 238 239 240 241 242 243 244 245 246 247 248 249 250 251 252 253 | """
A minimal SpeechLLM: Whisper encoder + connector + LLM with LoRA.
audio ──► Whisper encoder (trained or frozen) ──► connector (trained) ──┐
├──► LLM + LoRA ──► text
text ────────────────────────────────────────► embedding table ────────┘
Two connectors (`adapter_type`), both 50 fps -> 25 fps per downsampling step:
cnn strided Conv1d + GELU, then LayerNorm + Linear
concat_mlp concatenate adjacent frames, then LayerNorm + Linear + GELU + Linear
A frozen LLM does *not* make training cheap: the speech features enter at layer 0,
so gradients still flow back through every LLM layer to reach the connector, and
every LLM activation has to be kept for the backward pass. Turn on
`gradient_checkpointing` if that hurts.
"""
import logging
import torch
import torch.nn as nn
from transformers import AutoModelForCausalLM, WhisperModel
logger = logging.getLogger(__name__)
class SpeechConnector(nn.Module):
"""Whisper hidden states → LLM-width embeddings.
Takes the encoder's last hidden state and applies two steps:
1. **Strided Conv1d downsampling.** Whisper emits 50 frames/s, which is far
too many tokens to hand an LLM. One stride-2 convolution brings that to
25 frames/s -- 40 ms per token, the same rate as Qwen2-Audio. Each extra
layer halves the rate again.
2. **Linear projection to the LLM's width** (from Qwen2-Audio, whose entire
connector is ``AvgPool1d(2)`` + ``Linear``).
"""
def __init__(self, d_encoder: int, d_llm: int, n_downsample_layers: int = 1):
super().__init__()
layers = []
for _ in range(n_downsample_layers):
layers += [nn.Conv1d(d_encoder, d_encoder, kernel_size=3, stride=2, padding=1),
nn.GELU()]
self.downsample = nn.Sequential(*layers)
self.proj = nn.Sequential(nn.LayerNorm(d_encoder), nn.Linear(d_encoder, d_llm))
def forward(self, hidden_states): # (B, T, d_encoder)
hidden_states = self.downsample(hidden_states.transpose(1, 2)).transpose(1, 2)
return self.proj(hidden_states) # (B, T', d_llm)
class ConcatMLPConnector(nn.Module):
"""Whisper hidden states → LLM-width embeddings by frame concatenation.
Every `2 ** n_downsample_layers` adjacent frames are concatenated into one
feature (d_encoder * k), which a two-layer MLP maps to the LLM's width. This is
the SLAM-ASR / LLaVA-style projector: no convolution, all mixing happens in the
MLP. The sequence is right-padded to a multiple of k, so a length-L input gives
ceil(L / k) tokens -- the same count as the CNN, so `data.audio_token_count`
serves both.
"""
def __init__(self, d_encoder: int, d_llm: int, n_downsample_layers: int = 1,
hidden_dim: int = 0):
super().__init__()
self.k = 2 ** n_downsample_layers
hidden_dim = hidden_dim or d_llm
self.proj = nn.Sequential(
nn.LayerNorm(d_encoder * self.k),
nn.Linear(d_encoder * self.k, hidden_dim),
nn.GELU(),
nn.Linear(hidden_dim, d_llm),
)
def forward(self, hidden_states): # (B, T, d_encoder)
B, T, D = hidden_states.shape
pad = (-T) % self.k
if pad:
hidden_states = nn.functional.pad(hidden_states, (0, 0, 0, pad))
stacked = hidden_states.reshape(B, (T + pad) // self.k, D * self.k)
return self.proj(stacked) # (B, ceil(T/k), d_llm)
ADAPTERS = {"cnn": SpeechConnector, "concat_mlp": ConcatMLPConnector}
class SpeechLLM(nn.Module):
def __init__(self, llm_id: str, encoder_id: str,
n_downsample_layers: int = 1, dtype=torch.bfloat16,
adapter_type: str = "cnn", adapter_hidden_dim: int = 0,
gradient_checkpointing: bool = False, cache_dir=None,
freeze_encoder: bool = True, use_lora: bool = False,
lora_rank: int = 32, lora_alpha: int = 64, lora_dropout: float = 0.05,
lora_target_modules=("q_proj", "k_proj", "v_proj", "o_proj")):
super().__init__()
self.llm = AutoModelForCausalLM.from_pretrained(
llm_id, torch_dtype=dtype, cache_dir=cache_dir)
# We only ever need the encoder; loading WhisperModel and keeping .encoder
# lets the decoder weights fall out of scope immediately.
self.encoder = WhisperModel.from_pretrained(
encoder_id, torch_dtype=dtype, cache_dir=cache_dir).encoder
assert adapter_type in ADAPTERS, f"adapter_type must be one of {list(ADAPTERS)}"
extra = {"hidden_dim": adapter_hidden_dim} if adapter_type == "concat_mlp" else {}
self.connector = ADAPTERS[adapter_type](
d_encoder=self.encoder.config.d_model,
d_llm=self.llm.config.hidden_size,
n_downsample_layers=n_downsample_layers,
**extra,
)
logger.info("adapter %s: %.1fM params", adapter_type,
sum(p.numel() for p in self.connector.parameters()) / 1e6)
# --- what trains ---------------------------------------------------
# The connector always does. The LLM trains only through LoRA adapters,
# if at all. The encoder is a switch: leaving it frozen is much cheaper
# than it looks, because with every encoder parameter frozen and
# `input_features` not requiring grad, PyTorch builds no graph for the
# encoder at all -- its activations are never stored. Unfreezing it
# turns 32 Whisper layers back into stored activations, which costs far
# more memory than the 635M parameters themselves.
for param in self.llm.parameters():
param.requires_grad = False
self.use_lora = use_lora
if use_lora:
from peft import LoraConfig ,get_peft_model
lora_config = LoraConfig(
r=lora_rank, lora_alpha=lora_alpha, lora_dropout=lora_dropout,
target_modules=list(lora_target_modules), bias="none",
task_type="CAUSAL_LM",
)
# unwrap back to a plain module: we call the LLM directly (with
# inputs_embeds) rather than through the PeftModel wrapper
self.llm = get_peft_model(self.llm, lora_config).base_model.model
n_lora = sum(p.numel() for n, p in self.llm.named_parameters() if "lora_" in n)
assert n_lora > 0, f"LoRA matched no modules in {lora_target_modules}"
logger.info("LoRA r=%d on %s: %.1fM adapter params",
lora_rank, list(lora_target_modules), n_lora / 1e6)
self.freeze_encoder = freeze_encoder
for param in self.encoder.parameters():
param.requires_grad = not freeze_encoder
# Standard mixed precision: frozen weights stay in `dtype`, everything that
# trains is stored in float32 and autocast does the bf16 compute. bf16 has
# only 7 fraction bits (~0.8% between neighbouring values), so AdamW's
# small per-step updates would round to zero on bf16 weights.
# The connector and LoRA are created in float32 already; the encoder is
# loaded in `dtype`, so cast it when it trains.
if not freeze_encoder:
self.encoder.float()
if gradient_checkpointing:
# use_reentrant=False matters under DDP: the reentrant implementation
# runs the recomputation outside autograd's view, so DDP can mark a
# gradient ready twice and abort with "Expected to mark a variable
# ready only once". The non-reentrant path composes with DDP.
ckpt_kwargs = {"use_reentrant": False}
self.llm.gradient_checkpointing_enable(gradient_checkpointing_kwargs=ckpt_kwargs)
# deliberately no enable_input_require_grads(): the spliced speech
# features already make inputs_embeds require grad, and that hook
# would make it a leaf, breaking the splice (see build_inputs_embeds)
if not freeze_encoder:
self.encoder.gradient_checkpointing_enable(
gradient_checkpointing_kwargs=ckpt_kwargs)
dtypes = {}
for name, param in self.named_parameters():
if param.requires_grad:
part = name.split(".")[0]
dtypes.setdefault(part, set()).add(str(param.dtype).replace("torch.", ""))
logger.info("trainable dtypes: %s", {k: sorted(v) for k, v in sorted(dtypes.items())})
by_part = {}
for name, param in self.named_parameters():
if param.requires_grad:
by_part[name.split(".")[0]] = by_part.get(name.split(".")[0], 0) + param.numel()
trainable = sum(by_part.values())
total = sum(p.numel() for p in self.parameters())
logger.info("trainable %.1fM / %.1fM total (%.2f%%) -- %s",
trainable / 1e6, total / 1e6, 100 * trainable / total,
{k: f"{v/1e6:.1f}M" for k, v in sorted(by_part.items())})
def encode_audio(self, input_features):
"""mel features → speech embeddings in the LLM's width.
We run the encoder over the full 30 s window including padding, and crop
afterwards. Whisper is *trained* that way, so this is its native regime;
masking the padding would be off-distribution for a frozen encoder.
"""
input_features = input_features.to(self.encoder.conv1.weight.dtype)
return self.connector(self.encoder(input_features).last_hidden_state)
def build_inputs_embeds(self, input_ids, speech_features, audio_lengths, start_positions):
"""Overwrite each audio's placeholder embeddings with its speech features."""
embed = self.llm.get_input_embeddings()
# The embedding table is frozen, so `embed(input_ids)` comes back as a
# leaf with requires_grad=False -- and writing in-place into a leaf that
# requires grad is illegal, which is exactly what happens once gradient
# checkpointing turns that flag on. Clone to get a non-leaf we may write
# into. Splicing the speech features in is what makes the result require
# grad, which is also what checkpointing needs from the LLM's input.
inputs_embeds = embed(input_ids).clone()
for k, (row, start) in enumerate(start_positions):
segment = speech_features[k, :audio_lengths[k]]
end = start + segment.size(0)
assert end <= inputs_embeds.size(1), (
f"audio {k} needs slots [{start}:{end}] but the sequence is "
f"{inputs_embeds.size(1)} long")
# outside autocast (inference) the float32 connector output must be
# cast to the LLM's embedding dtype before the in-place write
inputs_embeds[row, start:end] = segment.to(inputs_embeds.dtype)
return inputs_embeds
def forward(self, input_ids, attention_mask, input_features, audio_lengths,
start_positions, labels=None):
speech_features = self.encode_audio(input_features)
inputs_embeds = self.build_inputs_embeds(
input_ids, speech_features, audio_lengths, start_positions)
return self.llm(inputs_embeds=inputs_embeds,
attention_mask=attention_mask,
labels=labels)
@torch.no_grad()
def generate(self, batch, tokenizer, **generation_kwargs):
"""Greedy/sampled decoding from the same batch dict used for training.
Note we pass `inputs_embeds`, so `generate` returns only the newly
generated ids -- there is no prompt prefix to strip.
"""
speech_features = self.encode_audio(batch["input_features"])
inputs_embeds = self.build_inputs_embeds(
batch["input_ids"], speech_features, batch["audio_lengths"],
batch["start_positions"])
return self.llm.generate(
inputs_embeds=inputs_embeds,
attention_mask=batch["attention_mask"],
pad_token_id=tokenizer.pad_token_id,
**generation_kwargs,
)
def trainable_state_dict(self):
"""Only the connector -- a few tens of MB instead of the full ~11 GB."""
return {name: param.detach().clone()
for name, param in self.named_parameters() if param.requires_grad}
|