""" 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 # Ensure T is divisible by downsample T_new = (T // self.downsample) * self.downsample x = x[:, :T_new] # Reshape: [B, T, D] -> [B, T // downsample, D * downsample] 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 # Determine attention implementation # Flash Attention 2 provides 10-20x memory savings on long sequences # Requires: Ampere/Ada/Hopper GPU (RTX 30xx, 40xx, A100, H100) attn_impl = "sdpa" # Default fallback if use_flash_attention and torch.cuda.is_available(): try: # Actually try to import flash_attn to verify it works import flash_attn # Check GPU capability (Flash Attention 2 requires SM 80+) major, _ = torch.cuda.get_device_capability() if major >= 8: # Ampere or newer (RTX 30xx, 40xx, A100, H100) 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") # Load model with optimizations try: llm = AutoModelForCausalLM.from_pretrained( model_path, torch_dtype=dtype, attn_implementation=attn_impl, ) except Exception as e: # Fallback if flash_attention_2 fails 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)