Add analysis plots for training progress (metrics.csv) #75

Closed
opened 2026-08-24 10:28:29 +02:00 by lars · 0 comments
Owner

Nothing currently plots training metrics. giant/training/metrics.py's MetricsCollector
writes one row per epoch to <run_dir>/metrics.csv (opened in __init__, write_epoch()
called from giant/training/loop.py), with columns declared per-stage via MetricSpec in
giant/training/trainers.py. Full column set observed:

  • Per-stage train/val losses: <stage>/train/loss, <stage>/val/loss, plus constituents
    loss_gen, loss_nsec, loss_stop, loss_type, and (routed stages) loss_balance,
    loss_proc, loss_entropy.
  • Accuracies: nsec_acc, stop_acc, type_acc (train + val).
  • WGAN stage (WGANStageTrainer): d_loss, g_loss, wasserstein, gp_loss,
    grad_norm_d, grad_norm_g, and (onehot particle-type target) grad_norm_type_slice,
    grad_norm_cont_slice. No val loss for this stage (supports_val_loss = False).
  • Router diagnostics (routed stages, computed at validation): <stage>/router/entropy,
    util_min, util_max, util_std.
  • Point-in-time: <stage>/lr (and critic_lr for WGAN).
  • Run-level: val/loss, val/marginal_kl, grad_norm, gpu_mem_mb, samples_per_sec,
    is_best, epoch_time_s; epoch is always column 0.

Because the exact column set varies by run (flow/ddpm vs wgan, routed vs not), the plotting
code should read the CSV header dynamically rather than hardcoding column names.

Proposed plots, one run dir in:

  1. Per-stage loss curves — train vs val, log-y, with constituent losses as separate panels/lines.
  2. LR schedule(s) — <stage>/lr (+ critic_lr) vs epoch.
  3. Accuracy curves — nsec_acc/stop_acc/type_acc, train vs val, vs epoch.
  4. Gradient-norm diagnostics — grad_norm (+ WGAN grad_norm_d/grad_norm_g/slice norms)
    vs epoch, for spotting instability.
  5. Router health over training — entropy/util_min/util_max/util_std vs epoch. Directly relevant
    to diagnosing the MoE expert-collapse issue described in the CLAUDE.md roadmap
    (lambda_balance=0 run).
  6. WGAN critic/generator balance — d_loss/g_loss/wasserstein/gp_loss vs epoch.
  7. Throughput/resource overview — samples_per_sec, gpu_mem_mb, epoch_time_s vs epoch,
    with is_best epochs marked.
  8. Best-epoch summary annotation (best_val_loss + epoch reached).

Suggested entry point: a new CLI command reading <run_dir>/metrics.csv directly (it's tiny —
no need for the reduce/chunk/condor machinery giant analyze uses for rollout data) and
rendering with the same plotstyle/matplotlib conventions as giant/analysis/render.py, for
visual consistency with the rollout-vs-reference plots.

As a natural follow-on (not in scope here, flag for later): plotting multiple runs'
metrics.csv together to compare hyperparameter choices — same "N series, one plot" shape as
the multi-rollout refactor issue.

Nothing currently plots training metrics. `giant/training/metrics.py`'s `MetricsCollector` writes one row per epoch to `<run_dir>/metrics.csv` (opened in `__init__`, `write_epoch()` called from `giant/training/loop.py`), with columns declared per-stage via `MetricSpec` in `giant/training/trainers.py`. Full column set observed: - Per-stage train/val losses: `<stage>/train/loss`, `<stage>/val/loss`, plus constituents `loss_gen`, `loss_nsec`, `loss_stop`, `loss_type`, and (routed stages) `loss_balance`, `loss_proc`, `loss_entropy`. - Accuracies: `nsec_acc`, `stop_acc`, `type_acc` (train + val). - WGAN stage (`WGANStageTrainer`): `d_loss`, `g_loss`, `wasserstein`, `gp_loss`, `grad_norm_d`, `grad_norm_g`, and (onehot particle-type target) `grad_norm_type_slice`, `grad_norm_cont_slice`. No val loss for this stage (`supports_val_loss = False`). - Router diagnostics (routed stages, computed at validation): `<stage>/router/entropy`, `util_min`, `util_max`, `util_std`. - Point-in-time: `<stage>/lr` (and `critic_lr` for WGAN). - Run-level: `val/loss`, `val/marginal_kl`, `grad_norm`, `gpu_mem_mb`, `samples_per_sec`, `is_best`, `epoch_time_s`; `epoch` is always column 0. Because the exact column set varies by run (flow/ddpm vs wgan, routed vs not), the plotting code should read the CSV header dynamically rather than hardcoding column names. Proposed plots, one run dir in: 1. Per-stage loss curves — train vs val, log-y, with constituent losses as separate panels/lines. 2. LR schedule(s) — `<stage>/lr` (+ `critic_lr`) vs epoch. 3. Accuracy curves — `nsec_acc`/`stop_acc`/`type_acc`, train vs val, vs epoch. 4. Gradient-norm diagnostics — `grad_norm` (+ WGAN `grad_norm_d`/`grad_norm_g`/slice norms) vs epoch, for spotting instability. 5. Router health over training — entropy/util_min/util_max/util_std vs epoch. Directly relevant to diagnosing the MoE expert-collapse issue described in the CLAUDE.md roadmap (`lambda_balance=0` run). 6. WGAN critic/generator balance — `d_loss`/`g_loss`/`wasserstein`/`gp_loss` vs epoch. 7. Throughput/resource overview — `samples_per_sec`, `gpu_mem_mb`, `epoch_time_s` vs epoch, with `is_best` epochs marked. 8. Best-epoch summary annotation (best_val_loss + epoch reached). Suggested entry point: a new CLI command reading `<run_dir>/metrics.csv` directly (it's tiny — no need for the reduce/chunk/condor machinery `giant analyze` uses for rollout data) and rendering with the same `plotstyle`/matplotlib conventions as `giant/analysis/render.py`, for visual consistency with the rollout-vs-reference plots. As a natural follow-on (not in scope here, flag for later): plotting multiple runs' `metrics.csv` together to compare hyperparameter choices — same "N series, one plot" shape as the multi-rollout refactor issue.
lars added the feature label 2026-08-24 10:28:29 +02:00
lars closed this issue 2026-08-24 11:44:14 +02:00
Sign in to join this conversation.
1 Participants
Notifications
Due Date
No due date set.
Dependencies

No dependencies set.

Reference: lars/giant#75