optimization_methods.py

"""
Optimization Methods Configuration

This module defines the available optimization methods for fine-tuning.
To add a new method, simply add it to the METHODS dictionary and implement
the required configuration functions.
"""

from typing import Dict, Any, Tuple
from transformers import AutoModelForCausalLM, AutoTokenizer
from peft import LoraConfig
import torch

class OptimizationMethod:
    """Base class for optimization methods."""
    
    def __init__(self, name: str, description: str):
        self.name = name
        self.description = description
    
    def load_model(self, pretrained_model: str, trust_remote_code: bool = True) -> Tuple[AutoModelForCausalLM, AutoTokenizer, str]:
        """Load model with this optimization method."""
        raise NotImplementedError
    
    def setup_lora_config(self, model_config: dict) -> LoraConfig:
        """Set up LoRA configuration for this method."""
        raise NotImplementedError
    
    def get_training_config(self, use_cuda: bool) -> Dict[str, Any]:
        """Get training configuration specific to this method."""
        raise NotImplementedError

class StandardLoRA(OptimizationMethod):
    """Standard LoRA fine-tuning method with low-rank matrix decomposition W = W₀ + BA."""
    
    def __init__(self):
        super().__init__("lora", "Standard LoRA fine-tuning with mixed precision")
    
    def load_model(self, pretrained_model: str, trust_remote_code: bool = True) -> Tuple[AutoModelForCausalLM, AutoTokenizer, str]:
        use_cuda = torch.cuda.is_available()
        
        tokenizer = AutoTokenizer.from_pretrained(pretrained_model, trust_remote_code=trust_remote_code)
        tokenizer.pad_token = tokenizer.eos_token or tokenizer.pad_token or "[PAD]"
        
        model = AutoModelForCausalLM.from_pretrained(
            pretrained_model,
            torch_dtype=torch.float16 if use_cuda else torch.float32,
            device_map="auto" if use_cuda else "cpu",
            trust_remote_code=trust_remote_code
        )
        
        return model, tokenizer, "lora"
    
    def setup_lora_config(self, model_config: dict) -> LoraConfig:
        return LoraConfig(
            r=model_config['lora_rank'],
            lora_alpha=model_config['lora_alpha'],
            lora_dropout=model_config['lora_dropout'],
            target_modules=["q_proj", "v_proj"],
            task_type="CAUSAL_LM"
        )
    
    def get_training_config(self, use_cuda: bool) -> Dict[str, Any]:
        return {
            "fp16": use_cuda,
            "bf16": False,
            "dataloader_pin_memory": use_cuda
        }

class QLoRA(OptimizationMethod):
    """QLoRA fine-tuning method with 4-bit quantization for memory efficiency."""
    
    def __init__(self):
        super().__init__("qlora", "QLoRA with 4-bit quantization and double quantization efficiency")
    
    def load_model(self, pretrained_model: str, trust_remote_code: bool = True) -> Tuple[AutoModelForCausalLM, AutoTokenizer, str]:
        from transformers import BitsAndBytesConfig
        
        use_cuda = torch.cuda.is_available()
        
        tokenizer = AutoTokenizer.from_pretrained(pretrained_model, trust_remote_code=trust_remote_code)
        tokenizer.pad_token = tokenizer.eos_token or tokenizer.pad_token or "[PAD]"
        
        if use_cuda:
            quantization_config = BitsAndBytesConfig(
                load_in_4bit=True,
                bnb_4bit_compute_dtype=torch.float16,
                bnb_4bit_quant_type="nf4",
                bnb_4bit_use_double_quant=True
            )
            model = AutoModelForCausalLM.from_pretrained(
                pretrained_model,
                quantization_config=quantization_config,
                torch_dtype=torch.float16,
                device_map="auto",
                trust_remote_code=trust_remote_code
            )
            return model, tokenizer, "qlora"
        else:
            # Fallback to standard LoRA if CUDA not available
            return StandardLoRA().load_model(pretrained_model, trust_remote_code)
    
    def setup_lora_config(self, model_config: dict) -> LoraConfig:
        return LoraConfig(
            r=model_config['lora_rank'],
            lora_alpha=model_config['lora_alpha'],
            lora_dropout=model_config['lora_dropout'],
            target_modules=["q_proj", "v_proj"],
            task_type="CAUSAL_LM",
            use_rslora=True
        )
    
    def get_training_config(self, use_cuda: bool) -> Dict[str, Any]:
        return {
            "fp16": False,  # Don't use FP16 with QLoRA
            "bf16": use_cuda,  # Use BF16 if available
            "dataloader_pin_memory": use_cuda
        }

