nested-learning 0.2.0__py3-none-any.whl

nested_learning/continual_streaming.py

@@ -0,0 +1,283 @@
+from __future__ import annotations
+
+import math
+from dataclasses import dataclass
+from typing import Any, Dict, List, Sequence
+
+import torch
+
+from .continual_classification import ClassificationExample, unique_labels
+from .memorize import MemorizeConfig, memorize_sequence
+from .tokenizer import SentencePieceTokenizer
+
+
+@dataclass(frozen=True)
+class StreamingTask:
+    task_id: int
+    labels: List[str]
+    train: List[ClassificationExample]
+    eval: List[ClassificationExample]
+
+
+@dataclass(frozen=True)
+class ContinualEvalConfig:
+    task_size: int = 10
+    seed: int = 0
+    train_per_label: int = 50
+    eval_per_label: int = 50
+    prompt_template: str = "Text: {text}\nLabel:"
+    label_template: str = "{label}"
+    task_aware: bool = True
+
+
+def _logprob_completion(
+    model,
+    tokenizer: SentencePieceTokenizer,
+    prompt: str,
+    completion: str,
+    device: torch.device,
+    *,
+    fast_state=None,
+) -> float:
+    prompt_ids = tokenizer.encode(prompt, add_bos=True)
+    completion_ids = tokenizer.encode(" " + completion, add_bos=False)
+    tokens = torch.cat([prompt_ids, completion_ids], dim=0).unsqueeze(0).to(device)
+    with torch.no_grad():
+        logits = model(tokens, fast_state=fast_state) if fast_state is not None else model(tokens)
+        log_probs = torch.log_softmax(logits[:, :-1, :], dim=-1)
+        target = tokens[:, 1:]
+        gathered = log_probs.gather(-1, target.unsqueeze(-1)).squeeze(-1)
+        prompt_len = prompt_ids.numel()
+        return float(gathered[0, prompt_len - 1 :].sum().item())
+
+
+def predict_label(
+    model,
+    tokenizer: SentencePieceTokenizer,
+    text: str,
+    candidates: Sequence[str],
+    device: torch.device,
+    *,
+    prompt_template: str,
+    label_template: str,
+    fast_state=None,
+) -> str:
+    if not candidates:
+        raise ValueError("predict_label requires at least one candidate label")
+    prompt = prompt_template.format(text=text)
+    best_label = candidates[0]
+    best_score = -math.inf
+    for label in candidates:
+        label_str = label_template.format(label=label)
+        score = _logprob_completion(
+            model, tokenizer, prompt, label_str, device, fast_state=fast_state
+        )
+        if score > best_score:
+            best_score = score
+            best_label = label
+    return best_label
+
+
+def _balanced_split(
+    examples: Sequence[ClassificationExample],
+    *,
+    labels: Sequence[str],
+    train_per_label: int,
+    eval_per_label: int,
+) -> tuple[List[ClassificationExample], List[ClassificationExample]]:
+    train: List[ClassificationExample] = []
+    eval_: List[ClassificationExample] = []
+    counts_train: Dict[str, int] = {lbl: 0 for lbl in labels}
+    counts_eval: Dict[str, int] = {lbl: 0 for lbl in labels}
+    for ex in examples:
+        lbl = ex.label
+        if lbl not in counts_train:
+            continue
+        if counts_train[lbl] < train_per_label:
+            train.append(ex)
+            counts_train[lbl] += 1
+        elif counts_eval[lbl] < eval_per_label:
+            eval_.append(ex)
+            counts_eval[lbl] += 1
+        if all(v >= train_per_label for v in counts_train.values()) and all(
+            v >= eval_per_label for v in counts_eval.values()
+        ):
+            break
+    return train, eval_
+
+
+def build_streaming_tasks(
+    examples: Sequence[ClassificationExample],
+    *,
+    cfg: ContinualEvalConfig,
+    label_order: Sequence[str] | None = None,
+) -> List[StreamingTask]:
+    labels = list(label_order) if label_order is not None else unique_labels(examples)
+    if label_order is None:
+        import random
+
+        rng = random.Random(cfg.seed)
+        rng.shuffle(labels)
+    if cfg.task_size <= 0:
+        raise ValueError("task_size must be positive")
+    tasks: List[StreamingTask] = []
+    for task_id, start in enumerate(range(0, len(labels), cfg.task_size)):
+        task_labels = labels[start : start + cfg.task_size]
+        if not task_labels:
+            break
+        task_examples = [ex for ex in examples if ex.label in set(task_labels)]
+        train, eval_ = _balanced_split(
+            task_examples,
+            labels=task_labels,
+            train_per_label=cfg.train_per_label,
+            eval_per_label=cfg.eval_per_label,
+        )
+        tasks.append(
+            StreamingTask(task_id=task_id, labels=list(task_labels), train=train, eval=eval_)
+        )
+    return tasks
+
+
+@dataclass(frozen=True)
+class ContinualEvalResult:
+    task_accuracy_matrix: List[List[float]]
+    per_task_forgetting: List[float]
+    avg_accuracy_final: float
+    avg_forgetting: float
+
+
+def evaluate_continual_classification(
+    model,
+    tokenizer: SentencePieceTokenizer,
+    tasks: Sequence[StreamingTask],
+    device: torch.device,
+    *,
+    cfg: ContinualEvalConfig,
+    memorize_cfg: MemorizeConfig,
+) -> tuple[ContinualEvalResult, Dict[str, Any]]:
+    """
+    Streaming class-incremental evaluation using generative classification + optional
+    test-time memorization.
+
+    - If `memorize_cfg.enabled`, each training example is memorized by appending the correct
+      label string.
+    - Accuracy is computed after each task on each task's eval set, producing a task-accuracy
+      matrix.
+    """
+    meta_snapshot: Dict[str, torch.Tensor] | None = None
+    if memorize_cfg.enabled and (not memorize_cfg.use_fast_state) and memorize_cfg.reset:
+        from .memorize import snapshot_state_dict  # local import to avoid cycles
+
+        meta_snapshot = snapshot_state_dict(model)
+
+    fast_state = None
+    if memorize_cfg.enabled and memorize_cfg.use_fast_state:
+        if not hasattr(model, "init_fast_state"):
+            raise RuntimeError("Model does not support fast state memorization")
+        fast_state = model.init_fast_state()
+
+    task_acc: List[List[float]] = [[float("nan") for _ in tasks] for _ in tasks]
+    best_acc: List[float] = [0.0 for _ in tasks]
+
+    memorize_stats_total: Dict[str, float] = {}
+
+    def _eval_task(task_idx: int) -> float:
+        task = tasks[task_idx]
+        candidates = (
+            task.labels
+            if cfg.task_aware
+            else [lbl for t in tasks[: current_task + 1] for lbl in t.labels]
+        )
+        if not task.eval:
+            return float("nan")
+        correct = 0
+        for ex in task.eval:
+            pred = predict_label(
+                model,
+                tokenizer,
+                ex.text,
+                candidates,
+                device,
+                prompt_template=cfg.prompt_template,
+                label_template=cfg.label_template,
+                fast_state=fast_state,
+            )
+            correct += int(pred == ex.label)
+        return correct / len(task.eval) if task.eval else float("nan")
+
+    for current_task, task in enumerate(tasks):
+        # Online "training" on this task's examples via optional memorization.
+        for ex in task.train:
+            candidates = (
+                task.labels
+                if cfg.task_aware
+                else [lbl for t in tasks[: current_task + 1] for lbl in t.labels]
+            )
+            _ = predict_label(
+                model,
+                tokenizer,
+                ex.text,
+                candidates,
+                device,
+                prompt_template=cfg.prompt_template,
+                label_template=cfg.label_template,
+                fast_state=fast_state,
+            )
+            if memorize_cfg.enabled:
+                prompt = cfg.prompt_template.format(text=ex.text)
+                target = cfg.label_template.format(label=ex.label)
+                memorize_text = f"{prompt} {target}"
+                if memorize_cfg.use_fast_state and memorize_cfg.reset:
+                    fast_state = model.init_fast_state()
+                stats = memorize_sequence(
+                    model, tokenizer, memorize_text, device, memorize_cfg, fast_state=fast_state
+                )
+                for k, v in stats.items():
+                    memorize_stats_total[k] = memorize_stats_total.get(k, 0.0) + v
+                if (
+                    (not memorize_cfg.use_fast_state)
+                    and memorize_cfg.reset
+                    and meta_snapshot is not None
+                ):
+                    from .memorize import restore_state_dict  # local import to avoid cycles
+
+                    restore_state_dict(model, meta_snapshot)
+
+        # Evaluate on all tasks seen so far.
+        for task_idx in range(current_task + 1):
+            acc = _eval_task(task_idx)
+            task_acc[task_idx][current_task] = acc
+            if not math.isnan(acc):
+                best_acc[task_idx] = max(best_acc[task_idx], acc)
+
+    final_accs = [task_acc[i][-1] for i in range(len(tasks)) if not math.isnan(task_acc[i][-1])]
+    avg_accuracy_final = sum(final_accs) / len(final_accs) if final_accs else float("nan")
+
+    per_task_forgetting: List[float] = []
+    for i in range(len(tasks)):
+        last = task_acc[i][-1]
+        if math.isnan(last):
+            per_task_forgetting.append(float("nan"))
+            continue
+        per_task_forgetting.append(best_acc[i] - last)
+    valid_forgetting = [f for f in per_task_forgetting if not math.isnan(f)]
+    avg_forgetting = (
+        sum(valid_forgetting) / len(valid_forgetting) if valid_forgetting else float("nan")
+    )
+
+    result = ContinualEvalResult(
+        task_accuracy_matrix=task_acc,
+        per_task_forgetting=per_task_forgetting,
+        avg_accuracy_final=avg_accuracy_final,
+        avg_forgetting=avg_forgetting,
+    )
+    meta = {
+        "task_size": cfg.task_size,
+        "train_per_label": cfg.train_per_label,
+        "eval_per_label": cfg.eval_per_label,
+        "task_aware": cfg.task_aware,
+        "prompt_template": cfg.prompt_template,
+        "label_template": cfg.label_template,
+        "memorize_stats": memorize_stats_total,
+    }
+    return result, meta

