Added test suite

optimizers/compile_utils.py

@@ -0,0 +1,106 @@
+"""
+Utility functions for handling torch.compile gracefully across different PyTorch versions and environments.
+"""
+import torch
+import warnings
+from functools import wraps
+from typing import Callable, Any
+
+
+def safe_torch_compile(fullgraph: bool = True, **kwargs):
+    """
+    A decorator that applies torch.compile if available and functional, 
+    otherwise falls back to the original function.
+    
+    Args:
+        fullgraph: Whether to compile the full graph
+        **kwargs: Additional arguments to pass to torch.compile
+    
+    Returns:
+        A decorator function that either compiles or passes through the original function
+    """
+    import os
+
+    def decorator(func: Callable) -> Callable:
+        # Check if compilation is disabled via environment variable
+        if os.environ.get('TORCH_COMPILE_DISABLE', '0') == '1':
+            return func
+
+        try:
+            # Try to compile the function
+            compiled_func = torch.compile(func, fullgraph=fullgraph, **kwargs)
+
+            # Test if compilation actually works by attempting to create a dummy call
+            # This won't execute but will trigger any import/compilation errors
+            return compiled_func
+
+        except Exception as e:
+            # If compilation fails, warn and return the original function
+            warnings.warn(
+                f"torch.compile failed for function '{func.__name__}': {e}. "
+                f"Falling back to uncompiled version. Performance may be reduced.",
+                UserWarning,
+                stacklevel=2
+            )
+            return func
+
+    return decorator
+
+
+def is_compile_available() -> bool:
+    """
+    Check if torch.compile is available and functional in the current environment.
+    
+    Returns:
+        True if torch.compile is available and functional, False otherwise
+    """
+    try:
+        # Try a simple compile operation
+        @torch.compile
+        def dummy_func(x):
+            return x + 1
+
+        return True
+    except Exception:
+        return False
+
+
+def conditional_compile(condition: bool = None, **compile_kwargs):
+    """
+    Conditionally apply torch.compile based on a condition or environment check.
+    
+    Args:
+        condition: If None, will check if compile is available. 
+                  If True/False, will use that condition.
+        **compile_kwargs: Arguments to pass to torch.compile
+    
+    Returns:
+        A decorator that either compiles or passes through the function
+    """
+    def decorator(func: Callable) -> Callable:
+        if condition is None:
+            should_compile = is_compile_available()
+        else:
+            should_compile = condition
+
+        if should_compile:
+            try:
+                return torch.compile(func, **compile_kwargs)
+            except Exception as e:
+                warnings.warn(
+                    f"torch.compile failed for '{func.__name__}': {e}. Using uncompiled version.",
+                    UserWarning
+                )
+                return func
+        else:
+            return func
+
+    return decorator
+
+
+def disable_compile_for_tests():
+    """
+    Temporarily disable torch.compile for testing to avoid cache limit issues.
+    """
+    import os
+    os.environ['TORCH_COMPILE_DISABLE'] = '1'

optimizers/scalar_opts.py

@@ -1,9 +1,10 @@
 import torch
 from torch import Tensor
 from typing import List