class DoRA(OptimizationMethod):
    """DoRA - Weight magnitude-direction decomposition for better performance."""
    
    def __init__(self):
        super().__init__("dora", "DoRA with weight magnitude-direction decomposition")
    
    def load_model(self, pretrained_model: str, trust_remote_code: bool = True) -> Tuple[AutoModelForCausalLM, AutoTokenizer, str]:
        use_cuda = torch.cuda.is_available()
        
        tokenizer = AutoTokenizer.from_pretrained(pretrained_model, trust_remote_code=trust_remote_code)
        tokenizer.pad_token = tokenizer.eos_token or tokenizer.pad_token or "[PAD]"
        
        model = AutoModelForCausalLM.from_pretrained(
            pretrained_model,
            torch_dtype=torch.float16 if use_cuda else torch.float32,
            device_map="auto" if use_cuda else "cpu",
            trust_remote_code=trust_remote_code
        )
        
        return model, tokenizer, "dora"
    
    def setup_lora_config(self, model_config: dict) -> LoraConfig:
        return LoraConfig(
            r=model_config['lora_rank'],
            lora_alpha=model_config['lora_alpha'],
            lora_dropout=model_config['lora_dropout'],
            target_modules=["q_proj", "v_proj", "k_proj", "o_proj"],
            task_type="CAUSAL_LM",
            use_dora=True  # Enable DoRA
        )
    
    def get_training_config(self, use_cuda: bool) -> Dict[str, Any]:
        return {
            "fp16": use_cuda,
            "bf16": False,
            "dataloader_pin_memory": use_cuda,
            "gradient_checkpointing": True
        }

class AdaLoRA(OptimizationMethod):
    """AdaLoRA with adaptive rank allocation using SVD-based dynamic rank adjustment."""
    
    def __init__(self):
        super().__init__("adalora", "AdaLoRA with adaptive rank allocation and SVD-based optimization")
    
    def load_model(self, pretrained_model: str, trust_remote_code: bool = True) -> Tuple[AutoModelForCausalLM, AutoTokenizer, str]:
        use_cuda = torch.cuda.is_available()
        
        tokenizer = AutoTokenizer.from_pretrained(pretrained_model, trust_remote_code=trust_remote_code)
        tokenizer.pad_token = tokenizer.eos_token or tokenizer.pad_token or "[PAD]"
        
        model = AutoModelForCausalLM.from_pretrained(
            pretrained_model,
            torch_dtype=torch.float16 if use_cuda else torch.float32,
            device_map="auto" if use_cuda else "cpu",
            trust_remote_code=trust_remote_code
        )
        
        return model, tokenizer, "adalora"
    
    def setup_lora_config(self, model_config: dict) -> LoraConfig:
        return LoraConfig(
            r=model_config['lora_rank'],
            lora_alpha=model_config['lora_alpha'],
            lora_dropout=model_config['lora_dropout'],
            target_modules=["q_proj", "v_proj", "k_proj", "o_proj"],
            task_type="CAUSAL_LM",
            target_modules_rank=model_config.get('target_modules_rank', model_config['lora_rank']),
            init_lora_weights=True
        )
    
    def get_training_config(self, use_cuda: bool) -> Dict[str, Any]:
        return {
            "fp16": use_cuda,
            "bf16": False,
            "dataloader_pin_memory": use_cuda,
            "gradient_checkpointing": True,
            "max_grad_norm": 1.0
        }

