diff --git a/examples/finetune_distilgpt2_colab.py b/examples/finetune_distilgpt2_colab.py new file mode 100644 index 0000000..f94909e --- /dev/null +++ b/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() diff --git a/pyproject.toml b/pyproject.toml index 556443f..979acf4 100644 --- a/pyproject.toml +++ b/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"] diff --git a/scao/__init__.py b/scao/__init__.py index ce72916..8c0538d 100644 --- a/scao/__init__.py +++ b/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" diff --git a/scao/cuda/__init__.py b/scao/cuda/__init__.py index a24ee9e..d9fd271 100644 --- a/scao/cuda/__init__.py +++ b/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( diff --git a/scao/cuda/setup.py b/scao/cuda/setup.py index e245159..a70174c 100644 --- a/scao/cuda/setup.py +++ b/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, }, ) ] diff --git a/scao/integrations/huggingface.py b/scao/integrations/huggingface.py index d644211..fbb4808 100644 --- a/scao/integrations/huggingface.py +++ b/scao/integrations/huggingface.py @@ -56,6 +56,22 @@ def get_scao_optimizer( training_args, scao_kwargs: dict | None = None, no_decay_names: tuple[str, ...] = ("bias", "LayerNorm.weight", "layer_norm.weight"), + no_precond_names: tuple[str, ...] = ( + "embed", + "embedding", + "embed_tokens", + "lm_head", + "wte", + "wpe", + ), + adapter_names: tuple[str, ...] = ( + "adapter", + "lora", + "ia3", + "prefix", + "prompt", + ), + preconditioner_policy: str = "all", num_training_steps: int | None = None, ): """ @@ -70,6 +86,17 @@ def get_scao_optimizer( training_args: ``transformers.TrainingArguments`` instance scao_kwargs: additional SCAO hyperparameters (override defaults) no_decay_names: parameter name substrings exempt from weight decay + no_precond_names: parameter name substrings that should skip SCAO + preconditioning while still using the optimizer moments. This is + useful for large embedding/output matrices in 100B+ models. + adapter_names: parameter name substrings used to identify PEFT/adapter + parameters for ``preconditioner_policy="adapters_only"``. + preconditioner_policy: one of: + ``"all"``: precondition all trainable parameters. This is the + backwards-compatible default. + ``"auto"``: precondition trainable non-embedding parameters. + ``"none"``: disable preconditioning for all parameters. + ``"adapters_only"``: precondition only adapter/LoRA-style params. num_training_steps: explicit scheduler length. Required for an accurate scheduler when ``training_args.max_steps`` is not set. @@ -97,20 +124,64 @@ def get_scao_optimizer( from scao.optimizer import SCAO + valid_policies = {"auto", "all", "none", "adapters_only"} + if preconditioner_policy not in valid_policies: + raise ValueError( + f"Invalid preconditioner_policy={preconditioner_policy!r}. " + f"Expected one of {sorted(valid_policies)}." + ) + + def should_precondition(name: str) -> bool: + if preconditioner_policy == "all": + return True + if preconditioner_policy == "none": + return False + if preconditioner_policy == "adapters_only": + return any(token in name for token in adapter_names) + return not any(token in name for token in no_precond_names) + # Separate parameters with / without weight decay (standard HF split) - decay_params, no_decay_params = [], [] + decay_precond_params, decay_plain_params = [], [] + no_decay_precond_params, no_decay_plain_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) + use_decay = not any(nd in name for nd in no_decay_names) + use_precond = should_precondition(name) + if use_decay and use_precond: + decay_precond_params.append(param) + elif use_decay: + decay_plain_params.append(param) + elif use_precond: + no_decay_precond_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}, - ] + no_decay_plain_params.append(param) + + param_groups = [] + if decay_precond_params: + param_groups.append({ + "params": decay_precond_params, + "weight_decay": training_args.weight_decay, + "preconditioner_enabled": True, + }) + if decay_plain_params: + param_groups.append({ + "params": decay_plain_params, + "weight_decay": training_args.weight_decay, + "preconditioner_enabled": False, + }) + if no_decay_precond_params: + param_groups.append({ + "params": no_decay_precond_params, + "weight_decay": 0.0, + "preconditioner_enabled": True, + }) + if no_decay_plain_params: + param_groups.append({ + "params": no_decay_plain_params, + "weight_decay": 0.0, + "preconditioner_enabled": False, + }) defaults: dict[str, Any] = { "lr": training_args.learning_rate, @@ -175,8 +246,32 @@ class SCAOTrainer(Trainer): trainer.train() """ - def __init__(self, *args, scao_kwargs: dict | None = None, **kwargs): + def __init__( + self, + *args, + scao_kwargs: dict | None = None, + preconditioner_policy: str = "all", + no_precond_names: tuple[str, ...] = ( + "embed", + "embedding", + "embed_tokens", + "lm_head", + "wte", + "wpe", + ), + adapter_names: tuple[str, ...] = ( + "adapter", + "lora", + "ia3", + "prefix", + "prompt", + ), + **kwargs, + ): self._scao_kwargs = scao_kwargs or {} + self._scao_preconditioner_policy = preconditioner_policy + self._scao_no_precond_names = no_precond_names + self._scao_adapter_names = adapter_names super().__init__(*args, **kwargs) def create_optimizer(self): @@ -185,6 +280,9 @@ def create_optimizer(self): self.model, self.args, scao_kwargs=self._scao_kwargs, + no_precond_names=self._scao_no_precond_names, + adapter_names=self._scao_adapter_names, + preconditioner_policy=self._scao_preconditioner_policy, ) self.optimizer = optimizer return self.optimizer diff --git a/scao/optimizer.py b/scao/optimizer.py index a639f0a..76e599d 100644 --- a/scao/optimizer.py +++ b/scao/optimizer.py @@ -253,6 +253,11 @@ class SCAO(Optimizer): adaptive_rank (bool): Adjusts preconditioner k based on per-layer grad norm ratio (R5). Default: False. + preconditioner_enabled (bool): + Whether to allocate and apply the Kronecker/diagonal preconditioner + for this parameter group. Set to False for very large embeddings or + output heads in 100B+ fine-tuning when the curvature overhead is not + worth the memory/latency tradeoff. """ def __init__( @@ -293,6 +298,7 @@ def __init__( use_gsnr_clip: bool = False, gsnr_threshold: float = 0.5, adaptive_rank: bool = False, + preconditioner_enabled: bool = True, ) -> None: if lr < 0: raise ValueError(f"Invalid lr: {lr}") @@ -311,6 +317,7 @@ def __init__( "min_precond_updates": min_precond_updates, "max_precond_dim": max_precond_dim, "use_newton_schulz": use_newton_schulz, "use_int8_ema": use_int8_ema, "beta3": beta3, "lars_coeff": lars_coeff, + "preconditioner_enabled": preconditioner_enabled, } super().__init__(params, defaults) @@ -422,16 +429,17 @@ def _init_state(self, p: Tensor, group: dict) -> None: memory_format=torch.preserve_format) state["prev_grad"] = torch.zeros_like(p, dtype=torch.float32, memory_format=torch.preserve_format) - state["preconditioner"] = SparsePreconditioner( - param=p, - epsilon_sparse=group["epsilon_sparse"], - k_min=group["k_min"], - k_max=group["k_max"], - rho=group["rho"], - max_precond_dim=group["max_precond_dim"], - use_newton_schulz=group["use_newton_schulz"], - use_int8_ema=group["use_int8_ema"], - ) + if group.get("preconditioner_enabled", True): + state["preconditioner"] = SparsePreconditioner( + param=p, + epsilon_sparse=group["epsilon_sparse"], + k_min=group["k_min"], + k_max=group["k_max"], + rho=group["rho"], + max_precond_dim=group["max_precond_dim"], + use_newton_schulz=group["use_newton_schulz"], + use_int8_ema=group["use_int8_ema"], + ) # ----------------------------------------------------------------------- # Grad norm (feeds R1, R2, R5) @@ -562,9 +570,10 @@ def _step_group(self, group: dict, avg_norm: float) -> None: exp_avg_sq = state["exp_avg_sq"] exp_avg_diff = state["exp_avg_diff"] prev_grad = state["prev_grad"] - precond: SparsePreconditioner = state["preconditioner"] + precond: SparsePreconditioner | None = state.get("preconditioner") - self._maybe_adjust_rank(precond, grad_norm, group) + if precond is not None: + self._maybe_adjust_rank(precond, grad_norm, group) # R1: adaptive warmup — use the per-step flag computed once in step(). # Non-adaptive path falls back to the per-param step counter. @@ -572,7 +581,7 @@ def _step_group(self, group: dict, avg_norm: float) -> None: self._in_warmup_global if self._adaptive_warmup else (step <= warmup_steps) - ) or (precond.precond_step < min_pu) + ) or (precond is not None and precond.precond_step < min_pu) if in_warmup: exp_avg.mul_(beta1).add_(grad, alpha=1.0 - beta1) @@ -581,7 +590,7 @@ def _step_group(self, group: dict, avg_norm: float) -> None: bc2 = 1.0 - beta2 ** step denom = (exp_avg_sq.sqrt() / math.sqrt(bc2)).add_(eps) # Accumulate curvature during warmup so Phase 2 starts warm - if step % precond_freq == 0: + if precond is not None and step % precond_freq == 0: self._update_precond_async(precond, p.grad) p.add_((exp_avg / denom).to(p.dtype), alpha=-(lr / bc1)) prev_grad.copy_(grad) @@ -609,8 +618,8 @@ def _step_group(self, group: dict, avg_norm: float) -> None: else: should_update = (step % precond_freq == 0) - g_precond = precond.precondition(p.grad) - if tau is not None: + g_precond = precond.precondition(p.grad) if precond is not None else grad + if tau is not None and precond is not None: g_precond = self._curvature_clip(g_precond, precond, p.grad, tau, eps) # Linear blend from raw gradient to preconditioned over blend_steps, @@ -654,7 +663,7 @@ def _step_group(self, group: dict, avg_norm: float) -> None: # Update curvature after the parameter update has consumed the # current factors. Async work can overlap the remaining optimizer # bookkeeping and is fenced at the start of the next step. - if should_update: + if should_update and precond is not None: self._update_precond_async(precond, p.grad) # ----------------------------------------------------------------------- @@ -953,6 +962,9 @@ def scao_40b(model_or_params: Any, lr: float = 5e-5, **kw: Any) -> SCAO: 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, @@ -963,6 +975,7 @@ def scao_40b(model_or_params: Any, lr: float = 5e-5, **kw: Any) -> SCAO: adaptive_rank=True, use_int8_ema=True, sparsity=0.75, + lookahead_k=0, ), ) @@ -975,7 +988,9 @@ def scao_125b(model_or_params: Any, lr: float = 2e-5, **kw: Any) -> SCAO: kw, warmup_steps=200, precond_freq=50, - max_precond_dim=2048, + k_min=4, + k_max=32, + max_precond_dim=1024, dynamic_sparsity=True, adaptive_warmup=True, lazy_precond=True, @@ -986,7 +1001,7 @@ def scao_125b(model_or_params: Any, lr: float = 2e-5, **kw: Any) -> SCAO: adaptive_rank=True, use_int8_ema=True, sparsity=0.80, - lookahead_k=10, + lookahead_k=0, async_precond=True, ), ) diff --git a/scao/preconditioner.py b/scao/preconditioner.py index 7d5850d..ffc1a88 100644 --- a/scao/preconditioner.py +++ b/scao/preconditioner.py @@ -63,7 +63,6 @@ adaptive_rank, dequantize_sym_int8, from_2d, - low_rank_matrix_power_neg_quarter, quantize_sym_int8, to_2d, ) @@ -233,6 +232,7 @@ def __init__( self.S_l = torch.ones(k_init, device=device, dtype=_PRECOND_DTYPE) self.U_r = torch.eye(n, k_init, device=device, dtype=_PRECOND_DTYPE) self.S_r = torch.ones(k_init, device=device, dtype=_PRECOND_DTYPE) + self._refresh_inverse_spectrum_cache() else: # Diagonal fallback: maintain per-element variance estimate in float32 self.diag_ema = torch.zeros(shape, device=device, dtype=_PRECOND_DTYPE) @@ -242,6 +242,15 @@ def __init__( # Step counter for this preconditioner (updated externally) self.precond_step: int = 0 + def _refresh_inverse_spectrum_cache(self, eps: float = 1e-8) -> None: + """Refresh derived inverse spectrum tensors after eigenfactor changes.""" + if not self.use_kronecker: + return + self.S_l_inv4 = self.S_l.clamp(min=eps).pow(-0.25) + self.S_r_inv4 = self.S_r.clamp(min=eps).pow(-0.25) + self.S_l_inv2 = self.S_l.clamp(min=eps).pow(-0.5) + self.S_r_inv2 = self.S_r.clamp(min=eps).pow(-0.5) + # ------------------------------------------------------------------ # R5: k property — safe external rank assignment # ------------------------------------------------------------------ @@ -285,6 +294,7 @@ def k(self, new_k: int) -> None: self.S_l = self.S_l[:new_k].contiguous() self.U_r = self.U_r[:, :new_k].contiguous() self.S_r = self.S_r[:new_k].contiguous() + self._refresh_inverse_spectrum_cache() elif new_k > self._k: # Rank increase: clamp to the current tensor width. # _update_eigenfactors will expand to new_k on the next update. @@ -453,6 +463,7 @@ def _update_eigenfactors(self, bias_factor: float = 1.0) -> None: self.S_l = S_l_full[:k_new].contiguous() self.U_r = U_r_full[:, :k_new].contiguous() self.S_r = S_r_full[:k_new].contiguous() + self._refresh_inverse_spectrum_cache() # ------------------------------------------------------------------ # Preconditioning (called every optimizer step) @@ -495,7 +506,7 @@ def precondition(self, grad: Tensor) -> Tensor: # Block-diagonal: precondition each slice independently, then reassemble. if self.use_block_diagonal: - result = torch.zeros_like(g2d) + result = torch.empty_like(g2d) offset = 0 for bs, blk in zip(self._block_sizes, self._blocks, strict=True): if self._block_dim == 0: @@ -513,14 +524,10 @@ def precondition(self, grad: Tensor) -> Tensor: result = g2d / denom return from_2d(result, orig_shape).to(orig_dtype) - eps = 1e-8 - _, S_l_inv4 = low_rank_matrix_power_neg_quarter(self.U_l, self.S_l, eps) - _, S_r_inv4 = low_rank_matrix_power_neg_quarter(self.U_r, self.S_r, eps) - - # Use fused CUDA kernel when available (k ≤ 128) — avoids materialising + # Use fused CUDA kernel when available (k <= 32) — avoids materialising # the (m,n) correction tensor and reduces memory bandwidth. # Falls back to pure PyTorch (inside _cuda_fused_precond) otherwise. - G_precond = _cuda_fused_precond(self.U_l, S_l_inv4, self.U_r, S_r_inv4, g2d) + G_precond = _cuda_fused_precond(self.U_l, self.S_l_inv4, self.U_r, self.S_r_inv4, g2d) return from_2d(G_precond, orig_shape).to(orig_dtype) @@ -556,14 +563,12 @@ def natural_grad_norm(self, grad: Tensor, eps: float = 1e-8) -> Tensor: return torch.stack(sq_norms).sum().sqrt() if not self.use_kronecker: - denom = self.diag_ema.reshape_as(g2d).pow(0.5).add_(eps) + bias = getattr(self, "_diag_bias_factor", 1.0) + denom = (self.diag_ema / bias).reshape_as(g2d).pow(0.5).add_(eps) return (g2d.pow(2) / denom).sum().sqrt() - S_l_inv2 = self.S_l.clamp(min=eps).pow(-0.5) - S_r_inv2 = self.S_r.clamp(min=eps).pow(-0.5) - G_proj = (self.U_l.T @ g2d) @ self.U_r # (k, k) - G_scaled = S_l_inv2.unsqueeze(1) * G_proj * S_r_inv2.unsqueeze(0) + G_scaled = self.S_l_inv2.unsqueeze(1) * G_proj * self.S_r_inv2.unsqueeze(0) return cast(Tensor, G_scaled.norm(p="fro")) @@ -618,6 +623,8 @@ def memory_bytes(self) -> int: total += self.R_ema.numel() * self.R_ema.element_size() for t in (self.U_l, self.S_l, self.U_r, self.S_r): total += t.numel() * t.element_size() + for t in (self.S_l_inv4, self.S_r_inv4, self.S_l_inv2, self.S_r_inv2): + total += t.numel() * t.element_size() else: total += self.diag_ema.numel() * self.diag_ema.element_size() return total @@ -644,6 +651,7 @@ def load_state_dict(self, state: dict) -> None: self.S_l = state["S_l"].to(device=self.device, dtype=_PRECOND_DTYPE) self.U_r = state["U_r"].to(device=self.device, dtype=_PRECOND_DTYPE) self.S_r = state["S_r"].to(device=self.device, dtype=_PRECOND_DTYPE) + self._refresh_inverse_spectrum_cache() else: self.diag_ema.copy_(state["diag_ema"]) @@ -717,6 +725,7 @@ def _broadcast_precond( dist.broadcast(precond.S_l, src=0, group=process_group) dist.broadcast(precond.U_r, src=0, group=process_group) dist.broadcast(precond.S_r, src=0, group=process_group) + precond._refresh_inverse_spectrum_cache() else: # Diagonal fallback diff --git a/scao/tests/test_huggingface_integration.py b/scao/tests/test_huggingface_integration.py new file mode 100644 index 0000000..53ca6e6 --- /dev/null +++ b/scao/tests/test_huggingface_integration.py @@ -0,0 +1,122 @@ +from __future__ import annotations + +import inspect +import sys +import types + +import pytest +import torch.nn as nn + +from scao.integrations.huggingface import SCAOTrainer, get_scao_optimizer + + +class _Args: + learning_rate = 1e-3 + warmup_steps = 0 + weight_decay = 0.01 + max_steps = 10 + lr_scheduler_type = "linear" + + +class _TinyPeftLikeModel(nn.Module): + def __init__(self) -> None: + super().__init__() + self.embed_tokens = nn.Embedding(8, 4) + self.base = nn.Linear(4, 4) + self.lora_A = nn.Linear(4, 2, bias=False) + self.lora_B = nn.Linear(2, 4, bias=False) + self.lm_head = nn.Linear(4, 8, bias=False) + + +def _install_transformers_stub(monkeypatch) -> None: + module = types.ModuleType("transformers") + + def get_scheduler(**kwargs): + return {"scheduler_kwargs": kwargs} + + module.get_scheduler = get_scheduler + monkeypatch.setitem(sys.modules, "transformers", module) + + +def _group_param_ids(opt, enabled: bool) -> set[int]: + ids: set[int] = set() + for group in opt.param_groups: + if group["preconditioner_enabled"] is enabled: + ids.update(id(param) for param in group["params"]) + return ids + + +def test_hf_auto_policy_skips_embeddings_and_lm_head(monkeypatch): + _install_transformers_stub(monkeypatch) + model = _TinyPeftLikeModel() + + opt, scheduler = get_scao_optimizer( + model, + _Args(), + preconditioner_policy="auto", + num_training_steps=10, + ) + + enabled = _group_param_ids(opt, True) + disabled = _group_param_ids(opt, False) + + assert id(model.base.weight) in enabled + assert id(model.lora_A.weight) in enabled + assert id(model.embed_tokens.weight) in disabled + assert id(model.lm_head.weight) in disabled + assert scheduler is not None + + +def test_hf_default_policy_preconditions_all_trainable_params(monkeypatch): + _install_transformers_stub(monkeypatch) + model = _TinyPeftLikeModel() + + opt, _ = get_scao_optimizer(model, _Args(), num_training_steps=10) + + enabled = _group_param_ids(opt, True) + assert id(model.base.weight) in enabled + assert id(model.lora_A.weight) in enabled + assert id(model.embed_tokens.weight) in enabled + assert id(model.lm_head.weight) in enabled + + +def test_hf_adapters_only_policy_preconditions_only_lora(monkeypatch): + _install_transformers_stub(monkeypatch) + model = _TinyPeftLikeModel() + + opt, _ = get_scao_optimizer( + model, + _Args(), + preconditioner_policy="adapters_only", + num_training_steps=10, + ) + + enabled = _group_param_ids(opt, True) + disabled = _group_param_ids(opt, False) + + assert id(model.lora_A.weight) in enabled + assert id(model.lora_B.weight) in enabled + assert id(model.base.weight) in disabled + assert id(model.embed_tokens.weight) in disabled + + +def test_hf_invalid_preconditioner_policy_raises(monkeypatch): + _install_transformers_stub(monkeypatch) + model = _TinyPeftLikeModel() + + try: + get_scao_optimizer(model, _Args(), preconditioner_policy="surprise") + except ValueError as exc: + assert "preconditioner_policy" in str(exc) + else: + raise AssertionError("expected invalid preconditioner_policy to raise") + + +def test_scao_trainer_exposes_preconditioner_policy_args(): + if SCAOTrainer is None: + pytest.skip("transformers is not installed") + + signature = inspect.signature(SCAOTrainer.__init__) + assert "preconditioner_policy" in signature.parameters + assert "no_precond_names" in signature.parameters + assert "adapter_names" in signature.parameters diff --git a/scao/tests/test_optimizer.py b/scao/tests/test_optimizer.py index c3e76bd..3af48fc 100644 --- a/scao/tests/test_optimizer.py +++ b/scao/tests/test_optimizer.py @@ -745,3 +745,104 @@ def test_sync_preconditioners_block_diagonal(self): opt.step() sync_preconditioners(opt) # dist not initialized → silent no-op + +# --------------------------------------------------------------------------- +# Large-scale presets +# --------------------------------------------------------------------------- + +class TestLargeScalePresets: + """Large presets should avoid accidental full-parameter extra state.""" + + def test_40b_preset_is_memory_conservative(self): + from scao import scao_40b + + model = nn.Linear(16, 16) + opt = scao_40b(model, lr=1e-4) + + assert opt._lookahead_k == 0 + assert opt.param_groups[0]["k_max"] == 64 + assert opt.param_groups[0]["max_precond_dim"] == 2048 + assert opt.param_groups[0]["use_int8_ema"] is True + + def test_125b_preset_is_memory_conservative(self): + from scao import scao_125b + + model = nn.Linear(16, 16) + opt = scao_125b(model, lr=1e-4) + + assert opt._lookahead_k == 0 + assert opt.param_groups[0]["k_max"] == 32 + assert opt.param_groups[0]["max_precond_dim"] == 1024 + assert opt.param_groups[0]["use_int8_ema"] is True + + def test_param_group_can_disable_preconditioner(self): + p_precond = nn.Parameter(torch.randn(8, 8)) + p_plain = nn.Parameter(torch.randn(8, 8)) + opt = SCAO( + [ + {"params": [p_precond], "preconditioner_enabled": True}, + {"params": [p_plain], "preconditioner_enabled": False}, + ], + lr=1e-3, + warmup_steps=0, + min_precond_updates=0, + sparsity=0.0, + noise_std_init=0.0, + lookahead_k=0, + ) + + loss = p_precond.pow(2).sum() + p_plain.pow(2).sum() + loss.backward() + opt.step() + + assert "preconditioner" in opt.state[p_precond] + assert "preconditioner" not in opt.state[p_plain] + + def test_state_dict_roundtrip_with_disabled_preconditioner_group(self): + p_precond = nn.Parameter(torch.randn(8, 8)) + p_plain = nn.Parameter(torch.randn(8, 8)) + opt = SCAO( + [ + {"params": [p_precond], "preconditioner_enabled": True}, + {"params": [p_plain], "preconditioner_enabled": False}, + ], + lr=1e-3, + warmup_steps=0, + min_precond_updates=0, + sparsity=0.0, + noise_std_init=0.0, + lookahead_k=0, + ) + + for _ in range(3): + loss = p_precond.pow(2).sum() + p_plain.pow(2).sum() + loss.backward() + opt.step() + opt.zero_grad() + + state = copy.deepcopy(opt.state_dict()) + + q_precond = nn.Parameter(torch.randn(8, 8)) + q_plain = nn.Parameter(torch.randn(8, 8)) + opt2 = SCAO( + [ + {"params": [q_precond], "preconditioner_enabled": True}, + {"params": [q_plain], "preconditioner_enabled": False}, + ], + lr=1e-3, + warmup_steps=0, + min_precond_updates=0, + sparsity=0.0, + noise_std_init=0.0, + lookahead_k=0, + ) + + # Initialize optimizer state before loading, matching real resume flows. + (q_precond.pow(2).sum() + q_plain.pow(2).sum()).backward() + opt2.step() + opt2.zero_grad() + opt2.load_state_dict(state) + + assert "preconditioner" in opt2.state[q_precond] + assert "preconditioner" not in opt2.state[q_plain] + diff --git a/scao/tests/test_scale_planner.py b/scao/tests/test_scale_planner.py new file mode 100644 index 0000000..9c23195 --- /dev/null +++ b/scao/tests/test_scale_planner.py @@ -0,0 +1,60 @@ +from __future__ import annotations + +from scao.tools.scale_planner import ( + LayerShape, + ScaoConfig, + _estimate_preconditioner_bytes, + _preconditioner_bytes_for_matrix, +) + + +def test_params_only_planner_estimate_is_conservative_and_warns(): + config = ScaoConfig( + k_max=32, + max_precond_dim=1024, + use_int8_ema=True, + lookahead_k=0, + ) + + estimate, warning = _estimate_preconditioner_bytes( + shapes=[], + config=config, + total_params=140_000_000_000, + precondition_embeddings=False, + ) + + assert estimate == 140_000_000_000 * 4 + assert warning is not None + assert "shape-agnostic" in warning + + +def test_shape_based_planner_can_skip_embeddings(): + config = ScaoConfig( + k_max=32, + max_precond_dim=1024, + use_int8_ema=True, + lookahead_k=0, + ) + shapes = [ + LayerShape("mlp.down", 1024, 4096, 2), + LayerShape("token_embedding", 32000, 1024, 1), + LayerShape("lm_head", 32000, 1024, 1), + ] + + without_embeddings, warning = _estimate_preconditioner_bytes( + shapes=shapes, + config=config, + total_params=sum(shape.params for shape in shapes), + precondition_embeddings=False, + ) + with_embeddings, _ = _estimate_preconditioner_bytes( + shapes=shapes, + config=config, + total_params=sum(shape.params for shape in shapes), + precondition_embeddings=True, + ) + + expected_without = _preconditioner_bytes_for_matrix(1024, 4096, config) * 2 + assert without_embeddings == expected_without + assert with_embeddings > without_embeddings + assert warning is None diff --git a/scao/tools/__init__.py b/scao/tools/__init__.py new file mode 100644 index 0000000..3602cb2 --- /dev/null +++ b/scao/tools/__init__.py @@ -0,0 +1 @@ +"""Command-line tools shipped with SCAO.""" diff --git a/scao/tools/preconditioner_microbench.py b/scao/tools/preconditioner_microbench.py new file mode 100644 index 0000000..ea45411 --- /dev/null +++ b/scao/tools/preconditioner_microbench.py @@ -0,0 +1,130 @@ +""" +Microbenchmark for the SCAO preconditioner hot path. + +Measures repeated precondition() calls after eigenfactors have already been +computed, which is the per-step overhead users pay during training. +""" + +from __future__ import annotations + +import argparse +import statistics +import time +import warnings + +import torch + +from scao.preconditioner import SparsePreconditioner + + +def _synchronize(device: torch.device) -> None: + if device.type == "cuda": + torch.cuda.synchronize(device) + + +def _build_preconditioner( + dim: int, + rank: int, + device: torch.device, +) -> tuple[SparsePreconditioner, torch.Tensor]: + param = torch.empty(dim, dim, device=device) + grad = torch.randn_like(param) + precond = SparsePreconditioner( + param, + k_min=min(rank, dim), + k_max=min(rank, dim), + max_precond_dim=dim, + rho=0.99, + use_int8_ema=False, + ) + precond.update_curvature(grad) + return precond, grad + + +def _bench_precondition( + precond: SparsePreconditioner, + grad: torch.Tensor, + warmup: int, + iters: int, +) -> list[float]: + device = grad.device + fallback_warned = False + with warnings.catch_warnings(record=True) as caught: + warnings.simplefilter("always", RuntimeWarning) + for _ in range(warmup): + out = precond.precondition(grad) + if not torch.isfinite(out).all(): + raise RuntimeError("precondition() produced non-finite values during warmup") + _synchronize(device) + fallback_warned = any("compiled CUDA extension" in str(item.message) for item in caught) + + times_ms: list[float] = [] + with warnings.catch_warnings(record=True) as caught: + warnings.simplefilter("always", RuntimeWarning) + for _ in range(iters): + start = time.perf_counter() + out = precond.precondition(grad) + _synchronize(device) + elapsed_ms = (time.perf_counter() - start) * 1000.0 + if not torch.isfinite(out).all(): + raise RuntimeError("precondition() produced non-finite values") + times_ms.append(elapsed_ms) + fallback_warned = fallback_warned or any( + "compiled CUDA extension" in str(item.message) for item in caught + ) + print("using_pytorch_fallback:", fallback_warned) + return times_ms + + +def main() -> None: + parser = argparse.ArgumentParser(description="SCAO preconditioner microbenchmark") + parser.add_argument("--device", default="cuda" if torch.cuda.is_available() else "cpu") + parser.add_argument("--dim", type=int, default=512) + parser.add_argument("--rank", type=int, default=32) + parser.add_argument("--warmup", type=int, default=10) + parser.add_argument("--iters", type=int, default=50) + args = parser.parse_args() + + if args.dim < 2: + raise SystemExit(f"--dim must be at least 2, got {args.dim}.") + if args.rank < 1: + raise SystemExit(f"--rank must be positive, got {args.rank}.") + if args.rank > args.dim: + raise SystemExit(f"--rank must be <= --dim, got rank={args.rank}, dim={args.dim}.") + if args.warmup < 0: + raise SystemExit(f"--warmup must be non-negative, got {args.warmup}.") + if args.iters < 1: + raise SystemExit(f"--iters must be positive, got {args.iters}.") + + device = torch.device(args.device) + if device.type == "cuda" and not torch.cuda.is_available(): + raise RuntimeError("CUDA requested but torch.cuda.is_available() is False") + + torch.manual_seed(42) + if device.type == "cuda": + torch.cuda.manual_seed_all(42) + torch.cuda.reset_peak_memory_stats(device) + + print("== SCAO preconditioner microbenchmark ==") + print("torch:", torch.__version__) + print("device:", device) + if device.type == "cuda": + print("gpu:", torch.cuda.get_device_name(device)) + print("dim:", args.dim) + print("rank:", args.rank) + + precond, grad = _build_preconditioner(args.dim, args.rank, device) + times_ms = _bench_precondition(precond, grad, args.warmup, args.iters) + + print("iters:", args.iters) + print("mean_ms:", f"{statistics.mean(times_ms):.4f}") + print("median_ms:", f"{statistics.median(times_ms):.4f}") + print("min_ms:", f"{min(times_ms):.4f}") + print("max_ms:", f"{max(times_ms):.4f}") + if device.type == "cuda": + peak_mb = torch.cuda.max_memory_allocated(device) / 1024**2 + print("peak_mb:", f"{peak_mb:.1f}") + + +if __name__ == "__main__": + main() diff --git a/scao/tools/scale_planner.py b/scao/tools/scale_planner.py new file mode 100644 index 0000000..7b6fac9 --- /dev/null +++ b/scao/tools/scale_planner.py @@ -0,0 +1,275 @@ +""" +Memory planner for large-scale SCAO fine-tuning. + +The planner estimates optimizer/preconditioner state for transformer-like +models before a run. It is meant for capacity planning on 40B-140B+ models, +where one accidental extra parameter copy can cost terabytes across a cluster. +""" + +from __future__ import annotations + +import argparse +import math +from dataclasses import dataclass + + +@dataclass(frozen=True) +class LayerShape: + name: str + rows: int + cols: int + count: int = 1 + + @property + def params(self) -> int: + return self.rows * self.cols * self.count + + +@dataclass(frozen=True) +class ScaoConfig: + k_max: int + max_precond_dim: int + use_int8_ema: bool + lookahead_k: int + + +def _gb(num_bytes: float) -> float: + return num_bytes / 1024**3 + + +def _format_gb(num_bytes: float) -> str: + return f"{_gb(num_bytes):,.2f} GB" + + +def _transformer_shapes( + hidden_size: int, + ffn_size: int, + layers: int, + vocab_size: int, + tie_embeddings: bool, +) -> list[LayerShape]: + shapes = [ + LayerShape("attn.qkv", 3 * hidden_size, hidden_size, layers), + LayerShape("attn.out", hidden_size, hidden_size, layers), + LayerShape("mlp.gate_up", 2 * ffn_size, hidden_size, layers), + LayerShape("mlp.down", hidden_size, ffn_size, layers), + LayerShape("token_embedding", vocab_size, hidden_size, 1), + ] + if not tie_embeddings: + shapes.append(LayerShape("lm_head", vocab_size, hidden_size, 1)) + return shapes + + +def _infer_params_from_shapes(shapes: list[LayerShape]) -> int: + dense_params = sum(shape.params for shape in shapes) + return int(dense_params * 1.002) + + +def _preconditioner_bytes_for_matrix(rows: int, cols: int, config: ScaoConfig) -> int: + if rows <= 1 or cols <= 1: + return rows * cols * 4 + + if max(rows, cols) > config.max_precond_dim: + total = 0 + if rows >= cols: + for offset in range(0, rows, config.max_precond_dim): + block_rows = min(config.max_precond_dim, rows - offset) + total += _preconditioner_bytes_for_matrix(block_rows, cols, config) + else: + for offset in range(0, cols, config.max_precond_dim): + block_cols = min(config.max_precond_dim, cols - offset) + total += _preconditioner_bytes_for_matrix(rows, block_cols, config) + return total + + k = min(config.k_max, rows, cols) + ema_bytes_per_elem = 1 if config.use_int8_ema else 4 + ema = (rows * rows + cols * cols) * ema_bytes_per_elem + eigen = (rows * k + cols * k + 2 * k) * 4 + inverse_cache = 4 * k * 4 + scale_overhead = 16 if config.use_int8_ema else 0 + return int(ema + eigen + inverse_cache + scale_overhead) + + +def _preconditioner_bytes_for_shape(shape: LayerShape, config: ScaoConfig) -> int: + return _preconditioner_bytes_for_matrix(shape.rows, shape.cols, config) * shape.count + + +def _estimate_preconditioner_bytes( + shapes: list[LayerShape], + config: ScaoConfig, + total_params: int, + precondition_embeddings: bool, +) -> tuple[int, str | None]: + if shapes: + total = 0 + for shape in shapes: + is_embedding_like = "embedding" in shape.name or shape.name == "lm_head" + if is_embedding_like and not precondition_embeddings: + continue + total += _preconditioner_bytes_for_shape(shape, config) + return total, None + + bytes_per_param = 4.0 if config.use_int8_ema else 8.0 + warning = ( + "params-only estimate is conservative and shape-agnostic; pass " + "--hidden-size --ffn-size --layers for capacity planning." + ) + return int(total_params * bytes_per_param), warning + + +def _recommended_config(params_b: float) -> ScaoConfig: + if params_b >= 100: + return ScaoConfig(k_max=32, max_precond_dim=1024, use_int8_ema=True, lookahead_k=0) + if params_b >= 40: + return ScaoConfig(k_max=64, max_precond_dim=2048, use_int8_ema=True, lookahead_k=0) + return ScaoConfig(k_max=96, max_precond_dim=2048, use_int8_ema=True, lookahead_k=0) + + +def _require_positive(name: str, value: int | float | None) -> None: + if value is not None and value <= 0: + raise SystemExit(f"{name} must be positive, got {value}.") + + +def _validate_args(args: argparse.Namespace) -> None: + _require_positive("--params-b", args.params_b) + _require_positive("--hidden-size", args.hidden_size) + _require_positive("--ffn-size", args.ffn_size) + _require_positive("--layers", args.layers) + _require_positive("--vocab-size", args.vocab_size) + _require_positive("--world-size", args.world_size) + _require_positive("--gpu-memory-gb", args.gpu_memory_gb) + _require_positive("--k-max", args.k_max) + _require_positive("--max-precond-dim", args.max_precond_dim) + if args.lookahead_k is not None and args.lookahead_k < 0: + raise SystemExit(f"--lookahead-k must be non-negative, got {args.lookahead_k}.") + + +def main() -> None: + parser = argparse.ArgumentParser(description="Plan SCAO memory for large fine-tuning") + parser.add_argument("--params-b", type=float, default=None) + parser.add_argument("--hidden-size", type=int, default=None) + parser.add_argument("--ffn-size", type=int, default=None) + parser.add_argument("--layers", type=int, default=None) + parser.add_argument("--vocab-size", type=int, default=32000) + parser.add_argument("--untied-embeddings", action="store_true") + parser.add_argument("--world-size", type=int, default=64) + parser.add_argument("--gpu-memory-gb", type=float, default=80.0) + parser.add_argument("--zero-stage", type=int, choices=(0, 1, 2, 3), default=3) + parser.add_argument( + "--precondition-embeddings", + action="store_true", + help="Include token embeddings/lm_head in SCAO preconditioner estimates.", + ) + parser.add_argument("--k-max", type=int, default=None) + parser.add_argument("--max-precond-dim", type=int, default=None) + parser.add_argument("--fp32-ema", action="store_true") + parser.add_argument("--lookahead-k", type=int, default=None) + args = parser.parse_args() + _validate_args(args) + + shape_args = (args.hidden_size, args.ffn_size, args.layers) + has_shapes = all(value is not None for value in shape_args) + if args.params_b is None and not has_shapes: + raise SystemExit("Provide either --params-b or --hidden-size/--ffn-size/--layers.") + if args.params_b is not None and has_shapes: + print("warning: both --params-b and explicit shapes were provided; using explicit shapes.") + + shapes: list[LayerShape] = [] + if has_shapes: + assert args.hidden_size is not None + assert args.ffn_size is not None + assert args.layers is not None + shapes = _transformer_shapes( + args.hidden_size, + args.ffn_size, + args.layers, + args.vocab_size, + tie_embeddings=not args.untied_embeddings, + ) + total_params = _infer_params_from_shapes(shapes) + else: + assert args.params_b is not None + total_params = int(args.params_b * 1_000_000_000) + + params_b = total_params / 1_000_000_000 + recommended = _recommended_config(params_b) + config = ScaoConfig( + k_max=args.k_max if args.k_max is not None else recommended.k_max, + max_precond_dim=( + args.max_precond_dim + if args.max_precond_dim is not None + else recommended.max_precond_dim + ), + use_int8_ema=not args.fp32_ema, + lookahead_k=args.lookahead_k if args.lookahead_k is not None else recommended.lookahead_k, + ) + + param_bytes = total_params * 2 + grad_bytes = total_params * 2 + adam_moments = total_params * 8 + master_weights = total_params * 4 + lookahead_bytes = total_params * 2 if config.lookahead_k > 0 else 0 + precond_bytes, estimate_warning = _estimate_preconditioner_bytes( + shapes, + config, + total_params, + precondition_embeddings=args.precondition_embeddings, + ) + + replicated_total = param_bytes + grad_bytes + adam_moments + master_weights + replicated_total += lookahead_bytes + precond_bytes + + shard_optimizer = args.zero_stage >= 1 + shard_grads = args.zero_stage >= 2 + shard_params = args.zero_stage >= 3 + + per_gpu = 0.0 + per_gpu += param_bytes / args.world_size if shard_params else param_bytes + per_gpu += grad_bytes / args.world_size if shard_grads else grad_bytes + optimizer_state = adam_moments + master_weights + precond_bytes + lookahead_bytes + per_gpu += optimizer_state / args.world_size if shard_optimizer else optimizer_state + + gpu_capacity = args.gpu_memory_gb * 1024**3 + headroom = gpu_capacity - per_gpu + + print("== SCAO large-scale planner ==") + print("estimated_params:", f"{params_b:,.2f}B") + print("world_size:", args.world_size) + print("zero_stage:", args.zero_stage) + print("gpu_memory:", f"{args.gpu_memory_gb:.1f} GB") + print("") + print("== Recommended SCAO knobs ==") + print("k_max:", config.k_max) + print("max_precond_dim:", config.max_precond_dim) + print("use_int8_ema:", config.use_int8_ema) + print("lookahead_k:", config.lookahead_k) + print("precondition_embeddings:", args.precondition_embeddings) + print("") + print("== Cluster-level state estimate ==") + print("params_bf16:", _format_gb(param_bytes)) + print("grads_bf16:", _format_gb(grad_bytes)) + print("adam_moments_fp32:", _format_gb(adam_moments)) + print("master_weights_fp32:", _format_gb(master_weights)) + print("scao_preconditioner:", _format_gb(precond_bytes)) + if estimate_warning: + print("estimate_warning:", estimate_warning) + print("lookahead_extra:", _format_gb(lookahead_bytes)) + print("total_no_sharding:", _format_gb(replicated_total)) + print("") + print("== Per-GPU estimate ==") + print("model_grad_optimizer:", _format_gb(per_gpu)) + print("activation_budget_remaining:", _format_gb(headroom)) + print("fits_before_activations:", headroom > 0) + if headroom <= 0: + min_world = math.ceil(per_gpu / gpu_capacity * args.world_size) + print("suggested_min_world_size_before_activations:", min_world) + elif headroom / gpu_capacity < 0.25: + print("warning: less than 25% memory remains for activations, dataloader, and fragmentation") + + if config.lookahead_k > 0 and params_b >= 40: + print("") + print("warning: lookahead creates an extra parameter copy; disable it for 40B+ runs") + + +if __name__ == "__main__": + main() diff --git a/scripts/colab_real_user_test.py b/scripts/colab_real_user_test.py new file mode 100644 index 0000000..640c788 --- /dev/null +++ b/scripts/colab_real_user_test.py @@ -0,0 +1,329 @@ +""" +Real-user Colab test for SCAO on a T4 GPU. + +This script runs a small Hugging Face fine-tuning job using SCAO exactly like +an end user would: model + dataset + Trainer + generation after training. + +Colab usage from the repository root: + + !pip install -U "transformers>=4.30" "datasets>=2.0" accelerate -q + !python scripts/colab_real_user_test.py + +Optional: + + !python scripts/colab_real_user_test.py --quick + !python scripts/colab_real_user_test.py --max-steps 200 --batch-size 2 + !python scripts/colab_real_user_test.py --compare-adamw +""" + +from __future__ import annotations + +import argparse +import inspect +import math +import os +import sys +import tempfile +import time +from pathlib import Path +from typing import Any + +import torch + + +def _require_cuda() -> torch.device: + print("== Environment ==") + print("python:", sys.version.split()[0]) + print("torch:", torch.__version__) + print("cuda_available:", torch.cuda.is_available()) + print("torch_cuda_version:", torch.version.cuda) + + if not torch.cuda.is_available(): + raise RuntimeError( + "CUDA is not available. In Colab, use Runtime > Change runtime type > GPU." + ) + + device = torch.device("cuda") + print("gpu:", torch.cuda.get_device_name(0)) + print("capability:", torch.cuda.get_device_capability(0)) + return device + + +def _load_dataset(tokenizer: Any, block_size: int, dataset_rows: int) -> Any: + try: + from datasets import load_dataset + except ImportError as exc: + raise ImportError( + "datasets is required. Install it with: " + 'pip install -U "datasets>=2.0"' + ) from exc + + dataset = load_dataset("wikitext", "wikitext-2-raw-v1", split="train") + dataset = dataset.filter(lambda row: bool(row["text"].strip())) + if dataset_rows > 0: + dataset = dataset.select(range(min(dataset_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, + desc="Tokenizing WikiText-2", + ) + return tokenized.train_test_split(test_size=0.05, seed=42) + + +def _build_model_and_tokenizer(model_name: str) -> tuple[Any, Any]: + try: + from transformers import AutoModelForCausalLM, AutoTokenizer + except ImportError as exc: + raise ImportError( + "transformers is required. Install it with: " + 'pip install -U "transformers>=4.30" accelerate' + ) from exc + + tokenizer = AutoTokenizer.from_pretrained(model_name) + if tokenizer.pad_token is None: + tokenizer.pad_token = tokenizer.eos_token + + model = AutoModelForCausalLM.from_pretrained(model_name) + model.config.pad_token_id = tokenizer.pad_token_id + return model, tokenizer + + +def _make_training_args(args: argparse.Namespace, output_dir: str) -> Any: + from transformers import TrainingArguments + + signature = inspect.signature(TrainingArguments) + kwargs: dict[str, Any] = { + "output_dir": 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": True, + "logging_steps": max(1, args.log_every), + "eval_steps": max(1, args.max_steps // 2), + "save_strategy": "no", + "report_to": "none", + "remove_unused_columns": False, + "dataloader_num_workers": 2, + "seed": args.seed, + } + + if "eval_strategy" in signature.parameters: + kwargs["eval_strategy"] = "steps" + elif "evaluation_strategy" in signature.parameters: + kwargs["evaluation_strategy"] = "steps" + + kwargs = { + name: value + for name, value in kwargs.items() + if name in signature.parameters + } + return TrainingArguments(**kwargs) + + +def _train_with_scao( + model: Any, + tokenizer: Any, + data: Any, + args: argparse.Namespace, + output_dir: str, +) -> dict[str, float]: + from transformers import DataCollatorForLanguageModeling, Trainer + + import scao + from scao.integrations.huggingface import get_scao_optimizer + + print("\n== SCAO package ==") + print("version:", getattr(scao, "__version__", "unknown")) + print("path:", getattr(scao, "__file__", "unknown")) + + training_args = _make_training_args(args, output_dir) + optimizer, scheduler = get_scao_optimizer( + model, + training_args, + scao_kwargs={ + "precond_freq": args.precond_freq, + "min_precond_updates": 2, + "max_precond_dim": args.max_precond_dim, + "k_min": args.k_min, + "k_max": args.k_max, + "noise_std_init": 0.0, + "sparsity": 0.0, + "lars_coeff": 0.0, + "lookahead_k": 0, + "async_precond": True, + }, + num_training_steps=args.max_steps, + ) + + trainer = Trainer( + model=model, + args=training_args, + train_dataset=data["train"], + eval_dataset=data["test"], + data_collator=DataCollatorForLanguageModeling(tokenizer=tokenizer, mlm=False), + optimizers=(optimizer, scheduler), + ) + + torch.cuda.reset_peak_memory_stats() + start = time.time() + train_result = trainer.train() + if hasattr(optimizer, "synchronize_precond"): + optimizer.synchronize_precond() + torch.cuda.synchronize() + + eval_metrics = trainer.evaluate() + train_loss = float(train_result.training_loss) + eval_loss = float(eval_metrics["eval_loss"]) + peak_mb = torch.cuda.max_memory_allocated() / 1024**2 + elapsed_s = time.time() - start + + if not math.isfinite(train_loss) or not math.isfinite(eval_loss): + raise RuntimeError(f"non-finite loss: train={train_loss}, eval={eval_loss}") + + return { + "train_loss": train_loss, + "eval_loss": eval_loss, + "peak_mb": peak_mb, + "elapsed_s": elapsed_s, + } + + +def _train_with_adamw( + model: Any, + tokenizer: Any, + data: Any, + args: argparse.Namespace, + output_dir: str, +) -> dict[str, float]: + from transformers import DataCollatorForLanguageModeling, Trainer + + training_args = _make_training_args(args, output_dir) + trainer = Trainer( + model=model, + args=training_args, + train_dataset=data["train"], + eval_dataset=data["test"], + data_collator=DataCollatorForLanguageModeling(tokenizer=tokenizer, mlm=False), + ) + + torch.cuda.reset_peak_memory_stats() + start = time.time() + train_result = trainer.train() + torch.cuda.synchronize() + eval_metrics = trainer.evaluate() + + return { + "train_loss": float(train_result.training_loss), + "eval_loss": float(eval_metrics["eval_loss"]), + "peak_mb": torch.cuda.max_memory_allocated() / 1024**2, + "elapsed_s": time.time() - start, + } + + +def _generate_sample(model: Any, tokenizer: Any, prompt: str) -> None: + print("\n== Generation sample ==") + model.eval() + inputs = tokenizer(prompt, return_tensors="pt").to(model.device) + with torch.no_grad(): + output = model.generate( + **inputs, + max_new_tokens=40, + 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)) + + +def _print_metrics(name: str, metrics: dict[str, float]) -> None: + print(f"\n== {name} result ==") + print("train_loss:", f"{metrics['train_loss']:.4f}") + print("eval_loss:", f"{metrics['eval_loss']:.4f}") + print("time_s:", f"{metrics['elapsed_s']:.1f}") + print("peak_mb:", f"{metrics['peak_mb']:.1f}") + + +def main() -> None: + parser = argparse.ArgumentParser(description="Real-user SCAO Colab T4 test") + parser.add_argument("--model", default="distilgpt2") + parser.add_argument("--max-steps", type=int, default=100) + parser.add_argument("--quick", action="store_true") + parser.add_argument("--batch-size", type=int, default=2) + parser.add_argument("--grad-accum", type=int, default=4) + parser.add_argument("--block-size", type=int, default=128) + parser.add_argument("--dataset-rows", type=int, default=1200) + 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("--precond-freq", type=int, default=10) + parser.add_argument("--max-precond-dim", type=int, default=1024) + parser.add_argument("--k-min", type=int, default=4) + parser.add_argument("--k-max", type=int, default=32) + parser.add_argument("--log-every", type=int, default=10) + parser.add_argument("--seed", type=int, default=42) + parser.add_argument("--compare-adamw", action="store_true") + parser.add_argument( + "--prompt", + default="Sparse curvature-aware optimization helps language models", + ) + args = parser.parse_args() + + if args.quick: + args.max_steps = min(args.max_steps, 20) + args.dataset_rows = min(args.dataset_rows, 300) + args.log_every = min(args.log_every, 5) + + os.environ.setdefault("TOKENIZERS_PARALLELISM", "false") + _require_cuda() + torch.manual_seed(args.seed) + + print("\n== Loading model and dataset ==") + model, tokenizer = _build_model_and_tokenizer(args.model) + model.to("cuda") + data = _load_dataset(tokenizer, args.block_size, args.dataset_rows) + + with tempfile.TemporaryDirectory(prefix="scao-real-user-") as tmp: + scao_metrics = _train_with_scao( + model, + tokenizer, + data, + args, + str(Path(tmp) / "scao"), + ) + _print_metrics("SCAO", scao_metrics) + _generate_sample(model, tokenizer, args.prompt) + + if args.compare_adamw: + print("\n== AdamW comparison ==") + adamw_model, adamw_tokenizer = _build_model_and_tokenizer(args.model) + adamw_model.to("cuda") + adamw_metrics = _train_with_adamw( + adamw_model, + adamw_tokenizer, + data, + args, + str(Path(tmp) / "adamw"), + ) + _print_metrics("AdamW", adamw_metrics) + + print("\nSCAO real-user Colab test: OK") + + +if __name__ == "__main__": + main() diff --git a/scripts/gpu_compat_check.py b/scripts/gpu_compat_check.py new file mode 100644 index 0000000..695bd13 --- /dev/null +++ b/scripts/gpu_compat_check.py @@ -0,0 +1,179 @@ +""" +SCAO GPU compatibility check. + +Run this first on a new machine before fine-tuning: + + python scripts/gpu_compat_check.py + +Optional: + + python scripts/gpu_compat_check.py --compile-cuda-ext + python scripts/gpu_compat_check.py --strict-cuda-ext + +The script reports the detected device, whether the optional CUDA extension is +loaded, whether SCAO can run through the PyTorch fallback, and a conservative +recommended config for the current GPU class. +""" + +from __future__ import annotations + +import argparse +import importlib +import subprocess +import sys +import warnings +from pathlib import Path + +import torch + +ROOT = Path(__file__).resolve().parents[1] + + +def _run(cmd: list[str], cwd: Path | None = None) -> bool: + print(f"\n$ {' '.join(cmd)}") + completed = subprocess.run(cmd, cwd=str(cwd or ROOT), check=False) + return completed.returncode == 0 + + +def _compile_cuda_extension() -> bool: + cuda_dir = ROOT / "scao" / "cuda" + if not cuda_dir.exists(): + print("cuda_extension_build: skipped, scao/cuda directory not found") + return False + return _run([sys.executable, "setup.py", "build_ext", "--inplace"], cwd=cuda_dir) + + +def _device_profile() -> tuple[str, str]: + if not torch.cuda.is_available(): + return "cpu", "CPU fallback" + + capability = torch.cuda.get_device_capability(0) + arch = f"sm_{capability[0]}{capability[1]}" + name = torch.cuda.get_device_name(0) + return arch, name + + +def _recommended_config(arch: str) -> dict[str, int | bool]: + base: dict[str, int | bool] = { + "precond_freq": 20, + "max_precond_dim": 1024, + "k_min": 4, + "k_max": 32, + "async_precond": True, + } + if arch in {"sm_80", "sm_89", "sm_90"}: + base["max_precond_dim"] = 2048 + if arch == "cpu": + base["precond_freq"] = 50 + base["max_precond_dim"] = 512 + base["async_precond"] = False + return base + + +def _support_status(arch: str) -> str: + statuses = { + "cpu": "supported via PyTorch fallback", + "sm_75": "validated target: T4", + "sm_80": "production target: A100", + "sm_86": "production target: A10/A10G/RTX 3090", + "sm_89": "production target: RTX 4090/L4/L40S", + "sm_90": "production target: H100", + } + return statuses.get(arch, "not yet validated; fallback may work if PyTorch supports it") + + +def _check_cuda_extension_loaded() -> bool: + try: + importlib.import_module("scao.cuda._scao_cuda") + except ImportError: + return False + return True + + +def _run_functional_check(strict_cuda_ext: bool) -> bool: + from scao import SCAO + from scao.cuda import fused_kronecker_precond + + device = torch.device("cuda" if torch.cuda.is_available() else "cpu") + ext_loaded = _check_cuda_extension_loaded() + if strict_cuda_ext and not ext_loaded: + print("functional_check: failed, CUDA extension was required but is not loaded") + return False + + with warnings.catch_warnings(record=True) as caught: + warnings.simplefilter("always", RuntimeWarning) + U = torch.eye(16, 4, device=device) + s = torch.ones(4, device=device) + G = torch.randn(16, 16, device=device) + out = fused_kronecker_precond(U, s, U, s, G) + + fallback = any("compiled CUDA extension" in str(item.message) for item in caught) + identity_ok = torch.allclose(out, G, atol=1e-6) + finite_ok = torch.isfinite(out).all().item() + + model = torch.nn.Linear(32, 16).to(device) + optimizer = SCAO( + model.parameters(), + lr=1e-3, + warmup_steps=2, + precond_freq=2, + k_max=8, + noise_std_init=0.0, + sparsity=0.0, + lars_coeff=0.0, + lookahead_k=0, + ) + for _ in range(3): + x = torch.randn(8, 32, device=device) + loss = model(x).pow(2).mean() + optimizer.zero_grad() + loss.backward() + optimizer.step() + if hasattr(optimizer, "synchronize_precond"): + optimizer.synchronize_precond() + + params_finite = all(torch.isfinite(param).all().item() for param in model.parameters()) + print("fused_precond_identity:", identity_ok) + print("fused_precond_finite:", finite_ok) + print("optimizer_step_finite:", params_finite) + print("cuda_extension_loaded:", ext_loaded) + print("using_pytorch_fallback:", fallback or not ext_loaded) + return bool(identity_ok and finite_ok and params_finite) + + +def main() -> None: + parser = argparse.ArgumentParser(description="Check SCAO GPU compatibility") + parser.add_argument("--compile-cuda-ext", action="store_true") + parser.add_argument("--strict-cuda-ext", action="store_true") + args = parser.parse_args() + + print("== Environment ==") + print("python:", sys.version.split()[0]) + print("torch:", torch.__version__) + print("cuda_available:", torch.cuda.is_available()) + print("torch_cuda_version:", torch.version.cuda) + + arch, device_name = _device_profile() + print("device:", device_name) + print("arch:", arch) + print("support_status:", _support_status(arch)) + + if args.compile_cuda_ext: + built = _compile_cuda_extension() + print("cuda_extension_build:", "ok" if built else "failed") + + print("\n== Functional check ==") + ok = _run_functional_check(args.strict_cuda_ext) + + print("\n== Recommended SCAO config ==") + for key, value in _recommended_config(arch).items(): + print(f"{key}: {value}") + + if ok: + print("\nSCAO compatibility check: OK") + else: + raise SystemExit("SCAO compatibility check: FAILED") + + +if __name__ == "__main__": + main() diff --git a/scripts/preconditioner_microbench.py b/scripts/preconditioner_microbench.py new file mode 100644 index 0000000..9ed389e --- /dev/null +++ b/scripts/preconditioner_microbench.py @@ -0,0 +1,8 @@ +"""Repo wrapper for the installable ``scao-preconditioner-bench`` command.""" + +from __future__ import annotations + +from scao.tools.preconditioner_microbench import main + +if __name__ == "__main__": + main() diff --git a/scripts/scao_scale_planner.py b/scripts/scao_scale_planner.py new file mode 100644 index 0000000..89d8c81 --- /dev/null +++ b/scripts/scao_scale_planner.py @@ -0,0 +1,8 @@ +"""Repo wrapper for the installable ``scao-scale-planner`` command.""" + +from __future__ import annotations + +from scao.tools.scale_planner import main + +if __name__ == "__main__": + main()