RL diagnostics: documents-seen clock, model-version broadcast, reward + per-token model_version, data-wait

[jlamypoirier]

Claude Opus 4.8, on behalf of @jlamypoirier.

Summary

Consolidates the remaining RL diagnostics on the Fast-LLM trainer side so the trainer is a single,
reliable source of truth for GSPO/GRPO runs — logged from the exact tensors each step trained on.

The config unification + GSPO/GRPO policy metrics landed separately in #555; this PR is rebased on
top of it and adds the document clock, the version broadcast, the reward / per-token-version data, and
the data-wait timer. Paired PipelineRL changes (raw-reward forwarding + document-count model version)
live in a separate PR against PipelineRL's fast-llm branch. Backward-compatible with the current producer.

Changes

• documents_seen / num_documents. Reimplemented on main, independent of docs_per_step, from
  the per-step DP-summed document count (the loss-normalization divisor). Persisted in checkpoint
  metadata. Offered as an alternative W&B x-axis via a lazily-declared define_metric (RL runs only;
  the global step is unchanged).
• Document-count model version. weights_ready now carries document_count (= documents_seen)
  alongside the existing step, so consumers can measure staleness in documents (DeepSpeed-aligned).
  Keeping step makes this backward-compatible.
• reward and per-token model_version. Two optional per-token fields threaded through the RL
  streaming schema and data pipeline (reusing TokenDataDocument/TokenDataBatch); both optional and
  guarded on presence. reward is the raw reward (distinct from the group-relative advantage). The
  shared policy-gradient loss logs mean/max/min of each when present (folded into the PolicyMetricsLevel
  metric set from Unify policy-loss metrics config and add GSPO metrics #555). Staleness is documents_seen − model_version, derivable from the logged stats.
• data_wait_time_ms. The schedule runner already timed data-loader waits but discarded the value;
  it is now logged as an input-starvation metric.

Tests

• The reward/version registration is checked against an independent plain-Python reference
  (test_policy_data_metrics), and the streaming schema round-trips the new fields (test_streaming_dataset).
• Rebased onto Unify policy-loss metrics config and add GSPO metrics #555 (now in main); CI validates the full suite.