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