class LoRAPlus(OptimizationMethod):
    """LoRA+ with differential learning rates for A/B matrices."""
    
    def __init__(self):
        super().__init__("lora_plus", "LoRA+ with differential learning rates λₐ ≠ λᵦ")
    
    def load_model(self, pretrained_model: str, trust_remote_code: bool = True) -> Tuple[AutoModelForCausalLM, AutoTokenizer, str]:
        use_cuda = torch.cuda.is_available()
        
        tokenizer = AutoTokenizer.from_pretrained(pretrained_model, trust_remote_code=trust_remote_code)
        tokenizer.pad_token = tokenizer.eos_token or tokenizer.pad_token or "[PAD]"
        
        model = AutoModelForCausalLM.from_pretrained(
            pretrained_model,
            torch_dtype=torch.float16 if use_cuda else torch.float32,
            device_map="auto" if use_cuda else "cpu",
            trust_remote_code=trust_remote_code
        )
        
        return model, tokenizer, "lora_plus"
    
    def setup_lora_config(self, model_config: dict) -> LoraConfig:
        return LoraConfig(
            r=model_config['lora_rank'],
            lora_alpha=model_config['lora_alpha'],
            lora_dropout=model_config['lora_dropout'],
            target_modules=["q_proj", "v_proj"],
            task_type="CAUSAL_LM",
            use_rslora=True
        )
    
    def get_training_config(self, use_cuda: bool) -> Dict[str, Any]:
        return {
            "fp16": use_cuda,
            "bf16": False,
            "dataloader_pin_memory": use_cuda,
            "optim": "adamw_torch",
            "learning_rate": 2e-4,
            "weight_decay": 0.01
        }

class DeltaLoRA(OptimizationMethod):
    """Delta-LoRA for continual learning optimization with parameter difference storage."""
    
    def __init__(self):
        super().__init__("delta_lora", "Delta-LoRA with continual learning and parameter difference storage")
    
    def load_model(self, pretrained_model: str, trust_remote_code: bool = True) -> Tuple[AutoModelForCausalLM, AutoTokenizer, str]:
        use_cuda = torch.cuda.is_available()
        
        tokenizer = AutoTokenizer.from_pretrained(pretrained_model, trust_remote_code=trust_remote_code)
        tokenizer.pad_token = tokenizer.eos_token or tokenizer.pad_token or "[PAD]"
        
        model = AutoModelForCausalLM.from_pretrained(
            pretrained_model,
            torch_dtype=torch.float16 if use_cuda else torch.float32,
            device_map="auto" if use_cuda else "cpu",
            trust_remote_code=trust_remote_code
        )
        
        return model, tokenizer, "delta_lora"
    
    def setup_lora_config(self, model_config: dict) -> LoraConfig:
        return LoraConfig(
            r=model_config['lora_rank'],
            lora_alpha=model_config['lora_alpha'],
            lora_dropout=model_config['lora_dropout'],
            target_modules=["q_proj", "v_proj", "k_proj", "o_proj"],
            task_type="CAUSAL_LM",
            inference_mode=False
        )
    
    def get_training_config(self, use_cuda: bool) -> Dict[str, Any]:
        return {
            "fp16": use_cuda,
            "bf16": False,
            "dataloader_pin_memory": use_cuda,
            "save_strategy": "steps",
            "save_steps": 100,
            "save_total_limit": 5
        }

class MultiLoRA(OptimizationMethod):
    """Multi-LoRA for multi-task adapter system with parallel LoRA modules."""
    
    def __init__(self):
        super().__init__("multi_lora", "Multi-LoRA with parallel adapter modules for different domains")
    
    def load_model(self, pretrained_model: str, trust_remote_code: bool = True) -> Tuple[AutoModelForCausalLM, AutoTokenizer, str]:
        use_cuda = torch.cuda.is_available()
        
        tokenizer = AutoTokenizer.from_pretrained(pretrained_model, trust_remote_code=trust_remote_code)
        tokenizer.pad_token = tokenizer.eos_token or tokenizer.pad_token or "[PAD]"
        
        model = AutoModelForCausalLM.from_pretrained(
            pretrained_model,
            torch_dtype=torch.float16 if use_cuda else torch.float32,
            device_map="auto" if use_cuda else "cpu",
            trust_remote_code=trust_remote_code
        )
        
        return model, tokenizer, "multi_lora"
    
    def setup_lora_config(self, model_config: dict) -> LoraConfig:
        return LoraConfig(
            r=model_config['lora_rank'],
            lora_alpha=model_config['lora_alpha'],
            lora_dropout=model_config['lora_dropout'],
            target_modules=["q_proj", "v_proj"],
            task_type="CAUSAL_LM",
            modules_to_save=["embed_tokens", "lm_head"]
        )
    
    def get_training_config(self, use_cuda: bool) -> Dict[str, Any]:
        return {
            "fp16": use_cuda,
            "bf16": False,
            "dataloader_pin_memory": use_cuda,
            "gradient_checkpointing": True
        }

