Log documents_seen / num_documents and offer it as a W&B x-axis

[jlamypoirier]

Claude Sonnet 5, on behalf of @jlamypoirier.

Summary

Extracted from #553 as the first, independent piece of the RL diagnostics umbrella: the
documents_seen document-clock counter itself. Split out ahead of the weights_ready
document-count model version + reward/model_version staleness metrics (separate PRs) because
those both depend on this counter existing (the model version is documents_seen, and
staleness is computed as documents_seen - model_version) — this PR has no dependents-side
prerequisites and stands alone.

• ScheduleRunner exposes the whole-step DP-summed document count (num_documents_in_batch,
  the loss-normalization divisor) as _num_documents_in_batch, also returned as part of
  run_step's return tuple.
• The trainer accumulates it into _documents_seen every step (persisted in checkpoint
  metadata, restored on resume) and logs num_documents + documents_seen in the step metrics
  (console + W&B) on every logging step.
• Wandb declares training/documents_seen as an alternative step_metric for training/* at
  setup, so it can be used as a cross-run x-axis. The global wandb.log(step=...) step is
  unchanged, so both axes remain available.
• Document counting is unconditional, not an opt-in loss preprocessing flag. It used to
  activate only as a side effect of GRPO/GSPO's loss preprocessing config requesting it (for
  their own, unrelated loss-normalization divisor) — an obscure, indirect enablement path for a
  training-level metric. Document counting is a real, DP-wide all_reduce per step, so this is a
  deliberate tradeoff (sync cost accepted for every run) rather than an oversight. The now-dead
  return_document_count flag and its now-empty TokenPreprocessingConfig host class were both
  removed, along with a couple of downstream now-always-true conditionals this uncovered
  (TokenModelInput.share_batch_data's guard, the lazy W&B metric-definition dance, the
  metrics-dict presence check).

Tests

• tests/data/test_streaming.py: 23 passed (CPU).
• Broader regression pass after making document counting unconditional: test_preprocessing.py,
  test_streaming.py, test_random.py, test_attention.py, test_varlen.py, test_ssm.py,
  test_lm_losses.py, test_lm_head.py — 1450 passed / 35 skipped, 0 failures.
• No dedicated unit test added for the documents_seen accumulator/checkpoint round-trip; would
  require new checkpoint-test infrastructure for GRPO/GSPO models (currently marked
  not_implemented in tests/utils/model_configs.py) — out of scope here, left as a known gap.
• The GPU end-to-end integration test (tests/models/test_streaming.py::test_model_streaming,
  which exercises documents_seen) was not run here (no GPU) — worth running in CI before
  merge.