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:
@@ -43,6 +43,8 @@ class _RecordingRunner:
|
||||
self.calls = 0
|
||||
self.step_calls = 0
|
||||
self._fresh = True
|
||||
self.policy_version = result.policy_version
|
||||
self.weight_updates = []
|
||||
|
||||
def __call__(self, batch):
|
||||
self.calls += 1
|
||||
@@ -53,6 +55,11 @@ class _RecordingRunner:
|
||||
def step(self):
|
||||
self.step_calls += 1
|
||||
|
||||
def update_weights(self, policy_version):
|
||||
self.policy_version = policy_version
|
||||
self.weight_updates.append(policy_version)
|
||||
return policy_version
|
||||
|
||||
def swap_result(self, result):
|
||||
self.result = result
|
||||
self._fresh = True
|
||||
@@ -302,6 +309,8 @@ def test_step_called_when_sync_gradients_true(device):
|
||||
strat({"input_ids": torch.randint(3, 200, (2, 4), device=device)})
|
||||
strat.on_optimizer_step()
|
||||
assert runner.step_calls == 1
|
||||
assert runner.weight_updates == [1]
|
||||
assert strat.policy_version == 1
|
||||
|
||||
|
||||
def test_loss_is_differentiable_dpo(device):
|
||||
|
||||
Reference in New Issue
Block a user