feat: version rollout weight updates

Track a monotonic policy version across optimizer steps, scheduler updates, and rollout results. Serialize synchronous generation with weight acknowledgements and invalidate reusable prefix KV entries so cached samples remain attributable to the behavior policy that generated them.
2026-09-02 19:01:41 +08:00
parent 1fad50d847
commit e58a728b80
13 changed files with 232 additions and 7 deletions
@@ -166,6 +166,13 @@ them with a `BaseRewardModel`. It refreshes cached rollouts every
 `rollout_interval` optimizer steps. `online_grpo` synchronizes `old_model` when
 a fresh rollout is produced.

+Every successful optimizer step advances a monotonic `policy_version` and
+acknowledges the shared-model weight update to the rollout scheduler. The
+scheduler invalidates reusable KV prefixes before accepting the new version.
+`RawRollout` and `RolloutResult` retain the version that actually generated
+their behavior log-probabilities, so cached rollout samples remain attributable
+even while later optimizer steps advance the live policy.
+
 Online strategies require `TrainConfig.reward_model_fn`. `train.py` exposes the
 rollout sampling parameters but does not yet offer a CLI argument for the reward
 model factory.