""" plot_training.py — Training Visualization Dashboard Reads train_log.jsonl and renders a clean, dark-mode training dashboard. Usage: # Static plot of completed/current run python plot_training.py --run_dir runs/run_001 # Live mode: refresh every 5 seconds while training runs python plot_training.py --run_dir runs/run_001 --live # Compare multiple runs python plot_training.py --run_dir runs/run_001 runs/run_002 Dashboard panels: 1. Training Loss (raw + EMA smoothed) 2. Validation Loss (if available) 3. Learning Rate schedule 4. Tokens / second (throughput) 5. VRAM usage (if logged) 6. Gradient norm (if logged) """ import os import sys import json import time import argparse from pathlib import Path import matplotlib import matplotlib.pyplot as plt import matplotlib.gridspec as gridspec import matplotlib.ticker as ticker import numpy as np # ------------------------------------------------------------------ # # STYLE # ------------------------------------------------------------------ # DARK_BG = "#0d1117" PANEL_BG = "#161b22" GRID_COLOR = "#21262d" TEXT_COLOR = "#c9d1d9" MUTED_COLOR = "#6e7681" ACCENT_BLUE = "#58a6ff" ACCENT_GREEN = "#3fb950" ACCENT_ORANGE= "#d29922" ACCENT_RED = "#f85149" ACCENT_PURPLE= "#bc8cff" ACCENT_TEAL = "#39d353" matplotlib.rcParams.update({ "figure.facecolor": DARK_BG, "axes.facecolor": PANEL_BG, "axes.edgecolor": GRID_COLOR, "axes.labelcolor": TEXT_COLOR, "axes.titlecolor": TEXT_COLOR, "xtick.color": MUTED_COLOR, "ytick.color": MUTED_COLOR, "grid.color": GRID_COLOR, "grid.linestyle": "--", "grid.linewidth": 0.5, "grid.alpha": 0.7, "legend.facecolor": PANEL_BG, "legend.edgecolor": GRID_COLOR, "legend.labelcolor": TEXT_COLOR, "text.color": TEXT_COLOR, "font.family": "DejaVu Sans", "font.size": 10, "axes.titlesize": 11, "axes.labelsize": 10, }) # ------------------------------------------------------------------ # # DATA LOADING # ------------------------------------------------------------------ # def load_log(log_path: str) -> dict: """ Loads train_log.jsonl and returns separate arrays for each metric. Returns dict of metric_name -> list of values, aligned by step. """ train_steps = [] train_loss = [] val_steps = [] val_loss = [] lr_steps = [] lr_vals = [] tok_steps = [] tok_vals = [] vram_steps = [] vram_vals = [] grad_steps = [] grad_vals = [] if not os.path.exists(log_path): return None with open(log_path, "r") as f: for line in f: line = line.strip() if not line: continue try: entry = json.loads(line) except json.JSONDecodeError: continue step = entry.get("step") if step is None: continue if "loss" in entry: train_steps.append(step) train_loss.append(entry["loss"]) if "val_loss" in entry: val_steps.append(step) val_loss.append(entry["val_loss"]) if "lr" in entry: lr_steps.append(step) lr_vals.append(entry["lr"]) if "tok_per_sec" in entry: tok_steps.append(step) tok_vals.append(entry["tok_per_sec"]) if "vram_gb" in entry: vram_steps.append(step) vram_vals.append(entry["vram_gb"]) if "grad_norm" in entry and entry["grad_norm"] is not None: grad_steps.append(step) grad_vals.append(entry["grad_norm"]) return { "train": (train_steps, train_loss), "val": (val_steps, val_loss), "lr": (lr_steps, lr_vals), "tok": (tok_steps, tok_vals), "vram": (vram_steps, vram_vals), "grad": (grad_steps, grad_vals), } def ema_smooth(values: list, alpha: float = 0.9) -> list: """Exponential moving average smoothing.""" if not values: return values smoothed = [values[0]] for v in values[1:]: smoothed.append(alpha * smoothed[-1] + (1 - alpha) * v) return smoothed # ------------------------------------------------------------------ # # PLOTTING # ------------------------------------------------------------------ # def make_dashboard(data_dict: dict, run_names: list, save_path: str = None): """ Renders a multi-panel training dashboard. Args: data_dict : dict of run_name -> metrics dict run_names : list of run display names save_path : if set, saves figure to this path instead of showing """ fig = plt.figure(figsize=(16, 10), facecolor=DARK_BG) fig.suptitle( "SLLM Training Dashboard", fontsize=16, fontweight="bold", color=TEXT_COLOR, y=0.98, ) # 3x2 grid of panels gs = gridspec.GridSpec(3, 2, figure=fig, hspace=0.45, wspace=0.3, left=0.06, right=0.97, top=0.93, bottom=0.06) ax_loss = fig.add_subplot(gs[0, 0]) ax_val = fig.add_subplot(gs[0, 1]) ax_lr = fig.add_subplot(gs[1, 0]) ax_tok = fig.add_subplot(gs[1, 1]) ax_vram = fig.add_subplot(gs[2, 0]) ax_grad = fig.add_subplot(gs[2, 1]) colors = [ACCENT_BLUE, ACCENT_GREEN, ACCENT_ORANGE, ACCENT_PURPLE] has_val = False has_vram = False has_grad = False for idx, (run_name, data) in enumerate(data_dict.items()): if data is None: continue color = colors[idx % len(colors)] # --- Train loss ------------------------------------------ # steps, loss = data["train"] if steps: smoothed = ema_smooth(loss, alpha=0.92) ax_loss.plot(steps, loss, color=color, alpha=0.25, linewidth=0.8) ax_loss.plot(steps, smoothed, color=color, alpha=1.0, linewidth=1.8, label=run_name) # Annotate final loss ax_loss.annotate( f"{smoothed[-1]:.4f}", xy=(steps[-1], smoothed[-1]), xytext=(5, 0), textcoords="offset points", color=color, fontsize=8, va="center", ) # --- Val loss -------------------------------------------- # vsteps, vloss = data["val"] if vsteps: has_val = True ax_val.plot(vsteps, vloss, color=color, linewidth=2, marker="o", markersize=4, label=run_name) ax_val.annotate( f"{vloss[-1]:.4f}", xy=(vsteps[-1], vloss[-1]), xytext=(5, 0), textcoords="offset points", color=color, fontsize=8, va="center", ) # --- LR -------------------------------------------------- # lsteps, lvals = data["lr"] if lsteps: ax_lr.plot(lsteps, lvals, color=color, linewidth=1.5, label=run_name) # --- Throughput ------------------------------------------ # tsteps, tvals = data["tok"] if tsteps: avg_tok = np.mean(tvals) ax_tok.plot(tsteps, tvals, color=color, alpha=0.6, linewidth=1.0) ax_tok.axhline(avg_tok, color=color, linewidth=1.5, linestyle="--", label=f"{run_name} (avg {avg_tok:.0f})") # --- VRAM ------------------------------------------------- # vsteps2, vvals = data["vram"] if vsteps2: has_vram = True ax_vram.plot(vsteps2, vvals, color=color, linewidth=1.5, label=run_name) # --- Grad norm ------------------------------------------- # gsteps, gvals = data["grad"] if gsteps: has_grad = True smoothed_g = ema_smooth(gvals, alpha=0.85) ax_grad.plot(gsteps, gvals, color=color, alpha=0.2, linewidth=0.8) ax_grad.plot(gsteps, smoothed_g, color=color, linewidth=1.5, label=run_name) # --- Style panels -------------------------------------------- # def _style(ax, title, xlabel, ylabel, legend=True): ax.set_title(title, fontweight="bold", pad=8) ax.set_xlabel(xlabel) ax.set_ylabel(ylabel) ax.grid(True) ax.tick_params(which="both", length=3) if legend and ax.get_legend_handles_labels()[0]: ax.legend(fontsize=8, loc="upper right") _style(ax_loss, "Training Loss (EMA smoothed)", "Step", "Loss") _style(ax_lr, "Learning Rate Schedule", "Step", "LR") _style(ax_tok, "Throughput", "Step", "Tokens / sec") if has_val: _style(ax_val, "Validation Loss", "Step", "Val Loss") else: ax_val.text(0.5, 0.5, "No validation data yet", ha="center", va="center", transform=ax_val.transAxes, color=MUTED_COLOR, fontsize=11) ax_val.set_title("Validation Loss", fontweight="bold", pad=8) if has_vram: _style(ax_vram, "VRAM Usage", "Step", "GB") ax_vram.axhline(4.0, color=ACCENT_RED, linewidth=1, linestyle=":", alpha=0.6, label="4 GB limit") ax_vram.legend(fontsize=8) else: ax_vram.text(0.5, 0.5, "No VRAM data\n(requires CUDA)", ha="center", va="center", transform=ax_vram.transAxes, color=MUTED_COLOR, fontsize=11) ax_vram.set_title("VRAM Usage", fontweight="bold", pad=8) if has_grad: _style(ax_grad, "Gradient Norm (EMA smoothed)", "Step", "Norm") else: ax_grad.text(0.5, 0.5, "No gradient norm data", ha="center", va="center", transform=ax_grad.transAxes, color=MUTED_COLOR, fontsize=11) ax_grad.set_title("Gradient Norm", fontweight="bold", pad=8) # LR scientific notation ax_lr.yaxis.set_major_formatter(ticker.ScalarFormatter(useMathText=True)) ax_lr.ticklabel_format(style="sci", axis="y", scilimits=(0, 0)) if save_path: plt.savefig(save_path, dpi=150, bbox_inches="tight", facecolor=DARK_BG) print(f"[PLOT] Saved to {save_path}") else: plt.show() # ------------------------------------------------------------------ # # CLI # ------------------------------------------------------------------ # def parse_args(): p = argparse.ArgumentParser(description="SLLM Training Dashboard") p.add_argument("--run_dir", nargs="+", default=["runs/run_001"], help="One or more run directories to plot") p.add_argument("--live", action="store_true", help="Refresh plot every --interval seconds (live mode)") p.add_argument("--interval", type=int, default=10, help="Refresh interval in seconds for --live mode") p.add_argument("--save", type=str, default=None, help="Save plot to this path instead of showing interactively") return p.parse_args() def main(): args = parse_args() run_dirs = args.run_dir run_names = [Path(d).name for d in run_dirs] def _reload_and_plot(): data_dict = {} for name, run_dir in zip(run_names, run_dirs): log_path = os.path.join(run_dir, "train_log.jsonl") data = load_log(log_path) if data is None: print(f"[WARN] No log found at: {log_path}") data_dict[name] = data # Check if any data was loaded total_steps = sum( len(d["train"][0]) for d in data_dict.values() if d ) if total_steps == 0: print("[PLOT] No data logged yet. Waiting...") return steps_info = {n: len(d["train"][0]) for n, d in data_dict.items() if d} print(f"[PLOT] Plotting {steps_info} train steps") plt.close("all") make_dashboard(data_dict, run_names, save_path=args.save) if args.live: print(f"[LIVE] Refreshing every {args.interval}s (Ctrl+C to stop)") matplotlib.use("TkAgg") if sys.platform == "win32" else None try: while True: _reload_and_plot() plt.pause(args.interval) except KeyboardInterrupt: print("\n[LIVE] Stopped.") else: _reload_and_plot() if __name__ == "__main__": main()