quarterbit 0.1.0__py3-none-any.whl

quarterbit/torch/optim.py

@@ -0,0 +1,728 @@
+"""
+QuarterBit Optimizers
+=====================
+
+Precision optimizers for PyTorch.
+Copyright (c) 2026 Clouthier Simulation Labs. All rights reserved.
+"""
+
+import torch
+from torch.optim import Optimizer
+from typing import List, Optional, Callable, Iterable
+import ctypes
+from .utils import get_lib, is_available
+
+def _ptr(tensor):
+    """Get ctypes pointer to tensor data."""
+    if tensor is None:
+        return None
+    return ctypes.cast(tensor.data_ptr(), ctypes.POINTER(ctypes.c_float))
+
+
+class SGD(Optimizer):
+    """
+    EFT-powered Stochastic Gradient Descent.
+
+    Drop-in replacement for torch.optim.SGD with precise gradient accumulation.
+
+    Args:
+        params: Iterable of parameters to optimize
+        lr: Learning rate
+        momentum: Momentum factor (default: 0)
+        dampening: Dampening for momentum (default: 0)
+        weight_decay: Weight decay (L2 penalty) (default: 0)
+        nesterov: Enables Nesterov momentum (default: False)
+
+    Example:
+        >>> from quarterbit.torch import SGD
+        >>> optimizer = SGD(model.parameters(), lr=0.01, momentum=0.9)
+    """
+
+    def __init__(
+        self,
+        params: Iterable[torch.nn.Parameter],
+        lr: float = 1e-3,
+        momentum: float = 0,
+        dampening: float = 0,
+        weight_decay: float = 0,
+        nesterov: bool = False
+    ):
+        if lr < 0.0:
+            raise ValueError(f"Invalid learning rate: {lr}")
+        if momentum < 0.0:
+            raise ValueError(f"Invalid momentum value: {momentum}")
+        if weight_decay < 0.0:
+            raise ValueError(f"Invalid weight_decay value: {weight_decay}")
+        if nesterov and (momentum <= 0 or dampening != 0):
+            raise ValueError("Nesterov momentum requires a momentum and zero dampening")
+
+        defaults = dict(
+            lr=lr, momentum=momentum, dampening=dampening,
+            weight_decay=weight_decay, nesterov=nesterov
+        )
+        super().__init__(params, defaults)
+
+        self._use_eft = is_available()
+        if self._use_eft:
+            self._lib = get_lib()
+
+        # Initialize compensation buffers
+        for group in self.param_groups:
+            for p in group['params']:
+                state = self.state[p]
+                state['weight_comp'] = torch.zeros_like(p.data)
+                if momentum != 0:
+                    state['momentum_buffer'] = torch.zeros_like(p.data)
+                    state['momentum_comp'] = torch.zeros_like(p.data)
+
+    @torch.no_grad()
+    def step(self, closure: Optional[Callable] = None):
+        """Performs a single optimization step."""
+        loss = None
+        if closure is not None:
+            with torch.enable_grad():
+                loss = closure()
+
+        for group in self.param_groups:
+            lr = group['lr']
+            momentum = group['momentum']
+            dampening = group['dampening']
+            weight_decay = group['weight_decay']
+            nesterov = group['nesterov']
+
+            for p in group['params']:
+                if p.grad is None:
+                    continue
+
+                grad = p.grad.data
+                state = self.state[p]
+
+                if self._use_eft and p.is_cuda and p.dtype == torch.float32:
+                    # Use EFT kernel
+                    n = p.numel()
+
+                    if momentum != 0:
+                        # TODO: Implement momentum with EFT
+                        # For now, fall back to standard for momentum
+                        self._standard_step(p, grad, state, group)
+                    else:
+                        # Pure SGD with EFT
+                        self._lib.eft_sgd_step(
+                            _ptr(p.data),
+                            _ptr(state['weight_comp']),
+                            _ptr(grad),
+                            ctypes.c_float(lr),
+                            ctypes.c_float(weight_decay),
+                            ctypes.c_int(n)
+                        )
+                else:
+                    # Fallback to standard PyTorch
+                    self._standard_step(p, grad, state, group)
+
+        return loss
+
+    def _standard_step(self, p, grad, state, group):
+        """Standard PyTorch SGD step (fallback)."""
+        weight_decay = group['weight_decay']
+        momentum = group['momentum']
+        dampening = group['dampening']
+        nesterov = group['nesterov']
+        lr = group['lr']
+
+        if weight_decay != 0:
+            grad = grad.add(p.data, alpha=weight_decay)
+
+        if momentum != 0:
+            buf = state.get('momentum_buffer')
+            if buf is None:
+                buf = torch.clone(grad).detach()
+                state['momentum_buffer'] = buf
+            else:
+                buf.mul_(momentum).add_(grad, alpha=1 - dampening)
+
+            if nesterov:
+                grad = grad.add(buf, alpha=momentum)
+            else:
+                grad = buf
+
+        p.data.add_(grad, alpha=-lr)
+
+
+class Adam(Optimizer):
+    """
+    EFT-powered Adam optimizer.
+
+    Drop-in replacement for torch.optim.Adam with precise moment accumulation.
+
+    The key insight: Adam's moment estimates (m and v) accumulate small updates
+    over millions of steps. Standard FP32 loses precision. EFT tracks the
+    exact error, ensuring stable training.
+
+    Args:
+        params: Iterable of parameters to optimize
+        lr: Learning rate (default: 1e-3)
+        betas: Coefficients for computing running averages (default: (0.9, 0.999))
+        eps: Term added to denominator for numerical stability (default: 1e-8)
+        weight_decay: Weight decay (L2 penalty) (default: 0)
+        amsgrad: Whether to use AMSGrad variant (default: False)
+
+    Example:
+        >>> from quarterbit.torch import Adam
+        >>> optimizer = Adam(model.parameters(), lr=1e-3)
+    """
+
+    def __init__(
+        self,
+        params: Iterable[torch.nn.Parameter],
+        lr: float = 1e-3,
+        betas: tuple = (0.9, 0.999),
+        eps: float = 1e-8,
+        weight_decay: float = 0,
+        amsgrad: bool = False
+    ):
+        if lr < 0.0:
+            raise ValueError(f"Invalid learning rate: {lr}")
+        if eps < 0.0:
+            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 weight_decay < 0.0:
+            raise ValueError(f"Invalid weight_decay value: {weight_decay}")
+
+        defaults = dict(lr=lr, betas=betas, eps=eps, weight_decay=weight_decay, amsgrad=amsgrad)
+        super().__init__(params, defaults)
+
+        self._use_eft = is_available()
+        if self._use_eft:
+            self._lib = get_lib()
+
+    def _init_state(self, p):
+        """Initialize state for a parameter."""
+        state = self.state[p]
+        state['step'] = 0
+        state['exp_avg'] = torch.zeros_like(p.data)
+        state['exp_avg_sq'] = torch.zeros_like(p.data)
+
+        # EFT compensation buffers
+        state['weight_comp'] = torch.zeros_like(p.data)
+        state['exp_avg_comp'] = torch.zeros_like(p.data)
+        state['exp_avg_sq_comp'] = torch.zeros_like(p.data)
+
+        return state
+
+    @torch.no_grad()
+    def step(self, closure: Optional[Callable] = None):
+        """Performs a single optimization step."""
+        loss = None
+        if closure is not None:
+            with torch.enable_grad():
+                loss = closure()
+
+        for group in self.param_groups:
+            beta1, beta2 = group['betas']
+            lr = group['lr']
+            eps = group['eps']
+            weight_decay = group['weight_decay']
+            amsgrad = group['amsgrad']
+
+            for p in group['params']:
+                if p.grad is None:
+                    continue
+
+                grad = p.grad.data
+                if grad.is_sparse:
+                    raise RuntimeError("Adam does not support sparse gradients")
+
+                # Initialize state if needed
+                state = self.state[p]
+                if len(state) == 0:
+                    state = self._init_state(p)
+
+                state['step'] += 1
+                step = state['step']
+
+                if self._use_eft and p.is_cuda and p.dtype == torch.float32 and not amsgrad:
+                    # Use EFT kernel
+                    n = p.numel()
+
+                    self._lib.eft_adam_step(
+                        _ptr(p.data),
+                        _ptr(state['weight_comp']),
+                        _ptr(grad),
+                        _ptr(state['exp_avg']),
+                        _ptr(state['exp_avg_sq']),
+                        _ptr(state['exp_avg_comp']),
+                        _ptr(state['exp_avg_sq_comp']),
+                        ctypes.c_float(lr),
+                        ctypes.c_float(beta1),
+                        ctypes.c_float(beta2),
+                        ctypes.c_float(eps),
+                        ctypes.c_float(weight_decay),
+                        ctypes.c_int(step),
+                        ctypes.c_int(n)
+                    )
+                else:
+                    # Fallback to standard PyTorch Adam
+                    self._standard_step(p, grad, state, group)
+
+        return loss
+
+    def _standard_step(self, p, grad, state, group):
+        """Standard PyTorch Adam step (fallback)."""
+        beta1, beta2 = group['betas']
+        lr = group['lr']
+        eps = group['eps']
+        weight_decay = group['weight_decay']
+        amsgrad = group['amsgrad']
+        step = state['step']
+
+        exp_avg = state['exp_avg']
+        exp_avg_sq = state['exp_avg_sq']
+
+        if weight_decay != 0:
+            grad = grad.add(p.data, alpha=weight_decay)
+
+        # Decay the first and second moment running average coefficient
+        exp_avg.mul_(beta1).add_(grad, alpha=1 - beta1)
+        exp_avg_sq.mul_(beta2).addcmul_(grad, grad, value=1 - beta2)
+
+        if amsgrad:
+            max_exp_avg_sq = state.get('max_exp_avg_sq')
+            if max_exp_avg_sq is None:
+                max_exp_avg_sq = torch.zeros_like(p.data)
+                state['max_exp_avg_sq'] = max_exp_avg_sq
+            torch.maximum(max_exp_avg_sq, exp_avg_sq, out=max_exp_avg_sq)
+            denom = (max_exp_avg_sq.sqrt() / (1 - beta2 ** step) ** 0.5).add_(eps)
+        else:
+            bias_correction1 = 1 - beta1 ** step
+            bias_correction2 = 1 - beta2 ** step
+            denom = (exp_avg_sq.sqrt() / (bias_correction2 ** 0.5)).add_(eps)
+
+        step_size = lr / (1 - beta1 ** step) if not amsgrad else lr / (1 - beta1 ** step)
+        p.data.addcdiv_(exp_avg, denom, value=-step_size)
+
+
+class AdamW(Adam):
+    """
+    EFT-powered AdamW optimizer (Adam with decoupled weight decay).
+
+    Same as Adam but with proper weight decay (not L2 regularization).
+
+    Example:
+        >>> from quarterbit.torch import AdamW
+        >>> optimizer = AdamW(model.parameters(), lr=1e-3, weight_decay=0.01)
+    """
+
+    def __init__(
+        self,
+        params: Iterable[torch.nn.Parameter],
+        lr: float = 1e-3,
+        betas: tuple = (0.9, 0.999),
+        eps: float = 1e-8,
+        weight_decay: float = 1e-2,
+        amsgrad: bool = False
+    ):
+        super().__init__(params, lr=lr, betas=betas, eps=eps, weight_decay=weight_decay, amsgrad=amsgrad)
+
+    def _standard_step(self, p, grad, state, group):
+        """AdamW step - weight decay applied directly to weights."""
+        beta1, beta2 = group['betas']
+        lr = group['lr']
+        eps = group['eps']
+        weight_decay = group['weight_decay']
+        amsgrad = group['amsgrad']
+        step = state['step']
+
+        exp_avg = state['exp_avg']
+        exp_avg_sq = state['exp_avg_sq']
+
+        # Decoupled weight decay (applied to weights, not gradients)
+        if weight_decay != 0:
+            p.data.mul_(1 - lr * weight_decay)
+
+        # Standard Adam update
+        exp_avg.mul_(beta1).add_(grad, alpha=1 - beta1)
+        exp_avg_sq.mul_(beta2).addcmul_(grad, grad, value=1 - beta2)
+
+        if amsgrad:
+            max_exp_avg_sq = state.get('max_exp_avg_sq')
+            if max_exp_avg_sq is None:
+                max_exp_avg_sq = torch.zeros_like(p.data)
+                state['max_exp_avg_sq'] = max_exp_avg_sq
+            torch.maximum(max_exp_avg_sq, exp_avg_sq, out=max_exp_avg_sq)
+            denom = (max_exp_avg_sq.sqrt() / (1 - beta2 ** step) ** 0.5).add_(eps)
+        else:
+            bias_correction1 = 1 - beta1 ** step
+            bias_correction2 = 1 - beta2 ** step
+            denom = (exp_avg_sq.sqrt() / (bias_correction2 ** 0.5)).add_(eps)
+
+        step_size = lr / (1 - beta1 ** step)
+        p.data.addcdiv_(exp_avg, denom, value=-step_size)
+
+
+class CompactAdam(Optimizer):
+    """
+    Memory-efficient Adam optimizer with compressed state storage.
+
+    Uses FP16+FP4 split weights and FP4+INT2 compressed optimizer states
+    for ~1.7x memory savings vs standard FP32 Adam, while maintaining
+    training quality through error feedback.
+
+    Memory budget per parameter:
+    - Standard FP32 Adam: 16 bytes/param
+    - CompactAdam (no EF): 9.25 bytes/param (1.7x savings)
+    - CompactAdam (with EF): 10.75 bytes/param (1.5x savings)
+
+    Args:
+        params: Iterable of parameters to optimize
+        lr: Learning rate (default: 1e-3)
+        betas: Coefficients for computing running averages (default: (0.9, 0.999))
+        eps: Term added to denominator for numerical stability (default: 1e-8)
+        weight_decay: Weight decay (L2 penalty) (default: 0)
+        use_error_feedback: Enable error feedback for better precision (default: True)
+
+    Example:
+        >>> from quarterbit.torch import CompactAdam
+        >>> optimizer = CompactAdam(model.parameters(), lr=1e-3)
+        >>> # For maximum memory savings (slight precision tradeoff):
+        >>> optimizer = CompactAdam(model.parameters(), lr=1e-3, use_error_feedback=False)
+    """
+
+    def __init__(
+        self,
+        params: Iterable[torch.nn.Parameter],
+        lr: float = 1e-3,
+        betas: tuple = (0.9, 0.999),
+        eps: float = 1e-8,
+        weight_decay: float = 0,
+        use_error_feedback: bool = True
+    ):
+        if lr < 0.0:
+            raise ValueError(f"Invalid learning rate: {lr}")
+        if eps < 0.0:
+            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 weight_decay < 0.0:
+            raise ValueError(f"Invalid weight_decay value: {weight_decay}")
+
+        defaults = dict(lr=lr, betas=betas, eps=eps, weight_decay=weight_decay)
+        super().__init__(params, defaults)
+
+        self._use_ef = 1 if use_error_feedback else 0
+        self._handles = {}  # Per-parameter handles
+        self._step_count = 0
+
+        # Try to load compact library
+        self._lib = None
+        self._use_compact = False
+        try:
+            from .utils import get_lib
+            lib = get_lib()
+            # Check if compact functions exist
+            if hasattr(lib, 'compact_adam_create'):
+                self._lib = lib
+                self._use_compact = True
+        except Exception:
+            pass
+
+    def _get_handle(self, p):
+        """Get or create compact adam handle for parameter."""
+        if id(p) not in self._handles:
+            if self._use_compact and p.is_cuda and p.dtype == torch.float32:
+                n = p.numel()
+                handle = self._lib.compact_adam_create(ctypes.c_int(n), ctypes.c_int(self._use_ef))
+                if handle:
+                    self._handles[id(p)] = handle
+                else:
+                    self._handles[id(p)] = None
+            else:
+                self._handles[id(p)] = None
+        return self._handles[id(p)]
+
+    @torch.no_grad()
+    def step(self, closure: Optional[Callable] = None):
+        """Performs a single optimization step."""
+        loss = None
+        if closure is not None:
+            with torch.enable_grad():
+                loss = closure()
+
+        self._step_count += 1
+
+        for group in self.param_groups:
+            beta1, beta2 = group['betas']
+            lr = group['lr']
+            eps = group['eps']
+            weight_decay = group['weight_decay']
+
+            for p in group['params']:
+                if p.grad is None:
+                    continue
+
+                grad = p.grad.data
+                if grad.is_sparse:
+                    raise RuntimeError("CompactAdam does not support sparse gradients")
+
+                handle = self._get_handle(p)
+
+                if handle is not None and self._use_compact:
+                    # Use compact CUDA kernel
+                    n = p.numel()
+
+                    # Apply weight decay to gradients (L2 regularization)
+                    if weight_decay != 0:
+                        grad = grad.add(p.data, alpha=weight_decay)
+
+                    self._lib.compact_adam_step(
+                        handle,
+                        _ptr(p.data),
+                        _ptr(grad),
+                        ctypes.c_float(lr),
+                        ctypes.c_float(beta1),
+                        ctypes.c_float(beta2),
+                        ctypes.c_float(eps),
+                        ctypes.c_int(self._step_count)
+                    )
+                else:
+                    # Fallback to standard Adam
+                    self._standard_step(p, grad, group)
+
+        return loss
+
+    def _standard_step(self, p, grad, group):
+        """Standard PyTorch Adam step (fallback for CPU or non-float32)."""
+        beta1, beta2 = group['betas']
+        lr = group['lr']
+        eps = group['eps']
+        weight_decay = group['weight_decay']
+
+        state = self.state[p]
+        if len(state) == 0:
+            state['step'] = 0
+            state['exp_avg'] = torch.zeros_like(p.data)
+            state['exp_avg_sq'] = torch.zeros_like(p.data)
+
+        state['step'] += 1
+        step = state['step']
+
+        exp_avg = state['exp_avg']
+        exp_avg_sq = state['exp_avg_sq']
+
+        if weight_decay != 0:
+            grad = grad.add(p.data, alpha=weight_decay)
+
+        exp_avg.mul_(beta1).add_(grad, alpha=1 - beta1)
+        exp_avg_sq.mul_(beta2).addcmul_(grad, grad, value=1 - beta2)
+
+        bias_correction1 = 1 - beta1 ** step
+        bias_correction2 = 1 - beta2 ** step
+        denom = (exp_avg_sq.sqrt() / (bias_correction2 ** 0.5)).add_(eps)
+
+        step_size = lr / bias_correction1
+        p.data.addcdiv_(exp_avg, denom, value=-step_size)
+
+    def __del__(self):
+        """Clean up handles on deletion."""
+        if hasattr(self, '_handles') and hasattr(self, '_lib') and self._lib is not None:
+            for handle in self._handles.values():
+                if handle is not None:
+                    try:
+                        self._lib.compact_adam_destroy(handle)
+                    except Exception:
+                        pass
+
+    def memory_savings(self):
+        """Report memory savings vs standard FP32 Adam."""
+        total_params = sum(p.numel() for group in self.param_groups for p in group['params'])
+        fp32_bytes = total_params * 16  # 4 bytes each for param, grad, m, v
+        compact_bytes = total_params * (10.75 if self._use_ef else 9.25)
+        savings = 1 - compact_bytes / fp32_bytes
+        return {
+            'total_params': total_params,
+            'fp32_bytes': fp32_bytes,
+            'compact_bytes': compact_bytes,
+            'savings_ratio': fp32_bytes / compact_bytes,
+            'savings_percent': savings * 100
+        }
+
+
+class CompactEFTAdam(Optimizer):
+    """
+    Production Adam optimizer: Memory-efficient + EFT precision.
+
+    Combines the best of both worlds:
+    - Compressed storage: FP16+FP4 weights, FP4+INT2 states (27% memory savings)
+    - EFT precision: TwoSum/TwoProduct for exact arithmetic in Adam updates
+
+    Memory budget per parameter:
+    - Standard FP32 Adam: 16 bytes/param
+    - CompactEFTAdam: ~13.25 bytes/param (17% savings with full precision)
+
+    This is the RECOMMENDED optimizer for production training.
+
+    Args:
+        params: Iterable of parameters to optimize
+        lr: Learning rate (default: 1e-3)
+        betas: Coefficients for computing running averages (default: (0.9, 0.999))
+        eps: Term added to denominator for numerical stability (default: 1e-8)
+        weight_decay: Weight decay (L2 penalty) (default: 0)
+
+    Example:
+        >>> from quarterbit.torch import CompactEFTAdam
+        >>> optimizer = CompactEFTAdam(model.parameters(), lr=1e-3)
+    """
+
+    def __init__(
+        self,
+        params: Iterable[torch.nn.Parameter],
+        lr: float = 1e-3,
+        betas: tuple = (0.9, 0.999),
+        eps: float = 1e-8,
+        weight_decay: float = 0,
+    ):
+        if lr < 0.0:
+            raise ValueError(f"Invalid learning rate: {lr}")
+        if eps < 0.0:
+            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 weight_decay < 0.0:
+            raise ValueError(f"Invalid weight_decay value: {weight_decay}")
+
+        defaults = dict(lr=lr, betas=betas, eps=eps, weight_decay=weight_decay)
+        super().__init__(params, defaults)
+
+        self._handles = {}
+        self._step_count = 0
+
+        # Try to load CompactEFT library
+        self._lib = None
+        self._use_compact_eft = False
+        try:
+            from .utils import _load_compact_eft_lib
+            self._lib = _load_compact_eft_lib()
+            if self._lib is not None:
+                self._use_compact_eft = True
+        except Exception:
+            pass
+
+    def _get_handle(self, p):
+        """Get or create CompactEFT handle for parameter."""
+        if id(p) not in self._handles:
+            if self._use_compact_eft and p.is_cuda and p.dtype == torch.float32:
+                handle = self._lib.compact_eft_adam_create(ctypes.c_int(p.numel()))
+                self._handles[id(p)] = handle if handle else None
+            else:
+                self._handles[id(p)] = None
+        return self._handles[id(p)]
+
+    @torch.no_grad()
+    def step(self, closure: Optional[Callable] = None):
+        """Performs a single optimization step."""
+        loss = None
+        if closure is not None:
+            with torch.enable_grad():
+                loss = closure()
+
+        self._step_count += 1
+
+        for group in self.param_groups:
+            beta1, beta2 = group['betas']
+            lr = group['lr']
+            eps = group['eps']
+            weight_decay = group['weight_decay']
+
+            for p in group['params']:
+                if p.grad is None:
+                    continue
+
+                grad = p.grad.data
+                if grad.is_sparse:
+                    raise RuntimeError("CompactEFTAdam does not support sparse gradients")
+
+                handle = self._get_handle(p)
+
+                if handle is not None and self._use_compact_eft:
+                    # Apply weight decay to gradients (L2 regularization)
+                    if weight_decay != 0:
+                        grad = grad.add(p.data, alpha=weight_decay)
+
+                    self._lib.compact_eft_adam_step(
+                        handle,
+                        _ptr(p.data),
+                        _ptr(grad),
+                        ctypes.c_float(lr),
+                        ctypes.c_float(beta1),
+                        ctypes.c_float(beta2),
+                        ctypes.c_float(eps),
+                        ctypes.c_int(self._step_count)
+                    )
+                else:
+                    # Fallback to standard Adam
+                    self._standard_step(p, grad, group)
+
+        return loss
+
+    def _standard_step(self, p, grad, group):
+        """Standard PyTorch Adam step (fallback for CPU or non-float32)."""
+        beta1, beta2 = group['betas']
+        lr = group['lr']
+        eps = group['eps']
+        weight_decay = group['weight_decay']
+
+        state = self.state[p]
+        if len(state) == 0:
+            state['step'] = 0
+            state['exp_avg'] = torch.zeros_like(p.data)
+            state['exp_avg_sq'] = torch.zeros_like(p.data)
+
+        state['step'] += 1
+        step = state['step']
+
+        exp_avg = state['exp_avg']
+        exp_avg_sq = state['exp_avg_sq']
+
+        if weight_decay != 0:
+            grad = grad.add(p.data, alpha=weight_decay)
+
+        exp_avg.mul_(beta1).add_(grad, alpha=1 - beta1)
+        exp_avg_sq.mul_(beta2).addcmul_(grad, grad, value=1 - beta2)
+
+        bias_correction1 = 1 - beta1 ** step
+        bias_correction2 = 1 - beta2 ** step
+        denom = (exp_avg_sq.sqrt() / (bias_correction2 ** 0.5)).add_(eps)
+
+        step_size = lr / bias_correction1
+        p.data.addcdiv_(exp_avg, denom, value=-step_size)
+
+    def __del__(self):
+        """Clean up handles on deletion."""
+        if hasattr(self, '_handles') and hasattr(self, '_lib') and self._lib is not None:
+            for handle in self._handles.values():
+                if handle is not None:
+                    try:
+                        self._lib.compact_eft_adam_destroy(handle)
+                    except Exception:
+                        pass
+
+    def memory_savings(self):
+        """Report memory savings vs standard FP32 Adam."""
+        total_params = sum(p.numel() for group in self.param_groups for p in group['params'])
+        fp32_bytes = total_params * 16
+        # Weights: 2.75, states: 2.5, EFT comp: 4, grads: 4 = 13.25 B/param
+        compact_eft_bytes = total_params * 13.25
+        savings = 1 - compact_eft_bytes / fp32_bytes
+        return {
+            'total_params': total_params,
+            'fp32_bytes': fp32_bytes,
+            'compact_eft_bytes': compact_eft_bytes,
+            'savings_ratio': fp32_bytes / compact_eft_bytes,
+            'savings_percent': savings * 100
+        }