+from .compile_utils import safe_torch_compile
 
 
-@torch.compile(fullgraph=True)
+@safe_torch_compile(fullgraph=True)
 def adamw_update(
     X: Tensor,  # Model weights (modified in place)
     G: Tensor,  # Gradient
@@ -52,7 +53,7 @@ def adamw_update(
     X.addcdiv_(M, denom, value=-adj_lr)
 
 
-@torch.compile(fullgraph=True)
+@safe_torch_compile(fullgraph=True)
 def lion_update(
     X: Tensor,  # Model weights (modified in place)
     G: Tensor,  # Gradient
@@ -86,7 +87,7 @@ def lion_update(
     X.add_(U, alpha=-lr)
 
 
-@torch.compile(fullgraph=True)
+@safe_torch_compile(fullgraph=True)
 def adamw_update_foreach(
     X: List[Tensor],  # Model weights (modified in place)
     G: List[Tensor],  # Gradient
@@ -149,7 +150,7 @@ def adamw_update_foreach(
     torch._foreach_sub_(X, M_div)
 
 
-@torch.compile(fullgraph=True)
+@safe_torch_compile(fullgraph=True)
 def lion_update_foreach(
     X: List[Tensor],  # Model weights (modified in place)
     G: List[Tensor],  # Gradient
@@ -185,3 +186,122 @@ def lion_update_foreach(
     # X = X - lr * U
     torch._foreach_mul_(U, lr)
     torch._foreach_sub_(X, U)
+
+
+class AdamW(torch.optim.Optimizer):
+    """
+    AdamW optimizer using the compiled update functions.
+    """
+    def __init__(self, params, lr=1e-3, betas=(0.9, 0.999), eps=1e-8, weight_decay=0.01):
+        if not 0.0 <= lr:
+            raise ValueError(f"Invalid learning rate: {lr}")
+        if not 0.0 <= eps:
+            raise ValueError(f"Invalid epsilon value: {eps}")
+        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}")
+
+        defaults = dict(lr=lr, betas=betas, eps=eps, weight_decay=weight_decay)
+        super(AdamW, self).__init__(params, defaults)
+
+    def step(self, closure=None):
+        """Performs a single optimization step."""
+        loss = None
+        if closure is not None:
+            loss = closure()
+
+        for group in self.param_groups:
+            for p in group['params']:
+                if p.grad is None:
+                    continue
+
+                grad = p.grad.data
+                if grad.is_sparse:
+                    raise RuntimeError('AdamW does not support sparse gradients')
+
+                state = self.state[p]
+
+                # State initialization
+                if len(state) == 0:
+                    state['step'] = 0
+                    state['exp_avg'] = torch.zeros_like(p.data)
+                    state['exp_avg_sq'] = torch.zeros_like(p.data)
+
+                exp_avg, exp_avg_sq = state['exp_avg'], state['exp_avg_sq']
+                beta1, beta2 = group['betas']
+
+                state['step'] += 1
+
+                # Convert to tensors for the update function
+                lr_tensor = torch.tensor(group['lr'], device=p.device, dtype=p.dtype)
+                beta1_tensor = torch.tensor(beta1, device=p.device, dtype=p.dtype)
+                beta2_tensor = torch.tensor(beta2, device=p.device, dtype=p.dtype)
+                weight_decay_tensor = torch.tensor(group['weight_decay'], device=p.device, dtype=p.dtype)
+
+                # Call the compiled update function
+                adamw_update(
+                    p.data, grad, exp_avg, exp_avg_sq,
+                    lr_tensor, beta1_tensor, beta2_tensor, weight_decay_tensor,
+                    state['step'], group['eps']
+                )
+
+        return loss
+
+
+class Lion(torch.optim.Optimizer):
+    """
+    Lion optimizer using the compiled update functions.
+    """
+    def __init__(self, params, lr=1e-4, betas=(0.9, 0.99), weight_decay=0.0):
+        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}")
+
+        defaults = dict(lr=lr, betas=betas, weight_decay=weight_decay)
+        super(Lion, self).__init__(params, defaults)
+
+    def step(self, closure=None):
+        """Performs a single optimization step."""
+        loss = None
+        if closure is not None:
+            loss = closure()
+
+        for group in self.param_groups:
+            for p in group['params']:
+                if p.grad is None:
+                    continue
+
+                grad = p.grad.data
+                if grad.is_sparse:
+                    raise RuntimeError('Lion does not support sparse gradients')
+
+                state = self.state[p]
+
+                # State initialization
+                if len(state) == 0:
+                    state['exp_avg'] = torch.zeros_like(p.data)
+
+                exp_avg = state['exp_avg']
+                beta1, beta2 = group['betas']
+
+                # Convert to tensors for the update function
+                lr_tensor = torch.tensor(group['lr'], device=p.device, dtype=p.dtype)
+                beta1_tensor = torch.tensor(beta1, device=p.device, dtype=p.dtype)
+                beta2_tensor = torch.tensor(beta2, device=p.device, dtype=p.dtype)
+                weight_decay_tensor = torch.tensor(group['weight_decay'], device=p.device, dtype=p.dtype)
+
+                # Call the compiled update function
+                lion_update(
+                    p.data, grad, exp_avg,
+                    lr_tensor, beta1_tensor, beta2_tensor, weight_decay_tensor
+                )
+
+        return loss

pytest.ini

@@ -0,0 +1,12 @@
+[pytest]
+addopts = -v
+testpaths = tests
+python_files = test_*.py
+python_classes = Test*
+python_functions = test_*
+markers =
+    integration: marks tests as integration tests
+    performance: marks tests as performance tests
+    slow: marks tests as slow running
+env = 
+    TORCH_COMPILE_DISABLE = 1