""" 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 from collections import defaultdict from pathlib import Path import statistics def write_screen_markdown_summary(output_dir: Path, rows: list[dict]) -> None: if not rows: return static_rows = [ row for row in rows if row["run_mode"] == "screen_static" and row["condition_kind"] == "static" ] if not static_rows: return by_model_prefix_rate: dict[tuple[str, int, float], list[dict]] = defaultdict(list) for row in rows: if row["run_mode"] == "screen_static" and row["condition_kind"] == "static": by_model_prefix_rate[ ( row["model_name"], int(row["token_limit"]), float(row["dropout_initial"]), ) ].append(row) aggregates: list[dict] = [] for (model_name, prefix, dropout), group_rows in by_model_prefix_rate.items(): first = group_rows[0] val_losses = [float(row["val_eval_loss"]) for row in group_rows] train_losses = [float(row["train_eval_loss"]) for row in group_rows] gaps = [float(row["generalization_gap"]) for row in group_rows] aggregates.append( { "model_name": model_name, "token_limit": prefix, "dropout_initial": dropout, "n": len(group_rows), "mean_val_eval_loss": statistics.fmean(val_losses), "std_val_eval_loss": statistics.stdev(val_losses) if len(val_losses) > 1 else 0.0, "mean_train_eval_loss": statistics.fmean(train_losses), "std_train_eval_loss": statistics.stdev(train_losses) if len(train_losses) > 1 else 0.0, "mean_generalization_gap": statistics.fmean(gaps), "std_generalization_gap": statistics.stdev(gaps) if len(gaps) > 1 else 0.0, "parameters": int(first["parameters"]), "n_layer": int(first["n_layer"]), "n_head": int(first["n_head"]), "n_embd": int(first["n_embd"]), "block_size": int(first["model_config"]["block_size"]), "vocab_size": int(first["model_config"]["vocab_size"]), "tokens_seen": int(first["tokens_seen"]), "seeds": sorted({int(row["seed"]) for row in group_rows}), } ) by_model: dict[str, list[dict]] = defaultdict(list) for row in aggregates: by_model[row["model_name"]].append(row) model_rows = [] for model_name, model_group in by_model.items(): first = model_group[0] seeds = sorted({seed for row in model_group for seed in row["seeds"]}) model_rows.append( { "model_name": model_name, "parameters": first["parameters"], "n_layer": first["n_layer"], "n_head": first["n_head"], "n_embd": first["n_embd"], "block_size": first["block_size"], "vocab_size": first["vocab_size"], "seeds": seeds, } ) lines = [ "# Static Dropout Screen Summary", "", f"Run directory: `{output_dir}`", "", "## Models", "", "| Model | Params | Layers | Heads | Embedding | Block | Vocab | Seeds |", "|---|---:|---:|---:|---:|---:|---:|---|", ] for model in sorted(model_rows, key=lambda item: item["parameters"]): lines.append( "| " f"`{model['model_name']}` | {model['parameters']:,} | " f"{model['n_layer']} | {model['n_head']} | {model['n_embd']} | " f"{model['block_size']} | {model['vocab_size']} | " f"{', '.join(str(seed) for seed in model['seeds'])} |" ) lines.extend( [ "", "## Best Dropout By Model And Prefix", "", "| Model | Prefix tokens | Effective epochs | Best dropout | Mean val loss | Val std | Mean train loss | Mean gap | Plateau/bracket note |", "|---|---:|---:|---:|---:|---:|---:|---:|---|", ] ) for model_name, model_group in sorted(by_model.items()): by_prefix: dict[int, list[dict]] = defaultdict(list) for row in model_group: by_prefix[int(row["token_limit"])].append(row) for prefix, prefix_rows in sorted(by_prefix.items()): best = min(prefix_rows, key=lambda row: row["mean_val_eval_loss"]) rates = [float(row["dropout_initial"]) for row in prefix_rows] eff_epochs = float(best["tokens_seen"]) / prefix if best["dropout_initial"] == max(rates): note = "not bracketed; best at top of tested grid" elif best["dropout_initial"] == min(rates): note = "not bracketed; best at bottom of tested grid" else: note = "bracketed by tested grid" lines.append( "| " f"`{model_name}` | {prefix:,} | {eff_epochs:.2f} | " f"{best['dropout_initial']:.2f} | " f"{best['mean_val_eval_loss']:.4f} | " f"{best['std_val_eval_loss']:.4f} | " f"{best['mean_train_eval_loss']:.4f} | " f"{best['mean_generalization_gap']:.4f} | {note} |" ) for model_name, model_group in sorted(by_model.items()): by_prefix = defaultdict(list) for row in model_group: by_prefix[int(row["token_limit"])].append(row) lines.extend( [ "", f"## Model `{model_name}`", ] ) for prefix, prefix_rows in sorted(by_prefix.items()): eff_epochs = float(prefix_rows[0]["tokens_seen"]) / prefix lines.extend( [ "", f"### Prefix {prefix:,} Tokens ({eff_epochs:.2f} Effective Epochs)", "", "| Dropout | N | Mean val loss | Val std | Mean train loss | Train std | Mean gap | Gap std | Sampled tokens | Params |", "|---:|---:|---:|---:|---:|---:|---:|---:|---:|---:|", ] ) for row in sorted(prefix_rows, key=lambda item: item["dropout_initial"]): lines.append( "| " f"{row['dropout_initial']:.2f} | {row['n']} | " f"{row['mean_val_eval_loss']:.4f} | " f"{row['std_val_eval_loss']:.4f} | " f"{row['mean_train_eval_loss']:.4f} | " f"{row['std_train_eval_loss']:.4f} | " f"{row['mean_generalization_gap']:.4f} | " f"{row['std_generalization_gap']:.4f} | " f"{int(row['tokens_seen']):,} | {int(row['parameters']):,} |" ) output = "\n".join(lines) + "\n" (output_dir / "RESULT_SUMMARY.md").write_text(output, encoding="utf-8") def svg_escape(value: object) -> str: return ( str(value) .replace("&", "&") .replace("<", "<") .replace(">", ">") .replace('"', """) ) def write_dropout_curve_svg(output_dir: Path, summary: list[dict]) -> None: rows = [ row for row in summary if row["run_mode"] == "screen_static" and row["condition_kind"] == "static" ] if not rows: return grouped: dict[tuple[str, int], list[dict]] = defaultdict(list) model_params: dict[str, int] = {} for row in rows: model_name = row["model_name"] grouped[(model_name, int(row["token_limit"]))].append(row) model_params[model_name] = int(row["parameters"]) models = sorted(model_params, key=lambda name: model_params[name]) prefixes = sorted({int(row["token_limit"]) for row in rows}) panel_w, panel_h = 230, 170 margin_l, margin_t = 58, 34 plot_w, plot_h = 142, 94 gap_x, gap_y = 18, 38 width = margin_l + len(prefixes) * panel_w + gap_x height = 70 + len(models) * (panel_h + gap_y) colors = ["#1f77b4", "#d62728", "#2ca02c", "#9467bd", "#ff7f0e"] parts = [ f'', "", '', 'Static dropout law: validation loss vs dropout', 'Each panel uses its own y-scale. Points are one-seed means unless N > 1.', ] for col, prefix in enumerate(prefixes): x = margin_l + col * panel_w + plot_w / 2 parts.append( f'{prefix:,} prefix tokens' ) for row_idx, model_name in enumerate(models): row_y = 92 + row_idx * (panel_h + gap_y) parts.append( f'' f'{svg_escape(model_name)} ({model_params[model_name] / 1_000_000:.1f}M)' ) for col, prefix in enumerate(prefixes): panel_x = margin_l + col * panel_w panel_y = row_y curve = sorted( grouped.get((model_name, prefix), []), key=lambda item: float(item["dropout_initial"]), ) if not curve: continue losses = [float(item["mean_val_eval_loss"]) for item in curve] min_loss, max_loss = min(losses), max(losses) pad = max(0.02, (max_loss - min_loss) * 0.08) y_min, y_max = min_loss - pad, max_loss + pad best = min(curve, key=lambda item: float(item["mean_val_eval_loss"])) def px(dropout: float) -> float: return panel_x + (dropout / 0.9) * plot_w def py(loss: float) -> float: scale = (loss - y_min) / (y_max - y_min) return panel_y + plot_h - scale * plot_h parts.extend( [ f'', f'', f'', f'{y_max:.2f}', f'{y_min:.2f}', f'0', f'0.9', ] ) points = " ".join( f"{px(float(item['dropout_initial'])):.1f},{py(float(item['mean_val_eval_loss'])):.1f}" for item in curve ) color = colors[row_idx % len(colors)] parts.append(f'') for item in curve: dropout = float(item["dropout_initial"]) loss = float(item["mean_val_eval_loss"]) radius = 4 if item is best else 2.7 fill = "#111827" if item is best else "#ffffff" parts.append( f'' ) parts.append( f'' f'best p={float(best["dropout_initial"]):.2f}' ) parts.append( f'' f'loss={float(best["mean_val_eval_loss"]):.3f}' ) parts.append("") (output_dir / "dropout_curves.svg").write_text("\n".join(parts), encoding="utf-8") def write_stream_markdown_summary(output_dir: Path, rows: list[dict]) -> None: stream_rows = [row for row in rows if row["run_mode"] == "locked_stream"] if not stream_rows: return by_condition_stage: dict[tuple[str, int], list[dict]] = defaultdict(list) by_condition: dict[str, list[dict]] = defaultdict(list) for row in stream_rows: condition = row["condition"] by_condition_stage[(condition, int(row["stage"]))].append(row) by_condition[condition].append(row) first = stream_rows[0] seeds = sorted({int(row["seed"]) for row in stream_rows}) conditions = sorted( by_condition, key=lambda name: ( by_condition[name][0]["condition_kind"] != "anchor_decay", by_condition[name][0]["dropout_initial"], name, ), ) stages = sorted({int(row["stage"]) for row in stream_rows}) lines = [ "# Locked Streaming Dropout Summary", "", f"Run directory: `{output_dir}`", "", ( f"Model: `{first['model_name']}` causal Transformer, " f"{int(first['parameters']):,} parameters, {first['n_layer']} layers, " f"{first['n_head']} heads, {first['n_embd']} embedding dim." ), ( f"Training per stage: {first['steps']:,} steps. " "Sampled tokens are cumulative in each stage row. " f"Seeds present: {', '.join(str(seed) for seed in seeds)}." ), "", "## Condition Ranking", "", "| Condition | Kind | Final dropout | Mean trajectory val loss | Final val loss | Final gap | Dropout path |", "|---|---|---:|---:|---:|---:|---|", ] ranking = [] for condition in conditions: stage_items = [] for stage in stages: group = by_condition_stage.get((condition, stage), []) if not group: continue stage_items.append( { "stage": stage, "token_limit": int(group[0]["token_limit"]), "mean_val": statistics.fmean( float(row["val_eval_loss"]) for row in group ), "mean_gap": statistics.fmean( float(row["generalization_gap"]) for row in group ), "mean_dropout": statistics.fmean( float(row["dropout_active_final"]) for row in group ), "kind": group[0]["condition_kind"], } ) if not stage_items: continue final = max(stage_items, key=lambda item: item["stage"]) ranking.append( { "condition": condition, "kind": stage_items[0]["kind"], "trajectory_val": statistics.fmean(item["mean_val"] for item in stage_items), "final_val": final["mean_val"], "final_gap": final["mean_gap"], "final_dropout": final["mean_dropout"], "dropout_path": " -> ".join( f"{item['mean_dropout']:.2f}" for item in stage_items ), } ) for item in sorted(ranking, key=lambda row: row["trajectory_val"]): lines.append( "| " f"`{item['condition']}` | {item['kind']} | " f"{item['final_dropout']:.2f} | {item['trajectory_val']:.4f} | " f"{item['final_val']:.4f} | {item['final_gap']:.4f} | " f"{item['dropout_path']} |" ) lines.extend(["", "## Stage Trajectory", ""]) for stage in stages: stage_groups = { condition: by_condition_stage[(condition, stage)] for condition in conditions if (condition, stage) in by_condition_stage } if not stage_groups: continue prefix = int(next(iter(stage_groups.values()))[0]["token_limit"]) lines.extend( [ f"### Stage {stage}: {prefix:,} Prefix Tokens", "", "| Condition | Dropout | Mean val loss | Mean train loss | Mean gap | N |", "|---|---:|---:|---:|---:|---:|", ] ) for condition, group in sorted( stage_groups.items(), key=lambda item: statistics.fmean( float(row["val_eval_loss"]) for row in item[1] ), ): val = statistics.fmean(float(row["val_eval_loss"]) for row in group) train = statistics.fmean(float(row["train_eval_loss"]) for row in group) gap = statistics.fmean(float(row["generalization_gap"]) for row in group) dropout = statistics.fmean( float(row["dropout_active_final"]) for row in group ) lines.append( "| " f"`{condition}` | {dropout:.2f} | {val:.4f} | " f"{train:.4f} | {gap:.4f} | {len(group)} |" ) lines.append("") (output_dir / "RESULT_SUMMARY.md").write_text( "\n".join(lines).rstrip() + "\n", encoding="utf-8", )