anigen/utils/optimizer_utils.py

"""
Optimizer utilities for PyTorch Lightning systems.
"""

import torch
import torch.nn as nn
from typing import Dict, List, Any


def create_muon_optimizer(models: Dict[str, nn.Module], optimizer_args: Dict[str, Any]):
    """
    Create Muon optimizer with different learning rates for different parameters.
    
    Args:
        models: Dictionary of models
        optimizer_args: Optimizer configuration with muon_lr, muon_weight_decay, 
                       other_lr, other_weight_decay
    
    Returns:
        Configured optimizer
    """
    muon_lr = optimizer_args.get('muon_lr', 0.02)
    muon_weight_decay = optimizer_args.get('muon_weight_decay', 0.01)
    other_lr = optimizer_args.get('other_lr', 1e-4)
    other_weight_decay = optimizer_args.get('other_weight_decay', 0.01)
    
    # Separate parameters for Muon and other optimizers
    muon_params = []
    other_params = []
    
    for name, model in models.items():
        for param_name, param in model.named_parameters():
            if not param.requires_grad:
                continue
                
            # Define which parameters should use Muon optimizer
            # This is a heuristic - adjust based on your model architecture
            if ('weight' in param_name and 
                param.dim() >= 2 and 
                param.numel() >= 1024):  # Large weight matrices
                muon_params.append(param)
            else:
                other_params.append(param)
    
    # Try to import and use Muon optimizer
    try:
        from muon import Muon
        
        # Create parameter groups
        param_groups = []
        
        if muon_params:
            param_groups.append({
                'params': muon_params,
                'lr': muon_lr,
                'weight_decay': muon_weight_decay,
                'momentum': 0.95,  # Muon-specific parameter
            })
        
        if other_params:
            param_groups.append({
                'params': other_params,
                'lr': other_lr,
                'weight_decay': other_weight_decay,
            })
        
        return Muon(param_groups)
        
    except ImportError:
        print("Warning: Muon optimizer not available. Falling back to AdamW.")
        # Fallback to AdamW with combined parameters
        all_params = muon_params + other_params
        return torch.optim.AdamW(all_params, lr=other_lr, weight_decay=other_weight_decay)


def create_optimizer(models: Dict[str, nn.Module], optimizer_config: Dict[str, Any]):
    """
    Create optimizer based on configuration.
    
    Args:
        models: Dictionary of models
        optimizer_config: Optimizer configuration
    
    Returns:
        Configured optimizer
    """
    optimizer_name = optimizer_config.get('name', 'Adam').lower()
    optimizer_args = optimizer_config.get('args', {})
    
    # Get all parameters
    params = []
    for model in models.values():
        params.extend([p for p in model.parameters() if p.requires_grad])
    
    if optimizer_name == 'adam':
        return torch.optim.Adam(params, **optimizer_args)
    elif optimizer_name == 'adamw':
        return torch.optim.AdamW(params, **optimizer_args)
    elif optimizer_name == 'sgd':
        return torch.optim.SGD(params, **optimizer_args)
    elif optimizer_name == 'muon':
        return create_muon_optimizer(models, optimizer_args)
    else:
        raise ValueError(f"Unknown optimizer: {optimizer_name}")


def create_lr_scheduler(optimizer, lr_scheduler_config: Dict[str, Any]):
    """
    Create learning rate scheduler based on configuration.
    
    Args:
        optimizer: The optimizer to schedule
        lr_scheduler_config: Scheduler configuration
    
    Returns:
        Configured scheduler
    """
    scheduler_name = lr_scheduler_config.get('name', 'CosineAnnealingLR')
    scheduler_args = lr_scheduler_config.get('args', {})
    
    if scheduler_name == 'CosineAnnealingLR':
        return torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, **scheduler_args)
    elif scheduler_name == 'LinearLR':
        return torch.optim.lr_scheduler.LinearLR(optimizer, **scheduler_args)
    elif scheduler_name == 'ExponentialLR':
        return torch.optim.lr_scheduler.ExponentialLR(optimizer, **scheduler_args)
    elif scheduler_name == 'StepLR':
        return torch.optim.lr_scheduler.StepLR(optimizer, **scheduler_args)
    elif scheduler_name == 'MultiStepLR':
        return torch.optim.lr_scheduler.MultiStepLR(optimizer, **scheduler_args)
    elif scheduler_name == 'SequentialLR':
        # Handle sequential scheduler
        schedulers = []
        for sched_config in lr_scheduler_config['schedulers']:
            sched = create_single_scheduler(sched_config, optimizer)
            schedulers.append(sched)
        return torch.optim.lr_scheduler.SequentialLR(
            optimizer, schedulers, **scheduler_args
        )
    else:
        raise ValueError(f"Unknown scheduler: {scheduler_name}")


def create_single_scheduler(scheduler_config: Dict[str, Any], optimizer):
    """Create a single scheduler for SequentialLR."""
    name = scheduler_config['name']
    args = scheduler_config['args']
    
    if name == 'LinearLR':
        return torch.optim.lr_scheduler.LinearLR(optimizer, **args)
    elif name == 'CosineAnnealingLR':
        return torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, **args)
    elif name == 'ExponentialLR':
        return torch.optim.lr_scheduler.ExponentialLR(optimizer, **args)
    elif name == 'StepLR':
        return torch.optim.lr_scheduler.StepLR(optimizer, **args)
    elif name == 'MultiStepLR':
        return torch.optim.lr_scheduler.MultiStepLR(optimizer, **args)
    else:
        raise ValueError(f"Unknown scheduler: {name}")


class AdaptiveGradClipper:
    """
    Adaptive gradient clipping based on percentile of gradient norms.
    """
    
    def __init__(self, max_norm: float = 1.0, clip_percentile: float = 95):
        self.max_norm = max_norm
        self.clip_percentile = clip_percentile
        self.grad_norms_history = []
        self.history_size = 1000
    
    def __call__(self, parameters):
        """Apply adaptive gradient clipping."""
        if isinstance(parameters, torch.Tensor):
            parameters = [parameters]
        parameters = [p for p in parameters if p.grad is not None]
        
        if not parameters:
            return torch.tensor(0.0)
        
        # Calculate gradient norm
        total_norm = torch.norm(
            torch.stack([torch.norm(p.grad.detach()) for p in parameters])
        )
        
        # Update history
        self.grad_norms_history.append(total_norm.item())
        if len(self.grad_norms_history) > self.history_size:
            self.grad_norms_history.pop(0)
        
        # Calculate adaptive threshold
        if len(self.grad_norms_history) >= 10:
            threshold = torch.quantile(
                torch.tensor(self.grad_norms_history), 
                self.clip_percentile / 100.0
            )
            clip_value = min(self.max_norm, threshold.item())
        else:
            clip_value = self.max_norm
        
        # Apply clipping
        if total_norm > clip_value:
            clip_coef = clip_value / (total_norm + 1e-6)
            for p in parameters:
                p.grad.detach().mul_(clip_coef)
        
        return total_norm


def create_grad_clipper(grad_clip_config: Dict[str, Any]):
    """Create gradient clipper based on configuration."""
    if grad_clip_config is None:
        return None
    
    name = grad_clip_config.get('name', 'norm')
    args = grad_clip_config.get('args', {})
    
    if name.lower() == 'norm':
        max_norm = args.get('max_norm', 1.0)
        return lambda params: torch.nn.utils.clip_grad_norm_(params, max_norm)
    elif name.lower() == 'adaptivegradclipper':
        return AdaptiveGradClipper(**args)
    else:
        raise ValueError(f"Unknown gradient clipper: {name}")