From f0b18973f58b100e6ece47922389dcc59d06987e Mon Sep 17 00:00:00 2001 From: Charles Martin Date: Tue, 11 Aug 2026 09:04:34 -0700 Subject: [PATCH 01/22] Add NGB v4 package metadata --- baseline/ngb/pyproject.toml | 36 ++++++++++++++++++++++++++++++++++++ 1 file changed, 36 insertions(+) create mode 100644 baseline/ngb/pyproject.toml diff --git a/baseline/ngb/pyproject.toml b/baseline/ngb/pyproject.toml new file mode 100644 index 0000000..fee7988 --- /dev/null +++ b/baseline/ngb/pyproject.toml @@ -0,0 +1,36 @@ +[build-system] +requires = ["setuptools>=68", "wheel"] +build-backend = "setuptools.build_meta" + +[project] +name = "rg-ngb" +version = "0.1.0" +description = "NGB v4 tuned nanoGPT optimizer baselines" +requires-python = ">=3.10" +dependencies = [ + "torch>=2.2", + "numpy>=1.24", + "pandas>=2.0", + "matplotlib>=3.7", + "pyyaml>=6.0", + "datasets>=2.19", + "tiktoken>=0.7", + "sacrebleu>=2.4", + "weightwatcher==0.7.7", + "jupyter>=1.0", + "ipykernel>=6.29", + "papermill>=2.6", +] + +[project.optional-dependencies] +dev = [ + "pytest>=8.0", + "nbformat>=5.10", +] + +[project.scripts] +ngb-prepare = "rg_ngb.data:main" +ngb-train = "rg_ngb.training:main" + +[tool.setuptools.packages.find] +where = ["src"] From 560538454fbf4ee8325fb98dc6bd8a36be65fe62 Mon Sep 17 00:00:00 2001 From: Charles Martin Date: Tue, 11 Aug 2026 09:05:12 -0700 Subject: [PATCH 02/22] Document NGB v4 protocols --- baseline/ngb/README.md | 118 +++++++++++++++++++++++++++++++++++++++++ 1 file changed, 118 insertions(+) create mode 100644 baseline/ngb/README.md diff --git a/baseline/ngb/README.md b/baseline/ngb/README.md new file mode 100644 index 0000000..924dc27 --- /dev/null +++ b/baseline/ngb/README.md @@ -0,0 +1,118 @@ +# NGB v4 nanoGPT baselines + +`baseline/ngb` is a separate experiment family for properly tuned small-language-model optimizer baselines. It does **not** reuse the v3 one-head result directories and does not alter the checked-in v3 protocol or outputs. + +All data, checkpoints, plots, and results default to: + +```text +/tmp/rg-ngb +``` + +No shell wrappers or project virtual environments are required. Use the currently activated conda environment. + +## Protocols + +| Configuration | Architecture | Horizon | Purpose | +|---|---|---:|---| +| `configs/v4_one_head.yaml` | 1 block, 1 head, width 128, context 256 | 2 corpus-equivalent epochs | Direct tuned successor to the v3 diagnostic model | +| `configs/v4_small_4x4.yaml` | 4 blocks, 4 heads, width 128, context 256 | 2 corpus-equivalent epochs | Distinct small but substantially more expressive language-model baseline | + +Both protocols use the same pinned, document-disjoint FineWeb-Edu 80M / 1M / 1M GPT-2-BPE corpus contract. The prepared corpus can be shared, while protocol fingerprints and result directories remain separate. + +## Tuned optimizer centers + +The v4 centers respond directly to the v3 instability observed in the checked-in comparison. + +### One-head v4 + +| Optimizer | Peak LR | Floor | Warm-up | Decay | +|---|---:|---:|---:|---:| +| SGD + Nesterov | `5e-2` | `5e-4` | 5% | `1e-2` | +| AdamW | `3e-4` | `1e-5` | 2.5% | `1e-1` | +| Muon matrices | `1e-2` | `2e-4` | 5% | `2e-2` | +| Muon auxiliary AdamW | `3e-4` | `1e-5` | 5% | `1e-1` | + +### Small 4x4 v4 + +The 4x4 architecture uses the same adaptive centers. Its SGD center is reduced to `3e-2` because four residual blocks produce a materially different optimization geometry. + +These are preregistered v4 centers, not a claim that a broad hyperparameter search has already been completed. Validation loss remains the only checkpoint-selection and qualification objective. + +## Install in the active conda environment + +From the repository root: + +```bash +cd baseline/ngb +python -m pip install -e ../nanogpt_one_head +python -m pip install -e '.[dev]' +``` + +The first editable install supplies shared, already-tested data, evaluation, checkpoint, runtime, and optimizer primitives. NGB owns its generalized model, protocol, spectral inventory, completion validation, result directories, and comparison logic. + +## Prepare the corpus + +```bash +export RG_NGB_ROOT="/tmp/rg-ngb" +export PYTORCH_ENABLE_MPS_FALLBACK=1 + +ngb-prepare --config configs/v4_one_head.yaml +``` + +The 4x4 protocol uses the same verified token files, so no second corpus download is required. + +## Run the tuned one-head v4 experiment + +```bash +ngb-train \ + --config configs/v4_one_head.yaml \ + --optimizer all \ + --device auto +``` + +## Run the distinct 4-layer / 4-head experiment + +```bash +ngb-train \ + --config configs/v4_small_4x4.yaml \ + --optimizer all \ + --device auto +``` + +The default canonical qualification seeds are `1337, 2027, 4099`. An explicit matched expansion can be run with, for example: + +```bash +ngb-train \ + --config configs/v4_one_head.yaml \ + --optimizer all \ + --seeds 1337,2027,4099,5003,6007,7013,8017,9011 \ + --device auto +``` + +Completed compatible runs are verified and reused. Stale or protocol-incompatible directories fail visibly rather than being silently accepted. + +## Notebooks + +```text +notebooks/01_v4_one_head_train.ipynb +notebooks/02_v4_one_head_compare.ipynb +notebooks/11_v4_small_4x4_train.ipynb +notebooks/12_v4_small_4x4_compare.ipynb +``` + +Run with the active conda kernel: + +```bash +papermill notebooks/01_v4_one_head_train.ipynb notebooks/01_v4_one_head_train.out.ipynb -k python3 +papermill notebooks/02_v4_one_head_compare.ipynb notebooks/02_v4_one_head_compare.out.ipynb -k python3 +``` + +The comparison notebooks discover the complete seed intersection shared by all three optimizers. They report run-level 95% Student-t intervals, matched-seed differences, validation-selected and final metrics, best-step drift, clipping frequency, maximum update-to-weight ratio, and WeightWatcher trajectories. Perplexity intervals are obtained by exponentiating the loss-space confidence interval, so they cannot have impossible negative lower bounds. + +## Validation + +```bash +python -m pytest -q tests +``` + +The tests cover both architectures, exact protocol settings, dynamic matrix inventories, `/tmp` path isolation, matched-seed discovery, transformed perplexity intervals, optimizer partitioning, notebook syntax, and a tiny CPU training path. From ffd664905bf10af1254c5ffe9684ebfb700d5c85 Mon Sep 17 00:00:00 2001 From: Charles Martin Date: Tue, 11 Aug 2026 09:05:32 -0700 Subject: [PATCH 03/22] Add tuned one-head v4 protocol --- baseline/ngb/configs/v4_one_head.yaml | 103 ++++++++++++++++++++++++++ 1 file changed, 103 insertions(+) create mode 100644 baseline/ngb/configs/v4_one_head.yaml diff --git a/baseline/ngb/configs/v4_one_head.yaml b/baseline/ngb/configs/v4_one_head.yaml new file mode 100644 index 0000000..e614d84 --- /dev/null +++ b/baseline/ngb/configs/v4_one_head.yaml @@ -0,0 +1,103 @@ +protocol: + name: ngb_v4_one_head_2epoch + version: 4 + architecture_id: one_head + description: Tuned two-epoch successor to the v3 one-block one-head FineWeb-Edu baseline. + +dataset: + name: HuggingFaceFW/fineweb-edu + config: sample-10BT + split: train + revision: 593b3a867298afb8ce42625a270ef20ddcad28f9 + tokenizer: gpt2 + train_tokens: 80000000 + val_tokens: 1000000 + test_tokens: 1000000 + +model: + vocab_size: 50257 + block_size: 256 + n_layer: 1 + n_head: 1 + n_embd: 128 + dropout: 0.0 + bias: false + tie_weights: true + +training: + seeds: [1337, 2027, 4099] + batch_size: 4 + grad_accum_steps: 8 + target_epochs: 2.0 + epoch_interval: 0.25 + eval_interval_steps: 500 + eval_batches: 64 + checkpoint_interval_steps: 500 + grad_clip: 1.0 + +optimizer_profiles: + sgd_momentum: + display_name: SGD + Nesterov momentum v4 + family: sgd + learning_rate: 0.05 + min_learning_rate: 0.0005 + warmup_fraction: 0.05 + schedule: warmup_cosine + momentum: 0.90 + dampening: 0.0 + nesterov: true + weight_decay: 0.01 + + adamw: + display_name: AdamW v4 tuned + family: adamw + learning_rate: 0.0003 + min_learning_rate: 0.00001 + warmup_fraction: 0.025 + schedule: warmup_cosine + beta1: 0.90 + beta2: 0.95 + epsilon: 1.0e-8 + weight_decay: 0.10 + + muon: + display_name: Muon + auxiliary AdamW v4 tuned + family: muon + matrix_learning_rate: 0.01 + matrix_min_learning_rate: 0.0002 + aux_learning_rate: 0.0003 + aux_min_learning_rate: 0.00001 + warmup_fraction: 0.05 + schedule: warmup_cosine + momentum: 0.95 + nesterov: true + newton_schulz_steps: 5 + muon_epsilon: 1.0e-7 + matrix_weight_decay: 0.02 + beta1: 0.90 + beta2: 0.95 + epsilon: 1.0e-8 + aux_weight_decay: 0.10 + +evaluation: + train_probe_seed: 21001 + validation_probe_seed: 22001 + test_probe_seed: 23001 + bleu_probe_seed: 24001 + bleu_examples: 64 + bleu_prompt_tokens: 64 + bleu_continuation_tokens: 32 + bleu_batch_size: 4 + +weightwatcher: + enabled: true + ERG: true + randomize: true + strict: true + min_evals: 20 + +runtime: + matmul_precision: high + mps_fallback: true + deterministic_algorithms: false + empty_mps_cache_after_weightwatcher: true From 7cc1d7cf0ee26484ffa2785d332a01ddbd17939f Mon Sep 17 00:00:00 2001 From: Charles Martin Date: Tue, 11 Aug 2026 09:05:57 -0700 Subject: [PATCH 04/22] Add expressive four-layer four-head v4 protocol --- baseline/ngb/configs/v4_small_4x4.yaml | 103 +++++++++++++++++++++++++ 1 file changed, 103 insertions(+) create mode 100644 baseline/ngb/configs/v4_small_4x4.yaml diff --git a/baseline/ngb/configs/v4_small_4x4.yaml b/baseline/ngb/configs/v4_small_4x4.yaml new file mode 100644 index 0000000..ca8eb34 --- /dev/null +++ b/baseline/ngb/configs/v4_small_4x4.yaml @@ -0,0 +1,103 @@ +protocol: + name: ngb_v4_small_4x4_2epoch + version: 4 + architecture_id: small_4x4 + description: Distinct four-block four-head width-128 FineWeb-Edu small-language-model baseline. + +dataset: + name: HuggingFaceFW/fineweb-edu + config: sample-10BT + split: train + revision: 593b3a867298afb8ce42625a270ef20ddcad28f9 + tokenizer: gpt2 + train_tokens: 80000000 + val_tokens: 1000000 + test_tokens: 1000000 + +model: + vocab_size: 50257 + block_size: 256 + n_layer: 4 + n_head: 4 + n_embd: 128 + dropout: 0.0 + bias: false + tie_weights: true + +training: + seeds: [1337, 2027, 4099] + batch_size: 4 + grad_accum_steps: 8 + target_epochs: 2.0 + epoch_interval: 0.25 + eval_interval_steps: 500 + eval_batches: 64 + checkpoint_interval_steps: 500 + grad_clip: 1.0 + +optimizer_profiles: + sgd_momentum: + display_name: SGD + Nesterov momentum v4 4x4 + family: sgd + learning_rate: 0.03 + min_learning_rate: 0.0003 + warmup_fraction: 0.05 + schedule: warmup_cosine + momentum: 0.90 + dampening: 0.0 + nesterov: true + weight_decay: 0.01 + + adamw: + display_name: AdamW v4 tuned 4x4 + family: adamw + learning_rate: 0.0003 + min_learning_rate: 0.00001 + warmup_fraction: 0.025 + schedule: warmup_cosine + beta1: 0.90 + beta2: 0.95 + epsilon: 1.0e-8 + weight_decay: 0.10 + + muon: + display_name: Muon + auxiliary AdamW v4 tuned 4x4 + family: muon + matrix_learning_rate: 0.01 + matrix_min_learning_rate: 0.0002 + aux_learning_rate: 0.0003 + aux_min_learning_rate: 0.00001 + warmup_fraction: 0.05 + schedule: warmup_cosine + momentum: 0.95 + nesterov: true + newton_schulz_steps: 5 + muon_epsilon: 1.0e-7 + matrix_weight_decay: 0.02 + beta1: 0.90 + beta2: 0.95 + epsilon: 1.0e-8 + aux_weight_decay: 0.10 + +evaluation: + train_probe_seed: 21001 + validation_probe_seed: 22001 + test_probe_seed: 23001 + bleu_probe_seed: 24001 + bleu_examples: 64 + bleu_prompt_tokens: 64 + bleu_continuation_tokens: 32 + bleu_batch_size: 4 + +weightwatcher: + enabled: true + ERG: true + randomize: true + strict: true + min_evals: 20 + +runtime: + matmul_precision: high + mps_fallback: true + deterministic_algorithms: false + empty_mps_cache_after_weightwatcher: true From f329b0bfecb6d2bff707e2e53de83b34630ae7cb Mon Sep 17 00:00:00 2001 From: Charles Martin Date: Tue, 11 Aug 2026 09:06:45 -0700 Subject: [PATCH 05/22] Add generalized NGB configuration --- baseline/ngb/src/rg_ngb/config.py | 224 ++++++++++++++++++++++++++++++ 1 file changed, 224 insertions(+) create mode 100644 baseline/ngb/src/rg_ngb/config.py diff --git a/baseline/ngb/src/rg_ngb/config.py b/baseline/ngb/src/rg_ngb/config.py new file mode 100644 index 0000000..e851216 --- /dev/null +++ b/baseline/ngb/src/rg_ngb/config.py @@ -0,0 +1,224 @@ +from __future__ import annotations + +from copy import deepcopy +import hashlib +import json +import os +from pathlib import Path +from typing import Any + +import yaml + +SUPPORTED_OPTIMIZERS = ("sgd_momentum", "adamw", "muon") +DEFAULT_ROOT = Path("/tmp/rg-ngb") + + +def protocol_name(cfg: dict[str, Any]) -> str: + value = str(cfg["protocol"]["name"]).strip() + if not value or any(character not in "abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ0123456789_-" for character in value): + raise ValueError("protocol.name must be a nonempty filesystem-safe slug") + return value + + +def roots(cfg: dict[str, Any] | None = None) -> dict[str, Path]: + root = Path(os.environ.get("RG_NGB_ROOT", DEFAULT_ROOT)) + data = Path(os.environ.get("RG_NGB_DATA_ROOT", root / "data")) + suffix = protocol_name(cfg) if cfg is not None else "unscoped" + results = Path(os.environ.get("RG_NGB_RESULTS_ROOT", root / "results" / suffix)) + plots = Path(os.environ.get("RG_NGB_PLOTS_ROOT", root / "plots" / suffix)) + return {"root": root, "data": data, "results": results, "plots": plots} + + +def load_config(path: str | Path) -> dict[str, Any]: + path = Path(path) + with path.open("r", encoding="utf-8") as handle: + cfg = yaml.safe_load(handle) + if not isinstance(cfg, dict): + raise ValueError("configuration root must be a mapping") + validate_config(cfg) + return cfg + + +def validate_config(cfg: dict[str, Any]) -> None: + required = ( + "protocol", + "dataset", + "model", + "training", + "optimizer_profiles", + "evaluation", + "weightwatcher", + "runtime", + ) + for section in required: + if section not in cfg: + raise ValueError(f"missing configuration section: {section}") + + protocol = cfg["protocol"] + if int(protocol.get("version", 0)) != 4: + raise ValueError("NGB requires protocol.version == 4") + protocol_name(cfg) + + model = cfg["model"] + for key in ("vocab_size", "block_size", "n_layer", "n_head", "n_embd"): + if int(model[key]) < 1: + raise ValueError(f"model.{key} must be positive") + if int(model["n_embd"]) % int(model["n_head"]) != 0: + raise ValueError("model.n_embd must be divisible by model.n_head") + if not 0.0 <= float(model.get("dropout", 0.0)) < 1.0: + raise ValueError("model.dropout must be in [0, 1)") + + dataset = cfg["dataset"] + for key in ("train_tokens", "val_tokens", "test_tokens"): + if int(dataset[key]) <= int(model["block_size"]) + 1: + raise ValueError(f"dataset.{key} is too small for the context length") + + training = cfg["training"] + for key in ( + "batch_size", + "grad_accum_steps", + "target_epochs", + "eval_interval_steps", + "eval_batches", + "checkpoint_interval_steps", + "epoch_interval", + ): + if float(training[key]) <= 0: + raise ValueError(f"training.{key} must be positive") + seeds = [int(seed) for seed in training["seeds"]] + if not seeds or len(seeds) != len(set(seeds)): + raise ValueError("training.seeds must be nonempty and unique") + if float(training["grad_clip"]) < 0: + raise ValueError("training.grad_clip must be nonnegative") + + profiles = cfg["optimizer_profiles"] + for name in SUPPORTED_OPTIMIZERS: + if name not in profiles: + raise ValueError(f"missing optimizer profile: {name}") + validate_optimizer_profile({**profiles[name], "name": name}) + + evaluation = cfg["evaluation"] + for key in ( + "bleu_examples", + "bleu_prompt_tokens", + "bleu_continuation_tokens", + "bleu_batch_size", + ): + if int(evaluation[key]) < 1: + raise ValueError(f"evaluation.{key} must be positive") + if int(evaluation["bleu_prompt_tokens"]) + int(evaluation["bleu_continuation_tokens"]) > int(model["block_size"]): + raise ValueError("BLEU prompt plus continuation must fit in block_size") + probe_keys = ( + "train_probe_seed", + "validation_probe_seed", + "test_probe_seed", + "bleu_probe_seed", + ) + probe_values = [int(evaluation[key]) for key in probe_keys] + if any(value < 0 for value in probe_values) or len(set(probe_values)) != len(probe_values): + raise ValueError("evaluation probe seeds must be distinct nonnegative values") + + ww = cfg["weightwatcher"] + if not bool(ww.get("ERG", False)) or not bool(ww.get("randomize", False)): + raise ValueError("NGB requires WeightWatcher ERG and randomize") + if int(ww["min_evals"]) < 5: + raise ValueError("weightwatcher.min_evals must be at least 5") + + +def validate_optimizer_profile(profile: dict[str, Any]) -> None: + family = str(profile.get("family", "")) + if family not in {"sgd", "adamw", "muon"}: + raise ValueError(f"unsupported optimizer family: {family}") + warmup_fraction = float(profile.get("warmup_fraction", -1.0)) + if not 0.0 <= warmup_fraction < 1.0: + raise ValueError("warmup_fraction must be in [0, 1)") + if str(profile.get("schedule")) != "warmup_cosine": + raise ValueError("NGB requires warmup_cosine schedules") + + if family in {"sgd", "adamw"}: + pairs = (("learning_rate", "min_learning_rate"),) + else: + pairs = ( + ("matrix_learning_rate", "matrix_min_learning_rate"), + ("aux_learning_rate", "aux_min_learning_rate"), + ) + for peak_key, floor_key in pairs: + peak = float(profile[peak_key]) + floor = float(profile[floor_key]) + if peak <= 0 or floor < 0 or floor > peak: + raise ValueError(f"inconsistent learning-rate pair: {peak_key}/{floor_key}") + if family == "muon" and int(profile["newton_schulz_steps"]) < 1: + raise ValueError("Muon newton_schulz_steps must be positive") + + +def optimizer_profile(cfg: dict[str, Any], optimizer: str) -> dict[str, Any]: + optimizer = str(optimizer).lower() + if optimizer not in SUPPORTED_OPTIMIZERS: + raise ValueError(f"unsupported optimizer {optimizer!r}") + profile = deepcopy(cfg["optimizer_profiles"][optimizer]) + profile["name"] = optimizer + validate_optimizer_profile(profile) + return profile + + +def canonical_seeds(cfg: dict[str, Any]) -> tuple[int, ...]: + return tuple(int(seed) for seed in cfg["training"]["seeds"]) + + +def tokens_per_step(cfg: dict[str, Any]) -> int: + return int(cfg["training"]["batch_size"]) * int(cfg["training"]["grad_accum_steps"]) * int(cfg["model"]["block_size"]) + + +def max_steps(cfg: dict[str, Any], train_tokens: int | None = None) -> int: + import math + + train_tokens = int(train_tokens or cfg["dataset"]["train_tokens"]) + target = float(cfg["training"]["target_epochs"]) * train_tokens + return max(1, int(math.ceil(target / tokens_per_step(cfg)))) + + +def warmup_steps(profile: dict[str, Any], total_steps: int) -> int: + if total_steps < 2: + return 0 + return min(total_steps - 1, max(1, int(round(total_steps * float(profile["warmup_fraction"]))))) + + +def epoch_step_map(cfg: dict[str, Any], train_tokens: int | None = None) -> dict[int, float]: + train_tokens = int(train_tokens or cfg["dataset"]["train_tokens"]) + total_steps = max_steps(cfg, train_tokens) + step_tokens = tokens_per_step(cfg) + target = float(cfg["training"]["target_epochs"]) + interval = float(cfg["training"]["epoch_interval"]) + points = [0.0] + current = interval + while current < target - 1e-12: + points.append(round(current, 12)) + current += interval + points.append(target) + result: dict[int, float] = {} + for epoch in points: + step = 0 if epoch == 0 else int(round(epoch * train_tokens / step_tokens)) + result[min(total_steps, max(0, step))] = float(epoch) + result[total_steps] = target + return dict(sorted(result.items())) + + +def expected_matrix_count(cfg: dict[str, Any]) -> int: + return 6 * int(cfg["model"]["n_layer"]) + + +def protocol_fingerprint(cfg: dict[str, Any], *, optimizer: str, seed: int, data_metadata: dict[str, Any]) -> str: + payload = { + "protocol": cfg["protocol"], + "dataset": cfg["dataset"], + "model": cfg["model"], + "training": cfg["training"], + "optimizer_profile": optimizer_profile(cfg, optimizer), + "evaluation": cfg["evaluation"], + "weightwatcher": cfg["weightwatcher"], + "optimizer": str(optimizer), + "seed": int(seed), + "data_metadata": data_metadata, + } + canonical = json.dumps(payload, sort_keys=True, separators=(",", ":"), default=str) + return hashlib.sha256(canonical.encode("utf-8")).hexdigest() From 806e69f84a709ea1c47ae0b72a328ff56346f482 Mon Sep 17 00:00:00 2001 From: Charles Martin Date: Tue, 11 Aug 2026 09:07:17 -0700 Subject: [PATCH 06/22] Add generalized NGB GPT model --- baseline/ngb/src/rg_ngb/model.py | 175 +++++++++++++++++++++++++++++++ 1 file changed, 175 insertions(+) create mode 100644 baseline/ngb/src/rg_ngb/model.py diff --git a/baseline/ngb/src/rg_ngb/model.py b/baseline/ngb/src/rg_ngb/model.py new file mode 100644 index 0000000..591d450 --- /dev/null +++ b/baseline/ngb/src/rg_ngb/model.py @@ -0,0 +1,175 @@ +from __future__ import annotations + +from dataclasses import dataclass +import math + +import torch +import torch.nn as nn +import torch.nn.functional as F + + +@dataclass(frozen=True) +class GPTConfig: + vocab_size: int = 50_257 + block_size: int = 256 + n_layer: int = 1 + n_head: int = 1 + n_embd: int = 128 + dropout: float = 0.0 + bias: bool = False + tie_weights: bool = True + + def __post_init__(self) -> None: + if self.n_layer < 1 or self.n_head < 1: + raise ValueError("n_layer and n_head must be positive") + if self.n_embd % self.n_head != 0: + raise ValueError("n_embd must be divisible by n_head") + if self.block_size < 2 or self.vocab_size < 2 or self.n_embd < 1: + raise ValueError("invalid GPT configuration") + if not 0.0 <= self.dropout < 1.0: + raise ValueError("dropout must be in [0, 1)") + + +class LayerNorm(nn.Module): + def __init__(self, width: int, bias: bool) -> None: + super().__init__() + self.weight = nn.Parameter(torch.ones(width)) + self.bias = nn.Parameter(torch.zeros(width)) if bias else None + + def forward(self, x: torch.Tensor) -> torch.Tensor: + return F.layer_norm(x, self.weight.shape, self.weight, self.bias, 1e-5) + + +class CausalSelfAttention(nn.Module): + def __init__(self, cfg: GPTConfig) -> None: + super().__init__() + self.n_head = cfg.n_head + self.n_embd = cfg.n_embd + self.dropout = float(cfg.dropout) + self.q_proj = nn.Linear(cfg.n_embd, cfg.n_embd, bias=cfg.bias) + self.k_proj = nn.Linear(cfg.n_embd, cfg.n_embd, bias=cfg.bias) + self.v_proj = nn.Linear(cfg.n_embd, cfg.n_embd, bias=cfg.bias) + self.out_proj = nn.Linear(cfg.n_embd, cfg.n_embd, bias=cfg.bias) + self.resid_dropout = nn.Dropout(cfg.dropout) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + batch, sequence, channels = x.shape + head_width = channels // self.n_head + q = self.q_proj(x).view(batch, sequence, self.n_head, head_width).transpose(1, 2) + k = self.k_proj(x).view(batch, sequence, self.n_head, head_width).transpose(1, 2) + v = self.v_proj(x).view(batch, sequence, self.n_head, head_width).transpose(1, 2) + y = F.scaled_dot_product_attention( + q, + k, + v, + attn_mask=None, + dropout_p=self.dropout if self.training else 0.0, + is_causal=True, + ) + y = y.transpose(1, 2).contiguous().view(batch, sequence, channels) + return self.resid_dropout(self.out_proj(y)) + + +class MLP(nn.Module): + def __init__(self, cfg: GPTConfig) -> None: + super().__init__() + self.fc = nn.Linear(cfg.n_embd, 4 * cfg.n_embd, bias=cfg.bias) + self.proj = nn.Linear(4 * cfg.n_embd, cfg.n_embd, bias=cfg.bias) + self.dropout = nn.Dropout(cfg.dropout) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + return self.dropout(self.proj(F.gelu(self.fc(x), approximate="tanh"))) + + +class Block(nn.Module): + def __init__(self, cfg: GPTConfig) -> None: + super().__init__() + self.ln1 = LayerNorm(cfg.n_embd, cfg.bias) + self.attn = CausalSelfAttention(cfg) + self.ln2 = LayerNorm(cfg.n_embd, cfg.bias) + self.mlp = MLP(cfg) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + x = x + self.attn(self.ln1(x)) + return x + self.mlp(self.ln2(x)) + + +class GPT(nn.Module): + def __init__(self, cfg: GPTConfig) -> None: + super().__init__() + self.cfg = cfg + self.token_embedding = nn.Embedding(cfg.vocab_size, cfg.n_embd) + self.position_embedding = nn.Embedding(cfg.block_size, cfg.n_embd) + self.drop = nn.Dropout(cfg.dropout) + self.blocks = nn.ModuleList([Block(cfg) for _ in range(cfg.n_layer)]) + self.ln_f = LayerNorm(cfg.n_embd, cfg.bias) + self.lm_head = nn.Linear(cfg.n_embd, cfg.vocab_size, bias=False) + if cfg.tie_weights: + self.lm_head.weight = self.token_embedding.weight + + self.apply(self._init_module) + residual_std = 0.02 / math.sqrt(2 * cfg.n_layer) + for block in self.blocks: + nn.init.normal_(block.attn.out_proj.weight, mean=0.0, std=residual_std) + nn.init.normal_(block.mlp.proj.weight, mean=0.0, std=residual_std) + + @staticmethod + def _init_module(module: nn.Module) -> None: + if isinstance(module, (nn.Linear, nn.Embedding)): + nn.init.normal_(module.weight, mean=0.0, std=0.02) + if isinstance(module, nn.Linear) and module.bias is not None: + nn.init.zeros_(module.bias) + + def hidden_states(self, idx: torch.Tensor) -> torch.Tensor: + _, sequence = idx.shape + if sequence > self.cfg.block_size: + raise ValueError("input sequence exceeds model.block_size") + positions = torch.arange(sequence, device=idx.device) + x = self.drop(self.token_embedding(idx) + self.position_embedding(positions)) + for block in self.blocks: + x = block(x) + return self.ln_f(x) + + def forward(self, idx: torch.Tensor, targets: torch.Tensor | None = None) -> tuple[torch.Tensor, torch.Tensor | None]: + logits = self.lm_head(self.hidden_states(idx)) + loss = None + if targets is not None: + loss = F.cross_entropy(logits.reshape(-1, logits.size(-1)), targets.reshape(-1)) + return logits, loss + + def next_token_logits(self, idx: torch.Tensor) -> torch.Tensor: + hidden = self.hidden_states(idx)[:, -1:, :] + return self.lm_head(hidden) + + @torch.inference_mode() + def generate_greedy(self, prompts: torch.Tensor, max_new_tokens: int) -> torch.Tensor: + if prompts.ndim != 2: + raise ValueError("prompts must be [batch, sequence]") + if max_new_tokens < 0: + raise ValueError("max_new_tokens must be nonnegative") + idx = prompts + for _ in range(int(max_new_tokens)): + idx_cond = idx[:, -self.cfg.block_size :] + logits = self.next_token_logits(idx_cond) + next_token = logits[:, -1, :].argmax(dim=-1, keepdim=True) + idx = torch.cat((idx, next_token), dim=1) + return idx + + def parameter_count(self) -> int: + return sum(parameter.numel() for parameter in self.parameters()) + + +def transformer_matrix_items(model: GPT) -> list[tuple[str, str, int, torch.Tensor]]: + items: list[tuple[str, str, int, torch.Tensor]] = [] + for block_index, block in enumerate(model.blocks): + matrices = ( + ("W_Q", block.attn.q_proj.weight), + ("W_K", block.attn.k_proj.weight), + ("W_V", block.attn.v_proj.weight), + ("W_O", block.attn.out_proj.weight), + ("W_MLP_IN", block.mlp.fc.weight), + ("W_MLP_OUT", block.mlp.proj.weight), + ) + for matrix_type, weight in matrices: + items.append((f"L{block_index:02d}_{matrix_type}", matrix_type, block_index, weight)) + return items From 756b0048630b57a81c7caca6438343554ea90ac0 Mon Sep 17 00:00:00 2001 From: Charles Martin Date: Tue, 11 Aug 2026 09:07:32 -0700 Subject: [PATCH 07/22] Add NGB data entry point --- baseline/ngb/src/rg_ngb/data.py | 37 +++++++++++++++++++++++++++++++++ 1 file changed, 37 insertions(+) create mode 100644 baseline/ngb/src/rg_ngb/data.py diff --git a/baseline/ngb/src/rg_ngb/data.py b/baseline/ngb/src/rg_ngb/data.py new file mode 100644 index 0000000..3187a42 --- /dev/null +++ b/baseline/ngb/src/rg_ngb/data.py @@ -0,0 +1,37 @@ +from __future__ import annotations + +import argparse +from pathlib import Path + +from rg_nanogpt_one_head.data import ( + TOKEN_DTYPE, + load_memmaps, + prepare_fineweb_edu, + validate_prepared_data, + write_token_splits, +) + +from .config import load_config, roots + +__all__ = [ + "TOKEN_DTYPE", + "load_memmaps", + "prepare_fineweb_edu", + "validate_prepared_data", + "write_token_splits", +] + + +def main() -> None: + parser = argparse.ArgumentParser(description="Prepare the pinned FineWeb-Edu corpus for NGB v4") + parser.add_argument("--config", required=True) + parser.add_argument("--output-dir") + parser.add_argument("--force", action="store_true") + args = parser.parse_args() + cfg = load_config(args.config) + output = Path(args.output_dir) if args.output_dir else roots(cfg)["data"] + prepare_fineweb_edu(cfg, output, force=args.force) + + +if __name__ == "__main__": + main() From b969c77d07ad5d58b723cb4673b7296134cc545d Mon Sep 17 00:00:00 2001 From: Charles Martin Date: Tue, 11 Aug 2026 09:08:16 -0700 Subject: [PATCH 08/22] Add dynamic NGB WeightWatcher measurements --- baseline/ngb/src/rg_ngb/spectral.py | 208 ++++++++++++++++++++++++++++ 1 file changed, 208 insertions(+) create mode 100644 baseline/ngb/src/rg_ngb/spectral.py diff --git a/baseline/ngb/src/rg_ngb/spectral.py b/baseline/ngb/src/rg_ngb/spectral.py new file mode 100644 index 0000000..d046f00 --- /dev/null +++ b/baseline/ngb/src/rg_ngb/spectral.py @@ -0,0 +1,208 @@ +from __future__ import annotations + +import json +from pathlib import Path +import random +from typing import Any + +import numpy as np +import pandas as pd +import torch +import torch.nn as nn + +from .model import GPT, transformer_matrix_items + +SPECTRAL_METRICS = ( + "alpha", + "alpha_weighted", + "ERG_gap", + "num_traps", + "detX_num", + "num_pl_spikes", + "num_ERG_spikes", + "D", + "stable_rank", + "mp_softrank", + "log_norm", + "log_spectral_norm", + "entropy", + "Lambda", + "rank_loss", +) + + +class WeightMatrixHolder(nn.Module): + def __init__(self, model: GPT) -> None: + super().__init__() + self.matrix_metadata: list[dict[str, object]] = [] + for name, matrix_type, block, weight in transformer_matrix_items(model): + layer = nn.Linear(weight.shape[1], weight.shape[0], bias=False) + layer.weight = nn.Parameter(weight.detach().float().cpu().clone(), requires_grad=False) + self.add_module(name, layer) + self.matrix_metadata.append({"matrix_name": name, "matrix_type": matrix_type, "block": int(block)}) + + +def _attach_matrix_metadata(frame: pd.DataFrame, metadata: list[dict[str, object]]) -> pd.DataFrame: + result = frame.copy().reset_index(drop=True) + names = [str(item["matrix_name"]) for item in metadata] + resolved: list[str | None] = [None] * len(result) + for row_index, row in result.iterrows(): + text = " ".join(str(row.get(column, "")) for column in ("longname", "name")) + for name in names: + if name in text: + resolved[row_index] = name + break + if any(value is None for value in resolved) and len(result) == len(metadata): + order = list(range(len(result))) + if "layer_id" in result.columns: + numeric = pd.to_numeric(result["layer_id"], errors="coerce") + if numeric.notna().all(): + order = list(numeric.sort_values().index) + for metadata_index, row_index in enumerate(order): + resolved[row_index] = names[metadata_index] + if any(value is None for value in resolved) or len(result) != len(metadata): + raise RuntimeError( + f"WeightWatcher returned {len(result)} rows for {len(metadata)} expected transformer matrices" + ) + by_name = {str(item["matrix_name"]): item for item in metadata} + result.insert(0, "matrix_name", resolved) + result.insert(1, "matrix_type", [by_name[str(name)]["matrix_type"] for name in resolved]) + result.insert(2, "block", [by_name[str(name)]["block"] for name in resolved]) + return result + + +def _atomic_csv(path: Path, frame: pd.DataFrame) -> None: + path.parent.mkdir(parents=True, exist_ok=True) + temporary = path.with_suffix(path.suffix + ".tmp") + frame.to_csv(temporary, index=False) + temporary.replace(path) + + +def _append_deduplicated(path: Path, frame: pd.DataFrame, keys: list[str]) -> None: + existing = pd.read_csv(path) if path.is_file() else pd.DataFrame() + combined = pd.concat([existing, frame], ignore_index=True, sort=False) if not existing.empty else frame.copy() + combined = combined.drop_duplicates(keys, keep="last").sort_values(keys) + _atomic_csv(path, combined) + + +def _capture_accelerator_rng() -> dict[str, Any]: + state: dict[str, Any] = {} + if torch.cuda.is_available(): + state["cuda"] = torch.cuda.get_rng_state_all() + if hasattr(torch, "mps") and hasattr(torch.mps, "get_rng_state") and torch.backends.mps.is_available(): + state["mps"] = torch.mps.get_rng_state() + return state + + +def _restore_accelerator_rng(state: dict[str, Any]) -> None: + if torch.cuda.is_available() and "cuda" in state: + torch.cuda.set_rng_state_all(state["cuda"]) + if "mps" in state and hasattr(torch, "mps") and hasattr(torch.mps, "set_rng_state") and torch.backends.mps.is_available(): + torch.mps.set_rng_state(state["mps"]) + + +def summarize_spectral_frame(frame: pd.DataFrame, *, step: int, tokens_seen: int, epoch: float) -> dict[str, Any]: + summary: dict[str, Any] = { + "step": int(step), + "tokens_seen": int(tokens_seen), + "epoch": float(epoch), + "n_matrices": int(len(frame)), + } + for metric in SPECTRAL_METRICS: + values = pd.to_numeric(frame[metric], errors="coerce") if metric in frame.columns else pd.Series(dtype=float) + array = values.to_numpy(dtype=float, na_value=np.nan) + finite = array[np.isfinite(array)] + summary[f"{metric}_n"] = int(finite.size) + for statistic in ("mean", "median", "std", "min", "max"): + summary[f"{metric}_{statistic}"] = float("nan") + if finite.size: + summary[f"{metric}_mean"] = float(np.mean(finite)) + summary[f"{metric}_median"] = float(np.median(finite)) + summary[f"{metric}_std"] = float(np.std(finite, ddof=1)) if finite.size > 1 else 0.0 + summary[f"{metric}_min"] = float(np.min(finite)) + summary[f"{metric}_max"] = float(np.max(finite)) + return summary + + +def run_weightwatcher( + model: GPT, + run_dir: str | Path, + *, + step: int, + tokens_seen: int, + train_tokens: int, + config: dict[str, Any], + seed: int, +) -> dict[str, Any]: + try: + import weightwatcher as ww + except ImportError as exc: + raise RuntimeError("WeightWatcher is required; install NGB dependencies in the active conda environment") from exc + + run_dir = Path(run_dir) + spectral_root = run_dir / "spectral" + raw_root = spectral_root / "raw" + raw_root.mkdir(parents=True, exist_ok=True) + raw_path = raw_root / f"weightwatcher_step_{int(step):07d}.csv" + if raw_path.is_file(): + frame = pd.read_csv(raw_path) + return summarize_spectral_frame(frame, step=step, tokens_seen=tokens_seen, epoch=tokens_seen / max(1, int(train_tokens))) + + py_state = random.getstate() + np_state = np.random.get_state() + torch_state = torch.random.get_rng_state() + accelerator_state = _capture_accelerator_rng() + diagnostic_seed = int(seed) + 1_000_003 + int(step) + random.seed(diagnostic_seed) + np.random.seed(diagnostic_seed % (2**32 - 1)) + torch.manual_seed(diagnostic_seed) + + try: + holder = WeightMatrixHolder(model) + details = ww.WeightWatcher(model=holder).analyze( + ERG=True, + randomize=True, + plot=False, + min_evals=int(config.get("min_evals", 20)), + ) + if details is None or len(details) == 0: + raise RuntimeError("WeightWatcher returned no transformer-matrix rows") + frame = _attach_matrix_metadata(pd.DataFrame(details), holder.matrix_metadata) + required = ("alpha", "ERG_gap", "num_traps") + missing = [column for column in required if column not in frame.columns] + if missing: + raise RuntimeError("WeightWatcher omitted required columns: " + ", ".join(missing)) + required_values = frame[list(required)].apply(pd.to_numeric, errors="coerce") + if not np.isfinite(required_values.to_numpy(dtype=float)).all(): + raise RuntimeError("WeightWatcher required alpha/ERG_gap/num_traps values are non-finite") + epoch = tokens_seen / max(1, int(train_tokens)) + frame.insert(0, "step", int(step)) + frame.insert(1, "tokens_seen", int(tokens_seen)) + frame.insert(2, "epoch", float(epoch)) + frame.insert(3, "diagnostic_seed", int(diagnostic_seed)) + _atomic_csv(raw_path, frame) + _append_deduplicated(spectral_root / "layers.csv", frame, keys=["step", "matrix_name"]) + summary = summarize_spectral_frame(frame, step=step, tokens_seen=tokens_seen, epoch=epoch) + _append_deduplicated(spectral_root / "summary.csv", pd.DataFrame([summary]), keys=["step"]) + status = { + "step": int(step), + "tokens_seen": int(tokens_seen), + "epoch": float(epoch), + "completed": True, + "matrix_count": len(holder.matrix_metadata), + "raw_path": str(raw_path), + } + (spectral_root / f"status_step_{int(step):07d}.json").write_text(json.dumps(status, indent=2, sort_keys=True), encoding="utf-8") + return summary + except Exception as exc: + status = {"step": int(step), "tokens_seen": int(tokens_seen), "completed": False, "error_type": type(exc).__name__, "error": str(exc)} + (spectral_root / f"status_step_{int(step):07d}.json").write_text(json.dumps(status, indent=2, sort_keys=True), encoding="utf-8") + if bool(config.get("strict", True)): + raise + print(f"[ngb-ww] WARNING step={step}: {type(exc).__name__}: {exc}", flush=True) + return status + finally: + random.setstate(py_state) + np.random.set_state(np_state) + torch.random.set_rng_state(torch_state) + _restore_accelerator_rng(accelerator_state) From 89918045fe8601912eccb4bee7601e7b21ab0bef Mon Sep 17 00:00:00 2001 From: Charles Martin Date: Tue, 11 Aug 2026 09:09:22 -0700 Subject: [PATCH 09/22] Add dynamic NGB completion validation --- baseline/ngb/src/rg_ngb/completion.py | 216 ++++++++++++++++++++++++++ 1 file changed, 216 insertions(+) create mode 100644 baseline/ngb/src/rg_ngb/completion.py diff --git a/baseline/ngb/src/rg_ngb/completion.py b/baseline/ngb/src/rg_ngb/completion.py new file mode 100644 index 0000000..a7f41fd --- /dev/null +++ b/baseline/ngb/src/rg_ngb/completion.py @@ -0,0 +1,216 @@ +"""Validation for completed NGB v4 experiment directories.""" + +from __future__ import annotations + +import json +import math +from pathlib import Path +from typing import Any, NoReturn + +import numpy as np +import pandas as pd +import torch + +_REQUIRED_FILES = ( + "run_complete.json", + "manifest.json", + "metrics.csv", + "epoch_metrics.csv", + "checkpoint_latest.pt", + "checkpoint_best.pt", + "checkpoint_final.pt", + "test_results.json", + "spectral/layers.csv", + "spectral/summary.csv", +) + + +class CompletedRunValidationError(RuntimeError): + pass + + +def _fail(message: str) -> NoReturn: + raise CompletedRunValidationError( + "completed NGB run is stale or inconsistent: " + + message + + ". Use the correct protocol results directory or rerun with explicit overwrite." + ) + + +def _read_json(path: Path) -> dict[str, Any]: + try: + payload = json.loads(path.read_text(encoding="utf-8")) + except (OSError, json.JSONDecodeError) as exc: + _fail(f"could not read valid JSON from {path}: {exc}") + if not isinstance(payload, dict): + _fail(f"{path} is not a JSON object") + return payload + + +def _read_csv(path: Path) -> pd.DataFrame: + try: + frame = pd.read_csv(path) + except Exception as exc: + _fail(f"could not read {path}: {exc}") + if frame.empty: + _fail(f"{path} is empty") + return frame + + +def _integer(value: Any, label: str) -> int: + try: + return int(value) + except (TypeError, ValueError): + _fail(f"{label} is not an integer: {value!r}") + + +def _expect(observed: Any, expected: Any, label: str) -> None: + if observed != expected: + _fail(f"{label} mismatch; observed {observed!r}, expected {expected!r}") + + +def _unique_steps(frame: pd.DataFrame, label: str) -> tuple[int, ...]: + if "step" not in frame.columns: + _fail(f"{label} has no step column") + values = pd.to_numeric(frame["step"], errors="coerce").to_numpy(dtype=float) + if not np.isfinite(values).all() or not np.allclose(values, np.rint(values)): + _fail(f"{label} contains invalid step values") + steps = tuple(int(value) for value in np.rint(values)) + if len(steps) != len(set(steps)): + _fail(f"{label} contains duplicate step rows") + return steps + + +def _load_checkpoint(path: Path) -> dict[str, Any]: + try: + payload = torch.load(path, map_location="cpu", weights_only=False) + except Exception as exc: + _fail(f"could not load checkpoint {path}: {exc}") + if not isinstance(payload, dict): + _fail(f"checkpoint {path} is not a mapping") + return payload + + +def validate_completed_run( + run_dir: str | Path, + *, + expected_fingerprint: str | None = None, + expected_optimizer: str | None = None, + expected_seed: int | None = None, + expected_total_steps: int | None = None, + expected_matrix_count: int | None = None, + verify_checkpoints: bool = True, +) -> dict[str, Any]: + root = Path(run_dir) + missing = [str(root / relative) for relative in _REQUIRED_FILES if not (root / relative).is_file() or (root / relative).stat().st_size == 0] + if missing: + _fail("missing required artifacts: " + ", ".join(missing)) + + completion = _read_json(root / "run_complete.json") + manifest = _read_json(root / "manifest.json") + test_results = _read_json(root / "test_results.json") + if completion.get("completed") is not True: + _fail("run_complete.json does not declare completed=true") + + recorded_fingerprint = str(completion.get("fingerprint", "")) + fingerprint = str(expected_fingerprint) if expected_fingerprint is not None else recorded_fingerprint + recorded_optimizer = str(completion.get("optimizer", "")) + optimizer = str(expected_optimizer) if expected_optimizer is not None else recorded_optimizer + recorded_seed = _integer(completion.get("seed"), "completion seed") + seed = int(expected_seed) if expected_seed is not None else recorded_seed + recorded_steps = _integer(completion.get("optimizer_steps"), "completion optimizer_steps") + total_steps = int(expected_total_steps) if expected_total_steps is not None else recorded_steps + best_step = _integer(completion.get("best_validation_step"), "completion best_validation_step") + if not fingerprint or not optimizer: + _fail("completion record lacks fingerprint or optimizer") + + model_cfg = manifest.get("model") + if not isinstance(model_cfg, dict): + _fail("manifest model section is missing") + recorded_matrix_count = 6 * _integer(model_cfg.get("n_layer"), "manifest model.n_layer") + matrix_count = int(expected_matrix_count) if expected_matrix_count is not None else recorded_matrix_count + + _expect(recorded_fingerprint, fingerprint, "completion fingerprint") + _expect(recorded_optimizer, optimizer, "completion optimizer") + _expect(recorded_seed, seed, "completion seed") + _expect(recorded_steps, total_steps, "completion optimizer_steps") + _expect(str(manifest.get("protocol_fingerprint", "")), fingerprint, "manifest fingerprint") + _expect(str(manifest.get("optimizer", "")), optimizer, "manifest optimizer") + _expect(_integer(manifest.get("seed"), "manifest seed"), seed, "manifest seed") + _expect(_integer(manifest.get("max_steps"), "manifest max_steps"), total_steps, "manifest max_steps") + _expect(recorded_matrix_count, matrix_count, "manifest transformer matrix count") + + final_test = test_results.get("final") + selected_test = test_results.get("validation_selected") + if not isinstance(final_test, dict) or not isinstance(selected_test, dict): + _fail("test_results.json lacks final or validation_selected results") + _expect(_integer(final_test.get("step"), "final test step"), total_steps, "final test step") + _expect(_integer(selected_test.get("step"), "selected test step"), best_step, "selected test step") + + metrics = _read_csv(root / "metrics.csv") + epoch_metrics = _read_csv(root / "epoch_metrics.csv") + layers = _read_csv(root / "spectral" / "layers.csv") + summary = _read_csv(root / "spectral" / "summary.csv") + metric_steps = _unique_steps(metrics, "metrics.csv") + epoch_steps = _unique_steps(epoch_metrics, "epoch_metrics.csv") + summary_steps = _unique_steps(summary, "spectral/summary.csv") + for label, steps in (("metrics.csv", metric_steps), ("epoch_metrics.csv", epoch_steps)): + if 0 not in steps or total_steps not in steps or max(steps) != total_steps: + _fail(f"{label} does not span step zero through {total_steps}") + + if "test_monitoring_only" not in epoch_metrics.columns: + _fail("epoch_metrics.csv lacks test_monitoring_only") + policy = pd.to_numeric(epoch_metrics["test_monitoring_only"], errors="coerce") + if policy.isna().any() or not policy.astype(int).eq(1).all(): + _fail("epoch_metrics.csv violates the monitoring-only test policy") + + required_layer_columns = {"step", "matrix_name", "alpha", "ERG_gap", "num_traps"} + missing_columns = required_layer_columns.difference(layers.columns) + if missing_columns: + _fail("spectral/layers.csv is missing columns " + ", ".join(sorted(missing_columns))) + if layers.duplicated(["step", "matrix_name"]).any(): + _fail("spectral/layers.csv has duplicate step/matrix rows") + layer_step_frame = layers[["step"]].drop_duplicates().sort_values("step") + layer_steps = _unique_steps(layer_step_frame, "spectral/layers.csv") + if set(summary_steps) != set(epoch_steps) or set(layer_steps) != set(epoch_steps): + _fail("spectral steps do not match epoch_metrics.csv") + if not layers.groupby("step")["matrix_name"].nunique().eq(matrix_count).all(): + _fail(f"spectral/layers.csv does not contain {matrix_count} matrices per epoch") + if "n_matrices" not in summary.columns: + _fail("spectral/summary.csv lacks n_matrices") + counts = pd.to_numeric(summary["n_matrices"], errors="coerce") + if counts.isna().any() or not counts.astype(int).eq(matrix_count).all(): + _fail(f"spectral/summary.csv does not report {matrix_count} matrices per epoch") + + if "checkpoint_path" not in epoch_metrics.columns: + _fail("epoch_metrics.csv lacks checkpoint_path") + resolved_paths: list[Path] = [] + for value in epoch_metrics["checkpoint_path"]: + recorded = Path(str(value)) + candidate = recorded if recorded.is_file() else root / "epoch_checkpoints" / recorded.name + if not candidate.is_file() or candidate.stat().st_size == 0: + _fail(f"missing epoch checkpoint: {recorded}") + resolved_paths.append(candidate.resolve()) + if len(resolved_paths) != len(set(resolved_paths)) or len(resolved_paths) != len(set(epoch_steps)): + _fail("epoch checkpoint inventory is duplicated or incomplete") + + if verify_checkpoints: + try: + best_loss = float(completion.get("best_validation_loss")) + except (TypeError, ValueError): + _fail("run_complete.json has invalid best_validation_loss") + for filename, expected_step in (("checkpoint_latest.pt", total_steps), ("checkpoint_final.pt", total_steps), ("checkpoint_best.pt", best_step)): + payload = _load_checkpoint(root / filename) + _expect(str(payload.get("fingerprint", "")), fingerprint, f"{filename} fingerprint") + _expect(str(payload.get("optimizer_name", "")), optimizer, f"{filename} optimizer") + _expect(_integer(payload.get("seed"), f"{filename} seed"), seed, f"{filename} seed") + _expect(_integer(payload.get("step"), f"{filename} step"), expected_step, f"{filename} step") + _expect(_integer(payload.get("best_validation_step"), f"{filename} best step"), best_step, f"{filename} best step") + try: + stored_loss = float(payload.get("best_validation_loss")) + except (TypeError, ValueError): + _fail(f"{filename} has invalid best_validation_loss") + if not math.isclose(stored_loss, best_loss, rel_tol=1e-12, abs_tol=1e-12): + _fail(f"{filename} best_validation_loss does not match completion") + + return completion From 375d1d926570a8f90b3a2d411d8ad9ba4b07bc93 Mon Sep 17 00:00:00 2001 From: Charles Martin Date: Tue, 11 Aug 2026 09:09:51 -0700 Subject: [PATCH 10/22] Add NGB run utilities --- baseline/ngb/src/rg_ngb/run_utils.py | 170 +++++++++++++++++++++++++++ 1 file changed, 170 insertions(+) create mode 100644 baseline/ngb/src/rg_ngb/run_utils.py diff --git a/baseline/ngb/src/rg_ngb/run_utils.py b/baseline/ngb/src/rg_ngb/run_utils.py new file mode 100644 index 0000000..78b33c9 --- /dev/null +++ b/baseline/ngb/src/rg_ngb/run_utils.py @@ -0,0 +1,170 @@ +from __future__ import annotations + +import csv +import json +from pathlib import Path + +import pandas as pd +import torch + +from rg_nanogpt_one_head.evaluation import evaluate_bleu, evaluate_probe + +from .completion import CompletedRunValidationError, validate_completed_run +from .config import expected_matrix_count, tokens_per_step +from .model import GPT + +METRIC_FIELDS = [ + "step", + "tokens_seen", + "epoch", + "elapsed_sec", + "tokens_per_sec", + "primary_lr", + "auxiliary_lr", + "train_loss", + "train_perplexity", + "train_accuracy", + "val_loss", + "val_perplexity", + "val_accuracy", + "test_loss", + "test_perplexity", + "test_accuracy", + "test_bleu", + "val_generalization_gap", + "test_generalization_gap", + "grad_norm_pre_clip", + "grad_norm_post_clip", + "gradient_clipped", + "weight_norm", + "update_norm_since_eval", + "update_to_weight_ratio", + "mps_current_allocated_mb", + "mps_driver_allocated_mb", +] +EPOCH_FIELDS = [*METRIC_FIELDS, "nominal_epoch", "checkpoint_path", "test_monitoring_only"] + + +def run_directory(results_root: str | Path, optimizer: str, seed: int) -> Path: + return Path(results_root) / str(optimizer) / f"seed_{int(seed)}" + + +def run_is_complete(results_root: str | Path, optimizer: str, seed: int) -> bool: + run_dir = run_directory(results_root, optimizer, seed) + try: + validate_completed_run( + run_dir, + expected_optimizer=str(optimizer), + expected_seed=int(seed), + verify_checkpoints=False, + ) + except (CompletedRunValidationError, OSError): + return False + return True + + +def prepare_csv(path: Path, fields: list[str], resume_step: int | None) -> None: + path.parent.mkdir(parents=True, exist_ok=True) + if path.is_file() and resume_step is not None: + frame = pd.read_csv(path) + if "step" in frame.columns: + frame = frame[pd.to_numeric(frame["step"], errors="coerce") < int(resume_step)] + temporary = path.with_suffix(path.suffix + ".tmp") + frame.to_csv(temporary, index=False) + temporary.replace(path) + if not path.is_file() or path.stat().st_size == 0: + with path.open("w", newline="", encoding="utf-8") as handle: + csv.DictWriter(handle, fieldnames=fields).writeheader() + + +def truncate_spectral_after(run_dir: Path, resume_step: int) -> None: + spectral_root = run_dir / "spectral" + for filename in ("layers.csv", "summary.csv"): + path = spectral_root / filename + if path.is_file(): + frame = pd.read_csv(path) + if "step" in frame.columns: + frame = frame[pd.to_numeric(frame["step"], errors="coerce") < int(resume_step)] + temporary = path.with_suffix(path.suffix + ".tmp") + frame.to_csv(temporary, index=False) + temporary.replace(path) + raw_root = spectral_root / "raw" + if raw_root.is_dir(): + for path in raw_root.glob("weightwatcher_step_*.csv"): + try: + step = int(path.stem.rsplit("_", 1)[-1]) + except ValueError: + continue + if step >= int(resume_step): + path.unlink(missing_ok=True) + for path in spectral_root.glob("status_step_*.json"): + try: + step = int(path.stem.rsplit("_", 1)[-1]) + except ValueError: + continue + if step >= int(resume_step): + path.unlink(missing_ok=True) + + +def write_manifest( + run_dir: Path, + *, + cfg: dict, + data_metadata: dict, + optimizer_name: str, + profile: dict, + seed: int, + device: torch.device, + total_steps: int, + warmup: int, + fingerprint: str, + model: GPT, +) -> None: + payload = { + "schema_version": 2, + "protocol": cfg["protocol"], + "optimizer": optimizer_name, + "optimizer_profile": profile, + "seed": int(seed), + "device": str(device), + "torch_version": torch.__version__, + "model": cfg["model"], + "parameter_count": model.parameter_count(), + "transformer_matrix_count": expected_matrix_count(cfg), + "data_metadata": data_metadata, + "training": cfg["training"], + "evaluation": cfg["evaluation"], + "weightwatcher": cfg["weightwatcher"], + "tokens_per_step": tokens_per_step(cfg), + "max_steps": int(total_steps), + "warmup_steps": int(warmup), + "planned_training_tokens": int(total_steps * tokens_per_step(cfg)), + "protocol_fingerprint": fingerprint, + "test_policy": "fixed test probes are monitoring-only and never select checkpoints or tune schedules", + "bleu_policy": "fixed greedy held-out continuation BLEU; secondary diagnostic", + } + temporary = run_dir / "manifest.json.tmp" + temporary.write_text(json.dumps(payload, indent=2, sort_keys=True, default=str), encoding="utf-8") + temporary.replace(run_dir / "manifest.json") + + +def checkpoint_eval( + checkpoint: Path, + *, + model: GPT, + test_probe, + bleu_probe, + device: torch.device, + bleu_batch_size: int, +) -> dict[str, float]: + payload = torch.load(checkpoint, map_location="cpu", weights_only=False) + model.load_state_dict(payload["model"]) + metrics = evaluate_probe(model, test_probe, device) + bleu = evaluate_bleu(model, bleu_probe, device=device, batch_size=bleu_batch_size) + return { + "step": int(payload["step"]), + "loss": float(metrics["loss"]), + "perplexity": float(metrics["perplexity"]), + "accuracy": float(metrics["accuracy"]), + "bleu": float(bleu["bleu"]), + } From 654e760cd2962c3c5fa7f92b314d20f3998088c6 Mon Sep 17 00:00:00 2001 From: Charles Martin Date: Tue, 11 Aug 2026 09:10:31 -0700 Subject: [PATCH 11/22] Add NGB training loop --- baseline/ngb/src/rg_ngb/train_loop.py | 268 ++++++++++++++++++++++++++ 1 file changed, 268 insertions(+) create mode 100644 baseline/ngb/src/rg_ngb/train_loop.py diff --git a/baseline/ngb/src/rg_ngb/train_loop.py b/baseline/ngb/src/rg_ngb/train_loop.py new file mode 100644 index 0000000..f794b50 --- /dev/null +++ b/baseline/ngb/src/rg_ngb/train_loop.py @@ -0,0 +1,268 @@ +from __future__ import annotations + +import csv +import math +from pathlib import Path +import time + +import torch + +from rg_nanogpt_one_head.checkpoints import save_epoch_model_checkpoint, save_training_checkpoint +from rg_nanogpt_one_head.evaluation import evaluate_bleu, evaluate_probe, random_batch +from rg_nanogpt_one_head.optimizers import optimizer_step, set_learning_rates, zero_grad +from rg_nanogpt_one_head.runtime import ( + empty_mps_cache, + gradient_norm, + model_weight_norm, + mps_memory_megabytes, + parameter_snapshot, + synchronize, + update_norm, +) + +from .spectral import run_weightwatcher + + +def execute_training_loop( + *, + cfg: dict, + model, + handles, + arrays: dict, + train_probe, + val_probe, + test_probe, + bleu_probe, + device: torch.device, + optimizer_name: str, + seed: int, + train_tokens: int, + total_steps: int, + warmup: int, + start_step: int, + best_validation_loss: float, + best_validation_step: int, + elapsed_offset: float, + fingerprint: str, + train_generator: torch.Generator, + epoch_steps: dict[int, float], + metrics_writer: csv.DictWriter, + metrics_handle, + epoch_writer: csv.DictWriter, + epoch_handle, + run_dir: Path, + latest_checkpoint: Path, + best_checkpoint: Path, + progress: bool, +) -> tuple[float, int, float]: + batch_size = int(cfg["training"]["batch_size"]) + grad_accum = int(cfg["training"]["grad_accum_steps"]) + block_size = int(cfg["model"]["block_size"]) + step_tokens = batch_size * grad_accum * block_size + eval_cfg = cfg["evaluation"] + + previous_snapshot = parameter_snapshot(model) + last_grad_pre = float("nan") + last_grad_post = float("nan") + last_clipped = False + started = time.time() + last_update_lrs = { + "primary": 0.0, + "auxiliary": 0.0 if any(handle.role == "auxiliary" for handle in handles) else float("nan"), + } + if start_step > 0: + for handle in handles: + last_update_lrs[handle.role] = float(handle.lr) + + for completed_steps in range(start_step, total_steps + 1): + schedule_index = min(completed_steps, total_steps - 1) + next_update_lrs = set_learning_rates( + handles, + update_index=schedule_index, + total_steps=total_steps, + warmup_steps=warmup, + ) + epoch_due = completed_steps in epoch_steps + evaluation_due = ( + completed_steps % int(cfg["training"]["eval_interval_steps"]) == 0 + or epoch_due + or completed_steps == total_steps + ) + + if evaluation_due: + synchronize(device) + train_metrics = evaluate_probe(model, train_probe, device) + val_metrics = evaluate_probe(model, val_probe, device) + elapsed = elapsed_offset + time.time() - started + if val_metrics["loss"] < best_validation_loss: + best_validation_loss = float(val_metrics["loss"]) + best_validation_step = int(completed_steps) + save_training_checkpoint( + best_checkpoint, + model=model, + handles=handles, + step=completed_steps, + best_validation_loss=best_validation_loss, + best_validation_step=best_validation_step, + elapsed_seconds=elapsed, + fingerprint=fingerprint, + cfg=cfg, + optimizer_name=optimizer_name, + seed=int(seed), + train_generator=train_generator, + ) + + test_metrics = {"loss": float("nan"), "perplexity": float("nan"), "accuracy": float("nan")} + bleu_metrics = {"bleu": float("nan")} + if epoch_due or completed_steps == total_steps: + test_metrics = evaluate_probe(model, test_probe, device) + bleu_metrics = evaluate_bleu( + model, + bleu_probe, + device=device, + batch_size=int(eval_cfg["bleu_batch_size"]), + ) + + tokens_seen = int(completed_steps * step_tokens) + actual_epoch = tokens_seen / max(1, train_tokens) + current_snapshot = parameter_snapshot(model) + delta_norm = update_norm(previous_snapshot, current_snapshot) + previous_snapshot = current_snapshot + weight_norm = model_weight_norm(model) + current_mps, driver_mps = mps_memory_megabytes(device) + row = { + "step": int(completed_steps), + "tokens_seen": tokens_seen, + "epoch": float(actual_epoch), + "elapsed_sec": float(elapsed), + "tokens_per_sec": tokens_seen / max(elapsed, 1e-9), + "primary_lr": float(last_update_lrs.get("primary", float("nan"))), + "auxiliary_lr": float(last_update_lrs.get("auxiliary", float("nan"))), + "train_loss": float(train_metrics["loss"]), + "train_perplexity": float(train_metrics["perplexity"]), + "train_accuracy": float(train_metrics["accuracy"]), + "val_loss": float(val_metrics["loss"]), + "val_perplexity": float(val_metrics["perplexity"]), + "val_accuracy": float(val_metrics["accuracy"]), + "test_loss": float(test_metrics["loss"]), + "test_perplexity": float(test_metrics["perplexity"]), + "test_accuracy": float(test_metrics["accuracy"]), + "test_bleu": float(bleu_metrics["bleu"]), + "val_generalization_gap": float(val_metrics["loss"] - train_metrics["loss"]), + "test_generalization_gap": float(test_metrics["loss"] - train_metrics["loss"]), + "grad_norm_pre_clip": float(last_grad_pre), + "grad_norm_post_clip": float(last_grad_post), + "gradient_clipped": int(last_clipped), + "weight_norm": float(weight_norm), + "update_norm_since_eval": float(delta_norm), + "update_to_weight_ratio": float(delta_norm / max(weight_norm, 1e-30)), + "mps_current_allocated_mb": float(current_mps), + "mps_driver_allocated_mb": float(driver_mps), + } + metrics_writer.writerow(row) + metrics_handle.flush() + + if epoch_due: + nominal_epoch = float(epoch_steps[completed_steps]) + checkpoint_path = save_epoch_model_checkpoint( + run_dir, + model=model, + step=completed_steps, + nominal_epoch=nominal_epoch, + actual_epoch=actual_epoch, + fingerprint=fingerprint, + cfg=cfg, + optimizer_name=optimizer_name, + seed=int(seed), + ) + epoch_writer.writerow({**row, "nominal_epoch": nominal_epoch, "checkpoint_path": str(checkpoint_path), "test_monitoring_only": 1}) + epoch_handle.flush() + ww_summary = run_weightwatcher( + model, + run_dir, + step=completed_steps, + tokens_seen=tokens_seen, + train_tokens=train_tokens, + config=cfg["weightwatcher"], + seed=int(seed), + ) + if progress: + print( + "[ngb-ww] " + f"protocol={cfg['protocol']['name']} optimizer={optimizer_name} seed={seed} " + f"epoch={nominal_epoch:.2f} alpha={ww_summary.get('alpha_median', float('nan')):.3f} " + f"ERG_gap={ww_summary.get('ERG_gap_median', float('nan')):.3f} " + f"num_traps={ww_summary.get('num_traps_mean', float('nan')):.2f}", + flush=True, + ) + if bool(cfg["runtime"].get("empty_mps_cache_after_weightwatcher", True)): + empty_mps_cache(device) + + if progress: + remaining = total_steps - completed_steps + rate = completed_steps / max(elapsed, 1e-9) + eta = remaining / rate if rate > 0 else float("nan") + eta_text = "unknown" if not math.isfinite(eta) else f"{eta / 60:.1f}m" + print( + "[ngb-train] " + f"protocol={cfg['protocol']['name']} optimizer={optimizer_name} seed={seed} " + f"step={completed_steps}/{total_steps} epoch={actual_epoch:.3f} " + f"last_lr={last_update_lrs.get('primary', float('nan')):.3e} " + f"next_lr={next_update_lrs.get('primary', float('nan')):.3e} " + f"train_loss={train_metrics['loss']:.4f} val_loss={val_metrics['loss']:.4f} " + f"val_ppl={val_metrics['perplexity']:.2f} val_acc={100 * val_metrics['accuracy']:.2f}% eta={eta_text}", + flush=True, + ) + + if completed_steps == total_steps: + break + + zero_grad(handles) + for _ in range(grad_accum): + x_cpu, y_cpu = random_batch( + arrays["train"], + batch_size=batch_size, + block_size=block_size, + generator=train_generator, + ) + x = x_cpu.to(device) + y = y_cpu.to(device) + _, loss = model(x, y) + if loss is None: + raise RuntimeError("training forward pass did not return loss") + (loss / grad_accum).backward() + + grad_pre_tensor = gradient_norm(model.parameters()) + last_grad_pre = float(grad_pre_tensor.detach().cpu()) + clip = float(cfg["training"]["grad_clip"]) + if clip > 0: + torch.nn.utils.clip_grad_norm_(model.parameters(), clip) + grad_post_tensor = gradient_norm(model.parameters()) + last_grad_post = float(grad_post_tensor.detach().cpu()) + last_clipped = bool(last_grad_pre > clip) if clip > 0 else False + optimizer_step(handles) + last_update_lrs = dict(next_update_lrs) + + new_step = completed_steps + 1 + checkpoint_due = ( + new_step % int(cfg["training"]["checkpoint_interval_steps"]) == 0 + or new_step in epoch_steps + or new_step == total_steps + ) + if checkpoint_due: + save_training_checkpoint( + latest_checkpoint, + model=model, + handles=handles, + step=new_step, + best_validation_loss=best_validation_loss, + best_validation_step=best_validation_step, + elapsed_seconds=elapsed_offset + time.time() - started, + fingerprint=fingerprint, + cfg=cfg, + optimizer_name=optimizer_name, + seed=int(seed), + train_generator=train_generator, + ) + + return float(best_validation_loss), int(best_validation_step), float(elapsed_offset + time.time() - started) From 126d242de6f1d6aeb9b3a0f41c64748514620940 Mon Sep 17 00:00:00 2001 From: Charles Martin Date: Tue, 11 Aug 2026 09:11:10 -0700 Subject: [PATCH 12/22] Add NGB run engine --- baseline/ngb/src/rg_ngb/engine.py | 258 ++++++++++++++++++++++++++++++ 1 file changed, 258 insertions(+) create mode 100644 baseline/ngb/src/rg_ngb/engine.py diff --git a/baseline/ngb/src/rg_ngb/engine.py b/baseline/ngb/src/rg_ngb/engine.py new file mode 100644 index 0000000..2323256 --- /dev/null +++ b/baseline/ngb/src/rg_ngb/engine.py @@ -0,0 +1,258 @@ +from __future__ import annotations + +import csv +import json +from pathlib import Path +import shutil + +import torch + +from rg_nanogpt_one_head.checkpoints import load_training_checkpoint, save_training_checkpoint +from rg_nanogpt_one_head.evaluation import fixed_bleu_probe, fixed_probe +from rg_nanogpt_one_head.optimizers import make_optimizer_handles +from rg_nanogpt_one_head.runtime import choose_device, configure_runtime, seed_everything + +from .completion import validate_completed_run +from .config import ( + SUPPORTED_OPTIMIZERS, + epoch_step_map, + expected_matrix_count, + max_steps, + optimizer_profile, + protocol_fingerprint, + tokens_per_step, + warmup_steps, +) +from .data import load_memmaps +from .model import GPT, GPTConfig +from .run_utils import ( + EPOCH_FIELDS, + METRIC_FIELDS, + checkpoint_eval, + prepare_csv, + run_directory, + truncate_spectral_after, + write_manifest, +) +from .train_loop import execute_training_loop + + +def run_one( + *, + cfg: dict, + data_root: str | Path, + results_root: str | Path, + optimizer_name: str, + seed: int, + device: str = "auto", + resume: bool = True, + overwrite: bool = False, + progress: bool = True, +) -> Path: + optimizer_name = str(optimizer_name).lower() + if optimizer_name not in SUPPORTED_OPTIMIZERS: + raise ValueError(f"unsupported optimizer: {optimizer_name}") + if resume and overwrite: + raise ValueError("resume and overwrite are mutually exclusive") + + data_root = Path(data_root) + results_root = Path(results_root) + run_dir = run_directory(results_root, optimizer_name, int(seed)) + completion_path = run_dir / "run_complete.json" + if run_dir.exists() and overwrite: + shutil.rmtree(run_dir) + if run_dir.exists() and not resume: + raise FileExistsError(f"run directory exists: {run_dir}; enable resume or overwrite") + run_dir.mkdir(parents=True, exist_ok=True) + + data_metadata, arrays = load_memmaps(data_root, cfg) + train_tokens = int(data_metadata["splits"]["train"]) + total_steps = max_steps(cfg, train_tokens) + profile = optimizer_profile(cfg, optimizer_name) + warmup = warmup_steps(profile, total_steps) + matrix_count = expected_matrix_count(cfg) + fingerprint = protocol_fingerprint( + cfg, + optimizer=optimizer_name, + seed=int(seed), + data_metadata=data_metadata, + ) + if completion_path.is_file(): + validate_completed_run( + run_dir, + expected_fingerprint=fingerprint, + expected_optimizer=optimizer_name, + expected_seed=int(seed), + expected_total_steps=total_steps, + expected_matrix_count=matrix_count, + verify_checkpoints=True, + ) + if progress: + print(f"[ngb-train] reuse verified completed {cfg['protocol']['name']} {optimizer_name} seed={seed}") + return run_dir + + resolved_device = choose_device(device) + configure_runtime(resolved_device, cfg) + seed_everything(int(seed)) + model = GPT(GPTConfig(**cfg["model"])).to(resolved_device) + handles = make_optimizer_handles(model, profile) + train_generator = torch.Generator(device="cpu").manual_seed(int(seed) + 11) + + batch_size = int(cfg["training"]["batch_size"]) + eval_batches = int(cfg["training"]["eval_batches"]) + block_size = int(cfg["model"]["block_size"]) + epoch_steps = epoch_step_map(cfg, train_tokens) + eval_cfg = cfg["evaluation"] + train_probe = fixed_probe(arrays["train"], batch_size=batch_size, block_size=block_size, n_batches=eval_batches, seed=int(eval_cfg["train_probe_seed"])) + val_probe = fixed_probe(arrays["val"], batch_size=batch_size, block_size=block_size, n_batches=eval_batches, seed=int(eval_cfg["validation_probe_seed"])) + test_probe = fixed_probe(arrays["test"], batch_size=batch_size, block_size=block_size, n_batches=eval_batches, seed=int(eval_cfg["test_probe_seed"])) + bleu_probe = fixed_bleu_probe( + arrays["test"], + examples=int(eval_cfg["bleu_examples"]), + prompt_tokens=int(eval_cfg["bleu_prompt_tokens"]), + continuation_tokens=int(eval_cfg["bleu_continuation_tokens"]), + seed=int(eval_cfg["bleu_probe_seed"]), + ) + + start_step = 0 + best_validation_loss = float("inf") + best_validation_step = 0 + elapsed_offset = 0.0 + latest_checkpoint = run_dir / "checkpoint_latest.pt" + best_checkpoint = run_dir / "checkpoint_best.pt" + final_checkpoint = run_dir / "checkpoint_final.pt" + if resume and latest_checkpoint.is_file(): + start_step, best_validation_loss, best_validation_step, elapsed_offset = load_training_checkpoint( + latest_checkpoint, + model=model, + handles=handles, + expected_fingerprint=fingerprint, + train_generator=train_generator, + ) + model.to(resolved_device) + truncate_spectral_after(run_dir, start_step) + if progress: + print(f"[ngb-train] resume {cfg['protocol']['name']} {optimizer_name} seed={seed} step={start_step}") + elif run_dir.exists() and any(run_dir.iterdir()) and resume: + nontrivial = [path for path in run_dir.iterdir() if path.name != "manifest.json"] + if nontrivial and not latest_checkpoint.is_file(): + raise FileNotFoundError(f"cannot resume {run_dir}: checkpoint_latest.pt is missing") + + write_manifest( + run_dir, + cfg=cfg, + data_metadata=data_metadata, + optimizer_name=optimizer_name, + profile=profile, + seed=int(seed), + device=resolved_device, + total_steps=total_steps, + warmup=warmup, + fingerprint=fingerprint, + model=model, + ) + + metrics_path = run_dir / "metrics.csv" + epoch_metrics_path = run_dir / "epoch_metrics.csv" + prepare_csv(metrics_path, METRIC_FIELDS, start_step if start_step else None) + prepare_csv(epoch_metrics_path, EPOCH_FIELDS, start_step if start_step else None) + with ( + metrics_path.open("a", newline="", encoding="utf-8") as metrics_handle, + epoch_metrics_path.open("a", newline="", encoding="utf-8") as epoch_handle, + ): + best_validation_loss, best_validation_step, elapsed_total = execute_training_loop( + cfg=cfg, + model=model, + handles=handles, + arrays=arrays, + train_probe=train_probe, + val_probe=val_probe, + test_probe=test_probe, + bleu_probe=bleu_probe, + device=resolved_device, + optimizer_name=optimizer_name, + seed=int(seed), + train_tokens=train_tokens, + total_steps=total_steps, + warmup=warmup, + start_step=start_step, + best_validation_loss=best_validation_loss, + best_validation_step=best_validation_step, + elapsed_offset=elapsed_offset, + fingerprint=fingerprint, + train_generator=train_generator, + epoch_steps=epoch_steps, + metrics_writer=csv.DictWriter(metrics_handle, fieldnames=METRIC_FIELDS), + metrics_handle=metrics_handle, + epoch_writer=csv.DictWriter(epoch_handle, fieldnames=EPOCH_FIELDS), + epoch_handle=epoch_handle, + run_dir=run_dir, + latest_checkpoint=latest_checkpoint, + best_checkpoint=best_checkpoint, + progress=progress, + ) + + for checkpoint in (final_checkpoint, latest_checkpoint): + save_training_checkpoint( + checkpoint, + model=model, + handles=handles, + step=total_steps, + best_validation_loss=best_validation_loss, + best_validation_step=best_validation_step, + elapsed_seconds=elapsed_total, + fingerprint=fingerprint, + cfg=cfg, + optimizer_name=optimizer_name, + seed=int(seed), + train_generator=train_generator, + ) + + final_state = torch.load(final_checkpoint, map_location="cpu", weights_only=False)["model"] + final_test = checkpoint_eval( + final_checkpoint, + model=model, + test_probe=test_probe, + bleu_probe=bleu_probe, + device=resolved_device, + bleu_batch_size=int(eval_cfg["bleu_batch_size"]), + ) + best_test = checkpoint_eval( + best_checkpoint, + model=model, + test_probe=test_probe, + bleu_probe=bleu_probe, + device=resolved_device, + bleu_batch_size=int(eval_cfg["bleu_batch_size"]), + ) + model.load_state_dict(final_state) + model.to(resolved_device) + + test_results = { + "policy": "test is monitoring-only; validation loss selects checkpoint_best.pt", + "final": final_test, + "validation_selected": best_test, + } + (run_dir / "test_results.json").write_text(json.dumps(test_results, indent=2, sort_keys=True), encoding="utf-8") + completion = { + "completed": True, + "protocol": cfg["protocol"]["name"], + "optimizer": optimizer_name, + "seed": int(seed), + "optimizer_steps": int(total_steps), + "train_epochs": float(total_steps * tokens_per_step(cfg) / train_tokens), + "elapsed_seconds": float(elapsed_total), + "best_validation_step": int(best_validation_step), + "best_validation_loss": float(best_validation_loss), + "final_test_loss": float(final_test["loss"]), + "final_test_perplexity": float(final_test["perplexity"]), + "final_test_accuracy": float(final_test["accuracy"]), + "final_test_bleu": float(final_test["bleu"]), + "fingerprint": fingerprint, + } + temporary = run_dir / "run_complete.json.tmp" + temporary.write_text(json.dumps(completion, indent=2, sort_keys=True), encoding="utf-8") + temporary.replace(completion_path) + if progress: + print(f"[ngb-train] complete {cfg['protocol']['name']} {optimizer_name} seed={seed}: {run_dir}") + return run_dir From f2235d19f78805e8a6ec415fa9b8741415cbb3cd Mon Sep 17 00:00:00 2001 From: Charles Martin Date: Tue, 11 Aug 2026 09:13:00 -0700 Subject: [PATCH 13/22] Add corrected NGB multi-seed analysis --- baseline/ngb/src/rg_ngb/analysis.py | 374 ++++++++++++++++++++++++++++ 1 file changed, 374 insertions(+) create mode 100644 baseline/ngb/src/rg_ngb/analysis.py diff --git a/baseline/ngb/src/rg_ngb/analysis.py b/baseline/ngb/src/rg_ngb/analysis.py new file mode 100644 index 0000000..487fd13 --- /dev/null +++ b/baseline/ngb/src/rg_ngb/analysis.py @@ -0,0 +1,374 @@ +from __future__ import annotations + +from itertools import combinations +import json +import math +from pathlib import Path +from typing import Iterable, Sequence + +import matplotlib.pyplot as plt +import numpy as np +import pandas as pd + +from .config import SUPPORTED_OPTIMIZERS +from .run_utils import run_directory, run_is_complete + +OPTIMIZER_LABELS = { + "sgd_momentum": "SGD + Nesterov", + "adamw": "AdamW", + "muon": "Muon + auxiliary AdamW", +} +OPTIMIZER_COLORS = { + "sgd_momentum": "#0072B2", + "adamw": "#D55E00", + "muon": "#009E73", +} +MATRIX_COLORS = { + "W_Q": "#0072B2", + "W_K": "#E69F00", + "W_V": "#009E73", + "W_O": "#D55E00", + "W_MLP_IN": "#CC79A7", + "W_MLP_OUT": "#56B4E9", +} +_T_975 = { + 1: 12.7062047364, + 2: 4.3026527297, + 3: 3.1824463053, + 4: 2.7764451052, + 5: 2.5705818356, + 6: 2.4469118511, + 7: 2.3646242510, + 8: 2.3060041352, + 9: 2.2621571629, + 10: 2.2281388520, + 11: 2.2009851601, + 12: 2.1788128297, + 13: 2.1603686565, + 14: 2.1447866879, + 15: 2.1314495456, + 16: 2.1199052992, + 17: 2.1098155778, + 18: 2.1009220402, + 19: 2.0930240544, + 20: 2.0859634473, + 21: 2.0796138447, + 22: 2.0738730679, + 23: 2.0686576104, + 24: 2.0638985616, + 25: 2.0595385528, + 26: 2.0555294386, + 27: 2.0518305165, + 28: 2.0484071418, + 29: 2.0452296421, + 30: 2.0422724563, +} + + +def mean_ci95(values: Iterable[float]) -> dict[str, float]: + array = np.asarray(list(values), dtype=float) + array = array[np.isfinite(array)] + n = int(array.size) + if n == 0: + return {key: np.nan for key in ("mean", "sd", "sem", "ci95_half_width", "ci95_lower", "ci95_upper")} | {"n": 0} + mean = float(array.mean()) + if n == 1: + return {"n": 1, "mean": mean, "sd": np.nan, "sem": np.nan, "ci95_half_width": np.nan, "ci95_lower": np.nan, "ci95_upper": np.nan} + sd = float(array.std(ddof=1)) + sem = sd / math.sqrt(n) + critical = _T_975.get(n - 1, 1.9599639845) + half = critical * sem + return {"n": n, "mean": mean, "sd": sd, "sem": sem, "ci95_half_width": half, "ci95_lower": mean - half, "ci95_upper": mean + half} + + +def completed_seed_sets(results_root: str | Path, optimizers: Sequence[str] = SUPPORTED_OPTIMIZERS) -> dict[str, set[int]]: + root = Path(results_root) + result: dict[str, set[int]] = {} + for optimizer in optimizers: + seeds: set[int] = set() + optimizer_root = root / optimizer + if optimizer_root.is_dir(): + for directory in optimizer_root.glob("seed_*"): + try: + seed = int(directory.name.split("_", 1)[1]) + except (IndexError, ValueError): + continue + if run_is_complete(root, optimizer, seed): + seeds.add(seed) + result[optimizer] = seeds + return result + + +def discover_matched_seeds( + results_root: str | Path, + *, + optimizers: Sequence[str] = SUPPORTED_OPTIMIZERS, + requested: Sequence[int] | None = None, + minimum: int = 2, +) -> tuple[int, ...]: + sets = completed_seed_sets(results_root, optimizers) + matched = set.intersection(*(sets[optimizer] for optimizer in optimizers)) if optimizers else set() + if requested is not None: + requested_set = {int(seed) for seed in requested} + missing = requested_set.difference(matched) + if missing: + raise FileNotFoundError(f"requested seeds are not complete for every optimizer: {sorted(missing)}") + matched = requested_set + result = tuple(sorted(matched)) + if len(result) < int(minimum): + detail = {key: sorted(value) for key, value in sets.items()} + raise FileNotFoundError(f"need at least {minimum} matched completed seeds; found {result}; per optimizer={detail}") + return result + + +def run_status_table(results_root: str | Path, *, optimizers: Sequence[str] = SUPPORTED_OPTIMIZERS, seeds: Sequence[int] | None = None) -> pd.DataFrame: + sets = completed_seed_sets(results_root, optimizers) + if seeds is None: + seeds = sorted(set().union(*sets.values())) if sets else [] + rows = [] + for optimizer in optimizers: + for seed in seeds: + run_dir = run_directory(results_root, optimizer, int(seed)) + completion_path = run_dir / "run_complete.json" + payload = json.loads(completion_path.read_text(encoding="utf-8")) if completion_path.is_file() else {} + rows.append({ + "optimizer": optimizer, + "optimizer_label": OPTIMIZER_LABELS.get(optimizer, optimizer), + "seed": int(seed), + "complete": int(seed) in sets.get(optimizer, set()), + "steps": payload.get("optimizer_steps", np.nan), + "best_validation_loss": payload.get("best_validation_loss", np.nan), + "final_test_loss": payload.get("final_test_loss", np.nan), + "run_dir": str(run_dir), + }) + return pd.DataFrame(rows) + + +def _load_csvs(results_root: str | Path, relative_path: str, *, optimizers: Sequence[str], seeds: Sequence[int]) -> pd.DataFrame: + frames: list[pd.DataFrame] = [] + for optimizer in optimizers: + for seed in seeds: + if not run_is_complete(results_root, optimizer, int(seed)): + raise FileNotFoundError(f"incomplete run: optimizer={optimizer} seed={seed}") + path = run_directory(results_root, optimizer, int(seed)) / relative_path + if not path.is_file(): + raise FileNotFoundError(path) + frame = pd.read_csv(path) + frame["optimizer"] = optimizer + frame["optimizer_label"] = OPTIMIZER_LABELS.get(optimizer, optimizer) + frame["seed"] = int(seed) + frames.append(frame) + return pd.concat(frames, ignore_index=True, sort=False) if frames else pd.DataFrame() + + +def load_metrics(results_root: str | Path, *, optimizers: Sequence[str], seeds: Sequence[int]) -> pd.DataFrame: + frame = _load_csvs(results_root, "metrics.csv", optimizers=optimizers, seeds=seeds) + return frame.sort_values(["optimizer", "seed", "step"]).drop_duplicates(["optimizer", "seed", "step"], keep="last") + + +def load_epoch_metrics(results_root: str | Path, *, optimizers: Sequence[str], seeds: Sequence[int]) -> pd.DataFrame: + frame = _load_csvs(results_root, "epoch_metrics.csv", optimizers=optimizers, seeds=seeds) + return frame.sort_values(["optimizer", "seed", "nominal_epoch"]).drop_duplicates(["optimizer", "seed", "nominal_epoch"], keep="last") + + +def load_layer_metrics(results_root: str | Path, *, optimizers: Sequence[str], seeds: Sequence[int]) -> pd.DataFrame: + frame = _load_csvs(results_root, "spectral/layers.csv", optimizers=optimizers, seeds=seeds) + return frame.sort_values(["optimizer", "seed", "epoch", "block", "matrix_type"]).drop_duplicates(["optimizer", "seed", "step", "matrix_name"], keep="last") + + +def load_spectral_summary(results_root: str | Path, *, optimizers: Sequence[str], seeds: Sequence[int]) -> pd.DataFrame: + frame = _load_csvs(results_root, "spectral/summary.csv", optimizers=optimizers, seeds=seeds) + return frame.sort_values(["optimizer", "seed", "epoch"]).drop_duplicates(["optimizer", "seed", "step"], keep="last") + + +def load_test_results(results_root: str | Path, *, optimizers: Sequence[str], seeds: Sequence[int]) -> pd.DataFrame: + rows = [] + for optimizer in optimizers: + for seed in seeds: + if not run_is_complete(results_root, optimizer, int(seed)): + raise FileNotFoundError(f"incomplete run: optimizer={optimizer} seed={seed}") + path = run_directory(results_root, optimizer, int(seed)) / "test_results.json" + payload = json.loads(path.read_text(encoding="utf-8")) + for checkpoint in ("final", "validation_selected"): + values = payload[checkpoint] + rows.append({ + "optimizer": optimizer, + "optimizer_label": OPTIMIZER_LABELS[optimizer], + "seed": int(seed), + "checkpoint": checkpoint, + "step": int(values["step"]), + "test_loss": float(values["loss"]), + "test_accuracy": float(values["accuracy"]), + "test_bleu": float(values["bleu"]), + }) + return pd.DataFrame(rows) + + +def summarize_by_epoch(frame: pd.DataFrame, metric: str, *, x: str = "nominal_epoch", group: Sequence[str] = ("optimizer",)) -> pd.DataFrame: + rows = [] + keys = [*group, x] + subset = frame[[*keys, "seed", metric]].copy() + subset[metric] = pd.to_numeric(subset[metric], errors="coerce") + for values, group_frame in subset.groupby(keys, sort=True): + values_tuple = values if isinstance(values, tuple) else (values,) + row = dict(zip(keys, values_tuple, strict=True)) + row.update(mean_ci95(group_frame[metric])) + rows.append(row) + return pd.DataFrame(rows) + + +def final_test_summary(test_results: pd.DataFrame) -> pd.DataFrame: + rows = [] + for (optimizer, checkpoint), group in test_results.groupby(["optimizer", "checkpoint"]): + loss_stats = mean_ci95(group["test_loss"]) + for metric in ("test_loss", "test_accuracy", "test_bleu"): + rows.append({ + "optimizer": optimizer, + "optimizer_label": OPTIMIZER_LABELS[optimizer], + "checkpoint": checkpoint, + "metric": metric, + "interval_method": "run_level_student_t", + **mean_ci95(group[metric]), + }) + rows.append({ + "optimizer": optimizer, + "optimizer_label": OPTIMIZER_LABELS[optimizer], + "checkpoint": checkpoint, + "metric": "test_perplexity", + "interval_method": "exp_of_loss_space_student_t_interval", + "n": loss_stats["n"], + "mean": math.exp(loss_stats["mean"]), + "sd": np.nan, + "sem": np.nan, + "ci95_half_width": np.nan, + "ci95_lower": math.exp(loss_stats["ci95_lower"]), + "ci95_upper": math.exp(loss_stats["ci95_upper"]), + }) + return pd.DataFrame(rows) + + +def paired_optimizer_differences(test_results: pd.DataFrame) -> pd.DataFrame: + rows = [] + optimizers = tuple(sorted(test_results["optimizer"].unique())) + for checkpoint in ("final", "validation_selected"): + checkpoint_frame = test_results[test_results["checkpoint"].eq(checkpoint)] + for left, right in combinations(optimizers, 2): + for metric in ("test_loss", "test_accuracy", "test_bleu"): + a = checkpoint_frame[checkpoint_frame["optimizer"].eq(left)][["seed", metric]].rename(columns={metric: "left"}) + b = checkpoint_frame[checkpoint_frame["optimizer"].eq(right)][["seed", metric]].rename(columns={metric: "right"}) + paired = a.merge(b, on="seed", validate="one_to_one") + stats = mean_ci95(paired["left"] - paired["right"]) + rows.append({ + "checkpoint": checkpoint, + "contrast": f"{OPTIMIZER_LABELS[left]} - {OPTIMIZER_LABELS[right]}", + "left_optimizer": left, + "right_optimizer": right, + "metric": metric, + **stats, + }) + return pd.DataFrame(rows) + + +def run_diagnostics(results_root: str | Path, *, optimizers: Sequence[str], seeds: Sequence[int]) -> pd.DataFrame: + rows = [] + for optimizer in optimizers: + for seed in seeds: + run_dir = run_directory(results_root, optimizer, int(seed)) + completion = json.loads((run_dir / "run_complete.json").read_text(encoding="utf-8")) + metrics = pd.read_csv(run_dir / "metrics.csv").sort_values("step") + final = metrics.iloc[-1] + clipped = pd.to_numeric(metrics["gradient_clipped"], errors="coerce").dropna() + ratios = pd.to_numeric(metrics["update_to_weight_ratio"], errors="coerce") + rows.append({ + "optimizer": optimizer, + "optimizer_label": OPTIMIZER_LABELS[optimizer], + "seed": int(seed), + "best_validation_step": int(completion["best_validation_step"]), + "best_validation_loss": float(completion["best_validation_loss"]), + "final_validation_loss": float(final["val_loss"]), + "final_minus_best_validation_loss": float(final["val_loss"] - completion["best_validation_loss"]), + "final_validation_accuracy": float(final["val_accuracy"]), + "gradient_clipped_fraction": float(clipped.mean()) if not clipped.empty else np.nan, + "max_update_to_weight_ratio": float(ratios.max(skipna=True)), + "elapsed_seconds": float(completion["elapsed_seconds"]), + }) + return pd.DataFrame(rows) + + +def diagnostic_summary(diagnostics: pd.DataFrame) -> pd.DataFrame: + metrics = ( + "best_validation_loss", + "final_validation_loss", + "final_minus_best_validation_loss", + "final_validation_accuracy", + "gradient_clipped_fraction", + "max_update_to_weight_ratio", + "elapsed_seconds", + ) + rows = [] + for optimizer, group in diagnostics.groupby("optimizer"): + for metric in metrics: + rows.append({ + "optimizer": optimizer, + "optimizer_label": OPTIMIZER_LABELS[optimizer], + "metric": metric, + **mean_ci95(group[metric]), + }) + return pd.DataFrame(rows) + + +def plot_epoch_metric(frame: pd.DataFrame, *, metric: str, x: str = "nominal_epoch", optimizers: Sequence[str] = SUPPORTED_OPTIMIZERS, title: str | None = None, output: str | Path | None = None): + figure, axis = plt.subplots(figsize=(10, 5.5)) + for optimizer in optimizers: + subset = frame[frame["optimizer"].eq(optimizer)] + if subset.empty: + continue + for _, seed_frame in subset.groupby("seed"): + axis.plot(seed_frame[x], seed_frame[metric], color=OPTIMIZER_COLORS[optimizer], alpha=0.16, linewidth=0.8) + summary = summarize_by_epoch(subset, metric, x=x) + axis.plot(summary[x], summary["mean"], color=OPTIMIZER_COLORS[optimizer], linewidth=2.2, label=OPTIMIZER_LABELS[optimizer]) + axis.fill_between(summary[x], summary["ci95_lower"], summary["ci95_upper"], color=OPTIMIZER_COLORS[optimizer], alpha=0.14) + axis.set_xlabel(x.replace("_", " ").title()) + axis.set_ylabel(metric.replace("_", " ").title()) + axis.set_title(title or f"{metric}: run-level 95% Student-t CI") + axis.grid(alpha=0.25) + axis.legend(frameon=False) + figure.tight_layout() + if output is not None: + output = Path(output) + output.parent.mkdir(parents=True, exist_ok=True) + figure.savefig(output, dpi=180, bbox_inches="tight") + return figure, axis + + +def plot_layer_metric(frame: pd.DataFrame, *, optimizer: str, metric: str, output: str | Path | None = None): + subset = frame[frame["optimizer"].eq(optimizer)].copy() + if subset.empty: + raise ValueError(f"no layer data for optimizer={optimizer}") + blocks = tuple(sorted(int(value) for value in subset["block"].unique())) + figure, axes = plt.subplots(len(blocks), 1, figsize=(11, max(5, 4 * len(blocks))), squeeze=False, sharex=True) + for row, block in enumerate(blocks): + axis = axes[row, 0] + block_frame = subset[subset["block"].eq(block)] + for matrix_type, color in MATRIX_COLORS.items(): + matrix = block_frame[block_frame["matrix_type"].eq(matrix_type)] + if matrix.empty: + continue + summary = summarize_by_epoch(matrix, metric, x="epoch", group=("matrix_type",)) + axis.plot(summary["epoch"], summary["mean"], color=color, linewidth=2.0, marker="o", markersize=3, label=matrix_type) + axis.fill_between(summary["epoch"], summary["ci95_lower"], summary["ci95_upper"], color=color, alpha=0.12) + if metric == "alpha": + axis.axhline(2.0, color="black", linestyle="--", linewidth=1.0, label="alpha = 2") + if metric == "ERG_gap": + axis.axhline(0.0, color="black", linestyle="--", linewidth=1.0) + axis.set_ylabel(metric) + axis.set_title(f"{OPTIMIZER_LABELS[optimizer]} block {block}") + axis.grid(alpha=0.25) + axes[-1, 0].set_xlabel("Epoch") + axes[0, 0].legend(frameon=False, ncol=3) + figure.tight_layout() + if output is not None: + output = Path(output) + output.parent.mkdir(parents=True, exist_ok=True) + figure.savefig(output, dpi=180, bbox_inches="tight") + return figure, axes From 7805edb88f3dc9d09cd357740c419b524c84988c Mon Sep 17 00:00:00 2001 From: Charles Martin Date: Tue, 11 Aug 2026 09:13:19 -0700 Subject: [PATCH 14/22] Export NGB v4 API --- baseline/ngb/src/rg_ngb/__init__.py | 62 +++++++++++++++++++++++++++++ 1 file changed, 62 insertions(+) create mode 100644 baseline/ngb/src/rg_ngb/__init__.py diff --git a/baseline/ngb/src/rg_ngb/__init__.py b/baseline/ngb/src/rg_ngb/__init__.py new file mode 100644 index 0000000..11437d4 --- /dev/null +++ b/baseline/ngb/src/rg_ngb/__init__.py @@ -0,0 +1,62 @@ +"""NGB v4 tuned nanoGPT baselines.""" + +from .analysis import ( + MATRIX_COLORS, + OPTIMIZER_COLORS, + OPTIMIZER_LABELS, + completed_seed_sets, + diagnostic_summary, + discover_matched_seeds, + final_test_summary, + load_epoch_metrics, + load_layer_metrics, + load_metrics, + load_spectral_summary, + load_test_results, + paired_optimizer_differences, + plot_epoch_metric, + plot_layer_metric, + run_diagnostics, + run_status_table, +) +from .config import ( + DEFAULT_ROOT, + SUPPORTED_OPTIMIZERS, + canonical_seeds, + load_config, + roots, +) +from .data import prepare_fineweb_edu +from .model import GPT, GPTConfig, transformer_matrix_items +from .training import run_all_replicates, run_optimizer_replicates + +__all__ = [ + "DEFAULT_ROOT", + "GPT", + "GPTConfig", + "MATRIX_COLORS", + "OPTIMIZER_COLORS", + "OPTIMIZER_LABELS", + "SUPPORTED_OPTIMIZERS", + "canonical_seeds", + "completed_seed_sets", + "diagnostic_summary", + "discover_matched_seeds", + "final_test_summary", + "load_config", + "load_epoch_metrics", + "load_layer_metrics", + "load_metrics", + "load_spectral_summary", + "load_test_results", + "paired_optimizer_differences", + "plot_epoch_metric", + "plot_layer_metric", + "prepare_fineweb_edu", + "roots", + "run_all_replicates", + "run_diagnostics", + "run_optimizer_replicates", + "run_status_table", + "transformer_matrix_items", +] From 55722833b19d29260097c5bbd7d5875d83e614ec Mon Sep 17 00:00:00 2001 From: Charles Martin Date: Tue, 11 Aug 2026 09:13:57 -0700 Subject: [PATCH 15/22] Add one-head v4 training notebook --- .../ngb/notebooks/01_v4_one_head_train.ipynb | 63 +++++++++++++++++++ 1 file changed, 63 insertions(+) create mode 100644 baseline/ngb/notebooks/01_v4_one_head_train.ipynb diff --git a/baseline/ngb/notebooks/01_v4_one_head_train.ipynb b/baseline/ngb/notebooks/01_v4_one_head_train.ipynb new file mode 100644 index 0000000..107baf0 --- /dev/null +++ b/baseline/ngb/notebooks/01_v4_one_head_train.ipynb @@ -0,0 +1,63 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "id": "title", + "metadata": {}, + "source": [ + "# NGB v4 — tuned one-head training\n", + "\n", + "Runs or resumes the separate two-epoch v4 SGD, AdamW, and Muon protocols. Results live under `/tmp/rg-ngb/results/ngb_v4_one_head_2epoch` and never mix with v3." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "setup", + "metadata": {}, + "outputs": [], + "source": [ + "from pathlib import Path\n", + "import sys\n", + "import pandas as pd\n", + "from IPython.display import display\n", + "\n", + "cwd = Path.cwd().resolve()\n", + "candidates = [cwd, cwd.parent, cwd / 'baseline' / 'ngb']\n", + "EXPERIMENT_ROOT = next((p for p in candidates if (p / 'configs' / 'v4_one_head.yaml').is_file()), None)\n", + "if EXPERIMENT_ROOT is None:\n", + " raise FileNotFoundError('Run from baseline/ngb, baseline/ngb/notebooks, or the repository root')\n", + "sys.path.insert(0, str(EXPERIMENT_ROOT / 'src'))\n", + "\n", + "from rg_ngb import canonical_seeds, load_config, roots, run_all_replicates, run_status_table\n", + "\n", + "pd.set_option('display.max_rows', None)\n", + "pd.set_option('display.max_columns', None)\n", + "pd.set_option('display.width', None)\n", + "CONFIG = load_config(EXPERIMENT_ROOT / 'configs' / 'v4_one_head.yaml')\n", + "SEEDS = canonical_seeds(CONFIG)\n", + "PATHS = roots(CONFIG)\n", + "display(pd.DataFrame([CONFIG['model'] | CONFIG['training']]))\n", + "display(run_status_table(PATHS['results'], seeds=SEEDS))" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "run", + "metadata": {}, + "outputs": [], + "source": [ + "RUNS = run_all_replicates(cfg=CONFIG, seeds=SEEDS, device='auto', resume=True, overwrite=False, progress=True)\n", + "print('completed or verified runs:', len(RUNS))\n", + "display(run_status_table(PATHS['results'], seeds=SEEDS))" + ] + } + ], + "metadata": { + "kernelspec": {"display_name": "Python 3", "language": "python", "name": "python3"}, + "language_info": {"name": "python", "version": "3.10"} + }, + "nbformat": 4, + "nbformat_minor": 5 +} From d33838b08967d3dcc230ebe5b5d76491c9596625 Mon Sep 17 00:00:00 2001 From: Charles Martin Date: Tue, 11 Aug 2026 09:14:44 -0700 Subject: [PATCH 16/22] Add corrected one-head v4 comparison notebook --- .../notebooks/02_v4_one_head_compare.ipynb | 157 ++++++++++++++++++ 1 file changed, 157 insertions(+) create mode 100644 baseline/ngb/notebooks/02_v4_one_head_compare.ipynb diff --git a/baseline/ngb/notebooks/02_v4_one_head_compare.ipynb b/baseline/ngb/notebooks/02_v4_one_head_compare.ipynb new file mode 100644 index 0000000..8dea41b --- /dev/null +++ b/baseline/ngb/notebooks/02_v4_one_head_compare.ipynb @@ -0,0 +1,157 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "id": "title", + "metadata": {}, + "source": [ + "# NGB v4 — tuned one-head optimizer comparison\n", + "\n", + "Uses the complete seed intersection shared by SGD, AdamW, and Muon. All uncertainty is run-level. Perplexity intervals are transformed from cross-entropy loss space." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "setup", + "metadata": {}, + "outputs": [], + "source": [ + "from pathlib import Path\n", + "import sys\n", + "import matplotlib.pyplot as plt\n", + "import pandas as pd\n", + "from IPython.display import display\n", + "\n", + "cwd = Path.cwd().resolve()\n", + "candidates = [cwd, cwd.parent, cwd / 'baseline' / 'ngb']\n", + "EXPERIMENT_ROOT = next((p for p in candidates if (p / 'configs' / 'v4_one_head.yaml').is_file()), None)\n", + "if EXPERIMENT_ROOT is None:\n", + " raise FileNotFoundError('Run from baseline/ngb, baseline/ngb/notebooks, or the repository root')\n", + "sys.path.insert(0, str(EXPERIMENT_ROOT / 'src'))\n", + "\n", + "from rg_ngb import (\n", + " SUPPORTED_OPTIMIZERS, diagnostic_summary, discover_matched_seeds,\n", + " final_test_summary, load_config, load_epoch_metrics, load_layer_metrics,\n", + " load_spectral_summary, load_test_results, paired_optimizer_differences,\n", + " plot_epoch_metric, plot_layer_metric, roots, run_diagnostics, run_status_table,\n", + ")\n", + "\n", + "pd.set_option('display.max_rows', None)\n", + "pd.set_option('display.max_columns', None)\n", + "pd.set_option('display.width', None)\n", + "CONFIG = load_config(EXPERIMENT_ROOT / 'configs' / 'v4_one_head.yaml')\n", + "PATHS = roots(CONFIG)\n", + "OPTIMIZERS = SUPPORTED_OPTIMIZERS\n", + "SEEDS = discover_matched_seeds(PATHS['results'], optimizers=OPTIMIZERS)\n", + "PLOT_DIR = PATHS['plots'] / 'comparison'\n", + "PLOT_DIR.mkdir(parents=True, exist_ok=True)\n", + "print('matched completed seeds:', SEEDS)\n", + "display(run_status_table(PATHS['results'], optimizers=OPTIMIZERS, seeds=SEEDS))" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "load", + "metadata": {}, + "outputs": [], + "source": [ + "epoch_metrics = load_epoch_metrics(PATHS['results'], optimizers=OPTIMIZERS, seeds=SEEDS)\n", + "layer_metrics = load_layer_metrics(PATHS['results'], optimizers=OPTIMIZERS, seeds=SEEDS)\n", + "spectral_summary = load_spectral_summary(PATHS['results'], optimizers=OPTIMIZERS, seeds=SEEDS)\n", + "test_results = load_test_results(PATHS['results'], optimizers=OPTIMIZERS, seeds=SEEDS)\n", + "diagnostics = run_diagnostics(PATHS['results'], optimizers=OPTIMIZERS, seeds=SEEDS)\n", + "display(epoch_metrics.sort_values(['optimizer', 'seed', 'nominal_epoch']))" + ] + }, + { + "cell_type": "markdown", + "id": "performance-heading", + "metadata": {}, + "source": ["## Full and post-transient task trajectories"] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "performance", + "metadata": {}, + "outputs": [], + "source": [ + "for metric in [\n", + " 'train_loss', 'val_loss', 'test_loss',\n", + " 'train_accuracy', 'val_accuracy', 'test_accuracy',\n", + " 'train_perplexity', 'val_perplexity', 'test_perplexity',\n", + " 'test_bleu', 'val_generalization_gap', 'test_generalization_gap',\n", + " 'weight_norm', 'update_to_weight_ratio', 'grad_norm_pre_clip',\n", + "]:\n", + " if metric not in epoch_metrics.columns:\n", + " continue\n", + " plot_epoch_metric(epoch_metrics, metric=metric, optimizers=OPTIMIZERS, output=PLOT_DIR / f'full_{metric}.png')\n", + " plt.show()\n", + " zoom = epoch_metrics[epoch_metrics['nominal_epoch'].ge(0.5)]\n", + " plot_epoch_metric(zoom, metric=metric, optimizers=OPTIMIZERS, title=f'{metric}: epoch 0.5 onward', output=PLOT_DIR / f'zoom_{metric}.png')\n", + " plt.show()" + ] + }, + { + "cell_type": "markdown", + "id": "spectral-heading", + "metadata": {}, + "source": ["## Optimizer-level and block-resolved WeightWatcher diagnostics"] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "spectral", + "metadata": {}, + "outputs": [], + "source": [ + "for metric in ['alpha_median', 'ERG_gap_median', 'num_traps_mean']:\n", + " plot_epoch_metric(spectral_summary, metric=metric, x='epoch', optimizers=OPTIMIZERS, output=PLOT_DIR / f'spectral_{metric}.png')\n", + " if metric == 'alpha_median':\n", + " plt.axhline(2.0, color='black', linestyle='--', linewidth=1.0, label='alpha = 2')\n", + " if metric == 'ERG_gap_median':\n", + " plt.axhline(0.0, color='black', linestyle='--', linewidth=1.0)\n", + " plt.show()\n", + "\n", + "for optimizer in OPTIMIZERS:\n", + " for metric in ['alpha', 'ERG_gap', 'num_traps']:\n", + " plot_layer_metric(layer_metrics, optimizer=optimizer, metric=metric, output=PLOT_DIR / f'{optimizer}_layer_{metric}.png')\n", + " plt.show()" + ] + }, + { + "cell_type": "markdown", + "id": "tables-heading", + "metadata": {}, + "source": ["## Final, validation-selected, paired, and stability tables"] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "tables", + "metadata": {}, + "outputs": [], + "source": [ + "terminal = final_test_summary(test_results)\n", + "paired = paired_optimizer_differences(test_results)\n", + "stability = diagnostic_summary(diagnostics)\n", + "display(terminal.sort_values(['checkpoint', 'metric', 'optimizer_label']))\n", + "display(paired.sort_values(['checkpoint', 'metric', 'contrast']))\n", + "display(diagnostics.sort_values(['optimizer', 'seed']))\n", + "display(stability.sort_values(['metric', 'optimizer_label']))\n", + "terminal.to_csv(PLOT_DIR / 'terminal_summary_95ci.csv', index=False)\n", + "paired.to_csv(PLOT_DIR / 'paired_optimizer_differences_95ci.csv', index=False)\n", + "diagnostics.to_csv(PLOT_DIR / 'run_diagnostics.csv', index=False)\n", + "stability.to_csv(PLOT_DIR / 'diagnostic_summary_95ci.csv', index=False)" + ] + } + ], + "metadata": { + "kernelspec": {"display_name": "Python 3", "language": "python", "name": "python3"}, + "language_info": {"name": "python", "version": "3.10"} + }, + "nbformat": 4, + "nbformat_minor": 5 +} From bbf005ddc1bd7547ed21d16143982bf18d1790e3 Mon Sep 17 00:00:00 2001 From: Charles Martin Date: Tue, 11 Aug 2026 09:15:03 -0700 Subject: [PATCH 17/22] Add four-layer four-head v4 training notebook --- .../ngb/notebooks/11_v4_small_4x4_train.ipynb | 63 +++++++++++++++++++ 1 file changed, 63 insertions(+) create mode 100644 baseline/ngb/notebooks/11_v4_small_4x4_train.ipynb diff --git a/baseline/ngb/notebooks/11_v4_small_4x4_train.ipynb b/baseline/ngb/notebooks/11_v4_small_4x4_train.ipynb new file mode 100644 index 0000000..1b1438a --- /dev/null +++ b/baseline/ngb/notebooks/11_v4_small_4x4_train.ipynb @@ -0,0 +1,63 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "id": "title", + "metadata": {}, + "source": [ + "# NGB v4 — small 4-layer / 4-head training\n", + "\n", + "Runs or resumes the distinct two-epoch 4×4 small-language-model protocol. Results live under `/tmp/rg-ngb/results/ngb_v4_small_4x4_2epoch`." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "setup", + "metadata": {}, + "outputs": [], + "source": [ + "from pathlib import Path\n", + "import sys\n", + "import pandas as pd\n", + "from IPython.display import display\n", + "\n", + "cwd = Path.cwd().resolve()\n", + "candidates = [cwd, cwd.parent, cwd / 'baseline' / 'ngb']\n", + "EXPERIMENT_ROOT = next((p for p in candidates if (p / 'configs' / 'v4_small_4x4.yaml').is_file()), None)\n", + "if EXPERIMENT_ROOT is None:\n", + " raise FileNotFoundError('Run from baseline/ngb, baseline/ngb/notebooks, or the repository root')\n", + "sys.path.insert(0, str(EXPERIMENT_ROOT / 'src'))\n", + "\n", + "from rg_ngb import canonical_seeds, load_config, roots, run_all_replicates, run_status_table\n", + "\n", + "pd.set_option('display.max_rows', None)\n", + "pd.set_option('display.max_columns', None)\n", + "pd.set_option('display.width', None)\n", + "CONFIG = load_config(EXPERIMENT_ROOT / 'configs' / 'v4_small_4x4.yaml')\n", + "SEEDS = canonical_seeds(CONFIG)\n", + "PATHS = roots(CONFIG)\n", + "display(pd.DataFrame([CONFIG['model'] | CONFIG['training']]))\n", + "display(run_status_table(PATHS['results'], seeds=SEEDS))" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "run", + "metadata": {}, + "outputs": [], + "source": [ + "RUNS = run_all_replicates(cfg=CONFIG, seeds=SEEDS, device='auto', resume=True, overwrite=False, progress=True)\n", + "print('completed or verified runs:', len(RUNS))\n", + "display(run_status_table(PATHS['results'], seeds=SEEDS))" + ] + } + ], + "metadata": { + "kernelspec": {"display_name": "Python 3", "language": "python", "name": "python3"}, + "language_info": {"name": "python", "version": "3.10"} + }, + "nbformat": 4, + "nbformat_minor": 5 +} From 63a5604f366e8dc40678d72edb02f9bd5ea865b8 Mon Sep 17 00:00:00 2001 From: Charles Martin Date: Tue, 11 Aug 2026 09:15:41 -0700 Subject: [PATCH 18/22] Add corrected four-layer four-head v4 comparison notebook --- .../notebooks/12_v4_small_4x4_compare.ipynb | 139 ++++++++++++++++++ 1 file changed, 139 insertions(+) create mode 100644 baseline/ngb/notebooks/12_v4_small_4x4_compare.ipynb diff --git a/baseline/ngb/notebooks/12_v4_small_4x4_compare.ipynb b/baseline/ngb/notebooks/12_v4_small_4x4_compare.ipynb new file mode 100644 index 0000000..56dd142 --- /dev/null +++ b/baseline/ngb/notebooks/12_v4_small_4x4_compare.ipynb @@ -0,0 +1,139 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "id": "title", + "metadata": {}, + "source": [ + "# NGB v4 — small 4-layer / 4-head optimizer comparison\n", + "\n", + "Compares the complete matched seed intersection for the distinct 4×4 architecture with run-level uncertainty, paired contrasts, and block-resolved WeightWatcher diagnostics." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "setup", + "metadata": {}, + "outputs": [], + "source": [ + "from pathlib import Path\n", + "import sys\n", + "import matplotlib.pyplot as plt\n", + "import pandas as pd\n", + "from IPython.display import display\n", + "\n", + "cwd = Path.cwd().resolve()\n", + "candidates = [cwd, cwd.parent, cwd / 'baseline' / 'ngb']\n", + "EXPERIMENT_ROOT = next((p for p in candidates if (p / 'configs' / 'v4_small_4x4.yaml').is_file()), None)\n", + "if EXPERIMENT_ROOT is None:\n", + " raise FileNotFoundError('Run from baseline/ngb, baseline/ngb/notebooks, or the repository root')\n", + "sys.path.insert(0, str(EXPERIMENT_ROOT / 'src'))\n", + "\n", + "from rg_ngb import (\n", + " SUPPORTED_OPTIMIZERS, diagnostic_summary, discover_matched_seeds,\n", + " final_test_summary, load_config, load_epoch_metrics, load_layer_metrics,\n", + " load_spectral_summary, load_test_results, paired_optimizer_differences,\n", + " plot_epoch_metric, plot_layer_metric, roots, run_diagnostics, run_status_table,\n", + ")\n", + "\n", + "pd.set_option('display.max_rows', None)\n", + "pd.set_option('display.max_columns', None)\n", + "pd.set_option('display.width', None)\n", + "CONFIG = load_config(EXPERIMENT_ROOT / 'configs' / 'v4_small_4x4.yaml')\n", + "PATHS = roots(CONFIG)\n", + "OPTIMIZERS = SUPPORTED_OPTIMIZERS\n", + "SEEDS = discover_matched_seeds(PATHS['results'], optimizers=OPTIMIZERS)\n", + "PLOT_DIR = PATHS['plots'] / 'comparison'\n", + "PLOT_DIR.mkdir(parents=True, exist_ok=True)\n", + "print('matched completed seeds:', SEEDS)\n", + "display(run_status_table(PATHS['results'], optimizers=OPTIMIZERS, seeds=SEEDS))" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "load", + "metadata": {}, + "outputs": [], + "source": [ + "epoch_metrics = load_epoch_metrics(PATHS['results'], optimizers=OPTIMIZERS, seeds=SEEDS)\n", + "layer_metrics = load_layer_metrics(PATHS['results'], optimizers=OPTIMIZERS, seeds=SEEDS)\n", + "spectral_summary = load_spectral_summary(PATHS['results'], optimizers=OPTIMIZERS, seeds=SEEDS)\n", + "test_results = load_test_results(PATHS['results'], optimizers=OPTIMIZERS, seeds=SEEDS)\n", + "diagnostics = run_diagnostics(PATHS['results'], optimizers=OPTIMIZERS, seeds=SEEDS)\n", + "display(epoch_metrics.sort_values(['optimizer', 'seed', 'nominal_epoch']))" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "performance", + "metadata": {}, + "outputs": [], + "source": [ + "for metric in [\n", + " 'train_loss', 'val_loss', 'test_loss',\n", + " 'train_accuracy', 'val_accuracy', 'test_accuracy',\n", + " 'train_perplexity', 'val_perplexity', 'test_perplexity',\n", + " 'test_bleu', 'val_generalization_gap', 'test_generalization_gap',\n", + " 'weight_norm', 'update_to_weight_ratio', 'grad_norm_pre_clip',\n", + "]:\n", + " if metric not in epoch_metrics.columns:\n", + " continue\n", + " plot_epoch_metric(epoch_metrics, metric=metric, optimizers=OPTIMIZERS, output=PLOT_DIR / f'full_{metric}.png')\n", + " plt.show()\n", + " zoom = epoch_metrics[epoch_metrics['nominal_epoch'].ge(0.5)]\n", + " plot_epoch_metric(zoom, metric=metric, optimizers=OPTIMIZERS, title=f'{metric}: epoch 0.5 onward', output=PLOT_DIR / f'zoom_{metric}.png')\n", + " plt.show()" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "spectral", + "metadata": {}, + "outputs": [], + "source": [ + "for metric in ['alpha_median', 'ERG_gap_median', 'num_traps_mean']:\n", + " plot_epoch_metric(spectral_summary, metric=metric, x='epoch', optimizers=OPTIMIZERS, output=PLOT_DIR / f'spectral_{metric}.png')\n", + " if metric == 'alpha_median':\n", + " plt.axhline(2.0, color='black', linestyle='--', linewidth=1.0, label='alpha = 2')\n", + " if metric == 'ERG_gap_median':\n", + " plt.axhline(0.0, color='black', linestyle='--', linewidth=1.0)\n", + " plt.show()\n", + "\n", + "for optimizer in OPTIMIZERS:\n", + " for metric in ['alpha', 'ERG_gap', 'num_traps']:\n", + " plot_layer_metric(layer_metrics, optimizer=optimizer, metric=metric, output=PLOT_DIR / f'{optimizer}_block_layer_{metric}.png')\n", + " plt.show()" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "tables", + "metadata": {}, + "outputs": [], + "source": [ + "terminal = final_test_summary(test_results)\n", + "paired = paired_optimizer_differences(test_results)\n", + "stability = diagnostic_summary(diagnostics)\n", + "display(terminal.sort_values(['checkpoint', 'metric', 'optimizer_label']))\n", + "display(paired.sort_values(['checkpoint', 'metric', 'contrast']))\n", + "display(diagnostics.sort_values(['optimizer', 'seed']))\n", + "display(stability.sort_values(['metric', 'optimizer_label']))\n", + "terminal.to_csv(PLOT_DIR / 'terminal_summary_95ci.csv', index=False)\n", + "paired.to_csv(PLOT_DIR / 'paired_optimizer_differences_95ci.csv', index=False)\n", + "diagnostics.to_csv(PLOT_DIR / 'run_diagnostics.csv', index=False)\n", + "stability.to_csv(PLOT_DIR / 'diagnostic_summary_95ci.csv', index=False)" + ] + } + ], + "metadata": { + "kernelspec": {"display_name": "Python 3", "language": "python", "name": "python3"}, + "language_info": {"name": "python", "version": "3.10"} + }, + "nbformat": 4, + "nbformat_minor": 5 +} From 4e5d44facfb49cc1bc5887cb80b73388901e0604 Mon Sep 17 00:00:00 2001 From: Charles Martin Date: Tue, 11 Aug 2026 09:16:22 -0700 Subject: [PATCH 19/22] Add NGB protocol and architecture tests --- baseline/ngb/tests/test_protocol.py | 134 ++++++++++++++++++++++++++++ 1 file changed, 134 insertions(+) create mode 100644 baseline/ngb/tests/test_protocol.py diff --git a/baseline/ngb/tests/test_protocol.py b/baseline/ngb/tests/test_protocol.py new file mode 100644 index 0000000..4614874 --- /dev/null +++ b/baseline/ngb/tests/test_protocol.py @@ -0,0 +1,134 @@ +from __future__ import annotations + +import ast +import json +from pathlib import Path + +import pytest +import torch + +from rg_nanogpt_one_head.optimizers import make_optimizer_handles, optimizer_step + +from rg_ngb.config import ( + DEFAULT_ROOT, + canonical_seeds, + epoch_step_map, + expected_matrix_count, + load_config, + max_steps, + optimizer_profile, + roots, +) +from rg_ngb.model import GPT, GPTConfig, transformer_matrix_items + +EXPERIMENT_ROOT = Path(__file__).resolve().parents[1] + + +def config(name: str) -> dict: + return load_config(EXPERIMENT_ROOT / "configs" / name) + + +def test_defaults_are_tmp_and_protocol_scoped(monkeypatch): + for name in ("RG_NGB_ROOT", "RG_NGB_DATA_ROOT", "RG_NGB_RESULTS_ROOT", "RG_NGB_PLOTS_ROOT"): + monkeypatch.delenv(name, raising=False) + cfg = config("v4_one_head.yaml") + assert DEFAULT_ROOT == Path("/tmp/rg-ngb") + assert roots(cfg) == { + "root": Path("/tmp/rg-ngb"), + "data": Path("/tmp/rg-ngb/data"), + "results": Path("/tmp/rg-ngb/results/ngb_v4_one_head_2epoch"), + "plots": Path("/tmp/rg-ngb/plots/ngb_v4_one_head_2epoch"), + } + + +def test_one_head_v4_contract(): + cfg = config("v4_one_head.yaml") + assert cfg["protocol"]["version"] == 4 + assert cfg["model"]["n_layer"] == 1 + assert cfg["model"]["n_head"] == 1 + assert cfg["training"]["target_epochs"] == 2.0 + assert cfg["training"]["epoch_interval"] == 0.25 + assert cfg["training"]["eval_interval_steps"] == 500 + assert max_steps(cfg) == 19_532 + assert list(epoch_step_map(cfg).values()) == [0.0, 0.25, 0.5, 0.75, 1.0, 1.25, 1.5, 1.75, 2.0] + assert expected_matrix_count(cfg) == 6 + assert canonical_seeds(cfg) == (1337, 2027, 4099) + sgd = optimizer_profile(cfg, "sgd_momentum") + adamw = optimizer_profile(cfg, "adamw") + muon = optimizer_profile(cfg, "muon") + assert sgd["learning_rate"] == pytest.approx(0.05) + assert sgd["min_learning_rate"] == pytest.approx(0.0005) + assert sgd["warmup_fraction"] == pytest.approx(0.05) + assert adamw["learning_rate"] == pytest.approx(3e-4) + assert adamw["min_learning_rate"] == pytest.approx(1e-5) + assert adamw["warmup_fraction"] == pytest.approx(0.025) + assert muon["matrix_learning_rate"] == pytest.approx(0.01) + assert muon["matrix_min_learning_rate"] == pytest.approx(2e-4) + assert muon["aux_weight_decay"] == pytest.approx(0.10) + + +def test_small_4x4_is_distinct_and_has_expected_inventory(): + cfg = config("v4_small_4x4.yaml") + assert cfg["protocol"]["architecture_id"] == "small_4x4" + assert cfg["model"]["n_layer"] == 4 + assert cfg["model"]["n_head"] == 4 + assert expected_matrix_count(cfg) == 24 + model = GPT(GPTConfig(**cfg["model"])) + matrices = transformer_matrix_items(model) + assert len(matrices) == 24 + assert {block for _, _, block, _ in matrices} == {0, 1, 2, 3} + assert model.parameter_count() == 7_253_248 + assert optimizer_profile(cfg, "sgd_momentum")["learning_rate"] == pytest.approx(0.03) + + +@pytest.mark.parametrize("configuration", ["v4_one_head.yaml", "v4_small_4x4.yaml"]) +@pytest.mark.parametrize("optimizer_name", ["sgd_momentum", "adamw", "muon"]) +def test_all_optimizer_paths_take_finite_step(configuration, optimizer_name): + cfg = config(configuration) + model_cfg = dict(cfg["model"]) + model_cfg.update({"vocab_size": 64, "block_size": 8, "n_embd": 16}) + model = GPT(GPTConfig(**model_cfg)) + handles = make_optimizer_handles(model, optimizer_profile(cfg, optimizer_name)) + x = torch.randint(0, 64, (2, 8)) + _, loss = model(x, x) + assert loss is not None + loss.backward() + optimizer_step(handles) + assert all(torch.isfinite(parameter).all() for parameter in model.parameters()) + assert len(handles) == (2 if optimizer_name == "muon" else 1) + + +def test_notebooks_are_python3_and_parse(): + expected = { + "01_v4_one_head_train.ipynb", + "02_v4_one_head_compare.ipynb", + "11_v4_small_4x4_train.ipynb", + "12_v4_small_4x4_compare.ipynb", + } + paths = sorted((EXPERIMENT_ROOT / "notebooks").glob("*.ipynb")) + assert {path.name for path in paths} == expected + for path in paths: + notebook = json.loads(path.read_text(encoding="utf-8")) + assert notebook["metadata"]["kernelspec"]["name"] == "python3" + for index, cell in enumerate(notebook["cells"]): + if cell["cell_type"] != "code": + continue + source = cell["source"] + if isinstance(source, list): + source = "".join(source) + ast.parse(source, filename=f"{path}:cell-{index}") + + +def test_ngb_contains_no_home_directory_defaults(): + forbidden = ("$HOME", "${HOME}", "Path.home(", ".expanduser(", "/home/", "~/") + paths = [EXPERIMENT_ROOT / "README.md"] + paths.extend((EXPERIMENT_ROOT / "src").rglob("*.py")) + paths.extend((EXPERIMENT_ROOT / "configs").glob("*.yaml")) + paths.extend((EXPERIMENT_ROOT / "notebooks").glob("*.ipynb")) + violations = [] + for path in paths: + text = path.read_text(encoding="utf-8") + for token in forbidden: + if token in text: + violations.append(f"{path.relative_to(EXPERIMENT_ROOT)}: {token}") + assert not violations, "forbidden path defaults found:\n" + "\n".join(violations) From 4d32b697c76c2710c699751c8042e3c09e0c0203 Mon Sep 17 00:00:00 2001 From: Charles Martin Date: Tue, 11 Aug 2026 09:16:48 -0700 Subject: [PATCH 20/22] Add NGB comparison regression tests --- baseline/ngb/tests/test_analysis.py | 89 +++++++++++++++++++++++++++++ 1 file changed, 89 insertions(+) create mode 100644 baseline/ngb/tests/test_analysis.py diff --git a/baseline/ngb/tests/test_analysis.py b/baseline/ngb/tests/test_analysis.py new file mode 100644 index 0000000..c7f2d22 --- /dev/null +++ b/baseline/ngb/tests/test_analysis.py @@ -0,0 +1,89 @@ +from __future__ import annotations + +import math +from pathlib import Path + +import pandas as pd +import pytest + +import rg_ngb.analysis as analysis + + +def test_matched_seed_discovery_uses_intersection(tmp_path, monkeypatch): + complete = { + "sgd_momentum": {1, 2, 3, 4}, + "adamw": {2, 3, 4, 5}, + "muon": {3, 4, 5, 6}, + } + for optimizer, seeds in complete.items(): + for seed in seeds: + (tmp_path / optimizer / f"seed_{seed}").mkdir(parents=True) + + monkeypatch.setattr( + analysis, + "run_is_complete", + lambda root, optimizer, seed: int(seed) in complete[str(optimizer)], + ) + assert analysis.discover_matched_seeds(tmp_path) == (3, 4) + assert analysis.discover_matched_seeds(tmp_path, requested=(4,)) == (4,) + with pytest.raises(FileNotFoundError, match="requested seeds"): + analysis.discover_matched_seeds(tmp_path, requested=(2, 3)) + + +def test_perplexity_interval_is_exponentiated_from_loss_space(): + frame = pd.DataFrame( + { + "optimizer": ["adamw"] * 3, + "optimizer_label": ["AdamW"] * 3, + "seed": [1, 2, 3], + "checkpoint": ["final"] * 3, + "step": [10] * 3, + "test_loss": [5.0, 5.2, 5.4], + "test_accuracy": [0.1, 0.2, 0.3], + "test_bleu": [0.0, 0.1, 0.2], + } + ) + summary = analysis.final_test_summary(frame) + loss = summary[summary["metric"].eq("test_loss")].iloc[0] + perplexity = summary[summary["metric"].eq("test_perplexity")].iloc[0] + assert perplexity["interval_method"] == "exp_of_loss_space_student_t_interval" + assert perplexity["mean"] == pytest.approx(math.exp(loss["mean"])) + assert perplexity["ci95_lower"] == pytest.approx(math.exp(loss["ci95_lower"])) + assert perplexity["ci95_upper"] == pytest.approx(math.exp(loss["ci95_upper"])) + assert perplexity["ci95_lower"] > 0 + + +def test_paired_differences_preserve_seed_matching(): + rows = [] + for optimizer, offset in (("sgd_momentum", 0.0), ("adamw", -0.2), ("muon", -0.3)): + for seed, base in ((1, 6.0), (2, 7.0), (3, 8.0)): + rows.append( + { + "optimizer": optimizer, + "optimizer_label": optimizer, + "seed": seed, + "checkpoint": "final", + "step": 10, + "test_loss": base + offset, + "test_accuracy": 0.1 - offset, + "test_bleu": 0.2 - offset, + } + ) + rows.append({**rows[-1], "checkpoint": "validation_selected", "step": 8}) + result = analysis.paired_optimizer_differences(pd.DataFrame(rows)) + contrast = result[ + result["checkpoint"].eq("final") + & result["metric"].eq("test_loss") + & result["left_optimizer"].eq("adamw") + & result["right_optimizer"].eq("sgd_momentum") + ].iloc[0] + assert contrast["n"] == 3 + assert contrast["mean"] == pytest.approx(-0.2) + assert contrast["ci95_half_width"] == pytest.approx(0.0, abs=1e-12) + + +def test_mean_ci95_uses_student_t_for_three_runs(): + result = analysis.mean_ci95([1.0, 2.0, 3.0]) + assert result["n"] == 3 + assert result["mean"] == pytest.approx(2.0) + assert result["ci95_half_width"] == pytest.approx(4.3026527297 / math.sqrt(3)) From 516f8263bb972ab0d7e8f44a5ab0bc804784838b Mon Sep 17 00:00:00 2001 From: Charles Martin Date: Tue, 11 Aug 2026 09:17:19 -0700 Subject: [PATCH 21/22] Add NGB v4 CI job --- .github/workflows/baseline-tests.yml | 26 ++++++++++++++++++++++++++ 1 file changed, 26 insertions(+) diff --git a/.github/workflows/baseline-tests.yml b/.github/workflows/baseline-tests.yml index e19b817..75338f4 100644 --- a/.github/workflows/baseline-tests.yml +++ b/.github/workflows/baseline-tests.yml @@ -71,3 +71,29 @@ jobs: PYTHONPATH: baseline/nanogpt_one_head/src MPLBACKEND: Agg run: pytest -q baseline/nanogpt_one_head/tests + + ngb-v4: + runs-on: ubuntu-latest + timeout-minutes: 25 + steps: + - uses: actions/checkout@v4 + - uses: actions/setup-python@v5 + with: + python-version: "3.11" + cache: pip + cache-dependency-path: | + baseline/nanogpt_one_head/pyproject.toml + baseline/ngb/pyproject.toml + - name: Install CPU PyTorch and NGB dependencies + run: | + python -m pip install --upgrade pip + python -m pip install torch --index-url https://download.pytorch.org/whl/cpu + python -m pip install -e './baseline/nanogpt_one_head' + python -m pip install -e './baseline/ngb[dev]' + python -m pip check + - name: Run NGB protocol, architecture, and analysis tests + env: + MPLBACKEND: Agg + run: | + python -m compileall -q baseline/ngb/src baseline/ngb/tests + pytest -q baseline/ngb/tests From e4a67ee0e0577b5770721d5a0ee18ab01239482c Mon Sep 17 00:00:00 2001 From: Charles Martin Date: Tue, 11 Aug 2026 09:21:40 -0700 Subject: [PATCH 22/22] Add bounded v4 validation qualification grid --- .../ngb/configs/v4_qualification_grid.yaml | 68 +++++++++++++++++++ 1 file changed, 68 insertions(+) create mode 100644 baseline/ngb/configs/v4_qualification_grid.yaml diff --git a/baseline/ngb/configs/v4_qualification_grid.yaml b/baseline/ngb/configs/v4_qualification_grid.yaml new file mode 100644 index 0000000..04f1bfb --- /dev/null +++ b/baseline/ngb/configs/v4_qualification_grid.yaml @@ -0,0 +1,68 @@ +qualification: + objective: mean_best_validation_loss + protected_metrics: [test_loss, test_perplexity, test_accuracy, test_bleu] + stage_1: + description: Reject unstable candidates before full-horizon qualification. + seeds: [1337, 2027] + target_epochs: 0.5 + reject_if: + final_minus_best_validation_loss_gt: 0.25 + nonfinite_metric: true + stage_2: + description: Qualify surviving centers over the complete two-epoch protocol. + seeds: [1337, 2027, 4099] + target_epochs: 2.0 + ranking: + - mean_best_validation_loss + - worst_seed_best_validation_loss + - validation_loss_standard_deviation + - final_minus_best_validation_loss + stage_3: + description: Freeze one recipe per optimizer, then expand the matched seed set. + seeds: [1337, 2027, 4099, 5003, 6007, 7013, 8017, 9011] + +one_head_candidates: + sgd_momentum: + learning_rate: [0.03, 0.05] + min_learning_rate_ratio: [0.01] + warmup_fraction: [0.05] + weight_decay: [0.01] + adamw: + learning_rate: [0.0002, 0.0003] + min_learning_rate: [0.00001] + warmup_fraction: [0.025] + beta1: [0.90] + beta2: [0.95] + weight_decay: [0.05, 0.10] + muon: + matrix_learning_rate: [0.005, 0.01] + matrix_min_learning_rate: [0.0001, 0.0002] + matrix_weight_decay: [0.01, 0.02] + aux_learning_rate: [0.0002, 0.0003] + aux_min_learning_rate: [0.00001] + aux_weight_decay: [0.01, 0.10] + momentum: [0.95] + newton_schulz_steps: [5] + +small_4x4_candidates: + sgd_momentum: + learning_rate: [0.02, 0.03] + min_learning_rate_ratio: [0.01] + warmup_fraction: [0.05] + weight_decay: [0.01] + adamw: + learning_rate: [0.0002, 0.0003] + min_learning_rate: [0.00001] + warmup_fraction: [0.025] + beta1: [0.90] + beta2: [0.95] + weight_decay: [0.05, 0.10] + muon: + matrix_learning_rate: [0.005, 0.01] + matrix_min_learning_rate: [0.0001, 0.0002] + matrix_weight_decay: [0.01, 0.02] + aux_learning_rate: [0.0002, 0.0003] + aux_min_learning_rate: [0.00001] + aux_weight_decay: [0.01, 0.10] + momentum: [0.95] + newton_schulz_steps: [5]