scao 0.1.0__py3-none-any.whl

scao/__init__.py

@@ -0,0 +1,49 @@
+"""
+SCAO - Sparse Curvature-Aware Adaptive Optimizer
+=================================================
+A second-order PyTorch optimizer that delivers Shampoo-quality
+preconditioned gradients at near-AdamW memory and throughput cost.
+
+Quickstart
+----------
+    import scao
+
+    optimizer = scao.SCAO(model.parameters(), lr=3.5e-4)
+
+    # or, explicit:
+    from scao import SCAO
+    optimizer = SCAO(
+        model.parameters(),
+        lr=3.5e-4,
+        weight_decay=0.1,
+        warmup_steps=100,
+        precond_freq=10,
+    )
+
+Paper
+-----
+SCAO: Sparse Curvature-Aware Adaptive Optimization for Large-Scale Models
+NeurIPS 2026 (under review)
+"""
+
+from .optimizer import SCAO
+from .preconditioner import SparsePreconditioner
+from .utils import matrix_power_neg_quarter, adaptive_rank
+from . import logging as scao_logging
+
+__version__ = "0.1.0"
+__author__ = "SCAO Authors"
+__license__ = "Apache-2.0"
+
+__all__ = [
+    # Main API — this is all most users need
+    "SCAO",
+    # Advanced / internals
+    "SparsePreconditioner",
+    "matrix_power_neg_quarter",
+    "adaptive_rank",
+    "scao_logging",
+    # Metadata
+    "__version__",
+]
+

scao/cuda/__init__.py

@@ -0,0 +1,130 @@
+"""
+SCAO CUDA acceleration layer (optional).
+=========================================
+Provides faster implementations of the critical hot paths.
+Falls back gracefully to pure-PyTorch if CUDA is not available or
+the CUDA extension has not been compiled.
+
+To compile the CUDA extension:
+    cd scao/cuda
+    python setup.py build_ext --inplace
+
+Or install from the project root:
+    pip install -e ".[cuda]"
+"""
+
+from __future__ import annotations
+
+import torch
+from torch import Tensor
+
+_cuda_ext = None  # Will hold the compiled C++/CUDA extension if available
+
+
+def _try_load_cuda_ext():
+    global _cuda_ext
+    if _cuda_ext is not None:
+        return _cuda_ext
+    try:
+        from scao.cuda import _scao_cuda  # type: ignore[import]
+        _cuda_ext = _scao_cuda
+    except ImportError:
+        _cuda_ext = None
+    return _cuda_ext
+
+
+# ---------------------------------------------------------------------------
+# Low-rank matrix-matrix multiply  (U diag(s) U^T) @ G
+# ---------------------------------------------------------------------------
+
+def low_rank_precond_mm(
+    U: Tensor,
+    s_inv_quarter: Tensor,
+    G: Tensor,
+    side: str = "left",
+) -> Tensor:
+    """
+    Compute  U * diag(s_inv_quarter) * (U^T * G)   if side='left'
+    or       G * U * diag(s_inv_quarter) * U^T      if side='right'
+
+    Uses fused CUDA kernel if available, otherwise pure PyTorch.
+
+    Args:
+        U: (m, k) orthonormal matrix
+        s_inv_quarter: (k,) vector of S^{-1/4} scaling factors
+        G: (m, n) gradient matrix
+        side: 'left' or 'right'
+
+    Returns:
+        Result tensor of same shape as G.
+    """
+    ext = _try_load_cuda_ext()
+    if ext is not None and G.is_cuda:
+        try:
+            return ext.low_rank_precond_mm(U, s_inv_quarter, G, side == "left")
+        except (AttributeError, RuntimeError):
+            pass  # Fall through to PyTorch fallback
+
+    # Pure PyTorch fallback
+    if side == "left":
+        # (k, n) = (m, k)^T @ (m, n)
+        proj = U.T @ G
+        # scale rows
+        proj = proj * s_inv_quarter.unsqueeze(1)
+        # (m, n) = (m, k) @ (k, n)
+        return U @ proj
+    else:
+        # (m, k) = (m, n) @ (n, k)
+        proj = G @ U
+        # scale cols
+        proj = proj * s_inv_quarter.unsqueeze(0)
+        # (m, n) = (m, k) @ (k, n)^T
+        return proj @ U.T
+
+
+# ---------------------------------------------------------------------------
+# Batched eigendecomposition with truncation
+# ---------------------------------------------------------------------------
+
+def truncated_eigh(
+    A: Tensor,
+    k: int,
+    eps: float = 1e-8,
+) -> tuple[Tensor, Tensor]:
+    """
+    Compute the top-k eigenvalues and eigenvectors of symmetric PSD matrix A.
+
+    Uses CUDA LOBPCG if available; falls back to full eigh + slice.
+
+    Args:
+        A: (m, m) symmetric PSD matrix
+        k: number of top eigenpairs to return
+
+    Returns:
+        (eigenvalues, eigenvectors) of shapes (k,) and (m, k),
+        eigenvalues in descending order.
+    """
+    ext = _try_load_cuda_ext()
+    m = A.shape[0]
+
+    if k >= m:
+        # Full decomposition
+        S, U = torch.linalg.eigh(A)
+        return S.flip(0).clamp(min=0.0), U.flip(1)
+
+    if ext is not None and A.is_cuda:
+        try:
+            S, U = ext.truncated_eigh(A, k)
+            return S, U
+        except (AttributeError, RuntimeError):
+            pass
+
+    # PyTorch fallback: full eigh + truncate (LOBPCG is less stable for small k)
+    try:
+        S, U = torch.linalg.eigh(A)
+    except torch.linalg.LinAlgError:
+        S, U = torch.linalg.eigh(A + eps * torch.eye(m, device=A.device, dtype=A.dtype))
+
+    S = S.flip(0).clamp(min=0.0)
+    U = U.flip(1)
+    return S[:k], U[:, :k]

