Single seed 20260830; 6 blocks; 8 heads per block; width 384; context 512; 512M training tokens. This report is a non-destructive snapshot and may be regenerated while training continues.
+{html.escape(summary)}
+{sections if sections else 'No plot-ready rows are available yet.
'} +""" + output.parent.mkdir(parents=True, exist_ok=True) + temporary = output.with_suffix(".html.tmp") + temporary.write_text(document, encoding="utf-8") + temporary.replace(output) + return output + + +def main() -> None: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--experiment-root", required=True, type=Path) + parser.add_argument("--output", required=True, type=Path) + args = parser.parse_args() + print(build_report(args.experiment_root, args.output)) + + +if __name__ == "__main__": + main() diff --git a/baseline/experiments/nanogpt_muonclip_large_2026_08_30/scripts/run_experiment.py b/baseline/experiments/nanogpt_muonclip_large_2026_08_30/scripts/run_experiment.py new file mode 100644 index 0000000..e4c130a --- /dev/null +++ b/baseline/experiments/nanogpt_muonclip_large_2026_08_30/scripts/run_experiment.py @@ -0,0 +1,354 @@ +#!/usr/bin/env python3 +"""Run, resume, inspect, and report the large MuonClip Apple-MPS experiment.""" + +from __future__ import annotations + +import argparse +import csv +import json +import math +import os +from pathlib import Path +import signal +import subprocess +import sys +from typing import Any + + +SCRIPT_PATH = Path(__file__).resolve() +EXPERIMENT_DIR = SCRIPT_PATH.parents[1] +REPOSITORY_ROOT = SCRIPT_PATH.parents[4] +PACKAGE_ROOT = REPOSITORY_ROOT / "baseline" / "nanogpt_one_head" +CONFIG_PATH = EXPERIMENT_DIR / "configs" / "muonclip_long_mps.yaml" +ROOT_ENV = "RG_NANOGPT_LARGE_EXPERIMENT_ROOT" +DEFAULT_ROOT = Path("/tmp/rg-nanogpt-muonclip-large-20260830") +OPTIMIZER = "muon_clip" +SEED = 20260830 + + +def experiment_root() -> Path: + value = Path(os.environ.get(ROOT_ENV, str(DEFAULT_ROOT))).expanduser() + if not value.is_absolute(): + raise ValueError(f"{ROOT_ENV} must be an absolute path") + return value.resolve(strict=False) + + +def paths() -> dict[str, Path]: + root = experiment_root() + result = { + "root": root, + "data": root / "data", + "results": root / "results", + "logs": root / "logs", + "report": root / "live_report", + } + for path in result.values(): + path.mkdir(parents=True, exist_ok=True) + return result + + +def run_dir() -> Path: + return paths()["results"] / OPTIMIZER / f"seed_{SEED}" + + +def _install_and_load() -> tuple[dict[str, Any], Any, Any, Any]: + from rg_nanogpt_one_head.muonclip import install_muonclip_extension + + install_muonclip_extension() + from rg_nanogpt_one_head.config import load_config, max_steps, tokens_per_step + from rg_nanogpt_one_head.model import GPT, GPTConfig + + cfg = load_config(CONFIG_PATH) + return cfg, GPT, GPTConfig, (max_steps, tokens_per_step) + + +def doctor(device_request: str, smoke_step: bool) -> None: + import datasets # noqa: F401 + import matplotlib # noqa: F401 + import pandas # noqa: F401 + import torch + import weightwatcher + from rg_nanogpt_one_head.runtime import choose_device + + cfg, GPT, GPTConfig, helpers = _install_and_load() + max_steps, tokens_per_step = helpers + model = GPT(GPTConfig(**cfg["model"])) + parameter_count = model.parameter_count() + step_tokens = tokens_per_step(cfg) + total_steps = max_steps(cfg) + processed_tokens = total_steps * step_tokens + resolved = choose_device(device_request) + payload = { + "config": str(CONFIG_PATH), + "experiment_root": str(experiment_root()), + "device_request": device_request, + "resolved_device": str(resolved), + "torch_version": torch.__version__, + "weightwatcher_version": weightwatcher.__version__, + "parameter_count": parameter_count, + "transformer_blocks": int(cfg["model"]["n_layer"]), + "attention_heads_per_block": int(cfg["model"]["n_head"]), + "transformer_matrix_count": 6 * int(cfg["model"]["n_layer"]), + "train_tokens": int(cfg["dataset"]["train_tokens"]), + "processed_tokens": processed_tokens, + "tokens_per_parameter": processed_tokens / parameter_count, + "tokens_per_optimizer_step": step_tokens, + "optimizer_steps": total_steps, + "warmup_steps": round( + total_steps + * float(cfg["optimizer_profiles"][OPTIMIZER]["warmup_fraction"]) + ), + "optimizer": OPTIMIZER, + "seed": SEED, + "estimated_dataset_gib": ( + 2 + * sum( + int(cfg["dataset"][name]) + for name in ("train_tokens", "val_tokens", "test_tokens") + ) + / 1024**3 + ), + } + if smoke_step: + from rg_nanogpt_one_head.config import optimizer_profile + from rg_nanogpt_one_head.optimizers import ( + make_optimizer_handles, + optimizer_step, + zero_grad, + ) + from rg_nanogpt_one_head.runtime import synchronize + + model.to(resolved) + model.train() + handles = make_optimizer_handles( + model, + optimizer_profile(cfg, OPTIMIZER), + ) + zero_grad(handles) + x = torch.randint( + 0, + int(cfg["model"]["vocab_size"]), + ( + int(cfg["training"]["batch_size"]), + int(cfg["model"]["block_size"]), + ), + device=resolved, + ) + _, loss = model(x, x) + if loss is None: + raise RuntimeError("large-model smoke step did not return a loss") + loss.backward() + torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0, foreach=False) + optimizer_step(handles) + synchronize(resolved) + loss_value = float(loss.detach().cpu()) + if not math.isfinite(loss_value): + raise FloatingPointError("large-model smoke loss is non-finite") + payload["smoke_optimizer_step"] = "passed" + payload["smoke_loss"] = loss_value + print(json.dumps(payload, indent=2, sort_keys=True)) + + +def prepare(force: bool) -> None: + cfg, _, _, _ = _install_and_load() + from rg_nanogpt_one_head.data import prepare_fineweb_edu + + prepare_fineweb_edu(cfg, paths()["data"], force=force) + + +def run_training(device: str, mps_retries: int, overwrite: bool) -> int: + resolved = paths() + command = [ + sys.executable, + "-u", + "-m", + "rg_nanogpt_one_head.muonclip", + "--config", + str(CONFIG_PATH), + "--optimizer", + OPTIMIZER, + "--seeds", + str(SEED), + "--data-root", + str(resolved["data"]), + "--results-root", + str(resolved["results"]), + "--device", + device, + "--mps-retries", + str(int(mps_retries)), + "--fail-fast", + ] + if overwrite: + command.append("--overwrite") + + log_path = resolved["logs"] / "train.log" + print("[large-muonclip] command:", " ".join(command), flush=True) + print(f"[large-muonclip] log: {log_path}", flush=True) + with log_path.open("a", encoding="utf-8") as log: + process = subprocess.Popen( + command, + cwd=EXPERIMENT_DIR, + stdout=subprocess.PIPE, + stderr=subprocess.STDOUT, + text=True, + bufsize=1, + ) + try: + assert process.stdout is not None + for line in process.stdout: + print(line, end="", flush=True) + log.write(line) + log.flush() + except KeyboardInterrupt: + process.send_signal(signal.SIGINT) + process.wait() + raise + return int(process.wait()) + + +def _last_csv_row(path: Path) -> dict[str, str] | None: + if not path.is_file() or path.stat().st_size == 0: + return None + for _ in range(2): + try: + with path.open("r", newline="", encoding="utf-8") as handle: + rows = list(csv.DictReader(handle)) + return rows[-1] if rows else None + except (csv.Error, OSError): + continue + return None + + +def _number(row: dict[str, str], name: str) -> float: + try: + return float(row[name]) + except (KeyError, TypeError, ValueError): + return float("nan") + + +def status() -> None: + cfg, _, _, helpers = _install_and_load() + total_steps = helpers[0](cfg) + directory = run_dir() + print(f"RUN: {directory}") + + completion = directory / "run_complete.json" + if completion.is_file(): + payload = json.loads(completion.read_text(encoding="utf-8")) + print("STATE: COMPLETE") + print(json.dumps(payload, indent=2, sort_keys=True)) + else: + print("STATE: INCOMPLETE OR RUNNING") + + row = _last_csv_row(directory / "metrics.csv") + if row is None: + print("TRAINING: waiting for the first evaluation row") + else: + step = int(_number(row, "step")) + elapsed = _number(row, "elapsed_sec") + remaining = max(0, total_steps - step) + eta = remaining * elapsed / step if step > 0 else float("nan") + print( + "TRAINING: " + f"step={step:,}/{total_steps:,} ({100 * step / total_steps:.2f}%) " + f"epoch={_number(row, 'epoch'):.4f} " + f"train_loss={_number(row, 'train_loss'):.4f} " + f"val_loss={_number(row, 'val_loss'):.4f} " + f"val_ppl={_number(row, 'val_perplexity'):.2f} " + f"val_acc={100 * _number(row, 'val_accuracy'):.2f}% " + f"tokens_per_sec={_number(row, 'tokens_per_sec'):,.0f} " + f"eta_hours={eta / 3600:.1f}" + ) + + spectral = _last_csv_row(directory / "spectral" / "summary.csv") + if spectral is None: + print("WEIGHTWATCHER: waiting for the first permanent state") + else: + print( + "WEIGHTWATCHER: " + f"step={int(_number(spectral, 'step')):,} " + f"epoch={_number(spectral, 'epoch'):.4f} " + f"matrices={int(_number(spectral, 'n_matrices'))} " + f"alpha_raw_median={_number(spectral, 'alpha_raw_median'):.3f} " + "alpha_clip_median=" + f"{_number(spectral, 'alpha_clip_xmax_median'):.3f} " + f"ERG_gap_median={_number(spectral, 'ERG_gap_median'):.3f} " + f"num_traps_mean={_number(spectral, 'num_traps_mean'):.2f}" + ) + + checkpoint = directory / "checkpoint_latest.pt" + print( + "CHECKPOINT: " + + ( + f"present ({checkpoint.stat().st_size / 1024**2:.1f} MiB)" + if checkpoint.is_file() + else "not written yet" + ) + ) + process = subprocess.run( + ["pgrep", "-fl", "rg_nanogpt_one_head.muonclip"], + text=True, + capture_output=True, + check=False, + ) + print("PROCESS:") + print(process.stdout.strip() or "no MuonClip process found") + + +def report(open_report: bool) -> None: + output = paths()["report"] / "report.html" + command = [ + sys.executable, + str(EXPERIMENT_DIR / "scripts" / "build_live_report.py"), + "--experiment-root", + str(experiment_root()), + "--output", + str(output), + ] + subprocess.run(command, check=True) + print(f"REPORT: {output}") + if open_report: + subprocess.run(["open", str(output)], check=True) + + +def main() -> None: + parser = argparse.ArgumentParser(description=__doc__) + subparsers = parser.add_subparsers(dest="command", required=True) + + doctor_parser = subparsers.add_parser("doctor") + doctor_parser.add_argument("--device", default="auto") + doctor_parser.add_argument("--smoke-step", action="store_true") + + prepare_parser = subparsers.add_parser("prepare") + prepare_parser.add_argument("--force", action="store_true") + + run_parser = subparsers.add_parser("run") + run_parser.add_argument("--device", default="auto") + run_parser.add_argument("--mps-retries", type=int, default=20) + run_parser.add_argument("--overwrite", action="store_true") + + subparsers.add_parser("status") + + report_parser = subparsers.add_parser("report") + report_parser.add_argument("--open", action="store_true") + + args = parser.parse_args() + if args.command == "doctor": + doctor(args.device, args.smoke_step) + elif args.command == "prepare": + prepare(args.force) + elif args.command == "run": + if args.mps_retries < 0: + parser.error("--mps-retries must be nonnegative") + raise SystemExit( + run_training(args.device, args.mps_retries, args.overwrite) + ) + elif args.command == "status": + status() + elif args.command == "report": + report(args.open) + + +if __name__ == "__main__": + main() diff --git a/baseline/nanogpt_one_head/src/rg_nanogpt_one_head/config.py b/baseline/nanogpt_one_head/src/rg_nanogpt_one_head/config.py index 2499d20..4359ded 100644 --- a/baseline/nanogpt_one_head/src/rg_nanogpt_one_head/config.py +++ b/baseline/nanogpt_one_head/src/rg_nanogpt_one_head/config.py @@ -225,8 +225,6 @@ def validate_config(cfg: dict[str, Any]) -> None: 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_layer"]) != 1: - raise ValueError("this experiment is fixed to one transformer block") 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: diff --git a/baseline/nanogpt_one_head/src/rg_nanogpt_one_head/model.py b/baseline/nanogpt_one_head/src/rg_nanogpt_one_head/model.py index d40fd38..029b426 100644 --- a/baseline/nanogpt_one_head/src/rg_nanogpt_one_head/model.py +++ b/baseline/nanogpt_one_head/src/rg_nanogpt_one_head/model.py @@ -20,8 +20,8 @@ class GPTConfig: tie_weights: bool = True def __post_init__(self) -> None: - if self.n_layer != 1: - raise ValueError("the reference architecture is fixed to one block") + if self.n_layer < 1: + raise ValueError("n_layer must be positive") if self.n_head < 1: raise ValueError("n_head must be positive") if self.n_embd % self.n_head != 0: diff --git a/baseline/nanogpt_one_head/src/rg_nanogpt_one_head/spectral.py b/baseline/nanogpt_one_head/src/rg_nanogpt_one_head/spectral.py index 1260d56..856a12e 100644 --- a/baseline/nanogpt_one_head/src/rg_nanogpt_one_head/spectral.py +++ b/baseline/nanogpt_one_head/src/rg_nanogpt_one_head/spectral.py @@ -4,7 +4,7 @@ import json from pathlib import Path import random -from typing import Any +from typing import Any, Sequence import numpy as np import pandas as pd @@ -44,7 +44,7 @@ class WeightMatrixHolder(nn.Module): - """CPU-only Linear view of the six one-block transformer matrices.""" + """CPU-only Linear views of every declared transformer matrix.""" def __init__(self, model: GPT) -> None: super().__init__() @@ -173,6 +173,7 @@ def _validate_weightwatcher_frame( frame: pd.DataFrame, *, finger_policy: str | bool, + expected_matrix_names: Sequence[str] | None = None, ) -> None: """Reject incomplete results and stale pre-finger-policy caches.""" @@ -207,10 +208,26 @@ def _validate_weightwatcher_frame( + ", ".join(missing) + ". Remove the stale run directory before restarting." ) - if len(frame) != 6 or frame["matrix_name"].nunique() != 6: + expected_names = ( + tuple(str(value) for value in expected_matrix_names) + if expected_matrix_names is not None + else () + ) + expected_count = len(expected_names) if expected_names else 6 + observed_names = tuple(frame["matrix_name"].astype(str)) + inventory_matches = ( + set(observed_names) == set(expected_names) + if expected_names + else len(set(observed_names)) == expected_count + ) + if ( + len(frame) != expected_count + or frame["matrix_name"].nunique() != expected_count + or not inventory_matches + ): raise RuntimeError( - "WeightWatcher must return exactly the six declared transformer " - "matrices" + "WeightWatcher must return exactly the declared transformer " + f"matrix inventory ({expected_count} matrices)" ) numeric = [ @@ -455,6 +472,9 @@ def run_weightwatcher( device = model_device(model) synchronize(device) current_model_hash = model_state_sha256(model.state_dict()) + expected_matrix_names = tuple( + name for name, _, _, _ in transformer_matrix_items(model) + ) diagnostic_seed = int(seed) + 1_000_003 + int(step) if raw_path.is_file(): frame = pd.read_csv(raw_path) @@ -462,6 +482,7 @@ def run_weightwatcher( _validate_weightwatcher_frame( frame, finger_policy=finger_policy, + expected_matrix_names=expected_matrix_names, ) expected_identities = { "run_seed": int(seed), @@ -519,8 +540,8 @@ def run_weightwatcher( try: # WeightWatcher remains deliberately CPU/NumPy based. For TPU/XLA, - # WeightMatrixHolder materializes only the six hidden matrices on the - # host and never exposes the live accelerator model to NumPy. + # WeightMatrixHolder materializes only the declared hidden matrices on + # the host and never exposes the live accelerator model to NumPy. holder = WeightMatrixHolder(model) watcher = ww.WeightWatcher(model=holder) analysis_kwargs: dict[str, Any] = { @@ -592,6 +613,7 @@ def run_weightwatcher( _validate_weightwatcher_frame( frame, finger_policy=finger_policy, + expected_matrix_names=expected_matrix_names, ) _atomic_csv(raw_path, frame) return _record_successful_frame( diff --git a/baseline/nanogpt_one_head/tests/test_muonclip_large_20260830.py b/baseline/nanogpt_one_head/tests/test_muonclip_large_20260830.py new file mode 100644 index 0000000..aa11fb9 --- /dev/null +++ b/baseline/nanogpt_one_head/tests/test_muonclip_large_20260830.py @@ -0,0 +1,118 @@ +from __future__ import annotations + +import importlib.util +from pathlib import Path +import sys + +import pandas as pd +import pytest + + +PACKAGE_ROOT = Path(__file__).resolve().parents[1] +REPOSITORY_ROOT = PACKAGE_ROOT.parents[1] +EXPERIMENT_ROOT = ( + REPOSITORY_ROOT + / "baseline" + / "experiments" + / "nanogpt_muonclip_large_2026_08_30" +) +CONFIG_PATH = EXPERIMENT_ROOT / "configs" / "muonclip_long_mps.yaml" +RUNNER_PATH = EXPERIMENT_ROOT / "scripts" / "run_experiment.py" +REPORT_PATH = EXPERIMENT_ROOT / "scripts" / "build_live_report.py" + +sys.path.insert(0, str(PACKAGE_ROOT / "src")) + +from rg_nanogpt_one_head.config import epoch_step_map, max_steps, tokens_per_step +from rg_nanogpt_one_head.model import GPT, GPTConfig, transformer_matrix_items +from rg_nanogpt_one_head.muonclip import install_muonclip_extension +from rg_nanogpt_one_head.spectral import _validate_weightwatcher_frame + + +def _load_config() -> dict: + install_muonclip_extension() + from rg_nanogpt_one_head.config import load_config + + return load_config(CONFIG_PATH) + + +def test_large_muonclip_protocol_has_the_declared_scale_and_schedule() -> None: + cfg = _load_config() + assert cfg["model"] == { + "vocab_size": 50_257, + "block_size": 512, + "n_layer": 6, + "n_head": 8, + "n_embd": 384, + "dropout": 0.0, + "bias": False, + "tie_weights": True, + } + assert cfg["dataset"]["train_tokens"] == 512_000_000 + assert cfg["training"]["seeds"] == [20260830] + assert tokens_per_step(cfg) == 8_192 + assert max_steps(cfg) == 62_500 + assert len(epoch_step_map(cfg)) == 26 + profile = cfg["optimizer_profiles"]["muon_clip"] + assert profile["learning_rate"] == pytest.approx(2e-4) + assert profile["min_learning_rate"] == pytest.approx(1e-5) + assert profile["warmup_fraction"] == pytest.approx(0.016) + + +def test_gpt_and_matrix_inventory_support_multiple_blocks() -> None: + model = GPT( + GPTConfig( + vocab_size=64, + block_size=8, + n_layer=3, + n_head=4, + n_embd=32, + ) + ) + matrices = transformer_matrix_items(model) + assert len(matrices) == 18 + assert {block for _, _, block, _ in matrices} == {0, 1, 2} + assert len({name for name, _, _, _ in matrices}) == 18 + + +def test_weightwatcher_validator_accepts_full_multiblock_inventory() -> None: + names = [f"L{block:02d}_{kind}" for block in range(2) for kind in ( + "W_Q", "W_K", "W_V", "W_O", "W_MLP_IN", "W_MLP_OUT" + )] + frame = pd.DataFrame( + { + "matrix_name": names, + "alpha": 3.0, + "alpha_raw": 3.0, + "ERG_gap": 1.0, + "num_traps": 0.0, + "rand_distance": 0.2, + "finger_policy": "none", + "primary_alpha_variant": "raw", + "weightwatcher_analysis_calls": 1, + "run_seed": 1, + "diagnostic_seed": 2, + "protocol_fingerprint": "fingerprint", + "model_state_sha256": "state", + } + ) + _validate_weightwatcher_frame( + frame, + finger_policy=False, + expected_matrix_names=names, + ) + with pytest.raises(RuntimeError, match="12 matrices"): + _validate_weightwatcher_frame( + frame.iloc[:-1], + finger_policy=False, + expected_matrix_names=names, + ) + + +def test_experiment_entrypoints_are_importable() -> None: + assert (EXPERIMENT_ROOT / "README.md").is_file() + assert CONFIG_PATH.is_file() + for path in (RUNNER_PATH, REPORT_PATH): + spec = importlib.util.spec_from_file_location(path.stem, path) + assert spec is not None and spec.loader is not None + module = importlib.util.module_from_spec(spec) + spec.loader.exec_module(module)