Add robust NaN loss diagnostics

[nightlessbaron]

What changed

• Add per-tensor finite/nonfinite diagnostics for policy loss inputs and outputs.
• Make empty and all-nonfinite tensors safe, including rank-asymmetric distributed reductions.
• Aggregate diagnostic minima, maxima, and counts with explicit distributed semantics.
• Keep scalar diagnostics out of token/domain reducers.
• Reuse the ratio already computed by policy loss and remove the full-vocabulary logits scan.
• Rate-limit long-batch warnings.
• Add 15 focused regression tests.

Why

NaN/Inf incidents in long-context policy training were difficult to localize. The original diagnostic branch also had correctness hazards around scalar TIS metrics, empty shards, context-parallel aggregation, and duplicate ratio computation.

Validation

• 15 passed in test_loss_nan_diagnostics.py using the Miles environment with heavyweight unrelated fixtures excluded.
• Ruff: passed.
• Black: passed.
• py_compile: passed.
• git diff --check: passed.

Merge risk

Diagnostics are currently unconditional. They add several tensor-to-host synchronizations per policy microbatch and three small distributed collectives per train step. Benchmark on the target long-context workload, or add a sampling/config gate, before marking this ready.