""" Derived from Andrej Karpathy's nanochat project. MIT License Copyright (c) 2025 Andrej Karpathy Permission is hereby granted, free of charge, to any person obtaining a copy of this software and associated documentation files (the "Software"), to deal in the Software without restriction, including without limitation the rights to use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies of the Software, and to permit persons to whom the Software is furnished to do so, subject to the following conditions: The above copyright notice and this permission notice shall be included in all copies or substantial portions of the Software. """ from __future__ import annotations import argparse from datetime import datetime import json from pathlib import Path import statistics import sys import numpy as np import torch from dropout_decay.datasets import ( encode_corpus, load_cached_splits, resolve_paths, train_or_load_tokenizer, ) from dropout_decay.experiments.artifacts import ( SELECTION_FIELDS, SUMMARY_FIELDS, build_model_selection, load_metrics, metric_key, planned_metric_key, summarize, write_csv, write_jsonl_row, ) from dropout_decay.experiments.device import assert_mps_only from dropout_decay.experiments.parsing import ( DEFAULT_DROPOUT_RATES, default_seeds, parse_anchor_decay_spec, parse_decay_spec, parse_model_spec, rate_label, ) from dropout_decay.experiments.progress import ProgressMeter from dropout_decay.experiments.reports import ( write_dropout_curve_svg, write_screen_markdown_summary, write_stream_markdown_summary, ) from dropout_decay.experiments.training import train_segment from dropout_decay.license import NANOCHAT_ATTRIBUTION from dropout_decay.specs import DropoutCondition, ModelSpec def static_conditions(dropout_rates: list[float]) -> list[DropoutCondition]: return [ DropoutCondition( name=f"static_dropout_{rate_label(rate)}", kind="static", initial=rate, final=rate, ) for rate in dropout_rates ] def run_fixed_static_sweep( *, args: argparse.Namespace, model_specs: list[ModelSpec], seeds: list[int], train_tokens: np.ndarray, val_tokens: np.ndarray, tokenizer_vocab_size: int, token_limits: list[int], device: torch.device, metrics_file, trace_file, completed_keys: set[tuple] | None = None, ) -> list[dict]: rows: list[dict] = [] completed_keys = completed_keys or set() conditions = static_conditions(sorted(set(args.dropout_rates))) planned = 0 for token_limit in token_limits: for model_spec in model_specs: for condition in conditions: for seed in seeds: key = planned_metric_key( mode=args.mode, condition=condition, model_spec=model_spec, seed=seed, token_limit=token_limit, ) if key not in completed_keys: planned += 1 progress = ProgressMeter(planned) for token_limit in token_limits: for model_spec in model_specs: best_val_loss = float("inf") worse_streak = 0 for condition in conditions: condition_rows: list[dict] = [] for seed in seeds: key = planned_metric_key( mode=args.mode, condition=condition, model_spec=model_spec, seed=seed, token_limit=token_limit, ) if key in completed_keys: write_jsonl_row( trace_file, { "event": "skipped_completed_condition", "run_mode": args.mode, "condition": condition.name, "model_name": model_spec.name, "seed": seed, "stage": None, "token_limit": int(token_limit), }, ) continue config = model_spec.config( tokenizer_vocab_size, args.block_size, condition.initial, ) model, optimizer, _, row = train_segment( run_mode=args.mode, condition=condition, model_spec=model_spec, config=config, train_tokens=train_tokens, val_tokens=val_tokens, token_limit=token_limit, steps=args.steps, seed=seed, args=args, device=device, dropout_fn=condition.make_fn( args.steps * args.batch_size * args.block_size ), metrics_file=metrics_file, trace_file=trace_file, ) rows.append(row) condition_rows.append(row) completed_keys.add(metric_key(row)) progress.mark_done(row) del model, optimizer torch.mps.empty_cache() if not condition_rows: continue mean_val_loss = statistics.fmean( float(row["val_eval_loss"]) for row in condition_rows ) if mean_val_loss < best_val_loss - args.screen_prune_min_delta: best_val_loss = mean_val_loss worse_streak = 0 elif mean_val_loss > best_val_loss + args.screen_prune_min_delta: worse_streak += 1 if ( args.mode == "screen_static" and args.screen_early_stop and worse_streak >= args.screen_prune_patience and condition.initial >= args.target_min_dropout ): write_jsonl_row( trace_file, { "event": "screen_pruned_model", "run_mode": args.mode, "model_name": model_spec.name, "token_limit": int(token_limit), "best_val_loss": best_val_loss, "pruned_after_dropout": condition.initial, "worse_streak": worse_streak, "remaining_dropouts": [ rate for rate in args.dropout_rates if rate > condition.initial ], }, ) break return rows def run_locked_stream( *, args: argparse.Namespace, model_specs: list[ModelSpec], seeds: list[int], train_tokens: np.ndarray, val_tokens: np.ndarray, tokenizer_vocab_size: int, stream_caps: list[int], device: torch.device, metrics_file, trace_file, completed_keys: set[tuple] | None = None, ) -> list[dict]: rows: list[dict] = [] completed_keys = completed_keys or set() conditions = args.anchor_decays + static_conditions(args.dropout_rates) + args.decays fallback_decay_tokens = ( args.decay_tokens or args.stage_steps * args.batch_size * args.block_size * len(stream_caps) ) planned = 0 for model_spec in model_specs: for condition in conditions: for seed in seeds: for stage, token_limit in enumerate(stream_caps): key = planned_metric_key( mode=args.mode, condition=condition, model_spec=model_spec, seed=seed, token_limit=token_limit, stage=stage, ) if key not in completed_keys: planned += 1 progress = ProgressMeter(planned) for model_spec in model_specs: for condition in conditions: for seed in seeds: model = None optimizer = None tokens_seen = 0 for stage, token_limit in enumerate(stream_caps): key = planned_metric_key( mode=args.mode, condition=condition, model_spec=model_spec, seed=seed, token_limit=token_limit, stage=stage, ) if key in completed_keys: write_jsonl_row( trace_file, { "event": "skipped_completed_condition", "run_mode": args.mode, "condition": condition.name, "model_name": model_spec.name, "seed": seed, "stage": stage, "token_limit": int(token_limit), }, ) continue config = model_spec.config( tokenizer_vocab_size, args.block_size, condition.initial, ) model, optimizer, tokens_seen, row = train_segment( run_mode=args.mode, condition=condition, model_spec=model_spec, config=config, train_tokens=train_tokens, val_tokens=val_tokens, token_limit=token_limit, steps=args.stage_steps, seed=seed, args=args, device=device, dropout_fn=condition.make_fn( fallback_decay_tokens, unique_tokens=token_limit, ), metrics_file=metrics_file, trace_file=trace_file, stage=stage, model=model, optimizer=optimizer, tokens_seen_start=tokens_seen, ) rows.append(row) completed_keys.add(metric_key(row)) progress.mark_done(row) del model, optimizer torch.mps.empty_cache() return rows def prepare_data(args: argparse.Namespace, output_dir: Path, required_train_tokens: int): cache_dir = Path(args.cache_dir) if args.cache_dir else output_dir / "cache" cache_dir.mkdir(parents=True, exist_ok=True) if args.use_cached_data: if args.force_retokenize: raise ValueError("--use-cached-data cannot be combined with --force-retokenize") return load_cached_splits( cache_dir=cache_dir, vocab_size=args.vocab_size, max_required_train_tokens=required_train_tokens, val_tokens=args.val_tokens, allow_short_corpus=args.allow_short_corpus, ) paths = resolve_paths(args.corpus, args.corpus_glob) tokenizer = train_or_load_tokenizer( paths=paths, output_dir=cache_dir, vocab_size=args.vocab_size, tokenizer_train_chars=args.tokenizer_train_chars, text_column=args.text_column, force_retrain=args.force_retokenize, ) splits = encode_corpus( paths=paths, tokenizer=tokenizer, output_dir=cache_dir, max_required_train_tokens=required_train_tokens, val_tokens=args.val_tokens, text_column=args.text_column, allow_short_corpus=args.allow_short_corpus, force_reencode=args.force_retokenize, ) return tokenizer, splits def run(args: argparse.Namespace) -> Path: device = assert_mps_only() seeds = default_seeds(args.mode, args.seeds) model_specs = [parse_model_spec(spec) for spec in args.models] if args.mode != "locked_stream" and (args.decays or args.anchor_decays): raise ValueError("--decays and --anchor-decays are only used with --mode locked_stream") if args.resume_from and args.mode == "locked_stream": raise ValueError("--resume-from currently supports fixed static sweeps only") if args.resume_from: output_dir = Path(args.resume_from) if not output_dir.exists(): raise FileNotFoundError(f"resume directory does not exist: {output_dir}") else: run_id = datetime.now().strftime("%Y%m%d-%H%M%S") output_dir = Path(args.output_dir) / args.mode / run_id output_dir.mkdir(parents=True, exist_ok=True) required_train_tokens = max( args.stream_token_caps if args.mode == "locked_stream" else args.token_limits ) tokenizer, splits = prepare_data(args, output_dir, required_train_tokens) token_limits = [min(limit, len(splits.train)) for limit in args.token_limits] stream_caps = [min(limit, len(splits.train)) for limit in args.stream_token_caps] args_payload = vars(args).copy() args_payload["decays"] = [condition.to_dict() for condition in args.decays] args_payload["anchor_decays"] = [ condition.to_dict() for condition in args.anchor_decays ] config_payload = { "args": args_payload, "mode": args.mode, "seeds": seeds, "models": [model.to_dict() for model in model_specs], "device": str(device), "torch": torch.__version__, "python": sys.version, "mps_available": torch.backends.mps.is_available(), "attribution": NANOCHAT_ATTRIBUTION, "tokenizer_path": str(splits.tokenizer_path), "encoded_path": str(splits.encoded_path), "train_tokens": int(len(splits.train)), "val_tokens": int(len(splits.val)), "effective_token_limits": [int(limit) for limit in token_limits], "effective_stream_token_caps": [int(limit) for limit in stream_caps], "resume_from": str(args.resume_from) if args.resume_from else None, } config_name = "config.resume.json" if args.resume_from else "config.json" (output_dir / config_name).write_text( json.dumps(config_payload, indent=2), encoding="utf-8", ) metrics_path = output_dir / "metrics.jsonl" trace_path = output_dir / "trace.jsonl" existing_rows = load_metrics(metrics_path) if args.resume_from else [] completed_keys = {metric_key(row) for row in existing_rows} with ( metrics_path.open("a" if args.resume_from else "w", encoding="utf-8") as metrics_file, trace_path.open("a" if args.resume_from else "w", encoding="utf-8") as trace_file, ): if args.mode in {"screen_static", "confirm_static"}: new_rows = run_fixed_static_sweep( args=args, model_specs=model_specs, seeds=seeds, train_tokens=splits.train, val_tokens=splits.val, tokenizer_vocab_size=tokenizer.vocab_size, token_limits=token_limits, device=device, metrics_file=metrics_file, trace_file=trace_file, completed_keys=completed_keys, ) else: new_rows = run_locked_stream( args=args, model_specs=model_specs, seeds=seeds, train_tokens=splits.train, val_tokens=splits.val, tokenizer_vocab_size=tokenizer.vocab_size, stream_caps=stream_caps, device=device, metrics_file=metrics_file, trace_file=trace_file, ) rows = existing_rows + new_rows summary = summarize(rows) (output_dir / "summary.json").write_text( json.dumps(summary, indent=2), encoding="utf-8", ) write_csv(output_dir / "summary.csv", summary, SUMMARY_FIELDS) if args.mode in {"screen_static", "confirm_static"}: selection = build_model_selection(summary, args) (output_dir / "model_selection.json").write_text( json.dumps(selection, indent=2), encoding="utf-8", ) write_csv(output_dir / "model_selection.csv", selection, SELECTION_FIELDS) write_screen_markdown_summary(output_dir, rows) write_dropout_curve_svg(output_dir, summary) elif args.mode == "locked_stream": write_stream_markdown_summary(output_dir, rows) print( json.dumps( { "output_dir": str(output_dir), "new_rows": len(new_rows), "total_metric_rows": len(rows), "summary_rows": len(summary), }, indent=2, ) ) return output_dir def build_parser() -> argparse.ArgumentParser: parser = argparse.ArgumentParser( description="MPS-only dropout/model selection experiments" ) parser.add_argument( "--mode", choices=["screen_static", "confirm_static", "locked_stream"], default="screen_static", ) parser.add_argument("--corpus", default=None, help="Text or parquet corpus path") parser.add_argument("--corpus-glob", default=None, help="Glob of text/parquet corpus paths") parser.add_argument("--text-column", default="text", help="Parquet text column") parser.add_argument( "--use-cached-data", action="store_true", help=( "Load tokenizer-v{vocab}.json and tokens-v{vocab}-*.npy from --cache-dir " "instead of requiring the original text/parquet corpus." ), ) parser.add_argument("--output-dir", default="runs") parser.add_argument( "--resume-from", default=None, help="Existing fixed-static run directory; completed metric rows are skipped", ) parser.add_argument("--cache-dir", default=".cache/dropout_decay") parser.add_argument( "--models", nargs="+", default=["8x8x256"], help="Model specs like 8x8x256 or name=8x8x256", ) parser.add_argument("--seeds", nargs="+", type=int, default=None) parser.add_argument("--token-limits", nargs="+", type=int, default=[5_000_000]) parser.add_argument( "--stream-token-caps", nargs="+", type=int, default=[5_000_000, 10_000_000, 20_000_000, 40_000_000], ) parser.add_argument("--val-tokens", type=int, default=500_000) parser.add_argument("--allow-short-corpus", action="store_true") parser.add_argument("--force-retokenize", action="store_true") parser.add_argument("--vocab-size", type=int, default=4096) parser.add_argument("--tokenizer-train-chars", type=int, default=10_000_000) parser.add_argument("--block-size", type=int, default=128) parser.add_argument("--batch-size", type=int, default=16) parser.add_argument("--steps", type=int, default=2000) parser.add_argument("--stage-steps", type=int, default=1000) parser.add_argument("--dropout-rates", nargs="*", type=float, default=DEFAULT_DROPOUT_RATES) parser.add_argument("--decays", nargs="*", type=parse_decay_spec, default=[]) parser.add_argument( "--anchor-decays", nargs="*", type=parse_anchor_decay_spec, default=[], help=( "Prefix-token anchor schedules like " "fit:250000=0.60,500000=0.40,1000000=0.30" ), ) parser.add_argument("--decay-tokens", type=int, default=None) parser.add_argument("--eval-batches", type=int, default=64) parser.add_argument("--train-eval-batches", type=int, default=32) parser.add_argument("--trace-eval-batches", type=int, default=8) parser.add_argument("--eval-every", type=int, default=0) parser.add_argument("--log-every", type=int, default=100) parser.add_argument("--lr", type=float, default=3e-4) parser.add_argument("--weight-decay", type=float, default=0.1) parser.add_argument("--grad-clip", type=float, default=1.0) parser.add_argument("--plateau-delta", type=float, default=0.01) parser.add_argument("--target-min-dropout", type=float, default=0.10) parser.add_argument("--min-nonzero-margin", type=float, default=0.01) parser.add_argument("--min-high-dropout-margin", type=float, default=0.03) parser.add_argument("--screen-early-stop", action="store_true") parser.add_argument("--screen-prune-patience", type=int, default=3) parser.add_argument("--screen-prune-min-delta", type=float, default=0.01) return parser def main() -> None: args = build_parser().parse_args() run(args) if __name__ == "__main__": main()