From 38eee6a76e28d17724f866b22fdf15349cb342e7 Mon Sep 17 00:00:00 2001 From: MichaelChung Date: Tue, 3 Nov 2026 01:28:20 +0800 Subject: [PATCH] Keep policy versions unique when a trainer resumes behind the store's releases - bridges: a buffer trainer whose round cut is behind the store's newest release continues the version account above it (resume_at, behind_round_cuts; LocalOnlySync.start). - Miles single trainer: on relaunch the publisher continues above the store's releases (_resume_single_trainer, ReleasingPublisher resumed_base) instead of reusing a version already released. Co-Authored-By: Claude Opus 5.5 --- tests/test_rl_sampler_island.py | 81 +++++++++++++++++++++++- tests/test_rl_verl_buffer_trainer.py | 34 ++++++++++ yeto/rl/adapters/miles/sampler_island.py | 49 +++++++++++++- yeto/rl/engine/bridges.py | 20 +++++- 4 files changed, 176 insertions(+), 8 deletions(-) diff --git a/tests/test_rl_sampler_island.py b/tests/test_rl_sampler_island.py index 2ef49361..60d6c7b5 100644 --- a/tests/test_rl_sampler_island.py +++ b/tests/test_rl_sampler_island.py @@ -256,9 +256,84 @@ def test_release_is_committed_by_latest_and_republish_is_a_no_op(tmp_path): assert releases.resolve(first) == tmp_path / RELEASE_DIR / "r1-v1" assert releases.release(1, "h1", _write_weights) is None # a restarted trainer republishes releases.release(2, "h2", _write_weights) - rewind = releases.release(1, "h1b", _write_weights) # restored from an older cut - assert (rewind.seq, rewind.version) == (3, 1) - assert sorted(p.name for p in (tmp_path / RELEASE_DIR).iterdir()) == [LATEST, "r2-v2", "r3-v1"] + assert releases.latest().version == 2 + # A trainer restored from an older cut used to release v1 again with other + # weights here (LATEST went back to v1, and v1 named two sets of weights). The + # trainer now continues above the store's releases and never releases a version + # the store holds (test_a_single_trainer_restored_from_an_older_cut_...). + assert sorted(p.name for p in (tmp_path / RELEASE_DIR).iterdir()) == [LATEST, "r1-v1", "r2-v2"] + + +def test_a_single_trainer_restored_from_an_older_cut_continues_above_the_store(tmp_path): + """Intended change of semantics: the store holds v1..v3 (released after the last + round cut, at v1). The relaunched single trainer continues its account at v3 (as + the verl buffer trainer, rl_round_cut_behind_releases), releases nothing for v3, + and its first trained round is v4: LATEST never goes back and every version + names one set of weights.""" + from yeto.rl.engine.bridges import LocalOnlySync + + releases = PolicyReleaseStore(CheckpointStore(tmp_path), keep=4) + for version in (1, 2, 3): + releases.release(version, f"h{version}", _write_weights) + hashes = {r.version: r.policy_hash for r in releases.records()} + events = [] + sync = LocalOnlySync(2) + sync.round_cuts = SimpleNamespace(resume=lambda d: {"next_rollout_id": 1}) # wired first + driver = SimpleNamespace( + publisher=SimpleNamespace(publish=lambda s: None), sync=sync, _max_policy_age=lambda: 2, + emit=lambda event, **fields: events.append((event, fields)), + export_local_resident=lambda *, policy_version: _State(policy_version, "cut"), + rollout=SimpleNamespace(_in_executor=lambda fn, *a: fn(None, *a))) + args = SimpleNamespace(yeto_rl_sample_buffer=sample_roles.SampleBufferConfig( + role="trainer", store=str(tmp_path), release_inline=True)) + exported = [] + actor = SimpleNamespace(export_hf=lambda rollout_id, path: exported.append(rollout_id) + or _write_weights(Path(path))) + wired = si.wire_policy_releases(driver, miles_args=args, actor_model=actor, + runner=SimpleNamespace(run=lambda value: value)) + br._ACTIVE.pop(br.TRAINER, None) + assert driver.published_policy_history() == [(1, "h1"), (2, "h2"), (3, "h3")] + + start = sync.start(driver) + assert (start.rollout_id, start.state.policy_version) == (3, 3) # not the cut's 1 + assert ("rl_round_cut_behind_releases", + {"cut_rollout_id": 1, "latest_release_version": 3}) in events + wired.publish(start.state) # the cut's weights: v3 is the store's, not released again + skipped = [f for e, f in events if e == "rl_policy_release_skipped"] + assert [(f["policy_version"], f["reason"]) for f in skipped] == [(3, "resumed_base_version")] + wired.publish(_State(4, "h4")) # the first trained round + assert exported == [4] + assert [f["policy_version"] for e, f in events if e == "rl_policy_released"] == [4] + records = releases.records() # the trainer's store keeps the newest two + assert [r.version for r in records] == [3, 4] # no version released twice + assert {r.version: r.policy_hash for r in records if r.version <= 3} == {3: hashes[3]} + assert releases.latest().version == 4 + + +def test_a_single_trainer_relaunched_without_a_cut_continues_above_the_store(tmp_path): + """No round cut (a plain relaunch): the account starts at the store's vK, not at 0.""" + from yeto.rl.engine.bridges import LocalOnlySync + + releases = PolicyReleaseStore(CheckpointStore(tmp_path), keep=4) + releases.release(1, "h1", _write_weights) + releases.release(2, "h2", _write_weights) + sync = LocalOnlySync(1) + driver = SimpleNamespace( + publisher=SimpleNamespace(publish=lambda s: None), sync=sync, _max_policy_age=lambda: 1, + emit=lambda event, **fields: None, + export_local_resident=lambda *, policy_version: _State(policy_version, "base"), + rollout=SimpleNamespace(_in_executor=lambda fn, *a: fn(None, *a))) + args = SimpleNamespace(yeto_rl_sample_buffer=sample_roles.SampleBufferConfig( + role="trainer", store=str(tmp_path), release_inline=True)) + wired = si.wire_policy_releases(driver, miles_args=args, actor_model=SimpleNamespace( + export_hf=lambda rollout_id, path: _write_weights(Path(path))), + runner=SimpleNamespace(run=lambda value: value)) + br._ACTIVE.pop(br.TRAINER, None) + start = sync.start(driver) + assert start.rollout_id == 2 and sync._end() == 3 + wired.publish(start.state) + wired.publish(_State(3, "h3")) + assert [(r.version, r.policy_hash) for r in releases.records()] == [(2, "h2"), (3, "h3")] def test_lora_release_records_the_file_digest_and_the_follower_checks_it(tmp_path): diff --git a/tests/test_rl_verl_buffer_trainer.py b/tests/test_rl_verl_buffer_trainer.py index a70400a0..dcac783b 100644 --- a/tests/test_rl_verl_buffer_trainer.py +++ b/tests/test_rl_verl_buffer_trainer.py @@ -1076,3 +1076,37 @@ def origin(state): assert [r["policy_version"] for r in log] == sorted(r["policy_version"] for r in log) assert [r["policy_version"] for r in log if r["released"] is True] == [1, 3] assert origin_threads and all(origin_threads) + + +def test_a_round_cut_behind_the_store_continues_above_its_releases(tmp_path): + """The store holds v1..v3 (released after the last round cut, at v1). Retraining + from the cut and calling its rounds v2, v3 would give those numbers other + weights than the store's; the account continues at v3 as a relaunch does, and + the first trained round releases v4: every version names one set of weights.""" + store_dir = tmp_path / "store" + store = CheckpointStore(store_dir) + releases = PolicyReleaseStore(store, keep=4) + base, _weights = _base(tmp_path, dtype=torch.float32) + for version in (1, 2, 3): + _released(releases, version) + hashes = {r.version: r.policy_hash for r in releases.records()} + writer = GroupWriter(FileSampleBuffer(store, batch_size=1), "s0") + for i, version in enumerate((2, 3)): + writer.write([_sample(version, index=0), _sample(version, index=1)], group_key=f"g{i}") + driver, island, _publisher = _trainer(tmp_path, store_dir, FakeBufferActor(), rounds=1, required=2, + base=base, releases=releases, resume_versions=True, + buffer=FileSampleBuffer(store, batch_size=1, create_index=False)) + driver.sync.round_cuts = SimpleNamespace( # the cut the run resumes from: v1 + resume=lambda d: {"next_rollout_id": 1}, due=lambda *a, **k: False, + save=lambda *a, **k: None, prepare=lambda *a, **k: None) + state = driver.run() + + assert driver.rounds_completed == 1 and state.policy_version == 4 + events = _lines(tmp_path / "tape.jsonl") + behind = [e for e in events if e["event"] == "rl_round_cut_behind_releases"] + assert [(e["cut_rollout_id"], e["latest_release_version"]) for e in behind] == [(1, 3)] + release_log = _lines(tmp_path / "out" / RELEASE_LOG.format(learner=0)) + assert [(r["policy_version"], r["released"]) for r in release_log] == [(3, False), (4, True)] + records = releases.records() + assert [r.version for r in records] == [1, 2, 3, 4] # no version released twice + assert {r.version: r.policy_hash for r in records if r.version <= 3} == hashes diff --git a/yeto/rl/adapters/miles/sampler_island.py b/yeto/rl/adapters/miles/sampler_island.py index 815d8220..6997f4a2 100644 --- a/yeto/rl/adapters/miles/sampler_island.py +++ b/yeto/rl/adapters/miles/sampler_island.py @@ -30,6 +30,14 @@ ``base_policy``). Synced trainers release it (a restarted trainer restores its published versions from the store's releases). + Relaunch of a single trainer (no lane, no syncer; as the verl buffer trainer, + ``merged_release.ResumedReleases``): the version account continues above the + highest release vK the store holds -- also when the round cut it resumes from + is behind vK (``LocalOnlySync.start``, ``rl_round_cut_behind_releases``). Its + first publication is vK and releases nothing (``resumed_base_version``); the + first trained round releases vK+1. A version number never names two sets of + weights and ``LATEST`` never goes back to an older version. + Engine weight load (the order of the publisher, ``publish.MilesPublisher``): under the controller's update lock, every engine loads the release, then the plain policy version is stamped and read back on every engine; any failure @@ -80,7 +88,12 @@ copy_with_digest, ReleaseNotLoaded, ) -from yeto.rl.engine.merged_release import MergedReleaseGuard, is_publisher, release_version +from yeto.rl.engine.merged_release import ( + MergedReleaseGuard, + is_publisher, + release_version, + resumed_releases, +) from yeto.rl.engine.policy_versions import engine_weight_version from yeto.rl.engine.release_delta import DeltaMirror, default_mirror from yeto.rl.engine.sample_roles import CONVERTIBLE, SAMPLER, TRAINER @@ -673,6 +686,7 @@ def __init__(self, inner: Any, releases: PolicyReleaseStore, read_position_from_thread: Callable[[], int | None] | None = None, read_shards_from_thread: Callable[[], dict[str, int] | None] | None = None, delta: DeltaChain | None = None, + resumed_base: int | None = None, timer: Callable[[], float] = time.perf_counter) -> None: """``delta``: staged releases go out as bf16 deltas along this chain (``policy_release`` module doc, delta releases). ``read_shards``: the @@ -683,7 +697,9 @@ def __init__(self, inner: Any, releases: PolicyReleaseStore, policy version), ``guard`` keeps the release versions increasing and ``origin(state)`` is the release record's ``origin``. ``read_position_from_thread`` / ``read_shards_from_thread``: the same reads, - callable from the upload thread (default: ``read_position`` / ``read_shards``).""" + callable from the upload thread (default: ``read_position`` / ``read_shards``). + ``resumed_base``: a relaunched single trainer's vK (module doc, relaunch): + versions up to it are the store's already and are not released again.""" if (staging is None) != (releaser is None): raise ValueError("staging and releaser go together") if delta is not None and staging is None: @@ -704,6 +720,7 @@ def __init__(self, inner: Any, releases: PolicyReleaseStore, self.staging = None if staging is None else Path(staging) self.releaser = releaser self.skip_base = bool(skip_base) + self.resumed_base = None if resumed_base is None else int(resumed_base) self.timer = timer def _engine_checksum(self) -> str | None: @@ -727,6 +744,13 @@ def publish(self, state: Any, **kwargs: Any) -> Any: latest_version=None if latest is None else latest.version, reason="not_newer_than_latest") return result + if self.resumed_base is not None and version <= self.resumed_base: + # the store holds vK already and nothing newer is trained yet: releasing it + # would give vK a second set of weights (module doc, relaunch) + self.emit("rl_policy_release_skipped", policy_version=version, + local_version=int(state.policy_version), latest_version=self.resumed_base, + reason="resumed_base_version") + return result if self.skip_base and version == BASE_POLICY_VERSION and self.releases.latest() is None: self.emit("rl_policy_release_skipped", policy_version=version, local_version=int(state.policy_version), latest_version=None, @@ -907,6 +931,8 @@ def on_gap(previous: int, version: int) -> None: if not publisher_island: return None synced = {"version_of": version_of, "guard": MergedReleaseGuard(on_gap), "origin": origin} + else: + synced = _resume_single_trainer(driver, release_store(config, environ)) lora = lora_mode(miles_args) # module doc: a trainer island stages its full releases unless --rl-release-inline @@ -946,6 +972,25 @@ def export(version: int, directory: Path) -> None: return publisher +def _resume_single_trainer(driver: Any, releases: PolicyReleaseStore) -> dict[str, Any]: + """A single trainer relaunched over a store that holds releases (module doc, + relaunch): its account continues at the highest release vK, and the versions + below it are restored with the hashes the store recorded. Returns the + ``ReleasingPublisher`` keywords ({} for a store without releases).""" + limit = getattr(driver, "_max_policy_age", None) + resumed = resumed_releases(releases, int(limit()) if callable(limit) else 0) + if resumed is None: + return {} + resume_at = getattr(getattr(driver, "sync", None), "resume_at", None) + if callable(resume_at): + resume_at(resumed.start_version, behind_round_cuts=True) + driver.published_policy_history = resumed.history + driver.emit("rl_buffer_resume_versions", base_version=resumed.base_version, + start_version=resumed.start_version, account_resumed=callable(resume_at), + known_versions=list(resumed.known_versions())) + return {"resumed_base": resumed.base_version} + + def start_trainer_releaser(miles_args: Any, config: Any, environ: Any = None) -> Any: """FP8 rollout (tier 2), trainer side: a thread of this driver converts every BF16 release into an FP8 release (``yeto.rl.quant.release.scoped_releaser``) for the sampler islands.""" diff --git a/yeto/rl/engine/bridges.py b/yeto/rl/engine/bridges.py index b780dd23..d5341190 100644 --- a/yeto/rl/engine/bridges.py +++ b/yeto/rl/engine/bridges.py @@ -177,17 +177,22 @@ def __init__(self, num_rollout: int) -> None: # (:meth:`resume_at`); 0 = this process starts the account at v0. self.resumed_at = 0 - def resume_at(self, version: int) -> None: + def resume_at(self, version: int, *, behind_round_cuts: bool = False) -> None: """Continue the version account at ``version`` K instead of at 0, for a buffer trainer that found releases up to vK in its store (``merged_release.ResumedReleases``): its first publication is vK, its rounds are K, K+1, ..., and ``num_rollout`` stays the rounds *this process* runs, counted from K. - Refused with a round-cut resume, which brings its own start round.""" + Refused when the round cuts are wired first. Wired before them (a buffer + trainer binds its publisher first), a round cut behind ``version`` continues + at ``version`` (:meth:`start`). ``behind_round_cuts``: wired after them (a + Miles single trainer wires its round cuts before its releases), the same + holds -- ``version`` is a floor under the round cut's start round, not a + second start round.""" if int(version) < 0: raise ValueError(f"the version account cannot start below 0, got {version!r}") - if self.round_cuts is not None: + if self.round_cuts is not None and not behind_round_cuts: raise ValueError("a resumed version account and a round-cut resume both set the " "start round; this island has round cuts configured") self.resumed_at = int(version) @@ -195,6 +200,15 @@ def resume_at(self, version: int) -> None: def start(self, driver: IslandDriver) -> SyncStart: resumed = self.round_cuts.resume(driver) if self.round_cuts is not None else None rollout_id = self.resumed_at if resumed is None else int(resumed["next_rollout_id"]) + if resumed is not None and self.resumed_at > rollout_id: + # The store already holds releases up to vK above this round cut (they were + # released after the cut). Retraining from the cut and naming its rounds + # v(cut)..vK would give those numbers other weights than the store's; the + # account continues at vK instead (as a relaunch without a cut does), so + # the first trained round releases vK+1 and every version names one set of weights. + driver.emit("rl_round_cut_behind_releases", cut_rollout_id=rollout_id, + latest_release_version=self.resumed_at) + rollout_id = self.resumed_at return SyncStart(_local_state(driver, rollout_id), rollout_id, finished=rollout_id >= self._end())