🚀 Refined BitTransformerLM: Organized codebase with best practices

bit_transformer/BTLM_Extensions/lion_optimizer.py

@@ -0,0 +1,430 @@
+"""
+Lion Optimizer for BitTransformerLM Extensions
+==============================================
+
+Implementation of the Lion optimizer (EvoLved Sign Momentum).
+Based on "Symbolic Discovery of Optimization Algorithms" research.
+
+Key features:
+- Sign-based momentum updates
+- Extremely memory efficient (only stores momentum)
+- Often outperforms Adam/AdamW with larger learning rates
+- Compatible with BitTransformerLM's training infrastructure
+"""
+
+import torch
+from torch.optim.optimizer import Optimizer
+from typing import Any, Dict, List, Optional, Tuple, Union
+
+
+class Lion(Optimizer):
+    """
+    Lion optimizer implementation.
+    
+    Lion uses the sign of the interpolated momentum for parameter updates,
+    making it very memory efficient while maintaining competitive performance.
+    
+    Args:
+        params: Iterable of parameters to optimize
+        lr: Learning rate (default: 1e-4, typically needs to be smaller than Adam)
+        betas: Coefficients for computing momentum (default: (0.9, 0.99))
+        weight_decay: Weight decay coefficient (default: 0.0)
+        eps: Small constant for numerical stability (default: 1e-8)
+        maximize: Whether to maximize the objective (default: False)
+    """
+    
+    def __init__(
+        self,
+        params,
+        lr: float = 1e-4,
+        betas: Tuple[float, float] = (0.9, 0.99),
+        weight_decay: float = 0.0,
+        eps: float = 1e-8,
+        maximize: bool = False,
+    ):
+        if not 0.0 <= lr:
+            raise ValueError(f"Invalid learning rate: {lr}")
+        if not 0.0 <= betas[0] < 1.0:
+            raise ValueError(f"Invalid beta parameter at index 0: {betas[0]}")
+        if not 0.0 <= betas[1] < 1.0:
+            raise ValueError(f"Invalid beta parameter at index 1: {betas[1]}")
+        if not 0.0 <= weight_decay:
+            raise ValueError(f"Invalid weight_decay value: {weight_decay}")
+        if not 0.0 <= eps:
+            raise ValueError(f"Invalid epsilon value: {eps}")
+            
+        defaults = dict(
+            lr=lr,
+            betas=betas,
+            weight_decay=weight_decay,
+            eps=eps,
+            maximize=maximize,
+        )
+        super().__init__(params, defaults)
+    
+    @torch.no_grad()
+    def step(self, closure=None):
+        """Perform a single optimization step."""
+        loss = None
+        if closure is not None:
+            with torch.enable_grad():
+                loss = closure()
+                
+        for group in self.param_groups:
+            for p in group["params"]:
+                if p.grad is None:
+                    continue
+                    
+                grad = p.grad
+                if group["maximize"]:
+                    grad = -grad
+                    
+                if grad.dtype in {torch.float16, torch.bfloat16}:
+                    grad = grad.float()
+                    
+                state = self.state[p]
+                
+                # State initialization
+                if len(state) == 0:
+                    state["momentum"] = torch.zeros_like(p, memory_format=torch.preserve_format)
+                    
+                momentum = state["momentum"]
+                beta1, beta2 = group["betas"]
+                
+                # Weight decay (applied to parameters, not gradients)
+                if group["weight_decay"] != 0:
+                    p.mul_(1 - group["lr"] * group["weight_decay"])
+                
+                # Interpolate between momentum and gradient
+                # c_t = beta1 * m_{t-1} + (1 - beta1) * g_t
+                interpolated = momentum.mul(beta1).add_(grad, alpha=1 - beta1)
+                
+                # Update parameters using sign of interpolated momentum
+                # theta_t = theta_{t-1} - lr * sign(c_t)
+                p.add_(torch.sign(interpolated), alpha=-group["lr"])
+                
+                # Update momentum
+                # m_t = beta2 * m_{t-1} + (1 - beta2) * g_t  
+                momentum.mul_(beta2).add_(grad, alpha=1 - beta2)
+                
+        return loss
+
+
+def configure_lion_optimizer(
+    model: torch.nn.Module,
+    lr: float = 1e-4,
+    betas: Tuple[float, float] = (0.9, 0.99),
+    weight_decay: float = 0.01,
+    total_steps: Optional[int] = None,
+    warmup_ratio: float = 0.1,
+    **lion_kwargs
+) -> Tuple[Lion, Optional[torch.optim.lr_scheduler._LRScheduler]]:
+    """
+    Configure Lion optimizer with OneCycle learning rate schedule.
+    
+    This function provides a drop-in replacement for BitTransformerLM's
+    configure_optimizer function, using Lion instead of AdamW.
+    
+    Note: Lion typically works well with learning rates about 3-10x smaller
+    than Adam/AdamW, but higher weight decay (0.01-0.1).
+    
+    Args:
+        model: PyTorch model to optimize
+        lr: Peak learning rate (typically smaller than Adam)
+        betas: Beta coefficients for momentum computation
+        weight_decay: Weight decay coefficient (can be higher than Adam)
+        total_steps: Total training steps for OneCycle schedule
+        warmup_ratio: Fraction of steps for warmup
+        **lion_kwargs: Additional arguments for Lion optimizer
+        
+    Returns:
+        Tuple of (optimizer, scheduler)
+    """
+    # Filter parameters that need weight decay
+    decay_params = []
+    no_decay_params = []
+    
+    for name, param in model.named_parameters():
+        if not param.requires_grad:
+            continue
+        # Apply weight decay to weights but not biases/norms
+        if param.dim() >= 2:
+            decay_params.append(param)
+        else:
+            no_decay_params.append(param)
+    
+    param_groups = [
+        {"params": decay_params, "weight_decay": weight_decay},
+        {"params": no_decay_params, "weight_decay": 0.0},
+    ]
+    
+    optimizer = Lion(
+        param_groups,
+        lr=lr,
+        betas=betas,
+        **lion_kwargs
+    )
+    
+    scheduler = None
+    if total_steps is not None and total_steps > 0:
+        scheduler = torch.optim.lr_scheduler.OneCycleLR(
+            optimizer,
+            max_lr=lr,
+            total_steps=total_steps,
+            pct_start=warmup_ratio,
+            anneal_strategy='cos',
+            cycle_momentum=False,  # Lion doesn't use cycling momentum
+            div_factor=25.0,
+            final_div_factor=1e4,
+        )
+    
+    return optimizer, scheduler
+
+
+def create_lion_training_config(
+    lr: float = 1e-4,
+    betas: Tuple[float, float] = (0.9, 0.99),
+    weight_decay: float = 0.01,
+    **kwargs
+) -> Dict[str, Any]:
+    """
+    Create a training configuration dictionary for Lion optimizer.
+    
+    This can be used with BitTransformerLM's training scripts by passing
+    the config to the training loop.
+    
+    Args:
+        lr: Learning rate
+        betas: Beta coefficients for momentum
+        weight_decay: Weight decay coefficient
+        **kwargs: Additional configuration options
+        
+    Returns:
+        Dictionary containing training configuration
+    """
+    config = {
+        "optimizer_type": "lion",
+        "optimizer_config": {
+            "lr": lr,
+            "betas": betas,
+            "weight_decay": weight_decay,
+            **kwargs
+        },
+        "scheduler_type": "onecycle",
+    }
+    
+    return config
+
+
+class AdaptiveLion(Lion):
+    """
+    Enhanced Lion optimizer with adaptive learning rate scaling.
+    
+    This variant automatically adjusts the learning rate based on the
+    magnitude of gradients and momentum, potentially improving stability.
+    """
+    
+    def __init__(
+        self,
+        params,
+        lr: float = 1e-4,
+        betas: Tuple[float, float] = (0.9, 0.99),
+        weight_decay: float = 0.0,
+        eps: float = 1e-8,
+        maximize: bool = False,
+        adaptive_scale: float = 0.1,
+        min_scale: float = 0.01,
+        max_scale: float = 10.0,
+    ):
+        """
+        Args:
+            adaptive_scale: Scaling factor for adaptive adjustment
+            min_scale: Minimum learning rate scale
+            max_scale: Maximum learning rate scale
+        """
+        self.adaptive_scale = adaptive_scale
+        self.min_scale = min_scale
+        self.max_scale = max_scale
+        
+        super().__init__(params, lr, betas, weight_decay, eps, maximize)
+    
+    @torch.no_grad()
+    def step(self, closure=None):
+        """Perform optimization step with adaptive scaling."""
+        loss = None
+        if closure is not None:
+            with torch.enable_grad():
+                loss = closure()
+                
+        for group in self.param_groups:
+            for p in group["params"]:
+                if p.grad is None:
+                    continue
+                    
+                grad = p.grad
+                if group["maximize"]:
+                    grad = -grad
+                    
+                if grad.dtype in {torch.float16, torch.bfloat16}:
+                    grad = grad.float()
+                    
+                state = self.state[p]
+                
+                if len(state) == 0:
+                    state["momentum"] = torch.zeros_like(p, memory_format=torch.preserve_format)
+                    state["step"] = 0
+                    
+                momentum = state["momentum"]
+                state["step"] += 1
+                beta1, beta2 = group["betas"]
+                
+                # Adaptive learning rate based on gradient magnitude
+                grad_norm = grad.norm().item()
+                momentum_norm = momentum.norm().item()
+                
+                # Scale learning rate based on gradient/momentum ratio
+                if momentum_norm > 1e-8:
+                    scale = 1.0 + self.adaptive_scale * (grad_norm / momentum_norm - 1.0)
+                    scale = torch.clamp(torch.tensor(scale), self.min_scale, self.max_scale).item()
+                else:
+                    scale = 1.0
+                
+                adaptive_lr = group["lr"] * scale
+                
+                # Weight decay
+                if group["weight_decay"] != 0:
+                    p.mul_(1 - adaptive_lr * group["weight_decay"])
+                
+                # Lion update with adaptive learning rate
+                interpolated = momentum.mul(beta1).add_(grad, alpha=1 - beta1)
+                p.add_(torch.sign(interpolated), alpha=-adaptive_lr)
+                momentum.mul_(beta2).add_(grad, alpha=1 - beta2)
+                
+        return loss
+
+
+def configure_adaptive_lion_optimizer(
+    model: torch.nn.Module,
+    lr: float = 1e-4,
+    adaptive_scale: float = 0.1,
+    **kwargs
+) -> Tuple[AdaptiveLion, Optional[torch.optim.lr_scheduler._LRScheduler]]:
+    """Configure AdaptiveLion optimizer with learning rate scheduling."""
+    # Similar to configure_lion_optimizer but with AdaptiveLion
+    decay_params = []
+    no_decay_params = []
+    
+    for name, param in model.named_parameters():
+        if not param.requires_grad:
+            continue
+        if param.dim() >= 2:
+            decay_params.append(param)
+        else:
+            no_decay_params.append(param)
+    
+    param_groups = [
+        {"params": decay_params, "weight_decay": kwargs.get("weight_decay", 0.01)},
+        {"params": no_decay_params, "weight_decay": 0.0},
+    ]
+    
+    optimizer = AdaptiveLion(
+        param_groups,
+        lr=lr,
+        adaptive_scale=adaptive_scale,
+        **{k: v for k, v in kwargs.items() if k != "weight_decay"}
+    )
+    
+    scheduler = None
+    total_steps = kwargs.get("total_steps")
+    if total_steps is not None and total_steps > 0:
+        scheduler = torch.optim.lr_scheduler.OneCycleLR(
+            optimizer,
+            max_lr=lr,
+            total_steps=total_steps,
+            pct_start=kwargs.get("warmup_ratio", 0.1),
+            anneal_strategy='cos',
+            cycle_momentum=False,
+            div_factor=25.0,
+            final_div_factor=1e4,
+        )
+    
+    return optimizer, scheduler
+
+
+# Example usage and integration helpers
+def integrate_with_bittransformerlm():
+    """
+    Example of how to integrate Lion optimizer with BitTransformerLM training.
+    
+    Usage:
+        from BTLM_Extensions.lion_optimizer import configure_lion_optimizer
+        
+        # Replace the standard optimizer configuration
+        # Note: Lion typically needs smaller learning rates than Adam
+        optimizer, scheduler = configure_lion_optimizer(
+            model, lr=1e-4, weight_decay=0.01, total_steps=1000
+        )
+        
+        # Use in training loop
+        train_loop(model, data, optimizer=optimizer, scheduler=scheduler)
+        
+        # For adaptive version:
+        from BTLM_Extensions.lion_optimizer import configure_adaptive_lion_optimizer
+        
+        optimizer, scheduler = configure_adaptive_lion_optimizer(
+            model, lr=1e-4, adaptive_scale=0.1, total_steps=1000
+        )
+    """
+    pass
+
+
+if __name__ == "__main__":
+    # Simple test of the optimizer
+    import torch.nn as nn
+    
+    model = nn.Sequential(
+        nn.Linear(10, 20),
+        nn.ReLU(),
+        nn.Linear(20, 1)
+    )
+    
+    print("Testing standard Lion optimizer...")
+    optimizer, scheduler = configure_lion_optimizer(model, lr=1e-4, total_steps=100)
+    
+    # Simple training step
+    x = torch.randn(32, 10)
+    y = torch.randn(32, 1)
+    
+    pred = model(x)
+    loss = nn.functional.mse_loss(pred, y)
+    initial_loss = loss.item()
+    loss.backward()
+    
+    optimizer.step()
+    if scheduler:
+        scheduler.step()
+    
+    print(f"Initial loss: {initial_loss:.4f}")
+    
+    # Test adaptive version
+    print("Testing Adaptive Lion optimizer...")
+    model2 = nn.Sequential(
+        nn.Linear(10, 20),
+        nn.ReLU(),
+        nn.Linear(20, 1)
+    )
+    
+    optimizer2, scheduler2 = configure_adaptive_lion_optimizer(
+        model2, lr=1e-4, adaptive_scale=0.1, total_steps=100
+    )
+    
+    pred2 = model2(x)
+    loss2 = nn.functional.mse_loss(pred2, y)
+    loss2.backward()
+    optimizer2.step()
+    if scheduler2:
+        scheduler2.step()
+    
+    print("Lion optimizers test completed successfully!")
+    print(f"Standard Lion loss: {initial_loss:.4f}")
+    print(f"Adaptive Lion loss: {loss2.item():.4f}")