feat(algo): add cispo_anchor to select old vs rollout IS-ratio anchor#1859
Draft
tbalestri-lila wants to merge 1 commit into
Draft
feat(algo): add cispo_anchor to select old vs rollout IS-ratio anchor#1859tbalestri-lila wants to merge 1 commit into
tbalestri-lila wants to merge 1 commit into
Conversation
CISPO currently anchors its clamped importance-sampling ratio on the recomputed old log-probs (pi_theta / pi_old). Under fully-async training the sampler lags the trainer, but pi_old is recomputed fresh at experience-prep, so at a single gradient step the ratio is ~1 and the CISPO clamp never bites -- CISPO is effectively inert async. Add `cispo.cispo_anchor: "old" | "rollout"`. With "rollout" the ratio is anchored on the rollout/sampler log-probs (pi_theta / pi_rollout), so it genuinely differs from 1 under async staleness and the clamped objective engages. This is the clamped (cap-not-zero) counterpart of the existing `rollout_is` loss, which anchors on pi_rollout but HARD-ZEROS out-of-range tokens; CISPO caps them so every token keeps a bounded gradient. The rollout-anchored ratio is itself the off-policy correction, so stacking TIS (off_policy_correction.tis_ratio_type) is rejected to avoid double-counting. Default "old" preserves existing behavior exactly. Adds tests for the rollout anchor, the rollout_logprobs requirement, the TIS guard, and config validation.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
What
Add
cispo.cispo_anchor: "old" | "rollout"to select which behavior policy the CISPO clamped importance-sampling ratio is anchored on. Default"old"preserves current behavior exactly.Why
CISPO currently anchors its clamped IS ratio on the recomputed old log-probs (π_θ / π_old), where π_old is recomputed fresh by the training backend at experience-prep. That ratio therefore only deviates from 1 once π_θ has moved away from π_old — i.e. only when the batch takes more than one gradient update (
update_epochs > 1, ortrain_batch_size > policy_mini_batch_sizeso there are multiple minibatch steps).Fully-async training takes exactly one gradient step per batch (it hard-asserts
train_batch_size == policy_mini_batch_sizeandupdate_epochs == 1), so at loss time π_θ == π_old exactly → ratio ≡ 1 → the CISPO clamp never bites. Stock CISPO is effectively inert under fully-async.With
cispo_anchor="rollout"the ratio is anchored on the rollout/sampler log-probs (π_θ / π_rollout) instead. Unlike π_old, π_rollout is not a fresh backend recompute, so the ratio deviates from 1 even at a single gradient step, from two async-intrinsic sources: (1) staleness — the rollout was sampled by an older policy version than the current trainer weights; and (2) sampler/trainer engine mismatch — vLLM logprobs differ from the training-backend logprobs even for identical weights. So the clamped objective genuinely engages under async. This is the clamped (cap-not-zero) counterpart of the existingrollout_isloss, which anchors on π_rollout but HARD-ZEROS out-of-range tokens; CISPO caps them so every token keeps a bounded gradient.Notes
"old"is a no-op — existing runs are unchanged.off_policy_correction.tis_ratio_type) is rejected to avoid double-counting."rollout"requiresrollout_logprobs(asserted); the trainer's skip-fwd path is made anchor-aware so the unused old-logprobs forward is skipped when the rollout anchor makes it unnecessary.Tests
test_losses.py: rollout-anchor loss value, therollout_logprobsrequirement, the TIS double-count guard, and config validation.test_skip_fwd_logprobs.py: the cispo+rollout skip-path.