training/config.py

"""
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,
        )