Release/0.2.2 production readiness

examples/finetune_distilgpt2_colab.py

@@ -0,0 +1,159 @@
+"""
+Lightweight SCAO fine-tuning example for Google Colab T4.
+
+Install dependencies:
+
+    pip install -U scao "transformers>=4.30" "datasets>=2.0" accelerate
+
+Run:
+
+    python examples/finetune_distilgpt2_colab.py
+
+This is intentionally small: it validates the real user path with Hugging Face
+Trainer, public data, fp16, and conservative SCAO defaults.
+"""
+
+from __future__ import annotations
+
+import argparse
+import inspect
+import os
+from typing import Any
+
+import torch
+from datasets import load_dataset
+from transformers import (
+    AutoModelForCausalLM,
+    AutoTokenizer,
+    DataCollatorForLanguageModeling,
+    Trainer,
+    TrainingArguments,
+)
+
+from scao.integrations.huggingface import get_scao_optimizer
+
+
+def _training_args(args: argparse.Namespace) -> TrainingArguments:
+    kwargs: dict[str, Any] = {
+        "output_dir": args.output_dir,
+        "overwrite_output_dir": True,
+        "max_steps": args.max_steps,
+        "per_device_train_batch_size": args.batch_size,
+        "per_device_eval_batch_size": args.batch_size,
+        "gradient_accumulation_steps": args.grad_accum,
+        "learning_rate": args.lr,
+        "warmup_steps": args.warmup_steps,
+        "weight_decay": args.weight_decay,
+        "fp16": torch.cuda.is_available(),
+        "logging_steps": args.logging_steps,
+        "eval_steps": max(1, args.max_steps // 2),
+        "save_strategy": "no",
+        "report_to": "none",
+        "remove_unused_columns": False,
+        "seed": args.seed,
+    }
+
+    signature = inspect.signature(TrainingArguments)
+    if "eval_strategy" in signature.parameters:
+        kwargs["eval_strategy"] = "steps"
+    elif "evaluation_strategy" in signature.parameters:
+        kwargs["evaluation_strategy"] = "steps"
+
+    kwargs = {key: value for key, value in kwargs.items() if key in signature.parameters}
+    return TrainingArguments(**kwargs)
+
+
+def _dataset(tokenizer: Any, rows: int, block_size: int) -> Any:
+    dataset = load_dataset("wikitext", "wikitext-2-raw-v1", split="train")
+    dataset = dataset.filter(lambda row: bool(row["text"].strip()))
+    dataset = dataset.select(range(min(rows, len(dataset))))
+
+    def tokenize(batch: dict[str, list[str]]) -> dict[str, Any]:
+        return tokenizer(
+            batch["text"],
+            truncation=True,
+            max_length=block_size,
+            padding="max_length",
+        )
+
+    tokenized = dataset.map(tokenize, batched=True, remove_columns=dataset.column_names)
+    return tokenized.train_test_split(test_size=0.05, seed=42)
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser(description="Fine-tune distilgpt2 with SCAO")
+    parser.add_argument("--model", default="distilgpt2")
+    parser.add_argument("--output-dir", default="/content/scao-distilgpt2")
+    parser.add_argument("--dataset-rows", type=int, default=1000)
+    parser.add_argument("--block-size", type=int, default=128)
+    parser.add_argument("--max-steps", type=int, default=150)
+    parser.add_argument("--batch-size", type=int, default=2)
+    parser.add_argument("--grad-accum", type=int, default=4)
+    parser.add_argument("--lr", type=float, default=3e-4)
+    parser.add_argument("--warmup-steps", type=int, default=10)
+    parser.add_argument("--weight-decay", type=float, default=0.01)
+    parser.add_argument("--logging-steps", type=int, default=10)
+    parser.add_argument("--seed", type=int, default=42)
+    parser.add_argument(
+        "--prompt",
+        default="Sparse optimization helps language models",
+    )
+    args = parser.parse_args()
+
+    os.environ.setdefault("TOKENIZERS_PARALLELISM", "false")
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    print("device:", device)
+    if device.type == "cuda":
+        print("gpu:", torch.cuda.get_device_name(0))
+
+    tokenizer = AutoTokenizer.from_pretrained(args.model)
+    if tokenizer.pad_token is None:
+        tokenizer.pad_token = tokenizer.eos_token
+
+    model = AutoModelForCausalLM.from_pretrained(args.model).to(device)
+    model.config.pad_token_id = tokenizer.pad_token_id
+
+    split = _dataset(tokenizer, args.dataset_rows, args.block_size)
+    training_args = _training_args(args)
+    optimizer, scheduler = get_scao_optimizer(
+        model,
+        training_args,
+        scao_kwargs={
+            "precond_freq": 20,
+            "max_precond_dim": 1024,
+            "k_min": 4,
+            "k_max": 32,
+            "async_precond": True,
+            "noise_std_init": 0.0,
+            "sparsity": 0.0,
+            "lars_coeff": 0.0,
+            "lookahead_k": 0,
+        },
+        num_training_steps=args.max_steps,
+    )
+
+    trainer = Trainer(
+        model=model,
+        args=training_args,
+        train_dataset=split["train"],
+        eval_dataset=split["test"],
+        data_collator=DataCollatorForLanguageModeling(tokenizer=tokenizer, mlm=False),
+        optimizers=(optimizer, scheduler),
+    )
+    trainer.train()
+    print(trainer.evaluate())
+
+    inputs = tokenizer(args.prompt, return_tensors="pt").to(device)
+    output = model.generate(
+        **inputs,
+        max_new_tokens=80,
+        do_sample=True,
+        temperature=0.8,
+        top_p=0.95,
+        pad_token_id=tokenizer.pad_token_id,
+    )
+    print(tokenizer.decode(output[0], skip_special_tokens=True))
+
+
+if __name__ == "__main__":
+    main()

pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "scao"
-version = "0.2.0"
+version = "0.2.2"
 description = "Sparse Curvature-Aware Adaptive Optimizer — second-order training at near-AdamW cost"
 readme = "README.md"
 requires-python = ">=3.10"
@@ -40,6 +40,10 @@ Documentation = "https://github.com/whispering3/scao#readme"
 Changelog = "https://github.com/whispering3/scao/blob/main/CHANGELOG.md"
 Paper = "https://doi.org/10.5281/zenodo.19839495"
 
+[project.scripts]
+scao-scale-planner = "scao.tools.scale_planner:main"
+scao-preconditioner-bench = "scao.tools.preconditioner_microbench:main"
+
 [project.optional-dependencies]
 dev = [
   "pytest>=7.0",
@@ -72,7 +76,7 @@ all = [
 [tool.setuptools.packages.find]
 where = ["."]
 include = ["scao*"]
-exclude = ["scao/tests*", "scao/benchmarks*"]
+exclude = ["scao.tests*", "scao.benchmarks*"]
 
 [tool.setuptools.package-data]
 scao = ["py.typed"]

scao/__init__.py

@@ -31,7 +31,7 @@
 from .preconditioner import SparsePreconditioner
 from .utils import adaptive_rank, matrix_power_neg_quarter
 
-__version__ = "0.2.0"
+__version__ = "0.2.2"
 __author__ = "Danilo Souza"
 __license__ = "Apache-2.0"
 

scao/cuda/__init__.py

@@ -113,11 +113,11 @@ def _fused_kronecker_precond_pytorch(
         delta = (s_l_inv4[:,None] * s_r_inv4[None,:] - 1) * P
         out   = G + U_l @ delta @ U_r^T
     """
-    P = (U_l.T @ G) @ U_r                                    # (k, k)
-    scale = s_l_inv4.unsqueeze(1) * s_r_inv4.unsqueeze(0)    # (k, k)
-    delta = (scale - 1.0) * P                                 # (k, k)
-    correction = U_l @ delta @ U_r.T                          # (m, n)
-    return G + correction
+    P = torch.mm(torch.mm(U_l.T, G), U_r)                     # (k, k)
+    scale = torch.outer(s_l_inv4, s_r_inv4)                   # (k, k)
+    delta = (scale - 1.0).mul_(P)                             # (k, k)
+    correction_left = torch.mm(U_l, delta)                    # (m, k)
+    return torch.addmm(G, correction_left, U_r.T)             # (m, n)
 
 
 def _low_rank_precond_mm_pytorch(

scao/cuda/setup.py

@@ -1,10 +1,31 @@
 """
 Optional CUDA extension build script for SCAO.
 Run:  python setup.py build_ext --inplace
+
+For production builds, prefer setting TORCH_CUDA_ARCH_LIST explicitly, e.g.:
+    TORCH_CUDA_ARCH_LIST="7.5;8.0;8.6;8.9;9.0" python setup.py build_ext --inplace
 """
 
+import os
+
 from setuptools import setup
 
+DEFAULT_NVCC_FLAGS = [
+    "-O3",
+    "--use_fast_math",
+]
+
+if not os.environ.get("TORCH_CUDA_ARCH_LIST"):
+    DEFAULT_NVCC_FLAGS.extend(
+        [
+            "-gencode=arch=compute_75,code=sm_75",   # T4
+            "-gencode=arch=compute_80,code=sm_80",   # A100
+            "-gencode=arch=compute_86,code=sm_86",   # A10/A10G/RTX 3090
+            "-gencode=arch=compute_89,code=sm_89",   # RTX 4090/L4/L40S
+            "-gencode=arch=compute_90,code=sm_90",   # H100
+        ]
+    )
+
 try:
     from torch.utils.cpp_extension import BuildExtension, CUDAExtension
     ext_modules = [
@@ -13,16 +34,7 @@
             sources=["low_rank_ops.cu"],
             extra_compile_args={
                 "cxx": ["-O3"],
-                "nvcc": [
-                    "-O3",
-                    "--use_fast_math",
-                    "-gencode=arch=compute_70,code=sm_70",   # V100
-                    "-gencode=arch=compute_75,code=sm_75",   # T4 / RTX 20xx
-                    "-gencode=arch=compute_80,code=sm_80",   # A100
-                    "-gencode=arch=compute_86,code=sm_86",   # RTX 30xx / A40
-                    "-gencode=arch=compute_89,code=sm_89",   # RTX 40xx / H100
-                    "-gencode=arch=compute_90,code=sm_90",   # H100 SXM
-                ],
+                "nvcc": DEFAULT_NVCC_FLAGS,
             },
         )
     ]