class OLoRA(OptimizationMethod):
    """OLoRA with orthogonal constraint adaptation using Gram-Schmidt optimization."""
    
    def __init__(self):
        super().__init__("olora", "OLoRA with orthogonal constraints and Gram-Schmidt optimization")
    
    def load_model(self, pretrained_model: str, trust_remote_code: bool = True) -> Tuple[AutoModelForCausalLM, AutoTokenizer, str]:
        use_cuda = torch.cuda.is_available()
        
        tokenizer = AutoTokenizer.from_pretrained(pretrained_model, trust_remote_code=trust_remote_code)
        tokenizer.pad_token = tokenizer.eos_token or tokenizer.pad_token or "[PAD]"
        
        model = AutoModelForCausalLM.from_pretrained(
            pretrained_model,
            torch_dtype=torch.float16 if use_cuda else torch.float32,
            device_map="auto" if use_cuda else "cpu",
            trust_remote_code=trust_remote_code
        )
        
        return model, tokenizer, "olora"
    
    def setup_lora_config(self, model_config: dict) -> LoraConfig:
        return LoraConfig(
            r=model_config['lora_rank'],
            lora_alpha=model_config['lora_alpha'],
            lora_dropout=model_config['lora_dropout'],
            target_modules=["q_proj", "v_proj"],
            task_type="CAUSAL_LM",
            use_rslora=True
        )
    
    def get_training_config(self, use_cuda: bool) -> Dict[str, Any]:
        return {
            "fp16": use_cuda,
            "bf16": False,
            "dataloader_pin_memory": use_cuda,
            "optim": "adamw_torch",
            "learning_rate": 1e-4
        }

class VeRA(OptimizationMethod):
    """VeRA with vector-based random adaptation and shared random matrices."""
    
    def __init__(self):
        super().__init__("vera", "VeRA with vector-based random adaptation and parameter efficiency")
    
    def load_model(self, pretrained_model: str, trust_remote_code: bool = True) -> Tuple[AutoModelForCausalLM, AutoTokenizer, str]:
        use_cuda = torch.cuda.is_available()
        
        tokenizer = AutoTokenizer.from_pretrained(pretrained_model, trust_remote_code=trust_remote_code)
        tokenizer.pad_token = tokenizer.eos_token or tokenizer.pad_token or "[PAD]"
        
        model = AutoModelForCausalLM.from_pretrained(
            pretrained_model,
            torch_dtype=torch.float16 if use_cuda else torch.float32,
            device_map="auto" if use_cuda else "cpu",
            trust_remote_code=trust_remote_code
        )
        
        return model, tokenizer, "vera"
    
    def setup_lora_config(self, model_config: dict) -> LoraConfig:
        return LoraConfig(
            r=model_config['lora_rank'],
            lora_alpha=model_config['lora_alpha'],
            lora_dropout=model_config['lora_dropout'],
            target_modules=["q_proj", "v_proj"],
            task_type="CAUSAL_LM",
            use_rslora=True
        )
    
    def get_training_config(self, use_cuda: bool) -> Dict[str, Any]:
        return {
            "fp16": use_cuda,
            "bf16": False,
            "dataloader_pin_memory": use_cuda,
            "learning_rate": 3e-4
        }