nested_learning/data.py

@@ -0,0 +1,153 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Iterator, List, Sequence
+
+import numpy as np
+import torch
+from torch.utils.data import Dataset, IterableDataset, get_worker_info
+
+
+@dataclass
+class SyntheticTextConfig:
+    vocab_size: int
+    seq_len: int
+    dataset_size: int
+
+
+class SyntheticTextDataset(Dataset[torch.Tensor]):
+    def __init__(self, config: SyntheticTextConfig):
+        self.config = config
+
+    def __len__(self) -> int:
+        return self.config.dataset_size
+
+    def __getitem__(self, idx: int) -> torch.Tensor:
+        g = torch.Generator().manual_seed(idx)
+        return torch.randint(0, self.config.vocab_size, (self.config.seq_len,), generator=g)
+
+
+class TokenShardDataset(Dataset[torch.Tensor]):
+    """Memory-mapped dataset over NumPy shards produced by shard_corpus.py."""
+
+    def __init__(self, shard_dir: str | Path):
+        self.shard_dir = Path(shard_dir)
+        if not self.shard_dir.exists():
+            msg = f"Shard directory {self.shard_dir} does not exist"
+            raise FileNotFoundError(msg)
+        self.paths = sorted(self.shard_dir.glob("*.npy"))
+        if not self.paths:
+            msg = f"No shard files found in {self.shard_dir}"
+            raise ValueError(msg)
+        self.metadata: List[tuple[int, int]] = []
+        self._cache: dict[int, np.memmap] = {}
+        total = 0
+        for idx, path in enumerate(self.paths):
+            arr = np.load(path, mmap_mode="r")
+            length = arr.shape[0]
+            self.metadata.append((total, length))
+            total += length
+        self.total_sequences = total
+
+    def __len__(self) -> int:
+        return self.total_sequences
+
+    def _load_array(self, shard_idx: int) -> np.memmap:
+        if shard_idx not in self._cache:
+            self._cache[shard_idx] = np.load(self.paths[shard_idx], mmap_mode="r")
+        return self._cache[shard_idx]
+
+    def __getitem__(self, idx: int) -> torch.Tensor:
+        if idx < 0 or idx >= self.total_sequences:
+            raise IndexError(idx)
+        shard_idx = self._find_shard(idx)
+        start_offset = self.metadata[shard_idx][0]
+        arr = self._load_array(shard_idx)
+        local_idx = idx - start_offset
+        tokens = torch.from_numpy(arr[local_idx])
+        return tokens.long()
+
+    def _find_shard(self, idx: int) -> int:
+        lo, hi = 0, len(self.metadata) - 1
+        while lo <= hi:
+            mid = (lo + hi) // 2
+            start, length = self.metadata[mid]
+            if idx < start:
+                hi = mid - 1
+            elif idx >= start + length:
+                lo = mid + 1
+            else:
+                return mid
+        return len(self.metadata) - 1
+
+
+@dataclass
+class ShardSourceConfig:
+    name: str
+    shards_dir: str
+    weight: float
+
+
+class ShardSource:
+    def __init__(self, config: ShardSourceConfig):
+        self.name = config.name
+        self.weight = config.weight
+        self.dir = Path(config.shards_dir)
+        if not self.dir.exists():
+            msg = f"Shard directory {self.dir} missing for source {self.name}"
+            raise FileNotFoundError(msg)
+        self.paths = sorted(self.dir.glob("*.npy"))
+        if not self.paths:
+            raise ValueError(f"No shard files in {self.dir}")
+        self._cache: dict[Path, np.memmap] = {}
+
+    def sample(self, rng: np.random.Generator) -> np.ndarray:
+        shard_path = self.paths[rng.integers(0, len(self.paths))]
+        if shard_path not in self._cache:
+            self._cache[shard_path] = np.load(shard_path, mmap_mode="r")
+        shard = self._cache[shard_path]
+        idx = rng.integers(0, shard.shape[0])
+        return shard[idx]
+
+
+class MixtureShardDataset(IterableDataset[torch.Tensor]):
+    def __init__(
+        self,
+        sources: Sequence[ShardSourceConfig],
+        *,
+        samples_per_epoch: int,
+        seed: int = 0,
+    ):
+        super().__init__()
+        self.sources = [ShardSource(cfg) for cfg in sources]
+        total_weight = sum(max(src.weight, 0.0) for src in self.sources)
+        if total_weight <= 0:
+            raise ValueError("Mixture weights must sum to > 0")
+        self.weights = np.array([max(src.weight, 0.0) / total_weight for src in self.sources])
+        self.samples_per_epoch = samples_per_epoch
+        self.seed = seed
+
+    def __len__(self) -> int:
+        return self.samples_per_epoch
+
+    def __iter__(self) -> Iterator[torch.Tensor]:
+        worker = get_worker_info()
+        if worker is None:
+            start = 0
+            end = self.samples_per_epoch
+            worker_seed = self.seed
+        else:
+            per_worker = (self.samples_per_epoch + worker.num_workers - 1) // worker.num_workers
+            start = worker.id * per_worker
+            end = min(start + per_worker, self.samples_per_epoch)
+            worker_seed = self.seed + worker.id
+        rng = np.random.default_rng(worker_seed)
+        for _ in range(start, end):
+            idx = rng.choice(len(self.sources), p=self.weights)
+            sample = np.array(self.sources[idx].sample(rng), copy=True)
+            yield torch.from_numpy(sample).long()
+
+
+def collate_batch(batch: list[torch.Tensor]) -> torch.Tensor:
+    return torch.stack(batch, dim=0)

