| """ |
| Model components for Speech-to-Speech training. |
| |
| Single Responsibility: Only defines model architectures. |
| Open/Closed: Can extend with new adapters without modifying existing code. |
| |
| Optimizations: |
| - Flash Attention 2 for memory-efficient attention (10-20x savings on long sequences) |
| - BFloat16 for better numerical stability than FP16 |
| - Gradient checkpointing for memory savings |
| """ |
|
|
| import torch |
| import torch.nn as nn |
| from typing import Optional |
| from .config import ( |
| DEFAULT_WHISPER_DIM, |
| DEFAULT_LLM_DIM, |
| DEFAULT_DOWNSAMPLE, |
| DEFAULT_INTERMEDIATE_DIM, |
| ) |
| from .utils import log |
|
|
|
|
| class SpeechAdapter(nn.Module): |
| """ |
| Speech adapter that maps Whisper features to LLM embedding space. |
| |
| Architecture: 5× downsampling + FFN with intermediate dim |
| |
| Based on LLaMA-Omni 2 design: |
| - Concatenates 5 consecutive Whisper frames |
| - Projects through 2-layer FFN |
| - Applies LayerNorm for stability |
| |
| Args: |
| whisper_dim: Dimension of Whisper features (default: 1280) |
| llm_dim: Dimension of LLM embeddings (default: 3072) |
| downsample: Downsampling factor (default: 5) |
| intermediate_dim: Hidden dimension of FFN (default: 2048) |
| """ |
|
|
| def __init__( |
| self, |
| whisper_dim: int = DEFAULT_WHISPER_DIM, |
| llm_dim: int = DEFAULT_LLM_DIM, |
| downsample: int = DEFAULT_DOWNSAMPLE, |
| intermediate_dim: int = DEFAULT_INTERMEDIATE_DIM |
| ): |
| super().__init__() |
| self.whisper_dim = whisper_dim |
| self.llm_dim = llm_dim |
| self.downsample = downsample |
| self.intermediate_dim = intermediate_dim |
|
|
| concat_dim = whisper_dim * downsample |
|
|
| self.ffn = nn.Sequential( |
| nn.Linear(concat_dim, intermediate_dim), |
| nn.GELU(), |
| nn.Linear(intermediate_dim, llm_dim), |
| nn.LayerNorm(llm_dim) |
| ) |
|
|
| def forward(self, x: torch.Tensor) -> torch.Tensor: |
| """ |
| Forward pass. |
| |
| Args: |
| x: Whisper features [B, T, D] |
| |
| Returns: |
| LLM embeddings [B, T // downsample, llm_dim] |
| """ |
| B, T, D = x.shape |
|
|
| |
| T_new = (T // self.downsample) * self.downsample |
| x = x[:, :T_new] |
|
|
| |
| x = x.reshape(B, T_new // self.downsample, D * self.downsample) |
|
|
| return self.ffn(x) |
|
|
| def get_num_params(self) -> int: |
| """Return total number of parameters.""" |
| return sum(p.numel() for p in self.parameters()) |
|
|
| def get_config(self) -> dict: |
| """Return configuration dict for serialization.""" |
| return { |
| "whisper_dim": self.whisper_dim, |
| "llm_dim": self.llm_dim, |
| "downsample": self.downsample, |
| "intermediate_dim": self.intermediate_dim, |
| } |
|
|
| @classmethod |
| def from_config(cls, config: dict) -> 'SpeechAdapter': |
| """Create adapter from configuration dict.""" |
| return cls(**config) |
|
|
|
|
| class ModelFactory: |
| """ |
| Factory for creating models with consistent settings. |
| |
| Single Responsibility: Only handles model instantiation. |
| Dependency Inversion: Depends on abstractions (config), not concretions. |
| """ |
|
|
| @staticmethod |
| def create_adapter( |
| whisper_dim: int = DEFAULT_WHISPER_DIM, |
| llm_dim: int = DEFAULT_LLM_DIM, |
| dtype: torch.dtype = torch.float32, |
| checkpoint_path: Optional[str] = None |
| ) -> SpeechAdapter: |
| """ |
| Create a SpeechAdapter, optionally loading from checkpoint. |
| |
| Args: |
| whisper_dim: Whisper feature dimension |
| llm_dim: LLM embedding dimension |
| dtype: Tensor dtype |
| checkpoint_path: Optional path to checkpoint |
| |
| Returns: |
| Initialized SpeechAdapter |
| """ |
| adapter = SpeechAdapter( |
| whisper_dim=whisper_dim, |
| llm_dim=llm_dim, |
| ).to(dtype=dtype) |
|
|
| if checkpoint_path: |
| ckpt = torch.load(checkpoint_path, map_location="cpu", weights_only=False) |
| if "adapter" in ckpt: |
| adapter.load_state_dict(ckpt["adapter"]) |
| elif "state_dict" in ckpt: |
| adapter.load_state_dict(ckpt["state_dict"]) |
| else: |
| adapter.load_state_dict(ckpt) |
|
|
| return adapter |
|
|
| @staticmethod |
| def create_llm( |
| model_path: str, |
| dtype: torch.dtype = torch.bfloat16, |
| freeze: bool = True, |
| gradient_checkpointing: bool = False, |
| use_flash_attention: bool = True, |
| verbose: bool = True, |
| ): |
| """ |
| Create and configure the LLM with memory optimizations. |
| |
| Args: |
| model_path: HuggingFace model path |
| dtype: Tensor dtype (BF16 recommended for stability) |
| freeze: Whether to freeze all parameters |
| gradient_checkpointing: Enable gradient checkpointing |
| use_flash_attention: Try to use Flash Attention 2 (10-20x memory savings) |
| verbose: Log configuration details |
| |
| Returns: |
| Configured LLM model |
| """ |
| from transformers import AutoModelForCausalLM |
|
|
| |
| |
| |
| attn_impl = "sdpa" |
|
|
| if use_flash_attention and torch.cuda.is_available(): |
| try: |
| |
| import flash_attn |
| |
| major, _ = torch.cuda.get_device_capability() |
| if major >= 8: |
| attn_impl = "flash_attention_2" |
| if verbose: |
| log(f"[LLM] Using Flash Attention 2 v{flash_attn.__version__} (10-20x memory savings)") |
| else: |
| if verbose: |
| log(f"[LLM] GPU SM {major}.x too old for Flash Attention 2, using SDPA") |
| except ImportError: |
| if verbose: |
| log("[LLM] flash_attn not installed, using SDPA") |
| except Exception as e: |
| if verbose: |
| log(f"[LLM] Flash Attention check failed: {e}, using SDPA") |
|
|
| |
| try: |
| llm = AutoModelForCausalLM.from_pretrained( |
| model_path, |
| torch_dtype=dtype, |
| attn_implementation=attn_impl, |
| ) |
| except Exception as e: |
| |
| if attn_impl == "flash_attention_2": |
| if verbose: |
| log(f"[LLM] Flash Attention failed ({e}), falling back to SDPA") |
| llm = AutoModelForCausalLM.from_pretrained( |
| model_path, |
| torch_dtype=dtype, |
| attn_implementation="sdpa", |
| ) |
| else: |
| raise |
|
|
| if freeze: |
| for p in llm.parameters(): |
| p.requires_grad = False |
| llm.eval() |
|
|
| if gradient_checkpointing: |
| llm.gradient_checkpointing_enable() |
| if verbose: |
| log("[LLM] Gradient checkpointing enabled") |
|
|
| return llm |
|
|
| @staticmethod |
| def apply_lora( |
| llm, |
| lora_config: 'LoRAConfig' |
| ): |
| """ |
| Apply LoRA to an LLM. |
| |
| Args: |
| llm: The LLM model |
| lora_config: LoRA configuration |
| |
| Returns: |
| LLM with LoRA applied |
| """ |
| from peft import get_peft_model |
| peft_config = lora_config.to_peft_config() |
| return get_peft_model(llm, peft_config) |
|
|
