diff --git a/demos/README.md b/demos/README.md index d97367cc30..57750098f6 100644 --- a/demos/README.md +++ b/demos/README.md @@ -22,6 +22,19 @@ 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. +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 +``` + ## 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..066f46c22c --- /dev/null +++ b/demos/lm_head_vocab_histogram_demo.sh @@ -0,0 +1,46 @@ +#!/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 \ + --export_lm_head_vocab_hist_html \ + --lm_head_vocab_hist_html_path "${OUT_DIR}/lm_head_vocab_histogram.html" + +cat < 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 _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] + 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"" + 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: + return + if len(self.lm_head_hist_snapshots) == 0: + print("Skipping lm_head vocab histogram HTML export: no snapshots captured.") + return + 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(self.lm_head_hist_snapshots) + html_text = f""" +lm_head vocab histogram + +

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 @@ -1485,6 +1631,8 @@ 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: @@ -1586,6 +1734,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) @@ -2304,6 +2453,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 7f7b2ace67..4de996fd0c 100644 --- a/train_args.py +++ b/train_args.py @@ -1451,6 +1451,10 @@ 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') + 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(