Add NGB v4 tuned and expressive nanoGPT baselines

.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

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.

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

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]