src/dropout_decay/experiments/runner.py

"""
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()