diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml new file mode 100644 index 0000000..dba0da0 --- /dev/null +++ b/.github/workflows/ci.yml @@ -0,0 +1,38 @@ +name: ci + +on: + push: + branches: [main] + pull_request: + +permissions: + contents: read + +jobs: + lint-and-smoke: + runs-on: ubuntu-latest + timeout-minutes: 15 + steps: + - uses: actions/checkout@v4 + + - name: Set up Python + uses: actions/setup-python@v5 + with: + python-version: "3.11" + cache: pip + + - name: Install CPU-only deps + run: | + python -m pip install --upgrade pip + # CPU-only torch + torchvision to keep CI cheap. + pip install --index-url https://download.pytorch.org/whl/cpu \ + torch torchvision + pip install -e .[dev] + + - name: ruff check + run: ruff check . + + - name: pytest (smoke only, CPU) + env: + TORCH_USE_CUDA_DSA: "0" + run: pytest -q diff --git a/.gitignore b/.gitignore new file mode 100644 index 0000000..721bc7b --- /dev/null +++ b/.gitignore @@ -0,0 +1,48 @@ +# Python +__pycache__/ +*.py[cod] +*$py.class +*.so +.Python + +# Build artifacts +build/ +dist/ +*.egg-info/ +*.egg +.eggs/ + +# Virtual environments +.venv/ +venv/ +env/ +ENV/ + +# Dataset cache +data/ +datasets/ +*.tar.gz +cifar-10-batches-py/ +cifar-100-python/ + +# Test / tooling caches +.pytest_cache/ +.ruff_cache/ +.mypy_cache/ +.tox/ +htmlcov/ +.coverage +.coverage.* + +# IDE / OS +.idea/ +.vscode/ +.DS_Store + +# Checkpoints & logs +checkpoints/ +runs/ +*.pt +*.pth +*.ckpt +logs/ diff --git a/README.md b/README.md index 840989b..91d3abd 100644 --- a/README.md +++ b/README.md @@ -5,26 +5,88 @@ Reference implementation anchor for: **"Hierarchical Abstraction is All You Need: A Mathematical Framework for Continuous AI Consciousness Through Sleep-Wake Cycles"** Jihyuk Im, Claude Opus 4. June 2025. -This repository will host: - -- Minimal sleep–wake cycle reference (L1 context / L2 LoRA / L3 small model / L4 large model hierarchy) -- Reproducibility scripts for the continual-learning benchmark reported in §3 of the paper -- Entropy-aware context compression reference (from the pseudocode in §2.4) +Paper: ## Status -**Stub.** Code and benchmarks will land in phases: +**Phase I MVP scaffolding.** Code runs, tests pass, benchmark numbers have not been measured yet (CI has no GPU). See [`results/README.md`](results/README.md). + +Vision-stack mapping of the paper's LLM-stack hierarchy (paper §2.3): + +| Layer | Paper (LLM stack) | This MVP (vision stack) | +|-------|-------------------|-------------------------| +| L1 | Context window | Task-batch ring buffer (`hal_sleep_wake/memory.py`) | +| L2 | LoRA weights | LoRA adapter on `Conv2d` + `fc` (`hal_sleep_wake/model.py`) | +| L3 | Small model (7B) | ResNet-18 base weights (`hal_sleep_wake/model.py`) | +| L4 | Large model (70B) | *Out of scope — Phase II+* | + +Task sequence (subset of paper §3.1): **CIFAR-10 → CIFAR-100**. Remaining datasets land in follow-up PRs. + +## Install + +```bash +pip install -e '.[dev]' +``` + +`torch` + `torchvision` are heavy dependencies; for CPU-only use the PyTorch CPU wheel index: + +```bash +pip install --index-url https://download.pytorch.org/whl/cpu torch torchvision +pip install -e '.[dev]' +``` + +## Run + +Dry-run (prints the planned loop without downloading data or building a model): + +```bash +python scripts/train_cifar10_cifar100.py --dry-run +``` + +Real training (needs a GPU + dataset download): + +```bash +python scripts/train_cifar10_cifar100.py \ + --data-root ./data --download --device cuda \ + --results-out results/runs/hal.json + +python scripts/eval.py results/runs/hal.json +``` + +`--no-sleep` disables the consolidation phase (baseline for the ablation row). `--replay-n N` draws N samples from the L1 buffer into each wake step (0 = off, the default). Both knobs share the same code path so baseline vs HAL runs are directly comparable. + +## Known limitations (Phase I MVP) + +Flagged explicitly so the follow-up PRs target the right gaps, and so the eventual measurement numbers aren't read as something they're not: + +- **Shared classifier head.** The model uses `max(num_classes_per_task) = 100` logits for both tasks. CIFAR-10 trains against the first 10 positions; the remaining 90 logits are not masked during eval. Per-task heads land in a follow-up PR. +- **`sleep_scale` is ad-hoc.** Paper §2.4 does not define this scalar. It's an ablation knob only; leaving it at `1.0` reproduces the paper's unscaled `W' = W + (α/r) · B @ A`. +- **LoRA targets omit residual projections.** `("conv1", "conv2", "fc")` misses ResNet-18's `downsample.0` 1×1 convs on the residual path. Documented choice, not a bug. +- **Forgetting rate formula.** Uses `max_{j ≤ T-1} A[i,j] − A[i,T-1]` (includes the final round in the max). Stays non-negative even if a task improves on its last round; see `tests/test_metrics.py::test_forgetting_zero_when_task_improves`. +- **No GPU numbers yet.** CI has no GPU; `results/baseline_vs_hal.md` is unpopulated until a follow-up measurement PR. + +## Test + +```bash +ruff check . +pytest -q +``` + +Smoke tests run on CPU and do not download datasets. CI (`.github/workflows/ci.yml`) runs exactly these two commands. + +## Sleep cycle, concretely + +Paper §2.4 gives the full pseudocode (wake → REM → NREM). This MVP implements the minimum viable subset: -1. Notebook-scale sleep–wake prototype on a single LLM + LoRA -2. Continual-learning benchmark suite (CIFAR-10/100 → permuted MNIST → Fashion-MNIST → SVHN → CelebA → TinyImageNet) -3. Ablations (no-REM, no-NREM, no-entropy-compression, fixed vs adaptive schedule) +1. **Wake** (`hal_sleep_wake/wake.py`) — SGD over LoRA-only params on the current task's `DataLoader`, pushing every batch into the L1 buffer. +2. **Sleep — NREM consolidation** (`hal_sleep_wake/sleep.py`) — at the end of each task, LoRA deltas are merged into the ResNet base weights (`W ← W + s · (α/r) · B @ A`) and then reset to a fresh Kaiming / zero init. Optimizer state is wiped. -## Paper +REM synthetic-dream generation and entropy-aware context compression (paper §2.4 lines 85–101) are explicit follow-up items. -See `paper/HAL/L1_Hierarchical_Abstraction_All_You_Need_v2.md` in the companion paper repository: +## Phase I goal - +Per paper §5 Implementation Roadmap: *"50% forgetting reduction on the CIFAR sequence vs a vanilla baseline."* This repository will report the raw numbers in [`results/baseline_vs_hal.md`](results/README.md) as soon as a GPU run is completed. ## License -TBD (will track the companion paper repo's license once that is finalized). +TBD (will track the companion [paper repository](https://github.com/labforadvancedstudy/paper) once that is finalized). diff --git a/hal_sleep_wake/__init__.py b/hal_sleep_wake/__init__.py new file mode 100644 index 0000000..fad9c4a --- /dev/null +++ b/hal_sleep_wake/__init__.py @@ -0,0 +1,7 @@ +"""hal-sleep-wake: Phase I MVP scaffolding. + +Implements a vision-stack mapping of the HAL paper's LLM-stack hierarchy +(L1 context buffer / L2 LoRA / L3 backbone). No L4 tier in Phase I. +""" + +__version__ = "0.0.1" diff --git a/hal_sleep_wake/config.py b/hal_sleep_wake/config.py new file mode 100644 index 0000000..4fd7cfd --- /dev/null +++ b/hal_sleep_wake/config.py @@ -0,0 +1,75 @@ +"""Hyperparameter dataclass for the Phase I continual-learning scaffold.""" + +from __future__ import annotations + +from dataclasses import dataclass + + +@dataclass +class HyperParams: + """Single source of truth for MVP hyperparameters. + + Sleep merge formula (paper §2.4): ``W' = W + (alpha / r) * B @ A``. + ``sleep_scale`` (``s``) is an ad-hoc scalar not from the paper; it + multiplies that delta so ablations can attenuate consolidation: + ``W' = W + s * (alpha / r) * B @ A``. Default 1.0 reproduces the + paper's unscaled formula. + """ + + # Task schedule (Phase I: CIFAR-10 then CIFAR-100). + tasks: tuple[str, ...] = ("cifar10", "cifar100") + num_classes_per_task: tuple[int, ...] = (10, 100) + + # Training. + epochs_per_task: int = 1 + batch_size: int = 64 + lr: float = 1e-3 + weight_decay: float = 0.0 + + # L2 LoRA adapter. + lora_rank: int = 4 + lora_alpha: int = 8 + lora_target_modules: tuple[str, ...] = ("conv1", "conv2", "fc") + + # Sleep / wake cadence. + sleep_every_n_tasks: int = 1 # fire sleep at the end of every task + sleep_scale: float = 1.0 # extra ad-hoc scalar on the merged delta + + # L1 working-memory buffer + replay. + buffer_capacity: int = 256 + replay_n: int = 0 # samples drawn from L1 per wake step; 0 = off + + # Runtime. + device: str = "cpu" + seed: int = 0 + + def __post_init__(self) -> None: + if len(self.num_classes_per_task) != len(self.tasks): + raise ValueError( + f"num_classes_per_task length ({len(self.num_classes_per_task)}) " + f"must match tasks length ({len(self.tasks)})" + ) + if self.replay_n < 0: + raise ValueError("replay_n must be >= 0") + if self.buffer_capacity <= 0: + raise ValueError("buffer_capacity must be positive") + + def plan(self) -> str: + """Return a compact plan string (used by --dry-run).""" + lines = [ + "HAL sleep-wake | Phase I MVP plan", + f" tasks : {list(self.tasks)}", + f" classes per task : {list(self.num_classes_per_task)}", + f" epochs per task : {self.epochs_per_task}", + f" batch size : {self.batch_size}", + f" lr / weight decay : {self.lr} / {self.weight_decay}", + f" LoRA rank / alpha : {self.lora_rank} / {self.lora_alpha}", + f" LoRA targets : {list(self.lora_target_modules)}", + f" sleep cadence : every {self.sleep_every_n_tasks} task(s)", + f" sleep merge scale (s) : {self.sleep_scale}", + f" L1 buffer capacity : {self.buffer_capacity}", + f" L1 replay per step : {self.replay_n}", + f" device : {self.device}", + f" seed : {self.seed}", + ] + return "\n".join(lines) diff --git a/hal_sleep_wake/data.py b/hal_sleep_wake/data.py new file mode 100644 index 0000000..dd28438 --- /dev/null +++ b/hal_sleep_wake/data.py @@ -0,0 +1,92 @@ +"""Dataset loaders for Phase I continual-learning scaffold. + +CIFAR-10 and CIFAR-100 from torchvision. Kept minimal on purpose — the +MVP only consumes two tasks in sequence. Additional datasets from the +paper §3.1 sequence will land in follow-up PRs. + +The loaders here never download data unless explicitly asked; the dry-run +mode in ``scripts/train_cifar10_cifar100.py`` must not hit the network. +""" + +from __future__ import annotations + +from collections.abc import Callable +from dataclasses import dataclass +from pathlib import Path +from typing import Any + +import torch +from torch.utils.data import DataLoader + + +@dataclass +class TaskDatasets: + """One (train, test) pair plus the task's label count.""" + + name: str + train: Any # torch.utils.data.Dataset but torchvision types vary + test: Any + num_classes: int + + +def _cifar_transform() -> Callable[[Any], torch.Tensor]: + # Kept inline so the module imports without torchvision at import time + # when only type hints / dataclasses are needed (e.g. dry-run plan). + from torchvision import transforms + + return transforms.Compose( + [ + transforms.ToTensor(), + transforms.Normalize( + mean=(0.4914, 0.4822, 0.4465), + std=(0.2470, 0.2435, 0.2616), + ), + ] + ) + + +def load_cifar10(root: str | Path, *, download: bool = False) -> TaskDatasets: + from torchvision import datasets + + tfm = _cifar_transform() + train = datasets.CIFAR10(root=str(root), train=True, transform=tfm, download=download) + test = datasets.CIFAR10(root=str(root), train=False, transform=tfm, download=download) + return TaskDatasets(name="cifar10", train=train, test=test, num_classes=10) + + +def load_cifar100(root: str | Path, *, download: bool = False) -> TaskDatasets: + from torchvision import datasets + + tfm = _cifar_transform() + train = datasets.CIFAR100(root=str(root), train=True, transform=tfm, download=download) + test = datasets.CIFAR100(root=str(root), train=False, transform=tfm, download=download) + return TaskDatasets(name="cifar100", train=train, test=test, num_classes=100) + + +_LOADERS: dict[str, Callable[..., TaskDatasets]] = { + "cifar10": load_cifar10, + "cifar100": load_cifar100, +} + + +def load_task(name: str, root: str | Path, *, download: bool = False) -> TaskDatasets: + """Dispatch to the correct loader by task name.""" + if name not in _LOADERS: + raise KeyError(f"unknown task '{name}'; known: {sorted(_LOADERS)}") + return _LOADERS[name](root, download=download) + + +def make_loaders( + ds: TaskDatasets, + *, + batch_size: int, + num_workers: int = 0, + shuffle_train: bool = True, +) -> tuple[DataLoader, DataLoader]: + train_loader = DataLoader( + ds.train, batch_size=batch_size, shuffle=shuffle_train, num_workers=num_workers + ) + test_loader = DataLoader( + ds.test, batch_size=batch_size, shuffle=False, num_workers=num_workers + ) + return train_loader, test_loader diff --git a/hal_sleep_wake/memory.py b/hal_sleep_wake/memory.py new file mode 100644 index 0000000..f469040 --- /dev/null +++ b/hal_sleep_wake/memory.py @@ -0,0 +1,68 @@ +"""Memory hierarchy for the Phase I MVP. + +Vision-stack mapping of paper section 2.3: + +- L1 = task-batch buffer (working memory analog): bounded ring buffer of + recent samples; used to stabilize the wake-phase optimization. +- L2 = LoRA adapter on the backbone (fast weights): see ``model.attach_lora``. +- L3 = backbone base weights (slow weights / long-term substrate): merged + into via the sleep phase. + +Only L1 and helpers for L2/L3 live here. The actual LoRA wrap lives in +``model.py``; the merge routine lives in ``sleep.py``. +""" + +from __future__ import annotations + +from collections import deque +from collections.abc import Iterable +from dataclasses import dataclass + +import torch + + +@dataclass +class BufferItem: + """One (x, y, task_id) sample stored in the L1 buffer.""" + + x: torch.Tensor + y: torch.Tensor + task_id: int + + +class TaskBatchBuffer: + """Bounded FIFO buffer of recent samples across tasks (L1).""" + + def __init__(self, capacity: int) -> None: + if capacity <= 0: + raise ValueError("capacity must be positive") + self.capacity = capacity + self._buf: deque[BufferItem] = deque(maxlen=capacity) + + def __len__(self) -> int: + return len(self._buf) + + def add_batch(self, xs: torch.Tensor, ys: torch.Tensor, task_id: int) -> None: + if xs.shape[0] != ys.shape[0]: + raise ValueError("xs and ys must have the same batch size") + for i in range(xs.shape[0]): + self._buf.append(BufferItem(x=xs[i].detach().cpu(), y=ys[i].detach().cpu(), + task_id=task_id)) + + def sample(self, n: int, generator: torch.Generator | None = None + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + if len(self._buf) == 0: + raise RuntimeError("buffer is empty") + n = min(n, len(self._buf)) + idx = torch.randint(0, len(self._buf), (n,), generator=generator).tolist() + picks = [self._buf[i] for i in idx] + xs = torch.stack([p.x for p in picks], dim=0) + ys = torch.stack([p.y for p in picks], dim=0) + tids = torch.tensor([p.task_id for p in picks], dtype=torch.long) + return xs, ys, tids + + def items(self) -> Iterable[BufferItem]: + return iter(self._buf) + + def clear(self) -> None: + self._buf.clear() diff --git a/hal_sleep_wake/metrics.py b/hal_sleep_wake/metrics.py new file mode 100644 index 0000000..dd374a5 --- /dev/null +++ b/hal_sleep_wake/metrics.py @@ -0,0 +1,97 @@ +"""Continual-learning metrics. + +The MVP tracks a standard **accuracy matrix** ``A[i, j]`` = accuracy on +task ``i`` measured after the model has finished training task ``j`` +(where ``j >= i``). From that matrix we derive: + +- ``avg_accuracy`` — mean of the last column (``A[i, T-1]`` for all ``i``). +- ``forgetting_rate`` — mean of ``max_j(A[i, j]) - A[i, T-1]`` across + ``i < T-1`` (the standard "forgetting" measure; higher = more + catastrophic forgetting). + +Both follow the convention used in GEM / AGEM / EWC reports so later PRs +can drop straight into the paper §3.2 comparison table. +""" + +from __future__ import annotations + +from dataclasses import dataclass, field + +import torch +import torch.nn as nn +from torch.utils.data import DataLoader + + +@dataclass +class AccuracyMatrix: + """Row per task, column per evaluation moment (after each task finished).""" + + num_tasks: int + matrix: list[list[float]] = field(default_factory=list) + + def record_round(self, accs: list[float]) -> None: + """Append one column of accuracies (length must equal ``num_tasks``).""" + if len(accs) != self.num_tasks: + raise ValueError(f"expected {self.num_tasks} accuracies, got {len(accs)}") + self.matrix.append(list(accs)) + + def as_tensor(self) -> torch.Tensor: + """Shape ``(T_rounds, num_tasks)``. Rows = rounds, cols = tasks.""" + if not self.matrix: + return torch.zeros(0, self.num_tasks) + return torch.tensor(self.matrix, dtype=torch.float32) + + +def average_accuracy(matrix: AccuracyMatrix) -> float: + """Mean accuracy on all tasks at the final round.""" + if not matrix.matrix: + return 0.0 + last = matrix.matrix[-1] + return float(sum(last) / len(last)) + + +def forgetting_rate(matrix: AccuracyMatrix) -> float: + """Average of max-past-accuracy minus final accuracy over prior tasks. + + Concretely: ``mean_{i < T-1} [ max_{j <= T-1} A[i, j] - A[i, T-1] ]``. + Returns 0.0 if there are fewer than 2 rounds recorded. + """ + m = matrix.as_tensor() + if m.shape[0] < 2: + return 0.0 + num_rounds = m.shape[0] + final = m[-1] # last recorded round + forget = [] + # Tasks that existed before the last round: task_id in [0, num_rounds-2] + # (assumes record_round appended once after training task_id). + for i in range(num_rounds - 1): + max_past = torch.max(m[:, i]).item() + forget.append(max_past - float(final[i])) + if not forget: + return 0.0 + return float(sum(forget) / len(forget)) + + +@torch.no_grad() +def eval_accuracy( + model: nn.Module, + loader: DataLoader, + device: torch.device, + limit_batches: int | None = None, +) -> float: + """Top-1 accuracy on the provided loader. ``limit_batches`` caps work.""" + model.eval() + correct = 0 + total = 0 + for bi, (x, y) in enumerate(loader): + if limit_batches is not None and bi >= limit_batches: + break + x = x.to(device) + y = y.to(device) + logits = model(x) + pred = logits.argmax(dim=1) + correct += int((pred == y).sum().item()) + total += int(y.numel()) + if total == 0: + return 0.0 + return correct / total diff --git a/hal_sleep_wake/model.py b/hal_sleep_wake/model.py new file mode 100644 index 0000000..beabb25 --- /dev/null +++ b/hal_sleep_wake/model.py @@ -0,0 +1,75 @@ +"""ResNet-18 backbone (L3) + LoRA injection (L2). + +Keeps ResNet-18 base weights frozen; only LoRA-A/B parameters are trainable. +We attach LoRA via :mod:`peft` on both ``Conv2d`` layers (``conv1``, ``conv2`` +inside ``layer1..layer4``) and the final ``fc`` classifier. peft >= 0.10 +supports Conv2d LoRA targets out of the box; the smoke tests assert that +only LoRA params are trainable regardless of backbone. +""" + +from __future__ import annotations + +from collections.abc import Iterable + +import torch.nn as nn +from peft import LoraConfig, PeftModel, get_peft_model +from torchvision.models import resnet18 + +from .config import HyperParams + + +def build_backbone(num_classes: int) -> nn.Module: + """Return a fresh (randomly initialized) ResNet-18 with a sized head.""" + m = resnet18(weights=None) + m.fc = nn.Linear(m.fc.in_features, num_classes) + return m + + +def _freeze(module: nn.Module) -> None: + for p in module.parameters(): + p.requires_grad_(False) + + +def attach_lora( + model: nn.Module, + r: int, + alpha: int, + target_modules: Iterable[str] = ("conv1", "conv2", "fc"), +) -> tuple[PeftModel, list[nn.Parameter]]: + """Attach LoRA adapters to ``model``; return ``(peft_model, lora_params)``. + + All non-LoRA parameters are frozen. The caller is expected to only pass + ``lora_params`` to the optimizer. The backbone base weights remain the + long-term substrate (L3); only LoRA deltas (L2) learn per task. + """ + _freeze(model) + cfg = LoraConfig( + r=r, + lora_alpha=alpha, + lora_dropout=0.0, + bias="none", + target_modules=list(target_modules), + ) + peft_model = get_peft_model(model, cfg) + lora_params = [p for n, p in peft_model.named_parameters() if "lora_" in n] + # Authoritative final requires_grad state: LoRA tensors train, everything + # else is frozen. peft may leave non-LoRA trainable in edge cases; this + # loop enforces the contract regardless. + for n, p in peft_model.named_parameters(): + p.requires_grad_("lora_" in n) + return peft_model, lora_params + + +def build_model(hp: HyperParams, num_classes: int) -> tuple[PeftModel, list[nn.Parameter]]: + """Convenience: backbone + LoRA in one call.""" + backbone = build_backbone(num_classes=num_classes) + return attach_lora( + backbone, + r=hp.lora_rank, + alpha=hp.lora_alpha, + target_modules=hp.lora_target_modules, + ) + + +def count_trainable(model: nn.Module) -> int: + return sum(p.numel() for p in model.parameters() if p.requires_grad) diff --git a/hal_sleep_wake/sleep.py b/hal_sleep_wake/sleep.py new file mode 100644 index 0000000..b07e84f --- /dev/null +++ b/hal_sleep_wake/sleep.py @@ -0,0 +1,88 @@ +"""Sleep phase: merge LoRA deltas into base weights, then re-init LoRA. + +Paper mapping: at end of each task (or after ``sleep_every_n_tasks``), +consolidate the fast weights (L2 LoRA) into the slow weights (L3 +backbone) by summing ``W' = W + (alpha/r) * B @ A`` into every target +module's weight tensor, then zero-reset the LoRA tensors so the next +task starts from a clean delta. + +``sleep_scale`` is NOT from the paper. It is an ad-hoc scalar knob that +multiplies the LoRA delta during merge (``W' = W + s * (alpha/r) * B @ A``) +so ablations can attenuate the consolidation aggressiveness. Default 1.0 +reproduces the paper's unscaled formula. +""" + +from __future__ import annotations + +import torch +import torch.nn as nn +from peft import PeftModel + + +def _iter_lora_linear(peft_model: PeftModel): + """Yield ``(name, module)`` for every LoRA-wrapped layer.""" + for name, mod in peft_model.named_modules(): + # peft LoRA layers have ``lora_A`` / ``lora_B`` as ModuleDicts. + has_A = hasattr(mod, "lora_A") and isinstance(mod.lora_A, nn.ModuleDict) + has_B = hasattr(mod, "lora_B") and isinstance(mod.lora_B, nn.ModuleDict) + if has_A and has_B and len(mod.lora_A) > 0: + yield name, mod + + +@torch.no_grad() +def merge_lora_into_base(peft_model: PeftModel, scale: float = 1.0, + adapter_name: str = "default") -> int: + """Merge LoRA deltas into the base weights. Returns # layers merged. + + Uses peft's native ``merge`` on each LoRA layer (which handles the + Conv2d vs Linear arithmetic for us), with an explicit scale applied + by temporarily patching the per-adapter ``scaling`` dict. After + merge, all LoRA A/B tensors are re-initialized (A=Kaiming, B=0) and + the layer's ``merged_adapters`` list is cleared so peft re-applies + the (fresh-zero) delta on the next forward — i.e. the next wake + phase sees ``W_consolidated + 0`` and can learn a new delta from + scratch without peft believing the adapter is already baked in. + """ + merged = 0 + for _name, layer in _iter_lora_linear(peft_model): + if adapter_name not in layer.scaling: + continue + original_scaling = layer.scaling[adapter_name] + layer.scaling[adapter_name] = original_scaling * float(scale) + try: + layer.merge(adapter_names=[adapter_name]) + finally: + # Always restore — never leak the scaled value even if merge raises. + layer.scaling[adapter_name] = original_scaling + merged += 1 + + # Re-init A to Kaiming-uniform, B to zeros, and clear merged_adapters so + # peft's forward() re-applies the adapter (now a zero delta) on top of + # the updated base weight. Without the ``merged_adapters.clear()`` call, + # subsequent training of A/B would be silently ignored — peft would skip + # the delta path because it believes the adapter is already merged. + for _name, layer in _iter_lora_linear(peft_model): + for key in list(layer.lora_A.keys()): + nn.init.kaiming_uniform_(layer.lora_A[key].weight, a=5**0.5) + nn.init.zeros_(layer.lora_B[key].weight) + if hasattr(layer, "merged_adapters"): + layer.merged_adapters.clear() + return merged + + +def reset_optimizer_state(optimizer: torch.optim.Optimizer) -> None: + """Wipe optimizer accumulators (momentum, second moments, ...). + + Uses ``state.clear()`` rather than ``state = {}`` so the underlying + ``defaultdict(dict)`` factory on recent torch versions is preserved. + """ + optimizer.state.clear() + + +def sleep_phase(peft_model: PeftModel, optimizer: torch.optim.Optimizer | None = None, + scale: float = 1.0) -> int: + """Run a full sleep cycle: merge LoRA -> base, reset LoRA + optimizer.""" + n = merge_lora_into_base(peft_model, scale=scale) + if optimizer is not None: + reset_optimizer_state(optimizer) + return n diff --git a/hal_sleep_wake/wake.py b/hal_sleep_wake/wake.py new file mode 100644 index 0000000..0dbd335 --- /dev/null +++ b/hal_sleep_wake/wake.py @@ -0,0 +1,93 @@ +"""Wake phase: train LoRA adapter on the current task's data. + +The wake loop mirrors a standard SGD epoch but explicitly: (a) only passes +LoRA parameters to the optimizer, (b) pushes each incoming batch into the +L1 buffer for future replay, and (c) if ``replay_n > 0`` and the buffer +contains samples from *past* batches, concats a random mini-batch drawn +from the buffer into the current batch before the forward pass. Replay +is off by default (``replay_n=0``) so the baseline run is a clean +LoRA-only SGD; turning it on is the Phase I knob for measuring whether +L1 replay alone already reduces forgetting vs the pure-sleep ablation. +""" + +from __future__ import annotations + +from collections.abc import Iterable + +import torch +import torch.nn as nn +import torch.nn.functional as F +from torch.utils.data import DataLoader + +from .memory import TaskBatchBuffer + + +def wake_step( + model: nn.Module, + batch: tuple[torch.Tensor, torch.Tensor], + optimizer: torch.optim.Optimizer, + device: torch.device, + buffer: TaskBatchBuffer | None = None, + replay_n: int = 0, +) -> float: + """One optimizer step. Returns scalar loss. + + If ``buffer`` is non-empty and ``replay_n > 0``, ``replay_n`` samples + are drawn and concatenated to the current batch. Replay uses the + current ``buffer`` contents *before* the caller has pushed the + current batch — i.e. the first batch of a brand-new task sees only + older task(s), not itself. + """ + model.train() + x, y = batch + x = x.to(device) + y = y.to(device) + if buffer is not None and replay_n > 0 and len(buffer) > 0: + rx, ry, _rtid = buffer.sample(replay_n) + x = torch.cat([x, rx.to(device)], dim=0) + y = torch.cat([y, ry.to(device)], dim=0) + optimizer.zero_grad(set_to_none=True) + logits = model(x) + loss = F.cross_entropy(logits, y) + loss.backward() + optimizer.step() + return float(loss.detach()) + + +def wake_epoch( + model: nn.Module, + loader: Iterable[tuple[torch.Tensor, torch.Tensor]], + optimizer: torch.optim.Optimizer, + device: torch.device, + task_id: int, + buffer: TaskBatchBuffer | None = None, + replay_n: int = 0, +) -> float: + """Run one wake-phase epoch; return mean per-batch loss.""" + losses: list[float] = [] + for x, y in loader: + # Sample replay first, then push this batch — so a sample drawn + # for this step never includes this step's own data. + losses.append(wake_step(model, (x, y), optimizer, device, + buffer=buffer, replay_n=replay_n)) + if buffer is not None: + buffer.add_batch(x, y, task_id=task_id) + return sum(losses) / max(len(losses), 1) + + +def wake_phase( + model: nn.Module, + loader: DataLoader, + optimizer: torch.optim.Optimizer, + device: torch.device, + task_id: int, + epochs: int, + buffer: TaskBatchBuffer | None = None, + replay_n: int = 0, +) -> list[float]: + """Full wake phase for one task; return per-epoch mean losses.""" + history: list[float] = [] + for _ in range(epochs): + history.append(wake_epoch(model, loader, optimizer, device, task_id, + buffer=buffer, replay_n=replay_n)) + return history diff --git a/pyproject.toml b/pyproject.toml new file mode 100644 index 0000000..83d0867 --- /dev/null +++ b/pyproject.toml @@ -0,0 +1,45 @@ +[build-system] +requires = ["setuptools>=68"] +build-backend = "setuptools.build_meta" + +[project] +name = "hal-sleep-wake" +version = "0.0.1" +description = "Phase I MVP scaffolding for the HAL sleep-wake continual-learning reference implementation." +readme = "README.md" +requires-python = ">=3.11" +license = { text = "TBD" } +authors = [ + { name = "Lab for Advanced Study" }, +] +dependencies = [ + "torch>=2.2", + "torchvision>=0.17", + "peft>=0.10", + "numpy", +] + +[project.optional-dependencies] +dev = [ + "pytest>=8", + "ruff>=0.6", +] + +[project.urls] +Homepage = "https://github.com/labforadvancedstudy/hal-sleep-wake" +Paper = "https://github.com/labforadvancedstudy/paper/blob/main/HAL/L1_Hierarchical_Abstraction_All_You_Need_v2.md" + +[tool.setuptools.packages.find] +include = ["hal_sleep_wake*"] +exclude = ["tests*", "scripts*", "results*"] + +[tool.ruff] +target-version = "py311" +line-length = 100 + +[tool.ruff.lint] +select = ["E", "F", "I", "B", "UP"] + +[tool.pytest.ini_options] +testpaths = ["tests"] +addopts = "-ra" diff --git a/results/README.md b/results/README.md new file mode 100644 index 0000000..033235d --- /dev/null +++ b/results/README.md @@ -0,0 +1,29 @@ +# results/ + +Placeholder for benchmark runs. + +## What lands here + +- `baseline_vs_hal.md` — table comparing `--no-sleep` baseline vs HAL sleep-wake on the CIFAR-10 → CIFAR-100 sequence (Phase I target: 50% forgetting reduction per paper §5). +- `runs/.json` — raw metric dumps produced by `scripts/train_cifar10_cifar100.py --results-out`. + +## Status + +**Not yet measured.** This MVP PR ships scaffolding + smoke tests only; CI has no GPU. Real runs will land in a follow-up PR. + +## How to run locally + +```bash +pip install -e '.[dev]' + +# baseline (LoRA only, no consolidation) +python scripts/train_cifar10_cifar100.py --data-root ./data --download \ + --device cuda --no-sleep --results-out results/runs/baseline.json + +# HAL sleep-wake +python scripts/train_cifar10_cifar100.py --data-root ./data --download \ + --device cuda --results-out results/runs/hal.json + +python scripts/eval.py results/runs/baseline.json +python scripts/eval.py results/runs/hal.json +``` diff --git a/scripts/eval.py b/scripts/eval.py new file mode 100644 index 0000000..8b07136 --- /dev/null +++ b/scripts/eval.py @@ -0,0 +1,62 @@ +#!/usr/bin/env python3 +"""Load a saved run JSON and emit the accuracy matrix + derived metrics. + +Usage: + python scripts/eval.py path/to/results.json + +Consumes the JSON produced by ``train_cifar10_cifar100.py --results-out``. +Does no model / data loading — pure numeric reporting. Useful for CI-time +smoke invocations that do not require GPU. +""" + +from __future__ import annotations + +import argparse +import json +import sys +from pathlib import Path + +_REPO_ROOT = Path(__file__).resolve().parents[1] +if str(_REPO_ROOT) not in sys.path: + sys.path.insert(0, str(_REPO_ROOT)) + +from hal_sleep_wake.metrics import AccuracyMatrix, average_accuracy, forgetting_rate # noqa: E402 + + +def build_arg_parser() -> argparse.ArgumentParser: + p = argparse.ArgumentParser(description=__doc__) + p.add_argument("results_json", type=str, help="Path to results JSON from training run.") + return p + + +def main(argv: list[str] | None = None) -> int: + args = build_arg_parser().parse_args(argv) + path = Path(args.results_json) + if not path.exists(): + print(f"ERROR: {path} does not exist", file=sys.stderr) + return 2 + + raw = json.loads(path.read_text()) + mat_list = raw.get("accuracy_matrix") or [] + if not mat_list: + print("ERROR: no accuracy_matrix in JSON", file=sys.stderr) + return 2 + + num_tasks = len(mat_list[0]) + am = AccuracyMatrix(num_tasks=num_tasks) + for row in mat_list: + am.record_round(row) + + print(f"accuracy matrix ({len(mat_list)} rounds × {num_tasks} tasks):") + for i, row in enumerate(mat_list): + print(f" round {i}: {[f'{v:.3f}' for v in row]}") + + print(f"\navg accuracy : {average_accuracy(am):.4f}") + print(f"forgetting rate : {forgetting_rate(am):.4f}") + if "sleep_enabled" in raw: + print(f"sleep enabled : {raw['sleep_enabled']}") + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/scripts/train_cifar10_cifar100.py b/scripts/train_cifar10_cifar100.py new file mode 100644 index 0000000..17fc95a --- /dev/null +++ b/scripts/train_cifar10_cifar100.py @@ -0,0 +1,172 @@ +#!/usr/bin/env python3 +"""Two-task continual-learning entry point (CIFAR-10 -> CIFAR-100). + +This is the MVP entry point. ``--dry-run`` prints the plan only (no data +download, no model build, no training). Without ``--dry-run`` it actually +trains — that requires GPU + dataset download and is OUT OF SCOPE for CI. + +Usage: + python scripts/train_cifar10_cifar100.py --dry-run + python scripts/train_cifar10_cifar100.py --data-root ./data --epochs 1 + +Phase I goal (paper §5): 50% forgetting reduction on a CIFAR sequence +vs. a LoRA-only baseline (no sleep merge). The script supports both +modes via ``--no-sleep`` so the same code path produces both rows of +``results/baseline_vs_hal.md``. +""" + +from __future__ import annotations + +import argparse +import json +import sys +from pathlib import Path + +# Allow running from repo root without installing. +_REPO_ROOT = Path(__file__).resolve().parents[1] +if str(_REPO_ROOT) not in sys.path: + sys.path.insert(0, str(_REPO_ROOT)) + +from hal_sleep_wake.config import HyperParams # noqa: E402 + + +def build_arg_parser() -> argparse.ArgumentParser: + p = argparse.ArgumentParser(description=__doc__) + p.add_argument("--dry-run", action="store_true", + help="Print the plan and exit (no downloads, no training).") + p.add_argument("--data-root", type=str, default="./data", + help="Where CIFAR datasets live / will be downloaded to.") + p.add_argument("--download", action="store_true", + help="Allow torchvision to download datasets if missing.") + p.add_argument("--epochs", type=int, default=None, + help="Override HyperParams.epochs_per_task.") + p.add_argument("--batch-size", type=int, default=None, + help="Override HyperParams.batch_size.") + p.add_argument("--no-sleep", action="store_true", + help="Disable sleep consolidation (baseline for ablation).") + p.add_argument("--replay-n", type=int, default=None, + help="Samples drawn from L1 buffer per wake step " + "(0 = replay off, default from HyperParams).") + p.add_argument("--device", type=str, default=None, + help="Override HyperParams.device (cpu / cuda / mps).") + p.add_argument("--seed", type=int, default=None, help="Override seed.") + p.add_argument("--results-out", type=str, default=None, + help="If set, write JSON metrics to this path after training.") + return p + + +def _apply_overrides(hp: HyperParams, args: argparse.Namespace) -> HyperParams: + if args.epochs is not None: + hp.epochs_per_task = args.epochs + if args.batch_size is not None: + hp.batch_size = args.batch_size + if args.device is not None: + hp.device = args.device + if args.seed is not None: + hp.seed = args.seed + if args.replay_n is not None: + hp.replay_n = args.replay_n + return hp + + +def main(argv: list[str] | None = None) -> int: + args = build_arg_parser().parse_args(argv) + hp = _apply_overrides(HyperParams(), args) + + print(hp.plan()) + print(f" sleep enabled : {not args.no_sleep}") + print(f" data root : {args.data_root}") + print(f" download allowed : {args.download}") + + if args.dry_run: + print("\n[dry-run] no downloads, no model build, no training. exiting.") + return 0 + + # Non-dry-run path. Guarded imports so the dry-run path never requires + # torch to be installed for a plan-only invocation. + import torch + + from hal_sleep_wake.data import load_task, make_loaders + from hal_sleep_wake.memory import TaskBatchBuffer + from hal_sleep_wake.metrics import ( + AccuracyMatrix, + average_accuracy, + eval_accuracy, + forgetting_rate, + ) + from hal_sleep_wake.model import build_model + from hal_sleep_wake.sleep import sleep_phase + from hal_sleep_wake.wake import wake_phase + + torch.manual_seed(hp.seed) + device = torch.device(hp.device) + + # Build datasets up-front; each task gets its own loaders + head size. + tasks = [] + for name in hp.tasks: + ds = load_task(name, args.data_root, download=args.download) + train_loader, test_loader = make_loaders(ds, batch_size=hp.batch_size) + tasks.append((ds, train_loader, test_loader)) + + # We use the max-over-tasks classifier head size (CIFAR-100 => 100) so + # the same model can score all tasks without head-switching gymnastics + # in the MVP. Per-task heads land in a follow-up PR. + num_classes = max(ds.num_classes for ds, _, _ in tasks) + + model, lora_params = build_model(hp, num_classes=num_classes) + model.to(device) + optimizer = torch.optim.AdamW(lora_params, lr=hp.lr, weight_decay=hp.weight_decay) + buffer = TaskBatchBuffer(capacity=hp.buffer_capacity) + + acc_mat = AccuracyMatrix(num_tasks=len(tasks)) + + for task_id, (ds, train_loader, _test_loader) in enumerate(tasks): + print(f"\n=== task {task_id} | {ds.name} ===") + losses = wake_phase( + model=model, + loader=train_loader, + optimizer=optimizer, + device=device, + task_id=task_id, + epochs=hp.epochs_per_task, + buffer=buffer, + replay_n=hp.replay_n, + ) + print(f" wake per-epoch mean loss: {losses}") + + if not args.no_sleep and ((task_id + 1) % hp.sleep_every_n_tasks == 0): + merged = sleep_phase(model, optimizer=optimizer, scale=hp.sleep_scale) + print(f" sleep merged layers: {merged}") + + # Evaluate on all tasks seen so far. + accs = [] + for eval_id, (_eds, _tl, test_loader) in enumerate(tasks): + if eval_id > task_id: + accs.append(0.0) + continue + accs.append(eval_accuracy(model, test_loader, device)) + acc_mat.record_round(accs) + print(f" round accuracies: {[f'{a:.3f}' for a in accs]}") + + avg = average_accuracy(acc_mat) + fgt = forgetting_rate(acc_mat) + print(f"\nfinal avg accuracy : {avg:.4f}") + print(f"forgetting rate : {fgt:.4f}") + + if args.results_out: + out = { + "avg_accuracy": avg, + "forgetting_rate": fgt, + "accuracy_matrix": acc_mat.matrix, + "sleep_enabled": not args.no_sleep, + "hyperparams": hp.__dict__, + } + Path(args.results_out).parent.mkdir(parents=True, exist_ok=True) + Path(args.results_out).write_text(json.dumps(out, indent=2)) + print(f"metrics written to {args.results_out}") + + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/tests/test_memory.py b/tests/test_memory.py new file mode 100644 index 0000000..95dffeb --- /dev/null +++ b/tests/test_memory.py @@ -0,0 +1,63 @@ +"""Smoke tests for the L1 working-memory buffer.""" + +from __future__ import annotations + +import pytest + +torch = pytest.importorskip("torch") # skips the module if torch is unavailable + +from hal_sleep_wake.memory import TaskBatchBuffer # noqa: E402 + + +def test_buffer_rejects_zero_capacity(): + with pytest.raises(ValueError): + TaskBatchBuffer(capacity=0) + + +def test_buffer_stores_and_samples(): + buf = TaskBatchBuffer(capacity=32) + xs = torch.randn(8, 3, 4, 4) + ys = torch.arange(8) + buf.add_batch(xs, ys, task_id=0) + assert len(buf) == 8 + + sx, sy, stid = buf.sample(4, generator=torch.Generator().manual_seed(0)) + assert sx.shape == (4, 3, 4, 4) + assert sy.shape == (4,) + assert stid.shape == (4,) + assert (stid == 0).all() + + +def test_buffer_respects_capacity(): + buf = TaskBatchBuffer(capacity=4) + xs = torch.randn(8, 2) + ys = torch.arange(8) + buf.add_batch(xs, ys, task_id=1) + assert len(buf) == 4 # FIFO drop + + +def test_buffer_shape_validation(): + buf = TaskBatchBuffer(capacity=4) + xs = torch.randn(2, 3) + ys = torch.arange(3) + with pytest.raises(ValueError): + buf.add_batch(xs, ys, task_id=0) + + +def test_buffer_clear(): + buf = TaskBatchBuffer(capacity=4) + buf.add_batch(torch.zeros(2, 2), torch.zeros(2, dtype=torch.long), task_id=0) + buf.clear() + assert len(buf) == 0 + with pytest.raises(RuntimeError): + buf.sample(1) + + +def test_buffer_mixes_multiple_task_ids(): + """Replay's whole point: draw across tasks, not just the current one.""" + buf = TaskBatchBuffer(capacity=32) + buf.add_batch(torch.zeros(8, 3, 4, 4), torch.zeros(8, dtype=torch.long), task_id=0) + buf.add_batch(torch.ones(8, 3, 4, 4), torch.ones(8, dtype=torch.long), task_id=1) + # Sample the whole buffer; both task IDs must appear. + _xs, _ys, tids = buf.sample(16, generator=torch.Generator().manual_seed(0)) + assert set(tids.tolist()) == {0, 1} diff --git a/tests/test_metrics.py b/tests/test_metrics.py new file mode 100644 index 0000000..7f9e06e --- /dev/null +++ b/tests/test_metrics.py @@ -0,0 +1,70 @@ +"""Unit tests for the continual-learning metrics.""" + +from __future__ import annotations + +import pytest + +# metrics.py imports torch; skip cleanly in torch-less environments. +pytest.importorskip("torch") + +from hal_sleep_wake.metrics import AccuracyMatrix, average_accuracy, forgetting_rate # noqa: E402 + + +def test_average_accuracy_empty_matrix_is_zero(): + am = AccuracyMatrix(num_tasks=2) + assert average_accuracy(am) == 0.0 + + +def test_record_round_validates_length(): + am = AccuracyMatrix(num_tasks=3) + with pytest.raises(ValueError): + am.record_round([0.1, 0.2]) # wrong length + + +def test_forgetting_zero_when_no_drop(): + am = AccuracyMatrix(num_tasks=2) + am.record_round([0.8, 0.0]) # after task 0 + am.record_round([0.8, 0.9]) # after task 1 — no drop on task 0 + assert forgetting_rate(am) == pytest.approx(0.0) + + +def test_forgetting_positive_when_drop(): + am = AccuracyMatrix(num_tasks=2) + am.record_round([0.9, 0.0]) # after task 0 + am.record_round([0.6, 0.9]) # task 0 dropped 0.3 + assert forgetting_rate(am) == pytest.approx(0.3) + + +def test_forgetting_single_round_is_zero(): + am = AccuracyMatrix(num_tasks=2) + am.record_round([0.5, 0.0]) + assert forgetting_rate(am) == 0.0 + + +def test_average_accuracy_uses_last_round(): + am = AccuracyMatrix(num_tasks=3) + am.record_round([0.9, 0.0, 0.0]) + am.record_round([0.6, 0.8, 0.0]) + am.record_round([0.5, 0.7, 0.9]) + assert average_accuracy(am) == pytest.approx((0.5 + 0.7 + 0.9) / 3) + + +def test_forgetting_three_tasks(): + am = AccuracyMatrix(num_tasks=3) + am.record_round([0.90, 0.0, 0.0]) # task 0: max 0.90 + am.record_round([0.70, 0.80, 0.0]) # task 0: 0.70, task 1 max 0.80 + am.record_round([0.50, 0.60, 0.95]) # task 0: 0.50, task 1: 0.60, task 2: 0.95 + # forgetting = mean[ (0.90 - 0.50), (0.80 - 0.60) ] = (0.40 + 0.20) / 2 = 0.30 + assert forgetting_rate(am) == pytest.approx(0.30) + + +def test_forgetting_zero_when_task_improves(): + """If a task's final accuracy beats its past best, forgetting stays at 0. + + Documents the ``max_{j <= T-1}`` (includes final round) convention — + our formula can never go negative for improvement cases. + """ + am = AccuracyMatrix(num_tasks=2) + am.record_round([0.50, 0.0]) # after task 0: task 0 at 0.50 + am.record_round([0.90, 0.80]) # after task 1: task 0 IMPROVED to 0.90 + assert forgetting_rate(am) == pytest.approx(0.0) diff --git a/tests/test_model.py b/tests/test_model.py new file mode 100644 index 0000000..1115e18 --- /dev/null +++ b/tests/test_model.py @@ -0,0 +1,57 @@ +"""Model + LoRA wiring tests. CPU-only, no dataset download.""" + +from __future__ import annotations + +import pytest + +torch = pytest.importorskip("torch") +pytest.importorskip("peft") +pytest.importorskip("torchvision") + +from hal_sleep_wake.config import HyperParams # noqa: E402 +from hal_sleep_wake.model import build_model, count_trainable # noqa: E402 + + +@pytest.fixture(scope="module") +def tiny_hp() -> HyperParams: + # Shrink the LoRA rank so the module imports + builds quickly on CPU. + return HyperParams( + tasks=("cifar10",), + num_classes_per_task=(10,), + lora_rank=2, + lora_alpha=4, + device="cpu", + seed=0, + ) + + +def test_build_model_only_lora_is_trainable(tiny_hp: HyperParams): + model, lora_params = build_model(tiny_hp, num_classes=10) + # Everything trainable must be a LoRA parameter. + trainable_names = {n for n, p in model.named_parameters() if p.requires_grad} + assert trainable_names, "no trainable parameters found — LoRA wiring broken" + for n in trainable_names: + assert "lora_" in n, f"non-LoRA param is trainable: {n}" + + # And the separately-returned lora_params must match that set exactly + # (both contain the same tensor objects). + lora_ids = {id(p) for p in lora_params} + trainable_ids = {id(p) for p in model.parameters() if p.requires_grad} + assert lora_ids == trainable_ids + + +def test_build_model_forward_shape(tiny_hp: HyperParams): + model, _ = build_model(tiny_hp, num_classes=10) + model.eval() + x = torch.randn(2, 3, 32, 32) + with torch.no_grad(): + y = model(x) + assert y.shape == (2, 10) + + +def test_trainable_count_is_small(tiny_hp: HyperParams): + """LoRA adds few params compared to ResNet-18 base (~11M).""" + model, _ = build_model(tiny_hp, num_classes=10) + trainable = count_trainable(model) + # We expect at most a few hundred thousand trainable LoRA params. + assert 0 < trainable < 2_000_000, f"unexpectedly high trainable count: {trainable}" diff --git a/tests/test_smoke.py b/tests/test_smoke.py new file mode 100644 index 0000000..7604acf --- /dev/null +++ b/tests/test_smoke.py @@ -0,0 +1,224 @@ +"""End-to-end smoke: dry-run entry point, config.plan(), sleep merge. + +CPU-only, no dataset download, no network. The goal is that CI confirms +the scaffold *runs* without importing torch for the dry-run path, and +that the sleep merge actually writes the documented delta into base +weights (not just "no exception"). +""" + +from __future__ import annotations + +import json +import subprocess +import sys +from pathlib import Path + +import pytest + +REPO_ROOT = Path(__file__).resolve().parents[1] + + +def test_config_plan_renders_without_torch(): + """``HyperParams.plan()`` must not import torch.""" + from hal_sleep_wake.config import HyperParams + text = HyperParams().plan() + assert "tasks" in text + assert "LoRA" in text + + +def test_config_rejects_mismatched_task_lengths(): + """``num_classes_per_task`` length must match ``tasks``.""" + from hal_sleep_wake.config import HyperParams + with pytest.raises(ValueError, match="num_classes_per_task"): + HyperParams(tasks=("a", "b"), num_classes_per_task=(10,)) + + +def test_dry_run_script_exits_zero(): + """``python scripts/train_cifar10_cifar100.py --dry-run`` must exit 0.""" + result = subprocess.run( + [sys.executable, str(REPO_ROOT / "scripts" / "train_cifar10_cifar100.py"), + "--dry-run"], + capture_output=True, + text=True, + timeout=60, + check=False, + ) + assert result.returncode == 0, f"stdout:\n{result.stdout}\nstderr:\n{result.stderr}" + assert "dry-run" in result.stdout + + +def test_eval_script_handles_missing_file(): + """``eval.py nonexistent`` must exit 2 with a clear error.""" + result = subprocess.run( + [sys.executable, str(REPO_ROOT / "scripts" / "eval.py"), + str(REPO_ROOT / "does_not_exist.json")], + capture_output=True, + text=True, + timeout=30, + check=False, + ) + assert result.returncode == 2, ( + f"expected exit 2, got {result.returncode}\n" + f"stdout:\n{result.stdout}\nstderr:\n{result.stderr}" + ) + + +def test_eval_script_rejects_empty_matrix(tmp_path: Path): + """``eval.py`` must refuse a JSON with an empty accuracy_matrix.""" + results_path = tmp_path / "empty.json" + results_path.write_text(json.dumps({ + "avg_accuracy": 0.0, + "forgetting_rate": 0.0, + "accuracy_matrix": [], + "sleep_enabled": False, + "hyperparams": {}, + })) + result = subprocess.run( + [sys.executable, str(REPO_ROOT / "scripts" / "eval.py"), str(results_path)], + capture_output=True, + text=True, + timeout=30, + check=False, + ) + assert result.returncode == 2, ( + f"expected exit 2 for empty matrix, got {result.returncode}\n" + f"stdout:\n{result.stdout}\nstderr:\n{result.stderr}" + ) + + +def test_sleep_merge_applies_correct_delta_to_fc_head(): + """The merge must write ``scale * (alpha/r) * B @ A`` into the base weight. + + This is the behavioral contract for the novel bit of the scaffold. + Asserting only ``merged >= 1`` would let a silent math regression + (wrong scaling, wrong matmul order, no-op merge) ship. + """ + torch = pytest.importorskip("torch") + pytest.importorskip("peft") + pytest.importorskip("torchvision") + + from hal_sleep_wake.config import HyperParams + from hal_sleep_wake.model import build_model + from hal_sleep_wake.sleep import sleep_phase + + torch.manual_seed(0) + hp = HyperParams(lora_rank=2, lora_alpha=4, device="cpu") + model, lora_params = build_model(hp, num_classes=10) + model.eval() + + # Find the fc LoRA-wrapped Linear (base_layer.weight is 2D). + target = None + for _name, mod in model.named_modules(): + if (hasattr(mod, "lora_A") and hasattr(mod, "base_layer") + and mod.base_layer.weight.ndim == 2): + target = mod + break + assert target is not None, "could not find LoRA-wrapped fc layer" + + # Give lora_B nonzero values so the delta is observable. + torch.nn.init.normal_(target.lora_B["default"].weight, mean=0.0, std=0.1) + + A = target.lora_A["default"].weight.detach().clone() # (r, in) + B = target.lora_B["default"].weight.detach().clone() # (out, r) + scaling = float(target.scaling["default"]) # alpha / r + base_before = target.base_layer.weight.detach().clone() + + scale = 0.5 + expected_delta = scale * scaling * (B @ A) + + # Populate optimizer state so we can verify sleep wipes it. + opt = torch.optim.SGD(lora_params, lr=1e-2) + for p in lora_params: + p.grad = torch.zeros_like(p) + opt.step() + assert len(opt.state) > 0 + + merged = sleep_phase(model, optimizer=opt, scale=scale) + assert merged >= 1 + + actual_delta = target.base_layer.weight.detach() - base_before + assert torch.allclose(actual_delta, expected_delta, atol=1e-6), ( + f"delta mismatch\n max abs diff: " + f"{(actual_delta - expected_delta).abs().max().item()}" + ) + + # B re-zeroed after merge. + assert torch.allclose( + target.lora_B["default"].weight, + torch.zeros_like(target.lora_B["default"].weight), + ), "lora_B should be zero after sleep" + + # Optimizer state wiped. + assert len(opt.state) == 0, "optimizer state must be cleared after sleep" + + # peft's merged_adapters list must be cleared so the next forward + # re-applies the (zero) delta — otherwise future wake steps are no-ops. + assert "default" not in target.merged_adapters, ( + "merged_adapters still contains 'default' — next wake step will be a no-op" + ) + + +def test_post_sleep_wake_actually_updates_lora_params(): + """After sleep, another wake step must move lora_A (guards against B2). + + If sleep_phase leaves peft's ``merged_adapters`` set, peft's forward + skips the LoRA branch and A's gradient is zero — the parameter never + moves. This test catches that regression. + """ + torch = pytest.importorskip("torch") + pytest.importorskip("peft") + pytest.importorskip("torchvision") + + from hal_sleep_wake.config import HyperParams + from hal_sleep_wake.model import build_model + from hal_sleep_wake.sleep import sleep_phase + from hal_sleep_wake.wake import wake_step + + torch.manual_seed(0) + hp = HyperParams(lora_rank=2, lora_alpha=4, device="cpu") + model, lora_params = build_model(hp, num_classes=10) + opt = torch.optim.SGD(lora_params, lr=1e-1) + x = torch.randn(4, 3, 32, 32) + y = torch.randint(0, 10, (4,)) + + # Pre-sleep wake step so lora_B becomes nonzero before we consolidate. + wake_step(model, (x, y), opt, torch.device("cpu")) + sleep_phase(model, optimizer=opt, scale=1.0) + + a_before = { + name: p.detach().clone() + for name, p in model.named_parameters() + if "lora_A" in name + } + wake_step(model, (x, y), opt, torch.device("cpu")) + + moved = any( + not torch.allclose(p.detach(), a_before[name]) + for name, p in model.named_parameters() + if "lora_A" in name + ) + assert moved, "post-sleep wake did not move any lora_A — merged_adapters bug" + + +def test_eval_script_parses_produced_json(tmp_path: Path): + """Round-trip: synthesize a results JSON, then run eval.py against it.""" + results_path = tmp_path / "results.json" + results_path.write_text(json.dumps({ + "avg_accuracy": 0.5, + "forgetting_rate": 0.1, + "accuracy_matrix": [[0.9, 0.0], [0.7, 0.8]], + "sleep_enabled": True, + "hyperparams": {}, + })) + result = subprocess.run( + [sys.executable, str(REPO_ROOT / "scripts" / "eval.py"), str(results_path)], + capture_output=True, + text=True, + timeout=30, + check=False, + ) + assert result.returncode == 0, ( + f"stdout:\n{result.stdout}\nstderr:\n{result.stderr}" + ) + assert "accuracy matrix" in result.stdout + assert "forgetting rate" in result.stdout