nested-learning 0.2.0__py3-none-any.whl

nested_learning/__init__.py

@@ -0,0 +1,12 @@
+"""Nested Learning (HOPE) reproduction package."""
+
+from importlib.metadata import PackageNotFoundError, version
+
+from .levels import LevelClock, LevelSpec  # noqa: F401
+
+try:
+    __version__ = version("nested-learning")
+except PackageNotFoundError:  # pragma: no cover - editable/local source tree
+    __version__ = "0.2.0"
+
+__all__ = ["LevelClock", "LevelSpec", "__version__"]

nested_learning/__main__.py

@@ -0,0 +1,12 @@
+from __future__ import annotations
+
+from .cli import app
+
+
+def main() -> None:
+    app()
+
+
+if __name__ == "__main__":
+    main()
+

nested_learning/assoc_memory.py

@@ -0,0 +1,23 @@
+from __future__ import annotations
+
+from typing import Protocol
+
+import torch
+import torch.nn as nn
+
+
+class AssocMemory(nn.Module):
+    """Base class for associative memories with explicit update hooks."""
+
+    def forward(self, query: torch.Tensor) -> torch.Tensor:  # type: ignore[override]
+        raise NotImplementedError
+
+    @torch.no_grad()
+    def update(
+        self, *, key: torch.Tensor, value: torch.Tensor, error_signal: torch.Tensor | None = None
+    ) -> None:
+        raise NotImplementedError
+
+
+class SupportsReset(Protocol):
+    def reset_state(self) -> None: ...

nested_learning/backbones.py

