diff --git a/openspec/changes/rl-resume-from-checkpoint/design.md b/openspec/changes/rl-resume-from-checkpoint/design.md index 3ebe7e801..4a3b7e19d 100644 --- a/openspec/changes/rl-resume-from-checkpoint/design.md +++ b/openspec/changes/rl-resume-from-checkpoint/design.md @@ -79,6 +79,7 @@ Modal 文档：退出处理函数有 30 秒宽限，被抢占时也调用；GPU - **单岛切点**（沿用格式）：adapter、优化器、学习率调度器、随机数、iteration、consumed_samples、pointer（下一轮号、策略版本、策略哈希、启动序号）+ 新增 `run_fingerprint`（算法、模型、LoRA 形状、并行布局、学习率配置、数据集指纹与洗牌种子）与 `lr_at_next_round`（下一轮应使用的学习率）与 `consumed_prompt_ids_digest`（已用题号集合的摘要，详细列表放账本）。 - **多岛**：每个岛存自己的切点（同上，路径 `/islands/<岛名>/`）；syncer 检查点同步到 `/syncer/`；切点 pointer 多记 `outer_version`。 - **频率**：`--rl-cut-every N`（默认：0.6B 为 1，FN 为 2）；最后一轮、收到"我方停机"时总是存一次。 +- **Default since 2026-10-21 (E1)**: `--rl-cut-every` defaults to 4; publication checks (policy hash, engine read-back) run on the first publication and every 4th round; `--rl-verify-every-round` restores per-round checks. See the requirement "Round-cut interval and publication checks" in specs/rl-resume-from-checkpoint/spec.md. - **保留**：`--rl-cut-keep K`（默认 2）；新切点 MANIFEST 落盘并提交后才删旧的。 - **目录**：`/<运行名>/cuts/<轮号>/{shards…, manifest.json, MANIFEST.sha256}` + `/<运行名>/LATEST`（最后写，指向最近完整切点）。 diff --git a/openspec/changes/rl-resume-from-checkpoint/specs/rl-resume-from-checkpoint/spec.md b/openspec/changes/rl-resume-from-checkpoint/specs/rl-resume-from-checkpoint/spec.md index 6f3cbcfb5..025a344b3 100644 --- a/openspec/changes/rl-resume-from-checkpoint/specs/rl-resume-from-checkpoint/spec.md +++ b/openspec/changes/rl-resume-from-checkpoint/specs/rl-resume-from-checkpoint/spec.md @@ -49,3 +49,10 @@ launcher SHALL 接受 `--rl-checkpoint-store modal-volume://<卷名>[/前缀]` #### Scenario: 续训后的 tape - **WHEN** 第二次启动重跑了第一次未进切点的第 4 轮 - **THEN** 第一次的第 4 轮记录仍在，`rl_resume.discarded_rounds` 含 4 + +### Requirement: Round-cut interval and publication checks (E1, 2026-10-21) +The product default round-cut interval SHALL be 4 rounds (`--rl-cut-every 4`; `--rl-cut-every 1` cuts every round), plus the last round and our own stop. A publication SHALL hash the policy and read the engine weight checksums back on the first publication, on every round whose number is a multiple of the cut interval (which includes every round that keeps a round cut), and on every round with `--rl-verify-every-round`. Every other publication SHALL still check that the trainer holds the published policy (trainer weights version) and SHALL be recorded with `sync/publication_verified: false` (verified publications carry no such field). A publication that was not verified MUST NOT be committed as a round cut. + +#### Scenario: an unverified round at a cut +- **WHEN** a cut is due at a safe point whose published policy was not verified +- **THEN** the cut is refused ("not verified") and the previous pointer stays the committed one diff --git a/tests/golden/decoupling/codex_harness.json b/tests/golden/decoupling/codex_harness.json index b957d6a5d..f74fb036c 100644 --- a/tests/golden/decoupling/codex_harness.json +++ b/tests/golden/decoupling/codex_harness.json @@ -81,6 +81,10 @@ "bridge", "--model-name", "SimpleNamespace", + "--attention-dropout", + "0.0", + "--hidden-dropout", + "0.0", "--num-layers", "2", "--hidden-size", @@ -229,7 +233,7 @@ "", "--group-query-attention" ], - "miles_argv_sha256": "279d3bae4c7c8381c9dfbea4e0e83ffdbfcd0afbeabb0f8931d7c3312984f771", + "miles_argv_sha256": "8a2934f5ac2b11fae39d1374ed834d0ab5ed4eb6c1d74f515a2f64bc2505830c", "name": "codex_harness", "placement": { "allow_cross_node_engine_tp": false, diff --git a/tests/golden/decoupling/decoupled.json b/tests/golden/decoupling/decoupled.json index eb40b112f..0d63eb515 100644 --- a/tests/golden/decoupling/decoupled.json +++ b/tests/golden/decoupling/decoupled.json @@ -63,6 +63,10 @@ "bridge", "--model-name", "SimpleNamespace", + "--attention-dropout", + "0.0", + "--hidden-dropout", + "0.0", "--num-layers", "2", "--hidden-size", @@ -195,7 +199,7 @@ "10", "--group-query-attention" ], - "miles_argv_sha256": "1108c08ed4e9ef28ec5ad2e37b3fe6ceb804db496824c9de4743b028a400a663", + "miles_argv_sha256": "6ff6da47c19cf9f9eac0ec0b9040e9e199220578244928625c5bec362379ccf6", "name": "decoupled", "placement": { "allow_cross_node_engine_tp": false, diff --git a/tests/golden/decoupling/drgrpo.json b/tests/golden/decoupling/drgrpo.json index fc7f4d60c..786cb2c18 100644 --- a/tests/golden/decoupling/drgrpo.json +++ b/tests/golden/decoupling/drgrpo.json @@ -110,6 +110,10 @@ "bridge", "--model-name", "SimpleNamespace", + "--attention-dropout", + "0.0", + "--hidden-dropout", + "0.0", "--num-layers", "2", "--hidden-size", @@ -245,7 +249,7 @@ "10", "--group-query-attention" ], - "miles_argv_sha256": "7ee222703a50f20c2c18b817596ddbca10140fe94645ca28acdcba05c871cc2b", + "miles_argv_sha256": "8545800520cd18d7ee951fbb8a06d012cbb6bbd39d70968e85f35c1b05d2c6e7", "name": "drgrpo", "placement": { "allow_cross_node_engine_tp": false, diff --git a/tests/golden/decoupling/elastic.json b/tests/golden/decoupling/elastic.json index 5664ac5bf..d37d4cd66 100644 --- a/tests/golden/decoupling/elastic.json +++ b/tests/golden/decoupling/elastic.json @@ -66,6 +66,10 @@ "bridge", "--model-name", "SimpleNamespace", + "--attention-dropout", + "0.0", + "--hidden-dropout", + "0.0", "--num-layers", "2", "--hidden-size", @@ -201,7 +205,7 @@ "120", "--group-query-attention" ], - "miles_argv_sha256": "c88e8996324d0261dc72f5b2949afc1e7eb81db916d3746ac76addea772020d1", + "miles_argv_sha256": "5aeca40e53d9f1b1a614a6a6ec814e5d037b354da3f448f04c1d10417751d135", "name": "elastic", "placement": { "allow_cross_node_engine_tp": false, diff --git a/tests/golden/decoupling/grpo_default.json b/tests/golden/decoupling/grpo_default.json index 61c611029..5235984f2 100644 --- a/tests/golden/decoupling/grpo_default.json +++ b/tests/golden/decoupling/grpo_default.json @@ -59,6 +59,10 @@ "bridge", "--model-name", "SimpleNamespace", + "--attention-dropout", + "0.0", + "--hidden-dropout", + "0.0", "--num-layers", "2", "--hidden-size", @@ -191,7 +195,7 @@ "10", "--group-query-attention" ], - "miles_argv_sha256": "cd9c51c2946a4e45bb63cd153ae64bc3ec0e3e998cf3d5e64fe6d4ed882c156b", + "miles_argv_sha256": "bb315aa0c82a03c2821899a3a376db84f09d30f1d32e2f6640833aec455d8634", "name": "grpo_default", "placement": { "allow_cross_node_engine_tp": false, diff --git a/tests/golden/decoupling/grpo_tis.json b/tests/golden/decoupling/grpo_tis.json index bca5a992a..9d65c992c 100644 --- a/tests/golden/decoupling/grpo_tis.json +++ b/tests/golden/decoupling/grpo_tis.json @@ -106,6 +106,10 @@ "bridge", "--model-name", "SimpleNamespace", + "--attention-dropout", + "0.0", + "--hidden-dropout", + "0.0", "--num-layers", "2", "--hidden-size", @@ -243,7 +247,7 @@ "10", "--group-query-attention" ], - "miles_argv_sha256": "06b407103acaff8218e5bab2309192be3af014b6ca635dd7c72f27cca28b1a2b", + "miles_argv_sha256": "70d04ad1b86c39f254ac2631b2be4a32558c00f4e0e7cc8c1767e03ea84b03ee", "name": "grpo_tis", "placement": { "allow_cross_node_engine_tp": false, diff --git a/tests/golden/decoupling/seq_adv_maxrl.json b/tests/golden/decoupling/seq_adv_maxrl.json index 801006738..29bd1d963 100644 --- a/tests/golden/decoupling/seq_adv_maxrl.json +++ b/tests/golden/decoupling/seq_adv_maxrl.json @@ -116,6 +116,10 @@ "bridge", "--model-name", "SimpleNamespace", + "--attention-dropout", + "0.0", + "--hidden-dropout", + "0.0", "--num-layers", "2", "--hidden-size", @@ -250,7 +254,7 @@ "10", "--group-query-attention" ], - "miles_argv_sha256": "8acfb1140db90c4d02fd4534526adcea04ca37a569125cee1952d3a697eb882d", + "miles_argv_sha256": "fd095a4b6d682aa8ba80b1953964b57575721859844217c78253847b7349a820", "name": "seq_adv_maxrl", "placement": { "allow_cross_node_engine_tp": false, diff --git a/tests/rl_cut_fakes.py b/tests/rl_cut_fakes.py index 8e9d097c3..1e74d4295 100644 --- a/tests/rl_cut_fakes.py +++ b/tests/rl_cut_fakes.py @@ -27,7 +27,7 @@ def named_parameters(self, model): def is_adapter(self, name): return "lora" in name - def export_optimizer(self, optimizer, named): + def export_optimizer(self, optimizer, named, *, on_device=False): # clones stay on their device by_id = {id(p): n for n, p in named} entries = {} for group in optimizer.param_groups: @@ -170,7 +170,7 @@ def __init__(self, coord=None, *, side_effect_free=True): super().__init__(coord) self.side_effect_free = side_effect_free - def export_optimizer(self, optimizer, named): + def export_optimizer(self, optimizer, named, *, on_device=False): # clones stay on their device from collections import defaultdict from yeto.rl.adapters.miles.cut_plugin import side_effect_free_state diff --git a/tests/rl_reshard_fakes.py b/tests/rl_reshard_fakes.py index 5d55ddbaf..cdf85a694 100644 --- a/tests/rl_reshard_fakes.py +++ b/tests/rl_reshard_fakes.py @@ -66,7 +66,7 @@ def named_parameters(self, model): def is_adapter(self, name): return "lora" in name - def export_optimizer(self, optimizer, named): + def export_optimizer(self, optimizer, named, *, on_device=False): # clones stay on their device optimizer.init_state() hyper = {k: v for k, v in optimizer.param_groups[0].items() if k != "params"} entries = {} diff --git a/tests/test_rl_cut_distopt_pp_ep.py b/tests/test_rl_cut_distopt_pp_ep.py index d687e8d38..308a20762 100644 --- a/tests/test_rl_cut_distopt_pp_ep.py +++ b/tests/test_rl_cut_distopt_pp_ep.py @@ -31,6 +31,9 @@ DPS_SOURCES = { "m5-c35702e": (os.environ.get("YETO_MILES_CHECKOUT", "/home/michael/work/miles-m3"), "c35702e"), "s25-distopt": (os.environ.get("YETO_MILES_DISTOPT_CHECKOUT", "/home/michael/work/miles-next"), "s25-distopt"), + # E1: + export_named_optimizer_state(device=None) on the image pin 44d41efd (not pushed) + "e1-device-export": (os.environ.get("YETO_MILES_DISTOPT_CHECKOUT", "/home/michael/work/miles-next"), + "yeto/s26-e1-device-export"), } @@ -355,3 +358,186 @@ def test_local_states_cover_with_fillers(): assert cp.range_problems(state, state) == [] other = {"entries": {"a": {**state["entries"]["a"], "start": 5}}} assert len(cp.range_problems(state, other)) == 2 + + +# -------------------------------------------------------------------------- +# E1: background cuts copy their state device-to-device first +# -------------------------------------------------------------------------- + + +@needs_dps +@pytest.mark.parametrize("layout", ["pp2dp1", "pp2dp2", "ep2"]) +def test_device_export_is_the_fork_host_export_once_on_the_host(layout): + ranks = _world(layout) + _train(ranks) + for rank in ranks: + named = list(rank.model[0].named_parameters()) + host = DPS.export_named_optimizer_state(rank.optimizer, named) + device = cp.to_host(cp._device_export(DPS)(rank.optimizer, named)) + assert cp.state_digest(device) == cp.state_digest(host) + assert device["format"] == host["format"] and list(device["entries"]) == list(host["entries"]) + for name, e in host["entries"].items(): + got = device["entries"][name] + assert {k: v for k, v in got.items() if k not in ("tensors", "scalars")} == \ + {k: v for k, v in e.items() if k not in ("tensors", "scalars")} + for part in ("tensors", "scalars"): + assert list(got[part]) == list(e[part]) + for key, t in e[part].items(): + assert got[part][key].dtype == t.dtype and torch.equal(got[part][key], t), (name, key) + + +def test_device_export_uses_the_fork_device_argument_when_it_has_one(): + calls = [] + + def export(optimizer, named, *, device="cpu"): + calls.append(device) + return {"entries": {}} + + cp._device_export(SimpleNamespace(export_named_optimizer_state=export))("opt", []) + assert calls == [None] + + +@needs_dps +@pytest.mark.parametrize("staged", [False, True], ids=["resident", "sleeps_after_cut"]) +@pytest.mark.parametrize("layout", ["pp2dp2", "ep2"]) +def test_background_device_capture_writes_the_inline_shard(tmp_path, layout, staged, monkeypatch): + """Same shard (file bytes, digests, N7.6 resume state) from the device path as from the + inline host path, although the trainer trains on before the writer copies to the host.""" + ranks = _world(layout) + _train(ranks) + inline = [cp.save_cut_shard(r, directory=str(tmp_path / "inline"), cut_id="c1") for r in ranks] + monkeypatch.setattr(cp, "_device_copy_plan", lambda actor, named: (True, "ok")) + if staged: # a sleeping trainer's copies go through the staging buffer (CPU stand-ins too) + real_stage = cp._stage + monkeypatch.setattr(cp, "_sleeps_after_cut", lambda actor: True) + monkeypatch.setattr(cp, "_stage", lambda snap: real_stage(snap, include_cpu=True)) + device = [] + for r, rank in enumerate(ranks): # one pending write per cut id per process: rank by rank + started = cp.save_cut_shard(rank, directory=str(tmp_path / "device"), cut_id="c1", background=True) + assert started["capture"] == {"mode": "device", "reason": "ok"}, started + for k, leaf in enumerate(_leaves(rank)): # the next step changes the live state first + leaf.train(2000 + 10 * r + k) + rank.opt_param_scheduler.num_steps = 16 + device.append(cp.finish_cut_shard(rank, cut_id="c1")) + assert [d.pop("writer") for d in device] == [{"mode": "process"}] * len(ranks) + assert cp._STAGING.free.is_set() and (cp._STAGING.buffer is not None or not staged) + assert device == inline + for name in sorted(p.name for p in (tmp_path / "inline").iterdir()): + assert (tmp_path / "device" / name).read_bytes() == (tmp_path / "inline" / name).read_bytes() + + +def test_background_capture_falls_back_to_the_host_and_says_why(tmp_path, monkeypatch): + from tests.rl_cut_fakes import make_rank, train_step + + rank = make_rank(0) + train_step(rank, torch.randn(3, 6, generator=torch.Generator().manual_seed(7))) + inline = cp.save_cut_shard(rank, directory=str(tmp_path / "inline"), cut_id="c1") + monkeypatch.setattr(cp, "_device_copy_plan", lambda actor, named: (False, "low_device_memory(free=1,need=2)")) + started = cp.save_cut_shard(rank, directory=str(tmp_path / "low"), cut_id="c1", background=True) + assert started["capture"] == {"mode": "host", "reason": "low_device_memory(free=1,need=2)"} + assert _without_writer(cp.finish_cut_shard(rank, cut_id="c1")) == inline + + real = cp._snapshot + + def oom(actor, backend, named, *, on_device=False): + if on_device: + raise torch.cuda.OutOfMemoryError("CUDA out of memory") + return real(actor, backend, named) + + monkeypatch.setattr(cp, "_device_copy_plan", lambda actor, named: (True, "ok")) + monkeypatch.setattr(cp, "_snapshot", oom) + monkeypatch.setattr(torch.cuda, "empty_cache", lambda: None) + started = cp.save_cut_shard(rank, directory=str(tmp_path / "oom"), cut_id="c1", background=True) + assert started["capture"] == {"mode": "host", "reason": "device_oom"} + assert _without_writer(cp.finish_cut_shard(rank, cut_id="c1")) == inline + + +def _without_writer(summary): + assert summary.pop("writer")["mode"] == "process", summary + return summary + + +def _one_rank_cut(tmp_path, name, *, background): + import random + + import numpy as np + + from tests.rl_cut_fakes import make_rank, train_step + + random.seed(5), np.random.seed(5), torch.manual_seed(5) # the shard holds the global RNG state + rank = make_rank(0) + train_step(rank, torch.randn(3, 6, generator=torch.Generator().manual_seed(7))) + if not background: + return cp.save_cut_shard(rank, directory=str(tmp_path / name), cut_id="c1") + cp.save_cut_shard(rank, directory=str(tmp_path / name), cut_id="c1", background=True) + return cp.finish_cut_shard(rank, cut_id="c1") + + +def test_a_dead_writer_process_falls_back_to_the_thread_and_says_why(tmp_path, monkeypatch): + """The writer process dies before it writes (its start fails): this rank's thread + writes the same shard and the summary records the fallback.""" + import multiprocessing + import os + from concurrent.futures import ProcessPoolExecutor + + inline = _one_rank_cut(tmp_path, "inline", background=False) + dead = ProcessPoolExecutor(max_workers=1, mp_context=multiprocessing.get_context("spawn"), + initializer=os._exit, initargs=(3,)) + monkeypatch.setattr(cp, "_WRITER", [dead]) + got = _one_rank_cut(tmp_path, "dead", background=True) + writer = got.pop("writer") + assert writer["mode"] == "thread" and "writer process failed" in writer["fallback"], writer + assert got == inline + assert cp._WRITER == [] # the broken pool is dropped; the next cut starts a new one + assert (tmp_path / "dead" / got["path"]).read_bytes() == (tmp_path / "inline" / got["path"]).read_bytes() + + +def test_a_partial_file_of_a_dead_writer_is_removed_before_the_fallback(tmp_path, monkeypatch): + inline = _one_rank_cut(tmp_path, "inline", background=False) + target = tmp_path / "partial" + target.mkdir() + (target / (inline["path"] + ".tmp")).write_bytes(b"half") + (target / inline["path"]).write_bytes(b"whole but unreported") + + class Broken: + def submit(self, *a, **k): + raise RuntimeError("writer died") + + def shutdown(self, **k): + pass + + monkeypatch.setattr(cp, "_WRITER", [Broken()]) + got = _one_rank_cut(tmp_path, "partial", background=True) + assert got.pop("writer") == {"mode": "thread", "fallback": "writer process failed: RuntimeError('writer died')"} + assert got == inline and not (target / (inline["path"] + ".tmp")).exists() + + +def test_the_writer_process_can_be_turned_off(tmp_path, monkeypatch): + inline = _one_rank_cut(tmp_path, "inline", background=False) + monkeypatch.setenv(cp.WRITER_PROCESS_ENV, "0") + got = _one_rank_cut(tmp_path, "thread", background=True) + assert got.pop("writer") == {"mode": "thread", "fallback": "disabled"} + assert got == inline + + +def test_device_copy_plan_needs_cuda_and_can_be_turned_off(monkeypatch): + named = [("lora_a", torch.nn.Parameter(torch.zeros(2)))] + assert cp._device_copy_plan(None, named) == (False, "not_cuda") + sleeper = SimpleNamespace(args=SimpleNamespace(offload_train=True)) # staged, below + assert cp._device_copy_plan(sleeper, named) == (False, "not_cuda") + monkeypatch.setenv(cp.DEVICE_COPY_ENV, "0") + assert cp._device_copy_plan(None, named) == (False, "disabled") + + +def test_pack_hands_the_state_over_in_one_buffer_and_unpacks_it_exactly(): + shared = torch.arange(6.0) + state = {"a": torch.randn(3, 4), "b": [torch.ones(2, dtype=torch.bfloat16), torch.tensor(7)], + "view": torch.arange(10.0)[2:5], "twice": (shared, shared), "empty": torch.zeros(0), "n": 3} + skeleton, layout, flat = cp._pack(state) + assert len(layout) == 3 # a, b[0], b[1]; the view, the shared storage and the empty one are not packed + assert flat.is_shared() and skeleton["n"] == 3 + back = cp._unpack(skeleton, layout, flat) + assert cp.state_digest(back) == cp.state_digest(state) + for key in ("a", "view", "empty"): + assert back[key].dtype == state[key].dtype and torch.equal(back[key], state[key]) + assert back["b"][0].untyped_storage().data_ptr() != flat.untyped_storage().data_ptr() # own copies diff --git a/tests/test_rl_engine_driver.py b/tests/test_rl_engine_driver.py index cf48504d1..b1bdefdb3 100644 --- a/tests/test_rl_engine_driver.py +++ b/tests/test_rl_engine_driver.py @@ -964,3 +964,52 @@ def test_first_optimizer_step_records_the_last_start_phase_once(tmp_path): assert names == list(START_PHASES) [summary] = [f for e, f in emitted if e == "trainer_start_summary"] assert summary["total_s"] == 6.0 and summary["missing"] == [] + + +def test_colocated_publish_offloaded_prepares_a_due_cut_before_offloading(tmp_path): + """E1: with round cuts, the trainer part of a due cut is written while the trainer + is still resident (before the offload for the publication), not after it.""" + engine = _engine() + engine.trainer.publish_offloaded = True + sync = LocalOnlySync(2) + order = [] + + class _Cuts: + def resume(self, driver): + return None + + def due(self, rollout_id, *, final=False): + return True + + def prepare(self, driver, *, rollout_id, state, final=False): + order.append(("prepare", rollout_id, state.policy_version, final, driver._trainer_offloaded)) + + def save(self, driver, *, rollout_id, final=False, **kw): + order.append(("save", rollout_id, final)) + return None + + sync.round_cuts = _Cuts() + driver = _driver(engine, sync, tmp_path) + driver.run() + if PROFILE == "ports-only": # publish resident: nothing to prepare + assert not any(o[0] == "prepare" for o in order) + return + assert [o for o in order if o[1] >= 1][:2] == [("prepare", 1, 1, False, False), ("save", 1, False)] + assert ("prepare", 2, 2, True, False) in order + phases = [e["phase"] for e in _events(tmp_path / "events.jsonl") if e["event"] == "rl_driver_phase"] + assert phases[4:8] == ["train", "sync", "checkpoint", "offload"] + + +def test_rollout_offload_after_generate_runs_in_its_own_phase(tmp_path): + """E1: Miles --offload-rollout frees the engine memory after generating; the + driver times it as its own phase so "generate" matches Miles' rollout_time.""" + engine = _engine() + calls = [] + engine.rollout.offloads_after_generate = True + engine.rollout.offload_after_generate = lambda: calls.append("offload_rollout") + driver = _driver(engine, LocalOnlySync(2), tmp_path) + driver.run() + assert calls == ["offload_rollout"] * 2 and engine.rollout.defer_offload is True + 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" diff --git a/tests/test_rl_fully_async_prompts.py b/tests/test_rl_fully_async_prompts.py index c92c68a39..1130228ff 100644 --- a/tests/test_rl_fully_async_prompts.py +++ b/tests/test_rl_fully_async_prompts.py @@ -87,7 +87,7 @@ def test_an_unexposed_stale_count_is_reported_unknown(executor, caplog): assert "reported unknown" in caplog.text -def test_cut_returns_every_untrained_prompt_and_stops_the_worker(executor): +def test_cut_records_every_untrained_prompt_and_keeps_the_worker_running(executor): source, ledger = executor.data_source, prompts.current() drawn = source.get_samples(5) # in flight / waiting in Miles' output buffer ledger.mark_trained(drawn[:2]) # trained this round @@ -102,18 +102,65 @@ def test_cut_returns_every_untrained_prompt_and_stops_the_worker(executor): _task_of(worker), loop).result() try: exported = export_in_flight(executor) + # E1: the worker keeps generating ahead across the cut + assert executor.generate_rollout._worker is not None + assert not executor.generate_rollout._worker.done() finally: loop.call_soon_threadsafe(loop.stop) - assert executor.generate_rollout._worker is None # the next call starts a fresh worker - assert source.offset == 5 # the cursor is not rewound: the prompts are in the buffer - assert sorted(g[0].group_index for g in source.buffer) == [2, 3, 4] - assert all(g[0].response == "" for g in source.buffer) # reset for a fresh generation - assert drawn[3][0].response == "partial" # a running task keeps its own sample + # the live run is untouched: same buffer, same in-flight samples + assert source.offset == 5 and [g[0].group_index for g in source.buffer] == [2] + assert drawn[3][0].response == "partial" + assert sorted(prompts.current().untrained) == [3, 4] + # the cut carries every drawn, untrained prompt, reset for a fresh generation assert exported["buffer_groups"] == 3 assert sorted({e["group_id"] for e in exported["entries"]}) == ["g2", "g3", "g4"] - # the resumed run draws exactly the untrained prompts again - assert sorted(g[0].group_index for g in source.get_samples(3)) == [2, 3, 4] + assert all(e["engine_state"]["miles_sample"]["response"] == "" for e in exported["entries"]) + + +def test_cut_reads_cursor_and_prompts_on_the_worker_loop(executor): + """The cursor and the untrained prompts come from one read on the loop the worker + draws on: a draw can never fall between them.""" + source = executor.data_source + for name in ("sample_offset", "epoch_id", "sample_group_index", "sample_index"): + setattr(source, name, 0) + real_get = prompts.current()._get_samples # the unwrapped DataSource.get_samples + + def get_samples(n): + groups = real_get(n) + source.sample_offset += n + source.sample_group_index += n + source.sample_index += n + return groups + + prompts.current()._get_samples = get_samples + loop = asyncio.new_event_loop() + thread = threading.Thread(target=loop.run_forever, daemon=True) + thread.start() + seen = [] + + async def drawing_worker(): + while True: + source.get_samples(1) # the ledger's wrapper (records the draw) + await asyncio.sleep(0) + + executor.generate_rollout._worker = asyncio.run_coroutine_threadsafe( + _start(drawing_worker), loop).result() + try: + for _ in range(20): + exported = export_in_flight(executor) + seen.append((exported["data_cursor"]["sample_offset"], exported["buffer_groups"])) + finally: + executor.generate_rollout._worker.get_loop().call_soon_threadsafe( + executor.generate_rollout._worker.cancel) + loop.call_soon_threadsafe(loop.stop) + # nothing is trained: every drawn prompt is in the cut, and the cursor is just past it + assert all(offset == groups for offset, groups in seen) + assert seen[-1][0] > seen[0][0] # the worker kept drawing between cuts + + +async def _start(fn): + return asyncio.create_task(fn()) async def _task_of(future): @@ -133,7 +180,7 @@ def test_cut_without_the_ledger_is_refused(): def test_overlap_round_report_carries_the_ledger_counts(monkeypatch): from yeto.rl.adapters.miles import rollout_meta_hook as hook - monkeypatch.setattr(hook, "build_metadata", lambda args, samples: { + monkeypatch.setattr(hook, "build_metadata", lambda args, samples, **_kw: { "completed": 2, "filtered": 0, "max_policy_age": 1}) sent = [] monkeypatch.setattr(hook, "put_to_sink", sent.append) diff --git a/tests/test_rl_meta_hook_sink_view.py b/tests/test_rl_meta_hook_sink_view.py new file mode 100644 index 000000000..d66646c35 --- /dev/null +++ b/tests/test_rl_meta_hook_sink_view.py @@ -0,0 +1,80 @@ +"""E1: one metadata-hook run reads each sink value once and looks the sink actor up once.""" + +from __future__ import annotations + +import sys +from types import SimpleNamespace + +import pytest + +from yeto.rl.adapters.miles import rollout_meta_hook as hook + + +class _Remote: + def __init__(self, fn): + self.remote = fn + + +class _FakeRay: + def __init__(self): + self.values = {"get_token": "yeto:4:" + "a" * 64, "get_limit": 1} + self.lookups = 0 + self.calls = [] + + def get_actor(self, name): + self.lookups += 1 + ray = self + + class Actor: + def __getattr__(self, method): + def call(*args): + ray.calls.append(method) + if method.startswith("set_"): + ray.values["get_" + method[4:]] = args[0] + return None + return ray.values.get(method) + return _Remote(call) + + return Actor() + + def get(self, value): + return value + + +@pytest.fixture() +def fake_ray(monkeypatch): + ray = _FakeRay() + monkeypatch.setitem(sys.modules, "ray", SimpleNamespace(get_actor=ray.get_actor, get=ray.get)) + monkeypatch.setenv(hook.META_SINK_ENV, "ray:yeto_rollout_meta") + return ray + + +def test_reads_inside_one_view_hit_the_actor_once(fake_ray): + with hook.one_sink_view(): + for _ in range(5): + assert hook.current_policy_token() == "yeto:4:" + "a" * 64 + assert hook.max_policy_age_override() == 1 + with hook.one_sink_view(): # nested: the same view + hook.current_policy_token() + assert fake_ray.calls == ["get_token", "get_limit"] and fake_ray.lookups == 1 + + +def test_outside_a_view_every_read_goes_to_the_actor(fake_ray): + hook.current_policy_token() + hook.current_policy_token() + assert fake_ray.calls == ["get_token", "get_token"] and fake_ray.lookups == 2 + + +def test_a_write_inside_the_view_is_read_back(fake_ray): + with hook.one_sink_view(): + assert hook.max_policy_age_override() == 1 + hook.put_max_policy_age(0) + assert hook.max_policy_age_override() == 0 + assert hook.max_policy_age_override() == 0 + + +def test_the_hook_limit_is_read_once_per_run(fake_ray): + args = SimpleNamespace(yeto_rl_max_policy_age=2) + with hook.one_sink_view(): + assert [hook._carry.max_policy_age(args) for _ in range(64)] == [1] * 64 + assert fake_ray.calls.count("get_limit") == 1 diff --git a/tests/test_rl_miles_adapter_config.py b/tests/test_rl_miles_adapter_config.py index 503fe2e28..7cbe30d98 100644 --- a/tests/test_rl_miles_adapter_config.py +++ b/tests/test_rl_miles_adapter_config.py @@ -795,3 +795,21 @@ def test_dynamic_batching_refused_for_bshd_gated_delta_net_models(): recompute_only = dataclasses.replace( cfg, model_recipe=recipe, train_memory=rc.TrainMemoryConfig(activation_recompute=True)) assert "--recompute-granularity" in mc.translate_run_config(recompute_only, AlgorithmSpec()).argv + + +def test_gradient_accumulation_fusion_is_opt_in(): + # S26 E0.1: the Miles recipes train with the fusion on; yeto keeps it off by default + off = " ".join(mc.translate_run_config(make_config(), AlgorithmSpec()).argv) + assert "--no-gradient-accumulation-fusion" in off + cfg = dataclasses.replace(make_config(), train_memory=rc.TrainMemoryConfig(gradient_accumulation_fusion=True)) + assert "--no-gradient-accumulation-fusion" not in mc.translate_run_config(cfg, AlgorithmSpec()).argv + assert rc._train_memory(SimpleNamespace(seq_len=8, rl_gradient_accumulation_fusion=True)) \ + .gradient_accumulation_fusion + + +def test_generic_recipe_trains_without_dropout(): + cfg = make_config() + if cfg.model_recipe.name != rc.RECIPE_GENERIC: + cfg = sub(cfg, "model_recipe", name=rc.RECIPE_GENERIC) + joined = " ".join(mc.translate_run_config(cfg, AlgorithmSpec()).argv) + assert "--attention-dropout 0.0 --hidden-dropout 0.0" in joined diff --git a/tests/test_rl_miles_adapter_trainer_publish.py b/tests/test_rl_miles_adapter_trainer_publish.py index 6e4ccbe36..25456531d 100644 --- a/tests/test_rl_miles_adapter_trainer_publish.py +++ b/tests/test_rl_miles_adapter_trainer_publish.py @@ -592,3 +592,22 @@ async def miles_update_weights(args, actor, executor, ctl, **kw): ) result = pub.publish(state()) assert calls == [] and result.members == frozenset() and pub.last_engine_checksums is None + + +# E1: publication checks at the round-cut interval +def test_an_unverified_publication_skips_only_the_engine_read_back(): + from dataclasses import replace as dc_replace + + class Unverified: # a verified state's identity, marked unverified + def __init__(self, s): + self._s, self.verified, self.policy_version = s, False, s.policy_version + + def __getattr__(self, name): + return getattr(self._s, name) + + ctl = Controller([Engine("http://e0")], checksum_ok=False) # a read-back would fail + pub, calls = publisher(ctl, export=lambda: state()) + pub.publish(Unverified(state())) # no read-back on an unverified round + assert pub.last_engine_checksums is None and calls == ["update_weights"] + with pytest.raises(PublicationError, match="checksum"): + pub.publish(state()) # a verified (spot-check) round still fails on a bad read-back diff --git a/tests/test_rl_miles_carry_over.py b/tests/test_rl_miles_carry_over.py index e94f3d267..01c13a9bb 100644 --- a/tests/test_rl_miles_carry_over.py +++ b/tests/test_rl_miles_carry_over.py @@ -348,3 +348,47 @@ def test_trajectory_records_carry_start_round_and_segments(sink): for r in recs: assert not [p for p in validate_trajectory_reward({**r, "rollout_id": 4, "policy_version": 4}) if "started_rollout_id" in p or "policy_versions" in p] + + +def test_deferred_cross_version_estimate_arrives_as_a_late_record(sink, tmp_path): + """E1 (fully_async): the estimate runs in the background and reaches the sink as + a late record of its rollout id; the round's fields only say it was deferred.""" + import glob + import json + import time + + engine = FakePartialEngine() + args = SimpleNamespace(yeto_rl_max_policy_age=1, tis_clip=2.0, tis_clip_low=0.0) + s = Sample(index=0, group_index=0) + engine.generate(s, 3, 4, finish=False) + engine.generate(s, 4, 2, finish=True, reward=1.0) + current = [0.5, -0.5, 0.0, 0.4, -0.5, -0.5] + hook.put_policy_token(tok(4), sink) + fields = hook.carry_over_fields(args, [s], sink, lambda _s: current, defer_estimate_for=4) + assert fields["cross_version_estimate"] == "deferred" + assert "cross_version_truncated_fraction" not in fields + handle = handle_from_metadata({**round_metadata(args, [[s]], sink), **fields}, rollout_id=4, + policy_version=4, policy_hash=HASHES[4], data_pack=None) + assert handle.cross_version_truncated_fraction is None + deadline = time.monotonic() + 5 + late = [] + while not late and time.monotonic() < deadline: + late = [json.loads(open(p).read()) for p in glob.glob(f"{tmp_path}/*.json")] + late = [r for r in late if r.get("schema") == hook.LATE_META_SCHEMA] + time.sleep(0.02) + assert late and late[0]["rollout_id"] == 4 + assert math.isclose(late[0]["cross_version_truncated_fraction"], 0.5) + + +def test_driver_governs_on_a_late_rollout_estimate(tmp_path): + """E1: a fully_async round's estimate arrives after its train step; the policy-age + governor still sees it.""" + from tests.test_rl_policy_age import _driver + + driver = _driver(tmp_path, 1) + calls, asked = [], [] + driver.rollout.set_max_policy_age = calls.append + driver.rollout.late_cross_version_fraction = lambda rid: asked.append(rid) or 0.9 + driver.run() + events = [json.loads(l) for l in (tmp_path / "e.jsonl").read_text().splitlines()] + assert asked and [e for e in events if e["event"] == "rl_policy_age_fallback"] and calls == [0] diff --git a/tests/test_rl_miles_train_profile.py b/tests/test_rl_miles_train_profile.py new file mode 100644 index 000000000..0325ad8ea --- /dev/null +++ b/tests/test_rl_miles_train_profile.py @@ -0,0 +1,85 @@ +"""E1: --rl-profile-trainer is observation only and reaches the trainer.""" + +from __future__ import annotations + +import gc +import json +import time + +from yeto.rl.adapters.miles import train_profile + + +def test_profiled_returns_the_call_result_and_prints_one_record(capsys): + def step(x, *, y): + junk = [[i] for i in range(20000)] + gc.collect() + time.sleep(0.03) + return x + y, len(junk) + + assert train_profile.profiled(step, 2, y=3) == (5, 20000) + (line,) = [l for l in capsys.readouterr().out.splitlines() if l.startswith("yeto train profile ")] + record = json.loads(line[len("yeto train profile "):]) + assert record["wall_s"] >= 0.03 and record["gc_pause_s"] > 0 and record["gc_collections"] + assert record["samples"] >= 1 and any("step" in label for label, _ in record["main_top"]) + + +def test_the_flag_reaches_the_launcher_args_and_the_eval_spec(): + from tools.eval.efficiency import spec + + sp = spec.get_spec("qwen3-0.6b", checkpoint="every-round", profile_trainer=True) + assert " --rl-profile-trainer" in spec.yeto_miles_launcher_args(sp, "async", "r1") + assert spec.control_fingerprint(sp) == spec.control_fingerprint(spec.get_spec("qwen3-0.6b")) + + +def test_the_flag_is_forwarded_to_the_island_and_set_on_the_miles_args(monkeypatch): + import argparse + + from yeto.rl.engine import phase_profile + + # the sampler would stay on for the rest of the test process and print on every + # driver phase (it broke the JSON output of test_rl_reconfig_e1 in the full suite) + monkeypatch.setattr(phase_profile, "enable", lambda: None) + from types import SimpleNamespace + + from yeto.launch import rl_ports + from yeto.rl.adapters.miles import island_entry + from yeto.rl.engine.run_flags import ISLAND, add_rl_flags + + _prelude, flags = rl_ports._ports_infra_flags(argparse.Namespace(rl_profile_trainer=True)) + assert " --rl-profile-trainer" in flags + parser = argparse.ArgumentParser() + add_rl_flags(parser, ISLAND) + for action in parser._actions: # parse just this flag + action.required = False + island_args, _ = parser.parse_known_args(["--rl-profile-trainer"]) + miles_args = SimpleNamespace() + island_entry.apply_ports_infra_switches(island_args, miles_args) + assert miles_args.yeto_rl_profile_trainer is True + + +def test_phase_profile_prints_each_finished_phase(capsys, monkeypatch): + from yeto.rl.engine import phase_profile + + monkeypatch.setitem(phase_profile._STATE, "thread", None) + phase_profile.on_phase("train", 0) # disabled: nothing + import threading + + phase_profile.enable() # from this thread; the phases come from another one + + def busy_driver(): + phase_profile.on_phase("train", 1) + end = time.perf_counter() + 0.05 + while time.perf_counter() < end: + sum(range(1000)) + phase_profile.on_phase("sync", 1) + + driver = threading.Thread(target=busy_driver) + driver.start() + driver.join() + (line,) = [l for l in capsys.readouterr().out.splitlines() if l.startswith("yeto phase profile ")] + record = json.loads(line[len("yeto phase profile "):]) + assert record["phase"] == "train" and record["rollout_id"] == 1 and record["wall_s"] >= 0.05 + assert record["samples"] >= 1 and any("busy_driver" in label for label, _ in record["top"]), record + sampler = phase_profile._STATE["thread"] + sampler.phase = None + monkeypatch.setitem(phase_profile._STATE, "thread", None) # off for later tests diff --git a/tests/test_rl_ports_image.py b/tests/test_rl_ports_image.py index 15585c2c3..c6b755fc7 100644 --- a/tests/test_rl_ports_image.py +++ b/tests/test_rl_ports_image.py @@ -631,3 +631,33 @@ def test_modal_public_image_config_has_no_login(monkeypatch, no_login): state = fake_modal(monkeypatch) mr.ModalOps("img-test").define(cfg) assert state["images"][0].calls[0] == ("from_registry", (cfg.image_ref,), {}) + + +def test_gpu_model_pin_refuses_an_a10g_for_an_a10(): + """E1: Modal "A10G" containers are an NVIDIA A10 or an A10G; --modal-gpu-name pins one.""" + from yeto.gpu_spec import parse_gpu_spec + + mr.check_gpu_model("NVIDIA A10", ["NVIDIA A10"] * 2, 2) + for names, count in ((["NVIDIA A10G"], 1), (["NVIDIA A10", "NVIDIA A10G"], 2), (["NVIDIA A10"], 2)): + with pytest.raises(RuntimeError, match="--modal-gpu-name"): + mr.check_gpu_model("NVIDIA A10", names, count) + cfg = _cfg(gpu="A10G", gpus_per_node=1, gpu_name="NVIDIA A10") + assert mr.ModalIslandConfig.from_json(cfg.to_json()).gpu_name == "NVIDIA A10" + old = json.loads(_cfg().to_json()) + old.pop("gpu_name") + assert mr.ModalIslandConfig.from_json(json.dumps(old)).gpu_name is None + args = _cli(("--modal-gpu-name", "NVIDIA A10")) + assert args.modal_gpu_name == "NVIDIA A10" + with pytest.raises(ValueError, match="--modal-gpu-name applies only to Modal"): + launcher.require_modal_for_gpu_exact(args, parse_gpu_spec("modal:1xa10g,aws:1xa100@us-east-1")) + + +def test_the_island_config_carries_the_gpu_model_pin(monkeypatch, login): + from yeto.gpu_spec import parse_gpu_spec + + args = _cli(("--modal-gpu-name", "NVIDIA A10")) + _prepare_rl_args(args) + task = _island_task(args, monkeypatch) + args.cluster_prefix = "img-test" + cfg = build_modal_island_config(args, parse_gpu_spec("modal:1xa10g")[0], 0, task, "1.2.3.4:5") + assert cfg.gpu_name == "NVIDIA A10" diff --git a/tests/test_rl_publish_fastpath.py b/tests/test_rl_publish_fastpath.py index bc88ee051..dbb99e4fb 100644 --- a/tests/test_rl_publish_fastpath.py +++ b/tests/test_rl_publish_fastpath.py @@ -257,3 +257,164 @@ def test_resident_state_type_guards(): d = PolicyDigest("a" * 64, "b" * 64, 8, ((NAMES[0], (2, 1)),)) with pytest.raises(ValueError): TrainerResidentState("full", REV, CFG, "c" * 64, 0, d, lambda v: None) + + +# E1: the digest path stages the tensors in one reusable host buffer +def _mixed_tensors(seed=0): + g = torch.Generator().manual_seed(seed) + return { + "b.bf16": torch.randn(8, 3, generator=g).to(torch.bfloat16), + "a.f32_transposed": torch.randn(5, 7, generator=g).t(), + "c.f16": torch.randn(4, generator=g).to(torch.float16), + "d.scalar": torch.tensor(1.5), + } + + +def test_staged_copies_are_the_per_tensor_copies_bit_for_bit(): + from yeto.rl.engine.policy_digest import HostStaging, stage_f32_on_host + + tensors = _mixed_tensors() + plain = stage_f32_on_host(tensors) + staging = HostStaging(pin=False) + staged = stage_f32_on_host(tensors, staging) + assert list(staged) == sorted(tensors) == list(plain) + for name, value in tensors.items(): + want = value.detach().to("cpu", torch.float32).contiguous() + for got in (plain[name], staged[name]): + assert got.dtype == torch.float32 and got.is_contiguous() and got.shape == want.shape + assert torch.equal(got.view(torch.int32) if got.dim() else got, want.view(torch.int32) if want.dim() else want) + kw = dict(base_model_revision=REV, lora_config_hash=CFG, layout_hash="0" * 64) + a, b = digest_canonical_tensors(staged, **kw), digest_canonical_tensors(plain, **kw) + assert (a.policy_tensor_hash, a.payload_hash, a.payload_bytes, a.specs) == \ + (b.policy_tensor_hash, b.payload_hash, b.payload_bytes, b.specs) + # the buffer is reused (and grown) by the next call; each call's values are its own + bigger = {**_mixed_tensors(1), "e.more": torch.randn(64)} + again = stage_f32_on_host(bigger, staging) + assert all(torch.equal(again[n], stage_f32_on_host(bigger)[n]) for n in bigger) + + +@pytest.mark.parametrize("staged", [False, True]) +def test_staging_names_the_first_non_finite_tensor(staged): + from yeto.rl.engine.policy_digest import HostStaging, NonFiniteTensor, stage_f32_on_host + + tensors = _mixed_tensors() + tensors["c.f16"][1] = float("inf") + tensors["b.bf16"][0, 0] = float("nan") + with pytest.raises(NonFiniteTensor, match="'b.bf16' contains NaN or Inf"): + stage_f32_on_host(tensors, HostStaging(pin=False) if staged else None) + + +def test_export_digest_is_stable_across_rounds_with_the_reused_buffer(): + ranks, group, ps = make() + for r in ranks: + r.train_step() + first = ps.export_digest(policy_version=2).policy_tensor_hash() + assert ps.export_digest(policy_version=2).policy_tensor_hash() == first == \ + ps.export(policy_version=2).policy_tensor_hash() + + +# E1: sync colocated mode hashes the staged policy while the trainer offloads +def test_deferred_digest_is_the_eager_digest_and_is_collected_once(): + ranks, group, ps = make() + for r in ranks: + r.train_step() + eager = ps.export_digest(policy_version=2) + group.calls.clear() + deferred = ps.export_digest(policy_version=2, deferred=True) + assert group.calls == [sp.EXPORT_DIGEST] and not deferred.digest.is_resolved + assert deferred.tensor_names == eager.tensor_names and deferred.layout_hash == eager.layout_hash + for r in ranks: # the trainer trains on (or offloads): the staged copy is hashed + r.train_step() + assert deferred.policy_tensor_hash() == eager.policy_tensor_hash() + assert _payload_of(deferred) == _payload_of(eager) + assert deferred.digest.to_wire() == {**eager.digest.to_wire(), "hash_seconds": deferred.digest.hash_seconds} + assert group.calls == [sp.EXPORT_DIGEST, sp.COLLECT_DIGEST] + deferred.policy_tensor_hash() + assert group.calls == [sp.EXPORT_DIGEST, sp.COLLECT_DIGEST] # resolved once + + +def test_a_new_export_waits_for_the_deferred_hash_of_the_reused_buffer(): + ranks, group, ps = make() + first = ps.export_digest(policy_version=1) + deferred = ps.export_digest(policy_version=1, deferred=True) + for r in ranks: + r.train_step() + other = ps.export_digest(policy_version=2) # overwrites the staging buffer + assert other.policy_tensor_hash() != first.policy_tensor_hash() + assert deferred.policy_tensor_hash() == first.policy_tensor_hash() + + +def test_collecting_another_version_is_refused(): + ranks, group, ps = make() + ps.export_digest(policy_version=1, deferred=True) + with pytest.raises(Exception, match="v1, not v3"): + ps._collect_digest(3) + + +def test_the_driver_defers_only_for_a_colocated_offloading_trainer(monkeypatch): + ranks, group, ps = make() + driver = _Driver(ps, fast=False) + from yeto.rl.engine.driver import defer_digest + + assert defer_digest(driver) is False + driver.colocated = driver.publish_offloaded = True + state = IslandDriver.export_local_resident(driver, policy_version=0) + assert not state.digest.is_resolved + monkeypatch.setenv("YETO_RL_DEFERRED_DIGEST", "0") + assert defer_digest(driver) is False + + +# E1: which publications are verified (policy hash + engine read-back) +def _verify_driver(ps, *, every, published, cut_due=lambda v: False, all_rounds=False): + from yeto.rl.engine import driver as driver_mod + + sync = SimpleNamespace(round_cuts=SimpleNamespace(every=every), cut_due=cut_due) + return SimpleNamespace(policy_state=ps, sync=sync, published_version=published, + verify_every_round=all_rounds), driver_mod.verify_publication + + +def test_cut_first_and_every_k_rounds_are_verified_and_the_switch_verifies_all(): + ranks, group, ps = make() + d, verify = _verify_driver(ps, every=4, published=None) + assert verify(d, 1) # the first publication + d, verify = _verify_driver(ps, every=4, published=1, cut_due=lambda v: v == 6) + assert [v for v in range(2, 10) if verify(d, v)] == [4, 6, 8] # every 4th and a due cut + d, verify = _verify_driver(ps, every=4, published=1, all_rounds=True) + assert all(verify(d, v) for v in range(2, 10)) # --rl-verify-every-round + d, verify = _verify_driver(ps, every=1, published=1) + assert all(verify(d, v) for v in range(2, 10)) # --rl-cut-every 1: every round + d, verify = _verify_driver(SimpleNamespace(), every=4, published=1) + assert verify(d, 3) # a policy state that cannot skip: every round + + +def test_an_unverified_round_still_checks_the_trainer_holds_its_policy(): + ranks, group, ps = make() + for r in ranks: + r.train_step() + verified = ps.export_digest(policy_version=2) + group.calls.clear() + unverified = ps.export_identity(policy_version=3) + assert sp.EXPORT_DIGEST not in group.calls # no export, no hash + assert unverified.verified is False and unverified.tensor_names == verified.tensor_names + assert unverified.policy_tensor_hash() != verified.policy_tensor_hash() + assert len(unverified.policy_tensor_hash()) == 64 + unverified.check_holds() + _step(ranks) + with pytest.raises(WeightsChanged): + unverified.check_holds() + + +def test_an_unverified_publication_is_never_committed_as_a_cut(): + from yeto.rl.engine.controller import RebuildRefused + from yeto.rl.engine.round_cut import CutSource + + ranks, group, ps = make() + ps.export_digest(policy_version=1) + unverified = ps.export_identity(policy_version=2) + source = CutSource.__new__(CutSource) + source.driver_ref = SimpleNamespace(published_version=2, published_state=unverified, at_safe_point=True, + expected_token="yeto:2:" + "0" * 64) + with pytest.raises(RebuildRefused, match="not verified"): + source.context("r000002-x") + with pytest.raises(RebuildRefused, match="not verified"): + source.context("r000002-x", pending=(2, unverified)) diff --git a/tests/test_rl_recovery_abstraction.py b/tests/test_rl_recovery_abstraction.py index 6064e8e1a..a6666639e 100644 --- a/tests/test_rl_recovery_abstraction.py +++ b/tests/test_rl_recovery_abstraction.py @@ -12,7 +12,7 @@ import torch from test_rl_multinode_round_cut import _Driver, _trainer from test_rl_reconfig_recovery import _ctl -from test_rl_resume_from_checkpoint import _cp, _Ledger, _rctl, _store, _trained_island, _Vol +from test_rl_resume_from_checkpoint import _cp, _cut, _Ledger, _rctl, _store, _trained_island, _Vol from tests.rl_cut_fakes import make_rank, params, train_step from yeto.cloud import modal_ckpt_store as mcs @@ -87,7 +87,7 @@ def test_preempt_periodic_and_manual_write_the_same_cut_format_and_resume(tmp_pa rank_a, ctl_a, driver_a, cp_a = _trained_island(tmp_path, store, 2) driver_a.sync = SimpleNamespace(outer_version=2) - periodic = cp_a.save(driver_a, rollout_id=2) + periodic = _cut(cp_a, driver_a, rollout_id=2) saver = s.InFlightSaver(lambda t: {"entries": [], "buffer_groups": 0}, str(tmp_path / "inflight"), store=_store(tmp_path / "inflight", vol), printer=lambda line: None) @@ -133,7 +133,7 @@ def test_resume_after_a_preemption_continues_from_the_last_periodic_cut(tmp_path the newest complete round cut (LATEST) with the same steps.""" store = _store(tmp_path / "store") rank_a, ctl_a, driver_a, cp_a = _trained_island(tmp_path, store, 2) - periodic = rc.ConsistentCut.from_record(cp_a.save(driver_a, rollout_id=2)[rc.CUT_KEY]) + periodic = rc.ConsistentCut.from_record(_cut(cp_a, driver_a, rollout_id=2)[rc.CUT_KEY]) saver = s.InFlightSaver(lambda t: {"entries": []}, str(tmp_path / "inflight"), printer=lambda line: None) saver.rehearse(2) _reclaim_save(saver) @@ -152,7 +152,7 @@ def test_same_save_code_on_another_backend_writes_the_same_cut(tmp_path): for name, vol in (("local", None), ("modal", _Vol())): store = _store(tmp_path / name / "store", vol) _, _, driver, cp = _trained_island(tmp_path / name, store, 2) - out = cp.save(driver, rollout_id=2) + out = _cut(cp, driver, rollout_id=2) record = dict(out[rc.CUT_KEY]) record["checkpoint_id"] = record["checkpoint_id"].split("-")[0] # incarnation differs records.append((out["store_synced"], record)) diff --git a/tests/test_rl_resume_from_checkpoint.py b/tests/test_rl_resume_from_checkpoint.py index cf6b44a44..a24480a58 100644 --- a/tests/test_rl_resume_from_checkpoint.py +++ b/tests/test_rl_resume_from_checkpoint.py @@ -5,6 +5,7 @@ from __future__ import annotations import json +from pathlib import Path from types import SimpleNamespace import pytest @@ -18,6 +19,7 @@ from yeto.rl.adapters.miles.rebuild_wiring import CutSource from yeto.rl.adapters.miles.round_cut import POINTER, RoundCutCheckpoint, RoundCutError from yeto.rl.engine.cut import CutError +from yeto.rl.engine.recovery import CUT_KEY from yeto.rl.engine import resume as rs from yeto.rl.engine.bridges import LocalOnlySync from yeto.rl.engine.controller import STORE_MANIFEST @@ -68,6 +70,19 @@ def _cp(ctl, driver, trainer, ledger, *, backend_fingerprint="miles@0af62f4d", * return RoundCutCheckpoint(controller=ctl, source=source, trainer=trainer, **kw) +def _cut(cp, driver, **kw): + """Save a cut and complete it: a periodic cut (also the store's first, E1) commits in + the background; the outcome of that commit is folded into the returned pointer.""" + info = cp.save(driver, **kw) + if info is None or info["store_commit"] != "deferred": + return info + done = cp.join_commit(driver) + assert done["ok"] is True, done + pointer = cp.pointer() + assert pointer["cut_id"] == info["cut_id"] + return {**info, "store_synced": True, "cut_bytes": done["cut_bytes"], CUT_KEY: pointer[CUT_KEY]} + + def _trained_island(tmp_path, store, rounds, *, name="a", seed=3): rank = make_rank(0) g = torch.Generator().manual_seed(seed) @@ -185,6 +200,101 @@ def test_cuts_without_latest_are_refused(tmp_path): _rctl(tmp_path, "a", store) +@pytest.mark.parametrize("layout", ["unmarked", "one_marked_one_not", "stray_file", "mark_of_another_cut"]) +def test_cuts_without_latest_are_refused_unless_all_are_a_marked_first_cut(tmp_path, layout): + """E1: only a store whose every cut carries the first-cut mark (never committed) + is started over; any other store with cuts and no LATEST is refused, untouched.""" + store = _store(tmp_path / "store") + cuts = store.root / rs.ROUND_CUTS + (cuts / "r000004-x").mkdir(parents=True) + if layout == "one_marked_one_not": + rs.mark_first_cut(store.root, "r000004-x") + (cuts / "r000008-x").mkdir() + elif layout == "stray_file": + rs.mark_first_cut(store.root, "r000004-x") + (cuts / "notes.txt").write_text("?") + elif layout == "mark_of_another_cut": + rs.mark_first_cut(store.root, "r000008-x") + before = sorted(p.name for p in cuts.iterdir()) + with pytest.raises(rs.StoreIntegrityError, match="no LATEST"): + _rctl(tmp_path, "a", store) + assert sorted(p.name for p in cuts.iterdir()) == before + + +def _first_cut_in_flight(tmp_path, store, every=4): + """An island at rollout ``every`` whose first (periodic) cut was captured, its + shards written, and then the process died before the store commit (no join).""" + rank, ctl, driver, cp = _trained_island(tmp_path, store, every) + cp.every = every + info = cp.save(driver, rollout_id=every) + assert info["store_commit"] == "deferred" + cp.trainer.finish_cut(cp._deferred["pending"]) # the writer finished; nothing committed + return rank, ctl, driver, cp, info + + +def test_first_periodic_cut_with_interval_n_commits_in_the_background(tmp_path): + """E1: with --rl-cut-every 4 the store's first cut (rollout 4) is in the steady + rounds: it is deferred like the others, marked until its LATEST is committed.""" + store = _store(tmp_path / "store", _Vol()) + rank, ctl, driver, cp = _trained_island(tmp_path, store, 4) + cp.every = 4 + calls = [] + start = cp.trainer.start_cut + cp.trainer.start_cut = lambda **kw: (calls.append(rs.first_cut_mark(store.root, kw["context"].cut_id) + .is_file()), start(**kw))[1] + info = cp.save(driver, rollout_id=4) + assert info["store_commit"] == "deferred" and info["total_s"] is not None + assert calls == [True] # marked before the shards were started + assert rs.read_latest(store.root) is None and cp.pointer() is None + done = cp.join_commit(driver) + assert done["ok"] is True and rs.read_latest(store.root)["cut_id"] == info["cut_id"] + assert not list((store.root / rs.ROUND_CUTS).glob(f"*{rs.FIRST_CUT_MARK}")) + assert _rctl(tmp_path, "b", store).restored["latest"]["cut_id"] == info["cut_id"] + + +def test_a_later_cut_is_not_marked(tmp_path): + store = _store(tmp_path / "store", _Vol()) + rank, ctl, driver, cp = _trained_island(tmp_path, store, 2) + _cut(cp, driver, rollout_id=2) + train_step(rank, torch.randn(3, 6, generator=torch.Generator().manual_seed(4))) + driver.local_step = 3 + driver.published_version = 3 + cp.source.ledger = _Ledger(3) + info = cp.save(driver, rollout_id=3) + assert info["store_commit"] == "deferred" + assert not rs.first_cut_mark(store.root, info["cut_id"]).exists() + assert cp.join_commit(driver)["ok"] is True + + +def test_a_run_dead_during_its_first_cut_commit_starts_over_at_zero(tmp_path): + store = _store(tmp_path / "store", _Vol()) + _rank, _ctl, _driver, _cp_a, info = _first_cut_in_flight(tmp_path, store) + cuts = store.root / rs.ROUND_CUTS + assert (cuts / info["cut_id"]).is_dir() and rs.first_cut_mark(store.root, info["cut_id"]).is_file() + assert rs.read_latest(store.root) is None + ctl_b = _rctl(tmp_path, "b", store) # a fresh machine + assert ctl_b.restored is None and sorted(ctl_b.discarded_first_cut) == sorted( + [f"{rs.ROUND_CUTS}/{info['cut_id']}", f"{rs.ROUND_CUTS}/{info['cut_id']}{rs.FIRST_CUT_MARK}"]) + assert list(cuts.iterdir()) == [] + journal = [json.loads(line) for line in (ctl_b.state_dir / rs.RESUME_JOURNAL).read_text().splitlines()] + assert journal[-1]["action"] == "discard_unfinished_first_cut" + rank_b = make_rank(9) + driver_b = _Driver(rank_b) + assert _cp(ctl_b, driver_b, _trainer(rank_b), _Ledger(0)).resume(driver_b) is None # rollout 0 + + +def test_an_in_place_restart_after_a_dead_first_cut_moves_the_state_dir_aside(tmp_path): + """The dead run's background commit may have pointed its local state dir at the + cut before dying; a restart on the same disk must not resume from that pointer.""" + store = _store(tmp_path / "store", _Vol()) + _rank, ctl, _driver, _cp_a, info = _first_cut_in_flight(tmp_path, store) + (ctl.state_dir / POINTER).write_text(json.dumps({"cut_id": info["cut_id"], "next_rollout_id": 4})) + ctl_b = rs.ResumeController(state_dir=ctl.state_dir, store=store) + assert ctl_b.discarded_first_cut and not (ctl_b.state_dir / POINTER).exists() + assert list(ctl.state_dir.parent.glob(f"{ctl.state_dir.name}.prev-*")) + assert ctl_b.incarnation["index"] == 0 # the dead run's incarnations went aside too + + # ---------------------------------------------------------------- round cut: save -> kill -> resume def test_save_kill_resume_without_elastic_continues_bitwise(tmp_path): """tasks 5.2: fake trainer, save -> the machine is gone -> a fresh trainer with @@ -192,8 +302,11 @@ def test_save_kill_resume_without_elastic_continues_bitwise(tmp_path): vol = _Vol() store = _store(tmp_path / "store", vol) rank_a, ctl_a, driver_a, cp_a = _trained_island(tmp_path, store, 2) - info = cp_a.save(driver_a, rollout_id=2) - assert info["store_synced"] and info["cut_bytes"] > 0 and info["lr_at_next_round"] == [5e-6] + info = _cut(cp_a, driver_a, rollout_id=2) + # E1: the store's first cut commits in the background too (marked until LATEST is committed) + assert info["store_synced"] is True and info["store_commit"] == "deferred" + assert not list((store.root / rs.ROUND_CUTS).glob(f"*{rs.FIRST_CUT_MARK}")) + assert info["cut_bytes"] > 0 and info["lr_at_next_round"] == [5e-6] assert info["run_fingerprint"] == FP and info["consumed_groups"] == 8 and info["incarnation_index"] == 0 assert vol.commits >= 1 and rs.read_latest(store.root)["cut_id"] == info["cut_id"] @@ -220,7 +333,7 @@ def test_save_kill_resume_without_elastic_continues_bitwise(tmp_path): def test_resume_refuses_a_changed_config_unless_allowed(tmp_path): store = _store(tmp_path / "store") _r, _c, driver_a, cp_a = _trained_island(tmp_path, store, 2) - cp_a.save(driver_a, rollout_id=2) + _cut(cp_a, driver_a, rollout_id=2) rank_b = make_rank(9) ctl_b = _rctl(tmp_path, "b", store) cp_b = _cp(ctl_b, _Driver(rank_b), _trainer(rank_b), _Ledger(2), fingerprint={**FP, "seed": 1}) @@ -240,7 +353,7 @@ def test_longer_run_changes_backend_fingerprint_and_lr_only_when_allowed(tmp_pat --rl-resume-allow-config-change is given.""" store = _store(tmp_path / "store") _r, _c, driver_a, cp_a = _trained_island(tmp_path, store, 2) - cp_a.save(driver_a, rollout_id=2) + _cut(cp_a, driver_a, rollout_id=2) rank_b = make_rank(9) cp_b = _cp(_rctl(tmp_path, "b", store), _Driver(rank_b), _trainer(rank_b), _Ledger(2), backend_fingerprint="miles@longer") @@ -261,7 +374,7 @@ def test_longer_run_changes_backend_fingerprint_and_lr_only_when_allowed(tmp_pat def test_resume_refuses_lr_jump_and_ledger_drift(tmp_path): store = _store(tmp_path / "store") _r, _c, driver_a, cp_a = _trained_island(tmp_path, store, 2) - cp_a.save(driver_a, rollout_id=2) + _cut(cp_a, driver_a, rollout_id=2) rank_b = make_rank(9) driver_b = _Driver(rank_b) cp_b = _cp(_rctl(tmp_path, "b", store), driver_b, _trainer(rank_b), _Ledger(2)) @@ -290,7 +403,7 @@ def test_cut_every_two_keeps_the_newest_two(tmp_path): driver.local_step = rid driver.published_state = SimpleNamespace(policy_tensor_hash=lambda: _hash(rank), policy_version=rid) cp.source.ledger = _Ledger(rid) - if cp.save(driver, rollout_id=rid) is not None: + if _cut(cp, driver, rollout_id=rid) is not None: saved.append(rid) assert saved == [2, 4, 6] assert cp.save(driver, rollout_id=7, final=True)["final"] is True # last round always cut @@ -541,3 +654,315 @@ def test_resume_wiring_carries_the_runtime_fingerprint(tmp_path): w = resume_wiring_for(SimpleNamespace(yeto_rl_resume={**cfg, "runtime_fingerprint": "miles@abc"}), sync=LocalOnlySync(3), environ={}) assert w.controller.runtime_fingerprint == "miles@abc" and w.every == 2 and w.ledger is not None + + +# ---------------------------------------------------------------- E1: cut before offload, background cut +class _CountingTrainer: + """The real MilesTrainerGroup cut path, counting the cut and onload/offload calls.""" + + def __init__(self, inner): + self.inner = inner + self.calls: list[str] = [] + + def save_cut(self, **kw): + self.calls.append("save_cut") + return self.inner.save_cut(**kw) + + def start_cut(self, **kw): + self.calls.append("start_cut") + return self.inner.start_cut(**kw) + + def finish_cut(self, pending): + self.calls.append("finish_cut") + return self.inner.finish_cut(pending) + + def collect_cut(self, pending): + self.calls.append("collect_cut") + return self.inner.collect_cut(pending) + + def onload(self): + self.calls.append("onload") + + def offload(self): + self.calls.append("offload") + + def __getattr__(self, name): + return getattr(self.inner, name) + + +def _publish(driver, version, policy_hash): + from yeto.rl.engine.driver import policy_token + + driver.published_version = version + driver.published_state = SimpleNamespace(policy_tensor_hash=lambda: policy_hash, policy_version=version) + driver.expected_token = policy_token(version, policy_hash) + + +def _island_with_first_cut(tmp_path, store, keep=2): + """One trained round, cut inline at v1 (the store's first cut); then a second + trained round whose cut (v2) is the periodic one under test.""" + rank = make_rank(0) + g = torch.Generator().manual_seed(3) + train_step(rank, torch.randn(3, 6, generator=g)) + ctl = _rctl(tmp_path, "a", store, keep=keep) + driver = _Driver(rank, published=None, local_step=1) + driver.last_applied_lrs = [5e-6] + trainer = _CountingTrainer(_trainer(rank)) + cp = _cp(ctl, driver, trainer, _Ledger(2)) + _publish(driver, 1, _hash(rank)) + first = _cut(cp, driver, rollout_id=1) + assert first["store_synced"] is True + train_step(rank, torch.randn(3, 6, generator=g)) + driver.local_step = 2 + trainer.calls.clear() + return rank, driver, trainer, cp, first + + +def _resume_fresh(tmp_path, store, rank): + rank_b = make_rank(9) + driver_b = _Driver(rank_b) + cp_b = _cp(_rctl(tmp_path, "b", store), driver_b, _trainer(rank_b), _Ledger(2)) + got = cp_b.resume(driver_b) + for n, v in params(rank).items(): + assert torch.equal(v, params(rank_b)[n]) + return got + + +def test_prepared_cut_is_written_in_the_background_without_waking_the_trainer(tmp_path): + store = _store(tmp_path / "store", _Vol()) + rank, driver, trainer, cp, first = _island_with_first_cut(tmp_path, store) + driver._trainer_offloaded = True + state = SimpleNamespace(policy_version=2, policy_tensor_hash=lambda: _hash(rank)) + assert cp.prepare(driver, rollout_id=2, state=state) is True + _publish(driver, 2, _hash(rank)) + info = cp.save(driver, rollout_id=2) + assert trainer.calls == ["start_cut"] # captured before the offload: no wake/sleep + assert info["prepared_before_offload"] is True and info["store_commit"] == "deferred" + assert cp.pointer()["cut_id"] == first["cut_id"] # not pointed at until complete + done = cp.join_commit(driver) + assert done["ok"] is True and done["cut_id"] == info["cut_id"] and trainer.calls[-1] == "collect_cut" + assert driver.events[-1][0] == "rl_cut_committed" and driver.events[-1][1]["ok"] is True + # E1: the event says where the shard was written (the rank's writer process) + assert driver.events[-1][1]["writer"] == [{"mode": "process"}] + assert cp.pointer()["cut_id"] == info["cut_id"] + assert rs.read_latest(store.root)["cut_id"] == info["cut_id"] + assert _resume_fresh(tmp_path, store, rank)["next_rollout_id"] == 2 + + +def test_prepared_cut_for_another_policy_is_dropped_and_cut_again(tmp_path): + store = _store(tmp_path / "store", _Vol()) + rank, driver, trainer, cp, _first = _island_with_first_cut(tmp_path, store) + driver._trainer_offloaded = True + cp.prepare(driver, rollout_id=2, state=SimpleNamespace(policy_version=2, policy_tensor_hash=lambda: "f" * 64)) + _publish(driver, 2, _hash(rank)) # the publication names other weights + info = cp.save(driver, rollout_id=2) + assert trainer.calls == ["start_cut", "finish_cut", "onload", "start_cut", "offload"] + assert info["prepared_before_offload"] is False and info["policy_hash"] == _hash(rank) + assert cp.join_commit(driver)["ok"] is True + assert _resume_fresh(tmp_path, store, rank)["policy_hash"] == _hash(rank) + + +def test_final_cut_completes_inline_after_the_pending_one(tmp_path): + store = _store(tmp_path / "store", _Vol()) + rank, driver, trainer, cp, _first = _island_with_first_cut(tmp_path, store) + _publish(driver, 2, _hash(rank)) + info = cp.save(driver, rollout_id=2, final=True) + assert info["store_commit"] == "inline" and info["store_synced"] is True + assert trainer.calls == ["save_cut"] and cp.join_commit(driver) is None + + +def test_failed_background_commit_keeps_the_previous_cut(tmp_path): + class _FailingVol(_Vol): + fail = False + + def commit(self): + if self.fail: + raise OSError("volume commit failed") + super().commit() + + vol = _FailingVol() + store = _store(tmp_path / "store", vol) + rank, driver, _trainer_c, cp, first = _island_with_first_cut(tmp_path, store, keep=1) + _publish(driver, 2, _hash(rank)) + vol.fail = True + cp.save(driver, rollout_id=2) + result = cp.join_commit(driver) + assert result["ok"] is False and "volume commit failed" in result["error"] + assert driver.events[-1][0] == "rl_cut_committed" and driver.events[-1][1]["ok"] is False + # nothing was pruned: the previous (committed) cut is still in the store + assert (store.root / rs.ROUND_CUTS / first["cut_id"]).is_dir() + + +def test_background_cut_can_be_turned_off(tmp_path, monkeypatch): + monkeypatch.setenv("YETO_RL_CUT_BACKGROUND_COMMIT", "0") + store = _store(tmp_path / "store", _Vol()) + rank, driver, trainer, cp, _first = _island_with_first_cut(tmp_path, store) + _publish(driver, 2, _hash(rank)) + info = cp.save(driver, rollout_id=2) + assert info["store_commit"] == "inline" and info["store_synced"] is True and trainer.calls == ["save_cut"] + + +def test_cut_context_takes_the_cursor_read_with_the_in_flight_entries(): + """E1: a rollout that keeps drawing during the cut (Miles fully_async) returns the + cursor read together with its entries; the context uses that one.""" + rank = make_rank(0) + driver = _Driver(rank, published=2, local_step=2) + driver.profile = SimpleNamespace(max_policy_age=1) + live = {"sample_offset": 9, "epoch_id": 0, "sample_group_index": 9, "sample_index": 9} + with_entries = {**live, "sample_offset": 12, "sample_group_index": 12, "sample_index": 12} + pool = SimpleNamespace(data_cursor=lambda: dict(live), export_in_flight=lambda: { + "entries": [], "buffer_groups": 0, "data_cursor": dict(with_entries)}) + source = CutSource(driver=lambda: driver, trainer=None, rollout=pool, ledger=_Ledger(2), + algorithm=SimpleNamespace(sha256=lambda: "a" * 64, to_legacy_runtime_attrs=lambda: {}), + backend_fingerprint="miles@0af62f4d", cut_root="", global_batch_size=ARGS.global_batch_size) + data = source.context("r000002-abc").data + assert data["sample_offset"] == 12 and data["buffer_length"] == 0 + + +def test_periodic_cut_is_skipped_while_the_previous_commit_runs(tmp_path): + import threading + + store = _store(tmp_path / "store", _Vol()) + rank, driver, trainer, cp, first = _island_with_first_cut(tmp_path, store) + gate = threading.Event() + busy = threading.Thread(target=gate.wait, daemon=True) + busy.start() + cp._commit = (busy, {"cut_id": first["cut_id"]}) + _publish(driver, 2, _hash(rank)) + assert cp.save(driver, rollout_id=2) is None and trainer.calls == [] + assert driver.events[-1][0] == "rl_cut_skipped" and cp.skipped_busy == 1 + gate.set() + busy.join() + info = cp.save(driver, rollout_id=2, final=True) # the final cut waits and is taken + assert info["store_commit"] == "inline" + + +def test_a_periodic_cut_never_waits_for_the_commit_it_starts(tmp_path): + import threading + + class _SlowVol(_Vol): + gate = threading.Event() + slow = False + + def commit(self): + if self.slow: + self.gate.wait(10) + super().commit() + + vol = _SlowVol() + store = _store(tmp_path / "store", vol) + rank, driver, _trainer_c, cp, _first = _island_with_first_cut(tmp_path, store) + g = torch.Generator().manual_seed(5) + _publish(driver, 2, _hash(rank)) + assert cp.save(driver, rollout_id=2)["store_commit"] == "deferred" + vol.slow = True + train_step(rank, torch.randn(3, 6, generator=g)) + driver.local_step = 3 + _publish(driver, 3, _hash(rank)) + info = cp.save(driver, rollout_id=3) # completes v2 and starts its (slow) commit + assert info["store_commit"] == "deferred" and cp._commit[0].is_alive() + train_step(rank, torch.randn(3, 6, generator=g)) + driver.local_step = 4 + _publish(driver, 4, _hash(rank)) + assert cp.save(driver, rollout_id=4) is None # v2 still committing: skipped + assert driver.events[-1][0] == "rl_cut_skipped" + vol.gate.set() + final = cp.save(driver, rollout_id=4, final=True) + assert final["store_commit"] == "inline" and rs.read_latest(store.root)["cut_id"] == final["cut_id"] + committed = [f["cut_id"] for e, f in driver.events if e == "rl_cut_committed"][1:] # after v1 + assert committed == ["r000002-" + cp.controller.incarnation["id"], info["cut_id"]] # v2, then v3 + + +def test_prune_keeps_a_cut_newer_than_every_kept_one(tmp_path): + """E1: a periodic cut is captured while the previous one commits in the background; + the commit's prune must not delete the newer cut being written.""" + root = tmp_path / "store" + for seq, cut_id in ((1, "r000002-aa"), (2, "r000004-aa")): + (root / rs.STATE_DIR / f"{seq:08d}").mkdir(parents=True) + (root / rs.STATE_DIR / f"{seq:08d}" / "cut").write_text(cut_id) + for cut_id in ("r000001-aa", "r000002-aa", "r000004-aa", "r000006-aa"): + (root / rs.ROUND_CUTS / cut_id).mkdir(parents=True) + (root / rs.LATEST).write_text(json.dumps({"seq": 2})) + deleted = rs.prune(root, keep=1, cut_ids_of=lambda p: (p / "cut").read_text()) + assert sorted(deleted) == [f"{rs.ROUND_CUTS}/r000001-aa", f"{rs.ROUND_CUTS}/r000002-aa", + f"{rs.STATE_DIR}/00000001"] + assert sorted(p.name for p in (root / rs.ROUND_CUTS).iterdir()) == ["r000004-aa", "r000006-aa"] + + +def test_a_prepared_cut_skipped_while_committing_is_collected_at_the_next_cut(tmp_path): + """E1: the safe point does not wait for the trainer writer of a prepared cut it + skips (busy commit); the next cut collects that writer and drops its files.""" + import threading + + from yeto.rl.engine.cut import cut_dir + + store = _store(tmp_path / "store", _Vol()) + rank, driver, trainer, cp, first = _island_with_first_cut(tmp_path, store) + driver._trainer_offloaded = True + state = SimpleNamespace(policy_version=2, policy_tensor_hash=lambda: _hash(rank)) + assert cp.prepare(driver, rollout_id=2, state=state) is True + gate = threading.Event() + busy = threading.Thread(target=gate.wait, daemon=True) + busy.start() + cp._commit = (busy, {"cut_id": first["cut_id"]}) + _publish(driver, 2, _hash(rank)) + assert cp.save(driver, rollout_id=2) is None and driver.events[-1][0] == "rl_cut_skipped" + assert trainer.calls == ["start_cut"] # the skipped cut's writer is not waited for + skipped = cut_dir(cp.root, cp._cut_id(2)) + gate.set() + busy.join() + info = cp.save(driver, rollout_id=2, final=True) + # the skipped writer is collected first, then the final cut is taken afresh + assert trainer.calls[:2] == ["start_cut", "finish_cut"] and cp._abandoned == [] + assert info["store_commit"] == "inline" and info["cut_id"] == skipped.name + assert rs.read_latest(store.root)["cut_id"] == info["cut_id"] + + +def test_a_cut_declined_at_prepare_does_not_wake_the_offloaded_trainer(tmp_path): + """E1: prepare met a running commit; when the commit is done by the safe point the + periodic cut is skipped instead of waking the offloaded trainer for it.""" + import threading + + store = _store(tmp_path / "store", _Vol()) + rank, driver, trainer, cp, first = _island_with_first_cut(tmp_path, store) + driver._trainer_offloaded = True + gate = threading.Event() + busy = threading.Thread(target=gate.wait, daemon=True) + busy.start() + cp._commit = (busy, {"cut_id": first["cut_id"]}) + state = SimpleNamespace(policy_version=2, policy_tensor_hash=lambda: _hash(rank)) + assert cp.prepare(driver, rollout_id=2, state=state) is False + gate.set() + busy.join() + _publish(driver, 2, _hash(rank)) + assert cp.save(driver, rollout_id=2) is None and trainer.calls == [] + assert driver.events[-1][0] == "rl_cut_skipped" + assert driver.events[-1][1]["reason"] == "not prepared before the trainer offloaded" + # a cut that was never declined still wakes the trainer (no silent loss of cuts) + train_step(rank, torch.randn(3, 6, generator=torch.Generator().manual_seed(7))) + driver.local_step = 3 + _publish(driver, 3, _hash(rank)) + assert cp.save(driver, rollout_id=3)["trainer_onloaded_for_cut"] is True + + +def test_a_deferred_cut_snapshots_the_state_dir_as_it_was_at_the_cut(tmp_path): + """E1 (Modal A10G resume e1c): a periodic cut is synced into the store after later + rounds changed the state dir (the ledger); the snapshot must hold the state dir as + it was at the cut, else a resume refuses the ledger (consumed groups differ).""" + store = _store(tmp_path / "store", _Vol()) + rank, driver, _trainer_c, cp, _first = _island_with_first_cut(tmp_path, store) + marker = Path(cp.controller.state_dir) / "ledger-e1-marker.txt" + marker.write_text("at the cut") + _publish(driver, 2, _hash(rank)) + info = cp.save(driver, rollout_id=2) + assert info["store_commit"] == "deferred" + marker.write_text("one round later") # the next round trains before the commit runs + done = cp.join_commit(driver) + assert done["ok"] is True + latest = rs.read_latest(store.root) + assert latest["cut_id"] == info["cut_id"] + snap = store.root / rs.STATE_DIR / f"{int(latest['seq']):08d}" + assert (snap / marker.name).read_text() == "at the cut" + assert json.loads((snap / POINTER).read_text())["cut_id"] == info["cut_id"] + assert cp.pointer()["cut_id"] == info["cut_id"] + assert not list(Path(cp.controller.state_dir).parent.glob(".*.stage-*")) # staging removed diff --git a/tests/test_rl_trainer_cut.py b/tests/test_rl_trainer_cut.py index a2e1c8b73..643f04fd2 100644 --- a/tests/test_rl_trainer_cut.py +++ b/tests/test_rl_trainer_cut.py @@ -284,8 +284,8 @@ def test_context_is_checked_before_any_rank_writes(tmp_path): def test_missing_optimizer_coverage_is_refused(tmp_path, monkeypatch): from yeto.rl.adapters.miles import cut_plugin - real = cut_plugin._save - monkeypatch.setattr(cut_plugin, "_save", lambda *a, **k: {**real(*a, **k), "optimizer_names": ["lora_A"]}) + real = cut_plugin._write + monkeypatch.setattr(cut_plugin, "_write", lambda *a, **k: {**real(*a, **k), "optimizer_names": ["lora_A"]}) with pytest.raises(CutError, match="no optimizer state"): _trainer(RankGroup([_trained_rank()])).save_cut(epoch=1, context=_context(tmp_path)) diff --git a/tests/test_rl_verl_adapter.py b/tests/test_rl_verl_adapter.py index 96c588451..1a417dcd1 100644 --- a/tests/test_rl_verl_adapter.py +++ b/tests/test_rl_verl_adapter.py @@ -447,3 +447,27 @@ def test_verl_data_cursor_refuses_what_it_cannot_honour(): cur.seek({"sample_offset": 2}) with pytest.raises(RuntimeError, match="twice"): VerlDataCursor(trainer) + + +def test_fully_async_runtime_env_keeps_nccl_host_buffers_on_shm(monkeypatch): + # S26 E0.1: a cuMem host-buffer sender and a /dev/shm receiver cannot attach on Modal + import sys + import types + + pkg = types.ModuleType("verl.trainer.constants_ppo") + pkg.get_ppo_ray_runtime_env = lambda _x: {"env_vars": {}} + for name in ("verl", "verl.trainer"): + monkeypatch.setitem(sys.modules, name, sys.modules.get(name) or types.ModuleType(name)) + monkeypatch.setitem(sys.modules, "verl.trainer.constants_ppo", pkg) + ray_stub = types.ModuleType("ray") + ray_stub.remote = lambda *a, **k: (lambda cls: cls) + monkeypatch.setitem(sys.modules, "ray", sys.modules.get("ray") or ray_stub) + monkeypatch.setitem(sys.modules, "yeto.rl.adapters.verl.verl_main", None) + monkeypatch.delitem(sys.modules, "yeto.rl.adapters.verl.verl_main") + from yeto.rl.adapters.verl import verl_main + + monkeypatch.delenv(verl_main.CUMEM_HOST_ENV, raising=False) + assert verl_main._runtime_env(fully_async=True)["env_vars"][verl_main.CUMEM_HOST_ENV] == "0" + assert verl_main.CUMEM_HOST_ENV not in verl_main._runtime_env()["env_vars"] + monkeypatch.setenv(verl_main.CUMEM_HOST_ENV, "1") + assert verl_main._runtime_env(fully_async=True)["env_vars"][verl_main.CUMEM_HOST_ENV] == "1" diff --git a/tests/test_rl_verl_buffer_trainer.py b/tests/test_rl_verl_buffer_trainer.py index 58dea8247..b37ad8a58 100644 --- a/tests/test_rl_verl_buffer_trainer.py +++ b/tests/test_rl_verl_buffer_trainer.py @@ -212,6 +212,12 @@ def yeto_param_version(self): def yeto_export_lora(self): return {k: v.clone() for k, v in self.tensors.items()}, {"count": 2} + def yeto_export_digest(self, kwargs): + from yeto.rl.adapters.verl.ports_impl import digest_lora_tensors + + self.digest_exports = getattr(self, "digest_exports", 0) + 1 + return digest_lora_tensors(dict(self.tensors), **kwargs) + def yeto_call_workers(self, fn, tensors, reset): for k in self.tensors: self.tensors[k] = tensors[k].clone().float() diff --git a/tests/test_rl_verl_fully_async_64b.py b/tests/test_rl_verl_fully_async_64b.py index 145cd8a2f..2ab12db40 100644 --- a/tests/test_rl_verl_fully_async_64b.py +++ b/tests/test_rl_verl_fully_async_64b.py @@ -256,6 +256,12 @@ def yeto_push_weights(self): def yeto_export_lora(self): return {k: v.clone() for k, v in self.tensors.items()}, {"count": 2} + def yeto_export_digest(self, kwargs): + from yeto.rl.adapters.verl.ports_impl import digest_lora_tensors + + self.digest_exports = getattr(self, "digest_exports", 0) + 1 + return digest_lora_tensors(dict(self.tensors), **kwargs) + def yeto_call_workers(self, fn, tensors, reset): for k in self.tensors: self.tensors[k] = tensors[k].clone().float() @@ -335,6 +341,22 @@ def test_driver_runs_fully_async_rounds_with_carry_over_and_discard(tmp_path): checks = [e for e in events if e["event"] == "rl_publication_check"] assert all(e["status"] == "VERIFIED" for e in checks) and len(checks) == 4 assert [e["param_version"] for e in events if e["event"] == "rl_verl_version_map"] == [0, 1, 2, 3] + # rl-publish-fastpath: LocalOnlySync takes digests from the trainer, no tensor export + assert actor.digest_exports == 4 + + +def test_fastpath_off_publishes_the_same_hashes(tmp_path, monkeypatch): + def run(sub): + actor = FakeTrainerActor(tmp_path / sub / "rb", [_sample("a", 0, 0, [(0, 4)])] * 2) + _driver(tmp_path / sub, actor, rounds=2, required=1)[0].run() + return [(e["policy_tensor_hash"], e["sent_checksum"]) for e in _events(tmp_path / sub / "tape.jsonl") + if e["event"] == "rl_publication_check"], getattr(actor, "digest_exports", 0) + + fast, fast_digests = run("fast") + monkeypatch.setenv("YETO_RL_PUBLISH_FASTPATH", "0") + slow, slow_digests = run("slow") + assert fast == slow and len(fast) == 3 + assert fast_digests == 3 and slow_digests == 0 def test_publication_mismatch_on_the_rollouter_replica_fails(tmp_path): diff --git a/tests/test_rl_verl_publish_fastpath.py b/tests/test_rl_verl_publish_fastpath.py new file mode 100644 index 000000000..186a93e06 --- /dev/null +++ b/tests/test_rl_verl_publish_fastpath.py @@ -0,0 +1,117 @@ +"""rl-publish-fastpath for verl: the trainer-side digest equals the driver-side +full-export values bit for bit, and the publishers use it without tensors.""" + +from __future__ import annotations + +import torch + +from yeto.rl.adapters.verl import publish as pub +from yeto.rl.adapters.verl.ports_impl import (VerlPolicyState, VerlResidentState, digest_lora_tensors, + sent_checksum) +from yeto.rl.core import CanonicalTensorSpec, canonical_layout_hash, canonical_state, policy_tensor_hash +from yeto.rl.engine.policy_digest import is_resident +from yeto.rl.engine.trainable_state import TrainableState + +NAMES = ("model.layers.0.q_proj.lora_A.weight", "model.layers.0.q_proj.lora_B.weight") +SPECS = (CanonicalTensorSpec(NAMES[0], (4, 8), "float32", 32), + CanonicalTensorSpec(NAMES[1], (8, 4), "float32", 32)) +IDS = {"base_model_revision": "a" * 40, "lora_config_hash": "b" * 64, "layout_hash": canonical_layout_hash(SPECS)} + + +def _tensors(seed: int = 0) -> dict: + gen = torch.Generator().manual_seed(seed) + # bf16 originals upcast to f32, as remote_export_lora returns them + return {s.name: torch.randn(s.shape, generator=gen).to(torch.bfloat16).float() for s in SPECS} + + +class _Island: + def __init__(self, tensors): + self.tensors = tensors + self.base_model_revision, self.lora_config_hash = IDS["base_model_revision"], IDS["lora_config_hash"] + self.layout_hash, self.expected_specs = IDS["layout_hash"], SPECS + self.exports = 0 + + def export_tensors(self): + self.exports += 1 + return {k: v.clone() for k, v in self.tensors.items()}, {} + + def export_digest(self, kwargs): + return digest_lora_tensors(dict(self.tensors), **kwargs) + + +def test_trainer_digest_equals_the_full_export_values(): + tensors = _tensors() + raw = digest_lora_tensors(tensors, **IDS, expected_specs=SPECS) + full = canonical_state(3, tensors, **IDS, expected_specs=SPECS) + assert raw["digest"]["policy_tensor_hash"] == policy_tensor_hash(full) + assert raw["lora_checksum"] == pub.lora_checksum(dict(full.tensors)) + + +def test_resident_state_crosses_no_tensors_and_publishes_the_same_checksum(): + island = _Island(_tensors(1)) + ps = VerlPolicyState(island) + resident = ps.export_digest(policy_version=5) + assert isinstance(resident, VerlResidentState) and is_resident(resident) + assert island.exports == 0 + full = ps.export() + assert resident.policy_tensor_hash() == full.policy_tensor_hash() + assert sent_checksum(resident) == sent_checksum(full) + assert island.exports == 1 + # a re-versioned handle keeps the checksum (bridges.with_version) + assert sent_checksum(resident.with_version(6))[0] == sent_checksum(full)[0] + + +def test_materialize_falls_back_to_the_full_export_at_the_handle_version(): + island = _Island(_tensors(2)) + resident = VerlPolicyState(island).export_digest(policy_version=7) + state = resident.materialize() + assert isinstance(state, TrainableState) and state.policy_version == 7 + assert state.policy_tensor_hash() == resident.policy_tensor_hash() + assert island.exports == 1 + + +def test_island_without_a_digest_hook_falls_back(): + class Plain: + pass + + assert VerlPolicyState(Plain()).export_digest(policy_version=0) is None + + +def test_driver_switch_turns_the_fast_path_off(monkeypatch): + from yeto.rl.engine.driver import IslandDriver + + driver = IslandDriver.__new__(IslandDriver) + driver.policy_state = VerlPolicyState(_Island(_tensors())) + assert is_resident(driver.export_local_resident(policy_version=0)) + monkeypatch.setenv("YETO_RL_PUBLISH_FASTPATH", "0") + assert driver.export_local_resident(policy_version=0) is None + + +def test_worker_digest_with_the_staged_export_equals_the_full_export(): + """E1: the worker stages its LoRA in one host buffer for the digest; the digest + and the bf16 send checksum equal those of the full export, round after round.""" + from types import SimpleNamespace + + from yeto.rl.adapters.verl import ports_impl as pi + + gen = torch.Generator().manual_seed(3) + params = { + "base_model.model.layers.0.q_proj.lora_A.default.weight": + torch.nn.Parameter(torch.randn(4, 8, generator=gen).to(torch.bfloat16)), + "base_model.model.layers.0.q_proj.lora_B.default.weight": + torch.nn.Parameter(torch.randn(4, 8, generator=gen).t()), # not contiguous + "base_model.model.layers.0.q_proj.base_layer.weight": torch.nn.Parameter(torch.randn(2, 2, generator=gen)), + } + peft = SimpleNamespace(named_parameters=lambda: list(params.items())) + worker = SimpleNamespace(actor=SimpleNamespace(engine=SimpleNamespace(module=peft))) + tensors, _info = pi.remote_export_lora(worker) + specs = tuple(CanonicalTensorSpec(n, tuple(t.shape), "float32", t.numel()) for n, t in sorted(tensors.items())) + kwargs = {**IDS, "layout_hash": canonical_layout_hash(specs), "expected_specs": specs} + def untimed(raw): # timings differ from call to call + return {**{k: v for k, v in raw.items() if k != "export_seconds"}, + "digest": {k: v for k, v in raw["digest"].items() if k != "hash_seconds"}} + + want = untimed(digest_lora_tensors(tensors, **kwargs)) + for _ in range(2): # the second round reuses the worker's buffer + assert untimed(pi.remote_export_digest(worker, kwargs)) == want + assert getattr(worker, "_yeto_digest_staging", None) is not None diff --git a/tests/test_tools_eval_efficiency.py b/tests/test_tools_eval_efficiency.py new file mode 100644 index 000000000..5f613bfd3 --- /dev/null +++ b/tests/test_tools_eval_efficiency.py @@ -0,0 +1,423 @@ +"""CPU tests of the efficiency evaluation (tools/eval/efficiency, S26 E0.1). + +The synthetic logs below follow the formats the parsers cite: + +* verl FileLogger rows ``{"step": N, "data": {...}}`` -- fork acad9875 + ``verl/utils/tracking.py`` ``FileLogger.log``; console ``step:N - k:v`` -- + ``verl/utils/logger/aggregate_logger.py`` ``concat_dict_to_str``; keys as in the + real T-EVAL file ``SESSION25/logs/T-EVAL/eval-e1-verl-c/.../verl-file-logger-0.jsonl`` + and ``metric_utils.compute_timing_metrics`` (``timing_s/*``); +* Miles ``perf N: {dict}`` lines -- copied in shape from the real T-EVAL log + ``SESSION25/logs/T-EVAL/eval-e1-miles-a/raw/launch.ts.log`` (rollout side + ``metrics.py:93``, train side ``train_metric_utils.py:56``); +* yeto tape events -- field names as in the T-EVAL tapes (``rl_round_wall``, + ``rl_local_round``, ``rl_round_trained``, ``rl_timeline_span``); +* NVML CSV -- ``nvidia-smi --query-gpu=... --format=csv`` output and the + ``miles_local.sh`` ``,`` prefixed form. +""" + +from __future__ import annotations + +import json +from pathlib import Path + +import pytest + +from tools.eval import __main__ as cli +from tools.eval.efficiency import eff_reward, metrics, parse, report, spec + +SPEC = spec.get_spec("qwen3-0.6b") + + +# ====================================================================== spec +def test_arms_cover_both_systems_frameworks_and_modes(): + assert len(spec.ARMS) == 8 + for name in spec.ARMS: + arm = spec.Arm.parse(name) + assert spec.Arm.parse(arm.peer).system != arm.system + with pytest.raises(ValueError): + spec.Arm.parse("bare-verl-fast") + + +def test_every_arm_carries_the_shared_values(): + plans = {a: spec.plan(SPEC, a, "r1") for a in spec.ARMS} + assert len({p["fingerprint"] for p in plans.values()}) == 1 + for name, p in plans.items(): + text = p.get("script") or " ".join(p.get("argv") or []) or p["launcher_args"] + assert p["cards"]["total"] == (1 if p["mode"] == "sync" else 2), name + if p["system"] == "bare": + assert "Qwen3-0.6B" in text and "/root/eff-data/train" in text, name + else: + assert SPEC.model_revision in text and SPEC.data_revision in text, name + for value in (SPEC.rounds, SPEC.groups_per_round, SPEC.samples_per_group, SPEC.max_response_len, + SPEC.lr, SPEC.lora_r, SPEC.seed): + assert str(value) in text, (name, value) + + +def test_fingerprint_changes_with_a_shared_value_only(): + base = spec.control_fingerprint(SPEC) + assert spec.control_fingerprint(spec.get_spec("qwen3-0.6b", repeats=5)) == base + assert spec.control_fingerprint(spec.get_spec("qwen3-0.6b", lr=2e-5)) != base + + +def test_bare_entry_points_follow_the_official_recipes(): + sync = spec.bare_script(SPEC, spec.Arm.parse("bare-verl-sync"), "r1") + async_ = spec.bare_script(SPEC, spec.Arm.parse("bare-verl-async"), "r1") + assert "-m verl.trainer.main_ppo" in sync and "fully_async_main" in async_ + assert "async_training.staleness_threshold=1" in async_ and "rollout.n_gpus_per_node=1" in async_ + assert "trainer.logger=[console,file]" in sync and "VERL_FILE_LOGGER_PATH" in sync + miles_sync = spec.bare_script(SPEC, spec.Arm.parse("bare-miles-sync"), "r1") + miles_async = spec.bare_script(SPEC, spec.Arm.parse("bare-miles-async"), "r1") + assert "/root/miles/train.py $MODEL_ARGS" in miles_sync and "--colocate" in miles_sync + assert "/root/miles/train_async.py" in miles_async and "--fully-async --max-weight-staleness 1" in miles_async + assert "scripts/models/qwen3-0.6B.py" in miles_sync # the image's own model flags + for script in (sync, async_, miles_sync, miles_async): + assert "yeto" not in script.replace("yeto-effeval", "").replace("yeto/", "") + + +def test_yeto_arms_use_the_same_policy_age_and_partition(): + assert "--rl-max-policy-age 1" in " ".join(spec.yeto_verl_argv(SPEC, "async", "/t")) + assert "--rl-max-policy-age" not in " ".join(spec.yeto_verl_argv(SPEC, "sync", "/t")) + # async needs an algorithm spec that tolerates the policy age (the default allows 0) + argv = spec.yeto_verl_argv(SPEC, "async", "/t") + algo = argv[argv.index("--rl-algorithm-spec") + 1] + assert algo.endswith(spec.ALGORITHM_SPEC) and (Path(__file__).parents[1] / spec.ALGORITHM_SPEC).is_file() + args = spec.yeto_miles_launcher_args(SPEC, "async", "r1") + assert "--gpu modal:2xa10g" in args and "--rl-placement fixed-partition --rl-rollout-gpus 1" in args + assert "--gpu modal:1xa10g" in spec.yeto_miles_launcher_args(SPEC, "sync", "r1") + # E1: on the bare arms' GPU model (an A10G container is refused) + assert "--modal-gpu-name 'NVIDIA A10' " in args + from tools.eval import modal_verl + + assert spec.GPU_MODELS == modal_verl.GPU_NAMES + # only the final round cut, as the bare arms keep no periodic checkpoint (E1) + assert f"--rl-cut-every {SPEC.rounds} " in spec.yeto_miles_launcher_args(SPEC, "sync", "r1") + every = spec.get_spec(SPEC.name, checkpoint="every-round") + assert "--rl-cut-every 1 " in spec.yeto_miles_launcher_args(every, "async", "r1") + assert spec.control_fingerprint(every) == spec.control_fingerprint(SPEC) # bare arms unchanged + default = spec.get_spec(SPEC.name, checkpoint="default") + assert "--rl-cut-every" not in spec.yeto_miles_launcher_args(default, "async", "r1") + with pytest.raises(ValueError): + spec.cut_every(spec.get_spec(SPEC.name, checkpoint="never")) + + +# overrides yeto's verl translation sets and the bare recipe leaves at the default (or the reverse); +# a new key here means the two arms drifted apart: decide, then extend the list +KNOWN_VERL_DIFF = { + "actor_rollout_ref.model.lora.merge", "actor_rollout_ref.rollout.enforce_eager", + "actor_rollout_ref.rollout.full_determinism", "actor_rollout_ref.rollout.load_format", + "ray_kwargs.ray_init.num_cpus", "trainer.project_name", "trainer.use_v1", "trainer.v1.trainer_mode", + "actor_rollout_ref.rollout.mode", "async_training.staleness_threshold", "data.return_raw_chat", + "rollout.total_rollout_steps", +} + + +@pytest.mark.parametrize("mode", spec.MODES) +def test_verl_knob_diff_is_the_known_set(mode): + diff = spec.verl_knob_diff(SPEC, mode) + assert set(diff) <= KNOWN_VERL_DIFF, set(diff) - KNOWN_VERL_DIFF + + +def test_4b_preset_is_ready_for_e02(): + s4 = spec.get_spec("qwen3-4b-instruct-2507") + assert s4.gpu == "H100!" and s4.repeats >= 2 and len(s4.model_revision) == 40 + assert "--gpu modal:2xh100" in spec.yeto_miles_launcher_args(s4, "async", "r") + assert "qwen3-4B-Instruct-2507.py" in spec.bare_script(s4, spec.Arm.parse("bare-miles-sync"), "r") + + +# ====================================================================== reward +@pytest.mark.parametrize("response,label", [ + ("so \\boxed{42}", "#### 42"), ("the answer is 1,234", "x #### 1234"), ("\\boxed{3.5} then 7", "#### 3.5"), + ("no number", "#### 1"), ("\\boxed{-2}", "#### 2"), ("", None), +]) +def test_bare_reward_equals_yeto_reward(response, label): + from yeto.rl.rewards.builtin import gsm8k_grade + + assert eff_reward.grade(response, label) == gsm8k_grade(response, label) + assert eff_reward.compute_score("d", response, label) == gsm8k_grade(response, label) + + +# ====================================================================== parse: verl +def _verl_rows(n: int, *, async_: bool = False) -> str: + rows = [] + for step in range(1, n + 1): + d = {"timing_s/gen": 10.0 + step, "timing_s/old_log_prob": 1.0, "timing_s/update_actor": 3.0, + "timing_s/step": 16.0 + step, "response_length/mean": 100.0, "perf/total_num_tokens": 9000, + "critic/score/mean": 0.1 * step, "rollout_corr/rollout_is_ratio_fraction_high": 0.01} + if async_: + d.pop("timing_s/gen") + d.update({"timing_s/param_sync": 0.5, "timing_s/wait_for_enough_samples": 2.0, + "fully_async/count/dropped_stale_samples": step, "fully_async/count/total_generated_samples": + 32 * step}) + else: + d["timing_s/update_weights"] = 0.7 + rows.append(json.dumps({"step": step, "data": d})) + return "\n".join(rows) + "\n" + + +def _verl_console(n: int, t0: float = 1000.0, dt: float = 20.0) -> str: + return "".join(f"{t0 + dt * s:.3f} (TaskRunner pid=1) step:{s} - timing_s/gen:{10 + s} - critic/score/mean:0.1\n" + for s in range(1, n + 1)) + + +def test_parse_verl_sync_rounds_and_sources(): + p = parse.parse_verl(_verl_rows(3), _verl_console(3), samples_per_round=32, trainer_cards=1, mode="sync") + r = p.rounds + assert [x["round"] for x in r] == [1, 2, 3] + assert r[1]["round_s"] == 20.0 and r[0]["round_s"] is None + assert r[0]["gen_tokens"] == 3200 and r[0]["train_s"] == 4.0 and r[0]["weight_sync_s"] == 0.7 + assert p.sources["gen_tokens"].startswith("derived:") and p.sources["wait_s"] == parse.NA + assert p.sources["sampler_update_s"] == parse.NA and p.sources["version_lag"] == parse.NA + assert p.segments_raw[0]["timing_s/step"] == 17.0 + + +def test_parse_verl_async_reads_fully_async_counters(): + p = parse.parse_verl(_verl_rows(2, async_=True), "", samples_per_round=32, trainer_cards=1, mode="async") + assert p.rounds[0]["wait_s"] == 2.0 and p.rounds[0]["weight_sync_s"] == 0.5 + assert p.rounds[1]["dropped_samples"] == 2 and p.sources["gen_s"] == parse.MISSING + assert p.sources["t_end_unix"] == parse.MISSING and p.sources["version_lag"] == parse.MISSING + + +# ====================================================================== parse: Miles +def _miles_log(n: int, *, async_: bool = False) -> str: + out = [] + for rid in range(n): + ro = {"rollout/raw_reward": 0.4 + 0.01 * rid, "perf/rollout_time": 12.0, + "perf/tokens_per_gpu_per_sec": 500.0, "rollout/num_training_samples": 32} + tr = {"perf/train_wait_time": 12.5, "perf/train_time": 20.0, "perf/actor_train_time": 15.0, + "perf/actor_train_tok_per_s": 800.0, "perf/update_weights_time": 0.15, "perf/sleep_time": 1.7} + out.append(f"2026-10-20T06:36:{10 + rid:02d}Z [p-l0-modal] \x1b[36m(RolloutExecutor pid=2765)\x1b[0m " + f"[2026-10-20 06:36:{10 + rid:02d}.055 rollout_executor] metrics.py:93 - perf {rid}: {ro!r}") + out.append(f"2026-10-20T06:37:{rid * 30 % 60:02d}Z [p-l0-modal] (MegatronTrainRayActor pid=3190) " + f"[2026-10-20 06:{37 + rid // 2}:{rid * 30 % 60:02d}.147 actor_cell00000_rank00000] " + f"train_metric_utils.py:56 - perf {rid}: {tr!r}") + return "\n".join(out) + "\nunrelated line\n" + + +def test_parse_miles_merges_rollout_and_train_lines(): + p = parse.parse_miles(_miles_log(3), rollout_cards=1, mode="sync") + r = p.rounds + assert [x["round"] for x in r] == [0, 1, 2] + assert r[0]["gen_s"] == 12.0 and r[0]["train_s"] == 20.0 and r[0]["weight_sync_s"] == 0.15 + assert r[0]["gen_tokens"] == 6000.0 and r[0]["train_tokens"] == 12000.0 and r[0]["samples"] == 32 + assert r[1]["round_s"] == 30.0 and r[0]["reward_mean"] == 0.4 + assert p.sources["wait_s"] == parse.NA and p.sources["is_clip_frac"] == parse.MISSING + + +def test_parse_miles_async_keeps_wait(): + p = parse.parse_miles(_miles_log(2), rollout_cards=1, mode="async") + assert p.rounds[0]["wait_s"] == 12.5 and p.sources["dropped_samples"] == parse.MISSING + + +# ====================================================================== parse: yeto tape +def _tape(n: int, t0: float = 2000.0) -> list[dict]: + ev = [] + for rid in range(n): + end = t0 + 25.0 * (rid + 1) + ev += [{"event": "rl_timeline_span", "rollout_id": rid, "task": "generate", "kind": "compute", + "start": 100.0 * rid, "end": 100.0 * rid + 11.0}, + {"event": "rl_timeline_span", "rollout_id": rid, "task": "train", "kind": "compute", + "start": 100.0 * rid + 11.0, "end": 100.0 * rid + 21.0}, + {"event": "rl_timeline_span", "rollout_id": rid, "task": "publish", "kind": "transfer", + "start": 100.0 * rid + 21.0, "end": 100.0 * rid + 22.0}, + {"event": "rl_timeline_span", "rollout_id": rid, "task": "offload", "kind": "transfer", + "start": 100.0 * rid + 22.0, "end": 100.0 * rid + 24.0}, + {"event": "rl_round_wall", "rollout_id": rid, "round_seconds": 25.0, "time_unix": end}, + {"event": "rl_local_round", "action_tokens": 3000 + rid, "reward_mean": 0.5, "time_unix": end - 1}, + {"event": "rl_round_trained", "rollout_id": rid, "trained_samples": 32, + "mismatch": {"tis_clipfrac": 0.0}, "time_unix": end}] + return ev + + +def test_parse_yeto_tape_with_framework_tokens(): + fw = parse.parse_miles(_miles_log(3), rollout_cards=1, mode="sync") + p = parse.parse_yeto(_tape(3), mode="sync", framework_parsed=fw) + r = p.rounds + assert r[0]["gen_s"] == 11.0 and r[0]["train_s"] == 10.0 and r[0]["weight_sync_s"] == 1.0 + assert r[1]["round_s"] == 25.0 and r[2]["gen_tokens"] == 3002 and r[0]["train_tokens"] == 12000.0 + assert p.sources["train_tokens"].startswith("framework log:") and p.segments_raw[0]["offload"] == 2.0 + + +def test_parse_yeto_driver_phase_fallback(): + ev = [{"event": "rl_driver_phase", "phase": "generate", "rollout_id": 0, "time_unix": 10.0}, + {"event": "rl_driver_phase", "phase": "train", "time_unix": 18.0}, + {"event": "rl_driver_phase", "phase": "publish", "time_unix": 21.0}, + {"event": "rl_driver_phase", "phase": "generate", "rollout_id": 1, "time_unix": 22.0}, + {"event": "rl_round_wall", "rollout_id": 0, "time_unix": 22.0}] + p = parse.parse_yeto(ev, mode="sync") + assert p.rounds[0]["gen_s"] == 8.0 and p.rounds[0]["train_s"] == 3.0 and p.rounds[0]["weight_sync_s"] == 1.0 + assert "rl_driver_phase" in p.sources["gen_s"] + + +def test_parse_yeto_verl_async_generation_from_rollouter_timing(): + ev = [{"event": "rl_publication_check", "policy_version": 0, + "rollouter_timing": {"fully_async/rollouter/active_time": 2.75}}, + {"event": "rl_publication_check", "policy_version": 1, + "rollouter_timing": {"fully_async/rollouter/active_time": 10.8}}, + {"event": "rl_round_wall", "rollout_id": 0, "time_unix": 21.3}, + {"event": "rl_publication_check", "policy_version": 2, + "rollouter_timing": {"fully_async/rollouter/active_time": 1.86}}, + {"event": "rl_round_wall", "rollout_id": 1, "time_unix": 26.5}] + p = parse.parse_yeto(ev, mode="async", skip=frozenset({"gen_s"}), rollouter_gen=True) + assert [r["gen_s"] for r in p.rounds] == [10.8, 1.86] + assert p.sources["gen_s"].startswith("rl_publication_check.rollouter_timing") + off = parse.parse_yeto(ev, mode="async", skip=frozenset({"gen_s"})) + assert off.rounds[0]["gen_s"] is None + + +def test_parse_nvml_both_formats(): + own = ("timestamp, index, name, utilization.gpu [%], memory.used [MiB], memory.total [MiB], power.draw [W]\n" + "2026/10/20 06:36:11.055, 0, NVIDIA A10, 45 %, 6077 MiB, 23028 MiB, 81.7 W\n" + "2026/10/20 06:36:12.055, 1, NVIDIA A10, 0 %, 100 MiB, 23028 MiB, 30.1 W\n") + polled = "ta-01ABC,2026/10/20 06:36:11.055, NVIDIA A10, 44 %, 6077 MiB, 81.76 W\n" + a, b = parse.parse_nvml(own), parse.parse_nvml(polled) + assert [x["util"] for x in a] == [45.0, 0.0] and a[1]["key"] == ":1" and a[0]["mem_total_mb"] == 23028.0 + assert b[0]["key"] == "ta-01ABC:0" and b[0]["mem_mb"] == 6077.0 and b[0]["power_w"] == 81.76 + + +# ====================================================================== run dirs, metrics, report +def _nvml_csv(t0: float, n: int, util: int) -> str: + import datetime as dt + + lines = ["timestamp, index, name, utilization.gpu [%], memory.used [MiB], memory.total [MiB], power.draw [W]"] + for i in range(n): + stamp = dt.datetime.fromtimestamp(t0 + i, dt.timezone.utc).strftime("%Y/%m/%d %H:%M:%S.%f")[:-3] + lines.append(f"{stamp}, 0, NVIDIA A10, {util if i % 4 else 0} %, 7000 MiB, 23028 MiB, 80 W") + return "\n".join(lines) + "\n" + + +def _run_dir(tmp: Path, arm: str, run_id: str, dt_round: float) -> Path: + plan = spec.plan(SPEC, arm, run_id) + d = tmp / run_id + spec.write_plan(d / "arm.json", plan) + raw = d / "raw" / run_id / "arm" + raw.mkdir(parents=True) + n, t0 = SPEC.rounds, 1000.0 + (raw / "meta.json").write_text(json.dumps({"rc": 0, "start_unix": t0 - 200, "end_unix": t0 + dt_round * n + 30})) + (raw / "nvml.csv").write_text(_nvml_csv(t0, int(dt_round * n) + 5, 60)) + if plan["system"] == "bare": + (raw / "main.log").write_text(_verl_console(n, t0, dt_round)) + (raw / "eff-out").mkdir() + (raw / "eff-out" / "verl-metrics.jsonl").write_text(_verl_rows(n)) + else: + out = raw / "yeto-output" + out.mkdir() + (out / "verl-file-logger-0.jsonl").write_text("".join( # steps 1..n = rollout ids 0..n-1 + json.dumps({"step": k, "data": {"global_seqlen/mean": 9000}}) + "\n" for k in range(1, n + 1))) + (out / "rl-island-0.jsonl").write_text("\n".join(json.dumps(e) for e in _tape(n, t0 - dt_round) + ) + "\n") + return d + + +def test_parse_run_metrics_and_report(tmp_path, monkeypatch): + dirs = [_run_dir(tmp_path, "bare-verl-sync", "b1", 20.0), _run_dir(tmp_path, "bare-verl-sync", "b2", 21.0), + _run_dir(tmp_path, "yeto-verl-sync", "y1", 25.0), _run_dir(tmp_path, "yeto-verl-sync", "y2", 25.0)] + for d in dirs: + assert cli.main(["eff-parse", str(d)]) == 0 + assert (d / "parsed.json").is_file() and (d / "summary.json").is_file() + s = json.loads((dirs[0] / "summary.json").read_text()) + assert s["rounds"] == SPEC.rounds and s["steady_wall_s"] == 20.0 * (SPEC.rounds - 1) + assert s["gen_tok_per_card_s"] == round(3200 / 20.0, 3) and s["bubble_frac"] == pytest.approx(0.25, abs=0.02) + assert s["usd_per_1k_samples_steady"] == round(2.0 * 20.0 * 7 / 3600 / (32 * 7) * 1000, 4) + assert s["gpu_mem_peak_mb"] == 7000.0 + assert s["missing"] == [] and "wait_s_median" in s["not_applicable"] and "drop_frac" in s["not_applicable"] + y = json.loads((dirs[2] / "summary.json").read_text()) + assert y["missing"] == [] and y["train_s_median"] == 10.0 and y["unaccounted_s_median"] == 3.0 + assert y["train_tok_per_card_s"] == round(9000 / 25.0, 3) + out = tmp_path / "report.md" + assert cli.main(["eff-report", *map(str, dirs), "--out", str(out)]) == 0 + text = out.read_text() + assert "### verl sync (bare n=2, yeto n=2)" in text + section = text.split("### verl sync")[1] + row = next(line for line in section.splitlines() if line.startswith("| round time (s)")) + assert "yeto worse" in row and "+4.5" in row + assert "| update sampler weights (s/round) | n/a | n/a | n/a | not applicable |" in text + assert "wait (s/round) | n/a" in text + + +def test_report_needs_two_repeats_and_one_spec(tmp_path): + one = [_run_dir(tmp_path, "bare-verl-sync", "b1", 20.0), _run_dir(tmp_path, "yeto-verl-sync", "y1", 25.0)] + for d in one: + cli.main(["eff-parse", str(d)]) + text = report.markdown(report.load(one), title="t") + assert "insufficient repeats (bare 1, yeto 1; need 2)" in text + runs = report.load(one) + runs[1]["fingerprint"] = "other" + with pytest.raises(ValueError, match="different specs"): + report.markdown(runs, title="t") + runs = report.load(one) + assert runs[0]["gpu_names"] # from the NVML samples + runs[1]["gpu_names"] = ["NVIDIA A10G"] + with pytest.raises(ValueError, match="different GPU models"): + report.markdown(runs, title="t") + + +def test_gpu_name_check_tells_a10_from_a10g(): + from tools.eval.modal_verl import gpu_name_matches + + assert gpu_name_matches("NVIDIA A10", "NVIDIA A10") + assert not gpu_name_matches("NVIDIA A10", "NVIDIA A10G") + assert gpu_name_matches("H100", "NVIDIA H100 80GB HBM3") and gpu_name_matches("L4", "NVIDIA L4") + + +def test_verdict_rules(): + assert report.verdict([10, 11], [12, 13], True) == "yeto worse" + assert report.verdict([10, 11], [8, 9], True) == "yeto better" + assert report.verdict([10, 12], [11, 13], True) == "within run-to-run range" + assert report.verdict([10, 11], [8, 9], False) == "yeto worse" + assert report.verdict([10, None], [8, 9], True) == "missing value" + + +def test_price_comes_from_the_config(tmp_path): + path = tmp_path / "p.json" + path.write_text(json.dumps({"A10G": {"usd_per_card_hour": 3.5}})) + assert metrics.price("A10G", path) == 3.5 + with pytest.raises(KeyError): + metrics.price("H100!", path) + + +def test_cli_plan_and_unregistered_run(tmp_path, monkeypatch, capsys): + assert cli.main(["eff-plan", "qwen3-0.6b", "--arm", "bare-verl-sync"]) == 0 + assert json.loads(capsys.readouterr().out)["arms"]["bare-verl-sync"]["cards"]["total"] == 1 + monkeypatch.setattr(cli, "LEDGER", tmp_path / "ledger.md") + (tmp_path / "ledger.md").write_text("nothing registered\n") + with pytest.raises(SystemExit, match="not pre-registered"): + cli.main(["eff-run", "qwen3-0.6b", "--arm", "bare-verl-sync", "--run-id", "eff-x", "--out", str(tmp_path)]) + + +def test_bare_verl_async_uses_the_fork_working_lora_sync(): + script = spec.bare_script(SPEC, spec.Arm.parse("bare-verl-async"), "r") + assert "actor_rollout_ref.model.lora.merge=True" in script and "export NCCL_CUMEM_HOST_ENABLE=0" in script + assert "NCCL_CUMEM_HOST_ENABLE" not in spec.bare_script(SPEC, spec.Arm.parse("bare-verl-sync"), "r") + + +def test_parse_bare_verl_async_rows_lag_one_version(): + # fork acad9875 fully_async: weight-sync row at version N, training row of step N at N+1, + # the last step's training row never written + rows = [{"step": 0, "data": {"fully_async/rollouter/active_time": 9.0}}] + for v in range(1, 4): + rows.append({"step": v, "data": {"fully_async/rollouter/active_time": 1.0}}) + if v > 1: + rows.append({"step": v, "data": {"training/global_step": v - 1, "response_length/mean": 10.0, + "critic/score/mean": 0.5}}) + log = "\n".join(f"{100 + v}.0 (FullyAsyncTrainer) step:{v} - x:1" for v in range(4)) + p = parse.parse_verl("\n".join(json.dumps(r) for r in rows), log, samples_per_round=4, trainer_cards=1, + mode="async") + assert [r["round"] for r in p.rounds] == [1, 2, 3] + assert [r["t_end_unix"] for r in p.rounds] == [101.0, 102.0, 103.0] + assert [r["reward_mean"] for r in p.rounds] == [0.5, 0.5, None] + assert [r["samples"] for r in p.rounds] == [4, 4, 4] + assert [r["gen_s"] for r in p.rounds] == [1.0, 1.0, 1.0] # rollouter active time, not the queue pull + + +def test_bare_miles_reads_prompts_in_file_order_like_yeto(): + for arm in ("bare-miles-sync", "bare-miles-async"): + assert "--rollout-shuffle" not in spec.bare_script(SPEC, spec.Arm.parse(arm), "r") + + +def test_tape_events_skip_the_repo_copy_and_the_store(tmp_path): + row = {"event": "rl_local_round", "local_round_id": 1, "time_unix": 1.0} + for sub in ("tape-direct/l0", "yeto/tests/fixtures", "store/x"): + (tmp_path / sub).mkdir(parents=True) + (tmp_path / sub / "rl-island-0.jsonl").write_text(json.dumps({**row, "src": sub}) + "\n") + assert [r.get("src") for r in parse._tape_events(tmp_path)] == ["tape-direct/l0"] diff --git a/tools/eval/__main__.py b/tools/eval/__main__.py index 1d047988e..775d2dbcb 100644 --- a/tools/eval/__main__.py +++ b/tools/eval/__main__.py @@ -5,6 +5,13 @@ python -m tools.eval judge DIR/eval-e3-verl-a python -m tools.eval report DIR/* --out report.md [--html report.html] +Efficiency evaluation (S26 E0.1, package :mod:`.efficiency`): + + python -m tools.eval eff-plan qwen3-0.6b [--arm ARM] + python -m tools.eval eff-run qwen3-0.6b --arm bare-verl-sync --run-id eff-... [--out DIR] [--dry-run] + python -m tools.eval eff-parse DIR/eff-... + python -m tools.eval eff-report DIR/eff-* --out report.md + ``run`` refuses a run id that is not pre-registered in the GPU ledger (the id must appear in ``gpu-spend.md``), checks the local thread count and Ray, writes ``cell.json``, runs the cell's harness (raw data into ``raw/``), then judges. @@ -123,6 +130,72 @@ def cmd_report(args) -> int: return 0 +EFF_OUT = Path("/home/michael/work/infra-drafts/SESSION26/logs/E0") + + +def cmd_eff_plan(args) -> int: + from .efficiency import spec + + sp = spec.get_spec(args.preset) + arms = [args.arm] if args.arm else list(spec.ARMS) + doc = {"spec": sp.to_dict(), "fingerprint": spec.control_fingerprint(sp), + "arms": {a: spec.plan(sp, a, "") for a in arms}, + "verl_knob_diff": {m: spec.verl_knob_diff(sp, m) for m in spec.MODES}} + print(json.dumps(doc, indent=1, default=str)) + return 0 + + +def cmd_eff_run(args) -> int: + from .efficiency import modal_eff, spec + + sp = spec.get_spec(args.preset, checkpoint=args.checkpoint, + **({"profile_trainer": True} if getattr(args, "profile_trainer", False) else {})) + plan = spec.plan(sp, args.arm, args.run_id) + run_dir = Path(args.out) / args.run_id + if not args.dry_run: + if (run_dir / "arm.json").exists(): + raise SystemExit(f"{run_dir} already used") + _preflight(args.run_id) + plan["code_sha"] = subprocess.run(["git", "-C", str(Path(__file__).resolve().parents[2]), "rev-parse", "HEAD"], + capture_output=True, text=True).stdout.strip() + plan["images"] = spec.images() + spec.write_plan(run_dir / "arm.json", plan) + if plan["harness"] == "miles_local": + summary = modal_eff.run_miles_launcher(plan, run_dir, dry_run=args.dry_run) + if args.dry_run: + print((run_dir / "dry-run.log").read_text()) + return summary["rc"] + elif args.dry_run: + print(plan.get("script") or json.dumps(plan.get("argv"))) + return 0 + else: + summary = modal_eff.run(plan, run_dir) + (run_dir / "run.json").write_text(json.dumps(summary, indent=1, default=str)) + return cmd_eff_parse(argparse.Namespace(run_dir=str(run_dir))) + + +def cmd_eff_parse(args) -> int: + from .efficiency import metrics, parse + + run_dir = Path(args.run_dir) + parsed = parse.parse_run(run_dir) + (run_dir / "parsed.json").write_text(json.dumps(parsed, indent=1, default=str)) + summary = metrics.summarize(parsed) + (run_dir / "summary.json").write_text(json.dumps(summary, indent=1, default=str)) + print(json.dumps({"run": parsed["run_id"], "arm": parsed["arm"], "rounds": summary["rounds"], + "missing": summary["missing"]})) + return 0 if summary["rounds"] == summary["rounds_expected"] else 1 + + +def cmd_eff_report(args) -> int: + from .efficiency import report as eff_report + + runs = eff_report.load([Path(d) for d in args.run_dirs]) + Path(args.out).write_text(eff_report.markdown(runs, title=args.title)) + print(f"{len(runs)} parsed run(s) -> {args.out}") + return 0 + + def main(argv: list[str] | None = None) -> int: p = argparse.ArgumentParser(prog="python -m tools.eval") sub = p.add_subparsers(dest="cmd", required=True) @@ -143,6 +216,32 @@ def main(argv: list[str] | None = None) -> int: rep.add_argument("--html", default=None) rep.add_argument("--title", default="yeto RL capability evaluation") rep.set_defaults(fn=cmd_report) + from .efficiency.spec import ARMS, PRESETS + + ep = sub.add_parser("eff-plan") + ep.add_argument("preset", choices=sorted(PRESETS)) + ep.add_argument("--arm", choices=ARMS, default=None) + ep.set_defaults(fn=cmd_eff_plan) + er = sub.add_parser("eff-run") + er.add_argument("preset", choices=sorted(PRESETS)) + er.add_argument("--arm", choices=ARMS, required=True) + er.add_argument("--run-id", required=True) + er.add_argument("--out", default=str(EFF_OUT)) + er.add_argument("--dry-run", action="store_true") + er.add_argument("--checkpoint", choices=("final", "every-round", "default"), default="final", + help="yeto+Miles round cuts: only at the end (fair control) or every round " + "or the product default interval (no --rl-cut-every)") + er.add_argument("--profile-trainer", action="store_true", + help="yeto+Miles arms: --rl-profile-trainer (observation only, E1)") + er.set_defaults(fn=cmd_eff_run) + eparse = sub.add_parser("eff-parse") + eparse.add_argument("run_dir") + eparse.set_defaults(fn=cmd_eff_parse) + erep = sub.add_parser("eff-report") + erep.add_argument("run_dirs", nargs="+") + erep.add_argument("--out", required=True) + erep.add_argument("--title", default="yeto RL efficiency evaluation") + erep.set_defaults(fn=cmd_eff_report) args = p.parse_args(argv) return args.fn(args) diff --git a/tools/eval/efficiency/__init__.py b/tools/eval/efficiency/__init__.py new file mode 100644 index 000000000..54197bef1 --- /dev/null +++ b/tools/eval/efficiency/__init__.py @@ -0,0 +1,24 @@ +"""yeto RL efficiency evaluation (S26 E0.1 / E0.2). + +Eight arms per model preset: {bare, yeto} x {verl, miles} x {sync, async}. The +bare arms run the framework's own entry point from the same image (no yeto +code on their path); the yeto arms run the same framework through yeto. Every +arm is generated from one :class:`.spec.EffSpec`, so model, data, order seed, +reward, hyperparameters, max tokens, card type and count and image are equal +by construction (:func:`.spec.control_fingerprint`). + + python -m tools.eval eff-plan qwen3-0.6b [--arm ARM] + python -m tools.eval eff-run qwen3-0.6b --arm bare-verl-sync --run-id ID [--dry-run] + python -m tools.eval eff-parse RUN_DIR + python -m tools.eval eff-report RUN_DIR ... --out report.md + +* :mod:`.spec` -- presets, arms, the per-arm command plans; +* :mod:`.parse` -- framework logs and yeto tapes -> one per-round record format; +* :mod:`.metrics` -- per-run summary (throughput, segments, cost, off-policy); +* :mod:`.report` -- Markdown report with yeto-vs-bare overhead and range verdicts; +* :mod:`.modal_eff` -- the Modal harness (bare arms and yeto+verl); + yeto+Miles runs through the yeto launcher (``tools/eval/miles_local.sh``). + +Raw data is written to ``/raw/`` before anything is computed; +``parsed.json`` and the report are derived from it and can be rebuilt. +""" diff --git a/tools/eval/efficiency/algo-eff-miles.json b/tools/eval/efficiency/algo-eff-miles.json new file mode 100644 index 000000000..f45d2e970 --- /dev/null +++ b/tools/eval/efficiency/algo-eff-miles.json @@ -0,0 +1,53 @@ +{ + "advantage": { + "estimator": "grpo", + "reward_binary": false, + "reward_postprocess": null, + "rewards_normalization": true, + "std_normalization": true, + "whiten": false + }, + "correction": { + "batch_normalize": false, + "function": null, + "level": "token", + "lower": 0.0, + "method": "tis", + "mismatch_metrics": false, + "opsm_delta": null, + "ratio_mask": "none", + "upper": 2.0, + "use_rollout_logprobs": false, + "veto": null + }, + "entropy_coef": 0.0, + "execution": { + "group_required": true, + "max_policy_staleness": 2, + "needs_critic": false, + "needs_rollout_logprobs": false + }, + "kl": { + "coef": null, + "estimator": null, + "placement": "none", + "unbiased": false + }, + "loss": { + "aggregation": "default", + "custom_loss": null, + "eps_clip": 0.2, + "eps_clip_c": null, + "eps_clip_high": 0.28, + "reducer": null, + "variant": "policy_loss" + }, + "plugins": [], + "sampling": { + "filter": null, + "max_replacements": null, + "over_sampling_batch_size": null, + "overlong_filter": false + }, + "schema": "yeto-rl-algorithm-spec-v2" +} \ No newline at end of file diff --git a/tools/eval/efficiency/eff_prep.py b/tools/eval/efficiency/eff_prep.py new file mode 100644 index 000000000..94fdee80e --- /dev/null +++ b/tools/eval/efficiency/eff_prep.py @@ -0,0 +1,65 @@ +"""Model and data preparation for the bare-framework arms (no yeto import). + +Run inside the container as ``python eff_prep.py ``. It downloads +the pinned model snapshot and the pinned dataset revision, then writes the +training prompts in the dataset's own row order: + +* ``/root/eff-data/train.parquet`` and ``test.parquet`` -- verl rows, the same + shape as ``yeto.rl.adapters.verl.island_entry.prepare_data`` (the yeto+verl + arm's data), so bare verl and yeto+verl read identical rows; +* ``/root/eff-data/train.jsonl`` -- Miles rows (``prompt`` = chat messages, + ``label``), for ``--input-key prompt --label-key label --apply-chat-template``; +* ``/root/eff-data/order.json`` -- sha256 of the label sequence (row order) and + the row count, so a report can show that two arms read the same rows. + +Shuffling is left to each framework (seeded by the spec's ``seed``), as in the +yeto arm of the same framework. +""" + +from __future__ import annotations + +import hashlib +import json +import sys +from pathlib import Path + +DATA_ROOT = Path("/root/eff-data") +MODEL_ROOT = Path("/root/eff-models") +VAL_ROWS = 64 + + +def main(spec_path: str) -> int: + import pandas as pd + from huggingface_hub import hf_hub_download, snapshot_download + + spec = json.loads(Path(spec_path).read_text()) + model_dir = MODEL_ROOT / spec["model"].split("/")[-1] + snapshot_download(spec["model"], revision=spec["model_revision"], local_dir=str(model_dir)) + DATA_ROOT.mkdir(parents=True, exist_ok=True) + order = {} + for split in ("train", "test"): + src = hf_hub_download(spec["data"], f"{split}.parquet", repo_type="dataset", revision=spec["data_revision"]) + frame = pd.read_parquet(src) + rows = [{ + "data_source": "yeto/" + spec["data"], + "prompt": [{"role": m["role"], "content": m["content"]} for m in row["messages"]], + "ability": "math", + "reward_model": {"style": "rule", "ground_truth": str(row["label"])}, + "extra_info": {"split": split, "index": int(i)}, + } for i, row in frame.iterrows()] + if split == "test": + rows = rows[:VAL_ROWS] + pd.DataFrame(rows).to_parquet(DATA_ROOT / f"{split}.parquet") + labels = [r["reward_model"]["ground_truth"] for r in rows] + order[split] = {"rows": len(rows), "label_sha256": hashlib.sha256("\n".join(labels).encode()).hexdigest()} + if split == "train": + with (DATA_ROOT / "train.jsonl").open("w") as fh: + for r in rows: + fh.write(json.dumps({"prompt": r["prompt"], "label": r["reward_model"]["ground_truth"]}) + "\n") + (DATA_ROOT / "order.json").write_text(json.dumps({**order, "model_dir": str(model_dir)}, indent=1)) + print(json.dumps({"model_dir": str(model_dir), **order})) + return 0 + + +if __name__ == "__main__": + sys.exit(main(sys.argv[1])) diff --git a/tools/eval/efficiency/eff_reward.py b/tools/eval/efficiency/eff_reward.py new file mode 100644 index 000000000..4358028ac --- /dev/null +++ b/tools/eval/efficiency/eff_reward.py @@ -0,0 +1,39 @@ +"""GSM8K exact-number reward for the bare-framework arms (no yeto import). + +The grading rule is a verbatim copy of ``yeto.rl.rewards.builtin.gsm8k_grade`` +(the reward of the yeto arms, ``--reward-function gsm8k_reward:score``); a CPU +test asserts both give the same value on the same cases. The file is copied +into the bare containers as ``/root/eff_reward.py``: + +* verl: ``reward.custom_reward_function.path=/root/eff_reward.py`` and + ``reward.custom_reward_function.name=compute_score``; +* Miles: ``--custom-rm-path eff_reward.score`` (``/root`` on ``PYTHONPATH``). +""" + +from __future__ import annotations + +import re + + +def _num(s): + s = str(s).replace(",", "") + m = re.findall(r"-?\d+(?:\.\d+)?", s) + return float(m[-1]) if m else None + + +def grade(response: str, label) -> float: + gold = _num(str(label).split("####")[-1]) if label is not None else None + text = response or "" + boxed = re.findall(r"\\boxed\{([^}]*)\}", text) + pred = _num(boxed[-1]) if boxed else _num(text[-200:]) + return 1.0 if (gold is not None and pred is not None and abs(pred - gold) < 1e-6) else 0.0 + + +def compute_score(data_source, solution_str, ground_truth, extra_info=None, **kwargs) -> float: + """verl custom reward signature.""" + return grade(solution_str, ground_truth) + + +async def score(args, sample, **kwargs) -> float: + """Miles custom reward signature (``--custom-rm-path``).""" + return grade(sample.response or "", sample.label) diff --git a/tools/eval/efficiency/metrics.py b/tools/eval/efficiency/metrics.py new file mode 100644 index 000000000..0df3bfa22 --- /dev/null +++ b/tools/eval/efficiency/metrics.py @@ -0,0 +1,141 @@ +"""Per-run summary from ``parsed.json`` (the definitions every arm shares). + +* Steady window: the rounds after the first (the first carries engine start-up, + CUDA-graph capture and caches); ``steady_wall_s`` = the sum of their + ``round_s`` (step end minus previous step end). +* Throughput per card: steady tokens / steady wall / every card of the arm + (rollout and trainer cards both count). +* Segments: per-round medians over the steady window; ``unaccounted_s`` = round + minus the named segments (offload/onload, reward, advantage, orchestration). +* GPU: NVML samples inside the steady window; ``bubble_frac`` = the share of + card samples with utilization below :data:`BUBBLE_UTIL_PCT`. +* Cost: ``usd_per_card_hour`` from ``prices.json`` x cards x container seconds + (whole run) and x steady wall (steady state), per 1k trained samples. + +A value whose inputs are missing is ``None`` and its name is listed in +``missing``; nothing is filled in. +""" + +from __future__ import annotations + +import json +import statistics +from pathlib import Path +from typing import Any + +PRICES = Path(__file__).with_name("prices.json") +BUBBLE_UTIL_PCT = 5.0 +SEGMENTS = ("gen_s", "train_s", "weight_sync_s", "sampler_update_s", "wait_s") + + +def price(gpu: str, path: Path = PRICES) -> float: + table = json.loads(path.read_text()) + if gpu not in table: + raise KeyError(f"no price for GPU {gpu!r} in {path}") + return float(table[gpu]["usd_per_card_hour"]) + + +def _vals(rows: list[dict], key: str) -> list[float]: + return [float(r[key]) for r in rows if r.get(key) is not None] + + +def _median(values: list[float]) -> float | None: + return round(statistics.median(values), 4) if values else None + + +def _sum_all(rows: list[dict], key: str) -> float | None: + """Sum over rows, or None when any row lacks the value (a partial sum would mislead).""" + vals = [r.get(key) for r in rows] + return None if not vals or any(v is None for v in vals) else float(sum(vals)) + + +def summarize(parsed: dict[str, Any], *, prices_path: Path = PRICES) -> dict[str, Any]: + rounds = sorted(parsed["rounds"], key=lambda r: r["round"]) + cards = int(parsed["cards"]["total"]) + steady = rounds[1:] + s: dict[str, Any] = {"arm": parsed["arm"], "run_id": parsed["run_id"], "rounds": len(rounds), + "rounds_expected": parsed["spec"]["rounds"], "cards": cards, "gpu": parsed["gpu"], + "rc": parsed.get("rc")} + wall = _sum_all(steady, "round_s") + s["steady_wall_s"] = round(wall, 3) if wall is not None else None + s["rate_rounds"] = {} + for key, out in (("gen_tokens", "gen_tok_per_card_s"), ("train_tokens", "train_tok_per_card_s"), + ("samples", "samples_per_s")): + # over the steady rounds that report the value (bare verl fully_async never + # logs its last step's batch); the round count is kept when it is fewer + known = [r for r in steady if r.get(key) is not None and r.get("round_s") is not None] + time_s = sum(r["round_s"] for r in known) * (cards if key != "samples" else 1) + s[out] = round(sum(r[key] for r in known) / time_s, 3) if known and time_s else None + if known and len(known) < len(steady): + s["rate_rounds"][out] = f"{len(known)}/{len(steady)}" + s["round_s_median"] = _median(_vals(steady, "round_s")) + s["round_s_range"] = [min(_vals(steady, "round_s")), max(_vals(steady, "round_s"))] if _vals(steady, "round_s") \ + else None + for seg in SEGMENTS: + s[f"{seg}_median"] = None if parsed["sources"].get(seg) in ("n/a",) else _median(_vals(steady, seg)) + named = [seg for seg in SEGMENTS if parsed["sources"].get(seg) != "n/a"] + rest = [r["round_s"] - sum(r[seg] for seg in named) for r in steady + if r.get("round_s") is not None and all(r.get(seg) is not None for seg in named)] + s["unaccounted_s_median"] = _median(rest) + s["segments_raw_median"] = _raw_medians(parsed.get("segments_raw") or []) + # reward against wall time and samples (every round; t from the first step end) + t0 = rounds[0].get("t_end_unix") if rounds else None + cum, curve = 0.0, [] + for r in rounds: + cum = cum + r["samples"] if r.get("samples") is not None and cum is not None else None + curve.append({"round": r["round"], "t_s": round(r["t_end_unix"] - t0, 3) + if r.get("t_end_unix") is not None and t0 is not None else None, + "samples": cum, "reward_mean": r.get("reward_mean")}) + s["reward_curve"] = curve + s["reward_last3_mean"] = _median(_vals(rounds[-3:], "reward_mean")) # median of the last three + # GPU (NVML samples inside the steady window) + lo = rounds[0].get("t_end_unix") if rounds else None + hi = rounds[-1].get("t_end_unix") if rounds else None + nv = [x for x in parsed.get("nvml") or [] if lo is not None and hi is not None and lo <= x["t_unix"] <= hi] + util = [x["util"] for x in nv if x.get("util") is not None] + s["nvml_samples"] = len(util) + s["nvml_interval_s"] = parsed.get("nvml_interval_s") + s["gpu_util_mean"] = round(statistics.fmean(util), 2) if util else None + s["bubble_frac"] = round(sum(u < BUBBLE_UTIL_PCT for u in util) / len(util), 4) if util else None + mem = [x["mem_mb"] for x in parsed.get("nvml") or [] if x.get("mem_mb") is not None] + s["gpu_mem_peak_mb"] = max(mem) if mem else None + # cost + rate = price(parsed["gpu"], prices_path) + s["usd_per_card_hour"] = rate + start, end = parsed.get("container_start_unix"), parsed.get("container_end_unix") + total_samples = _sum_all(rounds, "samples") + if start and end: + s["usd_run"] = round(rate * cards * (end - start) / 3600, 4) + s["usd_per_1k_samples_run"] = round(s["usd_run"] / total_samples * 1000, 4) if total_samples else None + else: + s["usd_run"] = s["usd_per_1k_samples_run"] = None + steady_samples = _sum_all(steady, "samples") + s["usd_per_1k_samples_steady"] = round(rate * cards * wall / 3600 / steady_samples * 1000, 4) \ + if wall and steady_samples else None + # off-policy + s["is_clip_frac_mean"] = round(statistics.fmean(_vals(rounds, "is_clip_frac")), 5) \ + if _vals(rounds, "is_clip_frac") else None + dropped, generated = _vals(rounds, "dropped_samples"), _vals(rounds, "generated_samples") + s["dropped_total"] = max(dropped) if dropped and "count/" in parsed["sources"].get("dropped_samples", "") \ + else (sum(dropped) if dropped else None) # verl's fully_async counters are cumulative + s["drop_frac"] = round(s["dropped_total"] / max(generated), 4) if s["dropped_total"] is not None and generated \ + and max(generated) else None + s["version_lag"] = [r.get("version_lag") for r in rounds if r.get("version_lag") is not None] or None + s["sources"] = parsed["sources"] + na = {k for k, v in parsed["sources"].items() if v == "n/a"} + s["not_applicable"] = sorted(k for k in s if s[k] is None and _field_of(k) in na) + s["missing"] = sorted(k for k, v in s.items() if v is None and k != "rc" and k not in s["not_applicable"]) + return s + + +# summary key -> the round field it is computed from (for the not-applicable list) +_DERIVED_FROM = {"dropped_total": "dropped_samples", "drop_frac": "dropped_samples"} + + +def _field_of(key: str) -> str: + return _DERIVED_FROM.get(key, key.removesuffix("_median")) + + +def _raw_medians(rows: list[dict[str, float]]) -> dict[str, float]: + keys = sorted({k for r in rows[1:] for k in r}) + return {k: _median([float(r[k]) for r in rows[1:] if r.get(k) is not None]) for k in keys} diff --git a/tools/eval/efficiency/modal_eff.py b/tools/eval/efficiency/modal_eff.py new file mode 100644 index 000000000..d67dd09e8 --- /dev/null +++ b/tools/eval/efficiency/modal_eff.py @@ -0,0 +1,182 @@ +"""Modal harness of the efficiency arms that run in one container: the four bare +arms and yeto+verl (yeto+Miles runs through the yeto launcher, ``miles_local.sh``). + +One GPU container per run, ``:``, in the framework's image (the +image the yeto arm of that framework uses; for verl the bare arms skip yeto's +source patch). The container: + +1. checks every visible GPU is the asked type (``modal_verl._gpu``); +2. samples NVML every second into ``nvml.csv`` (index, util, memory, power); +3. bare: writes the spec, ``eff_reward.py`` and ``eff_prep.py`` under + ``/root/eff`` and runs the arm's bash script (prep, then the framework entry); + yeto+verl: starts the island Ray and runs ``island_entry`` (as T-EVAL E1); +4. records ``meta.json`` (start/end unix time, rc, GPU names, argv/script) and + copies its outputs into the results Volume ``/arm/`` before returning. + +The local side spawns the call in a detached app, waits, pulls ``/`` into +``/raw/`` and stops only its own app (``yeto-eff-``). +""" + +from __future__ import annotations + +import json +import os +import subprocess +import sys +import time +from pathlib import Path +from typing import Any + +from .. import modal_verl as mv +from .spec import EFF_DIR, OUT_DIR + +HERE = Path(__file__).resolve().parent +APP_PREFIX = "yeto-eff-" +NVML_QUERY = "timestamp,index,name,utilization.gpu,memory.used,memory.total,power.draw" +NVML_INTERVAL_MS = 1000 + + +# ====================================================================== in-container +def _nvml(path: Path) -> subprocess.Popen: + return subprocess.Popen(["bash", "-c", f"nvidia-smi --query-gpu={NVML_QUERY} --format=csv " + f"-lms {NVML_INTERVAL_MS} >> {path}"]) + + +def _run_logged(argv: list[str], log: Path, timeout: float, **kwargs) -> int: + """Run ``argv``; every output line is written to ``log`` behind `` `` + (the step end times the parsers read). The process is killed at ``timeout``.""" + deadline = time.time() + timeout + with log.open("a") as fh, subprocess.Popen(argv, stdout=subprocess.PIPE, stderr=subprocess.STDOUT, text=True, + errors="replace", **kwargs) as proc: + for line in proc.stdout: + fh.write(f"{time.time():.3f} {line}") + fh.flush() + if time.time() > deadline: + proc.kill() + fh.write(f"{time.time():.3f} [eff harness] killed at the {timeout:.0f} s timeout\n") + break + return proc.wait() + + +def arm(run: str, plan: dict) -> dict: + """One run of one arm (see the module doc).""" + out = Path(mv.RESULTS) / run / "arm" + out.mkdir(parents=True, exist_ok=True) + meta: dict[str, Any] = {"arm": plan["arm"], "container": os.environ.get("MODAL_TASK_ID"), + "cards": plan["cards"]["total"]} + meta["gpu"] = mv._gpu(out, {"gpu": plan["gpu"]}) + nvml = _nvml(out / "nvml.csv") + meta["start_unix"] = time.time() + copies = {"eff-out": Path(OUT_DIR), "yeto-output": Path("/root/yeto-output")} + budget = plan["spec"]["timeout_s"] - 300 # leaves time to copy and commit + log = out / "main.log" + try: + if plan["system"] == "bare": + eff = Path(EFF_DIR) + eff.mkdir(parents=True, exist_ok=True) + for name, text in plan["files"].items(): + (eff / name).write_text(text) + (eff / "spec.json").write_text(json.dumps(plan["spec"])) + (out / "script.sh").write_text(plan["script"]) + meta["rc"] = _run_logged(["bash", str(out / "script.sh")], log, budget) + else: + sys.path.insert(0, mv.WORKDIR) + env = mv._island_env({"gpu": plan["gpu"]}, 0) + ray = str(Path(mv.VERL_PY).with_name("ray")) + _run_logged([ray, "start", "--head", "--node-ip-address=127.0.0.1", "--port=6379", "--num-cpus=32", + f"--num-gpus={plan['cards']['total']}", "--include-dashboard=false", + "--temp-dir=/root/eval-ray"], log, 300, env=env) + meta["argv"] = [mv.VERL_PY, *plan["argv"]] + meta["rc"] = _run_logged(meta["argv"], log, budget, env=env, cwd="/root") + except BaseException as exc: # noqa: BLE001 - recorded, then re-raised + meta["error"] = f"{type(exc).__name__}: {exc}" + raise + finally: + meta["end_unix"] = time.time() + nvml.terminate() + mv._finish(out, meta, meta["start_unix"], copies) + return meta + + +# ====================================================================== app +def build_app(run: str, plan: dict): + import modal + + sys.path.insert(0, str(mv.REPO)) + from modal._vendor import cloudpickle + + # the container functions travel by value: `tools` is not importable in the images + cloudpickle.register_pickle_by_value(sys.modules[__name__]) + cloudpickle.register_pickle_by_value(mv) + from yeto.rl.adapters.miles.pins import MILES_NEXT_IMAGE + from yeto.rl.adapters.verl.image import modal_image + + if plan["image"] == "verl": + # bare arms run the fork as shipped: yeto's patch replaces the LoRA weight-sync + # sender with an import of yeto (eff01-bare-verl-async-a/d) + image = modal_image(modal, patched=plan["system"] == "yeto").env({"PYTHONUNBUFFERED": "1"}) + if plan["system"] == "yeto": + image = image.add_local_dir(str(mv.REPO), mv.WORKDIR, copy=False, ignore=MOUNT_IGNORE) + else: + image = (modal.Image.from_registry(MILES_NEXT_IMAGE.removeprefix("docker:")) + .env({"HOME": "/root", "PYTHONUNBUFFERED": "1"})) + app = modal.App(APP_PREFIX + run) + results = modal.Volume.from_name(mv.RESULTS_VOLUME, create_if_missing=True, version=2) + gpu = f"{plan['gpu']}:{plan['cards']['total']}" + fn = app.function(image=image, gpu=gpu, cpu=8.0 * plan["cards"]["total"], + memory=49152 * plan["cards"]["total"], timeout=plan["spec"]["timeout_s"], retries=0, + volumes={mv.RESULTS: results}, single_use_containers=True, name="arm", serialized=True)(arm) + return app, fn + + +MOUNT_IGNORE = [".git", ".venv", ".claude", "**/__pycache__", "**/*.pyc", "s1-runs", "syncer/target", + "evidence*", "openspec", "docs", "tutorial", "data"] + + +def bare_files() -> dict[str, str]: + return {name: (HERE / name).read_text() for name in ("eff_reward.py", "eff_prep.py")} + + +def run(plan: dict, run_dir: Path) -> dict[str, Any]: + """Run one arm end to end (spawn, wait, pull, stop own app).""" + import modal + + run_id = plan["run_id"] + raw = run_dir / "raw" + raw.mkdir(parents=True, exist_ok=True) + log = run_dir / "driver.log" + if plan["system"] == "bare": + plan = {**plan, "files": bare_files()} + app, fn = build_app(run_id, plan) + began = time.time() + with modal.enable_output(), app.run(detach=True): # a killed driver never cancels a spawned call + app_id = app.app_id + (run_dir / "app_id.txt").write_text(f"{app_id}\n") + call = fn.spawn(run_id, plan) + mv._log(log, f"spawned arm {plan['arm']} app {APP_PREFIX + run_id}") + results = mv._wait_calls({"arm": call}, began + plan["spec"]["timeout_s"] + 600, log) + pulled = mv._modal_cli("volume", "get", "--force", mv.RESULTS_VOLUME, run_id, str(raw), timeout=1800) + stopped = mv._modal_cli("app", "stop", "-y", app_id, timeout=120) + summary = {"calls": results, "wall_s": round(time.time() - began, 1), "pull_rc": pulled.returncode, + "stop_rc": stopped.returncode} + (run_dir / "driver.json").write_text(json.dumps(summary, indent=1, default=str)) + return summary + + +# ====================================================================== yeto+Miles (launcher) +def run_miles_launcher(plan: dict, run_dir: Path, *, dry_run: bool = False) -> dict[str, Any]: + """The yeto+Miles arm: ``miles_local.sh`` with the plan's launcher arguments + (the T-EVAL Miles path; it pulls launch log, tapes and NVML into ``raw/``).""" + from .. import miles + + raw = run_dir / "raw" + env = dict(os.environ, P=plan["run_id"], R=str(raw), REPO=str(miles.REPO), PY=miles.PY, SKY=miles.SKY, + STOREVOL=miles.STORE_VOLUME, TAPEVOL=miles.TAPE_VOLUME, HARD=str(plan["spec"]["timeout_s"]), + HOME="/home/michael", ARGS=plan["launcher_args"], ISLANDS="1") + if dry_run: + env["DRY"] = "1" + log = run_dir / ("dry-run.log" if dry_run else "driver.log") + with log.open("w") as fh: + rc = subprocess.run(["bash", str(miles.HERE / "miles_local.sh")], env=env, stdout=fh, + stderr=subprocess.STDOUT, cwd="/tmp").returncode + return {"script": "miles_local.sh", "rc": rc, "log": str(log)} diff --git a/tools/eval/efficiency/parse.py b/tools/eval/efficiency/parse.py new file mode 100644 index 000000000..120d398b3 --- /dev/null +++ b/tools/eval/efficiency/parse.py @@ -0,0 +1,478 @@ +"""Framework logs and yeto tapes -> one per-round record format. + +Every parser returns a :class:`Parsed`: ``rounds`` (one dict per optimizer +step, the fields of :data:`ROUND_FIELDS`), ``sources`` (per field: the log key +it came from, ``"derived: ..."`` when it is computed from logged values by an +exact identity, ``MISSING`` when no source in the raw data carries it, or +``NA`` when the field does not apply to the arm), and ``segments_raw`` (every +native timing per round, for the overhead table). + +A field that a log does not carry stays ``None``; nothing is estimated. + +Formats read (each one is cited where it is parsed): + +* verl ``FileLogger`` JSONL (``VERL_FILE_LOGGER_PATH``): one + ``{"step": N, "data": {...}}`` per step (fork acad9875 ``verl/utils/tracking.py`` + ``FileLogger.log``); keys from ``metric_utils.compute_timing_metrics`` / + ``compute_throughout_metrics`` (``timing_s/*``, ``perf/*``) and the fully_async + trainer (``fully_async/*``, ``timing_s/param_sync``, ``timing_s/wait_for_enough_samples``); +* verl console lines ``step:N - key:value - ...`` (``LocalLogger``, + ``concat_dict_to_str``) behind the harness's `` `` line prefix, for the + step end time; +* Miles ``perf N: {python dict}`` log lines: rollout side + (``miles/ray/rollout/metrics.py:93``, keys ``rollout/*``, ``perf/rollout_time``, + ``perf/tokens_per_gpu_per_sec``) and train side + (``miles/utils/train_metric_utils.py:56``, keys ``perf/_time``, + ``perf/actor_train_tok_per_s``); the line's ``[YYYY-mm-dd HH:MM:SS.mmm`` stamp + (Miles' logger format) gives the time; +* yeto event tapes (JSONL ``{"event": ...}``): ``rl_round_wall``, + ``rl_local_round``, ``rl_round_trained``, ``rl_timeline_span``, + ``rl_driver_phase``, ``rl_rollout_carry_over``, ``rl_resource_sample``; +* NVML CSV (``nvidia-smi --query-gpu=timestamp,index,name,utilization.gpu, + memory.used,memory.total,power.draw --format=csv``), with or without a + ``,`` prefix per line (``miles_local.sh`` polls every 20 s). +""" + +from __future__ import annotations + +import ast +import csv +import datetime as _dt +import io +import json +import re +from collections import defaultdict +from dataclasses import dataclass, field +from pathlib import Path +from typing import Any, Iterable + +MISSING = "missing" +NA = "n/a" +ROUND_FIELDS = ( + "round", "t_end_unix", "round_s", "gen_s", "train_s", "weight_sync_s", "sampler_update_s", "wait_s", + "gen_tokens", "train_tokens", "samples", "reward_mean", "is_clip_frac", "dropped_samples", + "generated_samples", "version_lag", +) +ANSI = re.compile(r"\x1b\[[0-9;]*m") + + +@dataclass +class Parsed: + rounds: list[dict[str, Any]] = field(default_factory=list) + sources: dict[str, str] = field(default_factory=dict) + segments_raw: list[dict[str, float]] = field(default_factory=list) + + def to_dict(self) -> dict[str, Any]: + return {"rounds": self.rounds, "sources": self.sources, "segments_raw": self.segments_raw} + + +def _first(data: dict, keys: Iterable[str]) -> tuple[Any, str | None]: + for key in keys: + if data.get(key) is not None: + return data[key], key + return None, None + + +def _finish(parsed: Parsed, na: Iterable[str] = ()) -> Parsed: + """Fill ``sources`` for fields no round carries (missing or not applicable).""" + na = set(na) + parsed.sources["round"] = "step index (verl global step; Miles / yeto rollout id)" + for name in ROUND_FIELDS[1:]: + if name in na: + parsed.sources[name] = NA + for r in parsed.rounds: + r[name] = None + elif name not in parsed.sources or all(r.get(name) is None for r in parsed.rounds): + parsed.sources[name] = MISSING + return parsed + + +def _record(rnd: int) -> dict[str, Any]: + return {name: None for name in ROUND_FIELDS} | {"round": rnd} + + +def _set(parsed: Parsed, rec: dict, name: str, value: Any, source: str | None) -> None: + if value is None or source is None: + return + rec[name] = value + parsed.sources.setdefault(name, source) + + +def _step_ends(log_text: str, pattern: re.Pattern) -> dict[int, float]: + """Round -> unix time of its log line (lines carry a `` `` prefix).""" + out: dict[int, float] = {} + for line in log_text.splitlines(): + m = pattern.search(ANSI.sub("", line)) + if m: + out[int(m.group(2))] = float(m.group(1)) + return out + + +def _round_s_from_ends(parsed: Parsed) -> None: + """round_s = time between consecutive step ends (one definition for every arm).""" + prev = None + for r in parsed.rounds: + t = r.get("t_end_unix") + if t is not None and prev is not None: + r["round_s"] = round(t - prev, 3) + parsed.sources.setdefault("round_s", "derived: step end - previous step end") + prev = t + + +# ====================================================================== verl +VERL_STEP_LINE = re.compile(r"^(\d+(?:\.\d+)?) .*?\bstep:(\d+) - ") + + +def parse_verl(metrics_jsonl: str, log_text: str = "", *, samples_per_round: int, trainer_cards: int, + mode: str) -> Parsed: + """verl FileLogger rows (+ console step lines for end times).""" + p = Parsed() + ends = _step_ends(log_text, VERL_STEP_LINE) + recs: dict[int, dict[str, Any]] = {} + segs: dict[int, dict[str, float]] = {} + for line in metrics_jsonl.splitlines(): + if not line.strip(): + continue + row = json.loads(line) + d, step = row.get("data") or {}, int(row["step"]) + rnd = step + if mode == "async" and "training/global_step" in d: + # bare fully_async logs a step's training metrics one parameter version + # later (row step N+1, training/global_step N) and never logs the last + # step's; the rows without it carry the weight sync that ends version N. + rnd = int(d["training/global_step"]) + elif mode == "async" and step == 0: + continue # the weight sync before the first training step + rec = recs.setdefault(rnd, _record(rnd)) + if rnd == step: + _set(p, rec, "t_end_unix", ends.get(step), "console step line (harness time prefix)") + # fully_async: the trainer's timing_s/gen only pulls finished samples from the + # queue; the rollouter's active time in the version window is the generation + gen, key = _first(d, ("fully_async/rollouter/active_time",) if mode == "async" else ("timing_s/gen",)) + _set(p, rec, "gen_s", gen, key) + parts = [k for k in ("timing_s/old_log_prob", "timing_s/update_actor") if d.get(k) is not None] + if parts: + _set(p, rec, "train_s", sum(float(d[k]) for k in parts), " + ".join(parts)) + sync, key = _first(d, ("timing_s/update_weights", "timing_s/param_sync")) + _set(p, rec, "weight_sync_s", sync, key) + wait, key = _first(d, ("timing_s/wait_for_enough_samples",)) + _set(p, rec, "wait_s", wait, key) + if d.get("response_length/mean") is not None: # every step trains exactly samples_per_round responses + _set(p, rec, "gen_tokens", float(d["response_length/mean"]) * samples_per_round, + f"derived: response_length/mean x {samples_per_round} responses") + _set(p, rec, "samples", samples_per_round, "config: samples per step (verl trains a fixed batch)") + tokens, key = _first(d, ("perf/total_num_tokens",)) + if tokens is None and d.get("global_seqlen/mean") is not None: + tokens, key = float(d["global_seqlen/mean"]) * trainer_cards, \ + f"derived: global_seqlen/mean x {trainer_cards} trainer ranks" + _set(p, rec, "train_tokens", tokens, key) + reward, key = _first(d, ("critic/score/mean", "critic/rewards/mean")) + _set(p, rec, "reward_mean", reward, key) + clip, key = _first(d, ("rollout_corr/rollout_is_ratio_fraction_high", "yeto_mismatch/tis_clipfrac")) + _set(p, rec, "is_clip_frac", clip, key) + dropped, key = _first(d, ("fully_async/count/dropped_stale_samples",)) + _set(p, rec, "dropped_samples", dropped, key) + generated, key = _first(d, ("fully_async/count/total_generated_samples",)) + _set(p, rec, "generated_samples", generated, key) + segs.setdefault(rnd, {}).update({k: float(v) for k, v in d.items() if k.startswith("timing_s/") + and isinstance(v, (int, float))}) + for rec in recs.values(): # a step whose metrics verl never logged still trained one fixed batch + _set(p, rec, "samples", samples_per_round if rec["samples"] is None else rec["samples"], + "config: samples per step (verl trains a fixed batch)") + p.rounds = [recs[r] for r in sorted(recs)] + p.segments_raw = [segs[r] for r in sorted(recs)] + _round_s_from_ends(p) + na = {"sampler_update_s"} + if mode == "sync": + na |= {"wait_s", "dropped_samples", "generated_samples", "version_lag"} + return _finish(p, na) + + +# ====================================================================== Miles +MILES_DICT_LINE = re.compile(r"\b(perf|step|rollout) (\d+): (\{.*\})\s*$") +MILES_STAMP = re.compile(r"\[(\d{4}-\d\d-\d\d \d\d:\d\d:\d\d\.\d+)") + + +def _miles_stamp(line: str) -> float | None: + m = MILES_STAMP.search(line) + if not m: + return None + return _dt.datetime.strptime(m.group(1), "%Y-%m-%d %H:%M:%S.%f").replace(tzinfo=_dt.timezone.utc).timestamp() + + +def miles_dicts(log_text: str) -> dict[int, dict[str, Any]]: + """Rollout id -> {"rollout": {...}, "train": {...}, "step": {...}, "train_t": unix}.""" + out: dict[int, dict[str, Any]] = defaultdict(dict) + for raw in log_text.splitlines(): + line = ANSI.sub("", raw) + m = MILES_DICT_LINE.search(line) + if not m: + continue + try: + data = ast.literal_eval(m.group(3)) + except (ValueError, SyntaxError): + continue + rid = int(m.group(2)) + if m.group(1) == "perf" and ("perf/rollout_time" in data or any(k.startswith("rollout/") for k in data)): + out[rid].setdefault("rollout", {}).update(data) + elif m.group(1) == "perf": + out[rid].setdefault("train", {}).update(data) + out[rid]["train_t"] = _miles_stamp(line) + else: + out[rid].setdefault("step", {}).update(data) + return dict(out) + + +def parse_miles(log_text: str, *, rollout_cards: int, mode: str) -> Parsed: + p = Parsed() + for rid, parts in sorted(miles_dicts(log_text).items()): + ro, tr, st = parts.get("rollout", {}), parts.get("train", {}), parts.get("step", {}) + merged = {**ro, **tr, **st} + rec = _record(rid) + _set(p, rec, "t_end_unix", parts.get("train_t"), "Miles train perf line time stamp") + gen, key = _first(ro, ("perf/rollout_time",)) + _set(p, rec, "gen_s", gen, key) + train, key = _first(tr, ("perf/train_time",)) + _set(p, rec, "train_s", train, key) + sync, key = _first(tr, ("perf/update_weights_time",)) + _set(p, rec, "weight_sync_s", sync, key) + if mode == "async": + wait, key = _first(tr, ("perf/train_wait_time",)) + _set(p, rec, "wait_s", wait, key) + if ro.get("perf/tokens_per_gpu_per_sec") is not None and gen is not None: + _set(p, rec, "gen_tokens", float(ro["perf/tokens_per_gpu_per_sec"]) * float(gen) * rollout_cards, + f"derived: perf/tokens_per_gpu_per_sec x perf/rollout_time x {rollout_cards} rollout cards") + if tr.get("perf/actor_train_tok_per_s") is not None and tr.get("perf/actor_train_time") is not None: + _set(p, rec, "train_tokens", float(tr["perf/actor_train_tok_per_s"]) * float(tr["perf/actor_train_time"]), + "derived: perf/actor_train_tok_per_s x perf/actor_train_time") + samples, key = _first(merged, ("rollout/num_training_samples", "num_training_samples")) + _set(p, rec, "samples", samples, key) + reward, key = _first(merged, ("rollout/raw_reward", "rollout/episode_raw_reward", "rollout/rewards")) + _set(p, rec, "reward_mean", reward, key) + clip, key = _first(merged, ("train/tis_clipfrac", "train/tis_clip_frac", "tis_clipfrac")) + _set(p, rec, "is_clip_frac", clip, key) + lag = {k: v for k, v in merged.items() if "staleness" in k or "weight_version" in k} + _set(p, rec, "version_lag", lag or None, "Miles keys containing 'staleness' / 'weight_version'") + p.rounds.append(rec) + p.segments_raw.append({k: float(v) for k, v in merged.items() + if k.startswith("perf/") and k.endswith("_time") and isinstance(v, (int, float))}) + _round_s_from_ends(p) + na = {"sampler_update_s"} + if mode == "sync": + na |= {"wait_s", "dropped_samples", "generated_samples", "version_lag"} + return _finish(p, na) + + +# ====================================================================== yeto tape +def tape_rows(text: str) -> list[dict[str, Any]]: + rows = [] + for line in text.splitlines(): + line = line.strip() + start = line.find("{") + if start < 0: + continue + try: + row = json.loads(line[start:]) + except json.JSONDecodeError: + continue + if isinstance(row, dict) and row.get("event"): + rows.append(row) + return rows + + +# yeto task / phase names -> canonical segment (others stay in segments_raw only) +YETO_SEGMENT = {"generate": "gen_s", "train": "train_s", "publish": "weight_sync_s"} + + +def parse_yeto(events: list[dict[str, Any]], *, mode: str, framework_parsed: Parsed | None = None, + framework_round_offset: int = 0, skip: frozenset[str] = frozenset(), + rollouter_gen: bool = False) -> Parsed: + """yeto tape -> rounds; ``framework_parsed`` (the framework's own log in the + same run, e.g. Miles perf lines) fills fields the tape does not carry. Its + round ``rollout_id + framework_round_offset`` is the tape's ``rollout_id``: + 0 for Miles (perf ids are rollout ids), 1 for verl (FileLogger step = + global step, 1-based; T-EVAL eval-e1-verl-c: steps 1..4, rollout ids 0..3). + + ``rollouter_gen`` (yeto+verl fully_async): ``gen_s`` is the rollouter's + ``fully_async/rollouter/active_time`` from ``rl_publication_check.rollouter_timing``; + the check of policy version v closes round v - 1, the same window as the bare + arm's per-step value.""" + p = Parsed() + by = defaultdict(list) + for e in events: + by[e["event"]].append(e) + walls = {int(e["rollout_id"]): e for e in by["rl_round_wall"] if e.get("rollout_id") is not None} + trained = {int(e["rollout_id"]): e for e in by["rl_round_trained"] if e.get("rollout_id") is not None} + local = sorted(by["rl_local_round"], key=lambda e: float(e.get("time_unix") or 0)) + spans: dict[int, dict[str, float]] = defaultdict(lambda: defaultdict(float)) + for s in by["rl_timeline_span"]: + if s.get("rollout_id") is None or s.get("start") is None or s.get("end") is None: + continue + name = "wait" if s.get("kind") == "wait" else str(s.get("task")) + spans[int(s["rollout_id"])][name] += float(s["end"]) - float(s["start"]) + phases = _driver_phase_durations(by["rl_driver_phase"]) + drops = defaultdict(float) + for e in by["rl_rollout_carry_over"]: + if e.get("rollout_id") is not None and e.get("over_age_discarded_groups") is not None: + drops[int(e["rollout_id"])] += float(e["over_age_discarded_groups"]) + rollouter = _rollouter_active(by["rl_publication_check"]) if rollouter_gen else {} + rounds = sorted(set(walls) | set(trained)) + fw = {r["round"]: r for r in (framework_parsed.rounds if framework_parsed else [])} + for i, rid in enumerate(rounds): + rec = _record(rid) + w, t = walls.get(rid, {}), trained.get(rid, {}) + lr = local[i] if i < len(local) else {} + _set(p, rec, "t_end_unix", w.get("time_unix"), "rl_round_wall.time_unix") + seg = dict(spans.get(rid) or phases.get(rid) or {}) + src = "rl_timeline_span" if rid in spans else "rl_driver_phase (marker to next marker)" + for task, name in YETO_SEGMENT.items(): + if name not in skip: + _set(p, rec, name, seg.get(task), f"{src} {task}") + if rid in rollouter: + _set(p, rec, "gen_s", rollouter[rid], f"rl_publication_check.rollouter_timing {ROLLOUTER_ACTIVE}") + if mode == "async": + _set(p, rec, "wait_s", seg.get("wait"), f"{src} kind=wait") + _set(p, rec, "gen_tokens", lr.get("action_tokens"), "rl_local_round.action_tokens") + _set(p, rec, "samples", t.get("trained_samples"), "rl_round_trained.trained_samples") + _set(p, rec, "reward_mean", lr.get("reward_mean"), "rl_local_round.reward_mean") + _set(p, rec, "is_clip_frac", (t.get("mismatch") or {}).get("tis_clipfrac"), "rl_round_trained.mismatch.tis_clipfrac") + if mode == "async": + _set(p, rec, "dropped_samples", drops.get(rid), "rl_rollout_carry_over.over_age_discarded_groups (groups)") + f = fw.get(rid + framework_round_offset, {}) + for name in ROUND_FIELDS: # what the tape lacks, from the framework's own log of the same run + if rec[name] is None and f.get(name) is not None and name not in ("round", "t_end_unix", "round_s") \ + and name not in skip: + _set(p, rec, name, f[name], "framework log: " + framework_parsed.sources[name]) + p.rounds.append(rec) + p.segments_raw.append({k: round(v, 3) for k, v in seg.items()}) + _round_s_from_ends(p) + na = {"sampler_update_s"} # single-island arms: no separate sampler island + if mode == "sync": + na |= {"wait_s", "dropped_samples", "generated_samples", "version_lag"} + return _finish(p, na) + + +ROLLOUTER_ACTIVE = "fully_async/rollouter/active_time" + + +def _rollouter_active(checks: list[dict[str, Any]]) -> dict[int, float]: + """rl_publication_check events -> rollout id (policy version - 1) -> rollouter active s.""" + out: dict[int, float] = {} + for e in checks: + value = (e.get("rollouter_timing") or {}).get(ROLLOUTER_ACTIVE) + version = e.get("policy_version") + if isinstance(value, (int, float)) and version is not None and int(version) >= 1: + out[int(version) - 1] = float(value) + return out + + +def _driver_phase_durations(markers: list[dict[str, Any]]) -> dict[int, dict[str, float]]: + """rl_driver_phase markers (one per phase start) -> rollout id -> phase -> seconds + (marker to the next marker; a marker without rollout_id belongs to the round in progress).""" + out: dict[int, dict[str, float]] = defaultdict(lambda: defaultdict(float)) + rows = sorted(markers, key=lambda e: float(e.get("time_unix") or 0)) + current = 0 + for a, b in zip(rows, rows[1:]): + if a.get("rollout_id") is not None: + current = int(a["rollout_id"]) + out[current][str(a.get("phase"))] += float(b["time_unix"]) - float(a["time_unix"]) + return out + + +# ====================================================================== NVML +def parse_nvml(text: str) -> list[dict[str, Any]]: + """Rows {t_unix, key (container:index), name, util, mem_mb, mem_total_mb, power_w}; header lines skipped.""" + rows = [] + for parts in csv.reader(io.StringIO(text)): + parts = [x.strip() for x in parts] + if not parts or parts[0].startswith("timestamp") or "timestamp" in parts: + continue + container = "" + if not re.match(r"\d{4}/\d\d/\d\d", parts[0]): + container, parts = parts[0], parts[1:] + try: + t = _dt.datetime.strptime(parts[0], "%Y/%m/%d %H:%M:%S.%f").replace(tzinfo=_dt.timezone.utc).timestamp() + except (ValueError, IndexError): + continue + num = [re.sub(r"[^\d.]", "", x) for x in parts] + if len(parts) >= 7: # timestamp,index,name,util,mem.used,mem.total,power + index, util, mem, total, power = parts[1], num[3], num[4], num[5], num[6] + else: # timestamp,name,util,mem.used,power (T-EVAL format, no index) + index, util, mem, total, power = "0", num[2], num[3], "", num[4] if len(num) > 4 else "" + name = parts[2] if len(parts) >= 7 else parts[1] + rows.append({"t_unix": t, "key": f"{container}:{index}", "name": name, "util": float(util) if util else None, + "mem_mb": float(mem) if mem else None, "mem_total_mb": float(total) if total else None, + "power_w": float(power) if power else None}) + return rows + + +# ====================================================================== run directory +def _read(path: Path) -> str: + return path.read_text(errors="replace") if path.is_file() else "" + + +NOT_RUN_TAPES = ("yeto", "store") + + +def _tape_events(raw: Path) -> list[dict[str, Any]]: + seen, rows = set(), [] + for path in sorted(raw.rglob("*.jsonl")): + # miles_local.sh pulls a copy of the repo (raw/yeto, whose test fixtures are + # tapes of other runs) and the checkpoint store next to the run's own tapes + if "verl-" in path.name or path.relative_to(raw).parts[0] in NOT_RUN_TAPES: + continue + for row in tape_rows(_read(path)): + key = json.dumps(row, sort_keys=True, default=str) # the same event pulled from two sources + if key not in seen: + seen.add(key) + rows.append(row) + return rows + + +def parse_run(run_dir: Path) -> dict[str, Any]: + """Parse a run directory (``arm.json`` + ``raw/``) into ``parsed.json`` content.""" + run_dir = Path(run_dir) + plan = json.loads((run_dir / "arm.json").read_text()) + raw = run_dir / "raw" + mode, spec, c = plan["mode"], plan["spec"], plan["cards"] + arm_dir = next(iter(sorted(raw.glob("*/arm"))), raw / "arm") + nvml_text = _read(arm_dir / "nvml.csv") or _read(raw / "nvml.csv") + meta = json.loads(_read(arm_dir / "meta.json") or "{}") + if plan["framework"] == "verl": + metrics = next((p for p in sorted(raw.rglob("verl-metrics.jsonl"))), None) \ + or next((p for p in sorted(raw.rglob("verl-file-logger-*.jsonl"))), None) + fw = parse_verl(_read(metrics) if metrics else "", _read(arm_dir / "main.log"), + samples_per_round=spec["groups_per_round"] * spec["samples_per_group"], + trainer_cards=c["trainer"], mode=mode) + else: + log = _read(arm_dir / "main.log") or _read(raw / "launch.ts.log") + fw = parse_miles(log, rollout_cards=c["rollout"], mode=mode) + # yeto+verl fully_async: the driver's generate phase (and verl's timing_s/gen) only + # pull finished samples from the rollouter queue; generation time is the + # rollouter's own active time, carried by the tape's publication checks + verl_async = plan["framework"] == "verl" and mode == "async" + skip = frozenset({"gen_s"}) if verl_async else frozenset() + parsed = parse_yeto(_tape_events(raw), mode=mode, framework_parsed=fw, skip=skip, + framework_round_offset=1 if plan["framework"] == "verl" else 0, + rollouter_gen=verl_async) \ + if plan["system"] == "yeto" else fw + if meta.get("rc") is None and _read(raw / "rc").strip().lstrip("-").isdigit(): + meta["rc"] = int(_read(raw / "rc")) # miles_local.sh: the launcher's exit code + start = meta.get("start_unix") or _utc_file(raw / "start_utc.txt") + end = meta.get("end_unix") or _utc_file(raw / "end_utc.txt") + return {"arm": plan["arm"], "run_id": plan["run_id"], "fingerprint": plan["fingerprint"], "mode": mode, + "framework": plan["framework"], "system": plan["system"], "cards": c, "gpu": plan["gpu"], + "spec": spec, "container_start_unix": start, "container_end_unix": end, "rc": meta.get("rc"), + "framework_log": fw.to_dict() if plan["system"] == "yeto" else None, + "nvml": (nvml := parse_nvml(nvml_text)), + # E1: Modal "A10G" containers are NVIDIA A10 or (sometimes) A10G, which train slower + "gpu_names": sorted({r["name"] for r in nvml if r.get("name")}), + "nvml_interval_s": 20.0 if plan["harness"] == "miles_local" else 1.0, + **parsed.to_dict()} + + +def _utc_file(path: Path) -> float | None: + text = _read(path).strip() + if not text: + return None + return _dt.datetime.strptime(text, "%Y-%m-%dT%H:%M:%SZ").replace(tzinfo=_dt.timezone.utc).timestamp() diff --git a/tools/eval/efficiency/prices.json b/tools/eval/efficiency/prices.json new file mode 100644 index 000000000..75edc4eea --- /dev/null +++ b/tools/eval/efficiency/prices.json @@ -0,0 +1,6 @@ +{ + "_note": "USD per card-hour, all-in container estimate (GPU + the CPU and memory the harness requests). Modal list prices read 2026-10-20 for the GPU part; CPU/memory added as in tools/eval/cells.py. Estimates, not a bill: replace with the ledger's rate (infra-drafts/gpu-spend.md) when it differs.", + "A10G": {"usd_per_card_hour": 2.0, "gpu_list_usd_per_hour": 1.10}, + "L4": {"usd_per_card_hour": 1.6, "gpu_list_usd_per_hour": 0.80}, + "H100!": {"usd_per_card_hour": 5.3, "gpu_list_usd_per_hour": 3.95} +} diff --git a/tools/eval/efficiency/report.py b/tools/eval/efficiency/report.py new file mode 100644 index 000000000..92e469275 --- /dev/null +++ b/tools/eval/efficiency/report.py @@ -0,0 +1,213 @@ +"""Markdown report of a set of efficiency runs (``parsed.json`` in each run dir). + +Sections: the shared spec and its fingerprint; one row per run; per-arm +aggregates (median and min-max over repeats); yeto vs bare per framework and +mode, metric by metric, with a range verdict; yeto's own segments; reward +curves; and the fields each arm could not measure. + +Range verdict (repeats are not bit-reproducible on H100, so no single-run +comparison): with >= :data:`MIN_REPEATS` runs on both sides, yeto is "worse" +or "better" only when its min-max range and bare's do not overlap; otherwise +"within run-to-run range". Fewer runs, or a missing value, is said as such. +""" + +from __future__ import annotations + +import json +import statistics +from collections import defaultdict +from pathlib import Path +from typing import Any + +from .metrics import summarize +from .spec import ARMS, FRAMEWORKS, MODES, EffSpec, verl_knob_diff + +MIN_REPEATS = 2 +# metric -> (label, True when lower is better) +COMPARED = { + "round_s_median": ("round time (s)", True), + "gen_s_median": ("generation (s/round)", True), + "train_s_median": ("training (s/round)", True), + "weight_sync_s_median": ("weight sync trainer->rollout (s/round)", True), + "sampler_update_s_median": ("update sampler weights (s/round)", True), + "wait_s_median": ("wait (s/round)", True), + "unaccounted_s_median": ("other, not in the segments above (s/round)", True), + "gen_tok_per_card_s": ("generated tokens / card / s", False), + "train_tok_per_card_s": ("trained tokens / card / s", False), + "samples_per_s": ("trained samples / s", False), + "gpu_util_mean": ("GPU utilization mean (%)", False), + "bubble_frac": ("idle bubble share (util < 5 %)", True), + "gpu_mem_peak_mb": ("GPU memory peak (MiB)", True), + "usd_per_1k_samples_steady": ("USD / 1k samples (steady)", True), + "usd_per_1k_samples_run": ("USD / 1k samples (whole container run)", True), +} + + +def load(run_dirs: list[Path]) -> list[dict[str, Any]]: + out = [] + for d in run_dirs: + path = Path(d) / "parsed.json" + if path.is_file(): + out.append(json.loads(path.read_text())) + return out + + +def _fmt(v: Any) -> str: + if v is None: + return "missing" + if isinstance(v, float): + return f"{v:.4g}" + return str(v) + + +def _range(values: list[float]) -> str: + return "missing" if not values else f"{_fmt(statistics.median(values))} [{_fmt(min(values))}, {_fmt(max(values))}]" + + +def verdict(bare: list[float | None], yeto: list[float | None], lower_is_better: bool) -> str: + if any(v is None for v in bare + yeto) or not bare or not yeto: + return "missing value" + if len(bare) < MIN_REPEATS or len(yeto) < MIN_REPEATS: + return f"insufficient repeats (bare {len(bare)}, yeto {len(yeto)}; need {MIN_REPEATS})" + if max(yeto) < min(bare): + return "yeto better" if lower_is_better else "yeto worse" + if min(yeto) > max(bare): + return "yeto worse" if lower_is_better else "yeto better" + return "within run-to-run range" + + +def _delta(bare: list[float], yeto: list[float]) -> str: + if not bare or not yeto or any(v is None for v in bare + yeto): + return "missing" + b, y = statistics.median(bare), statistics.median(yeto) + if b == 0: + return f"{y - b:+.4g} (abs)" + return f"{y - b:+.4g} ({(y - b) / abs(b) * 100:+.1f} %)" + + +def markdown(parsed_runs: list[dict[str, Any]], *, title: str) -> str: + if not parsed_runs: + return f"# {title}\n\nNo parsed runs.\n" + prints = {p["fingerprint"] for p in parsed_runs} + if len(prints) > 1: + raise ValueError(f"runs come from different specs (fingerprints {sorted(prints)}); report them separately") + models = {tuple(p.get("gpu_names") or ()) for p in parsed_runs if p.get("gpu_names")} + if len(models) > 1: + # E1: e1a9 ran on an NVIDIA A10G, the others on an A10; its train step was ~12 % slower + raise ValueError(f"runs ran on different GPU models {sorted(models)}; report them separately") + spec = parsed_runs[0]["spec"] + summaries = [summarize(p) for p in parsed_runs] + by_arm: dict[str, list[dict]] = defaultdict(list) + for s in summaries: + by_arm[s["arm"]].append(s) + lines = [f"# {title}", "", + f"Spec `{spec['name']}`, fingerprint `{parsed_runs[0]['fingerprint'][:16]}`: model `{spec['model']}`@" + f"`{spec['model_revision'][:10]}`, data `{spec['data']}`@`{spec['data_revision'][:10]}`, " + f"{spec['rounds']} rounds x {spec['groups_per_round']} prompts x {spec['samples_per_group']} samples, " + f"max_tokens {spec['max_response_len']}, lr {spec['lr']}, LoRA r {spec['lora_r']}, seed {spec['seed']}, " + f"{spec['gpu']} (sync 1 card colocated, async 1 rollout + 1 trainer card), " + f"max_policy_age {spec['max_policy_age']} (async), TIS upper {spec['tis_upper']}.", "", + "Steady window = rounds after the first. Throughput counts every card of the arm. " + "Costs are list-price estimates (`prices.json`), not a bill.", ""] + lines += _runs_table(summaries) + _arms_table(by_arm) + _overhead(by_arm) + _yeto_segments(by_arm) + lines += _knobs(spec) + _rewards(summaries) + _missing(by_arm) + return "\n".join(lines) + "\n" + + +def _runs_table(summaries: list[dict]) -> list[str]: + head = ("arm", "run", "rc", "rounds", "steady wall s", "gen tok/card/s", "train tok/card/s", "round s med", + "util %", "bubble", "mem peak MiB", "$/1k steady", "$/1k run", "reward last3") + rows = ["## Runs", "", "| " + " | ".join(head) + " |", "|" + "---|" * len(head)] + for s in sorted(summaries, key=lambda s: (ARMS.index(s["arm"]), s["run_id"])): + vals = (s["arm"], f"`{s['run_id']}`", s["rc"], f"{s['rounds']}/{s['rounds_expected']}", s["steady_wall_s"], + s["gen_tok_per_card_s"], s["train_tok_per_card_s"], s["round_s_median"], s["gpu_util_mean"], + s["bubble_frac"], s["gpu_mem_peak_mb"], s["usd_per_1k_samples_steady"], s["usd_per_1k_samples_run"], + s["reward_last3_mean"]) + rows.append("| " + " | ".join(_fmt(v) for v in vals) + " |") + return rows + [""] + + +def _arms_table(by_arm: dict[str, list[dict]]) -> list[str]: + rows = ["## Per arm (median [min, max] over repeats)", "", + "| metric | " + " | ".join(a for a in ARMS if a in by_arm) + " |", "|---|" + "---|" * len(by_arm)] + for key, (label, _) in COMPARED.items(): + cells = [] + for a in ARMS: + if a in by_arm: + vals = [s[key] for s in by_arm[a]] + cells.append(_range([v for v in vals if v is not None]) if any(v is not None for v in vals) + else ("n/a" if all(s["sources"].get(key.removesuffix("_median")) == "n/a" + for s in by_arm[a]) else "missing")) + rows.append(f"| {label} | " + " | ".join(cells) + " |") + return rows + [""] + + +def _overhead(by_arm: dict[str, list[dict]]) -> list[str]: + rows = ["## yeto vs bare (same framework and mode)", "", + "Delta = median(yeto) - median(bare); positive time = yeto slower.", ""] + for fw in FRAMEWORKS: + for mode in MODES: + bare, yeto = by_arm.get(f"bare-{fw}-{mode}", []), by_arm.get(f"yeto-{fw}-{mode}", []) + if not bare and not yeto: + continue + rows += [f"### {fw} {mode} (bare n={len(bare)}, yeto n={len(yeto)})", "", + "| metric | bare | yeto | delta | verdict |", "|---|---|---|---|---|"] + for key, (label, lower) in COMPARED.items(): + bv, yv = [s[key] for s in bare], [s[key] for s in yeto] + if all(s["sources"].get(key.removesuffix("_median")) == "n/a" for s in bare + yeto): + rows.append(f"| {label} | n/a | n/a | n/a | not applicable |") + continue + rows.append(f"| {label} | {_range([v for v in bv if v is not None])} | " + f"{_range([v for v in yv if v is not None])} | {_delta(bv, yv)} | " + f"{verdict(bv, yv, lower)} |") + rows.append("") + return rows + + +def _yeto_segments(by_arm: dict[str, list[dict]]) -> list[str]: + rows = ["## Native segments per round (median, steady window)", "", + "yeto: tape task/phase names; bare: the framework's own timers. Where yeto adds work " + "(offload/onload, outer sync, publish), it shows here.", ""] + for a in ARMS: + for s in by_arm.get(a, []): + raw = s.get("segments_raw_median") or {} + rows.append(f"- {a} `{s['run_id']}`: " + (", ".join(f"{k} {_fmt(v)}" for k, v in raw.items()) + or "missing")) + return rows + [""] + + +def _knobs(spec: dict) -> list[str]: + rows = ["## verl overrides that differ between bare and yeto (same spec)", ""] + try: + for mode in MODES: + diff = verl_knob_diff(EffSpec(**spec), mode) + rows.append(f"- {mode}: " + ("; ".join(f"`{k}` bare={v['bare']} yeto={v['yeto']}" + for k, v in diff.items()) or "none")) + except ImportError as exc: + rows.append(f"- not computed ({exc})") + rows += ["- Miles: the yeto launcher builds the Miles argv at run time; compare the launch log's argv with " + "the bare script (`arm.json`) by hand (not automated in E0.1).", ""] + return rows + + +def _rewards(summaries: list[dict]) -> list[str]: + rows = ["## Reward against wall time and samples", "", "Per run: round: t_s / samples / reward_mean.", ""] + for s in sorted(summaries, key=lambda s: (ARMS.index(s["arm"]), s["run_id"])): + pts = " ".join(f"{c['round']}:{_fmt(c['t_s'])}/{_fmt(c['samples'])}/{_fmt(c['reward_mean'])}" + for c in s["reward_curve"]) + rows.append(f"- {s['arm']} `{s['run_id']}`: {pts or 'missing'}") + return rows + [""] + + +def _missing(by_arm: dict[str, list[dict]]) -> list[str]: + rows = ["## Field sources and gaps", "", "Per arm (first run): each round field's source; `missing` = no " + "source in the raw data, `n/a` = does not apply.", ""] + for a in ARMS: + if a in by_arm: + src = by_arm[a][0]["sources"] + rows.append(f"- {a}: " + "; ".join(f"{k}: {v}" for k, v in src.items())) + partial = {k: v for s in by_arm[a] for k, v in (s.get("rate_rounds") or {}).items()} + if partial: + rows.append(f" - {a}: rates over the steady rounds that report the value: " + + "; ".join(f"{k} {v}" for k, v in partial.items())) + return rows + [""] diff --git a/tools/eval/efficiency/spec.py b/tools/eval/efficiency/spec.py new file mode 100644 index 000000000..dcbcab0ec --- /dev/null +++ b/tools/eval/efficiency/spec.py @@ -0,0 +1,517 @@ +"""One shared spec per model preset, and the command plan of every arm. + +An arm is ``--``: + +* system ``bare`` (the framework's own entry point, no yeto code on its path) + or ``yeto`` (the same framework driven by yeto); +* framework ``verl`` (fork pin ``VERL_COMMIT``, the verl image yeto builds) or + ``miles`` (``MILES_NEXT_IMAGE``, the image yeto runs Miles in); +* mode ``sync`` (one card, rollout and trainer colocated) or ``async`` + (two cards, one rollout + one trainer, fully_async with a policy-age limit). + +Every value the comparison depends on comes from :class:`EffSpec`; +:func:`control_fingerprint` hashes them and every arm records it. + +Sources of the bare command lines (the "official recipe" each one follows): + +* bare verl sync: ``examples/grpo_trainer/run_qwen3_4b_fsdp.sh`` (fork + acad9875), ``python -m verl.trainer.main_ppo`` with its override groups; +* bare verl async: ``verl/experimental/fully_async_policy/shell/ + dapo_7b_math_fsdp2_4_4.sh``, ``python -m verl.experimental.fully_async_policy.fully_async_main``; +* bare Miles sync: ``scripts/run_qwen3_4b.py`` (Megatron branch, bridge mode) + -> ``train.py``; model flags from the image's own ``scripts/models/``; +* bare Miles async: ``examples/infra_features/fully_async/`` (``--fully-async``) + -> ``train_async.py``. + +Values the spec does not name stay at the recipe's value, or at the +framework default where the recipe sets none. Where yeto sets a value the +recipe also sets (LoRA, TIS, dtype), the bare arm takes the spec value, so the +two arms differ only by what yeto adds; :func:`verl_knob_diff` lists the verl +overrides that still differ. +""" + +from __future__ import annotations + +import hashlib +import json +import shlex +from dataclasses import asdict, dataclass, field, replace +from pathlib import Path +from typing import Any + +SYSTEMS = ("bare", "yeto") +FRAMEWORKS = ("verl", "miles") +MODES = ("sync", "async") +ARMS = tuple(f"{s}-{f}-{m}" for f in FRAMEWORKS for m in MODES for s in SYSTEMS) + +EFF_DIR = "/root/eff" # spec, reward and prep script inside a bare container +OUT_DIR = "/root/eff-out" # everything a container writes for the results Volume +DATA_DIR = "/root/eff-data" # eff_prep.py output +MODEL_DIR = "/root/eff-models" +VERL_ROOT = "/workspace/verl" +MILES_ROOT = "/root/miles" +MEGATRON_ROOT = "/root/Megatron-LM" # yeto.launch.remote_setup.PORTS_MEGATRON_PATH + + +@dataclass(frozen=True) +class EffSpec: + """Everything the arms of one comparison share.""" + + name: str + model: str + model_revision: str + miles_model_args: str # file under /scripts/models/ (the official model flags) + data: str = "zhuzilin/gsm8k" + data_revision: str = "0cbd9f31d91ac21a7613dcbc7fef992adac459ae" + gpu: str = "A10G" # Modal GPU type, the same for every arm + rounds: int = 8 # optimizer steps (= policy versions) per run + groups_per_round: int = 8 # prompts per step + samples_per_group: int = 4 # responses per prompt + max_prompt_len: int = 512 + max_response_len: int = 256 # "max_tokens" of the rollout + lr: float = 1e-5 + lora_r: int = 32 # yeto's verl backend is LoRA only, so every arm is LoRA + seed: int = 17 + temperature: float = 1.0 + top_p: float = 1.0 + mem_fraction: float = 0.5 # rollout engine static memory fraction + max_policy_age: int = 1 # async arms: versions a sample may lag + tis_upper: float = 2.0 # token-level truncated IS upper bound (all arms) + repeats: int = 2 # runs per arm (range verdicts need >= 2) + timeout_s: int = 3000 + chat_template_kwargs: dict[str, Any] = field(default_factory=dict) + # yeto arms only (the bare arms keep no checkpoint): "final" cuts once at the end + # (fair control), "every-round" is the product default (a round cut every round). + # Not in the control fingerprint: it does not change what the bare arms run. + checkpoint: str = "final" + # yeto+Miles only, observation (E1): --rl-profile-trainer (not in the fingerprint) + profile_trainer: bool = False + + @property + def samples_per_round(self) -> int: + return self.groups_per_round * self.samples_per_group + + @property + def seq_len(self) -> int: + return self.max_prompt_len + self.max_response_len + + def to_dict(self) -> dict[str, Any]: + return asdict(self) + + +PRESETS: dict[str, EffSpec] = { + # E0.1 acceptance (debug size): the T-EVAL model and data on the T-EVAL card + "qwen3-0.6b": EffSpec( + name="qwen3-0.6b", model="Qwen/Qwen3-0.6B", model_revision="c1899de289a04d12100db370d81485cdf75e47ca", + miles_model_args="qwen3-0.6B.py", chat_template_kwargs={"enable_thinking": False}), + # E0.2 (main comparison): revision read from the HF API 2026-10-20 + "qwen3-4b-instruct-2507": EffSpec( + name="qwen3-4b-instruct-2507", model="Qwen/Qwen3-4B-Instruct-2507", + model_revision="cdbee75f17c01a7cc42f958dc650907174af0554", miles_model_args="qwen3-4B-Instruct-2507.py", + gpu="H100!", rounds=20, groups_per_round=32, samples_per_group=8, max_prompt_len=512, + max_response_len=1024, lora_r=64, mem_fraction=0.6, timeout_s=5400), +} + + +@dataclass(frozen=True) +class Arm: + system: str + framework: str + mode: str + + @classmethod + def parse(cls, name: str) -> "Arm": + parts = name.split("-") + if len(parts) != 3 or name not in ARMS: + raise ValueError(f"unknown arm {name!r}; known: {', '.join(ARMS)}") + return cls(*parts) + + @property + def name(self) -> str: + return f"{self.system}-{self.framework}-{self.mode}" + + @property + def peer(self) -> str: + """The arm this one is compared with (same framework and mode).""" + return f"{'yeto' if self.system == 'bare' else 'bare'}-{self.framework}-{self.mode}" + + +def get_spec(preset: str, **overrides: Any) -> EffSpec: + if preset not in PRESETS: + raise ValueError(f"unknown preset {preset!r}; known: {', '.join(PRESETS)}") + return replace(PRESETS[preset], **overrides) if overrides else PRESETS[preset] + + +def cards(mode: str) -> dict[str, int]: + """Cards per role: sync colocates both on one card; async has one of each.""" + return {"total": 1, "rollout": 1, "trainer": 1, "colocated": True} if mode == "sync" else \ + {"total": 2, "rollout": 1, "trainer": 1, "colocated": False} + + +def control_fingerprint(spec: EffSpec) -> str: + """sha256 over every shared value except the repeat count (a report refuses + to compare runs whose fingerprints differ).""" + doc = {k: v for k, v in spec.to_dict().items() if k not in ("repeats", "timeout_s", "checkpoint", "profile_trainer")} + return hashlib.sha256(json.dumps(doc, sort_keys=True).encode()).hexdigest() + + +def images() -> dict[str, str]: + """The image pins every arm of a framework uses (same image for bare and yeto).""" + from yeto.rl.adapters.miles.pins import MILES_NEXT_IMAGE + from yeto.rl.adapters.verl.pins import VERL_COMMIT + + return {"verl": f"verl-build:{VERL_COMMIT}", "miles": MILES_NEXT_IMAGE} + + +def _model_path(spec: EffSpec) -> str: + return f"{MODEL_DIR}/{spec.model.split('/')[-1]}" + + +# ====================================================================== bare verl +def verl_bare_overrides(spec: EffSpec, mode: str, run_id: str) -> list[str]: + """Hydra overrides of the bare verl arm (recipe groups in the recipe's order).""" + c = cards(mode) + data = [ + "algorithm.adv_estimator=grpo", + "algorithm.use_kl_in_reward=False", + "algorithm.norm_adv_by_std_in_grpo=True", + f"data.train_files={DATA_DIR}/train.parquet", + f"data.val_files={DATA_DIR}/test.parquet", + f"data.max_prompt_length={spec.max_prompt_len}", + f"data.max_response_length={spec.max_response_len}", + "data.filter_overlong_prompts=True", + "data.truncation=error", + f"data.seed={spec.seed}", + *[f"+data.apply_chat_template_kwargs.{k}={v}" for k, v in sorted(spec.chat_template_kwargs.items())], + ] + model = [ + f"actor_rollout_ref.model.path={_model_path(spec)}", + "actor_rollout_ref.model.use_remove_padding=True", + "actor_rollout_ref.model.enable_gradient_checkpointing=True", + f"actor_rollout_ref.model.lora_rank={spec.lora_r}", + f"actor_rollout_ref.model.lora_alpha={spec.lora_r}", + "actor_rollout_ref.model.target_modules=all-linear", + ] + actor = [ + "actor_rollout_ref.actor.strategy=fsdp2", + "actor_rollout_ref.actor.fsdp_config.model_dtype=fp32", # yeto's VerlRunConfig default + f"actor_rollout_ref.actor.optim.lr={spec.lr!r}", + f"actor_rollout_ref.actor.ppo_mini_batch_size={spec.groups_per_round}", + "actor_rollout_ref.actor.use_dynamic_bsz=True", + "actor_rollout_ref.actor.ppo_max_token_len_per_gpu=16384", + "actor_rollout_ref.actor.use_kl_loss=False", + "actor_rollout_ref.actor.entropy_coeff=0", + "actor_rollout_ref.actor.fsdp_config.param_offload=False", + "actor_rollout_ref.actor.fsdp_config.optimizer_offload=False", + f"actor_rollout_ref.actor.data_loader_seed={spec.seed}", + ] + rollout = [ + "actor_rollout_ref.rollout.name=vllm", + "actor_rollout_ref.rollout.tensor_model_parallel_size=1", + f"actor_rollout_ref.rollout.gpu_memory_utilization={spec.mem_fraction}", + f"actor_rollout_ref.rollout.n={spec.samples_per_group}", + f"actor_rollout_ref.rollout.temperature={spec.temperature}", + f"actor_rollout_ref.rollout.top_p={spec.top_p}", + "actor_rollout_ref.rollout.top_k=-1", + "actor_rollout_ref.rollout.calculate_log_probs=True", + "actor_rollout_ref.rollout.log_prob_use_dynamic_bsz=True", + "actor_rollout_ref.rollout.log_prob_max_token_len_per_gpu=16384", + "algorithm.rollout_correction.rollout_is=token", + f"algorithm.rollout_correction.rollout_is_threshold={spec.tis_upper}", + "algorithm.rollout_correction.bypass_mode=False", + ] + reward = [f"reward.custom_reward_function.path={EFF_DIR}/eff_reward.py", + "reward.custom_reward_function.name=compute_score"] + trainer = [ + "trainer.critic_warmup=0", + "trainer.logger=[console,file]", + "trainer.project_name=yeto-effeval", + f"trainer.experiment_name={run_id}", + f"trainer.default_local_dir={OUT_DIR}/ckpt", + "trainer.nnodes=1", + "trainer.val_before_train=False", + "trainer.test_freq=-1", + "trainer.save_freq=-1", + "trainer.resume_mode=disable", + f"trainer.total_training_steps={spec.rounds}", + "trainer.total_epochs=1000", + "ray_kwargs.ray_init.runtime_env.py_executable=null", + "ray_kwargs.ray_init.num_cpus=32", # a Modal container reports too few Ray CPUs (S17 V2) + ] + if mode == "sync": + return [*data, f"data.train_batch_size={spec.groups_per_round}", *model, *actor, *rollout, *reward, + *trainer, f"trainer.n_gpus_per_node={c['total']}"] + fully_async = [ # dapo_7b_math_fsdp2_4_4.sh; one parameter version per spec round + "data.train_batch_size=0", + "data.gen_batch_size=1", + "data.return_raw_chat=True", + "actor_rollout_ref.hybrid_engine=False", + "actor_rollout_ref.rollout.mode=async", + "actor_rollout_ref.rollout.checkpoint_engine.backend=nccl", + # The fork's disaggregated sender drops the PEFT config of unmerged LoRA params + # (engine_workers.update_weights), so the rollout would load them as base + # weights; merged LoRA sends full weights, the fork's working LoRA path. + "actor_rollout_ref.model.lora.merge=True", + "rollout.nnodes=1", + f"rollout.n_gpus_per_node={c['rollout']}", + f"rollout.total_rollout_steps={spec.groups_per_round * spec.rounds}", + f"trainer.n_gpus_per_node={c['trainer']}", + f"async_training.staleness_threshold={spec.max_policy_age}", + "async_training.trigger_parameter_sync_step=1", + "async_training.require_batches=1", + "async_training.partial_rollout=True", + "async_training.use_trainer_do_validate=False", + "async_training.use_dynamic_resource_scheduling=False", + ] + return [*data, *model, *actor, *rollout, *reward, *trainer, *fully_async] + + +def verl_bare_entry(mode: str) -> str: + return "verl.trainer.main_ppo" if mode == "sync" else "verl.experimental.fully_async_policy.fully_async_main" + + +# ====================================================================== bare Miles +def miles_bare_argv(spec: EffSpec, mode: str) -> list[str]: + """Miles arguments after the model flags (the container prepends ``$MODEL_ARGS``).""" + c = cards(mode) + model = _model_path(spec) + argv = [ + # checkpoint: bridge mode loads the HF checkpoint directly (run_qwen3_4b.py --megatron-to-hf-mode bridge) + "--hf-checkpoint", model, "--ref-load", model, "--megatron-to-hf-mode", "bridge", + "--finetune", "--no-load-optim", "--no-load-rng", "--no-save-optim", "--no-save-rng", + # rollout + "--prompt-data", f"{DATA_DIR}/train.jsonl", "--input-key", "prompt", "--label-key", "label", + # no --rollout-shuffle (the recipe has it): yeto's Miles argv leaves rollout_shuffle + # False (eff01-yeto-miles-* Miles arg dump), and both arms must read prompts in one order + "--apply-chat-template", "--rollout-seed", str(spec.seed), + "--custom-rm-path", "eff_reward.score", + "--num-rollout", str(spec.rounds), "--rollout-batch-size", str(spec.groups_per_round), + "--n-samples-per-prompt", str(spec.samples_per_group), + "--rollout-max-response-len", str(spec.max_response_len), + "--rollout-max-context-len", str(spec.seq_len), + "--rollout-temperature", str(spec.temperature), "--rollout-top-p", str(spec.top_p), + "--global-batch-size", str(spec.samples_per_round), "--balance-data", + # GRPO + TIS (the correction every arm uses) + "--advantage-estimator", "grpo", "--entropy-coef", "0.00", "--eps-clip", "0.2", "--eps-clip-high", "0.28", + "--use-tis", "--tis-clip", str(spec.tis_upper), + # optimizer + "--optimizer", "adam", "--lr", str(spec.lr), "--lr-decay-style", "constant", + "--seed", str(spec.seed), + # LoRA (Miles' own flags; the same set yeto's Miles adapter emits) + "--lora-rank", str(spec.lora_r), "--lora-alpha", str(spec.lora_r), "--lora-dropout", "0", + "--lora-type", "canonical_lora", "--target-modules", "all-linear", + "--sglang-max-lora-rank", str(spec.lora_r), + # Megatron (run_qwen3_4b.py, Megatron branch) + "--train-backend", "megatron", "--seq-length", str(spec.seq_len), + "--tensor-model-parallel-size", "1", "--pipeline-model-parallel-size", "1", + "--context-parallel-size", "1", "--expert-model-parallel-size", "1", "--expert-tensor-parallel-size", "1", + "--recompute-granularity", "full", "--recompute-method", "uniform", "--recompute-num-layers", "1", + "--use-dynamic-batch-size", "--max-tokens-per-gpu", str(max_tokens_per_gpu(spec)), + "--attention-dropout", "0.0", "--hidden-dropout", "0.0", + "--accumulate-allreduce-grads-in-fp32", "--attention-softmax-in-fp32", + # yeto's generic recipe (Qwen3) trains with "unfused": TE's flash path picks the + # FlashAttention cute kernel, which fails to compile on A10G (sm86, eff01-bare-miles-sync-a). + "--attention-backend", "unfused", + "--bf16", + # SGLang + "--rollout-num-gpus-per-engine", "1", "--sglang-mem-fraction-static", str(spec.mem_fraction), + # placement + "--actor-num-nodes", "1", "--actor-num-gpus-per-node", str(c["trainer"]), + "--num-gpus-per-node", str(c["total"]), + ] + if spec.chat_template_kwargs: + argv += ["--apply-chat-template-kwargs", json.dumps(spec.chat_template_kwargs, separators=(",", ":"))] + if mode == "sync": + argv += ["--colocate", "--offload-train", "--lora-base-cpu-backup"] + else: + argv += ["--rollout-num-gpus", str(c["rollout"]), "--fully-async", + "--max-weight-staleness", str(spec.max_policy_age)] + return argv + + +def miles_bare_script(mode: str) -> str: + return "train.py" if mode == "sync" else "train_async.py" + + +def max_tokens_per_gpu(spec: EffSpec) -> int: + """Miles dynamic-batching budget (run_qwen3_4b.py: 9216), never below one sample.""" + return max(9216, spec.seq_len) + + +# ====================================================================== yeto arms +YETO_WORKDIR = "/root/sky_workdir" # modal_verl.WORKDIR: the repo inside the yeto+verl container +ALGORITHM_SPEC = "examples/rl_algorithms/decoupled-ppo-c2.json" # token TIS 2.0, max_policy_staleness 2 +# the same with the bare Miles recipe's clip range (--eps-clip 0.2 --eps-clip-high 0.28) +MILES_ALGORITHM_SPEC = "tools/eval/efficiency/algo-eff-miles.json" +def yeto_verl_argv(spec: EffSpec, mode: str, tape: str) -> list[str]: + """``island_entry`` arguments of the yeto+verl arm (the T-EVAL E1 shape).""" + argv = ["-m", "yeto.rl.adapters.verl.island_entry", "--model", spec.model, "--model-revision", + spec.model_revision, "--data", spec.data, "--data-revision", spec.data_revision, "--learner-id", "0", + "--reward-function", "gsm8k_reward:score", "--global-rounds", str(spec.rounds), + "--groups-per-round", str(spec.groups_per_round), "--samples-per-group", str(spec.samples_per_group), + "--rollout-max-response-len", str(spec.max_response_len), "--event-tape", tape, + "--inner-lr", str(spec.lr), "--lora-r", str(spec.lora_r), "--lora-targets", "all-linear", + "--seq-len", str(spec.seq_len), "--seed", str(spec.seed), "--rl-single-island-no-sync"] + if spec.chat_template_kwargs: + argv += ["--apply-chat-template-kwargs", json.dumps(spec.chat_template_kwargs)] + if mode == "async": + # The default algorithm spec allows no staleness (island_entry rejects age 1, + # eff01-yeto-verl-async-a); use the yeto+Miles arm's spec: token TIS 2.0, staleness 2. + argv += ["--rl-max-policy-age", str(spec.max_policy_age), + "--rl-algorithm-spec", f"{YETO_WORKDIR}/{ALGORITHM_SPEC}"] + return argv + + +CHECKPOINT_MODES = ("final", "every-round", "default") + + +def cut_every(spec: EffSpec) -> int | None: + """Round-cut interval of the yeto+Miles arm: ``final`` cuts only at the end (the + bare arms keep no checkpoint), ``every-round`` cuts every round, ``default`` passes + no interval (the product default, ``yeto.rl.engine.resume.DEFAULT_CUT_EVERY``). + The yeto+verl arm sets no checkpoint store and fully_async refuses cuts, so it + has no round cut in either mode.""" + if spec.checkpoint not in CHECKPOINT_MODES: + raise ValueError(f"checkpoint must be one of {CHECKPOINT_MODES}, got {spec.checkpoint!r}") + return {"final": spec.rounds, "every-round": 1, "default": None}[spec.checkpoint] + + +# E1: the GPU model every arm must run on, per spec.gpu (as tools/eval/modal_verl.GPU_NAMES: +# Modal "A10G" containers are an NVIDIA A10 or an A10G, and the A10G trains slower) +GPU_MODELS = {"A10G": "NVIDIA A10", "L4": "L4", "H100!": "H100"} + + +def gpu_name_arg(spec: EffSpec) -> str: + """The launcher option that pins the yeto+Miles arm to the bare arms' GPU model.""" + return f"--modal-gpu-name {shlex.quote(GPU_MODELS[spec.gpu])} " + + +def yeto_miles_launcher_args(spec: EffSpec, mode: str, run_id: str) -> str: + """The yeto launcher arguments of the yeto+Miles arm (``tools/eval/miles_local.sh`` ARGS).""" + from ..miles import DATA, STORE_VOLUME, TAPE_VOLUME + + card = {"A10G": "a10g", "H100!": "h100", "L4": "l4"}[spec.gpu] + c = cards(mode) + kwargs = json.dumps(spec.chat_template_kwargs) if spec.chat_template_kwargs else "" + args = (f"launch --training-mode rl --rl-engine ports --preflight-threads off --modal-gpu-exact " + f"{gpu_name_arg(spec)}" + # container resources as the bare arms' (modal_eff: 8 CPU + 48 GiB per card) + f"--modal-retries 0 --modal-cpu {8 * c['total']} --modal-memory {48 * c['total']} " + f"--modal-tape-volume {TAPE_VOLUME} " + f"--controller local --cluster-prefix {run_id} --tuning lora --lora-r {spec.lora_r} " + f"--rl-observe-timeline --model {spec.model} --model-revision {spec.model_revision} {DATA} " + f"--rl-checkpoint-store modal-volume://{STORE_VOLUME}/{run_id} " + f"{'' if cut_every(spec) is None else f'--rl-cut-every {cut_every(spec)} '}" + "--reward-function gsm8k_reward:score --fragments 1 --pipeline 1 " + f"--rollout-batch-size {spec.groups_per_round} --n-samples-per-prompt {spec.samples_per_group} " + f"--rollout-max-response-len {spec.max_response_len} --seq-len {spec.seq_len} " + f"--inner-lr {spec.lr} --seed {spec.seed} --sglang-mem-fraction-static {spec.mem_fraction} " + f"--gpu modal:{c['total']}x{card} --rl-single-island-no-sync --total-steps {spec.rounds} " + # Miles recipe settings of the bare arm (S26 GOAL §6: the bare framework's official + # settings win): no deterministic inference, constant LR, full recompute, dynamic + # batching, gradient-accumulation fusion; dropout 0 is the generic recipe's own + "--no-sglang-deterministic-inference --rl-lr-schedule constant --rl-activation-recompute " + f"--rl-max-tokens-per-gpu {max_tokens_per_gpu(spec)} --rl-gradient-accumulation-fusion " + # TIS (upper 2.0) and eps_clip_high 0.28 as in the bare arms; staleness 2 >= max_policy_age + f"--trust-remote-code --rl-algorithm-spec {MILES_ALGORITHM_SPEC}") + if spec.profile_trainer: + args += " --rl-profile-trainer" + if kwargs: + args += f" --apply-chat-template-kwargs {shlex.quote(kwargs)}" + if mode == "async": + args += (f" --rl-placement fixed-partition --rl-rollout-gpus {c['rollout']} " + f"--rl-max-policy-age {spec.max_policy_age}") + if spec.data != "zhuzilin/gsm8k" or spec.data_revision not in DATA: + raise ValueError("the yeto+Miles arm reads tools.eval.miles.DATA; keep the spec's data equal to it") + return args + + +# ====================================================================== plans +def bare_script(spec: EffSpec, arm: Arm, run_id: str) -> str: + """The bash script a bare container runs (after eff_prep.py staged model and data).""" + c = cards(arm.mode) + lines = ["set -euo pipefail", f"mkdir -p {OUT_DIR}", + f"python3 {EFF_DIR}/eff_prep.py {EFF_DIR}/spec.json > {OUT_DIR}/prep.json"] + if arm.framework == "verl": + lines += [f"cd {VERL_ROOT}", f"export VERL_FILE_LOGGER_PATH={OUT_DIR}/verl-metrics.jsonl", + f"export PYTHONPATH={EFF_DIR}", + # Modal: the trainer allocates NCCL host buffers with cuMem, the + # rollout-side receiver falls back to /dev/shm and cannot attach + # (eff01-bare-verl-async-d, NCCL_DEBUG=INFO); both use /dev/shm. + *(["export NCCL_CUMEM_HOST_ENABLE=0"] if arm.mode == "async" else []), + "exec python3 -m " + verl_bare_entry(arm.mode) + " " + shlex.join( + verl_bare_overrides(spec, arm.mode, run_id))] + else: + load = ("import importlib.util as u, sys; sys.path.insert(0, '{d}'); " + "s = u.spec_from_file_location('m', '{d}/{f}'); m = u.module_from_spec(s); " + "s.loader.exec_module(m); print(m.model_args())").format(d=f"{MILES_ROOT}/scripts/models", + f=spec.miles_model_args) + runtime_env = json.dumps({"env_vars": {"PYTHONPATH": f"{MEGATRON_ROOT}/:{EFF_DIR}", + "CUDA_DEVICE_MAX_CONNECTIONS": "1"}}) + lines += [f"cd {MILES_ROOT}", f"MODEL_ARGS=$(python3 -c {shlex.quote(load)})", + f"echo \"$MODEL_ARGS\" > {OUT_DIR}/model_args.txt", + f"ray start --head --node-ip-address 127.0.0.1 --num-gpus {c['total']} --num-cpus 32 " + "--disable-usage-stats --dashboard-host=127.0.0.1", + f"exec ray job submit --address=http://127.0.0.1:8265 --runtime-env-json={shlex.quote(runtime_env)} " + f"-- python3 {MILES_ROOT}/{miles_bare_script(arm.mode)} $MODEL_ARGS " + + shlex.join(miles_bare_argv(spec, arm.mode))] + return "\n".join(lines) + "\n" + + +def plan(spec: EffSpec, arm_name: str, run_id: str) -> dict[str, Any]: + """Everything one run of an arm needs (written to ``arm.json`` before the launch).""" + arm = Arm.parse(arm_name) + c = cards(arm.mode) + doc: dict[str, Any] = { + "arm": arm.name, "system": arm.system, "framework": arm.framework, "mode": arm.mode, "run_id": run_id, + "spec": spec.to_dict(), "fingerprint": control_fingerprint(spec), "cards": c, "gpu": spec.gpu, + } + if arm.system == "bare": + doc.update(harness="modal_eff", image=arm.framework, script=bare_script(spec, arm, run_id)) + elif arm.framework == "verl": + doc.update(harness="modal_eff", image="verl", + argv=yeto_verl_argv(spec, arm.mode, "/root/yeto-output/rl-island-0.jsonl")) + else: + doc.update(harness="miles_local", launcher_args=yeto_miles_launcher_args(spec, arm.mode, run_id)) + return doc + + +# ====================================================================== control check +def _override_map(overrides: list[str]) -> dict[str, str]: + out = {} + for o in overrides: + key, _, value = o.partition("=") + out[key.lstrip("+")] = value + return out + + +def verl_knob_diff(spec: EffSpec, mode: str) -> dict[str, dict[str, str | None]]: + """Hydra overrides where bare verl and yeto+verl differ (key -> {bare, yeto}). + + The yeto side is yeto's own translation (``build_overrides``, plus + ``fully_async_run_overrides`` for async) of the same spec; paths, names and + keys yeto handles in its own loop are compared too, so the list is the full + set of differences, not only the interesting ones.""" + from yeto.rl.adapters.verl.config import VerlRunConfig, build_overrides + from yeto.rl.adapters.verl.fully_async_round import fully_async_run_overrides + + cfg = VerlRunConfig( + model_path=_model_path(spec), train_file=f"{DATA_DIR}/train.parquet", val_file=f"{DATA_DIR}/test.parquet", + out_dir=OUT_DIR, lora_rank=spec.lora_r, lora_alpha=spec.lora_r, lr=spec.lr, + groups_per_round=spec.groups_per_round, samples_per_group=spec.samples_per_group, + max_prompt_len=spec.max_prompt_len, max_response_len=spec.max_response_len, seed=spec.seed, + total_rounds=spec.rounds, temperature=spec.temperature, top_p=spec.top_p, tis_upper=spec.tis_upper, + gpu_memory_utilization=spec.mem_fraction, reward_path=f"{EFF_DIR}/eff_reward.py", + chat_template_kwargs=dict(spec.chat_template_kwargs), experiment_name="x") + yeto = build_overrides(cfg) + if mode == "async": + yeto = fully_async_run_overrides(yeto, spec.max_policy_age, groups_per_round=spec.groups_per_round, + rounds=spec.rounds) + bare, ours = _override_map(verl_bare_overrides(spec, mode, "x")), _override_map(yeto) + return {k: {"bare": bare.get(k), "yeto": ours.get(k)} for k in sorted(set(bare) | set(ours)) + if bare.get(k) != ours.get(k)} + + +def write_plan(path: Path, doc: dict[str, Any]) -> None: + path.parent.mkdir(parents=True, exist_ok=True) + path.write_text(json.dumps(doc, indent=1, default=str)) diff --git a/tools/eval/modal_verl.py b/tools/eval/modal_verl.py index eadc46995..351bb1103 100644 --- a/tools/eval/modal_verl.py +++ b/tools/eval/modal_verl.py @@ -25,6 +25,7 @@ import json import os +import re import secrets import shutil import subprocess @@ -64,13 +65,19 @@ # ====================================================================== in-container helpers +def gpu_name_matches(want: str, name: str) -> bool: + """``name`` is the model ``want`` names: "NVIDIA A10" accepts "NVIDIA A10", not "NVIDIA A10G" + (Modal "A10G" containers report either; the A10G trains slower, E1 e1a9).""" + return re.search(re.escape(want) + r"(?![\w-])", name) is not None + + def _gpu(out: Path, cfg: dict) -> list[str]: names = subprocess.run(["nvidia-smi", "--query-gpu=name", "--format=csv,noheader"], capture_output=True, text=True, timeout=60).stdout.split("\n") names = [n.strip() for n in names if n.strip()] (out / "nvidia-smi.txt").write_text(subprocess.run(["nvidia-smi"], capture_output=True, text=True).stdout) want = GPU_NAMES[cfg["gpu"]] - if not names or not all(want in n for n in names): + if not names or not all(gpu_name_matches(want, n) for n in names): raise RuntimeError(f"asked for {want}, got {names}") return names diff --git a/yeto/cli.py b/yeto/cli.py index 7c68497f8..5e1c1d5c6 100644 --- a/yeto/cli.py +++ b/yeto/cli.py @@ -459,6 +459,14 @@ def int_or_auto(value: str): "may upgrade H100 to H200) and fail at container start unless " "nvidia-smi reports it; use for bitwise comparisons", ) + infra.add_argument( + "--modal-gpu-name", + default=None, + metavar="MODEL", + help="Modal islands: fail at container start unless every GPU nvidia-smi reports " + "is this model (e.g. 'NVIDIA A10': Modal A10G containers are an A10 or an A10G); " + "use when runs must compare on one GPU model", + ) infra.add_argument( "--modal-retries", type=int, diff --git a/yeto/launch/modal_islands.py b/yeto/launch/modal_islands.py index a7f12af6f..e6e7f86f5 100644 --- a/yeto/launch/modal_islands.py +++ b/yeto/launch/modal_islands.py @@ -104,6 +104,7 @@ def build_modal_island_config(args, spec: ClusterSpec, learner_id: int, task, sy envs={k: str(v) for k, v in envs.items() if v is not None}, region=spec.region, gpu_exact=bool(getattr(args, "modal_gpu_exact", False)), + gpu_name=getattr(args, "modal_gpu_name", None) or None, registry_login=bool(registry_login), registry_creds=registry_login, image_ref=image_ref_from_rl_image(args.rl_image) if rl else None, @@ -180,14 +181,13 @@ def pull_modal_tapes(modal_ops, modal_cfgs: dict, run_dir) -> dict[str, str]: def require_modal_for_gpu_exact(args, specs: list[ClusterSpec]) -> None: - """--modal-gpu-exact only means something on Modal islands; refuse it + """--modal-gpu-exact / --modal-gpu-name only mean something on Modal islands; refuse them rather than silently launching unpinned learners elsewhere.""" - if getattr(args, "modal_gpu_exact", False): - other = sorted({s.cloud for s in specs if s.cloud != "modal"}) + for flag, value in (("--modal-gpu-exact", getattr(args, "modal_gpu_exact", False)), + ("--modal-gpu-name", getattr(args, "modal_gpu_name", None))): + other = sorted({s.cloud for s in specs if s.cloud != "modal"}) if value else [] if other: - raise ValueError( - f"--modal-gpu-exact applies only to Modal islands; --gpu also has {other}" - ) + raise ValueError(f"{flag} applies only to Modal islands; --gpu also has {other}") def _tail_modal(modal_ops, call_id: str, prefix: str, collector=None, guard=None) -> int: diff --git a/yeto/launch/rl_island_task.py b/yeto/launch/rl_island_task.py index 1185f4c9a..ce1ceb50a 100644 --- a/yeto/launch/rl_island_task.py +++ b/yeto/launch/rl_island_task.py @@ -153,6 +153,8 @@ def make_miles_island_task( + f" --inner-lr {args.inner_lr}" f" --seq-len {args.seq_len}" + (" --rl-activation-recompute" if getattr(args, "rl_activation_recompute", False) else "") + + (" --rl-gradient-accumulation-fusion" + if getattr(args, "rl_gradient_accumulation_fusion", False) else "") + (f" --rl-max-tokens-per-gpu {int(args.rl_max_tokens_per_gpu)}" if getattr(args, "rl_max_tokens_per_gpu", None) is not None else "") + (f" --rl-optimizer-offload-fraction {float(args.rl_optimizer_offload_fraction)!r}" diff --git a/yeto/launch/rl_ports.py b/yeto/launch/rl_ports.py index 63b41f24f..99f608019 100644 --- a/yeto/launch/rl_ports.py +++ b/yeto/launch/rl_ports.py @@ -533,6 +533,10 @@ def _ports_infra_flags(args) -> tuple[str, str]: flags += " --rl-observe-timeline" if getattr(args, "rl_debug_dump", False): flags += " --rl-debug-dump" + if getattr(args, "rl_profile_trainer", False): + flags += " --rl-profile-trainer" + if getattr(args, "rl_verify_every_round", False): + flags += " --rl-verify-every-round" if getattr(args, "rl_resource_sample_interval", None) is not None: flags += f" --rl-resource-sample-interval {float(args.rl_resource_sample_interval)!r}" flags += _rl_backend_module(args, "elastic_hook").recommend_flags(args) diff --git a/yeto/modal_runner.py b/yeto/modal_runner.py index b25be3c4f..761b53c9a 100644 --- a/yeto/modal_runner.py +++ b/yeto/modal_runner.py @@ -204,6 +204,10 @@ class ModalIslandConfig: # Request exactly `gpu` ("H100!" -- no H200 upgrade) and fail the # container at start-up unless nvidia-smi reports that type. gpu_exact: bool = False + # E1: fail the container at start-up unless every GPU nvidia-smi reports is this + # model ("NVIDIA A10" refuses "NVIDIA A10G": Modal "A10G" containers are either, + # and the A10G trains about 12 % slower); None: any card of type `gpu`. + gpu_name: str | None = None # Pull image_ref with a private-registry login: ``registry_creds`` (the # SKYPILOT_DOCKER_* triple launcher.registry_login_for resolved) when # given, else -- with ``registry_login`` -- the launching process's @@ -394,6 +398,19 @@ def check_gpu_names(gpu: str, names: list[str], count: int) -> None: ) +def gpu_model_matches(want: str, name: str) -> bool: + """``name`` is the GPU model ``want`` names, not a longer model name that starts + with it ("NVIDIA A10" matches "NVIDIA A10", not "NVIDIA A10G").""" + return re.search(re.escape(want) + r"(?![\w-])", name) is not None + + +def check_gpu_model(want: str, names: list[str], count: int) -> None: + """Raise unless exactly ``count`` GPUs of the model ``want`` are visible.""" + wrong = [n for n in names if not gpu_model_matches(want, n)] + if wrong or len(names) != count: + raise RuntimeError(f"requested {count}x {want!r} (--modal-gpu-name) but the container sees {names!r}") + + def visible_gpu_names() -> list[str]: out = subprocess.run( ["nvidia-smi", "--query-gpu=name", "--format=csv,noheader"], @@ -883,7 +900,7 @@ def _island_body(cfg: ModalIslandConfig, rank: int, ips: list[str], all_ips: dic try: names = visible_gpu_names() except (OSError, subprocess.SubprocessError) as exc: - if cfg.gpu_exact: + if cfg.gpu_exact or cfg.gpu_name: raise RuntimeError(f"island {cfg.learner_id}: cannot read GPU names: {exc}") from exc names = [] print(f"[modal-island {cfg.learner_id}] requested {cfg.gpu_request}, got {names}", flush=True) @@ -896,6 +913,8 @@ def _island_body(cfg: ModalIslandConfig, rank: int, ips: list[str], all_ips: dic tape.start() if cfg.gpu_exact: check_gpu_names(cfg.gpu, names, cfg.gpus_per_node) + if cfg.gpu_name: + check_gpu_model(cfg.gpu_name, names, cfg.gpus_per_node) if cfg.setup_script: code = subprocess.call(container_command(cfg.setup_script), env=env) if code != 0: diff --git a/yeto/rl/adapters/miles/config.py b/yeto/rl/adapters/miles/config.py index dadbb341d..1af82e71c 100644 --- a/yeto/rl/adapters/miles/config.py +++ b/yeto/rl/adapters/miles/config.py @@ -472,6 +472,7 @@ def _check_sglang_tp(value, config): # translated below (Megatron full recompute / Miles dynamic batching) "train_memory.activation_recompute": _ok, "train_memory.max_tokens_per_gpu": _ok, + "train_memory.gradient_accumulation_fusion": _ok, "train_memory.optimizer_offload_fraction": _ok, **{f"sample_buffer.{name}": _ok for name in ("role", "store", "max_backlog_groups", "wait_s", "min_round_s")}, @@ -781,6 +782,7 @@ def translate_run_config( from yeto.rl.engine.run_config import ( RECIPE_DEEPSEEK_V4_FLASH, + RECIPE_GENERIC, RECIPE_GLM5_NEXT, RECIPE_QWEN3_5, RECIPE_QWEN3_8_NEXT, @@ -864,6 +866,10 @@ def translate_run_config( model_recipe_values = deepseek_v4.recipe_argv(trainable) else: model_name = recipe.provider_class + if recipe.name == RECIPE_GENERIC: + # Megatron's default dropout is 0.1; the Miles recipes train RL with + # none (the policy's train-time log-probs must match the rollout's) + model_recipe_values = ["--attention-dropout", "0.0", "--hidden-dropout", "0.0"] values: list[str] = [ "train.py", @@ -938,7 +944,7 @@ def translate_run_config( "--accumulate-allreduce-grads-in-fp32", "--attention-softmax-in-fp32", "--attention-backend", recipe.training_attention_backend, - "--no-gradient-accumulation-fusion", + *(() if config.train_memory.gradient_accumulation_fusion else ("--no-gradient-accumulation-fusion",)), "--bf16", "--no-load-optim", "--no-load-rng", diff --git a/yeto/rl/adapters/miles/cut_plugin.py b/yeto/rl/adapters/miles/cut_plugin.py index 50fe3afde..98bd9d03c 100644 --- a/yeto/rl/adapters/miles/cut_plugin.py +++ b/yeto/rl/adapters/miles/cut_plugin.py @@ -54,6 +54,7 @@ _MODULE = "yeto.rl.adapters.miles.cut_plugin" SAVE_CUT_SHARD = f"{_MODULE}.save_cut_shard" +FINISH_CUT_SHARD = f"{_MODULE}.finish_cut_shard" RESTORE_CUT_SHARD = f"{_MODULE}.restore_cut_shard" SHARD_SCHEMA = "yeto.cut_shard/v1" # Opt-in: read the optimizer back right after the load and report where a mismatch arises. @@ -99,11 +100,17 @@ def _dps(self): return dp_invariant_state - def export_optimizer(self, optimizer: Any, named: list) -> Any: + def export_optimizer(self, optimizer: Any, named: list, *, on_device: bool = False) -> Any: + """The fork-M5 named optimizer export. ``on_device`` (E1): the tensors are + device copies (same layout, not yet on the host); :func:`to_host` turns the + result into exactly the host export.""" + export = self._dps().export_named_optimizer_state + if on_device: + export = _device_export(self._dps()) if _UNSAFE_STATE_READS[0]: # diagnostic only: verify the fork fix without the yeto guard - return self._dps().export_named_optimizer_state(optimizer, named) + return export(optimizer, named) with side_effect_free_state(optimizer): - return self._dps().export_named_optimizer_state(optimizer, named) + return export(optimizer, named) def check_optimizer(self, optimizer: Any, named: list, states: list) -> Any: """Merge the DP shards of one (tp, pp) (DistOpt ranges) and validate against ``optimizer``.""" @@ -140,6 +147,51 @@ def set_megatron_counters(self, counters: Mapping[str, int]) -> None: setattr(margs, key, int(value)) +def _device_export(dps: Any) -> Any: + """``export_named_optimizer_state`` keeping each tensor on its device (E1). + + The Miles fork takes ``device=None`` for this (yeto/s26-e1-device-export); the + pinned image's fork does not yet, and then this module builds the same + entries from the fork's own ``optimizer_slots`` -- the fork function's body + without its ``.cpu()`` (tests pin the equality with the fork's host export).""" + import inspect + + export = dps.export_named_optimizer_state + if "device" in inspect.signature(export).parameters: + return lambda optimizer, named: export(optimizer, named, device=None) + + def export_on_device(optimizer: Any, named: list) -> dict[str, Any]: + import copy + + entries = {} + for slot in dps.optimizer_slots(optimizer, named): + tensors, scalars = {}, {} + for key, value in slot.get().items(): + value = value.detach() + if value.dim() == 0: + scalars[key] = value.clone() + continue + flat = value.reshape(-1).clone() + if flat.numel() != slot.end - slot.start: + raise dps.DpInvariantStateError( + f"{slot.name!r}: state {key!r} has {flat.numel()} elements but the owned range is " + f"[{slot.start}, {slot.end})" + ) + tensors[key] = flat + entries[slot.name] = { + "numel": slot.model_param.numel(), + "shape": tuple(slot.model_param.shape), + "start": slot.start, + "end": slot.end, + "tensors": tensors, + "scalars": scalars, + "hyper": {k: copy.deepcopy(v) for k, v in slot.group.items() if k != "params"}, + } + return {"format": dps.NAMED_STATE_FORMAT, "entries": entries} + + return export_on_device + + def _expert_coord(ps: Any) -> dict[str, int]: """EP coordinates for the same-shape check (only when EP>1: older cuts stay comparable).""" if int(ps.ep.size) <= 1: @@ -332,10 +384,12 @@ def _feed(digest: Any, value: Any) -> None: shape = tuple(value.shape) t = value.detach().to("cpu").contiguous().reshape(-1) digest.update(f"T{t.dtype}{shape}".encode()) - digest.update(t.view(torch.uint8).numpy().tobytes() if t.numel() else b"") + # a buffer, not a bytes copy: same digest, and hashlib releases the GIL while it + # hashes (a background cut writer must not stall the training thread, E1) + digest.update(t.view(torch.uint8).numpy() if t.numel() else b"") elif isinstance(value, np.ndarray): digest.update(f"N{value.dtype}{value.shape}".encode()) - digest.update(np.ascontiguousarray(value).tobytes()) + digest.update(np.ascontiguousarray(value).view(np.uint8) if value.size else b"") elif isinstance(value, Mapping): digest.update(b"{") for key in sorted(value, key=repr): @@ -538,18 +592,33 @@ def _entry_step(entry: Mapping[str, Any]) -> Any: return None -def resume_rank_state(actor: Any, optimizer_named: Any) -> dict[str, Any]: - """This rank's :class:`~yeto.rl.engine.resume_manifest.OptimizerResumeState` (as a dict), - read from the named optimizer export and the Megatron LR scheduler (S25 N7.6).""" - from yeto.rl.engine.resume_manifest import build_rank_state, scheduler_progress_of +def resume_inputs(actor: Any) -> dict[str, Any]: + """What :func:`resume_rank_state` reads from the live trainer: the optimizer + group lrs and the Megatron LR scheduler progress.""" + from yeto.rl.engine.resume_manifest import scheduler_progress_of + + return {"lrs": [float(g["lr"]) for g in getattr(actor.optimizer, "param_groups", ()) if "lr" in g], + "scheduler_progress": scheduler_progress_of(actor.opt_param_scheduler)} + +def resume_rank_state(actor: Any, optimizer_named: Any, *, lrs: list[float] | None = None, + scheduler_progress: int | None = None) -> dict[str, Any]: + """This rank's :class:`~yeto.rl.engine.resume_manifest.OptimizerResumeState` (as a dict), + read from the named optimizer export and the Megatron LR scheduler (S25 N7.6). + ``lrs`` / ``scheduler_progress`` (from :func:`resume_inputs`, taken at the step + boundary) replace the reads of ``actor`` (then ``None`` is allowed).""" + from yeto.rl.engine.resume_manifest import build_rank_state + + if lrs is None or scheduler_progress is None: + live = resume_inputs(actor) + lrs = live["lrs"] if lrs is None else lrs + scheduler_progress = live["scheduler_progress"] if scheduler_progress is None else scheduler_progress entries = (optimizer_named or {}).get("entries") or {} steps = [s for s in (_entry_step(e) for e in entries.values()) if s is not None] - lrs = [float(g["lr"]) for g in getattr(actor.optimizer, "param_groups", ()) if "lr" in g] return build_rank_state( steps_by_group={"named": steps}, - scheduler_progress=scheduler_progress_of(actor.opt_param_scheduler), - lrs=lrs, + scheduler_progress=int(scheduler_progress), + lrs=list(lrs), moments_by_name={n: _entry_moments(e) for n, e in entries.items()}, ).to_dict() @@ -573,16 +642,76 @@ def resume_state(actor: Any) -> dict[str, Any]: RESUME_STATE = f"{_MODULE}.resume_state" -def _snapshot(actor: Any, backend: Any, named: list) -> dict[str, Any]: +def _snapshot(actor: Any, backend: Any, named: list, *, on_device: bool = False) -> dict[str, Any]: + """The shard state at this step boundary. ``on_device`` (E1): the adapter and + optimizer tensors are device copies (device-to-device, no wait for the host); + :func:`to_host` later makes them exactly the host snapshot.""" + if on_device: + adapter = {n: p.detach().clone() for n, p in named} + optimizer_named = backend.export_optimizer(actor.optimizer, named, on_device=True) + else: + adapter = {n: p.detach().to("cpu").clone() for n, p in named} + optimizer_named = backend.export_optimizer(actor.optimizer, named) return { - "adapter": {n: p.detach().to("cpu").clone() for n, p in named}, - "optimizer_named": backend.export_optimizer(actor.optimizer, named), + "adapter": adapter, + "optimizer_named": optimizer_named, "scheduler": actor.opt_param_scheduler.state_dict(), "megatron_counters": backend.megatron_counters(), "miles_counters": miles_counters(actor), } +def to_host(value: Any) -> Any: + """``value`` with every non-CPU tensor copied to the host (containers rebuilt, + everything else as it is). A CPU tensor is returned as it is.""" + import torch + + if isinstance(value, torch.Tensor): + return value if value.device.type == "cpu" else value.to("cpu") + if isinstance(value, dict): + return {k: to_host(v) for k, v in value.items()} + if isinstance(value, list): + return [to_host(v) for v in value] + if isinstance(value, tuple): + return tuple(to_host(v) for v in value) + return value + + +# E1: device copies of a background cut need this much free device memory besides +# their own size (twice the copies: the allocator's rounding and what training +# allocates next); below it the capture stays on the host path +DEVICE_COPY_ENV = "YETO_RL_CUT_DEVICE_COPY" +DEVICE_COPY_HEADROOM_BYTES = 1 << 30 + + +def _device_copy_plan(actor: Any, named: list) -> tuple[bool, str]: + """(copy on the device?, why): the device path needs CUDA parameters, a trainer + that stays resident, and free memory for the copies (adapter + FP32 master + + two Adam moments, estimated).""" + import torch + + if os.environ.get(DEVICE_COPY_ENV, "1").strip().lower() in {"0", "false", "no", "off"}: + return False, "disabled" + devices = {p.device for _, p in named} + if len(devices) != 1 or next(iter(devices)).type != "cuda": + return False, "not_cuda" + need = sum(p.numel() * (p.element_size() + 3 * 4) for _, p in named) + free, _total = torch.cuda.mem_get_info(next(iter(devices))) + if free < 2 * need + DEVICE_COPY_HEADROOM_BYTES: + return False, f"low_device_memory(free={free},need={need})" + if _sleeps_after_cut(actor): + # --offload-train: the trainer sleeps right after the cut (colocated), so nothing + # may read the device copies later (they failed with "CUDA error: invalid + # argument" in GPU run e1g-yeto-miles-sync-1): they are staged on the host + # before the capture returns (_stage) + return True, "ok_staged" + return True, "ok" + + +def _sleeps_after_cut(actor: Any) -> bool: + return bool(getattr(getattr(actor, "args", None), "offload_train", False)) + + # -------------------------------------------------------------------------- # Plugins # -------------------------------------------------------------------------- @@ -596,32 +725,418 @@ def shard_name(coord: Mapping[str, int]) -> str: REFUSAL_KIND = "refusal_kind" # "refused" (deliberate) | "failed_before_write" -def save_cut_shard(actor: Any, *, directory: str, cut_id: str) -> dict[str, Any]: +# E1: shards being written by a background thread of this rank, by cut id +_PENDING_SHARDS: dict[str, Any] = {} + + +def save_cut_shard(actor: Any, *, directory: str, cut_id: str, background: bool = False) -> dict[str, Any]: """Write this rank's shard (tmp + fsync + rename) and return its summary. A refusal (unsupported configuration, not at a step boundary, ...) is RETURNED as ``{"refused": reason}`` instead of raised: Miles marks a cell errored on any ``run_plugin`` exception (GPU run 2026-09-30, C1), which would turn a clean "no cut" into a dead trainer. + + ``background`` (E1): take the state now (at this step boundary) and write, hash + and summarize it off the training thread (:func:`_write_background`), so + training goes on; returns ``{"pending": cut_id}``. :func:`finish_cut_shard` + returns the summary. """ try: with trainer_resident(actor): - return _save(actor, directory=directory, cut_id=cut_id) + captured = _capture(actor, cut_id=cut_id, device_copy=background) except CutPluginError as error: return {REFUSED: str(error), REFUSAL_KIND: "refused"} + if not background: + return _write_refusing(captured, directory=directory, cut_id=cut_id) + from concurrent.futures import ThreadPoolExecutor + + if cut_id in _PENDING_SHARDS: + return {REFUSED: f"cut {cut_id} is already being written", REFUSAL_KIND: "refused"} + pool = ThreadPoolExecutor(max_workers=1, thread_name_prefix=f"yeto-cut-{cut_id}") + _PENDING_SHARDS[cut_id] = (pool, pool.submit(_write_background, captured, directory=directory, cut_id=cut_id)) + pool.shutdown(wait=False) + return {"pending": cut_id, "coord": dict(captured["coord"]), "capture": dict(captured["capture"])} + + +def finish_cut_shard(actor: Any, *, cut_id: str) -> dict[str, Any]: + """The summary of a background :func:`save_cut_shard` (waits for its thread). + Its ``writer`` key says where the shard was serialized and hashed.""" + pending = _PENDING_SHARDS.pop(cut_id, None) + if pending is None: + return {REFUSED: f"no background write of cut {cut_id} on this rank", REFUSAL_KIND: "failed_before_write"} + try: + return pending[1].result() + except Exception as error: # noqa: BLE001 - reported like a refusal: the cut is not committed + return {REFUSED: f"background write of cut {cut_id} failed: {error!r}", REFUSAL_KIND: "failed"} -def _save(actor: Any, *, directory: str, cut_id: str) -> dict[str, Any]: +def _write_refusing(captured: dict[str, Any], *, directory: str, cut_id: str) -> dict[str, Any]: + try: + return _write(_with_resume(_on_host(captured)), directory=directory, cut_id=cut_id) + except CutPluginError as error: + return {REFUSED: str(error), REFUSAL_KIND: "refused"} + + +# E1 phase 5: a background cut is serialized (torch.save), hashed and committed in a +# writer process of this rank, not in a thread of the trainer process: the writer +# thread competed with training (each trained token about 8 % slower while it ran, +# GPU runs e1g/e1h). The trainer copies the state into shared memory and hands it +# over; the writer process writes the same file and summary. "0" keeps the thread. +WRITER_PROCESS_ENV = "YETO_RL_CUT_WRITER_PROCESS" +_WRITER: list[Any] = [] # this rank's writer process pool (one worker), started lazily + + +def _writer_enabled() -> bool: + return os.environ.get(WRITER_PROCESS_ENV, "1").strip().lower() not in {"0", "false", "no", "off"} + + +def _writer_init() -> None: + """Writer process start: below training in CPU priority, one intra-op thread.""" + import torch + + try: + os.nice(10) + except OSError: + pass + torch.set_num_threads(1) + + +def _writer_pool() -> Any: + if not _WRITER: + import multiprocessing + from concurrent.futures import ProcessPoolExecutor + + import torch.multiprocessing # noqa: F401 - registers the shared-memory tensor pickling + + # spawn: the trainer has CUDA initialized (fork would copy it); the child never touches CUDA + _WRITER.append(ProcessPoolExecutor(max_workers=1, mp_context=multiprocessing.get_context("spawn"), + initializer=_writer_init)) + import atexit + + atexit.register(_drop_writer_pool) # stop it before interpreter teardown + return _WRITER[0] + + +def _drop_writer_pool() -> None: + if _WRITER: + _WRITER.pop().shutdown(wait=False, cancel_futures=True) + + +def _remove_partial(directory: str, coord: Mapping[str, int]) -> None: + """Remove what a dead writer process left of this rank's shard (no summary was + returned, so no manifest names it).""" + final = os.path.join(directory, shard_name(coord)) + for path in (final + ".tmp", final): + if os.path.exists(path): + os.remove(path) + + +def _write_background(captured: dict[str, Any], *, directory: str, cut_id: str) -> dict[str, Any]: + """Thread of the trainer process: the state to the host and into one shared buffer + (:func:`_pack`), then the writer process serializes, hashes and summarizes it. + When the writer process fails (it died, the state did not pickle, ...) this thread + writes the same shard itself and the summary's ``writer`` says so (``mode=thread``, + ``fallback``).""" + try: + try: + skeleton, layout, flat = _pack(_on_host(captured)) + finally: + _STAGING.free.set() # a staged snapshot (sync cuts) is copied out or dropped + except CutPluginError as error: + return {REFUSED: str(error), REFUSAL_KIND: "refused"} + fallback = "disabled" + if _writer_enabled(): + try: + result = _writer_pool().submit(_write_packed, skeleton, layout, flat, + directory=directory, cut_id=cut_id).result() + return {**result, "writer": {"mode": "process"}} + except Exception as error: # noqa: BLE001 - fall back to this thread, say why + _drop_writer_pool() + fallback = f"writer process failed: {error!r}"[:500] + print("yeto cut writer fallback " + repr({"cut_id": cut_id, "reason": fallback}), flush=True) + _remove_partial(directory, captured["coord"]) + return {**_write_refusing(_unpack(skeleton, layout, flat), directory=directory, cut_id=cut_id), + "writer": {"mode": "thread", "fallback": fallback}} + + +class _Staging: + """The reused host buffer of :func:`_stage` (pinned when CUDA is available) and + whether a staged snapshot still lives in it (set again once :func:`_pack` copied + it out). One per trainer process.""" + + ALIGN = 64 + + def __init__(self) -> None: + import threading + + self.buffer = None + self.free = threading.Event() + self.free.set() + + def flat(self, nbytes: int): + import torch + + if self.buffer is None or self.buffer.numel() < nbytes: + self.buffer = None + self.buffer = torch.empty(nbytes, dtype=torch.uint8, pin_memory=torch.cuda.is_available()) + return self.buffer + + +_STAGING = _Staging() +STAGING_WAIT_S = 600.0 + + +def _tensors(value: Any, out: list) -> list: + import torch + + if isinstance(value, torch.Tensor): + out.append(value) + elif isinstance(value, Mapping): + for v in value.values(): + _tensors(v, out) + elif isinstance(value, (list, tuple)): + for v in value: + _tensors(v, out) + return out + + +def _stage(snap: Any, *, include_cpu: bool = False) -> Any: + """``snap`` with every contiguous non-CPU tensor replaced by a view of the + process's staging buffer holding its exact bytes (copied asynchronously, ONE + synchronization at the end); other tensors (CPU unless ``include_cpu``, a test + hook; non-contiguous) are copied to the host on their own as :func:`to_host` does. The views are valid until + :func:`_pack` copies them out (the writer thread then marks the buffer free).""" + import torch + + if not _STAGING.free.wait(STAGING_WAIT_S): + raise CutPluginError("the previous cut's staged state was never handed to its writer") + stageable = [t for t in _tensors(snap, []) + if (include_cpu or t.device.type != "cpu") and t.is_contiguous()] + if not stageable: + return to_host(snap) + _STAGING.free.clear() + align = _Staging.ALIGN + offsets, total = {}, 0 + for t in stageable: + if id(t) not in offsets: + offsets[id(t)] = total + total += -(-t.numel() * t.element_size() // align) * align + flat = _STAGING.flat(total) + views: dict[int, Any] = {} + + def staged(value: Any) -> Any: + if isinstance(value, torch.Tensor): + if id(value) in offsets: + if id(value) not in views: + start = offsets[id(value)] + view = flat[start:start + value.numel() * value.element_size()].view(value.dtype) + view = view.view(value.shape) + view.copy_(value, non_blocking=True) + views[id(value)] = view + return views[id(value)] + return value if value.device.type == "cpu" else value.to("cpu") + if isinstance(value, dict): + return {k: staged(v) for k, v in value.items()} + if isinstance(value, list): + return [staged(v) for v in value] + if isinstance(value, tuple): + return tuple(staged(v) for v in value) + return value + + try: + out = staged(snap) + devices = {t.device for t in stageable if t.device.type == "cuda"} + for device in devices: + torch.cuda.synchronize(device) # every staged copy is complete + except BaseException: + _STAGING.free.set() + raise + return out + + + +_PACKED = "__yeto_packed__" + + +def _pack(value: Any) -> tuple[Any, list, Any]: + """(skeleton, layout, flat): ``value`` with every packable CPU tensor replaced by + ``{_PACKED: index}`` and its bytes copied into ONE shared-memory buffer ``flat`` + (``layout[index]`` = offset, dtype, shape). Packable: contiguous and the sole + owner of its whole storage, so :func:`_unpack` rebuilds a tensor that saves to + the same bytes. Other tensors stay in the skeleton, moved to shared memory one by + one. One buffer is one file descriptor to hand over: per-tensor sharing passed + about 400 descriptors per cut through the trainer process's resource sharer + while the next round trained (E1 profile run e1j-yeto-miles-async-a).""" + import torch + + tensors: list = [] + _collect(value, tensors) + owners: dict[int, int] = {} + for t in tensors: + key = t.untyped_storage().data_ptr() + owners[key] = owners.get(key, 0) + 1 + + staging = _STAGING.buffer + staged_ptr = None if staging is None else staging.untyped_storage().data_ptr() + + def packable(t: Any) -> bool: + storage = t.untyped_storage() + if staged_ptr is not None and storage.data_ptr() == staged_ptr and t.device.type == "cpu": + return t.is_contiguous() and t.numel() > 0 # a view of the staging buffer: copied out + return (t.device.type == "cpu" and t.is_contiguous() and t.storage_offset() == 0 + and storage.nbytes() == t.numel() * t.element_size() and owners[storage.data_ptr()] == 1 + and t.numel() > 0) + + chosen = [t for t in tensors if packable(t)] + offsets, total = {}, 0 + for t in chosen: + offsets[id(t)] = total + total += -(-t.numel() * t.element_size() // 64) * 64 + flat = torch.empty(max(total, 1), dtype=torch.uint8).share_memory_() + layout: list = [] + + def walk(v: Any) -> Any: + if isinstance(v, torch.Tensor): + if id(v) not in offsets: + return v.share_memory_() if v.device.type == "cpu" else v + start, nbytes = offsets[id(v)], v.numel() * v.element_size() + flat[start:start + nbytes].copy_(v.reshape(-1).view(torch.uint8)) + layout.append((start, v.dtype, tuple(v.shape))) + return {_PACKED: len(layout) - 1} + if isinstance(v, dict): + return {k: walk(x) for k, x in v.items()} + if isinstance(v, list): + return [walk(x) for x in v] + if isinstance(v, tuple): + return tuple(walk(x) for x in v) + return v + + return walk(value), layout, flat + + +def _collect(value: Any, out: list) -> None: + import torch + + if isinstance(value, torch.Tensor): + out.append(value) + elif isinstance(value, Mapping): + for v in value.values(): + _collect(v, out) + elif isinstance(value, (list, tuple)): + for v in value: + _collect(v, out) + + +def _unpack(skeleton: Any, layout: list, flat: Any) -> Any: + """The value :func:`_pack` packed: every packed tensor is a fresh tensor holding + its own copy of the bytes.""" + import torch + + def walk(v: Any) -> Any: + if isinstance(v, dict) and _PACKED in v and len(v) == 1: + start, dtype, shape = layout[v[_PACKED]] + numel = 1 + for d in shape: + numel *= int(d) + nbytes = numel * torch.empty((), dtype=dtype).element_size() + return flat[start:start + nbytes].clone().view(dtype).view(shape) + if isinstance(v, dict): + return {k: walk(x) for k, x in v.items()} + if isinstance(v, list): + return [walk(x) for x in v] + if isinstance(v, tuple): + return tuple(walk(x) for x in v) + return v + + return walk(skeleton) + + +def _write_packed(skeleton: Any, layout: list, flat: Any, *, directory: str, cut_id: str) -> dict[str, Any]: + """Writer process: :func:`_write_refusing` of the unpacked state.""" + return _write_refusing(_unpack(skeleton, layout, flat), directory=directory, cut_id=cut_id) + + +def _capture(actor: Any, *, cut_id: str, device_copy: bool = False) -> dict[str, Any]: + """Everything the shard and its summary need, read at the step boundary. + + Host path: CPU copies, taken here. ``device_copy`` (E1, background cuts): the + adapter and optimizer tensors are copied device-to-device here and reach the + host in the writer thread (:func:`_on_host`), so the trainer waits only for + the device copies. The device path is skipped (host path, reason in + ``capture``) without CUDA parameters, with too little free device memory, or + when a device copy runs out of memory; the shard is the same either way.""" import torch backend = _backend(actor) coord = _require(actor, backend) if _STEP_GRAD_NORMS or _STEP_APPLIED_LRS or _STEP_LOSSES: raise CutPluginError("per-step records not drained: the cut is not at an optimizer-step boundary") + import json + import time + + marks = [time.monotonic()] named = _adapters(actor, backend) + on_device, reason = _device_copy_plan(actor, named) if device_copy else (False, "inline") with torch.no_grad(): - snap = _snapshot(actor, backend, named) + snap = None + if on_device: + try: + snap = _snapshot(actor, backend, named, on_device=True) + except torch.cuda.OutOfMemoryError: + snap, on_device, reason = None, False, "device_oom" + torch.cuda.empty_cache() + if snap is None: + snap = _snapshot(actor, backend, named) + ready = None + if on_device and _sleeps_after_cut(actor): + snap = _stage(snap) # host views; the device copies are dropped here + elif on_device and any(p.is_cuda for _, p in named): + ready = torch.cuda.Event() + ready.record() + marks.append(time.monotonic()) rng = backend.capture_rng() + capture = {"mode": "device" if on_device else "host", "reason": reason} + captured = {"coord": coord, "snap": snap, "rng": rng, "capture": capture, "ready": ready, + "scheduler_samples": int(actor.opt_param_scheduler.num_steps), + "tuning": _tuning(backend, named), "adapter_names": sorted(n for n, _ in named), + # what resume_rank_state reads from the live trainer, as it is now + "resume_inputs": resume_inputs(actor)} + marks.append(time.monotonic()) + if not on_device: + captured["resume"] = resume_rank_state(actor, snap["optimizer_named"], **captured["resume_inputs"]) + marks.append(time.monotonic()) + # E1 diagnosis: where the step-boundary capture of a cut goes (driver log only) + print("yeto cut capture timing " + json.dumps({**dict(zip( + ("snapshot_s", "rng_tuning_s", "resume_state_s"), + (round(b - a, 3) for a, b in zip(marks, marks[1:])))), "capture": capture}), flush=True) + return captured + + +def _on_host(captured: dict[str, Any]) -> dict[str, Any]: + """``captured`` with its device copies on the host (runs off the training + thread). Host captures are returned as they are.""" + ready = captured.get("ready") + if ready is None and "resume" in captured: + return captured + if ready is not None: + ready.synchronize() # the device copies are complete (whatever stream training uses) + return {**captured, "snap": to_host(captured["snap"]), "ready": None} + + +def _with_resume(captured: dict[str, Any]) -> dict[str, Any]: + """``captured`` with the N7.6 resume state a device capture defers (computed from + the host snapshot and the step-boundary resume inputs).""" + if "resume" in captured: + return captured + return {**captured, "resume": resume_rank_state(None, captured["snap"]["optimizer_named"], + **captured["resume_inputs"])} + + +def _write(captured: dict[str, Any], *, directory: str, cut_id: str) -> dict[str, Any]: + import torch + + coord, snap, rng = captured["coord"], captured["snap"], captured["rng"] shard = {"schema": SHARD_SCHEMA, "cut_id": cut_id, "coord": dict(coord), **snap, "rng": rng} os.makedirs(directory, exist_ok=True) name = shard_name(coord) @@ -637,13 +1152,13 @@ def _save(actor: Any, *, directory: str, cut_id: str) -> dict[str, Any]: import io buffer = io.BytesIO() - torch.save(to_safe(shard), buffer) + torch.save(_interned(to_safe(shard)), buffer) data = buffer.getvalue() fh.write(data[: len(data) // 2]) fh.flush() os.fsync(fh.fileno()) cut_injection.kill_now(f"save_cut shard {name} half written") - torch.save(to_safe(shard), fh) + torch.save(_interned(to_safe(shard)), fh) fh.flush() os.fsync(fh.fileno()) os.replace(tmp, final) @@ -652,7 +1167,7 @@ def _save(actor: Any, *, directory: str, cut_id: str) -> dict[str, Any]: "sha256": _sha256(final), "bytes": os.path.getsize(final), "coord": dict(coord), - "scheduler_samples": int(actor.opt_param_scheduler.num_steps), + "scheduler_samples": captured["scheduler_samples"], "has_optimizer_state": snap["optimizer_named"] is not None, "has_rng": rng is not None, "state_digest": state_digest(snap), @@ -660,14 +1175,33 @@ def _save(actor: Any, *, directory: str, cut_id: str) -> dict[str, Any]: "weight_version": (snap.get("miles_counters") or {}).get("weight_version"), "rng_digest": state_digest(rng), "components": component_digests({"state": snap, "rng": rng}), - "tuning": _tuning(backend, named), - "adapter_tensors": len(named), - "adapter_names": sorted(n for n, _ in named), + "tuning": captured["tuning"], + "adapter_tensors": len(captured["adapter_names"]), + "adapter_names": captured["adapter_names"], "optimizer_names": sorted((snap["optimizer_named"] or {}).get("entries", {})), - "resume": resume_rank_state(actor, snap["optimizer_named"]), + "resume": captured["resume"], } +def _interned(value: Any) -> Any: + """``value`` with every string interned. Pickle writes a string once per object + and back-references repeats of the same object; interning makes that depend on + the values only, so a shard serialized in the writer process (whose strings + arrived unpickled, as separate objects) has the same bytes as one serialized in + the trainer (E1 phase 5).""" + import sys + + if isinstance(value, str): + return sys.intern(value) + if isinstance(value, dict): + return {_interned(k): _interned(v) for k, v in value.items()} + if isinstance(value, list): + return [_interned(v) for v in value] + if isinstance(value, tuple): + return tuple(_interned(v) for v in value) + return value + + def _load_verified(directory: str, entry: Mapping[str, Any], cut_id: str) -> dict[str, Any]: import torch diff --git a/yeto/rl/adapters/miles/entry.py b/yeto/rl/adapters/miles/entry.py index f05f8ac4e..40745e744 100644 --- a/yeto/rl/adapters/miles/entry.py +++ b/yeto/rl/adapters/miles/entry.py @@ -26,6 +26,7 @@ from collections.abc import Awaitable, Callable, Mapping, Sequence from typing import Any +from yeto.rl.engine.resume import DEFAULT_CUT_EVERY from yeto.rl.engine.algorithm import ( BOUNDED_NONZERO_STD_FILTER, SECRLENV_GROUP_CHECK_FILTER, @@ -830,6 +831,8 @@ def expected_policy() -> tuple[int, str]: ), **({"evaluate_start": evaluate_start} if evaluate_start is not None else {}), ) + # E1: 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)) if elastic is not None: from .elastic_placement import ElasticPlacement @@ -1040,7 +1043,7 @@ def _wire_round_cuts(driver, *, elastic, miles_args, algorithm, base_model_revis checkpoint = wire_round_cuts( driver, controller=controller, source=source, every=int(getattr(elastic, "every", None) or config.get("every") - or getattr(miles_args, "yeto_rl_resume_every", None) or 1), + or getattr(miles_args, "yeto_rl_resume_every", None) or DEFAULT_CUT_EVERY), fingerprint=run_fingerprint_of(miles_args, algorithm), allow_config_change=bool(getattr(elastic, "allow_config_change", False) or config.get("allow_config_change", False)), diff --git a/yeto/rl/adapters/miles/fully_async_prompts.py b/yeto/rl/adapters/miles/fully_async_prompts.py index b448cdbda..bf1fbd95e 100644 --- a/yeto/rl/adapters/miles/fully_async_prompts.py +++ b/yeto/rl/adapters/miles/fully_async_prompts.py @@ -9,11 +9,13 @@ * round accounting: groups drawn since the last report and Miles' stale-group count (reported unknown, with a warning, when Miles does not expose it); -* cuts (decision P0-S-Q3): :func:`requeue_untrained` stops the background - worker and puts every drawn, untrained group back into the data source - buffer, reset for a fresh generation. The cut then carries them like any - buffered group, so the resumed run generates them again; nothing in flight - is saved. +* cuts (decision P0-S-Q3): background generation is never saved; every drawn, + untrained group goes into the cut reset for a fresh generation, so the + resumed run generates it again. :func:`snapshot_untrained` (E1, the cut path) + records them, the data source buffer and the data cursor in one step on the + worker's event loop and leaves the worker running: the live run keeps + generating ahead across the cut. :func:`requeue_untrained` (the older path) + stops the worker and puts the groups back into the live data source. Runs inside the rollout executor actor (``__ray_call__``). Import-light. """ @@ -117,6 +119,51 @@ def current() -> PromptLedger | None: return _LEDGER +def _reset_copies(groups: Iterable[Sequence[Any]]) -> list[list[Any]]: + copies = [copy.deepcopy(list(group)) for group in groups] + for group in copies: + for sample in group: + sample.reset_for_retry() + return copies + + +def snapshot_untrained(executor: Any, *, timeout_s: float = WORKER_STOP_TIMEOUT_S) -> dict[str, Any]: + """The cut's view of the fully_async prompts, taken without stopping the worker. + + Returns ``groups`` (the data source buffer, then every drawn, untrained group + as a reset copy, in group order) and ``data_cursor`` (the data source position, + ``rollout_meta_hook.data_cursor``). Both are read in ONE step on the worker's + event loop, where the worker draws prompts, so no group is both behind the + cursor and missing from ``groups``, or counted twice. A resume restores the + cursor and puts ``groups`` into the buffer: exactly the prompts the run had + drawn but not trained are generated again. The live data source and the + worker are left untouched.""" + from .rollout_meta_hook import data_cursor + + source = getattr(executor, "data_source", None) + source = getattr(source, "__self__", source) + ledger = getattr(source, _LEDGER_ATTR, None) + if ledger is None: + raise PromptLedgerError("fully_async prompts were not tracked; the cut would lose them") + + def read() -> dict[str, Any]: + cursor, _length = data_cursor(source) + buffered = [list(group) for group in (getattr(source, "buffer", None) or [])] + untrained = [group for _, group in sorted(ledger.untrained.items())] + return {"groups": _reset_copies(buffered + untrained), "data_cursor": cursor, + "untrained_groups": len(untrained)} + + worker = getattr(getattr(executor, "generate_rollout", None), "_worker", None) + loop = None if worker is None or worker.done() else worker.get_loop() + if loop is None or not loop.is_running() or getattr(loop, "_thread_id", None) == threading.get_ident(): + return read() # no concurrent draw: nothing else runs this loop now + + async def on_loop() -> dict[str, Any]: + return read() + + return asyncio.run_coroutine_threadsafe(on_loop(), loop).result(timeout=timeout_s) + + def requeue_untrained(executor: Any, *, timeout_s: float = WORKER_STOP_TIMEOUT_S) -> int: """Stop the background worker and return every drawn, untrained group to the data source, reset for a fresh generation. Returns the group count. @@ -134,10 +181,7 @@ def requeue_untrained(executor: Any, *, timeout_s: float = WORKER_STOP_TIMEOUT_S raise PromptLedgerError("fully_async prompts were not tracked; the cut would lose them") rollout_fn = getattr(executor, "generate_rollout", None) _stop_worker(rollout_fn, timeout_s) - groups = [copy.deepcopy(group) for _, group in sorted(ledger.untrained.items())] - for group in groups: - for sample in group: - sample.reset_for_retry() + groups = _reset_copies(group for _, group in sorted(ledger.untrained.items())) if groups: ledger.data_source.add_samples(groups) # the ledger forgets them if rollout_fn is not None: diff --git a/yeto/rl/adapters/miles/in_flight.py b/yeto/rl/adapters/miles/in_flight.py index 44a32207b..4e896ee22 100644 --- a/yeto/rl/adapters/miles/in_flight.py +++ b/yeto/rl/adapters/miles/in_flight.py @@ -132,18 +132,23 @@ def _codex_suspension() -> tuple[Mapping[str, Any], Mapping[str, Any]]: def export_in_flight(executor: Any) -> dict[str, Any]: """Runs INSIDE the rollout executor actor: the in-flight entries and the - data buffer length they account for.""" + data buffer length they account for (fully_async: also ``data_cursor``, the + data source position read together with them).""" + gates, inflight = _codex_suspension() if getattr(getattr(executor, "args", None), "fully_async", False): - # P0-S decision 3: background generation is discarded, never saved; - # its prompts go back into the data buffer (reset), so the cut carries - # them like any buffered group and the resumed run generates them again. - from .fully_async_prompts import requeue_untrained - - requeue_untrained(executor) + # P0-S decision 3: background generation is never saved; the cut carries + # every drawn, untrained prompt (reset) like a buffered group, so the + # resumed run generates it again. E1: read without stopping the worker + # (the live run keeps generating ahead); the cursor comes from the same + # read, so the cut's cursor and groups agree. + from .fully_async_prompts import snapshot_untrained + + snap = snapshot_untrained(executor) + return {"entries": buffered_entries(snap["groups"]) + agentic_entries(gates, inflight), + "buffer_groups": len(snap["groups"]), "data_cursor": snap["data_cursor"]} source = getattr(executor, "data_source", None) source = getattr(source, "__self__", source) buffer = list(getattr(source, "buffer", None) or []) - gates, inflight = _codex_suspension() entries = buffered_entries(buffer) + agentic_entries(gates, inflight) return {"entries": entries, "buffer_groups": len(buffer)} diff --git a/yeto/rl/adapters/miles/island_entry.py b/yeto/rl/adapters/miles/island_entry.py index 48585224f..a54105ec4 100644 --- a/yeto/rl/adapters/miles/island_entry.py +++ b/yeto/rl/adapters/miles/island_entry.py @@ -18,6 +18,7 @@ from yeto.rl.bridge import BridgeConfig from yeto.rl.decoupled import DecoupledBridgeConfig +from yeto.rl.engine.resume import DEFAULT_CUT_EVERY from yeto.rl.engine.run_flags import ISLAND, add_rl_flags, rl_option_pairs from yeto.rl.adapters.miles.models.deepseek_v4.expert_full import expert_full_specs from yeto.rl.export import adapter_targets, derive_peft_lora_specs @@ -282,6 +283,13 @@ def apply_ports_infra_switches(args, miles_args, environ=None) -> None: if getattr(args, "rl_overlap_eval", False): miles_args.yeto_rl_overlap_eval = True + if getattr(args, "rl_verify_every_round", False): # E1: read by entry (driver.verify_every_round) + miles_args.yeto_rl_verify_every_round = True + if getattr(args, "rl_profile_trainer", False): # E1: read by state_plugin.train_one_step + miles_args.yeto_rl_profile_trainer = True + from yeto.rl.engine import phase_profile + + phase_profile.enable() # samples the thread that drives the phases if getattr(args, "rl_debug_dump", False): out = os.path.dirname(os.path.abspath(os.path.expanduser( getattr(args, "event_tape", None) or "~/yeto-output/rl-island.jsonl"))) @@ -328,7 +336,7 @@ def apply_ports_infra_switches(args, miles_args, environ=None) -> None: miles_args.yeto_rl_resume = { "store": str(args.rl_resume_store), "state_dir": str(getattr(args, "rl_resume_state_dir", None) or "~/yeto-rl/resume-state"), - "every": int(getattr(args, "rl_cut_every", None) or 1), + "every": int(getattr(args, "rl_cut_every", None) or DEFAULT_CUT_EVERY), "keep": int(getattr(args, "rl_cut_keep", None) or 2), "allow_config_change": bool(getattr(args, "rl_resume_allow_config_change", False)), "stop_after": (None if getattr(args, "rl_stop_after_rounds", None) is None diff --git a/yeto/rl/adapters/miles/publish.py b/yeto/rl/adapters/miles/publish.py index 990843ef3..2bd8a1481 100644 --- a/yeto/rl/adapters/miles/publish.py +++ b/yeto/rl/adapters/miles/publish.py @@ -465,11 +465,19 @@ def _check_trainer_holds(self, state: Any) -> str: return tensor_hash def publish(self, state: TrainableState, *, token_rollout_id: int | None = None) -> PublicationResult: + started = time.monotonic() tensor_hash = self._check_trainer_holds(state) payload_hash, payload_bytes = _payload_of(state) version = state.policy_version if token_rollout_id is None else token_rollout_id token = policy_token(version, tensor_hash) - members, checksums = self._runner.run(self._publish(version, token)) + self.last_timing = {"check_s": time.monotonic() - started} + # E1: a round whose policy is not hashed (state.verified False) also skips the + # engine read-back; the next verified publication reads it back again + members, checksums = self._runner.run( + self._publish(version, token, read_back=getattr(state, "verified", True) is not False)) + # E1 diagnosis: where a publication's wall time goes (driver log only) + print("yeto publish timing " + json.dumps( + {k: round(v, 3) for k, v in self.last_timing.items()}, sort_keys=True), flush=True) self.last_engine_checksums = checksums bodies = list((checksums or {}).values()) self._reference = ( @@ -531,7 +539,7 @@ async def _stamp_version(self, version: int, label: str) -> tuple[list[Any], lis return engines, ids, members async def _publish( - self, version: int, token: str + self, version: int, token: str, *, read_back: bool = True ) -> tuple[frozenset[str], dict[str, Any] | None]: if getattr(self._args, "debug_train_only", False): # Trainer-only island (Miles --debug-train-only, samples from the store): @@ -547,18 +555,24 @@ async def _publish( # the first (engines start resident); KV comes back after it. if offload_rollout and self._published_once: await self._controller.onload_weights() + mark = time.monotonic() await update_weights(self._args, self._actor, self._executor, self._controller) + self._time("update_weights_s", mark) if offload_rollout: await self._controller.onload_kv() except Exception as exc: raise PublicationError(f"upstream update_weights failed: {exc}", cause=PublicationCause.UPDATE_FAILED) from exc self._published_once = True + mark = time.monotonic() engines, ids, members = await self._stamp_version(version, token) + self._time("stamp_s", mark) checksums = None - if self._verify_checksums: + if self._verify_checksums and read_back: + mark = time.monotonic() bodies = await read_engine_checksums(self._controller, self._flatten) + self._time("checksum_s", mark) if len(bodies) != len(engines): raise PublicationError( f"{len(engines) - len(bodies)} engines returned no weight checksum" @@ -567,6 +581,11 @@ async def _publish( return members, checksums + def _time(self, key: str, mark: float) -> None: + timing = getattr(self, "last_timing", None) + if timing is not None: + timing[key] = time.monotonic() - mark + # -- convertible node (rl-trainer-start-stop L7.2) -------------------------- @property def published_checksum(self) -> str | None: diff --git a/yeto/rl/adapters/miles/rollout.py b/yeto/rl/adapters/miles/rollout.py index df2080d54..27edd5b35 100644 --- a/yeto/rl/adapters/miles/rollout.py +++ b/yeto/rl/adapters/miles/rollout.py @@ -361,6 +361,8 @@ def __init__(self, name: str = DEFAULT_SINK_ACTOR, *, timeout_s: float = 30.0) - self.timeout_s = timeout_s self._ray = ray self._actor = _sink_actor_class().options(name=name, get_if_exists=True).remote() + self._late: dict[int, dict[str, Any]] = {} + self._held: list[dict[str, Any]] = [] @property def sink_spec(self) -> str: @@ -379,10 +381,36 @@ def set_max_policy_age(self, limit: int) -> None: def set_published_policies(self, policies: Mapping[str, str]) -> None: self._ray.get(self._actor.set_policies.remote(json.dumps(dict(policies), sort_keys=True))) + def _take_items(self) -> list[dict[str, Any]]: + """Every pushed record not yet taken; late round records (E1) are kept aside + by rollout id, the others are returned (and forgotten) in push order.""" + from .rollout_meta_hook import LATE_META_SCHEMA + + for raw in self._ray.get(self._actor.take_all.remote()): + item = json.loads(raw) + if item.get("schema") == LATE_META_SCHEMA: + self._late[int(item["rollout_id"])] = item + else: + self._held.append(item) + items, self._held = self._held, [] + return items + + def take_late(self, rollout_id: int, *, timeout_s: float | None = None) -> dict[str, Any] | None: + """The late record of ``rollout_id`` (the deferred cross-version estimate), + waiting up to ``timeout_s``; None when it did not arrive.""" + deadline = time.monotonic() + (self.timeout_s if timeout_s is None else timeout_s) + while True: + self._held = self._take_items() # a next round's records stay for take() + if rollout_id in self._late: + return self._late.pop(rollout_id) + if time.monotonic() > deadline: + return None + time.sleep(0.05) + def take(self, rollout_id: int) -> dict[str, Any]: deadline = time.monotonic() + self.timeout_s while True: - items = [json.loads(i) for i in self._ray.get(self._actor.take_all.remote())] + items = self._take_items() main = [i for i in items if i.get("schema") != ROUND_META_SCHEMA] extra = [i for i in items if i.get("schema") == ROUND_META_SCHEMA] if items: @@ -658,6 +686,7 @@ def __init__( self._controller = inference_controller self._executor = rollout_executor self._metadata = metadata + self._deferred_estimates: set[int] = set() # E1: late cross-version estimates self._expected_policy = expected_policy self._run = (runner or LoopRunner()).run self._last_cursor: dict[str, int] | None = None @@ -781,6 +810,17 @@ def _in_executor(self, fn: Any, *fn_args: Any, **fn_kwargs: Any) -> Any: raise RuntimeError("Ray actor handle %s has no __ray_call__" % type(target).__name__) return self._run(_awaited(remote(fn, *fn_args, **fn_kwargs))) + def late_cross_version_fraction(self, rollout_id: int) -> float | None: + """E1: the deferred cross-version IS estimate of ``rollout_id`` (fully_async), + None when the round deferred none or it did not arrive in time.""" + take_late = getattr(self._metadata, "take_late", None) + if not callable(take_late) or rollout_id not in self._deferred_estimates: + return None + self._deferred_estimates.discard(rollout_id) + late = take_late(rollout_id) + value = (late or {}).get("cross_version_truncated_fraction") + return None if value is None else float(value) + def export_in_flight(self) -> dict[str, Any]: """agentic-rollout-utilization 3.3 (limit > 0): the cut's in-flight entries (buffered groups, suspended agentic trajectories); raises when @@ -834,9 +874,12 @@ def generate( self._prompt_ledger_installed = True self._run(self._controller.prepare_rollout(rollout_id)) data_pack = self._run(self._executor.get(rollout_id)) - self._offload_after_rollout() + if not self.defer_offload: + self._offload_after_rollout() try: payload = self._metadata.take(rollout_id) + if payload.get("cross_version_estimate") == "deferred": + self._deferred_estimates.add(int(rollout_id)) handle = handle_from_metadata( payload, rollout_id=rollout_id, @@ -1045,6 +1088,17 @@ def set_admission(self, open_: bool) -> None: harness admission to this pool's running members (no-op without a board).""" self._harness_admission("open_admission" if open_ else "close_admission", self.members()) + # E1: set by a driver that runs the --offload-rollout step in its own phase + # (offload_after_generate); other callers keep the offload inside generate + defer_offload = False + + @property + def offloads_after_generate(self) -> bool: + return self._args is not None and bool(getattr(self._args, "offload_rollout", False)) + + def offload_after_generate(self) -> None: + self._offload_after_rollout() + def _offload_after_rollout(self) -> None: """Upstream ``train.py`` serial colocated branch (--offload-rollout).""" diff --git a/yeto/rl/adapters/miles/rollout_meta_hook.py b/yeto/rl/adapters/miles/rollout_meta_hook.py index d34e761d9..ea9e1f22a 100644 --- a/yeto/rl/adapters/miles/rollout_meta_hook.py +++ b/yeto/rl/adapters/miles/rollout_meta_hook.py @@ -35,12 +35,26 @@ import os import statistics import tempfile +import threading from collections.abc import Iterable, Mapping, Sequence +from contextlib import contextmanager from pathlib import Path from typing import Any META_SINK_ENV = "YETO_ROLLOUT_META_SINK" DEFAULT_SINK_ACTOR = "yeto_rollout_meta" +# E1: per-step seconds of the last metadata hook run (printed by record_trained_groups) +_HOOK_TIMES: dict[str, float] = {} + + +def _timed(name: str, fn: Any, *args: Any, **kwargs: Any) -> Any: + import time + + started = time.monotonic() + try: + return fn(*args, **kwargs) + finally: + _HOOK_TIMES[name] = _HOOK_TIMES.get(name, 0.0) + time.monotonic() - started DEFAULT_SINK = f"ray:{DEFAULT_SINK_ACTOR}" METADATA_SCHEMA = "yeto-rollout-meta-v1" _TRAINED_ATTR = "_yeto_trained_group_keys" @@ -54,6 +68,9 @@ # rollout's metadata by the driver-side reader, so they are never attributed # to the next round. ROUND_META_SCHEMA = "yeto-rl-round-metadata-v1" +# E1: the cross-version IS estimate of a fully_async round, sent after the round's +# metadata (the prefill scoring runs in the background, off the rollout's critical path) +LATE_META_SCHEMA = "yeto-rl-late-round-metadata-v1" # buffer_stale_groups: groups the sample-buffer reader passed over this round # because they were older than max_policy_age (consumed, not trained) ROUND_METADATA_KEYS = frozenset({"nonzero_advantages", "buffer_stale_groups"}) @@ -136,15 +153,24 @@ def record_trained_groups(args: Any, data: Sequence[Sequence[Any]]) -> None: to its buffer (``sglang_rollout.py:505-510``); only samples aborted under ``--partial-rollout`` go back. """ + with one_sink_view(): + _record_trained_groups(args, data) + + +def _record_trained_groups(args: Any, data: Sequence[Sequence[Any]]) -> None: + import time from yeto.rl.algos.sample_filters import apply_sample_filters from yeto.rl.harness.reasoning_loss_mask import apply_from_args as apply_reasoning_loss_policy + marks = [("start", time.monotonic())] apply_sample_filters(args, data) + marks.append(("sample_filters", time.monotonic())) # Side effect beyond filtering (S17 WP6): with --codex-exclude-reasoning-from-loss # this sets loss_mask 1 -> 0 on generated reasoning tokens of the kept samples # (before Miles converts them to train data). Default: returns at once, no sample touched. apply_reasoning_loss_policy(args, data, tokenizer_loader=_load_miles_tokenizer) + marks.append(("reasoning_mask", time.monotonic())) # S19 #8: a stop token forced by the engine's grammar mask (SGLang returns # logprob exactly 0.0 under the masked distribution) is not a model choice: # loss_mask 1 -> 0 before Miles converts the samples to train data. @@ -152,11 +178,13 @@ def record_trained_groups(args: Any, data: Sequence[Sequence[Any]]) -> None: setattr(args, _PLACEHOLDER_ATTR, mask_placeholder_logprobs(data, tokenizer_loader=_load_miles_tokenizer, args=args)) + marks.append(("placeholder_mask", time.monotonic())) # CompactionRL: the masks above may drop optimised tokens; the per-segment # counts (tokens_after, gae_length) must follow the final masks. from yeto.rl.harness.codex.codex_openenv_generate import refresh_segment_counts refresh_segment_counts([s for group in data for s in _flat(group)]) + marks.append(("segment_counts", time.monotonic())) setattr(args, _TRAINED_ATTR, {_group_key(group) for group in data}) # Full-parameter islands (P0-D8): bind the exact accepted batch, after the # final loss masks, for the role streams to read at the round boundary. @@ -176,7 +204,16 @@ def record_trained_groups(args: Any, data: Sequence[Sequence[Any]]) -> None: if ledger is None: raise RuntimeError("fully_async round without the yeto prompt ledger") ledger.mark_trained(data) - extract_rollout_metadata(args, data, round_counts=ledger.take_round_counts()) + counts = ledger.take_round_counts() + marks.append(("ledger", time.monotonic())) + extract_rollout_metadata(args, data, round_counts=counts) + marks.append(("round_metadata", time.monotonic())) + # E1: where the hook's time goes (it runs on the rollout's critical path) + detail = " ".join(f"{k}={v:.3f}" for k, v in sorted(_HOOK_TIMES.items())) + _HOOK_TIMES.clear() + print(f"[rl] sample filter hook: total={marks[-1][1] - marks[0][1]:.3f}s " + + " ".join(f"{name}={t - prev:.3f}" for (_, prev), (name, t) in zip(marks, marks[1:])) + + (f" [{detail}]" if detail else ""), flush=True) def _reward(args: Any, sample: Any) -> float: @@ -357,7 +394,8 @@ def batch_summary_enabled(args: Any) -> bool: def build_metadata( - args: Any, all_samples: Iterable[Sequence[Any]], sink: str | None = None + args: Any, all_samples: Iterable[Sequence[Any]], sink: str | None = None, *, + defer_estimate: bool = False, ) -> dict[str, Any]: trained = getattr(args, _TRAINED_ATTR, None) if trained is None: @@ -379,8 +417,8 @@ def build_metadata( if not samples: continue if rollout_id is None: - rollout_id = current_round_id(samples, sink) - record = group_record(args, group) + rollout_id = _timed("round_id", current_round_id, samples, sink) + record = _timed("group_record", group_record, args, group) tool_wait += sum(float(getattr(x, "non_generation_time", 0.0) or 0.0) for x in samples) aborted += int(record.pop("aborted")) key = tuple(record.pop("_key")) @@ -418,7 +456,9 @@ def build_metadata( if limit > 0: # agentic-rollout-utilization 4.1: carry-over accounting (limit > 0 only; # the default key set is unchanged). - payload.update(carry_over_fields(args, trained_samples, sink)) + payload.update(_timed("carry_over", carry_over_fields, args, trained_samples, sink, + defer_estimate_for=rollout_id if defer_estimate and rollout_id is not None + else None)) # 5.1/5.3: agentic suspension between turns (absent when not in use) payload.update(_carry.suspend_fields(args, all_samples)) if tool_wait > 0: @@ -453,10 +493,17 @@ def build_metadata( def carry_over_fields(args: Any, trained_samples: Sequence[Any], - sink: str | None = None, scorer: Any = None) -> dict[str, Any]: + sink: str | None = None, scorer: Any = None, *, + defer_estimate_for: int | None = None) -> dict[str, Any]: """4.1 rollout-metadata fields under a limit > 0: what the buffer filter carried in / discarded this round, the trained cross-version tokens and the estimated - cross-version IS truncated fraction (:mod:`.carry_over`).""" + cross-version IS truncated fraction (:mod:`.carry_over`). + + ``defer_estimate_for`` (E1, a rollout id): the estimate (one prefill of every + cross-version sample, about 1 s on A10G) runs in a background thread and goes + to the sink as a :data:`LATE_META_SCHEMA` record; the fields here then say + ``cross_version_estimate: "deferred"``. Its only reader, the policy-age + governor, runs after the train step.""" token = current_policy_token(sink) version = token_version(token) if token else None fields: dict[str, Any] = {"max_policy_age": _carry.max_policy_age(args)} @@ -467,12 +514,31 @@ def carry_over_fields(args: Any, trained_samples: Sequence[Any], fields["cross_version_tokens"] = crossed fields["trained_response_tokens"] = total if crossed: - fields.update(_carry.estimate_cross_version_truncation( - args, trained_samples, version, - scorer if scorer is not None else _carry.router_prefill_scorer(args))) + scorer = scorer if scorer is not None else _carry.router_prefill_scorer(args) + if defer_estimate_for is None: + fields.update(_carry.estimate_cross_version_truncation(args, trained_samples, version, scorer)) + else: + _estimate_in_background(args, list(trained_samples), version, scorer, sink, + rollout_id=int(defer_estimate_for)) + fields["cross_version_estimate"] = "deferred" return fields +def _estimate_in_background(args: Any, samples: list[Any], version: int, scorer: Any, + sink: str | None, *, rollout_id: int) -> None: + import threading + + def run() -> None: + try: + late = _carry.estimate_cross_version_truncation(args, samples, version, scorer) + except Exception as exc: # noqa: BLE001 - an estimate: unknown, never guessed + late = {"cross_version_truncated_fraction": None, + "cross_version_estimate_error": f"{type(exc).__name__}: {str(exc)[:200]}"} + put_to_sink({"schema": LATE_META_SCHEMA, "rollout_id": rollout_id, **late}, sink) + + threading.Thread(target=run, name=f"yeto-xv-estimate-{rollout_id}", daemon=True).start() + + # IR-3/IR-4 sample-metadata keys written by agentic generate code (codex-harness # codex_openenv_generate): summed per rollout into the metadata payload. # Neutral names and values live in the core port (decoupling 4.11); re-exported. @@ -751,9 +817,7 @@ def put_to_sink(payload: dict[str, Any], sink: str | None = None) -> None: os.replace(tmp, directory / name) return if kind == "ray": - import ray - - ray.get(ray.get_actor(target or DEFAULT_SINK_ACTOR).put.remote(encoded)) + _sink_call(target, "put", encoded) return raise ValueError(f"unknown rollout metadata sink {sink!r}") @@ -765,9 +829,50 @@ def _sink_target(sink: str | None) -> tuple[str, str]: return kind, target +# E1: inside one hook run (:func:`one_sink_view`) each sink value is read once and +# the sink actor is looked up once (ray.get_actor is a GCS round trip per call) +_SINK_VIEW = threading.local() + + +def _sink_call(target: str, actor_method: str, *call_args: Any) -> Any: + import ray + + name = target or DEFAULT_SINK_ACTOR + memo = getattr(_SINK_VIEW, "memo", None) + actor = None if memo is None else memo.get(("actor", name)) + if actor is None: + actor = ray.get_actor(name) + if memo is not None: + memo[("actor", name)] = actor + return ray.get(getattr(actor, actor_method).remote(*call_args)) + + +@contextmanager +def one_sink_view(): + """Within this block every sink value (policy token, published map, policy-age + limit) is read at most once per (sink, value) and reused (E1: the metadata + hook read the limit from the sink actor once per sample, ~70 actor calls a + round). A value another process writes while a block runs is seen from the + next block on (before, two reads inside one hook run could see two values; + now one run sees one view). A write from this thread inside the block drops + that value from the view. The sink actor handle is looked up once per block. + Nested blocks share the outer view.""" + if getattr(_SINK_VIEW, "memo", None) is not None: + yield + return + _SINK_VIEW.memo = {} + try: + yield + finally: + _SINK_VIEW.memo = None + + def _sink_write(sink: str | None, file_name: str, text: str, actor_method: str, value: Any) -> None: """Driver side: store one value (atomic file under ``dir:``, actor call under ``ray:``).""" kind, target = _sink_target(sink) + memo = getattr(_SINK_VIEW, "memo", None) + if memo is not None: + memo.pop((kind, target, file_name), None) if kind == "dir": directory = Path(target) directory.mkdir(parents=True, exist_ok=True) @@ -776,21 +881,24 @@ def _sink_write(sink: str | None, file_name: str, text: str, actor_method: str, handle.write(text) os.replace(tmp, directory / file_name) return - import ray - - actor = ray.get_actor(target or DEFAULT_SINK_ACTOR) - ray.get(getattr(actor, actor_method).remote(value)) + _sink_call(target, actor_method, value) def _sink_read(sink: str | None, file_name: str, actor_method: str) -> Any: """Rollout side: the stored value (file text under ``dir:``), None when unset.""" kind, target = _sink_target(sink) + memo = getattr(_SINK_VIEW, "memo", None) + key = (kind, target, file_name) + if memo is not None and key in memo: + return memo[key] if kind == "dir": path = Path(target) / file_name - return path.read_text(encoding="utf-8") if path.exists() else None - import ray - - return ray.get(getattr(ray.get_actor(target or DEFAULT_SINK_ACTOR), actor_method).remote()) + value = path.read_text(encoding="utf-8") if path.exists() else None + else: + value = _sink_call(target, actor_method) + if memo is not None: + memo[key] = value + return value def put_policy_token(token: str, sink: str | None = None) -> None: @@ -870,7 +978,11 @@ def policy_buffer_filter(args: Any, _rollout_id: Any, buffer: list, num_samples: Groups from any other policy (or incomplete ones) are dropped from the buffer, exactly as legacy filtered its completed-group queue. """ + with one_sink_view(): + return _policy_buffer_filter(args, _rollout_id, buffer, num_samples) + +def _policy_buffer_filter(args: Any, _rollout_id: Any, buffer: list, num_samples: int) -> list: token = current_policy_token() limit = _carry.max_policy_age(args) if limit > 0: @@ -975,15 +1087,24 @@ def extract_rollout_metadata( """``--rollout-all-samples-process-path`` hook (under fully_async: called by :func:`record_trained_groups` with the trained groups).""" + with one_sink_view(): + _extract_rollout_metadata(args, all_samples, data_source, round_counts=round_counts) + + +def _extract_rollout_metadata(args: Any, all_samples: Any, data_source: Any, *, + round_counts: Mapping[str, Any] | None) -> None: try: - payload = build_metadata(args, all_samples) + # fully_async (round_counts): the trainer is waiting for this batch; the + # cross-version estimate is read only after its train step (E1) + payload = _timed("build", build_metadata, args, all_samples, + defer_estimate=round_counts is not None) if elastic_metadata_enabled(args): cursor, buffer_length = data_cursor(data_source) if cursor is not None: payload["data_cursor"] = cursor if buffer_length is not None: payload["buffer_length"] = buffer_length - submitted = submitted_groups(args, data_source) + submitted = _timed("submitted", submitted_groups, args, data_source) if round_counts is not None: from yeto.rl.engine.overlap import ROLLOUT_OVERLAP_MODE @@ -1026,7 +1147,7 @@ def extract_rollout_metadata( payload["over_sampling"] = over_sampling if hasattr(args, "rollout_over_sampling_stats"): args.rollout_over_sampling_stats = None - put_to_sink(payload) + _timed("put", put_to_sink, payload) finally: # Reset per-rollout state: the bounded filter keys its memo on # ``yeto_rl_policy_version`` which legacy advanced per round; here the diff --git a/yeto/rl/adapters/miles/run_flags.py b/yeto/rl/adapters/miles/run_flags.py index 8a9f52027..6e4a5af0b 100644 --- a/yeto/rl/adapters/miles/run_flags.py +++ b/yeto/rl/adapters/miles/run_flags.py @@ -175,6 +175,12 @@ "at most this many tokens per GPU (must be >= --seq-len). Default: fixed " "micro batches", ), + shared_flag( + "--rl-gradient-accumulation-fusion", + action="store_true", + help="ports: Megatron gradient-accumulation fusion (the Miles recipes' " + "default). Off by default (--no-gradient-accumulation-fusion)", + ), shared_flag( "--rl-optimizer-offload-fraction", type=float, diff --git a/yeto/rl/adapters/miles/state.py b/yeto/rl/adapters/miles/state.py index 175296cf6..fcc1b7b47 100644 --- a/yeto/rl/adapters/miles/state.py +++ b/yeto/rl/adapters/miles/state.py @@ -20,6 +20,7 @@ from . import LoopRunner from .state_plugin import ( APPLY_STATE, + COLLECT_DIGEST, EXPORT_DIGEST, EXPORT_FULL_DIGEST, EXPORT_STATE, @@ -60,6 +61,10 @@ def require_run_plugin(controller_cls: Any = None) -> None: class MilesPolicyState: """Export/apply the LoRA trainable state of the single-cell actor group.""" + # E1: export_digest(deferred=True) is supported (LoRA digest path) + supports_deferred_digest = True + supports_unverified_rounds = True # E1: export_identity + def __init__( self, *, @@ -117,25 +122,59 @@ async def aexport(self, *, policy_version: int | None = None) -> TrainableState: results = await self._aplugin(EXPORT_STATE, {"policy_version": version}) return self._export_result(version, results) - def export_digest(self, *, policy_version: int | None = None): + def export_digest(self, *, policy_version: int | None = None, deferred: bool = False): """rl-publish-fastpath: the trainer exports and hashes; only the digests cross Ray. Returns a :class:`TrainerResidentState` whose tensors stay in the trainer - (``materialize`` falls back to :meth:`export`).""" + (``materialize`` falls back to :meth:`export`). + + ``deferred`` (E1): the trainer returns once the tensors are staged and hashes + them in the background (e.g. while it offloads); the state's digest is a + :class:`~yeto.rl.engine.policy_digest.DeferredDigest` collected on first read.""" import time version = self._export_version(policy_version) started = time.monotonic() - results = self._plugin(EXPORT_DIGEST, self._digest_kwargs(version)) + kwargs = {**self._digest_kwargs(version), **({"deferred": True} if deferred else {})} + results = self._plugin(EXPORT_DIGEST, kwargs) state = self._digest_result(version, results) export_s = next((r.get("export_seconds") for r in results if r is not None), None) # Timing for the next GPU run (log only; no tape field changes): total plugin # round trip vs the trainer-side hashing part. + hashed = "deferred" if deferred else f"{state.digest.hash_seconds:.1f}s" print(f"[rl] publish-fastpath digest export v{version}: " f"total={time.monotonic() - started:.1f}s trainer_export={export_s}s " - f"hash={state.digest.hash_seconds:.1f}s bytes={state.digest.payload_bytes}", flush=True) + f"hash={hashed} tensors={len(state.digest.specs)}", flush=True) return state + def export_identity(self, *, policy_version: int | None = None): + """E1: the resident state of a round whose publication is not verified: the + trainer reports only its weights version (no export, no hash); the identity is + :func:`~yeto.rl.engine.policy_digest.unverified_identity`. The layout is the last + verified digest's (None before any: the caller verifies that round).""" + from yeto.rl.engine.policy_digest import PolicyDigest, TrainerResidentState, unverified_identity + + last = getattr(self, "_last_verified", None) + if last is None: + return None + version = self._export_version(policy_version) + mark = self.weights_version() + identity = unverified_identity(version, mark) + digest = PolicyDigest(identity, identity, last.payload_bytes, last.specs) + return TrainerResidentState( + layout=self.layout, base_model_revision=self._revision, config_hash=self._config_hash, + layout_hash=self._expected_layout_hash, policy_version=version, digest=digest, + materialize_fn=lambda v: self.export(policy_version=v), weights_mark=mark, + holds_fn=self.check_holds, verified=False) + + def _collect_digest(self, version: int): + from yeto.rl.engine.policy_digest import PolicyDigest + + results = [r for r in self._plugin(COLLECT_DIGEST, {"policy_version": version}) if r is not None] + if len(results) != 1 or results[0].get("policy_version") != version: + raise PolicyStateError(f"expected exactly one main-rank digest of v{version}, got {results!r:.300}") + return PolicyDigest.from_wire(results[0]["digest"]) + def weights_version(self) -> dict[str, Any]: """The main trainer rank's weights version ``{process_id, version}`` (D3).""" @@ -194,7 +233,12 @@ def _digest_result(self, version: int, results: list[Any]): (result,) = results if result.get("policy_version") != version: raise PolicyStateError("exported policy version mismatch") - digest = PolicyDigest.from_wire(result["digest"]) + if result.get("deferred"): + from yeto.rl.engine.policy_digest import DeferredDigest + + digest = DeferredDigest(result["specs"], lambda: self._collect_digest(version)) + else: + digest = PolicyDigest.from_wire(result["digest"]) layout_hash = layout_hash_of(digest.specs) if self._expected_layout_hash is not None and layout_hash != self._expected_layout_hash: raise ValueError( @@ -204,6 +248,7 @@ def _digest_result(self, version: int, results: list[Any]): self._expected_layout_hash = layout_hash print(f"[rl] learned LoRA layout from first digest export: {len(digest.specs)} tensors, " f"hash={layout_hash}", flush=True) + self._last_verified = digest return TrainerResidentState( layout=self.layout, base_model_revision=self._revision, @@ -304,6 +349,9 @@ class MilesFullParameterLocalPolicyState(MilesPolicyState): apply. The layout hash is the full-parameter identity bound by run_miles. """ + supports_deferred_digest = False # its export_digest hashes every rank's shard eagerly + supports_unverified_rounds = False # E1: every publication hashes the shards + def __init__(self, *, actor_model: Any, base_model_revision: str, config_hash: str, expected_layout_hash: str | None, policy_version: int = 0, runner: LoopRunner | None = None) -> None: diff --git a/yeto/rl/adapters/miles/state_plugin.py b/yeto/rl/adapters/miles/state_plugin.py index 5125287ec..3e0684975 100644 --- a/yeto/rl/adapters/miles/state_plugin.py +++ b/yeto/rl/adapters/miles/state_plugin.py @@ -39,6 +39,7 @@ EXPORT_STATE = f"{_PLUGIN_MODULE}.export_state" # rl-publish-fastpath: export + hash on the trainer, return only the digests. EXPORT_DIGEST = f"{_PLUGIN_MODULE}.export_digest" +COLLECT_DIGEST = f"{_PLUGIN_MODULE}.collect_digest" WEIGHTS_VERSION = f"{_PLUGIN_MODULE}.weights_version" # A no-sync full-parameter island (L4.6): per-rank digest of every model parameter. EXPORT_FULL_DIGEST = f"{_PLUGIN_MODULE}.export_full_digest" @@ -648,33 +649,92 @@ def export_state(actor: Any, *, policy_version: int) -> dict[str, Any] | None: return _export_state(actor, policy_version=policy_version) +def _digest_staging(actor: Any) -> Any: + """The pinned host buffer this trainer's digests reuse every round (E1). A + deferred digest still hashing it is waited for first.""" + from yeto.rl.engine.policy_digest import HostStaging + + _wait_deferred_digest(actor) + staging = getattr(actor, "_yeto_digest_staging", None) + if staging is None: + staging = actor._yeto_digest_staging = HostStaging() + return staging + + def export_digest( actor: Any, *, policy_version: int, base_model_revision: str, lora_config_hash: str, + deferred: bool = False, ) -> dict[str, Any] | None: """rl-publish-fastpath: :func:`export_state` on the trainer, but return only the publication digests (policy tensor hash, payload hash/bytes, specs) instead of the tensors. The hash definitions are those of ``yeto.rl.core.policy_tensor_hash`` - and ``publish.payload_digest`` (see ``policy_digest``); only where they run moves.""" + and ``publish.payload_digest`` (see ``policy_digest``); only where they run moves. + + ``deferred`` (E1, sync colocated mode): return once the tensors are staged on the + host (layout only, ``"deferred": True``) and hash them in a thread of this + process, so the trainer can offload meanwhile; :func:`collect_digest` returns + the digest (the same values).""" from yeto.rl.engine.policy_digest import digest_canonical_tensors, layout_hash_of import time started = time.monotonic() - exported = export_state(actor, policy_version=policy_version) + install_grad_norm_recorder() + with trainer_resident(actor): + # E1: the tensors are only hashed here, so they may be staged in one pinned buffer + exported = _export_state(actor, policy_version=policy_version, staging=_digest_staging(actor)) export_seconds = round(time.monotonic() - started, 3) if exported is None: return None tensors = exported["tensors"] specs = [(name, tuple(int(d) for d in value.shape)) for name, value in sorted(tensors.items())] - digest = digest_canonical_tensors( - tensors, - base_model_revision=base_model_revision, - lora_config_hash=lora_config_hash, - layout_hash=layout_hash_of(specs), - ) - return {"policy_version": int(policy_version), "digest": digest.to_wire(), - "export_seconds": export_seconds, "weights": current_weights_version()} + + def digest() -> Any: + return digest_canonical_tensors( + tensors, + base_model_revision=base_model_revision, + lora_config_hash=lora_config_hash, + layout_hash=layout_hash_of(specs), + ) + + out = {"policy_version": int(policy_version), "export_seconds": export_seconds, + "weights": current_weights_version()} + if not deferred: + return {**out, "digest": digest().to_wire()} + from concurrent.futures import ThreadPoolExecutor + + pool = ThreadPoolExecutor(max_workers=1, thread_name_prefix="yeto-digest-deferred") + actor._yeto_deferred_digest = (int(policy_version), pool.submit(digest)) + pool.shutdown(wait=False) + return {**out, "deferred": True, "specs": [[n, list(shape)] for n, shape in specs]} + + +def collect_digest(actor: Any, *, policy_version: int) -> dict[str, Any] | None: + """The digest of the deferred :func:`export_digest` of ``policy_version`` (waits + for its hashing thread). Reads no trainer state (the trainer may be offloaded); + ``None`` on ranks that export nothing.""" + pending = getattr(actor, "_yeto_deferred_digest", None) + if pending is None: + if not _is_main_rank(actor): + return None + raise StatePluginError(f"no deferred digest of policy v{policy_version} on this rank") + version, future = pending + actor._yeto_deferred_digest = None + if version != int(policy_version): + raise StatePluginError(f"the deferred digest is of policy v{version}, not v{policy_version}") + return {"policy_version": version, "digest": future.result().to_wire()} + + +def _wait_deferred_digest(actor: Any) -> None: + """The staged tensors of a deferred digest are views of the reused buffer: let an + uncollected hashing finish before the buffer is written again.""" + pending = getattr(actor, "_yeto_deferred_digest", None) + if pending is not None: + try: + pending[1].result() + except Exception: # noqa: BLE001 - collect_digest reports it; only the buffer matters here + pass FULL_SHARD_PREFIX = b"yeto-rl-full-parameter-shard-v1\0" @@ -912,10 +972,15 @@ def _tp_ep_leader(actor: Any) -> bool: ) -def _export_state(actor: Any, *, policy_version: int) -> dict[str, Any] | None: +def _export_state(actor: Any, *, policy_version: int, staging: Any = None) -> dict[str, Any] | None: + """``staging`` (a :class:`~yeto.rl.engine.policy_digest.HostStaging`): copy the + tensors into its pinned buffer with one GPU synchronization; the result is then + only valid until the next call with it (digest path, E1). Same values either way.""" import torch + from yeto.rl.engine.policy_digest import NonFiniteTensor, stage_f32_on_host + if is_native_flash_next(actor): return _export_flash_next(actor, policy_version=policy_version) bindings = adapter_bindings(actor) @@ -931,12 +996,10 @@ def _export_state(actor: Any, *, policy_version: int) -> dict[str, Any] | None: tensors[b.name] = b.to_hf(master) if not _is_main_rank(actor): return None - out = {} - for name, value in sorted(tensors.items()): - value = value.detach().to(device="cpu", dtype=torch.float32).contiguous().clone() - if not torch.isfinite(value).all().item(): - raise StatePluginError(f"{name!r} contains NaN or Inf") - out[name] = value + try: + out = stage_f32_on_host(tensors, staging) + except NonFiniteTensor as error: + raise StatePluginError(str(error)) from None assert_grad_flow_intact(actor.model, params) return {"policy_version": int(policy_version), "tensors": out} @@ -1189,7 +1252,12 @@ def train_one_step(*args: Any, **kwargs: Any): _EV_STATS.clear() _POLICY_STATS.clear() try: - result = original(*args, **kwargs) + if getattr(args[0] if args else kwargs.get("args"), "yeto_rl_profile_trainer", False): + from .train_profile import profiled # E1 observation (--rl-profile-trainer) + + result = profiled(original, *args, **kwargs) + else: + result = original(*args, **kwargs) finally: bump_weights_version() # rl-publish-fastpath: any step may have written weights try: diff --git a/yeto/rl/adapters/miles/train_profile.py b/yeto/rl/adapters/miles/train_profile.py new file mode 100644 index 000000000..bfa12cc3f --- /dev/null +++ b/yeto/rl/adapters/miles/train_profile.py @@ -0,0 +1,121 @@ +"""E1 observation: where the trainer process spends one training step. + +Opt-in (``--rl-profile-trainer``): around each ``train_one_step`` call it records the +wall time, the main thread's CPU time, the whole process's CPU time, garbage +collector pauses, CUDA caching-allocator counters, and stack samples of every +Python thread (every ``SAMPLE_S``), and prints one ``yeto train profile {json}`` +line. It changes no value the step computes. +""" + +from __future__ import annotations + +import collections +import gc +import json +import os +import sys +import threading +import time +from typing import Any + +SAMPLE_S = 0.005 +TOP = 12 +_GC = {"pause_s": 0.0, "collections": collections.Counter(), "start": None} +_GC_INSTALLED = False + + +def _gc_callback(phase: str, info: dict) -> None: + if phase == "start": + _GC["start"] = time.perf_counter() + elif _GC["start"] is not None: + _GC["pause_s"] += time.perf_counter() - _GC["start"] + _GC["collections"][int(info.get("generation", -1))] += 1 + _GC["start"] = None + + +def _frame_label(frame: Any) -> str: + """The innermost frame outside the standard library and threading machinery.""" + first = None + while frame is not None: + code = frame.f_code + name = code.co_filename + label = f"{os.path.basename(name)}:{code.co_name}" + first = first or label + if "/lib/python3" not in name or "site-packages" in name: + return label + frame = frame.f_back + return first or "?" + + +class _Sampler(threading.Thread): + def __init__(self, main_id: int) -> None: + super().__init__(daemon=True, name="yeto-train-profile") + self.main_id = main_id + self.stop = threading.Event() + self.main = collections.Counter() + self.others = collections.Counter() + self.samples = 0 + + def run(self) -> None: + own = threading.get_ident() + names = {t.ident: t.name for t in threading.enumerate()} + while not self.stop.wait(SAMPLE_S): + self.samples += 1 + for ident, frame in sys._current_frames().items(): + if ident == own: + continue + label = _frame_label(frame) + if ident == self.main_id: + self.main[label] += 1 + elif not label.startswith(("threading.py:wait", "queue.py:get", "selectors.py")): + self.others[f"{names.get(ident, ident)}|{label}"] += 1 + + +def _cuda_counters() -> dict[str, int]: + try: + import torch + + if not torch.cuda.is_available(): + return {} + stats = torch.cuda.memory_stats() + except Exception: # noqa: BLE001 - observation only + return {} + keys = ("num_alloc_retries", "num_device_alloc", "num_device_free", "num_sync_all_streams") + out = {k: int(stats.get(k, 0)) for k in keys} + out["reserved_bytes"] = int(stats.get("reserved_bytes.all.current", 0)) + out["allocated_bytes"] = int(stats.get("allocated_bytes.all.current", 0)) + return out + + +def profiled(call: Any, *args: Any, **kwargs: Any) -> Any: + """``call(*args, **kwargs)`` with the profile line printed after it.""" + global _GC_INSTALLED + if not _GC_INSTALLED: + gc.callbacks.append(_gc_callback) + _GC_INSTALLED = True + sampler = _Sampler(threading.get_ident()) + gc_before, gens_before = _GC["pause_s"], collections.Counter(_GC["collections"]) + cuda_before = _cuda_counters() + wall, thread_cpu, process_cpu = time.perf_counter(), time.thread_time(), time.process_time() + sampler.start() + try: + return call(*args, **kwargs) + finally: + sampler.stop.set() + sampler.join() + cuda_after = _cuda_counters() + record = { + "wall_s": round(time.perf_counter() - wall, 4), + "main_thread_cpu_s": round(time.thread_time() - thread_cpu, 4), + "process_cpu_s": round(time.process_time() - process_cpu, 4), + "gc_pause_s": round(_GC["pause_s"] - gc_before, 4), + "gc_collections": dict(_GC["collections"] - gens_before), + "gc_objects": len(gc.get_objects()) if os.environ.get("YETO_RL_PROFILE_GC_OBJECTS") else None, + "cuda": {k: cuda_after[k] - cuda_before.get(k, 0) if k.startswith("num_") else cuda_after[k] + for k in cuda_after}, + "threads": threading.active_count(), + "samples": sampler.samples, + "main_top": sampler.main.most_common(TOP), + "others_top": sampler.others.most_common(TOP), + } + print("yeto train profile " + json.dumps(record), flush=True) diff --git a/yeto/rl/adapters/miles/trainer.py b/yeto/rl/adapters/miles/trainer.py index 54c65ed90..ab3602711 100644 --- a/yeto/rl/adapters/miles/trainer.py +++ b/yeto/rl/adapters/miles/trainer.py @@ -813,7 +813,15 @@ def save_cut(self, *, epoch: int, context: "CutContext") -> str: a non-zero ``carried_over``/ready-unconsumed ledger, an unsettled outer commit or a missing data cursor. """ - from yeto.rl.engine.cut import CutError, CutFile, CutManifest, commit_manifest, context_problems, cut_dir + return self.finish_cut(self.start_cut(epoch=epoch, context=context, background=False)) + + def start_cut(self, *, epoch: int, context: "CutContext", background: bool = True) -> dict[str, Any]: + """E1: the part of :meth:`save_cut` that must run at the step boundary: the + context checks and every rank's state capture. With ``background`` the ranks + write and hash their shards in a thread while training goes on; :meth:`finish_cut` + (later, before the next cut) collects them and commits the manifest. Nothing + names the cut until the manifest is committed. Returns the pending cut.""" + from yeto.rl.engine.cut import CutError, context_problems, cut_dir from .cut_plugin import SAVE_CUT_SHARD, config_problems cut_id = context.cut_id @@ -828,11 +836,42 @@ def save_cut(self, *, epoch: int, context: "CutContext") -> str: if problems: raise CutError("refusing to save a cut: " + "; ".join(problems)) directory = cut_dir(context.root, cut_id) - summaries = [ - dict(s) for s in self._run( - self._actor.run_plugin(SAVE_CUT_SHARD, {"directory": str(directory), "cut_id": cut_id}) - ) - ] + # the critic cut runs in finish_cut and reads the critic as it is then: inline only + background = bool(background) and self._critic is None + critic = self._save_critic_cut(directory / "critic", context.progress.policy_version) + kwargs = {"directory": str(directory), "cut_id": cut_id, **({"background": True} if background else {})} + results = [dict(s) for s in self._run(self._actor.run_plugin(SAVE_CUT_SHARD, kwargs))] + # E1: how each rank took its state (device copies or host path, and why) + capture = [dict(s["capture"]) for s in results if isinstance(s.get("capture"), Mapping)] + return {"epoch": int(epoch), "context": context, "runtime": runtime, "directory": directory, + "background": background, "results": results, "critic": critic, + **({"capture": capture} if capture else {})} + + def finish_cut(self, pending: Mapping[str, Any]) -> str: + """Collect the shards of :meth:`start_cut` and commit the manifest; returns the cut id.""" + return self.collect_cut(pending)() + + def collect_cut(self, pending: Mapping[str, Any]) -> Callable[[], str]: + """E1: the trainer part of :meth:`finish_cut` (collect and check every rank's + shard summary, build the manifest). Returns the commit of that manifest: plain + filesystem work (no trainer call), so the caller may run it on another thread. + Raises like :meth:`finish_cut` before anything is committed.""" + from yeto.rl.engine.cut import CutError, CutFile, CutManifest, commit_manifest + from .cut_plugin import FINISH_CUT_SHARD + + epoch, context, runtime = pending["epoch"], pending["context"], pending["runtime"] + cut_id = context.cut_id + summaries = list(pending["results"]) + if pending["background"]: + started = [s for s in summaries if "refused" not in s] + summaries = [s for s in summaries if "refused" in s] + if started: + summaries += [dict(s) for s in self._run( + self._actor.run_plugin(FINISH_CUT_SHARD, {"cut_id": cut_id}))] + # E1: where each rank's background shard was written (writer process, or this + # thread after a writer-process failure); reported with rl_cut_committed + self.last_cut_writer = [dict(s["writer"]) for s in summaries if isinstance(s.get("writer"), Mapping)] + problems: list[str] = [] refused = [f"[{s.get('refusal_kind', 'refused')}] {s['refused']}" for s in summaries if "refused" in s] if refused: raise CutError("rank refused the cut: " + "; ".join(sorted(set(refused)))) @@ -843,7 +882,7 @@ def save_cut(self, *, epoch: int, context: "CutContext") -> str: if problems: raise CutError("refusing an incomplete cut: " + "; ".join(problems)) resume = self._resume_manifest(context, summaries) - critic = self._save_critic_cut(directory / "critic", context.progress.policy_version) + critic = pending.get("critic") if critic is not None: # 4.3; absent for every critic-free trainer (manifest unchanged) runtime = {**runtime, "critic": critic} manifest = CutManifest( @@ -869,8 +908,13 @@ def save_cut(self, *, epoch: int, context: "CutContext") -> str: ) if problems: raise CutError("refusing an incomplete cut: " + "; ".join(problems)) - commit_manifest(context.root, manifest, check_files=context.shared_filesystem) - return cut_id + root, check_files = context.root, context.shared_filesystem + + def commit() -> str: + commit_manifest(root, manifest, check_files=check_files) + return cut_id + + return commit def _resume_manifest(self, context: "CutContext", summaries: list[Mapping[str, Any]]) -> Any: """Resume manifest of this cut: run counters plus every rank's optimizer state (S25 N7.6).""" diff --git a/yeto/rl/adapters/verl/fully_async_ports.py b/yeto/rl/adapters/verl/fully_async_ports.py index 061fb920c..f79505f91 100644 --- a/yeto/rl/adapters/verl/fully_async_ports.py +++ b/yeto/rl/adapters/verl/fully_async_ports.py @@ -15,8 +15,8 @@ * ``train_step`` -- old log-prob -> advantage -> actor update on the actor (verl's ``fit_step`` without ``_fit_update_weights``); mismatch criteria in the Miles convention on the trainer's DataProto. -* policy state -- ``ports_impl.VerlPolicyState`` unchanged: its two island hooks - (``export_tensors`` / ``call_workers``) run on the trainer actor's worker group. +* policy state -- ``ports_impl.VerlPolicyState`` unchanged: its island hooks + (``export_tensors`` / ``export_digest`` / ``call_workers``) run on the trainer actor's worker group. * ``publish`` -- the driver has applied the state; the actor pushes it through ``checkpoint_manager.update_weights(global_steps=param_version)``, resets the rollouter's staleness counter, and the pairing @@ -41,6 +41,7 @@ from .fully_async_round import (DISCARD_MECHANISM, FULLY_ASYNC_PLACEMENT, RoundCollector, SampleMeta, group_metadata) from .fully_async_translate import VersionMap +from .ports_impl import sent_checksum from .rollout_events import cutoff_fields MEMBER = "vllm-rollouter-0" @@ -84,6 +85,10 @@ def __init__(self, call: Callable[..., Any], *, learner_id: int, limit: int, req def export_tensors(self) -> tuple[dict, dict]: return self.call("yeto_export_lora") + def export_digest(self, kwargs: dict) -> dict: + """rl-publish-fastpath: digests + send checksum computed on the trainer actor.""" + return self.call("yeto_export_digest", kwargs) + def call_workers(self, fn, *args): return self.call("yeto_call_workers", fn, *args) @@ -251,8 +256,7 @@ def publish(self, state) -> PublicationResult: isl = self.island version = int(state.policy_version) - lora = state.to_lora() - sent = pub.lora_checksum(dict(lora.tensors)) + sent, payload_bytes = sent_checksum(state) (isl.readback_dir / "expected_version").write_text(str(version)) stale = isl.readback_dir / f"readback-v{version}.json" if stale.exists(): @@ -288,7 +292,6 @@ def publish(self, state) -> PublicationResult: digest = state.policy_tensor_hash() isl.published = (version, digest) isl.hashes[version] = digest - payload_bytes = sum(int(t.numel()) * 2 for t in lora.tensors.values()) manifest = InferencePublicationManifest( publication_mode="full", base_policy_version=None, target_policy_version=version, target_policy_hash=digest, target_manifest_hash=_sha(f"{version}:{digest}"), diff --git a/yeto/rl/adapters/verl/fully_async_runner.py b/yeto/rl/adapters/verl/fully_async_runner.py index 900ccf6d3..b3273048b 100644 --- a/yeto/rl/adapters/verl/fully_async_runner.py +++ b/yeto/rl/adapters/verl/fully_async_runner.py @@ -189,6 +189,11 @@ async def yeto_export_lora(self): return (await self.yeto_call_workers(remote_export_lora))[0] + async def yeto_export_digest(self, kwargs): + from .ports_impl import remote_export_digest + + return next(r for r in await self.yeto_call_workers(remote_export_digest, kwargs) if r is not None) + YetoFullyAsyncTrainer = ray.remote(num_cpus=10)(type("YetoFullyAsyncTrainer", (_YetoTrainerMixin, _TrainerBase), {})) diff --git a/yeto/rl/adapters/verl/image.py b/yeto/rl/adapters/verl/image.py index b096d1b18..29a522824 100644 --- a/yeto/rl/adapters/verl/image.py +++ b/yeto/rl/adapters/verl/image.py @@ -46,15 +46,18 @@ def image_env() -> dict[str, str]: } -def modal_image(modal, commit: str = VERL_COMMIT): - patch = Path(__file__).with_name("patch_verl.py") - return ( +def modal_image(modal, commit: str = VERL_COMMIT, patched: bool = True): + """The verl image; ``patched=False`` leaves the verl source as the fork ships it + (the bare-verl arms of the efficiency evaluation, which must not import yeto).""" + image = ( modal.Image.from_registry(VERL_BASE_IMAGE, add_python=VERL_PYTHON) .apt_install("git", "build-essential", "curl", "libnuma1") .pip_install(f"uv=={VERL_UV_VERSION}") .run_commands(*build_commands(commit)) .run_commands(*store_client_commands()) - .add_local_file(str(patch), "/workspace/patch_verl.py", copy=True) - .run_commands(f"{VENV_BIN}/python /workspace/patch_verl.py {VERL_ROOT}") - .env(image_env()) ) + if patched: + patch = Path(__file__).with_name("patch_verl.py") + image = (image.add_local_file(str(patch), "/workspace/patch_verl.py", copy=True) + .run_commands(f"{VENV_BIN}/python /workspace/patch_verl.py {VERL_ROOT}")) + return image.env(image_env()) diff --git a/yeto/rl/adapters/verl/ports_impl.py b/yeto/rl/adapters/verl/ports_impl.py index 4952a73c2..85f0e41c9 100644 --- a/yeto/rl/adapters/verl/ports_impl.py +++ b/yeto/rl/adapters/verl/ports_impl.py @@ -30,13 +30,17 @@ import json import os import time +from collections.abc import Mapping +from dataclasses import dataclass, field +from math import prod from pathlib import Path from typing import Any from yeto.rl.contracts import InferencePublicationManifest, LocalStepReceipt from yeto.rl.engine import mismatch_criteria, offpolicy_source +from yeto.rl.engine.policy_digest import PolicyDigest, TrainerResidentState, is_resident from yeto.rl.engine.ports import GroupMetadata, PlacementDescription, PublicationResult, RolloutBatchHandle -from yeto.rl.engine.trainable_state import TrainableState +from yeto.rl.engine.trainable_state import LAYOUT_LORA, TrainableState from . import publish as pub from .param_names import train_to_canonical @@ -59,22 +63,88 @@ def _peft_module(worker): return getattr(module, "_fsdp_wrapped_module", module) +def _lora_tensors(worker_dict) -> tuple[dict, set]: + """({canonical name: full LoRA tensor on its device}, parameter dtypes); collective + (``full_tensor``): every rank calls it.""" + peft = _peft_module(_inner_worker(worker_dict)) + out, dtypes = {}, set() + for name, param in peft.named_parameters(): + if ".lora_" not in name: + continue + dtypes.add(str(param.dtype)) + full = param.full_tensor() if hasattr(param, "full_tensor") else param + out[train_to_canonical(name)] = full.detach() + return out, dtypes + + def remote_export_lora(worker_dict): """(canonical f32 CPU tensors, info) from the actor's PEFT model.""" import torch - peft = _peft_module(_inner_worker(worker_dict)) - out, dtypes = {}, set() with torch.no_grad(): - for name, param in peft.named_parameters(): - if ".lora_" not in name: - continue - dtypes.add(str(param.dtype)) - full = param.full_tensor() if hasattr(param, "full_tensor") else param - out[train_to_canonical(name)] = full.detach().to("cpu", torch.float32).clone() + tensors, dtypes = _lora_tensors(worker_dict) + out = {name: value.to("cpu", torch.float32).clone() for name, value in tensors.items()} return out, {"param_dtypes": sorted(dtypes), "count": len(out)} +def digest_lora_tensors(tensors, *, base_model_revision: str, lora_config_hash: str, + layout_hash: str, expected_specs) -> dict: + """rl-publish-fastpath, trainer side: validate ``tensors`` exactly as the driver's + ``canonical_state`` would and return only what the driver reads -- the policy + digests and the bf16 send checksum of :func:`publish.lora_checksum` (small dicts, + no tensors). Same functions as the full-export path, so the values are equal + bit for bit.""" + + import time + + from concurrent.futures import ThreadPoolExecutor + + from yeto.rl.core import canonical_state_from_owned_tensors + from yeto.rl.engine.policy_digest import digest_canonical_tensors + + started = time.monotonic() + # E1: ``tensors`` are this call's own copies (the trainer's export), so the + # validation keeps them without a second clone; the values are the same + state = canonical_state_from_owned_tensors( + 0, tensors, base_model_revision=base_model_revision, lora_config_hash=lora_config_hash, + layout_hash=layout_hash, expected_specs=expected_specs) + # E1: the bf16 send checksum runs in its own thread beside the two digest streams + # (hashlib releases the GIL); each function is unchanged, so the values are too + with ThreadPoolExecutor(max_workers=1, thread_name_prefix="yeto-lora-checksum") as pool: + checksum = pool.submit(pub.lora_checksum, dict(state.tensors)) + digest = digest_canonical_tensors(state.tensors, base_model_revision=state.base_model_revision, + lora_config_hash=state.lora_config_hash, + layout_hash=state.layout_hash) + lora_checksum = checksum.result() + return {"digest": digest.to_wire(), "layout_hash": state.layout_hash, + "lora_checksum": lora_checksum, + "export_seconds": round(time.monotonic() - started, 3)} + + +def remote_export_digest(worker_dict, kwargs: dict): + """All ranks gather the LoRA (``full_tensor`` is collective); rank 0 hashes it. + + E1: rank 0 stages the tensors in one reusable pinned buffer (one GPU + synchronization, :func:`~yeto.rl.engine.policy_digest.stage_f32_on_host`); they + are hashed here and dropped, never returned.""" + import torch + import torch.distributed as dist + + from yeto.rl.engine.policy_digest import HostStaging, stage_f32_on_host + + with torch.no_grad(): + tensors, _dtypes = _lora_tensors(worker_dict) + if dist.is_available() and dist.is_initialized() and dist.get_rank() != 0: + return None + worker = _inner_worker(worker_dict) + staging = getattr(worker, "_yeto_digest_staging", None) + if staging is None: + staging = HostStaging() + setattr(worker, "_yeto_digest_staging", staging) + tensors = stage_f32_on_host(tensors, staging) + return digest_lora_tensors(tensors, **kwargs) + + def remote_apply_lora(worker_dict, tensors, reset_optimizer: bool): import torch @@ -161,6 +231,10 @@ def export_tensors(self) -> tuple[dict, dict]: results = self.call_workers(remote_export_lora) return results[0] + def export_digest(self, kwargs: dict) -> dict: + """rl-publish-fastpath: digests + send checksum computed on the workers.""" + return next(r for r in self.call_workers(remote_export_digest, kwargs) if r is not None) + def append_jsonl(self, name: str, record: dict) -> None: with (self.out / name).open("a") as handle: handle.write(json.dumps(record, sort_keys=True, default=str) + "\n") @@ -388,6 +462,25 @@ def round_metrics(self) -> dict: return out +@dataclass(frozen=True) +class VerlResidentState(TrainerResidentState): + """A verl policy kept in the trainer; carries the bf16 send checksum the + publisher compares with the vLLM read-back (no tensors on the driver).""" + + lora_checksum: Mapping[str, Any] | None = field(default=None, compare=False) + + +def sent_checksum(state) -> tuple[dict, int]: + """(send checksum, bf16 payload bytes) of the policy being published. A resident + state brings both from the trainer; any other state is hashed here as before.""" + + if is_resident(state) and getattr(state, "lora_checksum", None) is not None: + payload_bytes = sum(prod(shape) * 2 for _name, shape in state.digest.specs) + return dict(state.lora_checksum), int(payload_bytes) + tensors = dict(state.to_lora().tensors) + return pub.lora_checksum(tensors), sum(int(t.numel()) * 2 for t in tensors.values()) + + class VerlPolicyState: def __init__(self, island: VerlIsland): self.island = island @@ -403,6 +496,36 @@ def export(self) -> TrainableState: expected_specs=self.island.expected_specs) return TrainableState.from_lora(state) + def export_digest(self, *, policy_version: int) -> VerlResidentState | None: + """rl-publish-fastpath: the workers export and hash; only digests cross Ray. + None when the island has no digest hook (the driver then calls :meth:`export`).""" + + hook = getattr(self.island, "export_digest", None) + if not callable(hook): + return None + isl, started = self.island, time.monotonic() + raw = hook({"base_model_revision": isl.base_model_revision, + "lora_config_hash": isl.lora_config_hash, "layout_hash": isl.layout_hash, + "expected_specs": tuple(isl.expected_specs)}) + digest = PolicyDigest.from_wire(raw["digest"]) + print(f"[rl] verl publish-fastpath digest export v{policy_version}: " + f"total={time.monotonic() - started:.2f}s trainer={raw.get('export_seconds')}s " + f"hash={digest.hash_seconds:.2f}s bytes={digest.payload_bytes}", flush=True) + return VerlResidentState( + layout=LAYOUT_LORA, base_model_revision=isl.base_model_revision, + config_hash=isl.lora_config_hash, layout_hash=str(raw["layout_hash"]), + policy_version=int(policy_version), digest=digest, materialize_fn=self._materialize, + lora_checksum=raw["lora_checksum"]) + + def _materialize(self, policy_version: int) -> TrainableState: + from yeto.rl.core import canonical_state + + lora = self.export().to_lora() + return TrainableState.from_lora(canonical_state( + int(policy_version), lora.tensors, base_model_revision=lora.base_model_revision, + lora_config_hash=lora.lora_config_hash, layout_hash=lora.layout_hash, + expected_specs=lora.specs)) + def apply(self, state: TrainableState, *, optimizer: str, local_step: int) -> None: if optimizer not in ("reset", "preserve"): raise ValueError(optimizer) @@ -443,8 +566,7 @@ def publish(self, state: TrainableState) -> PublicationResult: isl, t = self.island, self.island.trainer version = int(state.policy_version) - lora = state.to_lora() - sent = pub.lora_checksum(dict(lora.tensors)) + sent, payload_bytes = sent_checksum(state) (isl.readback_dir / "expected_version").write_text(str(version)) stale = isl.readback_dir / f"readback-v{version}.json" if stale.exists(): @@ -470,7 +592,6 @@ def publish(self, state: TrainableState) -> PublicationResult: f"{list(check.mismatched[:3])}", cause=cause) digest = state.policy_tensor_hash() isl.published = (version, digest) - payload_bytes = sum(int(t_.numel()) * 2 for t_ in lora.tensors.values()) manifest = InferencePublicationManifest( publication_mode="full", base_policy_version=None, target_policy_version=version, target_policy_hash=digest, target_manifest_hash=_sha(f"{version}:{digest}"), diff --git a/yeto/rl/adapters/verl/round_cuts.py b/yeto/rl/adapters/verl/round_cuts.py index 735b282e2..aa7f4e341 100644 --- a/yeto/rl/adapters/verl/round_cuts.py +++ b/yeto/rl/adapters/verl/round_cuts.py @@ -31,6 +31,8 @@ from dataclasses import dataclass from typing import Any +from yeto.rl.engine.resume import DEFAULT_CUT_EVERY + RESUME_FLAGS = ("rl_cut_every", "rl_cut_keep", "rl_stop_after_rounds") DEFAULT_STATE_DIR = "~/yeto-rl/resume-state" # VerlRunConfig fields that may change between a cut and the launch resuming it: @@ -77,7 +79,7 @@ def resume_config(args: Any) -> dict[str, Any] | None: return { "store": str(store), "state_dir": str(getattr(args, "rl_resume_state_dir", None) or DEFAULT_STATE_DIR), - "every": int(getattr(args, "rl_cut_every", None) or 1), + "every": int(getattr(args, "rl_cut_every", None) or DEFAULT_CUT_EVERY), "keep": int(getattr(args, "rl_cut_keep", None) or 2), "allow_config_change": bool(getattr(args, "rl_resume_allow_config_change", False)), "stop_after": (None if getattr(args, "rl_stop_after_rounds", None) is None diff --git a/yeto/rl/adapters/verl/verl_main.py b/yeto/rl/adapters/verl/verl_main.py index 106a0093f..303ad9428 100644 --- a/yeto/rl/adapters/verl/verl_main.py +++ b/yeto/rl/adapters/verl/verl_main.py @@ -57,12 +57,21 @@ def run(self, config, plan: dict): tq.close() -def _runtime_env() -> dict: +# NCCL host buffers on /dev/shm in every process of the disaggregated weight sync. +# On Modal one process could use cuMem host buffers while its peer fell back to +# /dev/shm, and the peer then failed to attach (S26 E0.1 eff01-yeto-verl-async-e, +# eff01-bare-verl-async-d with NCCL_DEBUG=INFO). The caller's own value wins. +CUMEM_HOST_ENV = "NCCL_CUMEM_HOST_ENABLE" + + +def _runtime_env(fully_async: bool = False) -> dict: from verl.trainer.constants_ppo import get_ppo_ray_runtime_env env = get_ppo_ray_runtime_env(None) env.setdefault("env_vars", {}) env["env_vars"]["TRANSFER_QUEUE_ENABLE"] = "1" + if fully_async: + env["env_vars"][CUMEM_HOST_ENV] = os.environ.get(CUMEM_HOST_ENV, "0") for key, value in os.environ.items(): # YETO_ covers YETO_ISLAND_HMAC_KEY (HELLO HMAC in the Ray actor that # opens the syncer client) and YETO_SANDBOX_MODAL_TOKEN_*; no cloud @@ -90,7 +99,7 @@ def _keep_npu_visible_in_cpu_actors() -> None: os.environ.setdefault(NPU_NOSET_ENV, "1") -def _plan_and_ray() -> dict: +def _plan_and_ray(fully_async: bool = False) -> dict: from yeto.island_credential_guard import check_island_credentials check_island_credentials() # secret-handling-hardening D4 (driver, both paths) @@ -98,7 +107,7 @@ def _plan_and_ray() -> dict: if not ray.is_initialized(): _keep_npu_visible_in_cpu_actors() address = os.environ.get("RAY_ADDRESS") or None - ray.init(address=address, runtime_env=_runtime_env(), namespace="yeto-verl") + ray.init(address=address, runtime_env=_runtime_env(fully_async), namespace="yeto-verl") return plan @@ -125,7 +134,7 @@ def main_fully_async(argv: list[str] | None = None) -> None: import verl.trainer.main_ppo as main_ppo from verl.experimental.fully_async_policy import fully_async_main as fa_main - plan = _plan_and_ray() + plan = _plan_and_ray(fully_async=True) def run_ppo(config, task_runner_class=None): from yeto.rl.adapters.verl.fully_async_runner import YetoFullyAsyncTaskRunner @@ -147,7 +156,7 @@ def main_buffer_trainer(argv: list[str] | None = None) -> None: import verl.trainer.main_ppo as main_ppo from verl.experimental.fully_async_policy import fully_async_main as fa_main - plan = _plan_and_ray() + plan = _plan_and_ray(fully_async=True) def run_ppo(config, task_runner_class=None): from yeto.rl.adapters.verl.buffer_trainer_runner import YetoBufferTaskRunner diff --git a/yeto/rl/engine/bridges.py b/yeto/rl/engine/bridges.py index b1f955d32..47cb3b90d 100644 --- a/yeto/rl/engine/bridges.py +++ b/yeto/rl/engine/bridges.py @@ -138,7 +138,26 @@ def at_safe_point(self, driver: IslandDriver, *, rollout_id: int, final: bool = k: info.get(k) for k in ("cut_id", "cut_bytes", "save_s", "store_copy_s", "store_commit_s", "total_s", "write_bytes_per_s", "policy_hash", "store_synced", "latest_seq", "pruned", - "incarnation_index", "final", "trainer_onloaded_for_cut")}) + "incarnation_index", "final", "trainer_onloaded_for_cut", + "store_commit", "prepared_before_offload")}) + + def cut_due(self, rollout_id: int) -> bool: + """A round cut is due at the safe point before ``rollout_id``.""" + return self.round_cuts is not None and self.round_cuts.due( + rollout_id, final=rollout_id >= self._end()) + + def prepare_cut(self, driver: IslandDriver, *, rollout_id: int, state: Any) -> None: + """E1: write the trainer part of the cut due at the next safe point while the + trainer is still resident (the driver calls this right before it offloads a + colocated trainer for the publication). A failure is reported, not fatal: the + safe point then cuts the usual way.""" + if self.round_cuts is None: + return + try: + self.round_cuts.prepare(driver, rollout_id=rollout_id, state=state, + final=rollout_id >= self._end()) + except Exception as error: # noqa: BLE001 + driver.emit("rl_cut_prepare_failed", rollout_id=rollout_id, error=repr(error)[:2000]) def resume_data_cursor(self) -> dict[str, int] | None: """The data cursor recorded by the round cut this run resumed from (None: diff --git a/yeto/rl/engine/driver.py b/yeto/rl/engine/driver.py index d815a0a7e..f116459a1 100644 --- a/yeto/rl/engine/driver.py +++ b/yeto/rl/engine/driver.py @@ -74,6 +74,7 @@ ) from yeto.rl.elastic_client import LeaveReason +from . import phase_profile from .algorithm import AlgorithmSpec from .capabilities import BackendTraits, EngineCapabilities from .execution_profile import ( @@ -122,6 +123,8 @@ "publish": "trainer+rollout", "eval": "rollout", "offload": "trainer", + "checkpoint": "trainer", + "rollout_offload": "rollout", "onload": "trainer", } _PHASE_TASK = {"sync": "outer_sync"} @@ -366,6 +369,41 @@ def _pooled_reward(batch: RolloutBatchHandle) -> tuple[float, float]: return mean, math.sqrt(max(0.0, second - mean * mean)) +def defer_digest(driver: Any) -> bool: + """E1: a colocated trainer that offloads for the publication hashes the policy + while it offloads (``deferred`` digest, read at the publication). Needs a + policy state that supports it; ``YETO_RL_DEFERRED_DIGEST=0`` turns it off.""" + import os + + if os.environ.get("YETO_RL_DEFERRED_DIGEST", "1").strip().lower() in {"0", "false", "no", "off"}: + return False + return (bool(getattr(driver, "colocated", False)) and bool(getattr(driver, "publish_offloaded", False)) + and bool(getattr(driver.policy_state, "supports_deferred_digest", False))) + + +def verify_publication(driver: Any, version: int) -> bool: + """E1: does the publication of policy ``version`` hash the policy and read the + engine weight checksums back? Yes on the first publication, on a round whose safe + point keeps a round cut (resume relies on its policy hash), on every K-th version + (K = the round-cut interval) and, with ``--rl-verify-every-round``, always. Without + round cuts, or with a policy state that cannot skip it, every round is verified. + The other rounds still check that the trainer holds the policy (weights version).""" + if bool(getattr(driver, "verify_every_round", False)): + return True + if not bool(getattr(getattr(driver, "policy_state", None), "supports_unverified_rounds", False)): + return True + cuts = getattr(getattr(driver, "sync", None), "round_cuts", None) + every = int(getattr(cuts, "every", 0) or 0) + if cuts is None or every <= 1: + return True + if getattr(driver, "published_version", None) is None: + return True + cut_due = getattr(driver.sync, "cut_due", None) + if callable(cut_due) and cut_due(int(version)): + return True + return int(version) % every == 0 + + class IslandDriver: def __init__( self, @@ -466,6 +504,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")) + phase_profile.on_phase(name, fields.get("rollout_id")) # no-op unless --rl-profile-trainer if self.observe: self._close_span() self._open_span = (name, self.clock(), fields.get("rollout_id")) @@ -685,6 +724,12 @@ def export_local_resident(self, *, policy_version: int): export_digest = getattr(self.policy_state, "export_digest", None) if not callable(export_digest): return None + if not verify_publication(self, int(policy_version)): + state = self.policy_state.export_identity(policy_version=int(policy_version)) + if state is not None: + return state + if defer_digest(self): + return export_digest(policy_version=int(policy_version), deferred=True) return export_digest(policy_version=int(policy_version)) # -- per-round steps ------------------------------------------------- @@ -737,6 +782,9 @@ def publish(self, state: TrainableState, *, rollout_id: int) -> None: policy_version=rollout_id, **{ "rl/policy_token": self.expected_token, + # E1: only on a round whose policy content was not hashed (recorded tapes + # of verified rounds stay byte-identical) + **({"sync/publication_verified": False} if getattr(state, "verified", True) is False else {}), "sync/publication_payload_bytes": manifest.payload_bytes, "sync/publication_payload_hash": manifest.payload_hash, "sync/publication_members": sorted(result.members), @@ -889,11 +937,20 @@ def _generate(self, rollout_id: int) -> RolloutBatchHandle: f"generation of rollout {rollout_id}", ) self.phase("generate", rollout_id=rollout_id, policy_version=rollout_id) + offload_rollout = getattr(self.rollout, "offload_after_generate", None) + split_offload = callable(offload_rollout) and bool(getattr(self.rollout, "offloads_after_generate", False)) + if split_offload: + self.rollout.defer_offload = True with self._load_sampler(rollout_id): # IR-3: the rollout side learns the target policy token up front. batch = self.rollout.generate( rollout_id, expected_policy_version=self.expected_token ) + # E1: a colocated rollout that frees its memory after generating (Miles + # --offload-rollout) does it in its own phase, not inside "generate" + if split_offload: + self.phase("rollout_offload", rollout_id=rollout_id) + offload_rollout() if batch.rollout_id != rollout_id or batch.policy_version != rollout_id: raise PolicyIdentityError( f"rollout pool returned rollout {batch.rollout_id} " @@ -1167,20 +1224,35 @@ def loop() -> None: thread.join(timeout=self.load_sample_interval_s + 5) # S15 evidence-window lesson: one terminal probe after generate so a # round whose loads mostly land after the last tick still leaves a - # final ``rl_load_sample`` (``terminal=True``). Observe only. - try: - sample = probe() - except Exception as exc: # noqa: BLE001 - observation must not kill the round - import logging - - logging.getLogger(__name__).warning( - "yeto terminal load_sample probe raised %s: %s", type(exc).__name__, exc) - sample = None - if sample is not None: - self.emit("rl_load_sample", rollout_id=rollout_id, **dict(sample), - t=self.clock(), profile_hash=self.profile_hash, - epoch=self.config_epoch, - weight_transport=self.weight_transport, terminal=True) + # final ``rl_load_sample`` (``terminal=True``). Observe only, so it runs + # off the round's critical path (E1: a synchronous probe cost ~0.1 s/round). + self._join_terminal_probe() + self._terminal_probe = threading.Thread( + target=self._terminal_load_sample, args=(probe, rollout_id), + name="yeto-load-sample-terminal", daemon=True) + self._terminal_probe.start() + + def _join_terminal_probe(self) -> None: + """Wait (bounded) for the previous round's terminal load probe: at most one runs.""" + thread = getattr(self, "_terminal_probe", None) + if thread is not None: + thread.join(timeout=self.load_sample_interval_s + 5) + self._terminal_probe = None + + def _terminal_load_sample(self, probe: Callable[[], Any], rollout_id: int) -> None: + try: + sample = probe() + except Exception as exc: # noqa: BLE001 - observation must not kill the round + import logging + + logging.getLogger(__name__).warning( + "yeto terminal load_sample probe raised %s: %s", type(exc).__name__, exc) + return + if sample is not None: + self.emit("rl_load_sample", rollout_id=rollout_id, **dict(sample), + t=self.clock(), profile_hash=self.profile_hash, + epoch=self.config_epoch, + weight_transport=self.weight_transport, terminal=True) def _emit_round_labels(self, rollout_id, batch, metrics) -> None: """A5: per-round algorithm metrics carry the same profile/epoch/transport labels.""" @@ -1455,6 +1527,10 @@ def _train_round(self, rollout_id: int, batch: RolloutBatchHandle, raw = self.trainer.step_metrics() metrics = raw if isinstance(raw, TrainStepMetrics) else TrainStepMetrics(**dict(raw)) rollout_fraction = getattr(batch, "cross_version_truncated_fraction", None) + late = getattr(self.rollout, "late_cross_version_fraction", None) + if rollout_fraction is None and callable(late): + # E1: estimated in the background while this step trained (Miles fully_async) + rollout_fraction = late(rollout_id) if metrics.cross_version_truncated_fraction is None and rollout_fraction is not None: # agentic-rollout-utilization 4.1: an engine that estimates the # cross-version truncation on the rollout side (Miles carry-over) @@ -1531,6 +1607,12 @@ def _train_round(self, rollout_id: int, batch: RolloutBatchHandle, self.ledger.outer_recorded(rollout_id, next_policy_version=rollout_id + 1) self._join_eval() if self.colocated and self.publish_offloaded: + # E1: a round cut due at the next safe point writes the trainer state now, + # while it is resident, instead of waking the offloaded trainer for it + cut_due = getattr(self.sync, "cut_due", None) + if callable(cut_due) and cut_due(rollout_id + 1): + self.phase("checkpoint", rollout_id=rollout_id) + self.sync.prepare_cut(self, rollout_id=rollout_id + 1, state=boundary.state) self._offload_trainer(rollout_id) self.publish(boundary.state, rollout_id=rollout_id + 1) self._close_span() @@ -1617,6 +1699,11 @@ def safe_point(self, rollout_id: int) -> str | None: self.elastic_hook.at_safe_point(self, rollout_id) if not self.controller.has_pending() and not self.controller.recovery_required: return None + # E1: a reconfiguration syncs the checkpoint store itself; no round cut may still + # be committing it in the background + join_commit = getattr(getattr(self.sync, "round_cuts", None), "join_commit", None) + if callable(join_commit): + join_commit(self) epoch_before = self.config_epoch self.phase("reconfigure", rollout_id=rollout_id, config_epoch=epoch_before) from .controller import RecoveryRequired @@ -1885,6 +1972,7 @@ def _run(self) -> TrainableState: failure = error raise finally: + self._join_terminal_probe() # its rl_load_sample lands before the tape closes if self.eval_overlap is not None: self.eval_overlap.abort() # no orphan eval task on error/cancel paths self._close_sync(failure) diff --git a/yeto/rl/engine/phase_profile.py b/yeto/rl/engine/phase_profile.py new file mode 100644 index 000000000..4b342189d --- /dev/null +++ b/yeto/rl/engine/phase_profile.py @@ -0,0 +1,89 @@ +"""E1 observation (``--rl-profile-trainer``): where the island driver's main thread +spends each driver phase. + +When enabled, a daemon thread samples the stack of the thread that reports the +first phase (the driver's: it is not the thread that enables it, GPU run +e1n-yeto-miles-async-p) every ``SAMPLE_S`` and counts the samples per (phase, frame); at each +phase change the finished phase's wall time and top frames are printed as one +``yeto phase profile {json}`` line. Frames are labelled by the innermost frame in +yeto, miles or ray code (else the innermost frame). Nothing else changes. +""" + +from __future__ import annotations + +import collections +import json +import sys +import threading +import time +from typing import Any + +SAMPLE_S = 0.005 +TOP = 6 +_MARKERS = ("/yeto/", "/miles/", "/ray/") +_STATE: dict[str, Any] = {"thread": None} + + +def _label(frame: Any, depth: int = 4) -> str: + """The innermost ``depth`` yeto/miles/ray frames, innermost first (else the + innermost frame): enough of the call chain to tell who waits on what.""" + inner, chain = None, [] + while frame is not None and len(chain) < depth: + name = frame.f_code.co_filename + here = f"{name.rsplit('/', 1)[-1]}:{frame.f_code.co_name}" + inner = inner or here + if any(m in name for m in _MARKERS): + chain.append(here) + frame = frame.f_back + return " < ".join(chain) if chain else (inner or "?") + + +class _PhaseSampler(threading.Thread): + def __init__(self) -> None: + super().__init__(daemon=True, name="yeto-phase-profile") + self.target: int | None = None + self.lock = threading.Lock() + self.phase: tuple[str, Any] | None = None + self.started = time.perf_counter() + self.counts: collections.Counter = collections.Counter() + self.samples = 0 + + def run(self) -> None: + while True: + time.sleep(SAMPLE_S) + if self.target is None: + continue + frame = sys._current_frames().get(self.target) + if frame is None: + continue + with self.lock: + if self.phase is not None: + self.samples += 1 + self.counts[_label(frame)] += 1 + + def switch(self, phase: tuple[str, Any]) -> None: + now = time.perf_counter() + with self.lock: + if self.target is None: + self.target = threading.get_ident() # the thread that drives the phases + done, counts, samples = self.phase, self.counts, self.samples + self.phase, self.counts, self.samples, started, self.started = ( + phase, collections.Counter(), 0, self.started, now) + if done is not None: + print("yeto phase profile " + json.dumps({ + "phase": done[0], "rollout_id": done[1], "wall_s": round(now - started, 4), + "samples": samples, "top": counts.most_common(TOP)}), flush=True) + + +def enable() -> None: + """Start the sampler; it samples the thread of the first :func:`on_phase`.""" + if _STATE["thread"] is None: + sampler = _PhaseSampler() + sampler.start() + _STATE["thread"] = sampler + + +def on_phase(name: str, rollout_id: Any) -> None: + sampler = _STATE["thread"] + if sampler is not None: + sampler.switch((name, rollout_id)) diff --git a/yeto/rl/engine/policy_digest.py b/yeto/rl/engine/policy_digest.py index cd0702752..734fa765a 100644 --- a/yeto/rl/engine/policy_digest.py +++ b/yeto/rl/engine/policy_digest.py @@ -86,6 +86,49 @@ def from_wire(cls, raw: Mapping[str, Any]) -> "PolicyDigest": ) +class DeferredDigest: + """A :class:`PolicyDigest` whose hashes are still being computed in the trainer + (E1, sync colocated mode: the trainer hashes while it offloads). The layout + (``specs``) is known at once; the first read of a hash calls ``resolve`` once + (it waits for the trainer) and the values are those of the eager digest.""" + + def __init__(self, specs: Sequence[tuple[str, tuple[int, ...]]], resolve: Callable[[], "PolicyDigest"]): + self.specs = tuple((str(n), tuple(int(d) for d in s)) for n, s in specs) + self._resolve = resolve + self._done: PolicyDigest | None = None + + def resolved(self) -> "PolicyDigest": + if self._done is None: + done = self._resolve() + if done.specs != self.specs: + raise PolicyDigestError("the deferred digest's layout changed while it was hashed") + self._done = done + return self._done + + @property + def is_resolved(self) -> bool: + return self._done is not None + + @property + def policy_tensor_hash(self) -> str: + return self.resolved().policy_tensor_hash + + @property + def payload_hash(self) -> str: + return self.resolved().payload_hash + + @property + def payload_bytes(self) -> int: + return self.resolved().payload_bytes + + @property + def hash_seconds(self) -> float: + return self.resolved().hash_seconds + + def to_wire(self) -> dict[str, Any]: + return self.resolved().to_wire() + + def _raw(tensor: Any, name: str) -> memoryview: import torch @@ -159,6 +202,83 @@ def digest_canonical_tensors( hash_seconds=round(time.monotonic() - started, 3)) +class NonFiniteTensor(ValueError): + """A tensor staged by :func:`stage_f32_on_host` holds NaN or Inf (``name`` says which).""" + + def __init__(self, name: str): + super().__init__(f"{name!r} contains NaN or Inf") + self.name = name + + +class HostStaging: + """One reusable pinned host buffer for :func:`stage_f32_on_host` (E1). + + Allocating pinned memory costs milliseconds per call; the policy keeps the + same size every round, so the owner (the trainer process) keeps one of + these and the buffer is reallocated only when the policy grows.""" + + def __init__(self, *, pin: bool = True) -> None: + self._pin = pin + self._buffer = None + + def flat(self, numel: int): + import torch + + if self._buffer is None or self._buffer.numel() < numel: + pin = self._pin and torch.cuda.is_available() + self._buffer = torch.empty(numel, dtype=torch.float32, pin_memory=pin) + return self._buffer[:numel] + + +def stage_f32_on_host(tensors: Mapping[str, Any], staging: HostStaging | None = None) -> dict[str, Any]: + """Canonical CPU f32 contiguous copies of ``tensors`` (sorted by name), bit for bit + the values of ``t.detach().to("cpu", torch.float32).contiguous()`` per tensor. + + E1: with ``staging`` and every tensor on one device, the f32 conversion and the + finite check run on that device (the GPU) and the copies go to the pinned host + buffer asynchronously, with ONE synchronization for all of them (the per-tensor + path waits once per tensor: ~400 waits for a Qwen3-0.6B LoRA). No extra device + buffer is held: a tensor already f32 and contiguous is copied as it is. The + returned tensors are views of ``staging``'s buffer and are valid until the next + call with the same ``staging``: hash them, do not keep them (``torch.save`` of a + view would also write the whole buffer). Without ``staging`` (or with tensors on + several devices) each tensor is copied on its own, as before. Raises + :class:`NonFiniteTensor` naming the first non-finite tensor in name order, like + the per-tensor check.""" + + import torch + + names = sorted(tensors) + values = [tensors[n].detach() for n in names] + devices = {v.device for v in values} + if staging is None or len(devices) != 1: + out = {} + for name, value in zip(names, values): + value = value.to(device="cpu", dtype=torch.float32).contiguous().clone() + if not torch.isfinite(value).all().item(): + raise NonFiniteTensor(name) + out[name] = value + return out + flat = staging.flat(sum(v.numel() for v in values)) + out, finite, offset = {}, [], 0 + for name, value in zip(names, values): + source = value.to(dtype=torch.float32).contiguous() + finite.append(torch.isfinite(source).all()) + view = flat[offset:offset + source.numel()].view(source.shape) + view.copy_(source, non_blocking=True) + out[name] = view + offset += source.numel() + device = next(iter(devices)) + ok = bool(torch.stack(finite).all().item()) if finite else True + if device.type == "cuda": + torch.cuda.current_stream(device).synchronize() # the copies to the host are complete + if not ok: + for name, check in zip(names, finite): + if not bool(check.item()): + raise NonFiniteTensor(name) + return out + + def layout_hash_of(specs: Sequence[tuple[str, tuple[int, ...]]]) -> str: """The canonical layout hash rebuilt from (name, shape) only.""" @@ -194,6 +314,10 @@ class TrainerResidentState: weights_mark: Mapping[str, Any] | None = None # (state) -> None, raises when the trainer no longer holds this policy holds_fn: Callable[[Any], None] | None = field(default=None, repr=False, compare=False) + # E1: False for a round whose policy content is not hashed (see unverified_identity): + # ``digest`` then holds an identity of the trainer weights version, not a content hash, + # and the publisher skips the engine read-back + verified: bool = True _cache: dict = field(default_factory=dict, repr=False, compare=False) def __post_init__(self) -> None: @@ -229,7 +353,7 @@ def materialize(self) -> TrainableState: state = self._cache.get("state") if state is None: state = self.materialize_fn(self.policy_version) - if state.policy_tensor_hash() != self.digest.policy_tensor_hash: + if self.verified and state.policy_tensor_hash() != self.digest.policy_tensor_hash: raise PolicyDigestError( "trainer weights changed after the resident digest was taken") self._cache["state"] = state @@ -246,5 +370,17 @@ def policy_hash(self) -> str: return self.materialize().policy_hash() +UNVERIFIED_PREFIX = b"yeto-unverified-policy-v1\0" + + +def unverified_identity(policy_version: int, weights_mark: Mapping[str, Any]) -> str: + """The 64-hex policy identity of a publication whose content is not hashed (E1): + names the trainer process and its weights version, so two different weights never + share it; it is never equal to a content hash of the same weights.""" + raw = json.dumps({"policy_version": int(policy_version), "weights": dict(weights_mark)}, + sort_keys=True).encode("utf-8") + return hashlib.sha256(UNVERIFIED_PREFIX + raw).hexdigest() + + def is_resident(state: Any) -> bool: return isinstance(state, TrainerResidentState) diff --git a/yeto/rl/engine/resume.py b/yeto/rl/engine/resume.py index bfa3894e9..99b54ff38 100644 --- a/yeto/rl/engine/resume.py +++ b/yeto/rl/engine/resume.py @@ -52,6 +52,9 @@ SKIP_NAMES = frozenset({"journal.lock", STATE_MANIFEST, LATEST}) DEFAULT_CUT_KEEP = 2 DEFAULT_CUT_EVERY = 1 +# E1: next to a store's first round cut while it is written and committed in the +# background (``round-cuts/`` + this suffix, a file); see discard_unfinished_first_cut +FIRST_CUT_MARK = ".first-cut-in-progress" class StoreIntegrityError(RuntimeError): @@ -210,10 +213,13 @@ def read_latest(store_root: Path) -> dict[str, Any] | None: def write_snapshot(store: CheckpointStore, state_dir: Path, info: Mapping[str, Any], *, - skip: Iterable[str] = ()) -> dict[str, Any]: + skip: Iterable[str] = (), commit: bool = True) -> dict[str, Any]: """Copy ``state_dir`` into a NEW snapshot dir (every file hash-checked after the copy), write its manifest, then ``LATEST``, then commit. Returns the LATEST record - plus timings/bytes.""" + plus timings/bytes. + + ``commit=False`` leaves the commit to the caller (``commit_s`` is None): until it + commits, a reader of a committing store still sees the previous LATEST.""" root = store.root started = time.monotonic() latest = read_latest(root) @@ -233,7 +239,7 @@ def write_snapshot(store: CheckpointStore, state_dir: Path, info: Mapping[str, A record = {"seq": seq, "manifest_sha256": manifest_sha, "bytes": total, **{k: info[k] for k in ("cut_id", "next_rollout_id", "incarnation", "reason") if k in info}} _write_json_atomic(root / LATEST, record) - commit_s = store.commit() + commit_s = store.commit() if commit else None return {**record, "copy_s": copied_s, "commit_s": commit_s} @@ -366,15 +372,96 @@ def prune(store_root: Path, *, keep: int, cut_ids_of: Callable[[Path], str | Non shutil.rmtree(p, ignore_errors=True) deleted.append(f"{STATE_DIR}/{p.name}") wanted = {cid for cid in (cut_ids_of(p) for p in kept) if cid} + # E1: a cut newer than every kept one may be being written right now (a periodic + # cut is captured while the previous cut commits in the background): not pruned + newest = max((r for r in map(_cut_rollout, wanted) if r is not None), default=None) cuts = root / ROUND_CUTS if cuts.is_dir(): for p in sorted(cuts.iterdir()): + rollout = _cut_rollout(p.name) + if newest is not None and rollout is not None and rollout > newest: + continue if p.is_dir() and p.name not in wanted: shutil.rmtree(p, ignore_errors=True) deleted.append(f"{ROUND_CUTS}/{p.name}") return deleted +def _cut_rollout(cut_id: str) -> int | None: + """Rollout id of a round cut named ``r-`` (None: other name).""" + head, sep, _ = cut_id.partition("-") + if not sep or not head.startswith("r") or not head[1:].isdigit(): + return None + return int(head[1:]) + + +def first_cut_mark(store_root: Path, cut_id: str) -> Path: + """The "first cut in progress" mark of round cut ``cut_id`` (a file in round-cuts/).""" + return Path(store_root) / ROUND_CUTS / f"{cut_id}{FIRST_CUT_MARK}" + + +def mark_first_cut(store_root: Path, cut_id: str) -> Path: + """E1: written (fsynced) before the store's first round cut starts its shards, so a + run that dies before that cut's LATEST is committed can tell its files from the + cuts of an unfinished store (:func:`discard_unfinished_first_cut`).""" + mark = first_cut_mark(store_root, cut_id) + mark.parent.mkdir(parents=True, exist_ok=True) + with open(mark, "w", encoding="utf-8") as fh: + fh.write(json.dumps({"cut_id": cut_id})) + fh.flush() + os.fsync(fh.fileno()) + return mark + + +def clear_first_cut_marks(store_root: Path) -> list[str]: + """Remove every first-cut mark (after a LATEST is committed they mean nothing).""" + cuts = Path(store_root) / ROUND_CUTS + removed = [] + if cuts.is_dir(): + for p in sorted(cuts.glob(f"*{FIRST_CUT_MARK}")): + p.unlink(missing_ok=True) + removed.append(p.name) + return removed + + +def discard_unfinished_first_cut(store_root: Path) -> list[str] | None: + """For a store without LATEST that holds round-cut entries: when every cut directory + carries a first-cut mark (none was ever committed: the run died while writing its + first cut), delete them and the marks and return what was deleted (the run starts + at rollout 0). Any unmarked entry (a cut of an unfinished store, or an unknown + file) returns None: the caller refuses to start.""" + root = Path(store_root) + if read_latest(root) is not None: + raise ValueError(f"{root} has {LATEST}: its first cut is committed") + cuts = root / ROUND_CUTS + entries = sorted(cuts.iterdir()) if cuts.is_dir() else [] + marks = {p.name for p in entries if p.is_file() and p.name.endswith(FIRST_CUT_MARK)} + others = [p for p in entries if p.name not in marks] + if any(not p.is_dir() or f"{p.name}{FIRST_CUT_MARK}" not in marks for p in others): + return None + deleted = [] + for p in others: + shutil.rmtree(p) + deleted.append(f"{ROUND_CUTS}/{p.name}") + for name in sorted(marks): + (cuts / name).unlink() + deleted.append(f"{ROUND_CUTS}/{name}") + return deleted + + +def stage_state_dir(state_dir: Path, staged: Path) -> Path: + """Copy the snapshot files of ``state_dir`` (as :func:`write_snapshot` would take + them) into a fresh ``staged`` directory; returns it.""" + state_dir, staged = Path(state_dir), Path(staged) + if staged.exists(): + shutil.rmtree(staged) + for rel in _tree_files(state_dir): + (staged / rel).parent.mkdir(parents=True, exist_ok=True) + shutil.copy2(state_dir / rel, staged / rel) + staged.mkdir(parents=True, exist_ok=True) + return staged + + def directory_bytes(path: Path) -> int: return sum(p.stat().st_size for p in Path(path).rglob("*") if p.is_file()) @@ -423,6 +510,12 @@ def check_lr_continues(expected: list[float] | None, applied: list[float] | None return None +# E1 (product default since 2026-10-21): a round cut every 4 rounds (and always at the +# end); --rl-cut-every 1 cuts every round. The publication checks follow the same +# interval (driver.verify_publication). +DEFAULT_CUT_EVERY = 4 + + def should_cut(rollout_id: int, *, every: int, final: bool = False, stop: bool = False) -> bool: """Cut at the safe point before ``rollout_id`` every ``every`` rounds, and always at the end (last round, our own stop).""" @@ -455,7 +548,9 @@ class ResumeController: ``sync_checkpoint_store`` (snapshot + LATEST + commit + prune) and ``_record``. At construction, a store with ``LATEST`` is restored into the state dir (verified); - a store that holds cuts but no ``LATEST`` is refused (never a silent restart at 0).""" + a store that holds cuts but no ``LATEST`` is refused (never a silent restart at 0), + except when they are all the marked, never committed first cut + (:func:`discard_unfinished_first_cut`): those are deleted and the run starts at 0.""" def __init__(self, *, state_dir: str | os.PathLike[str], store: CheckpointStore, runtime_fingerprint: str = "", keep: int = DEFAULT_CUT_KEEP, @@ -474,6 +569,8 @@ def __init__(self, *, state_dir: str | os.PathLike[str], store: CheckpointStore, store.root.mkdir(parents=True, exist_ok=True) store.reload() self.restored: dict[str, Any] | None = None + # E1: the files of a first round cut that never committed, deleted at start + self.discarded_first_cut: list[str] | None = None latest = read_latest(store.root) started = time.monotonic() if latest is not None: @@ -481,8 +578,17 @@ def __init__(self, *, state_dir: str | os.PathLike[str], store: CheckpointStore, self.restored["latest"] = dict(latest) self.restored["seconds"] = time.monotonic() - started elif (store.root / ROUND_CUTS).is_dir() and any((store.root / ROUND_CUTS).iterdir()): - raise StoreIntegrityError(f"{store.root} holds round cuts but no {LATEST}: refusing to " - "start over at rollout 0 on top of an unfinished store") + discarded = discard_unfinished_first_cut(store.root) + if discarded is None: + raise StoreIntegrityError(f"{store.root} holds round cuts but no {LATEST}: refusing to " + "start over at rollout 0 on top of an unfinished store") + store.commit() + self.discarded_first_cut = discarded + # the run starts at 0: a local state dir of the dead run (its pointer may + # name the deleted cut) is moved aside, as a restore does + if self.state_dir.exists() and any(self.state_dir.iterdir()): + self.state_dir.rename(self.state_dir.with_name( + f"{self.state_dir.name}.prev-{int(time.time())}")) self.state_dir.mkdir(parents=True, exist_ok=True) previous = self._incarnations() self.incarnation = {"id": uuid.uuid4().hex[:12], "index": len(previous), "pid": os.getpid(), @@ -490,6 +596,9 @@ def __init__(self, *, state_dir: str | os.PathLike[str], store: CheckpointStore, with open(self.state_dir / INCARNATIONS, "a", encoding="utf-8") as fh: fh.write(json.dumps(self.incarnation, sort_keys=True) + "\n") self.previous_incarnations = previous + if self.discarded_first_cut is not None: + self._record("checkpoint_store", tx_id=None, action="discard_unfinished_first_cut", + deleted=self.discarded_first_cut, store=str(store.root)) if self.restored is not None: self._record("checkpoint_store", tx_id=None, action="restore", **{ k: v for k, v in self.restored.items() if k != "latest"}, @@ -509,15 +618,25 @@ def _record(self, kind: str, **fields: Any) -> dict[str, Any]: return record def sync_checkpoint_store(self, reason: str, *, skip_if_recovery: bool = False, - extra: Mapping[str, Any] | None = None) -> bool: + extra: Mapping[str, Any] | None = None, + defer_commit: bool = False, source_dir: Path | None = None) -> bool: """Snapshot the state dir into the store, then ``LATEST``, commit, prune. A failure is remembered (``last_store_sync``) and returned False: the previous - LATEST stays valid.""" + LATEST stays valid. + + ``defer_commit``: stop after ``LATEST``; :meth:`commit_checkpoint_store` commits + and prunes later (the round cut runs it off the round's critical path). Until + then the store's committed LATEST is the previous one, so a crash in between + resumes from the previous cut. + + ``source_dir``: snapshot this copy of the state dir instead (a round cut whose + store sync runs later stages the state dir as it was at the cut).""" if skip_if_recovery and self.recovery_required: return False + source = Path(source_dir) if source_dir is not None else self.state_dir info = {"incarnation": self.incarnation["id"], "incarnation_index": self.incarnation["index"], "reason": reason, "wall_time": self._wall(), **dict(extra or {})} - pointer = self.state_dir / "round-cut.json" + pointer = source / "round-cut.json" if pointer.is_file(): # LATEST names the cut the snapshot's pointer names try: p = json.loads(pointer.read_text(encoding="utf-8")) @@ -526,17 +645,44 @@ def sync_checkpoint_store(self, reason: str, *, skip_if_recovery: bool = False, except ValueError: pass try: - result = write_snapshot(self.store, self.state_dir, info) - deleted = prune(self.store.root, keep=self.keep, cut_ids_of=_snapshot_cut_id) - if deleted: - self.store.commit() + result = write_snapshot(self.store, source, info, commit=not defer_commit) + if defer_commit: + self.last_store_sync = {**info, **result, "ok": True, "committed": False, "pruned": None} + return True + deleted = self._prune() except Exception as exc: # noqa: BLE001 self.last_store_sync = {**info, "ok": False, "error": repr(exc)} logging.getLogger(__name__).warning("checkpoint store sync (%s) failed: %r", reason, exc) return False - self.last_store_sync = {**info, **result, "ok": True, "pruned": deleted} + self.last_store_sync = {**info, **result, "ok": True, "committed": True, "pruned": deleted} return True + def commit_checkpoint_store(self) -> dict[str, Any]: + """Commit a ``defer_commit`` sync, then prune (pruning only after the new LATEST + is committed, so a failed commit never loses the previous cut). Thread-safe + against the caller as long as nothing else writes to the store meanwhile. + Returns (and remembers in ``last_store_sync``) the outcome; never raises.""" + sync = dict(self.last_store_sync or {}) + try: + sync["commit_s"] = self.store.commit() + sync["committed"] = True + sync["pruned"] = self._prune() + except Exception as exc: # noqa: BLE001 + sync.update(ok=False, error=repr(exc)) + logging.getLogger(__name__).warning("checkpoint store commit (%s) failed: %r", + sync.get("reason"), exc) + self.last_store_sync = sync + return sync + + def _prune(self) -> list[str]: + """Housekeeping once a new LATEST is committed: old snapshots and cuts, and the + first-cut marks (the first cut is committed now; the next commit drops them).""" + clear_first_cut_marks(self.store.root) + deleted = prune(self.store.root, keep=self.keep, cut_ids_of=_snapshot_cut_id) + if deleted: + self.store.commit() + return deleted + def close(self) -> None: return None diff --git a/yeto/rl/engine/round_cut.py b/yeto/rl/engine/round_cut.py index cadc5047e..c0b1d010f 100644 --- a/yeto/rl/engine/round_cut.py +++ b/yeto/rl/engine/round_cut.py @@ -39,6 +39,7 @@ import json import os from collections.abc import Callable, Mapping +from dataclasses import dataclass from pathlib import Path from typing import Any @@ -92,15 +93,35 @@ def driver(self) -> Any: ref = self.driver_ref return ref() if callable(ref) and not hasattr(ref, "published_state") else ref - def context(self, cut_id: str) -> Any: + def context(self, cut_id: str, *, pending: tuple[int, Any] | None = None) -> Any: + """The cut context at the safe point after the current publication. + + ``pending`` = ``(version, state)``: build it before that publication, for + the state the driver is about to publish as ``version`` (the policy token + is the one the publication will set). Only valid when nothing but the + publication happens until the safe point (colocated, generation done); + :meth:`RoundCutCheckpoint.save` keeps such a cut only if the publication + then names the same policy.""" from yeto.rl.engine.controller import RebuildRefused from yeto.rl.engine.cut import CutProgress from yeto.rl.engine.cut import CutContext + from yeto.rl.engine.driver import policy_token driver = self.driver - state = driver.published_state - if state is None or driver.published_version is None: + if pending is not None: + # the token is read last (below): a deferred digest (E1) hashes meanwhile + version, state = int(pending[0]), pending[1] + settled, token = True, None + else: + version, state = driver.published_version, driver.published_state + settled, token = bool(driver.at_safe_point), driver.expected_token + if state is None or version is None: raise RebuildRefused("no published policy to cut") + if getattr(state, "verified", True) is False: + # E1: a cut carries the policy's content hash (resume checks it); a round + # whose publication was not hashed is never committed as a cut + raise RebuildRefused(f"the publication of policy v{version} was not verified " + "(no content hash); see --rl-verify-every-round") cursor = (self.rollout.data_cursor() if callable(getattr(self.rollout, "data_cursor", None)) else None) if cursor is None: @@ -122,6 +143,10 @@ def context(self, cut_id: str) -> Any: if problems: raise RebuildRefused("; ".join(problems)) in_flight = tuple(exported["entries"]) + if exported.get("data_cursor") is not None: + # read together with the entries (a rollout that draws while the + # cut is taken, e.g. Miles fully_async): keep cursor and entries consistent + data.update(exported["data_cursor"]) data["buffer_length"] = int(exported["buffer_groups"]) summary["carried_over"] = len(in_flight) summary["max_policy_age"] = limit @@ -129,17 +154,19 @@ def context(self, cut_id: str) -> Any: local_step=int(driver.local_step), scheduler_samples=int(driver.local_step) * self.global_batch_size, global_batch_size=self.global_batch_size, - next_rollout_id=int(driver.published_version), + next_rollout_id=int(version), policy_version=int(state.policy_version), policy_hash=state.policy_tensor_hash(), ) + if token is None: + token = policy_token(version, progress.policy_hash) return CutContext( root=self.cut_root, cut_id=cut_id, backend_fingerprint=self.backend_fingerprint, progress=progress, algorithm=self.identity, data=data, ledger=summary, # The safe point is after boundary() returned non-stop and every # member acknowledged the policy: the outer commit is settled. - outer={"settled": bool(driver.at_safe_point), "policy_version": state.policy_version, - "policy_token": driver.expected_token}, + outer={"settled": settled, "policy_version": state.policy_version, + "policy_token": token}, shared_filesystem=self.shared_filesystem, in_flight=in_flight, ) @@ -155,6 +182,26 @@ def expectation(self, layout: Mapping[str, int], *, epoch: int) -> Any: ) +@dataclass(frozen=True) +class _PreparedCut: + """The trainer part of a cut written before the publication (``prepare``).""" + + rollout_id: int + context: Any # CutContext the shards and manifest were written with + save_s: float + pending: Any = None # the trainer's pending cut (start_cut); None: already written + + def matches(self, rollout_id: int, context: Any) -> bool: + """The safe point's own context names the same cut: same id, progress + (version, policy hash, steps), data cursor, ledger summary, outer commit + (settled, policy token) and in-flight entries.""" + mine = self.context + return (self.rollout_id == rollout_id and mine.cut_id == context.cut_id + and mine.progress == context.progress and dict(mine.data) == dict(context.data) + and dict(mine.ledger) == dict(context.ledger) and dict(mine.outer) == dict(context.outer) + and tuple(mine.in_flight or ()) == tuple(context.in_flight or ())) + + class RoundCutCheckpoint: """``save`` at a safe point, ``resume`` at sync start (see the module doc). @@ -166,7 +213,7 @@ class RoundCutCheckpoint: def __init__(self, *, controller: Any, source: Any, trainer: Any, every: int = 1, fingerprint: dict[str, Any] | None = None, allow_config_change: bool = False, - next_lr: Any = None) -> None: + next_lr: Any = None, background_commit: bool | None = None) -> None: if getattr(controller, "checkpoint_store", None) is None: raise ValueError("round cuts need the checkpoint store") if int(every) < 1: @@ -191,6 +238,25 @@ def __init__(self, *, controller: Any, source: Any, trainer: Any, every: int = 1 # (buffer_rollout.wire_sample_buffer sets both). self.buffer_watermark: Any = None # callable(int) -> int | None self.buffer_resume: Any = None # callable(int) -> int + # E1: a cut prepared before the trainer offloads (see prepare()), consumed by save() + self._prepared: _PreparedCut | None = None + # E1: a periodic cut writes its shards and commits the store in the background + # (the controller must offer commit_checkpoint_store); completed before the next + # cut and at the final cut. YETO_RL_CUT_BACKGROUND_COMMIT=0 does it all inline. + if background_commit is None: + background_commit = os.environ.get("YETO_RL_CUT_BACKGROUND_COMMIT", "1").strip().lower() \ + not in {"0", "false", "no", "off"} + self.background_commit = bool(background_commit) and callable( + getattr(controller, "commit_checkpoint_store", None)) + self._commit: tuple[Any, dict[str, Any]] | None = None # (thread, result box) + # E1: a periodic cut captured at its safe point, completed by join_commit() + self._deferred: dict[str, Any] | None = None + self._commit_info: dict[str, Any] | None = None # the running commit's cut + self.skipped_busy = 0 # periodic cuts skipped while a commit was running + # E1: prepared cuts a busy commit made the safe point skip; their trainer-side + # writers are collected (and their files dropped) at the next cut, when done + self._abandoned: list[_PreparedCut] = [] + self._declined: int | None = None # rollout whose prepare met a running commit # ---------------------------------------------------------------- pointer @property @@ -203,12 +269,28 @@ def pointer(self) -> dict[str, Any] | None: return None return dict(json.loads(path.read_text(encoding="utf-8"))) - def _write_pointer(self, info: dict[str, Any]) -> None: - tmp = self.pointer_path.with_suffix(".json.tmp") + def _write_pointer(self, info: dict[str, Any], *, path: Path | None = None) -> None: + target = self.pointer_path if path is None else Path(path) + tmp = target.with_suffix(".json.tmp") tmp.write_text(json.dumps(info, sort_keys=True), encoding="utf-8") with open(tmp, "rb") as fh: os.fsync(fh.fileno()) - os.replace(tmp, self.pointer_path) + os.replace(tmp, target) + + def _stage(self, cut_id: str) -> Path: + """E1: a copy of the state dir as it is now (small: ledger, journal, pointer), + next to it, for a store sync that runs after later rounds.""" + from yeto.rl.engine.resume import stage_state_dir + + state_dir = Path(self.controller.state_dir) + return stage_state_dir(state_dir, state_dir.parent / f".{state_dir.name}.stage-{cut_id}") + + @staticmethod + def _unstage(staged: Path | None) -> None: + if staged is not None: + import shutil + + shutil.rmtree(staged, ignore_errors=True) def _epoch(self) -> int: return int(self.controller.journal.epochs.config_epoch) @@ -228,6 +310,216 @@ def due(self, rollout_id: int, *, final: bool = False) -> bool: return should_cut(int(rollout_id), every=self.every, final=final) + def _skip(self, driver: Any, rollout_id: int, *, final: bool) -> bool: + """No cut at this safe point: not due, island RECOVERY_REQUIRED, or just + resumed from (or already cut at) this rollout.""" + if self.controller.recovery_required or not self.due(rollout_id, final=final): + return True + current = self.pointer() + return current is not None and int(current["next_rollout_id"]) == int(rollout_id) + + def _cut_id(self, rollout_id: int) -> str: + return f"r{int(rollout_id):06d}-{self.controller.incarnation['id']}" + + def _start(self, context: Any, *, background: bool) -> Any: + """Start the trainer part of a cut: a pending cut whose shards are written in + the background (trainers with ``start_cut``), else the finished cut (``None``).""" + start = getattr(self.trainer, "start_cut", None) + if background and callable(start): + if self.pointer() is None: + # the store's first cut: until its LATEST is committed the store holds a + # cut and no LATEST; the mark lets a restart drop it and start at 0 + from yeto.rl.engine.resume import mark_first_cut + + mark_first_cut(self.controller.checkpoint_store, context.cut_id) + return start(epoch=self._epoch(), context=context, background=True) + self.trainer.save_cut(epoch=self._epoch(), context=context) + return None + + def prepare(self, driver: Any, *, rollout_id: int, state: Any, final: bool = False) -> bool: + """E1: take the trainer part of the cut due at the safe point before rollout + ``rollout_id`` NOW, while the trainer is still resident, for the ``state`` the + driver is about to publish as ``rollout_id``. Saves the wake/sleep a cut of an + offloaded colocated trainer costs. The pointer and the store sync stay at the + safe point (:meth:`save`), which keeps the cut only if the publication names + the same policy; otherwise the prepared files are never pointed at (dropped). + Returns True when a cut was prepared. Raises like :meth:`save`.""" + import time + + self._prepared = None + self._reap_abandoned() + if driver.published_version is None or self._skip(driver, rollout_id, final=final): + return False + if not final and self._commit is not None and self._commit[0].is_alive(): + self._declined = int(rollout_id) # save() skips it if still busy, else _unprepared_wake + return False # save() skips this periodic cut (_commit_busy) + if int(state.policy_version) != int(rollout_id): + raise RoundCutError(f"prepared cut {rollout_id} for policy v{state.policy_version}") + self.join_commit(driver, wait=False) # the previous cut is complete before this one starts + cut_id = self._cut_id(rollout_id) + self._drop_unpointed(cut_id) + started = time.monotonic() + context = self.source.context(cut_id, pending=(int(rollout_id), state)) + pending = self._start(context, background=self._defer(final)) + self._prepared = _PreparedCut(rollout_id=int(rollout_id), context=context, + save_s=time.monotonic() - started, pending=pending) + return True + + def _commit_busy(self, driver: Any, rollout_id: int, *, final: bool) -> bool: + """E1: a periodic cut is skipped while the previous cut's store commit still + runs (a slow store, e.g. a Modal volume commit of 2-10 s, would otherwise hold + the round); the next safe point cuts again. The final cut waits for it.""" + if final or self._commit is None or not self._commit[0].is_alive(): + return False + self.skipped_busy += 1 + emit = getattr(driver, "emit", None) + if callable(emit): + emit("rl_cut_skipped", rollout_id=int(rollout_id), reason="previous cut still committing", + committing=self._commit[1].get("cut_id"), skipped_total=self.skipped_busy) + return True + + def _unprepared_wake(self, driver: Any, rollout_id: int, prepared: "_PreparedCut | None", *, + final: bool) -> bool: + """E1: skip a periodic cut whose prepare was declined (it met a running commit) + when it would wake the offloaded trainer + (~2 s on the Modal A10G sync run e1a7) and a committed cut already exists; the + next safe point cuts again. The final cut and the first cut are always taken.""" + declined, self._declined = self._declined, None + if final or declined != int(rollout_id) or prepared is not None: + return False + if self.pointer() is None or not getattr(driver, "_trainer_offloaded", False): + return False + self.skipped_busy += 1 + emit = getattr(driver, "emit", None) + if callable(emit): + emit("rl_cut_skipped", rollout_id=int(rollout_id), reason="not prepared before the trainer offloaded", + skipped_total=self.skipped_busy) + return True + + def _defer(self, final: bool) -> bool: + """Write the trainer shards and commit the store in the background: periodic + cuts only (the final cut is complete when the run ends), and not when a sample + buffer commits the same store during the round (no two commits at once). The + store's first cut is deferred too (E1: with ``--rl-cut-every`` N > 1 it falls in + the steady rounds); :meth:`_start` marks it first, since a store holding cuts but + no LATEST refuses to start unless they are all such marked first cuts.""" + return self.background_commit and not final and self.buffer_watermark is None + + def _drop_unpointed(self, cut_id: str) -> None: + import shutil + + from yeto.rl.engine.cut import cut_dir + + directory = cut_dir(self.root, cut_id) + if directory.exists(): + shutil.rmtree(directory) + + def join_commit(self, driver: Any = None, *, wait: bool = True) -> dict[str, Any] | None: + """Complete the previous background cut: wait for a running store commit, + then collect the pending cut's shards, commit its manifest, point the state + dir at it and start its store commit (``wait``: and wait for that too). Each + finished commit is reported (``rl_cut_committed``); on any failure the + previous pointer stays the committed one. Returns the last outcome (None: + nothing was pending).""" + result = self._join_running(driver) + deferred, self._deferred = self._deferred, None + if deferred is not None: + done = self._complete_deferred(deferred) + if not done.get("ok", True): + self._report(driver, done) + return done + self._commit_info = done + result = self._join_running(driver) if wait else done + return result + + def _join_running(self, driver: Any) -> dict[str, Any] | None: + if self._commit is None: + return None + thread, box = self._commit + thread.join() + self._commit = None + result = {**(self._commit_info or {}), **dict(box.get("result") or { + "ok": False, "error": "commit thread ended without a result"})} + self._commit_info = None + result.setdefault("cut_id", box.get("cut_id")) + self._report(driver, result) + return result + + @staticmethod + def _report(driver: Any, result: Mapping[str, Any]) -> None: + emit = getattr(driver, "emit", None) + if callable(emit): + emit("rl_cut_committed", **{k: result.get(k) for k in ( + "cut_id", "rollout_id", "ok", "finish_s", "commit_s", "pruned", "seq", "cut_bytes")}, + **({"writer": result["writer"]} if result.get("writer") else {}), + **({"error": result["error"]} if result.get("error") else {})) + + def _complete_deferred(self, deferred: dict[str, Any]) -> dict[str, Any]: + """Collect the pending cut's shards from the trainer (this thread: the trainer + is driven from here), then commit its manifest, point the state dir at it and + commit the store in the background (:meth:`_commit_in_background`).""" + import time + + info = deferred["info"] + started = time.monotonic() + try: + collect = getattr(self.trainer, "collect_cut", None) + if callable(collect): + commit_manifest = collect(deferred["pending"]) + else: + self.trainer.finish_cut(deferred["pending"]) + commit_manifest = None + except Exception as error: # noqa: BLE001 - the previous cut stays the pointed one + self._drop_unpointed(info["cut_id"]) + self._unstage(deferred.get("staged")) + return {"cut_id": info["cut_id"], "rollout_id": deferred["rollout_id"], "ok": False, + "error": repr(error)[:2000]} + self._commit_in_background(deferred, commit_manifest) + writer = getattr(self.trainer, "last_cut_writer", None) + return {"cut_id": info["cut_id"], "rollout_id": deferred["rollout_id"], + "collect_s": round(time.monotonic() - started, 3), **({"writer": writer} if writer else {})} + + def _commit_in_background(self, deferred: dict[str, Any], commit_manifest: Any) -> None: + """E1: off the round's critical path: the cut manifest (``commit_manifest``, None + when the trainer committed it already), the pointer with the cut record, the + state dir copied into the store, the store commit. Until the pointer moves the + previous cut stays the pointed one; a failure before that drops the new cut's + files. The outcome lands in the result box, reported by :meth:`join_commit`.""" + import threading + import time + + from yeto.rl.engine.resume import directory_bytes + + info, cut = deferred["info"], deferred["cut"] + cut_id = info["cut_id"] + box: dict[str, Any] = {"cut_id": cut_id} + + def run() -> None: + started = time.monotonic() + try: + if commit_manifest is not None: + commit_manifest() + saved = recovery.save(cut, store=self.controller.store, + write=lambda record: self._write_and_sync( + info, record, defer=True, staged=deferred.get("staged"))) + except Exception as error: # noqa: BLE001 - reported at the next join + self._drop_unpointed(cut_id) + box["result"] = {"ok": False, "error": repr(error)[:2000]} + return + finally: + self._unstage(deferred.get("staged")) + done = {"finish_s": round(time.monotonic() - started, 3), + "cut_bytes": directory_bytes(Path(self.root) / cut_id)} + if not saved["synced"]: + box["result"] = {**done, "ok": False, "error": ( + getattr(self.controller, "last_store_sync", None) or {}).get("error")} + return + box["result"] = {**done, **dict(self.controller.commit_checkpoint_store() or {})} + + thread = threading.Thread(target=run, name=f"yeto-cut-commit-{cut_id}", daemon=True) + self._commit = (thread, box) + thread.start() + def save(self, driver: Any, *, rollout_id: int, final: bool = False, trigger: str | None = None) -> dict[str, Any] | None: """Cut the trainer at the safe point before rollout ``rollout_id`` (= the published @@ -235,6 +527,12 @@ def save(self, driver: Any, *, rollout_id: int, final: bool = False, when nothing was cut (not due, no publication yet, already cut at this rollout, island RECOVERY_REQUIRED). Raises on a refused/failed cut (the caller logs it). + A periodic cut (:meth:`_defer`) only captures the state here: the shards are + written, the pointer moved and the store committed in the background, completed + by :meth:`join_commit` before the next cut (``store_commit`` "deferred"; the + outcome is the ``rl_cut_committed`` event). Until then the previous cut stays + the committed one. + ``trigger`` names a final cut's trigger (default ``MANUAL``; ``SCHEDULED`` for a stop the trainer scheduler ordered); a periodic cut takes none.""" import time @@ -244,30 +542,51 @@ def save(self, driver: Any, *, rollout_id: int, final: bool = False, if trigger is not None and (not final or trigger not in FINAL_TRIGGERS): raise RoundCutError(f"cut trigger {trigger!r} needs final=True and one of {FINAL_TRIGGERS}") ctl = self.controller - if ctl.recovery_required or driver.published_version is None: + prepared, self._prepared = self._prepared, None + if driver.published_version is None or self._skip(driver, rollout_id, final=final): + self._abandon(prepared) + return None + if self._commit_busy(driver, rollout_id, final=final): + # its trainer may still be writing the shards: collect them at the next + # cut instead of waiting for the writer here (E1: up to ~0.9 s per round) + if prepared is not None: + self._abandoned.append(prepared) return None - if not self.due(rollout_id, final=final): + self._reap_abandoned() + if self._unprepared_wake(driver, rollout_id, prepared, final=final): + self._abandon(prepared) return None if int(driver.published_version) != int(rollout_id): + self._abandon(prepared) raise RoundCutError(f"safe point {rollout_id} != published version {driver.published_version}") - current = self.pointer() - if current is not None and int(current["next_rollout_id"]) == int(rollout_id): - return None # just resumed from (or already cut) this point - cut_id = f"r{int(rollout_id):06d}-{ctl.incarnation['id']}" + self.join_commit(driver, wait=final) # the final cut leaves nothing running started = time.monotonic() - context = self.source.context(cut_id) - # colocated + offloaded trainer (1 GPU): its weights/optimizer are not on the - # device between publish and the next train step; wake it for the cut only. - offloaded = bool(getattr(driver, "_trainer_offloaded", False)) - if offloaded: - self.trainer.onload() - try: - self.trainer.save_cut(epoch=self._epoch(), context=context) - finally: - if offloaded: - self.trainer.offload() - save_s = time.monotonic() - started + context = self.source.context(self._cut_id(rollout_id)) progress = context.progress + offloaded = False + defer = self._defer(final) + used_prepared = prepared is not None and prepared.matches(int(rollout_id), context) + if used_prepared: + save_s, pending = prepared.save_s, prepared.pending # taken before the trainer offloaded + else: + self._abandon(prepared) + # a failed prepare or an earlier attempt left files under this id; no pointer + # names them (_skip), and cuts are immutable: drop them + self._drop_unpointed(context.cut_id) + # colocated + offloaded trainer (1 GPU): its weights/optimizer are not on the + # device between publish and the next train step; wake it for the cut only. + offloaded = bool(getattr(driver, "_trainer_offloaded", False)) + if offloaded: + self.trainer.onload() + try: + pending = self._start(context, background=defer) + finally: + if offloaded: + self.trainer.offload() + save_s = time.monotonic() - started + if pending is None: + defer = False # the trainer wrote the cut already (no start_cut) + cut_id = context.cut_id consumed_n, consumed_sha = self._consumed() next_lr = self.next_lr(driver) if callable(self.next_lr) else None info = { @@ -291,31 +610,70 @@ def save(self, driver: Any, *, rollout_id: int, final: bool = False, completed_groups=consumed_n, completed_groups_digest=consumed_sha, buffer_watermark=(self.buffer_watermark(int(rollout_id)) if callable(self.buffer_watermark) else None)) + extra = {"prepared_before_offload": used_prepared, "trainer_onloaded_for_cut": offloaded, + "save_s": round(save_s, 3)} + capture = pending.get("capture") if isinstance(pending, Mapping) else None + if capture: + # E1: per rank, device copies or the host path (and why: e.g. low device memory) + extra["cut_capture"] = capture + if defer: + # the state dir (ledger, journal) as it is at this cut: its store sync runs + # in the background after later rounds have changed it + staged = self._stage(cut_id) + self._deferred = {"info": info, "cut": cut, "pending": pending, "rollout_id": int(rollout_id), + "staged": staged} + return {**info, **extra, "store_commit": "deferred", "store_synced": None, + "total_s": round(time.monotonic() - started, 3)} + if pending is not None: + self.trainer.finish_cut(pending) saved = recovery.save(cut, store=ctl.store, write=lambda record: self._write_and_sync(info, record)) info[CUT_KEY] = saved[CUT_KEY] synced = saved["synced"] sync = dict(getattr(ctl, "last_store_sync", None) or {}) cut_bytes = directory_bytes(Path(self.root) / cut_id) - total_s = time.monotonic() - started + info.update(extra) info["store_synced"] = bool(synced) + info["store_commit"] = "inline" info["cut_bytes"] = cut_bytes - info["save_s"] = round(save_s, 3) info["store_copy_s"] = sync.get("copy_s") info["store_commit_s"] = sync.get("commit_s") - info["total_s"] = round(total_s, 3) + info["total_s"] = round(time.monotonic() - started, 3) info["write_bytes_per_s"] = (cut_bytes / save_s) if save_s > 0 else None info["pruned"] = sync.get("pruned") info["latest_seq"] = sync.get("seq") - info["trainer_onloaded_for_cut"] = offloaded if not synced: info["store_error"] = sync.get("error") return info - def _write_and_sync(self, info: dict[str, Any], cut_record: dict[str, Any]) -> dict[str, Any]: + def _reap_abandoned(self) -> None: + abandoned, self._abandoned = self._abandoned, [] + for prepared in abandoned: + self._abandon(prepared) + + def _abandon(self, prepared: "_PreparedCut | None") -> None: + """A prepared cut the safe point does not keep: wait for its writer, drop its files.""" + if prepared is None: + return + if prepared.pending is not None: + try: + self.trainer.finish_cut(prepared.pending) + except Exception: # noqa: BLE001 - dropped below either way + pass + self._drop_unpointed(prepared.context.cut_id) + + def _write_and_sync(self, info: dict[str, Any], cut_record: dict[str, Any], *, + defer: bool = False, staged: Path | None = None) -> dict[str, Any]: """The round cut's ``recovery.save`` writer: the pointer (with the cut record) - into the state dir, then the state dir into the store (which commits).""" - self._write_pointer({**info, CUT_KEY: cut_record}) - synced = self.controller.sync_checkpoint_store("round_cut", skip_if_recovery=True) + into the state dir, then the state dir into the store (which commits, unless + ``defer``: the caller commits in the background). ``staged``: the state dir as + staged at the cut (:meth:`_stage`) is the one synced, with the same pointer.""" + pointer = {**info, CUT_KEY: cut_record} + self._write_pointer(pointer) + kwargs: dict[str, Any] = {"defer_commit": True} if defer else {} + if staged is not None: + self._write_pointer(pointer, path=Path(staged) / POINTER) + kwargs["source_dir"] = Path(staged) + synced = self.controller.sync_checkpoint_store("round_cut", skip_if_recovery=True, **kwargs) # the store sync commits itself (snapshot layout: in write_snapshot; flat layout: # in the controller, outside its lock), also when it failed: recovery.save must not return {"synced": bool(synced), "committed": True} diff --git a/yeto/rl/engine/run_config.py b/yeto/rl/engine/run_config.py index 616e588a5..90abd7dc4 100644 --- a/yeto/rl/engine/run_config.py +++ b/yeto/rl/engine/run_config.py @@ -255,6 +255,9 @@ class TrainMemoryConfig: max_tokens_per_gpu: int | None = None # Share of the optimizer state kept and updated on the host (0 = none). optimizer_offload_fraction: float = 0.0 + # Megatron gradient-accumulation fusion (fused wgrad into main_grad; needs + # APEX/TE fused kernels). Off keeps the historical argv. + gradient_accumulation_fusion: bool = False @dataclass(frozen=True) @@ -1094,6 +1097,7 @@ def _train_memory(args) -> TrainMemoryConfig: activation_recompute=bool(getattr(args, "rl_activation_recompute", False)), max_tokens_per_gpu=int(max_tokens) if max_tokens is not None else None, optimizer_offload_fraction=offload, + gradient_accumulation_fusion=bool(getattr(args, "rl_gradient_accumulation_fusion", False)), ) diff --git a/yeto/rl/engine/run_flags.py b/yeto/rl/engine/run_flags.py index b4de91487..e90e06bc3 100644 --- a/yeto/rl/engine/run_flags.py +++ b/yeto/rl/engine/run_flags.py @@ -495,6 +495,13 @@ def rl_forward_pairs(*dests: str | tuple[str, str]) -> tuple[tuple[str, str], .. # 1.7 observation: per-round timeline labels (entry observe=...), off by default. shared_flag("--rl-observe-timeline", action="store_true", help="ports: record per-round timeline labels (rl-infra-spec 1.7); off by default"), + shared_flag("--rl-verify-every-round", action="store_true", + help="ports: hash the policy and read the engine weight checksums back at every " + "publication (default: on the first publication and on every --rl-cut-every-th " + "round, which includes every round that keeps a round cut)"), + shared_flag("--rl-profile-trainer", action="store_true", + help="ports, observation only (E1): per training step, print where the trainer " + "process spends its time (stack samples, GC pauses, CUDA allocator counters)"), shared_flag("--rl-resource-sample-interval", type=float, default=None, metavar="SECONDS", help="ports: NVML rl_resource_sample period forwarded to the learner " "(default: learner default, 60 s with --rl-observe-timeline; 0 = off)"), @@ -551,7 +558,9 @@ def rl_forward_pairs(*dests: str | tuple[str, str]) -> tuple[tuple[str, str], .. # rl-resume-from-checkpoint (S17 C4) shared_flag("--rl-cut-every", type=int, default=None, help="with --rl-checkpoint-store: keep a round cut every N rounds (plus the " - "last round and our own stop); default 1"), + "last round and our own stop); default 4 (1: every round). The publication " + "checks (policy hash, engine weight read-back) run at the same interval and " + "on the first publication; see --rl-verify-every-round"), shared_flag("--rl-cut-keep", type=int, default=None, help="with --rl-checkpoint-store and no --rl-elastic: keep the newest K cuts " "(default 2); older ones are deleted only after the new LATEST is committed"),