nested_learning/device.py

@@ -0,0 +1,21 @@
+from __future__ import annotations
+
+import torch
+
+
+def resolve_device(device_str: str) -> torch.device:
+    normalized = str(device_str).strip().lower()
+    if normalized.startswith("cuda"):
+        if not torch.cuda.is_available():
+            return torch.device("cpu")
+        parts = normalized.split(":")
+        idx = int(parts[1]) if len(parts) > 1 else 0
+        if idx >= torch.cuda.device_count():
+            idx = max(torch.cuda.device_count() - 1, 0)
+        return torch.device(f"cuda:{idx}")
+    if normalized.startswith("mps"):
+        if not (hasattr(torch.backends, "mps") and torch.backends.mps.is_available()):
+            return torch.device("cpu")
+        return torch.device("mps")
+    return torch.device(device_str)
+

nested_learning/eval_state.py

@@ -0,0 +1,72 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+
+import torch
+
+
+@dataclass
+class EvalStreamingState:
+    fast_state: object | None = None
+    attention_cache: object | None = None
+
+
+def parse_eval_state_mode(mode: str) -> bool:
+    """
+    Returns True when eval state should be carried across samples.
+    """
+    normalized = str(mode).strip().lower()
+    if normalized in {"reset", "reset_per_sample", "isolated"}:
+        return False
+    if normalized in {"carry", "carry_across_samples", "stream"}:
+        return True
+    raise ValueError(
+        "Unsupported eval_state_mode={!r}; expected one of "
+        "['reset_per_sample', 'carry_across_samples']".format(mode)
+    )
+
+
+def init_eval_streaming_state(
+    model,
+    *,
+    use_fast_state: bool,
+    use_attention_cache: bool,
+) -> EvalStreamingState:
+    state = EvalStreamingState()
+    if use_fast_state:
+        init_fast_state = getattr(model, "init_fast_state", None)
+        if not callable(init_fast_state):
+            raise RuntimeError(
+                "Requested fast-state eval mode, but model.init_fast_state() is missing"
+            )
+        state.fast_state = init_fast_state()
+    if use_attention_cache:
+        init_attention_cache = getattr(model, "init_attention_cache", None)
+        if not callable(init_attention_cache):
+            raise RuntimeError(
+                "Requested attention-cache eval mode, but model.init_attention_cache() is missing"
+            )
+        state.attention_cache = init_attention_cache()
+    return state
+
+
+def forward_with_eval_state(
+    model,
+    tokens: torch.Tensor,
+    *,
+    state: EvalStreamingState | None = None,
+) -> tuple[torch.Tensor, EvalStreamingState | None]:
+    if state is None:
+        return model(tokens), None
+    if state.attention_cache is not None:
+        logits, next_cache = model(
+            tokens,
+            fast_state=state.fast_state,
+            attention_cache=state.attention_cache,
+            return_attention_cache=True,
+        )
+        state.attention_cache = next_cache
+        return logits, state
+    if state.fast_state is not None:
+        return model(tokens, fast_state=state.fast_state), state
+    return model(tokens), state

