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:
@@ -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
|
||||
}
|
||||
|
||||
|
||||
@@ -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)
|
||||
|
||||
```
|
||||
|
||||
@@ -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.
|
||||
|
||||
Reference in New Issue
Block a user