diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index e773f63..0084fad 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -8,12 +8,12 @@ on: jobs: test: - name: "Test · Python ${{ matrix.python-version }}" + name: "Test · Linux · Python ${{ matrix.python-version }}" runs-on: ubuntu-latest strategy: fail-fast: false matrix: - python-version: ["3.10", "3.11", "3.12"] + python-version: ["3.10", "3.11", "3.12", "3.13"] steps: - uses: actions/checkout@v4 @@ -31,10 +31,10 @@ jobs: pip install -e ".[dev]" - name: Lint with ruff - run: ruff check scao/ || true + run: ruff check . - name: Type-check with mypy - run: mypy scao/ --ignore-missing-imports --exclude 'scao/(benchmarks|tests|cuda)' || true + run: mypy scao/ --ignore-missing-imports --exclude 'scao/(benchmarks|tests|cuda)' - name: Run tests run: pytest scao/tests/ -v --tb=short --cov=scao --cov-report=xml --ignore=scao/tests/test_profiling.py --ignore=scao/tests/test_ddp.py @@ -48,6 +48,27 @@ jobs: files: coverage.xml fail_ci_if_error: false + test-windows: + name: "Test · Windows · Python 3.11" + runs-on: windows-latest + steps: + - uses: actions/checkout@v4 + + - name: Set up Python 3.11 + uses: actions/setup-python@v5 + with: + python-version: "3.11" + cache: pip + + - name: Install dependencies + run: | + python -m pip install --upgrade pip + pip install torch --index-url https://download.pytorch.org/whl/cpu + pip install -e ".[dev]" + + - name: Run unit & DDP tests + run: pytest scao/tests/ -v --tb=short --ignore=scao/tests/test_profiling.py + benchmark-smoke: name: "Benchmark smoke test" runs-on: ubuntu-latest @@ -63,9 +84,7 @@ jobs: pip install -e . - name: Run quick benchmark run: | - python scao/benchmarks/gpt_scale_benchmark.py \ - --scales tiny --steps 30 --optimizers "adamw,scao" \ - --seeds "42" --csv /tmp/smoke.csv + python scao/benchmarks/gpt_scale_benchmark.py --scales tiny --steps 30 --optimizers "adamw,scao" --seeds "42" --csv /tmp/smoke.csv build: name: "Build distribution" diff --git a/.gitignore b/.gitignore index acfcb66..399d337 100644 --- a/.gitignore +++ b/.gitignore @@ -26,6 +26,7 @@ MANIFEST .build/ *.whl *.tar.gz +scao-*/ # Testing .tox/ diff --git a/CHANGELOG.md b/CHANGELOG.md index 4de21cc..63129c6 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -5,6 +5,24 @@ Format follows [Keep a Changelog](https://keepachangelog.com/en/1.1.0/). --- +## [0.2.2] — 2026-10-04 + +### Added +- **Triton JIT Acceleration**: Added OpenAI Triton kernels (`scao/cuda/triton_ops.py`) for GPU acceleration without requiring a precompiled C++/CUDA extension or host `nvcc` compiler. +- **Axolotl Integration**: Official native support for Axolotl fine-tuning framework (`optimizer: scao`) documented in README and verified with E2E tests. +- **Cross-Platform CI**: Added Windows runner (`windows-latest`) to GitHub Actions CI matrix to test DDP and file paths natively on Windows. +- **Python 3.13 Support**: Added Python 3.13 to CI test matrix. + +### Changed +- **Modular Core Architecture**: Refactored `scao/optimizer.py` into focused, reusable modules while preserving 100% backward compatibility: + - `scao/filters.py`: Gradient filters (`_SparseGradFilter`, `_DynamicSparseFilter`), adaptive warmup, gSNR element-wise clipping, and lazy triggers. + - `scao/presets.py`: Model scale presets (`scao_sub1b` through `scao_125b`). + - `scao/optimizer.py`: Lean optimizer step loop and state management. +- **Strict CI Verification**: Removed `|| true` masking from `ruff` and `mypy` in CI workflows, enforcing zero-regression quality gates. +- **Repository Hygiene**: Removed legacy build directories and added `scao-*/` to `.gitignore`. + +--- + ## [0.2.0] — 2026-04-28 ### Added (SCAO v2) diff --git a/README.md b/README.md index f770657..712f413 100644 --- a/README.md +++ b/README.md @@ -6,10 +6,12 @@ [![DOI](https://zenodo.org/badge/DOI/10.5281/zenodo.19870556.svg)](https://doi.org/10.5281/zenodo.19870556) [![Python](https://img.shields.io/badge/python-3.10%2B-blue)](https://python.org) [![PyTorch](https://img.shields.io/badge/pytorch-2.0%2B-orange)](https://pytorch.org) +[![Axolotl](https://img.shields.io/badge/Axolotl-supported-blueviolet.svg)](https://github.com/axolotl-ai-cloud/axolotl/pull/3624) > **A second-order PyTorch optimizer that delivers Shampoo-quality preconditioned gradients at near-AdamW memory and throughput cost.** > Drop-in replacement for `AdamW`. One-line change. Real gains. -> **Now available on PyPI:** `pip install scao` +> **Now available on PyPI:** `pip install scao` +> **Native support:** [Axolotl](https://github.com/axolotl-ai-cloud/axolotl/pull/3624) (`optimizer: scao`) and HuggingFace `Trainer`. --- @@ -555,6 +557,26 @@ trainer = SCAOTrainer( trainer.train() ``` +### Axolotl (LLM Fine-Tuning) + +SCAO is integrated as a native custom optimizer in [Axolotl](https://github.com/axolotl-ai-cloud/axolotl) (PR [#3624](https://github.com/axolotl-ai-cloud/axolotl/pull/3624)). + +To use SCAO in your Axolotl YAML configuration: + +```yaml +# In your Axolotl config.yaml +optimizer: scao + +# Optional SCAO hyperparameters passed through optim_args +optim_args: + precond_freq: 20 # Curvature update frequency + k_min: 8 # Minimum rank + k_max: 64 # Maximum rank + sparsity: 0.7 # Base gradient sparsity + use_int8_ema: true # 4x memory reduction on curvature buffers (recommended for LoRA/QLoRA) + dynamic_sparsity: true # Adaptive per-layer gradient sparsity +``` + ### Monitoring and diagnostics ```python @@ -687,20 +709,26 @@ Open [`scripts/scao_colab_benchmark.ipynb`](scripts/scao_colab_benchmark.ipynb) ``` scao/ # Core library ├── optimizer.py # SCAO main class — drop-in for AdamW +├── filters.py # Gradient filters: dynamic sparsity, gSNR clipping, warmup +├── presets.py # Scale presets: scao_sub1b through scao_125b ├── preconditioner.py # SparsePreconditioner: Kronecker low-rank + int8 EMA ├── utils.py # adaptive_rank, quantize_sym_int8, dequantize_sym_int8 ├── distributed.py # ZeRO-3 / FSDP helpers ├── logging.py # ConsoleLogger, TensorBoardLogger, WandbLogger ├── integrations/ -│ └── huggingface.py # SCAOTrainer, SCAOMonitorCallback +│ └── huggingface.py # SCAOTrainer, get_scao_optimizer, SCAOMonitorCallback ├── benchmarks/ │ └── gpt_scale_benchmark.py # Multi-scale GPT: SCAO vs AdamW vs SCAO-int8 ├── tests/ -│ ├── test_optimizer.py # 40 optimizer correctness tests +│ ├── test_optimizer.py # 40+ optimizer correctness tests +│ ├── test_ddp.py # DistributedDataParallel tests (Gloo / CPU) +│ ├── test_huggingface_integration.py # HF Trainer & PEFT integration tests +│ ├── test_scale_planner.py # Memory & scaling verification tests │ └── test_profiling.py # 26 memory + timing profiling tests └── cuda/ ├── low_rank_ops.cu # Fused CUDA kernels: tiled GEMM, Kronecker precond, int8 EMA - ├── __init__.py # fused_kronecker_precond(), int8_ema_update(), truncated_eigh() + ├── triton_ops.py # Triton JIT kernels for GPU acceleration without nvcc + ├── __init__.py # fused_kronecker_precond(), int8_ema_update() └── setup.py # nvcc build (sm_70/75/80/86/89/90) scao_benchmarks_t4/ # Unified T4/Colab benchmark suite diff --git a/benchmark/inference.py b/benchmark/inference.py index f5bd86e..2bb855b 100644 --- a/benchmark/inference.py +++ b/benchmark/inference.py @@ -14,8 +14,8 @@ """ import torch -from transformers import AutoModelForCausalLM, AutoTokenizer from peft import PeftModel +from transformers import AutoModelForCausalLM, AutoTokenizer def main(): @@ -60,4 +60,4 @@ def main(): if __name__ == "__main__": - main() \ No newline at end of file + main() diff --git a/benchmark/scao_vs_adamw_bench/benchmark_4b.py b/benchmark/scao_vs_adamw_bench/benchmark_4b.py index 3e74606..5ac989a 100644 --- a/benchmark/scao_vs_adamw_bench/benchmark_4b.py +++ b/benchmark/scao_vs_adamw_bench/benchmark_4b.py @@ -1,15 +1,24 @@ import torch -from transformers import AutoModelForCausalLM, AutoTokenizer, TrainingArguments, Trainer, BitsAndBytesConfig, DataCollatorForLanguageModeling from datasets import load_dataset -from peft import get_peft_model, LoraConfig, prepare_model_for_kbit_training -from scao import SCAO # Your 2nd-order optimizer implementation +from peft import LoraConfig, get_peft_model, prepare_model_for_kbit_training +from transformers import ( + AutoModelForCausalLM, + AutoTokenizer, + BitsAndBytesConfig, + DataCollatorForLanguageModeling, + Trainer, + TrainingArguments, +) + +from scao import SCAO # Your 2nd-order optimizer implementation + def main(): print("🚀 Starting 4B-Scale Benchmark for SCAO...") - + # Using Qwen 2.5 3B model (optimal for the 4B category tests) model_id = "Qwen/Qwen2.5-3B" - + tokenizer = AutoTokenizer.from_pretrained(model_id) if tokenizer.pad_token is None: tokenizer.pad_token = tokenizer.eos_token @@ -30,7 +39,7 @@ def main(): print("🧠 Initializing LoRA adapters...") model = prepare_model_for_kbit_training(model) - + lora_config = LoraConfig( r=16, lora_alpha=32, @@ -48,10 +57,10 @@ def main(): print("📚 Loading dataset...") # Using wikitext-2 for consistency across benchmarks dataset = load_dataset("wikitext", "wikitext-2-raw-v1", split="train[:5%]") - + def tokenize(example): return tokenizer(example["text"], padding="max_length", truncation=True, max_length=256) - + tokenized_dataset = dataset.map(tokenize, batched=True) print("⚙️ Injecting SCAO Optimizer...") diff --git a/benchmark/scao_vs_shampoo_bench/scao_vs_shampoo_pro.py b/benchmark/scao_vs_shampoo_bench/scao_vs_shampoo_pro.py index 30795d5..6f2c04f 100644 --- a/benchmark/scao_vs_shampoo_bench/scao_vs_shampoo_pro.py +++ b/benchmark/scao_vs_shampoo_bench/scao_vs_shampoo_pro.py @@ -1,20 +1,23 @@ -import torch import argparse -import time import json import os +import time + +import torch +import torch_optimizer as optim +from datasets import load_dataset +from peft import LoraConfig, get_peft_model, prepare_model_for_kbit_training from transformers import ( - AutoModelForCausalLM, - AutoTokenizer, - TrainingArguments, - Trainer, - BitsAndBytesConfig, - DataCollatorForLanguageModeling + AutoModelForCausalLM, + AutoTokenizer, + BitsAndBytesConfig, + DataCollatorForLanguageModeling, + Trainer, + TrainingArguments, ) -from datasets import load_dataset -from peft import get_peft_model, LoraConfig, prepare_model_for_kbit_training + from scao import SCAO -import torch_optimizer as optim + def get_peak_vram(): if torch.cuda.is_available(): @@ -29,7 +32,7 @@ def main(): args_cmd = parser.parse_args() print(f"🚀 Starting Professional Benchmark: {args_cmd.optimizer.upper()} on {args_cmd.model_id}") - + # 1. Model Loading (4-bit) tokenizer = AutoTokenizer.from_pretrained(args_cmd.model_id) if tokenizer.pad_token is None: @@ -60,7 +63,7 @@ def main(): ) model = get_peft_model(model, lora_config) trainable_params = [p for p in model.parameters() if p.requires_grad] - + # 3. Dataset dataset = load_dataset("wikitext", "wikitext-2-raw-v1", split="train[:1%]") def tokenize(example): @@ -100,7 +103,7 @@ def tokenize(example): start_time = time.time() status = "Success" error_msg = "" - + try: print("⚡ Training active! Monitoring performance...") trainer.train() @@ -129,7 +132,7 @@ def tokenize(example): # Append to markdown summary summary_path = "benchmark_results.md" header_exists = os.path.exists(summary_path) and os.stat(summary_path).st_size > 0 - + with open(summary_path, "a") as f: if not header_exists: f.write("# SCAO vs Shampoo Comparison Results\n\n") diff --git a/benchmark/train_1m.py b/benchmark/train_1m.py index f691d29..71d42cc 100644 --- a/benchmark/train_1m.py +++ b/benchmark/train_1m.py @@ -16,15 +16,15 @@ python train_1m.py """ -import torch +from datasets import load_dataset from transformers import ( AutoModelForCausalLM, AutoTokenizer, - TrainingArguments, - Trainer, DataCollatorForLanguageModeling, + Trainer, + TrainingArguments, ) -from datasets import load_dataset + from scao import SCAO @@ -83,4 +83,4 @@ def tokenize(example): if __name__ == "__main__": - main() \ No newline at end of file + main() diff --git a/benchmark/train_local.py b/benchmark/train_local.py index 93522dc..ec53b7b 100644 --- a/benchmark/train_local.py +++ b/benchmark/train_local.py @@ -12,16 +12,16 @@ python train_local.py """ -import torch +from datasets import load_dataset +from peft import LoraConfig, TaskType, get_peft_model from transformers import ( AutoModelForCausalLM, AutoTokenizer, - TrainingArguments, - Trainer, DataCollatorForLanguageModeling, + Trainer, + TrainingArguments, ) -from datasets import load_dataset -from peft import LoraConfig, get_peft_model, TaskType + from scao import SCAO @@ -98,4 +98,4 @@ def tokenize(example): if __name__ == "__main__": - main() \ No newline at end of file + main() diff --git a/pyproject.toml b/pyproject.toml index 979acf4..aaaeec3 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -105,6 +105,7 @@ exclude = [ [tool.ruff] line-length = 100 target-version = "py310" +extend-exclude = ["*.ipynb", "dist", "build", "scao.egg-info"] [tool.ruff.lint] select = ["E", "F", "W", "I", "UP", "B", "C4", "T20"] @@ -116,5 +117,7 @@ ignore = [ [tool.ruff.lint.per-file-ignores] "scao/tests/*" = ["T201", "F841"] -"scao/benchmarks/*" = ["T201"] +"scao/benchmarks/*" = ["T201", "E402"] +"benchmark/*" = ["T201", "E402", "F401", "I001"] +"scripts/*" = ["T201", "E402", "E701", "E702", "B007", "F401", "I001"] diff --git a/scao/cuda/__init__.py b/scao/cuda/__init__.py index d9fd271..21585bc 100644 --- a/scao/cuda/__init__.py +++ b/scao/cuda/__init__.py @@ -287,4 +287,14 @@ def int8_ema_update( return ext.int8_ema_update(ema_q, ema_scale, new_val, rho) except (AttributeError, RuntimeError): pass + + if ema_q.is_cuda: + try: + from .triton_ops import int8_ema_update_triton + res = int8_ema_update_triton(ema_q, ema_scale, new_val, rho) + if res is not None: + return res + except Exception: + pass + return _int8_ema_update_pytorch(ema_q, ema_scale, new_val, rho) diff --git a/scao/cuda/triton_ops.py b/scao/cuda/triton_ops.py new file mode 100644 index 0000000..ca2f5db --- /dev/null +++ b/scao/cuda/triton_ops.py @@ -0,0 +1,104 @@ +""" +SCAO Triton Kernels +=================== +High-performance JIT-compiled kernels for SCAO operators using OpenAI Triton. +Allows GPU acceleration without requiring a precompiled C++/CUDA extension. +Falls back safely if Triton is not installed or CUDA is unavailable. +""" + +from __future__ import annotations + +from typing import Any + +import torch +from torch import Tensor + +_TRITON_AVAILABLE = False +try: + import triton + import triton.language as tl + _TRITON_AVAILABLE = torch.cuda.is_available() +except ImportError: + triton = None # type: ignore[assignment] + tl = None # type: ignore[assignment] + + +def is_triton_available() -> bool: + """Return True if Triton is available and can execute on CUDA.""" + return _TRITON_AVAILABLE + + +if _TRITON_AVAILABLE: + @triton.jit + def _int8_ema_kernel( + ema_q_ptr: Any, + new_val_ptr: Any, + out_q_ptr: Any, + ema_scale: Any, + rho: Any, + n_elements: Any, + inv_scale: Any, + BLOCK_SIZE: tl.constexpr, + ) -> None: + + pid = tl.program_id(axis=0) + block_start = pid * BLOCK_SIZE + offsets = block_start + tl.arange(0, BLOCK_SIZE) + mask = offsets < n_elements + + # Load int8, convert to float32, dequantize + q = tl.load(ema_q_ptr + offsets, mask=mask, other=0).to(tl.float32) + dequant = q * ema_scale + + # Load new_val and blend + nv = tl.load(new_val_ptr + offsets, mask=mask, other=0.0).to(tl.float32) + updated = dequant * rho + nv + + # Quantize to int8 + q_out = tl.math.clamp(tl.math.round(updated * inv_scale), -127.0, 127.0).to(tl.int8) + tl.store(out_q_ptr + offsets, q_out, mask=mask) + + +def int8_ema_update_triton( + ema_q: Tensor, + ema_scale: float, + new_val: Tensor, + rho: float, +) -> tuple[Tensor, float] | None: + """ + Fused int8 EMA update using Triton when available. + Returns (q_new, new_scale) or None if Triton execution fails. + """ + if not _TRITON_AVAILABLE or not ema_q.is_cuda: + return None + + try: + # Pre-pass for scale: estimate updated magnitude + # updated ≈ rho * (ema_q * scale) + new_val + n_elements = ema_q.numel() + fp32_approx = ema_q.to(torch.float32).mul_(ema_scale).mul_(rho).add_(new_val) + abs_max = fp32_approx.abs().max().item() + new_scale = abs_max / 127.0 if abs_max > 1e-30 else 1.0 + inv_scale = 1.0 / new_scale if new_scale > 1e-30 else 0.0 + + out_q = torch.empty_like(ema_q) + BLOCK_SIZE = 1024 + + def grid(meta: dict[str, int]) -> tuple[int]: + return ((n_elements + meta["BLOCK_SIZE"] - 1) // meta["BLOCK_SIZE"],) + + + _int8_ema_kernel[grid]( + + ema_q, + new_val, + out_q, + ema_scale, + rho, + n_elements, + inv_scale, + BLOCK_SIZE=BLOCK_SIZE, + ) + return out_q, new_scale + except Exception: + return None diff --git a/scao/filters.py b/scao/filters.py new file mode 100644 index 0000000..4ae9c1c --- /dev/null +++ b/scao/filters.py @@ -0,0 +1,189 @@ +""" +SCAO Gradient Filters & Schedulers +=================================== +Adaptive gradient sparsity, element-wise gSNR clipping, +adaptive warmup monitoring, and lazy preconditioner event triggers. +""" + +from __future__ import annotations + +import torch +from torch import Tensor + +# --------------------------------------------------------------------------- +# Gradient filters +# --------------------------------------------------------------------------- + +class _SparseGradFilter: + """EMA-magnitude adaptive sparse mask. Replaces static top-k.""" + + def __init__(self, sparsity: float = 0.7, ema: float = 0.99) -> None: + self.sparsity = sparsity + self.ema = ema + self.mag_ema: Tensor | None = None + + def __call__(self, grad: Tensor) -> Tensor: + mag = grad.abs().detach() + if self.mag_ema is None: + self.mag_ema = mag.clone() + else: + self.mag_ema.mul_(self.ema).add_((1.0 - self.ema) * mag) + threshold = torch.quantile(self.mag_ema.float(), self.sparsity) + return grad * (self.mag_ema >= threshold).to(grad.dtype) + + +class _DynamicSparseFilter(_SparseGradFilter): + """ + R2: Scales sparsity inversely with layer grad norm relative to the global EMA. + Active layers (high ‖g‖) receive lower sparsity, preserving curvature budget + where it matters. Critical for QLoRA where LoRA A/B norms differ by 10-100x + from full attention projections. + """ + + def __init__( + self, + base_sparsity: float = 0.7, + min_sparsity: float = 0.3, + max_sparsity: float = 0.9, + ema: float = 0.99, + norm_ema: float = 0.95, + ) -> None: + super().__init__(sparsity=base_sparsity, ema=ema) + self.base_sparsity = base_sparsity + self.min_sparsity = min_sparsity + self.max_sparsity = max_sparsity + self.norm_ema_coeff = norm_ema + self._norm_ema: float = 1.0 + self._global_norm_ema: float = 1.0 + + def set_global_norm_ref(self, global_norm: float) -> None: + self._global_norm_ema = max(global_norm, 1e-8) + + def __call__(self, grad: Tensor) -> Tensor: + layer_norm = float(grad.norm()) + self._norm_ema = ( + self.norm_ema_coeff * self._norm_ema + + (1.0 - self.norm_ema_coeff) * layer_norm + ) + # margin shrinks as ratio grows: high-activity layers keep more gradients + ratio = self._norm_ema / self._global_norm_ema + margin = 0.2 * (1.0 - min(ratio, 2.0) / 2.0) + self.sparsity = float( + max(self.min_sparsity, min(self.max_sparsity, self.base_sparsity + margin)) + ) + return super().__call__(grad) + + +# --------------------------------------------------------------------------- +# R1: Adaptive warmup scheduler +# --------------------------------------------------------------------------- + +class _AdaptiveWarmupScheduler: + """ + Monitors relative grad norm change as a curvature stability proxy. + Exits warmup early when |‖g‖_t - ‖g‖_{t-1}| / ‖g‖_{t-1} < threshold + for `patience` consecutive steps. + + Prevents small models from wasting preconditioned steps during the + steepest part of the loss curve by staying locked in Adam warmup. + """ + + def __init__( + self, + warmup_steps: int = 100, + stability_threshold: float = 0.05, + patience: int = 5, + min_warmup: int = 20, + ) -> None: + self.warmup_steps = warmup_steps + self.stability_threshold = stability_threshold + self.patience = patience + self.min_warmup = min_warmup + self._prev_norm: float = float("inf") + self._stable_count: int = 0 + self._early_exit_step: int | None = None + + def update(self, step: int, avg_grad_norm: float) -> bool: + """Returns True while still in warmup.""" + if step < self.min_warmup: + return True + + if self._prev_norm > 0: + rel_change = abs(avg_grad_norm - self._prev_norm) / (self._prev_norm + 1e-8) + self._stable_count = self._stable_count + 1 if rel_change < self.stability_threshold else 0 + + self._prev_norm = avg_grad_norm + + if self._stable_count >= self.patience: + self._early_exit_step = self._early_exit_step or step + return False + + return step < self.warmup_steps + + @property + def exited_early(self) -> bool: + return self._early_exit_step is not None + + @property + def actual_warmup_steps(self) -> int: + return self._early_exit_step or self.warmup_steps + + +# --------------------------------------------------------------------------- +# R4: gSNR element-wise clipping +# --------------------------------------------------------------------------- + +def _gsnr_clip( + grad: Tensor, + exp_avg: Tensor, + exp_avg_sq: Tensor, + eps: float, + clip_snr: float, +) -> Tensor: + """ + Masks elements where signal-to-noise ratio |m| / sqrt(v) < clip_snr. + Unlike global norm clipping, this preserves high-signal directions + while suppressing stochastic noise at the element level. Effective + with small per-device batch sizes in multi-GPU setups. + """ + snr = exp_avg.abs() / (exp_avg_sq.sqrt() + eps) + return grad * (snr >= clip_snr).to(grad.dtype) + + +# --------------------------------------------------------------------------- +# R3: Lazy preconditioner trigger +# --------------------------------------------------------------------------- + +class _LazyPrecondTrigger: + """ + Replaces the fixed `step % precond_freq == 0` schedule with an + event-driven policy: update only when ‖g_t - g_{t-1}‖ / ‖g_{t-1}‖ + exceeds delta_threshold, or when max_skip steps have elapsed. + + At 40B+ scale this cuts matrix inversion cost by 60-80% with + negligible impact on convergence quality. + """ + + def __init__(self, delta_threshold: float = 0.1, max_skip: int = 50) -> None: + self.delta_threshold = delta_threshold + self.max_skip = max_skip + self._prev_grad_norm: float = 0.0 + self._steps_since_update: int = 0 + + def should_update(self, grad_norm: float) -> bool: + self._steps_since_update += 1 + + if self._steps_since_update >= self.max_skip: + self._steps_since_update = 0 + self._prev_grad_norm = grad_norm + return True + + if self._prev_grad_norm > 0: + rel_delta = abs(grad_norm - self._prev_grad_norm) / (self._prev_grad_norm + 1e-8) + if rel_delta > self.delta_threshold: + self._steps_since_update = 0 + self._prev_grad_norm = grad_norm + return True + + self._prev_grad_norm = grad_norm + return False diff --git a/scao/optimizer.py b/scao/optimizer.py index 76e599d..528c342 100644 --- a/scao/optimizer.py +++ b/scao/optimizer.py @@ -27,196 +27,28 @@ import math import warnings from collections.abc import Callable, Iterable -from typing import Any, cast import torch from torch import Tensor from torch.optim import Optimizer +from .filters import ( + _AdaptiveWarmupScheduler, + _DynamicSparseFilter, + _gsnr_clip, + _LazyPrecondTrigger, + _SparseGradFilter, +) from .preconditioner import SparsePreconditioner, _broadcast_precond +from .presets import ( + scao_1b, + scao_3b, + scao_7b, + scao_40b, + scao_125b, + scao_sub1b, +) -# --------------------------------------------------------------------------- -# Gradient filters -# --------------------------------------------------------------------------- - -class _SparseGradFilter: - """EMA-magnitude adaptive sparse mask. Replaces static top-k.""" - - def __init__(self, sparsity: float = 0.7, ema: float = 0.99): - self.sparsity = sparsity - self.ema = ema - self.mag_ema: Tensor | None = None - - def __call__(self, grad: Tensor) -> Tensor: - mag = grad.abs().detach() - if self.mag_ema is None: - self.mag_ema = mag.clone() - else: - self.mag_ema.mul_(self.ema).add_((1.0 - self.ema) * mag) - threshold = torch.quantile(self.mag_ema.float(), self.sparsity) - return grad * (self.mag_ema >= threshold).to(grad.dtype) - - -class _DynamicSparseFilter(_SparseGradFilter): - """ - R2: Scales sparsity inversely with layer grad norm relative to the global EMA. - Active layers (high ‖g‖) receive lower sparsity, preserving curvature budget - where it matters. Critical for QLoRA where LoRA A/B norms differ by 10-100x - from full attention projections. - """ - - def __init__( - self, - base_sparsity: float = 0.7, - min_sparsity: float = 0.3, - max_sparsity: float = 0.9, - ema: float = 0.99, - norm_ema: float = 0.95, - ): - super().__init__(sparsity=base_sparsity, ema=ema) - self.base_sparsity = base_sparsity - self.min_sparsity = min_sparsity - self.max_sparsity = max_sparsity - self.norm_ema_coeff = norm_ema - self._norm_ema: float = 1.0 - self._global_norm_ema: float = 1.0 - - def set_global_norm_ref(self, global_norm: float) -> None: - self._global_norm_ema = max(global_norm, 1e-8) - - def __call__(self, grad: Tensor) -> Tensor: - layer_norm = float(grad.norm()) - self._norm_ema = ( - self.norm_ema_coeff * self._norm_ema - + (1.0 - self.norm_ema_coeff) * layer_norm - ) - # margin shrinks as ratio grows: high-activity layers keep more gradients - ratio = self._norm_ema / self._global_norm_ema - margin = 0.2 * (1.0 - min(ratio, 2.0) / 2.0) - self.sparsity = float( - max(self.min_sparsity, min(self.max_sparsity, self.base_sparsity + margin)) - ) - return super().__call__(grad) - - -# --------------------------------------------------------------------------- -# R1: Adaptive warmup scheduler -# --------------------------------------------------------------------------- - -class _AdaptiveWarmupScheduler: - """ - Monitors relative grad norm change as a curvature stability proxy. - Exits warmup early when |‖g‖_t - ‖g‖_{t-1}| / ‖g‖_{t-1} < threshold - for `patience` consecutive steps. - - Prevents small models from wasting preconditioned steps during the - steepest part of the loss curve by staying locked in Adam warmup. - """ - - def __init__( - self, - warmup_steps: int = 100, - stability_threshold: float = 0.05, - patience: int = 5, - min_warmup: int = 20, - ): - self.warmup_steps = warmup_steps - self.stability_threshold = stability_threshold - self.patience = patience - self.min_warmup = min_warmup - self._prev_norm: float = float("inf") - self._stable_count: int = 0 - self._early_exit_step: int | None = None - - def update(self, step: int, avg_grad_norm: float) -> bool: - """Returns True while still in warmup.""" - if step < self.min_warmup: - return True - - if self._prev_norm > 0: - rel_change = abs(avg_grad_norm - self._prev_norm) / (self._prev_norm + 1e-8) - self._stable_count = self._stable_count + 1 if rel_change < self.stability_threshold else 0 - - self._prev_norm = avg_grad_norm - - if self._stable_count >= self.patience: - self._early_exit_step = self._early_exit_step or step - return False - - return step < self.warmup_steps - - @property - def exited_early(self) -> bool: - return self._early_exit_step is not None - - @property - def actual_warmup_steps(self) -> int: - return self._early_exit_step or self.warmup_steps - - -# --------------------------------------------------------------------------- -# R4: gSNR element-wise clipping -# --------------------------------------------------------------------------- - -def _gsnr_clip( - grad: Tensor, - exp_avg: Tensor, - exp_avg_sq: Tensor, - eps: float, - clip_snr: float, -) -> Tensor: - """ - Masks elements where signal-to-noise ratio |m| / sqrt(v) < clip_snr. - Unlike global norm clipping, this preserves high-signal directions - while suppressing stochastic noise at the element level. Effective - with small per-device batch sizes in multi-GPU setups. - """ - snr = exp_avg.abs() / (exp_avg_sq.sqrt() + eps) - return grad * (snr >= clip_snr).to(grad.dtype) - - -# --------------------------------------------------------------------------- -# R3: Lazy preconditioner trigger -# --------------------------------------------------------------------------- - -class _LazyPrecondTrigger: - """ - Replaces the fixed `step % precond_freq == 0` schedule with an - event-driven policy: update only when ‖g_t - g_{t-1}‖ / ‖g_{t-1}‖ - exceeds delta_threshold, or when max_skip steps have elapsed. - - At 40B+ scale this cuts matrix inversion cost by 60-80% with - negligible impact on convergence quality. - """ - - def __init__(self, delta_threshold: float = 0.1, max_skip: int = 50): - self.delta_threshold = delta_threshold - self.max_skip = max_skip - self._prev_grad_norm: float = 0.0 - self._steps_since_update: int = 0 - - def should_update(self, grad_norm: float) -> bool: - self._steps_since_update += 1 - - if self._steps_since_update >= self.max_skip: - self._steps_since_update = 0 - self._prev_grad_norm = grad_norm - return True - - if self._prev_grad_norm > 0: - rel_delta = abs(grad_norm - self._prev_grad_norm) / (self._prev_grad_norm + 1e-8) - if rel_delta > self.delta_threshold: - self._steps_since_update = 0 - self._prev_grad_norm = grad_norm - return True - - self._prev_grad_norm = grad_norm - return False - - -# --------------------------------------------------------------------------- -# Main optimizer -# --------------------------------------------------------------------------- class SCAO(Optimizer): """ @@ -825,185 +657,18 @@ def precond_stats(self) -> dict[str, object]: } -# --------------------------------------------------------------------------- -# Scale presets -# --------------------------------------------------------------------------- - -def _params_from_model_or_iterable(model_or_params: Any) -> Iterable: - return cast(Iterable, ( - model_or_params.parameters() - if hasattr(model_or_params, "parameters") - else model_or_params - )) - - -def _merge_preset_defaults(overrides: dict[str, Any], **defaults: Any) -> dict[str, Any]: - merged = dict(defaults) - merged.update(overrides) - return merged - - -def scao_sub1b(model_or_params: Any, lr: float = 5e-4, **kw: Any) -> SCAO: - """ - <1B params (125M–999M). RTX 3060/T4, CPU offload. - Kronecker inversion is cheap at this scale — maximise update frequency. - Lookahead disabled: these models converge fast enough that slow weights - create net drag. k_min=4 handles narrow layers without falling back to - diagonal; verify SparsePreconditioner has this guard before use. - """ - return SCAO( - _params_from_model_or_iterable(model_or_params), lr=lr, - **_merge_preset_defaults( - kw, - warmup_steps=20, - min_precond_updates=3, - precond_freq=5, - k_min=4, - k_max=64, - max_precond_dim=1024, - epsilon_sparse=0.10, - sparsity=0.4, - dynamic_sparsity=True, - adaptive_warmup=True, - warmup_stability_threshold=0.08, - warmup_patience=3, - lazy_precond=False, - use_gsnr_clip=False, - adaptive_rank=False, - noise_std_init=0.005, - noise_anneal=0.995, - lars_coeff=5e-4, - lookahead_k=0, - beta3=0.98, - tau=0.8, - ), - ) - - -def scao_1b(model_or_params: Any, lr: float = 3e-4, **kw: Any) -> SCAO: - """ - 1B–3B params (Phi-2, Gemma-2B, TinyLlama). RTX 3090/A10, T4×2. - """ - return SCAO( - _params_from_model_or_iterable(model_or_params), lr=lr, - **_merge_preset_defaults( - kw, - warmup_steps=35, - min_precond_updates=5, - precond_freq=8, - k_min=4, - k_max=96, - max_precond_dim=2048, - epsilon_sparse=0.07, - sparsity=0.50, - dynamic_sparsity=True, - adaptive_warmup=True, - warmup_stability_threshold=0.06, - warmup_patience=4, - lazy_precond=False, - use_gsnr_clip=False, - adaptive_rank=False, - noise_std_init=0.008, - noise_anneal=0.997, - lars_coeff=8e-4, - lookahead_k=3, - lookahead_alpha=0.4, - beta3=0.985, - tau=0.9, - ), - ) - - -def scao_3b(model_or_params: Any, lr: float = 2e-4, **kw: Any) -> SCAO: - """3B params. T4/A10 16–24 GB, QLoRA.""" - return SCAO( - _params_from_model_or_iterable(model_or_params), lr=lr, - **_merge_preset_defaults( - kw, - warmup_steps=50, - blend_steps=30, - precond_freq=10, - dynamic_sparsity=True, - adaptive_warmup=True, - warmup_patience=3, - lazy_precond=False, - use_gsnr_clip=False, - adaptive_rank=False, - sparsity=0.6, - ), - ) - - -def scao_7b(model_or_params: Any, lr: float = 1e-4, **kw: Any) -> SCAO: - """7B params. A100 40 GB, QLoRA or full fine-tune.""" - return SCAO( - _params_from_model_or_iterable(model_or_params), lr=lr, - **_merge_preset_defaults( - kw, - warmup_steps=80, - precond_freq=15, - dynamic_sparsity=True, - adaptive_warmup=True, - lazy_precond=False, - use_gsnr_clip=True, - gsnr_threshold=0.4, - adaptive_rank=True, - sparsity=0.65, - ), - ) - - -def scao_40b(model_or_params: Any, lr: float = 5e-5, **kw: Any) -> SCAO: - """14B–70B params. 4–8×A100/H100.""" - return SCAO( - _params_from_model_or_iterable(model_or_params), lr=lr, - **_merge_preset_defaults( - kw, - warmup_steps=100, - blend_steps=50, - precond_freq=20, - k_min=4, - k_max=64, - max_precond_dim=2048, - dynamic_sparsity=True, - adaptive_warmup=True, - lazy_precond=True, - lazy_delta_threshold=0.08, - lazy_max_skip=40, - use_gsnr_clip=True, - gsnr_threshold=0.5, - adaptive_rank=True, - use_int8_ema=True, - sparsity=0.75, - lookahead_k=0, - ), - ) - - -def scao_125b(model_or_params: Any, lr: float = 2e-5, **kw: Any) -> SCAO: - """70B+ params. FSDP / Megatron, 32–64×H100.""" - return SCAO( - _params_from_model_or_iterable(model_or_params), lr=lr, - **_merge_preset_defaults( - kw, - warmup_steps=200, - precond_freq=50, - k_min=4, - k_max=32, - max_precond_dim=1024, - dynamic_sparsity=True, - adaptive_warmup=True, - lazy_precond=True, - lazy_delta_threshold=0.05, - lazy_max_skip=60, - use_gsnr_clip=True, - gsnr_threshold=0.6, - adaptive_rank=True, - use_int8_ema=True, - sparsity=0.80, - lookahead_k=0, - async_precond=True, - ), - ) - +__all__ = [ + "SCAO", + "scao_sub1b", + "scao_1b", + "scao_3b", + "scao_7b", + "scao_40b", + "scao_125b", + "_SparseGradFilter", + "_DynamicSparseFilter", + "_AdaptiveWarmupScheduler", + "_gsnr_clip", + "_LazyPrecondTrigger", +] diff --git a/scao/presets.py b/scao/presets.py new file mode 100644 index 0000000..657d5d7 --- /dev/null +++ b/scao/presets.py @@ -0,0 +1,204 @@ +""" +SCAO Scale Presets +=================== +Pre-tuned hyperparameter profiles tailored for model parameter scales +from sub-1B up to 125B+. +""" + +from __future__ import annotations + +from collections.abc import Iterable +from typing import TYPE_CHECKING, Any, cast + +if TYPE_CHECKING: + from .optimizer import SCAO + + +def _params_from_model_or_iterable(model_or_params: Any) -> Iterable: + return cast(Iterable, ( + model_or_params.parameters() + if hasattr(model_or_params, "parameters") + else model_or_params + )) + + +def _merge_preset_defaults(overrides: dict[str, Any], **defaults: Any) -> dict[str, Any]: + merged = dict(defaults) + merged.update(overrides) + return merged + + +def scao_sub1b(model_or_params: Any, lr: float = 5e-4, **kw: Any) -> SCAO: + """ + <1B params (125M–999M). RTX 3060/T4, CPU offload. + Kronecker inversion is cheap at this scale — maximise update frequency. + Lookahead disabled: these models converge fast enough that slow weights + create net drag. k_min=4 handles narrow layers without falling back to + diagonal; verify SparsePreconditioner has this guard before use. + """ + from .optimizer import SCAO + + return SCAO( + _params_from_model_or_iterable(model_or_params), lr=lr, + **_merge_preset_defaults( + kw, + warmup_steps=20, + min_precond_updates=3, + precond_freq=5, + k_min=4, + k_max=64, + max_precond_dim=1024, + epsilon_sparse=0.10, + sparsity=0.4, + dynamic_sparsity=True, + adaptive_warmup=True, + warmup_stability_threshold=0.08, + warmup_patience=3, + lazy_precond=False, + use_gsnr_clip=False, + adaptive_rank=False, + noise_std_init=0.005, + noise_anneal=0.995, + lars_coeff=5e-4, + lookahead_k=0, + beta3=0.98, + tau=0.8, + ), + ) + + +def scao_1b(model_or_params: Any, lr: float = 3e-4, **kw: Any) -> SCAO: + """ + 1B–3B params (Phi-2, Gemma-2B, TinyLlama). RTX 3090/A10, T4×2. + """ + from .optimizer import SCAO + + return SCAO( + _params_from_model_or_iterable(model_or_params), lr=lr, + **_merge_preset_defaults( + kw, + warmup_steps=35, + min_precond_updates=5, + precond_freq=8, + k_min=4, + k_max=96, + max_precond_dim=2048, + epsilon_sparse=0.07, + sparsity=0.50, + dynamic_sparsity=True, + adaptive_warmup=True, + warmup_stability_threshold=0.06, + warmup_patience=4, + lazy_precond=False, + use_gsnr_clip=False, + adaptive_rank=False, + noise_std_init=0.008, + noise_anneal=0.997, + lars_coeff=8e-4, + lookahead_k=3, + lookahead_alpha=0.4, + beta3=0.985, + tau=0.9, + ), + ) + + +def scao_3b(model_or_params: Any, lr: float = 2e-4, **kw: Any) -> SCAO: + """3B params. T4/A10 16–24 GB, QLoRA.""" + from .optimizer import SCAO + + return SCAO( + _params_from_model_or_iterable(model_or_params), lr=lr, + **_merge_preset_defaults( + kw, + warmup_steps=50, + blend_steps=30, + precond_freq=10, + dynamic_sparsity=True, + adaptive_warmup=True, + warmup_patience=3, + lazy_precond=False, + use_gsnr_clip=False, + adaptive_rank=False, + sparsity=0.6, + ), + ) + + +def scao_7b(model_or_params: Any, lr: float = 1e-4, **kw: Any) -> SCAO: + """7B params. A100 40 GB, QLoRA or full fine-tune.""" + from .optimizer import SCAO + + return SCAO( + _params_from_model_or_iterable(model_or_params), lr=lr, + **_merge_preset_defaults( + kw, + warmup_steps=80, + precond_freq=15, + dynamic_sparsity=True, + adaptive_warmup=True, + lazy_precond=False, + use_gsnr_clip=True, + gsnr_threshold=0.4, + adaptive_rank=True, + sparsity=0.65, + ), + ) + + +def scao_40b(model_or_params: Any, lr: float = 5e-5, **kw: Any) -> SCAO: + """14B–70B params. 4–8×A100/H100.""" + from .optimizer import SCAO + + return SCAO( + _params_from_model_or_iterable(model_or_params), lr=lr, + **_merge_preset_defaults( + kw, + warmup_steps=100, + blend_steps=50, + precond_freq=20, + k_min=4, + k_max=64, + max_precond_dim=2048, + dynamic_sparsity=True, + adaptive_warmup=True, + lazy_precond=True, + lazy_delta_threshold=0.08, + lazy_max_skip=40, + use_gsnr_clip=True, + gsnr_threshold=0.5, + adaptive_rank=True, + use_int8_ema=True, + sparsity=0.75, + lookahead_k=0, + ), + ) + + +def scao_125b(model_or_params: Any, lr: float = 2e-5, **kw: Any) -> SCAO: + """70B+ params. FSDP / Megatron, 32–64×H100.""" + from .optimizer import SCAO + + return SCAO( + _params_from_model_or_iterable(model_or_params), lr=lr, + **_merge_preset_defaults( + kw, + warmup_steps=200, + precond_freq=50, + k_min=4, + k_max=32, + max_precond_dim=1024, + dynamic_sparsity=True, + adaptive_warmup=True, + lazy_precond=True, + lazy_delta_threshold=0.05, + lazy_max_skip=60, + use_gsnr_clip=True, + gsnr_threshold=0.6, + adaptive_rank=True, + use_int8_ema=True, + sparsity=0.80, + lookahead_k=0, + async_precond=True, + ), + ) diff --git a/scripts/bench_125m_350m.py b/scripts/bench_125m_350m.py index e16fdec..2a91b80 100644 --- a/scripts/bench_125m_350m.py +++ b/scripts/bench_125m_350m.py @@ -30,8 +30,6 @@ import argparse import csv -import math -import os import statistics import sys import textwrap @@ -42,7 +40,7 @@ sys.path.insert(0, str(ROOT)) # Reuse the full benchmark infrastructure from gpt_scale_benchmark.py -from scao.benchmarks.gpt_scale_benchmark import ( # noqa: E402 +from scao.benchmarks.gpt_scale_benchmark import ( # noqa: E402 SCALE_CONFIGS, load_data, run_single, @@ -199,7 +197,7 @@ def main() -> None: w = 88 lines: list[str] = [] lines.append("=" * w) - lines.append(f" SCAO Large-Scale Benchmark Results") + lines.append(" SCAO Large-Scale Benchmark Results") lines.append(f" Scales: {scales} Seeds: {seeds} Device: {device}") lines.append("─" * w) lines.append( diff --git a/scripts/bench_ddp.py b/scripts/bench_ddp.py index 59a7032..ceb4e64 100644 --- a/scripts/bench_ddp.py +++ b/scripts/bench_ddp.py @@ -52,7 +52,6 @@ load_data, ) - # --------------------------------------------------------------------------- # Argument parsing # --------------------------------------------------------------------------- diff --git a/scripts/bench_largescale.py b/scripts/bench_largescale.py index 06ad35f..607b4e1 100644 --- a/scripts/bench_largescale.py +++ b/scripts/bench_largescale.py @@ -35,13 +35,11 @@ import csv import gc import math -import os import sys import time import tracemalloc -from dataclasses import dataclass, field +from dataclasses import dataclass from pathlib import Path -from typing import Iterator import torch import torch.nn as nn @@ -712,7 +710,7 @@ def main() -> None: device = torch.device(args.device) print(f"\n{'='*60}") - print(f" SCAO Large-Scale Benchmark") + print(" SCAO Large-Scale Benchmark") print(f" Device: {device}") print(f" Scales: {args.scales}") print(f" Optimizers: {args.optimizers}") diff --git a/scripts/kaggle_kernel_scao.py b/scripts/kaggle_kernel_scao.py index ece0b4f..7d5df9a 100644 --- a/scripts/kaggle_kernel_scao.py +++ b/scripts/kaggle_kernel_scao.py @@ -7,7 +7,11 @@ """ # ── 1. Setup ────────────────────────────────────────────────────────────────── -import subprocess, sys, os, pathlib +import os +import pathlib +import subprocess +import sys + def run(cmd, **kw): print(f"$ {cmd}") @@ -34,6 +38,7 @@ def run(cmd, **kw): # Check GPU import torch + print(f"\nPyTorch : {torch.__version__}") print(f"CUDA : {torch.cuda.is_available()}") if torch.cuda.is_available(): @@ -45,6 +50,7 @@ def run(cmd, **kw): print("WARNING: No GPU detected!") import scao + print(f"SCAO v{scao.__version__} imported OK") # ── 3. Compile CUDA kernels ─────────────────────────────────────────────────── @@ -96,6 +102,8 @@ def run(cmd, **kw): # ── 6. Print summary ────────────────────────────────────────────────────────── import csv + + def load_csv(path): if not os.path.exists(path): return [] with open(path) as f: return list(csv.DictReader(f)) diff --git a/scripts/push_to_kaggle.py b/scripts/push_to_kaggle.py index 56a8326..119dcee 100644 --- a/scripts/push_to_kaggle.py +++ b/scripts/push_to_kaggle.py @@ -11,8 +11,13 @@ C:\\Users\\\\.kaggle\\kaggle.json com {"username":"...","key":"..."} """ -import json, os, sys, zipfile, requests, pathlib, time -from base64 import b64encode +import json +import pathlib +import sys +import time +import zipfile + +import requests # ── Load credentials ────────────────────────────────────────────────────────── CRED_PATH = pathlib.Path.home() / ".kaggle" / "kaggle.json" @@ -192,10 +197,10 @@ def should_exclude(path_str: str) -> bool: if push_r.status_code in (200, 201): info = push_r.json() - print(f"Kernel pushed successfully!") + print("Kernel pushed successfully!") print(f"URL: https://www.kaggle.com/code/{USERNAME}/{KERNEL_SLUG}") print(f"Version: {info.get('versionNumber', '?')}") - print(f"\nThe kernel will start running automatically on a T4 GPU.") + print("\nThe kernel will start running automatically on a T4 GPU.") print(f"Check status at: https://www.kaggle.com/code/{USERNAME}/{KERNEL_SLUG}") else: print(f"ERROR pushing kernel: {push_r.status_code} {push_r.text}") diff --git a/scripts/run_experiment.py b/scripts/run_experiment.py index 40dff9f..6f91bcb 100644 --- a/scripts/run_experiment.py +++ b/scripts/run_experiment.py @@ -20,8 +20,6 @@ import argparse import csv import json -import math -import os import random import sys import time @@ -40,7 +38,6 @@ from scao import SCAO - # --------------------------------------------------------------------------- # Utilities # --------------------------------------------------------------------------- @@ -128,7 +125,7 @@ def main() -> None: out_root = Path(args.out_dir) / Path(args.config).stem print(f"\n{'='*60}") - print(f" SCAO Experiment Runner") + print(" SCAO Experiment Runner") print(f" Config : {args.config}") print(f" Steps : {steps}") print(f" Seeds : {args.seeds}") diff --git a/scripts/scao_colab_benchmark.ipynb b/scripts/scao_colab_benchmark.ipynb index 8e253dc..1e540aa 100644 --- a/scripts/scao_colab_benchmark.ipynb +++ b/scripts/scao_colab_benchmark.ipynb @@ -1,16 +1,4 @@ { - "nbformat": 4, - "nbformat_minor": 5, - "metadata": { - "kernelspec": {"display_name": "Python 3", "language": "python", "name": "python3"}, - "language_info": {"name": "python", "version": "3.10.0"}, - "accelerator": "GPU", - "colab": { - "provenance": [], - "gpuType": "T4", - "name": "SCAO v0.1.1 — GPU Benchmark" - } - }, "cells": [ { "cell_type": "markdown", @@ -38,7 +26,9 @@ "cell_type": "markdown", "id": "s1", "metadata": {}, - "source": ["## 1. GPU Check"] + "source": [ + "## 1. GPU Check" + ] }, { "cell_type": "code", @@ -47,7 +37,10 @@ "metadata": {}, "outputs": [], "source": [ - "import subprocess, torch\n", + "import subprocess\n", + "\n", + "import torch\n", + "\n", "r = subprocess.run(['nvidia-smi','--query-gpu=name,memory.total,driver_version','--format=csv,noheader'],\n", " capture_output=True, text=True)\n", "print('GPU :', r.stdout.strip() if r.returncode==0 else 'NOT FOUND — enable GPU runtime!')\n", @@ -66,7 +59,9 @@ "cell_type": "markdown", "id": "s2", "metadata": {}, - "source": ["## 2. Install Dependencies"] + "source": [ + "## 2. Install Dependencies" + ] }, { "cell_type": "code", @@ -105,8 +100,11 @@ "source": [ "# OPTION A: Upload ZIP\n", "# Create the zip first: cd 'BIG TEECH' && zip -r scao_project.zip . -x '*.pyc' -x '__pycache__/*'\n", + "import glob as _glob\n", + "import os\n", + "import zipfile\n", + "\n", "from google.colab import files\n", - "import zipfile, os, glob as _glob\n", "\n", "uploaded = files.upload()\n", "for fname in uploaded:\n", @@ -141,11 +139,12 @@ "outputs": [], "source": [ "import sys\n", + "\n", "if PROJECT_ROOT not in sys.path:\n", " sys.path.insert(0, PROJECT_ROOT)\n", "\n", "import scao\n", - "from scao import SCAO\n", + "\n", "print(f'SCAO v{scao.__version__} imported OK from {PROJECT_ROOT}')" ] }, @@ -167,7 +166,9 @@ "metadata": {}, "outputs": [], "source": [ - "import os, subprocess\n", + "import os\n", + "import subprocess\n", + "\n", "cuda_dir = os.path.join(PROJECT_ROOT, 'scao', 'cuda')\n", "print(f'Compiling CUDA kernels in: {cuda_dir}')\n", "result = subprocess.run(\n", @@ -198,7 +199,8 @@ "outputs": [], "source": [ "import torch\n", - "from scao.cuda import fused_kronecker_precond, low_rank_precond_mm, int8_ema_update\n", + "\n", + "from scao.cuda import fused_kronecker_precond, int8_ema_update\n", "\n", "device = 'cuda' if torch.cuda.is_available() else 'cpu'\n", "torch.manual_seed(0)\n", @@ -252,6 +254,7 @@ "outputs": [], "source": [ "import os\n", + "\n", "BENCH = os.path.join(PROJECT_ROOT, 'scao', 'benchmarks', 'gpt_scale_benchmark.py')\n", "!python {BENCH} \\\n", " --scale tiny \\\n", @@ -324,7 +327,9 @@ "cell_type": "markdown", "id": "s9", "metadata": {}, - "source": ["## 9. Results Summary"] + "source": [ + "## 9. Results Summary" + ] }, { "cell_type": "code", @@ -333,7 +338,9 @@ "metadata": {}, "outputs": [], "source": [ - "import csv, os\n", + "import csv\n", + "import os\n", + "\n", "\n", "def load_csv(path):\n", " if not os.path.exists(path): return []\n", @@ -365,7 +372,9 @@ "cell_type": "markdown", "id": "s10", "metadata": {}, - "source": ["## 10. Convergence Plot"] + "source": [ + "## 10. Convergence Plot" + ] }, { "cell_type": "code", @@ -374,9 +383,10 @@ "metadata": {}, "outputs": [], "source": [ - "import matplotlib.pyplot as plt\n", "from collections import defaultdict\n", "\n", + "import matplotlib.pyplot as plt\n", + "\n", "COLORS = {'adamw':'#2196F3','scao':'#F44336','scao_int8':'#FF9800','diag_shampoo':'#4CAF50'}\n", "LABELS = {'adamw':'AdamW','scao':'SCAO','scao_int8':'SCAO+int8','diag_shampoo':'Diag-Shampoo'}\n", "\n", @@ -413,7 +423,9 @@ "cell_type": "markdown", "id": "s11", "metadata": {}, - "source": ["## 11. Memory Savings — int8 EMA"] + "source": [ + "## 11. Memory Savings — int8 EMA" + ] }, { "cell_type": "code", @@ -424,6 +436,7 @@ "source": [ "if all_rows:\n", " from collections import defaultdict\n", + "\n", " import numpy as np\n", "\n", " by_opt = defaultdict(lambda: defaultdict(list))\n", @@ -463,7 +476,9 @@ "cell_type": "markdown", "id": "s12", "metadata": {}, - "source": ["## 12. Run Full Test Suite"] + "source": [ + "## 12. Run Full Test Suite" + ] }, { "cell_type": "code", @@ -479,7 +494,9 @@ "cell_type": "markdown", "id": "s13", "metadata": {}, - "source": ["## 13. Download Results"] + "source": [ + "## 13. Download Results" + ] }, { "cell_type": "code", @@ -488,9 +505,10 @@ "metadata": {}, "outputs": [], "source": [ - "from google.colab import files\n", "import os\n", "\n", + "from google.colab import files\n", + "\n", "to_download = [\n", " '/content/results_125m_350m.csv',\n", " '/content/results_125m_350m_curves.csv',\n", @@ -505,5 +523,24 @@ " print(f'Not found (run benchmarks first): {f}')" ] } - ] + ], + "metadata": { + "accelerator": "GPU", + "colab": { + "gpuType": "T4", + "name": "SCAO v0.1.1 — GPU Benchmark", + "provenance": [] + }, + "kernelspec": { + "display_name": "Python 3", + "language": "python", + "name": "python3" + }, + "language_info": { + "name": "python", + "version": "3.10.0" + } + }, + "nbformat": 4, + "nbformat_minor": 5 }