nested_learning/fast_state.py

@@ -0,0 +1,108 @@
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+from typing import Dict, cast
+
+import torch
+from torch import nn
+
+from .optim.manager import LevelConfig, LevelOptimizerManager
+from .titan.self_modifying import SelfModifyingTitansState
+
+ParamDict = Dict[str, torch.Tensor]
+
+
+@dataclass
+class CMSChunkBuffer:
+    """
+    Streaming CMS chunk buffer persisted across multiple model calls.
+
+    This is required to preserve update-period cadence when a logical sequence is
+    processed in several chunked forward/update calls.
+    """
+
+    inputs: list[torch.Tensor] = field(default_factory=list)
+    teach: list[torch.Tensor] = field(default_factory=list)
+    active: list[torch.Tensor] = field(default_factory=list)
+    count: int = 0
+
+
+def init_module_deltas(module: nn.Module) -> ParamDict:
+    """
+    Initialize a per-parameter "fast state" delta dict for meta+delta fast state.
+
+    The fast state stores *deltas* (initialized to 0) rather than detached parameter clones so that
+    forward passes can use `meta_param + delta`, allowing outer gradients to flow to meta params
+    while keeping online updates as stop-grad writes into the delta tensors.
+    """
+
+    return {name: torch.zeros_like(param).detach() for name, param in module.named_parameters()}
+
+
+@dataclass
+class BlockFastState:
+    titan_params: ParamDict | None
+    cms_params: Dict[str, ParamDict]
+    cms_online_buffers: Dict[str, CMSChunkBuffer]
+    level_manager: LevelOptimizerManager
+    selfmod_state: SelfModifyingTitansState | None = None
+
+
+def build_block_fast_state(
+    *,
+    titan_module: nn.Module | None,
+    cms_blocks: Dict[str, nn.Module],
+    selfmod_module: nn.Module | None = None,
+    specs,
+    optimizer_configs: Dict[str, dict],
+    default_lr: float,
+) -> BlockFastState:
+    titan_params = None
+    if titan_module is not None:
+        titan_params = init_module_deltas(titan_module)
+    cms_params = {name: init_module_deltas(block) for name, block in cms_blocks.items()}
+    cms_online_buffers = {name: CMSChunkBuffer() for name in cms_blocks}
+    level_cfg = LevelConfig(specs=specs, optimizer_configs=optimizer_configs, default_lr=default_lr)
+    level_manager = LevelOptimizerManager(level_cfg)
+    selfmod_state = None
+    if selfmod_module is not None:
+        init_fn = getattr(selfmod_module, "init_fast_state", None)
+        if callable(init_fn):
+            selfmod_state = cast(SelfModifyingTitansState, init_fn())
+    return BlockFastState(
+        titan_params=titan_params,
+        cms_params=cms_params,
+        cms_online_buffers=cms_online_buffers,
+        level_manager=level_manager,
+        selfmod_state=selfmod_state,
+    )
+
+
+@dataclass
+class ModelFastState:
+    blocks: list[BlockFastState]
+
+
+@dataclass
+class AttentionKVCache:
+    """
+    Per-layer autoregressive attention cache.
+
+    Shapes:
+    - key:   [batch, heads, cached_tokens, head_dim]
+    - value: [batch, heads, cached_tokens, head_dim]
+    """
+
+    key: torch.Tensor
+    value: torch.Tensor
+
+
+@dataclass
+class ModelAttentionCache:
+    """
+    Model-level container for per-block attention caches.
+
+    Blocks without attention store `None` entries.
+    """
+
+    blocks: list[AttentionKVCache | None]

