🚀 Refined BitTransformerLM: Organized codebase with best practices

bit_transformer/BTLM_Extensions/muon_optimizer.py

@@ -0,0 +1,332 @@
+"""
+Muon Optimizer for BitTransformerLM Extensions
+==============================================
+
+Implementation of the Muon optimizer with orthogonal momentum updates.
+Based on "Muon: Momentum Orthogonalized by Newton's method" research.
+
+Key features:
+- Orthogonal momentum updates
+- Better convergence properties than Adam/AdamW
+- Memory efficient implementation
+- Compatible with BitTransformerLM's training infrastructure
+"""
+
+import math
+import torch
+from torch.optim.optimizer import Optimizer
+from typing import Any, Dict, List, Optional, Tuple, Union
+import warnings
+
+
+class Muon(Optimizer):
+    """
+    Muon optimizer with orthogonal momentum updates.
+    
+    This implementation provides momentum updates that are orthogonalized using
+    Newton's method, leading to more stable training dynamics.
+    
+    Args:
+        params: Iterable of parameters to optimize
+        lr: Learning rate (default: 1e-3)
+        momentum: Momentum factor (default: 0.95)
+        nesterov: Enable Nesterov momentum (default: False)
+        backend: Backend for orthogonalization ('newtonschulz' or 'svd')
+        update_period: Period for updating orthogonalization (default: 1)
+        rank_deficiency_threshold: Threshold for rank deficiency detection
+        eps: Small constant for numerical stability (default: 1e-8)
+        weight_decay: Weight decay coefficient (default: 0.0)
+    """
+    
+    def __init__(
+        self,
+        params,
+        lr: float = 1e-3,
+        momentum: float = 0.95,
+        nesterov: bool = False,
+        backend: str = "newtonschulz",
+        update_period: int = 1,
+        rank_deficiency_threshold: float = 1e-6,
+        eps: float = 1e-8,
+        weight_decay: float = 0.0,
+    ):
+        if not 0.0 <= lr:
+            raise ValueError(f"Invalid learning rate: {lr}")
+        if not 0.0 <= momentum <= 1.0:
+            raise ValueError(f"Invalid momentum value: {momentum}")
+        if not 0.0 <= weight_decay:
+            raise ValueError(f"Invalid weight_decay value: {weight_decay}")
+        if backend not in ["newtonschulz", "svd"]:
+            raise ValueError(f"Invalid backend: {backend}")
+            
+        defaults = dict(
+            lr=lr,
+            momentum=momentum,
+            nesterov=nesterov,
+            backend=backend,
+            update_period=update_period,
+            rank_deficiency_threshold=rank_deficiency_threshold,
+            eps=eps,
+            weight_decay=weight_decay,
+        )
+        super().__init__(params, defaults)
+        
+    def _orthogonalize_newtonschulz(self, matrix: torch.Tensor, num_iterations: int = 5) -> torch.Tensor:
+        """Orthogonalize matrix using Newton-Schulz iteration."""
+        # Handle different shapes
+        original_shape = matrix.shape
+        if matrix.dim() > 2:
+            matrix = matrix.view(-1, matrix.shape[-1])
+            
+        if matrix.shape[0] >= matrix.shape[1]:
+            # Tall matrix - orthogonalize columns
+            X = matrix.clone()
+            for _ in range(num_iterations):
+                A = X.T @ X
+                X = X @ (1.5 * torch.eye(A.shape[0], device=A.device, dtype=A.dtype) - 0.5 * A)
+        else:
+            # Wide matrix - orthogonalize rows  
+            X = matrix.clone()
+            for _ in range(num_iterations):
+                A = X @ X.T
+                X = (1.5 * torch.eye(A.shape[0], device=A.device, dtype=A.dtype) - 0.5 * A) @ X
+                
+        return X.view(original_shape)
+    
+    def _orthogonalize_svd(self, matrix: torch.Tensor) -> torch.Tensor:
+        """Orthogonalize matrix using SVD decomposition."""
+        original_shape = matrix.shape
+        if matrix.dim() > 2:
+            matrix = matrix.view(-1, matrix.shape[-1])
+            
+        try:
+            U, _, Vt = torch.linalg.svd(matrix, full_matrices=False)
+            orthogonal = U @ Vt
+            return orthogonal.view(original_shape)
+        except torch._C._LinAlgError:
+            # Fallback to Newton-Schulz if SVD fails
+            warnings.warn("SVD failed, falling back to Newton-Schulz")
+            return self._orthogonalize_newtonschulz(matrix)
+    
+    @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 grad.dtype in {torch.float16, torch.bfloat16}:
+                    grad = grad.float()
+                    
+                state = self.state[p]
+                
+                # State initialization
+                if len(state) == 0:
+                    state["step"] = 0
+                    state["momentum_buffer"] = torch.zeros_like(p, memory_format=torch.preserve_format)
+                    
+                momentum_buffer = state["momentum_buffer"]
+                state["step"] += 1
+                
+                # Weight decay
+                if group["weight_decay"] != 0:
+                    grad = grad.add(p, alpha=group["weight_decay"])
+                
+                # Apply momentum
+                momentum_buffer.mul_(group["momentum"]).add_(grad)
+                
+                # Orthogonalize momentum every update_period steps
+                if state["step"] % group["update_period"] == 0 and momentum_buffer.numel() > 1:
+                    # Only orthogonalize if we have sufficient dimensions
+                    if momentum_buffer.dim() >= 2 and min(momentum_buffer.shape[-2:]) > 1:
+                        if group["backend"] == "newtonschulz":
+                            orthogonal_momentum = self._orthogonalize_newtonschulz(momentum_buffer)
+                        else:
+                            orthogonal_momentum = self._orthogonalize_svd(momentum_buffer)
+                        
+                        # Check for rank deficiency
+                        rank_ratio = torch.linalg.matrix_norm(orthogonal_momentum) / torch.linalg.matrix_norm(momentum_buffer)
+                        if rank_ratio < group["rank_deficiency_threshold"]:
+                            warnings.warn("Detected rank deficiency in momentum buffer")
+                        else:
+                            momentum_buffer.copy_(orthogonal_momentum)
+                
+                # Apply Nesterov acceleration if enabled
+                if group["nesterov"]:
+                    update = grad.add(momentum_buffer, alpha=group["momentum"])
+                else:
+                    update = momentum_buffer
+                
+                # Apply update
+                p.add_(update, alpha=-group["lr"])
+                
+        return loss
+
+
+def configure_muon_optimizer(
+    model: torch.nn.Module,
+    lr: float = 1e-3,
+    momentum: float = 0.95,
+    weight_decay: float = 0.01,
+    total_steps: Optional[int] = None,
+    warmup_ratio: float = 0.1,
+    nesterov: bool = False,
+    backend: str = "newtonschulz",
+    **muon_kwargs
+) -> Tuple[Muon, Optional[torch.optim.lr_scheduler._LRScheduler]]:
+    """
+    Configure Muon optimizer with OneCycle learning rate schedule.
+    
+    This function provides a drop-in replacement for BitTransformerLM's
+    configure_optimizer function, using Muon instead of AdamW.
+    
+    Args:
+        model: PyTorch model to optimize
+        lr: Peak learning rate
+        momentum: Momentum factor for Muon
+        weight_decay: Weight decay coefficient
+        total_steps: Total training steps for OneCycle schedule
+        warmup_ratio: Fraction of steps for warmup
+        nesterov: Enable Nesterov momentum
+        backend: Orthogonalization backend
+        **muon_kwargs: Additional arguments for Muon 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 = Muon(
+        param_groups,
+        lr=lr,
+        momentum=momentum,
+        nesterov=nesterov,
+        backend=backend,
+        **muon_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,  # Muon handles momentum internally
+            div_factor=25.0,
+            final_div_factor=1e4,
+        )
+    
+    return optimizer, scheduler
+
+
+def create_muon_training_config(
+    lr: float = 1e-3,
+    momentum: float = 0.95,
+    weight_decay: float = 0.01,
+    backend: str = "newtonschulz",
+    nesterov: bool = False,
+    **kwargs
+) -> Dict[str, Any]:
+    """
+    Create a training configuration dictionary for Muon optimizer.
+    
+    This can be used with BitTransformerLM's training scripts by passing
+    the config to the training loop.
+    
+    Args:
+        lr: Learning rate
+        momentum: Momentum factor
+        weight_decay: Weight decay coefficient  
+        backend: Orthogonalization backend
+        nesterov: Enable Nesterov momentum
+        **kwargs: Additional configuration options
+        
+    Returns:
+        Dictionary containing training configuration
+    """
+    config = {
+        "optimizer_type": "muon",
+        "optimizer_config": {
+            "lr": lr,
+            "momentum": momentum,
+            "weight_decay": weight_decay,
+            "backend": backend,
+            "nesterov": nesterov,
+            **kwargs
+        },
+        "scheduler_type": "onecycle",
+    }
+    
+    return config
+
+
+# Example usage and integration helpers
+def integrate_with_bittransformerlm():
+    """
+    Example of how to integrate Muon optimizer with BitTransformerLM training.
+    
+    Usage:
+        from BTLM_Extensions.muon_optimizer import configure_muon_optimizer
+        
+        # Replace the standard optimizer configuration
+        optimizer, scheduler = configure_muon_optimizer(
+            model, lr=1e-3, momentum=0.95, total_steps=1000
+        )
+        
+        # Use in training loop
+        train_loop(model, data, optimizer=optimizer, scheduler=scheduler)
+    """
+    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)
+    )
+    
+    optimizer, scheduler = configure_muon_optimizer(model, lr=1e-3, 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)
+    loss.backward()
+    
+    optimizer.step()
+    if scheduler:
+        scheduler.step()
+        
+    print("Muon optimizer test completed successfully!")
+    print(f"Loss: {loss.item():.4f}")