class LoRAFA(OptimizationMethod):
    """LoRA-FA with frozen-A initialization strategy and fixed A matrix."""
    
    def __init__(self):
        super().__init__("lora_fa", "LoRA-FA with frozen A matrix and trainable B optimization")
    
    def load_model(self, pretrained_model: str, trust_remote_code: bool = True) -> Tuple[AutoModelForCausalLM, AutoTokenizer, str]:
        use_cuda = torch.cuda.is_available()
        
        tokenizer = AutoTokenizer.from_pretrained(pretrained_model, trust_remote_code=trust_remote_code)
        tokenizer.pad_token = tokenizer.eos_token or tokenizer.pad_token or "[PAD]"
        
        model = AutoModelForCausalLM.from_pretrained(
            pretrained_model,
            torch_dtype=torch.float16 if use_cuda else torch.float32,
            device_map="auto" if use_cuda else "cpu",
            trust_remote_code=trust_remote_code
        )
        
        return model, tokenizer, "lora_fa"
    
    def setup_lora_config(self, model_config: dict) -> LoraConfig:
        return LoraConfig(
            r=model_config['lora_rank'],
            lora_alpha=model_config['lora_alpha'],
            lora_dropout=model_config['lora_dropout'],
            target_modules=["q_proj", "v_proj"],
            task_type="CAUSAL_LM",
            inference_mode=False
        )
    
    def get_training_config(self, use_cuda: bool) -> Dict[str, Any]:
        return {
            "fp16": use_cuda,
            "bf16": False,
            "dataloader_pin_memory": use_cuda,
            "learning_rate": 1e-4
        }

class LoRADrop(OptimizationMethod):
    """LoRA-drop with structured dropout patterns for regularization."""
    
    def __init__(self):
        super().__init__("lora_drop", "LoRA-drop with structured dropout and selective adapter deactivation")
    
    def load_model(self, pretrained_model: str, trust_remote_code: bool = True) -> Tuple[AutoModelForCausalLM, AutoTokenizer, str]:
        use_cuda = torch.cuda.is_available()
        
        tokenizer = AutoTokenizer.from_pretrained(pretrained_model, trust_remote_code=trust_remote_code)
        tokenizer.pad_token = tokenizer.eos_token or tokenizer.pad_token or "[PAD]"
        
        model = AutoModelForCausalLM.from_pretrained(
            pretrained_model,
            torch_dtype=torch.float16 if use_cuda else torch.float32,
            device_map="auto" if use_cuda else "cpu",
            trust_remote_code=trust_remote_code
        )
        
        return model, tokenizer, "lora_drop"
    
    def setup_lora_config(self, model_config: dict) -> LoraConfig:
        return LoraConfig(
            r=model_config['lora_rank'],
            lora_alpha=model_config['lora_alpha'],
            lora_dropout=model_config.get('lora_dropout', 0.3),  # Higher dropout for LoRA-drop
            target_modules=["q_proj", "v_proj"],
            task_type="CAUSAL_LM"
        )
    
    def get_training_config(self, use_cuda: bool) -> Dict[str, Any]:
        return {
            "fp16": use_cuda,
            "bf16": False,
            "dataloader_pin_memory": use_cuda,
            "gradient_checkpointing": True
        }

class TiedLoRA(OptimizationMethod):
    """Tied-LoRA with parameter sharing across layers for memory efficiency."""
    
    def __init__(self):
        super().__init__("tied_lora", "Tied-LoRA with parameter sharing and weight reuse")
    
    def load_model(self, pretrained_model: str, trust_remote_code: bool = True) -> Tuple[AutoModelForCausalLM, AutoTokenizer, str]:
        use_cuda = torch.cuda.is_available()
        
        tokenizer = AutoTokenizer.from_pretrained(pretrained_model, trust_remote_code=trust_remote_code)
        tokenizer.pad_token = tokenizer.eos_token or tokenizer.pad_token or "[PAD]"
        
        model = AutoModelForCausalLM.from_pretrained(
            pretrained_model,
            torch_dtype=torch.float16 if use_cuda else torch.float32,
            device_map="auto" if use_cuda else "cpu",
            trust_remote_code=trust_remote_code
        )
        
        return model, tokenizer, "tied_lora"
    
    def setup_lora_config(self, model_config: dict) -> LoraConfig:
        return LoraConfig(
            r=model_config['lora_rank'],
            lora_alpha=model_config['lora_alpha'],
            lora_dropout=model_config['lora_dropout'],
            target_modules=["q_proj", "v_proj"],
            task_type="CAUSAL_LM",
            inference_mode=False
        )
    
    def get_training_config(self, use_cuda: bool) -> Dict[str, Any]:
        return {
            "fp16": use_cuda,
            "bf16": False,
            "dataloader_pin_memory": use_cuda,
            "gradient_checkpointing": True
        }