scao/cuda/setup.py

@@ -0,0 +1,33 @@
+"""
+Optional CUDA extension build script for SCAO.
+Run:  python setup.py build_ext --inplace
+"""
+
+from setuptools import setup
+try:
+    from torch.utils.cpp_extension import CUDAExtension, BuildExtension
+    ext_modules = [
+        CUDAExtension(
+            name="scao.cuda._scao_cuda",
+            sources=["low_rank_ops.cu"],
+            extra_compile_args={
+                "cxx": ["-O3"],
+                "nvcc": [
+                    "-O3",
+                    "--use_fast_math",
+                    "-gencode=arch=compute_80,code=sm_80",   # A100
+                    "-gencode=arch=compute_89,code=sm_89",   # H100
+                ],
+            },
+        )
+    ]
+    cmdclass = {"build_ext": BuildExtension}
+except ImportError:
+    ext_modules = []
+    cmdclass = {}
+
+setup(
+    name="scao_cuda",
+    ext_modules=ext_modules,
+    cmdclass=cmdclass,
+)

scao/distributed.py

@@ -0,0 +1,114 @@
+"""
+SCAO distributed training utilities.
+=====================================
+Helpers for ZeRO-3 (DeepSpeed) and FSDP (PyTorch ≥ 2.0) compatibility.
+
+Usage with FSDP
+---------------
+    from scao.distributed import wrap_scao_for_fsdp
+
+    model = FSDP(model, ...)
+    optimizer = wrap_scao_for_fsdp(SCAO(model.parameters(), lr=1e-3))
+
+Usage with DeepSpeed ZeRO-3
+----------------------------
+SCAO is compatible with ZeRO-3 out of the box as long as you do NOT enable
+stage 3 optimizer state partitioning for preconditioner tensors (they must
+live on the rank that owns the corresponding parameters).
+
+In your DeepSpeed config, set:
+    "zero_optimization": {
+        "stage": 3,
+        "stage3_param_persistence_threshold": 1e4
+    }
+"""
+
+from __future__ import annotations
+
+import warnings
+from typing import TYPE_CHECKING
+
+import torch
+import torch.distributed as dist
+
+if TYPE_CHECKING:
+    from .optimizer import SCAO
+
+
+def sync_preconditioners(optimizer: "SCAO", process_group=None) -> None:
+    """
+    All-reduce preconditioner eigenfactors across all ranks.
+
+    In FSDP/ZeRO setups where parameters are sharded, each rank computes
+    curvature from its local gradient shard.  This function averages those
+    estimates so every rank has a globally consistent preconditioner.
+
+    Call this AFTER optimizer.step() and BEFORE the next forward pass, or
+    better: schedule it on the preconditioner update cadence (precond_freq).
+
+    Args:
+        optimizer:  SCAO optimizer instance.
+        process_group: optional distributed process group (default: global group).
+    """
+    if not dist.is_available() or not dist.is_initialized():
+        return
+
+    world_size = dist.get_world_size(group=process_group)
+    if world_size == 1:
+        return
+
+    handles = []
+    tensors_to_sync: list[torch.Tensor] = []
+
+    for state in optimizer.state.values():
+        if "preconditioner" not in state:
+            continue
+        prec = state["preconditioner"]
+        if prec.use_kronecker:
+            # Average eigenfactor matrices across ranks
+            for t in (prec.U_l, prec.S_l, prec.U_r, prec.S_r,
+                      prec.L_ema, prec.R_ema):
+                tensors_to_sync.append(t)
+        else:
+            tensors_to_sync.append(prec.diag_ema)
+
+    # Batch all-reduce into as few calls as possible
+    for t in tensors_to_sync:
+        h = dist.all_reduce(t, op=dist.ReduceOp.SUM, group=process_group, async_op=True)
+        handles.append((h, t))
+
+    for h, t in handles:
+        h.wait()
+        t.div_(world_size)
+
+
+def wrap_scao_for_fsdp(optimizer: "SCAO") -> "SCAO":
+    """
+    Register a post-step hook that synchronises preconditioners after each
+    curvature update in FSDP training.
+
+    Returns the same optimizer (mutation in place), for chaining:
+        optimizer = wrap_scao_for_fsdp(SCAO(...))
+    """
+    try:
+        from torch.distributed.fsdp import FullyShardedDataParallel as FSDP  # noqa: F401
+    except ImportError:
+        warnings.warn("torch.distributed.fsdp not available; FSDP wrapping skipped.")
+        return optimizer
+
+    original_step = optimizer.step
+
+    def patched_step(closure=None):
+        result = original_step(closure)
+        # Synchronise preconditioners on curvature-update steps
+        step = next(
+            (s.get("step", 0) for s in optimizer.state.values() if "step" in s),
+            0,
+        )
+        if step % optimizer.defaults["precond_freq"] == 0:
+            optimizer.synchronize_precond()
+            sync_preconditioners(optimizer)
+        return result
+
+    optimizer.step = patched_step  # type: ignore[method-assign]
+    return optimizer

