From 6cb0b47dac1a7f5f7b4e5df2d9b055ea69859c00 Mon Sep 17 00:00:00 2001 From: Kauna <16511995+klei22@users.noreply.github.com> Date: Tue, 16 Jun 2026 13:03:07 -0700 Subject: [PATCH] fix(viewer): restart playback from earliest snapshot --- README.md | 46 ++ analysis/min_angle_graph_plotly_viewer.html | 451 ++++++++++++++++++++ demos/min_angle_graph_export_demo.sh | 68 +++ explorations/min_angle_graph_export.yaml | 39 ++ train.py | 44 ++ train_args.py | 14 + utils/min_angle_graph_export.py | 152 +++++++ 7 files changed, 814 insertions(+) create mode 100644 analysis/min_angle_graph_plotly_viewer.html create mode 100755 demos/min_angle_graph_export_demo.sh create mode 100644 explorations/min_angle_graph_export.yaml create mode 100644 utils/min_angle_graph_export.py diff --git a/README.md b/README.md index 6a8182a68b..fdf66a08a9 100644 --- a/README.md +++ b/README.md @@ -258,6 +258,52 @@ python3 view_model_stats.py run1_stats.csv run2_stats.csv See [documentation/Model_Stats_Table.md](documentation/Model_Stats_Table.md) for more details. +## LM-head minimum-angle graph exports + +`train.py` can optionally export the same closest-neighbor graph used by the +LM-head angle explorer's minimum-angular-distance view after each validation +loss. The export treats each LM-head row as a token vector, streams +row/column blocks through the selected compute device, excludes each token's +self-distance, and records the closest non-self token by signed `0°–180°` +angular distance. The full `vocab_size × vocab_size` angle matrix is never +materialized. + +Enable the export by choosing a labelled output directory and turning on the +per-eval flag: + +```bash +python3 train.py \ + --export_min_angle_graph_dir out/min_angle_graphs \ + --export_min_angle_graph_each_eval \ + --export_min_angle_graph_label shakespeare_char_baseline +``` + +Each validation step writes a labelled CSV and JSON sidecar whose filename +includes the label, iteration, and validation loss: + +```text +out/min_angle_graphs/shakespeare_char_baseline_iter_00000250_val_1.234567.csv +out/min_angle_graphs/shakespeare_char_baseline_iter_00000250_val_1.234567.json +``` + +The CSV is an edge list with one directed nearest-neighbor edge per token: + +```text +token_id,nearest_token_id,min_angle_deg,cosine,token_vector_length,nearest_token_vector_length,min_angle_rank +``` + +Use `--export_min_angle_graph_block_size` to tune temporary matrix multiply +size, and `--export_min_angle_graph_device` to choose `auto`, `cpu`, or a CUDA +device such as `cuda:0`. In `auto` mode, the export uses the LM-head tensor's +current CUDA device when possible, otherwise streams blocks to `cuda:0` if CUDA +is available, and falls back to CPU. + +To review a sequence of exports visually, open +`analysis/min_angle_graph_plotly_viewer.html` in a browser and select the CSV +files from one export directory. The page sorts snapshots by the iteration in +their filenames and provides previous/next/play controls plus a slider to step +from the first validation export to the last. + ## TODO Section: TODO: Add links and descriptions to other Readme's and Demos. diff --git a/analysis/min_angle_graph_plotly_viewer.html b/analysis/min_angle_graph_plotly_viewer.html new file mode 100644 index 0000000000..1b57eda472 --- /dev/null +++ b/analysis/min_angle_graph_plotly_viewer.html @@ -0,0 +1,451 @@ + + + + + + LM-head Minimum-Angle Graph Viewer + + + + +
+

LM-head Minimum-Angle Graph Viewer

+

+ Load one or more CSV exports from --export_min_angle_graph_dir. + The viewer sorts snapshots by iteration and lets you step from the first + validation export to the last without uploading data anywhere. +

+
+ +
+
+
+
+ + +
+
+ + +
+
+ + +
+
+ + +
+
+
+ + + + +
+
Choose CSV files to begin.
+
+ +
+
+
+
+
+ +
+

Snapshot rows

+
+ + + + + + + + + + + + + +
ranktoken_idnearest_token_idmin_angle_degcosinetoken_vector_lengthnearest_token_vector_length
+
+
+
+ + + + diff --git a/demos/min_angle_graph_export_demo.sh b/demos/min_angle_graph_export_demo.sh new file mode 100755 index 0000000000..0b8fe11037 --- /dev/null +++ b/demos/min_angle_graph_export_demo.sh @@ -0,0 +1,68 @@ +#!/bin/bash + +# Demo the per-validation LM-head minimum-angle graph export on the small +# exploration config, then point the user at the local Plotly viewer. +# +# Optional environment variables: +# OUTPUT_DIR Base output directory for train.py runs. Default: out +# PREFIX Prefix for exploration run names. Default: timestamped demo prefix +# EXPORT_ROOT Directory where the exploration config writes CSV/JSON exports. +# Default: out/min_angle_graph_exports + +OUTPUT_DIR="${OUTPUT_DIR:-out}" +PREFIX="${PREFIX:-min_angle_graph_demo_$(date +%Y%m%d_%H%M%S)_}" +EXPORT_ROOT="${EXPORT_ROOT:-out/min_angle_graph_exports}" +CONFIG="explorations/min_angle_graph_export.yaml" +VIEWER="analysis/min_angle_graph_plotly_viewer.html" + +before_csv_count="$(find "${EXPORT_ROOT}" -type f -name "${PREFIX}*.csv" 2>/dev/null | wc -l | tr -d ' ')" + +cat </dev/null | wc -l | tr -d ' ')" +if [[ "${after_csv_count}" -le "${before_csv_count}" ]]; then + cat >&2 </ + +New CSV exports detected: +$(find "${EXPORT_ROOT}" -type f -name "${PREFIX}*.csv" 2>/dev/null | sort) + +Open the Plotly viewer in your browser: + ${VIEWER} + +Then select one or more exported CSV files from an export directory to step +through validation snapshots from first iteration to last. +EOF diff --git a/explorations/min_angle_graph_export.yaml b/explorations/min_angle_graph_export.yaml new file mode 100644 index 0000000000..dd8bad9818 --- /dev/null +++ b/explorations/min_angle_graph_export.yaml @@ -0,0 +1,39 @@ +# Smoke-test the optional LM-head minimum-angle graph export on a tiny +# validation cadence for both a character-tokenized and GPT-style dataset. +# +# Run with: +# python3 optimization_and_search/run_experiments.py \ +# --config explorations/min_angle_graph_export.yaml \ +# --config_format yaml \ +# --prefix min_angle_export_ +# +# The minipile run has a much larger vocabulary than shakespeare_char, so the +# default smoke settings use CUDA/float16 and a GPU-friendly export block size. +--- + +common_group: + n_layer: [3] + n_head: [3] + n_embd: [150] + block_size: [64] + batch_size: [4] + max_iters: [2000] + eval_interval: [20] + eval_iters: [20] + device: ["cuda:0"] + dtype: ["float16"] + compile: [true] + never_save_checkpoint: [true] + export_min_angle_graph_each_eval: [true] + export_min_angle_graph_device: ["cuda:0"] + +parameter_groups: + - dataset: ["shakespeare_char"] + export_min_angle_graph_block_size: [512] + export_min_angle_graph_label: ["${RUN_NAME}"] + export_min_angle_graph_dir: ["out/min_angle_graph_exports/${RUN_NAME}"] + + - dataset: ["minipile"] + export_min_angle_graph_block_size: [2048] + export_min_angle_graph_label: ["${RUN_NAME}"] + export_min_angle_graph_dir: ["out/min_angle_graph_exports/${RUN_NAME}"] diff --git a/train.py b/train.py index cccf0acba1..4b77dee589 100644 --- a/train.py +++ b/train.py @@ -12,6 +12,22 @@ from collections import deque from datetime import datetime, timedelta +# Some user worktrees may contain helper files named like Python stdlib modules +# (for example copy.py). Preload stdlib modules that Rich/dataclasses need +# before third-party imports so local helper files cannot shadow them during +# interpreter startup. +_repo_dir = os.path.dirname(os.path.abspath(__file__)) +_removed_sys_path_entries = [] +for _entry in ("", _repo_dir): + while _entry in sys.path: + sys.path.remove(_entry) + _removed_sys_path_entries.append(_entry) +import copy as _stdlib_copy +import dataclasses as _stdlib_dataclasses +for _entry in reversed(_removed_sys_path_entries): + sys.path.insert(0, _entry) +del _entry, _removed_sys_path_entries, _repo_dir, _stdlib_copy, _stdlib_dataclasses + from rich.console import Group from rich.console import Console from rich.text import Text @@ -27,6 +43,7 @@ from train_variations.distillation_loss_variants import build_distillation_loss from utils.gpu_monitoring import get_gpu_memory_info, get_process_gpu_memory_bytes +from utils.min_angle_graph_export import export_min_angle_graph as write_min_angle_graph_export from torch.cuda import reset_peak_memory_stats, max_memory_allocated, max_memory_reserved try: @@ -1797,6 +1814,26 @@ def save_checkpoint(self, filename): } torch.save(checkpoint, os.path.join(self.args.out_dir, filename)) + def export_min_angle_graph(self, losses): + """Export the current LM-head minimum-angle graph using the configured writer.""" + export_dir = getattr(self.args, "export_min_angle_graph_dir", None) + if not export_dir: + return + + weight = self.raw_model.apply_lm_head_norm(self.raw_model.lm_head.weight).detach() + label = getattr(self.args, "export_min_angle_graph_label", None) or "min_angle_graph" + val_loss = losses["val"].item() if hasattr(losses["val"], "item") else float(losses["val"]) + csv_path, _ = write_min_angle_graph_export( + weight=weight, + export_dir=export_dir, + label=label, + iter_num=self.iter_num, + val_loss=val_loss, + block_size=getattr(self.args, "export_min_angle_graph_block_size", 2048), + compute_device=getattr(self.args, "export_min_angle_graph_device", "auto"), + ) + print(f"Minimum-angle graph exported to {csv_path}") + def run_validation_step(self, running_mfu, current_epoch, current_dataset, num_steps_with_worse_loss, live=None): losses = self.estimate_loss() @@ -1888,6 +1925,13 @@ def run_validation_step(self, running_mfu, current_epoch, current_dataset, num_s print("Exiting with nan") file.write(str(self.iter_num)) + if self.args.export_min_angle_graph_each_eval: + if live: + live.stop() + self.export_min_angle_graph(losses) + if live: + live.start() + if (not self.args.never_save_checkpoint and self.args.save_major_ckpt_interval is not None): if self.iter_num % self.args.save_major_ckpt_interval == 0: diff --git a/train_args.py b/train_args.py index 835b24db06..424a1c7b39 100644 --- a/train_args.py +++ b/train_args.py @@ -80,6 +80,20 @@ def parse_args(): training_group.add_argument('--log_interval', default=10, type=int) training_group.add_argument('--eval_iters', default=200, type=int) training_group.add_argument('--eval_only', default=False, action=argparse.BooleanOptionalAction) + + # Validation-time LM-head minimum-angle graph export args + training_group.add_argument('--export_min_angle_graph_dir', default=None, type=str, + help='Optional directory for per-validation LM-head minimum-angle graph exports.') + training_group.add_argument('--export_min_angle_graph_each_eval', default=False, + action=argparse.BooleanOptionalAction, + help='Export the LM-head minimum-angle graph after every validation loss.') + training_group.add_argument('--export_min_angle_graph_block_size', default=2048, type=int, + help='Row/column block size used when exporting the minimum-angle graph.') + training_group.add_argument('--export_min_angle_graph_device', default='auto', type=str, + help="Compute device for minimum-angle graph export: 'auto', 'cpu', or a CUDA device such as 'cuda:0'.") + training_group.add_argument('--export_min_angle_graph_label', default=None, type=str, + help='Optional label inserted into minimum-angle graph export filenames.') + training_group.add_argument( '--mezo_epsilon', type=float, diff --git a/utils/min_angle_graph_export.py b/utils/min_angle_graph_export.py new file mode 100644 index 0000000000..d662ad5560 --- /dev/null +++ b/utils/min_angle_graph_export.py @@ -0,0 +1,152 @@ +# utils/min_angle_graph_export.py +import csv +import json +import os + +import torch +import torch.nn.functional as F + + +def resolve_min_angle_graph_device(weight, requested_device="auto"): + """Choose the compute device for blockwise minimum-angle graph export.""" + if requested_device == "auto": + if weight.is_cuda: + return weight.device + if torch.cuda.is_available(): + return torch.device("cuda:0") + return torch.device("cpu") + if requested_device.startswith("cuda") and not torch.cuda.is_available(): + print("CUDA requested for minimum-angle graph export but is unavailable; using CPU.") + return torch.device("cpu") + return torch.device(requested_device) + + +def compute_min_angle_graph(weight, block_size=2048, compute_device="auto"): + """Compute each row vector's closest non-self row by signed angular distance.""" + compute_device = resolve_min_angle_graph_device(weight, compute_device) + block_size = max(1, int(block_size)) + vocab_size = weight.shape[0] + + with torch.no_grad(): + weight = weight.detach() + norms = torch.linalg.vector_norm(weight, ord=2, dim=1).detach().cpu() + best_cosine = torch.full((vocab_size,), -float("inf"), device=compute_device) + best_other_token_id = torch.full((vocab_size,), -1, dtype=torch.long, device=compute_device) + + for row_start in range(0, vocab_size, block_size): + row_end = min(row_start + block_size, vocab_size) + row_block = weight[row_start:row_end].to(compute_device, non_blocking=True) + row_block = F.normalize(row_block, p=2, dim=1) + + for col_start in range(0, vocab_size, block_size): + col_end = min(col_start + block_size, vocab_size) + col_block = weight[col_start:col_end].to(compute_device, non_blocking=True) + col_block = F.normalize(col_block, p=2, dim=1) + cosine_block = row_block @ col_block.T + + if row_start < col_end and col_start < row_end: + diag_start = max(row_start, col_start) + diag_end = min(row_end, col_end) + diag_rows = torch.arange( + diag_start - row_start, + diag_end - row_start, + device=compute_device, + ) + diag_cols = torch.arange( + diag_start - col_start, + diag_end - col_start, + device=compute_device, + ) + cosine_block[diag_rows, diag_cols] = -float("inf") + + block_best_cosine, block_best_col = cosine_block.max(dim=1) + current_best = best_cosine[row_start:row_end] + current_other = best_other_token_id[row_start:row_end] + update_mask = block_best_cosine > current_best + current_best[update_mask] = block_best_cosine[update_mask] + current_other[update_mask] = block_best_col[update_mask] + col_start + + best_cosine = best_cosine.clamp(-1.0, 1.0).cpu() + best_other_token_id = best_other_token_id.cpu() + min_angles = torch.rad2deg(torch.acos(best_cosine)) + sorted_rank = torch.argsort(min_angles) + rank_by_token = torch.empty_like(sorted_rank) + rank_by_token[sorted_rank] = torch.arange(vocab_size, dtype=torch.long) + + return { + "vocab_size": vocab_size, + "block_size": block_size, + "compute_device": str(compute_device), + "norms": norms, + "best_cosine": best_cosine, + "best_other_token_id": best_other_token_id, + "min_angles": min_angles, + "rank_by_token": rank_by_token, + } + + +def safe_filename_label(label): + """Return a filesystem-friendly label while preserving readable separators.""" + return "".join(ch if ch.isalnum() or ch in ("-", "_") else "_" for ch in label) + + +def write_min_angle_graph_export(graph, export_dir, label, iter_num, val_loss): + """Write a minimum-angle graph CSV plus JSON sidecar and return both paths.""" + export_dir = os.path.expanduser(export_dir) + os.makedirs(export_dir, exist_ok=True) + + safe_label = safe_filename_label(label) + file_stem = f"{safe_label}_iter_{iter_num:08d}_val_{val_loss:.6f}" + csv_path = os.path.join(export_dir, f"{file_stem}.csv") + json_path = os.path.join(export_dir, f"{file_stem}.json") + + norms = graph["norms"] + best_cosine = graph["best_cosine"] + best_other_token_id = graph["best_other_token_id"] + min_angles = graph["min_angles"] + rank_by_token = graph["rank_by_token"] + vocab_size = graph["vocab_size"] + + with open(csv_path, "w", newline="") as csv_file: + writer = csv.writer(csv_file) + writer.writerow([ + "token_id", + "nearest_token_id", + "min_angle_deg", + "cosine", + "token_vector_length", + "nearest_token_vector_length", + "min_angle_rank", + ]) + for token_id in range(vocab_size): + other_id = int(best_other_token_id[token_id]) + writer.writerow([ + token_id, + other_id, + f"{float(min_angles[token_id]):.9f}", + f"{float(best_cosine[token_id]):.9f}", + f"{float(norms[token_id]):.9f}", + f"{float(norms[other_id]):.9f}" if other_id >= 0 else "", + int(rank_by_token[token_id]), + ]) + + metadata = { + "iter_num": iter_num, + "val_loss": val_loss, + "label": label, + "csv_path": csv_path, + "vocab_size": vocab_size, + "block_size": graph["block_size"], + "compute_device": graph["compute_device"], + "angle_definition": "signed 0-180 degrees, closest non-self token by maximum cosine", + } + with open(json_path, "w") as json_file: + json.dump(metadata, json_file, indent=2) + + return csv_path, json_path + + +def export_min_angle_graph(weight, export_dir, label, iter_num, val_loss, block_size=2048, compute_device="auto"): + """Compute and write the LM-head minimum-angle graph export.""" + graph = compute_min_angle_graph(weight, block_size=block_size, compute_device=compute_device) + return write_min_angle_graph_export(graph, export_dir, label, iter_num, val_loss)