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.
This commit is contained in:
0z5a
2026-09-02 19:01:41 +08:00
parent 1fad50d847
commit e58a728b80
13 changed files with 232 additions and 7 deletions
+10
View File
@@ -673,6 +673,7 @@ classDiagram
+Tensor responses
+Tensor response_mask
+Tensor logprobs_old
+int policy_version
+List[str] prompt_texts
+List[List[str]] response_texts
}
@@ -695,10 +696,14 @@ classDiagram
+float top_p
+float frequency_penalty
+int rep_window
+int policy_version
+update_weights(policy_version) int
+generate(batch) RawRollout
}
class RolloutRunner {
+int policy_version
+update_weights(policy_version) int
+step()
+clear_cache()
+__call__(batch) Tuple[RolloutResult, bool]
@@ -854,11 +859,13 @@ classDiagram
+int max_seq_len
+str device
+torch.dtype dtype
+int policy_version
+add_task(prompt, **kwargs) str
+remove_task(task_id)
+start()
+stop()
+get_stats() Dict
+update_weights(policy_version) int
+run_batch(prompt_ids_list, max_tokens, temperature, top_p, top_k, frequency_penalty, rep_window, return_logprobs) Union[List[List[int]], List[Tuple[List[int], List[float]]]]
}
@@ -871,6 +878,7 @@ classDiagram
+inc_ref(idx)
+touch(idx)
+ref_count(idx) int
+clear_cached() int
}
class RadixNode {
@@ -897,6 +905,7 @@ classDiagram
+extend(state, pos) bool
+write_indices(state, prompt_ids)
+record_hashes(state, prompt_ids, start_logical_page)
+invalidate_cache() int
}
class ContiguousStrategy {
@@ -960,6 +969,7 @@ classDiagram
+task_extend(task_id, pos) bool
+task_cached(task_id) int
+task_record_hashes(task_id, prompt_ids, start_logical_page)
+invalidate_cache() int
+bind(task_ids, workspace) KVCache
}
+6
View File
@@ -139,6 +139,12 @@ attention backends share the same rotary dispatch — it is backend-agnostic.
4. Decode → Run single-token forward for each same-position group
```
For in-process training rollout, `InferenceScheduler.update_weights(version)`
acknowledges that the shared model was updated in place. Versions are monotonic;
the scheduler rejects updates while requests are queued and invalidates reusable
prefix KV pages before exposing the new version. Synchronous `run_batch()` and
weight updates are serialized so a generation cannot straddle two versions.
## Sampling (Strategy Pattern)
```
+7
View File
@@ -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.