Skip to content

Carry reward and per-token model_version, log reward/version metrics#562

Open
jlamypoirier wants to merge 1 commit into
mainfrom
jlp_reward_model_version
Open

Carry reward and per-token model_version, log reward/version metrics#562
jlamypoirier wants to merge 1 commit into
mainfrom
jlp_reward_model_version

Conversation

@jlamypoirier

Copy link
Copy Markdown
Collaborator

Claude Sonnet 5, on behalf of @jlamypoirier.

Summary

Extracted from #553, as an independent piece with no dependency on the documents_seen /
weights_ready PRs (#559, #561) — verified by cherry-picking this commit directly onto main
with no involvement of the other two. The "staleness" framing (documents_seen - model_version)
is a usage note for whoever reads the logged metrics; the code here doesn't compute or need
documents_seen itself.

Was briefly combined with the weights_ready broadcast in #557 (now closed) before splitting
into fully independent pieces.

  • Two optional per-token fields added to the RL streaming schema and threaded through the data
    pipeline alongside advantages / old_log_probabilities (reusing
    TokenDataDocument/TokenDataBatch):
    • reward: the raw (un-normalized) reward, a per-rollout scalar broadcast per-token — distinct
      from the group-relative advantage.
    • model_version: the model version each token was generated under (documents-seen units),
      one per token, for staleness diagnostics.
  • Both are optional (absent when the producer does not send them), so the batch/target plumbing
    guards on presence.
  • The shared policy-gradient loss logs mean/max/min of each supplied field when metrics != none
    (GRPO and GSPO).

Paired PipelineRL change (raw-reward forwarding) is in a separate PR against PipelineRL's
fast-llm branch.

Tests

  • tests/data/test_streaming.py (schema round-trip) and tests/layers/test_lm_losses.py
    (metric registration, single-process and distributed): 584 passed / 21 skipped, 23 passed.
  • 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.

Add two optional per-token fields to the RL streaming schema and thread
them through the data pipeline alongside `advantages` /
`old_log_probabilities` (reusing `TokenDataDocument`/`TokenDataBatch`):
- `reward`: the raw (un-normalized) reward, a per-rollout scalar broadcast
  per-token — distinct from the group-relative `advantage`.
- `model_version`: the model version each token was generated under
  (documents-seen units), one per token, for staleness diagnostics.

Both are optional (absent when the producer does not send them), so the
batch/target plumbing guards on presence.

The shared policy-gradient loss logs mean/max/min of each supplied field
when `metrics != none` (GRPO and GSPO), registered only when the data is
present. Staleness is `documents_seen - model_version`, derivable from the
logged version stats and `documents_seen`.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant