From b5b65e651a705246a341eaf3333c22dbf9c6fd68 Mon Sep 17 00:00:00 2001 From: Kauna <16511995+klei22@users.noreply.github.com> Date: Mon, 11 May 2026 13:21:19 -0700 Subject: [PATCH 1/5] Add TensorBoard lm_head vocab magnitude histogram logging --- train.py | 30 ++++++++++++++++++++++++++++++ train_args.py | 2 ++ 2 files changed, 32 insertions(+) diff --git a/train.py b/train.py index d326db89c4..5084552471 100644 --- a/train.py +++ b/train.py @@ -1385,6 +1385,34 @@ def _safe_better_than_chance(self, vocab_size: float, loss_value: float) -> floa return 0.0 return vocab_size / math.exp(loss_value) + def _log_lm_head_vocab_magnitude_histogram(self, target_dataset: str) -> None: + if ( + not self.args.tensorboard_log + or self.writer is None + or not self.args.log_lm_head_vocab_hist + or self.iter_num % self.args.log_lm_head_vocab_hist_interval != 0 + ): + return + + lm_head = getattr(self.model, "lm_head", None) + if self.args.training_mode == "multicontext" and hasattr(self.model, "transformer"): + try: + dataset_idx = self.args.multicontext_datasets.index(target_dataset) + except ValueError: + dataset_idx = 0 + lm_head = self.model.transformer.get(f"lm_head_{dataset_idx}", lm_head) + + if lm_head is None or not hasattr(lm_head, "weight"): + return + + with torch.no_grad(): + magnitudes = lm_head.weight.detach().norm(dim=1).float().cpu() + self.writer.add_histogram( + f"{target_dataset}/lm_head_vocab_vector_magnitude", + magnitudes, + self.iter_num, + ) + def log_metrics(self, losses, running_mfu, epoch, tokens_trained, target_dataset, val_better_than_chance): compute_rankme = self.args.log_rankme or self.args.log_areq @@ -1485,6 +1513,7 @@ def log_metrics(self, losses, running_mfu, epoch, tokens_trained, target_dataset self._log_zeus_tensorboard(target_dataset, tokens_trained) self._log_bit_metrics(target_dataset, tokens_trained) + self._log_lm_head_vocab_magnitude_histogram(target_dataset) if self.args.csv_log: @@ -1586,6 +1615,7 @@ def log_metrics_non_validation(self, loss_training, running_mfu, epoch, tokens_t self._log_zeus_tensorboard(target_dataset, tokens_trained) self._log_bit_metrics(target_dataset, tokens_trained) + self._log_lm_head_vocab_magnitude_histogram(target_dataset) def write_to_csv(self, *args, prefix=""): args = list(args) diff --git a/train_args.py b/train_args.py index 7f7b2ace67..91ff774913 100644 --- a/train_args.py +++ b/train_args.py @@ -1451,6 +1451,8 @@ def parse_args(): logging_group.add_argument('--log_all_metrics', default=False, action=argparse.BooleanOptionalAction, help='Enable logging of all metrics including gns') logging_group.add_argument('--log_rankme', default=True, action=argparse.BooleanOptionalAction, help='Log RankMe representation metric during validation') logging_group.add_argument('--log_areq', default=True, action=argparse.BooleanOptionalAction, help='Log aReQ representation metric during validation') + logging_group.add_argument('--log_lm_head_vocab_hist', default=False, action=argparse.BooleanOptionalAction, help='Log TensorBoard histogram of per-token lm_head vector magnitudes over training') + logging_group.add_argument('--log_lm_head_vocab_hist_interval', default=100, type=int, help='Training-step interval for logging lm_head vocab magnitude histogram') # Turn activation/weight statistics off to save CPU RAM and wall time. training_group.add_argument( From 53c6d025adf18bd963888e6a85d46278cc5e5744 Mon Sep 17 00:00:00 2001 From: Kauna <16511995+klei22@users.noreply.github.com> Date: Mon, 11 May 2026 17:14:23 -0700 Subject: [PATCH 2/5] Add demo script for lm_head vocab histogram logging --- demos/README.md | 10 +++++++ demos/lm_head_vocab_histogram_demo.sh | 41 +++++++++++++++++++++++++++ 2 files changed, 51 insertions(+) create mode 100644 demos/lm_head_vocab_histogram_demo.sh diff --git a/demos/README.md b/demos/README.md index d97367cc30..fd20b42cdb 100644 --- a/demos/README.md +++ b/demos/README.md @@ -22,6 +22,16 @@ python3 demos/check_ckpt_for_gelu_shift.py \ `adam_vs_adamw.sh` trains two tiny Shakespeare models, one with Adam and one with AdamW, then compares their statistics using `view_model_stats.py`. +## lm_head Vocab Magnitude Histogram (TensorBoard Animation) + +`lm_head_vocab_histogram_demo.sh` runs a short `shakespeare_char` training job +with `--log_lm_head_vocab_hist` enabled so you can watch the distribution of +per-token `lm_head` vector magnitudes evolve over time in TensorBoard. + +```bash +bash demos/lm_head_vocab_histogram_demo.sh +``` + ## Grouped Asymmetric Vector PTQ Comparison `fake_ptq_asymmetric_grouped_vector_eval_demo_shakespeare_char.sh` runs a diff --git a/demos/lm_head_vocab_histogram_demo.sh b/demos/lm_head_vocab_histogram_demo.sh new file mode 100644 index 0000000000..27007eda51 --- /dev/null +++ b/demos/lm_head_vocab_histogram_demo.sh @@ -0,0 +1,41 @@ +#!/bin/bash +# demos/lm_head_vocab_histogram_demo.sh +# Quick demo: train a tiny Shakespeare-char model while logging a TensorBoard +# histogram of lm_head vocab-vector magnitudes over time. + +set -euo pipefail + +OUT_DIR="out/lm_head_vocab_hist_demo" +RUN_NAME="lm_head_vocab_hist_demo" + +echo "=== Training tiny model with lm_head vocab histogram logging enabled ===" +python3 train.py \ + --dataset shakespeare_char \ + --out_dir "${OUT_DIR}" \ + --n_layer 2 \ + --n_head 2 \ + --n_embd 128 \ + --block_size 128 \ + --batch_size 32 \ + --max_iters 200 \ + --eval_interval 50 \ + --eval_iters 20 \ + --log_interval 10 \ + --learning_rate 1e-3 \ + --tensorboard_log \ + --tensorboard_run_name "${RUN_NAME}" \ + --log_lm_head_vocab_hist \ + --log_lm_head_vocab_hist_interval 10 + +cat < Date: Mon, 11 May 2026 22:55:10 -0700 Subject: [PATCH 3/5] Export sortable lm_head vocab magnitude HTML report --- demos/README.md | 3 + demos/lm_head_vocab_histogram_demo.sh | 7 ++- train.py | 79 +++++++++++++++++++++++++++ train_args.py | 2 + 4 files changed, 90 insertions(+), 1 deletion(-) diff --git a/demos/README.md b/demos/README.md index fd20b42cdb..57750098f6 100644 --- a/demos/README.md +++ b/demos/README.md @@ -27,6 +27,9 @@ with AdamW, then compares their statistics using `view_model_stats.py`. `lm_head_vocab_histogram_demo.sh` runs a short `shakespeare_char` training job with `--log_lm_head_vocab_hist` enabled so you can watch the distribution of per-token `lm_head` vector magnitudes evolve over time in TensorBoard. +It also emits an interactive HTML file with sortable bars (by vocab id asc/desc +or magnitude asc/desc) and hover labels that include both vocab id and rendered +token text. ```bash bash demos/lm_head_vocab_histogram_demo.sh diff --git a/demos/lm_head_vocab_histogram_demo.sh b/demos/lm_head_vocab_histogram_demo.sh index 27007eda51..066f46c22c 100644 --- a/demos/lm_head_vocab_histogram_demo.sh +++ b/demos/lm_head_vocab_histogram_demo.sh @@ -25,7 +25,9 @@ python3 train.py \ --tensorboard_log \ --tensorboard_run_name "${RUN_NAME}" \ --log_lm_head_vocab_hist \ - --log_lm_head_vocab_hist_interval 10 + --log_lm_head_vocab_hist_interval 10 \ + --export_lm_head_vocab_hist_html \ + --lm_head_vocab_hist_html_path "${OUT_DIR}/lm_head_vocab_histogram.html" cat < None: self.iter_num, ) + def _get_vocab_label(self, token_id: int) -> str: + token_text = None + if hasattr(self, "itos") and self.itos is not None and token_id < len(self.itos): + token_text = self.itos[token_id] + elif hasattr(self, "decode") and callable(self.decode): + try: + token_text = self.decode([int(token_id)]) + except Exception: + token_text = None + if token_text is None: + token_text = f"" + escaped = html.escape(str(token_text)).replace("\n", "\\n") + return f"{token_id}: {escaped}" + + def _export_lm_head_vocab_histogram_html(self) -> None: + if not self.args.export_lm_head_vocab_hist_html: + return + lm_head = getattr(self.model, "lm_head", None) + if lm_head is None or not hasattr(lm_head, "weight"): + print("Skipping lm_head vocab histogram HTML export: lm_head not found.") + return + with torch.no_grad(): + magnitudes = lm_head.weight.detach().norm(dim=1).float().cpu().tolist() + vocab_data = [ + {"id": i, "magnitude": float(m), "label": self._get_vocab_label(i)} + for i, m in enumerate(magnitudes) + ] + out_path = self.args.lm_head_vocab_hist_html_path or os.path.join( + self.args.out_dir, "lm_head_vocab_histogram.html" + ) + os.makedirs(os.path.dirname(out_path), exist_ok=True) + payload = json.dumps(vocab_data) + html_text = f""" +lm_head vocab histogram + +

Final lm_head vocabulary vector magnitudes

+ + +
Bars include both token id and rendered token text in hover tooltip.
+
+ +""" + with open(out_path, "w", encoding="utf-8") as f: + f.write(html_text) + print(f"Exported lm_head vocab histogram HTML: {out_path}") + def log_metrics(self, losses, running_mfu, epoch, tokens_trained, target_dataset, val_better_than_chance): compute_rankme = self.args.log_rankme or self.args.log_areq @@ -2334,6 +2412,7 @@ def train(self): # End of training actions if self.iter_num > self.args.max_iters: + self._export_lm_head_vocab_histogram_html() print(self.best_val_loss, self.best_iter, self.best_tokens) if self.args.only_save_checkpoint_at_end: if not self.args.never_save_checkpoint: diff --git a/train_args.py b/train_args.py index 91ff774913..4de996fd0c 100644 --- a/train_args.py +++ b/train_args.py @@ -1453,6 +1453,8 @@ def parse_args(): logging_group.add_argument('--log_areq', default=True, action=argparse.BooleanOptionalAction, help='Log aReQ representation metric during validation') logging_group.add_argument('--log_lm_head_vocab_hist', default=False, action=argparse.BooleanOptionalAction, help='Log TensorBoard histogram of per-token lm_head vector magnitudes over training') logging_group.add_argument('--log_lm_head_vocab_hist_interval', default=100, type=int, help='Training-step interval for logging lm_head vocab magnitude histogram') + logging_group.add_argument('--export_lm_head_vocab_hist_html', default=False, action=argparse.BooleanOptionalAction, help='Export an interactive HTML report of final lm_head vocab-vector magnitudes') + logging_group.add_argument('--lm_head_vocab_hist_html_path', type=str, default=None, help='Optional output path for lm_head vocab histogram HTML (default: /lm_head_vocab_histogram.html)') # Turn activation/weight statistics off to save CPU RAM and wall time. training_group.add_argument( From 576ccdaa7ea6ba09c148e893e9980c2da8ad5635 Mon Sep 17 00:00:00 2001 From: Kauna <16511995+klei22@users.noreply.github.com> Date: Mon, 11 May 2026 23:45:53 -0700 Subject: [PATCH 4/5] Fix lm_head HTML sorting control and token-symbol x-axis --- train.py | 41 +++++++++++++++++++++++++---------------- 1 file changed, 25 insertions(+), 16 deletions(-) diff --git a/train.py b/train.py index b0b522124c..af6a0e192e 100644 --- a/train.py +++ b/train.py @@ -2,7 +2,6 @@ from contextlib import nullcontext import csv import json -import html import math import os import random @@ -1414,7 +1413,7 @@ def _log_lm_head_vocab_magnitude_histogram(self, target_dataset: str) -> None: self.iter_num, ) - def _get_vocab_label(self, token_id: int) -> str: + def _get_vocab_label_parts(self, token_id: int) -> tuple[str, str]: token_text = None if hasattr(self, "itos") and self.itos is not None and token_id < len(self.itos): token_text = self.itos[token_id] @@ -1425,8 +1424,9 @@ def _get_vocab_label(self, token_id: int) -> str: token_text = None if token_text is None: token_text = f"" - escaped = html.escape(str(token_text)).replace("\n", "\\n") - return f"{token_id}: {escaped}" + raw = str(token_text) + human = raw.encode("unicode_escape").decode("utf-8") + return raw, human def _export_lm_head_vocab_histogram_html(self) -> None: if not self.args.export_lm_head_vocab_hist_html: @@ -1437,10 +1437,17 @@ def _export_lm_head_vocab_histogram_html(self) -> None: return with torch.no_grad(): magnitudes = lm_head.weight.detach().norm(dim=1).float().cpu().tolist() - vocab_data = [ - {"id": i, "magnitude": float(m), "label": self._get_vocab_label(i)} - for i, m in enumerate(magnitudes) - ] + vocab_data = [] + for i, m in enumerate(magnitudes): + token_raw, token_display = self._get_vocab_label_parts(i) + vocab_data.append( + { + "id": i, + "magnitude": float(m), + "token_raw": token_raw, + "token_display": token_display, + } + ) out_path = self.args.lm_head_vocab_hist_html_path or os.path.join( self.args.out_dir, "lm_head_vocab_histogram.html" ) @@ -1457,7 +1464,7 @@ def _export_lm_head_vocab_histogram_html(self) -> None: -
Bars include both token id and rendered token text in hover tooltip.
+
X-axis is token symbol; hover shows token id + symbol.
""" with open(out_path, "w", encoding="utf-8") as f: From fd696c88dd884b87e71bb5bf84e60a9026b5ce99 Mon Sep 17 00:00:00 2001 From: Kauna <16511995+klei22@users.noreply.github.com> Date: Tue, 12 May 2026 00:29:51 -0700 Subject: [PATCH 5/5] Add time-snapshotted lm_head HTML playback across evals --- train.py | 76 ++++++++++++++++++++++++++++++++++++++++---------------- 1 file changed, 54 insertions(+), 22 deletions(-) diff --git a/train.py b/train.py index af6a0e192e..a4a9eab26d 100644 --- a/train.py +++ b/train.py @@ -108,6 +108,7 @@ def __init__(self, args, model_group, training_group, logging_group): self.peak_torch_allocated = 0.0 self.peak_torch_reserved = 0.0 self.peak_process_gpu_usage = 0.0 + self.lm_head_hist_snapshots = [] self.zeus_monitor = None self.zeus_enabled = False self.zeus_total_energy_j = 0.0 @@ -1431,32 +1432,18 @@ def _get_vocab_label_parts(self, token_id: int) -> tuple[str, str]: def _export_lm_head_vocab_histogram_html(self) -> None: if not self.args.export_lm_head_vocab_hist_html: return - lm_head = getattr(self.model, "lm_head", None) - if lm_head is None or not hasattr(lm_head, "weight"): - print("Skipping lm_head vocab histogram HTML export: lm_head not found.") + if len(self.lm_head_hist_snapshots) == 0: + print("Skipping lm_head vocab histogram HTML export: no snapshots captured.") return - with torch.no_grad(): - magnitudes = lm_head.weight.detach().norm(dim=1).float().cpu().tolist() - vocab_data = [] - for i, m in enumerate(magnitudes): - token_raw, token_display = self._get_vocab_label_parts(i) - vocab_data.append( - { - "id": i, - "magnitude": float(m), - "token_raw": token_raw, - "token_display": token_display, - } - ) out_path = self.args.lm_head_vocab_hist_html_path or os.path.join( self.args.out_dir, "lm_head_vocab_histogram.html" ) os.makedirs(os.path.dirname(out_path), exist_ok=True) - payload = json.dumps(vocab_data) + payload = json.dumps(self.lm_head_hist_snapshots) html_text = f""" lm_head vocab histogram -

Final lm_head vocabulary vector magnitudes

+

lm_head vocabulary vector magnitudes over time

+ + +
X-axis is token symbol; hover shows token id + symbol.
""" with open(out_path, "w", encoding="utf-8") as f: f.write(html_text) print(f"Exported lm_head vocab histogram HTML: {out_path}") + def _capture_lm_head_hist_snapshot(self, target_dataset: str, tokens_trained: float) -> None: + if not self.args.export_lm_head_vocab_hist_html: + return + lm_head = getattr(self.model, "lm_head", None) + if self.args.training_mode == "multicontext" and hasattr(self.model, "transformer"): + try: + dataset_idx = self.args.multicontext_datasets.index(target_dataset) + except ValueError: + dataset_idx = 0 + lm_head = self.model.transformer.get(f"lm_head_{dataset_idx}", lm_head) + if lm_head is None or not hasattr(lm_head, "weight"): + return + with torch.no_grad(): + magnitudes = lm_head.weight.detach().norm(dim=1).float().cpu().tolist() + vocab_data = [] + for i, m in enumerate(magnitudes): + token_raw, token_display = self._get_vocab_label_parts(i) + vocab_data.append({"id": i, "magnitude": float(m), "token_raw": token_raw, "token_display": token_display}) + self.lm_head_hist_snapshots.append({ + "iter_num": int(self.iter_num), + "tokens_trained": float(tokens_trained), + "dataset": target_dataset, + "vocab": vocab_data, + }) + def log_metrics(self, losses, running_mfu, epoch, tokens_trained, target_dataset, val_better_than_chance): compute_rankme = self.args.log_rankme or self.args.log_areq @@ -1601,6 +1632,7 @@ def log_metrics(self, losses, running_mfu, epoch, tokens_trained, target_dataset self._log_zeus_tensorboard(target_dataset, tokens_trained) self._log_bit_metrics(target_dataset, tokens_trained) self._log_lm_head_vocab_magnitude_histogram(target_dataset) + self._capture_lm_head_hist_snapshot(target_dataset, tokens_trained) if self.args.csv_log: