From 6ef9f6e0e96e5551f327fdee6e92118875f09a02 Mon Sep 17 00:00:00 2001 From: MichaelChung Date: Tue, 3 Nov 2026 01:28:20 +0800 Subject: [PATCH] verl buffer trainer: background releases, per-round phase timing and hybrid MFU accounting - --rl-release-background: the staged release is written on the training thread, then encoded, digested and uploaded on a worker thread while the next rounds train; at most one release in flight, its error raises at the next publication; the release log stays in version order and the release origin is computed on the training thread; the publisher is closed when the run raises. - rl_round_phases: per-round phase timing emitted on the trainer islands only; no clock reads for phase accounting while it is off. - hybrid_flops: Ascend MFU accounting and a CLI that prints step/wall MFU per trained round from a trainer event tape. Co-Authored-By: Claude Opus 5.5 --- docs/MILES_RL.md | 10 ++ tests/test_rl_engine_driver.py | 17 ++ tests/test_rl_hybrid_params.py | 35 ++++ tests/test_rl_verl_buffer_trainer.py | 141 +++++++++++++++- tests/test_verl_sample_role.py | 34 ++++ yeto/rl/adapters/miles/entry.py | 1 + yeto/rl/adapters/verl/buffer_ports.py | 158 +++++++++++++++--- .../rl/adapters/verl/buffer_trainer_runner.py | 3 + yeto/rl/adapters/verl/island_entry.py | 4 + yeto/rl/adapters/verl/trainer.py | 15 +- yeto/rl/engine/driver.py | 45 +++++ yeto/rl/hybrid_flops.py | 82 ++++++++- 12 files changed, 515 insertions(+), 30 deletions(-) diff --git a/docs/MILES_RL.md b/docs/MILES_RL.md index 064a0df7..ffc23eca 100644 --- a/docs/MILES_RL.md +++ b/docs/MILES_RL.md @@ -1015,6 +1015,12 @@ invariants in the [Development Guide](#development-guide) must hold. These options apply to a buffer trainer (`--rl-sample-role trainer`) that publishes policy releases to separate rollout islands. +- **Background releases** (`--rl-release-background`, island entry; needs + `--rl-release-delta-full-every N`): the trainer writes each staged release + on the training thread, then encodes, digests and uploads it on a worker + thread while the next rounds train. At most one release is in flight; the + next publication waits for it, and a failed background release raises + there. Single trainer only. - **Sparse weight updates** (sampler spec `engine.options.sparse_update`, `off` by default, or `auto`): when a release was rebuilt from the served one with change tracking (`YETO_RL_MIRROR_TRACK_CHANGES=1`), the SGLang @@ -1024,6 +1030,10 @@ publishes policy releases to separate rollout islands. elements, or touches a tensor the startup probe could not map, and reloads in full after a failed write or check. bf16 weights only; quantized weights and models whose `load_weights` derives tensors are not supported. +- **MFU of a hybrid model**: `python -m yeto.rl.hybrid_flops TAPE --config + config.json --gpus N --gpu CARD` prints one JSON line per trained round + (`step_mfu`, `wall_mfu`) from the trainer's event tape, for models whose + framework MFU counter does not cover them (Nemotron-H). Spot rollout islands on Nebius have four operator tools in `yeto/rl/spot/` (each with `--help`): diff --git a/tests/test_rl_engine_driver.py b/tests/test_rl_engine_driver.py index f8830bcd..13cb4c57 100644 --- a/tests/test_rl_engine_driver.py +++ b/tests/test_rl_engine_driver.py @@ -1013,3 +1013,20 @@ def test_rollout_offload_after_generate_runs_in_its_own_phase(tmp_path): phases = [e["phase"] for e in _events(tmp_path / "events.jsonl") if e["event"] == "rl_driver_phase"] i = phases.index("generate") assert phases[i + 1] == "rollout_offload" + + +def test_phase_accounting_off_reads_no_clock(): + """round_phases off (the default): phase changes do not read the clock, so a + driver's event stream (and a ticking test clock) is as before the breakdown.""" + from types import SimpleNamespace + + from yeto.rl.engine.driver import IslandDriver + + reads = [] + stub = SimpleNamespace(round_phases=False, clock=lambda: reads.append(1) or 0.0, + _phase_mark=None, _round_phase_seconds=None) + IslandDriver._account_phase(stub, "train") + assert reads == [] and stub._phase_mark is None + stub.round_phases = True + IslandDriver._account_phase(stub, "train") + assert reads == [1] and stub._phase_mark == ("train", 0.0) diff --git a/tests/test_rl_hybrid_params.py b/tests/test_rl_hybrid_params.py index e23d361e..44411ec3 100644 --- a/tests/test_rl_hybrid_params.py +++ b/tests/test_rl_hybrid_params.py @@ -213,3 +213,38 @@ def test_hybrid_flops_of_the_vl_config_counts_the_language_model(): assert hf.forward_flops(vl, [512, 7]) == hf.forward_flops(llm, [512, 7]) # ~12B active parameters per token (model card: 120B total, 12B active) assert 11e9 < hf.forward_flops(vl, [1]) / 2 < 14e9 + + +def test_ascend_peak_and_unknown_devices(): + from yeto.rl import hybrid_flops as hf + + assert hf.peak_flops("910B3") == pytest.approx(313e12) + assert hf.mfu(NANO, [1000], 1.0, 24, gpu="910B3") == pytest.approx( + hf.train_flops(NANO, [1000]) / (24 * 313e12)) + with pytest.raises(ValueError, match="no BF16 peak"): + hf.peak_flops("910B") + + +def test_round_mfu_separates_the_train_step_from_the_round_wall(tmp_path): + import json + + from yeto.rl import hybrid_flops as hf + + events = [ + {"event": "rl_round_trained", "rollout_id": 0, "step_seconds": 100.0, "tok_per_s": 3800.0}, + {"event": "rl_round_phases", "rollout_id": 0, "wall_seconds": 400.0, "trained_tokens": 380_000, + "phase_seconds": {"generate": 30.0, "train": 100.0, "publish": 270.0}}, + {"event": "rl_round_trained", "rollout_id": 1, "step_seconds": 50.0, "tok_per_s": 1000.0}, + ] + rows = hf.round_mfu(events, flops_per_token=7.34e10, n_gpus=24, gpu="910B3") + peak = 24 * 313e12 + assert rows[0]["step_mfu"] == pytest.approx(380_000 * 7.34e10 / (100 * peak)) + assert rows[0]["wall_mfu"] == pytest.approx(rows[0]["step_mfu"] / 4) + assert rows[0]["phase_seconds"]["publish"] == 270.0 + # an older tape without rl_round_phases: tokens from tok_per_s, no wall MFU + assert rows[1]["tokens"] == pytest.approx(50_000) and rows[1]["wall_mfu"] is None + tape = tmp_path / "tape.jsonl" + tape.write_text("\n".join(json.dumps(e) for e in events)) + cfg = tmp_path / "config.json" + cfg.write_text(json.dumps(NANO)) + assert hf.main([str(tape), "--config", str(cfg), "--gpus", "24", "--gpu", "910B3"]) == 0 diff --git a/tests/test_rl_verl_buffer_trainer.py b/tests/test_rl_verl_buffer_trainer.py index f768f613..a70400a0 100644 --- a/tests/test_rl_verl_buffer_trainer.py +++ b/tests/test_rl_verl_buffer_trainer.py @@ -235,7 +235,7 @@ def yeto_call_workers(self, fn, tensors, reset): def _trainer(tmp_path, store_dir, actor, *, rounds, required, base, buffer=None, limit=1, - releases=None, resume_versions=False, no_release=False, release_only=None): + releases=None, resume_versions=False, no_release=False, release_only=None, **ports): from yeto.rl.adapters.verl.entry import verl_capabilities from yeto.rl.adapters.verl.fully_async_ports import FullyAsyncIsland, FullyAsyncTrainerGroup from yeto.rl.adapters.verl.ports_impl import VerlPolicyState @@ -258,7 +258,7 @@ def _trainer(tmp_path, store_dir, actor, *, rounds, required, base, buffer=None, island, buffer=buffer, releases=releases or PolicyReleaseStore(store), shape=bf.RowShape(64, 16, 0), base_dir=base, tokenize=lambda text: [7, 8], poll_s=0.01, wait_s=30, resume_versions=resume_versions, - no_release=no_release, release_only=release_only) + no_release=no_release, release_only=release_only, **ports) profile = far.execution_profile_for(TOLERANT, 1, groups_per_round=required, samples_per_group=2, sync="none") driver = IslandDriver( learner_id=0, rollout=rollout, trainer=FullyAsyncTrainerGroup(island), policy_state=VerlPolicyState(island), @@ -266,6 +266,7 @@ def _trainer(tmp_path, store_dir, actor, *, rounds, required, base, buffer=None, capabilities=verl_capabilities("sha256:" + "0" * 64, fully_async_limit=1), algorithm=TOLERANT, sync=LocalOnlySync(rounds), events=tape, profile=profile) publisher.bind(driver) # as trainer.run_fully_async_island does for a buffer trainer + driver.round_phases = True # as trainer.run_fully_async_island does return driver, island, publisher @@ -769,6 +770,109 @@ def test_a_run_releases_only_the_listed_versions(tmp_path): assert beat["release"] is None # no release existed when it last read +def _staged_run(tmp_path, *, background): + tmp_path.mkdir(parents=True, exist_ok=True) + store_dir = tmp_path / "store" + store = CheckpointStore(store_dir) + releases = PolicyReleaseStore(store) + base, _weights = _base(tmp_path, dtype=torch.float32) + writer = GroupWriter(FileSampleBuffer(store, batch_size=1), "s0") + for i in range(4): + writer.write([_sample(0, index=0), _sample(0, index=1)], group_key=f"g{i}") + driver, island, publisher = _trainer( + tmp_path, store_dir, FakeBufferActor(), rounds=2, required=2, base=base, + buffer=FileSampleBuffer(store, batch_size=1, create_index=False), releases=releases, + staging=tmp_path / "staging", delta_full_every=4, release_background=background) + state = driver.run() + publisher.close() # as trainer.run_fully_async_island does after the driver returns + return driver, state, releases, _lines(tmp_path / "tape.jsonl"), _lines(island.out / RELEASE_LOG.format(learner=0)) + + +def test_background_releases_match_the_inline_ones(tmp_path): + """--rl-release-background: the training thread writes each staged release and + returns; the upload runs on a worker and its record lands at the next release + (or close). Same releases, same order, same log lines as the inline path.""" + inline = _staged_run(tmp_path / "inline", background=False) + background = _staged_run(tmp_path / "background", background=True) + for driver, state, releases, events, log in (inline, background): + assert driver.rounds_completed == 2 and state.policy_version == 2 + assert [r.version for r in releases.records()] == [1, 2] + assert [e["policy_version"] for e in events if e["event"] == "rl_policy_released"] == [1, 2] + assert [e["policy_version"] for e in events if e["event"] == "rl_policy_release_delta"] == [1, 2] + assert [r["released"] for r in log if r["policy_version"] > 0] == [True, True] + keys = ("policy_version", "policy_hash", "released", "seq", "path", "payload_bytes") + assert ([{k: r.get(k) for k in keys} for r in inline[4]] + == [{k: r.get(k) for k in keys} for r in background[4]]) + bg_events = background[3] + assert [e["policy_version"] for e in bg_events if e["event"] == "rl_policy_release_background"] == [1, 2] + assert not [e for e in inline[3] if e["event"] == "rl_policy_release_background"] + # v1's record is put when v2's publication waits for it, before v2 starts its upload + order = [(e["event"], e["policy_version"]) for e in bg_events + if e["event"] in ("rl_policy_released", "rl_policy_release_background")] + assert order == [("rl_policy_release_background", 1), ("rl_policy_released", 1), + ("rl_policy_release_background", 2), ("rl_policy_released", 2)] + assert all("wait_seconds" in r and "write_seconds" in r for r in background[4] if r["policy_version"] > 0) + + +def test_a_failed_background_release_raises_at_the_next_one(tmp_path): + from yeto.rl.engine.policy_release import PolicyReleaseError + + store_dir = tmp_path / "store" + store = CheckpointStore(store_dir) + releases = PolicyReleaseStore(store) + base, _weights = _base(tmp_path, dtype=torch.float32) + writer = GroupWriter(FileSampleBuffer(store, batch_size=1), "s0") + for i in range(4): + writer.write([_sample(0, index=0), _sample(0, index=1)], group_key=f"g{i}") + calls = [] + original = releases.release_staged + + def flaky(*args, **kwargs): + calls.append(args[0]) + if len(calls) == 1: + raise OSError("store unreachable") + return original(*args, **kwargs) + + releases.release_staged = flaky + driver, _island, _publisher = _trainer( + tmp_path, store_dir, FakeBufferActor(), rounds=2, required=2, base=base, + buffer=FileSampleBuffer(store, batch_size=1, create_index=False), releases=releases, + staging=tmp_path / "staging", delta_full_every=4, release_background=True) + with pytest.raises(PolicyReleaseError, match="store unreachable"): + driver.run() + assert calls == [1] # v2 was never uploaded over a failed v1 + assert releases.latest() is None + + +def test_background_releases_need_staging_on_a_single_trainer(tmp_path): + with pytest.raises(ValueError, match="--rl-release-background"): + BufferReleasePublisher(SimpleNamespace(), None, None, base_dir=tmp_path, background=True) + + +def test_release_background_is_planned_from_the_flag(): + from yeto.rl.adapters.verl.island_entry import sample_buffer_section + + common = ("--rl-sample-role", "trainer", "--rl-resume-store", "/s", "--rl-single-island-no-sync", + "--rl-max-policy-age", "1", "--rl-sample-wait-s", "60") + assert "release_background" not in sample_buffer_section(_entry_args(*common)) + assert sample_buffer_section(_entry_args(*common, "--rl-release-background"))["release_background"] is True + + +def test_every_round_puts_its_time_breakdown_on_the_tape(tmp_path): + """rl_round_phases: the wall seconds of each driver phase per round (the wait for + samples, train, sync, publish), and the trained tokens, for MFU per round.""" + _driver, _state, _releases, events, _log = _staged_run(tmp_path, background=False) + rounds = [e for e in events if e["event"] == "rl_round_phases"] + trained = {e["rollout_id"]: e for e in events if e["event"] == "rl_round_trained"} + assert [e["rollout_id"] for e in rounds] == sorted(trained) == [0, 1] + for row in rounds: + phases = row["phase_seconds"] + assert {"train", "publish"} <= set(phases) + assert row["wall_seconds"] == pytest.approx(sum(phases.values()), abs=0.01) + assert phases["train"] >= row["train_seconds"] - 0.01 + assert row["trained_tokens"] > 0 and 0 < row["train_fraction"] <= 1.0 + + def test_release_only_versions_is_planned_from_the_flag(): from yeto.rl.adapters.verl.island_entry import sample_buffer_section @@ -939,3 +1043,36 @@ def test_the_default_feed_emitter_prints_a_yeto_rl_event(capsys): bf.print_feed_event("rl_buffer_group_dropped", island_id=1, reason="short_group", count=1) line = capsys.readouterr().out.strip() assert line.startswith(PREFIX) and '"event":"rl_buffer_group_dropped"' in line + + +def test_background_release_log_stays_in_version_order_and_the_worker_never_reads_state(tmp_path, monkeypatch): + """A publication that releases nothing while a background release is in flight + puts its log line after that release's (flush); the release origin is computed + on the training thread.""" + import threading + + store_dir = tmp_path / "store" + store = CheckpointStore(store_dir) + releases = PolicyReleaseStore(store) + base, _weights = _base(tmp_path, dtype=torch.float32) + writer = GroupWriter(FileSampleBuffer(store, batch_size=1), "s0") + for i in range(6): # the last round's groups come from v1 (policy-age window 1) + writer.write([_sample(i // 4, index=0), _sample(i // 4, index=1)], group_key=f"g{i}") + driver, island, publisher = _trainer( + tmp_path, store_dir, FakeBufferActor(), rounds=3, required=2, base=base, + buffer=FileSampleBuffer(store, batch_size=1, create_index=False), releases=releases, + release_only=[1, 3], staging=tmp_path / "staging", delta_full_every=4, release_background=True) + origin_threads = [] + original = publisher._origin + + def origin(state): + origin_threads.append(threading.current_thread() is threading.main_thread()) + return original(state) + + monkeypatch.setattr(publisher, "_origin", origin) + driver.run() + publisher.close() + log = _lines(island.out / RELEASE_LOG.format(learner=0)) + 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) diff --git a/tests/test_verl_sample_role.py b/tests/test_verl_sample_role.py index 7706ec32..f84dda8c 100644 --- a/tests/test_verl_sample_role.py +++ b/tests/test_verl_sample_role.py @@ -405,3 +405,37 @@ def ports(island): assert not [e for e in _lines(tmp_path / "tape.jsonl") if e["event"] == "rl_verify_every_wired"] releases = _lines(tmp_path / "out" / RELEASE_LOG.format(learner=0)) assert all(r["policy_hash_verified"] for r in releases) + + +def test_a_failing_run_still_closes_the_buffer_publisher(tmp_path, verl_trainer, monkeypatch): + """A background release in flight is flushed (record, events) even when the + driver raises; a close error goes on the tape and the run's own error is raised.""" + from yeto.rl.adapters.verl import round_cuts + + actor = FullBufferActor(seed=0) + base = _full_base(tmp_path, actor.model) + store = CheckpointStore(tmp_path / "store") + plan = _plan(tmp_path, base, rounds=1, groups=1) + closed = [] + + def ports(island): + rollout, publisher, placement = install_buffer_ports( + island, buffer=FileSampleBuffer(store, batch_size=1, create_index=False), + releases=PolicyReleaseStore(store), shape=bf.RowShape(64, 16, 0), base_dir=base, + tokenize=lambda text: [7, 8], poll_s=0.01, wait_s=30, weights=release_weights_for(island, plan)) + + def close(): + closed.append(True) + raise RuntimeError("upload failed") + + publisher.close = close + return rollout, publisher, placement + + def boom(driver, resume, **kw): + raise KeyboardInterrupt("cut") + + monkeypatch.setattr(round_cuts, "run_driver", boom) + with pytest.raises(KeyboardInterrupt): + verl_trainer.run_fully_async_island(actor, None, plan, {}, buffer_ports=ports) + assert closed == [True] + assert [e for e in _lines(tmp_path / "tape.jsonl") if e["event"] == "rl_policy_release_close_failed"] diff --git a/yeto/rl/adapters/miles/entry.py b/yeto/rl/adapters/miles/entry.py index 6c3211ab..ed5811ec 100644 --- a/yeto/rl/adapters/miles/entry.py +++ b/yeto/rl/adapters/miles/entry.py @@ -823,6 +823,7 @@ def expected_policy() -> tuple[int, str]: ), **({"evaluate_start": evaluate_start} if evaluate_start is not None else {}), ) + driver.round_phases = True # per-round time breakdown (rl_round_phases) for MFU # publication checks on every round (default: cut, first and every-K rounds) driver.verify_every_round = bool(getattr(miles_args, "yeto_rl_verify_every_round", False)) # --rl-verify-every N (no-sync island): hash the policy only every N-th version diff --git a/yeto/rl/adapters/verl/buffer_ports.py b/yeto/rl/adapters/verl/buffer_ports.py index 96102bbc..44d11915 100644 --- a/yeto/rl/adapters/verl/buffer_ports.py +++ b/yeto/rl/adapters/verl/buffer_ports.py @@ -77,7 +77,7 @@ from yeto.rl.contracts import InferencePublicationManifest from yeto.rl.engine.merged_release import (BufferLane, MergedReleaseGuard, ResumedReleases, is_publisher, merged_version, resumed_releases) -from yeto.rl.engine.policy_release import FULL, PolicyReleaseStore +from yeto.rl.engine.policy_release import FULL, BackgroundReleaser, PolicyReleaseStore from yeto.rl.engine.sample_buffer import buffer_dir, write_read_position from yeto.rl.engine.ports import PlacementDescription, PublicationResult @@ -233,7 +233,7 @@ def __init__(self, island: FullyAsyncIsland, releases: PolicyReleaseStore, feed: weights: Callable[[Any], ReleaseWeights] | None = None, staging: str | Path | None = None, delta: Any = None, resumed: ResumedReleases | None = None, no_release: bool = False, - release_only: Iterable[int] | None = None) -> None: + release_only: Iterable[int] | None = None, background: bool = False) -> None: """``no_release`` (``--rl-measure-no-release``): a measurement run (memory, MFU) writes no policy release at all. One release of the 120B is about 225 GB written and a dozen minutes per round, which answers nothing such a run @@ -251,9 +251,20 @@ def __init__(self, island: FullyAsyncIsland, releases: PolicyReleaseStore, feed: writing it into the store (module doc, delta releases); ``delta``: the :class:`~yeto.rl.engine.policy_release.DeltaChain` those staged releases follow (:func:`delta_chain_for`), which needs ``staging``; ``resumed``: the - version account this process continues (module doc, relaunch).""" + version account this process continues (module doc, relaunch). + + ``background`` (``--rl-release-background``): a staged release writes its + weights on the training thread, as before (the next optimizer step changes + them), but encodes, digests and uploads them on a worker thread while the + next rounds train. One release is in flight at most: the next publication + that releases (or :meth:`close`) waits for it first and then puts its + record, events and heartbeat on the tape. A failed background release + raises there. Single trainer only (a synced trainer's peers wait for + the release).""" if delta is not None and staging is None: raise ValueError("delta releases need a staging directory") + if background and (staging is None or lane is not None): + raise ValueError("--rl-release-background needs staged releases on a single trainer") self.resumed = resumed self.staging = None if staging is None else Path(staging) self.delta = delta @@ -273,6 +284,12 @@ def __init__(self, island: FullyAsyncIsland, releases: PolicyReleaseStore, feed: "peers wait for its releases") self.no_release = bool(no_release) self.release_only = None if release_only is None else frozenset(int(v) for v in release_only) + self.background = bool(background) + self.releaser = BackgroundReleaser("verl-policy-release") if self.background else None + # the release in flight: (its record so far, its start, the worker's result) + self._pending: tuple[dict[str, Any], float, dict[str, Any]] | None = None + # publication records that came after the release in flight (committed by flush) + self._deferred: list[tuple[dict[str, Any], Any, dict[str, Any]]] = [] def bind(self, driver: Any) -> None: """The island's driver (synced trainers read its outer version; a resumed @@ -339,31 +356,35 @@ def _release_from_staging(self, version: int, digest: str, write: Callable[[Path """Write the release into a staging directory and publish it from there, as the delta against the chain's base when it has one (module doc). The staged directory stays only while it is the next delta's base.""" + staged = self._write_staged(version, write) + return self._upload_staged(version, digest, staged, position, shards, self._origin(state)) + + def _upload_staged(self, version: int, digest: str, staged: Path, position: int | None, + shards: dict[str, int] | None, origin: dict[str, Any] | None): + """The part of a staged release after the weights are written: the delta, + the digests and the upload (on the worker thread under ``background``). + Returns ``(release, delta event fields)``; the caller emits the event.""" import shutil - import tempfile - self.staging.mkdir(parents=True, exist_ok=True) - staged = Path(tempfile.mkdtemp(prefix=f"v{int(version)}-", dir=self.staging)) base = None if self.delta is None else self.delta.base() kept = False try: - write(staged) release = self.releases.release_staged( version, digest, staged, position=lambda: (position, shards), - weights=FULL, origin=self._origin(state), delta_base=base, + weights=FULL, origin=origin, delta_base=base, digest_full=self.delta is not None) kept = self.delta is not None and self.delta.advance(release, staged) finally: if not kept: shutil.rmtree(staged, ignore_errors=True) delta = None if release is None else release.delta - self.island.emit("rl_policy_release_delta", policy_version=int(version), - base_seq=None if delta is None else delta["base_seq"], - base_version=None if delta is None else delta["base_version"], - ratio=None if delta is None else delta["ratio"], - bytes=None if release is None else sum(int(v) for v in (release.files or {}).values()), - timings=self.releases.last_timings) - return release + event = dict(policy_version=int(version), + base_seq=None if delta is None else delta["base_seq"], + base_version=None if delta is None else delta["base_version"], + ratio=None if delta is None else delta["ratio"], + bytes=None if release is None else sum(int(v) for v in (release.files or {}).values()), + timings=dict(self.releases.last_timings)) + return release, event def _origin(self, state: Any) -> dict[str, Any] | None: # A release whose policy hash is an identity label (an unverified round @@ -392,16 +413,89 @@ def publish(self, state: Any) -> PublicationResult: position = self.feed.read_cursor() shards = self.feed.read_shards() released = self._release_version(state) - if released is not None and not weights.is_base(released): + releasing = released is not None and not weights.is_base(released) + if releasing: + self.flush() # the release in flight first: the store and the delta chain are serial def write(directory: Path) -> None: merge.update(weights.write(directory)) - if self.staging is not None: - release = self._release_from_staging(released, digest, write, position, shards, state) + if self.staging is not None and self.background: + staged = self._write_staged(released, write) + written = round(self.clock() - started, 3) + record = self._record(state, version, param_version, digest, position, shards, + None, merge, started) + record.update(released="pending", write_seconds=written) + result: dict[str, Any] = {} + origin = self._origin(state) # on this thread: the worker never reads ``state`` + + def upload() -> None: + result["release"], result["event"] = self._upload_staged( + released, digest, staged, position, shards, origin) + + self.releaser.submit(upload) + self._pending = (record, started, result) + isl.emit("rl_policy_release_background", policy_version=int(released), + write_seconds=written) + return self._finish_publication(version, param_version, digest, None, weights) + elif self.staging is not None: + release, event = self._release_from_staging(released, digest, write, position, shards, state) + isl.emit("rl_policy_release_delta", **event) else: release = self.releases.release(released, digest, write, buffer_selected=position, buffer_shards=shards, weights=FULL, origin=self._origin(state)) + record = self._record(state, version, param_version, digest, position, shards, release, + merge, started) + if self._pending is not None: # after the release in flight: the log stays in version order + self._deferred.append((record, release, merge)) + else: + self._commit_record(record, release, merge) + return self._finish_publication(version, param_version, digest, record.get("path"), weights, + record.get("payload_bytes")) + + # -- background releases ------------------------------------------------ + def _write_staged(self, version: int, write: Callable[[Path], None]) -> Path: + import shutil + import tempfile + + self.staging.mkdir(parents=True, exist_ok=True) + staged = Path(tempfile.mkdtemp(prefix=f"v{int(version)}-", dir=self.staging)) + try: + write(staged) + except BaseException: + shutil.rmtree(staged, ignore_errors=True) + raise + return staged + + def flush(self) -> None: + """Wait for the background release in flight (if any) and put its record, + events and heartbeat on the tape; its error, if it failed, raises here.""" + if self._pending is None: + return + record, started, result = self._pending + self._pending = None + waited = self.releaser.wait() + release = result["release"] + self.island.emit("rl_policy_release_delta", **result["event"]) + # seconds: publication start to upload end; wait_seconds: how long the + # training thread blocked on it (0 when the upload hid behind training) + record.update(released=release is not None, seconds=round(self.clock() - started, 3), + wait_seconds=round(waited, 3)) + self._commit_record(record, release, record.get("merge") or {}) + deferred, self._deferred = self._deferred, [] + for args in deferred: + self._commit_record(*args) + + def close(self) -> None: + """The end of the run: the release in flight goes out before the island reports done.""" + self.flush() + + # -- records ------------------------------------------------------------ + def _record(self, state: Any, version: int, param_version: int, digest: str, + position: int | None, shards: dict[str, int] | None, release: Any, + merge: dict[str, Any], started: float) -> dict[str, Any]: + """The publication log line (``RELEASE_LOG``) before the release fields.""" + isl = self.island # Under --rl-verify-every N a round the driver does not verify carries # an identity label (policy_digest.unverified_identity), not a content hash. Say # which one this is, so nothing downstream takes the label for the weights' hash. @@ -414,6 +508,14 @@ def write(directory: Path) -> None: if self.lane is not None: record.update(merged_version=merged_version(self.driver, state), merged_versions=self.feed.merged.to_dict()) + return record + + def _commit_record(self, record: dict[str, Any], release: Any, merge: dict[str, Any]) -> None: + """The release event, the publication log line and the heartbeat release.""" + isl = self.island + version, digest = int(record["policy_version"]), record["policy_hash"] + position = record["buffer_selected"] + verified = bool(record["policy_hash_verified"]) if release is not None: record.update(seq=release.seq, path=release.path, release_version=release.version, payload_bytes=sum(int(v) for v in (release.files or {}).values())) @@ -425,17 +527,23 @@ def write(directory: Path) -> None: max_relative_change=merge.get("max_relative_change"), seconds=record["seconds"]) isl.append_jsonl(RELEASE_LOG.format(learner=isl.learner_id), record) + if release is not None: + isl.heartbeat_release = {"seq": int(release.seq), "policy_hash": str(release.policy_hash)} + + def _finish_publication(self, version: int, param_version: int, digest: str, + path: str | None, weights: ReleaseWeights, + payload_bytes: int | None = None) -> PublicationResult: + isl = self.island isl.emit("rl_verl_version_map", param_version=param_version, outer_version=version) isl.published = (version, digest) isl.hashes[version] = digest - if release is not None: - isl.heartbeat_release = {"seq": int(release.seq), "policy_hash": str(release.policy_hash)} manifest = InferencePublicationManifest( publication_mode="full", base_policy_version=None, target_policy_version=version, target_policy_hash=digest, target_manifest_hash=_sha(f"{version}:{digest}"), - payload_hash=_sha(json.dumps({"release": record.get("path"), "hash": digest})), - # the release's files, or (nothing released) the weights' fallback size - payload_bytes=int(record.get("payload_bytes") or weights.fallback_bytes), + payload_hash=_sha(json.dumps({"release": path, "hash": digest})), + # the release's files, or (nothing released, or still uploading) the + # weights' fallback size + payload_bytes=int(payload_bytes or weights.fallback_bytes), complete=True) return PublicationResult(manifest, frozenset({MEMBER})) @@ -455,7 +563,7 @@ def install_buffer_ports(island: FullyAsyncIsland, *, buffer: Any, releases: Pol weights: Callable[[Any], ReleaseWeights] | None = None, staging: str | Path | None = None, delta_full_every: int = 0, resume_versions: bool = False, no_release: bool = False, - release_only: Iterable[int] | None = None, + release_only: Iterable[int] | None = None, release_background: bool = False, **feed_kwargs: Any) -> tuple[FullyAsyncRolloutPool, BufferReleasePublisher, BufferPlacement]: """Route the island's queue calls to a buffer feed and return the rollout @@ -530,5 +638,5 @@ def record(fields: dict[str, Any]) -> None: weights=weights, staging=staging if int(delta_full_every) > 0 else None, delta=delta, resumed=resumed, no_release=no_release, - release_only=release_only) + release_only=release_only, background=release_background) return FullyAsyncRolloutPool(island, queue_end_is_buffer_end=True), publisher, BufferPlacement() diff --git a/yeto/rl/adapters/verl/buffer_trainer_runner.py b/yeto/rl/adapters/verl/buffer_trainer_runner.py index 8424cdb9..fdb9d153 100644 --- a/yeto/rl/adapters/verl/buffer_trainer_runner.py +++ b/yeto/rl/adapters/verl/buffer_trainer_runner.py @@ -200,6 +200,9 @@ def ports(island): # A measurement run writes no release no_release=bool(section.get("no_release")), release_only=section.get("release_only"), + # encode and upload each staged release on a worker + # thread while the next rounds train + release_background=bool(section.get("release_background")), # This island counts its own # versions, so a new process of it continues the # account above the store's newest release diff --git a/yeto/rl/adapters/verl/island_entry.py b/yeto/rl/adapters/verl/island_entry.py index 6f7c5b4f..07c7c1bb 100644 --- a/yeto/rl/adapters/verl/island_entry.py +++ b/yeto/rl/adapters/verl/island_entry.py @@ -96,6 +96,9 @@ def parse_args(argv): # a memory / MFU measurement run on a buffer trainer writes no policy # release (one 120B release is about 225 GB and a dozen minutes per round) p.add_argument("--rl-measure-no-release", action="store_true") + # encode, digest and upload each staged (delta) release on a worker thread + # while the next rounds train; needs --rl-release-delta-full-every + p.add_argument("--rl-release-background", action="store_true") # a measurement run also skips the full policy hash of the first and the last # round (about 480 GB of FP32 masters, ~30 min each on the 120B island); needs # --rl-measure-no-release and --rl-verify-every (train_options.verify_plan) @@ -1361,6 +1364,7 @@ def sample_buffer_section(args) -> dict | None: # staged bf16 delta releases (0: the release goes into the store whole) "release_delta_full_every": int(config.release_delta_full_every or 0), **({"no_release": True} if getattr(args, "rl_measure_no_release", False) else {}), + **({"release_background": True} if getattr(args, "rl_release_background", False) else {}), **({"release_only": list(args.rl_release_only_versions)} if getattr(args, "rl_release_only_versions", None) is not None else {}), **({"shards": list(args.rl_buffer_shards)} if getattr(args, "rl_buffer_shards", None) else {})} diff --git a/yeto/rl/adapters/verl/trainer.py b/yeto/rl/adapters/verl/trainer.py index 78cf965c..b5dca273 100644 --- a/yeto/rl/adapters/verl/trainer.py +++ b/yeto/rl/adapters/verl/trainer.py @@ -469,6 +469,7 @@ def run_fully_async_island(call, start_rollouter, plan: dict, cfg: dict, *, buff capabilities=verl_capabilities(_fingerprint(cfg), fully_async_limit=limit), algorithm=spec, sync=sync, events=tape, profile=profile, ledger=None if resume is None else resume.ledger) driver.data_slot = plan.get("data_slot") # task assignment + driver.round_phases = True # per-round time breakdown (rl_round_phases) for MFU if full: # Observation only (exports and data positions are refused, .fully_async_full) from .driver_extras import wire_full_parameter_extras @@ -494,7 +495,19 @@ def run_fully_async_island(call, start_rollouter, plan: dict, cfg: dict, *, buff if getattr(driver, "verify_every", None) != int(plan["verify_every"]) or not getattr( policy_state, "supports_unverified_rounds", False): raise ValueError("--rl-verify-every was planned but the driver or policy state did not take it") - state = round_cuts.run_driver(driver, resume) + close = getattr(publisher, "close", None) + try: + state = round_cuts.run_driver(driver, resume) + except BaseException: + if callable(close): # a finished background upload still gets its record; the run's error wins + try: + close() + except Exception as error: # noqa: BLE001 - on the tape; the run's own error is raised + tape.append({"event": "rl_policy_release_close_failed", + "error": f"{type(error).__name__}: {error}"[:500]}) + raise + if callable(close): + close() # a background release still in flight goes out before the island reports done result = {"learner_id": learner_id, "rounds_completed": driver.rounds_completed, "final_policy_version": state.policy_version, "final_policy_tensor_hash": state.policy_tensor_hash(), diff --git a/yeto/rl/engine/driver.py b/yeto/rl/engine/driver.py index baea9ee2..d2eca66d 100644 --- a/yeto/rl/engine/driver.py +++ b/yeto/rl/engine/driver.py @@ -648,6 +648,7 @@ def emit(self, event: str, **fields: Any) -> None: def phase(self, name: str, **fields: Any) -> None: self._last_phase = (name, fields.get("rollout_id")) + self._account_phase(name) phase_profile.on_phase(name, fields.get("rollout_id")) # no-op unless --rl-profile-trainer if self.observe: self._close_span() @@ -679,6 +680,48 @@ def _labels(self) -> dict[str, Any]: "weight_transport": self.weight_transport, } + # -- per-round time breakdown ------------------------------------------ + #: emit ``rl_round_phases`` each round (off: the recorded reference tapes stay + #: byte-identical; the trainer islands turn it on) + round_phases: bool = False + _phase_mark: tuple[str, float] | None = None + _round_phase_seconds: dict[str, float] | None = None + + def _account_phase(self, name: str | None) -> None: + """Charge the time since the previous phase change to that phase; ``name`` + (None: stop the clock) starts the next one.""" + if not self.round_phases: # off: not even a clock read (the reference tapes) + return + now = self.clock() + if self._round_phase_seconds is None: + self._round_phase_seconds = {} + if self._phase_mark is not None: + previous, since = self._phase_mark + self._round_phase_seconds[previous] = ( + self._round_phase_seconds.get(previous, 0.0) + max(0.0, now - since)) + self._phase_mark = None if name is None else (name, now) + + def _emit_round_phases(self, rollout_id: int, *, train_seconds: float | None = None, + tokens: int | None = None) -> None: + """``rl_round_phases``: the wall seconds of every driver phase since the + previous round's publication (the wait for samples, train, sync, export, + publish, ...), so the share of a round the trainer computes is on the tape + without the observation spans. The clock keeps running into the next round.""" + if not self.round_phases: + return + current = None if self._phase_mark is None else self._phase_mark[0] + self._account_phase(current) + seconds = dict(self._round_phase_seconds or {}) + self._round_phase_seconds = {} + wall = sum(seconds.values()) + self.emit("rl_round_phases", rollout_id=rollout_id, + phase_seconds={k: round(v, 3) for k, v in sorted(seconds.items())}, + wall_seconds=round(wall, 3), + train_seconds=None if train_seconds is None else round(float(train_seconds), 3), + trained_tokens=tokens, + train_fraction=(round(float(train_seconds) / wall, 4) + if train_seconds is not None and wall > 0 else None)) + def _close_span(self) -> None: if not self.observe or self._open_span is None: return @@ -2138,6 +2181,8 @@ def _train_round(self, rollout_id: int, batch: RolloutBatchHandle, self._offload_trainer(rollout_id) self.publish(boundary.state, rollout_id=rollout_id + 1) self._close_span() + self._emit_round_phases(rollout_id, train_seconds=train_seconds, + tokens=sum(int(g.token_count) for g in batch.groups)) self.rounds_completed += 1 return boundary diff --git a/yeto/rl/hybrid_flops.py b/yeto/rl/hybrid_flops.py index c8e52589..8daffd6f 100644 --- a/yeto/rl/hybrid_flops.py +++ b/yeto/rl/hybrid_flops.py @@ -30,7 +30,9 @@ from .hybrid_params import hybrid_pattern, language_model_config # Dense BF16 peak, no sparsity (NVIDIA H100/H200 SXM datasheets: 989.4 TFLOPS). -PEAK_BF16_TFLOPS = {"H100": 989.4, "H200": 989.4} +# Ascend 910B3: 313 TFLOPS FP16/BF16 cube peak, the denominator every Ascend MFU +# figure of this project uses (SFT 18.6%, RL 11-14%), so the numbers stay comparable. +PEAK_BF16_TFLOPS = {"H100": 989.4, "H200": 989.4, "910B3": 313.0} def _get(cfg: Mapping[str, Any], key: str, default: Any = None) -> Any: @@ -83,6 +85,82 @@ def train_flops(cfg: Mapping[str, Any], seqlens: Iterable[int]) -> float: return 3.0 * forward_flops(cfg, seqlens) +def peak_flops(gpu: str) -> float: + """Dense BF16 peak FLOP/s of one device (``PEAK_BF16_TFLOPS``).""" + try: + return PEAK_BF16_TFLOPS[gpu] * 1e12 + except KeyError: + raise ValueError(f"no BF16 peak for device {gpu!r}; known: {sorted(PEAK_BF16_TFLOPS)}") from None + + def mfu(cfg: Mapping[str, Any], seqlens: Iterable[int], seconds: float, n_gpus: int, gpu: str = "H200") -> float: """Model FLOPs utilization of a training phase that processed ``seqlens`` in ``seconds``.""" - return train_flops(cfg, seqlens) / (seconds * n_gpus * PEAK_BF16_TFLOPS[gpu] * 1e12) + return train_flops(cfg, seqlens) / (seconds * n_gpus * peak_flops(gpu)) + + +def round_mfu(events: Iterable[Mapping[str, Any]], *, flops_per_token: float, n_gpus: int, + gpu: str) -> list[dict[str, Any]]: + """Per-round MFU of a trainer tape, two ways. + + ``step_mfu`` divides the round's training FLOPs by the train step alone + (``rl_round_trained.step_seconds``: forward, backward, the old-log-prob pass and + the optimizer); ``wall_mfu`` divides them by the whole round + (``rl_round_phases.wall_seconds``: also the wait for samples, the sync, the + export and the publication), which is what the cards actually delivered. + ``flops_per_token``: training FLOPs per token (``train_flops(cfg, [L]) / L`` at + the batch's typical length). A round without ``rl_round_phases`` (older + tapes) has ``wall_mfu`` None.""" + peak = peak_flops(gpu) * int(n_gpus) + trained: dict[int, Mapping[str, Any]] = {} + phases: dict[int, Mapping[str, Any]] = {} + for event in events: + name, rid = event.get("event"), event.get("rollout_id") + if rid is None: + continue + if name == "rl_round_trained": + trained[int(rid)] = event + elif name == "rl_round_phases": + phases[int(rid)] = event + out = [] + for rid in sorted(trained): + row = trained[rid] + step = row.get("step_seconds") + tokens = phases.get(rid, {}).get("trained_tokens") + if tokens is None and step and row.get("tok_per_s") is not None: + tokens = float(row["tok_per_s"]) * float(step) + wall = phases.get(rid, {}).get("wall_seconds") + flops = None if tokens is None else float(tokens) * float(flops_per_token) + out.append({ + "rollout_id": rid, "tokens": tokens, "step_seconds": step, "wall_seconds": wall, + "step_mfu": flops / (float(step) * peak) if flops is not None and step else None, + "wall_mfu": flops / (float(wall) * peak) if flops is not None and wall else None, + "phase_seconds": phases.get(rid, {}).get("phase_seconds"), + }) + return out + + +def main(argv: list[str] | None = None) -> int: + """``python -m yeto.rl.hybrid_flops TAPE --config config.json --gpus 24 --gpu 910B3``: + one JSON line per trained round (``round_mfu``).""" + import argparse + import json + from pathlib import Path + + p = argparse.ArgumentParser(description=main.__doc__) + p.add_argument("tape") + p.add_argument("--config", required=True, help="the model's HF config.json") + p.add_argument("--gpus", type=int, required=True) + p.add_argument("--gpu", required=True, choices=sorted(PEAK_BF16_TFLOPS)) + p.add_argument("--seq-len", type=int, default=16384, + help="typical sample length for the attention term (default 16384)") + args = p.parse_args(argv) + cfg = json.loads(Path(args.config).read_text()) + per_token = train_flops(cfg, [args.seq_len]) / args.seq_len + events = [json.loads(line) for line in Path(args.tape).read_text().splitlines() if line.strip()] + for row in round_mfu(events, flops_per_token=per_token, n_gpus=args.gpus, gpu=args.gpu): + print(json.dumps(row, sort_keys=True)) + return 0 + + +if __name__ == "__main__": + raise SystemExit(main())