""" Configuration management using dataclasses. Single Responsibility: Only handles configuration. """ from dataclasses import dataclass, field from typing import Optional, List import os import torch # ============================================================ # Constants (Single source of truth) # ============================================================ SNAC_BASE_OFFSET = 128266 SNAC_LAYERS_PER_FRAME = 7 SNAC_LAYER_OFFSET = 4096 EOS_TOKEN = 128009 # Model defaults DEFAULT_WHISPER_DIM = 1280 DEFAULT_LLM_DIM = 3072 DEFAULT_DOWNSAMPLE = 5 DEFAULT_INTERMEDIATE_DIM = 2048 DEFAULT_MODEL_PATH = "canopylabs/3b-es_it-ft-research_release" # LoRA defaults DEFAULT_LORA_R = 16 DEFAULT_LORA_ALPHA = 32 DEFAULT_LORA_DROPOUT = 0.05 DEFAULT_LORA_MODULES = [ "q_proj", "k_proj", "v_proj", "o_proj", "gate_proj", "up_proj", "down_proj" ] @dataclass class GPUConfig: """GPU configuration detected at runtime.""" name: str = "Unknown" vram_gb: int = 0 batch_size: int = 2 grad_accum: int = 16 dtype: torch.dtype = torch.float32 device_type: str = "cpu" @classmethod def auto_detect(cls) -> 'GPUConfig': """Detect GPU and return optimal configuration.""" config = cls() # Try CUDA (NVIDIA) if torch.cuda.is_available(): try: props = torch.cuda.get_device_properties(0) config.vram_gb = props.total_memory // (1024**3) config.name = props.name config.device_type = "cuda" config.dtype = torch.bfloat16 except Exception: pass # Try MPS (Apple Silicon) elif hasattr(torch.backends, 'mps') and torch.backends.mps.is_available(): config.name = "Apple Silicon (MPS)" config.device_type = "mps" config.dtype = torch.float32 try: import subprocess result = subprocess.run( ['sysctl', '-n', 'hw.memsize'], capture_output=True, text=True ) total_mem = int(result.stdout.strip()) // (1024**3) config.vram_gb = total_mem // 2 except Exception: config.vram_gb = 8 # Try ROCm (AMD) elif hasattr(torch, 'hip') or os.environ.get('ROCM_HOME'): try: if torch.cuda.is_available(): props = torch.cuda.get_device_properties(0) config.vram_gb = props.total_memory // (1024**3) config.name = f"AMD {props.name}" config.device_type = "cuda" config.dtype = torch.bfloat16 except Exception: config.name = "AMD ROCm" config.vram_gb = 16 # Fallback: nvidia-smi if config.vram_gb == 0: try: import subprocess result = subprocess.run( ['nvidia-smi', '--query-gpu=name,memory.total', '--format=csv,noheader,nounits'], capture_output=True, text=True ) lines = result.stdout.strip().split('\n') config.name, vram_mb = lines[0].split(', ') config.vram_gb = int(vram_mb) // 1024 config.device_type = "cuda" config.dtype = torch.bfloat16 except Exception: pass # Set batch size based on VRAM config.batch_size, config.grad_accum = cls._get_batch_config(config.vram_gb) return config @staticmethod def _get_batch_config(vram_gb: int) -> tuple: """Get optimal batch size and gradient accumulation based on VRAM.""" if vram_gb >= 140: # H200 (141GB) return 12, 3 # H200: batch=12, effective=36 elif vram_gb >= 80: return 6, 5 elif vram_gb >= 35: return 4, 8 elif vram_gb >= 16: return 2, 16 else: return 1, 32 @dataclass class TrainingConfig: """Training configuration with sensible defaults.""" # Data data_paths: List[str] = field(default_factory=list) output_dir: str = "./checkpoints" # Training hyperparameters learning_rate: float = 5e-5 epochs: int = 2 batch_size: Optional[int] = None grad_accum: Optional[int] = None warmup_ratio: float = 0.03 max_grad_norm: float = 1.0 label_smoothing: float = 0.1 weight_decay: float = 0.01 # Sequence limits max_audio_len: int = 500 max_seq_len: int = 2048 # Scheduled interleaving (IST-LM) initial_text_ratio: float = 0.9 decay_steps: int = 300 dynamic_decay: bool = False no_decay: bool = False # Stage 1: keep text_ratio fixed final_audio_portion: float = 0.2 # Model model_path: str = DEFAULT_MODEL_PATH # Checkpointing save_steps: int = 200 resume_from: Optional[str] = None # Memory vram_fraction: float = 0.80 ram_limit_gb: Optional[float] = None gradient_checkpointing: bool = False # Mode flags demo_mode: bool = False test_mode: bool = False def __post_init__(self): """Apply GPU auto-detection if batch_size not set.""" if self.batch_size is None or self.grad_accum is None: gpu_config = GPUConfig.auto_detect() if self.batch_size is None: self.batch_size = gpu_config.batch_size if self.grad_accum is None: self.grad_accum = gpu_config.grad_accum @classmethod def from_args(cls, args) -> 'TrainingConfig': """Create config from argparse namespace.""" return cls( data_paths=[p.strip() for p in args.data.split(",")], output_dir=args.output_dir, learning_rate=args.lr, epochs=args.epochs, batch_size=args.batch_size, grad_accum=args.grad_accum, warmup_ratio=args.warmup_ratio, max_grad_norm=args.max_grad_norm, label_smoothing=args.label_smoothing, max_seq_len=args.max_seq_len, initial_text_ratio=args.initial_text_ratio, decay_steps=args.decay_steps, dynamic_decay=getattr(args, 'dynamic_decay', False), no_decay=getattr(args, 'no_decay', False), final_audio_portion=getattr(args, 'final_audio_portion', 0.2), model_path=args.model_path, save_steps=args.save_steps, resume_from=args.resume, vram_fraction=args.vram_fraction, ram_limit_gb=args.ram_limit_gb, gradient_checkpointing=args.gradient_checkpointing, demo_mode=args.demo, test_mode=args.test, ) @dataclass class LoRAConfig: """LoRA configuration for Stage 2.""" r: int = DEFAULT_LORA_R alpha: int = DEFAULT_LORA_ALPHA dropout: float = DEFAULT_LORA_DROPOUT target_modules: List[str] = field(default_factory=lambda: DEFAULT_LORA_MODULES.copy()) bias: str = "none" def to_peft_config(self): """Convert to PEFT LoraConfig.""" from peft import LoraConfig as PeftLoraConfig, TaskType return PeftLoraConfig( r=self.r, lora_alpha=self.alpha, lora_dropout=self.dropout, target_modules=self.target_modules, bias=self.bias, task_type=TaskType.CAUSAL_LM, )