Fix SCAO optimizer stability and validation

pyproject.toml

@@ -93,6 +93,10 @@ disallow_untyped_defs = true
 disallow_incomplete_defs = true
 check_untyped_defs = true
 no_implicit_optional = true
+exclude = [
+  "scao/tests/",
+  "scao/benchmarks/",
+]
 
 [tool.ruff]
 line-length = 100

scao/__init__.py

@@ -26,12 +26,10 @@
 Zenodo 2026
 """
 
-from .optimizer import (
-    SCAO, scao_sub1b, scao_1b, scao_3b, scao_7b, scao_40b, scao_125b
-)
-from .preconditioner import SparsePreconditioner
-from .utils import matrix_power_neg_quarter, adaptive_rank
 from . import logging as scao_logging
+from .optimizer import SCAO, scao_1b, scao_3b, scao_7b, scao_40b, scao_125b, scao_sub1b
+from .preconditioner import SparsePreconditioner
+from .utils import adaptive_rank, matrix_power_neg_quarter
 
 __version__ = "0.2.0"
 __author__ = "Danilo Souza"

scao/benchmarks/compare_adamw_scao.py

@@ -16,8 +16,9 @@
 
 import argparse
 import math
+import os
+import sys
 import time
-import sys, os
 
 sys.path.insert(0, os.path.join(os.path.dirname(__file__), "..", ".."))
 
@@ -27,7 +28,6 @@
 
 from scao import SCAO
 
-
 # ---------------------------------------------------------------------------
 # Tiny GPT-like model
 # ---------------------------------------------------------------------------
@@ -229,7 +229,7 @@ def main():
         results.append(r)
 
     print(f"\n{'='*60}")
-    print(f"  SUMMARY")
+    print("  SUMMARY")
     print(f"{'='*60}")
     for r in results:
         print(f"\n  {r['optimizer'].upper()}")
@@ -244,7 +244,7 @@ def main():
         a, s = results[0], results[1]
         speedup = a["avg_loss_last_20"] / max(s["avg_loss_last_20"], 1e-9)
         print(f"\n  SCAO loss ratio vs AdamW: {speedup:.3f}x")
-        print(f"  (>1 means SCAO reached lower loss in same steps)")
+        print("  (>1 means SCAO reached lower loss in same steps)")
     print()
 
 

scao/benchmarks/gpt_comparison.py

@@ -29,7 +29,6 @@
 import os
 import sys
 import time
-from collections import defaultdict
 from typing import Any
 
 sys.path.insert(0, os.path.join(os.path.dirname(__file__), "..", ".."))
@@ -40,7 +39,6 @@
 
 from scao import SCAO
 
-
 # ===========================================================================
 # Baseline optimizers (no external deps)
 # ===========================================================================
@@ -63,7 +61,7 @@ def __init__(
         eps: float = 1e-8,
         weight_decay: float = 0.01,
     ) -> None:
-        defaults = dict(lr=lr, beta=beta, eps=eps, weight_decay=weight_decay)
+        defaults = {"lr": lr, "beta": beta, "eps": eps, "weight_decay": weight_decay}
         super().__init__(params, defaults)
 
     @torch.no_grad()
@@ -122,8 +120,12 @@ def __init__(
         weight_decay: float = 0.01,
         ns_steps: int = 5,
     ) -> None:
-        defaults = dict(lr=lr, momentum=momentum, weight_decay=weight_decay,
-                        ns_steps=ns_steps)
+        defaults = {
+            "lr": lr,
+            "momentum": momentum,
+            "weight_decay": weight_decay,
+            "ns_steps": ns_steps,
+        }
         super().__init__(params, defaults)
 
     @staticmethod
@@ -410,7 +412,7 @@ def _color(name: str, text: str) -> str:
 
 def ascii_loss_chart(results: list[dict], width: int = 72, height: int = 20) -> None:
     """Print a compact ASCII loss curve for all optimizers."""
-    all_losses = [l for r in results for l in r["losses"]]
+    all_losses = [loss for r in results for loss in r["losses"]]
     y_min = min(all_losses) * 0.98
     y_max = max(all_losses[: len(results[0]["losses"]) // 5]) * 1.02  # zoom on early
 
@@ -486,7 +488,7 @@ def print_summary(results: list[dict], threshold: float) -> None:
     if scao_r and adamw_r:
         auc_ratio     = adamw_r["auc"]      / max(scao_r["auc"],     1e-9)
         speed_ratio   = scao_r["steps_per_sec"] / max(adamw_r["steps_per_sec"], 1e-9)
-        print(f"\n  SCAO vs AdamW:")
+        print("\n  SCAO vs AdamW:")
         print(f"    AUC ratio      : {auc_ratio:.3f}x  (>1 means SCAO converges faster)")
         print(f"    Throughput     : {speed_ratio:.3f}x  (1 = same speed)")
         if scao_r.get("peak_memory_mb") and adamw_r.get("peak_memory_mb"):

scao/benchmarks/gpt_scale_benchmark.py

@@ -41,7 +41,6 @@
 import time
 import tracemalloc
 from dataclasses import dataclass
-from typing import Iterator
 
 sys.path.insert(0, os.path.join(os.path.dirname(__file__), "..", ".."))
 
@@ -251,7 +250,10 @@ class DiagonalShampoo(torch.optim.Optimizer):
 
     def __init__(self, params, lr: float = 1e-3, eps: float = 1e-8,
                  weight_decay: float = 0.0, rho: float = 0.999) -> None:
-        super().__init__(params, dict(lr=lr, eps=eps, weight_decay=weight_decay, rho=rho))
+        super().__init__(
+            params,
+            {"lr": lr, "eps": eps, "weight_decay": weight_decay, "rho": rho},
+        )
 
     @torch.no_grad()
     def step(self, closure=None):
@@ -439,7 +441,7 @@ def run_single(
         grad_bytes  = param_bytes  # one gradient tensor per param (same size)
         opt_bytes = 0
         for state in optimizer.state.values():
-            for k, v in state.items():
+            for _k, v in state.items():
                 if isinstance(v, torch.Tensor):
                     opt_bytes += v.numel() * v.element_size()
                 elif hasattr(v, "memory_bytes"):
@@ -532,7 +534,7 @@ def main() -> None:
         warmup_steps = args.warmup   or min(200, steps // 5)
 
         print(f"\n{'=' * 72}")
-        print(f"  SCAO Multi-Scale Benchmark")
+        print("  SCAO Multi-Scale Benchmark")
         print(f"  Scale: {scale_key} ({cfg.label} params)  "
               f"d={cfg.d_model}  L={cfg.n_layers}  H={cfg.n_head}")
         print(f"  device={device}  steps={steps}  batch={batch_size}  seeds={seeds}")
@@ -572,8 +574,8 @@ def main() -> None:
                       f"t={r['total_time_s']:.1f}s")
 
     # --- Summary table ---
-    from collections import defaultdict
     import statistics
+    from collections import defaultdict
 
     # Group by (scale, optimizer) across seeds
     groups: dict[tuple, list[dict]] = defaultdict(list)

scao/benchmarks/synthetic_benchmark.py

@@ -47,22 +47,20 @@
 import json
 import math
 import os
+import random
 import sys
 import time
 import tracemalloc
 from contextlib import contextmanager
-from dataclasses import dataclass, asdict, field
-from typing import Iterator, Optional
+from dataclasses import asdict, dataclass, field
 
 sys.path.insert(0, os.path.join(os.path.dirname(__file__), "..", ".."))
 
 import torch
 import torch.nn as nn
-import torch.nn.functional as F
 
 from scao import SCAO
 
-
 # ===========================================================================
 # Profiling helpers
 # ===========================================================================
@@ -180,7 +178,7 @@ class RunResult:
 
 def _seed_everything(seed: int) -> None:
     torch.manual_seed(seed)
-    import random; random.seed(seed)
+    random.seed(seed)
 
 
 def make_markov_text(
@@ -380,7 +378,10 @@ def forward(self, x: torch.Tensor) -> torch.Tensor:
 
 class DiagonalShampoo(torch.optim.Optimizer):
     def __init__(self, params, lr=1e-3, eps=1e-8, weight_decay=0.0, rho=0.999):
-        super().__init__(params, dict(lr=lr, eps=eps, weight_decay=weight_decay, rho=rho))
+        super().__init__(
+            params,
+            {"lr": lr, "eps": eps, "weight_decay": weight_decay, "rho": rho},
+        )
 
     @torch.no_grad()
     def step(self, closure=None):
@@ -413,7 +414,13 @@ def make_optimizer(name: str, params, lr: float, steps: int,
             warmup_steps=warmup_steps,
             precond_freq=max(5, steps // 80),
             min_precond_updates=5,
-            k_min=4, k_max=32, tau=None,
+            k_min=4, k_max=32,
+            tau=1.0,
+            sparsity=0.0,
+            noise_std_init=0.0,
+            lars_coeff=0.0,
+            lookahead_k=0,
+            adaptive_warmup=False,
         )
     if name == "diag_shampoo":
         return DiagonalShampoo(params, lr=lr * 3, weight_decay=0.01)
@@ -482,7 +489,7 @@ def run_lm(
     gc.collect()
     tracemalloc.start()
 
-    for step in range(1, steps + 1):
+    for _step in range(1, steps + 1):
         idx  = torch.randint(0, n_train, (batch_size,))
         xb   = train_x[idx]
         yb   = train_y[idx]
@@ -497,8 +504,7 @@ def run_lm(
         logits = model(xb)
         loss   = loss_fn(logits.reshape(-1, vocab), yb.reshape(-1))
         loss.backward()
-        if opt_name != "scao":
-            nn.utils.clip_grad_norm_(model.parameters(), 1.0)
+        nn.utils.clip_grad_norm_(model.parameters(), 1.0)
         opt.step()
         sched.step()
 
@@ -582,7 +588,7 @@ def run_regression(
     gc.collect()
     tracemalloc.start()
 
-    for step in range(1, steps + 1):
+    for _step in range(1, steps + 1):
         idx  = torch.randint(0, n_train, (batch_size,))
         xb, yb = train_x[idx], train_y[idx]
 
@@ -595,8 +601,7 @@ def run_regression(
         pred = model(xb)
         loss = loss_fn(pred, yb)
         loss.backward()
-        if opt_name != "scao":
-            nn.utils.clip_grad_norm_(model.parameters(), 1.0)
+        nn.utils.clip_grad_norm_(model.parameters(), 1.0)
         opt.step()
         sched.step()
 

scao/benchmarks/wikitext_benchmark.py

@@ -19,7 +19,6 @@
 from __future__ import annotations
 
 import argparse
-import csv
 import math
 import os
 import sys
@@ -29,18 +28,15 @@
 
 import torch
 import torch.nn as nn
-import torch.nn.functional as F
 
 from scao import SCAO
 from scao.benchmarks.gpt_comparison import (
-    TinyGPT,
     DiagonalShampoo,
+    TinyGPT,
     ascii_loss_chart,
-    print_summary,
     save_csv,
 )
 
-
 # ===========================================================================
 # Data pipeline
 # ===========================================================================
@@ -354,7 +350,7 @@ def main():
     if scao_r and adamw_r:
         val_ratio = adamw_r["final_val"] / max(scao_r["final_val"], 1e-9)
         ppl_delta = adamw_r["final_ppl"] - scao_r["final_ppl"]
-        print(f"\n  SCAO vs AdamW:")
+        print("\n  SCAO vs AdamW:")
         print(f"    Val loss ratio : {val_ratio:.4f}x")
         print(f"    PPL improvement: {ppl_delta:+.2f} (positive = SCAO better)")
     print(f"{'=' * w}\n")