| """ |
| 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'<svg xmlns="http://www.w3.org/2000/svg" width="{width}" height="{height}" viewBox="0 0 {width} {height}">', |
| "<style>", |
| "text{font-family:Arial,Helvetica,sans-serif;fill:#111827}", |
| ".small{font-size:10px}.label{font-size:11px}.title{font-size:15px;font-weight:700}", |
| ".axis{stroke:#374151;stroke-width:1}.grid{stroke:#e5e7eb;stroke-width:1}.line{fill:none;stroke-width:2}", |
| "</style>", |
| '<rect width="100%" height="100%" fill="#ffffff"/>', |
| '<text x="24" y="28" class="title">Static dropout law: validation loss vs dropout</text>', |
| '<text x="24" y="48" class="label">Each panel uses its own y-scale. Points are one-seed means unless N > 1.</text>', |
| ] |
|
|
| for col, prefix in enumerate(prefixes): |
| x = margin_l + col * panel_w + plot_w / 2 |
| parts.append( |
| f'<text x="{x:.1f}" y="70" text-anchor="middle" class="label">{prefix:,} prefix tokens</text>' |
| ) |
|
|
| for row_idx, model_name in enumerate(models): |
| row_y = 92 + row_idx * (panel_h + gap_y) |
| parts.append( |
| f'<text x="24" y="{row_y + 48}" class="label" transform="rotate(-90 24 {row_y + 48})">' |
| f'{svg_escape(model_name)} ({model_params[model_name] / 1_000_000:.1f}M)</text>' |
| ) |
| 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'<line x1="{panel_x:.1f}" y1="{panel_y:.1f}" x2="{panel_x:.1f}" y2="{panel_y + plot_h:.1f}" class="axis"/>', |
| f'<line x1="{panel_x:.1f}" y1="{panel_y + plot_h:.1f}" x2="{panel_x + plot_w:.1f}" y2="{panel_y + plot_h:.1f}" class="axis"/>', |
| f'<line x1="{panel_x:.1f}" y1="{panel_y:.1f}" x2="{panel_x + plot_w:.1f}" y2="{panel_y:.1f}" class="grid"/>', |
| f'<text x="{panel_x:.1f}" y="{panel_y - 6:.1f}" class="small">{y_max:.2f}</text>', |
| f'<text x="{panel_x:.1f}" y="{panel_y + plot_h + 13:.1f}" class="small">{y_min:.2f}</text>', |
| f'<text x="{panel_x:.1f}" y="{panel_y + plot_h + 28:.1f}" class="small">0</text>', |
| f'<text x="{panel_x + plot_w:.1f}" y="{panel_y + plot_h + 28:.1f}" text-anchor="end" class="small">0.9</text>', |
| ] |
| ) |
| 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'<polyline points="{points}" class="line" stroke="{color}"/>') |
| 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'<circle cx="{px(dropout):.1f}" cy="{py(loss):.1f}" r="{radius}" fill="{fill}" stroke="{color}" stroke-width="1.5"/>' |
| ) |
| parts.append( |
| f'<text x="{panel_x + plot_w + 8:.1f}" y="{panel_y + 14:.1f}" class="small">' |
| f'best p={float(best["dropout_initial"]):.2f}</text>' |
| ) |
| parts.append( |
| f'<text x="{panel_x + plot_w + 8:.1f}" y="{panel_y + 28:.1f}" class="small">' |
| f'loss={float(best["mean_val_eval_loss"]):.3f}</text>' |
| ) |
|
|
| parts.append("</svg>") |
| (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", |
| ) |
|
|