@@ -0,0 +1,147 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from .fast_state import AttentionKVCache
+
+
+@dataclass
+class AttentionConfig:
+    dim: int
+    heads: int
+    dropout: float = 0.0
+    use_flash: bool = True
+    causal: bool = True
+    qk_l2_norm: bool = False
+    qk_norm_eps: float = 1e-6
+    local_conv_window: int | None = None
+
+
+class SelfAttention(nn.Module):
+    def __init__(self, config: AttentionConfig):
+        super().__init__()
+        if config.dim % config.heads != 0:
+            msg = f"dim must be divisible by heads (got dim={config.dim}, heads={config.heads})"
+            raise ValueError(msg)
+        self.config = config
+        self.heads = config.heads
+        self.head_dim = config.dim // config.heads
+        self.qkv = nn.Linear(config.dim, config.dim * 3, bias=False)
+        self.out_proj = nn.Linear(config.dim, config.dim, bias=False)
+        self.resid_dropout = nn.Dropout(config.dropout)
+        self.norm = nn.LayerNorm(config.dim)
+        self.local_conv: nn.Conv1d | None = None
+        if config.local_conv_window is not None:
+            window = int(config.local_conv_window)
+            if window <= 0:
+                raise ValueError("local_conv_window must be positive")
+            self.local_conv = nn.Conv1d(
+                config.dim,
+                config.dim,
+                kernel_size=window,
+                groups=config.dim,
+                padding=0,
+                bias=False,
+            )
+
+    def forward(  # type: ignore[override]
+        self,
+        x: torch.Tensor,
+        *,
+        kv_cache: AttentionKVCache | None = None,
+        return_kv_cache: bool = False,
+    ) -> torch.Tensor | tuple[torch.Tensor, AttentionKVCache]:
+        residual = x
+        attn_inp = x
+        if kv_cache is not None and self.local_conv is not None:
+            raise RuntimeError(
+                "kv_cache with local_conv_window is not supported in this implementation."
+            )
+        if self.local_conv is not None:
+            kernel = self.local_conv.kernel_size[0]
+            attn_inp = attn_inp.transpose(1, 2)
+            # Causal depthwise conv: only attends to past tokens.
+            attn_inp = F.pad(attn_inp, (kernel - 1, 0))
+            attn_inp = self.local_conv(attn_inp).transpose(1, 2)
+        q, k, v = self._compute_qkv(attn_inp)
+        past_len = 0
+        k_all = k
+        v_all = v
+        if kv_cache is not None:
+            if kv_cache.key.size(0) != k.size(0):
+                raise ValueError("kv_cache batch dimension must match input batch dimension")
+            if kv_cache.key.size(1) != k.size(1) or kv_cache.key.size(-1) != k.size(-1):
+                raise ValueError("kv_cache shape is incompatible with attention heads/head_dim")
+            past_len = int(kv_cache.key.size(2))
+            k_all = torch.cat([kv_cache.key, k], dim=2)
+            v_all = torch.cat([kv_cache.value, v], dim=2)
+        attn_output = self._scaled_dot_product_attn(q, k_all, v_all, past_len=past_len)
+        attn_output = attn_output.transpose(1, 2).contiguous().view(x.size(0), x.size(1), -1)
+        attn_output = self.out_proj(attn_output)
+        attn_output = self.resid_dropout(attn_output)
+        out = self.norm(residual + attn_output)
+        if return_kv_cache:
+            return out, AttentionKVCache(key=k_all.detach(), value=v_all.detach())
+        return out
+
+    def _compute_qkv(self, x: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+        qkv = self.qkv(x)
+        q, k, v = qkv.chunk(3, dim=-1)
+        shape = (x.size(0), x.size(1), self.heads, self.head_dim)
+        q = q.view(*shape).transpose(1, 2)
+        k = k.view(*shape).transpose(1, 2)
+        v = v.view(*shape).transpose(1, 2)
+        if self.config.qk_l2_norm:
+            q = F.normalize(q, dim=-1, eps=self.config.qk_norm_eps)
+            k = F.normalize(k, dim=-1, eps=self.config.qk_norm_eps)
+        return q, k, v
+
+    def _scaled_dot_product_attn(
+        self,
+        q: torch.Tensor,
+        k: torch.Tensor,
+        v: torch.Tensor,
+        *,
+        past_len: int = 0,
+    ) -> torch.Tensor:
+        dropout_p = self.config.dropout if self.training else 0.0
+        attn_mask = None
+        if self.config.causal and past_len > 0:
+            query_len = int(q.size(-2))
+            key_len = int(k.size(-2))
+            key_positions = torch.arange(key_len, device=q.device)
+            query_positions = past_len + torch.arange(query_len, device=q.device)
+            attn_mask = key_positions.unsqueeze(0) <= query_positions.unsqueeze(1)
+        is_causal = self.config.causal and attn_mask is None
+        device_type = q.device.type
+        if (
+            device_type == "cuda"
+            and torch.cuda.is_available()
+            and hasattr(torch.backends, "cuda")
+            and hasattr(torch.backends.cuda, "sdp_kernel")
+        ):
+            with torch.backends.cuda.sdp_kernel(  # type: ignore[attr-defined]
+                enable_flash=self.config.use_flash,
+                enable_mem_efficient=True,
+                enable_math=not self.config.use_flash,
+            ):
+                return F.scaled_dot_product_attention(
+                    q,
+                    k,
+                    v,
+                    attn_mask=attn_mask,
+                    dropout_p=dropout_p,
+                    is_causal=is_causal,
+                )
+        return F.scaled_dot_product_attention(
+            q,
+            k,
+            v,
+            attn_mask=attn_mask,
+            dropout_p=dropout_p,
+            is_causal=is_causal,
+        )

nested_learning/capabilities.py

@@ -0,0 +1,104 @@
+from __future__ import annotations
+
+import platform
+import sys
+from dataclasses import asdict, dataclass, field
+from typing import Any
+
+import torch
+
+
+@dataclass
+class RuntimeCapabilities:
+    python_version: str
+    platform: str
+    machine: str
+    torch_version: str
+    cuda_available: bool
+    cuda_device_count: int
+    cuda_devices: list[str] = field(default_factory=list)
+    mps_available: bool = False
+    mps_built: bool = False
+    distributed_available: bool = False
+    compile_available: bool = False
+    sdpa_flash_available: bool = False
+    sdpa_mem_efficient_available: bool = False
+    sdpa_math_available: bool = True
+    bf16_supported: bool = False
+    fp16_supported: bool = False
+    default_device: str = "cpu"
+    warnings: list[str] = field(default_factory=list)
+
+    def to_dict(self) -> dict[str, Any]:
+        return asdict(self)
+
+
+def collect_runtime_capabilities() -> RuntimeCapabilities:
+    cuda_available = bool(torch.cuda.is_available())
+    cuda_device_count = int(torch.cuda.device_count() if cuda_available else 0)
+    cuda_devices: list[str] = []
+    warnings: list[str] = []
+
+    if cuda_available:
+        for idx in range(cuda_device_count):
+            try:
+                name = torch.cuda.get_device_name(idx)
+                cuda_devices.append(f"cuda:{idx} {name}")
+            except Exception as err:  # pragma: no cover - backend specific
+                warnings.append(f"failed to query cuda:{idx}: {err}")
+
+    mps_backend = getattr(torch.backends, "mps", None)
+    mps_available = bool(mps_backend and mps_backend.is_available())
+    mps_built = bool(mps_backend and mps_backend.is_built())
+
+    distributed_available = bool(torch.distributed.is_available())
+    compile_available = bool(hasattr(torch, "compile"))
+
+    flash_enabled = False
+    mem_eff_enabled = False
+    math_enabled = True
+    if hasattr(torch.backends, "cuda") and torch.backends.cuda.is_built():
+        try:
+            flash_enabled = bool(torch.backends.cuda.flash_sdp_enabled())
+            mem_eff_enabled = bool(torch.backends.cuda.mem_efficient_sdp_enabled())
+            math_enabled = bool(torch.backends.cuda.math_sdp_enabled())
+        except Exception as err:  # pragma: no cover - backend specific
+            warnings.append(f"failed to query SDPA backend flags: {err}")
+
+    bf16_supported = False
+    fp16_supported = False
+    if cuda_available:
+        try:
+            bf16_supported = bool(torch.cuda.is_bf16_supported())
+            fp16_supported = True
+        except Exception as err:  # pragma: no cover
+            warnings.append(f"failed to query CUDA dtype support: {err}")
+    elif mps_available:
+        fp16_supported = True
+
+    default_device = "cpu"
+    if cuda_available:
+        default_device = "cuda:0"
+    elif mps_available:
+        default_device = "mps"
+
+    return RuntimeCapabilities(
+        python_version=sys.version.split()[0],
+        platform=platform.platform(),
+        machine=platform.machine(),
+        torch_version=torch.__version__,
+        cuda_available=cuda_available,
+        cuda_device_count=cuda_device_count,
+        cuda_devices=cuda_devices,
+        mps_available=mps_available,
+        mps_built=mps_built,
+        distributed_available=distributed_available,
+        compile_available=compile_available,
+        sdpa_flash_available=flash_enabled,
+        sdpa_mem_efficient_available=mem_eff_enabled,
+        sdpa_math_available=math_enabled,
+        bf16_supported=bf16_supported,
+        fp16_supported=fp16_supported,
+        default_device=default_device,
+        warnings=warnings,
+    )

nested_learning/cli.py

@@ -0,0 +1,253 @@
+from __future__ import annotations
+
+import json
+from pathlib import Path
+from typing import Annotated
+
+import torch
+import typer
+from omegaconf import OmegaConf
+
+from .capabilities import collect_runtime_capabilities
+from .config_utils import compose_config
+from .device import resolve_device
+from .training import build_model_from_cfg
+
+app = typer.Typer(
+    add_completion=False,
+    no_args_is_help=True,
+    help="Nested Learning CLI (training, diagnostics, and smoke checks).",
+)
+
+
+def _resolve_cli_device(device: str) -> torch.device:
+    lowered = device.strip().lower()
+    if lowered == "auto":
+        caps = collect_runtime_capabilities()
+        return resolve_device(caps.default_device)
+    return resolve_device(device)
+
+
+@app.command("doctor")
+def doctor(
+    as_json: Annotated[
+        bool,
+        typer.Option("--json", help="Emit machine-readable JSON only."),
+    ] = False,
+    output: Annotated[
+        Path | None,
+        typer.Option(
+            "--output",
+            "-o",
+            help="Optional path for writing doctor output JSON.",
+            dir_okay=False,
+            writable=True,
+        ),
+    ] = None,
+) -> None:
+    """Inspect runtime capabilities for backend/device compatibility."""
+    payload = collect_runtime_capabilities().to_dict()
+    rendered = json.dumps(payload, indent=2, sort_keys=True)
+    if output is not None:
+        output.parent.mkdir(parents=True, exist_ok=True)
+        output.write_text(rendered + "\n", encoding="utf-8")
+    if as_json:
+        typer.echo(rendered)
+        return
+
+    typer.echo("Runtime Doctor")
+    typer.echo(f"python: {payload['python_version']}")
+    typer.echo(f"platform: {payload['platform']} ({payload['machine']})")
+    typer.echo(f"torch: {payload['torch_version']}")
+    typer.echo(f"default_device: {payload['default_device']}")
+    typer.echo(
+        "cuda_available: {available} ({count} device(s))".format(
+            available=payload["cuda_available"],
+            count=payload["cuda_device_count"],
+        )
+    )
+    for name in payload["cuda_devices"]:
+        typer.echo(f"  - {name}")
+    typer.echo(f"mps_available: {payload['mps_available']} (built={payload['mps_built']})")
+    typer.echo(f"distributed_available: {payload['distributed_available']}")
+    typer.echo(f"compile_available: {payload['compile_available']}")
+    typer.echo(
+        "sdpa backends: flash={flash} mem_efficient={mem} math={math}".format(
+            flash=payload["sdpa_flash_available"],
+            mem=payload["sdpa_mem_efficient_available"],
+            math=payload["sdpa_math_available"],
+        )
+    )
+    typer.echo(f"dtype support: bf16={payload['bf16_supported']} fp16={payload['fp16_supported']}")
+    if payload["warnings"]:
+        typer.echo("warnings:")
+        for warning in payload["warnings"]:
+            typer.echo(f"  - {warning}")
+
+
+@app.command("smoke")
+def smoke(
+    config_name: Annotated[
+        str,
+        typer.Option("--config-name", "-c", help="Hydra config name (e.g. pilot, hope/mid)."),
+    ] = "pilot_smoke",
+    override: Annotated[
+        list[str] | None,
+        typer.Option(
+            "--override",
+            "-O",
+            help="Hydra override(s), may be passed multiple times.",
+        ),
+    ] = None,
+    config_dir: Annotated[
+        Path | None,
+        typer.Option(
+            "--config-dir",
+            help="Optional explicit config directory.",
+            exists=True,
+            file_okay=False,
+            dir_okay=True,
+            readable=True,
+        ),
+    ] = None,
+    device: Annotated[
+        str,
+        typer.Option(
+            "--device",
+            help="Device string for smoke pass (cpu, cuda:0, mps, auto).",
+        ),
+    ] = "cpu",
+    batch_size: Annotated[
+        int,
+        typer.Option("--batch-size", help="Synthetic smoke batch size."),
+    ] = 1,
+    seq_len: Annotated[
+        int,
+        typer.Option("--seq-len", help="Synthetic smoke sequence length."),
+    ] = 32,
+) -> None:
+    """Run a lightweight forward-pass smoke test with composed config."""
+    cfg = compose_config(config_name, overrides=override or [], config_dir=config_dir)
+    model_cfg = cfg.model
+    torch_device = _resolve_cli_device(device)
+    model = build_model_from_cfg(model_cfg).to(torch_device)
+    model.eval()
+
+    vocab_size = int(model_cfg.vocab_size)
+    tokens = torch.randint(0, vocab_size, (batch_size, seq_len), device=torch_device)
+    with torch.no_grad():
+        outputs = model(tokens)
+    if isinstance(outputs, tuple):
+        logits = outputs[0]
+    else:
+        logits = outputs
+    typer.echo(
+        json.dumps(
+            {
+                "status": "ok",
+                "config_name": config_name,
+                "device": str(torch_device),
+                "batch_size": batch_size,
+                "seq_len": seq_len,
+                "logits_shape": list(logits.shape),
+                "dtype": str(logits.dtype),
+            },
+            sort_keys=True,
+        )
+    )
+
+
+@app.command("train")
+def train(
+    config_name: Annotated[
+        str,
+        typer.Option("--config-name", "-c", help="Hydra config name for training."),
+    ] = "pilot",
+    override: Annotated[
+        list[str] | None,
+        typer.Option("--override", "-O", help="Hydra override(s), may be passed multiple times."),
+    ] = None,
+    config_dir: Annotated[
+        Path | None,
+        typer.Option(
+            "--config-dir",
+            help="Optional explicit config directory.",
+            exists=True,
+            file_okay=False,
+            dir_okay=True,
+            readable=True,
+        ),
+    ] = None,
+    device: Annotated[
+        str | None,
+        typer.Option(
+            "--device",
+            help="Override cfg.train.device (e.g. cpu, cuda:1, auto).",
+        ),
+    ] = None,
+    dry_run: Annotated[
+        bool,
+        typer.Option("--dry-run", help="Print resolved config and exit."),
+    ] = False,
+) -> None:
+    """Launch a local (single-process) training loop."""
+    from .training import run_training_loop
+
+    cfg = compose_config(config_name, overrides=override or [], config_dir=config_dir)
+    if device is not None:
+        cfg.train.device = device
+    if dry_run:
+        typer.echo(OmegaConf.to_yaml(cfg))
+        return
+    train_device = _resolve_cli_device(str(cfg.train.device))
+    run_training_loop(cfg, device=train_device, distributed=False, dist_ctx=None)
+
+
+@app.command("audit")
+def audit(
+    config_name: Annotated[
+        str,
+        typer.Option("--config-name", "-c", help="Hydra config name to audit."),
+    ] = "pilot_paper_faithful",
+    override: Annotated[
+        list[str] | None,
+        typer.Option("--override", "-O", help="Hydra override(s), may be passed multiple times."),
+    ] = None,
+    config_dir: Annotated[
+        Path | None,
+        typer.Option(
+            "--config-dir",
+            help="Optional explicit config directory.",
+            exists=True,
+            file_okay=False,
+            dir_okay=True,
+            readable=True,
+        ),
+    ] = None,
+) -> None:
+    """Run static architecture checks on a composed config."""
+    cfg = compose_config(config_name, overrides=override or [], config_dir=config_dir)
+    model = build_model_from_cfg(cfg.model)
+    has_embed = hasattr(model, "embed")
+    has_lm_head = hasattr(model, "lm_head")
+    tied_weights = False
+    if has_embed and has_lm_head:
+        embed = getattr(model, "embed")
+        lm_head = getattr(model, "lm_head")
+        tied_weights = bool(embed.weight.data_ptr() == lm_head.weight.data_ptr())
+
+    report = {
+        "status": "ok",
+        "config_name": config_name,
+        "model_type": str(cfg.model.get("type", "hope")),
+        "block_variant": str(cfg.model.get("block_variant", "hope_hybrid")),
+        "surprise_metric": str(cfg.model.get("surprise_metric", "l2")),
+        "surprise_threshold": cfg.model.get("surprise_threshold"),
+        "teach_scale": float(cfg.model.get("teach_scale", 1.0)),
+        "teach_clip": float(cfg.model.get("teach_clip", 0.0)),
+        "freeze_backbone": bool(cfg.model.get("freeze_backbone", False)),
+        "has_embed": has_embed,
+        "has_lm_head": has_lm_head,
+        "lm_tied_to_embedding": tied_weights,
+    }
+    typer.echo(json.dumps(report, sort_keys=True))

nested_learning/cms.py

@@ -0,0 +1,92 @@
+from __future__ import annotations
+
+from typing import Dict, Sequence
+
+import torch
+import torch.nn as nn
+
+from .levels import LevelSpec, ensure_level_specs
+
+
+class CMSBlock(nn.Module):
+    def __init__(
+        self,
+        dim: int,
+        hidden_multiplier: int = 4,
+        activation: str = "gelu",
+        grad_clip: float = 1.0,
+        use_layernorm: bool = True,
+    ):
+        super().__init__()
+        hidden = dim * hidden_multiplier
+        act: nn.Module
+        if activation == "relu":
+            act = nn.ReLU()
+        elif activation == "silu":
+            act = nn.SiLU()
+        else:
+            act = nn.GELU()
+        norm: nn.Module = nn.LayerNorm(dim) if use_layernorm else nn.Identity()
+        self.net = nn.Sequential(
+            norm,
+            nn.Linear(dim, hidden),
+            act,
+            nn.Linear(hidden, dim),
+        )
+        self.grad_clip = grad_clip
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:  # type: ignore[override]
+        delta = self.net(x)
+        if self.training and self.grad_clip > 0:
+            with torch.no_grad():
+                norm = delta.norm(dim=-1, keepdim=True)
+                scale = torch.clamp(norm / self.grad_clip, min=1.0)
+            delta = delta / scale
+        return x + delta
+
+
+class CMS(nn.Module):
+    """Continuum Memory System with multi-frequency updates."""
+
+    def __init__(
+        self,
+        *,
+        dim: int,
+        levels: Sequence[LevelSpec],
+        hidden_multiplier: int = 4,
+        activation: str = "gelu",
+        use_layernorm: bool = True,
+    ) -> None:
+        super().__init__()
+        ordered = ensure_level_specs(levels)
+        self.level_specs: Sequence[LevelSpec] = tuple(ordered)
+        self.blocks = nn.ModuleDict(
+            {
+                spec.name: CMSBlock(
+                    dim,
+                    hidden_multiplier=hidden_multiplier,
+                    activation=activation,
+                    grad_clip=1.0,
+                    use_layernorm=use_layernorm,
+                )
+                for spec in self.level_specs
+            }
+        )
+
+    def forward(
+        self,
+        x: torch.Tensor,
+        *,
+        return_intermediates: bool = False,
+    ) -> torch.Tensor | tuple[torch.Tensor, Dict[str, torch.Tensor], Dict[str, torch.Tensor]]:
+        current = x
+        inputs: Dict[str, torch.Tensor] = {}
+        outputs: Dict[str, torch.Tensor] = {}
+        for spec in self.level_specs:
+            block = self.blocks[spec.name]
+            inputs[spec.name] = current
+            current = block(current)
+            outputs[spec.name] = current
+        if return_intermediates:
+            return current, inputs, outputs
+        return current

nested_learning/config_utils.py

@@ -0,0 +1,50 @@
+from __future__ import annotations
+
+from contextlib import contextmanager
+from importlib.resources import as_file, files
+from pathlib import Path
+from typing import Iterator
+
+from hydra import compose, initialize_config_dir
+from hydra.core.global_hydra import GlobalHydra
+from omegaconf import DictConfig
+
+from .training import unwrap_config
+
+
+def find_repo_root(start: Path | None = None) -> Path | None:
+    cursor = (start or Path.cwd()).resolve()
+    for candidate in (cursor, *cursor.parents):
+        if (candidate / ".git").exists() and (candidate / "configs").exists():
+            return candidate
+    return None
+
+
+@contextmanager
+def resolved_config_dir(config_dir: Path | None = None) -> Iterator[Path]:
+    if config_dir is not None:
+        yield config_dir.resolve()
+        return
+
+    module_path = Path(__file__).resolve()
+    repo_config_dir = module_path.parents[2] / "configs"
+    if repo_config_dir.exists():
+        yield repo_config_dir
+        return
+
+    package_configs = files("nested_learning").joinpath("configs")
+    with as_file(package_configs) as pkg_dir:
+        yield Path(pkg_dir)
+
+
+def compose_config(
+    config_name: str,
+    *,
+    overrides: list[str] | None = None,
+    config_dir: Path | None = None,
+) -> DictConfig:
+    with resolved_config_dir(config_dir) as cfg_dir:
+        GlobalHydra.instance().clear()
+        with initialize_config_dir(version_base=None, config_dir=str(cfg_dir)):
+            cfg = compose(config_name=config_name, overrides=overrides or [])
+    return unwrap_config(cfg)