Upload folder using huggingface_hub

9d45136 verified about 1 month ago

12.7 kB

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

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