From 5a01732c71b059b5bba89ba24e16032718396428 Mon Sep 17 00:00:00 2001 From: michael <78113554+chowhsu@users.noreply.github.com> Date: Tue, 20 Oct 2026 14:29:32 +0000 Subject: [PATCH 01/55] tools/eval: add efficiency evaluation (bare vs yeto, verl and Miles) Eight arms per model preset ({bare,yeto} x {verl,miles} x {sync,async}) generated from one shared spec with a control fingerprint. Bare arms run the frameworks' own entry points (verl main_ppo / fully_async_main, Miles train.py / train_async.py) in the same images yeto uses; yeto arms run verl island_entry and the yeto launcher. Parsers normalize verl FileLogger, Miles perf lines and yeto tapes into per-round records (missing fields are marked, never estimated); the report gives yeto-vs-bare deltas with min-max range verdicts. Presets: Qwen3-0.6B (E0.1) and Qwen3-4B-Instruct-2507 (E0.2). CLI: eff-plan, eff-run, eff-parse, eff-report. Co-Authored-By: Claude Opus 5.5 --- tests/test_tools_eval_efficiency.py | 338 +++++++++++++++++++++ tools/eval/__main__.py | 93 ++++++ tools/eval/efficiency/__init__.py | 24 ++ tools/eval/efficiency/eff_prep.py | 65 ++++ tools/eval/efficiency/eff_reward.py | 39 +++ tools/eval/efficiency/metrics.py | 136 +++++++++ tools/eval/efficiency/modal_eff.py | 179 +++++++++++ tools/eval/efficiency/parse.py | 420 ++++++++++++++++++++++++++ tools/eval/efficiency/prices.json | 6 + tools/eval/efficiency/report.py | 205 +++++++++++++ tools/eval/efficiency/spec.py | 451 ++++++++++++++++++++++++++++ 11 files changed, 1956 insertions(+) create mode 100644 tests/test_tools_eval_efficiency.py create mode 100644 tools/eval/efficiency/__init__.py create mode 100644 tools/eval/efficiency/eff_prep.py create mode 100644 tools/eval/efficiency/eff_reward.py create mode 100644 tools/eval/efficiency/metrics.py create mode 100644 tools/eval/efficiency/modal_eff.py create mode 100644 tools/eval/efficiency/parse.py create mode 100644 tools/eval/efficiency/prices.json create mode 100644 tools/eval/efficiency/report.py create mode 100644 tools/eval/efficiency/spec.py diff --git a/tests/test_tools_eval_efficiency.py b/tests/test_tools_eval_efficiency.py new file mode 100644 index 00000000..a7ff78f6 --- /dev/null +++ b/tests/test_tools_eval_efficiency.py @@ -0,0 +1,338 @@ +"""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")) + 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") + + +# 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_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") + + +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)]) diff --git a/tools/eval/__main__.py b/tools/eval/__main__.py index 1d047988..32ed6571 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,71 @@ 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) + 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 +215,27 @@ 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.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 00000000..54197bef --- /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/eff_prep.py b/tools/eval/efficiency/eff_prep.py new file mode 100644 index 00000000..94fdee80 --- /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 00000000..4358028a --- /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 00000000..6fe298c0 --- /dev/null +++ b/tools/eval/efficiency/metrics.py @@ -0,0 +1,136 @@ +"""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 + for key, out in (("gen_tokens", "gen_tok_per_card_s"), ("train_tokens", "train_tok_per_card_s"), + ("samples", "samples_per_s")): + total = _sum_all(steady, key) + denom = (wall * (cards if key != "samples" else 1)) if wall else None + s[out] = round(total / denom, 3) if total is not None and denom else None + 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 00000000..0bcf5719 --- /dev/null +++ b/tools/eval/efficiency/modal_eff.py @@ -0,0 +1,179 @@ +"""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 +same image the yeto arm of that framework uses). 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": + image = modal_image(modal).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 00000000..167d0bab --- /dev/null +++ b/tools/eval/efficiency/parse.py @@ -0,0 +1,420 @@ +"""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) + for line in metrics_jsonl.splitlines(): + if not line.strip(): + continue + row = json.loads(line) + d, step = row.get("data") or {}, int(row["step"]) + rec = _record(step) + _set(p, rec, "t_end_unix", ends.get(step), "console step line (harness time prefix)") + gen, key = _first(d, ("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) + p.rounds.append(rec) + p.segments_raw.append({k: float(v) for k, v in d.items() if k.startswith("timing_s/") + and isinstance(v, (int, float))}) + p.rounds.sort(key=lambda r: r["round"]) + _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) -> 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).""" + 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"]) + 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(): + _set(p, rec, name, seg.get(task), f"{src} {task}") + 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"): + _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) + + +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), 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 "" + rows.append({"t_unix": t, "key": f"{container}:{index}", "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 "" + + +def _tape_events(raw: Path) -> list[dict[str, Any]]: + seen, rows = set(), [] + for path in sorted(raw.rglob("*.jsonl")): + if "verl-" in path.name: + 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) + parsed = parse_yeto(_tape_events(raw), mode=mode, framework_parsed=fw, + framework_round_offset=1 if plan["framework"] == "verl" else 0) \ + if plan["system"] == "yeto" else fw + 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": parse_nvml(nvml_text), "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 00000000..75edc4ee --- /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 00000000..38c6a9f0 --- /dev/null +++ b/tools/eval/efficiency/report.py @@ -0,0 +1,205 @@ +"""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") + 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())) + return rows + [""] diff --git a/tools/eval/efficiency/spec.py b/tools/eval/efficiency/spec.py new file mode 100644 index 00000000..be0e6a0a --- /dev/null +++ b/tools/eval/efficiency/spec.py @@ -0,0 +1,451 @@ +"""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) + + @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")} + 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", + "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", + "--apply-chat-template", "--rollout-shuffle", "--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(9216, spec.seq_len)), + "--attention-dropout", "0.0", "--hidden-dropout", "0.0", + "--accumulate-allreduce-grads-in-fp32", "--attention-softmax-in-fp32", "--attention-backend", "flash", + "--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" + + +# ====================================================================== yeto arms +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": + argv += ["--rl-max-policy-age", str(spec.max_policy_age)] + return argv + + +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, SPEC, 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"--modal-retries 0 --modal-cpu 8 --modal-memory 64 --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} " + "--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} " + # TIS (upper 2.0) in both modes, as in the bare arms; tolerates staleness 2 >= max_policy_age + f"--trust-remote-code {SPEC}") + 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}", + "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)) From 261ba2e1695783aff28df5a220dd66dccc4a5465 Mon Sep 17 00:00:00 2001 From: michael <78113554+chowhsu@users.noreply.github.com> Date: Tue, 20 Oct 2026 14:52:50 +0000 Subject: [PATCH 02/55] tools/eval: bare Miles trains with unfused attention (match yeto generic recipe; TE flash fails on sm86) Co-Authored-By: Claude Opus 5.5 --- tools/eval/efficiency/spec.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/tools/eval/efficiency/spec.py b/tools/eval/efficiency/spec.py index be0e6a0a..ea0f79cb 100644 --- a/tools/eval/efficiency/spec.py +++ b/tools/eval/efficiency/spec.py @@ -296,7 +296,10 @@ def miles_bare_argv(spec: EffSpec, mode: str) -> list[str]: "--recompute-granularity", "full", "--recompute-method", "uniform", "--recompute-num-layers", "1", "--use-dynamic-batch-size", "--max-tokens-per-gpu", str(max(9216, spec.seq_len)), "--attention-dropout", "0.0", "--hidden-dropout", "0.0", - "--accumulate-allreduce-grads-in-fp32", "--attention-softmax-in-fp32", "--attention-backend", "flash", + "--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), From c3c9d09d8184a2ea3c5b5e5dce12be005337bc49 Mon Sep 17 00:00:00 2001 From: michael <78113554+chowhsu@users.noreply.github.com> Date: Tue, 20 Oct 2026 15:20:51 +0000 Subject: [PATCH 03/55] tools/eval: yeto+verl async arm passes the staleness-2 TIS algorithm spec Co-Authored-By: Claude Opus 5.5 --- tests/test_tools_eval_efficiency.py | 4 ++++ tools/eval/efficiency/spec.py | 7 ++++++- 2 files changed, 10 insertions(+), 1 deletion(-) diff --git a/tests/test_tools_eval_efficiency.py b/tests/test_tools_eval_efficiency.py index a7ff78f6..b1ca57df 100644 --- a/tests/test_tools_eval_efficiency.py +++ b/tests/test_tools_eval_efficiency.py @@ -78,6 +78,10 @@ def test_bare_entry_points_follow_the_official_recipes(): 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") diff --git a/tools/eval/efficiency/spec.py b/tools/eval/efficiency/spec.py index ea0f79cb..b10acc48 100644 --- a/tools/eval/efficiency/spec.py +++ b/tools/eval/efficiency/spec.py @@ -322,6 +322,8 @@ def miles_bare_script(mode: str) -> str: # ====================================================================== 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 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", @@ -334,7 +336,10 @@ def yeto_verl_argv(spec: EffSpec, mode: str, tape: str) -> list[str]: if spec.chat_template_kwargs: argv += ["--apply-chat-template-kwargs", json.dumps(spec.chat_template_kwargs)] if mode == "async": - argv += ["--rl-max-policy-age", str(spec.max_policy_age)] + # 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 From 63dc064482fa86b0ae42a2518433150471a6e75d Mon Sep 17 00:00:00 2001 From: michael <78113554+chowhsu@users.noreply.github.com> Date: Tue, 20 Oct 2026 16:02:58 +0000 Subject: [PATCH 04/55] tools/eval: bare verl arms run the unpatched fork; async uses merged LoRA and /dev/shm NCCL host buffers The verl image applies yeto's patch_verl, whose sender hook imports yeto, so the bare arms now build the image without it. The fork's disaggregated checkpoint engine drops the PEFT config of unmerged LoRA params; the bare async arm sends merged weights instead. On Modal the trainer allocated NCCL host buffers with cuMem while the receiver fell back to /dev/shm and failed to attach (NCCL_DEBUG=INFO run); NCCL_CUMEM_HOST_ENABLE=0 keeps both on /dev/shm. Co-Authored-By: Claude Opus 5.5 --- tests/test_tools_eval_efficiency.py | 6 ++++++ tools/eval/efficiency/modal_eff.py | 7 +++++-- tools/eval/efficiency/spec.py | 8 ++++++++ yeto/rl/adapters/verl/image.py | 15 +++++++++------ 4 files changed, 28 insertions(+), 8 deletions(-) diff --git a/tests/test_tools_eval_efficiency.py b/tests/test_tools_eval_efficiency.py index b1ca57df..44356fbf 100644 --- a/tests/test_tools_eval_efficiency.py +++ b/tests/test_tools_eval_efficiency.py @@ -340,3 +340,9 @@ def test_cli_plan_and_unregistered_run(tmp_path, monkeypatch, capsys): (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") diff --git a/tools/eval/efficiency/modal_eff.py b/tools/eval/efficiency/modal_eff.py index 0bcf5719..d67dd09e 100644 --- a/tools/eval/efficiency/modal_eff.py +++ b/tools/eval/efficiency/modal_eff.py @@ -2,7 +2,8 @@ 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 -same image the yeto arm of that framework uses). The container: +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); @@ -111,7 +112,9 @@ def build_app(run: str, plan: dict): from yeto.rl.adapters.verl.image import modal_image if plan["image"] == "verl": - image = modal_image(modal).env({"PYTHONUNBUFFERED": "1"}) + # 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: diff --git a/tools/eval/efficiency/spec.py b/tools/eval/efficiency/spec.py index b10acc48..5b158ae5 100644 --- a/tools/eval/efficiency/spec.py +++ b/tools/eval/efficiency/spec.py @@ -242,6 +242,10 @@ def verl_bare_overrides(spec: EffSpec, mode: str, run_id: str) -> list[str]: "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}", @@ -381,6 +385,10 @@ def bare_script(spec: EffSpec, arm: Arm, run_id: str) -> str: 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: diff --git a/yeto/rl/adapters/verl/image.py b/yeto/rl/adapters/verl/image.py index b096d1b1..29a52282 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()) From 052c31e167f946f323dd4476521dc8cfa923ad66 Mon Sep 17 00:00:00 2001 From: michael <78113554+chowhsu@users.noreply.github.com> Date: Tue, 20 Oct 2026 16:36:04 +0000 Subject: [PATCH 05/55] tools/eval: parse bare verl fully_async rows by training step; rates over rounds that report the value The fork logs a step's training metrics one parameter version late and never logs the last step's, so rows are keyed by training/global_step and the step without metrics keeps its trained batch from the config. Throughput rates use the steady rounds that report the value; the report lists partial counts. Co-Authored-By: Claude Opus 5.5 --- tests/test_tools_eval_efficiency.py | 18 ++++++++++++++++++ tools/eval/efficiency/metrics.py | 11 ++++++++--- tools/eval/efficiency/parse.py | 26 ++++++++++++++++++++------ tools/eval/efficiency/report.py | 4 ++++ 4 files changed, 50 insertions(+), 9 deletions(-) diff --git a/tests/test_tools_eval_efficiency.py b/tests/test_tools_eval_efficiency.py index 44356fbf..cdd9df04 100644 --- a/tests/test_tools_eval_efficiency.py +++ b/tests/test_tools_eval_efficiency.py @@ -346,3 +346,21 @@ 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] diff --git a/tools/eval/efficiency/metrics.py b/tools/eval/efficiency/metrics.py index 6fe298c0..0df3bfa2 100644 --- a/tools/eval/efficiency/metrics.py +++ b/tools/eval/efficiency/metrics.py @@ -58,11 +58,16 @@ def summarize(parsed: dict[str, Any], *, prices_path: Path = PRICES) -> dict[str "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")): - total = _sum_all(steady, key) - denom = (wall * (cards if key != "samples" else 1)) if wall else None - s[out] = round(total / denom, 3) if total is not None and denom else None + # 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 diff --git a/tools/eval/efficiency/parse.py b/tools/eval/efficiency/parse.py index 167d0bab..066f8d6f 100644 --- a/tools/eval/efficiency/parse.py +++ b/tools/eval/efficiency/parse.py @@ -128,13 +128,24 @@ def parse_verl(metrics_jsonl: str, log_text: str = "", *, samples_per_round: int """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"]) - rec = _record(step) - _set(p, rec, "t_end_unix", ends.get(step), "console step line (harness time prefix)") + 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)") gen, key = _first(d, ("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] @@ -161,10 +172,13 @@ def parse_verl(metrics_jsonl: str, log_text: str = "", *, samples_per_round: int _set(p, rec, "dropped_samples", dropped, key) generated, key = _first(d, ("fully_async/count/total_generated_samples",)) _set(p, rec, "generated_samples", generated, key) - p.rounds.append(rec) - p.segments_raw.append({k: float(v) for k, v in d.items() if k.startswith("timing_s/") - and isinstance(v, (int, float))}) - p.rounds.sort(key=lambda r: r["round"]) + 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": diff --git a/tools/eval/efficiency/report.py b/tools/eval/efficiency/report.py index 38c6a9f0..be571fe7 100644 --- a/tools/eval/efficiency/report.py +++ b/tools/eval/efficiency/report.py @@ -202,4 +202,8 @@ def _missing(by_arm: dict[str, list[dict]]) -> list[str]: 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 + [""] From c9245e5825193e336aa910e11cc11745f6087aab Mon Sep 17 00:00:00 2001 From: michael <78113554+chowhsu@users.noreply.github.com> Date: Tue, 20 Oct 2026 16:38:44 +0000 Subject: [PATCH 06/55] tools/eval: bare Miles arms do not shuffle prompts (yeto's Miles argv does not) Co-Authored-By: Claude Opus 5.5 --- tests/test_tools_eval_efficiency.py | 5 +++++ tools/eval/efficiency/spec.py | 4 +++- 2 files changed, 8 insertions(+), 1 deletion(-) diff --git a/tests/test_tools_eval_efficiency.py b/tests/test_tools_eval_efficiency.py index cdd9df04..2cd54e26 100644 --- a/tests/test_tools_eval_efficiency.py +++ b/tests/test_tools_eval_efficiency.py @@ -364,3 +364,8 @@ def test_parse_bare_verl_async_rows_lag_one_version(): 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] + + +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") diff --git a/tools/eval/efficiency/spec.py b/tools/eval/efficiency/spec.py index 5b158ae5..bf40111a 100644 --- a/tools/eval/efficiency/spec.py +++ b/tools/eval/efficiency/spec.py @@ -275,7 +275,9 @@ def miles_bare_argv(spec: EffSpec, mode: str) -> list[str]: "--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", - "--apply-chat-template", "--rollout-shuffle", "--rollout-seed", str(spec.seed), + # 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), From 73442ad84a79d71c881bc14358175f3fadf6fad3 Mon Sep 17 00:00:00 2001 From: michael <78113554+chowhsu@users.noreply.github.com> Date: Tue, 20 Oct 2026 17:01:51 +0000 Subject: [PATCH 07/55] verl adapter: fully_async keeps NCCL host buffers on /dev/shm in every Ray worker On Modal one weight-sync process could allocate NCCL host buffers with cuMem while its peer fell back to /dev/shm and then failed to attach (eff01-yeto-verl-async-e; NCCL_DEBUG=INFO in eff01-bare-verl-async-d). NCCL_CUMEM_HOST_ENABLE=0 unless the caller sets it. Co-Authored-By: Claude Opus 5.5 --- tests/test_rl_verl_adapter.py | 24 ++++++++++++++++++++++++ yeto/rl/adapters/verl/verl_main.py | 19 ++++++++++++++----- 2 files changed, 38 insertions(+), 5 deletions(-) diff --git a/tests/test_rl_verl_adapter.py b/tests/test_rl_verl_adapter.py index 96c58845..1a417dcd 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/yeto/rl/adapters/verl/verl_main.py b/yeto/rl/adapters/verl/verl_main.py index 106a0093..303ad942 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 From 0122f4e7aab4b445f9251374a9f039df6e4f1a95 Mon Sep 17 00:00:00 2001 From: michael <78113554+chowhsu@users.noreply.github.com> Date: Tue, 20 Oct 2026 17:23:56 +0000 Subject: [PATCH 08/55] tools/eval: read the yeto+Miles launcher exit code into the parsed run Co-Authored-By: Claude Opus 5.5 --- tools/eval/efficiency/parse.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/tools/eval/efficiency/parse.py b/tools/eval/efficiency/parse.py index 066f8d6f..e7cbef96 100644 --- a/tools/eval/efficiency/parse.py +++ b/tools/eval/efficiency/parse.py @@ -417,6 +417,8 @@ def parse_run(run_dir: Path) -> dict[str, Any]: parsed = parse_yeto(_tape_events(raw), mode=mode, framework_parsed=fw, framework_round_offset=1 if plan["framework"] == "verl" else 0) \ 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, From e0f833dccd1bfed3524756fcf2c71183f6443e19 Mon Sep 17 00:00:00 2001 From: michael <78113554+chowhsu@users.noreply.github.com> Date: Tue, 20 Oct 2026 17:34:51 +0000 Subject: [PATCH 09/55] miles adapter: opt-in gradient-accumulation fusion; generic recipe trains without dropout --rl-gradient-accumulation-fusion lets a run keep Megatron's fused wgrad (the Miles recipes' default); off keeps the historical argv. The generic recipe passed no dropout, so Megatron's 0.1 applied during RL training while the Miles recipes pass 0.0; it now passes 0.0 as the Qwen3.5 recipe does. Co-Authored-By: Claude Opus 5.5 --- tests/test_rl_miles_adapter_config.py | 18 ++++++++++++++++++ yeto/launch/rl_island_task.py | 2 ++ yeto/rl/adapters/miles/config.py | 8 +++++++- yeto/rl/adapters/miles/run_flags.py | 6 ++++++ yeto/rl/engine/run_config.py | 4 ++++ 5 files changed, 37 insertions(+), 1 deletion(-) diff --git a/tests/test_rl_miles_adapter_config.py b/tests/test_rl_miles_adapter_config.py index 503fe2e2..7cbe30d9 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/yeto/launch/rl_island_task.py b/yeto/launch/rl_island_task.py index 1185f4c9..ce1ceb50 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/rl/adapters/miles/config.py b/yeto/rl/adapters/miles/config.py index dadbb341..1af82e71 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/run_flags.py b/yeto/rl/adapters/miles/run_flags.py index 8a9f5202..6e4a5af0 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/engine/run_config.py b/yeto/rl/engine/run_config.py index 616e588a..90abd7dc 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)), ) From 6a4263e24d62eabbe7cc991a5422375e340ac246 Mon Sep 17 00:00:00 2001 From: michael <78113554+chowhsu@users.noreply.github.com> Date: Tue, 20 Oct 2026 17:35:46 +0000 Subject: [PATCH 10/55] tools/eval: yeto+Miles arm takes the bare Miles recipe settings MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit No deterministic inference, constant LR, full recompute, dynamic batching, gradient-accumulation fusion, and an algorithm spec with the recipe's clip range (0.2 / 0.28); S26 GOAL §6 makes the bare framework's settings the reference. Co-Authored-By: Claude Opus 5.5 --- tools/eval/efficiency/algo-eff-miles.json | 53 +++++++++++++++++++++++ tools/eval/efficiency/spec.py | 20 +++++++-- 2 files changed, 69 insertions(+), 4 deletions(-) create mode 100644 tools/eval/efficiency/algo-eff-miles.json diff --git a/tools/eval/efficiency/algo-eff-miles.json b/tools/eval/efficiency/algo-eff-miles.json new file mode 100644 index 00000000..f45d2e97 --- /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/spec.py b/tools/eval/efficiency/spec.py index bf40111a..1d1de74b 100644 --- a/tools/eval/efficiency/spec.py +++ b/tools/eval/efficiency/spec.py @@ -300,7 +300,7 @@ def miles_bare_argv(spec: EffSpec, mode: str) -> list[str]: "--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(9216, spec.seq_len)), + "--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 @@ -327,9 +327,16 @@ 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", @@ -351,7 +358,7 @@ def yeto_verl_argv(spec: EffSpec, mode: str, tape: str) -> list[str]: 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, SPEC, STORE_VOLUME, TAPE_VOLUME + from ..miles import DATA, STORE_VOLUME, TAPE_VOLUME card = {"A10G": "a10g", "H100!": "h100", "L4": "l4"}[spec.gpu] c = cards(mode) @@ -366,8 +373,13 @@ def yeto_miles_launcher_args(spec: EffSpec, mode: str, run_id: str) -> str: 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} " - # TIS (upper 2.0) in both modes, as in the bare arms; tolerates staleness 2 >= max_policy_age - f"--trust-remote-code {SPEC}") + # 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 kwargs: args += f" --apply-chat-template-kwargs {shlex.quote(kwargs)}" if mode == "async": From ff893bcd146969ed20458cbeeb72b199c93e55c9 Mon Sep 17 00:00:00 2001 From: michael <78113554+chowhsu@users.noreply.github.com> Date: Tue, 20 Oct 2026 17:38:14 +0000 Subject: [PATCH 11/55] tools/eval: read only the run's own tapes (skip the repo copy and the store) miles_local.sh pulls a copy of the repo whose test fixtures are tapes of other runs; their rl_local_round rows replaced the run's rewards (all four yeto+Miles runs showed the same fixture curve with reward 0 in round 0). Co-Authored-By: Claude Opus 5.5 --- tests/test_tools_eval_efficiency.py | 8 ++++++++ tools/eval/efficiency/parse.py | 7 ++++++- 2 files changed, 14 insertions(+), 1 deletion(-) diff --git a/tests/test_tools_eval_efficiency.py b/tests/test_tools_eval_efficiency.py index 2cd54e26..c5f19024 100644 --- a/tests/test_tools_eval_efficiency.py +++ b/tests/test_tools_eval_efficiency.py @@ -369,3 +369,11 @@ def test_parse_bare_verl_async_rows_lag_one_version(): 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/efficiency/parse.py b/tools/eval/efficiency/parse.py index e7cbef96..59f9ee15 100644 --- a/tools/eval/efficiency/parse.py +++ b/tools/eval/efficiency/parse.py @@ -383,10 +383,15 @@ 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")): - if "verl-" in path.name: + # 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 From ba179531cb1012c47464825ff39944f28371aafa Mon Sep 17 00:00:00 2001 From: michael <78113554+chowhsu@users.noreply.github.com> Date: Tue, 20 Oct 2026 17:39:05 +0000 Subject: [PATCH 12/55] tools/eval: verl fully_async generation time from the rollouter, else missing The trainer's timing_s/gen and yeto's driver generate phase only pull finished samples from the rollouter queue. Bare verl logs the rollouter's active time per parameter version; yeto logs none per round, so its value stays missing. Co-Authored-By: Claude Opus 5.5 --- tests/test_tools_eval_efficiency.py | 1 + tools/eval/efficiency/parse.py | 18 +++++++++++++----- 2 files changed, 14 insertions(+), 5 deletions(-) diff --git a/tests/test_tools_eval_efficiency.py b/tests/test_tools_eval_efficiency.py index c5f19024..2cd3821c 100644 --- a/tests/test_tools_eval_efficiency.py +++ b/tests/test_tools_eval_efficiency.py @@ -364,6 +364,7 @@ def test_parse_bare_verl_async_rows_lag_one_version(): 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(): diff --git a/tools/eval/efficiency/parse.py b/tools/eval/efficiency/parse.py index 59f9ee15..f3b2f8ee 100644 --- a/tools/eval/efficiency/parse.py +++ b/tools/eval/efficiency/parse.py @@ -146,7 +146,9 @@ def parse_verl(metrics_jsonl: str, log_text: str = "", *, samples_per_round: int rec = recs.setdefault(rnd, _record(rnd)) if rnd == step: _set(p, rec, "t_end_unix", ends.get(step), "console step line (harness time prefix)") - gen, key = _first(d, ("timing_s/gen",)) + # 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: @@ -283,7 +285,7 @@ def tape_rows(text: str) -> list[dict[str, Any]]: def parse_yeto(events: list[dict[str, Any]], *, mode: str, framework_parsed: Parsed | None = None, - framework_round_offset: int = 0) -> Parsed: + framework_round_offset: int = 0, skip: frozenset[str] = frozenset()) -> 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``: @@ -317,7 +319,8 @@ def parse_yeto(events: list[dict[str, Any]], *, mode: str, framework_parsed: Par 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(): - _set(p, rec, name, seg.get(task), f"{src} {task}") + if name not in skip: + _set(p, rec, name, seg.get(task), f"{src} {task}") 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") @@ -328,7 +331,8 @@ def parse_yeto(events: list[dict[str, Any]], *, mode: str, framework_parsed: Par _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"): + 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()}) @@ -419,7 +423,11 @@ def parse_run(run_dir: Path) -> dict[str, Any]: else: log = _read(arm_dir / "main.log") or _read(raw / "launch.ts.log") fw = parse_miles(log, rollout_cards=c["rollout"], mode=mode) - parsed = parse_yeto(_tape_events(raw), mode=mode, framework_parsed=fw, + # yeto+verl fully_async: the driver's generate phase (and verl's timing_s/gen) only + # pull finished samples from the rollouter queue, and yeto logs no per-round + # rollouter time, so generation time stays missing + skip = frozenset({"gen_s"}) if plan["framework"] == "verl" and mode == "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) \ if plan["system"] == "yeto" else fw if meta.get("rc") is None and _read(raw / "rc").strip().lstrip("-").isdigit(): From a4b11047ceaf856ad7f2fc4af11dae8b35cc4cc1 Mon Sep 17 00:00:00 2001 From: michael <78113554+chowhsu@users.noreply.github.com> Date: Wed, 21 Oct 2026 00:39:19 +0000 Subject: [PATCH 13/55] eff eval: final round cut only for yeto+Miles; verl async generation from the rollouter E0.1 found yeto+Miles 77 % (sync) and 252 % (async) slower than bare Miles. The tape shows the gap is the per-round round cut: the arm sets a checkpoint store, so a cut runs after every publish (5.7-12.9 s: trainer wake, export, sleep, Modal volume commit), and in fully_async mode the cut also stops Miles' background worker, so generation never runs ahead of training. The bare arms keep no checkpoint; cut only at the end (--rl-cut-every rounds). yeto+verl async generation time was missing: the tape's rl_publication_check.rollouter_timing carries the rollouter's active time per policy version, the same metric the bare arm reports. Read it (version v closes round v - 1). Co-Authored-By: Claude Opus 5.5 --- tests/test_tools_eval_efficiency.py | 18 ++++++++++++++ tools/eval/efficiency/parse.py | 37 ++++++++++++++++++++++++----- tools/eval/efficiency/spec.py | 4 ++++ 3 files changed, 53 insertions(+), 6 deletions(-) diff --git a/tests/test_tools_eval_efficiency.py b/tests/test_tools_eval_efficiency.py index 2cd3821c..9be2abe0 100644 --- a/tests/test_tools_eval_efficiency.py +++ b/tests/test_tools_eval_efficiency.py @@ -85,6 +85,8 @@ def test_yeto_arms_use_the_same_policy_age_and_partition(): 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") + # 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") # overrides yeto's verl translation sets and the bare recipe leaves at the default (or the reverse); @@ -235,6 +237,22 @@ def test_parse_yeto_driver_phase_fallback(): 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" diff --git a/tools/eval/efficiency/parse.py b/tools/eval/efficiency/parse.py index f3b2f8ee..98cd55b9 100644 --- a/tools/eval/efficiency/parse.py +++ b/tools/eval/efficiency/parse.py @@ -285,12 +285,18 @@ def tape_rows(text: str) -> list[dict[str, Any]]: def parse_yeto(events: list[dict[str, Any]], *, mode: str, framework_parsed: Parsed | None = None, - framework_round_offset: int = 0, skip: frozenset[str] = frozenset()) -> Parsed: + 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).""" + 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: @@ -309,6 +315,7 @@ def parse_yeto(events: list[dict[str, Any]], *, mode: str, framework_parsed: Par 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): @@ -321,6 +328,8 @@ def parse_yeto(events: list[dict[str, Any]], *, mode: str, framework_parsed: Par 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") @@ -343,6 +352,20 @@ def parse_yeto(events: list[dict[str, Any]], *, mode: str, framework_parsed: Par 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).""" @@ -424,11 +447,13 @@ def parse_run(run_dir: Path) -> dict[str, Any]: 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, and yeto logs no per-round - # rollouter time, so generation time stays missing - skip = frozenset({"gen_s"}) if plan["framework"] == "verl" and mode == "async" else frozenset() + # 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) \ + 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 diff --git a/tools/eval/efficiency/spec.py b/tools/eval/efficiency/spec.py index 1d1de74b..56d29d45 100644 --- a/tools/eval/efficiency/spec.py +++ b/tools/eval/efficiency/spec.py @@ -368,6 +368,10 @@ def yeto_miles_launcher_args(spec: EffSpec, mode: str, run_id: str) -> str: 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} " + # Only the final round cut: the bare arms save no checkpoint. A periodic cut + # (default every round) is synchronous (~6-9 s on A10G) and stops Miles' + # fully_async worker, so the async arm would generate nothing ahead (E1). + f"--rl-cut-every {spec.rounds} " "--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} " From 77ce9cc3308fe8c8580a40f166ca1f475382c996 Mon Sep 17 00:00:00 2001 From: michael <78113554+chowhsu@users.noreply.github.com> Date: Wed, 21 Oct 2026 00:45:21 +0000 Subject: [PATCH 14/55] verl: hash the LoRA on the trainer for the single-island publish fast path VerlPolicyState.export_digest returns a trainer-resident handle: the workers validate the LoRA like canonical_state, compute the policy digests and the bf16 send checksum, and only those cross Ray. Both verl publishers take the send checksum from the handle; other states are hashed on the driver as before. materialize falls back to the full export; YETO_RL_PUBLISH_FASTPATH=0 keeps the old path. Tests pin hash and checksum equality with the full export. Co-Authored-By: Claude Opus 5.5 --- tests/test_rl_verl_buffer_trainer.py | 6 ++ tests/test_rl_verl_fully_async_64b.py | 22 +++++ tests/test_rl_verl_publish_fastpath.py | 87 ++++++++++++++++++ yeto/rl/adapters/verl/fully_async_ports.py | 13 +-- yeto/rl/adapters/verl/fully_async_runner.py | 5 ++ yeto/rl/adapters/verl/ports_impl.py | 98 ++++++++++++++++++++- 6 files changed, 222 insertions(+), 9 deletions(-) create mode 100644 tests/test_rl_verl_publish_fastpath.py diff --git a/tests/test_rl_verl_buffer_trainer.py b/tests/test_rl_verl_buffer_trainer.py index 58dea824..b37ad8a5 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 145cd8a2..2ab12db4 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 00000000..4428a65f --- /dev/null +++ b/tests/test_rl_verl_publish_fastpath.py @@ -0,0 +1,87 @@ +"""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 diff --git a/yeto/rl/adapters/verl/fully_async_ports.py b/yeto/rl/adapters/verl/fully_async_ports.py index 061fb920..f79505f9 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 900ccf6d..b3273048 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/ports_impl.py b/yeto/rl/adapters/verl/ports_impl.py index 4952a73c..c7f32936 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 @@ -75,6 +79,41 @@ def remote_export_lora(worker_dict): 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 yeto.rl.core import canonical_state + from yeto.rl.engine.policy_digest import digest_canonical_tensors + + started = time.monotonic() + state = canonical_state(0, tensors, base_model_revision=base_model_revision, + lora_config_hash=lora_config_hash, layout_hash=layout_hash, + expected_specs=expected_specs) + 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) + return {"digest": digest.to_wire(), "layout_hash": state.layout_hash, + "lora_checksum": pub.lora_checksum(dict(state.tensors)), + "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.""" + import torch.distributed as dist + + tensors, _info = remote_export_lora(worker_dict) + if dist.is_available() and dist.is_initialized() and dist.get_rank() != 0: + return None + return digest_lora_tensors(tensors, **kwargs) + + def remote_apply_lora(worker_dict, tensors, reset_optimizer: bool): import torch @@ -161,6 +200,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 +431,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 +465,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 +535,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 +561,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}"), From 74016eb0c434e123e436364dbd9b5b63c77df4b8 Mon Sep 17 00:00:00 2001 From: michael <78113554+chowhsu@users.noreply.github.com> Date: Wed, 21 Oct 2026 00:50:47 +0000 Subject: [PATCH 15/55] round cut: write before the trainer offloads; commit the store in the background E0.1 measured a periodic round cut at 5.7-12.9 s per round on the critical path of a colocated Miles island: wake the offloaded trainer (0.7 s), export (1.8 s), sleep it again (1.0 s), then a synchronous store commit (1.5-8 s). - prepare(): when a cut is due at the next safe point, the driver writes the trainer part right before it offloads a colocated trainer for the publication (phase "checkpoint"), for the state it is about to publish. The safe point keeps it only if its own context names the same cut (id, progress incl. policy hash, data cursor, ledger, outer token, in-flight); otherwise the unpointed files are dropped and the cut is taken as before. - Background commit: a periodic cut writes the snapshot and LATEST, then a thread commits and prunes (ResumeController.commit_checkpoint_store). The next cut joins it first and reports rl_cut_committed. Pruning only follows a successful commit, so a failed or lost commit leaves the previous cut as the committed one. The final cut, and a run with a sample buffer that commits the same store, stay inline. YETO_RL_CUT_BACKGROUND_COMMIT=0 turns it off. Co-Authored-By: Claude Opus 5.5 --- tests/test_rl_engine_driver.py | 34 ++++ tests/test_rl_resume_from_checkpoint.py | 133 +++++++++++++++- yeto/rl/engine/bridges.py | 21 ++- yeto/rl/engine/driver.py | 7 + yeto/rl/engine/resume.py | 53 +++++-- yeto/rl/engine/round_cut.py | 196 ++++++++++++++++++++---- 6 files changed, 401 insertions(+), 43 deletions(-) diff --git a/tests/test_rl_engine_driver.py b/tests/test_rl_engine_driver.py index cf48504d..738529b8 100644 --- a/tests/test_rl_engine_driver.py +++ b/tests/test_rl_engine_driver.py @@ -964,3 +964,37 @@ 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"] diff --git a/tests/test_rl_resume_from_checkpoint.py b/tests/test_rl_resume_from_checkpoint.py index cf6b44a4..2b64ffe2 100644 --- a/tests/test_rl_resume_from_checkpoint.py +++ b/tests/test_rl_resume_from_checkpoint.py @@ -193,8 +193,12 @@ def test_save_kill_resume_without_elastic_continues_bitwise(tmp_path): 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] + # E1: a periodic cut commits the store in the background; LATEST is written already + assert info["store_synced"] is None and info["store_commit"] == "background" + 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 cp_a.join_commit(driver_a)["ok"] is True + assert driver_a.events[-1][0] == "rl_cut_committed" and driver_a.events[-1][1]["cut_id"] == info["cut_id"] assert vol.commits >= 1 and rs.read_latest(store.root)["cut_id"] == info["cut_id"] rank_b = make_rank(9) @@ -541,3 +545,130 @@ 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 commit +class _CountingTrainer: + """The real MilesTrainerGroup cut path, counting save/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 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 _prepared_island(tmp_path, store, **kw): + """Two trained rounds, policy v1 published, v2 about to be published by an offloaded + colocated trainer.""" + rank = make_rank(0) + g = torch.Generator().manual_seed(3) + for _ in range(2): + train_step(rank, torch.randn(3, 6, generator=g)) + ctl = _rctl(tmp_path, "a", store) + driver = _Driver(rank, published=1, local_step=2) + driver.last_applied_lrs = [5e-6] + driver._trainer_offloaded = True + trainer = _CountingTrainer(_trainer(rank)) + return rank, driver, trainer, _cp(ctl, driver, trainer, _Ledger(2), **kw) + + +def test_cut_prepared_before_offload_is_kept_without_waking_the_trainer(tmp_path): + store = _store(tmp_path / "store", _Vol()) + rank, driver, trainer, cp = _prepared_island(tmp_path, store) + 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 == ["save_cut"] # no wake/sleep for the cut + assert info["prepared_before_offload"] is True and info["trainer_onloaded_for_cut"] is False + assert cp.join_commit(driver)["ok"] is True + # the prepared cut is a complete cut: a fresh island resumes from it bitwise + 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)) + assert cp_b.resume(driver_b)["next_rollout_id"] == 2 + for n, v in params(rank).items(): + assert torch.equal(v, params(rank_b)[n]) + + +def test_prepared_cut_for_another_policy_is_dropped_and_cut_again(tmp_path): + store = _store(tmp_path / "store", _Vol()) + rank, driver, trainer, cp = _prepared_island(tmp_path, store) + 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 == ["save_cut", "onload", "save_cut", "offload"] + assert info["prepared_before_offload"] is False and info["policy_hash"] == _hash(rank) + cp.join_commit(driver) + 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)) + assert cp_b.resume(driver_b)["policy_hash"] == _hash(rank) + + +def test_final_cut_commits_inline(tmp_path): + store = _store(tmp_path / "store", _Vol()) + _rank, _ctl_a, driver, cp = _trained_island(tmp_path, store, 2) + info = cp.save(driver, rollout_id=2, final=True) + assert info["store_commit"] == "inline" and info["store_synced"] is True + assert 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 = make_rank(0) + g = torch.Generator().manual_seed(3) + ctl = _rctl(tmp_path, "a", store, keep=1) + driver = _Driver(rank, published=0, local_step=0) + driver.last_applied_lrs = [5e-6] + cp = _cp(ctl, driver, _trainer(rank), _Ledger(2)) + train_step(rank, torch.randn(3, 6, generator=g)) + driver.published_version, driver.local_step = 1, 1 + first = cp.save(driver, rollout_id=1) + assert cp.join_commit(driver)["ok"] is True + train_step(rank, torch.randn(3, 6, generator=g)) + driver.published_version, driver.local_step = 2, 2 + 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_commit_can_be_turned_off(tmp_path, monkeypatch): + monkeypatch.setenv("YETO_RL_CUT_BACKGROUND_COMMIT", "0") + store = _store(tmp_path / "store", _Vol()) + _rank, _ctl_a, driver, cp = _trained_island(tmp_path, store, 2) + info = cp.save(driver, rollout_id=2) + assert info["store_commit"] == "inline" and info["store_synced"] is True diff --git a/yeto/rl/engine/bridges.py b/yeto/rl/engine/bridges.py index b1f955d3..47cb3b90 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 d815a0a7..f27ea005 100644 --- a/yeto/rl/engine/driver.py +++ b/yeto/rl/engine/driver.py @@ -122,6 +122,7 @@ "publish": "trainer+rollout", "eval": "rollout", "offload": "trainer", + "checkpoint": "trainer", "onload": "trainer", } _PHASE_TASK = {"sync": "outer_sync"} @@ -1531,6 +1532,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() diff --git a/yeto/rl/engine/resume.py b/yeto/rl/engine/resume.py index bfa3894e..6778ee34 100644 --- a/yeto/rl/engine/resume.py +++ b/yeto/rl/engine/resume.py @@ -210,10 +210,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 +236,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} @@ -509,10 +512,16 @@ 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) -> 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.""" if skip_if_recovery and self.recovery_required: return False info = {"incarnation": self.incarnation["id"], "incarnation_index": self.incarnation["index"], @@ -526,17 +535,41 @@ 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, self.state_dir, 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]: + 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 cadc5047..e2eee98c 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,14 +93,28 @@ 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: + version, state = int(pending[0]), pending[1] + settled, token = True, policy_token(version, state.policy_tensor_hash()) + 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") cursor = (self.rollout.data_cursor() if callable(getattr(self.rollout, "data_cursor", None)) else None) @@ -129,7 +144,7 @@ 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(), ) @@ -138,8 +153,8 @@ def context(self, cut_id: str) -> Any: 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 +170,25 @@ 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 + + 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 +200,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 +225,17 @@ 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's store commit runs in the background (the controller must + # offer commit_checkpoint_store); joined before the next cut writes to the store + # and at the final cut. YETO_RL_CUT_BACKGROUND_COMMIT=0 commits 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) # ---------------------------------------------------------------- pointer @property @@ -228,6 +273,78 @@ 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 prepare(self, driver: Any, *, rollout_id: int, state: Any, final: bool = False) -> bool: + """E1: write 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 (pruned). + Returns True when a cut was prepared. Raises like :meth:`save`.""" + import time + + self._prepared = None + if driver.published_version is None or self._skip(driver, rollout_id, final=final): + return False + 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) # the shards go into the store: no commit may run + 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)) + self.trainer.save_cut(epoch=self._epoch(), context=context) + self._prepared = _PreparedCut(rollout_id=int(rollout_id), context=context, + save_s=time.monotonic() - started) + return True + + 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) -> dict[str, Any] | None: + """Wait for a background store commit and report it (``rl_cut_committed``).""" + if self._commit is None: + return None + thread, box = self._commit + thread.join() + self._commit = None + result = dict(box.get("result") or {"ok": False, "error": "commit thread ended without a result"}) + emit = getattr(driver, "emit", None) + if callable(emit): + emit("rl_cut_committed", cut_id=box.get("cut_id"), ok=bool(result.get("ok")), + commit_s=result.get("commit_s"), pruned=result.get("pruned"), + latest_seq=result.get("seq"), **({"error": result["error"]} if result.get("error") else {})) + return result + + def _commit_in_background(self, cut_id: str) -> None: + import threading + + box: dict[str, Any] = {"cut_id": cut_id} + + def run() -> None: + box["result"] = self.controller.commit_checkpoint_store() + + 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 @@ -244,30 +361,38 @@ 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: - return None - if not self.due(rollout_id, final=final): + prepared, self._prepared = self._prepared, None + if driver.published_version is None or self._skip(driver, rollout_id, final=final): return None if int(driver.published_version) != int(rollout_id): 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) 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 + used_prepared = prepared is not None and prepared.matches(int(rollout_id), context) + if used_prepared: + save_s = prepared.save_s # written before the trainer offloaded (prepare) + else: + # a prepared cut that does not match (or a failed prepare) 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: + self.trainer.save_cut(epoch=self._epoch(), context=context) + finally: + if offloaded: + self.trainer.offload() + save_s = time.monotonic() - started + cut_id = context.cut_id + # inline at the end (the run leaves a committed cut) and when a sample buffer + # commits the same store during the round (no two commits at once) + background = self.background_commit and not final and self.buffer_watermark is None consumed_n, consumed_sha = self._consumed() next_lr = self.next_lr(driver) if callable(self.next_lr) else None info = { @@ -291,13 +416,19 @@ 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)) - saved = recovery.save(cut, store=ctl.store, write=lambda record: self._write_and_sync(info, record)) + saved = recovery.save(cut, store=ctl.store, + write=lambda record: self._write_and_sync(info, record, defer=background)) 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) + if background and synced: + self._commit_in_background(cut_id) # after the last read of last_store_sync total_s = time.monotonic() - started - info["store_synced"] = bool(synced) + # background: written (LATEST) but not yet committed; rl_cut_committed reports it + info["store_synced"] = None if background and synced else bool(synced) + info["store_commit"] = "background" if background and synced else "inline" + info["prepared_before_offload"] = used_prepared info["cut_bytes"] = cut_bytes info["save_s"] = round(save_s, 3) info["store_copy_s"] = sync.get("copy_s") @@ -311,11 +442,14 @@ def save(self, driver: Any, *, rollout_id: int, final: bool = False, 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 _write_and_sync(self, info: dict[str, Any], cut_record: dict[str, Any], *, + defer: bool = False) -> 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).""" + into the state dir, then the state dir into the store (which commits, unless + ``defer``: the caller commits in the background).""" self._write_pointer({**info, CUT_KEY: cut_record}) - synced = self.controller.sync_checkpoint_store("round_cut", skip_if_recovery=True) + kwargs = {"defer_commit": True} if defer else {} + 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} From 3b00fe50732e373e69524970f3adf53125266b03 Mon Sep 17 00:00:00 2001 From: michael <78113554+chowhsu@users.noreply.github.com> Date: Wed, 21 Oct 2026 00:54:25 +0000 Subject: [PATCH 16/55] fully_async cut: record untrained prompts without stopping the worker A round cut on a Miles fully_async island stopped the background worker and requeued every untrained prompt into the live data source (P0-S decision 3). With a cut every round the worker restarted each round, so generation never ran ahead of training (E0.1: 8 worker starts in 8 rounds, policy age 0, 2.15 s of serial generation per round). snapshot_untrained() reads the data source buffer, every drawn untrained group (reset copies) and the data cursor in one step on the worker's event loop, where the worker draws; the cut carries them and the cursor read with them, and the live run is untouched. A resume still regenerates exactly the drawn, untrained prompts (decision 3 unchanged); the resume manifest's data watermark is that cursor. requeue_untrained() stays for its other callers. Co-Authored-By: Claude Opus 5.5 --- tests/test_rl_fully_async_prompts.py | 63 ++++++++++++++++--- tests/test_rl_resume_from_checkpoint.py | 17 +++++ yeto/rl/adapters/miles/fully_async_prompts.py | 62 +++++++++++++++--- yeto/rl/adapters/miles/in_flight.py | 21 ++++--- yeto/rl/engine/round_cut.py | 4 ++ 5 files changed, 142 insertions(+), 25 deletions(-) diff --git a/tests/test_rl_fully_async_prompts.py b/tests/test_rl_fully_async_prompts.py index c92c68a3..eea87337 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): diff --git a/tests/test_rl_resume_from_checkpoint.py b/tests/test_rl_resume_from_checkpoint.py index 2b64ffe2..82b03b28 100644 --- a/tests/test_rl_resume_from_checkpoint.py +++ b/tests/test_rl_resume_from_checkpoint.py @@ -672,3 +672,20 @@ def test_background_commit_can_be_turned_off(tmp_path, monkeypatch): _rank, _ctl_a, driver, cp = _trained_island(tmp_path, store, 2) info = cp.save(driver, rollout_id=2) assert info["store_commit"] == "inline" and info["store_synced"] is True + + +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 diff --git a/yeto/rl/adapters/miles/fully_async_prompts.py b/yeto/rl/adapters/miles/fully_async_prompts.py index b448cdbd..bf1fbd95 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 44a32207..4e896ee2 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/engine/round_cut.py b/yeto/rl/engine/round_cut.py index e2eee98c..c7fe3e3c 100644 --- a/yeto/rl/engine/round_cut.py +++ b/yeto/rl/engine/round_cut.py @@ -137,6 +137,10 @@ def context(self, cut_id: str, *, pending: tuple[int, Any] | None = None) -> 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 From 84d43ce86fb4a0a73a76f659aa44992bc974bf00 Mon Sep 17 00:00:00 2001 From: michael <78113554+chowhsu@users.noreply.github.com> Date: Wed, 21 Oct 2026 00:55:59 +0000 Subject: [PATCH 17/55] golden: regenerate decoupling samples for the explicit Miles dropout 0.0 e0f833dc made the generic Miles recipe pass --attention-dropout 0.0 and --hidden-dropout 0.0 (the bare Miles recipe's values); the seven decoupling golden samples were not regenerated and failed on miles_argv. The only difference is those four argv items (and the argv hash); no algorithm or contract hash changes. Co-Authored-By: Claude Opus 5.5 --- tests/golden/decoupling/codex_harness.json | 6 +++++- tests/golden/decoupling/decoupled.json | 6 +++++- tests/golden/decoupling/drgrpo.json | 6 +++++- tests/golden/decoupling/elastic.json | 6 +++++- tests/golden/decoupling/grpo_default.json | 6 +++++- tests/golden/decoupling/grpo_tis.json | 6 +++++- tests/golden/decoupling/seq_adv_maxrl.json | 6 +++++- 7 files changed, 35 insertions(+), 7 deletions(-) diff --git a/tests/golden/decoupling/codex_harness.json b/tests/golden/decoupling/codex_harness.json index b957d6a5..f74fb036 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 eb40b112..0d63eb51 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 fc7f4d60..786cb2c1 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 5664ac5b..d37d4cd6 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 61c61102..5235984f 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 bca5a992..9d65c992 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 80100673..29bd1d96 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, From 36eb805d1e76a97341c66f332240a694088dd206 Mon Sep 17 00:00:00 2001 From: michael <78113554+chowhsu@users.noreply.github.com> Date: Wed, 21 Oct 2026 00:59:34 +0000 Subject: [PATCH 18/55] driver: time the rollout offload after generation as its own phase E0.1: the yeto+Miles sync generate span was 0.46 s longer than Miles' perf/rollout_time. 0.44 s of it is --offload-rollout (release_memory_occupation), which Miles times outside rollout_time but MilesRolloutPool.generate ran inside the driver's generate phase. The driver now asks the pool to defer it and runs it as phase "rollout_offload"; callers that use the pool directly keep the old in-generate offload. Co-Authored-By: Claude Opus 5.5 --- tests/test_rl_engine_driver.py | 15 +++++++++++++++ yeto/rl/adapters/miles/rollout.py | 14 +++++++++++++- yeto/rl/engine/driver.py | 10 ++++++++++ 3 files changed, 38 insertions(+), 1 deletion(-) diff --git a/tests/test_rl_engine_driver.py b/tests/test_rl_engine_driver.py index 738529b8..b1bdefdb 100644 --- a/tests/test_rl_engine_driver.py +++ b/tests/test_rl_engine_driver.py @@ -998,3 +998,18 @@ def save(self, driver, *, rollout_id, final=False, **kw): 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/yeto/rl/adapters/miles/rollout.py b/yeto/rl/adapters/miles/rollout.py index df2080d5..f2589c6f 100644 --- a/yeto/rl/adapters/miles/rollout.py +++ b/yeto/rl/adapters/miles/rollout.py @@ -834,7 +834,8 @@ 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) handle = handle_from_metadata( @@ -1045,6 +1046,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/engine/driver.py b/yeto/rl/engine/driver.py index f27ea005..c66a2405 100644 --- a/yeto/rl/engine/driver.py +++ b/yeto/rl/engine/driver.py @@ -123,6 +123,7 @@ "eval": "rollout", "offload": "trainer", "checkpoint": "trainer", + "rollout_offload": "rollout", "onload": "trainer", } _PHASE_TASK = {"sync": "outer_sync"} @@ -890,11 +891,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} " From dfe8b914ede4d1452a3168224c0c1d4f65ac98d0 Mon Sep 17 00:00:00 2001 From: michael <78113554+chowhsu@users.noreply.github.com> Date: Wed, 21 Oct 2026 01:14:22 +0000 Subject: [PATCH 19/55] eff eval: --checkpoint final|every-round for the yeto+Miles arm final (default) cuts once at the end like the bare arms keep no checkpoint; every-round is the product default. Kept out of the control fingerprint so the bare runs stay comparable. The yeto+verl arm has no checkpoint store in either mode. Co-Authored-By: Claude Opus 5.5 --- tests/test_tools_eval_efficiency.py | 5 +++++ tools/eval/__main__.py | 5 ++++- tools/eval/efficiency/spec.py | 24 +++++++++++++++++++----- 3 files changed, 28 insertions(+), 6 deletions(-) diff --git a/tests/test_tools_eval_efficiency.py b/tests/test_tools_eval_efficiency.py index 9be2abe0..6da76604 100644 --- a/tests/test_tools_eval_efficiency.py +++ b/tests/test_tools_eval_efficiency.py @@ -87,6 +87,11 @@ def test_yeto_arms_use_the_same_policy_age_and_partition(): assert "--gpu modal:1xa10g" in spec.yeto_miles_launcher_args(SPEC, "sync", "r1") # 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 + 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); diff --git a/tools/eval/__main__.py b/tools/eval/__main__.py index 32ed6571..cbb998b6 100644 --- a/tools/eval/__main__.py +++ b/tools/eval/__main__.py @@ -148,7 +148,7 @@ def cmd_eff_plan(args) -> int: def cmd_eff_run(args) -> int: from .efficiency import modal_eff, spec - sp = spec.get_spec(args.preset) + sp = spec.get_spec(args.preset, checkpoint=args.checkpoint) plan = spec.plan(sp, args.arm, args.run_id) run_dir = Path(args.out) / args.run_id if not args.dry_run: @@ -227,6 +227,9 @@ def main(argv: list[str] | None = None) -> int: 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="final", + help="yeto+Miles round cuts: only at the end (fair control) or every round " + "(product default)") er.set_defaults(fn=cmd_eff_run) eparse = sub.add_parser("eff-parse") eparse.add_argument("run_dir") diff --git a/tools/eval/efficiency/spec.py b/tools/eval/efficiency/spec.py index 56d29d45..cf3864b2 100644 --- a/tools/eval/efficiency/spec.py +++ b/tools/eval/efficiency/spec.py @@ -80,6 +80,10 @@ class EffSpec: 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" @property def samples_per_round(self) -> int: @@ -145,7 +149,7 @@ def cards(mode: str) -> dict[str, int]: 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")} + doc = {k: v for k, v in spec.to_dict().items() if k not in ("repeats", "timeout_s", "checkpoint")} return hashlib.sha256(json.dumps(doc, sort_keys=True).encode()).hexdigest() @@ -356,6 +360,19 @@ def yeto_verl_argv(spec: EffSpec, mode: str, tape: str) -> list[str]: return argv +CHECKPOINT_MODES = ("final", "every-round") + + +def cut_every(spec: EffSpec) -> int: + """Round-cut interval of the yeto+Miles arm: ``final`` cuts only at the end (the + bare arms keep no checkpoint), ``every-round`` is the product default. + 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 spec.rounds if spec.checkpoint == "final" else 1 + + 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 @@ -368,10 +385,7 @@ def yeto_miles_launcher_args(spec: EffSpec, mode: str, run_id: str) -> str: 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} " - # Only the final round cut: the bare arms save no checkpoint. A periodic cut - # (default every round) is synchronous (~6-9 s on A10G) and stops Miles' - # fully_async worker, so the async arm would generate nothing ahead (E1). - f"--rl-cut-every {spec.rounds} " + 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} " From 8ef1018a1c327f3388ff42ef8bec0db9f27dd9f6 Mon Sep 17 00:00:00 2001 From: michael <78113554+chowhsu@users.noreply.github.com> Date: Wed, 21 Oct 2026 01:34:09 +0000 Subject: [PATCH 20/55] round cut: write the trainer shards in the background for periodic cuts GPU canary e1a-yeto-miles-async-1 (every-round cuts, after 74016eb0 and 3b00fe50): overlap is back (one fully_async worker start), but each cut still held the round for 2.1 s writing and hashing the 283 MB shard, so a round took 5.3 s against 4.15 s bare (78 %). MilesTrainerGroup.start_cut captures every rank's state at the step boundary (CPU copies) and the ranks write, fsync and hash their shards in a thread (cut_plugin.save_cut_shard background=True); finish_cut collects the summaries and commits the manifest. A periodic cut now only captures at its safe point; the shard write, manifest, pointer, store snapshot and commit complete before the next cut (join_commit, rl_cut_committed). Until then the previous cut stays the pointed and committed one; a failed completion drops the unpointed files. The store's first cut, the final cut, runs with a sample buffer or a critic, and YETO_RL_CUT_BACKGROUND_COMMIT=0 stay inline. Co-Authored-By: Claude Opus 5.5 --- tests/test_rl_resume_from_checkpoint.py | 115 +++++++++--------- tests/test_rl_trainer_cut.py | 4 +- yeto/rl/adapters/miles/cut_plugin.py | 64 +++++++++-- yeto/rl/adapters/miles/trainer.py | 40 +++++-- yeto/rl/engine/round_cut.py | 147 ++++++++++++++++++------ 5 files changed, 265 insertions(+), 105 deletions(-) diff --git a/tests/test_rl_resume_from_checkpoint.py b/tests/test_rl_resume_from_checkpoint.py index 82b03b28..13bd0bb5 100644 --- a/tests/test_rl_resume_from_checkpoint.py +++ b/tests/test_rl_resume_from_checkpoint.py @@ -193,12 +193,10 @@ def test_save_kill_resume_without_elastic_continues_bitwise(tmp_path): 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) - # E1: a periodic cut commits the store in the background; LATEST is written already - assert info["store_synced"] is None and info["store_commit"] == "background" + # E1: the store's first cut is complete inline (a store with cuts but no LATEST refuses to start) + assert info["store_synced"] is True and info["store_commit"] == "inline" 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 cp_a.join_commit(driver_a)["ok"] is True - assert driver_a.events[-1][0] == "rl_cut_committed" and driver_a.events[-1][1]["cut_id"] == info["cut_id"] assert vol.commits >= 1 and rs.read_latest(store.root)["cut_id"] == info["cut_id"] rank_b = make_rank(9) @@ -547,9 +545,9 @@ def test_resume_wiring_carries_the_runtime_fingerprint(tmp_path): assert w.controller.runtime_fingerprint == "miles@abc" and w.every == 2 and w.ledger is not None -# ---------------------------------------------------------------- E1: cut before offload, background commit +# ---------------------------------------------------------------- E1: cut before offload, background cut class _CountingTrainer: - """The real MilesTrainerGroup cut path, counting save/onload/offload calls.""" + """The real MilesTrainerGroup cut path, counting the cut and onload/offload calls.""" def __init__(self, inner): self.inner = inner @@ -559,6 +557,14 @@ 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 onload(self): self.calls.append("onload") @@ -577,61 +583,75 @@ def _publish(driver, version, policy_hash): driver.expected_token = policy_token(version, policy_hash) -def _prepared_island(tmp_path, store, **kw): - """Two trained rounds, policy v1 published, v2 about to be published by an offloaded - colocated trainer.""" +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) - for _ in range(2): - train_step(rank, torch.randn(3, 6, generator=g)) - ctl = _rctl(tmp_path, "a", store) - driver = _Driver(rank, published=1, local_step=2) + 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] - driver._trainer_offloaded = True trainer = _CountingTrainer(_trainer(rank)) - return rank, driver, trainer, _cp(ctl, driver, trainer, _Ledger(2), **kw) + cp = _cp(ctl, driver, trainer, _Ledger(2)) + _publish(driver, 1, _hash(rank)) + first = cp.save(driver, rollout_id=1) + assert first["store_commit"] == "inline" and 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 test_cut_prepared_before_offload_is_kept_without_waking_the_trainer(tmp_path): - store = _store(tmp_path / "store", _Vol()) - rank, driver, trainer, cp = _prepared_island(tmp_path, store) - 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 == ["save_cut"] # no wake/sleep for the cut - assert info["prepared_before_offload"] is True and info["trainer_onloaded_for_cut"] is False - assert cp.join_commit(driver)["ok"] is True - # the prepared cut is a complete cut: a fresh island resumes from it bitwise +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)) - assert cp_b.resume(driver_b)["next_rollout_id"] == 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] == "finish_cut" + assert driver.events[-1][0] == "rl_cut_committed" and driver.events[-1][1]["ok"] is True + 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 = _prepared_island(tmp_path, store) + 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 == ["save_cut", "onload", "save_cut", "offload"] + assert trainer.calls == ["start_cut", "finish_cut", "onload", "start_cut", "offload"] assert info["prepared_before_offload"] is False and info["policy_hash"] == _hash(rank) - cp.join_commit(driver) - 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)) - assert cp_b.resume(driver_b)["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_commits_inline(tmp_path): +def test_final_cut_completes_inline_after_the_pending_one(tmp_path): store = _store(tmp_path / "store", _Vol()) - _rank, _ctl_a, driver, cp = _trained_island(tmp_path, store, 2) + 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 cp.join_commit(driver) is None + assert trainer.calls == ["save_cut"] and cp.join_commit(driver) is None def test_failed_background_commit_keeps_the_previous_cut(tmp_path): @@ -645,18 +665,8 @@ def commit(self): vol = _FailingVol() store = _store(tmp_path / "store", vol) - rank = make_rank(0) - g = torch.Generator().manual_seed(3) - ctl = _rctl(tmp_path, "a", store, keep=1) - driver = _Driver(rank, published=0, local_step=0) - driver.last_applied_lrs = [5e-6] - cp = _cp(ctl, driver, _trainer(rank), _Ledger(2)) - train_step(rank, torch.randn(3, 6, generator=g)) - driver.published_version, driver.local_step = 1, 1 - first = cp.save(driver, rollout_id=1) - assert cp.join_commit(driver)["ok"] is True - train_step(rank, torch.randn(3, 6, generator=g)) - driver.published_version, driver.local_step = 2, 2 + 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) @@ -666,12 +676,13 @@ def commit(self): assert (store.root / rs.ROUND_CUTS / first["cut_id"]).is_dir() -def test_background_commit_can_be_turned_off(tmp_path, monkeypatch): +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, _ctl_a, driver, cp = _trained_island(tmp_path, store, 2) + 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 + 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(): diff --git a/tests/test_rl_trainer_cut.py b/tests/test_rl_trainer_cut.py index a2e1c8b7..643f04fd 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/yeto/rl/adapters/miles/cut_plugin.py b/yeto/rl/adapters/miles/cut_plugin.py index 50fe3afd..ddc7fed7 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. @@ -596,22 +597,59 @@ 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 (CPU copies, at this step boundary) and + write, hash and summarize it in a thread of this rank, 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) 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_refusing, captured, directory=directory, cut_id=cut_id)) + pool.shutdown(wait=False) + return {"pending": cut_id, "coord": dict(captured["coord"])} + + +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).""" + 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(captured, directory=directory, cut_id=cut_id) + except CutPluginError as error: + return {REFUSED: str(error), REFUSAL_KIND: "refused"} + + +def _capture(actor: Any, *, cut_id: str) -> dict[str, Any]: + """Everything the shard and its summary need, read at the step boundary (CPU copies).""" import torch backend = _backend(actor) @@ -622,6 +660,16 @@ def _save(actor: Any, *, directory: str, cut_id: str) -> dict[str, Any]: with torch.no_grad(): snap = _snapshot(actor, backend, named) rng = backend.capture_rng() + return {"coord": coord, "snap": snap, "rng": rng, + "scheduler_samples": int(actor.opt_param_scheduler.num_steps), + "tuning": _tuning(backend, named), "adapter_names": sorted(n for n, _ in named), + "resume": resume_rank_state(actor, snap["optimizer_named"])} + + +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) @@ -652,7 +700,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,11 +708,11 @@ 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"], } diff --git a/yeto/rl/adapters/miles/trainer.py b/yeto/rl/adapters/miles/trainer.py index 54c65ed9..63259e9f 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,29 @@ 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))] + return {"epoch": int(epoch), "context": context, "runtime": runtime, "directory": directory, + "background": background, "results": results, "critic": critic} + + def finish_cut(self, pending: Mapping[str, Any]) -> str: + """Collect the shards of :meth:`start_cut` and commit the manifest; returns the cut id.""" + 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}))] + 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 +869,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( diff --git a/yeto/rl/engine/round_cut.py b/yeto/rl/engine/round_cut.py index c7fe3e3c..d2988ce7 100644 --- a/yeto/rl/engine/round_cut.py +++ b/yeto/rl/engine/round_cut.py @@ -181,6 +181,7 @@ class _PreparedCut: 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 @@ -231,15 +232,17 @@ def __init__(self, *, controller: Any, source: Any, trainer: Any, every: int = 1 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's store commit runs in the background (the controller must - # offer commit_checkpoint_store); joined before the next cut writes to the store - # and at the final cut. YETO_RL_CUT_BACKGROUND_COMMIT=0 commits inline. + # 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 # ---------------------------------------------------------------- pointer @property @@ -288,13 +291,22 @@ def _skip(self, driver: Any, rollout_id: int, *, final: bool) -> bool: 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): + 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: write the trainer part of the cut due at the safe point before rollout + """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 (pruned). + 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 @@ -303,16 +315,25 @@ def prepare(self, driver: Any, *, rollout_id: int, state: Any, final: bool = Fal return False 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) # the shards go into the store: no commit may run + self.join_commit(driver) # 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)) - self.trainer.save_cut(epoch=self._epoch(), context=context) + pending = self._start(context, background=self._defer(final)) self._prepared = _PreparedCut(rollout_id=int(rollout_id), context=context, - save_s=time.monotonic() - started) + save_s=time.monotonic() - started, pending=pending) 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), and + not before the first cut of the store is committed (a store holding cuts but + no LATEST refuses to start).""" + return (self.background_commit and not final and self.buffer_watermark is None + and self.pointer() is not None) + def _drop_unpointed(self, cut_id: str) -> None: import shutil @@ -323,20 +344,53 @@ def _drop_unpointed(self, cut_id: str) -> None: shutil.rmtree(directory) def join_commit(self, driver: Any = None) -> dict[str, Any] | None: - """Wait for a background store commit and report it (``rl_cut_committed``).""" - if self._commit is None: + """Complete the previous background cut: collect its shards, commit its + manifest, point the state dir at it and commit the store; then report it + (``rl_cut_committed``). On any failure the previous pointer stays the + committed one. Returns the outcome (None: nothing was pending).""" + deferred, self._deferred = self._deferred, None + result: dict[str, Any] | None = None + if deferred is not None: + result = self._complete_deferred(deferred) + if self._commit is not None: + thread, box = self._commit + thread.join() + self._commit = None + result = {**(result or {}), **dict(box.get("result") or { + "ok": False, "error": "commit thread ended without a result"})} + result.setdefault("cut_id", box.get("cut_id")) + if result is None: return None - thread, box = self._commit - thread.join() - self._commit = None - result = dict(box.get("result") or {"ok": False, "error": "commit thread ended without a result"}) emit = getattr(driver, "emit", None) if callable(emit): - emit("rl_cut_committed", cut_id=box.get("cut_id"), ok=bool(result.get("ok")), - commit_s=result.get("commit_s"), pruned=result.get("pruned"), - latest_seq=result.get("seq"), **({"error": result["error"]} if result.get("error") else {})) + emit("rl_cut_committed", **{k: result.get(k) for k in ( + "cut_id", "rollout_id", "ok", "finish_s", "commit_s", "pruned", "seq", "cut_bytes")}, + **({"error": result["error"]} if result.get("error") else {})) return result + def _complete_deferred(self, deferred: dict[str, Any]) -> dict[str, Any]: + import time + + from yeto.rl.engine.resume import directory_bytes + + info, cut = deferred["info"], deferred["cut"] + started = time.monotonic() + try: + self.trainer.finish_cut(deferred["pending"]) + saved = recovery.save(cut, store=self.controller.store, + write=lambda record: self._write_and_sync(info, record, defer=True)) + except Exception as error: # noqa: BLE001 - the previous cut stays the pointed one + self._drop_unpointed(info["cut_id"]) + return {"cut_id": info["cut_id"], "rollout_id": deferred["rollout_id"], "ok": False, + "error": repr(error)[:2000]} + out = {"cut_id": info["cut_id"], "rollout_id": deferred["rollout_id"], + "finish_s": round(time.monotonic() - started, 3), + "cut_bytes": directory_bytes(Path(self.root) / info["cut_id"])} + if saved["synced"]: + self._commit_in_background(info["cut_id"]) + return out + return {**out, "ok": False, "error": (getattr(self.controller, "last_store_sync", None) or {}).get("error")} + def _commit_in_background(self, cut_id: str) -> None: import threading @@ -356,6 +410,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 @@ -367,20 +427,24 @@ def save(self, driver: Any, *, rollout_id: int, final: bool = False, ctl = self.controller 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 int(driver.published_version) != int(rollout_id): + self._abandon(prepared) raise RoundCutError(f"safe point {rollout_id} != published version {driver.published_version}") self.join_commit(driver) started = time.monotonic() 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 = prepared.save_s # written before the trainer offloaded (prepare) + save_s, pending = prepared.save_s, prepared.pending # taken before the trainer offloaded else: - # a prepared cut that does not match (or a failed prepare) left files under - # this id; no pointer names them (_skip), and cuts are immutable: drop them + 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. @@ -388,15 +452,14 @@ def save(self, driver: Any, *, rollout_id: int, final: bool = False, if offloaded: self.trainer.onload() try: - self.trainer.save_cut(epoch=self._epoch(), context=context) + 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 - # inline at the end (the run leaves a committed cut) and when a sample buffer - # commits the same store during the round (no two commits at once) - background = self.background_commit and not final and self.buffer_watermark is None consumed_n, consumed_sha = self._consumed() next_lr = self.next_lr(driver) if callable(self.next_lr) else None info = { @@ -420,32 +483,44 @@ 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)) - saved = recovery.save(cut, store=ctl.store, - write=lambda record: self._write_and_sync(info, record, defer=background)) + extra = {"prepared_before_offload": used_prepared, "trainer_onloaded_for_cut": offloaded, + "save_s": round(save_s, 3)} + if defer: + self._deferred = {"info": info, "cut": cut, "pending": pending, "rollout_id": int(rollout_id)} + 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) - if background and synced: - self._commit_in_background(cut_id) # after the last read of last_store_sync - total_s = time.monotonic() - started - # background: written (LATEST) but not yet committed; rl_cut_committed reports it - info["store_synced"] = None if background and synced else bool(synced) - info["store_commit"] = "background" if background and synced else "inline" - info["prepared_before_offload"] = used_prepared + 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 _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) -> dict[str, Any]: """The round cut's ``recovery.save`` writer: the pointer (with the cut record) From de275ef57738932f7a17578fec240509825656b0 Mon Sep 17 00:00:00 2001 From: michael <78113554+chowhsu@users.noreply.github.com> Date: Wed, 21 Oct 2026 01:58:15 +0000 Subject: [PATCH 21/55] round cut: skip a periodic cut while the previous commit runs; time the sample filter hook GPU e1a2-yeto-miles-async-1: with the shards written in the background a cut held the safe point 0.2 s, but joining the previous Modal volume commit (1.8-9.9 s) at the next cut still held the round. A periodic cut is now skipped (rl_cut_skipped) while that commit runs; the final cut waits for it. The same run spent 0.9 s per round between Miles' 'Finish rollout' and the train data (0.04 s in E0.1): the yeto sample filter hook now prints its per-step timing to find it. Co-Authored-By: Claude Opus 5.5 --- tests/test_rl_resume_from_checkpoint.py | 18 ++++++++++++++++++ yeto/rl/adapters/miles/rollout_meta_hook.py | 15 ++++++++++++++- yeto/rl/engine/round_cut.py | 19 +++++++++++++++++++ 3 files changed, 51 insertions(+), 1 deletion(-) diff --git a/tests/test_rl_resume_from_checkpoint.py b/tests/test_rl_resume_from_checkpoint.py index 13bd0bb5..a08112d3 100644 --- a/tests/test_rl_resume_from_checkpoint.py +++ b/tests/test_rl_resume_from_checkpoint.py @@ -700,3 +700,21 @@ def test_cut_context_takes_the_cursor_read_with_the_in_flight_entries(): 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" diff --git a/yeto/rl/adapters/miles/rollout_meta_hook.py b/yeto/rl/adapters/miles/rollout_meta_hook.py index d34e761d..450936f8 100644 --- a/yeto/rl/adapters/miles/rollout_meta_hook.py +++ b/yeto/rl/adapters/miles/rollout_meta_hook.py @@ -137,14 +137,19 @@ def record_trained_groups(args: Any, data: Sequence[Sequence[Any]]) -> None: ``--partial-rollout`` go back. """ + 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 +157,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 +183,13 @@ 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) + 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:])), flush=True) def _reward(args: Any, sample: Any) -> float: diff --git a/yeto/rl/engine/round_cut.py b/yeto/rl/engine/round_cut.py index d2988ce7..09ad8c2a 100644 --- a/yeto/rl/engine/round_cut.py +++ b/yeto/rl/engine/round_cut.py @@ -243,6 +243,7 @@ def __init__(self, *, controller: Any, source: Any, trainer: Any, every: int = 1 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.skipped_busy = 0 # periodic cuts skipped while a commit was running # ---------------------------------------------------------------- pointer @property @@ -313,6 +314,8 @@ def prepare(self, driver: Any, *, rollout_id: int, state: Any, final: bool = Fal self._prepared = None 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(): + 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) # the previous cut is complete before this one starts @@ -325,6 +328,19 @@ def prepare(self, driver: Any, *, rollout_id: int, state: Any, final: bool = Fal 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 _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 @@ -429,6 +445,9 @@ def save(self, driver: Any, *, rollout_id: int, final: bool = False, 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): + 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}") From 77b460cdbde709794cf29fbc5257129350920b07 Mon Sep 17 00:00:00 2001 From: michael <78113554+chowhsu@users.noreply.github.com> Date: Wed, 21 Oct 2026 02:02:52 +0000 Subject: [PATCH 22/55] cut digests: hash tensor buffers without a bytes copy state_digest copied every tensor into a bytes object before hashing, holding the GIL for the copy. Hashing the buffer gives the same digest (checked for float32, bfloat16, empty tensors and the numpy RNG arrays) and lets hashlib release the GIL, so a background cut writer does not stall the trainer's own thread (E1 GPU: a colocated offload took 4.0 s instead of 1.3 s while the shard was written). Co-Authored-By: Claude Opus 5.5 --- yeto/rl/adapters/miles/cut_plugin.py | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/yeto/rl/adapters/miles/cut_plugin.py b/yeto/rl/adapters/miles/cut_plugin.py index ddc7fed7..68bef508 100644 --- a/yeto/rl/adapters/miles/cut_plugin.py +++ b/yeto/rl/adapters/miles/cut_plugin.py @@ -333,10 +333,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): From 86d8872fd5e94899feb2d4d7dfbc5b5d257c786d Mon Sep 17 00:00:00 2001 From: michael <78113554+chowhsu@users.noreply.github.com> Date: Wed, 21 Oct 2026 02:12:36 +0000 Subject: [PATCH 23/55] round cut: a periodic cut does not wait for the store commit it starts e1a3-yeto-miles-async-1 still held each round 5-6 s: join_commit completed the pending cut, started its store commit and then joined that same commit. A periodic cut now completes the pending cut and starts its commit without waiting (the commit still running at the next safe point skips that cut); only the final cut waits. Co-Authored-By: Claude Opus 5.5 --- tests/test_rl_resume_from_checkpoint.py | 36 +++++++++++++++++ yeto/rl/engine/round_cut.py | 51 ++++++++++++++++--------- 2 files changed, 69 insertions(+), 18 deletions(-) diff --git a/tests/test_rl_resume_from_checkpoint.py b/tests/test_rl_resume_from_checkpoint.py index a08112d3..79b0fea4 100644 --- a/tests/test_rl_resume_from_checkpoint.py +++ b/tests/test_rl_resume_from_checkpoint.py @@ -718,3 +718,39 @@ def test_periodic_cut_is_skipped_while_the_previous_commit_runs(tmp_path): 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"] + assert committed == ["r000002-" + cp.controller.incarnation["id"], info["cut_id"]] # v2, then v3 diff --git a/yeto/rl/engine/round_cut.py b/yeto/rl/engine/round_cut.py index 09ad8c2a..d40c7533 100644 --- a/yeto/rl/engine/round_cut.py +++ b/yeto/rl/engine/round_cut.py @@ -243,6 +243,7 @@ def __init__(self, *, controller: Any, source: Any, trainer: Any, every: int = 1 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 # ---------------------------------------------------------------- pointer @@ -318,7 +319,7 @@ def prepare(self, driver: Any, *, rollout_id: int, state: Any, final: bool = Fal 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) # the previous cut is complete before this one starts + 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() @@ -359,30 +360,44 @@ def _drop_unpointed(self, cut_id: str) -> None: if directory.exists(): shutil.rmtree(directory) - def join_commit(self, driver: Any = None) -> dict[str, Any] | None: - """Complete the previous background cut: collect its shards, commit its - manifest, point the state dir at it and commit the store; then report it - (``rl_cut_committed``). On any failure the previous pointer stays the - committed one. Returns the outcome (None: nothing was pending).""" + 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 - result: dict[str, Any] | None = None if deferred is not None: - result = self._complete_deferred(deferred) - if self._commit is not None: - thread, box = self._commit - thread.join() - self._commit = None - result = {**(result or {}), **dict(box.get("result") or { - "ok": False, "error": "commit thread ended without a result"})} - result.setdefault("cut_id", box.get("cut_id")) - if result is 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")}, **({"error": result["error"]} if result.get("error") else {})) - return result def _complete_deferred(self, deferred: dict[str, Any]) -> dict[str, Any]: import time @@ -451,7 +466,7 @@ def save(self, driver: Any, *, rollout_id: int, final: bool = False, if int(driver.published_version) != int(rollout_id): self._abandon(prepared) raise RoundCutError(f"safe point {rollout_id} != published version {driver.published_version}") - self.join_commit(driver) + self.join_commit(driver, wait=final) # the final cut leaves nothing running started = time.monotonic() context = self.source.context(self._cut_id(rollout_id)) progress = context.progress From 0f615ab819864035640f752b11a26efc6839011f Mon Sep 17 00:00:00 2001 From: michael <78113554+chowhsu@users.noreply.github.com> Date: Wed, 21 Oct 2026 02:13:37 +0000 Subject: [PATCH 24/55] rollout metadata hook: time its steps (E1 diagnosis) The fully_async round metadata took 1.1 s per round on GPU (0.09 s in round 0). Print the seconds of round id, group records, carry-over fields, submitted groups and the sink put with the hook timing line. Co-Authored-By: Claude Opus 5.5 --- yeto/rl/adapters/miles/rollout_meta_hook.py | 29 ++++++++++++++++----- 1 file changed, 22 insertions(+), 7 deletions(-) diff --git a/yeto/rl/adapters/miles/rollout_meta_hook.py b/yeto/rl/adapters/miles/rollout_meta_hook.py index 450936f8..25dc5ef0 100644 --- a/yeto/rl/adapters/miles/rollout_meta_hook.py +++ b/yeto/rl/adapters/miles/rollout_meta_hook.py @@ -41,6 +41,18 @@ 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" @@ -188,8 +200,11 @@ def record_trained_groups(args: Any, data: Sequence[Sequence[Any]]) -> None: 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:])), flush=True) + + " ".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: @@ -392,8 +407,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")) @@ -431,7 +446,7 @@ 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)) # 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: @@ -989,14 +1004,14 @@ def extract_rollout_metadata( :func:`record_trained_groups` with the trained groups).""" try: - payload = build_metadata(args, all_samples) + payload = _timed("build", build_metadata, args, all_samples) 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 @@ -1039,7 +1054,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 From 427d142741a453067e70d97be4acedfe3c58a4b3 Mon Sep 17 00:00:00 2001 From: michael <78113554+chowhsu@users.noreply.github.com> Date: Wed, 21 Oct 2026 02:26:53 +0000 Subject: [PATCH 25/55] fully_async: estimate the cross-version IS truncation off the critical path e1a4-yeto-miles-async-1: with overlap restored every trained token is one version old, so each round's metadata hook re-scored all 32 samples by an SGLang prefill (carry_over_fields: 0.7-1.2 s) before Miles returned the batch, while the trainer waited. Its only reader is the policy-age governor, after the train step. Under fully_async the estimate now runs in a background thread of the rollout process and reaches the sink as a late record of its rollout id; MilesRolloutPool keeps it aside and the driver reads it after the train step (late_cross_version_fraction). Other modes estimate inline as before. Co-Authored-By: Claude Opus 5.5 --- tests/test_rl_miles_carry_over.py | 44 ++++++++++++++++++ yeto/rl/adapters/miles/rollout.py | 44 +++++++++++++++++- yeto/rl/adapters/miles/rollout_meta_hook.py | 51 +++++++++++++++++---- yeto/rl/engine/driver.py | 4 ++ 4 files changed, 134 insertions(+), 9 deletions(-) diff --git a/tests/test_rl_miles_carry_over.py b/tests/test_rl_miles_carry_over.py index e94f3d26..01c13a9b 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/yeto/rl/adapters/miles/rollout.py b/yeto/rl/adapters/miles/rollout.py index f2589c6f..27edd5b3 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 @@ -838,6 +878,8 @@ def generate( 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, diff --git a/yeto/rl/adapters/miles/rollout_meta_hook.py b/yeto/rl/adapters/miles/rollout_meta_hook.py index 25dc5ef0..877b5308 100644 --- a/yeto/rl/adapters/miles/rollout_meta_hook.py +++ b/yeto/rl/adapters/miles/rollout_meta_hook.py @@ -66,6 +66,9 @@ def _timed(name: str, fn: Any, *args: Any, **kwargs: Any) -> Any: # 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"}) @@ -385,7 +388,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: @@ -446,7 +450,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(_timed("carry_over", 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: @@ -481,10 +487,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)} @@ -495,12 +508,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. @@ -1004,7 +1036,10 @@ def extract_rollout_metadata( :func:`record_trained_groups` with the trained groups).""" try: - payload = _timed("build", 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: diff --git a/yeto/rl/engine/driver.py b/yeto/rl/engine/driver.py index c66a2405..9ec6a4f6 100644 --- a/yeto/rl/engine/driver.py +++ b/yeto/rl/engine/driver.py @@ -1466,6 +1466,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) From 4015cc50f9078c2525bbd71eeb0bc5e4cd786aa6 Mon Sep 17 00:00:00 2001 From: michael <78113554+chowhsu@users.noreply.github.com> Date: Wed, 21 Oct 2026 02:37:00 +0000 Subject: [PATCH 26/55] test: the fully_async round report fake takes build_metadata's new keyword Co-Authored-By: Claude Opus 5.5 --- tests/test_rl_fully_async_prompts.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/test_rl_fully_async_prompts.py b/tests/test_rl_fully_async_prompts.py index eea87337..1130228f 100644 --- a/tests/test_rl_fully_async_prompts.py +++ b/tests/test_rl_fully_async_prompts.py @@ -180,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) From 1502e7d2d53feec7390411a7822c50863a3ba2a8 Mon Sep 17 00:00:00 2001 From: michael <78113554+chowhsu@users.noreply.github.com> Date: Wed, 21 Oct 2026 03:28:51 +0000 Subject: [PATCH 27/55] round cuts: commit the deferred cut fully off the critical path The completion of a periodic cut (manifest commit with shard re-hash, pointer, state-dir copy into the store) ran on the driver thread at the next safe point (0.33-0.36 s per cut on the Modal A10G run e1a5). Only the trainer call that collects the rank summaries stays on the driver thread (collect_cut); the rest runs in the commit thread. The store prune no longer deletes a round cut newer than every kept one: such a cut may be being captured while the previous cut commits in the background. A reconfiguration at a safe point first joins a running cut commit. The Miles publisher logs a per-step timing line (E1 diagnosis). Co-Authored-By: Claude Opus 5.5 --- tests/test_rl_resume_from_checkpoint.py | 22 +++++++++- yeto/rl/adapters/miles/publish.py | 16 +++++++ yeto/rl/adapters/miles/trainer.py | 16 ++++++- yeto/rl/engine/driver.py | 5 +++ yeto/rl/engine/resume.py | 14 ++++++ yeto/rl/engine/round_cut.py | 58 ++++++++++++++++++------- 6 files changed, 112 insertions(+), 19 deletions(-) diff --git a/tests/test_rl_resume_from_checkpoint.py b/tests/test_rl_resume_from_checkpoint.py index 79b0fea4..4e1b37e5 100644 --- a/tests/test_rl_resume_from_checkpoint.py +++ b/tests/test_rl_resume_from_checkpoint.py @@ -565,6 +565,10 @@ 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") @@ -625,7 +629,7 @@ def test_prepared_cut_is_written_in_the_background_without_waking_the_trainer(tm 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] == "finish_cut" + 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 assert cp.pointer()["cut_id"] == info["cut_id"] assert rs.read_latest(store.root)["cut_id"] == info["cut_id"] @@ -754,3 +758,19 @@ def commit(self): 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"] 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"] diff --git a/yeto/rl/adapters/miles/publish.py b/yeto/rl/adapters/miles/publish.py index 990843ef..329b66bf 100644 --- a/yeto/rl/adapters/miles/publish.py +++ b/yeto/rl/adapters/miles/publish.py @@ -465,11 +465,16 @@ 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) + self.last_timing = {"check_s": time.monotonic() - started} members, checksums = self._runner.run(self._publish(version, token)) + # 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 = ( @@ -547,18 +552,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: + 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 +578,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/trainer.py b/yeto/rl/adapters/miles/trainer.py index 63259e9f..206aae56 100644 --- a/yeto/rl/adapters/miles/trainer.py +++ b/yeto/rl/adapters/miles/trainer.py @@ -846,6 +846,13 @@ def start_cut(self, *, epoch: int, context: "CutContext", background: bool = Tru 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 @@ -895,8 +902,13 @@ def finish_cut(self, pending: Mapping[str, Any]) -> 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/engine/driver.py b/yeto/rl/engine/driver.py index 9ec6a4f6..930dcc98 100644 --- a/yeto/rl/engine/driver.py +++ b/yeto/rl/engine/driver.py @@ -1638,6 +1638,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 diff --git a/yeto/rl/engine/resume.py b/yeto/rl/engine/resume.py index 6778ee34..b96d36f7 100644 --- a/yeto/rl/engine/resume.py +++ b/yeto/rl/engine/resume.py @@ -369,15 +369,29 @@ 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 directory_bytes(path: Path) -> int: return sum(p.stat().st_size for p in Path(path).rglob("*") if p.is_file()) diff --git a/yeto/rl/engine/round_cut.py b/yeto/rl/engine/round_cut.py index d40c7533..00220df0 100644 --- a/yeto/rl/engine/round_cut.py +++ b/yeto/rl/engine/round_cut.py @@ -400,35 +400,61 @@ def _report(driver: Any, result: Mapping[str, Any]) -> None: **({"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 - from yeto.rl.engine.resume import directory_bytes - - info, cut = deferred["info"], deferred["cut"] + info = deferred["info"] started = time.monotonic() try: - self.trainer.finish_cut(deferred["pending"]) - saved = recovery.save(cut, store=self.controller.store, - write=lambda record: self._write_and_sync(info, record, defer=True)) + 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"]) return {"cut_id": info["cut_id"], "rollout_id": deferred["rollout_id"], "ok": False, "error": repr(error)[:2000]} - out = {"cut_id": info["cut_id"], "rollout_id": deferred["rollout_id"], - "finish_s": round(time.monotonic() - started, 3), - "cut_bytes": directory_bytes(Path(self.root) / info["cut_id"])} - if saved["synced"]: - self._commit_in_background(info["cut_id"]) - return out - return {**out, "ok": False, "error": (getattr(self.controller, "last_store_sync", None) or {}).get("error")} - - def _commit_in_background(self, cut_id: str) -> None: + self._commit_in_background(deferred, commit_manifest) + return {"cut_id": info["cut_id"], "rollout_id": deferred["rollout_id"], + "collect_s": round(time.monotonic() - started, 3)} + + 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: - box["result"] = self.controller.commit_checkpoint_store() + 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)) + 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 + 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) From ae82993f3dfc894c41d87b1f02faa34d3c1c662c Mon Sep 17 00:00:00 2001 From: michael <78113554+chowhsu@users.noreply.github.com> Date: Wed, 21 Oct 2026 04:15:32 +0000 Subject: [PATCH 28/55] driver: take the terminal load probe off the round's critical path The terminal rl_load_sample probe after generate ran synchronously (~0.1 s per round on the Modal A10G fully_async run e1a6). It is observation only: it now runs in a thread; the next round and the end of the run wait for it (bounded), so its event still lands before the tape closes. Co-Authored-By: Claude Opus 5.5 --- yeto/rl/engine/driver.py | 44 +++++++++++++++++++++++++++------------- 1 file changed, 30 insertions(+), 14 deletions(-) diff --git a/yeto/rl/engine/driver.py b/yeto/rl/engine/driver.py index 930dcc98..6ae1a5b7 100644 --- a/yeto/rl/engine/driver.py +++ b/yeto/rl/engine/driver.py @@ -1178,20 +1178,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.""" @@ -1911,6 +1926,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) From 31a951a4addaef656b4bc1882bb6a07dc9fc6904 Mon Sep 17 00:00:00 2001 From: michael <78113554+chowhsu@users.noreply.github.com> Date: Wed, 21 Oct 2026 04:15:32 +0000 Subject: [PATCH 29/55] round cuts: collect a skipped prepared cut's writer at the next cut A periodic cut prepared before the trainer offload and then skipped because the previous cut was still committing waited for its trainer-side shard writer at the safe point (~0.9 s every other round, Modal A10G sync run e1a6). The writer is now collected, and the files dropped, at the next cut. The Miles cut capture logs a per-step timing line (E1 diagnosis). Co-Authored-By: Claude Opus 5.5 --- tests/test_rl_resume_from_checkpoint.py | 29 +++++++++++++++++++++++++ yeto/rl/adapters/miles/cut_plugin.py | 20 +++++++++++++---- yeto/rl/engine/round_cut.py | 15 ++++++++++++- 3 files changed, 59 insertions(+), 5 deletions(-) diff --git a/tests/test_rl_resume_from_checkpoint.py b/tests/test_rl_resume_from_checkpoint.py index 4e1b37e5..fddb7fe8 100644 --- a/tests/test_rl_resume_from_checkpoint.py +++ b/tests/test_rl_resume_from_checkpoint.py @@ -774,3 +774,32 @@ def test_prune_keeps_a_cut_newer_than_every_kept_one(tmp_path): 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"] diff --git a/yeto/rl/adapters/miles/cut_plugin.py b/yeto/rl/adapters/miles/cut_plugin.py index 68bef508..85e095c1 100644 --- a/yeto/rl/adapters/miles/cut_plugin.py +++ b/yeto/rl/adapters/miles/cut_plugin.py @@ -658,14 +658,26 @@ def _capture(actor: Any, *, cut_id: str) -> dict[str, Any]: 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) with torch.no_grad(): snap = _snapshot(actor, backend, named) + marks.append(time.monotonic()) rng = backend.capture_rng() - return {"coord": coord, "snap": snap, "rng": rng, - "scheduler_samples": int(actor.opt_param_scheduler.num_steps), - "tuning": _tuning(backend, named), "adapter_names": sorted(n for n, _ in named), - "resume": resume_rank_state(actor, snap["optimizer_named"])} + captured = {"coord": coord, "snap": snap, "rng": rng, + "scheduler_samples": int(actor.opt_param_scheduler.num_steps), + "tuning": _tuning(backend, named), "adapter_names": sorted(n for n, _ in named)} + marks.append(time.monotonic()) + captured["resume"] = resume_rank_state(actor, snap["optimizer_named"]) + 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:]))))), flush=True) + return captured def _write(captured: dict[str, Any], *, directory: str, cut_id: str) -> dict[str, Any]: diff --git a/yeto/rl/engine/round_cut.py b/yeto/rl/engine/round_cut.py index 00220df0..7a3006ef 100644 --- a/yeto/rl/engine/round_cut.py +++ b/yeto/rl/engine/round_cut.py @@ -245,6 +245,9 @@ def __init__(self, *, controller: Any, source: Any, trainer: Any, every: int = 1 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] = [] # ---------------------------------------------------------------- pointer @property @@ -313,6 +316,7 @@ def prepare(self, driver: Any, *, rollout_id: int, state: Any, final: bool = Fal 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(): @@ -487,8 +491,12 @@ def save(self, driver: Any, *, rollout_id: int, final: bool = False, self._abandon(prepared) return None if self._commit_busy(driver, rollout_id, final=final): - self._abandon(prepared) + # 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 + self._reap_abandoned() if int(driver.published_version) != int(rollout_id): self._abandon(prepared) raise RoundCutError(f"safe point {rollout_id} != published version {driver.published_version}") @@ -570,6 +578,11 @@ def save(self, driver: Any, *, rollout_id: int, final: bool = False, info["store_error"] = sync.get("error") return info + 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: From 0a11aa163b1271d42ea8d2051052af0ab1dfd416 Mon Sep 17 00:00:00 2001 From: michael <78113554+chowhsu@users.noreply.github.com> Date: Wed, 21 Oct 2026 04:26:46 +0000 Subject: [PATCH 30/55] round cuts: do not wake an offloaded trainer for a declined periodic cut When the prepare of a periodic cut met a running commit, and the commit was done by the safe point, the cut was taken there by waking the offloaded colocated trainer (~2 s, Modal A10G sync run e1a7). Such a cut is now skipped (rl_cut_skipped) when a committed cut exists; the next safe point cuts again. Cuts whose prepare was not declined, the first cut and the final cut are taken as before. Co-Authored-By: Claude Opus 5.5 --- tests/test_rl_resume_from_checkpoint.py | 27 +++++++++++++++++++++++++ yeto/rl/engine/round_cut.py | 23 +++++++++++++++++++++ 2 files changed, 50 insertions(+) diff --git a/tests/test_rl_resume_from_checkpoint.py b/tests/test_rl_resume_from_checkpoint.py index fddb7fe8..f940719c 100644 --- a/tests/test_rl_resume_from_checkpoint.py +++ b/tests/test_rl_resume_from_checkpoint.py @@ -803,3 +803,30 @@ def test_a_prepared_cut_skipped_while_committing_is_collected_at_the_next_cut(tm 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 diff --git a/yeto/rl/engine/round_cut.py b/yeto/rl/engine/round_cut.py index 7a3006ef..81054489 100644 --- a/yeto/rl/engine/round_cut.py +++ b/yeto/rl/engine/round_cut.py @@ -248,6 +248,7 @@ def __init__(self, *, controller: Any, source: Any, trainer: Any, every: int = 1 # 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 @@ -320,6 +321,7 @@ def prepare(self, driver: Any, *, rollout_id: int, state: Any, final: bool = Fal 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}") @@ -346,6 +348,24 @@ def _commit_busy(self, driver: Any, rollout_id: int, *, final: bool) -> bool: 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 @@ -497,6 +517,9 @@ def save(self, driver: Any, *, rollout_id: int, final: bool = False, self._abandoned.append(prepared) return None 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}") From f19a05549e84205ec8110906290786fc264a00a0 Mon Sep 17 00:00:00 2001 From: michael <78113554+chowhsu@users.noreply.github.com> Date: Wed, 21 Oct 2026 05:26:01 +0000 Subject: [PATCH 31/55] eval: give the yeto+Miles arm the bare arms' container resources The yeto+Miles arm ran in one 8-CPU / 64 GiB container whatever its card count; the bare arms (modal_eff) get 8 CPU and 48 GiB per card, i.e. 16 CPU for the two-card fully_async arm. The yeto fully_async trainer took 6-11 % more time per token than the bare one on the Modal A10G E1 runs while the one-card sync arms matched, with and without periodic cuts. The arm now asks for 8 CPU and 48 GiB per card. Co-Authored-By: Claude Opus 5.5 --- tools/eval/efficiency/spec.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/tools/eval/efficiency/spec.py b/tools/eval/efficiency/spec.py index cf3864b2..f4ef5450 100644 --- a/tools/eval/efficiency/spec.py +++ b/tools/eval/efficiency/spec.py @@ -381,7 +381,9 @@ def yeto_miles_launcher_args(spec: EffSpec, mode: str, run_id: str) -> str: 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"--modal-retries 0 --modal-cpu 8 --modal-memory 64 --modal-tape-volume {TAPE_VOLUME} " + # 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} " From 19e595ae7e6655035696aeff65191ef3676184ea Mon Sep 17 00:00:00 2001 From: michael <78113554+chowhsu@users.noreply.github.com> Date: Wed, 21 Oct 2026 06:47:50 +0000 Subject: [PATCH 32/55] round cuts: sync a deferred cut's state dir as it was at the cut A periodic cut whose store sync is deferred copied the state dir into the store only at the next safe point, after one more round had changed it: the snapshot held the ledger of round N+1 with the pointer of the cut at round N, and a resume from it was refused ('the restored ledger's consumed groups (64, ...) differ from the pointer's (56, ...)', Modal A10G resume run e1c, cut written by e1a7). The deferred cut now stages a copy of the state dir at the cut (a few KB: ledger, journal) and the background sync snapshots that copy, with the same pointer. Co-Authored-By: Claude Opus 5.5 --- tests/test_rl_resume_from_checkpoint.py | 24 +++++++++++++ yeto/rl/engine/resume.py | 25 ++++++++++--- yeto/rl/engine/round_cut.py | 47 ++++++++++++++++++++----- 3 files changed, 83 insertions(+), 13 deletions(-) diff --git a/tests/test_rl_resume_from_checkpoint.py b/tests/test_rl_resume_from_checkpoint.py index f940719c..f06809c8 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 @@ -830,3 +831,26 @@ def test_a_cut_declined_at_prepare_does_not_wake_the_offloaded_trainer(tmp_path) 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/yeto/rl/engine/resume.py b/yeto/rl/engine/resume.py index b96d36f7..389d45ec 100644 --- a/yeto/rl/engine/resume.py +++ b/yeto/rl/engine/resume.py @@ -392,6 +392,19 @@ def _cut_rollout(cut_id: str) -> int | None: return int(head[1:]) +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()) @@ -527,7 +540,7 @@ def _record(self, kind: str, **fields: Any) -> dict[str, Any]: def sync_checkpoint_store(self, reason: str, *, skip_if_recovery: bool = False, extra: Mapping[str, Any] | None = None, - defer_commit: bool = False) -> bool: + 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. @@ -535,12 +548,16 @@ def sync_checkpoint_store(self, reason: str, *, skip_if_recovery: bool = False, ``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.""" + 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")) @@ -549,7 +566,7 @@ 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, commit=not defer_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 diff --git a/yeto/rl/engine/round_cut.py b/yeto/rl/engine/round_cut.py index 81054489..c8b222ab 100644 --- a/yeto/rl/engine/round_cut.py +++ b/yeto/rl/engine/round_cut.py @@ -261,12 +261,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) @@ -440,6 +456,7 @@ def _complete_deferred(self, deferred: dict[str, Any]) -> dict[str, Any]: 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) @@ -467,11 +484,14 @@ def run() -> None: 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)) + 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"]: @@ -577,7 +597,11 @@ def save(self, driver: Any, *, rollout_id: int, final: bool = False, extra = {"prepared_before_offload": used_prepared, "trainer_onloaded_for_cut": offloaded, "save_s": round(save_s, 3)} if defer: - self._deferred = {"info": info, "cut": cut, "pending": pending, "rollout_id": int(rollout_id)} + # 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: @@ -618,12 +642,17 @@ def _abandon(self, prepared: "_PreparedCut | None") -> None: self._drop_unpointed(prepared.context.cut_id) def _write_and_sync(self, info: dict[str, Any], cut_record: dict[str, Any], *, - defer: bool = False) -> 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, unless - ``defer``: the caller commits in the background).""" - self._write_pointer({**info, CUT_KEY: cut_record}) - kwargs = {"defer_commit": True} if defer else {} + ``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 From 8154509dfc57543d92f3c970a24c99f8bd9b12a0 Mon Sep 17 00:00:00 2001 From: michael <78113554+chowhsu@users.noreply.github.com> Date: Wed, 21 Oct 2026 07:54:58 +0000 Subject: [PATCH 33/55] round cuts: copy a background cut's state on the device first A background cut took its adapter and optimizer copies to the host at the step boundary (0.12-0.14 s per round on A10, one wait per tensor). It now clones them device-to-device there and the writer thread moves them to the host, then computes the N7.6 resume state from the lrs and scheduler progress read at the boundary. The shard is byte-identical to the inline host path. The device path needs CUDA parameters and free memory for twice the copies plus 1 GiB; otherwise, or on a device OOM, the capture takes the host path. Each rank reports its mode and reason; rl_round_cut carries them as cut_capture. YETO_RL_CUT_DEVICE_COPY=0 turns the device path off. The Miles fork takes export_named_optimizer_state(device=None) on branch yeto/s26-e1-device-export (466b89b92, not pushed); for the pinned image the same entries are built from the fork's optimizer_slots. Tests compare both with the fork's host export. Co-Authored-By: Claude Opus 5.5 --- tests/rl_cut_fakes.py | 4 +- tests/rl_reshard_fakes.py | 2 +- tests/test_rl_cut_distopt_pp_ep.py | 93 +++++++++++++ yeto/rl/adapters/miles/cut_plugin.py | 200 +++++++++++++++++++++++---- yeto/rl/adapters/miles/trainer.py | 5 +- yeto/rl/engine/round_cut.py | 4 + 6 files changed, 280 insertions(+), 28 deletions(-) diff --git a/tests/rl_cut_fakes.py b/tests/rl_cut_fakes.py index 8e9d097c..1e74d429 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 5d55ddba..cdf85a69 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 d687e8d3..8ae45c6b 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,93 @@ 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("layout", ["pp2dp2", "ep2"]) +def test_background_device_capture_writes_the_inline_shard(tmp_path, layout, 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")) + 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 device == inline + + +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 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 cp.finish_cut_shard(rank, cut_id="c1") == 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") + monkeypatch.setenv(cp.DEVICE_COPY_ENV, "0") + assert cp._device_copy_plan(None, named) == (False, "disabled") diff --git a/yeto/rl/adapters/miles/cut_plugin.py b/yeto/rl/adapters/miles/cut_plugin.py index 85e095c1..0dd6ed61 100644 --- a/yeto/rl/adapters/miles/cut_plugin.py +++ b/yeto/rl/adapters/miles/cut_plugin.py @@ -100,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``.""" @@ -141,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: @@ -541,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() @@ -576,16 +642,65 @@ 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 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})" + return True, "ok" + + # -------------------------------------------------------------------------- # Plugins # -------------------------------------------------------------------------- @@ -617,7 +732,7 @@ def save_cut_shard(actor: Any, *, directory: str, cut_id: str, background: bool """ try: with trainer_resident(actor): - captured = _capture(actor, 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: @@ -629,7 +744,7 @@ def save_cut_shard(actor: Any, *, directory: str, cut_id: str, background: bool pool = ThreadPoolExecutor(max_workers=1, thread_name_prefix=f"yeto-cut-{cut_id}") _PENDING_SHARDS[cut_id] = (pool, pool.submit(_write_refusing, captured, directory=directory, cut_id=cut_id)) pool.shutdown(wait=False) - return {"pending": cut_id, "coord": dict(captured["coord"])} + return {"pending": cut_id, "coord": dict(captured["coord"]), "capture": dict(captured["capture"])} def finish_cut_shard(actor: Any, *, cut_id: str) -> dict[str, Any]: @@ -645,13 +760,20 @@ def finish_cut_shard(actor: Any, *, cut_id: str) -> dict[str, Any]: def _write_refusing(captured: dict[str, Any], *, directory: str, cut_id: str) -> dict[str, Any]: try: - return _write(captured, directory=directory, cut_id=cut_id) + return _write(_on_host(captured), directory=directory, cut_id=cut_id) except CutPluginError as error: return {REFUSED: str(error), REFUSAL_KIND: "refused"} -def _capture(actor: Any, *, cut_id: str) -> dict[str, Any]: - """Everything the shard and its summary need, read at the step boundary (CPU copies).""" +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) @@ -663,23 +785,53 @@ def _capture(actor: Any, *, cut_id: str) -> dict[str, Any]: 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 any(p.is_cuda for _, p in named): + ready = torch.cuda.Event() + ready.record() marks.append(time.monotonic()) rng = backend.capture_rng() - captured = {"coord": coord, "snap": snap, "rng": 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)} + "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()) - captured["resume"] = resume_rank_state(actor, snap["optimizer_named"]) + 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( + 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:]))))), flush=True) + (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 and the resume state it + defers (runs in the writer thread). Host captures are returned as they are.""" + if "resume" in captured: + return captured + ready = captured.get("ready") + if ready is not None: + ready.synchronize() # the device copies are complete (whatever stream training uses) + snap = to_host(captured["snap"]) + return {**captured, "snap": snap, "ready": None, + "resume": resume_rank_state(None, snap["optimizer_named"], **captured["resume_inputs"])} + + def _write(captured: dict[str, Any], *, directory: str, cut_id: str) -> dict[str, Any]: import torch diff --git a/yeto/rl/adapters/miles/trainer.py b/yeto/rl/adapters/miles/trainer.py index 206aae56..4ccb37fa 100644 --- a/yeto/rl/adapters/miles/trainer.py +++ b/yeto/rl/adapters/miles/trainer.py @@ -841,8 +841,11 @@ def start_cut(self, *, epoch: int, context: "CutContext", background: bool = Tru 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} + "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.""" diff --git a/yeto/rl/engine/round_cut.py b/yeto/rl/engine/round_cut.py index c8b222ab..bcda15d6 100644 --- a/yeto/rl/engine/round_cut.py +++ b/yeto/rl/engine/round_cut.py @@ -596,6 +596,10 @@ def save(self, driver: Any, *, rollout_id: int, final: bool = False, 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 From 13d53ab25266fb54c47d401fd100fb487b0b46f2 Mon Sep 17 00:00:00 2001 From: michael <78113554+chowhsu@users.noreply.github.com> Date: Wed, 21 Oct 2026 07:54:58 +0000 Subject: [PATCH 34/55] publish fast path: stage the digest's tensors with one GPU wait The trainer-side policy digest copied each LoRA tensor to the host on its own (about 400 waits for Qwen3-0.6B). stage_f32_on_host converts and checks the tensors on the GPU and copies them into one reusable pinned buffer with one synchronization. The values, and so both digests, are bit for bit those of the per-tensor copies; a non-finite tensor is named as before. verl: the worker digest stages the same way, validates its own copies without a second clone, and computes the bf16 send checksum in a thread beside the two digest streams. Only the digest path stages; full exports keep their own tensors. Co-Authored-By: Claude Opus 5.5 --- tests/test_rl_publish_fastpath.py | 52 +++++++++++++++++ tests/test_rl_verl_publish_fastpath.py | 28 ++++++++++ yeto/rl/adapters/miles/state_plugin.py | 32 ++++++++--- yeto/rl/adapters/verl/ports_impl.py | 71 +++++++++++++++++------- yeto/rl/engine/policy_digest.py | 77 ++++++++++++++++++++++++++ 5 files changed, 232 insertions(+), 28 deletions(-) diff --git a/tests/test_rl_publish_fastpath.py b/tests/test_rl_publish_fastpath.py index bc88ee05..e755e70a 100644 --- a/tests/test_rl_publish_fastpath.py +++ b/tests/test_rl_publish_fastpath.py @@ -257,3 +257,55 @@ 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) + assert digest_canonical_tensors(staged, **kw) == digest_canonical_tensors(plain, **kw) + # 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() diff --git a/tests/test_rl_verl_publish_fastpath.py b/tests/test_rl_verl_publish_fastpath.py index 4428a65f..bf654321 100644 --- a/tests/test_rl_verl_publish_fastpath.py +++ b/tests/test_rl_verl_publish_fastpath.py @@ -85,3 +85,31 @@ def test_driver_switch_turns_the_fast_path_off(monkeypatch): 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} + want = digest_lora_tensors(tensors, **kwargs) + for _ in range(2): # the second round reuses the worker's buffer + got = pi.remote_export_digest(worker, kwargs) + assert {k: v for k, v in got.items() if k != "export_seconds"} == \ + {k: v for k, v in want.items() if k != "export_seconds"} + assert getattr(worker, "_yeto_digest_staging", None) is not None diff --git a/yeto/rl/adapters/miles/state_plugin.py b/yeto/rl/adapters/miles/state_plugin.py index 5125287e..570e3a0e 100644 --- a/yeto/rl/adapters/miles/state_plugin.py +++ b/yeto/rl/adapters/miles/state_plugin.py @@ -648,6 +648,16 @@ 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).""" + from yeto.rl.engine.policy_digest import HostStaging + + 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, ) -> dict[str, Any] | None: @@ -661,7 +671,10 @@ def export_digest( 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 @@ -912,10 +925,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 +949,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} diff --git a/yeto/rl/adapters/verl/ports_impl.py b/yeto/rl/adapters/verl/ports_impl.py index c7f32936..85f0e41c 100644 --- a/yeto/rl/adapters/verl/ports_impl.py +++ b/yeto/rl/adapters/verl/ports_impl.py @@ -63,19 +63,27 @@ 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)} @@ -89,28 +97,51 @@ def digest_lora_tensors(tensors, *, base_model_revision: str, lora_config_hash: import time - from yeto.rl.core import canonical_state + 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() - state = canonical_state(0, tensors, base_model_revision=base_model_revision, - lora_config_hash=lora_config_hash, layout_hash=layout_hash, - expected_specs=expected_specs) - 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) + # 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": pub.lora_checksum(dict(state.tensors)), + "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.""" + """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 - tensors, _info = remote_export_lora(worker_dict) - if dist.is_available() and dist.is_initialized() and dist.get_rank() != 0: - return None + 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) diff --git a/yeto/rl/engine/policy_digest.py b/yeto/rl/engine/policy_digest.py index cd070275..7f8477cf 100644 --- a/yeto/rl/engine/policy_digest.py +++ b/yeto/rl/engine/policy_digest.py @@ -159,6 +159,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.""" From 19f2648257af6b0ad52af9e5fd428f7c6a30ae20 Mon Sep 17 00:00:00 2001 From: michael <78113554+chowhsu@users.noreply.github.com> Date: Wed, 21 Oct 2026 07:54:58 +0000 Subject: [PATCH 35/55] rollout metadata hook: read each sink value once per hook run Under max_policy_age > 0 the hook read the policy-age limit from the sink actor once per group and twice per trajectory record, each time with a ray.get_actor lookup: about 70 actor calls a round, 0.03 s of the 0.045 s hook on A10. one_sink_view memoizes sink reads and the actor handle within one run of record_trained_groups, the all-samples hook and the buffer filter. A write from the same thread drops the value from the view; other writers are seen from the next run on. Co-Authored-By: Claude Opus 5.5 --- tests/test_rl_meta_hook_sink_view.py | 80 +++++++++++++++++++++ yeto/rl/adapters/miles/rollout_meta_hook.py | 80 ++++++++++++++++++--- 2 files changed, 149 insertions(+), 11 deletions(-) create mode 100644 tests/test_rl_meta_hook_sink_view.py 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 00000000..d66646c3 --- /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/yeto/rl/adapters/miles/rollout_meta_hook.py b/yeto/rl/adapters/miles/rollout_meta_hook.py index 877b5308..ea9e1f22 100644 --- a/yeto/rl/adapters/miles/rollout_meta_hook.py +++ b/yeto/rl/adapters/miles/rollout_meta_hook.py @@ -35,7 +35,9 @@ 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 @@ -151,7 +153,11 @@ 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 @@ -811,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}") @@ -825,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) @@ -836,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: @@ -930,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: @@ -1035,6 +1087,12 @@ 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: # fully_async (round_counts): the trainer is waiting for this batch; the # cross-version estimate is read only after its train step (E1) From 4a1c2457d74b7d8fd994c653ca69022805e446a7 Mon Sep 17 00:00:00 2001 From: michael <78113554+chowhsu@users.noreply.github.com> Date: Wed, 21 Oct 2026 07:54:58 +0000 Subject: [PATCH 36/55] eval: tell an NVIDIA A10 from an A10G Modal A10G containers report NVIDIA A10 or, sometimes, NVIDIA A10G; the A10G trained about 12 % slower (E1 run e1a9). The bare and yeto+verl containers now refuse an A10G when the spec asks for the A10 the other runs had, parsed runs record the GPU names from NVML, and the report refuses to compare runs on different GPU models. Co-Authored-By: Claude Opus 5.5 --- tests/test_tools_eval_efficiency.py | 13 +++++++++++++ tools/eval/efficiency/parse.py | 10 +++++++--- tools/eval/efficiency/report.py | 4 ++++ tools/eval/modal_verl.py | 9 ++++++++- 4 files changed, 32 insertions(+), 4 deletions(-) diff --git a/tests/test_tools_eval_efficiency.py b/tests/test_tools_eval_efficiency.py index 6da76604..c66cc884 100644 --- a/tests/test_tools_eval_efficiency.py +++ b/tests/test_tools_eval_efficiency.py @@ -338,6 +338,19 @@ def test_report_needs_two_repeats_and_one_spec(tmp_path): 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(): diff --git a/tools/eval/efficiency/parse.py b/tools/eval/efficiency/parse.py index 98cd55b9..120d398b 100644 --- a/tools/eval/efficiency/parse.py +++ b/tools/eval/efficiency/parse.py @@ -381,7 +381,7 @@ def _driver_phase_durations(markers: list[dict[str, Any]]) -> dict[int, dict[str # ====================================================================== NVML def parse_nvml(text: str) -> list[dict[str, Any]]: - """Rows {t_unix, key (container:index), util, mem_mb, mem_total_mb, power_w}; header lines skipped.""" + """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] @@ -399,7 +399,8 @@ def parse_nvml(text: str) -> list[dict[str, Any]]: 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 "" - rows.append({"t_unix": t, "key": f"{container}:{index}", "util": float(util) if util else None, + 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 @@ -463,7 +464,10 @@ def parse_run(run_dir: Path) -> dict[str, Any]: "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": parse_nvml(nvml_text), "nvml_interval_s": 20.0 if plan["harness"] == "miles_local" else 1.0, + "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()} diff --git a/tools/eval/efficiency/report.py b/tools/eval/efficiency/report.py index be571fe7..92e46927 100644 --- a/tools/eval/efficiency/report.py +++ b/tools/eval/efficiency/report.py @@ -91,6 +91,10 @@ def markdown(parsed_runs: list[dict[str, Any]], *, title: str) -> str: 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) diff --git a/tools/eval/modal_verl.py b/tools/eval/modal_verl.py index eadc4699..351bb110 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 From 1fde82f4b1ad30b11a82f0d6565346821e81f316 Mon Sep 17 00:00:00 2001 From: michael <78113554+chowhsu@users.noreply.github.com> Date: Wed, 21 Oct 2026 08:03:40 +0000 Subject: [PATCH 37/55] test: compare staged digests without their timing fields hash_seconds and export_seconds are wall times and may differ between two calls on the same tensors (seen on yeto-cpu-test). Co-Authored-By: Claude Opus 5.5 --- tests/test_rl_publish_fastpath.py | 4 +++- tests/test_rl_verl_publish_fastpath.py | 10 ++++++---- 2 files changed, 9 insertions(+), 5 deletions(-) diff --git a/tests/test_rl_publish_fastpath.py b/tests/test_rl_publish_fastpath.py index e755e70a..d35d5bd3 100644 --- a/tests/test_rl_publish_fastpath.py +++ b/tests/test_rl_publish_fastpath.py @@ -284,7 +284,9 @@ def test_staged_copies_are_the_per_tensor_copies_bit_for_bit(): 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) - assert digest_canonical_tensors(staged, **kw) == digest_canonical_tensors(plain, **kw) + 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) diff --git a/tests/test_rl_verl_publish_fastpath.py b/tests/test_rl_verl_publish_fastpath.py index bf654321..186a93e0 100644 --- a/tests/test_rl_verl_publish_fastpath.py +++ b/tests/test_rl_verl_publish_fastpath.py @@ -107,9 +107,11 @@ def test_worker_digest_with_the_staged_export_equals_the_full_export(): 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} - want = digest_lora_tensors(tensors, **kwargs) + 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 - got = pi.remote_export_digest(worker, kwargs) - assert {k: v for k, v in got.items() if k != "export_seconds"} == \ - {k: v for k, v in want.items() if k != "export_seconds"} + assert untimed(pi.remote_export_digest(worker, kwargs)) == want assert getattr(worker, "_yeto_digest_staging", None) is not None From d583c5972f181ca3a5fb6f1fc97d33b3f070df6f Mon Sep 17 00:00:00 2001 From: michael <78113554+chowhsu@users.noreply.github.com> Date: Wed, 21 Oct 2026 08:28:30 +0000 Subject: [PATCH 38/55] round cuts: keep a sleeping trainer's cut on the host path With --offload-train the colocated trainer sleeps right after the cut is prepared, and the writer thread then failed to read the device copies (CUDA error: invalid argument; 5 of 6 cuts refused in GPU run e1g-yeto-miles-sync-1, the run itself went on). Such a trainer now takes the host path, reported as reason offload_train. Co-Authored-By: Claude Opus 5.5 --- tests/test_rl_cut_distopt_pp_ep.py | 2 ++ yeto/rl/adapters/miles/cut_plugin.py | 10 ++++++++-- 2 files changed, 10 insertions(+), 2 deletions(-) diff --git a/tests/test_rl_cut_distopt_pp_ep.py b/tests/test_rl_cut_distopt_pp_ep.py index 8ae45c6b..4a8e150f 100644 --- a/tests/test_rl_cut_distopt_pp_ep.py +++ b/tests/test_rl_cut_distopt_pp_ep.py @@ -446,5 +446,7 @@ def oom(actor, backend, named, *, on_device=False): 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)) + assert cp._device_copy_plan(sleeper, named) == (False, "offload_train") monkeypatch.setenv(cp.DEVICE_COPY_ENV, "0") assert cp._device_copy_plan(None, named) == (False, "disabled") diff --git a/yeto/rl/adapters/miles/cut_plugin.py b/yeto/rl/adapters/miles/cut_plugin.py index 0dd6ed61..179418cb 100644 --- a/yeto/rl/adapters/miles/cut_plugin.py +++ b/yeto/rl/adapters/miles/cut_plugin.py @@ -685,12 +685,18 @@ def to_host(value: Any) -> Any: def _device_copy_plan(actor: Any, named: list) -> tuple[bool, str]: - """(copy on the device?, why): the device path needs CUDA parameters and free - memory for the copies (adapter + FP32 master + two Adam moments, estimated).""" + """(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" + if getattr(getattr(actor, "args", None), "offload_train", False): + # --offload-train: the trainer sleeps right after the cut (colocated) and the + # writer's reads of the device copies then fail ("CUDA error: invalid + # argument", 5 of 6 cuts of GPU run e1g-yeto-miles-sync-1) + return False, "offload_train" devices = {p.device for _, p in named} if len(devices) != 1 or next(iter(devices)).type != "cuda": return False, "not_cuda" From 2f3d79c36efa61b249497ce28493c356170fa14d Mon Sep 17 00:00:00 2001 From: michael <78113554+chowhsu@users.noreply.github.com> Date: Wed, 21 Oct 2026 08:59:04 +0000 Subject: [PATCH 39/55] round cuts: serialize and hash background shards in a writer process The background cut writer ran torch.save and four sha256 passes in a thread of the trainer process while the next step trained; each trained token was about 8 % slower in those rounds (GPU runs e1g/e1h). The trainer thread now copies the device snapshot into shared host memory and a per-rank spawn writer process (nice 10, one intra-op thread) serializes, hashes and summarizes it. A writer-process failure removes its partial file and the thread writes the same shard; the summary's writer field and rl_cut_committed record the mode and the fallback. YETO_RL_CUT_WRITER_PROCESS=0 keeps the thread. Strings are interned before torch.save so the shard bytes depend on values only: a shard written in the writer process is byte-identical to the inline one (tests compare file bytes, digests and the N7.6 state). Co-Authored-By: Claude Opus 5.5 --- tests/test_rl_cut_distopt_pp_ep.py | 75 +++++++++++- tests/test_rl_resume_from_checkpoint.py | 2 + yeto/rl/adapters/miles/cut_plugin.py | 151 +++++++++++++++++++++--- yeto/rl/adapters/miles/trainer.py | 3 + yeto/rl/engine/round_cut.py | 4 +- 5 files changed, 217 insertions(+), 18 deletions(-) diff --git a/tests/test_rl_cut_distopt_pp_ep.py b/tests/test_rl_cut_distopt_pp_ep.py index 4a8e150f..849e288e 100644 --- a/tests/test_rl_cut_distopt_pp_ep.py +++ b/tests/test_rl_cut_distopt_pp_ep.py @@ -414,7 +414,10 @@ def test_background_device_capture_writes_the_inline_shard(tmp_path, layout, mon 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 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): @@ -426,7 +429,7 @@ def test_background_capture_falls_back_to_the_host_and_says_why(tmp_path, monkey 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 cp.finish_cut_shard(rank, cut_id="c1") == inline + assert _without_writer(cp.finish_cut_shard(rank, cut_id="c1")) == inline real = cp._snapshot @@ -440,7 +443,75 @@ def oom(actor, backend, named, *, on_device=False): 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 cp.finish_cut_shard(rank, cut_id="c1") == inline + 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): diff --git a/tests/test_rl_resume_from_checkpoint.py b/tests/test_rl_resume_from_checkpoint.py index f06809c8..c3e43461 100644 --- a/tests/test_rl_resume_from_checkpoint.py +++ b/tests/test_rl_resume_from_checkpoint.py @@ -632,6 +632,8 @@ def test_prepared_cut_is_written_in_the_background_without_waking_the_trainer(tm 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 diff --git a/yeto/rl/adapters/miles/cut_plugin.py b/yeto/rl/adapters/miles/cut_plugin.py index 179418cb..1e7d0b80 100644 --- a/yeto/rl/adapters/miles/cut_plugin.py +++ b/yeto/rl/adapters/miles/cut_plugin.py @@ -732,9 +732,10 @@ def save_cut_shard(actor: Any, *, directory: str, cut_id: str, background: bool 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 (CPU copies, at this step boundary) and - write, hash and summarize it in a thread of this rank, so training goes on; - returns ``{"pending": cut_id}``. :func:`finish_cut_shard` returns the summary. + ``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): @@ -748,13 +749,14 @@ def save_cut_shard(actor: Any, *, directory: str, cut_id: str, background: bool 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_refusing, captured, directory=directory, cut_id=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).""" + """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"} @@ -766,11 +768,104 @@ def finish_cut_shard(actor: Any, *, cut_id: str) -> dict[str, Any]: def _write_refusing(captured: dict[str, Any], *, directory: str, cut_id: str) -> dict[str, Any]: try: - return _write(_on_host(captured), directory=directory, cut_id=cut_id) + 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: device copies to shared host memory, then the + writer process serializes, hashes and summarizes them. 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: + host = _share(_on_host(captured)) + except CutPluginError as error: + return {REFUSED: str(error), REFUSAL_KIND: "refused"} + fallback = "disabled" + if _writer_enabled(): + try: + result = _writer_pool().submit(_write_refusing, host, 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, host["coord"]) + return {**_write_refusing(host, directory=directory, cut_id=cut_id), + "writer": {"mode": "thread", "fallback": fallback}} + + +def _share(value: Any) -> Any: + """``value`` with every CPU tensor moved into shared memory (in place), so the + writer process maps it instead of receiving a copy through the pipe.""" + import torch + + if isinstance(value, torch.Tensor): + return value.share_memory_() if value.device.type == "cpu" else value + if isinstance(value, dict): + return {k: _share(v) for k, v in value.items()} + if isinstance(value, list): + return [_share(v) for v in value] + if isinstance(value, tuple): + return tuple(_share(v) for v in value) + return value + + 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. @@ -826,16 +921,23 @@ def _capture(actor: Any, *, cut_id: str, device_copy: bool = False) -> dict[str, def _on_host(captured: dict[str, Any]) -> dict[str, Any]: - """``captured`` with its device copies on the host and the resume state it - defers (runs in the writer thread). Host captures are returned as they are.""" - if "resume" in captured: - return captured + """``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) - snap = to_host(captured["snap"]) - return {**captured, "snap": snap, "ready": None, - "resume": resume_rank_state(None, snap["optimizer_named"], **captured["resume_inputs"])} + 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]: @@ -857,13 +959,13 @@ def _write(captured: dict[str, Any], *, directory: str, cut_id: str) -> dict[str 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) @@ -888,6 +990,25 @@ def _write(captured: dict[str, Any], *, directory: str, cut_id: str) -> dict[str } +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/trainer.py b/yeto/rl/adapters/miles/trainer.py index 4ccb37fa..ab360271 100644 --- a/yeto/rl/adapters/miles/trainer.py +++ b/yeto/rl/adapters/miles/trainer.py @@ -868,6 +868,9 @@ def collect_cut(self, pending: Mapping[str, Any]) -> Callable[[], str]: 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: diff --git a/yeto/rl/engine/round_cut.py b/yeto/rl/engine/round_cut.py index bcda15d6..fb1143ff 100644 --- a/yeto/rl/engine/round_cut.py +++ b/yeto/rl/engine/round_cut.py @@ -437,6 +437,7 @@ def _report(driver: Any, result: Mapping[str, Any]) -> 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]: @@ -460,8 +461,9 @@ def _complete_deferred(self, deferred: dict[str, Any]) -> dict[str, Any]: 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)} + "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 From 27f66dd8759022c61a80e4c103225c594d5c8bcb Mon Sep 17 00:00:00 2001 From: michael <78113554+chowhsu@users.noreply.github.com> Date: Wed, 21 Oct 2026 09:01:10 +0000 Subject: [PATCH 40/55] launcher: pin a Modal island to one GPU model with --modal-gpu-name Modal A10G containers are an NVIDIA A10 or an A10G, and the A10G trains about 12 % slower (E1 run e1a9). --modal-gpu-name MODEL fails an island at container start unless every GPU nvidia-smi reports is that model ("NVIDIA A10" refuses "NVIDIA A10G"). The efficiency eval passes it for the yeto+Miles arm with the model the bare arms require, so every arm runs on the same card. Co-Authored-By: Claude Opus 5.5 --- tests/test_rl_ports_image.py | 30 +++++++++++++++++++++++++++++ tests/test_tools_eval_efficiency.py | 5 +++++ tools/eval/efficiency/spec.py | 11 +++++++++++ yeto/cli.py | 8 ++++++++ yeto/launch/modal_islands.py | 12 ++++++------ yeto/modal_runner.py | 21 +++++++++++++++++++- 6 files changed, 80 insertions(+), 7 deletions(-) diff --git a/tests/test_rl_ports_image.py b/tests/test_rl_ports_image.py index 15585c2c..c6b755fc 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_tools_eval_efficiency.py b/tests/test_tools_eval_efficiency.py index c66cc884..94911c11 100644 --- a/tests/test_tools_eval_efficiency.py +++ b/tests/test_tools_eval_efficiency.py @@ -85,6 +85,11 @@ def test_yeto_arms_use_the_same_policy_age_and_partition(): 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") diff --git a/tools/eval/efficiency/spec.py b/tools/eval/efficiency/spec.py index f4ef5450..174e9eb6 100644 --- a/tools/eval/efficiency/spec.py +++ b/tools/eval/efficiency/spec.py @@ -373,6 +373,16 @@ def cut_every(spec: EffSpec) -> int: return spec.rounds if spec.checkpoint == "final" else 1 +# 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 @@ -381,6 +391,7 @@ def yeto_miles_launcher_args(spec: EffSpec, mode: str, run_id: str) -> str: 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} " diff --git a/yeto/cli.py b/yeto/cli.py index 7c68497f..5e1c1d5c 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 a7f12af6..e6e7f86f 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/modal_runner.py b/yeto/modal_runner.py index b25be3c4..761b53c9 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: From d8d159c1a20a7b0d2649e1c1e88274b64782e59f Mon Sep 17 00:00:00 2001 From: michael <78113554+chowhsu@users.noreply.github.com> Date: Wed, 21 Oct 2026 09:05:49 +0000 Subject: [PATCH 41/55] publish fastpath: hash the policy while a colocated trainer offloads In sync colocated mode the trainer hashed the staged policy (two sha256 streams, about 0.1 s) before it could offload. export_digest now takes deferred=True: the trainer returns once the tensors are staged in its pinned buffer and hashes them in a thread; the state's DeferredDigest collects the result (collect_digest plugin) on first read, at the publication. The values are those of the eager digest; a new export waits for an uncollected hash before it reuses the buffer. The cut context reads the policy hash last. The driver defers only for a colocated trainer that offloads to publish (YETO_RL_DEFERRED_DIGEST=0 turns it off). Co-Authored-By: Claude Opus 5.5 --- tests/test_rl_publish_fastpath.py | 51 ++++++++++++++++++++ yeto/rl/adapters/miles/state.py | 35 ++++++++++++-- yeto/rl/adapters/miles/state_plugin.py | 67 ++++++++++++++++++++++---- yeto/rl/engine/driver.py | 14 ++++++ yeto/rl/engine/policy_digest.py | 43 +++++++++++++++++ yeto/rl/engine/round_cut.py | 5 +- 6 files changed, 199 insertions(+), 16 deletions(-) diff --git a/tests/test_rl_publish_fastpath.py b/tests/test_rl_publish_fastpath.py index d35d5bd3..54877fd8 100644 --- a/tests/test_rl_publish_fastpath.py +++ b/tests/test_rl_publish_fastpath.py @@ -311,3 +311,54 @@ def test_export_digest_is_stable_across_rounds_with_the_reused_buffer(): 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 diff --git a/yeto/rl/adapters/miles/state.py b/yeto/rl/adapters/miles/state.py index 175296cf..49120156 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,9 @@ 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 + def __init__( self, *, @@ -117,25 +121,39 @@ 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 _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 +212,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( @@ -304,6 +327,8 @@ 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 + 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 570e3a0e..312eaed7 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" @@ -649,9 +650,11 @@ def export_state(actor: Any, *, policy_version: int) -> dict[str, Any] | None: def _digest_staging(actor: Any) -> Any: - """The pinned host buffer this trainer's digests reuse every round (E1).""" + """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() @@ -660,11 +663,17 @@ def _digest_staging(actor: Any) -> Any: 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 @@ -680,14 +689,52 @@ def export_digest( 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" diff --git a/yeto/rl/engine/driver.py b/yeto/rl/engine/driver.py index 6ae1a5b7..f6ea1afb 100644 --- a/yeto/rl/engine/driver.py +++ b/yeto/rl/engine/driver.py @@ -368,6 +368,18 @@ 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))) + + class IslandDriver: def __init__( self, @@ -687,6 +699,8 @@ 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 defer_digest(self): + return export_digest(policy_version=int(policy_version), deferred=True) return export_digest(policy_version=int(policy_version)) # -- per-round steps ------------------------------------------------- diff --git a/yeto/rl/engine/policy_digest.py b/yeto/rl/engine/policy_digest.py index 7f8477cf..7d58afbd 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 diff --git a/yeto/rl/engine/round_cut.py b/yeto/rl/engine/round_cut.py index fb1143ff..ce207118 100644 --- a/yeto/rl/engine/round_cut.py +++ b/yeto/rl/engine/round_cut.py @@ -109,8 +109,9 @@ def context(self, cut_id: str, *, pending: tuple[int, Any] | None = None) -> Any driver = self.driver 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, policy_token(version, state.policy_tensor_hash()) + settled, token = True, None else: version, state = driver.published_version, driver.published_state settled, token = bool(driver.at_safe_point), driver.expected_token @@ -152,6 +153,8 @@ def context(self, cut_id: str, *, pending: tuple[int, Any] | None = None) -> Any 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, From 1c0df83e7ad860227579285523284abb094b4c96 Mon Sep 17 00:00:00 2001 From: michael <78113554+chowhsu@users.noreply.github.com> Date: Wed, 21 Oct 2026 10:09:46 +0000 Subject: [PATCH 42/55] eval: opt-in trainer step profile (--rl-profile-trainer) Observation only, for E1: around each train_one_step the trainer process prints the step's wall time, main-thread and process CPU time, garbage collector pauses, CUDA caching-allocator counters and stack samples of every Python thread. The launcher forwards the flag to the island, the island sets it on the Miles args, and eff-run --profile-trainer passes it for the yeto+Miles arms (not in the control fingerprint). Co-Authored-By: Claude Opus 5.5 --- tests/test_rl_miles_train_profile.py | 51 ++++++++++ tools/eval/__main__.py | 5 +- tools/eval/efficiency/spec.py | 6 +- yeto/launch/rl_ports.py | 2 + yeto/rl/adapters/miles/island_entry.py | 2 + yeto/rl/adapters/miles/state_plugin.py | 7 +- yeto/rl/adapters/miles/train_profile.py | 121 ++++++++++++++++++++++++ yeto/rl/engine/run_flags.py | 3 + 8 files changed, 194 insertions(+), 3 deletions(-) create mode 100644 tests/test_rl_miles_train_profile.py create mode 100644 yeto/rl/adapters/miles/train_profile.py diff --git a/tests/test_rl_miles_train_profile.py b/tests/test_rl_miles_train_profile.py new file mode 100644 index 00000000..1ef45ba7 --- /dev/null +++ b/tests/test_rl_miles_train_profile.py @@ -0,0 +1,51 @@ +"""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(): + import argparse + 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 diff --git a/tools/eval/__main__.py b/tools/eval/__main__.py index cbb998b6..8160439f 100644 --- a/tools/eval/__main__.py +++ b/tools/eval/__main__.py @@ -148,7 +148,8 @@ def cmd_eff_plan(args) -> int: def cmd_eff_run(args) -> int: from .efficiency import modal_eff, spec - sp = spec.get_spec(args.preset, checkpoint=args.checkpoint) + 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: @@ -230,6 +231,8 @@ def main(argv: list[str] | None = None) -> int: er.add_argument("--checkpoint", choices=("final", "every-round"), default="final", help="yeto+Miles round cuts: only at the end (fair control) or every round " "(product default)") + 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") diff --git a/tools/eval/efficiency/spec.py b/tools/eval/efficiency/spec.py index 174e9eb6..1b582f28 100644 --- a/tools/eval/efficiency/spec.py +++ b/tools/eval/efficiency/spec.py @@ -84,6 +84,8 @@ class EffSpec: # (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: @@ -149,7 +151,7 @@ def cards(mode: str) -> dict[str, int]: 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")} + 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() @@ -411,6 +413,8 @@ def yeto_miles_launcher_args(spec: EffSpec, mode: str, run_id: str) -> str: 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": diff --git a/yeto/launch/rl_ports.py b/yeto/launch/rl_ports.py index 63b41f24..9e78b20b 100644 --- a/yeto/launch/rl_ports.py +++ b/yeto/launch/rl_ports.py @@ -533,6 +533,8 @@ 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_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/rl/adapters/miles/island_entry.py b/yeto/rl/adapters/miles/island_entry.py index 48585224..50b053d9 100644 --- a/yeto/rl/adapters/miles/island_entry.py +++ b/yeto/rl/adapters/miles/island_entry.py @@ -282,6 +282,8 @@ 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_profile_trainer", False): # E1: read by state_plugin.train_one_step + miles_args.yeto_rl_profile_trainer = True 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"))) diff --git a/yeto/rl/adapters/miles/state_plugin.py b/yeto/rl/adapters/miles/state_plugin.py index 312eaed7..3e068497 100644 --- a/yeto/rl/adapters/miles/state_plugin.py +++ b/yeto/rl/adapters/miles/state_plugin.py @@ -1252,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 00000000..bfa12cc3 --- /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/engine/run_flags.py b/yeto/rl/engine/run_flags.py index b4de9148..38491c91 100644 --- a/yeto/rl/engine/run_flags.py +++ b/yeto/rl/engine/run_flags.py @@ -495,6 +495,9 @@ 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-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)"), From fe5e612ef37a64f45d5c187c71acf6070e27ec01 Mon Sep 17 00:00:00 2001 From: michael <78113554+chowhsu@users.noreply.github.com> Date: Wed, 21 Oct 2026 10:54:50 +0000 Subject: [PATCH 43/55] round cuts: stage a background cut's device copies on the host at the step The writer thread copied the device snapshot to the host tensor by tensor while the next round trained: about 400 synchronous device-to-host copies interleaved with the training kernels, the suspected cause of the slower training (profile run e1j-yeto-miles-async-a: the step grew while GC, allocator and main-thread CPU stayed as in the final-only run). The capture now copies the device snapshot into one reused pinned host buffer with a single synchronization, drops the device copies, and the writer thread copies the staged views into shared memory (the buffer is reused only after that). Since nothing reads device memory after the capture returns, a trainer that sleeps right after the cut (--offload-train) takes the device path too. Tests compare the shard bytes and summaries of staged, unstaged and inline cuts. Co-Authored-By: Claude Opus 5.5 --- tests/test_rl_cut_distopt_pp_ep.py | 11 +- yeto/rl/adapters/miles/cut_plugin.py | 144 +++++++++++++++++++++++++-- 2 files changed, 142 insertions(+), 13 deletions(-) diff --git a/tests/test_rl_cut_distopt_pp_ep.py b/tests/test_rl_cut_distopt_pp_ep.py index 849e288e..57bf2da7 100644 --- a/tests/test_rl_cut_distopt_pp_ep.py +++ b/tests/test_rl_cut_distopt_pp_ep.py @@ -398,14 +398,18 @@ def export(optimizer, named, *, device="cpu"): @needs_dps +@pytest.mark.parametrize("staged", [False, True], ids=["unstaged", "staged"]) @pytest.mark.parametrize("layout", ["pp2dp2", "ep2"]) -def test_background_device_capture_writes_the_inline_shard(tmp_path, layout, monkeypatch): +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: # CPU stand-ins for device tensors go through the staging buffer too + real_stage = cp._stage + 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) @@ -415,6 +419,7 @@ def test_background_device_capture_writes_the_inline_shard(tmp_path, layout, mon 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() 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() @@ -517,7 +522,7 @@ def test_the_writer_process_can_be_turned_off(tmp_path, monkeypatch): 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)) - assert cp._device_copy_plan(sleeper, named) == (False, "offload_train") + sleeper = SimpleNamespace(args=SimpleNamespace(offload_train=True)) # staged before the sleep + 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") diff --git a/yeto/rl/adapters/miles/cut_plugin.py b/yeto/rl/adapters/miles/cut_plugin.py index 1e7d0b80..6b253106 100644 --- a/yeto/rl/adapters/miles/cut_plugin.py +++ b/yeto/rl/adapters/miles/cut_plugin.py @@ -692,11 +692,10 @@ def _device_copy_plan(actor: Any, named: list) -> tuple[bool, str]: if os.environ.get(DEVICE_COPY_ENV, "1").strip().lower() in {"0", "false", "no", "off"}: return False, "disabled" - if getattr(getattr(actor, "args", None), "offload_train", False): - # --offload-train: the trainer sleeps right after the cut (colocated) and the - # writer's reads of the device copies then fail ("CUDA error: invalid - # argument", 5 of 6 cuts of GPU run e1g-yeto-miles-sync-1) - return False, "offload_train" + # --offload-train (the trainer sleeps right after the cut) is fine since E1 phase 6: + # the device copies reach the host buffer before the capture returns (_stage); the + # writer used to read them after the sleep ("CUDA error: invalid argument", GPU run + # e1g-yeto-miles-sync-1) devices = {p.device for _, p in named} if len(devices) != 1 or next(iter(devices)).type != "cuda": return False, "not_cuda" @@ -836,6 +835,8 @@ def _write_background(captured: dict[str, Any], *, directory: str, cut_id: str) host = _share(_on_host(captured)) except CutPluginError as error: return {REFUSED: str(error), REFUSAL_KIND: "refused"} + finally: + _STAGING.free.set() # also when the hand-over failed fallback = "disabled" if _writer_enabled(): try: @@ -850,11 +851,112 @@ def _write_background(captured: dict[str, Any], *, directory: str, cut_id: str) "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:`_share` 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:`_share` copies them out (it marks the buffer free again).""" + 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) + 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) + _STAGING.free.clear() + 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 {"__yeto_staged__": True, "snap": out} + + def _share(value: Any) -> Any: - """``value`` with every CPU tensor moved into shared memory (in place), so the - writer process maps it instead of receiving a copy through the pipe.""" + """``value`` with every CPU tensor in shared memory, so the writer process maps it + instead of receiving a copy through the pipe: moved in place, or, for a snapshot + staged by :func:`_stage`, copied out of the staging buffer (which is then free).""" import torch + if isinstance(value, Mapping) and value.get("__yeto_staged__"): + try: + return _share_copy(value["snap"]) + finally: + _STAGING.free.set() if isinstance(value, torch.Tensor): return value.share_memory_() if value.device.type == "cpu" else value if isinstance(value, dict): @@ -866,6 +968,20 @@ def _share(value: Any) -> Any: return value +def _share_copy(value: Any) -> Any: + import torch + + if isinstance(value, torch.Tensor): + return torch.empty_like(value, device="cpu").share_memory_().copy_(value) + if isinstance(value, dict): + return {k: _share_copy(v) for k, v in value.items()} + if isinstance(value, list): + return [_share_copy(v) for v in value] + if isinstance(value, tuple): + return tuple(_share_copy(v) for v in value) + return value + + 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. @@ -898,9 +1014,17 @@ def _capture(actor: Any, *, cut_id: str, device_copy: bool = False) -> dict[str, if snap is None: snap = _snapshot(actor, backend, named) ready = None - if on_device and any(p.is_cuda for _, p in named): - ready = torch.cuda.Event() - ready.record() + if on_device: + # E1 phase 6: one copy of the device snapshot into a reused (pinned) host + # buffer, one synchronization, here at the step boundary; the device copies + # are dropped. The writer thread used to copy them tensor by tensor while the + # next round trained (~400 synchronous device-to-host copies interleaved + # with training kernels, profile run e1j-yeto-miles-async-a). + snap = _stage(snap) + if getattr(getattr(actor, "args", None), "offload_train", False) and torch.cuda.is_available(): + # the trainer sleeps next: hand the dropped copies' blocks back to the + # device (the rollout engine shares it) + torch.cuda.empty_cache() marks.append(time.monotonic()) rng = backend.capture_rng() capture = {"mode": "device" if on_device else "host", "reason": reason} From 40cf73a4162d1ee7ae8503f90b60f3845ea098d1 Mon Sep 17 00:00:00 2001 From: michael <78113554+chowhsu@users.noreply.github.com> Date: Wed, 21 Oct 2026 11:41:33 +0000 Subject: [PATCH 44/55] Revert "round cuts: stage a background cut's device copies on the host at the step" This reverts commit fe5e612e. GPU runs e1k-yeto-miles-async-2 and e1k-yeto-miles-sync-1: the trained tokens were not faster (0.553 ms per token, as before) and the step-boundary capture got slower (async 0.033 -> 0.08 s, sync host path 0.16-0.28 -> 0.30 s). The interleaved device-to-host copies were not the cause of the slower training. Co-Authored-By: Claude Opus 5.5 --- tests/test_rl_cut_distopt_pp_ep.py | 11 +- yeto/rl/adapters/miles/cut_plugin.py | 144 ++------------------------- 2 files changed, 13 insertions(+), 142 deletions(-) diff --git a/tests/test_rl_cut_distopt_pp_ep.py b/tests/test_rl_cut_distopt_pp_ep.py index 57bf2da7..849e288e 100644 --- a/tests/test_rl_cut_distopt_pp_ep.py +++ b/tests/test_rl_cut_distopt_pp_ep.py @@ -398,18 +398,14 @@ def export(optimizer, named, *, device="cpu"): @needs_dps -@pytest.mark.parametrize("staged", [False, True], ids=["unstaged", "staged"]) @pytest.mark.parametrize("layout", ["pp2dp2", "ep2"]) -def test_background_device_capture_writes_the_inline_shard(tmp_path, layout, staged, monkeypatch): +def test_background_device_capture_writes_the_inline_shard(tmp_path, layout, 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: # CPU stand-ins for device tensors go through the staging buffer too - real_stage = cp._stage - 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) @@ -419,7 +415,6 @@ def test_background_device_capture_writes_the_inline_shard(tmp_path, layout, sta 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() 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() @@ -522,7 +517,7 @@ def test_the_writer_process_can_be_turned_off(tmp_path, monkeypatch): 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 before the sleep - assert cp._device_copy_plan(sleeper, named) == (False, "not_cuda") + sleeper = SimpleNamespace(args=SimpleNamespace(offload_train=True)) + assert cp._device_copy_plan(sleeper, named) == (False, "offload_train") monkeypatch.setenv(cp.DEVICE_COPY_ENV, "0") assert cp._device_copy_plan(None, named) == (False, "disabled") diff --git a/yeto/rl/adapters/miles/cut_plugin.py b/yeto/rl/adapters/miles/cut_plugin.py index 6b253106..1e7d0b80 100644 --- a/yeto/rl/adapters/miles/cut_plugin.py +++ b/yeto/rl/adapters/miles/cut_plugin.py @@ -692,10 +692,11 @@ def _device_copy_plan(actor: Any, named: list) -> tuple[bool, str]: if os.environ.get(DEVICE_COPY_ENV, "1").strip().lower() in {"0", "false", "no", "off"}: return False, "disabled" - # --offload-train (the trainer sleeps right after the cut) is fine since E1 phase 6: - # the device copies reach the host buffer before the capture returns (_stage); the - # writer used to read them after the sleep ("CUDA error: invalid argument", GPU run - # e1g-yeto-miles-sync-1) + if getattr(getattr(actor, "args", None), "offload_train", False): + # --offload-train: the trainer sleeps right after the cut (colocated) and the + # writer's reads of the device copies then fail ("CUDA error: invalid + # argument", 5 of 6 cuts of GPU run e1g-yeto-miles-sync-1) + return False, "offload_train" devices = {p.device for _, p in named} if len(devices) != 1 or next(iter(devices)).type != "cuda": return False, "not_cuda" @@ -835,8 +836,6 @@ def _write_background(captured: dict[str, Any], *, directory: str, cut_id: str) host = _share(_on_host(captured)) except CutPluginError as error: return {REFUSED: str(error), REFUSAL_KIND: "refused"} - finally: - _STAGING.free.set() # also when the hand-over failed fallback = "disabled" if _writer_enabled(): try: @@ -851,112 +850,11 @@ def _write_background(captured: dict[str, Any], *, directory: str, cut_id: str) "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:`_share` 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:`_share` copies them out (it marks the buffer free again).""" - 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) - 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) - _STAGING.free.clear() - 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 {"__yeto_staged__": True, "snap": out} - - def _share(value: Any) -> Any: - """``value`` with every CPU tensor in shared memory, so the writer process maps it - instead of receiving a copy through the pipe: moved in place, or, for a snapshot - staged by :func:`_stage`, copied out of the staging buffer (which is then free).""" + """``value`` with every CPU tensor moved into shared memory (in place), so the + writer process maps it instead of receiving a copy through the pipe.""" import torch - if isinstance(value, Mapping) and value.get("__yeto_staged__"): - try: - return _share_copy(value["snap"]) - finally: - _STAGING.free.set() if isinstance(value, torch.Tensor): return value.share_memory_() if value.device.type == "cpu" else value if isinstance(value, dict): @@ -968,20 +866,6 @@ def _share(value: Any) -> Any: return value -def _share_copy(value: Any) -> Any: - import torch - - if isinstance(value, torch.Tensor): - return torch.empty_like(value, device="cpu").share_memory_().copy_(value) - if isinstance(value, dict): - return {k: _share_copy(v) for k, v in value.items()} - if isinstance(value, list): - return [_share_copy(v) for v in value] - if isinstance(value, tuple): - return tuple(_share_copy(v) for v in value) - return value - - 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. @@ -1014,17 +898,9 @@ def _capture(actor: Any, *, cut_id: str, device_copy: bool = False) -> dict[str, if snap is None: snap = _snapshot(actor, backend, named) ready = None - if on_device: - # E1 phase 6: one copy of the device snapshot into a reused (pinned) host - # buffer, one synchronization, here at the step boundary; the device copies - # are dropped. The writer thread used to copy them tensor by tensor while the - # next round trained (~400 synchronous device-to-host copies interleaved - # with training kernels, profile run e1j-yeto-miles-async-a). - snap = _stage(snap) - if getattr(getattr(actor, "args", None), "offload_train", False) and torch.cuda.is_available(): - # the trainer sleeps next: hand the dropped copies' blocks back to the - # device (the rollout engine shares it) - torch.cuda.empty_cache() + if 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} From 74c2f6fbbb23187c10b5f99572e03976d5ddc176 Mon Sep 17 00:00:00 2001 From: michael <78113554+chowhsu@users.noreply.github.com> Date: Wed, 21 Oct 2026 11:43:08 +0000 Subject: [PATCH 45/55] round cuts: hand a background cut to its writer in one shared buffer Moving each tensor to shared memory passed one file descriptor per tensor (about 400 per cut) through the trainer process's resource sharer, and its queue feeder and server threads ran while the next round trained (profile run e1j-yeto-miles-async-a: about 15 % of the step's stack samples; none in the final-only run). The writer thread now copies every contiguous tensor that owns its whole storage into one shared uint8 buffer and sends offsets; the writer process rebuilds each tensor as its own copy. Other tensors are shared one by one as before. The shard bytes and summaries are unchanged (tests compare them with the inline cut). Co-Authored-By: Claude Opus 5.5 --- tests/test_rl_cut_distopt_pp_ep.py | 14 ++++ yeto/rl/adapters/miles/cut_plugin.py | 108 +++++++++++++++++++++++---- 2 files changed, 109 insertions(+), 13 deletions(-) diff --git a/tests/test_rl_cut_distopt_pp_ep.py b/tests/test_rl_cut_distopt_pp_ep.py index 849e288e..0285d2e7 100644 --- a/tests/test_rl_cut_distopt_pp_ep.py +++ b/tests/test_rl_cut_distopt_pp_ep.py @@ -521,3 +521,17 @@ def test_device_copy_plan_needs_cuda_and_can_be_turned_off(monkeypatch): assert cp._device_copy_plan(sleeper, named) == (False, "offload_train") 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/yeto/rl/adapters/miles/cut_plugin.py b/yeto/rl/adapters/miles/cut_plugin.py index 1e7d0b80..849a08ff 100644 --- a/yeto/rl/adapters/miles/cut_plugin.py +++ b/yeto/rl/adapters/miles/cut_plugin.py @@ -833,13 +833,15 @@ def _write_background(captured: dict[str, Any], *, directory: str, cut_id: str) 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: - host = _share(_on_host(captured)) + host = _on_host(captured) except CutPluginError as error: return {REFUSED: str(error), REFUSAL_KIND: "refused"} fallback = "disabled" if _writer_enabled(): try: - result = _writer_pool().submit(_write_refusing, host, directory=directory, cut_id=cut_id).result() + skeleton, layout, flat = _pack(host) + 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() @@ -850,20 +852,100 @@ def _write_background(captured: dict[str, Any], *, directory: str, cut_id: str) "writer": {"mode": "thread", "fallback": fallback}} -def _share(value: Any) -> Any: - """``value`` with every CPU tensor moved into shared memory (in place), so the - writer process maps it instead of receiving a copy through the pipe.""" +_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 + + def packable(t: Any) -> bool: + storage = t.untyped_storage() + 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): - return value.share_memory_() if value.device.type == "cpu" else value - if isinstance(value, dict): - return {k: _share(v) for k, v in value.items()} - if isinstance(value, list): - return [_share(v) for v in value] - if isinstance(value, tuple): - return tuple(_share(v) for v in value) - return value + 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]: From b77491c9e659d2f305750b413ca9e86d4a6f3d03 Mon Sep 17 00:00:00 2001 From: michael <78113554+chowhsu@users.noreply.github.com> Date: Wed, 21 Oct 2026 12:34:20 +0000 Subject: [PATCH 46/55] eval: sample the island driver per phase under --rl-profile-trainer Observation only, for E1: a daemon thread samples the driver thread's stack every 5 ms and, at each driver phase change, prints the finished phase's wall time and its most frequent yeto/miles/ray frames. Without the flag the phase hook does nothing. Co-Authored-By: Claude Opus 5.5 --- tests/test_rl_miles_train_profile.py | 16 +++++ yeto/rl/adapters/miles/island_entry.py | 3 + yeto/rl/engine/driver.py | 2 + yeto/rl/engine/phase_profile.py | 82 ++++++++++++++++++++++++++ 4 files changed, 103 insertions(+) create mode 100644 yeto/rl/engine/phase_profile.py diff --git a/tests/test_rl_miles_train_profile.py b/tests/test_rl_miles_train_profile.py index 1ef45ba7..883bb45a 100644 --- a/tests/test_rl_miles_train_profile.py +++ b/tests/test_rl_miles_train_profile.py @@ -49,3 +49,19 @@ def test_the_flag_is_forwarded_to_the_island_and_set_on_the_miles_args(): 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 + phase_profile.enable() + phase_profile.on_phase("train", 1) + time.sleep(0.05) + phase_profile.on_phase("sync", 1) + (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 + phase_profile._STATE["thread"].phase = None diff --git a/yeto/rl/adapters/miles/island_entry.py b/yeto/rl/adapters/miles/island_entry.py index 50b053d9..4bb7af54 100644 --- a/yeto/rl/adapters/miles/island_entry.py +++ b/yeto/rl/adapters/miles/island_entry.py @@ -284,6 +284,9 @@ def apply_ports_infra_switches(args, miles_args, environ=None) -> None: miles_args.yeto_rl_overlap_eval = 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() # this thread drives the island (asyncio.run below) 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"))) diff --git a/yeto/rl/engine/driver.py b/yeto/rl/engine/driver.py index f6ea1afb..c704f4e3 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 ( @@ -480,6 +481,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")) diff --git a/yeto/rl/engine/phase_profile.py b/yeto/rl/engine/phase_profile.py new file mode 100644 index 00000000..056f4673 --- /dev/null +++ b/yeto/rl/engine/phase_profile.py @@ -0,0 +1,82 @@ +"""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 called +:func:`enable` 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) -> str: + inner = None + while frame is not None: + 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): + return here + frame = frame.f_back + return inner or "?" + + +class _PhaseSampler(threading.Thread): + def __init__(self, target: int) -> None: + super().__init__(daemon=True, name="yeto-phase-profile") + self.target = target + 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) + 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: + 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 sampling the calling thread (the driver's).""" + if _STATE["thread"] is None: + sampler = _PhaseSampler(threading.get_ident()) + 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)) From 92f33bc05c90f4fc66c067bfc6691f728414df0d Mon Sep 17 00:00:00 2001 From: michael <78113554+chowhsu@users.noreply.github.com> Date: Wed, 21 Oct 2026 14:07:32 +0000 Subject: [PATCH 47/55] eval: sample the thread that drives the phases, not the enabling one GPU run e1n-yeto-miles-async-p: the island enables the sampler from the thread that starts asyncio, but the driver runs its phases on another thread, so every sample showed the event loop. The sampler now binds to the thread of the first phase report. Co-Authored-By: Claude Opus 5.5 --- tests/test_rl_miles_train_profile.py | 20 +++++++++++++++----- yeto/rl/adapters/miles/island_entry.py | 2 +- yeto/rl/engine/phase_profile.py | 17 +++++++++++------ 3 files changed, 27 insertions(+), 12 deletions(-) diff --git a/tests/test_rl_miles_train_profile.py b/tests/test_rl_miles_train_profile.py index 883bb45a..139f64dc 100644 --- a/tests/test_rl_miles_train_profile.py +++ b/tests/test_rl_miles_train_profile.py @@ -56,12 +56,22 @@ def test_phase_profile_prints_each_finished_phase(capsys, monkeypatch): monkeypatch.setitem(phase_profile._STATE, "thread", None) phase_profile.on_phase("train", 0) # disabled: nothing - phase_profile.enable() - phase_profile.on_phase("train", 1) - time.sleep(0.05) - phase_profile.on_phase("sync", 1) + 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 + assert record["samples"] >= 1 and any("busy_driver" in label for label, _ in record["top"]), record phase_profile._STATE["thread"].phase = None diff --git a/yeto/rl/adapters/miles/island_entry.py b/yeto/rl/adapters/miles/island_entry.py index 4bb7af54..6fe8d2c0 100644 --- a/yeto/rl/adapters/miles/island_entry.py +++ b/yeto/rl/adapters/miles/island_entry.py @@ -286,7 +286,7 @@ def apply_ports_infra_switches(args, miles_args, environ=None) -> None: miles_args.yeto_rl_profile_trainer = True from yeto.rl.engine import phase_profile - phase_profile.enable() # this thread drives the island (asyncio.run below) + 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"))) diff --git a/yeto/rl/engine/phase_profile.py b/yeto/rl/engine/phase_profile.py index 056f4673..abcaa415 100644 --- a/yeto/rl/engine/phase_profile.py +++ b/yeto/rl/engine/phase_profile.py @@ -1,8 +1,9 @@ """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 called -:func:`enable` every ``SAMPLE_S`` and counts the samples per (phase, frame); at each +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. @@ -36,9 +37,9 @@ def _label(frame: Any) -> str: class _PhaseSampler(threading.Thread): - def __init__(self, target: int) -> None: + def __init__(self) -> None: super().__init__(daemon=True, name="yeto-phase-profile") - self.target = target + self.target: int | None = None self.lock = threading.Lock() self.phase: tuple[str, Any] | None = None self.started = time.perf_counter() @@ -48,6 +49,8 @@ def __init__(self, target: int) -> None: 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 @@ -59,6 +62,8 @@ def run(self) -> None: 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) @@ -69,9 +74,9 @@ def switch(self, phase: tuple[str, Any]) -> None: def enable() -> None: - """Start sampling the calling thread (the driver's).""" + """Start the sampler; it samples the thread of the first :func:`on_phase`.""" if _STATE["thread"] is None: - sampler = _PhaseSampler(threading.get_ident()) + sampler = _PhaseSampler() sampler.start() _STATE["thread"] = sampler From 0dec43de0a37e00125584d19e8dd3a2a2b374c7b Mon Sep 17 00:00:00 2001 From: michael <78113554+chowhsu@users.noreply.github.com> Date: Wed, 21 Oct 2026 15:06:08 +0000 Subject: [PATCH 48/55] eval: phase profile labels carry the caller chain A sample is labelled by its innermost four yeto/miles/ray frames, so a wait inside Ray shows which driver call waits (GPU run e1o-yeto-miles-sync-p only showed the Ray frame). Co-Authored-By: Claude Opus 5.5 --- yeto/rl/engine/phase_profile.py | 12 +++++++----- 1 file changed, 7 insertions(+), 5 deletions(-) diff --git a/yeto/rl/engine/phase_profile.py b/yeto/rl/engine/phase_profile.py index abcaa415..4b342189 100644 --- a/yeto/rl/engine/phase_profile.py +++ b/yeto/rl/engine/phase_profile.py @@ -24,16 +24,18 @@ _STATE: dict[str, Any] = {"thread": None} -def _label(frame: Any) -> str: - inner = None - while frame is not 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): - return here + chain.append(here) frame = frame.f_back - return inner or "?" + return " < ".join(chain) if chain else (inner or "?") class _PhaseSampler(threading.Thread): From 2ec20f76ded4278b714409a942edb2cc09ab8846 Mon Sep 17 00:00:00 2001 From: michael <78113554+chowhsu@users.noreply.github.com> Date: Wed, 21 Oct 2026 15:18:53 +0000 Subject: [PATCH 49/55] WIP: verify publications only on cut, first and every-K rounds Untested work in progress (E1 phase 8): verify_publication, the unverified identity of a skipped round, --rl-verify-every-round. Co-Authored-By: Claude Opus 5.5 --- yeto/launch/rl_ports.py | 2 ++ yeto/rl/adapters/miles/entry.py | 2 ++ yeto/rl/adapters/miles/island_entry.py | 2 ++ yeto/rl/adapters/miles/state.py | 23 ++++++++++++++++++++++ yeto/rl/engine/driver.py | 27 ++++++++++++++++++++++++++ yeto/rl/engine/policy_digest.py | 18 ++++++++++++++++- yeto/rl/engine/run_flags.py | 4 ++++ 7 files changed, 77 insertions(+), 1 deletion(-) diff --git a/yeto/launch/rl_ports.py b/yeto/launch/rl_ports.py index 9e78b20b..99f60801 100644 --- a/yeto/launch/rl_ports.py +++ b/yeto/launch/rl_ports.py @@ -535,6 +535,8 @@ def _ports_infra_flags(args) -> tuple[str, str]: 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/rl/adapters/miles/entry.py b/yeto/rl/adapters/miles/entry.py index f05f8ac4..df9aa11b 100644 --- a/yeto/rl/adapters/miles/entry.py +++ b/yeto/rl/adapters/miles/entry.py @@ -830,6 +830,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 diff --git a/yeto/rl/adapters/miles/island_entry.py b/yeto/rl/adapters/miles/island_entry.py index 6fe8d2c0..6d287518 100644 --- a/yeto/rl/adapters/miles/island_entry.py +++ b/yeto/rl/adapters/miles/island_entry.py @@ -282,6 +282,8 @@ 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 diff --git a/yeto/rl/adapters/miles/state.py b/yeto/rl/adapters/miles/state.py index 49120156..fcc1b7b4 100644 --- a/yeto/rl/adapters/miles/state.py +++ b/yeto/rl/adapters/miles/state.py @@ -63,6 +63,7 @@ class MilesPolicyState: # E1: export_digest(deferred=True) is supported (LoRA digest path) supports_deferred_digest = True + supports_unverified_rounds = True # E1: export_identity def __init__( self, @@ -146,6 +147,26 @@ def export_digest(self, *, policy_version: int | None = None, deferred: bool = F 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 @@ -227,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, @@ -328,6 +350,7 @@ class MilesFullParameterLocalPolicyState(MilesPolicyState): """ 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, diff --git a/yeto/rl/engine/driver.py b/yeto/rl/engine/driver.py index c704f4e3..55cd89c4 100644 --- a/yeto/rl/engine/driver.py +++ b/yeto/rl/engine/driver.py @@ -381,6 +381,29 @@ def defer_digest(driver: Any) -> bool: 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, @@ -701,6 +724,10 @@ 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)) diff --git a/yeto/rl/engine/policy_digest.py b/yeto/rl/engine/policy_digest.py index 7d58afbd..734fa765 100644 --- a/yeto/rl/engine/policy_digest.py +++ b/yeto/rl/engine/policy_digest.py @@ -314,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: @@ -349,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 @@ -366,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/run_flags.py b/yeto/rl/engine/run_flags.py index 38491c91..5bd72feb 100644 --- a/yeto/rl/engine/run_flags.py +++ b/yeto/rl/engine/run_flags.py @@ -495,6 +495,10 @@ 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: at round-cut rounds, the first publication and every " + "--rl-cut-every rounds)"), 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)"), From b664b9da4eb04b75915d6180bee1dba30b0e097c Mon Sep 17 00:00:00 2001 From: michael <78113554+chowhsu@users.noreply.github.com> Date: Wed, 21 Oct 2026 15:20:37 +0000 Subject: [PATCH 50/55] round cuts: device-copy capture for a trainer that sleeps after the cut Under --offload-train the cut capture took host copies tensor by tensor (0.16-0.28 s at the step boundary, GPU runs e1i/e1l sync). It now copies the state device-to-device and stages the copies into one reused pinned host buffer with a single synchronization before the capture returns, so nothing reads device memory after the trainer sleeps. The writer thread copies the staged views into the shared hand-over buffer and frees the staging buffer. Shard bytes and summaries equal the inline cut's (tests). Co-Authored-By: Claude Opus 5.5 --- tests/test_rl_cut_distopt_pp_ep.py | 12 ++- yeto/rl/adapters/miles/cut_plugin.py | 139 ++++++++++++++++++++++++--- 2 files changed, 134 insertions(+), 17 deletions(-) diff --git a/tests/test_rl_cut_distopt_pp_ep.py b/tests/test_rl_cut_distopt_pp_ep.py index 0285d2e7..308a2076 100644 --- a/tests/test_rl_cut_distopt_pp_ep.py +++ b/tests/test_rl_cut_distopt_pp_ep.py @@ -398,14 +398,19 @@ def export(optimizer, named, *, device="cpu"): @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, monkeypatch): +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) @@ -415,6 +420,7 @@ def test_background_device_capture_writes_the_inline_shard(tmp_path, layout, mon 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() @@ -517,8 +523,8 @@ def test_the_writer_process_can_be_turned_off(tmp_path, monkeypatch): 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)) - assert cp._device_copy_plan(sleeper, named) == (False, "offload_train") + 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") diff --git a/yeto/rl/adapters/miles/cut_plugin.py b/yeto/rl/adapters/miles/cut_plugin.py index 849a08ff..98bd9d03 100644 --- a/yeto/rl/adapters/miles/cut_plugin.py +++ b/yeto/rl/adapters/miles/cut_plugin.py @@ -692,11 +692,6 @@ def _device_copy_plan(actor: Any, named: list) -> tuple[bool, str]: if os.environ.get(DEVICE_COPY_ENV, "1").strip().lower() in {"0", "false", "no", "off"}: return False, "disabled" - if getattr(getattr(actor, "args", None), "offload_train", False): - # --offload-train: the trainer sleeps right after the cut (colocated) and the - # writer's reads of the device copies then fail ("CUDA error: invalid - # argument", 5 of 6 cuts of GPU run e1g-yeto-miles-sync-1) - return False, "offload_train" devices = {p.device for _, p in named} if len(devices) != 1 or next(iter(devices)).type != "cuda": return False, "not_cuda" @@ -704,9 +699,19 @@ def _device_copy_plan(actor: Any, named: list) -> tuple[bool, str]: 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 # -------------------------------------------------------------------------- @@ -828,18 +833,21 @@ def _remove_partial(directory: str, coord: Mapping[str, int]) -> None: def _write_background(captured: dict[str, Any], *, directory: str, cut_id: str) -> dict[str, Any]: - """Thread of the trainer process: device copies to shared host memory, then the - writer process serializes, hashes and summarizes them. 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``).""" + """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: - host = _on_host(captured) + 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: - skeleton, layout, flat = _pack(host) result = _writer_pool().submit(_write_packed, skeleton, layout, flat, directory=directory, cut_id=cut_id).result() return {**result, "writer": {"mode": "process"}} @@ -847,11 +855,107 @@ def _write_background(captured: dict[str, Any], *, directory: str, cut_id: str) _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, host["coord"]) - return {**_write_refusing(host, directory=directory, cut_id=cut_id), + _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__" @@ -873,8 +977,13 @@ def _pack(value: Any) -> tuple[Any, list, Any]: 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) @@ -980,7 +1089,9 @@ def _capture(actor: Any, *, cut_id: str, device_copy: bool = False) -> dict[str, if snap is None: snap = _snapshot(actor, backend, named) ready = None - if on_device and any(p.is_cuda for _, p in named): + 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()) From e79520dfd1f3423590aca7ea2797d5d85d539881 Mon Sep 17 00:00:00 2001 From: michael <78113554+chowhsu@users.noreply.github.com> Date: Wed, 21 Oct 2026 15:25:29 +0000 Subject: [PATCH 51/55] round cuts: default interval 4; verify publications at that interval The product default round-cut interval is now 4 (--rl-cut-every 1 cuts every round). A publication hashes the policy and reads the engine weight checksums back only on the first publication, on every round that is a multiple of the cut interval (every round that keeps a cut) and, with --rl-verify-every-round, on every round. Other publications skip both but still check, by weights version, that the trainer holds the policy; their identity is a 64-hex value derived from the trainer weights version (never a content hash), and the publication event says sync/publication_verified false. A cut refuses an unverified publication, so resume always sees a content hash. Help texts and the rl-resume-from-checkpoint spec and design say so; the efficiency eval gains --checkpoint default (no --rl-cut-every). Tests: which rounds verify, the switch, the skipped read-back, the weights-version check on an unverified round, and the cut refusal. Co-Authored-By: Claude Opus 5.5 --- .../rl-resume-from-checkpoint/design.md | 1 + .../specs/rl-resume-from-checkpoint/spec.md | 7 +++ .../test_rl_miles_adapter_trainer_publish.py | 19 +++++++ tests/test_rl_publish_fastpath.py | 56 +++++++++++++++++++ tests/test_tools_eval_efficiency.py | 2 + tools/eval/__main__.py | 4 +- tools/eval/efficiency/spec.py | 11 ++-- yeto/rl/adapters/miles/entry.py | 3 +- yeto/rl/adapters/miles/island_entry.py | 3 +- yeto/rl/adapters/miles/publish.py | 9 ++- yeto/rl/adapters/verl/round_cuts.py | 4 +- yeto/rl/engine/driver.py | 2 + yeto/rl/engine/resume.py | 6 ++ yeto/rl/engine/round_cut.py | 5 ++ yeto/rl/engine/run_flags.py | 8 ++- 15 files changed, 124 insertions(+), 16 deletions(-) diff --git a/openspec/changes/rl-resume-from-checkpoint/design.md b/openspec/changes/rl-resume-from-checkpoint/design.md index 3ebe7e80..4a3b7e19 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 6f3cbcfb..0a63bda6 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`. 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/test_rl_miles_adapter_trainer_publish.py b/tests/test_rl_miles_adapter_trainer_publish.py index 6e4ccbe3..25456531 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_publish_fastpath.py b/tests/test_rl_publish_fastpath.py index 54877fd8..dbb99e4f 100644 --- a/tests/test_rl_publish_fastpath.py +++ b/tests/test_rl_publish_fastpath.py @@ -362,3 +362,59 @@ def test_the_driver_defers_only_for_a_colocated_offloading_trainer(monkeypatch): 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_tools_eval_efficiency.py b/tests/test_tools_eval_efficiency.py index 94911c11..5f613bfd 100644 --- a/tests/test_tools_eval_efficiency.py +++ b/tests/test_tools_eval_efficiency.py @@ -95,6 +95,8 @@ def test_yeto_arms_use_the_same_policy_age_and_partition(): 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")) diff --git a/tools/eval/__main__.py b/tools/eval/__main__.py index 8160439f..775d2dbc 100644 --- a/tools/eval/__main__.py +++ b/tools/eval/__main__.py @@ -228,9 +228,9 @@ def main(argv: list[str] | None = None) -> int: 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="final", + 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 " - "(product default)") + "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) diff --git a/tools/eval/efficiency/spec.py b/tools/eval/efficiency/spec.py index 1b582f28..dcbcab0e 100644 --- a/tools/eval/efficiency/spec.py +++ b/tools/eval/efficiency/spec.py @@ -362,17 +362,18 @@ def yeto_verl_argv(spec: EffSpec, mode: str, tape: str) -> list[str]: return argv -CHECKPOINT_MODES = ("final", "every-round") +CHECKPOINT_MODES = ("final", "every-round", "default") -def cut_every(spec: EffSpec) -> int: +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`` is the product default. + 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 spec.rounds if spec.checkpoint == "final" else 1 + 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: @@ -400,7 +401,7 @@ def yeto_miles_launcher_args(spec: EffSpec, mode: str, run_id: str) -> str: 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"--rl-cut-every {cut_every(spec)} " + 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} " diff --git a/yeto/rl/adapters/miles/entry.py b/yeto/rl/adapters/miles/entry.py index df9aa11b..40745e74 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, @@ -1042,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/island_entry.py b/yeto/rl/adapters/miles/island_entry.py index 6d287518..a54105ec 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 @@ -335,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 329b66bf..2bd8a148 100644 --- a/yeto/rl/adapters/miles/publish.py +++ b/yeto/rl/adapters/miles/publish.py @@ -471,7 +471,10 @@ def publish(self, state: TrainableState, *, token_rollout_id: int | None = None) version = state.policy_version if token_rollout_id is None else token_rollout_id token = policy_token(version, tensor_hash) self.last_timing = {"check_s": time.monotonic() - started} - members, checksums = self._runner.run(self._publish(version, token)) + # 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) @@ -536,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): @@ -566,7 +569,7 @@ async def _publish( 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) diff --git a/yeto/rl/adapters/verl/round_cuts.py b/yeto/rl/adapters/verl/round_cuts.py index 735b282e..aa7f4e34 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/engine/driver.py b/yeto/rl/engine/driver.py index 55cd89c4..959f4059 100644 --- a/yeto/rl/engine/driver.py +++ b/yeto/rl/engine/driver.py @@ -782,6 +782,8 @@ def publish(self, state: TrainableState, *, rollout_id: int) -> None: policy_version=rollout_id, **{ "rl/policy_token": self.expected_token, + # E1: False on a round whose policy content was not hashed + "sync/publication_verified": getattr(state, "verified", True) is not False, "sync/publication_payload_bytes": manifest.payload_bytes, "sync/publication_payload_hash": manifest.payload_hash, "sync/publication_members": sorted(result.members), diff --git a/yeto/rl/engine/resume.py b/yeto/rl/engine/resume.py index 389d45ec..9a68f22d 100644 --- a/yeto/rl/engine/resume.py +++ b/yeto/rl/engine/resume.py @@ -453,6 +453,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).""" diff --git a/yeto/rl/engine/round_cut.py b/yeto/rl/engine/round_cut.py index ce207118..2fa8378e 100644 --- a/yeto/rl/engine/round_cut.py +++ b/yeto/rl/engine/round_cut.py @@ -117,6 +117,11 @@ def context(self, cut_id: str, *, pending: tuple[int, Any] | None = None) -> Any 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: diff --git a/yeto/rl/engine/run_flags.py b/yeto/rl/engine/run_flags.py index 5bd72feb..e90e06bc 100644 --- a/yeto/rl/engine/run_flags.py +++ b/yeto/rl/engine/run_flags.py @@ -497,8 +497,8 @@ def rl_forward_pairs(*dests: str | tuple[str, str]) -> tuple[tuple[str, str], .. 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: at round-cut rounds, the first publication and every " - "--rl-cut-every rounds)"), + "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)"), @@ -558,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"), From 096bdf6159b67b22ac0d05da749a8d4e7b8686ef Mon Sep 17 00:00:00 2001 From: michael <78113554+chowhsu@users.noreply.github.com> Date: Wed, 21 Oct 2026 15:33:32 +0000 Subject: [PATCH 52/55] driver: mark only unverified publications in the publication event Verified publications keep the event as it was, so recorded golden tapes stay byte-identical (test_decoupling_golden, test_rl_driver_profiles). Co-Authored-By: Claude Opus 5.5 --- .../specs/rl-resume-from-checkpoint/spec.md | 2 +- yeto/rl/engine/driver.py | 5 +++-- 2 files changed, 4 insertions(+), 3 deletions(-) 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 0a63bda6..025a344b 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 @@ -51,7 +51,7 @@ launcher SHALL 接受 `--rl-checkpoint-store modal-volume://<卷名>[/前缀]` - **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`. A publication that was not verified MUST NOT be committed as a round cut. +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 diff --git a/yeto/rl/engine/driver.py b/yeto/rl/engine/driver.py index 959f4059..f116459a 100644 --- a/yeto/rl/engine/driver.py +++ b/yeto/rl/engine/driver.py @@ -782,8 +782,9 @@ def publish(self, state: TrainableState, *, rollout_id: int) -> None: policy_version=rollout_id, **{ "rl/policy_token": self.expected_token, - # E1: False on a round whose policy content was not hashed - "sync/publication_verified": getattr(state, "verified", True) is not False, + # 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), From becb3a39a8f8f91e8909c750e3625fd23f26715f Mon Sep 17 00:00:00 2001 From: michael <78113554+chowhsu@users.noreply.github.com> Date: Wed, 21 Oct 2026 15:41:31 +0000 Subject: [PATCH 53/55] test: keep the phase profiler off after the profile tests A test enabled the process-wide driver phase sampler and left it on; it then printed on every later driver phase and broke the JSON output that test_rl_reconfig_e1::test_command_inbox_and_cli parses (failed in the full suite only). Co-Authored-By: Claude Opus 5.5 --- tests/test_rl_miles_train_profile.py | 12 ++++++++++-- 1 file changed, 10 insertions(+), 2 deletions(-) diff --git a/tests/test_rl_miles_train_profile.py b/tests/test_rl_miles_train_profile.py index 139f64dc..0325ad8e 100644 --- a/tests/test_rl_miles_train_profile.py +++ b/tests/test_rl_miles_train_profile.py @@ -31,8 +31,14 @@ def test_the_flag_reaches_the_launcher_args_and_the_eval_spec(): 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(): +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 @@ -74,4 +80,6 @@ def busy_driver(): 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 - phase_profile._STATE["thread"].phase = None + sampler = phase_profile._STATE["thread"] + sampler.phase = None + monkeypatch.setitem(phase_profile._STATE, "thread", None) # off for later tests From 2f7e2882530aa8adb527da16d447d51e31c597d7 Mon Sep 17 00:00:00 2001 From: michael <78113554+chowhsu@users.noreply.github.com> Date: Wed, 21 Oct 2026 16:51:04 +0000 Subject: [PATCH 54/55] round cuts: commit the store's first cut in the background too A first periodic cut was written inline because a store that holds cuts but no LATEST refuses to start. With --rl-cut-every 4 that cut falls in the steady rounds (Modal A10 e1q: 2.3-9.8 s on the critical path). The first cut is now deferred like the others. Before its shards start, a fsynced mark round-cuts/.first-cut-in-progress is written; it is cleared once a LATEST is committed. At start, a store without LATEST whose every cut carries the mark (never committed) is cleaned and the run starts at rollout 0 (the local state dir is moved aside, journaled). Any other store with cuts and no LATEST is still refused. Co-Authored-By: Claude Opus 5.5 --- tests/test_rl_resume_from_checkpoint.py | 127 ++++++++++++++++++++++-- yeto/rl/engine/resume.py | 82 ++++++++++++++- yeto/rl/engine/round_cut.py | 16 ++- 3 files changed, 207 insertions(+), 18 deletions(-) diff --git a/tests/test_rl_resume_from_checkpoint.py b/tests/test_rl_resume_from_checkpoint.py index c3e43461..d3b275c8 100644 --- a/tests/test_rl_resume_from_checkpoint.py +++ b/tests/test_rl_resume_from_checkpoint.py @@ -69,6 +69,17 @@ 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 + return {**info, "store_synced": True, "cut_bytes": done["cut_bytes"]} + + def _trained_island(tmp_path, store, rounds, *, name="a", seed=3): rank = make_rank(0) g = torch.Generator().manual_seed(seed) @@ -186,6 +197,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 @@ -193,9 +299,10 @@ 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) - # E1: the store's first cut is complete inline (a store with cuts but no LATEST refuses to start) - assert info["store_synced"] is True and info["store_commit"] == "inline" + 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"] @@ -223,7 +330,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}) @@ -243,7 +350,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") @@ -264,7 +371,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)) @@ -293,7 +400,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 @@ -600,8 +707,8 @@ def _island_with_first_cut(tmp_path, store, keep=2): trainer = _CountingTrainer(_trainer(rank)) cp = _cp(ctl, driver, trainer, _Ledger(2)) _publish(driver, 1, _hash(rank)) - first = cp.save(driver, rollout_id=1) - assert first["store_commit"] == "inline" and first["store_synced"] is True + 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() @@ -759,7 +866,7 @@ def commit(self): 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"] + 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 diff --git a/yeto/rl/engine/resume.py b/yeto/rl/engine/resume.py index 9a68f22d..99b54ff3 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): @@ -392,6 +395,60 @@ def _cut_rollout(cut_id: str) -> int | 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.""" @@ -491,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, @@ -510,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: @@ -517,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(), @@ -526,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"}, @@ -602,6 +675,9 @@ def commit_checkpoint_store(self) -> dict[str, Any]: 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() diff --git a/yeto/rl/engine/round_cut.py b/yeto/rl/engine/round_cut.py index 2fa8378e..c0b1d010 100644 --- a/yeto/rl/engine/round_cut.py +++ b/yeto/rl/engine/round_cut.py @@ -326,6 +326,12 @@ def _start(self, context: Any, *, background: bool) -> Any: 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 @@ -393,11 +399,11 @@ def _unprepared_wake(self, driver: Any, rollout_id: int, prepared: "_PreparedCut 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), and - not before the first cut of the store is committed (a store holding cuts but - no LATEST refuses to start).""" - return (self.background_commit and not final and self.buffer_watermark is None - and self.pointer() is not None) + 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 233f1390216a20cbf78176f29fc6b4b4540588ef Mon Sep 17 00:00:00 2001 From: michael <78113554+chowhsu@users.noreply.github.com> Date: Wed, 21 Oct 2026 16:59:21 +0000 Subject: [PATCH 55/55] test: complete deferred first cuts in the recovery abstraction tests Co-Authored-By: Claude Opus 5.5 --- tests/test_rl_recovery_abstraction.py | 8 ++++---- tests/test_rl_resume_from_checkpoint.py | 5 ++++- 2 files changed, 8 insertions(+), 5 deletions(-) diff --git a/tests/test_rl_recovery_abstraction.py b/tests/test_rl_recovery_abstraction.py index 6064e8e1..a6666639 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 d3b275c8..a24480a5 100644 --- a/tests/test_rl_resume_from_checkpoint.py +++ b/tests/test_rl_resume_from_checkpoint.py @@ -19,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 @@ -77,7 +78,9 @@ def _cut(cp, driver, **kw): return info done = cp.join_commit(driver) assert done["ok"] is True, done - return {**info, "store_synced": True, "cut_bytes": done["cut_bytes"]} + 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):