class LoRAXS(OptimizationMethod):
    """LoRA-XS with extreme parameter compression maintaining performance."""
    
    def __init__(self):
        super().__init__("lora_xs", "LoRA-XS with <0.01% parameter ratio and performance maintenance")
    
    def load_model(self, pretrained_model: str, trust_remote_code: bool = True) -> Tuple[AutoModelForCausalLM, AutoTokenizer, str]:
        use_cuda = torch.cuda.is_available()
        
        tokenizer = AutoTokenizer.from_pretrained(pretrained_model, trust_remote_code=trust_remote_code)
        tokenizer.pad_token = tokenizer.eos_token or tokenizer.pad_token or "[PAD]"
        
        model = AutoModelForCausalLM.from_pretrained(
            pretrained_model,
            torch_dtype=torch.float16 if use_cuda else torch.float32,
            device_map="auto" if use_cuda else "cpu",
            trust_remote_code=trust_remote_code
        )
        
        return model, tokenizer, "lora_xs"
    
    def setup_lora_config(self, model_config: dict) -> LoraConfig:
        # Use very low rank for extreme compression
        rank = min(model_config.get('lora_rank', 4), 4)
        return LoraConfig(
            r=rank,
            lora_alpha=rank * 2,  # Smaller alpha for XS
            lora_dropout=model_config.get('lora_dropout', 0.1),
            target_modules=["q_proj", "v_proj"],
            task_type="CAUSAL_LM"
        )
    
    def get_training_config(self, use_cuda: bool) -> Dict[str, Any]:
        return {
            "fp16": use_cuda,
            "bf16": False,
            "dataloader_pin_memory": use_cuda,
            "learning_rate": 5e-4,  # Higher LR for XS
            "gradient_accumulation_steps": 8
        }

class HLoRA(OptimizationMethod):
    """HLoRA with hierarchical low-rank structures for complex tasks."""
    
    def __init__(self):
        super().__init__("hlora", "HLoRA with hierarchical low-rank structures and multi-level adaptation")
    
    def load_model(self, pretrained_model: str, trust_remote_code: bool = True) -> Tuple[AutoModelForCausalLM, AutoTokenizer, str]:
        use_cuda = torch.cuda.is_available()
        
        tokenizer = AutoTokenizer.from_pretrained(pretrained_model, trust_remote_code=trust_remote_code)
        tokenizer.pad_token = tokenizer.eos_token or tokenizer.pad_token or "[PAD]"
        
        model = AutoModelForCausalLM.from_pretrained(
            pretrained_model,
            torch_dtype=torch.float16 if use_cuda else torch.float32,
            device_map="auto" if use_cuda else "cpu",
            trust_remote_code=trust_remote_code
        )
        
        return model, tokenizer, "hlora"
    
    def setup_lora_config(self, model_config: dict) -> LoraConfig:
        return LoraConfig(
            r=model_config['lora_rank'],
            lora_alpha=model_config['lora_alpha'],
            lora_dropout=model_config['lora_dropout'],
            target_modules=["q_proj", "v_proj", "k_proj", "o_proj", "gate_proj", "up_proj", "down_proj"],
            task_type="CAUSAL_LM",
            inference_mode=False
        )
    
    def get_training_config(self, use_cuda: bool) -> Dict[str, Any]:
        return {
            "fp16": use_cuda,
            "bf16": False,
            "dataloader_pin_memory": use_cuda,
            "gradient_checkpointing": True,
            "max_grad_norm": 0.5
        }

