Add rich diagnostics to close the trainee feedback loop

[blueberryvertigo]

Summary

• Training loss curve: Records smoothed loss every 50 steps, reports convergence analysis (early vs late loss, improvement rate, whether model was still improving at cutoff)
• Per-position loss: Mean loss bucketed by sequence position (0-64, 64-256, 256-512, 512-1024, 1024-2048) to reveal long-range dependency issues
• Attention pattern analysis: Captures Q/K before flash attention via a lightweight _diag_store flag, reconstructs attention matrices for a single example, reports per-head entropy, mean attention distance, and peak weight per layer
• Model text samples: 5 unconditional generations of 200 tokens each showing what the model actually learned

All diagnostics written to diagnostics.log after each run. The existing run.log grep workflow is unchanged.

program.md updated with instructions to read diagnostics every run and an interpretation guide mapping diagnostic patterns (plateaued early, high position loss at late positions, dead attention heads, repetitive samples, etc.) to concrete architectural hypotheses.

Motivation

The agent was deciding keep/revert based on a single scalar (val_bpb). It never saw what the model generates, where in the sequence it struggles, whether attention heads are dead or redundant, or whether training was still improving at cutoff. This closes that feedback loop.

Test plan

• Run uv run train.py and verify diagnostics.log is produced with all 5 sections
• Verify training behavior is unchanged (_diag_store is a no-op during training)
• Verify existing grep workflow on run.log still works

Made with Cursor