verl buffer trainer: background releases, per-round phase timing and hybrid MFU accounting

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`):

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)

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

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,14 +258,15 @@ 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),
         publisher=publisher, placement=placement,
         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)

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"]

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