class GLoRA(OptimizationMethod):
    """GLoRA with group-wise low-rank adaptation and clustered parameter updates."""
    
    def __init__(self):
        super().__init__("glora", "GLoRA with group-wise adaptation and clustered parameter updates")
    
    def load_model(self, pretrained_model: str, trust_remote_code: bool = True) -> Tuple[AutoModelForCausalLM, AutoTokenizer, str]:
        use_cuda = torch.cuda.is_available()
        
        tokenizer = AutoTokenizer.from_pretrained(pretrained_model, trust_remote_code=trust_remote_code)
        tokenizer.pad_token = tokenizer.eos_token or tokenizer.pad_token or "[PAD]"
        
        model = AutoModelForCausalLM.from_pretrained(
            pretrained_model,
            torch_dtype=torch.float16 if use_cuda else torch.float32,
            device_map="auto" if use_cuda else "cpu",
            trust_remote_code=trust_remote_code
        )
        
        return model, tokenizer, "glora"
    
    def setup_lora_config(self, model_config: dict) -> LoraConfig:
        return LoraConfig(
            r=model_config['lora_rank'],
            lora_alpha=model_config['lora_alpha'],
            lora_dropout=model_config['lora_dropout'],
            target_modules=["q_proj", "v_proj", "k_proj", "o_proj"],
            task_type="CAUSAL_LM",
            inference_mode=False
        )
    
    def get_training_config(self, use_cuda: bool) -> Dict[str, Any]:
        return {
            "fp16": use_cuda,
            "bf16": False,
            "dataloader_pin_memory": use_cuda,
            "gradient_checkpointing": True,
            "learning_rate": 2e-4
        }

# Registry of available optimization methods
METHODS = {
    "lora": StandardLoRA(),
    "qlora": QLoRA(),
    "dora": DoRA(),
    "adalora": AdaLoRA(),
    "lora_plus": LoRAPlus(),
    "delta_lora": DeltaLoRA(),
    "multi_lora": MultiLoRA(),
    "olora": OLoRA(),
    "vera": VeRA(),
    "lora_fa": LoRAFA(),
    "lora_drop": LoRADrop(),
    "tied_lora": TiedLoRA(),
    "lora_xs": LoRAXS(),
    "hlora": HLoRA(),
    "glora": GLoRA()
}

def get_available_methods() -> list:
    """Get list of available optimization method names."""
    return list(METHODS.keys())

def get_method(method_name: str) -> OptimizationMethod:
    """Get optimization method by name."""
    if method_name.lower() not in METHODS:
        raise ValueError(f"Unknown optimization method: {method_name}. Available: {', '.join(get_available_methods())}")
    return METHODS[method_name.lower()]

def add_method(method_name: str, method: OptimizationMethod):
    """Add a new optimization method to the registry."""
    METHODS[method_name.lower()] = method

def get_method_info(method_name: str) -> Dict[str, str]:
    """Get detailed information about a specific method."""
    method = get_method(method_name)
    return {
        "name": method.name,
        "description": method.description,
        "key_features": get_method_key_features(method_name)
    }

def get_method_key_features(method_name: str) -> list:
    """Get key features and characteristics of each method."""
    features = {
        "lora": ["Low-rank decomposition", "Mixed precision", "Standard approach"],
        "qlora": ["4-bit quantization", "Memory efficient", "CUDA optimized"],
        "dora": ["Magnitude-direction decomposition", "Better performance", "Enhanced adaptation"],
        "adalora": ["Adaptive rank allocation", "SVD-based optimization", "Dynamic adjustment"],
        "lora_plus": ["Differential learning rates", "A/B matrix optimization", "Enhanced training"],
        "delta_lora": ["Continual learning", "Parameter difference storage", "Task sequences"],
        "multi_lora": ["Multi-task adapters", "Parallel modules", "Domain specialization"],
        "olora": ["Orthogonal constraints", "Gram-Schmidt optimization", "Reduced interference"],
        "vera": ["Vector-based adaptation", "Shared random matrices", "Parameter efficiency"],
        "lora_fa": ["Frozen-A strategy", "Fixed A matrix", "B optimization only"],
        "lora_drop": ["Structured dropout", "Selective deactivation", "Regularization"],
        "tied_lora": ["Parameter sharing", "Weight reuse", "Memory efficiency"],
        "lora_xs": ["Extreme compression", "<0.01% parameters", "Performance maintenance"],
        "hlora": ["Hierarchical structures", "Multi-level adaptation", "Complex tasks"],
        "glora": ["Group-wise adaptation", "Clustered updates", "Efficient training"]
    }
    return features.get(method_name.lower(), ["No specific features documented"])