scao/integrations/huggingface.py

@@ -0,0 +1,246 @@
+"""
+HuggingFace Transformers integration for SCAO
+==============================================
+
+Provides three integration paths:
+
+Path A — Manual (recommended for control)
+-----------------------------------------
+    from scao import SCAO
+    from scao.integrations.huggingface import get_scao_optimizer
+
+    optimizer, scheduler = get_scao_optimizer(model, training_args)
+    trainer = Trainer(model=model, args=training_args,
+                      optimizers=(optimizer, scheduler))
+
+Path B — SCAOTrainer drop-in
+-----------------------------
+    from scao.integrations.huggingface import SCAOTrainer
+
+    trainer = SCAOTrainer(
+        model=model,
+        args=training_args,
+        train_dataset=train_dataset,
+        scao_kwargs=dict(precond_freq=20, k_min=8, k_max=64),
+    )
+    trainer.train()
+
+Path C — Trainer callback for monitoring
+-----------------------------------------
+    from scao.integrations.huggingface import SCAOMonitorCallback
+    from scao import SCAO
+
+    optimizer = SCAO(model.parameters(), lr=training_args.learning_rate)
+    trainer = Trainer(model=model, args=training_args,
+                      optimizers=(optimizer, None))
+    trainer.add_callback(SCAOMonitorCallback(optimizer))
+"""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Any
+
+import torch
+import torch.nn as nn
+
+if TYPE_CHECKING:
+    pass
+
+
+# ---------------------------------------------------------------------------
+# Path A: get_scao_optimizer
+# ---------------------------------------------------------------------------
+
+def get_scao_optimizer(
+    model: nn.Module,
+    training_args,
+    scao_kwargs: dict | None = None,
+    no_decay_names: tuple[str, ...] = ("bias", "LayerNorm.weight", "layer_norm.weight"),
+):
+    """
+    Build a SCAO optimizer + learning-rate scheduler compatible with
+    HuggingFace ``Trainer``.
+
+    Applies weight decay only to parameters whose names do NOT match
+    ``no_decay_names`` (standard HF practice).
+
+    Args:
+        model: the model to optimise
+        training_args: ``transformers.TrainingArguments`` instance
+        scao_kwargs: additional SCAO hyperparameters (override defaults)
+        no_decay_names: parameter name substrings exempt from weight decay
+
+    Returns:
+        (optimizer, scheduler) tuple ready for ``Trainer(optimizers=...)``
+
+    Example::
+
+        from transformers import TrainingArguments, Trainer
+        from scao.integrations.huggingface import get_scao_optimizer
+
+        args = TrainingArguments(output_dir="out", num_train_epochs=3, ...)
+        optimizer, scheduler = get_scao_optimizer(model, args)
+        trainer = Trainer(model=model, args=args,
+                          optimizers=(optimizer, scheduler))
+    """
+    try:
+        from transformers import get_scheduler  # type: ignore[import]
+    except ImportError as e:
+        raise ImportError(
+            "transformers is required for get_scao_optimizer. "
+            "Install it with: pip install transformers"
+        ) from e
+
+    from scao.optimizer import SCAO
+
+    # Separate parameters with / without weight decay (standard HF split)
+    decay_params, no_decay_params = [], []
+    for name, param in model.named_parameters():
+        if not param.requires_grad:
+            continue
+        if any(nd in name for nd in no_decay_names):
+            no_decay_params.append(param)
+        else:
+            decay_params.append(param)
+
+    param_groups = [
+        {"params": decay_params,    "weight_decay": training_args.weight_decay},
+        {"params": no_decay_params, "weight_decay": 0.0},
+    ]
+
+    defaults: dict[str, Any] = dict(
+        lr=training_args.learning_rate,
+        warmup_steps=int(getattr(training_args, "warmup_steps", 100)),
+    )
+    if scao_kwargs:
+        defaults.update(scao_kwargs)
+
+    optimizer = SCAO(param_groups, **defaults)
+
+    # Build HF-compatible LR scheduler
+    num_training_steps = getattr(training_args, "max_steps", -1)
+    if num_training_steps < 0:
+        # Estimate from epochs and dataset size if available
+        num_training_steps = 10_000  # safe default; user should override
+
+    scheduler = get_scheduler(
+        name=getattr(training_args, "lr_scheduler_type", "cosine"),
+        optimizer=optimizer,
+        num_warmup_steps=getattr(training_args, "warmup_steps", 0),
+        num_training_steps=num_training_steps,
+    )
+
+    return optimizer, scheduler
+
+
+# ---------------------------------------------------------------------------
+# Path B: SCAOTrainer
+# ---------------------------------------------------------------------------
+
+def _make_scao_trainer_class():
+    """
+    Lazily build SCAOTrainer to avoid importing transformers at module load.
+    """
+    try:
+        from transformers import Trainer  # type: ignore[import]
+    except ImportError:
+        return None
+
+    class SCAOTrainer(Trainer):
+        """
+        Drop-in ``Trainer`` subclass that uses SCAO as the optimizer.
+
+        Args:
+            scao_kwargs: dict of extra SCAO hyperparameters (e.g.
+                ``precond_freq``, ``k_min``, ``k_max``, ``tau``).
+
+        All other arguments are passed through to ``transformers.Trainer``.
+
+        Example::
+
+            trainer = SCAOTrainer(
+                model=model,
+                args=training_args,
+                train_dataset=dataset,
+                scao_kwargs=dict(precond_freq=20, k_max=64),
+            )
+            trainer.train()
+        """
+
+        def __init__(self, *args, scao_kwargs: dict | None = None, **kwargs):
+            self._scao_kwargs = scao_kwargs or {}
+            super().__init__(*args, **kwargs)
+
+        def create_optimizer(self):
+            if self.optimizer is None:
+                optimizer, _ = get_scao_optimizer(
+                    self.model,
+                    self.args,
+                    scao_kwargs=self._scao_kwargs,
+                )
+                self.optimizer = optimizer
+            return self.optimizer
+
+        def create_optimizer_and_scheduler(self, num_training_steps: int):
+            self.create_optimizer()
+            if self.lr_scheduler is None:
+                try:
+                    from transformers import get_scheduler  # type: ignore[import]
+                    self.lr_scheduler = get_scheduler(
+                        name=self.args.lr_scheduler_type,
+                        optimizer=self.optimizer,
+                        num_warmup_steps=self.args.get_warmup_steps(num_training_steps),
+                        num_training_steps=num_training_steps,
+                    )
+                except Exception:
+                    super().create_optimizer_and_scheduler(num_training_steps)
+
+    return SCAOTrainer
+
+
+# Expose SCAOTrainer — will be None if transformers not installed
+SCAOTrainer = _make_scao_trainer_class()
+
+
+# ---------------------------------------------------------------------------
+# Path C: SCAOMonitorCallback
+# ---------------------------------------------------------------------------
+
+def _make_monitor_callback():
+    try:
+        from transformers import TrainerCallback  # type: ignore[import]
+    except ImportError:
+        return None
+
+    class SCAOMonitorCallback(TrainerCallback):
+        """
+        HuggingFace ``TrainerCallback`` that logs SCAO rank and curvature
+        health metrics to the Trainer's log at every logging step.
+
+        Args:
+            optimizer: the SCAO optimizer instance
+            log_every: override log frequency (default: use Trainer's setting)
+
+        Example::
+
+            trainer.add_callback(SCAOMonitorCallback(optimizer))
+        """
+
+        def __init__(self, optimizer, log_every: int | None = None) -> None:
+            self.optimizer = optimizer
+            self.log_every = log_every
+
+        def on_log(self, args, state, control, logs=None, **kwargs):
+            from scao.logging import collect_metrics
+            metrics = collect_metrics(self.optimizer)
+            if logs is not None:
+                # Inject numeric metrics into the Trainer's log dict
+                logs.update(
+                    {k: v for k, v in metrics.items()
+                     if isinstance(v, (int, float)) and k != "step"}
+                )
+
+    return SCAOMonitorCallback
+
+
+SCAOMonitorCallback = _make_monitor_callback()