Forward raw reward, document-count model version, and per-token model version

[jlamypoirier]

Claude Opus 4.8, on behalf of @jlamypoirier.

Summary

Producer-side changes paired with the Fast-LLM trainer PR (ServiceNow/Fast-LLM#553), which consolidates
RL diagnostics on the Fast-LLM side. Both trainer-clock changes are backward-compatible.

Targets the fast-llm branch (which carries the Fast-LLM streaming producer).

Changes

• Use the Fast-LLM document count as the model version. Fast-LLM now emits document_count
  (cumulative documents seen) alongside step in the weights_ready event. TrainerState stamps
  document_count as propagated_weight_version so rollout staleness is measured in documents
  (aligning with DeepSpeed's clock), and keeps the raw step in completed_step for logging. Falls
  back to step when document_count is absent, so it works against older trainers too.
• Forward the raw reward. convert_to_fast_llm_format now includes the per-rollout raw reward
  (already on the entry, previously dropped) so the trainer can log reward statistics. It is a
  diagnostic and does not affect the loss — the group-relative advantage is unchanged.
• Capture per-token model_version. Weight swaps can happen mid-generation (pause mode keep
  retains in-flight tokens), so a rollout's tokens may span more than one version. The vLLM v1 output
  path is monkeypatched (API-server process) to tag each generated token with the version active when
  it is committed: a holder set at swap completion before generation resumes, LogprobsProcessor
  annotates each committed position's Logprob, and OpenAIServingChat._create_chat_logprobs encodes
  it into the existing token string (token_id:<id> → token_id:<id>:v<version>, no response-schema
  change). parse_token_id_and_version carries it on TokenLogprob.version →
  TrainingText.token_versions → convert_to_fast_llm_format, which left-pads it to the full sequence
  like old_log_probabilities. When the server reports no per-token version, it falls back to
  broadcasting the per-rollout version, so the wire format is identical either way. The Fast-LLM
  consumer already accepts the per-token array.

Caveats

• The reward + document-count-version changes are small and backward-compatible but not validated
  end-to-end here (requires a live streaming RL run); a smoke run is advisable before relying on the
  new trainer-side metrics.
• The per-token model_version producer exercises vLLM 0.18.1 v1 internals (LogprobsProcessor,
  OpenAIServingChat) that cannot be run without a GPU/vLLM. It is grounded in the real v0.18.1 source
  but is not runtime-tested — validate on a live streaming run: confirm the per-token version varies
  across an in-flight weight swap and that the trainer's model_version/staleness metrics populate. The
  install is defensive: if the vLLM seam has moved (different build), per-token capture is disabled and
  it falls back to the per-rollout version rather than crashing the server. The pure parse/pad logic and
  the Fast-LLM consumer round-trip are verified.

🤖 Generated with Claude Code