nested_learning/functional.py

@@ -0,0 +1,69 @@
+from __future__ import annotations
+
+from typing import Any, Dict, Mapping, Tuple
+
+import torch
+from torch import nn
+from torch.func import functional_call
+
+ParamDict = Dict[str, torch.Tensor]
+
+
+def params_with_deltas(module: nn.Module, deltas: ParamDict) -> ParamDict:
+    params: ParamDict = {}
+    missing: list[str] = []
+    for name, param in module.named_parameters():
+        delta = deltas.get(name)
+        if delta is None:
+            missing.append(name)
+            continue
+        params[name] = param + delta
+    if missing:
+        raise KeyError(
+            f"Missing fast-state delta(s) for {module.__class__.__name__}: {sorted(missing)[:10]}"
+        )
+    return params
+
+
+def module_buffers(module: nn.Module) -> ParamDict:
+    return {name: buf for name, buf in module.named_buffers()}
+
+
+def call_with_params(
+    module: nn.Module,
+    params: ParamDict,
+    *args: Any,
+    **kwargs: Any,
+) -> Any:
+    buffers = module_buffers(module)
+    return functional_call(module, (params, buffers), args, kwargs, strict=True)
+
+
+def call_with_deltas(
+    module: nn.Module,
+    deltas: ParamDict,
+    *args: Any,
+    **kwargs: Any,
+) -> Any:
+    return call_with_params(module, params_with_deltas(module, deltas), *args, **kwargs)
+
+
+def require_grad_params(
+    params: Mapping[str, torch.Tensor], *, detach: bool = True
+) -> ParamDict:
+    out: ParamDict = {}
+    for name, value in params.items():
+        if detach:
+            out[name] = value.detach().requires_grad_(True)
+        else:
+            out[name] = value.requires_grad_(True)
+    return out
+
+
+def grads_to_dict(params: ParamDict, grads: Tuple[torch.Tensor | None, ...]) -> ParamDict:
+    out: ParamDict = {}
+    for (name, _), grad in zip(params.items(), grads, strict=True):
+        if grad is None:
+            continue
+        out[name] = grad
+    return out