feat(train): add gradient statistics

[rej55]

Summary

This PR adds per-step gradient statistics to the supervised training loop so that vanishing and exploding gradients can be observed during training. The statistics are logged to Weights & Biases through the existing logging path with no additional configuration.

What's added

A new helper compute_grad_stats() in utils/train_utils.py computes the following over the concatenation of all parameter gradients (the global gradient vector):

• $L_1$ norm (grad/l1_norm)
• $L_2$ norm (grad/l2_norm)
• $L_\infty$ norm (grad/linf_norm, max absolute value)
• mean (grad/mean)
• standard deviation (grad/std)

How it works

In train_epoch.py, the statistics are computed right after loss.backward() and before clip_grad_norm_(), so that exploding gradients are not masked by gradient clipping:

loss["loss"].backward()

# Gradient statistics (computed before clipping so that exploding
# gradients are not masked by clip_grad_norm_).
loss.update(compute_grad_stats(model.parameters()))

nn.utils.clip_grad_norm_(model.parameters(), 5)

The values are merged into the per-batch loss dict, so they automatically flow through the existing pipeline: get_epoch_mean_loss() averages them per epoch, and they are logged to wandb as train_loss/grad/* via the existing train_loss/{k} logging in train_predictor.py.

Notes

No new dependencies and no API/config changes; existing wandb dashboards pick up the new metrics automatically.
The GRPO training path (grpo_epoch.py) is intentionally left unchanged because it uses a keyed all-reduce across ranks that requires a consistent metric-key set; it can be addressed separately if needed.

Testing

Verified that train_epoch.py and train_utils.py parse without errors.

[go-sakayori]

Confirmed it works