RL diagnostics: documents_seen, weights_ready document-count version, reward/model_version staleness metrics

[jlamypoirier]

Claude Sonnet 5, on behalf of @jlamypoirier.

Summary

Extracted from #553. Stacked on #559 (documents_seen / num_documents tracking), which this
PR depends on: the model version broadcast below is documents_seen, and staleness is
computed as documents_seen - model_version. Base will retarget to main once #559 merges.

• 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.
  TrainerCallback.run_begin / step_end gain documents_seen alongside step.
• 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 (GRPO and
  GSPO, via the shared _data_metric_definitions/_register_data_metrics). Staleness is
  documents_seen − model_version, derivable from the logged version stats.

Paired PipelineRL changes (raw-reward forwarding + document-count model version) are in a
separate PR against PipelineRL's fast-llm branch.

Tests

• tests/data/test_streaming.py (schema round-trip for the new fields) and
  tests/layers/test_lm_losses.py (reward/model_version metric registration, single-process and
  distributed) pass on CPU: 584 passed / 21 skipped (lm_losses), 87 passed (streaming + lm_head).

Caveats

• The GPU end-to-end integration test (tests/models/test_streaming.py::test_model_streaming,
  which exercises documents_seen + the weights_ready document_count path) was not run
  here (no GPU) — worth running in CI before merge.
• The per-token model_version consumer and its metric are in place, but the producer
  (tagging tokens with the active version inside vLLM's output path) is a separate,
  cluster-validated follow-up; until it lands, the model_version/staleness metrics simply don't
  populate.

[jlamypoirier]

Claude Sonnet 5, on behalf of @jlamypoirier: closing this — it bundled two commits
(weights_ready document-count broadcast + reward/per-token model_version) that turn out to have
no code dependency on each other. Split into two fully independent PRs instead:

• Broadcast documents_seen as the model version in weights_ready #561 — weights_ready document-count model version (stacked on Log documents_seen / num_documents and offer it as a W&B x-axis #559)
• Carry reward and per-token model_version, log reward/version metrics #562 — reward + per-token model_version + staleness metrics (standalone, based on main)