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/62] 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 8358c1b13a69ba8b6d43df4735a01e4f4bc2ba73 Mon Sep 17 00:00:00 2001 From: michael <78113554+chowhsu@users.noreply.github.com> Date: Tue, 20 Oct 2026 14:31:17 +0000 Subject: [PATCH 02/62] syncer: pipeline policy-sweep fragments and async per-round checkpoints Allow --pipeline 1..P with --policy-sweep-fragments so fragment k+1 uploads while fragment k merges, checkpoints and broadcasts; commits stay in strict step order. Add --async-checkpoint (policy sweep only): the committed state is snapshotted and written on a blocking task, and the round is broadcast only after the checkpoint is durable. Checkpoint bytes, versions and the semantic profile hash are unchanged; the tape gains checkpoint_ms. Launcher flags --rl-sync-fragment-pipeline and --rl-sync-async-checkpoint (default off). Adds bit-equality, failure and real-syncer tests and a CPU micro-benchmark (tools/bench_sync_boundary.py). Co-Authored-By: Claude Opus 5.5 --- syncer/src/main.rs | 5 + syncer/src/server.rs | 496 +++++++++++++++++++++-- syncer/src/state.rs | 118 ++++-- tests/test_rl_dense_sweep_real_syncer.py | 87 +++- tests/test_rl_sao_streaming_runtime.py | 2 +- tests/test_run_flags.py | 2 +- tests/test_wandb_tape.py | 4 +- tools/bench_sync_boundary.py | 163 ++++++++ yeto/launch/rl_args.py | 7 +- yeto/launch/syncer.py | 3 +- yeto/rl/engine/run_flags.py | 14 + yeto/syncer_profile.py | 2 +- yeto/wandb_tape.py | 2 + 13 files changed, 830 insertions(+), 75 deletions(-) create mode 100644 tools/bench_sync_boundary.py diff --git a/syncer/src/main.rs b/syncer/src/main.rs index 513caabe..8243ab29 100644 --- a/syncer/src/main.rs +++ b/syncer/src/main.rs @@ -98,6 +98,10 @@ struct Args { /// Write the snapshot every N outer steps (0 disables). #[arg(long, default_value_t = 8)] checkpoint_every: u64, + /// Policy sweep: write per-round checkpoints off the scheduler; each + /// round is broadcast only after its checkpoint is durable. + #[arg(long, default_value_t = false)] + async_checkpoint: bool, /// Resume from --checkpoint-path if it exists. #[arg(long, default_value_t = false)] resume: bool, @@ -273,6 +277,7 @@ fn main() -> anyhow::Result<()> { final_state: args.final_state, checkpoint_path: args.checkpoint_path, checkpoint_every: args.checkpoint_every, + async_checkpoint: args.async_checkpoint, resume: args.resume, mark_final_checkpoint: args.mark_final_checkpoint, learner_budget_steps: args.learner_budget_steps, diff --git a/syncer/src/server.rs b/syncer/src/server.rs index a5ed5b67..e7b42b3e 100644 --- a/syncer/src/server.rs +++ b/syncer/src/server.rs @@ -91,6 +91,12 @@ pub struct Config { /// the quiescent cut between rounds, resumed from when `resume` is set. pub checkpoint_path: Option, pub checkpoint_every: u64, + /// Policy-sweep only: persist each per-round checkpoint on a blocking + /// writer instead of the scheduler, and broadcast that round only after + /// the write is durable. Merge results, versions, and checkpoint bytes + /// are unchanged, so this is deliberately not part of the semantic + /// profile hash. + pub async_checkpoint: bool, pub resume: bool, /// Create the adjacent final marker only for a completed terminal cut. pub mark_final_checkpoint: bool, @@ -702,8 +708,11 @@ fn validate_config(cfg: &Config) -> Result<()> { if cfg.total_steps == 0 { bail!("--policy-sweep-fragments requires positive --total-steps"); } - if cfg.pipeline != 1 { - bail!("--policy-sweep-fragments requires --pipeline 1"); + // Fragment rounds of one sweep are independent (H=1, plain average, + // no momentum), so up to P of them may gather concurrently; commits + // stay in strict step order. + if cfg.pipeline == 0 || cfg.pipeline > fragments { + bail!("--policy-sweep-fragments requires --pipeline between 1 and the fragment count"); } if !cfg.total_steps.is_multiple_of(u64::from(fragments)) { bail!("--total-steps must be divisible by --policy-sweep-fragments"); @@ -733,6 +742,14 @@ fn validate_config(cfg: &Config) -> Result<()> { ); } } + if cfg.async_checkpoint { + if cfg.policy_sweep_fragments.is_none() { + bail!("--async-checkpoint requires --policy-sweep-fragments"); + } + if cfg.learner_budget_steps.is_some() { + bail!("--async-checkpoint does not support --learner-budget-steps"); + } + } Ok(()) } @@ -2116,12 +2133,15 @@ async fn scheduler( let mut computing = tokio::task::JoinSet::>::new(); let mut ready: BTreeMap = BTreeMap::new(); let mut busy_fragments: HashSet = HashSet::new(); + // --async-checkpoint: the committed round whose checkpoint is in flight. + let mut publishing: Option = None; let mut cutoff_established = false; let mut first_budget_report: Option<(Member, u64)> = None; 'outer: while next_t <= cfg.total_steps || !inflight.is_empty() || !computing.is_empty() || !ready.is_empty() + || publishing.is_some() { if budget_cutoff.deadline().is_some() { cutoff_established = true; @@ -2138,17 +2158,37 @@ async fn scheduler( let completed = take_next_commit(&mut ready, next_commit_t) .context("ready commit disappeared after linearization")?; let p = completed.round.p; - commit_round( - &cfg, - &mut st, - ®istry, - &mut last_sync_secs, - completed, - &budget_cutoff, - ) - .await?; - if !busy_fragments.remove(&p) { - bail!("fragment {p} lost its busy ownership at commit"); + if cfg.async_checkpoint { + // At most one checkpoint write is outstanding: the snapshot + // of t must be durable before t+1 is snapshotted, which also + // bounds host memory to one extra copy of the state. + if let Some(pending) = publishing.take() { + settle_pending_publish( + &cfg, + &st, + ®istry, + &mut last_sync_secs, + pending, + &budget_cutoff, + &mut busy_fragments, + ) + .await?; + } + let committed = apply_round_commit(&cfg, &mut st, completed)?; + publishing = Some(start_async_checkpoint(&cfg, &st, committed)?); + } else { + commit_round( + &cfg, + &mut st, + ®istry, + &mut last_sync_secs, + completed, + &budget_cutoff, + ) + .await?; + if !busy_fragments.remove(&p) { + bail!("fragment {p} lost its busy ownership at commit"); + } } next_commit_t = next_commit_t .checked_add(1) @@ -2349,7 +2389,19 @@ async fn scheduler( if next_fragment_available { earliest = Some(earliest.map_or(next_launch, |d| d.min(next_launch))); } - let wake = match earliest { + let wake = if let Some(pending) = publishing.as_mut() { + // Same event-first bias; a durable checkpoint outranks compute + // so the oldest committed round is exposed first. + let sleep_until = earliest.map(tokio::time::Instant::from_std); + tokio::select! { + biased; + event = events.recv() => SchedulerWake::Event(event), + written = &mut pending.write => SchedulerWake::Durable(written), + result = computing.join_next(), if !computing.is_empty() => SchedulerWake::Computed(result), + _ = sleep_until_opt(sleep_until) => SchedulerWake::Deadline, + } + } else { + match earliest { Some(deadline) if !computing.is_empty() => tokio::select! { biased; wake = prefer_event_over_compute(events.recv(), computing.join_next()) => wake, @@ -2364,9 +2416,27 @@ async fn scheduler( prefer_event_over_compute(events.recv(), computing.join_next()).await } None => continue, + } }; match wake { SchedulerWake::Deadline => continue, + SchedulerWake::Durable(written) => { + let pending = publishing + .take() + .context("durable checkpoint without a pending publish")?; + finish_pending_publish( + &cfg, + &st, + ®istry, + &mut last_sync_secs, + pending.committed, + pending.write_started, + written, + &budget_cutoff, + &mut busy_fragments, + ) + .await?; + } SchedulerWake::Computed(Some(result)) => { let completed = result.context("iso compute task panicked")??; let t = completed.round.t; @@ -2469,6 +2539,20 @@ async fn scheduler( // slots in existing rounds. The helper revalidates the // generation because catch-up awaits socket writers and // a newer HELLO may supersede this one meanwhile. + // Catch-up must never expose a round whose checkpoint + // is not yet durable. + if let Some(pending) = publishing.take() { + settle_pending_publish( + &cfg, + &st, + ®istry, + &mut last_sync_secs, + pending, + &budget_cutoff, + &mut busy_fragments, + ) + .await?; + } if st.global_step < cfg.total_steps { send_all_fragments(&st, &group, &budget_cutoff).await; } @@ -2553,6 +2637,18 @@ async fn scheduler( // Catch-up must precede a replay even when the // disconnect event wins the race with the queued // replacement HELLO. + if let Some(pending) = publishing.take() { + settle_pending_publish( + &cfg, + &st, + ®istry, + &mut last_sync_secs, + pending, + &budget_cutoff, + &mut busy_fragments, + ) + .await?; + } if st.global_step < cfg.total_steps { send_all_fragments(&st, &group, &budget_cutoff).await; } @@ -2640,6 +2736,7 @@ async fn scheduler( if next_commit_t != cfg.total_steps.saturating_add(1) || st.global_step != cfg.total_steps || !busy_fragments.is_empty() + || publishing.is_some() { bail!( "non-quiescent terminal cut: next_commit_t={next_commit_t} global_step={} busy={:?}", @@ -2935,6 +3032,14 @@ enum SchedulerWake { Deadline, Event(Option), Computed(Option, tokio::task::JoinError>>), + Durable(std::result::Result, tokio::task::JoinError>), +} + +async fn sleep_until_opt(deadline: Option) { + match deadline { + Some(deadline) => tokio::time::sleep_until(deadline).await, + None => std::future::pending().await, + } } type ComputeJoin = Option, tokio::task::JoinError>>; @@ -3189,6 +3294,24 @@ fn prepare_round_compute(cfg: &Config, st: &GlobalState, round: &Round) -> Resul st.prepare_merge(round.p, round.base_version, &outer_gradients, &weights) } +/// A round whose merge is applied to `GlobalState` but not yet externally +/// visible: no BCAST, no tape record, no log line. +struct CommittedRound { + t: u64, + p: usize, + base_version: u64, + attempt: u32, + started: Instant, + expected_members: Vec, + quorum_size: usize, + quorum_ms: Option, + grace_ms: Option, + pushes: HashMap, + gnorm: f64, + compute_secs: f64, + sync_start: Instant, +} + /// Commit an already-computed round. Called only for the contiguous t prefix, /// so Nesterov, versions, ledger, event tape, checkpoint, and broadcast all /// share one deterministic order even when SVD completion is out of order. @@ -3200,6 +3323,26 @@ async fn commit_round( completed: ComputedRound, budget_cutoff: &BudgetCutoff, ) -> Result<()> { + let committed = apply_round_commit(cfg, st, completed)?; + let mut checkpoint_ms = None; + if checkpoint_before_broadcast(cfg) { + if let Some(path) = cfg.checkpoint_path.as_ref() { + let checkpoint_start = Instant::now(); + st.save_checkpoint(path)?; + checkpoint_ms = Some(checkpoint_start.elapsed().as_millis() as u64); + info!(step = committed.t, path = %path.display(), "checkpoint committed"); + } + } + publish_round(cfg, st, registry, last_sync_secs, committed, checkpoint_ms, budget_cutoff).await +} + +/// Apply the merge, version, step, and ledger updates of the next contiguous +/// round. Nothing leaves the process here. +fn apply_round_commit( + cfg: &Config, + st: &mut GlobalState, + completed: ComputedRound, +) -> Result { let ComputedRound { round, merge, @@ -3237,7 +3380,6 @@ async fn commit_round( st.global_step = t; let mut ordered_pushes: Vec<_> = pushes.iter().collect(); ordered_pushes.sort_unstable_by_key(|(member, _)| **member); - let responders: Vec = ordered_pushes.iter().map(|(member, _)| **member).collect(); for (_, push) in ordered_pushes { if let Some(fragments) = cfg.policy_sweep_fragments { st.record_fragment_merge( @@ -3252,15 +3394,52 @@ async fn commit_round( st.record_merge(push.learner_id, push.c_steps, push.c_tokens); } } + Ok(CommittedRound { + t, + p, + base_version, + attempt, + started, + expected_members, + quorum_size, + quorum_ms, + grace_ms, + pushes, + gnorm, + compute_secs, + sync_start, + }) +} - let checkpoint_each_round = checkpoint_before_broadcast(cfg); - if checkpoint_each_round { - if let Some(path) = cfg.checkpoint_path.as_ref() { - st.save_checkpoint(path)?; - info!(step = t, path = %path.display(), "checkpoint committed"); - } - } - +/// Expose a committed round: BCAST its fragment, log it, append its tape +/// record, and write the periodic (non-per-round) checkpoint. A per-round +/// checkpoint must already be durable when this is called. +async fn publish_round( + cfg: &Config, + st: &GlobalState, + registry: &Registry, + last_sync_secs: &mut f64, + committed: CommittedRound, + checkpoint_ms: Option, + budget_cutoff: &BudgetCutoff, +) -> Result<()> { + let CommittedRound { + t, + p, + base_version, + attempt, + started, + expected_members, + quorum_size, + quorum_ms, + grace_ms, + pushes, + gnorm, + compute_secs, + sync_start, + } = committed; + let mut responders: Vec = pushes.keys().copied().collect(); + responders.sort_unstable(); // A disconnected learner can catch up after reconnecting. Deterministic // local encoding/allocation failures must terminate the coordinator // instead of silently committing a cut that was never broadcastable. @@ -3318,6 +3497,7 @@ async fn commit_round( quorum_ms, grace_ms, sync_ms, + checkpoint_ms, &pushes, gnorm, ms, @@ -3330,7 +3510,7 @@ async fn commit_round( // have also been broadcast; the terminal cut is delivered by FINAL. // A crash-resume loses those gathers; their fragments simply merge on // a later cycle, which the quorum design already tolerates. - if !checkpoint_each_round { + if !checkpoint_before_broadcast(cfg) { if let Some(path) = &cfg.checkpoint_path { if cfg.checkpoint_every > 0 && t % cfg.checkpoint_every == 0 { st.save_checkpoint(path)?; @@ -3341,6 +3521,97 @@ async fn commit_round( Ok(()) } +/// A committed round whose per-round checkpoint is being written by a +/// blocking task (`--async-checkpoint`). Its fragment stays busy and its +/// BCAST is withheld until the write is durable, so no learner can observe +/// a version a crash could lose. +struct PendingPublish { + committed: CommittedRound, + write_started: Instant, + write: tokio::task::JoinHandle>, +} + +/// Snapshot the just-committed state and start persisting it off the +/// scheduler thread. +fn start_async_checkpoint( + cfg: &Config, + st: &GlobalState, + committed: CommittedRound, +) -> Result { + let path = cfg + .checkpoint_path + .clone() + .context("--async-checkpoint requires --checkpoint-path")?; + let write_started = Instant::now(); + let bytes = st.encode_checkpoint()?; + let write = tokio::task::spawn_blocking(move || { + crate::state::write_checkpoint_bytes(&path, &bytes) + }); + Ok(PendingPublish { + committed, + write_started, + write, + }) +} + +/// Wait for a pending per-round checkpoint, then publish its round. A failed +/// write fails the coordinator before the round becomes visible; the +/// previous checkpoint stays the resume point. +async fn settle_pending_publish( + cfg: &Config, + st: &GlobalState, + registry: &Registry, + last_sync_secs: &mut f64, + pending: PendingPublish, + budget_cutoff: &BudgetCutoff, + busy_fragments: &mut HashSet, +) -> Result<()> { + let PendingPublish { + committed, + write_started, + write, + } = pending; + finish_pending_publish( + cfg, + st, + registry, + last_sync_secs, + committed, + write_started, + write.await, + budget_cutoff, + busy_fragments, + ) + .await +} + +#[allow(clippy::too_many_arguments)] +async fn finish_pending_publish( + cfg: &Config, + st: &GlobalState, + registry: &Registry, + last_sync_secs: &mut f64, + committed: CommittedRound, + write_started: Instant, + written: std::result::Result, tokio::task::JoinError>, + budget_cutoff: &BudgetCutoff, + busy_fragments: &mut HashSet, +) -> Result<()> { + let (t, p) = (committed.t, committed.p); + written + .context("async checkpoint task panicked")? + .with_context(|| format!("async checkpoint for step {t} failed; step {t} not broadcast"))?; + let checkpoint_ms = Some(write_started.elapsed().as_millis() as u64); + if let Some(path) = cfg.checkpoint_path.as_ref() { + info!(step = t, path = %path.display(), "checkpoint committed"); + } + publish_round(cfg, st, registry, last_sync_secs, committed, checkpoint_ms, budget_cutoff).await?; + if !busy_fragments.remove(&p) { + bail!("fragment {p} lost its busy ownership at publish"); + } + Ok(()) +} + #[cfg(test)] async fn complete_round( cfg: &Config, @@ -3361,6 +3632,30 @@ async fn complete_round( commit_round(cfg, st, registry, last_sync_secs, completed, &cutoff).await } +/// Test twin of `complete_round` through the `--async-checkpoint` path. +#[cfg(test)] +async fn complete_round_async( + cfg: &Config, + st: &mut GlobalState, + registry: &Registry, + last_sync_secs: &mut f64, + round: Round, +) -> Result<()> { + let prepared = prepare_round_compute(cfg, st, &round)?; + let compute_start = Instant::now(); + let merge = prepared.compute().await?; + let completed = ComputedRound { + round, + merge, + compute_secs: compute_start.elapsed().as_secs_f64(), + }; + let cutoff = BudgetCutoff::new(false, Duration::ZERO); + let mut busy = HashSet::from([completed.round.p]); + let committed = apply_round_commit(cfg, st, completed)?; + let pending = start_async_checkpoint(cfg, st, committed)?; + settle_pending_publish(cfg, st, registry, last_sync_secs, pending, &cutoff, &mut busy).await +} + fn new_state_for(group: &Arc, cfg: &Config) -> Result { if let Some(fragments) = cfg.policy_sweep_fragments { if group.dtype != DTYPE_F32 { @@ -3856,6 +4151,7 @@ fn append_tape( quorum_ms: Option, grace_ms: Option, sync_ms: u64, + checkpoint_ms: Option, pushes: &HashMap, gnorm: f64, ms: u64, @@ -3932,9 +4228,10 @@ fn append_tape( .map(|byte| format!("{byte:02x}")) .collect(); let merge_seconds = sync_ms as f64 / 1000.0; + let checkpoint_field = checkpoint_tape_field(checkpoint_ms); let sweep_fields = policy_sweep_event_fields(step, policy_sweep_fragments); let line = format!( - "{{\"protocol_version\":{PROTOCOL_VERSION},\"delta_semantics\":\"local_minus_raw_anchor\",\"sync/layout_hash\":\"{layout_hash}\",\"sync/base_version\":{launch_base_version},\"sync/responders\":{},\"sync/quorum\":{quorum},\"sync/rejected_stale_updates\":0,\"sync/merge_seconds\":{merge_seconds},\"sync/global_delta_norm\":{gnorm},\"step\":{step},\"fragment\":{fragment}{sweep_fields},\"launch_base_version\":{launch_base_version},\"attempt\":{attempt},\"gnorm\":{gnorm},\"ms\":{ms},\"quorum\":{quorum},\"expected\":{},\"expected_members\":{},\"responded\":{},\"responded_members\":{},\"missed_grace\":{},\"missed_members\":{},\"quorum_ms\":{quorum_ms},\"grace_ms\":{grace_ms},\"sync_ms\":{sync_ms},\"time_unix\":{:?},\"responders\":[{}]}}\n", + "{{\"protocol_version\":{PROTOCOL_VERSION},\"delta_semantics\":\"local_minus_raw_anchor\",\"sync/layout_hash\":\"{layout_hash}\",\"sync/base_version\":{launch_base_version},\"sync/responders\":{},\"sync/quorum\":{quorum},\"sync/rejected_stale_updates\":0,\"sync/merge_seconds\":{merge_seconds},\"sync/global_delta_norm\":{gnorm},\"step\":{step},\"fragment\":{fragment}{sweep_fields},\"launch_base_version\":{launch_base_version},\"attempt\":{attempt},\"gnorm\":{gnorm},\"ms\":{ms},\"quorum\":{quorum},\"expected\":{},\"expected_members\":{},\"responded\":{},\"responded_members\":{},\"missed_grace\":{},\"missed_members\":{},\"quorum_ms\":{quorum_ms},\"grace_ms\":{grace_ms},\"sync_ms\":{sync_ms}{checkpoint_field},\"time_unix\":{:?},\"responders\":[{}]}}\n", responded.len(), json_ids(&expected_learners), json_members(expected_members), @@ -3954,6 +4251,12 @@ fn append_tape( Ok(()) } +/// Optional top-level tape field: wall time of the per-round checkpoint +/// (synchronous save, or snapshot + async write until durable). +fn checkpoint_tape_field(checkpoint_ms: Option) -> String { + checkpoint_ms.map_or_else(String::new, |ms| format!(",\"checkpoint_ms\":{ms}")) +} + fn json_ids(ids: &[u32]) -> String { let mut ids = ids.to_vec(); ids.sort_unstable(); @@ -4354,6 +4657,7 @@ mod tests { final_state: None, checkpoint_path: None, checkpoint_every: 0, + async_checkpoint: false, resume: false, mark_final_checkpoint: false, learner_budget_steps: None, @@ -4618,10 +4922,30 @@ mod tests { assert!(format!("{:#}", validate_config(&invalid).unwrap_err()) .contains("requires positive --total-steps")); + let mut streamed = sweep_test_config(2, 4); + streamed.pipeline = 2; + streamed.async_checkpoint = true; + assert!(validate_config(&streamed).is_ok()); + for pipeline in [0, 3] { + invalid = sweep_test_config(2, 4); + invalid.pipeline = pipeline; + assert!(format!("{:#}", validate_config(&invalid).unwrap_err()) + .contains("requires --pipeline between 1 and the fragment count")); + } + let mut legacy_async = round_test_config(3); + legacy_async.async_checkpoint = true; + assert!(format!("{:#}", validate_config(&legacy_async).unwrap_err()) + .contains("--async-checkpoint requires --policy-sweep-fragments")); invalid = sweep_test_config(2, 4); - invalid.pipeline = 2; + invalid.async_checkpoint = true; + invalid.learner_budget_steps = Some(1); assert!(format!("{:#}", validate_config(&invalid).unwrap_err()) - .contains("requires --pipeline 1")); + .contains("does not support --learner-budget-steps")); + // Async persistence is not semantic: the profile hash ignores it. + let mut hashed = sweep_test_config(2, 4); + let before = hashed.semantic_profile_hash(); + hashed.async_checkpoint = true; + assert_eq!(before, hashed.semantic_profile_hash()); invalid = sweep_test_config(3, 4); assert!( @@ -5244,6 +5568,120 @@ mod tests { std::fs::remove_dir_all(&directory).ok(); } + fn unique_test_directory(tag: &str) -> std::path::PathBuf { + let directory = std::env::temp_dir().join(format!( + "yeto-{tag}-{}-{}", + std::process::id(), + std::time::SystemTime::now() + .duration_since(std::time::UNIX_EPOCH) + .unwrap() + .as_nanos() + )); + std::fs::create_dir_all(&directory).unwrap(); + directory + } + + #[tokio::test] + async fn async_checkpoint_matches_the_synchronous_path_bit_for_bit() { + let directory = unique_test_directory("async-checkpoint-equal"); + let learner = member(0, 30); + let mut outputs = Vec::new(); + for asynchronous in [false, true] { + let checkpoint = directory.join(format!("state-{asynchronous}.ckpt")); + let tape = directory.join(format!("tape-{asynchronous}.jsonl")); + let (group, mut frames) = streaming_test_group(DTYPE_F32, false); + let registry = registry_with_group(group); + let mut config = round_test_config(4); + config.pipeline = 2; + config.checkpoint_path = Some(checkpoint.clone()); + config.checkpoint_every = 1; + config.async_checkpoint = asynchronous; + config.event_tape = Some(tape.clone()); + let mut state = two_fragment_streaming_state(); + let mut sync_seconds = 0.0; + for step in 1..=2u64 { + let round = sweep_test_round(step, (step - 1) as usize, 1, step, learner); + if asynchronous { + complete_round_async(&config, &mut state, ®istry, &mut sync_seconds, round) + .await + .unwrap(); + } else { + complete_round(&config, &mut state, ®istry, &mut sync_seconds, round) + .await + .unwrap(); + } + // Every non-terminal round is broadcast exactly once. + assert_eq!( + chunked_inner_type(&frames.try_recv().unwrap()), + MSG_BCAST_FRAGMENT + ); + } + let tape_text = std::fs::read_to_string(&tape).unwrap(); + assert_eq!(tape_text.lines().count(), 2); + assert!(tape_text.lines().all(|line| line.contains("\"checkpoint_ms\":"))); + outputs.push((state.params.clone(), std::fs::read(&checkpoint).unwrap())); + // The snapshot encoder writes exactly the streaming file format. + let mut on_disk = two_fragment_streaming_state(); + on_disk.load_checkpoint(&checkpoint).unwrap(); + assert_eq!(on_disk.encode_checkpoint().unwrap(), std::fs::read(&checkpoint).unwrap()); + } + assert_eq!(outputs[0].0, outputs[1].0); + assert_eq!(outputs[0].1, outputs[1].1); + std::fs::remove_dir_all(&directory).ok(); + } + + #[tokio::test] + async fn async_checkpoint_failure_withholds_broadcast_and_keeps_previous_checkpoint() { + let directory = unique_test_directory("async-checkpoint-failure"); + let checkpoint = directory.join("state.ckpt"); + let learner = member(0, 30); + let (group, mut frames) = streaming_test_group(DTYPE_F32, false); + let registry = registry_with_group(group); + let mut config = round_test_config(4); + config.pipeline = 2; + config.checkpoint_path = Some(checkpoint.clone()); + config.checkpoint_every = 1; + config.async_checkpoint = true; + let mut state = two_fragment_streaming_state(); + let mut sync_seconds = 0.0; + complete_round_async( + &config, + &mut state, + ®istry, + &mut sync_seconds, + sweep_test_round(1, 0, 1, 1, learner), + ) + .await + .unwrap(); + assert_eq!( + chunked_inner_type(&frames.try_recv().unwrap()), + MSG_BCAST_FRAGMENT + ); + // The atomic writer creates this exact temporary path; a directory + // there makes the next write fail deterministically. + std::fs::create_dir(checkpoint.with_extension("tmp")).unwrap(); + let error = complete_round_async( + &config, + &mut state, + ®istry, + &mut sync_seconds, + sweep_test_round(2, 1, 1, 1, learner), + ) + .await + .unwrap_err(); + assert!(format!("{error:#}").contains("step 2 not broadcast")); + assert!(matches!( + frames.try_recv(), + Err(tokio::sync::mpsc::error::TryRecvError::Empty) + )); + // The resume point is still the last durable, broadcast cut. + let mut recovered = two_fragment_streaming_state(); + recovered.load_checkpoint(&checkpoint).unwrap(); + assert_eq!(recovered.global_step, 1); + assert_eq!(recovered.versions, vec![1, 0]); + std::fs::remove_dir_all(&directory).ok(); + } + #[tokio::test] async fn policy_sweep_checkpoint_resume_mid_sweep_does_not_double_account() { let directory = std::env::temp_dir().join(format!( @@ -6227,6 +6665,7 @@ mod tests { Some(11), Some(22), 33, + None, &pushes, 0.5, 44, @@ -6275,6 +6714,7 @@ mod tests { Some(1), Some(0), 2, + None, &pushes, 0.5, 3, @@ -6303,6 +6743,7 @@ mod tests { Some(1), Some(0), 2, + None, &pushes, 0.5, 3, @@ -6410,6 +6851,7 @@ mod tests { Some(1), Some(0), 2, + None, &pushes, 0.5, 3, diff --git a/syncer/src/state.rs b/syncer/src/state.rs index f19ce4c9..6d441561 100644 --- a/syncer/src/state.rs +++ b/syncer/src/state.rs @@ -628,6 +628,36 @@ pub fn remove_final_marker(path: &std::path::Path) -> Result<()> { } } +/// Little-endian f32 values in 64 KiB blocks (same bytes as per-value +/// `to_le_bytes`, without one writer call per element). +fn write_f32s(f: &mut W, values: &[f32]) -> Result<()> { + let mut block = Vec::with_capacity(16 * 1024 * 4); + for chunk in values.chunks(16 * 1024) { + block.clear(); + for v in chunk { + block.extend_from_slice(&v.to_le_bytes()); + } + f.write_all(&block)?; + } + Ok(()) +} + +/// Atomically persist checkpoint bytes from `encode_checkpoint`: write a +/// sibling temporary file, fsync it, rename it over `path`, fsync the +/// directory. A failure at any point leaves the previous checkpoint intact. +pub fn write_checkpoint_bytes(path: &std::path::Path, bytes: &[u8]) -> Result<()> { + use std::io::Write; + let tmp = path.with_extension("tmp"); + { + let mut f = std::fs::File::create(&tmp) + .with_context(|| format!("create checkpoint temporary {}", tmp.display()))?; + f.write_all(bytes)?; + f.sync_all()?; + } + std::fs::rename(&tmp, path)?; + sync_parent_directory(path) +} + fn sync_parent_directory(path: &std::path::Path) -> Result<()> { let parent = path.parent().unwrap_or_else(|| std::path::Path::new(".")); std::fs::File::open(parent) @@ -660,47 +690,11 @@ impl GlobalState { /// Written to `.tmp` then renamed, so a crash mid-write never /// corrupts the previous checkpoint. pub fn save_checkpoint(&self, path: &std::path::Path) -> Result<()> { - use std::io::Write; - let layout_fingerprint = self.layout_fingerprint; let tmp = path.with_extension("tmp"); { let mut f = std::io::BufWriter::new(std::fs::File::create(&tmp)?); - f.write_all(&CKPT_MAGIC_V3.to_le_bytes())?; - f.write_all(&[self.iso_backend_kind() as u8])?; - f.write_all(&layout_fingerprint)?; - f.write_all(&self.global_step.to_le_bytes())?; - f.write_all(&(self.params.len() as u32).to_le_bytes())?; - for p in 0..self.params.len() { - f.write_all(&self.versions[p].to_le_bytes())?; - f.write_all(&(self.params[p].len() as u64).to_le_bytes())?; - for v in &self.params[p] { - f.write_all(&v.to_le_bytes())?; - } - for v in &self.momentum[p] { - f.write_all(&v.to_le_bytes())?; - } - } - f.write_all(&(self.ledger.len() as u32).to_le_bytes())?; - for (id, l) in &self.ledger { - f.write_all(&id.to_le_bytes())?; - f.write_all(&l.merges.to_le_bytes())?; - f.write_all(&l.steps.to_le_bytes())?; - f.write_all(&l.tokens.to_le_bytes())?; - } - f.write_all(&self.layout_fingerprint)?; - if let Some(fragments) = self.policy_sweep_fragments { - f.write_all(&POLICY_SWEEP_CKPT_MAGIC.to_le_bytes())?; - f.write_all(&fragments.to_le_bytes())?; - f.write_all( - &self - .session_contract_hash - .context("policy-sweep checkpoint requires a session contract hash")?, - )?; - } else if let Some(contract_hash) = self.session_contract_hash { - f.write_all(&SESSION_CONTRACT_CKPT_MAGIC.to_le_bytes())?; - f.write_all(&contract_hash)?; - } - f.flush()?; + self.write_checkpoint_to(&mut f)?; + std::io::Write::flush(&mut f)?; f.get_ref().sync_all()?; } std::fs::rename(&tmp, path)?; @@ -708,6 +702,54 @@ impl GlobalState { Ok(()) } + /// The exact bytes `save_checkpoint` writes, as an owned snapshot. + /// + /// The asynchronous checkpoint path encodes on the scheduler (one memory + /// copy of the consistent cut) and hands the bytes to a blocking writer, + /// so later commits may mutate `self` while the snapshot is persisted. + pub fn encode_checkpoint(&self) -> Result> { + let state_floats: usize = self.params.iter().map(Vec::len).sum::() * 2; + let mut bytes = Vec::with_capacity(state_floats * 4 + 4096); + self.write_checkpoint_to(&mut bytes)?; + Ok(bytes) + } + + fn write_checkpoint_to(&self, f: &mut W) -> Result<()> { + let layout_fingerprint = self.layout_fingerprint; + f.write_all(&CKPT_MAGIC_V3.to_le_bytes())?; + f.write_all(&[self.iso_backend_kind() as u8])?; + f.write_all(&layout_fingerprint)?; + f.write_all(&self.global_step.to_le_bytes())?; + f.write_all(&(self.params.len() as u32).to_le_bytes())?; + for p in 0..self.params.len() { + f.write_all(&self.versions[p].to_le_bytes())?; + f.write_all(&(self.params[p].len() as u64).to_le_bytes())?; + write_f32s(f, &self.params[p])?; + write_f32s(f, &self.momentum[p])?; + } + f.write_all(&(self.ledger.len() as u32).to_le_bytes())?; + for (id, l) in &self.ledger { + f.write_all(&id.to_le_bytes())?; + f.write_all(&l.merges.to_le_bytes())?; + f.write_all(&l.steps.to_le_bytes())?; + f.write_all(&l.tokens.to_le_bytes())?; + } + f.write_all(&self.layout_fingerprint)?; + if let Some(fragments) = self.policy_sweep_fragments { + f.write_all(&POLICY_SWEEP_CKPT_MAGIC.to_le_bytes())?; + f.write_all(&fragments.to_le_bytes())?; + f.write_all( + &self + .session_contract_hash + .context("policy-sweep checkpoint requires a session contract hash")?, + )?; + } else if let Some(contract_hash) = self.session_contract_hash { + f.write_all(&SESSION_CONTRACT_CKPT_MAGIC.to_le_bytes())?; + f.write_all(&contract_hash)?; + } + Ok(()) + } + /// Restore params/momentum/versions/step/ledger from a snapshot. /// /// V3 binds the snapshot to both the ISO backend and exact semantic HELLO diff --git a/tests/test_rl_dense_sweep_real_syncer.py b/tests/test_rl_dense_sweep_real_syncer.py index fd14e997..0976992a 100644 --- a/tests/test_rl_dense_sweep_real_syncer.py +++ b/tests/test_rl_dense_sweep_real_syncer.py @@ -41,23 +41,25 @@ def _hmac_key(monkeypatch): monkeypatch.setenv("YETO_ISLAND_HMAC_KEY", "dense-sweep-test-key") -def _dense_args(): +def _dense_args(*extra: str, total_steps: int = 1): from yeto.launcher import _prepare_rl_args args = _launcher_args("ports", ( "--rl-sync-preset", "full-parameter", "--rl-placement", "fixed-partition", "--rl-rollout-gpus", "1", "--rl-parameter-layout-sha256", "e" * 64, - "--fragments", str(FRAGMENTS), "--total-steps", "1"), gpu="modal:2xh100") + "--fragments", str(FRAGMENTS), "--total-steps", str(total_steps), *extra), + gpu="modal:2xh100") args.controller = "local" _prepare_rl_args(args) return args -def _start_syncer(binary: Path, tmp_path: Path) -> tuple[subprocess.Popen, int, str]: +def _start_syncer(binary: Path, tmp_path: Path, *extra: str, + total_steps: int = 1) -> tuple[subprocess.Popen, int, str]: """The launcher's dense actor syncer, relocated to a free port and tmp_path.""" from yeto import launcher - args = _dense_args() + args = _dense_args(*extra, total_steps=total_steps) command = launcher.role_syncer_commands(args, 1, binary=str(binary))["actor"] with socket.socket() as probe: probe.bind(("127.0.0.1", 0)) @@ -153,3 +155,80 @@ def test_launcher_dense_syncer_refuses_another_profile_hash(syncer_binary, tmp_p client.close() log = _stop(process) assert "syncer semantic profile does not match" in log + + +STREAMED = ("--rl-sync-fragment-pipeline", str(FRAGMENTS), "--rl-sync-async-checkpoint") + + +def _run_sweeps(binary: Path, tmp_path: Path, extra: tuple[str, ...], sweeps: int): + """Drive ``sweeps`` strict sweeps; return (max PULLs seen at once, finals, checkpoint).""" + process, port, profile_hash = _start_syncer(binary, tmp_path, *extra, total_steps=sweeps) + client = _client(port, profile_hash) + concurrent = 0 + try: + client.start() + for fid in range(FRAGMENTS): + client.send_init(fid, _f32([float(fid) + 0.125 * i for i in range(NUMEL)])) + for sweep in range(1, sweeps + 1): + permits: dict[int, object] = {} + + def gather(): + client.check_health() + for permit in client.drain_pulls(): + permits[permit.fragment_id] = permit + return any(not getattr(p, "pushed", False) for p in permits.values()) + + while len(permits) < FRAGMENTS: + _wait(gather, f"no PULL in sweep {sweep}") + time.sleep(0.2) # let every PULL the profile permits arrive + gather() + concurrent = max(concurrent, sum(not getattr(p, "pushed", False) + for p in permits.values())) + # Push only what is permitted now; the serial profile then + # releases the next fragment's PULL after this commit. + for fid, permit in sorted(permits.items()): + if getattr(permit, "pushed", False): + continue + base = 0 if sweep == 1 else (sweep - 2) * FRAGMENTS + fid + 1 + delta = [0.5 + fid + 0.03125 * sweep + 0.001 * i for i in range(NUMEL)] + client.push_fragment(fid, permit.global_step, permit.round_attempt, + base, sweep, 1, 7, _f32(delta)) + permit.pushed = True + manifest, finals = client.wait_for_final_fragments(timeout=30) + result = {f.fragment_id: bytes(f.data) for f in finals} + client.acknowledge_finalization(manifest) + assert process.wait(timeout=15) == 0 + finally: + client.close() + log = _stop(process) + assert "syncer semantic profile" not in log, log + return concurrent, list(manifest.versions), result, (tmp_path / "state.ckpt").read_bytes(), log + + +def test_streamed_sweep_matches_the_serial_sweep_bit_for_bit(syncer_binary, tmp_path): + """E3: fragment pipelining + async checkpoint changes timing, not results.""" + serial_dir, streamed_dir = tmp_path / "serial", tmp_path / "streamed" + serial_dir.mkdir() + streamed_dir.mkdir() + serial = _run_sweeps(syncer_binary, serial_dir, (), sweeps=3) + streamed = _run_sweeps(syncer_binary, streamed_dir, STREAMED, sweeps=3) + assert serial[0] == 1, "the serial profile must permit one fragment at a time" + assert streamed[0] == FRAGMENTS, "the streamed profile must permit a whole sweep at once" + assert serial[1] == streamed[1] == [5, 6] + assert serial[2] == streamed[2] + assert serial[3] == streamed[3] + assert "--async-checkpoint" not in serial[4] + + +def test_streamed_flags_bind_the_profile_and_are_full_parameter_only(): + serial = _dense_args() + streamed = _dense_args(*STREAMED) + assert serial.pipeline == 1 and streamed.pipeline == FRAGMENTS + from yeto import launcher + + command = launcher.role_syncer_commands(streamed, 1, binary="yeto-syncer")["actor"] + assert f"--pipeline {FRAGMENTS}" in command and "--async-checkpoint" in command + assert (launcher.rl_syncer_profile(serial, 1).sha256 + != launcher.rl_syncer_profile(streamed, 1).sha256) + with pytest.raises(ValueError, match="--rl-sync-fragment-pipeline between 1 and --fragments"): + _dense_args("--rl-sync-fragment-pipeline", str(FRAGMENTS + 1)) diff --git a/tests/test_rl_sao_streaming_runtime.py b/tests/test_rl_sao_streaming_runtime.py index ced8e08e..88ec57b2 100644 --- a/tests/test_rl_sao_streaming_runtime.py +++ b/tests/test_rl_sao_streaming_runtime.py @@ -279,7 +279,7 @@ def test_load_or_role_pairing_rejects_role_or_lockstep_drift(tmp_path, mutation, ("checkpoint_every", 2, "checkpointing every round"), ("resume", False, "checkpointing every round"), ("max_base_lag", 1, "max_base_lag=0"), - ("policy_sweep_fragments", 4, "policy sweep"), + ("policy_sweep_fragments", 4, "policy.sweep"), ("learner_budget_steps", 2, "learner-budget cutoff mode"), ), ) diff --git a/tests/test_run_flags.py b/tests/test_run_flags.py index 9ef5c8ea..99b5ea01 100644 --- a/tests/test_run_flags.py +++ b/tests/test_run_flags.py @@ -250,7 +250,7 @@ def test_forward_tables_name_flags_of_both_layers(): # flag renamed in the schema but still spelled the old way in the launcher fails. _NON_RL_LITERALS = frozenset({ "--abbrev-ref", "--adapter-dir", "--add", "--address", "--allow-unattested-legacy-adapter", - "--allow-unauthenticated-islands", "--bind", "--cap-add", "--carry-gamma", + "--allow-unauthenticated-islands", "--async-checkpoint", "--bind", "--cap-add", "--carry-gamma", "--checkpoint-every", "--checkpoint-path", "--codex-harness-binary", "--config-json", "--create", "--dashboard-host", "--dearmor", "--depth", "--detach", "--dtype", "--entrypoint", "--expected-island-contract", "--final-ack-timeout-s", "--final-grace-s", "--follow", "--force", diff --git a/tests/test_wandb_tape.py b/tests/test_wandb_tape.py index 0a9abf08..2305d7fc 100644 --- a/tests/test_wandb_tape.py +++ b/tests/test_wandb_tape.py @@ -176,13 +176,15 @@ def test_every_rust_tape_field_is_consumed_or_deliberately_skipped(): # The responders literal is built first, the record literal last. responder_fmt, record_fmt = body.split("let line = format!(", 1) top_level = set(re.findall(r'\\"(\w+)\\":', record_fmt)) + optional = source.split("fn checkpoint_tape_field(", 1)[1].split("\nfn ", 1)[0] + top_level |= set(re.findall(r'\\"(\w+)\\":', optional)) responder_keys = set(re.findall(r'\\"(\w+)\\":', responder_fmt)) assert responder_keys, "could not find the responder literal in append_tape" consumed_top = { "step", "fragment", "attempt", "gnorm", "ms", "sync_ms", "quorum", "expected", "responded", "missed_grace", "launch_base_version", - "quorum_ms", "grace_ms", "responders", + "quorum_ms", "grace_ms", "responders", "checkpoint_ms", } # Constants and the member-identity mirrors of the id lists: nothing a # time series can be drawn from. diff --git a/tools/bench_sync_boundary.py b/tools/bench_sync_boundary.py new file mode 100644 index 00000000..3b449632 --- /dev/null +++ b/tools/bench_sync_boundary.py @@ -0,0 +1,163 @@ +"""CPU micro-benchmark of the full-parameter sync boundary (S26 E3). + +Runs the real Rust syncer on localhost with two in-process learners and a +synthetic f32 model, and times each strict sweep from the first push to the +last updated fragment received by every learner. It compares the serial +profile (pipeline 1, checkpoint on the scheduler) with the streamed profile +(pipeline P, ``--async-checkpoint``). + +These are local CPU numbers: loopback transfer and a local NVMe checkpoint. +They show where the boundary time goes on this machine; they do not predict +GPU-cluster or WAN timings. + + python tools/bench_sync_boundary.py --floats-per-fragment 33554432 --sweeps 4 +""" + +from __future__ import annotations + +import argparse +import json +import socket +import statistics +import subprocess +import tempfile +import threading +import time +from pathlib import Path + +import numpy as np + +from yeto.fragments import MERGE_RDA, Fragment, FragmentLayout +from yeto.protocol import DTYPE_F32, SyncerClient + +ROOT = Path(__file__).resolve().parent.parent + + +def _free_port() -> int: + with socket.socket() as probe: + probe.bind(("127.0.0.1", 0)) + return probe.getsockname()[1] + + +def _syncer(binary: Path, workdir: Path, *, port: int, learners: int, fragments: int, + sweeps: int, streamed: bool) -> subprocess.Popen: + command = [ + str(binary), "--port", str(port), "--learners", str(learners), + "--quorum", str(learners), "--grace-ms", "0", + "--pipeline", str(fragments if streamed else 1), + "--sync-interval-steps", "0", "--delta-correction", "none", + "--total-steps", str(sweeps * fragments), "--policy-sweep-fragments", str(fragments), + "--outer-lr", "1", "--outer-momentum", "0", "--max-base-lag", "0", + "--learner-weight", "equal", "--checkpoint-path", str(workdir / "state.ckpt"), + "--checkpoint-every", "1", "--event-tape", str(workdir / "tape.jsonl"), + "--allow-unauthenticated-islands", # loopback-only benchmark + ] + if streamed: + command.append("--async-checkpoint") + return subprocess.Popen(command, stdout=open(workdir / "syncer.log", "w"), + stderr=subprocess.STDOUT) + + +def _learner(learner_id: int, port: int, layout: FragmentLayout, numel: int, sweeps: int, + timings: list[list[float]], errors: list[BaseException]) -> None: + fragments = len(layout.fragments) + client = SyncerClient(("127.0.0.1", port), learner_id, layout, dtype=DTYPE_F32, + num_streams=1, connect_timeout=30, max_reconnects=0, + finalization_timeout=600) + rng = np.random.default_rng(learner_id) + try: + client.start() + if learner_id == 0: + for fid in range(fragments): + client.send_init(fid, rng.standard_normal(numel, dtype=np.float32).tobytes()) + deltas = [rng.standard_normal(numel, dtype=np.float32).tobytes() for _ in range(fragments)] + carried: list = [] # a PULL of the next sweep can arrive before this sweep's BCAST + for sweep in range(1, sweeps + 1): + pushed: set[int] = set() + updated: set[int] = set() + first_push = None + while len(updated) < fragments or len(pushed) < fragments: + client.check_health() + permits, carried = carried + client.drain_pulls(), [] + for permit in permits: + if permit.fragment_id in pushed: + carried.append(permit) + continue + base = 0 if sweep == 1 else (sweep - 2) * fragments + permit.fragment_id + 1 + first_push = first_push or time.monotonic() + client.push_fragment(permit.fragment_id, permit.global_step, + permit.round_attempt, base, sweep, 1, 1, + deltas[permit.fragment_id]) + pushed.add(permit.fragment_id) + for update in client.drain_updates(): + if update.version > (sweep - 1) * fragments: + updated.add(update.fragment_id) + if sweep == sweeps and len(pushed) == fragments: + break # the terminal sweep is delivered by FINAL, not BCAST + time.sleep(0.001) + if sweep == sweeps: + manifest, _ = client.wait_for_final_fragments(timeout=600) + timings[learner_id].append(time.monotonic() - first_push) + client.acknowledge_finalization(manifest) + else: + timings[learner_id].append(time.monotonic() - first_push) + except BaseException as error: # noqa: BLE001 - reported by the caller + errors.append(error) + finally: + client.close() + + +def run(binary: Path, *, streamed: bool, learners: int, fragments: int, numel: int, + sweeps: int) -> dict: + layout = FragmentLayout([Fragment(MERGE_RDA, [(f"w{i}", numel)]) for i in range(fragments)]) + with tempfile.TemporaryDirectory(prefix="e3-bench-", dir="/tmp") as tmp: + workdir = Path(tmp) + port = _free_port() + process = _syncer(binary, workdir, port=port, learners=learners, fragments=fragments, + sweeps=sweeps, streamed=streamed) + timings: list[list[float]] = [[] for _ in range(learners)] + errors: list[BaseException] = [] + threads = [threading.Thread(target=_learner, args=(i, port, layout, numel, sweeps, + timings, errors)) + for i in range(learners)] + for thread in threads: + thread.start() + for thread in threads: + thread.join() + code = process.wait(timeout=120) + if errors or code != 0: + raise RuntimeError(f"benchmark failed (syncer exit {code}): {errors!r}; " + f"log: {(workdir / 'syncer.log').read_text()[-2000:]}") + tape = [json.loads(line) for line in (workdir / "tape.jsonl").read_text().splitlines() + if '"sync_ms"' in line] + # Sweep 1 includes connection set-up; the last sweep ends with FINAL. + steady = [max(per[i] for per in timings) for i in range(1, sweeps - 1)] or \ + [max(per[i] for per in timings) for i in range(sweeps)] + return { + "profile": "streamed" if streamed else "serial", + "boundary_s_per_sweep": [round(x, 3) for x in steady], + "boundary_s_median": round(statistics.median(steady), 3), + "checkpoint_ms_median": statistics.median(r["checkpoint_ms"] for r in tape), + "sync_ms_median": statistics.median(r["sync_ms"] for r in tape), + } + + +def main() -> None: + parser = argparse.ArgumentParser(description=__doc__.splitlines()[0]) + parser.add_argument("--binary", type=Path, default=ROOT / "syncer/target/release/yeto-syncer") + parser.add_argument("--learners", type=int, default=2) + parser.add_argument("--fragments", type=int, default=2) + parser.add_argument("--floats-per-fragment", type=int, default=32 * 1024 * 1024) + parser.add_argument("--sweeps", type=int, default=4) + args = parser.parse_args() + results = [run(args.binary, streamed=streamed, learners=args.learners, + fragments=args.fragments, numel=args.floats_per_fragment, sweeps=args.sweeps) + for streamed in (False, True)] + print(json.dumps({"note": "CPU localhost numbers; not GPU or WAN timings", + "learners": args.learners, "fragments": args.fragments, + "floats_per_fragment": args.floats_per_fragment, "results": results}, + indent=2)) + + +if __name__ == "__main__": + main() diff --git a/yeto/launch/rl_args.py b/yeto/launch/rl_args.py index 0c21669b..e1ce4ff2 100644 --- a/yeto/launch/rl_args.py +++ b/yeto/launch/rl_args.py @@ -169,11 +169,16 @@ def _prepare_dense_full_sweep(args, *, num_islands: int) -> None: problems.append("--rl-island-scheduling legacy (no fragment frames in elastic mode)") if rl_needs_critic(args): problems.append("a critic-free algorithm (the strict sweep trains the actor only)") + fragment_pipeline = int(getattr(args, "rl_sync_fragment_pipeline", 1)) + if not 1 <= fragment_pipeline <= max(args.fragments, 1): + problems.append("--rl-sync-fragment-pipeline between 1 and --fragments") if problems: raise ValueError(f"--rl-sync-preset {FULL_PARAMETER_SYNC_PRESET} requires " + "; ".join(problems)) args.quorum = num_islands args.grace_ms = 0 - args.pipeline = 1 + # Fragment rounds of one sweep are independent (H=1, plain average), so the + # syncer may gather several at once; it still commits them in step order. + args.pipeline = fragment_pipeline args.sync_interval_steps = 0.0 args.delta_correction = "none" args.outer_lr = 1.0 diff --git a/yeto/launch/syncer.py b/yeto/launch/syncer.py index 93b1aa6a..7d951ae9 100644 --- a/yeto/launch/syncer.py +++ b/yeto/launch/syncer.py @@ -161,7 +161,8 @@ def _syncer_profile_binding(args) -> str: full-parameter adds strict P-fragment policy sweeps; sao-streaming-full adds the frozen contract's timeouts and the publishable final checkpoint.""" if rl_dense_sweep(args): - return f" --policy-sweep-fragments {args.fragments} --require-profile-binding" + async_checkpoint = " --async-checkpoint" if getattr(args, "rl_sync_async_checkpoint", False) else "" + return f" --policy-sweep-fragments {args.fragments} --require-profile-binding{async_checkpoint}" if rl_sao_streaming(args): profile = args.rl_sao_manifest.profile return (f" --quorum-timeout-s {profile.quorum_timeout_s}" diff --git a/yeto/rl/engine/run_flags.py b/yeto/rl/engine/run_flags.py index b4de9148..f9be06f6 100644 --- a/yeto/rl/engine/run_flags.py +++ b/yeto/rl/engine/run_flags.py @@ -841,6 +841,20 @@ def rl_forward_pairs(*dests: str | tuple[str, str]) -> tuple[tuple[str, str], .. "sao-streaming-full: full-parameter SAO, actor and critic streams " "(ports engine, --rl-placement fixed-partition, --rl-sao-contracts)", ), + cli_flag( + "--rl-sync-fragment-pipeline", + type=int, + default=1, + help="full-parameter only: fragment rounds the syncer gathers at once " + "(1..--fragments). Above 1, fragment k+1 uploads while fragment k merges " + "and downloads; merge results are unchanged (default 1 = serial)", + ), + cli_flag( + "--rl-sync-async-checkpoint", + action="store_true", + help="full-parameter only: the syncer writes each per-round checkpoint off " + "its scheduler and broadcasts that round only once the checkpoint is durable", + ), cli_flag( "--rl-initial-adapter", default=None, diff --git a/yeto/syncer_profile.py b/yeto/syncer_profile.py index 390d2bc4..c5e011b0 100644 --- a/yeto/syncer_profile.py +++ b/yeto/syncer_profile.py @@ -259,7 +259,7 @@ def _validate_general(self) -> None: raise ValueError("syncer profile resume requires checkpointing") if self.policy_sweep_fragments is not None: fragments = self.policy_sweep_fragments - if self.pipeline != 1 or self.total_steps % fragments: + if not 1 <= self.pipeline <= fragments or self.total_steps % fragments: raise ValueError("syncer profile policy sweep schedule is invalid") def validate_sao(self) -> None: diff --git a/yeto/wandb_tape.py b/yeto/wandb_tape.py index 125afa62..59c4e584 100644 --- a/yeto/wandb_tape.py +++ b/yeto/wandb_tape.py @@ -110,6 +110,8 @@ def tape_metrics(rec: dict) -> dict: ): if source in rec: metrics[metric] = int(rec[source]) + if "checkpoint_ms" in rec: # per-round checkpoint wall time (S26 E3) + metrics["sync/checkpoint_ms"] = float(rec["checkpoint_ms"]) if "sweep_complete" in rec: metrics["sync/sweep_complete"] = int(bool(rec["sweep_complete"])) if expected: 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 03/62] 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 eb9e22905aea7ba61f659c94f838dcd991df18ab Mon Sep 17 00:00:00 2001 From: michael <78113554+chowhsu@users.noreply.github.com> Date: Tue, 20 Oct 2026 14:56:04 +0000 Subject: [PATCH 04/62] syncer: bf16 sweep pushes and sweep-aligned checkpoint intervals Add session dtype 4 (f32 session, bf16 PUSH deltas): INIT/BCAST/FINAL, masters and momentum stay f32; the client rounds f32 delta bytes to bf16 (round-to-nearest-even, streamed across part boundaries). The dtype is bound into the role-stream session contract. Launcher --rl-sync-push-dtype, island --sweep-push-dtype (default f32). Allow --checkpoint-every N (1 or a multiple of the fragment count) for policy sweeps: only checkpointed steps are durable before their broadcast, on both the synchronous and async paths; the interval is part of the semantic profile hash. Launcher --rl-sync-checkpoint-every (default 1). E3 launcher options are refused outside the full-parameter preset. Adds bf16 error-bound, interval/resume and real-syncer tests, and a third micro-benchmark profile. Co-Authored-By: Claude Opus 5.5 --- docs/PROTOCOL.md | 9 +- syncer/src/protocol.rs | 19 ++- syncer/src/server.rs | 150 ++++++++++++++++-- tests/test_e3_sync_options.py | 63 ++++++++ tests/test_rl_dense_sweep_real_syncer.py | 48 +++++- tools/bench_sync_boundary.py | 40 +++-- yeto/launch/rl_args.py | 20 +++ yeto/launch/rl_island_task.py | 2 + yeto/launch/syncer.py | 9 +- yeto/protocol.py | 53 ++++++- .../miles/full_parameter/role_stream.py | 7 +- yeto/rl/adapters/miles/island_entry.py | 1 + yeto/rl/adapters/miles/run_flags.py | 6 + yeto/rl/engine/run_flags.py | 16 ++ yeto/syncer_profile.py | 4 + 15 files changed, 404 insertions(+), 43 deletions(-) create mode 100644 tests/test_e3_sync_options.py diff --git a/docs/PROTOCOL.md b/docs/PROTOCOL.md index 579db093..3d08e098 100644 --- a/docs/PROTOCOL.md +++ b/docs/PROTOCOL.md @@ -32,7 +32,7 @@ or allocation overflows are errors rather than panics. | type | name | direction | payload | |---:|---|---|---| -| 1 | HELLO | learner → syncer | protocol_version:u16 (=4), learner_id:u32, connection_generation:u64, dtype:u8 (1=f32, 2=bf16, 3=q4), num_fragments:u32, per-fragment layout, layout_fingerprint:[u8;32], session_contract_hash:[u8;32], optional syncer_profile_hash:[u8;32], num_streams:u16, optional compat_group trailer (`b"YCG1"`, u32 length, 1–64 bytes utf-8) | +| 1 | HELLO | learner → syncer | protocol_version:u16 (=4), learner_id:u32, connection_generation:u64, dtype:u8 (1=f32, 2=bf16, 3=q4, 4=f32 with bf16 pushes), num_fragments:u32, per-fragment layout, layout_fingerprint:[u8;32], session_contract_hash:[u8;32], optional syncer_profile_hash:[u8;32], num_streams:u16, optional compat_group trailer (`b"YCG1"`, u32 length, 1–64 bytes utf-8) | | 2 | INIT_PARAMS | learner → syncer | fragment_id:u32, full tensor bytes; only learner 0 may initialize | | 3 | PULL_REQ | syncer → learner | fragment_id:u32, global_step:u64, round_attempt:u32 | | 4 | PUSH_FRAGMENT | learner → syncer | learner_id:u32, fragment_id:u32, global_step:u64, round_attempt:u32, base_version:u64, local_step:u64, c_steps:u32, c_tokens:u64, base-relative learner-delta bytes | @@ -223,6 +223,13 @@ The coordinator is one sequential scheduler actor. Pipelined rounds target distinct fragments, and merge + version update + broadcast for a completion is serialized. No two completions mutate one fragment concurrently. +## bf16 pushes over an f32 session (dtype 4) + +Dtype 4 keeps INIT, BCAST and FINAL exact f32 and sends only PUSH deltas as +bf16 (round-to-nearest-even). The syncer decodes and negates them like dtype 2 +and keeps f32 masters and momentum. The full-parameter policy sweep accepts +dtype 1 or 4 (S26 E3, `--rl-sync-push-dtype bf16`). + ## Q4 pushes Q4 uses the same learner-delta semantics. INIT and BCAST stay bf16; PUSH is diff --git a/syncer/src/protocol.rs b/syncer/src/protocol.rs index bf4b530e..2103ce3b 100644 --- a/syncer/src/protocol.rs +++ b/syncer/src/protocol.rs @@ -57,6 +57,10 @@ pub const DTYPE_BF16: u8 = 2; /// base-relative learner deltas; INIT_PARAMS and BCAST_FRAGMENT travel as /// bf16 (see `bulk_dtype`, docs/PROTOCOL.md v4). pub const DTYPE_Q4: u8 = 3; +/// Session dtype 4 (S26 E3): PUSH_FRAGMENT deltas travel as bf16 while +/// INIT_PARAMS, BCAST_FRAGMENT and FINAL stay exact f32. The syncer keeps f32 +/// masters and momentum; only the learner's outer gradient is rounded. +pub const DTYPE_F32_PUSH_BF16: u8 = 4; /// Values per q4 scale block (f32 absmax scale + 128 nibble bytes each). pub const Q4_BLOCK: usize = 256; @@ -64,7 +68,17 @@ pub const Q4_BLOCK: usize = 256; /// Dtype of INIT_PARAMS/BCAST_FRAGMENT tensors for a session dtype: q4 /// applies only to push deltas; full parameter payloads stay bf16. pub fn bulk_dtype(dtype: u8) -> u8 { - if dtype == DTYPE_Q4 { + match dtype { + DTYPE_Q4 => DTYPE_BF16, + DTYPE_F32_PUSH_BF16 => DTYPE_F32, + other => other, + } +} + +/// Dtype of PUSH_FRAGMENT deltas for an unquantized session dtype (q4 has +/// its own block codec). +pub fn push_dtype(dtype: u8) -> u8 { + if dtype == DTYPE_F32_PUSH_BF16 { DTYPE_BF16 } else { dtype @@ -417,6 +431,9 @@ mod tests { assert_eq!(bulk_dtype(DTYPE_Q4), DTYPE_BF16); assert_eq!(bulk_dtype(DTYPE_BF16), DTYPE_BF16); assert_eq!(bulk_dtype(DTYPE_F32), DTYPE_F32); + assert_eq!(bulk_dtype(DTYPE_F32_PUSH_BF16), DTYPE_F32); + assert_eq!(push_dtype(DTYPE_F32_PUSH_BF16), DTYPE_BF16); + assert_eq!(push_dtype(DTYPE_F32), DTYPE_F32); } #[test] diff --git a/syncer/src/server.rs b/syncer/src/server.rs index e7b42b3e..6d0b4ba1 100644 --- a/syncer/src/server.rs +++ b/syncer/src/server.rs @@ -735,10 +735,14 @@ fn validate_config(cfg: &Config) -> Result<()> { // --resume is not required here: a fresh sweep has no checkpoint yet and // --resume refuses a missing file (see check_resume_checkpoint). The // launcher adds --resume only when the checkpoint exists. - if cfg.checkpoint_path.is_none() || cfg.checkpoint_every != 1 { + if cfg.checkpoint_path.is_none() + || cfg.checkpoint_every == 0 + || (cfg.checkpoint_every != 1 + && !cfg.checkpoint_every.is_multiple_of(u64::from(fragments))) + { bail!( "--policy-sweep-fragments requires --checkpoint-path and \ - --checkpoint-every 1 for crash-safe sweeps" + --checkpoint-every 1 or a multiple of the fragment count" ); } } @@ -760,6 +764,20 @@ fn checkpoint_before_broadcast(cfg: &Config) -> bool { cfg.checkpoint_path.is_some() && cfg.checkpoint_every == 1 } +/// Whether step `t` must be durable before its BCAST. Per-round checkpoints +/// always are; a policy sweep with `--checkpoint-every N` (N a multiple of +/// the fragment count) persists only every N-th step, at sweep boundaries, +/// and broadcasts the other steps without a durable record (a crash resumes +/// at the newest checkpoint; islands ahead of it fail closed on their next +/// push's local-step check). +fn persist_before_broadcast(cfg: &Config, t: u64) -> bool { + checkpoint_before_broadcast(cfg) + || (cfg.policy_sweep_fragments.is_some() + && cfg.checkpoint_path.is_some() + && cfg.checkpoint_every > 0 + && t.is_multiple_of(cfg.checkpoint_every)) +} + fn validate_policy_sweep_push(push: &Push, fragments: u32) -> Result<()> { if fragments == 0 { bail!("policy sweep fragment count must be positive"); @@ -1000,7 +1018,7 @@ fn negotiated_payload_limits(layout: &Layout, dtype: u8) -> Result<(u64, u64)> { let delta_bytes = if dtype == DTYPE_Q4 { q4_nbytes(numel)? } else { - tensor_nbytes(dtype, numel)? + tensor_nbytes(push_dtype(dtype), numel)? } as u64; max_init = max_init.max( 4u64.checked_add(init_bytes) @@ -1146,7 +1164,7 @@ fn parse_hello(payload: &[u8], expected_learners: u32) -> Result { bail!("connection generation must be nonzero"); } let dtype = r.u8()?; - if !matches!(dtype, DTYPE_F32 | DTYPE_BF16 | DTYPE_Q4) { + if !matches!(dtype, DTYPE_F32 | DTYPE_BF16 | DTYPE_Q4 | DTYPE_F32_PUSH_BF16) { bail!("unsupported session dtype {dtype}"); } let num_fragments = r.u32()?; @@ -1720,7 +1738,7 @@ async fn dispatch_inner( let expected = if group.dtype == DTYPE_Q4 { q4_nbytes(numel)? } else { - tensor_nbytes(group.dtype, numel)? + tensor_nbytes(push_dtype(group.dtype), numel)? }; if bytes.len() != expected { bail!( @@ -1732,7 +1750,7 @@ async fn dispatch_inner( if group.dtype == DTYPE_Q4 { decode_q4(bytes, numel, &mut outer_gradient)?; } else { - decode_tensor(group.dtype, bytes, &mut outer_gradient)?; + decode_tensor(push_dtype(group.dtype), bytes, &mut outer_gradient)?; } // Every wire dtype carries learner_delta = local - raw_anchor. // Merge math consumes the corresponding outer gradient with the @@ -2102,7 +2120,7 @@ async fn scheduler( // Persist a fresh version zero before its first BCAST. A resumed cut is // already committed and only needs to be rebroadcast. - if checkpoint_each_round { + if checkpoint_each_round || cfg.policy_sweep_fragments.is_some() { if let Some(path) = cfg.checkpoint_path.as_ref() { if !cfg.resume || !path.exists() { st.save_checkpoint(path)?; @@ -2175,7 +2193,23 @@ async fn scheduler( .await?; } let committed = apply_round_commit(&cfg, &mut st, completed)?; - publishing = Some(start_async_checkpoint(&cfg, &st, committed)?); + if persist_before_broadcast(&cfg, committed.t) { + publishing = Some(start_async_checkpoint(&cfg, &st, committed)?); + } else { + publish_round( + &cfg, + &st, + ®istry, + &mut last_sync_secs, + committed, + None, + &budget_cutoff, + ) + .await?; + if !busy_fragments.remove(&p) { + bail!("fragment {p} lost its busy ownership at commit"); + } + } } else { commit_round( &cfg, @@ -3325,7 +3359,7 @@ async fn commit_round( ) -> Result<()> { let committed = apply_round_commit(cfg, st, completed)?; let mut checkpoint_ms = None; - if checkpoint_before_broadcast(cfg) { + if persist_before_broadcast(cfg, committed.t) { if let Some(path) = cfg.checkpoint_path.as_ref() { let checkpoint_start = Instant::now(); st.save_checkpoint(path)?; @@ -3510,7 +3544,8 @@ async fn publish_round( // have also been broadcast; the terminal cut is delivered by FINAL. // A crash-resume loses those gathers; their fragments simply merge on // a later cycle, which the quorum design already tolerates. - if !checkpoint_before_broadcast(cfg) { + // Policy sweeps persist only before broadcast (see persist_before_broadcast). + if cfg.policy_sweep_fragments.is_none() && !checkpoint_before_broadcast(cfg) { if let Some(path) = &cfg.checkpoint_path { if cfg.checkpoint_every > 0 && t % cfg.checkpoint_every == 0 { st.save_checkpoint(path)?; @@ -3652,14 +3687,18 @@ async fn complete_round_async( let cutoff = BudgetCutoff::new(false, Duration::ZERO); let mut busy = HashSet::from([completed.round.p]); let committed = apply_round_commit(cfg, st, completed)?; + if !persist_before_broadcast(cfg, committed.t) { + return publish_round(cfg, st, registry, last_sync_secs, committed, None, &cutoff).await; + } let pending = start_async_checkpoint(cfg, st, committed)?; settle_pending_publish(cfg, st, registry, last_sync_secs, pending, &cutoff, &mut busy).await } fn new_state_for(group: &Arc, cfg: &Config) -> Result { if let Some(fragments) = cfg.policy_sweep_fragments { - if group.dtype != DTYPE_F32 { - bail!("--policy-sweep-fragments requires an FP32 HELLO session"); + // Masters and every BCAST/FINAL stay f32; only push deltas may be bf16. + if !matches!(group.dtype, DTYPE_F32 | DTYPE_F32_PUSH_BF16) { + bail!("--policy-sweep-fragments requires an FP32 HELLO session (optionally with bf16 pushes)"); } if group.layout.fragments.len() != fragments as usize { bail!( @@ -4978,13 +5017,21 @@ mod tests { for mutate in [ |config: &mut Config| config.checkpoint_path = None, - |config: &mut Config| config.checkpoint_every = 2, + |config: &mut Config| config.checkpoint_every = 0, + |config: &mut Config| config.checkpoint_every = 3, ] { invalid = sweep_test_config(2, 4); mutate(&mut invalid); assert!(format!("{:#}", validate_config(&invalid).unwrap_err()) .contains("requires --checkpoint-path")); } + // Sweep-aligned checkpoint intervals are accepted. + let mut interval = sweep_test_config(2, 4); + interval.checkpoint_every = 4; + assert!(validate_config(&interval).is_ok()); + assert!(persist_before_broadcast(&interval, 4)); + assert!(!persist_before_broadcast(&interval, 2)); + assert!(persist_before_broadcast(&sweep_test_config(2, 4), 1)); } #[test] @@ -5581,6 +5628,83 @@ mod tests { directory } + #[test] + fn bf16_push_session_rounds_only_the_delta_within_bf16_error() { + // dtype 4 decodes pushes as bf16 and keeps every bulk payload f32. + assert_eq!(bulk_dtype(DTYPE_F32_PUSH_BF16), DTYPE_F32); + let delta: Vec = (0..4096) + .map(|i| ((i as f32) * 0.37).sin() * 10f32.powi((i % 9) - 4)) + .collect(); + let mut wire = Vec::new(); + encode_tensor(push_dtype(DTYPE_F32_PUSH_BF16), &delta, &mut wire).unwrap(); + assert_eq!(wire.len(), delta.len() * 2); + let mut decoded = Vec::new(); + decode_tensor(push_dtype(DTYPE_F32_PUSH_BF16), &wire, &mut decoded).unwrap(); + for (exact, rounded) in delta.iter().zip(&decoded) { + // bf16 keeps 8 significant bits: round-to-nearest error <= 2^-8 relative. + assert!((exact - rounded).abs() <= exact.abs() * (1.0 / 256.0)); + } + // A merge of bf16 deltas into f32 masters stays within the same bound + // per learner, so a two-learner average is off by at most that much. + let master = 1.0f32; + for (exact, rounded) in delta.iter().zip(&decoded) { + let merged_exact = master + exact; + let merged = master + rounded; + assert!((merged_exact - merged).abs() <= exact.abs() / 256.0 + f32::EPSILON * 2.0); + } + } + + #[tokio::test] + async fn sweep_checkpoint_interval_persists_only_sweep_aligned_steps() { + let directory = unique_test_directory("sweep-checkpoint-interval"); + let learner = member(0, 30); + for asynchronous in [false, true] { + let checkpoint = directory.join(format!("state-{asynchronous}.ckpt")); + let (group, mut frames) = streaming_test_group(DTYPE_F32, false); + let registry = registry_with_group(group); + let mut config = sweep_test_config(2, 8); + config.checkpoint_path = Some(checkpoint.clone()); + config.checkpoint_every = 4; + config.async_checkpoint = asynchronous; + let mut state = two_fragment_sweep_state(); + state.save_checkpoint(&checkpoint).unwrap(); // version zero + let mut sync_seconds = 0.0; + for step in 1..=6u64 { + let fragment = ((step - 1) % 2) as usize; + let mut round = sweep_test_round(step, fragment, step.div_ceil(2), 1, learner); + let base = state.versions[fragment]; + round.base_version = base; + for push in round.pushes.values_mut() { + push.base_version = base; + } + if asynchronous { + complete_round_async(&config, &mut state, ®istry, &mut sync_seconds, round) + .await + .unwrap(); + } else { + complete_round(&config, &mut state, ®istry, &mut sync_seconds, round) + .await + .unwrap(); + } + assert_eq!( + chunked_inner_type(&frames.try_recv().unwrap()), + MSG_BCAST_FRAGMENT + ); + let mut on_disk = two_fragment_sweep_state(); + on_disk.load_checkpoint(&checkpoint).unwrap(); + // Only step 4 (a sweep boundary and a multiple of 4) is durable. + let expected = if step >= 4 { 4 } else { 0 }; + assert_eq!(on_disk.global_step, expected, "step {step}"); + } + // The newest checkpoint resumes exactly at its sweep boundary. + let mut resumed = two_fragment_sweep_state(); + resumed.load_checkpoint(&checkpoint).unwrap(); + validate_resumed_policy_sweep(&config, &resumed).unwrap(); + assert_eq!(resumed.versions, vec![3, 4]); + } + std::fs::remove_dir_all(&directory).ok(); + } + #[tokio::test] async fn async_checkpoint_matches_the_synchronous_path_bit_for_bit() { let directory = unique_test_directory("async-checkpoint-equal"); diff --git a/tests/test_e3_sync_options.py b/tests/test_e3_sync_options.py new file mode 100644 index 00000000..0e8492ef --- /dev/null +++ b/tests/test_e3_sync_options.py @@ -0,0 +1,63 @@ +"""S26 E3 sync options: bf16 push transcoding, session binding, preset scoping. CPU only.""" + +from __future__ import annotations + +from dataclasses import replace + +import pytest +import torch + +from test_rl_algorithm_provenance import _fake_sky, _launcher_args, _no_modal_listing # noqa: F401 +from test_rl_full_parameter_sync import _role_stream_config +from yeto.protocol import ( + DTYPE_BF16, + DTYPE_F32, + DTYPE_F32_PUSH_BF16, + _f32_parts_to_bf16, + bulk_dtype, + f32_bytes_to_bf16, + push_wire_dtype, +) + + +def test_push_bf16_session_keeps_bulk_payloads_f32(): + assert bulk_dtype(DTYPE_F32_PUSH_BF16) == DTYPE_F32 + assert push_wire_dtype(DTYPE_F32_PUSH_BF16) == DTYPE_BF16 + assert push_wire_dtype(DTYPE_F32) == DTYPE_F32 + + +def test_f32_to_bf16_matches_torch_round_to_nearest_even_across_split_parts(): + values = torch.randn(50_001) * torch.logspace(-20, 20, 50_001) + values[:3] = torch.tensor([0.0, -0.0, float("inf")]) + raw = values.numpy().tobytes() + want = values.to(torch.bfloat16).view(torch.int16).numpy().astype(" FragmentLayout: return FragmentLayout([Fragment(MERGE_RDA, [(f"actor.w{i}", NUMEL)]) for i in range(FRAGMENTS)]) -def _client(port: int, profile_hash: str) -> SyncerClient: - return SyncerClient(("127.0.0.1", port), 0, _layout(), dtype=DTYPE_F32, num_streams=1, +def _client(port: int, profile_hash: str, dtype: int = DTYPE_F32) -> SyncerClient: + return SyncerClient(("127.0.0.1", port), 0, _layout(), dtype=dtype, num_streams=1, connect_timeout=10, max_reconnects=0, finalization_timeout=20, syncer_profile_hash=bytes.fromhex(profile_hash)) @@ -160,10 +160,11 @@ def test_launcher_dense_syncer_refuses_another_profile_hash(syncer_binary, tmp_p STREAMED = ("--rl-sync-fragment-pipeline", str(FRAGMENTS), "--rl-sync-async-checkpoint") -def _run_sweeps(binary: Path, tmp_path: Path, extra: tuple[str, ...], sweeps: int): +def _run_sweeps(binary: Path, tmp_path: Path, extra: tuple[str, ...], sweeps: int, + dtype: int = DTYPE_F32): """Drive ``sweeps`` strict sweeps; return (max PULLs seen at once, finals, checkpoint).""" process, port, profile_hash = _start_syncer(binary, tmp_path, *extra, total_steps=sweeps) - client = _client(port, profile_hash) + client = _client(port, profile_hash, dtype) concurrent = 0 try: client.start() @@ -232,3 +233,40 @@ def test_streamed_flags_bind_the_profile_and_are_full_parameter_only(): != launcher.rl_syncer_profile(streamed, 1).sha256) with pytest.raises(ValueError, match="--rl-sync-fragment-pipeline between 1 and --fragments"): _dense_args("--rl-sync-fragment-pipeline", str(FRAGMENTS + 1)) + + +def test_bf16_pushes_stay_within_bf16_error_of_the_f32_sweep(syncer_binary, tmp_path): + """E3: bf16 outer deltas change results only by bf16 rounding of each delta.""" + exact_dir, bf16_dir = tmp_path / "f32", tmp_path / "bf16" + exact_dir.mkdir() + bf16_dir.mkdir() + exact = _run_sweeps(syncer_binary, exact_dir, STREAMED, sweeps=3) + rounded = _run_sweeps(syncer_binary, bf16_dir, (*STREAMED, "--rl-sync-push-dtype", "bf16"), + sweeps=3, dtype=DTYPE_F32_PUSH_BF16) + assert exact[1] == rounded[1] == [5, 6] # versions are unchanged + for fid, data in exact[2].items(): + want = torch.frombuffer(bytearray(data), dtype=torch.float32)[-NUMEL:] + got = torch.frombuffer(bytearray(rounded[2][fid]), dtype=torch.float32)[-NUMEL:] + # Three sweeps, each delta < 2 in magnitude, bf16 error <= 2^-8 relative. + assert torch.allclose(got, want, rtol=0, atol=3 * 2.0 / 256), fid + assert not torch.equal(got, want), "bf16 pushes must actually round" + + +def test_sweep_checkpoint_interval_reaches_the_syncer_and_its_profile(): + every = 2 * FRAGMENTS + args = _dense_args("--rl-sync-checkpoint-every", str(every), total_steps=4) + from yeto import launcher + + command = launcher.role_syncer_commands(args, 1, binary="yeto-syncer")["actor"] + assert f"--checkpoint-every {every}" in command + assert launcher.rl_syncer_profile(args, 1).checkpoint_every == every + with pytest.raises(ValueError, match="--rl-sync-checkpoint-every 1 or a multiple"): + _dense_args("--rl-sync-checkpoint-every", str(FRAGMENTS + 1)) + + +def test_bf16_push_flag_reaches_the_island_command(): + from yeto.launch.rl_island_task import _rl_trainable_island_flags + + args = _dense_args("--rl-sync-push-dtype", "bf16") + assert "--sweep-push-dtype bf16" in _rl_trainable_island_flags(args, 0, 1) + assert "--sweep-push-dtype" not in _rl_trainable_island_flags(_dense_args(), 0, 1) diff --git a/tools/bench_sync_boundary.py b/tools/bench_sync_boundary.py index 3b449632..a55fb906 100644 --- a/tools/bench_sync_boundary.py +++ b/tools/bench_sync_boundary.py @@ -3,8 +3,9 @@ Runs the real Rust syncer on localhost with two in-process learners and a synthetic f32 model, and times each strict sweep from the first push to the last updated fragment received by every learner. It compares the serial -profile (pipeline 1, checkpoint on the scheduler) with the streamed profile -(pipeline P, ``--async-checkpoint``). +profile (pipeline 1, checkpoint on the scheduler), the streamed profile +(pipeline P, ``--async-checkpoint``), and streamed plus bf16 pushes and a +checkpoint every second sweep. These are local CPU numbers: loopback transfer and a local NVMe checkpoint. They show where the boundary time goes on this machine; they do not predict @@ -28,7 +29,7 @@ import numpy as np from yeto.fragments import MERGE_RDA, Fragment, FragmentLayout -from yeto.protocol import DTYPE_F32, SyncerClient +from yeto.protocol import DTYPE_F32, DTYPE_F32_PUSH_BF16, SyncerClient ROOT = Path(__file__).resolve().parent.parent @@ -40,7 +41,7 @@ def _free_port() -> int: def _syncer(binary: Path, workdir: Path, *, port: int, learners: int, fragments: int, - sweeps: int, streamed: bool) -> subprocess.Popen: + sweeps: int, streamed: bool, checkpoint_every: int) -> subprocess.Popen: command = [ str(binary), "--port", str(port), "--learners", str(learners), "--quorum", str(learners), "--grace-ms", "0", @@ -49,7 +50,7 @@ def _syncer(binary: Path, workdir: Path, *, port: int, learners: int, fragments: "--total-steps", str(sweeps * fragments), "--policy-sweep-fragments", str(fragments), "--outer-lr", "1", "--outer-momentum", "0", "--max-base-lag", "0", "--learner-weight", "equal", "--checkpoint-path", str(workdir / "state.ckpt"), - "--checkpoint-every", "1", "--event-tape", str(workdir / "tape.jsonl"), + "--checkpoint-every", str(checkpoint_every), "--event-tape", str(workdir / "tape.jsonl"), "--allow-unauthenticated-islands", # loopback-only benchmark ] if streamed: @@ -59,9 +60,9 @@ def _syncer(binary: Path, workdir: Path, *, port: int, learners: int, fragments: def _learner(learner_id: int, port: int, layout: FragmentLayout, numel: int, sweeps: int, - timings: list[list[float]], errors: list[BaseException]) -> None: + timings: list[list[float]], errors: list[BaseException], dtype: int) -> None: fragments = len(layout.fragments) - client = SyncerClient(("127.0.0.1", port), learner_id, layout, dtype=DTYPE_F32, + client = SyncerClient(("127.0.0.1", port), learner_id, layout, dtype=dtype, num_streams=1, connect_timeout=30, max_reconnects=0, finalization_timeout=600) rng = np.random.default_rng(learner_id) @@ -107,18 +108,29 @@ def _learner(learner_id: int, port: int, layout: FragmentLayout, numel: int, swe client.close() -def run(binary: Path, *, streamed: bool, learners: int, fragments: int, numel: int, +PROFILES = { + "serial": {"streamed": False, "dtype": DTYPE_F32, "sweeps_per_checkpoint": 0}, + "streamed": {"streamed": True, "dtype": DTYPE_F32, "sweeps_per_checkpoint": 0}, + "streamed+bf16+ckpt/2sweeps": {"streamed": True, "dtype": DTYPE_F32_PUSH_BF16, + "sweeps_per_checkpoint": 2}, +} + + +def run(binary: Path, *, profile: str, learners: int, fragments: int, numel: int, sweeps: int) -> dict: + options = PROFILES[profile] + streamed, dtype = options["streamed"], options["dtype"] + every = options["sweeps_per_checkpoint"] * fragments or 1 layout = FragmentLayout([Fragment(MERGE_RDA, [(f"w{i}", numel)]) for i in range(fragments)]) with tempfile.TemporaryDirectory(prefix="e3-bench-", dir="/tmp") as tmp: workdir = Path(tmp) port = _free_port() process = _syncer(binary, workdir, port=port, learners=learners, fragments=fragments, - sweeps=sweeps, streamed=streamed) + sweeps=sweeps, streamed=streamed, checkpoint_every=every) timings: list[list[float]] = [[] for _ in range(learners)] errors: list[BaseException] = [] threads = [threading.Thread(target=_learner, args=(i, port, layout, numel, sweeps, - timings, errors)) + timings, errors, dtype)) for i in range(learners)] for thread in threads: thread.start() @@ -134,10 +146,10 @@ def run(binary: Path, *, streamed: bool, learners: int, fragments: int, numel: i steady = [max(per[i] for per in timings) for i in range(1, sweeps - 1)] or \ [max(per[i] for per in timings) for i in range(sweeps)] return { - "profile": "streamed" if streamed else "serial", + "profile": profile, "boundary_s_per_sweep": [round(x, 3) for x in steady], "boundary_s_median": round(statistics.median(steady), 3), - "checkpoint_ms_median": statistics.median(r["checkpoint_ms"] for r in tape), + "checkpoints_written": sum("checkpoint_ms" in r for r in tape), "sync_ms_median": statistics.median(r["sync_ms"] for r in tape), } @@ -150,9 +162,9 @@ def main() -> None: parser.add_argument("--floats-per-fragment", type=int, default=32 * 1024 * 1024) parser.add_argument("--sweeps", type=int, default=4) args = parser.parse_args() - results = [run(args.binary, streamed=streamed, learners=args.learners, + results = [run(args.binary, profile=profile, learners=args.learners, fragments=args.fragments, numel=args.floats_per_fragment, sweeps=args.sweeps) - for streamed in (False, True)] + for profile in PROFILES] print(json.dumps({"note": "CPU localhost numbers; not GPU or WAN timings", "learners": args.learners, "fragments": args.fragments, "floats_per_fragment": args.floats_per_fragment, "results": results}, diff --git a/yeto/launch/rl_args.py b/yeto/launch/rl_args.py index e1ce4ff2..df0c0efb 100644 --- a/yeto/launch/rl_args.py +++ b/yeto/launch/rl_args.py @@ -148,6 +148,21 @@ def sao_contract_mounts(args, learner_id: int) -> dict[str, str]: for shipped in (island.runtime, island.layout_attestation)} +_FULL_PARAMETER_SYNC_DEFAULTS = ( + ("rl_sync_fragment_pipeline", 1, "--rl-sync-fragment-pipeline"), + ("rl_sync_async_checkpoint", False, "--rl-sync-async-checkpoint"), + ("rl_sync_push_dtype", "f32", "--rl-sync-push-dtype"), + ("rl_sync_checkpoint_every", 1, "--rl-sync-checkpoint-every"), +) + + +def _refuse_full_parameter_sync_options(args) -> None: + """The S26 E3 sync options exist only for the full-parameter preset.""" + for dest, default, option in _FULL_PARAMETER_SYNC_DEFAULTS: + if getattr(args, dest, default) != default: + raise ValueError(f"{option} only applies to --rl-sync-preset {FULL_PARAMETER_SYNC_PRESET}") + + def _prepare_dense_full_sweep(args, *, num_islands: int) -> None: """full-parameter preset (P0-D8): one actor role swept as whole-model fragments. @@ -172,6 +187,9 @@ def _prepare_dense_full_sweep(args, *, num_islands: int) -> None: fragment_pipeline = int(getattr(args, "rl_sync_fragment_pipeline", 1)) if not 1 <= fragment_pipeline <= max(args.fragments, 1): problems.append("--rl-sync-fragment-pipeline between 1 and --fragments") + checkpoint_every = int(getattr(args, "rl_sync_checkpoint_every", 1)) + if checkpoint_every != 1 and (checkpoint_every < 1 or checkpoint_every % max(args.fragments, 1)): + problems.append("--rl-sync-checkpoint-every 1 or a multiple of --fragments") if problems: raise ValueError(f"--rl-sync-preset {FULL_PARAMETER_SYNC_PRESET} requires " + "; ".join(problems)) args.quorum = num_islands @@ -728,6 +746,8 @@ def _prepare_rl_args( f"--rl-sync-preset {FULL_PARAMETER_SYNC_PRESET}") if args.rl_sync_preset != SAO_STREAMING_SYNC_PRESET and getattr(args, "rl_sao_contracts", None): raise ValueError("--rl-sao-contracts only applies to --rl-sync-preset sao-streaming-full") + if args.rl_sync_preset != FULL_PARAMETER_SYNC_PRESET: + _refuse_full_parameter_sync_options(args) if args.rl_sync_preset == "decoupled": if args.experimental_rl_sync: raise ValueError("decoupled RL does not accept --experimental-rl-sync") diff --git a/yeto/launch/rl_island_task.py b/yeto/launch/rl_island_task.py index 1185f4c9..a5cd883f 100644 --- a/yeto/launch/rl_island_task.py +++ b/yeto/launch/rl_island_task.py @@ -72,6 +72,8 @@ def _rl_trainable_island_flags(args, learner_id: int, num_learners: int) -> str: sweep = (f" --parameter-layout-sha256 {args.rl_parameter_layout_sha256}" f" --num-learners {num_learners}" f" --syncer-profile-sha256 {rl_syncer_profile(args, num_learners).sha256}") + if getattr(args, "rl_sync_push_dtype", "f32") != "f32": + sweep += f" --sweep-push-dtype {args.rl_sync_push_dtype}" if getattr(args, "expert_full_count", 0): return lora + sweep return " --parameter-mode full" + sweep diff --git a/yeto/launch/syncer.py b/yeto/launch/syncer.py index 7d951ae9..91587c07 100644 --- a/yeto/launch/syncer.py +++ b/yeto/launch/syncer.py @@ -145,6 +145,11 @@ def _syncer_quorum_timeout(args) -> str: return "" if value is None else f" --quorum-timeout-s {value}" +def _syncer_checkpoint_every(args) -> int: + """Per-round checkpoints (1) unless the full-parameter preset sets an interval.""" + return int(getattr(args, "rl_sync_checkpoint_every", 1)) if rl_dense_sweep(args) else 1 + + def rl_dense_sweep(args) -> bool: return (getattr(args, "training_mode", "sft") == "rl" and getattr(args, "rl_sync_preset", None) == FULL_PARAMETER_SYNC_PRESET) @@ -206,7 +211,7 @@ def rl_syncer_profile(args, num_learners: int): "outer_lr": args.outer_lr, "outer_momentum": args.outer_momentum, "checkpoint_enabled": True, - "checkpoint_every": 1, + "checkpoint_every": _syncer_checkpoint_every(args), "resume": False, # not hashed (v2): fresh start and restart agree "mark_final_checkpoint": mark_final, "learner_budget_steps": None, @@ -447,7 +452,7 @@ def _rl_role_syncer_command(args, num_learners: int, binary: str, f" --outer-momentum {args.outer_momentum}" " --max-base-lag 0 --learner-weight equal" f" --checkpoint-path {syncer.checkpoint}" - " --checkpoint-every 1" + f" --checkpoint-every {_syncer_checkpoint_every(args)}" f" {_resume_if_exists(syncer.checkpoint)}" f" --event-tape {syncer.event_tape}" ) diff --git a/yeto/protocol.py b/yeto/protocol.py index 7beb2254..28331475 100644 --- a/yeto/protocol.py +++ b/yeto/protocol.py @@ -34,7 +34,7 @@ import struct import threading import time -from collections.abc import Callable, Iterable +from collections.abc import Callable, Iterable, Iterator from dataclasses import dataclass, field from .fragments import MERGE_ISO, FragmentLayout @@ -66,6 +66,11 @@ # base-relative learner deltas; INIT_PARAMS and BCAST_FRAGMENT stay bf16 # (see bulk_dtype and docs/PROTOCOL.md v4). DTYPE_Q4 = 3 +# Session dtype 4 (S26 E3): PUSH_FRAGMENT deltas travel as bf16; INIT_PARAMS, +# BCAST_FRAGMENT and FINAL stay exact f32. Callers still hand the client f32 +# delta bytes; the client rounds them (round-to-nearest-even, as the syncer's +# ``half::bf16::from_f32``) on the way out. +DTYPE_F32_PUSH_BF16 = 4 CHUNK_SIZE = 4 * 1024 * 1024 MAX_HANDSHAKE_FRAME = 16 * 1024 * 1024 @@ -93,7 +98,40 @@ def bulk_dtype(dtype: int) -> int: Q4 applies only to push deltas (small dynamic range); full parameter payloads would not survive 4 bits, so they travel as bf16. """ - return DTYPE_BF16 if dtype == DTYPE_Q4 else dtype + if dtype == DTYPE_Q4: + return DTYPE_BF16 + if dtype == DTYPE_F32_PUSH_BF16: + return DTYPE_F32 + return dtype + + +def push_wire_dtype(dtype: int) -> int: + """Dtype of unquantized PUSH_FRAGMENT deltas on the wire for a session dtype.""" + return DTYPE_BF16 if dtype == DTYPE_F32_PUSH_BF16 else dtype + + +def f32_bytes_to_bf16(data: bytes | bytearray | memoryview) -> bytes: + """Round little-endian f32 values to bf16 (nearest, ties to even).""" + import numpy as np # noqa: PLC0415 - keep the protocol import light + + bits = np.frombuffer(data, dtype="> 16) & 1)) >> 16 + nan = (bits & 0x7FFF_FFFF) > 0x7F80_0000 + rounded[nan] = (bits[nan] >> 16) | 0x40 # keep NaN quiet, never round to inf + return rounded.astype(" Iterator[bytes]: + """Stream ``f32_bytes_to_bf16`` over parts whose boundaries may split a value.""" + carry = b"" + for part in parts: + data = carry + bytes(part) + whole = len(data) - len(data) % 4 + carry = data[whole:] + if whole: + yield f32_bytes_to_bf16(data[:whole]) + if carry: + raise ValueError("f32 push delta ends mid-value") def layout_fingerprint(layout: FragmentLayout) -> bytes: @@ -480,7 +518,7 @@ def __init__( ): if not 0 <= learner_id <= 0xFFFF_FFFF: raise ValueError(f"learner_id must fit u32, got {learner_id}") - if dtype not in (DTYPE_F32, DTYPE_BF16, DTYPE_Q4): + if dtype not in (DTYPE_F32, DTYPE_BF16, DTYPE_Q4, DTYPE_F32_PUSH_BF16): raise ValueError(f"unsupported session dtype {dtype}") if layout.num_fragments < 1: raise ValueError("layout must contain at least one fragment") @@ -833,7 +871,8 @@ def push_fragment( expected = ( _q4_nbytes(numel) if self.dtype == DTYPE_Q4 - else _tensor_nbytes(self.dtype, numel) + # dtype 4 callers pass f32 bytes; the client rounds them below + else _tensor_nbytes(DTYPE_F32 if self.dtype == DTYPE_F32_PUSH_BF16 else self.dtype, numel) ) if len(tensor_bytes) != expected: raise ValueError( @@ -851,6 +890,8 @@ def push_fragment( c_steps, c_tokens, ) + if self.dtype == DTYPE_F32_PUSH_BF16: + tensor_bytes = f32_bytes_to_bf16(tensor_bytes) self._send_large(MSG_PUSH_FRAGMENT, head + tensor_bytes) def push_fragment_parts( @@ -885,8 +926,10 @@ def push_fragment_parts( expected = ( _q4_nbytes(numel) if self.dtype == DTYPE_Q4 - else _tensor_nbytes(self.dtype, numel) + else _tensor_nbytes(push_wire_dtype(self.dtype), numel) ) + if self.dtype == DTYPE_F32_PUSH_BF16: + tensor_parts = _f32_parts_to_bf16(tensor_parts) head = struct.pack( " None: if not self.algorithm.isidentifier(): @@ -219,7 +222,7 @@ def role_stream_session_contract_hash( "pipeline_depth": config.pipeline_depth, "local_horizon": config.local_horizon, "optimizer_steps_per_round": config.optimizer_steps_per_round, - "wire_dtype": "fp32", + "wire_dtype": "fp32-push-bf16" if config.push_bf16 else "fp32", }, } encoded = json.dumps(payload, sort_keys=True, separators=(",", ":")).encode() @@ -334,7 +337,7 @@ async def initialize(self, group) -> None: self.config.syncer_addr, self.config.learner_id, adapter.layout.fragments, - dtype=DTYPE_F32, + dtype=DTYPE_F32_PUSH_BF16 if self.config.push_bf16 else DTYPE_F32, num_streams=self.config.wan_streams, max_reconnects=None, session_contract_hash=role_stream_session_contract_hash( diff --git a/yeto/rl/adapters/miles/island_entry.py b/yeto/rl/adapters/miles/island_entry.py index 48585224..5a9d8874 100644 --- a/yeto/rl/adapters/miles/island_entry.py +++ b/yeto/rl/adapters/miles/island_entry.py @@ -838,6 +838,7 @@ def _dense_full_sync_config( pipeline_depth=1, algorithm_spec_sha256=algorithm_spec_sha256, wan_streams=args.wan_streams, + push_bf16=getattr(args, "sweep_push_dtype", "f32") == "bf16", ) return FullParameterSyncConfig((actor,), sweep=True) diff --git a/yeto/rl/adapters/miles/run_flags.py b/yeto/rl/adapters/miles/run_flags.py index 8a9f5202..adbf93ca 100644 --- a/yeto/rl/adapters/miles/run_flags.py +++ b/yeto/rl/adapters/miles/run_flags.py @@ -313,6 +313,12 @@ ), island_flag("--fragments", type=int, default=1), island_flag("--parameter-layout-sha256", default=None), + island_flag( + "--sweep-push-dtype", + choices=("f32", "bf16"), + default="f32", + help="full-parameter sweep: wire dtype of pushed outer deltas (syncer masters stay f32)", + ), island_flag( "--syncer-profile-sha256", default=None, diff --git a/yeto/rl/engine/run_flags.py b/yeto/rl/engine/run_flags.py index f9be06f6..6e6b95c1 100644 --- a/yeto/rl/engine/run_flags.py +++ b/yeto/rl/engine/run_flags.py @@ -849,6 +849,22 @@ def rl_forward_pairs(*dests: str | tuple[str, str]) -> tuple[tuple[str, str], .. "(1..--fragments). Above 1, fragment k+1 uploads while fragment k merges " "and downloads; merge results are unchanged (default 1 = serial)", ), + cli_flag( + "--rl-sync-push-dtype", + choices=("f32", "bf16"), + default="f32", + help="full-parameter only: wire dtype of the islands' pushed outer deltas " + "(pseudo-gradients). bf16 halves upload bytes; the syncer keeps f32 masters, " + "momentum and f32 BCAST/FINAL (default f32 = exact)", + ), + cli_flag( + "--rl-sync-checkpoint-every", + type=int, + default=1, + help="full-parameter only: syncer checkpoint interval in fragment steps " + "(1, or a multiple of --fragments). Above 1, only checkpointed steps are " + "persisted before their broadcast; a crash resumes at the newest checkpoint", + ), cli_flag( "--rl-sync-async-checkpoint", action="store_true", diff --git a/yeto/syncer_profile.py b/yeto/syncer_profile.py index c5e011b0..ce23fa99 100644 --- a/yeto/syncer_profile.py +++ b/yeto/syncer_profile.py @@ -261,6 +261,10 @@ def _validate_general(self) -> None: fragments = self.policy_sweep_fragments if not 1 <= self.pipeline <= fragments or self.total_steps % fragments: raise ValueError("syncer profile policy sweep schedule is invalid") + if self.checkpoint_every != 1 and ( + self.checkpoint_every < 1 or self.checkpoint_every % fragments): + raise ValueError("syncer profile policy sweep checkpoint interval must be " + "1 or a multiple of the fragment count") def validate_sao(self) -> None: """Enforce the paired actor/critic fail-closed server profile.""" 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 05/62] 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 1a783d09529d711b01cba7478cce188312141239 Mon Sep 17 00:00:00 2001 From: michael <78113554+chowhsu@users.noreply.github.com> Date: Tue, 20 Oct 2026 15:48:24 +0000 Subject: [PATCH 06/62] syncer: int8 sweep pushes, bf16 downlink deltas, sync/train overlap Add session dtype 5 (f32 session, block-scaled int8 PUSH deltas): u32 block size, then per block an f32 absmax/127 scale and one i8 per value; error is at most absmax/254 per element. Launcher --rl-sync-push-int8-block (default 128) is bound into the role-stream contract. Add session flag 0x80 (bf16 downlink deltas) for f32-bulk sessions. Syncer and learners hold a bit-identical f32 anchor per fragment; BCAST is FULL f32 or bf16(theta - anchor), with a FULL refresh every N versions (--bcast-delta-full-every), FULL re-sends on HELLO/reconnect, and a SHA-256 check of the anchor after every update. A mismatch drops the connection and the reconnect restores the anchor. SyncerClient decodes before adapters see the payload: the Miles streaming sink is wrapped, the strict bridge (verl and Miles LoRA) gets f32 bytes. Masters, checkpoints and FINAL stay exact. Add --rl-sync-overlap for the full-parameter sweep: an island pushes round k, trains round k+1 at once, and at the next boundary rebases in place, theta <- G_k + (local - base), with G_k as the anchor so no delta is counted twice. The launcher raises --rl-max-policy-age to at least 1 before its checks; the island refuses overlap below that age. Adds Rust and Python codec/anchor tests, real-syncer int8/downlink/resync tests, a strict-bridge downlink test, and coordinator overlap tests. Co-Authored-By: Claude Opus 5.5 --- docs/PROTOCOL.md | 26 +- syncer/src/downlink.rs | 289 +++++++++++++++ syncer/src/main.rs | 7 + syncer/src/protocol.rs | 197 ++++++++++- syncer/src/server.rs | 113 ++++-- syncer/src/state.rs | 3 + tests/test_e3_sync_options.py | 261 +++++++++++++- tests/test_e3_sync_overlap.py | 331 ++++++++++++++++++ tests/test_rl_dense_sweep_real_syncer.py | 176 +++++++++- tests/test_run_flags.py | 2 +- yeto/downlink.py | 235 +++++++++++++ yeto/launch/rl_args.py | 36 +- yeto/launch/rl_island_task.py | 5 + yeto/launch/syncer.py | 7 + yeto/protocol.py | 170 +++++++-- .../miles/full_parameter/role_stream.py | 140 +++++++- yeto/rl/adapters/miles/island_entry.py | 12 +- .../miles/models/chunked_full_parameter.py | 54 +++ yeto/rl/adapters/miles/run_flags.py | 20 +- yeto/rl/adapters/verl/island_entry.py | 3 + yeto/rl/adapters/verl/trainer.py | 3 +- yeto/rl/bridge.py | 18 +- yeto/rl/engine/full_parameter_sync.py | 75 +++- yeto/rl/engine/run_flags.py | 32 +- yeto/rl/engine/sync_overlap.py | 41 +++ 25 files changed, 2144 insertions(+), 112 deletions(-) create mode 100644 syncer/src/downlink.rs create mode 100644 tests/test_e3_sync_overlap.py create mode 100644 yeto/downlink.py create mode 100644 yeto/rl/engine/sync_overlap.py diff --git a/docs/PROTOCOL.md b/docs/PROTOCOL.md index 3d08e098..db8e47e9 100644 --- a/docs/PROTOCOL.md +++ b/docs/PROTOCOL.md @@ -32,7 +32,7 @@ or allocation overflows are errors rather than panics. | type | name | direction | payload | |---:|---|---|---| -| 1 | HELLO | learner → syncer | protocol_version:u16 (=4), learner_id:u32, connection_generation:u64, dtype:u8 (1=f32, 2=bf16, 3=q4, 4=f32 with bf16 pushes), num_fragments:u32, per-fragment layout, layout_fingerprint:[u8;32], session_contract_hash:[u8;32], optional syncer_profile_hash:[u8;32], num_streams:u16, optional compat_group trailer (`b"YCG1"`, u32 length, 1–64 bytes utf-8) | +| 1 | HELLO | learner → syncer | protocol_version:u16 (=4), learner_id:u32, connection_generation:u64, dtype:u8 (1=f32, 2=bf16, 3=q4, 4=f32 with bf16 pushes, 5=f32 with int8 pushes; flag 0x80 = bf16 downlink deltas), num_fragments:u32, per-fragment layout, layout_fingerprint:[u8;32], session_contract_hash:[u8;32], optional syncer_profile_hash:[u8;32], num_streams:u16, optional compat_group trailer (`b"YCG1"`, u32 length, 1–64 bytes utf-8) | | 2 | INIT_PARAMS | learner → syncer | fragment_id:u32, full tensor bytes; only learner 0 may initialize | | 3 | PULL_REQ | syncer → learner | fragment_id:u32, global_step:u64, round_attempt:u32 | | 4 | PUSH_FRAGMENT | learner → syncer | learner_id:u32, fragment_id:u32, global_step:u64, round_attempt:u32, base_version:u64, local_step:u64, c_steps:u32, c_tokens:u64, base-relative learner-delta bytes | @@ -230,6 +230,30 @@ bf16 (round-to-nearest-even). The syncer decodes and negates them like dtype 2 and keeps f32 masters and momentum. The full-parameter policy sweep accepts dtype 1 or 4 (S26 E3, `--rl-sync-push-dtype bf16`). +## int8 pushes over an f32 session (dtype 5) + +Dtype 5 sends PUSH deltas as block-scaled int8: u32 block size B +(32..1048576), then per block of B values (the last block may be shorter) an +f32 scale (block absmax / 127) and one i8 per value (`-128` is invalid). +A value decodes to `q * scale`, so its error is at most absmax / 254. The +learner picks B (launcher `--rl-sync-push-int8-block`, default 128); it is +bound into the role-stream session contract. INIT, BCAST and FINAL stay f32. + +## bf16 downlink deltas (dtype flag 0x80) + +On an f32-bulk session (dtype 1, 4 or 5) the flag 0x80 changes BCAST_FRAGMENT +to `fragment_id:u32, version:u64, kind:u8, base_version:u64, +anchor_sha256:[u8;32]`, then the tensor: kind 0 (FULL) carries f32 values, +kind 1 (DELTA) carries `bf16(theta - anchor)`. Syncer and learner hold the +same f32 anchor per fragment: FULL sets it, DELTA adds the decoded delta in +f32; the digest is checked after every update. Every N-th version of a +fragment is FULL (`--bcast-delta-full-every N`, N >= 2; the syncer refuses a +flagged session without it and the reverse). Initial and reconnect re-sends +are always FULL with the shared anchor. A learner with a missing or wrong +anchor drops its connection, and the reconnect re-send restores it. The +syncer's masters, momentum, checkpoints and FINAL are unchanged and exact; +anchors are not checkpointed, so a resumed syncer starts with FULL. + ## Q4 pushes Q4 uses the same learner-delta semantics. INIT and BCAST stay bf16; PUSH is diff --git a/syncer/src/downlink.rs b/syncer/src/downlink.rs new file mode 100644 index 00000000..e303e119 --- /dev/null +++ b/syncer/src/downlink.rs @@ -0,0 +1,289 @@ +//! bf16 downlink deltas for BCAST_FRAGMENT (S26 E3, session flag +//! `DTYPE_FLAG_BCAST_DELTA`). +//! +//! Every learner of a session holds, per fragment, an f32 *anchor* that is +//! bit-identical to the syncer's copy here. A broadcast of a new version is +//! either +//! +//! * `FULL`: the f32 values; the anchor becomes exactly those values, or +//! * `DELTA`: `d = bf16(θ − anchor)`; both sides set `anchor += d` in f32 +//! (IEEE addition is exact-rounded, so both sides get the same bits). +//! +//! Each delta is taken against the current anchor, so the rounding error does +//! not accumulate: `|θ − anchor| <= |θ − anchor_prev| * 2^-9` per element +//! after every delta. The learner trains from the anchor; its next push is +//! relative to the anchor, and the syncer adds that push to its exact θ, so +//! θ stays the exact sum of all pushes. +//! +//! Drift guard: every `full_every`-th broadcast of a fragment is FULL; a +//! re-sent fragment (HELLO/reconnect recovery) is always FULL; and every +//! header carries the SHA-256 of the f32 anchor after the update. A learner +//! whose anchor version or digest does not match drops its connection; the +//! reconnect recovery then sends FULL. +//! +//! Anchors live only in memory: a resumed syncer starts with none, so its +//! first broadcast of every fragment is FULL. +//! +//! Header after the usual `fragment_id u32, version u64` prefix: +//! `kind u8 (0 FULL, 1 DELTA), base_version u64, anchor_sha256 [32]`. +//! The tensor that follows is f32 (FULL) or bf16 (DELTA). + +use std::sync::Arc; + +use anyhow::{bail, Result}; +use sha2::{Digest, Sha256}; + +use crate::protocol::{DTYPE_BF16, DTYPE_F32}; + +pub const KIND_FULL: u8 = 0; +pub const KIND_DELTA: u8 = 1; +/// fragment_id + version + kind + base_version + digest. +pub const PREFIX_LEN: usize = 4 + 8 + 1 + 8 + 32; + +struct Anchor { + version: u64, + values: Arc>, + digest: [u8; 32], + deltas_since_full: u64, +} + +/// One encoded BCAST: header prefix plus tensor values in `dtype` (DELTA +/// values are already bf16-representable, so encoding them is exact). +pub struct EncodedBcast { + pub prefix: Vec, + pub dtype: u8, + pub values: Arc>, +} + +impl EncodedBcast { + #[cfg(test)] + pub fn kind(&self) -> u8 { + self.prefix[12] + } +} + +pub struct DownlinkAnchors { + full_every: u64, + anchors: Vec>, +} + +pub fn digest_f32(values: &[f32]) -> [u8; 32] { + let mut hasher = Sha256::new(); + // Hash in bounded slices to avoid a second full-size byte buffer. + let mut buf = Vec::with_capacity(4 * 65536); + for chunk in values.chunks(65536) { + buf.clear(); + for v in chunk { + buf.extend_from_slice(&v.to_le_bytes()); + } + hasher.update(&buf); + } + hasher.finalize().into() +} + +fn prefix(fragment: usize, version: u64, kind: u8, base_version: u64, digest: &[u8; 32]) -> Vec { + let mut out = Vec::with_capacity(PREFIX_LEN); + out.extend_from_slice(&(fragment as u32).to_le_bytes()); + out.extend_from_slice(&version.to_le_bytes()); + out.push(kind); + out.extend_from_slice(&base_version.to_le_bytes()); + out.extend_from_slice(digest); + out +} + +impl DownlinkAnchors { + pub fn new(fragments: usize, full_every: u64) -> Result { + if full_every < 2 { + bail!("--bcast-delta-full-every must be at least 2"); + } + Ok(Self { + full_every, + anchors: (0..fragments).map(|_| None).collect(), + }) + } + + /// Encode the broadcast of a newly committed `version` of fragment `p` + /// and advance the shared anchor. Call exactly once per new version. + pub fn publish(&mut self, p: usize, version: u64, params: &[f32]) -> Result { + let slot = self + .anchors + .get_mut(p) + .ok_or_else(|| anyhow::anyhow!("downlink anchor for unknown fragment {p}"))?; + let delta_ok = matches!( + slot, + Some(a) if a.deltas_since_full + 1 < self.full_every + && a.values.len() == params.len() + && a.version < version + ); + if !delta_ok { + return Ok(Self::reset(slot, p, version, params)); + } + let anchor = slot.as_mut().unwrap(); + let base_version = anchor.version; + let values = Arc::make_mut(&mut anchor.values); + let mut delta = Vec::new(); + delta + .try_reserve_exact(params.len()) + .map_err(|error| anyhow::anyhow!("cannot allocate downlink delta: {error}"))?; + for (a, t) in values.iter_mut().zip(params) { + let d = half::bf16::from_f32(t - *a).to_f32(); + if !d.is_finite() { + // bf16 overflow of a huge step: fall back to FULL. + return Ok(Self::reset(slot, p, version, params)); + } + delta.push(d); + *a += d; + } + anchor.version = version; + anchor.digest = digest_f32(values); + anchor.deltas_since_full += 1; + Ok(EncodedBcast { + prefix: prefix(p, version, KIND_DELTA, base_version, &anchor.digest), + dtype: DTYPE_BF16, + values: Arc::new(delta), + }) + } + + /// FULL re-send for recovery (HELLO / reconnect): the anchor every other + /// learner holds, when it is at `version`; otherwise the exact params. + pub fn resend(&self, p: usize, version: u64, params: &[f32]) -> EncodedBcast { + match self.anchors.get(p).and_then(Option::as_ref) { + Some(anchor) if anchor.version == version => EncodedBcast { + prefix: prefix(p, version, KIND_FULL, version, &anchor.digest), + dtype: DTYPE_F32, + values: Arc::clone(&anchor.values), + }, + _ => EncodedBcast { + prefix: prefix(p, version, KIND_FULL, version, &digest_f32(params)), + dtype: DTYPE_F32, + values: Arc::new(params.to_vec()), + }, + } + } + + fn reset(slot: &mut Option, p: usize, version: u64, params: &[f32]) -> EncodedBcast { + let values = Arc::new(params.to_vec()); + let digest = digest_f32(&values); + *slot = Some(Anchor { + version, + values: Arc::clone(&values), + digest, + deltas_since_full: 0, + }); + EncodedBcast { + prefix: prefix(p, version, KIND_FULL, version, &digest), + dtype: DTYPE_F32, + values, + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + + /// A learner-side mirror: applies broadcasts exactly as yeto/protocol.py. + fn apply(anchor: &mut Option<(u64, Vec)>, enc: &EncodedBcast) -> Result<()> { + let version = u64::from_le_bytes(enc.prefix[4..12].try_into().unwrap()); + let base = u64::from_le_bytes(enc.prefix[13..21].try_into().unwrap()); + let digest: [u8; 32] = enc.prefix[21..53].try_into().unwrap(); + let next = match enc.kind() { + KIND_FULL => enc.values.to_vec(), + KIND_DELTA => { + let Some((held, values)) = anchor.as_ref() else { bail!("no anchor") }; + if *held != base { + bail!("anchor version {held} != base {base}"); + } + values.iter().zip(enc.values.iter()).map(|(a, d)| a + d).collect() + } + _ => bail!("bad kind"), + }; + if digest_f32(&next) != digest { + bail!("digest mismatch"); + } + *anchor = Some((version, next)); + Ok(()) + } + + fn params(seed: u64, n: usize) -> Vec { + (0..n) + .map(|i| ((i as f64 * 0.37 + seed as f64 * 1.3).sin() * 3.0) as f32) + .collect() + } + + #[test] + fn deltas_track_params_within_bf16_error_and_refresh_every_m() { + let mut syncer = DownlinkAnchors::new(1, 4).unwrap(); + let mut learner = None; + let mut theta = params(0, 1000); + let mut kinds = vec![]; + for version in 1..=9u64 { + let step = params(version, 1000); + let prev_anchor = learner.as_ref().map(|(_, v): &(u64, Vec)| v.clone()); + for (t, s) in theta.iter_mut().zip(&step) { + *t += 0.01 * s; + } + let enc = syncer.publish(0, version, &theta).unwrap(); + kinds.push(enc.kind()); + apply(&mut learner, &enc).unwrap(); + let (_, held) = learner.as_ref().unwrap(); + if enc.kind() == KIND_FULL { + assert_eq!(held, &theta); + } else { + let prev = prev_anchor.unwrap(); + for ((h, t), a) in held.iter().zip(&theta).zip(&prev) { + // one bf16 rounding of (θ − anchor_prev): rel. 2^-8 bound + assert!((h - t).abs() <= (t - a).abs() / 256.0 + 2.0 * f32::EPSILON * t.abs()); + } + } + } + assert_eq!( + kinds, + vec![KIND_FULL, KIND_DELTA, KIND_DELTA, KIND_DELTA, KIND_FULL, KIND_DELTA, KIND_DELTA, KIND_DELTA, KIND_FULL] + ); + } + + #[test] + fn resend_is_the_shared_anchor_and_lets_a_reconnected_learner_continue() { + let mut syncer = DownlinkAnchors::new(1, 100).unwrap(); + let mut a = None; + let theta1 = params(1, 64); + apply(&mut a, &syncer.publish(0, 1, &theta1).unwrap()).unwrap(); + let theta2: Vec = theta1.iter().map(|v| v + 0.123).collect(); + apply(&mut a, &syncer.publish(0, 2, &theta2).unwrap()).unwrap(); + // A fresh learner (reconnect) gets the anchor, not θ, as FULL. + let mut b = None; + let full = syncer.resend(0, 2, &theta2); + assert_eq!(full.kind(), KIND_FULL); + apply(&mut b, &full).unwrap(); + assert_eq!(a, b); + let theta3: Vec = theta2.iter().map(|v| v - 0.05).collect(); + let enc = syncer.publish(0, 3, &theta3).unwrap(); + assert_eq!(enc.kind(), KIND_DELTA); + apply(&mut a, &enc).unwrap(); + apply(&mut b, &enc).unwrap(); + assert_eq!(a, b); + } + + #[test] + fn a_learner_with_a_wrong_anchor_detects_it() { + let mut syncer = DownlinkAnchors::new(1, 100).unwrap(); + let mut a = None; + apply(&mut a, &syncer.publish(0, 1, ¶ms(1, 32)).unwrap()).unwrap(); + // Corrupt the learner's anchor bits: the digest check must fail. + if let Some((_, v)) = a.as_mut() { + v[3] += 1e-3; + } + let enc = syncer.publish(0, 2, ¶ms(2, 32)).unwrap(); + assert!(apply(&mut a, &enc).unwrap_err().to_string().contains("digest")); + // A learner that skipped a version fails the base-version check. + let mut stale = Some((0u64, params(1, 32))); + let enc = syncer.publish(0, 3, ¶ms(3, 32)).unwrap(); + assert!(apply(&mut stale, &enc).unwrap_err().to_string().contains("base")); + } + + #[test] + fn full_every_below_two_is_refused() { + assert!(DownlinkAnchors::new(1, 1).is_err()); + } +} diff --git a/syncer/src/main.rs b/syncer/src/main.rs index 8243ab29..58ddc964 100644 --- a/syncer/src/main.rs +++ b/syncer/src/main.rs @@ -2,6 +2,7 @@ mod elastic; mod elastic_server; mod iso_worker; mod merge; +mod downlink; mod protocol; mod server; mod state; @@ -102,6 +103,11 @@ struct Args { /// round is broadcast only after its checkpoint is durable. #[arg(long, default_value_t = false)] async_checkpoint: bool, + /// Learners whose session asks for bf16 downlink deltas get a full f32 + /// BCAST every N-th version of a fragment (N >= 2). Unset: such + /// sessions are refused. + #[arg(long)] + bcast_delta_full_every: Option, /// Resume from --checkpoint-path if it exists. #[arg(long, default_value_t = false)] resume: bool, @@ -278,6 +284,7 @@ fn main() -> anyhow::Result<()> { checkpoint_path: args.checkpoint_path, checkpoint_every: args.checkpoint_every, async_checkpoint: args.async_checkpoint, + bcast_delta_full_every: args.bcast_delta_full_every, resume: args.resume, mark_final_checkpoint: args.mark_final_checkpoint, learner_budget_steps: args.learner_budget_steps, diff --git a/syncer/src/protocol.rs b/syncer/src/protocol.rs index 2103ce3b..d9bd34bc 100644 --- a/syncer/src/protocol.rs +++ b/syncer/src/protocol.rs @@ -62,26 +62,80 @@ pub const DTYPE_Q4: u8 = 3; /// masters and momentum; only the learner's outer gradient is rounded. pub const DTYPE_F32_PUSH_BF16: u8 = 4; +/// Session dtype 5 (S26 E3): PUSH_FRAGMENT deltas travel as block-scaled +/// int8 (see `decode_i8`); INIT_PARAMS, BCAST_FRAGMENT and FINAL stay f32. +pub const DTYPE_F32_PUSH_I8: u8 = 5; +/// Session-dtype flag (S26 E3): BCAST_FRAGMENT carries a 41-byte downlink +/// header and, between periodic f32 refreshes, a bf16 delta against the +/// anchor every learner holds (see `crate::downlink`). Valid only on f32-bulk +/// sessions (dtypes 1, 4, 5). +pub const DTYPE_FLAG_BCAST_DELTA: u8 = 0x80; + /// Values per q4 scale block (f32 absmax scale + 128 nibble bytes each). pub const Q4_BLOCK: usize = 256; +/// Bounds of the int8 push block size the learner declares per payload. +pub const I8_BLOCK_MIN: usize = 32; +pub const I8_BLOCK_MAX: usize = 1 << 20; + +/// The session dtype without its flags. +pub fn session_base(dtype: u8) -> u8 { + dtype & !DTYPE_FLAG_BCAST_DELTA +} + +/// Whether the session receives bf16 downlink deltas. +pub fn bcast_delta(dtype: u8) -> bool { + dtype & DTYPE_FLAG_BCAST_DELTA != 0 +} + +/// Whether `dtype` is a session dtype this syncer accepts. +pub fn valid_session_dtype(dtype: u8) -> bool { + match session_base(dtype) { + DTYPE_F32 | DTYPE_F32_PUSH_BF16 | DTYPE_F32_PUSH_I8 => true, + DTYPE_BF16 | DTYPE_Q4 => !bcast_delta(dtype), + _ => false, + } +} /// Dtype of INIT_PARAMS/BCAST_FRAGMENT tensors for a session dtype: q4 /// applies only to push deltas; full parameter payloads stay bf16. pub fn bulk_dtype(dtype: u8) -> u8 { - match dtype { + match session_base(dtype) { DTYPE_Q4 => DTYPE_BF16, - DTYPE_F32_PUSH_BF16 => DTYPE_F32, + DTYPE_F32_PUSH_BF16 | DTYPE_F32_PUSH_I8 => DTYPE_F32, other => other, } } -/// Dtype of PUSH_FRAGMENT deltas for an unquantized session dtype (q4 has -/// its own block codec). +/// Dtype of PUSH_FRAGMENT deltas for an unquantized session dtype (q4 and +/// int8 have their own block codecs). pub fn push_dtype(dtype: u8) -> u8 { - if dtype == DTYPE_F32_PUSH_BF16 { - DTYPE_BF16 - } else { - dtype + match session_base(dtype) { + DTYPE_F32_PUSH_BF16 => DTYPE_BF16, + other => other, + } +} + +/// Largest PUSH_FRAGMENT delta payload a session dtype may carry. +pub fn max_push_nbytes(dtype: u8, numel: usize) -> Result { + match session_base(dtype) { + DTYPE_Q4 => q4_nbytes(numel), + DTYPE_F32_PUSH_I8 => i8_nbytes(numel, I8_BLOCK_MIN), + other => tensor_nbytes(push_dtype(other), numel), + } +} + +/// Decode a PUSH_FRAGMENT delta payload of any session dtype into f32s. +pub fn decode_push(dtype: u8, bytes: &[u8], numel: usize, out: &mut Vec) -> Result<()> { + match session_base(dtype) { + DTYPE_Q4 => decode_q4(bytes, numel, out), + DTYPE_F32_PUSH_I8 => decode_i8(bytes, numel, out), + other => { + let expected = tensor_nbytes(push_dtype(other), numel)?; + if bytes.len() != expected { + bail!("push delta has {} bytes, expected {expected}", bytes.len()); + } + decode_tensor(push_dtype(other), bytes, out) + } } } @@ -271,6 +325,80 @@ pub fn encode_tensor(dtype: u8, vals: &[f32], out: &mut Vec) -> Result<()> { Ok(()) } +/// Exact encoded byte length of an int8 push payload with block size `block`. +pub fn i8_nbytes(numel: usize, block: usize) -> Result { + if !(I8_BLOCK_MIN..=I8_BLOCK_MAX).contains(&block) { + bail!("int8 block size {block} outside {I8_BLOCK_MIN}..={I8_BLOCK_MAX}"); + } + numel + .div_ceil(block) + .checked_mul(4) + .and_then(|scales| scales.checked_add(numel)) + .and_then(|body| body.checked_add(4)) + .ok_or_else(|| anyhow::anyhow!("int8 tensor byte length overflow")) +} + +/// Decode an int8 delta payload (yeto/protocol.py `f32_bytes_to_i8`). +/// +/// Layout: u32 block size B, then per block of B values (the last block may +/// be shorter): f32 scale (absmax / 127), then one i8 per value. A value +/// decodes to `q as f32 * scale`, so |error| <= scale / 2 = absmax / 254 per +/// element (plus one f32 rounding). The scale is non-negative and finite. +pub fn decode_i8(bytes: &[u8], numel: usize, out: &mut Vec) -> Result<()> { + if bytes.len() < 4 { + bail!("int8 payload is shorter than its block header"); + } + let block = u32::from_le_bytes(bytes[..4].try_into().unwrap()) as usize; + let expected = i8_nbytes(numel, block)?; + if bytes.len() != expected { + bail!( + "int8 payload has {} bytes, expected {expected} for {numel} values in blocks of {block}", + bytes.len() + ); + } + out.clear(); + out.try_reserve(numel) + .map_err(|error| anyhow::anyhow!("cannot allocate int8 tensor: {error}"))?; + let mut at = 4usize; + let mut remaining = numel; + let mut b = 0usize; + while remaining > 0 { + let len = remaining.min(block); + let scale = f32::from_le_bytes(bytes[at..at + 4].try_into().unwrap()); + if !scale.is_finite() || scale < 0.0 { + bail!("int8 block {b} has invalid scale {scale}"); + } + at += 4; + for &q in &bytes[at..at + len] { + let q = q as i8; + if q == i8::MIN { + bail!("int8 block {b} uses the reserved value -128"); + } + out.push(q as f32 * scale); + } + at += len; + remaining -= len; + b += 1; + } + Ok(()) +} + +/// Encode f32s as an int8 push payload (test counterpart of the Python codec). +#[cfg(test)] +pub fn encode_i8(vals: &[f32], block: usize) -> Vec { + let mut out = (block as u32).to_le_bytes().to_vec(); + for chunk in vals.chunks(block) { + let absmax = chunk.iter().fold(0f32, |m, v| m.max(v.abs())); + let scale = absmax / 127.0; + out.extend_from_slice(&scale.to_le_bytes()); + for v in chunk { + let q = if scale > 0.0 { (v / scale).round_ties_even().clamp(-127.0, 127.0) } else { 0.0 }; + out.push(q as i8 as u8); + } + } + out +} + /// Decode a q4 delta payload (yeto/tensor_io.py `quantize_q4`) into f32s. /// /// Per block of `Q4_BLOCK` values: f32 absmax scale, then 128 bytes of @@ -326,6 +454,59 @@ mod tests { use super::*; use tokio::io::AsyncWriteExt; + #[test] + fn int8_push_round_trips_within_half_a_scale_step() { + let vals: Vec = (0..1000) + .map(|i| ((i as f32) * 0.731).sin() * if i % 300 < 150 { 1e-3 } else { 2.5 }) + .collect(); + for block in [I8_BLOCK_MIN, 128, 333, 4096] { + let wire = encode_i8(&vals, block); + assert_eq!(wire.len(), i8_nbytes(vals.len(), block).unwrap()); + assert!(wire.len() <= max_push_nbytes(DTYPE_F32_PUSH_I8, vals.len()).unwrap()); + let mut out = Vec::new(); + decode_push(DTYPE_F32_PUSH_I8, &wire, vals.len(), &mut out).unwrap(); + for (chunk, decoded) in vals.chunks(block).zip(out.chunks(block)) { + let absmax = chunk.iter().fold(0f32, |m, v| m.max(v.abs())); + for (v, d) in chunk.iter().zip(decoded) { + // |error| <= scale/2 = absmax/254, plus f32 rounding. + assert!((v - d).abs() <= absmax / 254.0 * (1.0 + 1e-5), "{v} {d} {absmax}"); + } + } + } + // All-zero blocks decode to exact zeros. + let mut out = Vec::new(); + decode_i8(&encode_i8(&[0.0; 40], 32), 40, &mut out).unwrap(); + assert!(out.iter().all(|v| *v == 0.0)); + } + + #[test] + fn int8_push_rejects_malformed_payloads() { + let vals = vec![1.0f32; 64]; + let mut out = Vec::new(); + let wire = encode_i8(&vals, 32); + assert!(decode_i8(&wire[..wire.len() - 1], 64, &mut out).is_err()); + assert!(decode_i8(&encode_i8(&vals, 16), 64, &mut out).is_err()); + let mut bad_scale = wire.clone(); + bad_scale[4..8].copy_from_slice(&(-1.0f32).to_le_bytes()); + assert!(decode_i8(&bad_scale, 64, &mut out).is_err()); + let mut reserved = wire; + reserved[8] = 0x80; + assert!(decode_i8(&reserved, 64, &mut out).is_err()); + } + + #[test] + fn session_dtype_flags_compose_only_with_f32_bulk() { + assert!(valid_session_dtype(DTYPE_F32_PUSH_I8)); + assert!(valid_session_dtype(DTYPE_F32 | DTYPE_FLAG_BCAST_DELTA)); + assert!(valid_session_dtype(DTYPE_F32_PUSH_BF16 | DTYPE_FLAG_BCAST_DELTA)); + assert!(valid_session_dtype(DTYPE_F32_PUSH_I8 | DTYPE_FLAG_BCAST_DELTA)); + assert!(!valid_session_dtype(DTYPE_BF16 | DTYPE_FLAG_BCAST_DELTA)); + assert!(!valid_session_dtype(DTYPE_Q4 | DTYPE_FLAG_BCAST_DELTA)); + assert!(!valid_session_dtype(6)); + assert_eq!(bulk_dtype(DTYPE_F32_PUSH_I8 | DTYPE_FLAG_BCAST_DELTA), DTYPE_F32); + assert_eq!(push_dtype(DTYPE_F32_PUSH_BF16 | DTYPE_FLAG_BCAST_DELTA), DTYPE_BF16); + } + #[test] fn reserved_message_ids_are_stable() { assert_eq!(MSG_ERROR, 10); diff --git a/syncer/src/server.rs b/syncer/src/server.rs index 6d0b4ba1..46d4cb52 100644 --- a/syncer/src/server.rs +++ b/syncer/src/server.rs @@ -97,6 +97,10 @@ pub struct Config { /// are unchanged, so this is deliberately not part of the semantic /// profile hash. pub async_checkpoint: bool, + /// S26 E3: full-f32 refresh period of bf16 downlink deltas (sessions + /// with `DTYPE_FLAG_BCAST_DELTA`). Wire encoding only: merges, versions + /// and checkpoints are unchanged, so it is not in the profile hash. + pub bcast_delta_full_every: Option, pub resume: bool, /// Create the adjacent final marker only for a completed terminal cut. pub mark_final_checkpoint: bool, @@ -746,6 +750,9 @@ fn validate_config(cfg: &Config) -> Result<()> { ); } } + if matches!(cfg.bcast_delta_full_every, Some(n) if n < 2) { + bail!("--bcast-delta-full-every must be at least 2"); + } if cfg.async_checkpoint { if cfg.policy_sweep_fragments.is_none() { bail!("--async-checkpoint requires --policy-sweep-fragments"); @@ -1015,11 +1022,7 @@ fn negotiated_payload_limits(layout: &Layout, dtype: u8) -> Result<(u64, u64)> { for fragment in &layout.fragments { let numel = fragment.numel()?; let init_bytes = tensor_nbytes(bulk_dtype(dtype), numel)? as u64; - let delta_bytes = if dtype == DTYPE_Q4 { - q4_nbytes(numel)? - } else { - tensor_nbytes(push_dtype(dtype), numel)? - } as u64; + let delta_bytes = max_push_nbytes(dtype, numel)? as u64; max_init = max_init.max( 4u64.checked_add(init_bytes) .context("INIT_PARAMS payload limit overflow")?, @@ -1164,7 +1167,7 @@ fn parse_hello(payload: &[u8], expected_learners: u32) -> Result { bail!("connection generation must be nonzero"); } let dtype = r.u8()?; - if !matches!(dtype, DTYPE_F32 | DTYPE_BF16 | DTYPE_Q4 | DTYPE_F32_PUSH_BF16) { + if !valid_session_dtype(dtype) { bail!("unsupported session dtype {dtype}"); } let num_fragments = r.u32()?; @@ -1735,23 +1738,9 @@ async fn dispatch_inner( .with_context(|| format!("PUSH_FRAGMENT for unknown fragment {fragment_id}"))?; let numel = fragment.numel()?; let bytes = r.rest(); - let expected = if group.dtype == DTYPE_Q4 { - q4_nbytes(numel)? - } else { - tensor_nbytes(push_dtype(group.dtype), numel)? - }; - if bytes.len() != expected { - bail!( - "PUSH_FRAGMENT fragment {fragment_id} has {} delta bytes, expected {expected}", - bytes.len() - ); - } let mut outer_gradient = Vec::new(); - if group.dtype == DTYPE_Q4 { - decode_q4(bytes, numel, &mut outer_gradient)?; - } else { - decode_tensor(push_dtype(group.dtype), bytes, &mut outer_gradient)?; - } + decode_push(group.dtype, bytes, numel, &mut outer_gradient) + .with_context(|| format!("PUSH_FRAGMENT fragment {fragment_id}"))?; // Every wire dtype carries learner_delta = local - raw_anchor. // Merge math consumes the corresponding outer gradient with the // opposite sign and never reconstructs a full learner parameter. @@ -1897,6 +1886,17 @@ async fn send_fragment_until_cutoff( dtype: u8, budget_cutoff: &BudgetCutoff, ) -> Result { + if let Some(encoded) = recovery_bcast(st, p, msg_type) { + return send_tensor_until_cutoff( + group, + msg_type, + &encoded.prefix, + encoded.dtype, + &encoded.values, + budget_cutoff, + ) + .await; + } let mut prefix = [0u8; 12]; prefix[..4].copy_from_slice(&(p as u32).to_le_bytes()); prefix[4..].copy_from_slice(&st.versions[p].to_le_bytes()); @@ -1911,6 +1911,17 @@ async fn send_fragment_until_cutoff( .await } +/// A BCAST outside publication (initial cut, HELLO/reconnect recovery) on a +/// bf16-downlink session: always FULL, carrying the shared anchor. +fn recovery_bcast(st: &GlobalState, p: usize, msg_type: u8) -> Option { + if msg_type != MSG_BCAST_FRAGMENT { + return None; + } + st.downlink + .as_ref() + .map(|anchors| anchors.lock().unwrap().resend(p, st.versions[p], &st.params[p])) +} + async fn send_pull_until_cutoff( groups: &[Arc], pull: &bytes::Bytes, @@ -3481,16 +3492,31 @@ async fn publish_round( // its exact f32 cut immediately after the scheduler reaches quiescence. if t < cfg.total_steps { let dtype = bulk_dtype(st.wire_dtype); + // bf16 downlink: encode this version once (advancing the shared + // anchor) and send the same bytes to every learner. + let published = match &st.downlink { + Some(anchors) => Some(anchors.lock().unwrap().publish(p, st.versions[p], &st.params[p])?), + None => None, + }; for group in current_groups(registry) { - match send_fragment_until_cutoff( - st, - &group, - p, - MSG_BCAST_FRAGMENT, - dtype, - budget_cutoff, - ) - .await + let sent = match &published { + Some(encoded) => { + send_tensor_until_cutoff( + &group, + MSG_BCAST_FRAGMENT, + &encoded.prefix, + encoded.dtype, + &encoded.values, + budget_cutoff, + ) + .await + } + None => { + send_fragment_until_cutoff(st, &group, p, MSG_BCAST_FRAGMENT, dtype, budget_cutoff) + .await + } + }; + match sent { Ok(true) => {} Ok(false) => break, @@ -3697,8 +3723,8 @@ async fn complete_round_async( fn new_state_for(group: &Arc, cfg: &Config) -> Result { if let Some(fragments) = cfg.policy_sweep_fragments { // Masters and every BCAST/FINAL stay f32; only push deltas may be bf16. - if !matches!(group.dtype, DTYPE_F32 | DTYPE_F32_PUSH_BF16) { - bail!("--policy-sweep-fragments requires an FP32 HELLO session (optionally with bf16 pushes)"); + if bulk_dtype(group.dtype) != DTYPE_F32 { + bail!("--policy-sweep-fragments requires an FP32 HELLO session (optionally with bf16/int8 pushes)"); } if group.layout.fragments.len() != fragments as usize { bail!( @@ -3717,6 +3743,17 @@ fn new_state_for(group: &Arc, cfg: &Config) -> Result { )?; st.policy_sweep_fragments = cfg.policy_sweep_fragments; st.session_contract_hash = Some(group.session_contract_hash); + match (bcast_delta(group.dtype), cfg.bcast_delta_full_every) { + (true, Some(full_every)) => { + st.downlink = Some(Mutex::new(crate::downlink::DownlinkAnchors::new( + group.layout.fragments.len(), + full_every, + )?)); + } + (true, None) => bail!("session asks for bf16 downlink deltas, but --bcast-delta-full-every is unset"), + (false, Some(_)) => bail!("--bcast-delta-full-every is set, but the session has no bf16 downlink deltas"), + (false, None) => {} + } if cfg.delta_correction { st.delta_correction = Some(crate::merge::Heloco::default()); } @@ -4055,6 +4092,11 @@ async fn send_fragment( msg_type: u8, dtype: u8, ) -> Result<()> { + if let Some(encoded) = recovery_bcast(st, p, msg_type) { + return group + .send_tensor_large(msg_type, &encoded.prefix, encoded.dtype, &encoded.values) + .await; + } let mut prefix = [0u8; 12]; prefix[..4].copy_from_slice(&(p as u32).to_le_bytes()); prefix[4..].copy_from_slice(&st.versions[p].to_le_bytes()); @@ -4697,6 +4739,7 @@ mod tests { checkpoint_path: None, checkpoint_every: 0, async_checkpoint: false, + bcast_delta_full_every: None, resume: false, mark_final_checkpoint: false, learner_budget_steps: None, @@ -5972,13 +6015,13 @@ mod tests { #[tokio::test] async fn updated_broadcast_propagates_local_errors_but_tolerates_closed_streams() { - let invalid_state = broadcast_test_state(255); + let invalid_state = broadcast_test_state(9); let (connected, _receiver) = streaming_test_group(DTYPE_F32, false); let error = broadcast_updated_fragment(&invalid_state, ®istry_with_group(connected), 0) .await .unwrap_err(); assert!( - format!("{error:#}").contains("unknown tensor dtype 255"), + format!("{error:#}").contains("unknown tensor dtype 9"), "unexpected error: {error:#}" ); diff --git a/syncer/src/state.rs b/syncer/src/state.rs index 6d441561..3b54cb9e 100644 --- a/syncer/src/state.rs +++ b/syncer/src/state.rs @@ -146,6 +146,8 @@ pub struct GlobalState { pub delta_correction: Option, /// Spectrum flattener. Torch mode owns one persistent child process. iso_backend: IsoBackend, + /// bf16 downlink anchors (S26 E3); in memory only, never checkpointed. + pub downlink: Option>, } /// A merge whose deterministic learner reduction has completed but whose @@ -344,6 +346,7 @@ impl GlobalState { wire_dtype, delta_correction: None, iso_backend: IsoBackend::start(iso_backend_config)?, + downlink: None, }) } diff --git a/tests/test_e3_sync_options.py b/tests/test_e3_sync_options.py index 0e8492ef..a3256227 100644 --- a/tests/test_e3_sync_options.py +++ b/tests/test_e3_sync_options.py @@ -2,21 +2,44 @@ from __future__ import annotations +import hashlib +import itertools from dataclasses import replace +import numpy as np import pytest import torch - -from test_rl_algorithm_provenance import _fake_sky, _launcher_args, _no_modal_listing # noqa: F401 +from test_rl_algorithm_provenance import ( # noqa: F401 + _fake_sky, + _island_run, + _launcher_args, + _no_modal_listing, +) +from test_rl_dense_sweep_real_syncer import syncer_binary # noqa: F401 from test_rl_full_parameter_sync import _role_stream_config + +from yeto.downlink import ( + HEADER, + KIND_DELTA, + KIND_FULL, + DownlinkAnchors, + DownlinkChunkSink, + DownlinkResync, +) from yeto.protocol import ( DTYPE_BF16, DTYPE_F32, DTYPE_F32_PUSH_BF16, + DTYPE_F32_PUSH_I8, + DTYPE_FLAG_BCAST_DELTA, _f32_parts_to_bf16, + _f32_parts_to_i8, bulk_dtype, f32_bytes_to_bf16, + f32_bytes_to_i8, + i8_nbytes, push_wire_dtype, + valid_session_dtype, ) @@ -44,7 +67,7 @@ def test_push_dtype_is_bound_into_the_role_stream_session_contract(): ) exact = _role_stream_config("actor", 30100, sweep=True) - rounded = replace(exact, push_bf16=True) + rounded = replace(exact, push_dtype="bf16") assert (role_stream_session_contract_hash("1" * 64, exact) != role_stream_session_contract_hash("1" * 64, rounded)) @@ -61,3 +84,235 @@ def test_e3_options_only_apply_to_the_full_parameter_preset(flag, value): args.controller = "local" with pytest.raises(ValueError, match=f"{flag} only applies to --rl-sync-preset full-parameter"): _prepare_rl_args(args) + + +# -- S26 E3 third batch: int8 pushes, bf16 downlink --------------------------- + +def _decode_i8(wire: bytes, numel: int) -> np.ndarray: + block = int(np.frombuffer(wire, "= 2"), + (("--rl-sync-preset", "decoupled", "--local-rl-rounds-per-sync", "2", + "--rl-sync-bcast-delta-full-every", "4"), "only applies to --rl-sync-preset strict-avg"), + (("--rl-sync-push-dtype", "int8"), "--rl-sync-push-dtype only applies"), +)) +def test_e3_downlink_and_int8_scope(extra, match): + from yeto.launcher import _prepare_rl_args + + args = _launcher_args("ports", extra, gpu="modal:2xh100") + args.controller = "local" + with pytest.raises(ValueError, match=match): + _prepare_rl_args(args) + + +def _run_bridge_pair(binary, tmp_path, rounds, *, full_every): + """Two fake islands through ``StrictRlBridge`` (the verl/Miles LoRA path).""" + import subprocess + + from test_rl_integration import _BridgeThread, _FakeMiles, _port, _start + + from yeto.rl.bridge import BridgeConfig, StrictRlBridge + + port = _port() + process = _start(binary, port, tmp_path / "state.ckpt", rounds=rounds) + if full_every: + process.kill() + process.wait() + argv = [*process.args, "--bcast-delta-full-every", str(full_every)] + process = subprocess.Popen(argv, stdout=subprocess.PIPE, stderr=subprocess.STDOUT, text=True) + runtimes = (_FakeMiles([0.1, 0.3], 0), _FakeMiles([0.3, 0.5], 1)) + threads = [] + try: + for learner_id, runtime in enumerate(runtimes): + config = BridgeConfig( + syncer_addr=("127.0.0.1", port), learner_id=learner_id, global_rounds=rounds, + groups_per_round=1, samples_per_group=2, local_optimizer_steps=1, wan_streams=0, + expected_specs=runtime.current.specs, + base_model_revision=runtime.current.base_model_revision, + lora_config_hash=runtime.current.lora_config_hash, + layout_hash=runtime.current.layout_hash, + event_tape=str(tmp_path / f"island-{learner_id}.jsonl"), + bcast_delta=bool(full_every)) + threads.append(_BridgeThread(StrictRlBridge(runtime, config))) + for thread in threads: + thread.start() + for thread in threads: + thread.join(timeout=30) + assert not thread.is_alive() and thread.error is None, thread.error + assert process.wait(timeout=10) == 0 + return [next(iter(t.result.tensors.values())) for t in threads], runtimes + finally: + for thread in threads: + thread.bridge.client.close() + if process.poll() is None: + process.kill() + process.wait() + + +def test_strict_bridge_with_bf16_downlink_finishes_on_the_exact_f32_cut(syncer_binary, tmp_path): # noqa: F811 + """E3, verl/LoRA path: islands train from the bf16-delta anchor, the syncer + keeps exact masters, and the final cut matches the exact f32 run (up to the + f32 rounding of ``local - anchor`` on a different anchor).""" + binary = syncer_binary + exact_dir, delta_dir = tmp_path / "f32", tmp_path / "delta" + exact_dir.mkdir() + delta_dir.mkdir() + exact, _ = _run_bridge_pair(binary, exact_dir, 4, full_every=0) + delta, runtimes = _run_bridge_pair(binary, delta_dir, 4, full_every=3) + for want, got in zip(exact, delta, strict=True): + assert torch.allclose(want, got, rtol=0, atol=1e-6), (want, got) + assert all(runtime.applied[-1] == 4 for runtime in runtimes) + + +def test_verl_island_accepts_the_downlink_flag(): + from yeto.rl.adapters.verl.island_entry import parse_args as verl_parse_args + + args = verl_parse_args([ + "--model", "m", "--model-revision", "r", "--data", "d", "--learner-id", "0", + "--reward-function", "f", "--global-rounds", "1", "--groups-per-round", "1", + "--samples-per-group", "1", "--rollout-max-response-len", "8", "--event-tape", "t", + "--inner-lr", "1e-6", "--seq-len", "8", "--seed", "0", "--bcast-delta"]) + assert args[0].bcast_delta is True diff --git a/tests/test_e3_sync_overlap.py b/tests/test_e3_sync_overlap.py new file mode 100644 index 00000000..9b9fbc4d --- /dev/null +++ b/tests/test_e3_sync_overlap.py @@ -0,0 +1,331 @@ +"""S26 E3 item 3: sync/train overlap of the full-parameter sweep (CPU only). + +A numeric fake role stream and an in-process strict-sweep syncer drive the +real ``FullParameterSync``/``RolePairing``: the island trains round k+1 while +sweep k is in flight and rebases ``theta <- G_k + (local - base)``. The tests +check that every round's progress reaches the global weights exactly once +(no double counting), that the trainer keeps its progress across a rebase, +and the launcher/island policy-age rule. +""" + +from __future__ import annotations + +import asyncio +from types import SimpleNamespace + +import numpy as np +import pytest +from test_rl_algorithm_provenance import ( # noqa: F401 + _fake_sky, + _launcher_args, + _no_modal_listing, +) + +from yeto.protocol import PullRequest +from yeto.rl.adapters.miles.full_parameter.role_stream import RoleSubmissionCandidate +from yeto.rl.engine.full_parameter_sync import ( + FullParameterSync, + FullParameterSyncProfile, + RolePairing, +) +from yeto.rl.engine.sync_overlap import rebase_values, required_policy_age + +P = 2 # fragments +NUMEL = 4 + + +class _SweepSyncer: + """One-island strict sweep: PULLs a sweep once the previous one committed; + ``G_f += pushed delta``; the terminal step publishes FINAL.""" + + def __init__(self, initial: list[np.ndarray], sweeps: int) -> None: + self.global_ = [value.copy() for value in initial] + self.versions = [0] * P + self.sweeps = sweeps + self.step = 0 # last committed global step + self.pushed: list[tuple[int, np.ndarray]] = [] + self.bcasts: list[tuple[int, int, np.ndarray]] = [] + self.finalizing = False + self._since_commit = 0 + + def permits(self) -> list[PullRequest]: + sweep = self.step // P # completed sweeps + if self.finalizing or sweep >= self.sweeps or self.step % P: + return [] + pushed = {step for step, _ in self.pushed} + return [PullRequest(f, sweep * P + f + 1, 1) for f in range(P) if sweep * P + f + 1 not in pushed] + + def push(self, step: int, fragment_id: int, delta: np.ndarray) -> None: + assert self.step < step <= self.step + P and (step - 1) % P == fragment_id + self.pushed.append((step, delta.copy())) + self.global_[fragment_id] = self.global_[fragment_id] + delta + self.versions[fragment_id] = step + self._since_commit += 1 + if self._since_commit < P: + return + self._since_commit = 0 + self.step += P + if self.step == self.sweeps * P: + self.finalizing = True + else: + self.bcasts.extend((f, self.versions[f], self.global_[f].copy()) for f in range(P)) + + +class _NumericRole: + """A RoleStream whose cuts are lists of per-fragment f32 arrays.""" + + role = "actor" + + def __init__(self, syncer: _SweepSyncer, progress) -> None: + self.syncer = syncer + self.progress = progress # rollout_id -> list of per-fragment progress + self.config = SimpleNamespace(overlap=True) + self.current = [g.copy() for g in syncer.global_] + self.anchors = [g.copy() for g in syncer.global_] + self._versions = [0] * P + self.submission = None + self.submission_versions = None + self.local_rounds = self.submission_rounds = 0 + self.rounds_at_anchor = [0] * P + self.round_base = None + self.rounds_at_base = 0 + self._pending: dict[int, PullRequest] = {} + self._staged: dict[tuple[int, int], np.ndarray] = {} + self.terminal_submitted = False + self.final_acknowledged = False + self.closed = False + self.trained_from: list[list[np.ndarray]] = [] + + @property + def fragment_versions(self): + return tuple(self._versions) + + @property + def finalizing(self): + return self.syncer.finalizing + + async def initialize(self, group): + del group + + def resident_state(self, policy_version): + return SimpleNamespace(policy_version=policy_version, policy_tensor_hash=lambda: "cut") + + async def complete_local_round(self, *, rollout_id, trajectories, hold_submission=False): + del trajectories + self.trained_from.append([c.copy() for c in self.current]) + self.current = [c + d for c, d in zip(self.current, self.progress(rollout_id), strict=True)] + self.local_rounds += 1 + if not hold_submission: + self.submission = [c.copy() for c in self.current] + self.submission_versions = self.fragment_versions + self.submission_rounds = self.local_rounds + + def ready_submission_candidates(self): + for permit in self.syncer.permits(): + self._pending.setdefault(permit.fragment_id, permit) + out = [] + for f, permit in sorted(self._pending.items()): + base = max(0, permit.global_step - P) + if (self.submission is None or self._versions[f] != base + or self.submission_versions[f] != base + or self.submission_rounds - self.rounds_at_anchor[f] < 1): + continue + out.append(RoleSubmissionCandidate( + "actor", permit, base, self.submission_rounds, 1, 1, NUMEL * 4, + self.submission, self.anchors[f])) + return tuple(out) + + def retain_candidate(self, candidate): + del candidate + + def release_candidate(self, candidate): + del candidate + + def push_candidate(self, candidate): + f = candidate.permit.fragment_id + self.syncer.push(candidate.permit.global_step, f, + candidate.submission_cut[f] - candidate.anchor) + self._pending.pop(f) + self.terminal_submitted |= candidate.permit.global_step == self.syncer.sweeps * P + return SimpleNamespace(role="actor", fragment_id=f, global_step=candidate.permit.global_step) + + async def stage_available_broadcasts(self): + for f, version, value in self.syncer.bcasts: + if version > self._versions[f]: + self._staged[(f, version)] = value + self.syncer.bcasts.clear() + return tuple(self._staged) + + @property + def staged_fragment_versions(self): + return tuple(sorted(self._staged)) + + def has_staged_broadcast(self, fragment_id, version): + return (fragment_id, version) in self._staged + + async def apply_staged_broadcast(self, fragment_id, version, *, rebase=False): + assert rebase, "overlap applies every sweep by rebasing" + g = self._staged.pop((fragment_id, version)) + self.current[fragment_id] = rebase_values( + g, self.current[fragment_id], self.round_base[fragment_id]) + self.anchors[fragment_id] = g + self._versions[fragment_id] = version + self.rounds_at_anchor[fragment_id] = self.rounds_at_base + return (fragment_id,) + + def finish_rebase(self): + self.submission = [c.copy() for c in self.current] + self.submission_versions = self.fragment_versions + self.submission_rounds = self.local_rounds + + def mark_round_base(self): + self.round_base = [c.copy() for c in self.current] + self.rounds_at_base = self.local_rounds + + async def wait_and_apply_final(self): + self.current = [g.copy() for g in self.syncer.global_] + self._versions = [self.syncer.sweeps * P - P + f + 1 for f in range(P)] + + def acknowledge_final(self): + self.final_acknowledged = True + + def discard_detached_submission_cut(self): + pass + + def close(self): + self.closed = True + + +class _Driver: + def __init__(self): + self.events = [] + + def emit(self, event, **fields): + self.events.append((event, fields)) + + +def _run(sweeps: int, overlap: bool = True): + rng = np.random.default_rng(7) + initial = [rng.standard_normal(NUMEL).astype(np.float32) for _ in range(P)] + steps = {k: [rng.standard_normal(NUMEL).astype(np.float32) * 0.1 for _ in range(P)] + for k in range(sweeps)} + syncer = _SweepSyncer(initial, sweeps) + role = _NumericRole(syncer, lambda rollout_id: steps[rollout_id]) + profile = FullParameterSyncProfile(roles=("actor",), total_fragment_steps=sweeps * P, fragments=P, + local_horizon=1, pipeline_depth=1, wait_timeout=5.0, + poll_seconds=0.001, sweep=True, overlap=overlap) + session = FullParameterSync(RolePairing({"actor": role}, profile), groups={"actor": "g"}, + run=asyncio.run, trajectories=lambda rollout_id: rollout_id) + driver = _Driver() + session.start(driver) + stops = [] + for rollout_id in range(sweeps): + boundary = session.boundary(driver, rollout_id=rollout_id, stats=None) + stops.append(boundary.stop) + if boundary.stop: + session.published(driver, rollout_id=rollout_id, policy_hash="x") + break + return initial, steps, syncer, role, stops + + +def test_overlap_counts_every_round_exactly_once(): + sweeps = 5 + initial, steps, syncer, role, stops = _run(sweeps) + assert stops == [False] * (sweeps - 1) + [True] + # Pushed deltas are exactly the per-round progress, in round order. + pushed = sorted(syncer.pushed, key=lambda row: row[0]) + for k in range(sweeps): + for f in range(P): + np.testing.assert_allclose(pushed[k * P + f][1], steps[k][f], rtol=0, atol=1e-6) + # Global = initial + sum of every round's progress, once. + for f in range(P): + want = initial[f] + sum(steps[k][f] for k in range(sweeps)) + np.testing.assert_allclose(syncer.global_[f], want, rtol=0, atol=1e-5) + np.testing.assert_allclose(role.current[f], want, rtol=0, atol=1e-5) + assert role.final_acknowledged + + +def test_overlap_trains_one_outer_step_stale_and_keeps_local_progress(): + sweeps = 4 + initial, steps, _syncer, role, _ = _run(sweeps) + # Round 0 starts from the initial cut, round 1 from its own round-0 result + # (sweep 0 still in flight), round k>=2 from G_{k-2} + progress of round k-1. + for f in range(P): + np.testing.assert_allclose(role.trained_from[0][f], initial[f]) + np.testing.assert_allclose(role.trained_from[1][f], initial[f] + steps[0][f], atol=1e-6) + for k in range(2, sweeps): + g_prev = initial[f] + sum(steps[j][f] for j in range(k - 1)) + np.testing.assert_allclose(role.trained_from[k][f], g_prev + steps[k - 1][f], atol=1e-5) + + +def test_rebase_values_is_anchor_plus_local_progress(): + anchor = np.array([1.0, -2.0], dtype=np.float32) + current = np.array([0.5, 0.25], dtype=np.float32) + base = np.array([0.25, 0.5], dtype=np.float32) + np.testing.assert_array_equal(rebase_values(anchor, current, base), [1.25, -2.25]) + with pytest.raises(ValueError, match="float32"): + rebase_values(anchor.astype(np.float64), current, base) + with pytest.raises(ValueError, match="non-finite"): + rebase_values(np.array([np.inf, 0], np.float32), current, base) + + +def test_overlap_profile_needs_the_strict_sweep(): + with pytest.raises(ValueError, match="overlap needs the dense strict sweep"): + FullParameterSyncProfile(roles=("actor",), total_fragment_steps=4, fragments=2, local_horizon=2, + pipeline_depth=1, wait_timeout=1.0, poll_seconds=0.01, overlap=True) + + +def _dense(*extra): + from yeto.launcher import _prepare_rl_args + + args = _launcher_args("ports", ( + "--rl-sync-preset", "full-parameter", "--rl-placement", "fixed-partition", + "--rl-rollout-gpus", "1", "--rl-parameter-layout-sha256", "e" * 64, + "--fragments", "2", "--total-steps", "2", *extra), gpu="modal:2xh100") + args.controller = "local" + _prepare_rl_args(args) + return args + + +def test_overlap_raises_the_minimum_policy_age_and_reaches_the_island(): + from yeto.launch.rl_island_task import _rl_trainable_island_flags + + assert required_policy_age(0) == 1 + args = _dense("--rl-sync-overlap") + assert args.rl_max_policy_age == 1 + assert "--sweep-overlap" in _rl_trainable_island_flags(args, 0, 1) + assert _dense("--rl-sync-overlap", "--rl-max-policy-age", "3").rl_max_policy_age == 3 + plain = _dense() + assert int(plain.rl_max_policy_age or 0) == 0 + assert "--sweep-overlap" not in _rl_trainable_island_flags(plain, 0, 1) + + +def test_overlap_launch_fails_when_the_algorithm_spec_does_not_tolerate_the_age(): + from yeto.launch_preflight import PreflightError, check_policy_age_spec + + args = _dense("--rl-sync-overlap") + with pytest.raises(PreflightError, match="--rl-max-policy-age 1 needs an algorithm spec"): + check_policy_age_spec(args) + + +def test_island_refuses_overlap_without_the_raised_policy_age(): + from yeto.rl.adapters.miles.island_entry import _dense_full_sync_config + + args = SimpleNamespace(sweep_overlap=True, rl_max_policy_age=0) + with pytest.raises(ValueError, match="--sweep-overlap .* needs --rl-max-policy-age >= 1"): + _dense_full_sync_config(args, model_config_sha256="0" * 64, + training_contract_sha256="0" * 64, algorithm_spec_sha256="0" * 64) + + +def test_overlap_is_bound_into_the_role_stream_contract(): + from dataclasses import replace + + from test_rl_full_parameter_sync import _role_stream_config + + from yeto.rl.adapters.miles.full_parameter.role_stream import ( + role_stream_session_contract_hash, + ) + + exact = _role_stream_config("actor", 30100, sweep=True) + assert (role_stream_session_contract_hash("1" * 64, exact) + != role_stream_session_contract_hash("1" * 64, replace(exact, overlap=True))) diff --git a/tests/test_rl_dense_sweep_real_syncer.py b/tests/test_rl_dense_sweep_real_syncer.py index 5d01ad19..31589d67 100644 --- a/tests/test_rl_dense_sweep_real_syncer.py +++ b/tests/test_rl_dense_sweep_real_syncer.py @@ -21,7 +21,13 @@ from test_rl_algorithm_provenance import _fake_sky, _launcher_args, _no_modal_listing # noqa: F401 from yeto.fragments import MERGE_RDA, Fragment, FragmentLayout -from yeto.protocol import DTYPE_F32, DTYPE_F32_PUSH_BF16, SyncerClient +from yeto.protocol import ( + DTYPE_F32, + DTYPE_F32_PUSH_BF16, + DTYPE_F32_PUSH_I8, + DTYPE_FLAG_BCAST_DELTA, + SyncerClient, +) ROOT = Path(__file__).resolve().parent.parent CARGO = shutil.which("cargo") or str(Path.home() / ".cargo/bin/cargo") @@ -78,10 +84,11 @@ def _layout() -> FragmentLayout: return FragmentLayout([Fragment(MERGE_RDA, [(f"actor.w{i}", NUMEL)]) for i in range(FRAGMENTS)]) -def _client(port: int, profile_hash: str, dtype: int = DTYPE_F32) -> SyncerClient: +def _client(port: int, profile_hash: str, dtype: int = DTYPE_F32, *, + max_reconnects: int = 0, push_i8_block: int = 128) -> SyncerClient: return SyncerClient(("127.0.0.1", port), 0, _layout(), dtype=dtype, num_streams=1, - connect_timeout=10, max_reconnects=0, finalization_timeout=20, - syncer_profile_hash=bytes.fromhex(profile_hash)) + connect_timeout=10, max_reconnects=max_reconnects, finalization_timeout=20, + syncer_profile_hash=bytes.fromhex(profile_hash), push_i8_block=push_i8_block) def _f32(values: list[float]) -> bytes: @@ -161,10 +168,18 @@ def test_launcher_dense_syncer_refuses_another_profile_hash(syncer_binary, tmp_p def _run_sweeps(binary: Path, tmp_path: Path, extra: tuple[str, ...], sweeps: int, - dtype: int = DTYPE_F32): - """Drive ``sweeps`` strict sweeps; return (max PULLs seen at once, finals, checkpoint).""" + dtype: int = DTYPE_F32, *, after_sweep=None, max_reconnects: int = 0, + push_i8_block: int = 128, updates: dict | None = None, sink=None): + """Drive ``sweeps`` strict sweeps; return (max PULLs seen at once, finals, checkpoint). + + ``updates`` collects every BCAST the client accepted, as + ``{(fragment_id, version): f32 bytes}``; ``after_sweep(client, sweep)`` + runs once all pushes of a sweep are queued.""" process, port, profile_hash = _start_syncer(binary, tmp_path, *extra, total_steps=sweeps) - client = _client(port, profile_hash, dtype) + client = _client(port, profile_hash, dtype, max_reconnects=max_reconnects, + push_i8_block=push_i8_block) + if sink is not None: + client.install_inbound_chunk_sink(sink) concurrent = 0 try: client.start() @@ -195,7 +210,11 @@ def gather(): client.push_fragment(fid, permit.global_step, permit.round_attempt, base, sweep, 1, 7, _f32(delta)) permit.pushed = True + if after_sweep is not None: + after_sweep(client, sweep) manifest, finals = client.wait_for_final_fragments(timeout=30) + if updates is not None: + updates.update({(u.fragment_id, u.version): bytes(u.data) for u in client.drain_updates()}) result = {f.fragment_id: bytes(f.data) for f in finals} client.acknowledge_finalization(manifest) assert process.wait(timeout=15) == 0 @@ -270,3 +289,146 @@ def test_bf16_push_flag_reaches_the_island_command(): args = _dense_args("--rl-sync-push-dtype", "bf16") assert "--sweep-push-dtype bf16" in _rl_trainable_island_flags(args, 0, 1) assert "--sweep-push-dtype" not in _rl_trainable_island_flags(_dense_args(), 0, 1) + + +def _tail(data: bytes) -> torch.Tensor: + return torch.frombuffer(bytearray(data), dtype=torch.float32) + + +def test_int8_pushes_stay_within_their_block_error_of_the_f32_sweep(syncer_binary, tmp_path): + """E3: int8 outer deltas change results only by half a scale step per delta.""" + exact_dir, int8_dir = tmp_path / "f32", tmp_path / "int8" + exact_dir.mkdir() + int8_dir.mkdir() + exact = _run_sweeps(syncer_binary, exact_dir, STREAMED, sweeps=3) + quantized = _run_sweeps(syncer_binary, int8_dir, + (*STREAMED, "--rl-sync-push-dtype", "int8", "--rl-sync-push-int8-block", "32"), + sweeps=3, dtype=DTYPE_F32_PUSH_I8, push_i8_block=32) + assert exact[1] == quantized[1] == [5, 6] + for fid, data in exact[2].items(): + want, got = _tail(data)[-NUMEL:], _tail(quantized[2][fid])[-NUMEL:] + # Three sweeps; each delta's absmax < 2, so error <= 3 * 2 / 254. + assert torch.allclose(got, want, rtol=0, atol=3 * 2.0 / 254), fid + assert not torch.equal(got, want), "int8 pushes must actually quantize" + + +def _downlink_run(binary, tmp_path, full_every: int, sweeps: int, **kwargs): + updates: dict = {} + result = _run_sweeps(binary, tmp_path, + (*STREAMED, "--rl-sync-bcast-delta-full-every", str(full_every)), + sweeps=sweeps, dtype=DTYPE_F32 | DTYPE_FLAG_BCAST_DELTA, + updates=updates, **kwargs) + return result, updates + + +def test_bf16_downlink_tracks_the_exact_broadcasts_and_leaves_the_merge_exact(syncer_binary, tmp_path): + """E3: BCAST deltas change only what islands hold between f32 refreshes; + the syncer's masters, versions, FINAL and checkpoint stay bit-identical.""" + exact_dir, delta_dir = tmp_path / "f32", tmp_path / "delta" + exact_dir.mkdir() + delta_dir.mkdir() + exact_updates: dict = {} + exact = _run_sweeps(syncer_binary, exact_dir, STREAMED, sweeps=4, updates=exact_updates) + delta, delta_updates = _downlink_run(syncer_binary, delta_dir, 3, sweeps=4) + assert exact[1] == delta[1] == [7, 8] + assert exact[2] == delta[2], "FINAL must be the exact f32 masters" + assert exact[3] == delta[3], "checkpoints must not depend on the downlink" + assert set(exact_updates) == set(delta_updates) + # Fragment 0 publishes versions 1, 3, 5, 7: FULL, DELTA, DELTA, FULL (M=3). + kinds = {} + for key, data in sorted(delta_updates.items()): + want, got = _tail(exact_updates[key]), _tail(data) + kinds[key] = "full" if torch.equal(got, want) else "delta" + # one bf16 rounding of (theta - anchor), |theta - anchor| < 4 here + assert torch.allclose(got, want, rtol=0, atol=4.0 / 256), key + assert [kinds[(0, v)] for v in (0, 1, 3, 5, 7)] == ["full", "full", "delta", "delta", "full"] + assert [kinds[(1, v)] for v in (0, 2, 4, 6)] == ["full", "full", "delta", "delta"] + + +def test_bf16_downlink_anchor_mismatch_reconnects_and_resyncs_with_f32(syncer_binary, tmp_path): + """E3 fallback: a learner whose anchor drifted fails the digest check, + reconnects, receives the shared anchor as FULL, and the run finishes + with the same exact result.""" + exact_dir, delta_dir = tmp_path / "f32", tmp_path / "delta" + exact_dir.mkdir() + delta_dir.mkdir() + exact = _run_sweeps(syncer_binary, exact_dir, STREAMED, sweeps=4) + corrupted, reconnects = [], [] + + def corrupt(client, sweep): + if sweep == 4: + reconnects.append(client._reconnects_used) + if sweep != 1: + return + _wait(lambda: client._downlink.version(0) == 1, "fragment 0 version 1 never arrived") + held = client._downlink._held[0] + held.values = held.values.copy() + held.values[0] += 1e-3 + corrupted.append(client.connection_generation) + + delta, _ = _downlink_run(syncer_binary, delta_dir, 100, sweeps=4, after_sweep=corrupt, + max_reconnects=5) + assert corrupted and reconnects and reconnects[0] >= 1, "the digest mismatch must reconnect" + assert exact[1] == delta[1] and exact[2] == delta[2] and exact[3] == delta[3] + + +def test_bf16_downlink_and_int8_flags_reach_syncer_and_island_commands(): + from yeto import launcher + from yeto.launch.rl_island_task import _rl_trainable_island_flags + + args = _dense_args("--rl-sync-push-dtype", "int8", "--rl-sync-push-int8-block", "64", + "--rl-sync-bcast-delta-full-every", "4") + island = _rl_trainable_island_flags(args, 0, 1) + assert "--sweep-push-dtype int8" in island and "--sweep-push-int8-block 64" in island + command = launcher.role_syncer_commands(args, 1, binary="yeto-syncer")["actor"] + assert "--bcast-delta-full-every 4" in command + plain = _dense_args() + assert "--bcast-delta" not in launcher.role_syncer_commands(plain, 1, binary="yeto-syncer")["actor"] + with pytest.raises(ValueError, match="--rl-sync-bcast-delta-full-every 0"): + _dense_args("--rl-sync-bcast-delta-full-every", "1") + with pytest.raises(ValueError, match="--rl-sync-push-int8-block only with"): + _dense_args("--rl-sync-push-int8-block", "64") + + +class _BytesSink: + """Minimal transactional inbound sink: assembles the f32 fragment bytes.""" + + def begin_message(self, message_id, total_bytes): + return {"buf": bytearray(), "bound": None} + + def bind_fragment(self, tx, msg_type, fragment_id, version, payload_bytes): + tx["bound"] = payload_bytes + + def stage_chunk(self, tx, payload): + return bytes(payload) + + def consume_staged_chunk(self, tx, offset, token, nbytes): + self.consume_chunk(tx, offset, token) + + def consume_chunk(self, tx, offset, payload): + assert offset == len(tx["buf"]) + tx["buf"] += bytes(payload) + + def finish_message(self, tx): + import hashlib + + from yeto.protocol import StreamedInboundPayload + + assert len(tx["buf"]) == tx["bound"] + return StreamedInboundPayload(bytes(tx["buf"]), hashlib.sha256(tx["buf"]).digest()) + + def abort_message(self, tx, staged_tokens): + tx["buf"].clear() + + +def test_bf16_downlink_through_a_streaming_sink_matches_the_buffered_path(syncer_binary, tmp_path): + """E3: the Miles path streams BCAST chunks into a transactional sink; the + downlink wrapper must hand it the same f32 anchor bytes.""" + buffered_dir, streamed_dir = tmp_path / "buffered", tmp_path / "streamed" + buffered_dir.mkdir() + streamed_dir.mkdir() + buffered, buffered_updates = _downlink_run(syncer_binary, buffered_dir, 3, sweeps=3) + streamed, streamed_updates = _downlink_run(syncer_binary, streamed_dir, 3, sweeps=3, + sink=_BytesSink()) + assert buffered[2] == streamed[2] + assert buffered_updates == streamed_updates diff --git a/tests/test_run_flags.py b/tests/test_run_flags.py index 99b5ea01..9cdba9d4 100644 --- a/tests/test_run_flags.py +++ b/tests/test_run_flags.py @@ -250,7 +250,7 @@ def test_forward_tables_name_flags_of_both_layers(): # flag renamed in the schema but still spelled the old way in the launcher fails. _NON_RL_LITERALS = frozenset({ "--abbrev-ref", "--adapter-dir", "--add", "--address", "--allow-unattested-legacy-adapter", - "--allow-unauthenticated-islands", "--async-checkpoint", "--bind", "--cap-add", "--carry-gamma", + "--allow-unauthenticated-islands", "--async-checkpoint", "--bcast-delta-full-every", "--bind", "--cap-add", "--carry-gamma", "--checkpoint-every", "--checkpoint-path", "--codex-harness-binary", "--config-json", "--create", "--dashboard-host", "--dearmor", "--depth", "--detach", "--dtype", "--entrypoint", "--expected-island-contract", "--final-ack-timeout-s", "--final-grace-s", "--follow", "--force", diff --git a/yeto/downlink.py b/yeto/downlink.py new file mode 100644 index 00000000..d3a613c0 --- /dev/null +++ b/yeto/downlink.py @@ -0,0 +1,235 @@ +"""Learner side of bf16 downlink deltas (S26 E3; syncer ``src/downlink.rs``). + +A session whose dtype carries ``DTYPE_FLAG_BCAST_DELTA`` receives every +BCAST_FRAGMENT with a 41-byte downlink header after ``fragment_id, version``:: + + kind u8 (0 FULL, 1 DELTA) | base_version u64 | anchor_sha256 [32] + +FULL carries f32 values; DELTA carries ``bf16(theta - anchor)``. The client +keeps one f32 *anchor* per fragment, bit-identical to the syncer's: FULL sets +it, DELTA adds the decoded delta in f32 (numpy float32 addition rounds like +Rust's). The digest is checked after every update. Callers above the client +(``drain_updates``, inbound chunk sinks) only ever see the reconstructed f32 +anchor, so adapters, N7.6 cuts and manifests are unchanged. + +A missing or mismatched anchor raises :class:`DownlinkResync` (a +``ConnectionError``): the client drops its connection group, and the +syncer's reconnect recovery re-sends every fragment as FULL. +""" + +from __future__ import annotations + +import hashlib +import struct +import threading +from dataclasses import dataclass, field +from typing import TYPE_CHECKING + +if TYPE_CHECKING: # numpy is imported lazily: yeto.protocol stays import-light + import numpy as np + +KIND_FULL = 0 +KIND_DELTA = 1 +HEADER = struct.Struct(" int: + if kind == KIND_FULL: + return 4 * numel + if kind == KIND_DELTA: + return 2 * numel + raise ValueError(f"unknown downlink kind {kind}") + + +@dataclass +class _Held: + version: int + values: np.ndarray # float32, never mutated after commit + + +class DownlinkDecode: + """Streams one BCAST's tensor bytes into the next f32 anchor.""" + + def __init__(self, fragment_id: int, version: int, kind: int, base: np.ndarray | None, + digest: bytes, numel: int) -> None: + import numpy as np + + self.fragment_id = fragment_id + self.version = version + self.kind = kind + self._base = base + self._digest = digest + self.values = np.empty(numel, dtype=np.float32) + self._width = 4 if kind == KIND_FULL else 2 + self._carry = b"" + self._filled = 0 + self._hasher = hashlib.sha256() + + def feed(self, data: bytes | bytearray | memoryview) -> bytes: + """Decode the next wire bytes; return the f32 bytes they complete.""" + import numpy as np + + raw = self._carry + bytes(data) + whole = len(raw) - len(raw) % self._width + self._carry = raw[whole:] + if not whole: + return b"" + count = whole // self._width + start, end = self._filled, self._filled + count + if end > self.values.size: + raise ValueError("downlink BCAST has more values than its fragment") + out = self.values[start:end] + if self.kind == KIND_FULL: + out[:] = np.frombuffer(raw, dtype=" None: + if self._carry or self._filled != self.values.size: + raise ValueError("downlink BCAST ended before its fragment was complete") + if self._hasher.digest() != self._digest: + raise DownlinkResync( + f"downlink anchor digest mismatch for fragment {self.fragment_id} " + f"version {self.version}") + + +class DownlinkAnchors: + """Per-fragment f32 anchors shared (bit for bit) with the syncer.""" + + def __init__(self, num_fragments: int) -> None: + self._held: list[_Held | None] = [None] * num_fragments + self._lock = threading.Lock() + + def version(self, fragment_id: int) -> int | None: + with self._lock: + held = self._held[fragment_id] + return None if held is None else held.version + + def begin(self, fragment_id: int, version: int, header: bytes, numel: int) -> DownlinkDecode: + kind, base_version, digest = HEADER.unpack(header) + if kind == KIND_FULL: + return DownlinkDecode(fragment_id, version, kind, None, digest, numel) + if kind != KIND_DELTA: + raise ValueError(f"unknown downlink kind {kind}") + with self._lock: + held = self._held[fragment_id] + if held is None or held.version != base_version or held.values.size != numel: + raise DownlinkResync( + f"downlink delta for fragment {fragment_id} version {version} needs anchor " + f"version {base_version}; learner holds " + f"{None if held is None else held.version}") + return DownlinkDecode(fragment_id, version, kind, held.values, digest, numel) + + def commit(self, decode: DownlinkDecode) -> None: + """Adopt a finished decode as the anchor unless a newer one is held.""" + with self._lock: + held = self._held[decode.fragment_id] + if held is None or decode.version >= held.version: + self._held[decode.fragment_id] = _Held(decode.version, decode.values) + + def decode_whole(self, fragment_id: int, version: int, payload: bytes | memoryview, + numel: int) -> bytes: + """Non-streamed BCAST: header + tensor -> committed f32 anchor bytes.""" + if len(payload) < HEADER_LEN: + raise ValueError("downlink BCAST header truncated") + decode = self.begin(fragment_id, version, bytes(payload[:HEADER_LEN]), numel) + if len(payload) - HEADER_LEN != tensor_nbytes(decode.kind, numel): + raise ValueError(f"downlink BCAST {fragment_id} has {len(payload) - HEADER_LEN} " + f"tensor bytes, expected {tensor_nbytes(decode.kind, numel)}") + data = decode.feed(payload[HEADER_LEN:]) + decode.finish() + self.commit(decode) + return data + + +@dataclass +class _Transaction: + inner: object + passthrough: bool = True + fragment_id: int | None = None + version: int | None = None + numel: int = 0 + header: bytearray = field(default_factory=bytearray) + decode: DownlinkDecode | None = None + out_offset: int = 0 + + +class DownlinkChunkSink: + """Wraps a transactional inbound chunk sink for a downlink session. + + BCAST payload bytes (downlink header + f32/bf16 tensor) are decoded in + order and handed to the wrapped sink as the plain f32 fragment; FINAL + passes through. Out-of-order chunks are staged here as bytes (the + wrapped sink sees only in-order f32 bytes).""" + + def __init__(self, inner: object, anchors: DownlinkAnchors, numels: list[int], + bcast_type: int) -> None: + self.inner = inner + self.anchors = anchors + self.numels = numels + self.bcast_type = bcast_type + + def begin_message(self, message_id: int, total_bytes: int) -> _Transaction: + return _Transaction(self.inner.begin_message(message_id, total_bytes)) + + def bind_fragment(self, tx: _Transaction, msg_type: int, fragment_id: int, version: int, + payload_bytes: int) -> None: + if msg_type != self.bcast_type: + self.inner.bind_fragment(tx.inner, msg_type, fragment_id, version, payload_bytes) + return + tx.passthrough = False + tx.fragment_id, tx.version = fragment_id, version + tx.numel = self.numels[fragment_id] + self.inner.bind_fragment(tx.inner, msg_type, fragment_id, version, 4 * tx.numel) + + def stage_chunk(self, tx: _Transaction, payload) -> bytes: + del tx + return bytes(payload) + + def consume_staged_chunk(self, tx: _Transaction, offset: int, token: bytes, nbytes: int) -> None: + if len(token) != nbytes: + raise ValueError("staged downlink chunk is malformed") + self.consume_chunk(tx, offset, token) + + def consume_chunk(self, tx: _Transaction, offset: int, payload) -> None: + if tx.passthrough: + self.inner.consume_chunk(tx.inner, offset, payload) + return + data = memoryview(payload).cast("B") + if tx.decode is None: + take = min(HEADER_LEN - len(tx.header), data.nbytes) + tx.header += data[:take] + data = data[take:] + if len(tx.header) < HEADER_LEN: + return + tx.decode = self.anchors.begin(tx.fragment_id, tx.version, bytes(tx.header), tx.numel) + out = tx.decode.feed(data) + if out: + self.inner.consume_chunk(tx.inner, tx.out_offset, out) + tx.out_offset += len(out) + + def finish_message(self, tx: _Transaction): + if tx.passthrough: + return self.inner.finish_message(tx.inner) + if tx.decode is None: + raise ValueError("downlink BCAST ended inside its header") + tx.decode.finish() + committed = self.inner.finish_message(tx.inner) + self.anchors.commit(tx.decode) + return committed + + def abort_message(self, tx: _Transaction, staged_tokens) -> None: + del staged_tokens # plain bytes held here + self.inner.abort_message(tx.inner, ()) diff --git a/yeto/launch/rl_args.py b/yeto/launch/rl_args.py index df0c0efb..60e63f88 100644 --- a/yeto/launch/rl_args.py +++ b/yeto/launch/rl_args.py @@ -153,16 +153,35 @@ def sao_contract_mounts(args, learner_id: int) -> dict[str, str]: ("rl_sync_async_checkpoint", False, "--rl-sync-async-checkpoint"), ("rl_sync_push_dtype", "f32", "--rl-sync-push-dtype"), ("rl_sync_checkpoint_every", 1, "--rl-sync-checkpoint-every"), + ("rl_sync_push_int8_block", 128, "--rl-sync-push-int8-block"), + ("rl_sync_overlap", False, "--rl-sync-overlap"), ) def _refuse_full_parameter_sync_options(args) -> None: - """The S26 E3 sync options exist only for the full-parameter preset.""" + """The S26 E3 sync options exist only for the full-parameter preset + (the bf16 downlink also for strict-avg, see ``_check_bcast_delta``).""" for dest, default, option in _FULL_PARAMETER_SYNC_DEFAULTS: if getattr(args, dest, default) != default: raise ValueError(f"{option} only applies to --rl-sync-preset {FULL_PARAMETER_SYNC_PRESET}") +def _check_bcast_delta(args) -> None: + """S26 E3 bf16 downlink: a legacy-scheduled strict-avg or full-parameter + syncer only (the elastic server and SAO have no downlink anchors).""" + full_every = int(getattr(args, "rl_sync_bcast_delta_full_every", 0) or 0) + if full_every == 0: + return + if full_every < 2: + raise ValueError("--rl-sync-bcast-delta-full-every 0 (off) or >= 2") + if args.rl_sync_preset not in ("strict-avg", FULL_PARAMETER_SYNC_PRESET): + raise ValueError("--rl-sync-bcast-delta-full-every only applies to --rl-sync-preset " + f"strict-avg or {FULL_PARAMETER_SYNC_PRESET}") + if island_scheduling_mode(args) != "legacy" or args.experimental_rl_sync: + raise ValueError("--rl-sync-bcast-delta-full-every needs --rl-island-scheduling legacy " + "and no --experimental-rl-sync") + + def _prepare_dense_full_sweep(args, *, num_islands: int) -> None: """full-parameter preset (P0-D8): one actor role swept as whole-model fragments. @@ -190,6 +209,11 @@ def _prepare_dense_full_sweep(args, *, num_islands: int) -> None: checkpoint_every = int(getattr(args, "rl_sync_checkpoint_every", 1)) if checkpoint_every != 1 and (checkpoint_every < 1 or checkpoint_every % max(args.fragments, 1)): problems.append("--rl-sync-checkpoint-every 1 or a multiple of --fragments") + int8_block = int(getattr(args, "rl_sync_push_int8_block", 128)) + if not 32 <= int8_block <= 1 << 20: + problems.append("--rl-sync-push-int8-block between 32 and 1048576") + elif int8_block != 128 and getattr(args, "rl_sync_push_dtype", "f32") != "int8": + problems.append("--rl-sync-push-int8-block only with --rl-sync-push-dtype int8") if problems: raise ValueError(f"--rl-sync-preset {FULL_PARAMETER_SYNC_PRESET} requires " + "; ".join(problems)) args.quorum = num_islands @@ -339,6 +363,15 @@ def _prepare_rl_args( # what the selected backend supports before any cloud or GPU work. from ..rl.engine.policy_age import validate_limit + if args.rl_sync_preset == FULL_PARAMETER_SYNC_PRESET and getattr(args, "rl_sync_overlap", False): + # S26 E3: sync/train overlap adds one outer step of staleness; raise the + # limit to the new minimum here so every check below (backend support, + # and the preflight's algorithm-spec staleness) validates it at launch. + from ..rl.engine.sync_overlap import required_policy_age + + required = required_policy_age(0) + if int(getattr(args, "rl_max_policy_age", 0) or 0) < required: + args.rl_max_policy_age = required max_policy_age = validate_limit(getattr(args, "rl_max_policy_age", 0) or 0) policy_age_module = _rl_backend_module(args, "policy_age") policy_age_module.SUPPORT.check(max_policy_age) @@ -748,6 +781,7 @@ def _prepare_rl_args( raise ValueError("--rl-sao-contracts only applies to --rl-sync-preset sao-streaming-full") if args.rl_sync_preset != FULL_PARAMETER_SYNC_PRESET: _refuse_full_parameter_sync_options(args) + _check_bcast_delta(args) if args.rl_sync_preset == "decoupled": if args.experimental_rl_sync: raise ValueError("decoupled RL does not accept --experimental-rl-sync") diff --git a/yeto/launch/rl_island_task.py b/yeto/launch/rl_island_task.py index a5cd883f..532fff20 100644 --- a/yeto/launch/rl_island_task.py +++ b/yeto/launch/rl_island_task.py @@ -74,6 +74,10 @@ def _rl_trainable_island_flags(args, learner_id: int, num_learners: int) -> str: f" --syncer-profile-sha256 {rl_syncer_profile(args, num_learners).sha256}") if getattr(args, "rl_sync_push_dtype", "f32") != "f32": sweep += f" --sweep-push-dtype {args.rl_sync_push_dtype}" + if getattr(args, "rl_sync_overlap", False): + sweep += " --sweep-overlap" + if getattr(args, "rl_sync_push_dtype", "f32") == "int8": + sweep += f" --sweep-push-int8-block {int(getattr(args, 'rl_sync_push_int8_block', 128))}" if getattr(args, "expert_full_count", 0): return lora + sweep return " --parameter-mode full" + sweep @@ -152,6 +156,7 @@ def make_miles_island_task( f" --rollout-num-gpus-per-engine {args.rollout_num_gpus_per_engine}" f" --sglang-mem-fraction-static {args.sglang_mem_fraction_static}" + _rl_trainable_island_flags(args, learner_id, num_learners) + + (" --bcast-delta" if int(getattr(args, "rl_sync_bcast_delta_full_every", 0) or 0) else "") + f" --inner-lr {args.inner_lr}" f" --seq-len {args.seq_len}" + (" --rl-activation-recompute" if getattr(args, "rl_activation_recompute", False) else "") diff --git a/yeto/launch/syncer.py b/yeto/launch/syncer.py index 91587c07..2836e1b0 100644 --- a/yeto/launch/syncer.py +++ b/yeto/launch/syncer.py @@ -150,6 +150,12 @@ def _syncer_checkpoint_every(args) -> int: return int(getattr(args, "rl_sync_checkpoint_every", 1)) if rl_dense_sweep(args) else 1 +def _syncer_bcast_delta(args) -> str: + """S26 E3 bf16 downlink deltas (checked by rl_args._check_bcast_delta).""" + full_every = int(getattr(args, "rl_sync_bcast_delta_full_every", 0) or 0) + return f" --bcast-delta-full-every {full_every}" if full_every else "" + + def rl_dense_sweep(args) -> bool: return (getattr(args, "training_mode", "sft") == "rl" and getattr(args, "rl_sync_preset", None) == FULL_PARAMETER_SYNC_PRESET) @@ -448,6 +454,7 @@ def _rl_role_syncer_command(args, num_learners: int, binary: str, f"{_syncer_island_scheduling(args)}" f"{_syncer_expected_contract(args)}" f"{_syncer_profile_binding(args)}" + f"{_syncer_bcast_delta(args)}" f" --outer-lr {args.outer_lr}" f" --outer-momentum {args.outer_momentum}" " --max-base-lag 0 --learner-weight equal" diff --git a/yeto/protocol.py b/yeto/protocol.py index 28331475..fc2b6d73 100644 --- a/yeto/protocol.py +++ b/yeto/protocol.py @@ -37,6 +37,9 @@ from collections.abc import Callable, Iterable, Iterator from dataclasses import dataclass, field +from .downlink import HEADER_LEN as DOWNLINK_HEADER_LEN +from .downlink import DownlinkAnchors, DownlinkChunkSink +from .downlink import tensor_nbytes as downlink_tensor_nbytes from .fragments import MERGE_ISO, FragmentLayout MAGIC = 0xD170C0DE @@ -71,6 +74,15 @@ # delta bytes; the client rounds them (round-to-nearest-even, as the syncer's # ``half::bf16::from_f32``) on the way out. DTYPE_F32_PUSH_BF16 = 4 +# Session dtype 5 (S26 E3): PUSH_FRAGMENT deltas travel as block-scaled int8 +# (``f32_bytes_to_i8``); INIT/BCAST/FINAL stay f32. Callers hand f32 bytes. +DTYPE_F32_PUSH_I8 = 5 +# Session-dtype flag (S26 E3): BCAST_FRAGMENT carries bf16 downlink deltas +# between f32 refreshes (yeto/downlink.py). Only on f32-bulk sessions. +DTYPE_FLAG_BCAST_DELTA = 0x80 +I8_BLOCK_MIN = 32 +I8_BLOCK_MAX = 1 << 20 +DEFAULT_I8_BLOCK = 128 CHUNK_SIZE = 4 * 1024 * 1024 MAX_HANDSHAKE_FRAME = 16 * 1024 * 1024 @@ -92,22 +104,103 @@ _BUDGET_DONE = struct.Struct(" int: + """The session dtype without its flags.""" + return dtype & ~DTYPE_FLAG_BCAST_DELTA + + +def bcast_delta(dtype: int) -> bool: + """Whether the session receives bf16 downlink deltas.""" + return bool(dtype & DTYPE_FLAG_BCAST_DELTA) + + +def valid_session_dtype(dtype: int) -> bool: + base = session_base(dtype) + if base in (DTYPE_F32, DTYPE_F32_PUSH_BF16, DTYPE_F32_PUSH_I8): + return 0 <= dtype <= 0xFF + return base in (DTYPE_BF16, DTYPE_Q4) and not bcast_delta(dtype) + + def bulk_dtype(dtype: int) -> int: """Dtype of INIT_PARAMS/BCAST_FRAGMENT tensors for a session dtype. Q4 applies only to push deltas (small dynamic range); full parameter - payloads would not survive 4 bits, so they travel as bf16. + payloads would not survive 4 bits, so they travel as bf16. A downlink + session's BCAST is reconstructed to f32 before callers see it. """ - if dtype == DTYPE_Q4: + base = session_base(dtype) + if base == DTYPE_Q4: return DTYPE_BF16 - if dtype == DTYPE_F32_PUSH_BF16: + if base in (DTYPE_F32_PUSH_BF16, DTYPE_F32_PUSH_I8): return DTYPE_F32 - return dtype + return base def push_wire_dtype(dtype: int) -> int: """Dtype of unquantized PUSH_FRAGMENT deltas on the wire for a session dtype.""" - return DTYPE_BF16 if dtype == DTYPE_F32_PUSH_BF16 else dtype + base = session_base(dtype) + return DTYPE_BF16 if base == DTYPE_F32_PUSH_BF16 else base + + +def i8_nbytes(numel: int, block: int) -> int: + """Encoded size of an int8 push: u32 block, then per block f32 scale + i8s.""" + if not I8_BLOCK_MIN <= block <= I8_BLOCK_MAX: + raise ValueError(f"int8 block size {block} outside {I8_BLOCK_MIN}..{I8_BLOCK_MAX}") + return 4 + 4 * -(-numel // block) + numel + + +def f32_bytes_to_i8(data: bytes | bytearray | memoryview, block: int, *, header: bool = True) -> bytes: + """Quantize little-endian f32 values to block-scaled int8 (syncer ``decode_i8``). + + Per block of ``block`` values (the last may be shorter): f32 scale = + absmax / 127, then ``round_half_even(x / scale)`` clipped to [-127, 127]. + Decoding ``q * scale`` errs by at most absmax / 254 per element (plus f32 + rounding). Blocks must start at multiples of ``block`` values; with + ``header`` the u32 block size is prepended. + """ + import numpy as np + + i8_nbytes(0, block) # validates the block size + values = np.frombuffer(data, dtype=" None: # chunk: (blocks, width) + absmax = np.abs(chunk).max(axis=1) + scale = (absmax / np.float32(127)).astype(np.float32) + safe = np.where(scale > 0, scale, np.float32(1)) + q = np.clip(np.rint(chunk / safe[:, None]), -127, 127).astype(np.int8) + q[scale == 0] = 0 + rows = np.empty((chunk.shape[0], 4 + chunk.shape[1]), dtype=np.uint8) + rows[:, :4] = scale.astype(" full: + encode(values[full:].reshape(1, -1)) + return b"".join(parts) + + +def _f32_parts_to_i8(parts: Iterable[bytes | bytearray | memoryview], block: int) -> Iterator[bytes]: + """Stream ``f32_bytes_to_i8`` over parts; blocks may span part boundaries.""" + yield struct.pack(" bytes: @@ -206,8 +299,8 @@ def seal_hello(body: bytes, key: bytes | None) -> bytes: The legacy syncer with a key refuses a HELLO without this trailer.""" if not key: return body - import hmac import hashlib + import hmac return body + hmac.new(key, HELLO_MAC_DOMAIN + body, hashlib.sha256).digest() @@ -515,11 +608,13 @@ def __init__( syncer_profile_hash: bytes | None = None, hmac_key: bytes | None = None, compat_group: str | None = None, + push_i8_block: int = DEFAULT_I8_BLOCK, ): if not 0 <= learner_id <= 0xFFFF_FFFF: raise ValueError(f"learner_id must fit u32, got {learner_id}") - if dtype not in (DTYPE_F32, DTYPE_BF16, DTYPE_Q4, DTYPE_F32_PUSH_BF16): + if not valid_session_dtype(dtype): raise ValueError(f"unsupported session dtype {dtype}") + i8_nbytes(0, push_i8_block) # validates the int8 block size if layout.num_fragments < 1: raise ValueError("layout must contain at least one fragment") if not 0 <= num_streams <= 256: @@ -537,6 +632,11 @@ def __init__( self.learner_id = learner_id self.layout = layout self.dtype = dtype + self._push_base = session_base(dtype) + self.push_i8_block = push_i8_block + # S26 E3: f32 anchors of a bf16-downlink session (yeto/downlink.py). + self._downlink = DownlinkAnchors(layout.num_fragments) if bcast_delta(dtype) else None + self._bcast_prefix = 12 + (DOWNLINK_HEADER_LEN if self._downlink is not None else 0) self.num_streams = num_streams self.session_contract_hash = session_contract_hash self.syncer_profile_hash = syncer_profile_hash @@ -548,7 +648,7 @@ def __init__( self.max_reconnects = max_reconnects self.finalization_timeout = finalization_timeout self._max_bcast_payload = max( - 12 + _tensor_nbytes(bulk_dtype(dtype), fragment.numel) + self._bcast_prefix + _tensor_nbytes(bulk_dtype(dtype), fragment.numel) for fragment in layout.fragments ) self._max_final_payload = max( @@ -870,9 +970,9 @@ def push_fragment( numel = self.layout.fragments[fragment_id].numel expected = ( _q4_nbytes(numel) - if self.dtype == DTYPE_Q4 - # dtype 4 callers pass f32 bytes; the client rounds them below - else _tensor_nbytes(DTYPE_F32 if self.dtype == DTYPE_F32_PUSH_BF16 else self.dtype, numel) + if self._push_base == DTYPE_Q4 + # dtypes 4/5 callers pass f32 bytes; the client encodes them below + else _tensor_nbytes(bulk_dtype(self.dtype), numel) ) if len(tensor_bytes) != expected: raise ValueError( @@ -890,8 +990,10 @@ def push_fragment( c_steps, c_tokens, ) - if self.dtype == DTYPE_F32_PUSH_BF16: + if self._push_base == DTYPE_F32_PUSH_BF16: tensor_bytes = f32_bytes_to_bf16(tensor_bytes) + elif self._push_base == DTYPE_F32_PUSH_I8: + tensor_bytes = f32_bytes_to_i8(tensor_bytes, self.push_i8_block) self._send_large(MSG_PUSH_FRAGMENT, head + tensor_bytes) def push_fragment_parts( @@ -923,12 +1025,14 @@ def push_fragment_parts( if c_steps < 1: raise ValueError("c_steps must be positive") numel = self.layout.fragments[fragment_id].numel - expected = ( - _q4_nbytes(numel) - if self.dtype == DTYPE_Q4 - else _tensor_nbytes(push_wire_dtype(self.dtype), numel) - ) - if self.dtype == DTYPE_F32_PUSH_BF16: + if self._push_base == DTYPE_Q4: + expected = _q4_nbytes(numel) + elif self._push_base == DTYPE_F32_PUSH_I8: + expected = i8_nbytes(numel, self.push_i8_block) + tensor_parts = _f32_parts_to_i8(tensor_parts, self.push_i8_block) + else: + expected = _tensor_nbytes(push_wire_dtype(self.dtype), numel) + if self._push_base == DTYPE_F32_PUSH_BF16: tensor_parts = _f32_parts_to_bf16(tensor_parts) head = struct.pack( " None: raise TypeError( "inbound chunk sink does not implement the transaction API" ) + if sink is not None and self._downlink is not None: + sink = DownlinkChunkSink( + sink, self._downlink, [f.numel for f in self.layout.fragments], MSG_BCAST_FRAGMENT) with self._reasm_lock: if self._reasm: raise RuntimeError("cannot replace inbound chunk sink mid-message") @@ -1542,6 +1649,13 @@ def _stream_chunk_locked( dtype, self.layout.fragments[fragment_id].numel, ) + if msg_type == MSG_BCAST_FRAGMENT and self._downlink is not None: + # The downlink header (kind, base version, digest) is part of + # the payload the wrapping sink decodes; its kind fixes the size. + if len(data) < tensor_start + DOWNLINK_HEADER_LEN: + raise ValueError("first downlink fragment chunk is header-truncated") + expected_payload = DOWNLINK_HEADER_LEN + downlink_tensor_nbytes( + data[tensor_start], self.layout.fragments[fragment_id].numel) if length != 12 + expected_payload: raise ValueError( f"streamed fragment {fragment_id} has {length - 12} tensor " @@ -1810,15 +1924,19 @@ def _dispatch_live( fid, version = struct.unpack_from("= self.layout.num_fragments: raise ValueError(f"BCAST_FRAGMENT for unknown fragment {fid}") - expected = _tensor_nbytes( - bulk_dtype(self.dtype), self.layout.fragments[fid].numel - ) - if len(payload) - 12 != expected: - raise ValueError( - f"BCAST_FRAGMENT {fid} has {len(payload) - 12} tensor bytes, " - f"expected {expected}" + if self._downlink is not None: + data = self._downlink.decode_whole( + fid, version, payload[12:], self.layout.fragments[fid].numel) + else: + expected = _tensor_nbytes( + bulk_dtype(self.dtype), self.layout.fragments[fid].numel ) - data = payload[12:] + if len(payload) - 12 != expected: + raise ValueError( + f"BCAST_FRAGMENT {fid} has {len(payload) - 12} tensor bytes, " + f"expected {expected}" + ) + data = payload[12:] digest = hashlib.sha256(data).digest() with self._bcast_lock: seen = self._bcast_seen[fid] diff --git a/yeto/rl/adapters/miles/full_parameter/role_stream.py b/yeto/rl/adapters/miles/full_parameter/role_stream.py index 7ea7d475..4b2a5199 100644 --- a/yeto/rl/adapters/miles/full_parameter/role_stream.py +++ b/yeto/rl/adapters/miles/full_parameter/role_stream.py @@ -20,19 +20,28 @@ from typing import Any from yeto.protocol import ( + DEFAULT_I8_BLOCK, DTYPE_F32, DTYPE_F32_PUSH_BF16, + DTYPE_F32_PUSH_I8, + DTYPE_FLAG_BCAST_DELTA, + I8_BLOCK_MAX, + I8_BLOCK_MIN, FinalManifest, PullRequest, SyncerClient, ) +from yeto.rl.adapters.miles.full_parameter.inbound_sink import MilesRayInboundChunkSink from yeto.rl.adapters.miles.models.chunked_full_parameter import ( MilesChunkedFullParameterAdapter, ReferencedPolicyCut, StoredAuthoritativeFragment, ) -from yeto.rl.adapters.miles.full_parameter.inbound_sink import MilesRayInboundChunkSink -from yeto.rl.engine.policy_digest import PolicyDigest, TrainerResidentState, WeightsChanged +from yeto.rl.engine.policy_digest import ( + PolicyDigest, + TrainerResidentState, + WeightsChanged, +) from yeto.rl.engine.trainable_state import LAYOUT_FRAGMENTS, UnsupportedLayoutError from yeto.rl.local_learner import ComponentIdentity, layout_stream_role from yeto.rl.trajectory_evidence import TrajectoryBatchEvidence @@ -83,8 +92,27 @@ class RoleStreamConfig: poll_seconds: float = 0.01 max_fragment_bytes: int = 2 << 30 max_chunk_bytes: int = 256 << 20 - # S26 E3: push outer deltas as bf16 (session dtype 4); BCAST/FINAL stay f32. - push_bf16: bool = False + # S26 E3 wire options (default exact f32): push outer deltas as bf16 + # (session dtype 4) or block-scaled int8 (dtype 5); receive BCAST as bf16 + # deltas against a shared f32 anchor. FINAL always stays f32. + push_dtype: str = "f32" + push_int8_block: int = DEFAULT_I8_BLOCK + bcast_delta: bool = False + # S26 E3: train round k+1 while sweep k syncs, then rebase the trainer + # (see FullParameterSync._overlap_boundary); one more outer step stale. + overlap: bool = False + + @property + def session_dtype(self) -> int: + base = {"f32": DTYPE_F32, "bf16": DTYPE_F32_PUSH_BF16, "int8": DTYPE_F32_PUSH_I8}[self.push_dtype] + return base | (DTYPE_FLAG_BCAST_DELTA if self.bcast_delta else 0) + + @property + def wire_dtype_label(self) -> str: + """Session-contract label; plain f32 keeps the historical ``fp32``.""" + label = {"f32": "fp32", "bf16": "fp32-push-bf16", + "int8": f"fp32-push-int8-b{self.push_int8_block}"}[self.push_dtype] + return label + ("+bcast-bf16-delta" if self.bcast_delta else "") def __post_init__(self) -> None: if not self.algorithm.isidentifier(): @@ -129,6 +157,10 @@ def __post_init__(self) -> None: value = getattr(self, name) if type(value) is not int or value < 1: raise ValueError(f"role-stream {name} must be positive") + if self.push_dtype not in ("f32", "bf16", "int8"): + raise ValueError("role-stream push dtype must be f32, bf16 or int8") + if type(self.push_int8_block) is not int or not I8_BLOCK_MIN <= self.push_int8_block <= I8_BLOCK_MAX: + raise ValueError("role-stream int8 block size is invalid") if self.wait_timeout <= 0 or not 0 < self.poll_seconds <= 1: raise ValueError("role-stream wait bounds are invalid") if ( @@ -222,7 +254,8 @@ def role_stream_session_contract_hash( "pipeline_depth": config.pipeline_depth, "local_horizon": config.local_horizon, "optimizer_steps_per_round": config.optimizer_steps_per_round, - "wire_dtype": "fp32-push-bf16" if config.push_bf16 else "fp32", + "wire_dtype": config.wire_dtype_label, + **({"sync_train_overlap": True} if config.overlap else {}), }, } encoded = json.dumps(payload, sort_keys=True, separators=(",", ":")).encode() @@ -264,6 +297,11 @@ def __init__(self, config: RoleStreamConfig, *, ray_module=None) -> None: self._staged: dict[tuple[int, int], StoredAuthoritativeFragment] = {} self._pinned_cuts: dict[int, tuple[ReferencedPolicyCut, int]] = {} self._pinned_anchors: dict[int, tuple[StoredAuthoritativeFragment, int]] = {} + # Overlap: the cut the current local round started from, and the + # round/token counters at that point (the next push's c_steps base). + self._round_base: ReferencedPolicyCut | None = None + self._rounds_at_base = 0 + self._tokens_at_base = 0 self.terminal_manifest: FinalManifest | None = None self.terminal_submitted = False self.final_applied = False @@ -337,7 +375,8 @@ async def initialize(self, group) -> None: self.config.syncer_addr, self.config.learner_id, adapter.layout.fragments, - dtype=DTYPE_F32_PUSH_BF16 if self.config.push_bf16 else DTYPE_F32, + dtype=self.config.session_dtype, + push_i8_block=self.config.push_int8_block, num_streams=self.config.wan_streams, max_reconnects=None, session_contract_hash=role_stream_session_contract_hash( @@ -474,7 +513,14 @@ async def complete_local_round( *, rollout_id: int, trajectories: TrajectoryBatchEvidence, + hold_submission: bool = False, ) -> RoleRoundResult: + """Record one trained round and capture its cut. + + ``hold_submission`` (overlap): the previous sweep is still in flight, + so the new cut is not yet submittable; it becomes the submission only + after :meth:`apply_staged_broadcast` rebased it and + :meth:`finish_rebase`. The round base stays pinned until then.""" if ( self.adapter is None or self.current is None @@ -503,14 +549,19 @@ async def complete_local_round( previous = self.current previous_submission = self.submission_cut self.current = local - self.submission_cut = local - self.submission_fragment_versions = self.fragment_versions self.local_generation = expected_generation self.local_rounds += 1 self.action_tokens += trajectories.trained_tokens - self.submission_rounds = self.local_rounds - self.submission_tokens = self.action_tokens - self._release_cuts(previous, previous_submission, keep=(local,)) + if hold_submission: + if not self.config.overlap or self._round_base is None: + raise RuntimeError(f"{self.config.role} holds a submission without an overlap round base") + self._release_cuts(previous, keep=(local, self._round_base, previous_submission)) + else: + self.submission_cut = local + self.submission_fragment_versions = self.fragment_versions + self.submission_rounds = self.local_rounds + self.submission_tokens = self.action_tokens + self._release_cuts(previous, previous_submission, keep=(local, self._round_base)) # The policy-level coordinator owns permit matching and release. A # role must never answer a PULL here: the peer role may not yet have @@ -562,8 +613,15 @@ async def apply_staged_broadcast( self, fragment_id: int, version: int, + *, + rebase: bool = False, ) -> tuple[int, ...]: - """Apply one coordinator-approved staged fragment at a safe boundary.""" + """Apply one coordinator-approved staged fragment at a safe boundary. + + ``rebase`` (overlap): the trainer keeps its progress since the round + base, ``theta_f <- G_f + (current_f - base_f)``; ``G_f`` becomes the + anchor, and the counters at the round base become the anchor's, so + the next push carries exactly that progress.""" key = (fragment_id, version) stored = self._staged.get(key) @@ -576,14 +634,56 @@ async def apply_staged_broadcast( raise RuntimeError( f"{self.config.algorithm} {self.config.role} staged fragment skipped its base version" ) - await self._apply_authoritative({fragment_id: stored}) + if not rebase: + await self._apply_authoritative({fragment_id: stored}) + else: + if not self.config.overlap or self._round_base is None or self.current is None: + raise RuntimeError(f"{self.config.role} rebase without an overlap round base") + rebased = self.adapter.rebase_fragment( + stored, self.current, self._round_base, ray_module=self.ray_module) + try: + await self._apply_authoritative( + {fragment_id: rebased}, + anchors={fragment_id: stored}, + counters=(self._rounds_at_base, self._tokens_at_base), + ) + finally: + rebased.release() # the applied target holds its own references self._staged.pop(key, None) return (fragment_id,) + def finish_rebase(self) -> None: + """Overlap: the rebased cut is the next submission (``G + progress``).""" + + if not self.config.overlap or self.current is None: + raise RuntimeError(f"{self.config.role} finish_rebase outside overlap") + previous_submission = self.submission_cut + self.submission_cut = self.current + self.submission_fragment_versions = self.fragment_versions + self.submission_rounds = self.local_rounds + self.submission_tokens = self.action_tokens + self._release_cuts(previous_submission, keep=(self.current, self._round_base)) + + def mark_round_base(self) -> None: + """Overlap: pin the cut the next local round starts from.""" + + if not self.config.overlap or self.current is None: + raise RuntimeError(f"{self.config.role} mark_round_base outside overlap") + previous = self._round_base + self._round_base = self.current + self._rounds_at_base = self.local_rounds + self._tokens_at_base = self.action_tokens + self._release_cuts(previous, keep=(self.current, self.submission_cut)) + async def _apply_authoritative( self, selected: Mapping[int, StoredAuthoritativeFragment], + *, + anchors: Mapping[int, StoredAuthoritativeFragment] | None = None, + counters: tuple[int, int] | None = None, ) -> None: + """Apply ``selected`` fragments; ``anchors``/``counters`` override the + recorded anchor and its round/token counters (overlap rebase).""" if self.adapter is None or self.current is None or self.group is None: raise RuntimeError(f"{self.config.algorithm} {self.config.role} stream is unavailable") target_version = self.transport_version + 1 @@ -613,16 +713,18 @@ async def _apply_authoritative( self.local_generation = 0 self._release_cuts( previous_cut, - keep=(self.current, self.submission_cut), + keep=(self.current, self.submission_cut, self._round_base), ) - for fragment_id, stored in selected.items(): + rounds, tokens = counters if counters is not None else (self.local_rounds, self.action_tokens) + for fragment_id, applied in selected.items(): + stored = applied if anchors is None else anchors[fragment_id] previous_anchor = self._anchors[fragment_id] if previous_anchor is not None and previous_anchor is not stored: self._release_anchor(previous_anchor) self._anchors[fragment_id] = stored self._fragment_versions[fragment_id] = stored.version - self._rounds_at_anchor[fragment_id] = self.local_rounds - self._tokens_at_anchor[fragment_id] = self.action_tokens + self._rounds_at_anchor[fragment_id] = rounds + self._tokens_at_anchor[fragment_id] = tokens def _apply_token( self, @@ -914,8 +1016,9 @@ def close(self) -> None: retained_cuts = tuple(value for value, _count in self._pinned_cuts.values()) self._pinned_cuts.clear() self._pinned_anchors.clear() - self._release_cuts(self.current, self.submission_cut, *retained_cuts) + self._release_cuts(self.current, self.submission_cut, self._round_base, *retained_cuts) self._staged.clear() + self._round_base = None self.current = None self.submission_cut = None self.submission_fragment_versions = None @@ -937,6 +1040,7 @@ def _unpin_cut(self, cut: ReferencedPolicyCut) -> None: remaining == 0 and cut is not self.current and cut is not self.submission_cut + and cut is not self._round_base and self.adapter is not None ): self.adapter.release(cut) diff --git a/yeto/rl/adapters/miles/island_entry.py b/yeto/rl/adapters/miles/island_entry.py index 5a9d8874..c382f8cb 100644 --- a/yeto/rl/adapters/miles/island_entry.py +++ b/yeto/rl/adapters/miles/island_entry.py @@ -814,8 +814,14 @@ def _dense_full_sync_config( from yeto.rl.adapters.miles.full_parameter.role_stream import RoleStreamConfig from yeto.rl.adapters.miles.full_parameter.session import FullParameterSyncConfig + from yeto.rl.engine.sync_overlap import required_policy_age from yeto.rl.local_learner import ComponentIdentity + overlap = bool(getattr(args, "sweep_overlap", False)) + if overlap and int(getattr(args, "rl_max_policy_age", 0) or 0) < required_policy_age(0): + raise ValueError("--sweep-overlap rollouts are one outer step stale: it needs " + f"--rl-max-policy-age >= {required_policy_age(0)}") + actor = RoleStreamConfig( algorithm="grpo", role="actor", @@ -838,7 +844,10 @@ def _dense_full_sync_config( pipeline_depth=1, algorithm_spec_sha256=algorithm_spec_sha256, wan_streams=args.wan_streams, - push_bf16=getattr(args, "sweep_push_dtype", "f32") == "bf16", + push_dtype=getattr(args, "sweep_push_dtype", "f32"), + push_int8_block=int(getattr(args, "sweep_push_int8_block", 128)), + bcast_delta=bool(getattr(args, "bcast_delta", False)), + overlap=overlap, ) return FullParameterSyncConfig((actor,), sweep=True) @@ -2099,6 +2108,7 @@ def run_miles( algorithm_semantic_sha256=algorithm_semantic_sha256, code_pack_digest=code_pack_digest, compat_group=_island_compat_group(args), + bcast_delta=bool(getattr(args, "bcast_delta", False)), ) if getattr(args, "rl_island_scheduling", "legacy") == "elastic": # rl-inter-island-scheduling 0.15 (legacy sets nothing) diff --git a/yeto/rl/adapters/miles/models/chunked_full_parameter.py b/yeto/rl/adapters/miles/models/chunked_full_parameter.py index 9b7e6c99..f16a6ec5 100644 --- a/yeto/rl/adapters/miles/models/chunked_full_parameter.py +++ b/yeto/rl/adapters/miles/models/chunked_full_parameter.py @@ -786,6 +786,60 @@ def assemble_mixed_target( ) return self.validate_cut(target, target_policy_version, 0) + def rebase_fragment( + self, + stored: StoredAuthoritativeFragment, + current: ReferencedPolicyCut, + base: ReferencedPolicyCut, + *, + ray_module=None, + ) -> StoredAuthoritativeFragment: + """Sync/train overlap: ``stored + (current - base)`` for one fragment. + + The result is stored like an inbound authoritative fragment (same + id, version and provenance, fresh Ray refs) so ``assemble_mixed_target`` + validates and applies it unchanged. The caller keeps ``stored`` as + the anchor and releases the result after the apply.""" + + import numpy as np + from miles.backends.megatron_utils.full_parameter_state import ( + _resolve_ray_module, + ) + + from yeto.rl.engine.sync_overlap import rebase_values + + fragment_id = stored.fragment_id + for cut in (current, base): + if cut.layout_hash != self.layout.layout_hash: + raise ValueError("Miles rebase cut layout changed") + owner = _fragment_owner_shard(current.transport_cut, fragment_id) + if ( + stored.parameter_layout_hash != self.layout.layout_hash + or stored.topology != owner.topology + or stored.plan_hash != owner.plan_hash + ): + raise ValueError("authoritative fragment provenance changed") + ray = _resolve_ray_module(ray_module) + chunks = sorted(zip(stored.descriptor.chunks, stored.refs, strict=True), + key=lambda row: row[0].flat_offset) + anchor_parts, offset = [], 0 + for descriptor, reference in chunks: + if reference is None or descriptor.flat_offset != offset: + raise ValueError("authoritative fragment chunks are released or not contiguous") + value = ray.get(reference) + anchor_parts.append(value.detach().to("cpu").reshape(-1).numpy()) + offset += descriptor.numel + + def flat(cut: ReferencedPolicyCut) -> np.ndarray: + raw = b"".join(bytes(part) for part in self.fragment_parts( + cut, fragment_id, ray_module=ray_module)()) + return np.frombuffer(raw, dtype="= 1)", + ), + island_flag( + "--bcast-delta", + action="store_true", + help="strict-avg / full-parameter sweep: accept bf16 downlink deltas against the " + "shared f32 anchor (the syncer runs --bcast-delta-full-every)", + ), island_flag( "--syncer-profile-sha256", default=None, diff --git a/yeto/rl/adapters/verl/island_entry.py b/yeto/rl/adapters/verl/island_entry.py index ab6f640f..5fec3a02 100644 --- a/yeto/rl/adapters/verl/island_entry.py +++ b/yeto/rl/adapters/verl/island_entry.py @@ -60,6 +60,8 @@ def parse_args(argv): p.add_argument("--seq-len", type=int, required=True) p.add_argument("--seed", type=int, required=True) p.add_argument("--wan-streams", type=int, default=4) + # S26 E3: bf16 downlink deltas (yeto/downlink.py; the syncer runs --bcast-delta-full-every) + p.add_argument("--bcast-delta", action="store_true") p.add_argument("--apply-chat-template-kwargs", default=None) p.add_argument("--rl-algorithm-spec", default=None) p.add_argument("--rl-island-scheduling", default="legacy") @@ -357,6 +359,7 @@ def main(argv=None) -> int: "model_path": model_path, "model_revision": args.model_revision, "lora_rank": args.lora_r, "lora_targets": args.lora_targets, "event_tape": os.path.expanduser(args.event_tape), "out_dir": str(OUTPUT), "tis_upper": algo["tis_upper"], "wan_streams": args.wan_streams, + "bcast_delta": bool(args.bcast_delta), "thresholds_key": ["verl", "fsdp2", "vllm-" + str(manifest["versions"].get("vllm")), _threshold_card(manifest)], "algorithm_spec": spec_dict, "syncer_epoch": args.rl_syncer_epoch, "island_lease_s": args.rl_island_lease_s, diff --git a/yeto/rl/adapters/verl/trainer.py b/yeto/rl/adapters/verl/trainer.py index 6be18c9e..08c72e83 100644 --- a/yeto/rl/adapters/verl/trainer.py +++ b/yeto/rl/adapters/verl/trainer.py @@ -186,7 +186,8 @@ def build_sync(plan: dict, specs, *, lora_hash: str, layout_hash: str, identity_ expected_specs=tuple(specs), base_model_revision=plan["model_revision"], lora_config_hash=lora_hash, layout_hash=layout_hash, event_tape=plan["event_tape"], wan_streams=int(plan.get("wan_streams", 4)), - backend_identity_sha256=identity_sha256, compat_group=compat_group) + backend_identity_sha256=identity_sha256, compat_group=compat_group, + bcast_delta=bool(plan.get("bcast_delta", False))) if mode == "strict": sync = StrictAvgSync(bridge, progress=None) elif mode == "elastic": diff --git a/yeto/rl/bridge.py b/yeto/rl/bridge.py index 9cf0e388..7cae9256 100644 --- a/yeto/rl/bridge.py +++ b/yeto/rl/bridge.py @@ -13,8 +13,13 @@ import torch -from .engine.backend_identity import config_island_contract_sha256 -from ..protocol import DTYPE_F32, FinalManifest, PullRequest, SyncerClient +from ..protocol import ( + DTYPE_F32, + DTYPE_FLAG_BCAST_DELTA, + FinalManifest, + PullRequest, + SyncerClient, +) from ..tensor_io import pack_tensor, unpack_fragment from .core import ( CanonicalLoraState, @@ -29,6 +34,7 @@ policy_delta, tensors_from_flat_owned, ) +from .engine.backend_identity import config_island_contract_sha256 class IslandRuntime(Protocol): @@ -77,6 +83,12 @@ class BridgeConfig: # 7.7c: "-" (yeto.hw.catalog); sent in HELLO / JOIN so the # syncer refuses another card type by name. None = not declared. compat_group: str | None = None + # S26 E3: receive BCAST as bf16 deltas against the shared f32 anchor + # (session flag; the syncer runs --bcast-delta-full-every). The client + # hands this bridge the reconstructed f32 anchor, so nothing else changes. + # An anchor mismatch drops the connection; with max_reconnects=0 the + # island exits and its restart re-reads the committed cut as f32. + bcast_delta: bool = False def _write_round_audit( @@ -278,7 +290,7 @@ def __init__(self, runtime: IslandRuntime, config: BridgeConfig) -> None: config.syncer_addr, config.learner_id, self.layout, - dtype=DTYPE_F32, + dtype=DTYPE_F32 | (DTYPE_FLAG_BCAST_DELTA if config.bcast_delta else 0), num_streams=config.wan_streams, session_contract_hash=_session_contract( self.layout, config_island_contract_sha256(config)), diff --git a/yeto/rl/engine/full_parameter_sync.py b/yeto/rl/engine/full_parameter_sync.py index f4b2141d..7c9d0921 100644 --- a/yeto/rl/engine/full_parameter_sync.py +++ b/yeto/rl/engine/full_parameter_sync.py @@ -81,6 +81,11 @@ async def stage_available_broadcasts(self) -> tuple[tuple[int, int], ...]: ... def staged_fragment_versions(self) -> tuple[tuple[int, int], ...]: ... def has_staged_broadcast(self, fragment_id: int, version: int) -> bool: ... async def apply_staged_broadcast(self, fragment_id: int, version: int) -> tuple[int, ...]: ... + # S26 E3 sync/train overlap (profile.overlap) only: + # complete_local_round(..., hold_submission=True) keeps the round base and + # the previous submission; apply_staged_broadcast(..., rebase=True) applies + # G + (current - base) with G as the anchor; finish_rebase() makes the + # rebased cut the submission; mark_round_base() pins the next round base. async def wait_and_apply_final(self) -> None: ... def acknowledge_final(self) -> None: ... def discard_detached_submission_cut(self) -> None: ... @@ -99,6 +104,8 @@ class FullParameterSyncProfile: wait_timeout: float poll_seconds: float sweep: bool = False + # S26 E3: train round k+1 while sweep k syncs; rebase at the next boundary. + overlap: bool = False def __post_init__(self) -> None: if not self.roles or len(set(self.roles)) != len(self.roles): @@ -107,6 +114,13 @@ def __post_init__(self) -> None: raise ValueError("the actor is the published role and must come first") if self.sweep and (self.roles != ("actor",) or self.local_horizon != 1): raise ValueError("a dense strict sweep has one actor role and H=1") + if self.overlap and not self.sweep: + raise ValueError("sync/train overlap needs the dense strict sweep") + + @property + def total_rounds(self) -> int: + """Local rounds of a sweep profile (one sweep per round).""" + return self.total_fragment_steps // self.fragments @property def effective_pipeline_depth(self) -> int: @@ -131,6 +145,7 @@ def from_configs( wait_timeout=min(config.wait_timeout for config in configs.values()), poll_seconds=min(config.poll_seconds for config in configs.values()), sweep=sweep, + overlap=bool(getattr(first, "overlap", False)), ) @@ -389,9 +404,12 @@ async def _apply_pair( step, fragment_id, _base = key if not pair.complete: raise RuntimeError("broadcast arrived before every paired PUSH was retained") - results = await asyncio.gather( - *(stream.apply_staged_broadcast(fragment_id, step) for stream in self.streams.values()) - ) + if self.profile.overlap: + calls = (stream.apply_staged_broadcast(fragment_id, step, rebase=True) + for stream in self.streams.values()) + else: + calls = (stream.apply_staged_broadcast(fragment_id, step) for stream in self.streams.values()) + results = await asyncio.gather(*calls) for role, fragments in zip(self.streams, results, strict=True): applied[role].extend(fragments) if self.streams[role].fragment_versions[fragment_id] != step: @@ -509,6 +527,8 @@ def __init__( self.closed = False self._driver: IslandDriver | None = None self._timing = SyncTiming(None, None) # L5.6, set by each boundary + # Overlap: fragment versions before the sweep still in flight. + self._pending_sweep: tuple[int, ...] | None = None @property def actor(self) -> RoleStream: @@ -539,6 +559,8 @@ def boundary(self, driver: IslandDriver, *, rollout_id: int, stats: Any) -> Sync return SyncBoundary(state=state, stop=self.terminal) async def _boundary(self, rollout_id: int) -> bool: + if self.pairing.profile.overlap: + return await self._overlap_boundary(rollout_id) pairing = self.pairing trajectories = self.trajectories(rollout_id) await pairing.gather_roles( @@ -564,6 +586,53 @@ async def _boundary(self, rollout_id: int) -> bool: stream.discard_detached_submission_cut() return False + async def _overlap_boundary(self, rollout_id: int) -> bool: + """Streaming-DiLoCo overlap (S26 E3; one more outer step of staleness). + + Round k's delta (``current - base``) is pushed at boundary k and the + island trains round k+1 at once. At boundary k+1 the island first + waits for sweep k (``G_k``), then rebases its trainer in place: + ``theta <- G_k + (theta_local_{k+1} - base_k)`` with ``G_k`` as the + anchor, so the next push is exactly round k+1's progress and round + k's delta is never counted twice. The terminal round waits for the + final cut like the strict sweep.""" + + pairing = self.pairing + trajectories = self.trajectories(rollout_id) + hold = self._pending_sweep is not None + await pairing.gather_roles( + lambda _role, stream: stream.complete_local_round( + rollout_id=rollout_id, trajectories=trajectories, hold_submission=hold + ) + ) + wait_ts = unix_now() + if hold: + await self._wait_for_sweep(self._pending_sweep) + self._pending_sweep = None + if not pairing.any_finalizing: + for stream in pairing.streams.values(): + stream.finish_rebase() + rebased_ts = unix_now() + if pairing.any_finalizing: + await pairing.finish_without_more_training() + return True + before = pairing.fragment_versions() + pairing.submit_ready() + submit_ts = unix_now() + if rollout_id + 1 >= pairing.profile.total_rounds: + # Terminal round: nothing trains after it, so wait for the cut. + await pairing.commit_staged_broadcasts() + await self._wait_for_sweep(before) + self._timing = SyncTiming(submit_ts, unix_now()) + await pairing.finish_without_more_training() + return True + self._pending_sweep = before + for stream in pairing.streams.values(): + stream.mark_round_base() + # This boundary blocked only on the previous sweep (wait + rebase). + self._timing = SyncTiming(wait_ts, rebased_ts if hold else None) + return False + async def _wait_for_sweep(self, before: tuple[int, ...]) -> None: """Strict sweep: block until all P fragments of this local step are applied, i.e. every fragment version moved past ``before``. An empty diff --git a/yeto/rl/engine/run_flags.py b/yeto/rl/engine/run_flags.py index 6e6b95c1..a69e67fd 100644 --- a/yeto/rl/engine/run_flags.py +++ b/yeto/rl/engine/run_flags.py @@ -851,11 +851,37 @@ def rl_forward_pairs(*dests: str | tuple[str, str]) -> tuple[tuple[str, str], .. ), cli_flag( "--rl-sync-push-dtype", - choices=("f32", "bf16"), + choices=("f32", "bf16", "int8"), default="f32", help="full-parameter only: wire dtype of the islands' pushed outer deltas " - "(pseudo-gradients). bf16 halves upload bytes; the syncer keeps f32 masters, " - "momentum and f32 BCAST/FINAL (default f32 = exact)", + "(pseudo-gradients). bf16 halves upload bytes, int8 (block absmax scale, see " + "--rl-sync-push-int8-block) quarters them; the syncer keeps f32 masters, " + "momentum and f32 FINAL (default f32 = exact)", + ), + cli_flag( + "--rl-sync-overlap", + action="store_true", + help="full-parameter only (E2 option): train the next round while the " + "sweep syncs, then rebase the trainer onto the new global weights " + "(theta <- G_k + local progress). Rollouts are one outer step staler: " + "--rl-max-policy-age is raised to at least 1 and the algorithm spec must " + "tolerate it (checked before launch)", + ), + cli_flag( + "--rl-sync-push-int8-block", + type=int, + default=128, + help="full-parameter only, with --rl-sync-push-dtype int8: values per f32 " + "scale (32..1048576; error <= block absmax / 254 per value)", + ), + cli_flag( + "--rl-sync-bcast-delta-full-every", + type=int, + default=0, + help="full-parameter only: above 0, the syncer sends each new fragment " + "version as a bf16 delta against the islands' shared f32 anchor, with an " + "exact f32 refresh every N-th version (N >= 2; digest-checked, a mismatch " + "reconnects for f32). 0 = always f32", ), cli_flag( "--rl-sync-checkpoint-every", diff --git a/yeto/rl/engine/sync_overlap.py b/yeto/rl/engine/sync_overlap.py new file mode 100644 index 00000000..4c47f344 --- /dev/null +++ b/yeto/rl/engine/sync_overlap.py @@ -0,0 +1,41 @@ +"""Sync/train overlap for the full-parameter sweep (S26 E3 item 3, E2 option). + +Round k's outer delta is pushed at boundary k and the island trains round +k+1 immediately on its own weights. At boundary k+1 it waits for the global +result of sweep k (``G_k``) and rebases in place:: + + theta <- G_k + (theta_local_{k+1} - base_k) + +``base_k`` is the cut round k+1 started from. ``G_k`` becomes the anchor, so +the next push ``theta - G_k`` equals round k+1's progress exactly once: round +k's delta is already inside ``G_k`` and is not counted again. + +Rollouts of round k+1 use weights one outer step older than the global +policy, so the run's minimum ``--rl-max-policy-age`` grows by +:data:`EXTRA_POLICY_AGE`. +""" + +from __future__ import annotations + +import numpy as np + +EXTRA_POLICY_AGE = 1 + + +def required_policy_age(minimum: int) -> int: + """The smallest accepted ``--rl-max-policy-age`` once overlap is on.""" + return int(minimum) + EXTRA_POLICY_AGE + + +def rebase_values(anchor: np.ndarray, current: np.ndarray, base: np.ndarray) -> np.ndarray: + """``anchor + (current - base)`` in f32 (the trainer keeps its progress).""" + if not (anchor.shape == current.shape == base.shape): + raise ValueError("rebase operands have different shapes") + for name, value in (("anchor", anchor), ("current", current), ("base", base)): + if value.dtype != np.float32: + raise ValueError(f"rebase {name} must be float32") + out = np.subtract(current, base, dtype=np.float32) + np.add(anchor, out, out=out) + if not np.isfinite(out).all(): + raise ValueError("rebased fragment has a non-finite value") + return out 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 07/62] 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 8af78faf152180d156bc6cf6d996d9339e253e99 Mon Sep 17 00:00:00 2001 From: michael <78113554+chowhsu@users.noreply.github.com> Date: Tue, 20 Oct 2026 16:04:45 +0000 Subject: [PATCH 08/62] overlap: keep rollout policy age 0, relabel rebased cuts, push whole sweep Pre-GPU check of --rl-sync-overlap on the Miles role stream found three problems: - The launcher raised --rl-max-policy-age to 1, but the dense sweep island refuses any age other than 0, so an overlap run could never start; and an age >= 1 also turns on Miles' fully_async rollout (partial rollout when colocated). Overlap is not rollout staleness: the driver publishes the boundary's cut before the next generation, so rollouts stay on-policy and only the local weights lag the global policy by one outer step. Keep the age at 0 (refuse others at launch) and bind the lag into the session contract only. - After a rebase the applied cut has local step generation 0, and Miles pushes only cuts with progress, so the second sweep's push failed. The submission is now the rebased cut relabeled with the rounds trained since the round base (same values, hashes and refs). - The boundary submitted once; a PULL that arrived later waited a whole local round. The overlap boundary now submits the whole sweep before training on, and the launcher requires --rl-sync-fragment-pipeline equal to --fragments so the syncer sends all PULLs at once. Adds Miles role-stream tests (real adapter and role stream, CPU tensors, fake trainer/Ray/syncer): exact-once rounds through rebase_fragment, and syncer restarts while the overlapped round trains (before commit, after commit with a lost broadcast, partly committed sweep). Co-Authored-By: Claude Opus 5.5 --- tests/test_e3_overlap_miles_role_stream.py | 495 ++++++++++++++++++ tests/test_e3_sync_overlap.py | 81 ++- yeto/launch/rl_args.py | 15 +- .../miles/full_parameter/role_stream.py | 15 +- yeto/rl/adapters/miles/island_entry.py | 10 +- .../miles/models/chunked_full_parameter.py | 36 ++ yeto/rl/adapters/miles/run_flags.py | 2 +- yeto/rl/engine/full_parameter_sync.py | 36 +- yeto/rl/engine/run_flags.py | 6 +- yeto/rl/engine/sync_overlap.py | 32 +- 10 files changed, 670 insertions(+), 58 deletions(-) create mode 100644 tests/test_e3_overlap_miles_role_stream.py diff --git a/tests/test_e3_overlap_miles_role_stream.py b/tests/test_e3_overlap_miles_role_stream.py new file mode 100644 index 00000000..6d9f0aa1 --- /dev/null +++ b/tests/test_e3_overlap_miles_role_stream.py @@ -0,0 +1,495 @@ +"""S26 E3: sync/train overlap on the real Miles role stream (CPU tensors, no GPU). + +``tests/test_e3_sync_overlap.py`` drives ``FullParameterSync`` with a numeric +fake role. Here the coordinator drives the real ``MilesFullParameterRoleStream`` +and ``MilesChunkedFullParameterAdapter`` (``rebase_fragment``, held +submissions, ``finish_rebase``, ``mark_round_base``, retained replays). Only +the Megatron trainer group, Ray and the syncer connection are fakes: + +* the trainer keeps one f32 tensor per fragment and adds a fixed progress per + round, so every pushed delta must equal that round's progress exactly; +* the syncer merges the island's delta with a second, virtual island's delta + (plain average) so the global weights differ from the local ones and a + wrong rebase shows up in the numbers. + +The syncer-restart tests cut the run while the island trains the overlapped +round: the syncer resumes from its last checkpoint (uncommitted pushes lost, +PULLs re-sent with a new attempt, committed fragments re-broadcast) and the +run must still count every round exactly once. + +Needs Miles (``importorskip``): run it in the Miles image. +""" + +from __future__ import annotations + +# ruff: noqa: I001 -- the Miles test helpers install a bridge stub before import. + +import asyncio +import threading +from types import SimpleNamespace + +import pytest +import torch + +pytest.importorskip("miles") + +from test_rl_miles_chunked_full_parameter import FakeRay, _evidence, component, manifest # noqa: E402 + +from miles.backends.megatron_utils.full_parameter_state import ( # noqa: E402 + FullParameterChunkDescriptor, + FullParameterChunkedShardState, + FullParameterFragmentDescriptor, + _chunk_hash, + _fragment_hasher, + _tensor_bytes, +) +from miles.ray.full_parameter_transport import FullParameterChunkedCut # noqa: E402 +from yeto.protocol import BcastFragment, FinalFragment, FinalManifest, PullRequest # noqa: E402 +from yeto.rl.adapters.miles.full_parameter import role_stream as role_stream_module # noqa: E402 +from yeto.rl.adapters.miles.full_parameter.role_stream import ( # noqa: E402 + MilesFullParameterRoleStream, + RoleStreamConfig, +) +from yeto.rl.adapters.miles.models.chunked_full_parameter import ( # noqa: E402 + MilesChunkedFullParameterAdapter, +) +from yeto.rl.engine.full_parameter_sync import ( # noqa: E402 + FullParameterSync, + FullParameterSyncProfile, + RolePairing, +) + +P = 2 +SIZES = (3, 2) # one parameter per owner shard; one fragment per shard +OTHER_WEIGHT = 0.5 # plain average of this island and one virtual peer + + +def _adapter() -> MilesChunkedFullParameterAdapter: + return MilesChunkedFullParameterAdapter.create( + (manifest(0, (SIZES[0],)), manifest(1, (SIZES[1],))), + algorithm="grpo", + components=(component(),), + minimum_fragments=P, + max_fragment_bytes=64, + max_chunk_bytes=8, # two floats per chunk: fragment 0 spans two chunks + ) + + +def _build_cut(adapter, values, *, policy_version, generation, ray) -> FullParameterChunkedCut: + """A reference-backed Miles cut holding ``values[fragment_id]``.""" + shards = [] + for plan in adapter.plans: + descriptors, refs = [], [] + for fragment in plan.fragments: + value = values[fragment.fragment_id].to(torch.float32).contiguous() + assert value.numel() == fragment.numel + digest = _fragment_hasher(fragment.fragment_id, adapter.layout.layout_hash, plan.plan_hash) + chunks, chunk_refs = [], [] + step = adapter.max_chunk_bytes // 4 + for index, start in enumerate(range(0, value.numel(), step)): + chunk = value[start:start + step].clone().contiguous() + digest.update(_tensor_bytes(chunk)) + chunks.append(FullParameterChunkDescriptor( + index, start, chunk.numel(), + _chunk_hash(fragment.fragment_id, index, start, chunk, + parameter_layout_hash=adapter.layout.layout_hash, + plan_hash=plan.plan_hash))) + chunk_refs.append(ray.put(chunk)) + descriptors.append(FullParameterFragmentDescriptor( + fragment.fragment_id, value.numel(), digest.hexdigest(), tuple(chunks))) + refs.append(chunk_refs) + shards.append(FullParameterChunkedShardState( + policy_version, generation, plan.topology, adapter.layout.layout_hash, + plan.plan_hash, tuple(descriptors), refs)) + return FullParameterChunkedCut(policy_version, generation, adapter.layout.layout_hash, + tuple(sorted(shards, key=lambda state: state.topology))) + + +class _Trainer: + """The Megatron trainer group: one f32 tensor per fragment, fixed progress per round.""" + + def __init__(self, adapter, ray, initial, progress) -> None: + self.adapter = adapter + self.ray = ray + self.values = {f: value.clone() for f, value in initial.items()} + self.progress = progress + self.installed = 0 + self.generation = 0 + self.scheduler_steps = 0 + self.trained_from: list[dict[int, torch.Tensor]] = [] + self.on_train = None # hook: runs while a round trains (cut point) + + async def record_full_parameter_local_step(self, *, base_policy_version, rollout_id): + assert base_policy_version == self.installed + self.trained_from.append({f: v.clone() for f, v in self.values.items()}) + for f, delta in self.progress(rollout_id).items(): + self.values[f] = self.values[f] + delta + self.generation += 1 + start, self.scheduler_steps = self.scheduler_steps, self.scheduler_steps + 1 + if self.on_train is not None: + self.on_train(rollout_id) + return tuple( + SimpleNamespace(topology=shard.topology, role="actor", base_policy_version=base_policy_version, + local_step_generation=self.generation, rollout_id=rollout_id, + optimizer_steps=1, scheduler_start_steps=start, + scheduler_end_steps=self.scheduler_steps) + for shard in self.adapter.manifests) + + async def export_full_parameter_chunked_cut(self, policy_version, local_step_generation=0, *, + max_chunk_bytes): + assert max_chunk_bytes == self.adapter.max_chunk_bytes + return _build_cut(self.adapter, self.values, policy_version=policy_version, + generation=local_step_generation, ray=self.ray) + + async def apply_full_parameter_chunked_cut(self, cut, *, commit_token=None, max_commit_attempts=3): + assert commit_token + for shard in cut.shards: + for descriptor, refs in zip(shard.fragments, shard.chunk_refs, strict=True): + assert all(ref is not None for ref in refs), "applied a released chunk" + self.values[descriptor.fragment_id] = torch.cat([self.ray.get(ref) for ref in refs]).clone() + self.installed = cut.policy_version + self.generation = 0 + return self.adapter.expected_parameter_tensor_count + + async def full_parameter_optimizer_states(self): + return tuple(SimpleNamespace(role="actor", installed_policy_version=self.installed, + local_step_generation=self.generation, + scheduler_num_steps=self.scheduler_steps) + for _ in self.adapter.manifests) + + +class _Syncer: + """In-process strict sweep (``--pipeline P``, checkpoint every round). + + ``G_f <- G_f + OTHER_WEIGHT * (island delta + peer delta)``; committed state + is checkpointed at once, so a restart loses only uncommitted pushes.""" + + def __init__(self, initial, sweeps, peer) -> None: + self.global_ = {f: v.clone() for f, v in initial.items()} + self.versions = [0] * P + self.sweeps = sweeps + self.peer = peer # rollout_id -> {fragment: peer delta} + self.received: dict[int, torch.Tensor] = {} # uncommitted pushes by step + self.pushed: list[tuple[int, int, torch.Tensor]] = [] # (step, attempt, delta), every PUSH seen + self.committed: list[int] = [] + self.attempt = 1 + self.client = None + + def commit_ready(self) -> None: + """Commit fragments in strict step order once their PUSH arrived.""" + while True: + step = len(self.committed) + 1 + if step > self.sweeps * P or step not in self.received: + return + f, k = (step - 1) % P, (step - 1) // P + delta = self.received.pop(step) + self.global_[f] = self.global_[f] + OTHER_WEIGHT * (delta + self.peer(k)[f]) + self.versions[f] = step + self.committed.append(step) + if step == self.sweeps * P: + self.client.publish_final() + elif step % P == 0: + # The sweep is committed: broadcast it, then PULL the next one. + for fragment in range(P): + self.client.broadcast(fragment, self.versions[fragment], self.global_[fragment]) + self.client.pull_sweep(step // P, self.attempt) + + def restart_from_checkpoint(self) -> None: + """Crash + ``--resume``: uncommitted pushes are lost; the reconnect + re-broadcasts every committed fragment and re-sends the open PULLs.""" + self.received.clear() + self.attempt += 1 + self.client.reconnect() + for fragment in range(P): + if self.versions[fragment]: + self.client.broadcast(fragment, self.versions[fragment], self.global_[fragment]) + if len(self.committed) % P == 0 and len(self.committed) < self.sweeps * P: + self.client.pull_sweep(len(self.committed) // P, self.attempt) + else: + # A partly committed sweep: only its uncommitted fragments are open. + sweep = len(self.committed) // P + for step in range(len(self.committed) + 1, (sweep + 1) * P + 1): + self.client.pulls.append(PullRequest((step - 1) % P, step, self.attempt)) + + +class _Client: + """The ``SyncerClient`` surface the role stream uses, backed by ``_Syncer``.""" + + def __init__(self, syncer: _Syncer) -> None: + self.syncer = syncer + syncer.client = self + self.connection_generation = 1 + self.finalizing = threading.Event() + self.pulls: list[PullRequest] = [] + self.updates: list[BcastFragment] = [] + self.inbound = None + self.final: tuple[FinalManifest, list[FinalFragment]] | None = None + self.acknowledged: FinalManifest | None = None + self.closed = False + + # -- syncer side + def _stored(self, fragment_id, version, value): + return self.inbound.sink(fragment_id, version, _tensor_bytes(value.to(torch.float32).contiguous())) + + def broadcast(self, fragment_id, version, value) -> None: + self.updates.append(BcastFragment(fragment_id, version, self._stored(fragment_id, version, value), + stored=True)) + + def pull_sweep(self, sweep, attempt) -> None: + self.pulls.extend(PullRequest(f, sweep * P + f + 1, attempt) for f in range(P)) + + def reconnect(self) -> None: + self.connection_generation += 1 + self.pulls.clear() + + def publish_final(self) -> None: + last = self.syncer.sweeps * P + versions = tuple(range(last - P + 1, last + 1)) + fragments = [FinalFragment(f, versions[f], self._stored(f, versions[f], self.syncer.global_[f]), + stored=True) for f in range(P)] + self.final = (FinalManifest(last, versions), fragments) + self.finalizing.set() + + # -- island side (SyncerClient API) + def install_inbound_chunk_sink(self, sink) -> None: + self.inbound = sink + + def start(self) -> None: + pass + + def send_init_parts(self, fragment_id, parts) -> bool: + payload = b"".join(bytes(part) for part in parts) + value = torch.frombuffer(bytearray(payload), dtype=torch.float32) + assert torch.equal(value, self.syncer.global_[fragment_id]) + self.broadcast(fragment_id, 0, value) + if fragment_id == P - 1: + self.pull_sweep(0, 1) + return True + + def check_health(self) -> None: + assert not self.closed + + def drain_updates(self): + out, self.updates = self.updates, [] + return out + + def drain_pulls(self): + out, self.pulls = self.pulls, [] + return out + + def push_fragment_parts(self, fragment_id, global_step, round_attempt, base_version, local_step, + c_steps, c_tokens, parts) -> bool: + assert (global_step - 1) % P == fragment_id and c_steps == 1 and c_tokens >= 1 + assert base_version == self.syncer.versions[fragment_id], "PUSH on a stale base" + delta = torch.frombuffer(bytearray(b"".join(bytes(part) for part in parts)), dtype=torch.float32) + self.syncer.pushed.append((global_step, round_attempt, delta.clone())) + self.syncer.received[global_step] = delta.clone() + self.syncer.commit_ready() + return True + + def wait_for_final_fragments(self, timeout): + assert self.final is not None, "no FINAL before the terminal barrier" + return self.final + + def acknowledge_finalization(self, manifest) -> None: + self.acknowledged = manifest + + def close(self) -> None: + self.closed = True + + +class _Driver: + def __init__(self) -> None: + self.events = [] + + def emit(self, event, **fields) -> None: + self.events.append((event, fields)) + + +def _config(sweeps: int) -> RoleStreamConfig: + adapter = _adapter() + return RoleStreamConfig( + algorithm="grpo", role="actor", component=component(), syncer_addr=("127.0.0.1", 30100), + learner_id=0, learner_generation=0, learner_generations=(0, 0), + total_fragment_steps=sweeps * P, expected_fragments=P, + expected_layout_hash=adapter.layout.layout_hash, local_horizon=1, optimizer_steps_per_round=1, + training_contract_hash="2" * 64, syncer_profile_hash="3" * 64, pipeline_depth=1, + algorithm_spec_sha256="a1" * 32, wait_timeout=5.0, poll_seconds=0.001, + max_fragment_bytes=64, max_chunk_bytes=8, overlap=True) + + +def _run(monkeypatch, sweeps: int, *, cut=None): + """Run ``sweeps`` overlapped rounds; ``cut(rollout_id, syncer)`` runs while + that round trains (the mid-overlap cut point).""" + generator = torch.Generator().manual_seed(11) + + def rand(scale=1.0): + return {f: torch.randn(n, generator=generator) * scale for f, n in enumerate(SIZES)} + + initial = rand() + steps = {k: rand(0.1) for k in range(sweeps)} + peer = {k: rand(0.1) for k in range(sweeps)} + ray = FakeRay() + adapter = _adapter() + trainer = _Trainer(adapter, ray, initial, lambda k: steps[k]) + syncer = _Syncer(initial, sweeps, lambda k: peer[k]) + client = _Client(syncer) + + async def initialize(cls, group, *, policy_version, **_kwargs): + cut = await group.export_full_parameter_chunked_cut(policy_version, 0, + max_chunk_bytes=adapter.max_chunk_bytes) + return adapter, adapter.validate_cut(cut, policy_version, 0) + + monkeypatch.setattr(MilesChunkedFullParameterAdapter, "initialize", classmethod(initialize)) + monkeypatch.setattr(role_stream_module, "SyncerClient", lambda *_args, **_kwargs: client) + + stream = MilesFullParameterRoleStream(_config(sweeps), ray_module=ray) + profile = FullParameterSyncProfile.from_configs({"actor": stream.config}, sweep=True) + assert profile.overlap + session = FullParameterSync( + RolePairing({"actor": stream}, profile), groups={"actor": trainer}, run=asyncio.run, + trajectories=lambda rollout_id: _evidence(rollout_id=rollout_id, + behavior_policy_hash=stream.current.policy_hash)) + if cut is not None: + trainer.on_train = lambda rollout_id: cut(rollout_id, syncer) + driver = _Driver() + session.start(driver) + stops = [] + for rollout_id in range(sweeps): + boundary = session.boundary(driver, rollout_id=rollout_id, stats=None) + stops.append(boundary.stop) + if boundary.stop: + session.published(driver, rollout_id=rollout_id, policy_hash=boundary.state.policy_tensor_hash()) + break + session.finish(driver) + return SimpleNamespace(initial=initial, steps=steps, peer=peer, trainer=trainer, syncer=syncer, + client=client, stream=stream, stops=stops, ray=ray) + + +def _expected_global(run, sweeps): + return {f: run.initial[f] + sum(OTHER_WEIGHT * (run.steps[k][f] + run.peer[k][f]) + for k in range(sweeps)) for f in range(P)} + + +def _committed_deltas(run): + """The delta of each committed step (the last PUSH seen for that step).""" + last = {} + for step, attempt, delta in run.syncer.pushed: + if attempt >= last.get(step, (0, None))[0]: + last[step] = (attempt, delta) + return {step: delta for step, (_attempt, delta) in last.items()} + + +def _check_exactly_once(run, sweeps): + assert run.stops == [False] * (sweeps - 1) + [True] + assert run.syncer.committed == list(range(1, sweeps * P + 1)) + deltas = _committed_deltas(run) + for k in range(sweeps): + for f in range(P): + # Each push is exactly that round's progress: no round is pushed + # twice, none is lost, the rebase kept the local progress. + torch.testing.assert_close(deltas[k * P + f + 1], run.steps[k][f], rtol=0, atol=1e-6) + want = _expected_global(run, sweeps) + for f in range(P): + torch.testing.assert_close(run.syncer.global_[f], want[f], rtol=0, atol=1e-5) + torch.testing.assert_close(run.trainer.values[f], want[f], rtol=0, atol=1e-5) + assert run.stream.final_applied and run.stream.final_acknowledged + assert run.client.acknowledged == run.stream.terminal_manifest + + +def test_overlap_on_the_miles_role_stream_rebases_and_counts_every_round_once(monkeypatch): + sweeps = 4 + run = _run(monkeypatch, sweeps) + _check_exactly_once(run, sweeps) + trained = run.trainer.trained_from + for f in range(P): + # Round 0 starts at the initial weights, round 1 at its own round-0 + # result (sweep 0 in flight), round k >= 2 at G_{k-2} + progress of + # round k-1: the rebase ran in the trainer through rebase_fragment. + torch.testing.assert_close(trained[0][f], run.initial[f]) + torch.testing.assert_close(trained[1][f], run.initial[f] + run.steps[0][f]) + for k in range(2, sweeps): + g = run.initial[f] + sum(OTHER_WEIGHT * (run.steps[j][f] + run.peer[j][f]) for j in range(k - 1)) + torch.testing.assert_close(trained[k][f], g + run.steps[k - 1][f], rtol=0, atol=1e-5) + # Exactly one PUSH per step without a restart. + assert sorted(step for step, _attempt, _delta in run.syncer.pushed) == list(range(1, sweeps * P + 1)) + + +def test_overlap_role_stream_releases_every_superseded_cut(monkeypatch): + run = _run(monkeypatch, 3) + stream = run.stream + assert not stream._pinned_cuts and not stream._pinned_anchors + assert not stream._staged + + +@pytest.mark.parametrize("cut_round", [1, 2]) +def test_syncer_restart_before_the_commit_mid_overlap_replays_the_retained_push(monkeypatch, cut_round): + """Cut while round ``cut_round`` trains: sweep ``cut_round - 1`` was pushed + but not committed. The resumed syncer re-PULLs it with attempt 2; the + island replays the retained bytes at its next boundary.""" + sweeps = 4 + + def cut(rollout_id, syncer): + if rollout_id == cut_round and syncer.attempt == 1: + # Undo the commit of the in-flight sweep: the crash came before it. + sweep = cut_round - 1 + for step in range(sweep * P + 1, (sweep + 1) * P + 1): + syncer.committed.remove(step) + syncer.global_ = {f: g.clone() for f, g in checkpoints[sweep].items()} + syncer.versions = [sweep * P - P + f + 1 if sweep else 0 for f in range(P)] + syncer.client.updates.clear() # its BCAST never left the syncer + syncer.restart_from_checkpoint() + + checkpoints = {} + original = _Syncer.commit_ready + + def commit_ready(self): + if len(self.committed) % P == 0: + checkpoints.setdefault(len(self.committed) // P, {f: g.clone() for f, g in self.global_.items()}) + original(self) + + monkeypatch.setattr(_Syncer, "commit_ready", commit_ready) + run = _run(monkeypatch, sweeps, cut=cut) + _check_exactly_once(run, sweeps) + replayed = [(step, attempt) for step, attempt, _ in run.syncer.pushed if attempt == 2] + first = (cut_round - 1) * P + 1 + assert [step for step, _ in replayed][:P] == list(range(first, first + P)) + + +@pytest.mark.parametrize("cut_round", [1, 2]) +def test_syncer_restart_after_the_commit_mid_overlap_rebases_on_the_resent_broadcast(monkeypatch, cut_round): + """Cut while round ``cut_round`` trains: sweep ``cut_round - 1`` committed + and checkpointed, but its BCAST was lost. The resumed syncer re-sends it + (FULL) and PULLs the next sweep; the island rebases on the re-sent bytes.""" + sweeps = 4 + + def cut(rollout_id, syncer): + if rollout_id == cut_round and syncer.attempt == 1: + syncer.client.updates.clear() # the BCAST was lost with the connection + syncer.restart_from_checkpoint() + + run = _run(monkeypatch, sweeps, cut=cut) + _check_exactly_once(run, sweeps) + assert all(attempt == 1 for step, attempt, _ in run.syncer.pushed if step <= cut_round * P) + + +def test_syncer_restart_with_a_partly_committed_sweep_mid_overlap(monkeypatch): + """Cut while round 2 trains: fragment 0 of sweep 1 committed (and is + re-broadcast), fragment 1 was lost. Only fragment 1 is re-PULLed and + replayed; the island applies the sweep only once both fragments arrived.""" + sweeps = 3 + lost = 2 * P # sweep 1, fragment 1 + + def cut(rollout_id, syncer): + if rollout_id == 2 and syncer.attempt == 1: + f = (lost - 1) % P + pushed = [delta for step, _attempt, delta in syncer.pushed if step == lost][-1] + syncer.committed.remove(lost) + syncer.global_[f] = syncer.global_[f] - OTHER_WEIGHT * (pushed + syncer.peer(1)[f]) + syncer.versions[f] = lost - P + syncer.client.updates.clear() + syncer.restart_from_checkpoint() + + run = _run(monkeypatch, sweeps, cut=cut) + _check_exactly_once(run, sweeps) + assert [step for step, attempt, _ in run.syncer.pushed if attempt == 2][0] == lost diff --git a/tests/test_e3_sync_overlap.py b/tests/test_e3_sync_overlap.py index 9b9fbc4d..6421d6ce 100644 --- a/tests/test_e3_sync_overlap.py +++ b/tests/test_e3_sync_overlap.py @@ -5,7 +5,7 @@ sweep k is in flight and rebases ``theta <- G_k + (local - base)``. The tests check that every round's progress reaches the global weights exactly once (no double counting), that the trainer keeps its progress across a rebase, -and the launcher/island policy-age rule. +and the launcher rules (rollout policy age 0, fragment pipeline P). """ from __future__ import annotations @@ -28,7 +28,7 @@ FullParameterSyncProfile, RolePairing, ) -from yeto.rl.engine.sync_overlap import rebase_values, required_policy_age +from yeto.rl.engine.sync_overlap import OUTER_STEP_LAG, overlap_option_problems, rebase_values P = 2 # fragments NUMEL = 4 @@ -95,6 +95,7 @@ def __init__(self, syncer: _SweepSyncer, progress) -> None: self.final_acknowledged = False self.closed = False self.trained_from: list[list[np.ndarray]] = [] + self.push_rounds: list[tuple[int, int]] = [] # (global step, local rounds trained then) @property def fragment_versions(self): @@ -143,6 +144,7 @@ def release_candidate(self, candidate): def push_candidate(self, candidate): f = candidate.permit.fragment_id + self.push_rounds.append((candidate.permit.global_step, self.local_rounds)) self.syncer.push(candidate.permit.global_step, f, candidate.submission_cut[f] - candidate.anchor) self._pending.pop(f) @@ -204,12 +206,31 @@ def emit(self, event, **fields): self.events.append((event, fields)) -def _run(sweeps: int, overlap: bool = True): +class _LatePullSyncer(_SweepSyncer): + """Shows the last fragment's PULL of each sweep only after a few polls.""" + + LATE_POLLS = 3 + + def __init__(self, *args) -> None: + super().__init__(*args) + self._polls: dict[int, int] = {} + + def permits(self): + permits = super().permits() + late = [p for p in permits if p.fragment_id == P - 1] + if late: + seen = self._polls[late[0].global_step] = self._polls.get(late[0].global_step, 0) + 1 + if seen <= self.LATE_POLLS: + return [p for p in permits if p.fragment_id != P - 1] + return permits + + +def _run(sweeps: int, overlap: bool = True, syncer_type=_SweepSyncer): rng = np.random.default_rng(7) initial = [rng.standard_normal(NUMEL).astype(np.float32) for _ in range(P)] steps = {k: [rng.standard_normal(NUMEL).astype(np.float32) * 0.1 for _ in range(P)] for k in range(sweeps)} - syncer = _SweepSyncer(initial, sweeps) + syncer = syncer_type(initial, sweeps) role = _NumericRole(syncer, lambda rollout_id: steps[rollout_id]) profile = FullParameterSyncProfile(roles=("actor",), total_fragment_steps=sweeps * P, fragments=P, local_horizon=1, pipeline_depth=1, wait_timeout=5.0, @@ -245,6 +266,18 @@ def test_overlap_counts_every_round_exactly_once(): assert role.final_acknowledged +def test_overlap_boundary_pushes_the_whole_sweep_even_when_a_pull_arrives_late(): + sweeps = 3 + _initial, _steps, syncer, role, stops = _run(sweeps, syncer_type=_LatePullSyncer) + assert stops == [False] * (sweeps - 1) + [True] + # Every sweep was complete before the next round trained: sweep k's last + # fragment is pushed before round k+1's progress exists (no sweep spans + # two boundaries). + assert role.push_rounds == [(k * P + f + 1, k + 1) for k in range(sweeps) for f in range(P)] + assert role.final_acknowledged + assert all(polls > _LatePullSyncer.LATE_POLLS for polls in syncer._polls.values()) + + def test_overlap_trains_one_outer_step_stale_and_keeps_local_progress(): sweeps = 4 initial, steps, _syncer, role, _ = _run(sweeps) @@ -287,34 +320,40 @@ def _dense(*extra): return args -def test_overlap_raises_the_minimum_policy_age_and_reaches_the_island(): +def test_overlap_keeps_the_rollout_policy_age_zero_and_reaches_the_island(): from yeto.launch.rl_island_task import _rl_trainable_island_flags + from yeto.launch.rl_ports import _ports_infra_flags - assert required_policy_age(0) == 1 - args = _dense("--rl-sync-overlap") - assert args.rl_max_policy_age == 1 + args = _dense("--rl-sync-overlap", "--rl-sync-fragment-pipeline", "2") + # Overlap is not rollout staleness: no policy age, so Miles gets neither + # partial rollout (carry-over) nor fully_async generation. + assert int(args.rl_max_policy_age or 0) == 0 + assert OUTER_STEP_LAG == 1 assert "--sweep-overlap" in _rl_trainable_island_flags(args, 0, 1) - assert _dense("--rl-sync-overlap", "--rl-max-policy-age", "3").rl_max_policy_age == 3 + assert "--rl-max-policy-age" not in " ".join(_ports_infra_flags(args)) plain = _dense() - assert int(plain.rl_max_policy_age or 0) == 0 assert "--sweep-overlap" not in _rl_trainable_island_flags(plain, 0, 1) -def test_overlap_launch_fails_when_the_algorithm_spec_does_not_tolerate_the_age(): - from yeto.launch_preflight import PreflightError, check_policy_age_spec - - args = _dense("--rl-sync-overlap") - with pytest.raises(PreflightError, match="--rl-max-policy-age 1 needs an algorithm spec"): - check_policy_age_spec(args) +def test_overlap_launch_refuses_a_policy_age_or_a_partial_fragment_pipeline(): + with pytest.raises(ValueError, match="--rl-max-policy-age 0 with --rl-sync-overlap"): + _dense("--rl-sync-overlap", "--rl-sync-fragment-pipeline", "2", "--rl-max-policy-age", "1") + with pytest.raises(ValueError, match="--rl-sync-fragment-pipeline equal to --fragments"): + _dense("--rl-sync-overlap") + assert overlap_option_problems(max_policy_age=0, fragment_pipeline=4, fragments=4) == [] -def test_island_refuses_overlap_without_the_raised_policy_age(): +def test_island_builds_overlap_without_a_policy_age(): from yeto.rl.adapters.miles.island_entry import _dense_full_sync_config - args = SimpleNamespace(sweep_overlap=True, rl_max_policy_age=0) - with pytest.raises(ValueError, match="--sweep-overlap .* needs --rl-max-policy-age >= 1"): - _dense_full_sync_config(args, model_config_sha256="0" * 64, - training_contract_sha256="0" * 64, algorithm_spec_sha256="0" * 64) + args = SimpleNamespace( + sweep_overlap=True, rl_max_policy_age=0, model_revision="A" * 40, syncer="127.0.0.1:30100", + learner_id=0, num_learners=2, total_fragment_steps=4, fragments=2, + parameter_layout_sha256="1" * 64, optimizer_steps=1, syncer_profile_sha256="3" * 64, + wan_streams=4) + config = _dense_full_sync_config(args, model_config_sha256="0" * 64, + training_contract_sha256="0" * 64, algorithm_spec_sha256="0" * 64) + assert config.roles[0].overlap def test_overlap_is_bound_into_the_role_stream_contract(): diff --git a/yeto/launch/rl_args.py b/yeto/launch/rl_args.py index 60e63f88..e623e6dc 100644 --- a/yeto/launch/rl_args.py +++ b/yeto/launch/rl_args.py @@ -214,6 +214,12 @@ def _prepare_dense_full_sweep(args, *, num_islands: int) -> None: problems.append("--rl-sync-push-int8-block between 32 and 1048576") elif int8_block != 128 and getattr(args, "rl_sync_push_dtype", "f32") != "int8": problems.append("--rl-sync-push-int8-block only with --rl-sync-push-dtype int8") + if getattr(args, "rl_sync_overlap", False): + from ..rl.engine.sync_overlap import overlap_option_problems + + problems += overlap_option_problems( + max_policy_age=getattr(args, "rl_max_policy_age", 0) or 0, + fragment_pipeline=fragment_pipeline, fragments=args.fragments) if problems: raise ValueError(f"--rl-sync-preset {FULL_PARAMETER_SYNC_PRESET} requires " + "; ".join(problems)) args.quorum = num_islands @@ -363,15 +369,6 @@ def _prepare_rl_args( # what the selected backend supports before any cloud or GPU work. from ..rl.engine.policy_age import validate_limit - if args.rl_sync_preset == FULL_PARAMETER_SYNC_PRESET and getattr(args, "rl_sync_overlap", False): - # S26 E3: sync/train overlap adds one outer step of staleness; raise the - # limit to the new minimum here so every check below (backend support, - # and the preflight's algorithm-spec staleness) validates it at launch. - from ..rl.engine.sync_overlap import required_policy_age - - required = required_policy_age(0) - if int(getattr(args, "rl_max_policy_age", 0) or 0) < required: - args.rl_max_policy_age = required max_policy_age = validate_limit(getattr(args, "rl_max_policy_age", 0) or 0) policy_age_module = _rl_backend_module(args, "policy_age") policy_age_module.SUPPORT.check(max_policy_age) diff --git a/yeto/rl/adapters/miles/full_parameter/role_stream.py b/yeto/rl/adapters/miles/full_parameter/role_stream.py index 4b2a5199..1855ad78 100644 --- a/yeto/rl/adapters/miles/full_parameter/role_stream.py +++ b/yeto/rl/adapters/miles/full_parameter/role_stream.py @@ -99,7 +99,8 @@ class RoleStreamConfig: push_int8_block: int = DEFAULT_I8_BLOCK bcast_delta: bool = False # S26 E3: train round k+1 while sweep k syncs, then rebase the trainer - # (see FullParameterSync._overlap_boundary); one more outer step stale. + # (see FullParameterSync._overlap_boundary); the local weights lag the + # global policy by one outer step (yeto.rl.engine.sync_overlap). overlap: bool = False @property @@ -653,16 +654,20 @@ async def apply_staged_broadcast( return (fragment_id,) def finish_rebase(self) -> None: - """Overlap: the rebased cut is the next submission (``G + progress``).""" + """Overlap: the rebased cut is the next submission (``G + progress``). - if not self.config.overlap or self.current is None: + The applied cut has generation 0; the submission is the same weights + labeled with the rounds trained since the round base.""" + + if not self.config.overlap or self.current is None or self.adapter is None: raise RuntimeError(f"{self.config.role} finish_rebase outside overlap") previous_submission = self.submission_cut - self.submission_cut = self.current + self.submission_cut = self.adapter.relabel_local_progress( + self.current, local_rounds=self.local_rounds - self._rounds_at_base) self.submission_fragment_versions = self.fragment_versions self.submission_rounds = self.local_rounds self.submission_tokens = self.action_tokens - self._release_cuts(previous_submission, keep=(self.current, self._round_base)) + self._release_cuts(previous_submission, keep=(self.current, self.submission_cut, self._round_base)) def mark_round_base(self) -> None: """Overlap: pin the cut the next local round starts from.""" diff --git a/yeto/rl/adapters/miles/island_entry.py b/yeto/rl/adapters/miles/island_entry.py index c382f8cb..6ff6a2e8 100644 --- a/yeto/rl/adapters/miles/island_entry.py +++ b/yeto/rl/adapters/miles/island_entry.py @@ -814,14 +814,8 @@ def _dense_full_sync_config( from yeto.rl.adapters.miles.full_parameter.role_stream import RoleStreamConfig from yeto.rl.adapters.miles.full_parameter.session import FullParameterSyncConfig - from yeto.rl.engine.sync_overlap import required_policy_age from yeto.rl.local_learner import ComponentIdentity - overlap = bool(getattr(args, "sweep_overlap", False)) - if overlap and int(getattr(args, "rl_max_policy_age", 0) or 0) < required_policy_age(0): - raise ValueError("--sweep-overlap rollouts are one outer step stale: it needs " - f"--rl-max-policy-age >= {required_policy_age(0)}") - actor = RoleStreamConfig( algorithm="grpo", role="actor", @@ -847,7 +841,9 @@ def _dense_full_sync_config( push_dtype=getattr(args, "sweep_push_dtype", "f32"), push_int8_block=int(getattr(args, "sweep_push_int8_block", 128)), bcast_delta=bool(getattr(args, "bcast_delta", False)), - overlap=overlap, + # The rollout policy age stays 0 (``_check_dense_full_sweep``); the + # overlap's outer-step lag is bound into the session contract. + overlap=bool(getattr(args, "sweep_overlap", False)), ) return FullParameterSyncConfig((actor,), sweep=True) diff --git a/yeto/rl/adapters/miles/models/chunked_full_parameter.py b/yeto/rl/adapters/miles/models/chunked_full_parameter.py index f16a6ec5..14a675d6 100644 --- a/yeto/rl/adapters/miles/models/chunked_full_parameter.py +++ b/yeto/rl/adapters/miles/models/chunked_full_parameter.py @@ -840,6 +840,42 @@ def flat(cut: ReferencedPolicyCut) -> np.ndarray: sink = AuthoritativeFragmentSink(current, ray_module=ray_module) return sink(fragment_id, stored.version, values.tobytes()) + def relabel_local_progress( + self, + cut: ReferencedPolicyCut, + *, + local_rounds: int, + ) -> ReferencedPolicyCut: + """Sync/train overlap: label a rebased cut with its local progress. + + ``apply`` resets the local step generation to 0, but a rebased cut + (``G + (current - base)``) still holds ``local_rounds`` rounds of + progress over its anchor ``G``, and Miles pushes only cuts with + progress (generation >= 1). Values, chunk hashes and refs are shared + with ``cut`` (fresh ref lists, so releasing one leaves the other); + only the generation label changes.""" + + from dataclasses import replace + + from miles.ray.full_parameter_transport import FullParameterChunkedCut + + if isinstance(local_rounds, bool) or not isinstance(local_rounds, int) or local_rounds < 1: + raise ValueError("a rebased cut needs at least one local round") + if cut.layout_hash != self.layout.layout_hash or cut.local_step_generation != 0: + raise ValueError("only an applied (generation 0) cut is relabeled") + transport = cut.transport_cut + relabeled = FullParameterChunkedCut( + policy_version=transport.policy_version, + local_step_generation=local_rounds, + parameter_layout_hash=transport.parameter_layout_hash, + shards=tuple( + replace(shard, local_step_generation=local_rounds, + chunk_refs=[list(references) for references in shard.chunk_refs]) + for shard in transport.shards + ), + ) + return self.validate_cut(relabeled, cut.policy_version, local_rounds) + async def apply( self, group: MilesChunkedFullParameterGroup, diff --git a/yeto/rl/adapters/miles/run_flags.py b/yeto/rl/adapters/miles/run_flags.py index 79b6da70..ff2830bc 100644 --- a/yeto/rl/adapters/miles/run_flags.py +++ b/yeto/rl/adapters/miles/run_flags.py @@ -329,7 +329,7 @@ "--sweep-overlap", action="store_true", help="full-parameter sweep: train the next round while the sweep syncs, then " - "rebase onto the new global weights (needs --rl-max-policy-age >= 1)", + "rebase onto the new global weights (rollout policy age stays 0)", ), island_flag( "--bcast-delta", diff --git a/yeto/rl/engine/full_parameter_sync.py b/yeto/rl/engine/full_parameter_sync.py index 7c9d0921..744f3053 100644 --- a/yeto/rl/engine/full_parameter_sync.py +++ b/yeto/rl/engine/full_parameter_sync.py @@ -612,27 +612,53 @@ async def _overlap_boundary(self, rollout_id: int) -> bool: if not pairing.any_finalizing: for stream in pairing.streams.values(): stream.finish_rebase() - rebased_ts = unix_now() if pairing.any_finalizing: await pairing.finish_without_more_training() return True before = pairing.fragment_versions() - pairing.submit_ready() - submit_ts = unix_now() if rollout_id + 1 >= pairing.profile.total_rounds: # Terminal round: nothing trains after it, so wait for the cut. + pairing.submit_ready() + submit_ts = unix_now() await pairing.commit_staged_broadcasts() await self._wait_for_sweep(before) self._timing = SyncTiming(submit_ts, unix_now()) await pairing.finish_without_more_training() return True + await self._push_whole_sweep(before) self._pending_sweep = before for stream in pairing.streams.values(): stream.mark_round_base() - # This boundary blocked only on the previous sweep (wait + rebase). - self._timing = SyncTiming(wait_ts, rebased_ts if hold else None) + # This boundary blocked on the previous sweep (wait + rebase) and on + # handing every fragment of this sweep to the wire. + self._timing = SyncTiming(wait_ts, unix_now()) return False + async def _push_whole_sweep(self, before: tuple[int, ...]) -> None: + """Overlap: submit every fragment of the sweep after ``before`` before + training on. A PULL that arrives later would otherwise wait for the + next boundary, a whole local round, and stall the sweep. The syncer + sends all P PULLs at once with ``--pipeline P`` (required by the + launcher), so this waits only for their arrival.""" + + pairing = self.pairing + fragments = pairing.profile.fragments + steps = {version + fragments if version else fragment_id + 1 + for fragment_id, version in enumerate(before)} + deadline = time.monotonic() + pairing.profile.wait_timeout + + def pushed() -> bool: + return steps <= {key[0] for key, pair in pairing.inflight.items() if pair.complete} + + pairing.submit_ready() + while not pushed(): + if time.monotonic() >= deadline: + error = TimeoutError("overlap could not submit the whole sweep before the wait timeout") + pairing.fail_stop(error) + raise error + await asyncio.sleep(pairing.profile.poll_seconds) + pairing.submit_ready() + async def _wait_for_sweep(self, before: tuple[int, ...]) -> None: """Strict sweep: block until all P fragments of this local step are applied, i.e. every fragment version moved past ``before``. An empty diff --git a/yeto/rl/engine/run_flags.py b/yeto/rl/engine/run_flags.py index a69e67fd..2af06a5c 100644 --- a/yeto/rl/engine/run_flags.py +++ b/yeto/rl/engine/run_flags.py @@ -863,9 +863,9 @@ def rl_forward_pairs(*dests: str | tuple[str, str]) -> tuple[tuple[str, str], .. action="store_true", help="full-parameter only (E2 option): train the next round while the " "sweep syncs, then rebase the trainer onto the new global weights " - "(theta <- G_k + local progress). Rollouts are one outer step staler: " - "--rl-max-policy-age is raised to at least 1 and the algorithm spec must " - "tolerate it (checked before launch)", + "(theta <- G_k + local progress). The local weights lag the global policy " + "by one outer step; rollouts stay on-policy, so --rl-max-policy-age stays 0. " + "Needs --rl-sync-fragment-pipeline equal to --fragments", ), cli_flag( "--rl-sync-push-int8-block", diff --git a/yeto/rl/engine/sync_overlap.py b/yeto/rl/engine/sync_overlap.py index 4c47f344..3a7f17ec 100644 --- a/yeto/rl/engine/sync_overlap.py +++ b/yeto/rl/engine/sync_overlap.py @@ -10,21 +10,39 @@ the next push ``theta - G_k`` equals round k+1's progress exactly once: round k's delta is already inside ``G_k`` and is not counted again. -Rollouts of round k+1 use weights one outer step older than the global -policy, so the run's minimum ``--rl-max-policy-age`` grows by -:data:`EXTRA_POLICY_AGE`. +Staleness is configured apart from ``--rl-max-policy-age``. The driver +publishes the boundary's cut to the island's engines before the next +generation (boundary, then publish), so every rollout is still generated by +the weights the trainer then trains: the rollout policy age stays 0 and the +sweep keeps refusing any other value. What overlap adds is a lag of the local +weights behind the global policy of :data:`OUTER_STEP_LAG` outer step (like +one extra DiLoCo local step); it is bound into the role-stream session +contract (``sync_train_overlap``), not into the rollout policy age, which in +Miles would switch on partial rollout (carry-over) or fully_async generation. """ from __future__ import annotations import numpy as np -EXTRA_POLICY_AGE = 1 +OUTER_STEP_LAG = 1 -def required_policy_age(minimum: int) -> int: - """The smallest accepted ``--rl-max-policy-age`` once overlap is on.""" - return int(minimum) + EXTRA_POLICY_AGE +def overlap_option_problems(*, max_policy_age: int, fragment_pipeline: int, fragments: int) -> list[str]: + """Launcher requirements of ``--rl-sync-overlap`` (empty when satisfied). + + * rollout policy age 0: see the module doc (overlap is not rollout staleness); + * syncer fragment pipeline = P: the syncer then PULLs every fragment of a + sweep at once, so the island pushes the whole sweep at the boundary. With + a smaller pipeline the later PULLs arrive while the next round trains and + wait for the next boundary, which serializes the sweep again.""" + problems = [] + if int(max_policy_age or 0) != 0: + problems.append("--rl-max-policy-age 0 with --rl-sync-overlap (overlap lags the global " + "policy by one outer step; rollouts stay on-policy)") + if int(fragment_pipeline) != int(fragments): + problems.append("--rl-sync-fragment-pipeline equal to --fragments with --rl-sync-overlap") + return problems def rebase_values(anchor: np.ndarray, current: np.ndarray, base: np.ndarray) -> np.ndarray: 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 09/62] 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 10/62] 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 11/62] 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 12/62] 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 152374b226bc86aa46fad3a0764e039e5c047342 Mon Sep 17 00:00:00 2001 From: michael <78113554+chowhsu@users.noreply.github.com> Date: Tue, 20 Oct 2026 17:30:04 +0000 Subject: [PATCH 13/62] test: overlap island with bf16 downlink deltas against the real syncer The overlap rebase uses the bf16-tracked anchor and pushes theta - anchor against that same anchor, so the syncer's f32 masters still receive every round's progress exactly once and no DownlinkResync happens. Co-Authored-By: Claude Opus 5.5 --- tests/test_rl_dense_sweep_real_syncer.py | 89 ++++++++++++++++++++++++ 1 file changed, 89 insertions(+) diff --git a/tests/test_rl_dense_sweep_real_syncer.py b/tests/test_rl_dense_sweep_real_syncer.py index 31589d67..ec4a66ec 100644 --- a/tests/test_rl_dense_sweep_real_syncer.py +++ b/tests/test_rl_dense_sweep_real_syncer.py @@ -432,3 +432,92 @@ def test_bf16_downlink_through_a_streaming_sink_matches_the_buffered_path(syncer sink=_BytesSink()) assert buffered[2] == streamed[2] assert buffered_updates == streamed_updates + + +def _overlap_island_run(binary: Path, tmp_path: Path, extra: tuple[str, ...], dtype: int, + sweeps: int, progress: dict[tuple[int, int], list[float]]): + """One overlap island (E3 item 3) against the real syncer. + + Round s trains before sweep s-1 comes back; at boundary s the island waits + for G_{s-1}, rebases ``theta <- G_{s-1} + (theta - base)`` and pushes + ``theta - G_{s-1}``. Returns (FINAL bytes, the island's anchors per sweep, + redials used).""" + process, port, profile_hash = _start_syncer(binary, tmp_path, *extra, total_steps=sweeps) + client = _client(port, profile_hash, dtype) + init = {fid: torch.tensor([float(fid) + 0.125 * i for i in range(NUMEL)]) for fid in range(FRAGMENTS)} + anchors: dict[int, dict[int, torch.Tensor]] = {} + held: dict[tuple[int, int], torch.Tensor] = {} + try: + client.start() + for fid in range(FRAGMENTS): + client.send_init(fid, _f32(init[fid].tolist())) + anchor = dict(init) # what the last push was measured against + theta = {fid: init[fid] + torch.tensor(progress[(1, fid)]) for fid in range(FRAGMENTS)} + for sweep in range(1, sweeps + 1): + if sweep > 1: + # Round `sweep` trained from `base` while sweep-1 was in flight. + base = dict(theta) + theta = {fid: base[fid] + torch.tensor(progress[(sweep, fid)]) for fid in range(FRAGMENTS)} + + def got_sweep(): + client.check_health() + held.update({(u.fragment_id, u.version): _tail(bytes(u.data)) + for u in client.drain_updates()}) + return all((fid, (sweep - 2) * FRAGMENTS + fid + 1) in held for fid in range(FRAGMENTS)) + + _wait(got_sweep, f"sweep {sweep - 1} never came back") + anchor = {fid: held[(fid, (sweep - 2) * FRAGMENTS + fid + 1)] for fid in range(FRAGMENTS)} + anchors[sweep - 1] = anchor + theta = {fid: anchor[fid] + (theta[fid] - base[fid]) for fid in range(FRAGMENTS)} + permits: dict[int, object] = {} + + def gather(): + client.check_health() + permits.update({p.fragment_id: p for p in client.drain_pulls()}) + return len(permits) == FRAGMENTS + + _wait(gather, f"no PULLs for sweep {sweep}") + for fid, permit in sorted(permits.items()): + base_version = 0 if sweep == 1 else (sweep - 2) * FRAGMENTS + fid + 1 + client.push_fragment(fid, permit.global_step, permit.round_attempt, base_version, + sweep, 1, 7, _f32((theta[fid] - anchor[fid]).tolist())) + manifest, finals = client.wait_for_final_fragments(timeout=30) + result = {f.fragment_id: _tail(bytes(f.data)) for f in finals} + client.acknowledge_finalization(manifest) + redials = client._reconnects_used + assert process.wait(timeout=15) == 0 + finally: + client.close() + log = _stop(process) + assert "syncer semantic profile" not in log, log + return init, result, anchors, redials + + +def test_overlap_with_bf16_downlink_counts_each_round_once_without_a_resync(syncer_binary, tmp_path): + """E3: overlap + bf16 BCAST deltas. The island rebases on a bf16-tracked + anchor, but it pushes ``theta - anchor`` against that same anchor, so the + syncer's f32 masters still get every round's progress exactly once; the + downlink error stays out of the masters and never triggers a resync.""" + sweeps = 4 + progress = {(s, fid): [0.5 + fid + 0.03125 * s + 0.001 * i for i in range(NUMEL)] + for s in range(1, sweeps + 1) for fid in range(FRAGMENTS)} + exact_dir, delta_dir = tmp_path / "f32", tmp_path / "delta" + exact_dir.mkdir() + delta_dir.mkdir() + init, exact, exact_anchors, exact_redials = _overlap_island_run( + syncer_binary, exact_dir, STREAMED, DTYPE_F32, sweeps, progress) + _, delta, delta_anchors, delta_redials = _overlap_island_run( + syncer_binary, delta_dir, (*STREAMED, "--rl-sync-bcast-delta-full-every", "3"), + DTYPE_F32 | DTYPE_FLAG_BCAST_DELTA, sweeps, progress) + assert exact_redials == delta_redials == 0, "no DownlinkResync expected" + for fid in range(FRAGMENTS): + want = init[fid] + sum(torch.tensor(progress[(s, fid)]) for s in range(1, sweeps + 1)) + # (anchor + p) - anchor in f32 is p up to an ulp of the anchor: no bf16 error. + assert torch.allclose(exact[fid], want, rtol=0, atol=1e-5), fid + assert torch.allclose(delta[fid], want, rtol=0, atol=1e-5), fid + for sweep, anchors in exact_anchors.items(): + got = delta_anchors[sweep][fid] + # one bf16 rounding of (G - previous anchor); |G - previous anchor| < 4 + assert torch.allclose(got, anchors[fid], rtol=0, atol=4.0 / 256), (sweep, fid) + assert any(not torch.equal(delta_anchors[s][f], exact_anchors[s][f]) + for s in delta_anchors for f in range(FRAGMENTS)), "the downlink must actually round" 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 14/62] 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 15/62] 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 16/62] 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 17/62] 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 18/62] 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 19/62] 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 20/62] 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 21/62] 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 22/62] 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 23/62] 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 24/62] 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 bc8f53b373ad2a61ddc762c8cbe133affadc2d51 Mon Sep 17 00:00:00 2001 From: michael <78113554+chowhsu@users.noreply.github.com> Date: Wed, 21 Oct 2026 01:24:08 +0000 Subject: [PATCH 25/62] syncer: send each fragment's FINAL version as an exact FULL broadcast With bf16 downlink deltas, the last BCAST of a fragment can carry the same version as its FINAL. Islands store one copy per (fragment, version) and refuse a second copy with other bytes, so the delta-tracked BCAST and the exact FINAL collided ('authoritative fragment replay changed bytes'), the islands redialed in a loop and finalization timed out (s26-e3-p06-D8). Broadcasts of a fragment's last version now always go out FULL. Co-Authored-By: Claude Opus 5.5 --- syncer/src/downlink.rs | 13 +++++++++++++ syncer/src/server.rs | 14 +++++++++++++- tests/test_rl_dense_sweep_real_syncer.py | 16 ++++++++++++++++ 3 files changed, 42 insertions(+), 1 deletion(-) diff --git a/syncer/src/downlink.rs b/syncer/src/downlink.rs index e303e119..32e6f82d 100644 --- a/syncer/src/downlink.rs +++ b/syncer/src/downlink.rs @@ -144,6 +144,19 @@ impl DownlinkAnchors { }) } + /// Encode a newly committed `version` of fragment `p` as FULL (exact f32) + /// and reset the shared anchor to it. + pub fn publish_full(&mut self, p: usize, version: u64, params: &[f32]) -> EncodedBcast { + match self.anchors.get_mut(p) { + Some(slot) => Self::reset(slot, p, version, params), + None => EncodedBcast { + prefix: prefix(p, version, KIND_FULL, version, &digest_f32(params)), + dtype: DTYPE_F32, + values: Arc::new(params.to_vec()), + }, + } + } + /// FULL re-send for recovery (HELLO / reconnect): the anchor every other /// learner holds, when it is at `version`; otherwise the exact params. pub fn resend(&self, p: usize, version: u64, params: &[f32]) -> EncodedBcast { diff --git a/syncer/src/server.rs b/syncer/src/server.rs index 46d4cb52..269c0b44 100644 --- a/syncer/src/server.rs +++ b/syncer/src/server.rs @@ -3495,7 +3495,19 @@ async fn publish_round( // bf16 downlink: encode this version once (advancing the shared // anchor) and send the same bytes to every learner. let published = match &st.downlink { - Some(anchors) => Some(anchors.lock().unwrap().publish(p, st.versions[p], &st.params[p])?), + Some(anchors) => { + // The last version of each fragment is also its FINAL version. + // Islands key stored fragments by (fragment, version) and + // refuse a second copy with other bytes, so that broadcast + // must carry the exact f32 masters (FULL), not a bf16 delta. + let last_version = t + st.params.len() as u64 > cfg.total_steps; + let mut anchors = anchors.lock().unwrap(); + Some(if last_version { + anchors.publish_full(p, st.versions[p], &st.params[p]) + } else { + anchors.publish(p, st.versions[p], &st.params[p])? + }) + } None => None, }; for group in current_groups(registry) { diff --git a/tests/test_rl_dense_sweep_real_syncer.py b/tests/test_rl_dense_sweep_real_syncer.py index ec4a66ec..eb9f22d3 100644 --- a/tests/test_rl_dense_sweep_real_syncer.py +++ b/tests/test_rl_dense_sweep_real_syncer.py @@ -521,3 +521,19 @@ def test_overlap_with_bf16_downlink_counts_each_round_once_without_a_resync(sync assert torch.allclose(got, anchors[fid], rtol=0, atol=4.0 / 256), (sweep, fid) assert any(not torch.equal(delta_anchors[s][f], exact_anchors[s][f]) for s in delta_anchors for f in range(FRAGMENTS)), "the downlink must actually round" + + +def test_bf16_downlink_sends_each_fragments_final_version_exact(syncer_binary, tmp_path): + """E3 (s26-e3-p06-D8): the last BCAST of a fragment has the same version as + its FINAL. Islands keep one copy per (fragment, version) and refuse a + second copy with other bytes, so that BCAST must be the exact f32 masters + even when the FULL cadence (here M=100) would send a bf16 delta.""" + delta, updates = _downlink_run(syncer_binary, tmp_path, 100, sweeps=4) + final_versions, finals = delta[1], delta[2] + assert final_versions == [7, 8] + # Fragment 0's FINAL version 7 was also broadcast; fragment 1's version 8 + # is the terminal step, which only FINAL carries. + assert updates[(0, 7)] == finals[0] + assert (1, 8) not in updates + # Earlier versions still travel as bf16 deltas. + assert any(updates[(0, v)] != finals[0] for v in (3, 5)) 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 26/62] 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 27/62] 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 28/62] 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 29/62] 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 30/62] 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 31/62] 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 32/62] 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 33/62] 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 34/62] 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 35/62] 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 36/62] 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 37/62] 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 38/62] 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 3dcbec86a95ececa6dd6170f7da6b7eac1dc0523 Mon Sep 17 00:00:00 2001 From: michael <78113554+chowhsu@users.noreply.github.com> Date: Wed, 21 Oct 2026 07:31:58 +0000 Subject: [PATCH 39/62] wan: per-stream stall timeout, rate-aware chunk striping, sync diagnostics E3 showed sweeps waiting on the slowest island's upload and one C32 run where a stream stopped moving with every socket still open. Island-side SyncerClient changes (all default off): - --wan-stall-timeout S: a watchdog fails the connection group when an outbound stream holds bytes, or an inbound chunked message stops growing, for S seconds; sockets also get TCP_USER_TIMEOUT=S. Recovery uses the existing redial path: the syncer replays its PULLs and the island pushes the fragment again at the same (global_step, round_attempt). - --wan-stream-schedule adaptive: each 4 MiB chunk goes to the data stream with the earliest estimated finish (backlog / EWMA rate) instead of round robin, so slow or stalled streams carry less. No wire or syncer change. - rl_full_parameter_round now records per-stream upload rates, the wait after this island's upload drained, and stall counts; scalars go to W&B on rl/step. A real-syncer test stalls one stream mid-sweep (rr and adaptive) and checks the redial keeps versions and finals bit for bit. Co-Authored-By: Claude Opus 5.5 --- tests/test_rl_dense_sweep_real_syncer.py | 96 ++++++++ tests/test_wan_streams.py | 229 ++++++++++++++++++ yeto/cli.py | 16 ++ yeto/launch/rl_island_task.py | 17 +- yeto/protocol.py | 144 ++++++++++- .../miles/full_parameter/role_stream.py | 10 + yeto/rl/adapters/miles/island_entry.py | 2 + yeto/rl/adapters/miles/run_flags.py | 2 + yeto/rl/engine/full_parameter_sync.py | 25 +- yeto/rl/engine/sync_timing.py | 39 +++ yeto/rl/wandb_rl.py | 30 +++ yeto/wan_streams.py | 190 +++++++++++++++ 12 files changed, 789 insertions(+), 11 deletions(-) create mode 100644 tests/test_wan_streams.py create mode 100644 yeto/wan_streams.py diff --git a/tests/test_rl_dense_sweep_real_syncer.py b/tests/test_rl_dense_sweep_real_syncer.py index eb9f22d3..92bf9b97 100644 --- a/tests/test_rl_dense_sweep_real_syncer.py +++ b/tests/test_rl_dense_sweep_real_syncer.py @@ -13,6 +13,7 @@ import shutil import socket import subprocess +import threading import time from pathlib import Path @@ -537,3 +538,98 @@ def test_bf16_downlink_sends_each_fragments_final_version_exact(syncer_binary, t assert (1, 8) not in updates # Earlier versions still travel as bf16 deltas. assert any(updates[(0, v)] != finals[0] for v in (3, 5)) + + +class _StallingSocket: + """Proxy whose ``send`` stops making progress once ``armed`` is set, + until the socket is closed.""" + + def __init__(self, sock: socket.socket, armed: threading.Event) -> None: + self._sock = sock + self._armed = armed + + def send(self, data) -> int: + if not self._armed.is_set(): + return self._sock.send(data) + while self._sock.fileno() != -1: + time.sleep(0.02) + raise OSError("stalled stream closed") + + def __getattr__(self, name): + return getattr(self._sock, name) + + +class _StallOnceClient(SyncerClient): + """Data stream 1 of the first connection group stalls once ``armed``.""" + + stalled_generation: int | None = None + + def __init__(self, *args, **kwargs) -> None: + super().__init__(*args, **kwargs) + self.armed = threading.Event() + + def _send_loop(self, gen, sock, q, meter): + if meter is self._meters[1] and self.stalled_generation in (None, gen): + self.stalled_generation = gen + sock = _StallingSocket(sock, self.armed) + return super()._send_loop(gen, sock, q, meter) + + +def _sweep_until_final(client: SyncerClient, sweeps: int): + """Answer every PULL (also PULLs the syncer replays after a redial) with + the deterministic delta of its (fragment, global_step) until FINAL.""" + answered: set[tuple[int, int, int]] = set() + deadline = time.monotonic() + 60 + while not client.finalizing.is_set(): + assert time.monotonic() < deadline, "sweeps did not finish" + client.check_health() + for permit in client.drain_pulls(): + key = (client.connection_generation, permit.global_step, permit.round_attempt) + if key in answered: + continue + fid = permit.fragment_id + sweep = (permit.global_step - 1) // FRAGMENTS + 1 + assert permit.global_step == (sweep - 1) * FRAGMENTS + fid + 1 and sweep <= sweeps + base = 0 if sweep == 1 else permit.global_step - FRAGMENTS + delta = [0.5 + fid + 0.03125 * sweep + 0.001 * i for i in range(NUMEL)] + client.push_fragment(fid, permit.global_step, permit.round_attempt, + base, sweep, 1, 7, _f32(delta)) + answered.add(key) + time.sleep(0.01) + return client.wait_for_final_fragments(timeout=30) + + +@pytest.mark.parametrize("schedule", ["rr", "adaptive"]) +def test_stalled_stream_redials_and_the_replayed_push_keeps_the_exact_result( + syncer_binary, tmp_path, schedule): + """E3b: a data stream that stops moving bytes (C32 first run) is dropped + after the stall timeout; the syncer replays its PULL on the new group and + the island pushes that fragment again, so versions and finals match a run + without the stall bit for bit.""" + exact = _run_sweeps(syncer_binary, tmp_path, STREAMED, sweeps=3) + stall_dir = tmp_path / "stall" + stall_dir.mkdir() + process, port, profile_hash = _start_syncer(syncer_binary, stall_dir, *STREAMED, total_steps=3) + client = _StallOnceClient(("127.0.0.1", port), 0, _layout(), dtype=DTYPE_F32, num_streams=2, + connect_timeout=10, max_reconnects=5, finalization_timeout=20, + syncer_profile_hash=bytes.fromhex(profile_hash), + stream_schedule=schedule, stream_stall_timeout=1.0) + try: + client.start() + for fid in range(FRAGMENTS): + client.send_init(fid, _f32([float(fid) + 0.125 * i for i in range(NUMEL)])) + # Stall only the pushes: INIT is sent once and not replayed after a redial. + _wait(lambda: client.send_progress().queued_chunks == 0 and client._bcasts.qsize() >= FRAGMENTS, + "initial broadcast") + client.armed.set() + manifest, finals = _sweep_until_final(client, sweeps=3) + result = {f.fragment_id: bytes(f.data) for f in finals} + client.acknowledge_finalization(manifest) + assert process.wait(timeout=15) == 0 + progress = client.send_progress() + finally: + client.close() + _stop(process) + assert progress.stalls >= 1 and progress.reconnects >= 1 + assert list(manifest.versions) == exact[1] + assert result == exact[2] diff --git a/tests/test_wan_streams.py b/tests/test_wan_streams.py new file mode 100644 index 00000000..4e08efb0 --- /dev/null +++ b/tests/test_wan_streams.py @@ -0,0 +1,229 @@ +"""S26 E3b: per-stream accounting, stall watchdog and rate-aware striping.""" + +from __future__ import annotations + +import queue +import socket +import threading +import time + +import pytest + +from yeto.protocol import SyncerClient +from yeto.rl.engine.sync_timing import SyncTiming, upload_fields +from yeto.wan_streams import ( + StreamMeter, + StreamSnapshot, + UploadWindow, + pick_stream, + upload_window, +) + + +class _Clock: + def __init__(self) -> None: + self.now = 100.0 + + def __call__(self) -> float: + return self.now + + +def _meter_with_rate(clock: _Clock, rate_bps: float, backlog: int = 0) -> StreamMeter: + meter = StreamMeter(clock) + meter.enqueued(int(rate_bps)) + meter.begin() + clock.now += 1.0 + meter.progress(int(rate_bps)) + meter.end() + if backlog: + meter.enqueued(backlog) + return meter + + +def test_round_robin_ignores_rates(): + clock = _Clock() + meters = [_meter_with_rate(clock, 1e6), _meter_with_rate(clock, 1e3)] + assert [pick_stream("rr", meters, ticket) for ticket in range(4)] == [0, 1, 0, 1] + + +def test_adaptive_spreads_an_idle_start_over_every_stream(): + meters = [StreamMeter() for _ in range(4)] + picked = [] + for ticket in range(4): + index = pick_stream("adaptive", meters, ticket) + meters[index].enqueued(4 << 20) + picked.append(index) + assert sorted(picked) == [0, 1, 2, 3] + + +def test_adaptive_gives_a_slow_stream_fewer_chunks_in_proportion_to_its_rate(): + clock = _Clock() + fast, slow = 40e6, 5e6 + meters = [_meter_with_rate(clock, fast), _meter_with_rate(clock, slow)] + counts = [0, 0] + for ticket in range(90): + index = pick_stream("adaptive", meters, ticket) + meters[index].enqueued(4 << 20) + counts[index] += 1 + # Equal estimated finish time: chunk counts follow the 8:1 rate ratio. + assert counts[1] == pytest.approx(90 / 9, abs=1) + assert counts[0] == pytest.approx(90 * 8 / 9, abs=1) + + +def test_adaptive_starves_a_stream_whose_backlog_is_stuck(): + clock = _Clock() + meters = [_meter_with_rate(clock, 10e6), _meter_with_rate(clock, 10e6, backlog=64 << 20)] + assert {pick_stream("adaptive", meters, ticket) for ticket in range(10)} == {0} + + +def test_unknown_schedule_is_refused(): + with pytest.raises(ValueError, match="schedule"): + pick_stream("random", [StreamMeter()], 0) + with pytest.raises(ValueError, match="stream_schedule"): + SyncerClient(("127.0.0.1", 1), 0, _layout(), stream_schedule="random") + with pytest.raises(ValueError, match="stall_timeout"): + SyncerClient(("127.0.0.1", 1), 0, _layout(), stream_stall_timeout=0) + + +def test_meter_reports_a_stall_only_while_bytes_wait(): + clock = _Clock() + meter = StreamMeter(clock) + clock.now += 50 + assert meter.stalled_for() == 0.0 # idle is not stalled + meter.enqueued(10) + meter.begin() + meter.progress(4) + clock.now += 7 + assert meter.stalled_for() == pytest.approx(7) + meter.progress(6) + meter.end() + clock.now += 30 + assert meter.stalled_for() == 0.0 and meter.idle + meter.enqueued(5) + meter.drop_pending() + assert meter.idle + + +def test_upload_window_and_fields_report_rates_and_the_wait_after_upload(): + before = (StreamSnapshot(100, 1.0, 0, None), StreamSnapshot(0, 0.0, 0, None)) + after = (StreamSnapshot(100 + 40_000_000, 3.0, 0, 2e7), StreamSnapshot(5_000_000, 1.0, 0, 5e6)) + window = upload_window(before, after) + assert window == UploadWindow((40_000_000, 5_000_000), (2.0, 1.0)) + assert window.stream_mbps() == [20.0, 5.0] + fields = upload_fields(window, drained_ts=1010.0, timing=SyncTiming(1000.0, 1025.0), stalls=1) + assert fields["sync/upload_mbps"] == pytest.approx(22.5) + assert fields["sync/upload_min_stream_mbps"] == 5.0 + assert fields["sync/upload_active_streams"] == 2 + assert fields["sync/wait_after_upload_seconds"] == pytest.approx(15.0) + assert fields["sync/stream_stalls"] == 1 + # The upload drained before this boundary's submit: only the submit counts. + early = upload_fields(window, drained_ts=990.0, timing=SyncTiming(1000.0, 1025.0), stalls=0) + assert early["sync/wait_after_upload_seconds"] == pytest.approx(25.0) + unknown = upload_fields(window, drained_ts=None, timing=SyncTiming(1000.0, None), stalls=0) + assert unknown["sync/wait_after_upload_seconds"] is None + + +def test_sync_round_diagnostics_reach_wandb_on_the_policy_axis(): + from yeto.rl.wandb_rl import event_metrics + + metrics = event_metrics({ + "event": "rl_full_parameter_round", "policy_version": 3, + "sync/sync_wait_seconds": 40.0, "sync/upload_mbps": 22.5, + "sync/upload_min_stream_mbps": 5.0, "sync/upload_stream_mbps": [20.0, 5.0], + "sync/wait_after_upload_seconds": 15.0, "sync/global_policy_hash": "ab", + }) + assert metrics == {"rl/step": 3, "sync/sync_wait_seconds": 40.0, "sync/upload_mbps": 22.5, + "sync/upload_min_stream_mbps": 5.0, "sync/wait_after_upload_seconds": 15.0} + assert event_metrics({"event": "rl_full_parameter_round", "policy_version": 3}) is None + + +# -- socket-level behaviour --------------------------------------------------- + + +def _layout(): + from yeto.fragments import MERGE_RDA, Fragment, FragmentLayout + + return FragmentLayout([Fragment(MERGE_RDA, [("w", 4)])]) + + +def test_send_loop_meters_bytes_and_marks_the_upload_drained(): + left, right = socket.socketpair() + client = SyncerClient(("127.0.0.1", 1), 0, _layout(), num_streams=1) + q: queue.Queue = queue.Queue() + meter = client._meters[1] + worker = threading.Thread(target=client._send_loop, args=(client._gen, left, q, meter)) + worker.start() + try: + meter.enqueued(3 << 20) + q.put(type("Item", (), {"data": b"x" * (3 << 20), "sent": None})()) + received = 0 + while received < 3 << 20: + received += len(right.recv(1 << 20)) + q.put(None) + worker.join(timeout=5) + snap = client.stream_snapshots()[0] + assert snap.bytes_sent == 3 << 20 and snap.backlog_bytes == 0 + assert client.upload_drained_at is not None + finally: + left.close() + right.close() + + +def test_watchdog_fails_the_group_when_a_stream_stops_moving(): + client = SyncerClient(("127.0.0.1", 1), 0, _layout(), num_streams=2, stream_stall_timeout=0.2) + a, b = socket.socketpair() + try: + client._socks = [a] + client._connected.set() + meter = client._meters[2] + meter.enqueued(10) + meter.begin() + assert client._stalled_stream() is None + time.sleep(0.3) + assert "outbound stream 2" in client._stalled_stream() + client._watchdog = threading.Thread(target=client._watch_stalls, daemon=True) + client._watchdog.start() + assert client._failure.wait(3), "the watchdog must fail the live group" + assert client.send_progress().stalls == 1 + assert "no progress" in str(client._last_err) + finally: + client._closed.set() + a.close() + b.close() + + +def test_watchdog_notices_an_inbound_message_that_stops_growing(): + client = SyncerClient(("127.0.0.1", 1), 0, _layout(), num_streams=1, stream_stall_timeout=0.2) + client._reasm_touched[5] = time.monotonic() - 1.0 + assert "inbound message 5" in client._stalled_stream() + client._clear_partial_reassemblies() + assert client._stalled_stream() is None + + +def test_stall_timeout_sets_tcp_user_timeout(): + if not hasattr(socket, "TCP_USER_TIMEOUT"): + pytest.skip("Linux-only socket option") + from yeto.protocol import set_tcp_user_timeout + + with socket.socket() as sock: + set_tcp_user_timeout(sock, 2.5) + assert sock.getsockopt(socket.IPPROTO_TCP, socket.TCP_USER_TIMEOUT) == 2500 + + +def test_transport_flags_reach_the_island_command_only_when_set(): + from argparse import Namespace + + from yeto.launch.rl_island_task import _wan_transport_flags + + assert _wan_transport_flags(Namespace()) == "" + assert _wan_transport_flags(Namespace(wan_stream_schedule="adaptive", wan_stall_timeout=60.0)) == ( + " --wan-stream-schedule adaptive --wan-stall-timeout 60.0") + with pytest.raises(ValueError, match="wan-stall-timeout"): + _wan_transport_flags(Namespace(wan_stall_timeout=-1.0)) + + +def test_island_parser_accepts_the_transport_flags(): + from yeto.rl.engine.run_flags import ISLAND, rl_option + + assert rl_option("wan_stream_schedule", ISLAND) == "--wan-stream-schedule" + assert rl_option("wan_stall_timeout", ISLAND) == "--wan-stall-timeout" diff --git a/yeto/cli.py b/yeto/cli.py index 7c68497f..40abdc1f 100644 --- a/yeto/cli.py +++ b/yeto/cli.py @@ -438,6 +438,22 @@ def int_or_auto(value: str): "deltas (~4x less learner egress; broadcasts stay bf16)", ) sync.add_argument("--wan-streams", type=int, default=4, help="parallel TCP streams per learner") + sync.add_argument( + "--wan-stream-schedule", + choices=("rr", "adaptive"), + default="rr", + help="full-parameter RL islands: rr = round-robin chunks over the streams; " + "adaptive = send each chunk to the stream with the earliest estimated " + "finish (backlog / measured rate), so slow streams carry less", + ) + sync.add_argument( + "--wan-stall-timeout", + type=float, + default=0.0, + help="full-parameter RL islands: drop and redial the connection group when " + "one stream holds bytes with no progress for this many seconds (also sets " + "TCP_USER_TIMEOUT); the fragment is pushed again on the replayed PULL. 0 = off", + ) infra = p.add_argument_group("infrastructure") infra.add_argument( diff --git a/yeto/launch/rl_island_task.py b/yeto/launch/rl_island_task.py index 532fff20..a51d6ea1 100644 --- a/yeto/launch/rl_island_task.py +++ b/yeto/launch/rl_island_task.py @@ -51,6 +51,20 @@ def submitted_unix() -> float: return time.time() +def _wan_transport_flags(args) -> str: + """S26 E3b island transport options; emitted only when not the default.""" + flags = "" + schedule = getattr(args, "wan_stream_schedule", "rr") + if schedule != "rr": + flags += f" --wan-stream-schedule {schedule}" + stall = float(getattr(args, "wan_stall_timeout", 0.0) or 0.0) + if stall < 0: + raise ValueError("--wan-stall-timeout must be >= 0 (0 = off)") + if stall > 0: + flags += f" --wan-stall-timeout {stall!r}" + return flags + + def _rl_trainable_island_flags(args, learner_id: int, num_learners: int) -> str: """LoRA rank/targets; for the full-parameter preset the sweep identity (the probed layout hash, the frozen roster size and the syncer profile hash); @@ -166,7 +180,8 @@ def make_miles_island_task( if float(getattr(args, "rl_optimizer_offload_fraction", 0.0) or 0.0) > 0 else "") + f" --seed {args.seed}" f" --wan-streams {args.wan_streams}" - " --miles-root ~/miles" + + _wan_transport_flags(args) + + " --miles-root ~/miles" ) if getattr(args, "wandb", False): flags += " --wandb" diff --git a/yeto/protocol.py b/yeto/protocol.py index fc2b6d73..42cfd3f8 100644 --- a/yeto/protocol.py +++ b/yeto/protocol.py @@ -41,6 +41,7 @@ from .downlink import DownlinkAnchors, DownlinkChunkSink from .downlink import tensor_nbytes as downlink_tensor_nbytes from .fragments import MERGE_ISO, FragmentLayout +from .wan_streams import STREAM_SCHEDULES, StreamMeter, StreamSnapshot, StreamStalledError, pick_stream MAGIC = 0xD170C0DE PROTOCOL_VERSION = 4 @@ -571,6 +572,7 @@ class SendProgress: queued_chunks: int reconnects: int connected: bool + stalls: int = 0 @dataclass @@ -594,6 +596,10 @@ class SyncerClient: forever); once exhausted, check_health raises. """ + # Transport defaults (S26 E3b); __init__ overrides them per client. + stream_schedule: str = "rr" + stream_stall_timeout: float | None = None + def __init__( self, addr: tuple[str, int], @@ -609,7 +615,13 @@ def __init__( hmac_key: bytes | None = None, compat_group: str | None = None, push_i8_block: int = DEFAULT_I8_BLOCK, + stream_schedule: str = "rr", + stream_stall_timeout: float | None = None, ): + if stream_schedule not in STREAM_SCHEDULES: + raise ValueError(f"stream_schedule must be one of {STREAM_SCHEDULES}, got {stream_schedule!r}") + if stream_stall_timeout is not None and not stream_stall_timeout > 0: + raise ValueError("stream_stall_timeout must be positive or None") if not 0 <= learner_id <= 0xFFFF_FFFF: raise ValueError(f"learner_id must fit u32, got {learner_id}") if not valid_session_dtype(dtype): @@ -638,6 +650,15 @@ def __init__( self._downlink = DownlinkAnchors(layout.num_fragments) if bcast_delta(dtype) else None self._bcast_prefix = 12 + (DOWNLINK_HEADER_LEN if self._downlink is not None else 0) self.num_streams = num_streams + # S26 E3b: chunk placement over the data streams and the per-stream + # no-progress timeout (None = off, the historical behaviour). + self.stream_schedule = stream_schedule + self.stream_stall_timeout = stream_stall_timeout + # Index 0 is the control socket; the meters outlive connection groups + # so their counters are cumulative for the whole session. + self._meters = [StreamMeter() for _ in range(num_streams + 1)] + self._upload_drained_at: float | None = None + self._stalls = 0 self.session_contract_hash = session_contract_hash self.syncer_profile_hash = syncer_profile_hash # secret-handling-hardening: HELLO carries an HMAC when the island has @@ -682,6 +703,7 @@ def __init__( | _StreamedReassembly, ] = {} self._reasm_lock = threading.Lock() + self._reasm_touched: dict[int, float] = {} # msg_id -> last chunk (monotonic) self._inbound_chunk_sink: object | None = None self._msg_id = itertools.count() self._rr = itertools.count() @@ -702,6 +724,7 @@ def __init__( self._bytes_sent_lock = threading.Lock() self._reset_bcasts_on_reconnect = False self._supervisor: threading.Thread | None = None + self._watchdog: threading.Thread | None = None # -- lifecycle ------------------------------------------------------------- @@ -710,6 +733,8 @@ def _dial(self, timeout: float) -> socket.socket: sock.settimeout(None) sock.setsockopt(socket.IPPROTO_TCP, socket.TCP_NODELAY, 1) enable_tcp_keepalive(sock) + if self.stream_stall_timeout is not None: + set_tcp_user_timeout(sock, self.stream_stall_timeout) return sock def _connect_one(self) -> socket.socket: @@ -740,6 +765,9 @@ def start(self) -> None: self._connect_group(patient=True) self._supervisor = threading.Thread(target=self._supervise, name="yeto-supervisor", daemon=True) self._supervisor.start() + if self.stream_stall_timeout is not None: + self._watchdog = threading.Thread(target=self._watch_stalls, name="yeto-stall-watchdog", daemon=True) + self._watchdog.start() def _connect_group(self, patient: bool) -> None: """Dial control + data sockets, send HELLO/DATA_HELLO, start workers. @@ -797,7 +825,7 @@ def _connect_group(self, patient: bool) -> None: q: queue.Queue[_Outbound | None] = queue.Queue(maxsize=256) self._queues.append(q) ts = threading.Thread( - target=self._send_loop, args=(gen, s, q), name=f"yeto-send-{i}", daemon=True + target=self._send_loop, args=(gen, s, q, self._meters[i]), name=f"yeto-send-{i}", daemon=True ) tr = threading.Thread(target=self._recv_loop, args=(gen, s), name=f"yeto-recv-{i}", daemon=True) self._threads += [ts, tr] @@ -822,6 +850,8 @@ def _teardown_group(self) -> None: pass for s in socks: _close_socket(s) + for meter in self._meters: + meter.drop_pending() self._clear_partial_reassemblies() self._drain(self._pulls) # stale pull requests; answering them is pointless with self._lock: @@ -914,8 +944,21 @@ def send_progress(self) -> SendProgress: queued_chunks=queued, reconnects=self._reconnects_used, connected=connected, + stalls=self._stalls, ) + def stream_snapshots(self) -> tuple[StreamSnapshot, ...]: + """Cumulative counters of each data stream (control stream excluded; + with ``num_streams == 0`` the control stream carries the data).""" + meters = self._meters[1:] if self.num_streams else self._meters[:1] + return tuple(meter.snapshot() for meter in meters) + + @property + def upload_drained_at(self) -> float | None: + """Unix time when the data streams last became idle with nothing + queued, i.e. when this island's latest upload left the host.""" + return self._upload_drained_at + # -- learner-facing API ------------------------------------------------------ def send_init(self, fragment_id: int, tensor_bytes: bytes) -> None: @@ -1279,6 +1322,8 @@ def _enqueue( if gen is not None and gen != self._gen: return False q = self._queues[stream] + meter = self._meters[stream] + meter.enqueued(len(data)) # before put: the sender may drain it at once try: q.put( _Outbound(data, sent), @@ -1286,8 +1331,16 @@ def _enqueue( ) # bounded wait, then re-check the group return True except queue.Full: + meter.enqueued(-len(data)) continue + def _next_data_stream(self) -> int: + """Socket index (1-based; 0 only without data streams) for the next + CHUNK envelope, by ``stream_schedule``.""" + if self.num_streams == 0: + return 0 + return 1 + pick_stream(self.stream_schedule, self._meters[1:], next(self._rr)) + def _send_large(self, msg_type: int, payload: bytes) -> None: inner = self._frame(msg_type, payload) if self.num_streams == 0: @@ -1302,7 +1355,7 @@ def _send_large(self, msg_type: int, payload: bytes) -> None: for offset in range(0, total, CHUNK_SIZE): chunk = inner[offset : offset + CHUNK_SIZE] envelope = self._frame(MSG_CHUNK, _CHUNK_HEAD.pack(msg_id, total, offset) + chunk) - stream = 1 + next(self._rr) % self.num_streams + stream = self._next_data_stream() if not self._enqueue(stream, envelope, gen=gen): return # group died mid-message; drop the remainder @@ -1423,11 +1476,7 @@ def _send_large_parts( before_last_enqueue = None callback() - stream = ( - 0 - if self.num_streams == 0 - else 1 + next(self._rr) % self.num_streams - ) + stream = self._next_data_stream() if not self._enqueue(stream, envelope, gen=gen): if enqueued_chunks: raise PartialMessageGenerationLost( @@ -1487,20 +1536,85 @@ def _protocol_failed(self, gen: int, exc: BaseException) -> None: with self._final_cond: self._final_cond.notify_all() - def _send_loop(self, gen: int, sock: socket.socket, q: queue.Queue) -> None: + def _send_loop( + self, + gen: int, + sock: socket.socket, + q: queue.Queue, + meter: StreamMeter | None = None, + ) -> None: + if meter is None: # a sender outside a connection group (tests) + meter = StreamMeter() try: while True: item = q.get() if item is None: return - sock.sendall(item.data) + # send() instead of sendall() so the meter sees byte progress + # inside one 4 MiB envelope (stall watchdog, per-stream rate). + view = memoryview(item.data) + meter.begin() + try: + while view: + sent = sock.send(view) + meter.progress(sent) + view = view[sent:] + finally: + meter.end() + view.release() with self._bytes_sent_lock: self._bytes_sent += len(item.data) + if all(m.idle for m in self._meters): + self._upload_drained_at = time.time() if item.sent is not None: item.sent.set() except BaseException as e: self._socket_failed(gen, e) + def _stalled_stream(self) -> str | None: + """Describe the first stream without progress for the stall timeout.""" + limit = self.stream_stall_timeout + now = time.monotonic() + for index, meter in enumerate(self._meters): + idle_for = meter.stalled_for(now) + if idle_for > limit: + return f"outbound stream {index} made no progress for {idle_for:.0f}s" + with self._reasm_lock: + touched = dict(self._reasm_touched) + for msg_id, last in touched.items(): + if now - last > limit: + return f"inbound message {msg_id} got no chunk for {now - last:.0f}s" + return None + + def _watch_stalls(self) -> None: + """Fail the live group when one stream stops moving bytes. + + A stalled TCP stream can sit for many minutes before the kernel gives + up (C32: 900 s with every socket still open). The failure goes + through the normal reconnect path, so the syncer replays its PULLs + and the island pushes the interrupted fragment again.""" + period = min(1.0, self.stream_stall_timeout / 4) + while not self._closed.wait(period): + if self.shutdown.is_set(): + return + with self._lock: + gen = self._gen + live = self._connected.is_set() + if not live: + continue + reason = self._stalled_stream() + if reason is None: + continue + with self._lock: + if gen != self._gen: + continue + self._stalls += 1 + socks = list(self._socks) + _connect_log(f"{reason} (learner {self.learner_id}); dropping the connection group") + self._socket_failed(gen, StreamStalledError(reason)) + for sock in socks: # unblock senders stuck in send() + _close_socket(sock) + def _recv_loop(self, gen: int, sock: socket.socket) -> None: try: while True: @@ -1569,6 +1683,7 @@ def _reassemble( sink=sink, transaction=sink.begin_message(msg_id, total), ) + self._reasm_touched[msg_id] = time.monotonic() partial = self._reasm[msg_id] if isinstance(partial, _StreamedReassembly): if partial.total != total: @@ -1588,6 +1703,7 @@ def _reassemble( raise if completed is not None: self._reasm.pop(msg_id, None) + self._reasm_touched.pop(msg_id, None) return completed buf, filled, ranges = partial if len(buf) != total: @@ -1602,6 +1718,7 @@ def _reassemble( if filled[0] != total: raise ValueError("chunk byte count exceeds frame length") del self._reasm[msg_id] + self._reasm_touched.pop(msg_id, None) magic, msg_type, length = _HEADER.unpack_from(buf) if magic != MAGIC or length != total - _HEADER.size: raise ValueError("corrupt reassembled frame") @@ -1769,6 +1886,7 @@ def _clear_partial_reassemblies(self) -> None: with self._reasm_lock: partials = list(self._reasm.values()) self._reasm.clear() + self._reasm_touched.clear() for partial in partials: if isinstance(partial, _StreamedReassembly): try: @@ -2043,6 +2161,14 @@ def enable_tcp_keepalive(sock: socket.socket) -> None: sock.setsockopt(socket.IPPROTO_TCP, option, value) +def set_tcp_user_timeout(sock: socket.socket, seconds: float) -> None: + """Bound how long unacknowledged data may sit before the kernel resets + the connection (Linux TCP_USER_TIMEOUT); other platforms keep defaults.""" + option = getattr(socket, "TCP_USER_TIMEOUT", None) + if option is not None: + sock.setsockopt(socket.IPPROTO_TCP, option, max(1, int(seconds * 1000))) + + def _close_socket(sock: socket.socket) -> None: try: sock.shutdown(socket.SHUT_RDWR) diff --git a/yeto/rl/adapters/miles/full_parameter/role_stream.py b/yeto/rl/adapters/miles/full_parameter/role_stream.py index 1855ad78..fa746fcb 100644 --- a/yeto/rl/adapters/miles/full_parameter/role_stream.py +++ b/yeto/rl/adapters/miles/full_parameter/role_stream.py @@ -88,6 +88,10 @@ class RoleStreamConfig: # P0-4.8: the island's AlgorithmSpec hash; placement capabilities stay out. algorithm_spec_sha256: str wan_streams: int = 4 + # S26 E3b transport options (yeto/wan_streams.py); not part of the session + # contract: they change how bytes travel, never which bytes. + wan_stream_schedule: str = "rr" + wan_stall_timeout: float = 0.0 # seconds; 0 = off wait_timeout: float = 900.0 poll_seconds: float = 0.01 max_fragment_bytes: int = 2 << 30 @@ -158,6 +162,10 @@ def __post_init__(self) -> None: value = getattr(self, name) if type(value) is not int or value < 1: raise ValueError(f"role-stream {name} must be positive") + if self.wan_stream_schedule not in ("rr", "adaptive"): + raise ValueError("role-stream wan stream schedule must be rr or adaptive") + if not self.wan_stall_timeout >= 0: + raise ValueError("role-stream wan stall timeout must be >= 0") if self.push_dtype not in ("f32", "bf16", "int8"): raise ValueError("role-stream push dtype must be f32, bf16 or int8") if type(self.push_int8_block) is not int or not I8_BLOCK_MIN <= self.push_int8_block <= I8_BLOCK_MAX: @@ -379,6 +387,8 @@ async def initialize(self, group) -> None: dtype=self.config.session_dtype, push_i8_block=self.config.push_int8_block, num_streams=self.config.wan_streams, + stream_schedule=self.config.wan_stream_schedule, + stream_stall_timeout=self.config.wan_stall_timeout or None, max_reconnects=None, session_contract_hash=role_stream_session_contract_hash( adapter.layout.layout_hash, diff --git a/yeto/rl/adapters/miles/island_entry.py b/yeto/rl/adapters/miles/island_entry.py index 6ff6a2e8..01367ff3 100644 --- a/yeto/rl/adapters/miles/island_entry.py +++ b/yeto/rl/adapters/miles/island_entry.py @@ -838,6 +838,8 @@ def _dense_full_sync_config( pipeline_depth=1, algorithm_spec_sha256=algorithm_spec_sha256, wan_streams=args.wan_streams, + wan_stream_schedule=getattr(args, "wan_stream_schedule", "rr"), + wan_stall_timeout=float(getattr(args, "wan_stall_timeout", 0.0) or 0.0), push_dtype=getattr(args, "sweep_push_dtype", "f32"), push_int8_block=int(getattr(args, "sweep_push_int8_block", 128)), bcast_delta=bool(getattr(args, "bcast_delta", False)), diff --git a/yeto/rl/adapters/miles/run_flags.py b/yeto/rl/adapters/miles/run_flags.py index ff2830bc..30778809 100644 --- a/yeto/rl/adapters/miles/run_flags.py +++ b/yeto/rl/adapters/miles/run_flags.py @@ -416,6 +416,8 @@ island_flag("--seq-len", type=int, required=True), island_flag("--seed", type=int, required=True), island_flag("--wan-streams", type=int, default=4), + island_flag("--wan-stream-schedule", choices=("rr", "adaptive"), default="rr"), + island_flag("--wan-stall-timeout", type=float, default=0.0), island_flag("--miles-root", required=True), island_flag( "--miles-commit", diff --git a/yeto/rl/engine/full_parameter_sync.py b/yeto/rl/engine/full_parameter_sync.py index 744f3053..af2a1a1d 100644 --- a/yeto/rl/engine/full_parameter_sync.py +++ b/yeto/rl/engine/full_parameter_sync.py @@ -33,9 +33,10 @@ from yeto.protocol import PartialMessageGenerationLost from yeto.rl.elastic_client import LeaveReason +from yeto.wan_streams import upload_window from .driver import SyncBoundary, SyncStart -from .sync_timing import SyncTiming, unix_now +from .sync_timing import SyncTiming, unix_now, upload_fields if TYPE_CHECKING: from .driver import IslandDriver @@ -529,6 +530,8 @@ def __init__( self._timing = SyncTiming(None, None) # L5.6, set by each boundary # Overlap: fragment versions before the sweep still in flight. self._pending_sweep: tuple[int, ...] | None = None + # S26 E3b: actor data-stream counters at the previous boundary. + self._upload_before: tuple | None = None @property def actor(self) -> RoleStream: @@ -537,6 +540,7 @@ def actor(self) -> RoleStream: def start(self, driver: IslandDriver) -> SyncStart: self._driver = driver self.run(self.pairing.gather_roles(lambda role, stream: stream.initialize(self.groups[role]))) + self._upload_before = self._stream_snapshots() return SyncStart(state=self.actor.resident_state(0), rollout_id=0) def boundary(self, driver: IslandDriver, *, rollout_id: int, stats: Any) -> SyncBoundary: @@ -554,10 +558,29 @@ def boundary(self, driver: IslandDriver, *, rollout_id: int, stats: Any) -> Sync "sync/global_policy_hash": state.policy_tensor_hash(), "sync/inflight_pairs": len(self.pairing.inflight), **self._timing.fields(), + **self._upload_fields(), }, ) return SyncBoundary(state=state, stop=self.terminal) + def _stream_snapshots(self) -> tuple | None: + snapshots = getattr(getattr(self.actor, "client", None), "stream_snapshots", None) + return snapshots() if callable(snapshots) else None + + def _upload_fields(self) -> dict[str, object]: + """Actor upload rates of this boundary and the wait after the upload.""" + after = self._stream_snapshots() + before, self._upload_before = self._upload_before, after + if after is None or before is None: + return {} + client = self.actor.client + return upload_fields( + upload_window(before, after), + drained_ts=client.upload_drained_at, + timing=self._timing, + stalls=client.send_progress().stalls, + ) + async def _boundary(self, rollout_id: int) -> bool: if self.pairing.profile.overlap: return await self._overlap_boundary(rollout_id) diff --git a/yeto/rl/engine/sync_timing.py b/yeto/rl/engine/sync_timing.py index eaa4ad9b..eae22025 100644 --- a/yeto/rl/engine/sync_timing.py +++ b/yeto/rl/engine/sync_timing.py @@ -44,3 +44,42 @@ def fields(self) -> dict[str, float | None]: def unix_now() -> float: return time.time() + + +def upload_fields( + window, + *, + drained_ts: float | None, + timing: SyncTiming, + stalls: int, +) -> dict[str, object]: + """Per-boundary upload diagnostics of one island (S26 E3b). + + ``window`` is a :class:`yeto.wan_streams.UploadWindow` covering this + boundary. ``sync/upload_stream_mbps`` lists each data stream's rate over + the time it was busy; ``sync/upload_mbps`` is the bytes over the longest + busy stream (the streams run in parallel). ``sync/wait_after_upload_seconds`` + is the time from this island's last byte leaving to the global policy + being ready: waiting for slower islands, the merge and the broadcast.""" + rates = window.stream_mbps() + busy = max(window.busy_seconds_per_stream, default=0.0) + measured = [rate for rate in rates if rate is not None] + wait_after = None + if ( + drained_ts is not None + and timing.submit_ts is not None + and timing.global_ready_ts is not None + and drained_ts <= timing.global_ready_ts + ): + wait_after = timing.global_ready_ts - max(drained_ts, timing.submit_ts) + return { + "sync/upload_bytes": window.total_bytes, + "sync/upload_busy_seconds": busy, + "sync/upload_mbps": (window.total_bytes / busy / 1e6) if busy > 0 else None, + "sync/upload_stream_mbps": [None if rate is None else round(rate, 3) for rate in rates], + "sync/upload_min_stream_mbps": min(measured) if measured else None, + "sync/upload_active_streams": len(measured), + "sync/upload_drained_ts": drained_ts, + "sync/wait_after_upload_seconds": wait_after, + "sync/stream_stalls": stalls, + } diff --git a/yeto/rl/wandb_rl.py b/yeto/rl/wandb_rl.py index 3467c503..60a08c7b 100644 --- a/yeto/rl/wandb_rl.py +++ b/yeto/rl/wandb_rl.py @@ -28,6 +28,7 @@ "train/*": TRAIN_STEP_KEY, "rl/*": STEP_KEY, "eval/*": STEP_KEY, + "sync/*": STEP_KEY, } # Closed output surface. Input aliases are limited to namespaced scalar fields @@ -115,11 +116,40 @@ def _non_negative_step(event: dict[str, Any], aliases: tuple[str, ...]) -> int | return None +# S26 E3b: per-sweep sync diagnostics of a full-parameter island, on the +# global policy-version axis (rl/step). +_SYNC_ROUND_OUTPUTS = ( + "sync/sync_wait_seconds", + "sync/upload_mbps", + "sync/upload_min_stream_mbps", + "sync/upload_active_streams", + "sync/wait_after_upload_seconds", + "sync/stream_stalls", +) + + +def _sync_round_metrics(event: dict[str, Any]) -> dict[str, Any] | None: + step = _non_negative_step(event, ("policy_version",)) + if step is None: + return None + metrics = { + key: value + for key in _SYNC_ROUND_OUTPUTS + if (value := _finite_scalar(event, (key,))) is not None + } + if not metrics: + return None + metrics[STEP_KEY] = step + return metrics + + def event_metrics(event: dict[str, Any]) -> dict[str, Any] | None: """Project one event onto the explicit scalar-only W&B surface.""" global _derived_step_logged event_name = event.get("event") + if event_name == "rl_full_parameter_round": + return _sync_round_metrics(event) if event_name not in {"rl_local_round", "rl_eval_result", "rl_round_trained"}: return None diff --git a/yeto/wan_streams.py b/yeto/wan_streams.py new file mode 100644 index 00000000..8d2cb6ab --- /dev/null +++ b/yeto/wan_streams.py @@ -0,0 +1,190 @@ +"""Per-stream accounting, stall detection and rate-aware striping for the +island side of a striped syncer connection (S26 E3b). + +A ``SyncerClient`` stripes one logical message over ``num_streams`` TCP data +sockets as 4 MiB CHUNK envelopes. The syncer reassembles a message by +``(msg_id, offset)``, so any chunk may travel on any data stream of the same +connection group. This module holds the parts of that transport that do not +need a socket: + +* :class:`StreamMeter` counts bytes, busy time and the last byte progress of + one stream. The send loop feeds it; diagnostics and the stall watchdog read + it. +* :func:`pick_stream` chooses the data stream for the next chunk. ``"rr"`` + is the historical round robin. ``"adaptive"`` sends each chunk to the + stream with the earliest estimated finish time (backlog / measured rate), + so a slow stream gets fewer chunks and a stalled stream gets none. The + number of open streams is fixed by HELLO; the number of streams that carry + data follows the measured rates. +* :class:`StreamStalledError` is the failure a watchdog reports when a + stream holds unsent bytes and makes no byte progress for the stall timeout. + The client handles it like any socket error: the whole group is torn down + and redialed, and the syncer replays its PULLs, so the island pushes the + interrupted fragment again at the same (global_step, round_attempt). +""" + +from __future__ import annotations + +import threading +import time +from collections.abc import Sequence +from dataclasses import dataclass + +STREAM_SCHEDULES = ("rr", "adaptive") + +# EWMA weight of the newest completed send in the per-stream rate estimate. +RATE_EWMA_ALPHA = 0.3 +# Sends shorter than this carry no useful rate signal (kernel buffer absorbs +# them), so they do not update the estimate. +MIN_RATE_SAMPLE_SECONDS = 0.01 + + +class StreamStalledError(ConnectionError): + """One stream held unsent (or partially received) bytes with no progress + for longer than the stall timeout.""" + + +@dataclass(frozen=True) +class StreamSnapshot: + """Cumulative counters of one stream at one instant.""" + + bytes_sent: int + busy_seconds: float + backlog_bytes: int + rate_bps: float | None + + +class StreamMeter: + """Counters of one outbound stream. Thread-safe; cheap enough to update + on every ``send`` call.""" + + def __init__(self, clock=time.monotonic) -> None: + self._clock = clock + self._lock = threading.Lock() + self.bytes_sent = 0 + self.busy_seconds = 0.0 + self.queued_bytes = 0 # enqueued, not yet handed to the kernel + self.rate_bps: float | None = None + self._item_started: float | None = None + self._item_bytes = 0 + self._last_progress = clock() + + def enqueued(self, nbytes: int) -> None: + with self._lock: + self.queued_bytes += nbytes + + def begin(self) -> None: + with self._lock: + now = self._clock() + self._item_started = now + self._item_bytes = 0 + self._last_progress = now + + def progress(self, nbytes: int) -> None: + with self._lock: + self.bytes_sent += nbytes + self.queued_bytes -= nbytes + self._item_bytes += nbytes + self._last_progress = self._clock() + + def end(self) -> None: + with self._lock: + if self._item_started is None: + return + elapsed = self._clock() - self._item_started + self.busy_seconds += elapsed + if elapsed >= MIN_RATE_SAMPLE_SECONDS and self._item_bytes: + sample = self._item_bytes / elapsed + self.rate_bps = ( + sample + if self.rate_bps is None + else (1 - RATE_EWMA_ALPHA) * self.rate_bps + RATE_EWMA_ALPHA * sample + ) + self._item_started = None + + def drop_pending(self) -> None: + """The connection group died: whatever was queued is gone.""" + with self._lock: + self.queued_bytes = 0 + self._item_started = None + self._item_bytes = 0 + + @property + def idle(self) -> bool: + with self._lock: + return self._item_started is None and self.queued_bytes <= 0 + + def stalled_for(self, now: float | None = None) -> float: + """Seconds without byte progress while bytes are waiting; 0 when idle.""" + with self._lock: + if self._item_started is None and self.queued_bytes <= 0: + return 0.0 + return max(0.0, (self._clock() if now is None else now) - self._last_progress) + + def snapshot(self) -> StreamSnapshot: + with self._lock: + busy = self.busy_seconds + if self._item_started is not None: + busy += self._clock() - self._item_started + return StreamSnapshot(self.bytes_sent, busy, max(0, self.queued_bytes), self.rate_bps) + + +def pick_stream( + schedule: str, + meters: Sequence[StreamMeter], + rr_ticket: int, +) -> int: + """Index into ``meters`` (the data streams) for the next chunk. + + ``adaptive`` minimizes ``backlog / rate``. A stream without a rate sample + yet uses the mean measured rate (or 1 B/s when nothing is measured, which + reduces the rule to least backlog). Ties go to the round-robin ticket so + an idle start still spreads chunks over every stream.""" + count = len(meters) + if count == 0: + raise ValueError("no data streams") + if schedule == "rr": + return rr_ticket % count + if schedule != "adaptive": + raise ValueError(f"unknown stream schedule {schedule!r}") + snaps = [meter.snapshot() for meter in meters] + known = [snap.rate_bps for snap in snaps if snap.rate_bps] + default_rate = sum(known) / len(known) if known else 1.0 + best = None + for step in range(count): + index = (rr_ticket + step) % count + snap = snaps[index] + finish = snap.backlog_bytes / (snap.rate_bps or default_rate) + if best is None or finish < best[0]: + best = (finish, index) + return best[1] + + +@dataclass(frozen=True) +class UploadWindow: + """Data-stream traffic between two snapshots (one sync boundary).""" + + bytes_per_stream: tuple[int, ...] + busy_seconds_per_stream: tuple[float, ...] + + @property + def total_bytes(self) -> int: + return sum(self.bytes_per_stream) + + def stream_mbps(self) -> list[float | None]: + """MB/s (1e6 B) of each stream over the time it was busy.""" + return [ + (b / s / 1e6) if s > 0 and b > 0 else None + for b, s in zip(self.bytes_per_stream, self.busy_seconds_per_stream, strict=True) + ] + + +def upload_window( + before: Sequence[StreamSnapshot], after: Sequence[StreamSnapshot] +) -> UploadWindow: + if len(before) != len(after): # the group was redialed with another width + before = [StreamSnapshot(0, 0.0, 0, None)] * len(after) + return UploadWindow( + tuple(max(0, a.bytes_sent - b.bytes_sent) for b, a in zip(before, after, strict=True)), + tuple(max(0.0, a.busy_seconds - b.busy_seconds) for b, a in zip(before, after, strict=True)), + ) 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 40/62] 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 41/62] 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 42/62] 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 43/62] 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 44/62] 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 45/62] 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 46/62] 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 47/62] 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 48/62] 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 2509fdefc62af5e6f41c06084ec4584b249eccd1 Mon Sep 17 00:00:00 2001 From: michael <78113554+chowhsu@users.noreply.github.com> Date: Wed, 21 Oct 2026 09:17:17 +0000 Subject: [PATCH 49/62] wan: replay INIT after a mid-init redial, link probe with container relaunch s26-e3b-p06-C32a: learner 0's connection group died before the syncer applied its INIT, the client never sent INIT again, and no island got the initial cut. s26-e3b runs also showed that a slow island is slow on every stream (container level), which no stream scheduling fixes. - SyncerClient.send_initial_fragments/init_replay_due/ mark_initial_cut_received: learner 0 replays INIT on every new connection group until the initial cut arrives (the syncer ignores INIT it already applied). TCP_USER_TIMEOUT is set only after the initial cut, so the multi-GB INIT and first broadcast are never reset by the kernel. - MSG_BW_PROBE (26): an HMAC-sealed probe connection to the legacy syncer uploads N bytes, gets an ack, downloads M bytes; never part of a session, at most 64 concurrent, 1 GiB per direction, 300 s. - --wan-probe-min-mbps / --wan-probe-min-down-mbps (default off): Miles islands probe the actor syncer before loading the model, record rl_wan_probe, and exit 75 below the minimum. The fleet controller requests a new container up to --wan-probe-relaunches times (default 2), also under --no-island-relaunch, since the island never joined the session. Real-syncer test: INIT interrupted between fragments is replayed and the finals match an undisturbed run bit for bit; the probe works on the same port and a wrong key is refused. Co-Authored-By: Claude Opus 5.5 --- syncer/src/protocol.rs | 4 + syncer/src/server.rs | 122 +++++++++++++- tests/test_e3_overlap_miles_role_stream.py | 9 ++ tests/test_rl_dense_sweep_real_syncer.py | 64 ++++++++ tests/test_wan_probe.py | 126 +++++++++++++++ tests/test_wan_streams.py | 32 ++++ yeto/cli.py | 23 +++ yeto/launch/fleet.py | 47 +++++- yeto/launch/modal_islands.py | 12 ++ yeto/launch/rl_island_task.py | 9 ++ yeto/launcher.py | 7 +- yeto/protocol.py | 59 ++++++- .../miles/full_parameter/role_stream.py | 31 ++-- yeto/rl/adapters/miles/island_entry.py | 35 ++++ yeto/rl/adapters/miles/run_flags.py | 3 + yeto/rl/event_echo.py | 8 +- yeto/wan_probe.py | 149 ++++++++++++++++++ 17 files changed, 716 insertions(+), 24 deletions(-) create mode 100644 tests/test_wan_probe.py create mode 100644 yeto/wan_probe.py diff --git a/syncer/src/protocol.rs b/syncer/src/protocol.rs index d9bd34bc..e214faf2 100644 --- a/syncer/src/protocol.rs +++ b/syncer/src/protocol.rs @@ -45,6 +45,10 @@ pub const MSG_ELASTIC_BASE: u8 = 23; pub const MSG_SAMPLE_VERDICT: u8 = 24; /// Elastic: run finished (final outer_version + base policy hash). pub const MSG_FINISHED: u8 = 25; +/// S26 E3b: link-bandwidth probe. First and only frame of its own connection +/// (never part of a learner session): the island uploads N bytes, the syncer +/// acknowledges with the u64 byte count, then streams M bytes back. +pub const MSG_BW_PROBE: u8 = 26; /// Version of the final-artifact handshake carried inside FINAL_MANIFEST and /// FINAL_ACK. Keeping this in the payload makes incompatible peers fail with diff --git a/syncer/src/server.rs b/syncer/src/server.rs index 269c0b44..c03011a3 100644 --- a/syncer/src/server.rs +++ b/syncer/src/server.rs @@ -1260,8 +1260,8 @@ async fn handle_connection( stream.set_nodelay(true)?; let (mut rd, mut wr) = stream.into_split(); let first = match read_frame_limited(&mut rd, |msg_type| match msg_type { - MSG_HELLO | MSG_DATA_HELLO => Ok(MAX_HELLO_FRAME), - other => bail!("first frame must be HELLO/DATA_HELLO, got {other}"), + MSG_HELLO | MSG_DATA_HELLO | MSG_BW_PROBE => Ok(MAX_HELLO_FRAME), + other => bail!("first frame must be HELLO/DATA_HELLO/BW_PROBE, got {other}"), }) .await { @@ -1466,8 +1466,98 @@ async fn handle_connection( tokio::spawn(writer_task(wr, rx)); read_loop(&mut rd, &group, &event_tx, &budget_cutoff).await } - t => bail!("first frame must be HELLO/DATA_HELLO, got {t}"), + MSG_BW_PROBE => { + let request = open_hello(hello_auth.key.as_deref(), &first.payload) + .and_then(parse_bw_probe); + let (upload, download) = match request { + Ok(sizes) => sizes, + Err(error) => { + let message = format!("invalid BW_PROBE: {error:#}"); + let _ = send_direct(&mut wr, MSG_ERROR, message.as_bytes()).await; + return Err(error); + } + }; + let Some(_slot) = BwProbeSlot::acquire() else { + let message = "too many concurrent BW_PROBE connections"; + send_direct(&mut wr, MSG_ERROR, message.as_bytes()).await?; + bail!(message); + }; + tokio::time::timeout(BW_PROBE_TIMEOUT, serve_bw_probe(&mut rd, &mut wr, upload, download)) + .await + .context("BW_PROBE timed out")? + } + t => bail!("first frame must be HELLO/DATA_HELLO/BW_PROBE, got {t}"), + } +} + +/// Largest upload or download one BW_PROBE connection may ask for. +const BW_PROBE_MAX_BYTES: u64 = 1 << 30; +const BW_PROBE_MAX_CONNECTIONS: usize = 64; +const BW_PROBE_TIMEOUT: Duration = Duration::from_secs(300); +static BW_PROBE_ACTIVE: AtomicUsize = AtomicUsize::new(0); + +/// One of at most BW_PROBE_MAX_CONNECTIONS concurrent probe connections. +struct BwProbeSlot; + +impl BwProbeSlot { + fn acquire() -> Option { + let previous = BW_PROBE_ACTIVE.fetch_add(1, Ordering::AcqRel); + if previous >= BW_PROBE_MAX_CONNECTIONS { + BW_PROBE_ACTIVE.fetch_sub(1, Ordering::AcqRel); + return None; + } + Some(BwProbeSlot) + } +} + +impl Drop for BwProbeSlot { + fn drop(&mut self) { + BW_PROBE_ACTIVE.fetch_sub(1, Ordering::AcqRel); + } +} + +/// BW_PROBE body (after the HELLO-style HMAC trailer is removed): +/// u64 upload bytes, u64 download bytes, both at most BW_PROBE_MAX_BYTES. +fn parse_bw_probe(body: &[u8]) -> Result<(u64, u64)> { + if body.len() != 16 { + bail!("BW_PROBE body has {} bytes, expected 16", body.len()); + } + let upload = u64::from_le_bytes(body[0..8].try_into().unwrap()); + let download = u64::from_le_bytes(body[8..16].try_into().unwrap()); + if upload > BW_PROBE_MAX_BYTES || download > BW_PROBE_MAX_BYTES { + bail!("BW_PROBE asks for more than {BW_PROBE_MAX_BYTES} bytes per direction"); } + Ok((upload, download)) +} + +/// Read and discard `upload` raw bytes, acknowledge with their count, then +/// write `download` zero bytes and close. Nothing touches session state. +async fn serve_bw_probe( + rd: &mut (impl tokio::io::AsyncReadExt + Unpin), + wr: &mut OwnedWriteHalf, + upload: u64, + download: u64, +) -> Result<()> { + let mut buf = vec![0u8; 1 << 20]; + let mut remaining = upload; + while remaining > 0 { + let want = remaining.min(buf.len() as u64) as usize; + let got = rd.read(&mut buf[..want]).await?; + if got == 0 { + bail!("BW_PROBE upload ended after {} of {upload} bytes", upload - remaining); + } + remaining -= got as u64; + } + wr.write_all(&upload.to_le_bytes()).await?; + buf.fill(0); + let mut remaining = download; + while remaining > 0 { + let take = remaining.min(buf.len() as u64) as usize; + wr.write_all(&buf[..take]).await?; + remaining -= take as u64; + } + wr.shutdown().await?; + Ok(()) } async fn send_direct(wr: &mut OwnedWriteHalf, msg_type: u8, payload: &[u8]) -> Result<()> { @@ -4410,6 +4500,28 @@ mod hello_auth_tests { out } + #[test] + fn bw_probe_body_is_two_bounded_sizes() { + let mut body = Vec::new(); + body.extend_from_slice(&5u64.to_le_bytes()); + body.extend_from_slice(&7u64.to_le_bytes()); + assert_eq!(parse_bw_probe(&body).unwrap(), (5, 7)); + assert!(parse_bw_probe(&body[..15]).is_err()); + let mut too_big = Vec::new(); + too_big.extend_from_slice(&(BW_PROBE_MAX_BYTES + 1).to_le_bytes()); + too_big.extend_from_slice(&0u64.to_le_bytes()); + assert!(parse_bw_probe(&too_big).is_err()); + } + + #[test] + fn bw_probe_slots_are_bounded() { + let slots: Vec<_> = (0..BW_PROBE_MAX_CONNECTIONS).map(|_| BwProbeSlot::acquire()).collect(); + assert!(slots.iter().all(Option::is_some)); + assert!(BwProbeSlot::acquire().is_none()); + drop(slots); + assert!(BwProbeSlot::acquire().is_some()); + } + #[test] fn open_hello_accepts_the_right_key_only() { let body = b"hello body bytes"; @@ -4951,7 +5063,7 @@ mod tests { } let mut stream = stream.expect("legacy syncer did not start"); // A correctly sealed elastic payload (epoch 0) still must not be - // accepted: legacy only knows HELLO/DATA_HELLO as a first frame. + // accepted: legacy only knows HELLO/DATA_HELLO/BW_PROBE as a first frame. let payload = crate::elastic::seal(&key, msg_type, 0u64.to_le_bytes().to_vec()); let mut frame = MAGIC.to_le_bytes().to_vec(); frame.push(msg_type); @@ -4968,7 +5080,7 @@ mod tests { assert_eq!(reply.get(4), Some(&MSG_ERROR), "type {msg_type}"); let text = String::from_utf8_lossy(&reply[13..]); assert!( - text.contains(&format!("first frame must be HELLO/DATA_HELLO, got {msg_type}")), + text.contains(&format!("first frame must be HELLO/DATA_HELLO/BW_PROBE, got {msg_type}")), "type {msg_type}: {text}" ); assert!(!server.is_finished(), "type {msg_type}: server exited"); diff --git a/tests/test_e3_overlap_miles_role_stream.py b/tests/test_e3_overlap_miles_role_stream.py index 6d9f0aa1..6a69f85f 100644 --- a/tests/test_e3_overlap_miles_role_stream.py +++ b/tests/test_e3_overlap_miles_role_stream.py @@ -266,6 +266,15 @@ def send_init_parts(self, fragment_id, parts) -> bool: self.pull_sweep(0, 1) return True + def send_initial_fragments(self, parts_for) -> bool: + return all(self.send_init_parts(f, parts_for(f)) for f in range(P)) + + def init_replay_due(self) -> bool: + return False + + def mark_initial_cut_received(self) -> None: + pass + def check_health(self) -> None: assert not self.closed diff --git a/tests/test_rl_dense_sweep_real_syncer.py b/tests/test_rl_dense_sweep_real_syncer.py index 92bf9b97..7fd428eb 100644 --- a/tests/test_rl_dense_sweep_real_syncer.py +++ b/tests/test_rl_dense_sweep_real_syncer.py @@ -633,3 +633,67 @@ def test_stalled_stream_redials_and_the_replayed_push_keeps_the_exact_result( assert progress.stalls >= 1 and progress.reconnects >= 1 assert list(manifest.versions) == exact[1] assert result == exact[2] + + +def _initial_values(fid: int) -> list[float]: + return [float(fid) + 0.125 * i for i in range(NUMEL)] + + +def test_init_interrupted_by_a_redial_is_replayed_and_the_result_stays_exact(syncer_binary, tmp_path): + """E3b (s26-e3b-p06-C32a): learner 0's group died after INIT of fragment 0 + and before fragment 1, so the syncer never finished initializing and no + island got the initial cut. The client now replays INIT on the next group + until the initial cut arrives; the result matches an undisturbed run. A + link probe on the same port (yeto.wan_probe) leaves the session alone.""" + from yeto.protocol import island_hmac_key_from_env + from yeto.wan_probe import probe_link + + exact = _run_sweeps(syncer_binary, tmp_path, STREAMED, sweeps=3) + run_dir = tmp_path / "replay" + run_dir.mkdir() + process, port, profile_hash = _start_syncer(syncer_binary, run_dir, *STREAMED, total_steps=3) + client = SyncerClient(("127.0.0.1", port), 0, _layout(), dtype=DTYPE_F32, num_streams=2, + connect_timeout=10, max_reconnects=5, finalization_timeout=20, + syncer_profile_hash=bytes.fromhex(profile_hash), stream_stall_timeout=30.0) + killed = [] + + def parts_for(fid: int): + if fid == 1 and not killed: + killed.append(client.connection_generation) + client._poison_outbound_generation(client._gen, ConnectionError("test: group died mid-INIT")) + return [_f32(_initial_values(fid))] + + try: + def listening() -> bool: + try: + socket.create_connection(("127.0.0.1", port), timeout=1).close() + except OSError: + return False + return True + + _wait(listening, "syncer listening") + good = probe_link(("127.0.0.1", port), streams=2, bytes_per_direction=8 << 20, + key=island_hmac_key_from_env(), timeout=30) + assert good.error is None and good.upload_mbps > 0 and good.download_mbps > 0 + bad = probe_link(("127.0.0.1", port), streams=1, bytes_per_direction=1 << 20, + key=b"wrong key", timeout=10) + assert bad.error is not None and not bad.accepted(0.0, 0.0) + client.start() + assert client.send_initial_fragments(parts_for) is False + replays = 0 + while client._bcasts.qsize() < FRAGMENTS: + client.check_health() + if client.init_replay_due(): + replays += client.send_initial_fragments(parts_for) + time.sleep(0.02) + client.mark_initial_cut_received() + assert killed and replays == 1 and client.connection_generation != killed[0] + manifest, finals = _sweep_until_final(client, sweeps=3) + result = {f.fragment_id: bytes(f.data) for f in finals} + client.acknowledge_finalization(manifest) + assert process.wait(timeout=15) == 0 + finally: + client.close() + _stop(process) + assert list(manifest.versions) == exact[1] + assert result == exact[2] diff --git a/tests/test_wan_probe.py b/tests/test_wan_probe.py new file mode 100644 index 00000000..9ce848ab --- /dev/null +++ b/tests/test_wan_probe.py @@ -0,0 +1,126 @@ +"""S26 E3b: island link probe, its island hook, and the launcher's relaunch of +an island whose probe rejected the container's link.""" + +from __future__ import annotations + +import json +from argparse import Namespace + +import pytest + +from yeto.wan_probe import EXIT_SLOW_LINK, ProbeResult, probe_link + + +def _fake_stream(up_s: float, down_s: float): + def run(addr, upload, download, key, timeout, out): + out["up"], out["down"] = up_s, down_s + return run + + +def test_probe_rates_use_the_slowest_stream_per_direction(): + result = probe_link(("h", 1), streams=4, bytes_per_direction=40_000_000, key=None, + stream_fn=_fake_stream(2.0, 0.5)) + assert (result.upload_mbps, result.download_mbps) == (20.0, 80.0) + assert result.accepted(20.0, 80.0) and not result.accepted(20.1, 0.0) + assert result.fields()["wan_probe/upload_mbps"] == 20.0 + + +def test_a_failed_stream_rejects_the_link(): + def broken(addr, upload, download, key, timeout, out): + raise ConnectionError("reset") + result = probe_link(("h", 1), streams=2, bytes_per_direction=2 << 20, key=None, stream_fn=broken) + assert result.error == "ConnectionError: reset" and not result.accepted(0.0, 0.0) + with pytest.raises(ValueError): + probe_link(("h", 1), streams=0, bytes_per_direction=1, key=None) + + +def _island_args(tmp_path, **kw): + return Namespace(syncer="10.0.0.1:29400", wan_streams=8, wan_probe_mib=16, + event_tape=str(tmp_path / "tape.jsonl"), learner_id=1, **kw) + + +def test_island_probe_is_off_by_default_and_exits_75_on_a_slow_link(tmp_path, capsys): + from yeto.rl.adapters.miles.island_entry import probe_wan_link + + calls = [] + + def probe(addr, **kw): + calls.append((addr, kw["streams"], kw["bytes_per_direction"])) + return ProbeResult(8, 16 << 20, 16 << 20, 4.0, 1.0) + + probe_wan_link(_island_args(tmp_path), probe=probe) + assert calls == [] + probe_wan_link(_island_args(tmp_path, wan_probe_min_mbps=2.0, wan_probe_min_down_mbps=10.0), probe=probe) + with pytest.raises(SystemExit) as stop: + probe_wan_link(_island_args(tmp_path, wan_probe_min_mbps=8.0), probe=probe) + assert stop.value.code == EXIT_SLOW_LINK + assert calls == [(("10.0.0.1", 29400), 8, 16 << 20)] * 2 + records = [json.loads(line) for line in (tmp_path / "tape.jsonl").read_text().splitlines()] + assert [r["accepted"] for r in records if r["event"] == "rl_wan_probe"] == [True, False] + assert "[yeto-wan-probe] rejected: up 4.2 MB/s" in capsys.readouterr().out + + +def test_launcher_flags_reach_the_island_and_bound_the_relaunches(): + from yeto.launch.modal_islands import wan_probe_relaunches + from yeto.launch.rl_island_task import _wan_transport_flags + + on = Namespace(wan_probe_min_mbps=20.0, wan_probe_min_down_mbps=0.0, wan_probe_mib=128, + wan_probe_relaunches=3) + assert _wan_transport_flags(on) == (" --wan-probe-min-mbps 20.0 --wan-probe-min-down-mbps 0.0" + " --wan-probe-mib 128") + assert wan_probe_relaunches(on) == 3 + assert wan_probe_relaunches(Namespace(wan_probe_relaunches=3)) == 0 # probe off + with pytest.raises(ValueError): + _wan_transport_flags(Namespace(wan_probe_min_mbps=1.0, wan_probe_mib=0)) + + +def test_echo_collector_counts_rejected_probes(tmp_path): + from yeto.rl.event_echo import PREFIX, TapeCollector + + collector = TapeCollector(tmp_path / "tape.jsonl", fresh=True) + for accepted, t in ((False, 1), (True, 2), (False, 3)): + collector.feed(PREFIX + json.dumps({"event": "rl_wan_probe", "accepted": accepted, "island_id": 1, "time_unix": t}) + "\n") + assert collector.wan_rejects == 2 + + +def _fleet(ops, rejects, relaunches, **kw): + from test_controller import RUNNING, SYNCER, ImmediateThread + + from yeto.launcher import FleetController + + ops.status_seq.setdefault(SYNCER, [RUNNING]) + return FleetController(learners={"l0": ("task-l0", 1), "l1": ("task-l1", 2)}, + syncer=(SYNCER, "task-syncer", 1), sky_ops=ops, poll_interval=30, + recover_timeout=0, thread_cls=ImmediateThread, fixed_roster=True, + no_island_relaunch=True, wan_reject_probe=lambda name: rejects[name], + wan_probe_relaunches=relaunches, **kw) + + +def test_rejected_link_is_relaunched_even_without_island_relaunch(): + from test_controller import FAILED, RUNNING, SUCCEEDED, FakeOps + + ops = FakeOps() + ops.status_seq["l0"] = [RUNNING, RUNNING, RUNNING, SUCCEEDED] + ops.status_seq["l1"] = [FAILED] + ops.relaunch_results["l1"] = [7] + ops.after_relaunch["l1"] = [RUNNING, SUCCEEDED] + rejects = {"l0": 0, "l1": 1} + ctl = _fleet(ops, rejects, relaunches=2) + ctl.run() + assert ops.relaunch_calls == ["l1"] + assert ctl.learners["l1"]["wan_relaunches"] == 1 and ctl.learners["l1"]["state"] == "done" + + +def test_rejected_link_budget_is_bounded(): + from test_controller import FAILED, RUNNING, FakeOps + + from yeto.launcher import FixedRosterIslandAbandoned + + ops = FakeOps() + ops.status_seq["l0"] = [RUNNING] + ops.status_seq["l1"] = [FAILED] + rejects = {"l0": 0, "l1": 1} + ctl = _fleet(ops, rejects, relaunches=0) + with pytest.raises(FixedRosterIslandAbandoned): + ctl.run() + assert ops.relaunch_calls == [] diff --git a/tests/test_wan_streams.py b/tests/test_wan_streams.py index 4e08efb0..99f3d4dc 100644 --- a/tests/test_wan_streams.py +++ b/tests/test_wan_streams.py @@ -227,3 +227,35 @@ def test_island_parser_accepts_the_transport_flags(): assert rl_option("wan_stream_schedule", ISLAND) == "--wan-stream-schedule" assert rl_option("wan_stall_timeout", ISLAND) == "--wan-stall-timeout" + + +def test_tcp_user_timeout_waits_for_the_initial_cut(): + if not hasattr(socket, "TCP_USER_TIMEOUT"): + pytest.skip("Linux-only socket option") + with socket.create_server(("127.0.0.1", 0)) as server: + addr = server.getsockname() + client = SyncerClient(addr, 0, _layout(), num_streams=1, stream_stall_timeout=2.0) + early = client._dial(5) + try: + assert early.getsockopt(socket.IPPROTO_TCP, socket.TCP_USER_TIMEOUT) == 0 + client._socks = [early] + client.mark_initial_cut_received() + assert early.getsockopt(socket.IPPROTO_TCP, socket.TCP_USER_TIMEOUT) == 2000 + late = client._dial(5) + assert late.getsockopt(socket.IPPROTO_TCP, socket.TCP_USER_TIMEOUT) == 2000 + late.close() + finally: + early.close() + + +def test_init_replay_is_due_only_on_a_new_group_before_the_initial_cut(): + client = SyncerClient(("127.0.0.1", 1), 0, _layout(), num_streams=1) + client._connected.set() + client.connection_generation = 11 + assert client.init_replay_due() + client._init_generation = 11 + assert not client.init_replay_due() + client.connection_generation = 12 + assert client.init_replay_due() + client.mark_initial_cut_received() + assert not client.init_replay_due() diff --git a/yeto/cli.py b/yeto/cli.py index 40abdc1f..ab16fb4b 100644 --- a/yeto/cli.py +++ b/yeto/cli.py @@ -454,6 +454,29 @@ def int_or_auto(value: str): "one stream holds bytes with no progress for this many seconds (also sets " "TCP_USER_TIMEOUT); the fragment is pushed again on the replayed PULL. 0 = off", ) + sync.add_argument( + "--wan-probe-min-mbps", + type=float, + default=0.0, + help="Miles RL islands: before loading the model, measure the upload rate to the " + "actor syncer over --wan-streams connections; below this MB/s the island exits " + "(code 75) and the launcher requests another container (--wan-probe-relaunches). 0 = off", + ) + sync.add_argument( + "--wan-probe-min-down-mbps", + type=float, + default=0.0, + help="like --wan-probe-min-mbps, for the download rate from the syncer. 0 = off", + ) + sync.add_argument("--wan-probe-mib", type=int, default=256, + help="bytes per direction of the link probe, in MiB") + sync.add_argument( + "--wan-probe-relaunches", + type=int, + default=2, + help="containers the launcher may request again per island after a rejected link " + "probe (also under --no-island-relaunch: the island never joined the session)", + ) infra = p.add_argument_group("infrastructure") infra.add_argument( diff --git a/yeto/launch/fleet.py b/yeto/launch/fleet.py index b7051a46..d1d9ce63 100644 --- a/yeto/launch/fleet.py +++ b/yeto/launch/fleet.py @@ -48,6 +48,8 @@ def __init__(self): # time spent recovering in the window. FIXED_ROSTER_MAX_RELAUNCHES = 2 RECOVERY_STABLE_S = 300.0 +# S26 E3b: how long one rejected-link relaunch may take to get a new container. +WAN_RECOVERY_S = 900.0 # How long a failed-looking island is given for its rl_learner_finalized record # to arrive over the log stream before recovery starts. FINALIZE_GRACE_S = 60.0 @@ -154,6 +156,8 @@ def __init__( elastic: bool = False, no_island_relaunch: bool = False, replacement_planner=None, + wan_reject_probe=None, + wan_probe_relaunches: int = 0, ): """`learners` maps cluster name -> (task, job_id); `syncer` is (name, task, job_id) for a cluster syncer, or None with @@ -169,6 +173,12 @@ def __init__( self.left_pool_probe = left_pool_probe self.elastic = elastic self.no_island_relaunch = no_island_relaunch + # S26 E3b: name -> number of rejected link probes the island logged + # (rl_wan_probe accepted=false, exit 75). Such an island never joined + # the syncer session, so it is relaunched in a fresh container up to + # wan_probe_relaunches times, also under --no-island-relaunch. + self.wan_reject_probe = wan_reject_probe + self.wan_probe_relaunches = wan_probe_relaunches # rl-spot-cost-saving 2.3: yeto.cloud.replace.ReplacementPlanner or None. # Asked once per recovery window before the same-place relaunch. Auto # launch off (default): it only writes spot_replace_advice and the @@ -372,6 +382,8 @@ def _poll(self, rec, is_syncer: bool) -> None: self._fleet("island_done", rec["name"], reason="syncer finished") print(f"[launcher] {rec['name']}: {verdict} after the syncer finished; " "counted as a normal end, not relaunched", file=sys.stderr) + elif not is_syncer and self._rejected_link(rec): + self._relaunch_rejected_link(rec) elif not is_syncer and self.elastic and self._left_pool(rec): # 0.27: re-JOIN exhausted (exit 7): the island left the elastic pool. # Not the fixed-roster abandon path (exit 4): relaunch unless @@ -418,6 +430,36 @@ def _syncer_finished(self) -> bool: except Exception: return False + def _rejected_link(self, rec) -> bool: + """A new rejected link probe since the island's last (re)launch.""" + if self.wan_reject_probe is None: + return False + try: + seen = int(self.wan_reject_probe(rec["name"])) + except Exception: + return False + return seen > rec.get("wan_rejects_handled", 0) + + def _relaunch_rejected_link(self, rec) -> None: + rec["wan_rejects_handled"] = rec.get("wan_rejects_handled", 0) + 1 + used = rec.get("wan_relaunches", 0) + if used >= self.wan_probe_relaunches: + self._fleet("island_lost", rec["name"], reason="link probe rejected; relaunch budget used") + print(f"[launcher] {rec['name']}: link probe rejected again after {used} new " + "container(s); not relaunched", file=sys.stderr) + self._abandon(rec, 0.0, reason="link probe rejected") + return + rec["wan_relaunches"] = used + 1 + self._fleet("island_lost", rec["name"], + reason=f"link probe rejected; new container {used + 1}/{self.wan_probe_relaunches}") + print(f"[launcher] {rec['name']}: link probe rejected; requesting a new container " + f"({used + 1}/{self.wan_probe_relaunches})", file=sys.stderr) + rec["state"] = RECOVERING + rec["failed_at"] = self.ops.now() + rec["wan_recovery"] = True + rec.pop("recovered_at", None) + self._drive_recovery(rec, False) + def _left_pool(self, rec) -> bool: if self.left_pool_probe is None: return False @@ -577,6 +619,7 @@ def _drive_recovery(self, rec, is_syncer: bool) -> None: self.ops.now() - rec["failed_at"]) rec["failed_at"] = None rec["replace_planned"] = False + rec.pop("wan_recovery", None) rec["recovered_at"] = self.ops.now() print( f"[launcher] {rec['name']} recovered: relaunched as job " @@ -597,7 +640,9 @@ def _drive_recovery(self, rec, is_syncer: bool) -> None: if rec.get("failures", 0) > FIXED_ROSTER_MAX_RELAUNCHES: self._abandon(rec, elapsed) return - if self.recover_timeout <= 0 or elapsed > self.recover_timeout: + # A rejected-link relaunch has its own budget (the island never joined). + timeout = WAN_RECOVERY_S if rec.get("wan_recovery") else self.recover_timeout + if timeout <= 0 or elapsed > timeout: if is_syncer: # The syncer is never abandoned: without it no learner can # make outer progress, so keep trying and complain loudly. diff --git a/yeto/launch/modal_islands.py b/yeto/launch/modal_islands.py index a7f12af6..976211e3 100644 --- a/yeto/launch/modal_islands.py +++ b/yeto/launch/modal_islands.py @@ -315,3 +315,15 @@ def effective_recover_timeout(args) -> float: if getattr(args, "modal_retries", None) == 0 and not getattr(args, "modal_launcher_relaunch", False): return 0 return args.recover_timeout + + +def wan_probe_relaunches(args) -> int: + """New containers per island after a rejected link probe (S26 E3b); + 0 when no island runs the probe.""" + if not (float(getattr(args, "wan_probe_min_mbps", 0.0) or 0.0) > 0 + or float(getattr(args, "wan_probe_min_down_mbps", 0.0) or 0.0) > 0): + return 0 + count = int(getattr(args, "wan_probe_relaunches", 2)) + if count < 0: + raise ValueError("--wan-probe-relaunches must be >= 0") + return count diff --git a/yeto/launch/rl_island_task.py b/yeto/launch/rl_island_task.py index a51d6ea1..0254a603 100644 --- a/yeto/launch/rl_island_task.py +++ b/yeto/launch/rl_island_task.py @@ -62,6 +62,15 @@ def _wan_transport_flags(args) -> str: raise ValueError("--wan-stall-timeout must be >= 0 (0 = off)") if stall > 0: flags += f" --wan-stall-timeout {stall!r}" + min_up = float(getattr(args, "wan_probe_min_mbps", 0.0) or 0.0) + min_down = float(getattr(args, "wan_probe_min_down_mbps", 0.0) or 0.0) + if min_up < 0 or min_down < 0: + raise ValueError("--wan-probe-min-mbps / --wan-probe-min-down-mbps must be >= 0 (0 = off)") + if min_up > 0 or min_down > 0: + mib = int(getattr(args, "wan_probe_mib", 256)) + if not 1 <= mib <= 1024: + raise ValueError("--wan-probe-mib must be between 1 and 1024") + flags += f" --wan-probe-min-mbps {min_up!r} --wan-probe-min-down-mbps {min_down!r} --wan-probe-mib {mib}" return flags diff --git a/yeto/launcher.py b/yeto/launcher.py index 0abd2f4b..7819fbc5 100644 --- a/yeto/launcher.py +++ b/yeto/launcher.py @@ -61,7 +61,7 @@ from .launch.modal_islands import ( _modal_tape_run_dir, _tail_modal, _verify_modal_app_stopped, build_modal_island_config, container_set_poller_for, effective_recover_timeout, pull_modal_tapes, - require_modal_for_gpu_exact, + require_modal_for_gpu_exact, wan_probe_relaunches, ) from .launch.fleet import ( CONTAINER_CHANGED_EXIT, DEFAULT_RL_STALL_TIMEOUT_S, FAILED_RUN_DRAIN_S, ISLAND_FAILED_EXIT, @@ -660,6 +660,11 @@ def container_changed() -> bool: ), elastic=island_scheduling_mode(args) == "elastic", no_island_relaunch=bool(getattr(args, "no_island_relaunch", False)), + wan_reject_probe=( + (lambda name: event_collectors[name].wan_rejects if name in event_collectors else 0) + if echo_names else None + ), + wan_probe_relaunches=wan_probe_relaunches(args), ) def drain_tapes(limit: float = NO_SYNC_EVENT_DRAIN_S) -> None: if echo_names: diff --git a/yeto/protocol.py b/yeto/protocol.py index 42cfd3f8..e55444b5 100644 --- a/yeto/protocol.py +++ b/yeto/protocol.py @@ -60,6 +60,8 @@ MSG_FINAL_ACK = 12 MSG_FINAL_FRAGMENT = 13 MSG_BUDGET_DONE = 14 +# S26 E3b: link-bandwidth probe connection (yeto/wan_probe.py; syncer serve_bw_probe). +MSG_BW_PROBE = 26 FINALIZATION_REVISION = 1 FINALIZATION_TIMEOUT = 900.0 @@ -659,6 +661,12 @@ def __init__( self._meters = [StreamMeter() for _ in range(num_streams + 1)] self._upload_drained_at: float | None = None self._stalls = 0 + # S26 E3b: set once the initial cut arrived. Before that a redial must + # replay INIT_PARAMS (learner 0), and TCP_USER_TIMEOUT stays unset: + # the multi-GB INIT upload and initial broadcast must not be reset by + # the kernel on a slow link (s26-e3b-p06-C32a). + self._initial_cut_received = threading.Event() + self._init_generation: int | None = None self.session_contract_hash = session_contract_hash self.syncer_profile_hash = syncer_profile_hash # secret-handling-hardening: HELLO carries an HMAC when the island has @@ -733,7 +741,7 @@ def _dial(self, timeout: float) -> socket.socket: sock.settimeout(None) sock.setsockopt(socket.IPPROTO_TCP, socket.TCP_NODELAY, 1) enable_tcp_keepalive(sock) - if self.stream_stall_timeout is not None: + if self.stream_stall_timeout is not None and self._initial_cut_received.is_set(): set_tcp_user_timeout(sock, self.stream_stall_timeout) return sock @@ -961,6 +969,55 @@ def upload_drained_at(self) -> float | None: # -- learner-facing API ------------------------------------------------------ + def send_initial_fragments( + self, + parts_for: Callable[[int], Iterable[bytes | bytearray | memoryview]], + ) -> bool: + """Learner 0: stream INIT_PARAMS of every fragment on the live group. + + ``parts_for(fragment_id)`` returns a fresh part iterable each call, so + the same INIT can be replayed after a redial. Returns True when every + fragment was queued on one connection generation; False on an outage + (nothing to undo: the syncer ignores INIT it already applied, and a + partial INIT of a dead generation is dropped by the syncer).""" + generation = self.connection_generation + try: + for fragment_id in range(self.layout.num_fragments): + if self.send_init_parts(fragment_id, parts_for(fragment_id)) is not True: + return False + except PartialMessageGenerationLost: + return False + if generation != self.connection_generation: + return False + self._init_generation = generation + return True + + def init_replay_due(self) -> bool: + """True when the live group never carried a complete INIT and the + initial cut has not arrived yet (learner 0 must send INIT again).""" + with self._lock: + live = self._connected.is_set() + generation = self.connection_generation + return ( + live + and not self._initial_cut_received.is_set() + and generation != self._init_generation + ) + + def mark_initial_cut_received(self) -> None: + """The initial cut arrived: INIT is no longer replayed, and the stall + timeout now also bounds unacknowledged data in the kernel.""" + self._initial_cut_received.set() + if self.stream_stall_timeout is None: + return + with self._lock: + socks = list(self._socks) + for sock in socks: + try: + set_tcp_user_timeout(sock, self.stream_stall_timeout) + except OSError: + pass # a dying socket; the group is being replaced anyway + def send_init(self, fragment_id: int, tensor_bytes: bytes) -> None: if not 0 <= fragment_id < self.layout.num_fragments: raise ValueError(f"INIT_PARAMS for unknown fragment {fragment_id}") diff --git a/yeto/rl/adapters/miles/full_parameter/role_stream.py b/yeto/rl/adapters/miles/full_parameter/role_stream.py index fa746fcb..549743f7 100644 --- a/yeto/rl/adapters/miles/full_parameter/role_stream.py +++ b/yeto/rl/adapters/miles/full_parameter/role_stream.py @@ -400,27 +400,23 @@ async def initialize(self, group) -> None: client.install_inbound_chunk_sink(MilesRayInboundChunkSink(sink)) try: client.start() + init_parts = None if self.config.learner_id == 0: - for fragment_id in range(adapter.layout.fragments.num_fragments): - queued = client.send_init_parts( - fragment_id, - adapter.fragment_parts( - initial, - fragment_id, - ray_module=self.ray_module, - )(), - ) - if queued is not True: - raise RuntimeError( - f"{self.config.algorithm} {self.config.role} initial fragment was not queued" - ) + def init_parts(fragment_id: int): + return adapter.fragment_parts( + initial, fragment_id, ray_module=self.ray_module + )() + + # An outage here is not fatal: the wait below replays INIT on + # the next connection group until the initial cut arrives. + client.send_initial_fragments(init_parts) # Wait off the event loop: an island initializes its roles # concurrently, and each role's initial cut needs every island's # INIT for that role. A blocking wait here would stop this # island's other role from connecting, and two islands blocked # on different roles deadlock (s23-g3-sao-main-20261014a). authoritative = await asyncio.to_thread( - self._wait_for_initial_fragments, client, adapter + self._wait_for_initial_fragments, client, adapter, init_parts ) adapter.verify_initial_fragments(initial, authoritative) except BaseException: @@ -443,6 +439,7 @@ def _wait_for_initial_fragments( self, client, adapter: MilesChunkedFullParameterAdapter, + init_parts=None, ) -> tuple[StoredAuthoritativeFragment, ...]: started = time.monotonic() deadline = started + self.config.wait_timeout @@ -452,6 +449,11 @@ def _wait_for_initial_fragments( self._report_initial_cut(client, observed, count, started, "waiting") while len(observed) < count: client.check_health() + if init_parts is not None and client.init_replay_due(): + # S26 E3b (s26-e3b-p06-C32a): the group carrying INIT died + # before the syncer applied it; send it again on this group. + self._report_initial_cut(client, observed, count, started, "replaying INIT") + client.send_initial_fragments(init_parts) for update in client.drain_updates(): stored = self._stored_update(update) if update.version != 0 or stored.fragment_id in observed: @@ -471,6 +473,7 @@ def _wait_for_initial_fragments( ) time.sleep(self.config.poll_seconds) self._report_initial_cut(client, observed, count, started, "received") + client.mark_initial_cut_received() return tuple(observed[index] for index in range(count)) def _report_initial_cut( diff --git a/yeto/rl/adapters/miles/island_entry.py b/yeto/rl/adapters/miles/island_entry.py index 01367ff3..cf98a201 100644 --- a/yeto/rl/adapters/miles/island_entry.py +++ b/yeto/rl/adapters/miles/island_entry.py @@ -1633,6 +1633,40 @@ def probe_syncers(args, *, connect=None, attempts: int = 10, delay: float = 3.0, return addrs +def probe_wan_link(args, *, probe=None) -> None: + """S26 E3b: measure this container's link to the actor syncer before any + model work; exit EXIT_SLOW_LINK below the configured rates so the launcher + can request another container. Off unless a minimum rate is set.""" + min_up = float(getattr(args, "wan_probe_min_mbps", 0.0) or 0.0) + min_down = float(getattr(args, "wan_probe_min_down_mbps", 0.0) or 0.0) + if min_up <= 0 and min_down <= 0: + return + from yeto.protocol import island_hmac_key_from_env + from yeto.wan_probe import EXIT_SLOW_LINK, PROBE_EVENT, PROBE_MARKER, probe_link + + result = (probe or probe_link)( + _syncer_address(args.syncer), + streams=max(1, int(args.wan_streams)), + bytes_per_direction=int(getattr(args, "wan_probe_mib", 256)) << 20, + key=island_hmac_key_from_env(), + ) + accepted = result.accepted(min_up, min_down) + record = {"event": PROBE_EVENT, "accepted": accepted, + "wan_probe/min_upload_mbps": min_up, "wan_probe/min_download_mbps": min_down, + **result.fields()} + print(f"{PROBE_MARKER} {'accepted' if accepted else 'rejected'}: up {result.upload_mbps:.1f} MB/s " + f"(min {min_up}), down {result.download_mbps:.1f} MB/s (min {min_down}), " + f"{result.streams} stream(s){'' if result.error is None else '; ' + result.error}", flush=True) + if getattr(args, "event_tape", None): + from yeto.rl.engine.driver import EventTape + + if getattr(args, "rl_echo_events", False): + install_event_echo() # the launcher reads the verdict from the log stream + EventTape(args.event_tape, int(getattr(args, "learner_id", 0) or 0)).append(record) + if not accepted: + raise SystemExit(EXIT_SLOW_LINK) + + def require_ports_router_mode(miles_args) -> None: """Upstream Miles owns the SGLang router. @@ -2394,6 +2428,7 @@ def main(argv=None) -> None: verify_miles_revision(miles_root, expected=pins) _preflight_codex_harness(args) probe_syncers(args) + probe_wan_link(args) from miles.utils.misc import load_function diff --git a/yeto/rl/adapters/miles/run_flags.py b/yeto/rl/adapters/miles/run_flags.py index 30778809..803b19e1 100644 --- a/yeto/rl/adapters/miles/run_flags.py +++ b/yeto/rl/adapters/miles/run_flags.py @@ -418,6 +418,9 @@ island_flag("--wan-streams", type=int, default=4), island_flag("--wan-stream-schedule", choices=("rr", "adaptive"), default="rr"), island_flag("--wan-stall-timeout", type=float, default=0.0), + island_flag("--wan-probe-min-mbps", type=float, default=0.0), + island_flag("--wan-probe-min-down-mbps", type=float, default=0.0), + island_flag("--wan-probe-mib", type=int, default=256), island_flag("--miles-root", required=True), island_flag( "--miles-commit", diff --git a/yeto/rl/event_echo.py b/yeto/rl/event_echo.py index b3f2a33e..b9abdadc 100644 --- a/yeto/rl/event_echo.py +++ b/yeto/rl/event_echo.py @@ -133,6 +133,7 @@ def __init__(self, path, *, fresh: bool = True) -> None: self.discarded = 0 self.finalized = False self.left_pool = False # 0.27: elastic_left_pool seen (re-JOIN exhausted) + self.wan_rejects = 0 # S26 E3b: rl_wan_probe records with accepted=false self.closed = False @property @@ -179,8 +180,11 @@ def _record_line(self, line: str) -> None: with self.path.open("a", encoding="utf-8") as handle: handle.write(raw + "\n") self.count += 1 - event = json.loads(raw).get("event") - if event == FINALIZED_EVENT: + record = json.loads(raw) + event = record.get("event") + if event == "rl_wan_probe" and record.get("accepted") is False: + self.wan_rejects += 1 + elif event == FINALIZED_EVENT: self.finalized = True elif event == "elastic_left_pool": self.left_pool = True diff --git a/yeto/wan_probe.py b/yeto/wan_probe.py new file mode 100644 index 00000000..4ccb91fe --- /dev/null +++ b/yeto/wan_probe.py @@ -0,0 +1,149 @@ +"""Island-to-syncer link-bandwidth probe (S26 E3b). + +E3/E3b showed that the sync share follows the island *container*: on a slow +container every TCP stream is slow at once (1-5 MB/s each, both directions), +so no stream scheduling can help. An island can measure its link to the head +before it loads a model and exit early when the link is too slow; the +launcher then requests a fresh container (``--wan-probe-relaunches``). + +Each probe connection sends one ``MSG_BW_PROBE`` frame (HMAC-sealed like +HELLO), uploads ``upload_bytes`` raw bytes, reads the syncer's 8-byte +acknowledgement, then reads ``download_bytes`` raw bytes. The syncer never +connects a probe to a learner session. +""" + +from __future__ import annotations + +import socket +import struct +import threading +import time +from collections.abc import Callable +from dataclasses import dataclass + +from .protocol import MAGIC, MSG_BW_PROBE, ProtocolError, seal_hello, write_frame + +# Exit code of an island that rejected its own link (the launcher may relaunch it). +EXIT_SLOW_LINK = 75 +PROBE_EVENT = "rl_wan_probe" +PROBE_MARKER = "[yeto-wan-probe]" +_BLOCK = memoryview(bytes(1 << 20)) + + +@dataclass(frozen=True) +class ProbeResult: + streams: int + upload_bytes: int + download_bytes: int + upload_seconds: float + download_seconds: float + error: str | None = None + + @property + def upload_mbps(self) -> float: + return self.upload_bytes / self.upload_seconds / 1e6 if self.upload_seconds > 0 else 0.0 + + @property + def download_mbps(self) -> float: + return self.download_bytes / self.download_seconds / 1e6 if self.download_seconds > 0 else 0.0 + + def accepted(self, min_upload_mbps: float, min_download_mbps: float) -> bool: + return ( + self.error is None + and self.upload_mbps >= min_upload_mbps + and self.download_mbps >= min_download_mbps + ) + + def fields(self) -> dict[str, object]: + return { + "wan_probe/streams": self.streams, + "wan_probe/upload_bytes": self.upload_bytes, + "wan_probe/download_bytes": self.download_bytes, + "wan_probe/upload_mbps": round(self.upload_mbps, 3), + "wan_probe/download_mbps": round(self.download_mbps, 3), + "wan_probe/error": self.error, + } + + +def encode_probe(upload_bytes: int, download_bytes: int, key: bytes | None) -> bytes: + return seal_hello(struct.pack(" None: + sock = socket.create_connection(addr, timeout=timeout) + try: + sock.setsockopt(socket.IPPROTO_TCP, socket.TCP_NODELAY, 1) + write_frame(sock, MSG_BW_PROBE, encode_probe(upload, download, key)) + start = time.monotonic() + left = upload + while left: + take = min(left, len(_BLOCK)) + sock.sendall(_BLOCK[:take]) + left -= take + ack = b"" + while len(ack) < 8: + part = sock.recv(8 - len(ack)) + if not part: + # The syncer refused the probe: it answers with an ERROR frame + # and closes; the ERROR body is the reason. + raise ConnectionError("syncer closed the probe before acknowledging the upload") + ack += part + if ack.startswith(struct.pack(" ProbeResult: + """Measure aggregate upload and download rates over ``streams`` parallel + connections; ``bytes_per_direction`` is split evenly over the streams. + A rate is total bytes over the slowest stream's time in that direction.""" + if streams < 1 or bytes_per_direction < streams: + raise ValueError("probe needs at least one stream and one byte per stream") + per = bytes_per_direction // streams + results = [dict() for _ in range(streams)] + errors: list[str] = [] + + def run(index: int) -> None: + try: + stream_fn(addr, per, per, key, timeout, results[index]) + except (OSError, ProtocolError) as error: + errors.append(f"{type(error).__name__}: {error}") + + threads = [threading.Thread(target=run, args=(i,), daemon=True) for i in range(streams)] + for thread in threads: + thread.start() + for thread in threads: + thread.join(timeout + 5) + done = [r for r in results if "up" in r and "down" in r] + if len(done) < streams and not errors: + errors.append(f"{streams - len(done)} probe stream(s) did not finish in {timeout:.0f}s") + return ProbeResult( + streams=streams, + upload_bytes=per * len(done), + download_bytes=per * len(done), + upload_seconds=max((r["up"] for r in done), default=0.0), + download_seconds=max((r["down"] for r in done), default=0.0), + error=errors[0] if errors else None, + ) + From 87c298a02b81c94c087202acfcd0c720ba720a41 Mon Sep 17 00:00:00 2001 From: michael <78113554+chowhsu@users.noreply.github.com> Date: Wed, 21 Oct 2026 09:55:26 +0000 Subject: [PATCH 50/62] eval efficiency: spec.cards sets the per-run card total (E11 whole-cluster arms) 0 keeps the E0 shape and its control fingerprint; async splits the total evenly between rollout and trainer. Co-Authored-By: Claude Opus 5.5 --- tests/test_tools_eval_efficiency.py | 16 ++++++++++++++++ tools/eval/efficiency/spec.py | 29 ++++++++++++++++++++--------- 2 files changed, 36 insertions(+), 9 deletions(-) diff --git a/tests/test_tools_eval_efficiency.py b/tests/test_tools_eval_efficiency.py index 94911c11..6620e234 100644 --- a/tests/test_tools_eval_efficiency.py +++ b/tests/test_tools_eval_efficiency.py @@ -419,3 +419,19 @@ def test_tape_events_skip_the_repo_copy_and_the_store(tmp_path): (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"] + + +def test_cards_whole_cluster_split(): + from tools.eval.efficiency import spec as S + + assert S.cards("sync") == {"total": 1, "rollout": 1, "trainer": 1, "colocated": True} + assert S.cards("async") == {"total": 2, "rollout": 1, "trainer": 1, "colocated": False} + assert S.cards("async", 4) == {"total": 4, "rollout": 2, "trainer": 2, "colocated": False} + assert S.cards("sync", 4)["trainer"] == 4 + with pytest.raises(ValueError): + S.cards("async", 3) + base = S.get_spec("qwen3-4b-instruct-2507") + # cards=0 keeps the fingerprint of runs made before the field existed + assert S.control_fingerprint(base) != S.control_fingerprint(S.get_spec("qwen3-4b-instruct-2507", cards=4)) + plan = S.plan(S.get_spec("qwen3-4b-instruct-2507", cards=4), "bare-miles-async", "e11-x") + assert "--rollout-num-gpus 2" in plan["script"] and "--actor-num-gpus-per-node 2" in plan["script"] diff --git a/tools/eval/efficiency/spec.py b/tools/eval/efficiency/spec.py index 174e9eb6..97c30d3a 100644 --- a/tools/eval/efficiency/spec.py +++ b/tools/eval/efficiency/spec.py @@ -84,6 +84,9 @@ 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" + # cards per run: 0 keeps the E0 shape (sync 1, async 1 rollout + 1 trainer); E11 + # (whole-cluster comparison) sets the total, split evenly between rollout and trainer in async + cards: int = 0 @property def samples_per_round(self) -> int: @@ -140,16 +143,24 @@ def get_spec(preset: str, **overrides: Any) -> EffSpec: 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 cards(mode: str, total: int = 0) -> dict[str, int]: + """Cards per role: sync colocates both on every card; async splits them evenly + (``total`` 0: one card for sync, one of each role for async).""" + if mode == "sync": + n = total or 1 + return {"total": n, "rollout": n, "trainer": n, "colocated": True} + n = total or 2 + if n % 2: + raise ValueError(f"async needs an even card count (rollout + trainer halves), got {n}") + return {"total": n, "rollout": n // 2, "trainer": n // 2, "colocated": False} def control_fingerprint(spec: EffSpec) -> str: """sha256 over every shared value except the repeat count (a report refuses to compare runs whose fingerprints differ).""" doc = {k: v for k, v in spec.to_dict().items() if k not in ("repeats", "timeout_s", "checkpoint")} + if not spec.cards: + doc.pop("cards") # the E0 shape: fingerprints from before the field existed stay valid return hashlib.sha256(json.dumps(doc, sort_keys=True).encode()).hexdigest() @@ -168,7 +179,7 @@ def _model_path(spec: EffSpec) -> str: # ====================================================================== 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) + c = cards(mode, spec.cards) data = [ "algorithm.adv_estimator=grpo", "algorithm.use_kl_in_reward=False", @@ -271,7 +282,7 @@ def verl_bare_entry(mode: str) -> str: # ====================================================================== 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) + c = cards(mode, spec.cards) model = _model_path(spec) argv = [ # checkpoint: bridge mode loads the HF checkpoint directly (run_qwen3_4b.py --megatron-to-hf-mode bridge) @@ -388,7 +399,7 @@ def yeto_miles_launcher_args(spec: EffSpec, mode: str, run_id: str) -> str: from ..miles import DATA, STORE_VOLUME, TAPE_VOLUME card = {"A10G": "a10g", "H100!": "h100", "L4": "l4"}[spec.gpu] - c = cards(mode) + c = cards(mode, spec.cards) 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)}" @@ -424,7 +435,7 @@ def yeto_miles_launcher_args(spec: EffSpec, mode: str, run_id: str) -> str: # ====================================================================== 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) + c = cards(arm.mode, spec.cards) 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": @@ -456,7 +467,7 @@ def bare_script(spec: EffSpec, arm: Arm, run_id: str) -> str: 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) + c = cards(arm.mode, spec.cards) 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, From 823b2e82e6bbfbcce8b4ca904707a283e758fa90 Mon Sep 17 00:00:00 2001 From: michael <78113554+chowhsu@users.noreply.github.com> Date: Wed, 21 Oct 2026 10:02:09 +0000 Subject: [PATCH 51/62] wan: refuse the link probe with the elastic syncer; test a lost socket mid-INIT Only the legacy syncer serves MSG_BW_PROBE. With --rl-elastic every probe would fail, every island would be judged slow and relaunched until its budget ran out, so the launch now refuses the probe flags there. Tests: - real syncer: one data socket of learner 0's group is shut down between the INIT of fragment 0 and fragment 1; the client replaces the group, replays INIT, and the final versions and bytes equal an undisturbed run. - slow-container decision: threshold equality accepts, upload just below rejects, slow download rejects, a zero download minimum is not checked, and a stream error rejects a fast link. Co-Authored-By: Claude Opus 5.5 --- tests/test_rl_dense_sweep_real_syncer.py | 46 ++++++++++++++++++++++++ tests/test_wan_probe.py | 24 +++++++++++++ yeto/launch/rl_island_task.py | 5 +++ 3 files changed, 75 insertions(+) diff --git a/tests/test_rl_dense_sweep_real_syncer.py b/tests/test_rl_dense_sweep_real_syncer.py index 7fd428eb..76351037 100644 --- a/tests/test_rl_dense_sweep_real_syncer.py +++ b/tests/test_rl_dense_sweep_real_syncer.py @@ -697,3 +697,49 @@ def listening() -> bool: _stop(process) assert list(manifest.versions) == exact[1] assert result == exact[2] + + +def test_one_socket_lost_mid_init_is_redialed_and_the_result_stays_exact(syncer_binary, tmp_path): + """E3c: a single data socket of learner 0's group fails between INIT of + fragment 0 and fragment 1 (as in s26-e3b-p06-C32a, where one connection + reported BrokenPipe). The client replaces the whole group, replays INIT, + and the final versions and bytes equal an undisturbed run.""" + exact = _run_sweeps(syncer_binary, tmp_path, STREAMED, sweeps=3) + run_dir = tmp_path / "one-socket" + run_dir.mkdir() + process, port, profile_hash = _start_syncer(syncer_binary, run_dir, *STREAMED, total_steps=3) + client = SyncerClient(("127.0.0.1", port), 0, _layout(), dtype=DTYPE_F32, num_streams=2, + connect_timeout=10, max_reconnects=5, finalization_timeout=20, + syncer_profile_hash=bytes.fromhex(profile_hash)) + broken = [] + + def parts_for(fid: int): + if fid == 1 and not broken: + with client._lock: + sock = client._socks[-1] + broken.append(client.connection_generation) + sock.shutdown(socket.SHUT_RDWR) # one real socket dies; the others stay open + return [_f32(_initial_values(fid))] + + try: + client.start() + client.send_initial_fragments(parts_for) + replays = 0 + deadline = time.monotonic() + 60 + while client._bcasts.qsize() < FRAGMENTS: + assert time.monotonic() < deadline, "initial cut never arrived" + client.check_health() + if client.init_replay_due(): + replays += client.send_initial_fragments(parts_for) + time.sleep(0.02) + client.mark_initial_cut_received() + assert broken and replays >= 1 and client.connection_generation != broken[0] + manifest, finals = _sweep_until_final(client, sweeps=3) + result = {f.fragment_id: bytes(f.data) for f in finals} + client.acknowledge_finalization(manifest) + assert process.wait(timeout=15) == 0 + finally: + client.close() + _stop(process) + assert list(manifest.versions) == exact[1] + assert result == exact[2] diff --git a/tests/test_wan_probe.py b/tests/test_wan_probe.py index 9ce848ab..ad8a19df 100644 --- a/tests/test_wan_probe.py +++ b/tests/test_wan_probe.py @@ -124,3 +124,27 @@ def test_rejected_link_budget_is_bounded(): with pytest.raises(FixedRosterIslandAbandoned): ctl.run() assert ops.relaunch_calls == [] + + +@pytest.mark.parametrize( + ("up_s", "down_s", "error", "min_up", "min_down", "accepted"), + [ + (1.0, 1.0, None, 100.0, 100.0, True), # 100 MB in 1 s each way: exactly at the minimum + (1.0, 1.0, None, 100.1, 0.0, False), # upload just below the minimum + (1.0, 4.0, None, 50.0, 50.0, False), # fast upload, slow download + (1.0, 4.0, None, 50.0, 0.0, True), # download not checked when its minimum is 0 + (0.1, 0.1, "OSError: x", 1.0, 1.0, False), # any stream error rejects a fast link + ], +) +def test_slow_container_decision(up_s, down_s, error, min_up, min_down, accepted): + result = ProbeResult(streams=8, upload_bytes=100_000_000, download_bytes=100_000_000, + upload_seconds=up_s, download_seconds=down_s, error=error) + assert result.accepted(min_up, min_down) is accepted + + +def test_probe_is_refused_with_the_elastic_syncer(): + from yeto.launch.rl_island_task import _wan_transport_flags + + with pytest.raises(ValueError, match="elastic"): + _wan_transport_flags(Namespace(wan_probe_min_mbps=20.0, rl_elastic=True)) + assert _wan_transport_flags(Namespace(rl_elastic=True)) == "" # probe off: no error diff --git a/yeto/launch/rl_island_task.py b/yeto/launch/rl_island_task.py index 0254a603..ce0ecea7 100644 --- a/yeto/launch/rl_island_task.py +++ b/yeto/launch/rl_island_task.py @@ -67,6 +67,11 @@ def _wan_transport_flags(args) -> str: if min_up < 0 or min_down < 0: raise ValueError("--wan-probe-min-mbps / --wan-probe-min-down-mbps must be >= 0 (0 = off)") if min_up > 0 or min_down > 0: + if getattr(args, "rl_elastic", False): + # Only the legacy syncer serves MSG_BW_PROBE; on the elastic server + # every probe would fail and every island would be rejected. + raise ValueError("--wan-probe-min-mbps / --wan-probe-min-down-mbps need the " + "legacy syncer; the elastic syncer has no link probe") mib = int(getattr(args, "wan_probe_mib", 256)) if not 1 <= mib <= 1024: raise ValueError("--wan-probe-mib must be between 1 and 1024") From 140aefa299890cd86b2a01be4ca61e256d61714b Mon Sep 17 00:00:00 2001 From: michael <78113554+chowhsu@users.noreply.github.com> Date: Wed, 21 Oct 2026 10:03:30 +0000 Subject: [PATCH 52/62] eval efficiency: E11 preset and the 2-island yeto+Miles launcher args Co-Authored-By: Claude Opus 5.5 --- tools/eval/efficiency/spec.py | 56 +++++++++++++++++++++++++++++++++++ 1 file changed, 56 insertions(+) diff --git a/tools/eval/efficiency/spec.py b/tools/eval/efficiency/spec.py index 97c30d3a..5a316a9a 100644 --- a/tools/eval/efficiency/spec.py +++ b/tools/eval/efficiency/spec.py @@ -111,6 +111,13 @@ def to_dict(self) -> dict[str, Any]: 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), + # E11 (whole cluster): the E0.2 spec on 4 cards; async (2 rollout + 2 trainer, or 2 + # islands of 1 + 1), so the 2-island yeto arm and the single-island arms share one shape + "e11-qwen3-4b": EffSpec( + name="e11-qwen3-4b", 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, cards=4), } @@ -432,6 +439,55 @@ def yeto_miles_launcher_args(spec: EffSpec, mode: str, run_id: str) -> str: return args +HEAD_REGION = "nebius/eu-north1" # the multi-island syncer head: a CPU-only Nebius VM + + +def yeto_miles_islands_launcher_args(spec: EffSpec, islands: int, run_id: str, *, data: str, + island_tasks: str) -> str: + """E11: the yeto+Miles arm split into ``islands`` async islands that sync through a + head syncer (LoRA strict-avg), with the card total and the prompts per round of the + single-island arm. ``data`` is a local JSONL of the spec's rows and ``island_tasks`` + the ``--rl-island-tasks`` spec that gives each island its share of every round.""" + c = cards("async", spec.cards) + if c["total"] % islands or (c["total"] // islands) % 2 or spec.groups_per_round % islands: + raise ValueError(f"{islands} islands need an even card share and an equal prompt share " + f"({c['total']} cards, {spec.groups_per_round} prompts per round)") + per = c["total"] // islands + card = {"A10G": "a10g", "H100!": "h100", "L4": "l4"}[spec.gpu] + kwargs = json.dumps(spec.chat_template_kwargs) if spec.chat_template_kwargs else "" + args = (f"launch --controller head --syncer-region {HEAD_REGION} --on-demand --training-mode rl " + f"--rl-engine ports --preflight-threads wait --preflight-thread-start 2950 --modal-gpu-exact " + f"{gpu_name_arg(spec)}--modal-retries 0 --no-island-relaunch " + f"--modal-cpu {8 * per} --modal-memory-gib {48 * per} " + f"--cluster-prefix {run_id} --tuning lora --lora-r {spec.lora_r} --lora-targets all-linear " + f"--rl-observe-timeline --model {spec.model} --model-revision {spec.model_revision} " + f"--rl-allow-local-data --data {data} --rl-island-tasks {shlex.quote(island_tasks)} " + "--reward-function gsm8k_reward:score --fragments 1 --pipeline 1 " + f"--rollout-batch-size {spec.groups_per_round // islands} " + f"--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 {','.join([f'modal:{per}x{card}'] * islands)} --total-steps {spec.rounds} " + # the single-island arm's Miles recipe settings (yeto_miles_launcher_args) + "--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 " + f"--trust-remote-code --rl-algorithm-spec {MILES_ALGORITHM_SPEC} " + f"--rl-placement fixed-partition --rl-rollout-gpus {per // 2} " + f"--rl-max-policy-age {spec.max_policy_age}") + if kwargs: + args += f" --apply-chat-template-kwargs {shlex.quote(kwargs)}" + return args + + +def island_task_split(groups_per_round: int, rounds: int, islands: int) -> dict[str, list[str]]: + """Round-major split of the first ``rounds * groups_per_round`` rows (task ids + ``gsm``): island k gets the k-th equal slice of every round's prompts, so the + islands together read the single-island arm's prompts of each round.""" + share = groups_per_round // islands + return {f"l{k}": [f"gsm{r * groups_per_round + k * share + i + 1}" for r in range(rounds) + for i in range(share)] for k in range(islands)} + + # ====================================================================== 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).""" From b0dfb1cf18caca6e5de66a983ba344c30977f7cd Mon Sep 17 00:00:00 2001 From: michael <78113554+chowhsu@users.noreply.github.com> Date: Wed, 21 Oct 2026 10:22:29 +0000 Subject: [PATCH 53/62] eval efficiency: JSONL datasets, a train-row cap and local task files for the yeto arms E11 trains on DAPO-Math-17k (gsm8k saturates Qwen3-4B-Instruct-2507). eff_prep reads a JSONL data_file and keeps the first train_rows rows; the yeto+Miles arms read the same rows as a local task file (EXTRA_FILES in miles_local.sh). The 2-island arm splits every round's prompts between the islands. Co-Authored-By: Claude Opus 5.5 --- tests/test_tools_eval_efficiency.py | 34 +++++++++++++++++++ tools/eval/efficiency/eff_prep.py | 50 ++++++++++++++++++++++----- tools/eval/efficiency/modal_eff.py | 5 +++ tools/eval/efficiency/spec.py | 52 ++++++++++++++++++++++++----- tools/eval/miles_local.sh | 7 ++-- 5 files changed, 129 insertions(+), 19 deletions(-) diff --git a/tests/test_tools_eval_efficiency.py b/tests/test_tools_eval_efficiency.py index 6620e234..431a80c1 100644 --- a/tests/test_tools_eval_efficiency.py +++ b/tests/test_tools_eval_efficiency.py @@ -435,3 +435,37 @@ def test_cards_whole_cluster_split(): assert S.control_fingerprint(base) != S.control_fingerprint(S.get_spec("qwen3-4b-instruct-2507", cards=4)) plan = S.plan(S.get_spec("qwen3-4b-instruct-2507", cards=4), "bare-miles-async", "e11-x") assert "--rollout-num-gpus 2" in plan["script"] and "--actor-num-gpus-per-node 2" in plan["script"] + + +def test_jsonl_source_rows_and_yeto_task_file(tmp_path): + from tools.eval.efficiency import eff_prep + + src = tmp_path / "d.jsonl" + src.write_text("".join(json.dumps({"prompt": [{"role": "user", "content": f"q{i}"}], "label": str(i)}) + "\n" + for i in range(10))) + sp = spec.get_spec("e11-qwen3-4b", train_rows=4).to_dict() + rows = eff_prep.source_rows(sp, lambda *a, **k: str(src)) + assert [lab for _, lab in rows["train"]] == ["0", "1", "2", "3"] + assert [lab for _, lab in rows["test"]] == [str(i) for i in range(4, 10)] + tasks = eff_prep.yeto_task_rows(sp, lambda *a, **k: str(src)) + assert tasks[2] == {"messages": [{"role": "user", "content": "q2"}], "label": "2", "metadata": {"task_id": "gsm3"}} + with pytest.raises(ValueError): + eff_prep.source_rows({**sp, "train_rows": 0}, lambda *a, **k: str(src)) + + +def test_e11_arms_read_one_row_set(): + e11 = spec.get_spec("e11-qwen3-4b") + assert e11.train_rows == e11.rounds * e11.groups_per_round + single = spec.yeto_miles_launcher_args(e11, "async", "e11-x") + assert f"--rl-allow-local-data --data {spec.YETO_TASK_FILE}" in single and "--gpu modal:4xh100" in single + multi = spec.yeto_miles_islands_launcher_args(e11, 2, "e11-y", data="t.jsonl", island_tasks="t.json") + assert "--gpu modal:2xh100,modal:2xh100" in multi and "--rollout-batch-size 16" in multi + assert "--rl-rollout-gpus 1" in multi and "--controller head" in multi + split = spec.island_task_split(e11.groups_per_round, e11.rounds, 2) + ids = sorted(int(t[3:]) for v in split.values() for t in v) + assert ids == list(range(1, e11.train_rows + 1)) + assert split["l1"][:2] == ["gsm17", "gsm18"] + with pytest.raises(ValueError): + spec.yeto_verl_argv(e11, "async", "/t") + # the E0 presets keep their fingerprints (new fields at their defaults are left out) + assert spec.control_fingerprint(spec.get_spec("qwen3-0.6b")).startswith("59ed2ab081afb9f9") diff --git a/tools/eval/efficiency/eff_prep.py b/tools/eval/efficiency/eff_prep.py index 94fdee80..90e9bb73 100644 --- a/tools/eval/efficiency/eff_prep.py +++ b/tools/eval/efficiency/eff_prep.py @@ -28,6 +28,42 @@ VAL_ROWS = 64 +def source_rows(spec: dict, download) -> dict[str, list[tuple[list, str]]]: + """Split -> ``(chat messages, label)`` in the dataset's own row order. + + Two layouts: ``train.parquet``/``test.parquet`` with ``messages`` and ``label`` + (zhuzilin/gsm8k), or one JSONL file ``spec["data_file"]`` with ``prompt`` (chat + messages) and ``label`` (zhuzilin/dapo-math-17k), whose test rows are the + ``VAL_ROWS`` rows after the training rows. ``spec["train_rows"]`` (0: all) keeps + the first rows only, so every arm trains on the same prompt set in one epoch. + """ + import pandas as pd + + n = int(spec.get("train_rows") or 0) + if spec.get("data_file"): + path = download(spec["data"], spec["data_file"], repo_type="dataset", revision=spec["data_revision"]) + with open(path) as fh: + records = [(r["prompt"], str(r["label"])) for r in map(json.loads, fh)] + if not n: + raise ValueError("a JSONL data_file needs train_rows (its test rows follow them)") + return {"train": records[:n], "test": records[n:n + VAL_ROWS]} + out = {} + for split in ("train", "test"): + frame = pd.read_parquet(download(spec["data"], f"{split}.parquet", repo_type="dataset", + revision=spec["data_revision"])) + records = [(list(row["messages"]), str(row["label"])) for _, row in frame.iterrows()] + out[split] = records[:n] if (split == "train" and n) else records[:VAL_ROWS] if split == "test" else records + return out + + +def yeto_task_rows(spec: dict, download) -> list[dict]: + """The training rows as a yeto local task file (``messages``, ``label``, task id + ``gsm``): the yeto arms read the same rows as the bare arms' ``train.jsonl``.""" + return [{"messages": [{"role": m["role"], "content": m["content"]} for m in messages], "label": label, + "metadata": {"task_id": f"gsm{i + 1}"}} + for i, (messages, label) in enumerate(source_rows(spec, download)["train"])] + + def main(spec_path: str) -> int: import pandas as pd from huggingface_hub import hf_hub_download, snapshot_download @@ -37,18 +73,14 @@ def main(spec_path: str) -> int: 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) + for split, records in source_rows(spec, hf_hub_download).items(): rows = [{ "data_source": "yeto/" + spec["data"], - "prompt": [{"role": m["role"], "content": m["content"]} for m in row["messages"]], + "prompt": [{"role": m["role"], "content": m["content"]} for m in 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] + "reward_model": {"style": "rule", "ground_truth": str(label)}, + "extra_info": {"split": split, "index": i}, + } for i, (messages, label) in enumerate(records)] 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()} diff --git a/tools/eval/efficiency/modal_eff.py b/tools/eval/efficiency/modal_eff.py index d67dd09e..d063046f 100644 --- a/tools/eval/efficiency/modal_eff.py +++ b/tools/eval/efficiency/modal_eff.py @@ -28,6 +28,7 @@ from typing import Any from .. import modal_verl as mv +from . import spec from .spec import EFF_DIR, OUT_DIR HERE = Path(__file__).resolve().parent @@ -175,6 +176,10 @@ def run_miles_launcher(plan: dict, run_dir: Path, *, dry_run: bool = False) -> d HOME="/home/michael", ARGS=plan["launcher_args"], ISLANDS="1") if dry_run: env["DRY"] = "1" + sp = spec.EffSpec(**plan["spec"]) + if spec.uses_task_file(sp): + run_dir.mkdir(parents=True, exist_ok=True) + env["EXTRA_FILES"] = str(spec.write_task_file(sp, run_dir / spec.YETO_TASK_FILE)) 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, diff --git a/tools/eval/efficiency/spec.py b/tools/eval/efficiency/spec.py index 5a316a9a..5a91d6ce 100644 --- a/tools/eval/efficiency/spec.py +++ b/tools/eval/efficiency/spec.py @@ -63,6 +63,8 @@ class EffSpec: miles_model_args: str # file under /scripts/models/ (the official model flags) data: str = "zhuzilin/gsm8k" data_revision: str = "0cbd9f31d91ac21a7613dcbc7fef992adac459ae" + data_file: str = "" # one JSONL file of the dataset repo (prompt, label); "" = train/test.parquet + train_rows: int = 0 # keep the first rows only (0 = all); needed with data_file 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 @@ -111,12 +113,16 @@ def to_dict(self) -> dict[str, Any]: 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), - # E11 (whole cluster): the E0.2 spec on 4 cards; async (2 rollout + 2 trainer, or 2 - # islands of 1 + 1), so the 2-island yeto arm and the single-island arms share one shape + # E11 (whole cluster): the E0.2 model and hyperparameters on 4 cards; async (2 rollout + + # 2 trainer, or 2 islands of 1 + 1), so the 2-island yeto arm and the single-island arms + # share one shape. Data: DAPO-Math-17k (gsm8k is saturated for this model: reward 0.97 + # in round 1 of e02-bare-verl-sync-1), the first rounds x groups rows (one epoch). "e11-qwen3-4b": EffSpec( name="e11-qwen3-4b", 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, + data="zhuzilin/dapo-math-17k", data_revision="2e65612930298bde4c5d58fd97b3f23a483aaff9", + data_file="dapo-math-17k.jsonl", train_rows=12 * 32, + gpu="H100!", rounds=12, 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, cards=4), } @@ -166,8 +172,9 @@ 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")} - if not spec.cards: - doc.pop("cards") # the E0 shape: fingerprints from before the field existed stay valid + for key in ("cards", "data_file", "train_rows"): # fields added after E0.1: a default value + if not doc[key]: # keeps the fingerprints of earlier runs valid + doc.pop(key) return hashlib.sha256(json.dumps(doc, sort_keys=True).encode()).hexdigest() @@ -361,6 +368,8 @@ def max_tokens_per_gpu(spec: EffSpec) -> int: 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).""" + if uses_task_file(spec): + raise ValueError("the yeto+verl arm reads --data from the Hub; it has no task-file path") 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), @@ -403,7 +412,7 @@ def gpu_name_arg(spec: EffSpec) -> 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, STORE_VOLUME, TAPE_VOLUME + from ..miles import STORE_VOLUME, TAPE_VOLUME card = {"A10G": "a10g", "H100!": "h100", "L4": "l4"}[spec.gpu] c = cards(mode, spec.cards) @@ -414,7 +423,7 @@ def yeto_miles_launcher_args(spec: EffSpec, mode: str, run_id: str) -> str: 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-observe-timeline --model {spec.model} --model-revision {spec.model_revision} {_yeto_data(spec)} " f"--rl-checkpoint-store modal-volume://{STORE_VOLUME}/{run_id} " f"--rl-cut-every {cut_every(spec)} " "--reward-function gsm8k_reward:score --fragments 1 --pipeline 1 " @@ -434,9 +443,36 @@ def yeto_miles_launcher_args(spec: EffSpec, mode: str, run_id: str) -> str: if mode == "async": args += (f" --rl-placement fixed-partition --rl-rollout-gpus {c['rollout']} " f"--rl-max-policy-age {spec.max_policy_age}") + return args + + +YETO_TASK_FILE = "eff-tasks.jsonl" # the local task file in the launcher's checkout (local_task_file) + + +def uses_task_file(spec: EffSpec) -> bool: + """The yeto+Miles arms read a local task file (eff_prep.yeto_task_rows) unless the + spec is the full gsm8k train split, which they read from the Hub as before.""" + return bool(spec.data_file or spec.train_rows) + + +def write_task_file(spec: EffSpec, path: Path) -> Path: + """Write the yeto arms' local task file (the bare arms' training rows, in order).""" + from huggingface_hub import hf_hub_download + + from .eff_prep import yeto_task_rows + + path.write_text("".join(json.dumps(r) + "\n" for r in yeto_task_rows(spec.to_dict(), hf_hub_download))) + return path + + +def _yeto_data(spec: EffSpec) -> str: + from ..miles import DATA + + if uses_task_file(spec): + return f"--rl-allow-local-data --data {YETO_TASK_FILE}" 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 + return DATA HEAD_REGION = "nebius/eu-north1" # the multi-island syncer head: a CPU-only Nebius VM diff --git a/tools/eval/miles_local.sh b/tools/eval/miles_local.sh index 6a24cba1..0c634b04 100644 --- a/tools/eval/miles_local.sh +++ b/tools/eval/miles_local.sh @@ -9,7 +9,8 @@ # ARGS (launcher args; the store is modal-volume://$STOREVOL/$P), SARGS (optional), # ISLANDS (containers before the GPU lock is released), HARD (s), STOREVOL, TAPEVOL, # RECLAIM_LABEL (e.g. "l0") + RECLAIM_ROUND (rollout id written before the stop), -# PY (python with modal + the launcher deps), DRY (non-empty: --dry-run only) +# PY (python with modal + the launcher deps), DRY (non-empty: --dry-run only), +# EXTRA_FILES (optional: files copied into the checkout root and committed) set -u : "${P:?}" "${R:?}" "${REPO:?}" "${ARGS:?}" "${STOREVOL:?}" "${PY:?}" HARD=${HARD:-3000}; ISLANDS=${ISLANDS:-1}; TAPEVOL=${TAPEVOL:-yeto-event-tapes} @@ -22,7 +23,9 @@ mkdir -p $R/tapes $R/store # a real clone: the code pack needs a git checkout git clone -q --no-local $REPO $CLONE && git -C $REPO rev-parse HEAD > $R/yeto_sha.txt cp $REPO/yeto/rl/gsm8k_reward.py $CLONE/gsm8k_reward.py -(cd $CLONE && git add -f gsm8k_reward.py && git -c user.email=eval@yeto -c user.name=t-eval commit -qm 'run inputs') +# EXTRA_FILES: run inputs the launcher reads from its checkout (e.g. a local task file) +for f in ${EXTRA_FILES:-}; do cp "$f" $CLONE/; done +(cd $CLONE && git add -f gsm8k_reward.py $(for f in ${EXTRA_FILES:-}; do basename "$f"; done) && git -c user.email=eval@yeto -c user.name=t-eval commit -qm 'run inputs') eval "$(/usr/bin/python3 - <<'PYC' import shlex, tomllib t = tomllib.load(open("/home/michael/.modal.toml", "rb")) From 392368b4604a6555dcdea2ecac0d4da3e991fc28 Mon Sep 17 00:00:00 2001 From: michael <78113554+chowhsu@users.noreply.github.com> Date: Wed, 21 Oct 2026 10:46:26 +0000 Subject: [PATCH 54/62] eval efficiency: derived Miles model scripts and the yeto+verl rollout memory fraction The bare Miles loader puts model_args_utils on the path (qwen3-4B-Instruct-2507.py imports it: e11-bare-miles-async-1). The yeto+verl arm passes the spec's mem_fraction as YETO_VERL_GPU_MEM; at the 0.5 default vLLM cannot fit the 4B model's 262144-token context (e02-yeto-verl-sync-1). Co-Authored-By: Claude Opus 5.5 --- tools/eval/efficiency/modal_eff.py | 2 +- tools/eval/efficiency/spec.py | 6 +++++- 2 files changed, 6 insertions(+), 2 deletions(-) diff --git a/tools/eval/efficiency/modal_eff.py b/tools/eval/efficiency/modal_eff.py index d063046f..3aa0f21c 100644 --- a/tools/eval/efficiency/modal_eff.py +++ b/tools/eval/efficiency/modal_eff.py @@ -82,7 +82,7 @@ def arm(run: str, plan: dict) -> dict: 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) + env = {**mv._island_env({"gpu": plan["gpu"]}, 0), **plan.get("env", {})} 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", diff --git a/tools/eval/efficiency/spec.py b/tools/eval/efficiency/spec.py index 5a91d6ce..60ab9b48 100644 --- a/tools/eval/efficiency/spec.py +++ b/tools/eval/efficiency/spec.py @@ -540,7 +540,10 @@ def bare_script(spec: EffSpec, arm: Arm, run_id: str) -> str: "exec python3 -m " + verl_bare_entry(arm.mode) + " " + shlex.join( verl_bare_overrides(spec, arm.mode, run_id))] else: + # model scripts derived from another (qwen3-4B-Instruct-2507.py) import Miles' + # model_args_utils (miles/utils/external_utils), which is not on the scripts' path load = ("import importlib.util as u, sys; sys.path.insert(0, '{d}'); " + f"sys.path.insert(0, '{MILES_ROOT}/miles/utils/external_utils'); " "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) @@ -567,7 +570,8 @@ def plan(spec: EffSpec, arm_name: str, run_id: str) -> dict[str, Any]: 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", + # island_entry reads the rollout memory fraction from YETO_VERL_GPU_MEM (default 0.5) + doc.update(harness="modal_eff", image="verl", env={"YETO_VERL_GPU_MEM": str(spec.mem_fraction)}, 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)) From 5593787fe9c534749b59285741c9f913cb263571 Mon Sep 17 00:00:00 2001 From: michael <78113554+chowhsu@users.noreply.github.com> Date: Wed, 21 Oct 2026 11:06:38 +0000 Subject: [PATCH 55/62] eval efficiency: the E11 2-island arm uses elastic island scheduling Legacy scheduling with LoRA Miles islands hung in the first policy apply (e11-yeto-2i-1); elastic is the multi-island LoRA path with real-card runs. Co-Authored-By: Claude Opus 5.5 --- tests/test_tools_eval_efficiency.py | 2 +- tools/eval/efficiency/spec.py | 3 +++ 2 files changed, 4 insertions(+), 1 deletion(-) diff --git a/tests/test_tools_eval_efficiency.py b/tests/test_tools_eval_efficiency.py index 431a80c1..17eaee15 100644 --- a/tests/test_tools_eval_efficiency.py +++ b/tests/test_tools_eval_efficiency.py @@ -460,7 +460,7 @@ def test_e11_arms_read_one_row_set(): assert f"--rl-allow-local-data --data {spec.YETO_TASK_FILE}" in single and "--gpu modal:4xh100" in single multi = spec.yeto_miles_islands_launcher_args(e11, 2, "e11-y", data="t.jsonl", island_tasks="t.json") assert "--gpu modal:2xh100,modal:2xh100" in multi and "--rollout-batch-size 16" in multi - assert "--rl-rollout-gpus 1" in multi and "--controller head" in multi + assert "--rl-rollout-gpus 1" in multi and "--controller head" in multi and "--rl-island-scheduling elastic" in multi split = spec.island_task_split(e11.groups_per_round, e11.rounds, 2) ids = sorted(int(t[3:]) for v in split.values() for t in v) assert ids == list(range(1, e11.train_rows + 1)) diff --git a/tools/eval/efficiency/spec.py b/tools/eval/efficiency/spec.py index 60ab9b48..53173396 100644 --- a/tools/eval/efficiency/spec.py +++ b/tools/eval/efficiency/spec.py @@ -508,6 +508,9 @@ def yeto_miles_islands_launcher_args(spec: EffSpec, islands: int, run_id: str, * "--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 " f"--trust-remote-code --rl-algorithm-spec {MILES_ALGORITHM_SPEC} " + # elastic: the multi-island LoRA path with real-card runs (S24 d4, S25 c3); legacy + # scheduling hung in the first apply on both islands (e11-yeto-2i-1, no round in 13 min) + "--rl-island-scheduling elastic " f"--rl-placement fixed-partition --rl-rollout-gpus {per // 2} " f"--rl-max-policy-age {spec.max_policy_age}") if kwargs: From ae1a26e7ba3d54babb4d8c0c3222c7406be5b4b7 Mon Sep 17 00:00:00 2001 From: michael <78113554+chowhsu@users.noreply.github.com> Date: Wed, 21 Oct 2026 11:08:43 +0000 Subject: [PATCH 56/62] eval efficiency: cluster.py, reward against samples and time for the E11 arms Co-Authored-By: Claude Opus 5.5 --- tests/test_tools_eval_efficiency.py | 14 +++ tools/eval/efficiency/cluster.py | 161 ++++++++++++++++++++++++++++ 2 files changed, 175 insertions(+) create mode 100644 tools/eval/efficiency/cluster.py diff --git a/tests/test_tools_eval_efficiency.py b/tests/test_tools_eval_efficiency.py index 17eaee15..6e8b841e 100644 --- a/tests/test_tools_eval_efficiency.py +++ b/tests/test_tools_eval_efficiency.py @@ -469,3 +469,17 @@ def test_e11_arms_read_one_row_set(): spec.yeto_verl_argv(e11, "async", "/t") # the E0 presets keep their fingerprints (new fields at their defaults are left out) assert spec.control_fingerprint(spec.get_spec("qwen3-0.6b")).startswith("59ed2ab081afb9f9") + + +def test_cluster_target_and_smoothing(): + from tools.eval.efficiency import cluster as C + + pts = [C.Point(100.0 + 10 * i, 128, r) for i, r in enumerate([0.1, 0.1, 0.3, 0.5, 0.5, 0.5])] + c = C.Curve("r", "yeto-miles-2islands", 4, 0.0, 200.0, pts) + sm = C.smoothed(pts, 256) + assert sm[0] == 0.1 and sm[2] == pytest.approx(0.2) and sm[4] == pytest.approx(0.5) + assert C.start_reward([c], 256) == pytest.approx(0.1) + m = C.metrics(c, target=0.4, window=256, usd_per_card_hour=5.0) + assert m["reached"] and m["to_target_s"] == 130.0 and m["to_target_samples"] == 4 * 128 + assert m["to_target_gpu_h"] == pytest.approx(130 / 3600 * 4) + assert not C.metrics(c, target=0.9, window=256, usd_per_card_hour=5.0)["reached"] diff --git a/tools/eval/efficiency/cluster.py b/tools/eval/efficiency/cluster.py new file mode 100644 index 00000000..6ed2b040 --- /dev/null +++ b/tools/eval/efficiency/cluster.py @@ -0,0 +1,161 @@ +"""E11 whole-cluster comparison: reward against samples and time on one card total. + +Every arm becomes one list of :class:`Point` (one per trained batch, in end-time +order). A single-island arm gives one point per round (``parsed.json`` of an +``eff-run`` directory); the 2-island yeto arm (``logs/E11/run2i.sh``) gives one +point per island round, read from each island's tape, so its two islands +interleave by time. + +Definitions (SESSION26/gpu-review/E11.md §3, fixed before the runs): + +* reward of a point = mean reward of the batch's rollout samples; +* smoothed reward = sample-weighted mean over the latest ``window`` samples + (3 single-island rounds), evaluated at every point; +* target T = sample-weighted mean reward of the first ``2 * round_samples`` + samples of every run, pooled, plus ``margin``; +* time to T = first point whose smoothed reward >= T, measured from the run's + launch (``t0``: the local launch command) and from the end of its first batch + (no cold start); GPU hours = that time x the run's card total. +""" + +from __future__ import annotations + +import json +from dataclasses import asdict, dataclass +from datetime import datetime, timezone +from pathlib import Path +from typing import Any + +from .parse import _tape_events, _utc_file, parse_yeto + + +@dataclass(frozen=True) +class Point: + t_end: float + samples: int + reward: float + + +@dataclass +class Curve: + run_id: str + arm: str + cards: int + t0: float # local launch start (unix s) + t_last: float # the run's GPU containers ended (or its last point) + points: list[Point] + usd_extra: float = 0.0 # e.g. the 2-island arm's head VM + + +def _from_parsed(run_dir: Path) -> Curve: + parsed = json.loads((run_dir / "parsed.json").read_text()) + pts = [Point(float(r["t_end_unix"]), int(r["samples"]), float(r["reward_mean"])) + for r in sorted(parsed["rounds"], key=lambda r: r["round"]) + if r.get("t_end_unix") is not None and r.get("samples") and r.get("reward_mean") is not None] + t0 = _launch_start(run_dir) or float(parsed["container_start_unix"]) + end = parsed.get("container_end_unix") or (pts[-1].t_end if pts else t0) + return Curve(parsed["run_id"], parsed["arm"], int(parsed["cards"]["total"]), t0, float(end), pts) + + +def _launch_start(run_dir: Path) -> float | None: + """The local launch time: miles_local.sh's start_utc.txt, else the driver's first log line.""" + utc = _utc_file(run_dir / "raw" / "start_utc.txt") + if utc: + return utc + log = run_dir / "driver.log" + if not log.is_file(): + return None + try: + first = log.read_text().split(" ", 1)[0] + return datetime.strptime(first, "%Y-%m-%dT%H:%M:%SZ").replace(tzinfo=timezone.utc).timestamp() + except ValueError: + return None + + +def _from_islands(run_dir: Path, *, cards: int, head_usd_per_hour: float) -> Curve: + """A run2i.sh directory: one tape per island under ``tape-direct/``.""" + pts: list[Point] = [] + tapes = sorted((run_dir / "tape-direct").rglob("rl-island-*.jsonl")) + if not tapes: + raise ValueError(f"{run_dir}: no island tapes under tape-direct/") + for tape in tapes: + island = parse_yeto(_tape_events(tape.parent), mode="async") + pts += [Point(float(r["t_end_unix"]), int(r["samples"]), float(r["reward_mean"])) + for r in island.rounds + if r.get("t_end_unix") is not None and r.get("samples") and r.get("reward_mean") is not None] + pts.sort(key=lambda p: p.t_end) + t0 = _utc_file(run_dir / "start_utc.txt") + end = _utc_file(run_dir / "end_utc.txt") or (pts[-1].t_end if pts else t0) + hours = (end - t0) / 3600 if t0 and end else 0.0 + return Curve(run_dir.name, "yeto-miles-2islands", cards, float(t0), float(end), pts, + usd_extra=hours * head_usd_per_hour) + + +def load(run_dir: Path, *, cards: int = 4, head_usd_per_hour: float = 0.0) -> Curve: + run_dir = Path(run_dir) + if (run_dir / "parsed.json").is_file(): + return _from_parsed(run_dir) + return _from_islands(run_dir, cards=cards, head_usd_per_hour=head_usd_per_hour) + + +def smoothed(points: list[Point], window: int) -> list[float]: + """Sample-weighted mean reward over the latest ``window`` samples at every point.""" + out = [] + for i in range(len(points)): + n = s = 0.0 + for p in reversed(points[:i + 1]): + take = min(p.samples, window - n) + s += p.reward * take + n += take + if n >= window: + break + out.append(s / n if n else float("nan")) + return out + + +def start_reward(curves: list[Curve], samples: int) -> float: + """Pooled sample-weighted mean reward of every run's first ``samples`` samples.""" + n = s = 0.0 + for c in curves: + left = samples + for p in c.points: + take = min(p.samples, left) + s += p.reward * take + n += take + left -= take + if left <= 0: + break + return s / n if n else float("nan") + + +def metrics(c: Curve, *, target: float, window: int, usd_per_card_hour: float) -> dict[str, Any]: + sm = smoothed(c.points, window) + hit = next((i for i, v in enumerate(sm) if v >= target), None) + total = sum(p.samples for p in c.points) + run_s = c.t_last - c.t0 + out: dict[str, Any] = { + "run_id": c.run_id, "arm": c.arm, "cards": c.cards, "points": len(c.points), "samples": total, + "first_batch_s": c.points[0].t_end - c.t0 if c.points else None, + "train_s": c.points[-1].t_end - c.points[0].t_end if len(c.points) > 1 else None, + "samples_per_s": ((total - c.points[0].samples) / (c.points[-1].t_end - c.points[0].t_end) + if len(c.points) > 1 else None), + "reward_last3": sm[-1] if sm else None, + "run_s": run_s, "usd_run": run_s / 3600 * c.cards * usd_per_card_hour + c.usd_extra, + "target": target, "reached": hit is not None, + } + if hit is not None: + p = c.points[hit] + out.update(to_target_s=p.t_end - c.t0, to_target_after_first_s=p.t_end - c.points[0].t_end, + to_target_samples=sum(q.samples for q in c.points[:hit + 1])) + out["to_target_gpu_h"] = out["to_target_s"] / 3600 * c.cards + out["to_target_usd"] = out["to_target_gpu_h"] * usd_per_card_hour + return out + + +def curve_table(c: Curve, window: int) -> list[dict[str, Any]]: + sm, cum = smoothed(c.points, window), 0 + rows = [] + for p, v in zip(c.points, sm): + cum += p.samples + rows.append({**asdict(p), "t_from_launch_s": p.t_end - c.t0, "cum_samples": cum, "smoothed": v}) + return rows From fb520a6065dd8577545d59fa4d60a2bfcfe1608f Mon Sep 17 00:00:00 2001 From: michael <78113554+chowhsu@users.noreply.github.com> Date: Wed, 21 Oct 2026 13:23:47 +0000 Subject: [PATCH 57/62] Full-parameter no-sync single island and the E11 0.6B full-parameter preset A single --rl-single-island-no-sync island with both roles may now train every parameter (--tuning full); before, only sampler/trainer islands could. The efficiency spec gains tuning and tensor_parallel fields and the e11-qwen3-0.6b-full preset (trainer TP 2, so DP stays 1 as Miles full-parameter islands with engines require). Co-Authored-By: Claude Opus 5.5 --- tests/test_rl_sampler_trainer_pair_launch.py | 11 ++++- tools/eval/efficiency/spec.py | 49 ++++++++++++++++---- yeto/launch/rl_args.py | 3 +- yeto/rl/adapters/miles/island_entry.py | 25 +++++----- yeto/rl/engine/sample_roles.py | 10 ++-- 5 files changed, 68 insertions(+), 30 deletions(-) diff --git a/tests/test_rl_sampler_trainer_pair_launch.py b/tests/test_rl_sampler_trainer_pair_launch.py index 0cbd7e69..df7d661a 100644 --- a/tests/test_rl_sampler_trainer_pair_launch.py +++ b/tests/test_rl_sampler_trainer_pair_launch.py @@ -109,11 +109,20 @@ def test_pair_island_commands_pass_the_island_parser_and_no_sync_checks(): assert (config.role, config.store) == (role, "/root/yeto-checkpoint-store/run-a") -def test_full_parameter_without_a_buffer_role_is_still_refused_on_the_island(): +def test_full_parameter_on_a_no_sync_island_with_both_roles_is_accepted(): + """S26 E11: a single no-sync island with both roles trains every parameter.""" from yeto.rl.adapters.miles import island_entry args = types.SimpleNamespace(parameter_mode="full", rl_single_island_no_sync=True, rl_sample_role="both") + assert island_entry._full_parameter_mode(args, "strict-avg") is True + + +def test_full_parameter_with_a_syncer_still_needs_the_full_parameter_preset(): + from yeto.rl.adapters.miles import island_entry + + args = types.SimpleNamespace(parameter_mode="full", rl_single_island_no_sync=False, + rl_sample_role="both") with pytest.raises(ValueError, match="full-parameter sync preset"): island_entry._full_parameter_mode(args, "strict-avg") diff --git a/tools/eval/efficiency/spec.py b/tools/eval/efficiency/spec.py index 53173396..0df7b29b 100644 --- a/tools/eval/efficiency/spec.py +++ b/tools/eval/efficiency/spec.py @@ -89,6 +89,10 @@ class EffSpec: # cards per run: 0 keeps the E0 shape (sync 1, async 1 rollout + 1 trainer); E11 # (whole-cluster comparison) sets the total, split evenly between rollout and trainer in async cards: int = 0 + # "lora" or "full" (every parameter; E11 0.6B). Miles full-parameter islands with + # engines need DP=1, so a 2-card trainer runs tensor parallel 2 in every arm. + tuning: str = "lora" + tensor_parallel: int = 1 @property def samples_per_round(self) -> int: @@ -124,6 +128,16 @@ def to_dict(self) -> dict[str, Any]: data_file="dapo-math-17k.jsonl", train_rows=12 * 32, gpu="H100!", rounds=12, 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, cards=4), + # E11 full parameter: the E11 data, rounds and lengths on Qwen3-0.6B, every parameter + # trained (lora_r is unused). 2 rollout + 2 trainer cards, trainer TP 2 (DP 1). + "e11-qwen3-0.6b-full": EffSpec( + name="e11-qwen3-0.6b-full", model="Qwen/Qwen3-0.6B", model_revision="c1899de289a04d12100db370d81485cdf75e47ca", + miles_model_args="qwen3-0.6B.py", + data="zhuzilin/dapo-math-17k", data_revision="2e65612930298bde4c5d58fd97b3f23a483aaff9", + data_file="dapo-math-17k.jsonl", train_rows=12 * 32, + gpu="H100!", rounds=12, groups_per_round=32, samples_per_group=8, max_prompt_len=512, + max_response_len=1024, mem_fraction=0.6, timeout_s=3600, cards=4, tuning="full", + tensor_parallel=2, chat_template_kwargs={"enable_thinking": False}), } @@ -172,8 +186,10 @@ 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")} - for key in ("cards", "data_file", "train_rows"): # fields added after E0.1: a default value - if not doc[key]: # keeps the fingerprints of earlier runs valid + # fields added after E0.1: a default value keeps the fingerprints of earlier runs valid + for key, default in (("cards", 0), ("data_file", ""), ("train_rows", 0), ("tuning", "lora"), + ("tensor_parallel", 1)): + if doc[key] == default: doc.pop(key) return hashlib.sha256(json.dumps(doc, sort_keys=True).encode()).hexdigest() @@ -320,13 +336,10 @@ def miles_bare_argv(spec: EffSpec, mode: str) -> list[str]: # 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), + *_miles_lora_flags(spec), # 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", + "--tensor-model-parallel-size", str(spec.tensor_parallel), "--pipeline-model-parallel-size", "1", "--context-parallel-size", "1", "--expert-model-parallel-size", "1", "--expert-tensor-parallel-size", "1", "--recompute-granularity", "full", "--recompute-method", "uniform", "--recompute-num-layers", "1", "--use-dynamic-batch-size", "--max-tokens-per-gpu", str(max_tokens_per_gpu(spec)), @@ -345,13 +358,31 @@ def miles_bare_argv(spec: EffSpec, mode: str) -> list[str]: 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"] + argv += ["--colocate", "--offload-train"] + if spec.tuning == "lora": + argv += ["--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_lora_flags(spec: EffSpec) -> list[str]: + """Miles' own LoRA flags (the same set yeto's Miles adapter emits); none for full parameter.""" + if spec.tuning == "full": + return [] + return ["--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)] + + +def _yeto_tuning(spec: EffSpec) -> str: + """The yeto launcher's tuning options (full parameter: a no-sync single island).""" + if spec.tuning == "full": + return f"--tuning full --tensor-parallel {spec.tensor_parallel} " + return f"--tuning lora --lora-r {spec.lora_r} " + + def miles_bare_script(mode: str) -> str: return "train.py" if mode == "sync" else "train_async.py" @@ -422,7 +453,7 @@ def yeto_miles_launcher_args(spec: EffSpec, mode: str, run_id: str) -> str: # 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"--controller local --cluster-prefix {run_id} {_yeto_tuning(spec)}" f"--rl-observe-timeline --model {spec.model} --model-revision {spec.model_revision} {_yeto_data(spec)} " f"--rl-checkpoint-store modal-volume://{STORE_VOLUME}/{run_id} " f"--rl-cut-every {cut_every(spec)} " diff --git a/yeto/launch/rl_args.py b/yeto/launch/rl_args.py index e623e6dc..bdb4ce16 100644 --- a/yeto/launch/rl_args.py +++ b/yeto/launch/rl_args.py @@ -320,8 +320,7 @@ def _prepare_rl_args( from ..rl.engine import sample_roles if args.tuning != "lora" and not sample_roles.buffer_full_parameter(args): - raise ValueError("RL v0 requires --tuning lora (--tuning full: only sampler / trainer " - "islands, --rl-sample-role or --rl-island-sample-role, with " + raise ValueError("RL v0 requires --tuning lora (--tuning full: only with " "--rl-single-island-no-sync)") if args.lora_r <= 0: raise ValueError("RL v0 requires a positive LoRA rank") diff --git a/yeto/rl/adapters/miles/island_entry.py b/yeto/rl/adapters/miles/island_entry.py index cf98a201..dbafe785 100644 --- a/yeto/rl/adapters/miles/island_entry.py +++ b/yeto/rl/adapters/miles/island_entry.py @@ -75,11 +75,11 @@ def parse_args(argv=None): ) if (args.parameter_mode == "full" and args.sync_preset not in FULL_PARAMETER_SYNC_PRESETS - and not _buffer_island(args)): + and not _no_sync_island(args)): parser.error( "--parameter-mode full requires a full-parameter --sync-preset " - f"({', '.join(sorted(FULL_PARAMETER_SYNC_PRESETS))}) or a no-sync " - "--rl-sample-role sampler / trainer island" + f"({', '.join(sorted(FULL_PARAMETER_SYNC_PRESETS))}) or " + "--rl-single-island-no-sync" ) if args.sync_preset == FULL_PARAMETER_SYNC_PRESET and not _trains_base_parameters(args): parser.error( @@ -151,7 +151,7 @@ def build_ports_launch(args, run_config, extra_argv=()): ) return translate_run_config(run_config, base_algorithm, extra_argv=tuple(extra_argv), max_policy_age=int(getattr(args, "rl_max_policy_age", 0) or 0), - full_parameter_sync=not _buffer_island(args)) + full_parameter_sync=not _no_sync_island(args)) _ELASTIC_COMPANIONS = rl_option_pairs( @@ -696,12 +696,11 @@ def _full_parameter_identity(model_config_sha256: str) -> tuple[str, str]: ) -def _buffer_island(args) -> bool: - """A no-sync sampler / trainer island (L4.6): it has no outer sync, so full - parameter mode runs without the full-parameter sync preset.""" - from yeto.rl.engine.sample_roles import BOTH, sample_role - - return bool(getattr(args, "rl_single_island_no_sync", False)) and sample_role(args) != BOTH +def _no_sync_island(args) -> bool: + """A no-sync island (a sampler / trainer island, L4.6, or a single island + with both roles, S26 E11): it has no outer sync, so full parameter mode runs + without the full-parameter sync preset.""" + return bool(getattr(args, "rl_single_island_no_sync", False)) def _full_parameter_mode(args, sync_preset: str) -> bool: @@ -716,9 +715,9 @@ def _full_parameter_mode(args, sync_preset: str) -> bool: full_mode = getattr(args, "parameter_mode", "lora") == "full" full_preset = sync_preset in FULL_PARAMETER_SYNC_PRESETS - if full_mode and not full_preset and not _buffer_island(args): + if full_mode and not full_preset and not _no_sync_island(args): raise ValueError("full parameter mode needs a full-parameter sync preset " - "(or a no-sync --rl-sample-role sampler / trainer island)") + "(or --rl-single-island-no-sync)") if full_preset and not _trains_base_parameters(args): raise ValueError(f"the {sync_preset} sync needs full parameter mode or expert-full tuning") if full_preset and not full_mode and sync_preset != FULL_PARAMETER_SYNC_PRESET: @@ -1793,7 +1792,7 @@ def run_miles( # full-parameter strict sweep also carries DeepSeek V4 expert-full (P0 1.7b.2). full_parameter_sync = sync_preset in FULL_PARAMETER_SYNC_PRESETS dense_full = sync_preset == FULL_PARAMETER_SYNC_PRESET - if _trains_base_parameters(args) and not full_parameter_sync and not _buffer_island(args): + if _trains_base_parameters(args) and not full_parameter_sync and not _no_sync_island(args): raise ValueError(f"expert-full tuning needs the {FULL_PARAMETER_SYNC_PRESET} sync") sao_runtime = None if dense_full: diff --git a/yeto/rl/engine/sample_roles.py b/yeto/rl/engine/sample_roles.py index de65b83f..ca766c3e 100644 --- a/yeto/rl/engine/sample_roles.py +++ b/yeto/rl/engine/sample_roles.py @@ -240,12 +240,12 @@ def island_learner(args: Any, island: int, num_learners: int) -> tuple[int, int] def buffer_full_parameter(args: Any) -> bool: - """``--tuning full`` on sampler/trainer islands (single or pair, no-sync): the - island trains every parameter without the full-parameter sync, since a buffer - island has no syncer; the policy reaches the sampler as a store release.""" + """``--tuning full`` on a no-sync island: the island trains every parameter + without the full-parameter sync, since it has no syncer. A sampler/trainer + island (single or pair) passes the policy to the sampler as a store release; + a single island with both roles (S26 E11) updates its own engines in place.""" return (getattr(args, "tuning", "lora") == "full" - and bool(getattr(args, "rl_single_island_no_sync", False)) - and (is_pair(args) or sample_role(args) != BOTH)) + and bool(getattr(args, "rl_single_island_no_sync", False))) def check_pair(args: Any, num_islands: int) -> None: From 8783b2524d297863955d4e933dfcb8d8791a8767 Mon Sep 17 00:00:00 2001 From: michael <78113554+chowhsu@users.noreply.github.com> Date: Wed, 21 Oct 2026 13:26:39 +0000 Subject: [PATCH 58/62] eval efficiency: no-progress watchdog and the E11 full-parameter 2-island arm eff-run kills a run when no training step line arrives for 15 min (30 min before the first step), in the container for the bare and yeto+verl arms and in miles_local.sh for the yeto+Miles arm. yeto_miles_full_islands_launcher_args builds the 2-island full-parameter arm with the E3 C sync configuration. Co-Authored-By: Claude Opus 5.5 --- tests/test_tools_eval_efficiency.py | 52 +++++++++++++++++++++++++++ tools/eval/efficiency/modal_eff.py | 55 ++++++++++++++++++++++++----- tools/eval/efficiency/spec.py | 45 +++++++++++++++++++++++ tools/eval/miles_local.sh | 19 +++++++++- 4 files changed, 162 insertions(+), 9 deletions(-) diff --git a/tests/test_tools_eval_efficiency.py b/tests/test_tools_eval_efficiency.py index 6e8b841e..c4eeb1f2 100644 --- a/tests/test_tools_eval_efficiency.py +++ b/tests/test_tools_eval_efficiency.py @@ -483,3 +483,55 @@ def test_cluster_target_and_smoothing(): assert m["reached"] and m["to_target_s"] == 130.0 and m["to_target_samples"] == 4 * 128 assert m["to_target_gpu_h"] == pytest.approx(130 / 3600 * 4) assert not C.metrics(c, target=0.9, window=256, usd_per_card_hour=5.0)["reached"] + + +# ---------------------------------------------------------------- E11 full parameter +def test_e11_full_preset_runs_every_parameter_with_trainer_tp2(): + sp = spec.get_spec("e11-qwen3-0.6b-full") + argv = spec.miles_bare_argv(sp, "async") + assert not any(a.startswith(("--lora", "--sglang-max-lora", "--target-modules")) for a in argv) + assert argv[argv.index("--tensor-model-parallel-size") + 1] == "2" + assert argv[argv.index("--rollout-num-gpus") + 1] == "2" + args = spec.yeto_miles_launcher_args(sp, "async", "r") + assert "--tuning full --tensor-parallel 2 " in args and "--lora-r" not in args + assert sp.samples_per_round == 256 and sp.rounds * sp.groups_per_round == sp.train_rows + + +def test_tuning_fields_keep_earlier_fingerprints(): + sp = spec.get_spec("e11-qwen3-4b") + assert spec.control_fingerprint(sp) == spec.control_fingerprint( + spec.replace(sp, tuning="lora", tensor_parallel=1)) + assert spec.control_fingerprint(sp) != spec.control_fingerprint(spec.replace(sp, tuning="full")) + + +def test_full_islands_args_use_the_e3_c_sync_and_split_each_round(): + sp = spec.get_spec("e11-qwen3-0.6b-full") + args = spec.yeto_miles_full_islands_launcher_args(sp, 2, "r", data="t.jsonl", island_tasks="s.json") + for flag in ("--rl-sync-preset full-parameter", "--fragments 2", "--rl-sync-fragment-pipeline 2", + "--rl-sync-async-checkpoint", "--rl-sync-push-dtype bf16", "--rl-sync-checkpoint-every 8", + "--wan-streams 8", "--rollout-batch-size 16", "--gpu modal:2xh100,modal:2xh100", + "--rl-max-policy-age 0", "--total-steps 12"): + assert flag in args + with pytest.raises(ValueError): + spec.yeto_miles_full_islands_launcher_args(spec.get_spec("e11-qwen3-4b"), 2, "r", data="t", island_tasks="s") + + +def test_run_logged_kills_a_run_without_training_steps(tmp_path, monkeypatch): + from tools.eval.efficiency import modal_eff + + monkeypatch.setattr(modal_eff, "WATCH_POLL_S", 0.2) + monkeypatch.setattr(modal_eff, "FIRST_STEP_S", 1) + monkeypatch.setattr(modal_eff, "STALL_S", 1) + log = tmp_path / "main.log" + script = "echo \"step 1: {'train/step': 0}\"; sleep 0.5; echo \"{'train/step': 1}\"; sleep 30" + rc = modal_eff._run_logged(["bash", "-c", script], log, 60, watch_progress=True) + assert rc != 0 + assert "killed: no training step" in log.read_text() + + +def test_run_logged_without_the_watch_waits_for_the_process(tmp_path): + from tools.eval.efficiency import modal_eff + + log = tmp_path / "main.log" + assert modal_eff._run_logged(["bash", "-c", "sleep 1; echo done"], log, 60) == 0 + assert "killed" not in log.read_text() diff --git a/tools/eval/efficiency/modal_eff.py b/tools/eval/efficiency/modal_eff.py index 3aa0f21c..5e7d0478 100644 --- a/tools/eval/efficiency/modal_eff.py +++ b/tools/eval/efficiency/modal_eff.py @@ -21,8 +21,10 @@ import json import os +import re import subprocess import sys +import threading import time from pathlib import Path from typing import Any @@ -37,24 +39,60 @@ NVML_INTERVAL_MS = 1000 +# No-progress watchdog (E11: a hung bare verl run burned 52 min): a training step +# line is progress; the run is killed when none arrives for STALL_S, or before the +# first step for FIRST_STEP_S (cold start: model load, engine start, first rollout). +# Miles logs "'train/step': N", verl "step:N -". +PROGRESS = re.compile(r"'train/step'|\bstep:\d+ ") +STALL_S = 900 +FIRST_STEP_S = 1800 +WATCH_POLL_S = 10 + + # ====================================================================== 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: +def _run_logged(argv: list[str], log: Path, timeout: float, *, watch_progress: bool = False, + **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 + (the step end times the parsers read). The process is killed at ``timeout`` and, + with ``watch_progress``, when no step line arrives in time (``PROGRESS``).""" + started = time.time() + deadline = started + timeout + last_step: list[float] = [] # time of the latest step line (empty: none yet) + lock = threading.Lock() with log.open("a") as fh, subprocess.Popen(argv, stdout=subprocess.PIPE, stderr=subprocess.STDOUT, text=True, errors="replace", **kwargs) as proc: + def note(text: str) -> None: + with lock: + fh.write(f"{time.time():.3f} [eff harness] {text}\n") + fh.flush() + + def watch() -> None: + while proc.poll() is None: + time.sleep(WATCH_POLL_S) + quiet = time.time() - (last_step[-1] if last_step else started) + limit = STALL_S if last_step else FIRST_STEP_S + if quiet > limit and proc.poll() is None: + note(f"killed: no training step for {quiet:.0f} s (limit {limit} s)") + proc.kill() + return + + if watch_progress: + threading.Thread(target=watch, daemon=True).start() for line in proc.stdout: - fh.write(f"{time.time():.3f} {line}") - fh.flush() + if watch_progress and PROGRESS.search(line): + last_step.append(time.time()) + del last_step[:-1] + with lock: + 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") + note(f"killed at the {timeout:.0f} s timeout") break return proc.wait() @@ -79,7 +117,7 @@ def arm(run: str, plan: dict) -> dict: (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) + meta["rc"] = _run_logged(["bash", str(out / "script.sh")], log, budget, watch_progress=True) else: sys.path.insert(0, mv.WORKDIR) env = {**mv._island_env({"gpu": plan["gpu"]}, 0), **plan.get("env", {})} @@ -88,7 +126,7 @@ def arm(run: str, plan: dict) -> dict: 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") + meta["rc"] = _run_logged(meta["argv"], log, budget, watch_progress=True, env=env, cwd="/root") except BaseException as exc: # noqa: BLE001 - recorded, then re-raised meta["error"] = f"{type(exc).__name__}: {exc}" raise @@ -172,6 +210,7 @@ def run_miles_launcher(plan: dict, run_dir: Path, *, dry_run: bool = False) -> d 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, + STALL_S=str(STALL_S), FIRST_STEP_S=str(FIRST_STEP_S), 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: diff --git a/tools/eval/efficiency/spec.py b/tools/eval/efficiency/spec.py index 0df7b29b..6d1a932b 100644 --- a/tools/eval/efficiency/spec.py +++ b/tools/eval/efficiency/spec.py @@ -549,6 +549,51 @@ def yeto_miles_islands_launcher_args(spec: EffSpec, islands: int, run_id: str, * return args +# E3 "C" sync configuration (s1-runs/s26-e3/run.sh GROUP=C): 2 fragments, fragment pipeline, +# background checkpoint every 4 sweeps, bf16 upload, 8 adaptive WAN streams (E3d X4a/X4b) +FULL_PARAMETER_FRAGMENTS = 2 +QWEN3_06B_LAYOUT_SHA256 = "3d5a0936038d8c9587afffc4bc043769d094695c8916ebef8696c0758d83a3b5" + + +def yeto_miles_full_islands_launcher_args(spec: EffSpec, islands: int, run_id: str, *, data: str, + island_tasks: str) -> str: + """E11 full parameter: ``islands`` islands of 1 rollout + 1 trainer card that sync every + parameter through the head syncer with the E3 C configuration. The islands together + read the single-island arm's prompts of each round (``island_task_split``) and run its + Miles recipe settings; the full-parameter sweep needs policy age 0 (E3).""" + if spec.tuning != "full" or spec.model != "Qwen/Qwen3-0.6B": + raise ValueError("the full-parameter island arm is pinned to Qwen3-0.6B (layout hash)") + c = cards("async", spec.cards) + if c["total"] != 2 * islands or spec.groups_per_round % islands: + raise ValueError(f"{islands} islands of 2 cards need {2 * islands} cards and an equal prompt share") + card = {"A10G": "a10g", "H100!": "h100", "L4": "l4"}[spec.gpu] + frags = FULL_PARAMETER_FRAGMENTS + kwargs = json.dumps(spec.chat_template_kwargs) if spec.chat_template_kwargs else "" + args = (f"launch --controller head --syncer-region {HEAD_REGION} --on-demand --training-mode rl " + f"--rl-engine ports --preflight-threads wait --preflight-thread-start 2950 --modal-gpu-exact " + f"{gpu_name_arg(spec)}--modal-retries 0 --no-island-relaunch --modal-cpu 8 --modal-memory-gib 64 " + f"--cluster-prefix {run_id} --rl-observe-timeline " + f"--model {spec.model} --model-revision {spec.model_revision} " + f"--rl-allow-local-data --data {data} --rl-island-tasks {shlex.quote(island_tasks)} " + "--reward-function gsm8k_reward:score --pipeline 1 " + f"--rollout-batch-size {spec.groups_per_round // islands} " + f"--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 {','.join([f'modal:2x{card}'] * islands)} --total-steps {spec.rounds} " + "--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 " + f"--trust-remote-code --rl-algorithm-spec {MILES_ALGORITHM_SPEC} " + f"--rl-sync-preset full-parameter --rl-parameter-layout-sha256 {QWEN3_06B_LAYOUT_SHA256} " + f"--fragments {frags} --rl-sync-fragment-pipeline {frags} --rl-sync-async-checkpoint " + f"--rl-sync-push-dtype bf16 --rl-sync-checkpoint-every {4 * frags} " + "--wan-streams 8 --wan-stream-schedule adaptive --wan-stall-timeout 90 " + "--rl-placement fixed-partition --rl-rollout-gpus 1 --rl-max-policy-age 0") + if kwargs: + args += f" --apply-chat-template-kwargs {shlex.quote(kwargs)}" + return args + + def island_task_split(groups_per_round: int, rounds: int, islands: int) -> dict[str, list[str]]: """Round-major split of the first ``rounds * groups_per_round`` rows (task ids ``gsm``): island k gets the k-th equal slice of every round's prompts, so the diff --git a/tools/eval/miles_local.sh b/tools/eval/miles_local.sh index 0c634b04..34d5458f 100644 --- a/tools/eval/miles_local.sh +++ b/tools/eval/miles_local.sh @@ -10,7 +10,9 @@ # ISLANDS (containers before the GPU lock is released), HARD (s), STOREVOL, TAPEVOL, # RECLAIM_LABEL (e.g. "l0") + RECLAIM_ROUND (rollout id written before the stop), # PY (python with modal + the launcher deps), DRY (non-empty: --dry-run only), -# EXTRA_FILES (optional: files copied into the checkout root and committed) +# EXTRA_FILES (optional: files copied into the checkout root and committed), +# STALL_S / FIRST_STEP_S (optional: stop the apps when no Miles training step line +# reaches the launch log for STALL_S s, or before the first one for FIRST_STEP_S s) set -u : "${P:?}" "${R:?}" "${REPO:?}" "${ARGS:?}" "${STOREVOL:?}" "${PY:?}" HARD=${HARD:-3000}; ISLANDS=${ISLANDS:-1}; TAPEVOL=${TAPEVOL:-yeto-event-tapes} @@ -71,6 +73,21 @@ for i in $(seq 1 120); do kill -0 $SUB 2>/dev/null || break; sleep 10 done exec 9>&- +# no-progress watchdog (tools/eval/efficiency/modal_eff.py STALL_S): stop the apps when the +# count of Miles training step lines stops growing +if [ -n "${STALL_S:-}" ]; then +( cd /tmp; last=0; since=$(date +%s); limit=${FIRST_STEP_S:-1800} + while kill -0 $SUB 2>/dev/null; do + sleep 15 + n=$(grep -ac "'train/step'" $R/launch.ts.log 2>/dev/null) + if [ "${n:-0}" -gt "$last" ]; then last=$n; since=$(date +%s); limit=$STALL_S; fi + if [ $(( $(date +%s) - since )) -gt "$limit" ]; then + echo "$(date -u +%FT%TZ) no training step for $(( $(date +%s) - since )) s (limit $limit s, steps $last): stopping apps" > $R/STALL_FIRED + for a in $APP $SAPP; do timeout 120 $MODAL app stop -y $a; done >> $R/STALL_FIRED 2>&1 + exit 0 + fi + done ) 9>&- & +fi # one reclaim: island $RECLAIM_LABEL wrote rollout $RECLAIM_ROUND -> stop its container if [ -n "${RECLAIM_LABEL:-}" ]; then ( cd /tmp From 189df102ae484170178c51caa896e92f363ea5e3 Mon Sep 17 00:00:00 2001 From: michael <78113554+chowhsu@users.noreply.github.com> Date: Wed, 21 Oct 2026 15:42:10 +0000 Subject: [PATCH 59/62] E5: time the policy release phases; optional fused LoRA stage Trainer side: PolicyReleaseStore.release records the seconds of each phase (read LATEST, export, digest, weight commit, LATEST commit, prune); ReleasingPublisher adds the publish to its own engines, released_at and the release size to rl_policy_released. Sampler side: PolicyFollower records its phases and when it first read each release (seen_at); EngineReleaseLoader records stage, lock, engine load, stamp, unlock, warm-up, checksum and probe seconds; rl_sampler_policy_loaded carries them with released_at. rl_sampler_round carries generate_s. --rl-sampler-fused-stage (default off, LoRA sampler only): the staging copy of the adapter computes the release digest, so a release is read from the store once instead of twice. The LoRA copy now happens before the update lock in both modes; engine calls and their order are unchanged. tools/eval/sampler_islands.py: E4 scale report (rounds, rates, preemption recovery, speedup vs a baseline run) and E5 latency report (per release and sampler island) from island tapes. Co-Authored-By: Claude Opus 5.5 --- tests/test_rl_release_latency.py | 261 ++++++++++++++++++++ tests/test_tools_sampler_islands.py | 116 +++++++++ tools/eval/sampler_islands.py | 288 +++++++++++++++++++++++ yeto/rl/adapters/miles/config.py | 3 +- yeto/rl/adapters/miles/sampler_island.py | 152 ++++++++---- yeto/rl/engine/policy_release.py | 161 ++++++++++--- yeto/rl/engine/run_flags.py | 4 + yeto/rl/engine/sample_roles.py | 18 +- yeto/rl/engine/sampler_loop.py | 12 +- 9 files changed, 937 insertions(+), 78 deletions(-) create mode 100644 tests/test_rl_release_latency.py create mode 100644 tests/test_tools_sampler_islands.py create mode 100644 tools/eval/sampler_islands.py diff --git a/tests/test_rl_release_latency.py b/tests/test_rl_release_latency.py new file mode 100644 index 00000000..f0ca9425 --- /dev/null +++ b/tests/test_rl_release_latency.py @@ -0,0 +1,261 @@ +"""S26 E5 (CPU, no Ray, no Miles): the phase timings of a policy release (trainer +side and sampler side) and the fused LoRA stage (``--rl-sampler-fused-stage``).""" + +from __future__ import annotations + +import argparse +import itertools +from pathlib import Path +from types import SimpleNamespace + +import pytest + +from tests.test_rl_sampler_island import ( + _Engine, + _launch_args, + _lora_loader, + _State, + _write_adapter, + _write_weights, +) +from yeto.rl.adapters.miles import sampler_island as si +from yeto.rl.engine import sample_roles +from yeto.rl.engine.policy_release import ( + LORA, + RECORD, + PhaseTimer, + PolicyFollower, + PolicyReleaseError, + PolicyReleaseStore, + copy_with_digest, + directory_digest, +) +from yeto.rl.engine.resume import CheckpointStore +from yeto.rl.engine.sample_buffer import FileSampleBuffer +from yeto.rl.engine.sampler_loop import SamplerLoop + + +def _ticks(step: float = 0.5): + """A clock that advances ``step`` seconds per read.""" + counter = itertools.count() + return lambda: next(counter) * step + + +def test_phase_timer_adds_each_phase_in_first_use_order(): + timer = PhaseTimer(_ticks()) + with timer.phase("b"): + pass + with timer.phase("a"): + pass + with timer.phase("b"): + pass + timer.add("c", 0.25) + assert timer.as_dict() == {"b": 1.0, "a": 0.5, "c": 0.25} + with pytest.raises(RuntimeError), timer.phase("d"): + raise RuntimeError("boom") + assert timer.as_dict()["d"] == 0.5 # a failed phase still counts + + +def test_copy_with_digest_copies_everything_and_equals_the_directory_digest(tmp_path): + source = tmp_path / "src" + (source / "sub").mkdir(parents=True) + (source / "adapter_model.safetensors").write_bytes(b"\x00\x01" * 70_000) # > one chunk + (source / "sub" / "extra.json").write_text("{}") + (source / RECORD).write_text('{"seq": 1}') + target = tmp_path / "dst" + digest = copy_with_digest(source, target) + assert digest == directory_digest(source) == directory_digest(target) + copied = sorted(p.relative_to(target).as_posix() for p in target.rglob("*") if p.is_file()) + assert copied == [RECORD, "adapter_model.safetensors", "sub/extra.json"] + assert (target / "adapter_model.safetensors").read_bytes() == ( + source / "adapter_model.safetensors").read_bytes() + (source / RECORD).write_text('{"seq": 2}') # the record is not digested + assert copy_with_digest(source, tmp_path / "dst2") == digest + + +def test_release_records_the_trainer_phases(tmp_path): + releases = PolicyReleaseStore(CheckpointStore(tmp_path), clock=lambda: 7.0, timer=_ticks()) + assert releases.last_timings == {} + lora = releases.release(1, "h1", _write_adapter, weights=LORA) + assert lora.released_at == 7.0 + assert list(releases.last_timings) == ["read_latest_s", "export_s", "digest_s", + "commit_weights_s", "commit_latest_s", "prune_s"] + assert all(value == 0.5 for value in releases.last_timings.values()) + before = dict(releases.last_timings) + assert releases.release(1, "h1", _write_adapter, weights=LORA) is None # republish + assert releases.last_timings == before # only a written release sets the timings + + +def test_follower_records_its_phases_and_when_it_first_saw_a_release(tmp_path): + releases = PolicyReleaseStore(CheckpointStore(tmp_path)) + wall = iter([100.0, 200.0, 300.0]) + loaded = [] + follower = PolicyFollower(releases, lambda r, d: loaded.append(r.seq), timer=_ticks(), + wall=lambda: next(wall)) + assert follower.refresh() is None and follower.seen_at is None + releases.release(1, "h1", _write_adapter, weights=LORA) + assert follower.refresh().seq == 1 and follower.seen_at == 100.0 + assert list(follower.last_timings) == ["read_latest_s", "resolve_s", "verify_s", "load_s"] + assert follower.refresh() is None and follower.seen_at == 100.0 # same release: same time + releases.release(2, "h2", _write_adapter, weights=LORA) + assert follower.refresh().seq == 2 and follower.seen_at == 200.0 and loaded == [1, 2] + + +def test_follower_leaves_the_file_check_to_a_hook_that_verifies_files(tmp_path): + releases = PolicyReleaseStore(CheckpointStore(tmp_path)) + release = releases.release(1, "h1", _write_adapter, weights=LORA) + weights = releases.resolve(release) / "adapter_model.safetensors" + weights.write_bytes(b"x" * weights.stat().st_size) # same size, other bytes + with pytest.raises(PolicyReleaseError, match="weight files sha256"): + PolicyFollower(releases, lambda r, d: None).refresh() + + class Hook: + verifies_files = True + + def __init__(self): + self.calls = [] + + def __call__(self, release, directory): + self.calls.append(release.seq) + + hook = Hook() + follower = PolicyFollower(releases, hook) + assert follower.refresh().seq == 1 and hook.calls == [1] + assert "verify_s" not in follower.last_timings + + +def _fused(log, staging, fused=True): + body = {"model.w": "base", "lora:miles_lora:q_proj:0:A": "a1"} + loader = _lora_loader(log, [_Engine(log, "e0"), _Engine(log, "e1")], [body, dict(body)], + staging=str(staging)) + loader.fused_stage = fused + return loader + + +def test_fused_stage_loads_like_the_plain_copy_and_records_the_load_phases(tmp_path): + staging = tmp_path / "staging" + staging.mkdir() + releases = PolicyReleaseStore(CheckpointStore(tmp_path / "store")) + plain_log, fused_log = [], [] + plain, fused = _fused(plain_log, staging, fused=False), _fused(fused_log, staging) + assert (plain.verifies_files, fused.verifies_files) == (False, True) + release = releases.release(0, "h0", _write_adapter, weights=LORA) + PolicyFollower(releases, fused).refresh() + (path,) = [str(p) for p in staging.iterdir()] + assert (Path(path) / "adapter_model.safetensors").read_text() == "a" + plain(release, releases.resolve(release)) + plain_path = str(plain._staged) + # the same engine calls in the same order; only the local copy's name differs + assert [step[:2] for step in fused_log] == [step[:2] for step in plain_log] + assert ("load_lora_adapter", "e0", path) in fused_log + assert ("load_lora_adapter", "e0", plain_path) in plain_log + assert set(fused.last_timings) == {"stage_s", "lock_s", "engine_load_s", "stamp_s", + "unlock_s", "warmup_s", "checksum_s", "probe_s"} + assert fused.last_checksum == plain.last_checksum is not None + + +def test_fused_stage_refuses_changed_bytes_before_any_engine_call(tmp_path): + staging = tmp_path / "staging" + staging.mkdir() + releases = PolicyReleaseStore(CheckpointStore(tmp_path / "store")) + release = releases.release(0, "h0", _write_adapter, weights=LORA) + weights = releases.resolve(release) / "adapter_model.safetensors" + weights.write_bytes(b"x" * weights.stat().st_size) + log = [] + loader = _fused(log, staging) + follower = PolicyFollower(releases, loader) + with pytest.raises(PolicyReleaseError, match="weight files sha256"): + follower.refresh() + assert log == [] and list(staging.iterdir()) == [] # no lock, no engine, no copy left + assert follower.release is None and "stage_s" in loader.last_timings + + +def test_full_parameter_loader_ignores_the_fused_stage(tmp_path): + loader = si.EngineReleaseLoader(None, None, None, fused_stage=True) + assert (loader.fused_stage, loader.verifies_files) == (False, False) + + +def test_a_failed_engine_load_removes_the_unused_copy(tmp_path): + staging = tmp_path / "staging" + staging.mkdir() + releases = PolicyReleaseStore(CheckpointStore(tmp_path / "store")) + release = releases.release(0, "h0", _write_adapter, weights=LORA) + log = [] + loader = _fused(log, staging) + + async def failing_post(engine, endpoint, payload): + raise RuntimeError("engine down") + + loader.post = failing_post + with pytest.raises(RuntimeError, match="engine down"): + loader(release, releases.resolve(release)) + assert list(staging.iterdir()) == [] and log[-1] == ("abort",) + + +def test_sampler_loaded_event_carries_the_latency_fields(tmp_path): + releases = PolicyReleaseStore(CheckpointStore(tmp_path), clock=lambda: 50.0) + follower = PolicyFollower(releases, lambda r, d: None, wall=lambda: 60.0) + events = [] + loop = SamplerLoop(follower=follower, buffer=FileSampleBuffer(CheckpointStore(tmp_path), batch_size=2), + generate=lambda r: None, num_rollouts=1, max_backlog_groups=4, + emit=lambda e, **f: events.append((e, f)), + loaded_fields=lambda: {"load_timings": {"stage_s": 0.1}}) + releases.release(1, "h1", _write_weights) + loop.run() + (loaded,) = [f for e, f in events if e == "rl_sampler_policy_loaded"] + assert (loaded["released_at"], loaded["seen_at"]) == (50.0, 60.0) + assert set(loaded["follower_timings"]) == {"read_latest_s", "resolve_s", "verify_s", "load_s"} + assert loaded["load_timings"] == {"stage_s": 0.1} + + +def test_releasing_publisher_reports_the_publish_and_release_phases(tmp_path): + events = [] + publisher = si.ReleasingPublisher( + SimpleNamespace(publish=lambda state: None), PolicyReleaseStore(CheckpointStore(tmp_path)), + export=lambda v, d: _write_adapter(d), read_position=lambda: 0, + emit=lambda e, **f: events.append((e, f)), weights=LORA) + publisher.publish(_State(0, "h0")) + (released,) = [f for e, f in events if e == "rl_policy_released"] + assert released["publish_s"] >= 0 and released["released_at"] > 0 + assert "export_s" in released["release_timings"] + assert released["release_bytes"] == len('{"r": 8}') + len("a") + + +# ---------------------------------------------------------------- the launch flag +def test_fused_stage_flag_is_sampler_and_lora_only(): + sample_roles.check_launch(_launch_args(rl_sampler_fused_stage=True)) + for kw, message in ( + ({"rl_sample_role": "trainer", "rl_max_policy_age": 1}, "needs --rl-sample-role sampler"), + ({"tuning": "full"}, "--tuning lora")): + with pytest.raises(ValueError, match=message): + sample_roles.check_launch(_launch_args(rl_sampler_fused_stage=True, **kw)) + with pytest.raises(ValueError, match="needs --rl-sample-role sampler"): + sample_roles.check_launch(SimpleNamespace(rl_sampler_fused_stage=True)) + + +def test_fused_stage_flag_is_forwarded_to_sampler_islands_only(): + args = _launch_args(rl_sampler_fused_stage=True) + assert sample_roles.island_flags(args).endswith(" --rl-sampler-fused-stage") + assert "--rl-sampler-fused-stage" not in sample_roles.island_flags(_launch_args()) + island = SimpleNamespace(rl_sample_role="sampler", rl_resume_store="/mnt/store", + rl_sampler_fused_stage=True) + assert sample_roles.resolve(island, groups_per_round=8).fused_stage is True + assert sample_roles.resolve(SimpleNamespace(rl_sample_role="sampler", rl_resume_store="/s"), + groups_per_round=8).fused_stage is False + pair = SimpleNamespace(rl_island_sample_role=["0:sampler", "1:trainer"], + rl_sampler_fused_stage=True, rl_sampler_min_round_seconds=None, + rl_sample_wait_s=None, rl_sampler_max_backlog_rounds=None) + assert sample_roles.island_args(pair, 0, 2).rl_sampler_fused_stage is True + assert sample_roles.island_args(pair, 1, 2).rl_sampler_fused_stage is False + + +def test_fused_stage_flag_exists_on_both_layers(): + from yeto.rl.engine.run_flags import CLI, ISLAND, flags_for + + for layer in (CLI, ISLAND): + parser = argparse.ArgumentParser() + for flag in flags_for(layer): + if flag.options[0] == "--rl-sampler-fused-stage": + parser.add_argument(*flag.options, **dict(flag.kwargs)) + assert parser.parse_args([]).rl_sampler_fused_stage is False + assert parser.parse_args(["--rl-sampler-fused-stage"]).rl_sampler_fused_stage is True diff --git a/tests/test_tools_sampler_islands.py b/tests/test_tools_sampler_islands.py new file mode 100644 index 00000000..3bbdffbf --- /dev/null +++ b/tests/test_tools_sampler_islands.py @@ -0,0 +1,116 @@ +"""S26 E4/E5 measurement tool (tools/eval/sampler_islands.py) on synthetic tapes.""" + +from __future__ import annotations + +import json +from pathlib import Path + +from tools.eval import sampler_islands as tool + + +def _tape(run: Path, island: int, name: str, rows: list[dict]) -> None: + path = run / "tape-direct" / "app" / f"l{island}" / "rank0" / name + path.parent.mkdir(parents=True, exist_ok=True) + lines = [json.dumps(r) for r in rows] + path.write_text("\n".join(lines) + '\n{"event": "rl_cut', encoding="utf-8") # a cut last line + + +def _sampler_round(t, rollout_id, generate_s=10.0, groups=32): + return {"event": "rl_sampler_round", "time_unix": t, "rollout_id": rollout_id, + "groups": groups, "generate_s": generate_s} + + +def _trained(t, groups=32): + return {"event": "rl_round_trained", "time_unix": t, "trained_groups": groups} + + +def _run(tmp_path: Path, name: str, *, samplers: int, preempt: bool = False) -> Path: + run = tmp_path / name + for m in range(samplers): + rows = [{"event": "rl_sampler_join", "time_unix": 90.0, "incarnation": f"a{m}"}, + _sampler_round(110.0, 0), _sampler_round(120.0, 1), + {"event": "rl_sampler_admitted", "time_unix": 125.0, "waited_s": 5.0}] + if preempt and m == 0: + rows.append({"event": "rl_sampler_leave", "time_unix": 121.0, "reason": "PREEMPTED", + "incarnation": "a0", "in_flight_rollout_id": 2}) + _tape(run, m, "rl-island-0.jsonl", rows) + rows = [{"event": "rl_sampler_join", "time_unix": 181.0, "incarnation": "b0"}, + _sampler_round(201.0, 2)] + _tape(run, m, "rl-island-0.inc1.jsonl", rows) + else: + _tape(run, m, "rl-island-0.jsonl", rows) + step = 20.0 / samplers if not preempt else 40.0 + _tape(run, samplers, "rl-island-0.jsonl", + [_trained(130.0), _trained(130.0 + step), _trained(130.0 + 2 * step)]) + return run + + +def test_scale_report_counts_rounds_rates_and_preemptions(tmp_path): + report = tool.scale_report(_run(tmp_path, "n2k", samplers=2, preempt=True), samples_per_group=4) + assert report["sampler_islands"] == 2 and report["preemptions"] == 1 + s0 = report["samplers"][0] + assert s0["incarnations"] == ["a0", "b0"] and s0["rounds"] == 3 and s0["groups"] == 96 + assert s0["generate_s"] == 30.0 and s0["backlog_wait_s"] == 5.0 + (p,) = s0["preemptions"] + assert (p["rejoin_s"], p["first_new_round_s"], p["in_flight_rollout_id"]) == (60.0, 80.0, 2) + # production: 160 groups from the first round start (100) to the last end (201) + assert report["production_groups_per_s"] == round(160 / 101, 4) + assert report["trainer"]["steady_groups_per_s"] == round(64 / 80, 4) + assert report["consumption_samples_per_s"] == round(64 / 80 * 4, 4) + + +def test_compare_gives_speedup_for_more_islands_and_a_ratio_for_the_same_size(tmp_path): + one = tool.scale_report(_run(tmp_path, "n1", samplers=1)) + two = tool.scale_report(_run(tmp_path, "n2", samplers=2)) + killed = tool.scale_report(_run(tmp_path, "n2k", samplers=2, preempt=True)) + up = tool.compare(two, one) + assert up["speedup"] == 2.0 and up["efficiency"] == 1.0 + same = tool.compare(killed, two) + assert same["effective_ratio"] == round((64 / 80) / (64 / 20), 4) and "speedup" not in same + + +def test_missing_generate_seconds_give_no_production_rate(tmp_path): + run = tmp_path / "old" + _tape(run, 0, "t.jsonl", [{"event": "rl_sampler_round", "time_unix": 5.0, "groups": 8}]) + report = tool.scale_report(run) + assert report["samplers"][0]["generate_s"] is None and report["production_groups_per_s"] is None + assert report["trainer"] is None + + +def test_latency_report_joins_releases_and_loads_by_seq(tmp_path): + run = tmp_path / "lat" + released = [{"event": "rl_policy_released", "time_unix": 100.0 * seq, "seq": seq, + "publish_s": 1.0, "release_bytes": 10, + "release_timings": {"export_s": 2.0, "commit_weights_s": 0.5}} for seq in (1, 2, 3)] + _tape(run, 1, "t.jsonl", [_trained(1.0), _trained(2.0)] + released) + loads = [{"event": "rl_sampler_policy_loaded", "time_unix": 100.0 * seq + 9.0, "seq": seq, + "weights": "lora", "seen_at": 100.0 * seq + 4.0, "fused_stage": True, + "follower_timings": {"read_latest_s": 0.1, "load_s": 4.9}, + "load_timings": {"stage_s": 1.0}} for seq in (1, 2, 3)] + loads.append({"event": "rl_sampler_policy_loaded", "time_unix": 450.0, "seq": 4}) # no release + _tape(run, 0, "t.jsonl", loads + [_sampler_round(5.0, 0)]) + report = tool.latency_report(run) + assert report["loads"] == 3 + row = report["rows"][1] + assert (row["trainer_release_s"], row["wait_s"], row["sampler_load_s"]) == (2.5, 4.0, 5.0) + assert (row["released_to_loaded_s"], row["end_to_end_s"]) == (9.0, 12.5) + assert report["summary_after_first"]["end_to_end_s"]["n"] == 2 + assert report["phases"]["sampler.load_s"]["median"] == 4.9 + assert report["phases"]["trainer.export_s"]["n"] == 3 + + +def test_latency_report_of_old_tapes_keeps_unknown_phases_none(tmp_path): + run = tmp_path / "old" + _tape(run, 1, "t.jsonl", [_trained(1.0), {"event": "rl_policy_released", "time_unix": 10.0, "seq": 2}]) + _tape(run, 0, "t.jsonl", [{"event": "rl_sampler_policy_loaded", "time_unix": 13.5, "seq": 2}, + _sampler_round(5.0, 0)]) + (row,) = tool.latency_report(run)["rows"] + assert row["released_to_loaded_s"] == 3.5 + assert row["wait_s"] is None and row["end_to_end_s"] is None + + +def test_cli_prints_json(tmp_path, capsys): + run = _run(tmp_path, "n1", samplers=1) + assert tool.main(["scale", str(run), "--baseline", str(run)]) == 0 + out = json.loads(capsys.readouterr().out) + assert out["compare"]["effective_ratio"] == 1.0 diff --git a/tools/eval/sampler_islands.py b/tools/eval/sampler_islands.py new file mode 100644 index 00000000..5ab79a8f --- /dev/null +++ b/tools/eval/sampler_islands.py @@ -0,0 +1,288 @@ +"""Sampler-island measurements from island event tapes (S26 E4, E5). + +Input: one run directory per call. The tapes are read from +``/tape-direct//l/rank0/*.jsonl`` (one file per incarnation, the +layout ``modal volume get`` of the tape volume gives), else from +``/tapes/l/*.jsonl``. ``l`` is the island; its role comes from its +events (``rl_sampler_*``: sampler; ``rl_round_trained``: trainer). + +Two reports (JSON on stdout): + +* ``scale`` (E4): per sampler island the rounds, groups, generation seconds, + waits and preemptions (LEAVE PREEMPTED to the next JOIN of the same sampler and + to its first new round); for the run the sampler production rate and the + trainer's consumption rate (groups per second). ``--baseline RUN`` adds the + speedup and efficiency against a baseline run (``N`` sampler islands vs the + baseline's), or the effective-throughput ratio of a preempted run vs an + unpreempted run of the same size. +* ``latency`` (E5): per release and sampler island the phases of one policy + update: the trainer's publish to its own engines and its release phases + (export, digest, commits), the wait until the sampler first read the release, + and the sampler's load phases. ``wait_s`` compares two hosts' wall clocks + (released on the trainer, seen on the sampler); it includes the clock offset. + +Rates use the trainer's rounds after its first one (the first round includes +the start-up). Every number is computed from the tapes; a missing field is +reported as None, never estimated. + +CLI:: + + python -m tools.eval.sampler_islands scale RUN [--baseline RUN] [--samples-per-group K] + python -m tools.eval.sampler_islands latency RUN +""" + +from __future__ import annotations + +import argparse +import json +import re +import statistics +import sys +from collections import defaultdict +from collections.abc import Iterable +from pathlib import Path +from typing import Any + +SAMPLER = "sampler" +TRAINER = "trainer" +_ISLAND_DIR = re.compile(r"l(\d+)") + + +# ---------------------------------------------------------------- reading +def read_jsonl(path: Path) -> list[dict[str, Any]]: + """Every JSON object line of ``path`` (a line cut by a live pull is skipped).""" + rows = [] + for line in path.read_text(errors="replace").splitlines(): + try: + row = json.loads(line) + except ValueError: + continue + if isinstance(row, dict) and "event" in row: + rows.append(row) + return rows + + +def _island_dirs(run: Path) -> dict[int, list[Path]]: + found: dict[int, list[Path]] = defaultdict(list) + for base in sorted((run / "tape-direct").glob("*")) + [run / "tapes"]: + if not base.is_dir(): + continue + for directory in sorted(base.iterdir()): + match = _ISLAND_DIR.fullmatch(directory.name) + if match and directory.is_dir(): + found[int(match.group(1))].extend(sorted(directory.rglob("*.jsonl"))) + if found: + break + return found + + +def load_run(run: Path) -> dict[int, list[dict[str, Any]]]: + """{island: events sorted by time}, duplicates (same event and time) dropped.""" + out: dict[int, list[dict[str, Any]]] = {} + for island, files in _island_dirs(Path(run)).items(): + seen, rows = set(), [] + for path in files: + for row in read_jsonl(path): + key = (row.get("event"), row.get("time_unix"), row.get("incarnation")) + if key not in seen: + seen.add(key) + rows.append(row) + out[island] = sorted(rows, key=lambda r: float(r.get("time_unix") or 0.0)) + return out + + +def island_role(events: Iterable[dict[str, Any]]) -> str | None: + names = {e.get("event") for e in events} + if "rl_round_trained" in names: + return TRAINER + if any(str(n).startswith("rl_sampler_") for n in names): + return SAMPLER + return None + + +def _of(events: list[dict[str, Any]], name: str) -> list[dict[str, Any]]: + return [e for e in events if e.get("event") == name] + + +def _spread(values: list[float]) -> dict[str, float] | None: + values = [float(v) for v in values if v is not None] + if not values: + return None + return {"n": len(values), "median": round(statistics.median(values), 3), + "min": round(min(values), 3), "max": round(max(values), 3)} + + +# ---------------------------------------------------------------- E4: scale +def sampler_summary(events: list[dict[str, Any]]) -> dict[str, Any]: + rounds = _of(events, "rl_sampler_round") + generate = [r.get("generate_s") for r in rounds] + starts = [float(r["time_unix"]) - float(r["generate_s"]) for r in rounds + if r.get("generate_s") is not None] + joins, leaves = _of(events, "rl_sampler_join"), _of(events, "rl_sampler_leave") + preemptions = [] + for leave in (x for x in leaves if x.get("reason") == "PREEMPTED"): + t = float(leave["time_unix"]) + rejoin = next((j for j in joins if float(j["time_unix"]) > t), None) + first_round = next((r for r in rounds if float(r["time_unix"]) > t), None) + preemptions.append({ + "leave_unix": round(t, 3), "incarnation": leave.get("incarnation"), + "in_flight_rollout_id": leave.get("in_flight_rollout_id"), + "rejoin_s": None if rejoin is None else round(float(rejoin["time_unix"]) - t, 3), + "first_new_round_s": None if first_round is None else round(float(first_round["time_unix"]) - t, 3), + }) + return { + "incarnations": sorted({str(j.get("incarnation")) for j in joins}), + "rounds": len(rounds), + "groups": sum(int(r.get("groups") or 0) for r in rounds), + "generate_s": None if any(g is None for g in generate) or not generate + else round(sum(float(g) for g in generate), 3), + "round_generate_s": _spread(generate), + "backlog_wait_s": round(sum(float(e.get("waited_s") or 0.0) + for e in _of(events, "rl_sampler_admitted")), 3), + "policy_wait_s": round(sum(float(e.get("waited_s") or 0.0) + for e in _of(events, "rl_sampler_policy_ready")), 3), + "first_round_start_unix": round(min(starts), 3) if starts else None, + "last_round_end_unix": round(float(rounds[-1]["time_unix"]), 3) if rounds else None, + "preemptions": preemptions, + } + + +def trainer_summary(events: list[dict[str, Any]]) -> dict[str, Any]: + trained = _of(events, "rl_round_trained") + groups = [int(e.get("trained_groups") or 0) for e in trained] + out: dict[str, Any] = {"rounds": len(trained), "groups": sum(groups), + "steady_groups_per_s": None, "steady_window_s": None} + if len(trained) >= 2: + window = float(trained[-1]["time_unix"]) - float(trained[0]["time_unix"]) + if window > 0: + out["steady_window_s"] = round(window, 3) + out["steady_groups_per_s"] = round(sum(groups[1:]) / window, 4) + return out + + +def scale_report(run: Path, *, samples_per_group: int | None = None) -> dict[str, Any]: + islands = load_run(run) + samplers = {m: sampler_summary(ev) for m, ev in islands.items() if island_role(ev) == SAMPLER} + trainers = {m: trainer_summary(ev) for m, ev in islands.items() if island_role(ev) == TRAINER} + starts = [s["first_round_start_unix"] for s in samplers.values() if s["first_round_start_unix"]] + ends = [s["last_round_end_unix"] for s in samplers.values() if s["last_round_end_unix"]] + groups = sum(s["groups"] for s in samplers.values()) + production = None + if starts and ends and max(ends) > min(starts): + production = round(groups / (max(ends) - min(starts)), 4) + trainer = next(iter(trainers.values()), None) + report = { + "run": Path(run).name, "sampler_islands": len(samplers), "samplers": samplers, + "trainer": trainer, + "sampler_groups": groups, + "production_groups_per_s": production, + "consumption_groups_per_s": None if trainer is None else trainer["steady_groups_per_s"], + "preemptions": sum(len(s["preemptions"]) for s in samplers.values()), + } + if samples_per_group: + k = int(samples_per_group) + for key in ("production_groups_per_s", "consumption_groups_per_s"): + value = report[key] + report[key.replace("groups", "samples")] = None if value is None else round(value * k, 4) + return report + + +def compare(run: dict[str, Any], baseline: dict[str, Any]) -> dict[str, Any]: + """Speedup of ``run`` over ``baseline`` (consumption and production rates). + Same sampler count: ``effective_ratio`` (a preempted run vs an unpreempted + one); more samplers: ``speedup`` and ``efficiency`` = speedup / island ratio.""" + out: dict[str, Any] = {"baseline": baseline["run"]} + for key in ("consumption_groups_per_s", "production_groups_per_s"): + a, b = run.get(key), baseline.get(key) + out[key.replace("_groups_per_s", "_ratio")] = None if not a or not b else round(a / b, 4) + n, n0 = run["sampler_islands"], baseline["sampler_islands"] + ratio = out["consumption_ratio"] + if n == n0: + out["effective_ratio"] = ratio + elif n0: + out["speedup"] = ratio + out["efficiency"] = None if ratio is None else round(ratio / (n / n0), 4) + return out + + +# ---------------------------------------------------------------- E5: latency +def _sum(timings: dict[str, Any] | None) -> float | None: + if not timings: + return None + return round(sum(float(v) for v in timings.values()), 3) + + +def latency_report(run: Path) -> dict[str, Any]: + islands = load_run(run) + released: dict[int, dict[str, Any]] = {} + for events in (ev for ev in islands.values() if island_role(ev) == TRAINER): + for e in _of(events, "rl_policy_released"): + released[int(e["seq"])] = e + rows = [] + for island, events in sorted(islands.items()): + if island_role(events) != SAMPLER: + continue + for loaded in _of(events, "rl_sampler_policy_loaded"): + rel = released.get(int(loaded.get("seq", -1))) + if rel is None: + continue + t_rel, t_load = float(rel["time_unix"]), float(loaded["time_unix"]) + seen = loaded.get("seen_at") + follower = loaded.get("follower_timings") or None + row = { + "island": island, "seq": int(loaded["seq"]), + "policy_version": loaded.get("policy_version"), "weights": loaded.get("weights"), + "release_bytes": rel.get("release_bytes"), + "trainer_publish_s": rel.get("publish_s"), + "trainer_release_s": _sum(rel.get("release_timings")), + "trainer_release_phases": rel.get("release_timings"), + "wait_s": None if seen is None else round(float(seen) - t_rel, 3), + "sampler_load_s": _sum(follower), + "sampler_follower_phases": follower, + "sampler_load_phases": loaded.get("load_timings"), + "fused_stage": loaded.get("fused_stage"), + "released_to_loaded_s": round(t_load - t_rel, 3), + } + parts = (row["trainer_publish_s"], row["trainer_release_s"], row["released_to_loaded_s"]) + row["end_to_end_s"] = None if any(p is None for p in parts) else round(sum(parts), 3) + rows.append(row) + keys = ("trainer_publish_s", "trainer_release_s", "wait_s", "sampler_load_s", + "released_to_loaded_s", "end_to_end_s") + phases: dict[str, list[float]] = defaultdict(list) + for row in rows: + for group in ("trainer_release_phases", "sampler_follower_phases", "sampler_load_phases"): + for name, value in (row.get(group) or {}).items(): + phases[f"{group.split('_')[0]}.{name}"].append(value) + # the first release is loaded before the first round (start-up, not a policy update) + later = [r for r in rows if r["seq"] > 1] + return { + "run": Path(run).name, "loads": len(rows), "rows": rows, + "summary_after_first": {k: _spread([r[k] for r in later]) for k in keys}, + "phases": {name: _spread(values) for name, values in sorted(phases.items())}, + } + + +def main(argv: list[str] | None = None) -> int: + parser = argparse.ArgumentParser(prog="python -m tools.eval.sampler_islands") + sub = parser.add_subparsers(dest="command", required=True) + scale = sub.add_parser("scale") + scale.add_argument("run", type=Path) + scale.add_argument("--baseline", type=Path, default=None) + scale.add_argument("--samples-per-group", type=int, default=None) + latency = sub.add_parser("latency") + latency.add_argument("run", type=Path) + args = parser.parse_args(argv) + if args.command == "scale": + report = scale_report(args.run, samples_per_group=args.samples_per_group) + if args.baseline is not None: + report["compare"] = compare(report, scale_report(args.baseline, + samples_per_group=args.samples_per_group)) + else: + report = latency_report(args.run) + json.dump(report, sys.stdout, indent=1, sort_keys=False) + print() + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/yeto/rl/adapters/miles/config.py b/yeto/rl/adapters/miles/config.py index 1af82e71..528dc6bd 100644 --- a/yeto/rl/adapters/miles/config.py +++ b/yeto/rl/adapters/miles/config.py @@ -475,7 +475,8 @@ def _check_sglang_tp(value, config): "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")}, + for name in ("role", "store", "max_backlog_groups", "wait_s", "min_round_s", + "fused_stage")}, } diff --git a/yeto/rl/adapters/miles/sampler_island.py b/yeto/rl/adapters/miles/sampler_island.py index 22da4c7d..546725b4 100644 --- a/yeto/rl/adapters/miles/sampler_island.py +++ b/yeto/rl/adapters/miles/sampler_island.py @@ -27,6 +27,13 @@ (:attr:`EngineReleaseLoader.last_checksum`); they must equal the trainer's (``PolicyRelease.engine_checksum``) when the release carries it. +LoRA staging (S26 E5): the adapter is copied to local disk before the update +lock. With ``fused_stage`` (``--rl-sampler-fused-stage``, default off) the copy +also computes the release digest, so the release is read from the store once +instead of twice (the follower's check, then the copy); a digest mismatch fails +the load before any engine sees the files, as the follower's check does. +``last_timings`` holds the seconds of each load phase. + This module imports Miles only inside :func:`run_sampler_island`. """ @@ -38,16 +45,20 @@ import shutil import tempfile import time -from collections.abc import Callable +from collections.abc import Callable, Iterator +from contextlib import contextmanager from pathlib import Path from typing import Any from yeto.rl.engine.policy_release import ( FULL, LORA, + PhaseTimer, PolicyFollower, PolicyRelease, PolicyReleaseStore, + check_release_digest, + copy_with_digest, ) from yeto.rl.engine.merged_release import MergedReleaseGuard, is_publisher, merged_version from yeto.rl.engine.policy_versions import engine_weight_version @@ -132,7 +143,8 @@ class EngineReleaseLoader: must then carry ``weights == "lora"``. ``adapter_loaded``: the engines already serve the adapter (a convertible node: its trainer pushed it at boot), so the first load unloads it first. ``warmup`` and ``flatten`` replace the warm-up - generate and the checksum flattening (tests).""" + generate and the checksum flattening (tests). ``fused_stage``: the LoRA staging + copy checks the release digest (module doc; ``verifies_files``).""" def __init__(self, controller: Any, executor: Any, run: Callable[[Any], Any], *, lora: bool = False, post: Callable[..., Any] = _post_engine, @@ -140,7 +152,8 @@ def __init__(self, controller: Any, executor: Any, run: Callable[[Any], Any], *, flatten: Callable[[Any], list[dict[str, Any]]] | None = None, probe: Callable[[Any, str], Any] | None = probe_generate, warmup_timeout_s: float = 120.0, staging_root: str | None = None, - adapter_loaded: bool = False) -> None: + adapter_loaded: bool = False, fused_stage: bool = False, + timer: Callable[[], float] = time.perf_counter) -> None: self.controller = controller self.executor = executor self.run = run @@ -158,23 +171,72 @@ def __init__(self, controller: Any, executor: Any, run: Callable[[Any], Any], *, self.staging_root = staging_root # local copies of loaded adapters (None: system temp) self._staged: Path | None = None # the local copy the engines serve self.last_checksum: str | None = None # LoRA: the engines' lora:* read-back digest + self.fused_stage = bool(fused_stage) and self.lora + self.timer = timer + self._timing: PhaseTimer | None = None # the load in progress + self.last_timings: dict[str, float] = {} # seconds per phase of the last load @property def weights(self) -> str: return LORA if self.lora else FULL + @property + def verifies_files(self) -> bool: + """True: this loader checks the release digest (``PolicyFollower`` skips its check).""" + return self.fused_stage + + @contextmanager + def _phase(self, name: str) -> Iterator[None]: + if self._timing is None: + yield + return + with self._timing.phase(name): + yield + def __call__(self, release: PolicyRelease, directory: Path) -> None: if release.weights != self.weights: raise SampleBufferRolloutError( f"release r{release.seq}-v{release.version} holds {release.weights} weights; " f"this sampler serves {self.weights}") - engines = self.run(self._load(str(directory), int(release.version))) - if self.lora: - self.last_checksum = self.run(self._read_back_adapter(engines)) - if release.engine_checksum is not None and self.last_checksum != release.engine_checksum: - raise SampleBufferRolloutError( - f"release r{release.seq}-v{release.version}: the engines serve LoRA checksum " - f"{self.last_checksum}, the trainer published {release.engine_checksum}") + self._timing = PhaseTimer(self.timer) + try: + self._call(release, directory) + finally: + self.last_timings = self._timing.as_dict() + self._timing = None + + def _call(self, release: PolicyRelease, directory: Path) -> None: + if not self.lora: + self.run(self._load(str(directory), int(release.version))) + return + staged = self._stage(release, directory) + try: + engines = self.run(self._load(str(staged), int(release.version))) + except BaseException: + if self._staged != staged: # no engine took this copy + shutil.rmtree(staged, ignore_errors=True) + raise + self.last_checksum = self.run(self._read_back_adapter(engines)) + if release.engine_checksum is not None and self.last_checksum != release.engine_checksum: + raise SampleBufferRolloutError( + f"release r{release.seq}-v{release.version}: the engines serve LoRA checksum " + f"{self.last_checksum}, the trainer published {release.engine_checksum}") + + def _stage(self, release: PolicyRelease, directory: Path) -> Path: + """A local copy of the adapter (SGLang keeps the adapter files open: open + files on a Modal Volume make every later reload fail, T-LSMP g). With + ``fused_stage`` the copy also checks the release digest.""" + staged = Path(tempfile.mkdtemp(prefix=f"v{int(release.version)}-", dir=self.staging_root)) + try: + with self._phase("stage_s"): + if self.fused_stage: + check_release_digest(release, copy_with_digest(Path(directory), staged)) + else: + shutil.copytree(directory, staged, dirs_exist_ok=True) + except BaseException: + shutil.rmtree(staged, ignore_errors=True) + raise + return staged def stamp(self, version: int) -> None: """Stamp ``version`` on the engines without loading weights (the base policy).""" @@ -186,26 +248,18 @@ async def _load_weights(self, engines: list[Any], path: str, version: int) -> No await asyncio.gather(*(e.update_weights_from_disk(model_path=path, weight_version=stamp) for e in engines)) return - # SGLang keeps the adapter files open (safetensors mmap); open files on a Modal - # Volume make every later reload fail, so no new release is ever seen - # (T-LSMP g). The engines load a local copy; the previous copy is removed - # once its adapter is unloaded. - staged = Path(tempfile.mkdtemp(prefix=f"v{int(version)}-", dir=self.staging_root)) - shutil.copytree(path, staged, dirs_exist_ok=True) - try: - if self._adapter_loaded: - await asyncio.gather(*(self.post(e, "unload_lora_adapter", {"lora_name": self.adapter}) - for e in engines)) - self._adapter_loaded = False - self._drop_staged() - await asyncio.gather(*(self.post(e, "load_lora_adapter", - {"lora_name": self.adapter, "lora_path": str(staged)}) + # ``path`` is the local copy (``_stage``); the previous copy is removed once its + # adapter is unloaded. On a failure before the load the caller removes ``path``. + if self._adapter_loaded: + await asyncio.gather(*(self.post(e, "unload_lora_adapter", {"lora_name": self.adapter}) for e in engines)) - except BaseException: - shutil.rmtree(staged, ignore_errors=True) - raise + self._adapter_loaded = False + self._drop_staged() + await asyncio.gather(*(self.post(e, "load_lora_adapter", + {"lora_name": self.adapter, "lora_path": path}) + for e in engines)) self._adapter_loaded = True - self._staged = staged + self._staged = Path(path) def _drop_staged(self) -> None: if self._staged is not None: @@ -214,35 +268,42 @@ def _drop_staged(self) -> None: async def _read_back_adapter(self, engines: list[Any]) -> str | None: """Warm the adapter into every engine's LoRA pool, then read the checksums.""" - await asyncio.gather(*(self.warmup(e, self.adapter, self.warmup_timeout_s) for e in engines)) - bodies = await read_engine_checksums(self.controller, self.flatten) + with self._phase("warmup_s"): + await asyncio.gather(*(self.warmup(e, self.adapter, self.warmup_timeout_s) for e in engines)) + with self._phase("checksum_s"): + bodies = await read_engine_checksums(self.controller, self.flatten) if self.probe is not None: - self.last_probe = dict(await self.probe(engines[0], self.adapter)) + with self._phase("probe_s"): + self.last_probe = dict(await self.probe(engines[0], self.adapter)) return weight_view_digest(bodies, lora=True) async def _load(self, path: str | None, version: int) -> list[Any]: """Load and stamp under the update lock; returns the engines.""" - info = await self.controller.start_update_weights() + with self._phase("lock_s"): + info = await self.controller.start_update_weights() ok = False try: engines = list(info.rollout_engines) if not engines: raise SampleBufferRolloutError("no rollout engines to load the policy into") if path is not None: - await self._load_weights(engines, path, version) + with self._phase("engine_load_s"): + await self._load_weights(engines, path, version) ids = [_engine_id(e, i) for i, e in enumerate(engines)] - failed, cause = await _stamp_engine_version(engines, ids, version) + with self._phase("stamp_s"): + failed, cause = await _stamp_engine_version(engines, ids, version) if failed: raise SampleBufferRolloutError( f"{len(failed)}/{len(engines)} rollout engines did not take policy " f"v{version} ({cause.value}): {sorted(failed)}") ok = True finally: - if ok: - await self.controller.end_update_weights( - snapshot_cell_id_to_hashes=info.snapshot_cell_id_to_hashes) - else: - await self.controller.abort_update_weights() + with self._phase("unlock_s"): + if ok: + await self.controller.end_update_weights( + snapshot_cell_id_to_hashes=info.snapshot_cell_id_to_hashes) + else: + await self.controller.abort_update_weights() await _maybe_await(self.executor.set_weight_version(int(version))) return engines @@ -352,7 +413,8 @@ def leave_reclaimed() -> None: try: controller, executor = runner.run(init()) lora = lora_mode(miles_args) - loader = EngineReleaseLoader(controller, executor, runner.run, lora=lora) + loader = EngineReleaseLoader(controller, executor, runner.run, lora=lora, + fused_stage=bool(config.fused_stage)) if not lora: # LoRA: no adapter is served before the first release (require_release) loader.stamp(BASE_POLICY_VERSION) # until the first release is loaded follower = PolicyFollower(release_store(config, environ), loader, @@ -364,7 +426,9 @@ def leave_reclaimed() -> None: max_backlog_groups=int(config.max_backlog_groups), min_round_s=config.min_round_s, require_release=lora, should_stop=reclaimed, loaded_fields=lambda: {"served_checksum": loader.last_checksum, - "probe": loader.last_probe}) + "probe": loader.last_probe, + "load_timings": loader.last_timings, + "fused_stage": loader.fused_stage}) latest = follower.releases.latest() membership.join(latest_release=None if latest is None else {"seq": latest.seq, "policy_version": latest.version}) @@ -450,7 +514,9 @@ def _engine_checksum(self) -> str | None: return weight_view_digest(list(checksums.values()), lora=self.weights == LORA) def publish(self, state: Any, **kwargs: Any) -> Any: + started = time.perf_counter() result = self.inner.publish(state, **kwargs) + publish_s = round(time.perf_counter() - started, 3) # the trainer's own engines version = int(self.version_of(state)) if self.guard is not None: latest = self.releases.latest() @@ -476,7 +542,9 @@ def publish(self, state: Any, **kwargs: Any) -> Any: local_version=int(state.policy_version), origin=release.origin, policy_hash=release.policy_hash, buffer_selected=release.buffer_selected, weights=release.weights, sha256=release.sha256, - engine_checksum=release.engine_checksum) + engine_checksum=release.engine_checksum, released_at=release.released_at, + publish_s=publish_s, release_timings=self.releases.last_timings, + release_bytes=sum((release.files or {}).values())) return result def __getattr__(self, name: str) -> Any: diff --git a/yeto/rl/engine/policy_release.py b/yeto/rl/engine/policy_release.py index 45aeef17..cc279b16 100644 --- a/yeto/rl/engine/policy_release.py +++ b/yeto/rl/engine/policy_release.py @@ -50,6 +50,14 @@ :class:`PolicyFollower` is the sampler side: it loads a release when its ``seq`` differs from the one its engines serve. When the trainer stops, no new release appears and the sampler keeps the last version. + +Latency phases (S26 E5): the trainer side (:meth:`PolicyReleaseStore.release`) +and the sampler side (:meth:`PolicyFollower.refresh`) time their phases with +:class:`PhaseTimer` (``last_timings``, seconds); the islands write them in their +release and load events. They only measure: no phase changes its order or work. +A load hook that checks the weight files itself while it copies them +(``verifies_files``, :func:`copy_with_digest`) reads a LoRA release once instead +of twice; the follower then skips its own check. """ from __future__ import annotations @@ -60,7 +68,8 @@ import re import shutil import time -from collections.abc import Callable +from collections.abc import Callable, Iterator +from contextlib import contextmanager from dataclasses import asdict, dataclass from pathlib import Path from typing import Any @@ -134,22 +143,75 @@ def check_release_scheme(release: PolicyRelease, expected: str) -> None: f"this reader serves {expected!r} (it follows the wrong release directory?)") +class PhaseTimer: + """Wall seconds per named phase (``phase(name)`` adds to the name's total). + ``as_dict()``: the totals in milliseconds precision, in first-use order.""" + + def __init__(self, clock: Callable[[], float] = time.perf_counter) -> None: + self.clock = clock + self._totals: dict[str, float] = {} + + @contextmanager + def phase(self, name: str) -> Iterator[None]: + started = self.clock() + try: + yield + finally: + self._totals[name] = self._totals.get(name, 0.0) + (self.clock() - started) + + def add(self, name: str, seconds: float) -> None: + self._totals[name] = self._totals.get(name, 0.0) + float(seconds) + + def as_dict(self) -> dict[str, float]: + return {name: round(value, 3) for name, value in self._totals.items()} + + +def _digested_files(root: Path) -> list[Path]: + """The files :func:`directory_digest` covers, in its order (all but the record).""" + return sorted(p for p in root.rglob("*") if p.is_file() and p.relative_to(root) != Path(RECORD)) + + +def _digest_header(digest: Any, root: Path, path: Path) -> None: + relative = path.relative_to(root).as_posix().encode("utf-8") + digest.update(len(relative).to_bytes(8, "big") + relative) + digest.update(os.path.getsize(path).to_bytes(8, "big")) + + def directory_digest(directory: Path) -> str: """sha256 over every file under ``directory`` except the release record: relative path, size and bytes, in sorted path order.""" digest = hashlib.sha256() root = Path(directory) - files = sorted(p for p in root.rglob("*") if p.is_file() and p.relative_to(root) != Path(RECORD)) - for path in files: - relative = path.relative_to(root).as_posix().encode("utf-8") - digest.update(len(relative).to_bytes(8, "big") + relative) - digest.update(os.path.getsize(path).to_bytes(8, "big")) + for path in _digested_files(root): + _digest_header(digest, root, path) with path.open("rb") as handle: while chunk := handle.read(_DIGEST_CHUNK): digest.update(chunk) return digest.hexdigest() +def copy_with_digest(source: Path, target: Path) -> str: + """Copy every file under ``source`` into ``target`` (created; existing files are + overwritten) and return :func:`directory_digest` of ``source``, reading each + file once. The release record is copied but not digested, as there.""" + root, out = Path(source), Path(target) + out.mkdir(parents=True, exist_ok=True) + digested = set(_digested_files(root)) + digest = hashlib.sha256() + for path in sorted(p for p in root.rglob("*") if p.is_file()): + destination = out / path.relative_to(root) + destination.parent.mkdir(parents=True, exist_ok=True) + hashed = path in digested + if hashed: + _digest_header(digest, root, path) + with path.open("rb") as reader, destination.open("wb") as writer: + while chunk := reader.read(_DIGEST_CHUNK): + if hashed: + digest.update(chunk) + writer.write(chunk) + return digest.hexdigest() + + class ReleaseIncomplete(PolicyReleaseError): """A release whose files are not all present with their recorded sizes yet.""" @@ -195,10 +257,14 @@ def check_complete(release: PolicyRelease, directory: Path) -> None: def check_release_files(release: PolicyRelease, directory: Path) -> None: """Raise when the weight files differ from the ones the trainer released.""" - if release.sha256 is None: - return - actual = directory_digest(directory) - if actual != release.sha256: + if release.sha256 is not None: + check_release_digest(release, directory_digest(directory)) + + +def check_release_digest(release: PolicyRelease, actual: str) -> None: + """Raise when ``actual`` (a :func:`directory_digest`) is not the released one + (no check for a release without ``sha256``).""" + if release.sha256 is not None and actual != release.sha256: raise PolicyReleaseError(f"release r{release.seq}-v{release.version}: weight files " f"sha256 {actual} != released {release.sha256}") @@ -212,7 +278,8 @@ class PolicyReleaseStore: def __init__(self, store: CheckpointStore, *, keep: int = 2, clock: Callable[[], float] = time.time, cache_dir: str | os.PathLike[str] = DEFAULT_CACHE_DIR, - release_dir: str = RELEASE_DIR) -> None: + release_dir: str = RELEASE_DIR, + timer: Callable[[], float] = time.perf_counter) -> None: if int(keep) < 1: raise ValueError("keep must be >= 1") self.store = store @@ -220,6 +287,9 @@ def __init__(self, store: CheckpointStore, *, keep: int = 2, self.clock = clock self.release_dir = str(release_dir) self.root = Path(store.root) / self.release_dir + self.timer = timer + # seconds per phase of the last release() that wrote a release (module doc) + self.last_timings: dict[str, float] = {} self.cache_dir = Path(cache_dir).expanduser() if self.release_dir != RELEASE_DIR: # its own cache: the prune below drops other directories self.cache_dir /= self.release_dir @@ -293,30 +363,39 @@ def release(self, version: int, policy_hash: str, write: Callable[[Path], None], raise ValueError("policy_hash must be non-empty") if weights not in WEIGHT_FORMATS: raise ValueError(f"weights must be one of {WEIGHT_FORMATS}, not {weights!r}") - latest = self.latest() + timer = PhaseTimer(self.timer) + with timer.phase("read_latest_s"): + latest = self.latest() if latest is not None and (latest.version, latest.policy_hash) == (version, policy_hash): return None seq = 1 if latest is None else latest.seq + 1 name = f"r{seq}-v{version}" directory = self.root / name - if directory.exists(): # a partial release of an earlier attempt (no LATEST names it) - shutil.rmtree(directory) - directory.mkdir(parents=True) - write(directory) + with timer.phase("export_s"): + if directory.exists(): # a partial release of an earlier attempt (no LATEST names it) + shutil.rmtree(directory) + directory.mkdir(parents=True) + write(directory) + with timer.phase("digest_s"): + sha256 = directory_digest(directory) if weights == LORA else None + files = file_sizes(directory) record = PolicyRelease(seq=seq, version=version, policy_hash=str(policy_hash), path=name, buffer_selected=None if buffer_selected is None else int(buffer_selected), - released_at=round(self.clock(), 3), weights=weights, - sha256=directory_digest(directory) if weights == LORA else None, - engine_checksum=engine_checksum or None, files=file_sizes(directory), + released_at=round(self.clock(), 3), weights=weights, sha256=sha256, + engine_checksum=engine_checksum or None, files=files, buffer_shards=None if buffer_shards is None else {str(k): int(v) for k, v in sorted(buffer_shards.items())}, derived_from=None if derived_from is None else dict(derived_from), origin=None if origin is None else dict(origin)) - _write_json_atomic(directory / RECORD, asdict(record)) - self.store.commit() # the weights and the record are durable before LATEST names them - _write_json_atomic(self.root / LATEST, asdict(record)) - self.store.commit() - self._prune() + with timer.phase("commit_weights_s"): + _write_json_atomic(directory / RECORD, asdict(record)) + self.store.commit() # the weights and the record are durable before LATEST names them + with timer.phase("commit_latest_s"): + _write_json_atomic(self.root / LATEST, asdict(record)) + self.store.commit() + with timer.phase("prune_s"): + self._prune() + self.last_timings = timer.as_dict() return record def _prune(self) -> list[str]: @@ -337,15 +416,22 @@ class PolicyFollower: ``load(release, directory)`` loads the weights and stamps ``release.version`` on every engine (raises on any failure; the served release is then unchanged). - The weight files are checked against the record's ``sha256`` first. + The weight files are checked against the record's ``sha256`` first, unless + the hook has a true ``verifies_files`` attribute: it then checks them itself + (with :func:`check_release_digest`) before any engine sees them. ``version`` is the version the engines serve before any release (the base policy, 0). ``scheme``: the weight scheme the engines serve; a release of another scheme raises :class:`SchemeMismatch` before any file is read - (None: no check, for callers that check in ``load``).""" + (None: no check, for callers that check in ``load``). + + ``last_timings``: seconds per phase of the last refresh that loaded a + release; ``seen_at``: the wall time this follower first read the newest + release (``last_seen``) in ``LATEST``.""" def __init__(self, releases: PolicyReleaseStore, load: Callable[[PolicyRelease, Path], None], *, version: int = 0, - scheme: str | None = None) -> None: + scheme: str | None = None, timer: Callable[[], float] = time.perf_counter, + wall: Callable[[], float] = time.time) -> None: self.releases = releases self.scheme = scheme self._load = load @@ -353,25 +439,38 @@ def __init__(self, releases: PolicyReleaseStore, self.release: PolicyRelease | None = None # the release the engines serve self.last_seen: PolicyRelease | None = None # the newest release read self.incomplete: str | None = None # why the newest release is not loaded yet + self.timer = timer + self.wall = wall + self.seen_at: float | None = None + self.last_timings: dict[str, float] = {} def refresh(self) -> PolicyRelease | None: """The release loaded by this call, or None when there is no new one, or when the newest one is not complete yet (``incomplete`` says why; the engines keep the release they serve).""" - latest = self.releases.latest() + timer = PhaseTimer(self.timer) + with timer.phase("read_latest_s"): + latest = self.releases.latest() + if latest is not None and (self.last_seen is None or latest.seq != self.last_seen.seq): + self.seen_at = round(self.wall(), 3) self.last_seen = latest if latest is None or (self.release is not None and latest.seq == self.release.seq): return None if self.scheme is not None: check_release_scheme(latest, self.scheme) try: - directory = self.releases.resolve(latest) + with timer.phase("resolve_s"): + directory = self.releases.resolve(latest) except ReleaseIncomplete as error: self.incomplete = str(error) return None self.incomplete = None - check_release_files(latest, directory) - self._load(latest, directory) + if not getattr(self._load, "verifies_files", False): + with timer.phase("verify_s"): + check_release_files(latest, directory) + with timer.phase("load_s"): + self._load(latest, directory) self.version = latest.version self.release = latest + self.last_timings = timer.as_dict() return latest diff --git a/yeto/rl/engine/run_flags.py b/yeto/rl/engine/run_flags.py index 2af06a5c..7711e3d5 100644 --- a/yeto/rl/engine/run_flags.py +++ b/yeto/rl/engine/run_flags.py @@ -585,6 +585,10 @@ def rl_forward_pairs(*dests: str | tuple[str, str]) -> tuple[tuple[str, str], .. shared_flag("--rl-sampler-min-round-seconds", type=float, default=None, metavar="S", help="--rl-sample-role sampler: start one round at most every S seconds (a " "throttle; a check run uses it to make the trainer idle); default no throttle"), + shared_flag("--rl-sampler-fused-stage", action="store_true", default=False, + help="--rl-sample-role sampler with --tuning lora: check each release's digest " + "while copying the adapter to local disk (one read of the store instead of " + "two); default off"), shared_flag("--rl-sample-wait-s", type=float, default=None, metavar="S", help="--rl-sample-role trainer: max seconds to wait for one full round from " "the buffer (default 3600)"), diff --git a/yeto/rl/engine/sample_roles.py b/yeto/rl/engine/sample_roles.py index ca766c3e..d9fcf673 100644 --- a/yeto/rl/engine/sample_roles.py +++ b/yeto/rl/engine/sample_roles.py @@ -58,13 +58,16 @@ class SampleBufferConfig: groups); ``wait_s``: trainer, max seconds to wait for one full round; ``min_round_s``: sampler, min seconds from one round start to the next (a - throttle, so a test can make the sampler slower than the trainer).""" + throttle, so a test can make the sampler slower than the trainer); + ``fused_stage``: LoRA sampler, check the release digest while copying the + adapter to local disk (``--rl-sampler-fused-stage``, S26 E5).""" role: str store: str max_backlog_groups: int | None = None wait_s: float | None = None min_round_s: float | None = None + fused_stage: bool = False def sample_role(args: Any) -> str: @@ -229,6 +232,8 @@ def island_args(args: Any, island: int, num_islands: int) -> Any: else: out.rl_sampler_max_backlog_rounds = None out.rl_sampler_min_round_seconds = None + if out.rl_sample_role != SAMPLER: + out.rl_sampler_fused_stage = False return out @@ -321,6 +326,8 @@ def check_launch(args: Any) -> None: ("--rl-sample-wait-s", "rl_sample_wait_s")): if getattr(args, dest, None) is not None: raise ValueError(f"{flag} needs --rl-sample-role sampler or trainer") + if getattr(args, "rl_sampler_fused_stage", False): + raise ValueError("--rl-sampler-fused-stage needs --rl-sample-role sampler") return if not getattr(args, "rl_single_island_no_sync", False) and role != TRAINER: raise ValueError(f"--rl-sample-role {role} needs --rl-single-island-no-sync: a sampler " @@ -349,6 +356,10 @@ def check_launch(args: Any) -> None: if min_round_s is not None and (role != SAMPLER or not float(min_round_s) > 0): raise ValueError("--rl-sampler-min-round-seconds needs --rl-sample-role sampler " "and a value > 0") + if getattr(args, "rl_sampler_fused_stage", False) and ( + role != SAMPLER or (getattr(args, "tuning", None) or "lora") != "lora"): + raise ValueError("--rl-sampler-fused-stage needs --rl-sample-role sampler and " + "--tuning lora (it fuses the LoRA adapter copy with its digest)") wait_s = getattr(args, "rl_sample_wait_s", None) if wait_s is not None and (role not in (TRAINER, CONVERTIBLE) or not float(wait_s) > 0): raise ValueError("--rl-sample-wait-s needs --rl-sample-role trainer and a value > 0") @@ -390,7 +401,8 @@ def resolve(args: Any, *, groups_per_round: int) -> SampleBufferConfig | None: max_backlog_groups=None if rounds is None else int(groups_per_round) * rounds, wait_s=None if wait_s is None else float(wait_s), min_round_s=(float(min_round_s) if role == SAMPLER and min_round_s is not None - else None)) + else None), + fused_stage=role == SAMPLER and bool(getattr(args, "rl_sampler_fused_stage", False))) def island_env(args: Any) -> dict[str, str]: @@ -420,6 +432,8 @@ def island_flags(args: Any) -> str: flags += f" --rl-sampler-min-round-seconds {float(args.rl_sampler_min_round_seconds)!r}" if getattr(args, "rl_sample_wait_s", None) is not None: flags += f" --rl-sample-wait-s {float(args.rl_sample_wait_s)!r}" + if role == SAMPLER and getattr(args, "rl_sampler_fused_stage", False): + flags += " --rl-sampler-fused-stage" if getattr(args, "rl_sampler_node", None): flags += f" --rl-sampler-node {args.rl_sampler_node}" if getattr(args, "rl_buffer_shards", None): diff --git a/yeto/rl/engine/sampler_loop.py b/yeto/rl/engine/sampler_loop.py index a4cca960..7fb88b4a 100644 --- a/yeto/rl/engine/sampler_loop.py +++ b/yeto/rl/engine/sampler_loop.py @@ -19,6 +19,10 @@ weights that failed a check) it waits instead of generating, so no sample of unchecked weights reaches the buffer. +``rl_sampler_policy_loaded`` carries the latency fields of S26 E5: the trainer's +``released_at``, the follower's ``seen_at`` (both wall clock, of two hosts) and +``follower_timings`` (:attr:`~.policy_release.PolicyFollower.last_timings`). + After the last round it marks the data end, so the trainer stops after reading the rest. The loop is engine-neutral: ``generate(rollout_id)`` runs one round. """ @@ -93,7 +97,9 @@ def _load_newest(self) -> None: if loaded is not None: self.emit("rl_sampler_policy_loaded", seq=loaded.seq, policy_version=loaded.version, policy_hash=loaded.policy_hash, buffer_selected=loaded.buffer_selected, - weights=loaded.weights, sha256=loaded.sha256, **self.loaded_fields()) + weights=loaded.weights, sha256=loaded.sha256, + released_at=loaded.released_at, seen_at=self.follower.seen_at, + follower_timings=self.follower.last_timings, **self.loaded_fields()) def _admit(self, rollout_id: int) -> bool: """True when the round may start; False when asked to stop while waiting.""" @@ -150,9 +156,11 @@ def run(self, start: int = 0) -> int: started = self.clock() self.generate(rollout_id) self.buffer.refresh() + generate_s = self.clock() - started self.emit("rl_sampler_round", rollout_id=rollout_id, policy_version=self.follower.version, - groups=len(self.buffer) - before, buffer_end=len(self.buffer)) + groups=len(self.buffer) - before, buffer_end=len(self.buffer), + generate_s=round(generate_s, 3)) done += 1 self._throttle(started, last=rollout_id == self.num_rollouts - 1) self.buffer.mark_data_end() From f183d759d53610baae773e63acd306a287fbaf0e Mon Sep 17 00:00:00 2001 From: michael <78113554+chowhsu@users.noreply.github.com> Date: Wed, 21 Oct 2026 15:49:48 +0000 Subject: [PATCH 60/62] E4: trainer driver phase seconds and buffer wait share in the scale report Co-Authored-By: Claude Opus 5.5 --- tests/test_tools_sampler_islands.py | 10 ++++++++++ tools/eval/sampler_islands.py | 19 ++++++++++++++++++- 2 files changed, 28 insertions(+), 1 deletion(-) diff --git a/tests/test_tools_sampler_islands.py b/tests/test_tools_sampler_islands.py index 3bbdffbf..34d8e07b 100644 --- a/tests/test_tools_sampler_islands.py +++ b/tests/test_tools_sampler_islands.py @@ -114,3 +114,13 @@ def test_cli_prints_json(tmp_path, capsys): assert tool.main(["scale", str(run), "--baseline", str(run)]) == 0 out = json.loads(capsys.readouterr().out) assert out["compare"]["effective_ratio"] == 1.0 + + +def test_trainer_phase_seconds_and_buffer_wait_share(tmp_path): + run = tmp_path / "ph" + marks = [("publish", 0.0), ("generate", 1.0), ("train", 7.0), ("publish", 9.0), ("finish", 10.0)] + rows = [{"event": "rl_driver_phase", "phase": name, "time_unix": t} for name, t in marks] + _tape(run, 0, "t.jsonl", rows + [_trained(8.0), _trained(9.5)]) + trainer = tool.scale_report(run)["trainer"] + assert trainer["phase_s"] == {"publish": 2.0, "generate": 6.0, "train": 2.0} + assert trainer["buffer_wait_share"] == 0.6 diff --git a/tools/eval/sampler_islands.py b/tools/eval/sampler_islands.py index 5ab79a8f..591a1249 100644 --- a/tools/eval/sampler_islands.py +++ b/tools/eval/sampler_islands.py @@ -147,11 +147,28 @@ def sampler_summary(events: list[dict[str, Any]]) -> dict[str, Any]: } +def driver_phase_seconds(events: list[dict[str, Any]]) -> dict[str, float]: + """Seconds per driver phase: ``rl_driver_phase`` marks a phase start; a phase + lasts until the next mark (the last mark has no end and is not counted). On a + buffer trainer ``generate`` is the wait for a round from the buffer.""" + marks = _of(events, "rl_driver_phase") + out: dict[str, float] = defaultdict(float) + for mark, following in zip(marks, marks[1:]): + out[str(mark.get("phase"))] += float(following["time_unix"]) - float(mark["time_unix"]) + return {name: round(value, 3) for name, value in out.items()} + + def trainer_summary(events: list[dict[str, Any]]) -> dict[str, Any]: trained = _of(events, "rl_round_trained") groups = [int(e.get("trained_groups") or 0) for e in trained] + phases = driver_phase_seconds(events) + total = sum(phases.values()) out: dict[str, Any] = {"rounds": len(trained), "groups": sum(groups), - "steady_groups_per_s": None, "steady_window_s": None} + "steady_groups_per_s": None, "steady_window_s": None, + "phase_s": phases, + # share of the driver's time spent waiting for samples + "buffer_wait_share": round(phases.get("generate", 0.0) / total, 4) + if total > 0 else None} if len(trained) >= 2: window = float(trained[-1]["time_unix"]) - float(trained[0]["time_unix"]) if window > 0: From 72adc7cb24f584d8e070875ac51c692bfe56cd29 Mon Sep 17 00:00:00 2001 From: michael <78113554+chowhsu@users.noreply.github.com> Date: Thu, 22 Oct 2026 02:47:19 +0000 Subject: [PATCH 61/62] E4 tool: merge launch-log events; a preempted round with no groups is not a round Co-Authored-By: Claude Opus 5.5 --- tests/test_tools_sampler_islands.py | 21 ++++++++++++++ tools/eval/sampler_islands.py | 45 +++++++++++++++++++++++++---- 2 files changed, 60 insertions(+), 6 deletions(-) diff --git a/tests/test_tools_sampler_islands.py b/tests/test_tools_sampler_islands.py index 34d8e07b..39b8de85 100644 --- a/tests/test_tools_sampler_islands.py +++ b/tests/test_tools_sampler_islands.py @@ -124,3 +124,24 @@ def test_trainer_phase_seconds_and_buffer_wait_share(tmp_path): trainer = tool.scale_report(run)["trainer"] assert trainer["phase_s"] == {"publish": 2.0, "generate": 6.0, "train": 2.0} assert trainer["buffer_wait_share"] == 0.6 + + +def test_events_only_in_the_launch_log_are_merged(tmp_path): + run = _run(tmp_path, "log", samplers=1) + leave = {"event": "rl_sampler_leave", "time_unix": 121.0, "reason": "PREEMPTED", + "incarnation": "a0", "in_flight_rollout_id": 2} + rejoin = {"event": "rl_sampler_join", "time_unix": 150.0, "incarnation": "b0"} + dup = _sampler_round(110.0, 0) + (run / "launch.ts.log").write_text("\n".join( + f"2026-10-22T00:00:00Z [p-l0-modal] YETO_RL_EVENT {json.dumps(r)}" for r in (leave, rejoin, dup))) + s0 = tool.scale_report(run)["samplers"][0] + assert s0["rounds"] == 2 # the echoed round is not counted twice + assert s0["preemptions"][0]["rejoin_s"] == 29.0 + + +def test_a_round_cut_by_a_preemption_is_not_a_round(tmp_path): + run = _run(tmp_path, "cut", samplers=1) + cut = {"event": "rl_sampler_round", "time_unix": 122.0, "rollout_id": 2, "groups": 0, "generate_s": 3.0} + (run / "launch.ts.log").write_text(f"t [p-l0-modal] YETO_RL_EVENT {json.dumps(cut)}") + s0 = tool.scale_report(run)["samplers"][0] + assert (s0["rounds"], s0["cut_rounds"], s0["cut_generate_s"], s0["generate_s"]) == (2, 1, 3.0, 20.0) diff --git a/tools/eval/sampler_islands.py b/tools/eval/sampler_islands.py index 591a1249..b366db8a 100644 --- a/tools/eval/sampler_islands.py +++ b/tools/eval/sampler_islands.py @@ -76,13 +76,40 @@ def _island_dirs(run: Path) -> dict[int, list[Path]]: return found +_LOG_EVENT = re.compile(r"\[[^\]]*-l(\d+)-modal\] YETO_RL_EVENT (\{.*\})\s*$") + + +def launch_log_events(run: Path) -> dict[int, list[dict[str, Any]]]: + """Events the launch log echoed (``/launch.ts.log``), by island. A container + stopped right after an event may echo it here before its tape file is saved.""" + out: dict[int, list[dict[str, Any]]] = defaultdict(list) + log = Path(run) / "launch.ts.log" + if not log.is_file(): + return out + for line in log.read_text(errors="replace").splitlines(): + match = _LOG_EVENT.search(line) + if not match: + continue + try: + row = json.loads(match.group(2)) + except ValueError: + continue + if isinstance(row, dict) and "event" in row: + out[int(match.group(1))].append(row) + return out + + def load_run(run: Path) -> dict[int, list[dict[str, Any]]]: - """{island: events sorted by time}, duplicates (same event and time) dropped.""" + """{island: events sorted by time} from the tapes and the launch log, + duplicates (same event, time and incarnation) dropped.""" out: dict[int, list[dict[str, Any]]] = {} - for island, files in _island_dirs(Path(run)).items(): + echoed = launch_log_events(Path(run)) + dirs = _island_dirs(Path(run)) + for island in sorted(set(dirs) | set(echoed)): + files = dirs.get(island, []) seen, rows = set(), [] - for path in files: - for row in read_jsonl(path): + for path in [*files, None]: + for row in (read_jsonl(path) if path is not None else echoed.get(island, [])): key = (row.get("event"), row.get("time_unix"), row.get("incarnation")) if key not in seen: seen.add(key) @@ -114,7 +141,9 @@ def _spread(values: list[float]) -> dict[str, float] | None: # ---------------------------------------------------------------- E4: scale def sampler_summary(events: list[dict[str, Any]]) -> dict[str, Any]: - rounds = _of(events, "rl_sampler_round") + # a round cut by a preemption ends with groups 0 (nothing written): not a round + every = _of(events, "rl_sampler_round") + rounds = [r for r in every if int(r.get("groups") or 0) > 0] generate = [r.get("generate_s") for r in rounds] starts = [float(r["time_unix"]) - float(r["generate_s"]) for r in rounds if r.get("generate_s") is not None] @@ -123,7 +152,8 @@ def sampler_summary(events: list[dict[str, Any]]) -> dict[str, Any]: for leave in (x for x in leaves if x.get("reason") == "PREEMPTED"): t = float(leave["time_unix"]) rejoin = next((j for j in joins if float(j["time_unix"]) > t), None) - first_round = next((r for r in rounds if float(r["time_unix"]) > t), None) + after = t if rejoin is None else float(rejoin["time_unix"]) + first_round = next((r for r in rounds if float(r["time_unix"]) > after), None) preemptions.append({ "leave_unix": round(t, 3), "incarnation": leave.get("incarnation"), "in_flight_rollout_id": leave.get("in_flight_rollout_id"), @@ -133,6 +163,9 @@ def sampler_summary(events: list[dict[str, Any]]) -> dict[str, Any]: return { "incarnations": sorted({str(j.get("incarnation")) for j in joins}), "rounds": len(rounds), + "cut_rounds": len(every) - len(rounds), + "cut_generate_s": round(sum(float(r.get("generate_s") or 0.0) for r in every + if int(r.get("groups") or 0) == 0), 3), "groups": sum(int(r.get("groups") or 0) for r in rounds), "generate_s": None if any(g is None for g in generate) or not generate else round(sum(float(g) for g in generate), 3), From dd88b26d3347facb4506addaf1e60e11de2a9604 Mon Sep 17 00:00:00 2001 From: michael <78113554+chowhsu@users.noreply.github.com> Date: Thu, 22 Oct 2026 04:30:54 +0000 Subject: [PATCH 62/62] E5: a release writer keeps the LATEST it wrote and commits once PolicyReleaseStore(cached_latest=True, the default): the writer of a release directory remembers the release it wrote last instead of reloading the store (a Modal Volume reload took 2-27 s per release in E4/E5) to read LATEST. It reloads for its first release in a process (a resumed trainer continues the seq), after a failed release, and when the LATEST it sees without a reload is not the one it wrote (another writer). Weights, record and LATEST go out in one commit; readers already load only complete releases (check_complete), so an early-visible LATEST only delays the load. --rl-release-reload-latest (trainer or convertible node) restores the earlier behavior. rl_policy_released carries latest_source. Co-Authored-By: Claude Opus 5.5 --- tests/test_rl_release_latency.py | 89 +++++++++++++++++++++++- yeto/rl/adapters/miles/config.py | 2 +- yeto/rl/adapters/miles/sampler_island.py | 4 +- yeto/rl/engine/policy_release.py | 77 +++++++++++++++++--- yeto/rl/engine/run_flags.py | 4 ++ yeto/rl/engine/sample_roles.py | 19 ++++- 6 files changed, 181 insertions(+), 14 deletions(-) diff --git a/tests/test_rl_release_latency.py b/tests/test_rl_release_latency.py index f0ca9425..0457161e 100644 --- a/tests/test_rl_release_latency.py +++ b/tests/test_rl_release_latency.py @@ -74,7 +74,8 @@ def test_copy_with_digest_copies_everything_and_equals_the_directory_digest(tmp_ def test_release_records_the_trainer_phases(tmp_path): - releases = PolicyReleaseStore(CheckpointStore(tmp_path), clock=lambda: 7.0, timer=_ticks()) + releases = PolicyReleaseStore(CheckpointStore(tmp_path), clock=lambda: 7.0, timer=_ticks(), + cached_latest=False) assert releases.last_timings == {} lora = releases.release(1, "h1", _write_adapter, weights=LORA) assert lora.released_at == 7.0 @@ -259,3 +260,89 @@ def test_fused_stage_flag_exists_on_both_layers(): parser.add_argument(*flag.options, **dict(flag.kwargs)) assert parser.parse_args([]).rl_sampler_fused_stage is False assert parser.parse_args(["--rl-sampler-fused-stage"]).rl_sampler_fused_stage is True + + +# ---------------------------------------------------------------- the writer's known LATEST +class _CountingStore(CheckpointStore): + def __init__(self, root): + super().__init__(root) + self.reloads = self.commits = 0 + + def reload(self): + self.reloads += 1 + + def commit(self): + self.commits += 1 + return 0.0 + + +def test_one_writer_reloads_only_for_its_first_release_and_commits_once(tmp_path): + store = _CountingStore(tmp_path) + releases = PolicyReleaseStore(store) + first = releases.release(1, "h1", _write_adapter, weights=LORA) + assert (releases.last_latest_source, store.reloads, store.commits) == ("reloaded", 1, 1) + second = releases.release(2, "h2", _write_adapter, weights=LORA) + assert (releases.last_latest_source, store.reloads, store.commits) == ("cached", 1, 2) + assert "commit_s" in releases.last_timings and "commit_weights_s" not in releases.last_timings + assert releases.release(2, "h2", _write_adapter, weights=LORA) is None # republish: still known + assert store.reloads == 1 + assert (first.seq, second.seq) == (1, 2) and PolicyReleaseStore(CheckpointStore(tmp_path)).latest() == second + + +def test_a_resumed_writer_reads_latest_again_and_continues_the_seq(tmp_path): + PolicyReleaseStore(CheckpointStore(tmp_path)).release(1, "h1", _write_weights) + store = _CountingStore(tmp_path) + resumed = PolicyReleaseStore(store) # a new process over the same store + assert resumed.release(1, "h1", _write_weights) is None # same version and hash: nothing + assert (resumed.last_latest_source, store.reloads) == ("reloaded", 1) + nxt = resumed.release(2, "h2", _write_weights) + assert nxt.seq == 2 and resumed.last_latest_source == "cached" + + +def test_another_writer_is_detected_and_latest_is_read_again(tmp_path): + store = _CountingStore(tmp_path) + mine = PolicyReleaseStore(store) + mine.release(1, "h1", _write_weights) + PolicyReleaseStore(CheckpointStore(tmp_path)).release(2, "other", _write_weights) # seq 2 + third = mine.release(3, "h3", _write_weights) + assert (mine.last_latest_source, store.reloads) == ("other_writer", 2) + assert third.seq == 3 # never reuses the other writer's seq + + +def test_a_failed_release_makes_the_next_one_read_latest_again(tmp_path): + store = _CountingStore(tmp_path) + releases = PolicyReleaseStore(store) + releases.release(1, "h1", _write_weights) + with pytest.raises(RuntimeError): + releases.release(2, "h2", lambda d: (_ for _ in ()).throw(RuntimeError("export failed"))) + releases.release(2, "h2", _write_weights) + assert releases.last_latest_source == "reloaded" and store.reloads == 2 + + +def test_cached_latest_off_is_the_earlier_behavior(tmp_path): + store = _CountingStore(tmp_path) + releases = PolicyReleaseStore(store, cached_latest=False) + releases.release(1, "h1", _write_weights) + releases.release(2, "h2", _write_weights) + assert (releases.last_latest_source, store.reloads, store.commits) == ("reloaded", 2, 4) + + +def test_reload_latest_flag_is_trainer_only_and_reaches_the_release_store(tmp_path): + trainer = _launch_args(rl_sample_role="trainer", rl_max_policy_age=1, rl_release_reload_latest=True) + sample_roles.check_launch(trainer) + assert sample_roles.island_flags(trainer).endswith(" --rl-release-reload-latest") + with pytest.raises(ValueError, match="only a trainer writes"): + sample_roles.check_launch(_launch_args(rl_release_reload_latest=True)) + island = SimpleNamespace(rl_sample_role="trainer", rl_resume_store=str(tmp_path), + rl_release_reload_latest=True) + config = sample_roles.resolve(island, groups_per_round=8) + assert config.release_legacy is True + assert si.release_store(config).cached_latest is False + default = sample_roles.resolve(SimpleNamespace(rl_sample_role="trainer", rl_resume_store=str(tmp_path)), + groups_per_round=8) + assert si.release_store(default).cached_latest is True + pair = SimpleNamespace(rl_island_sample_role=["0:sampler", "1:trainer"], rl_release_reload_latest=True, + rl_sampler_min_round_seconds=None, rl_sample_wait_s=None, + rl_sampler_max_backlog_rounds=None) + assert sample_roles.island_args(pair, 0, 2).rl_release_reload_latest is False + assert sample_roles.island_args(pair, 1, 2).rl_release_reload_latest is True diff --git a/yeto/rl/adapters/miles/config.py b/yeto/rl/adapters/miles/config.py index 528dc6bd..2e9b2f5b 100644 --- a/yeto/rl/adapters/miles/config.py +++ b/yeto/rl/adapters/miles/config.py @@ -476,7 +476,7 @@ def _check_sglang_tp(value, config): "train_memory.optimizer_offload_fraction": _ok, **{f"sample_buffer.{name}": _ok for name in ("role", "store", "max_backlog_groups", "wait_s", "min_round_s", - "fused_stage")}, + "fused_stage", "release_legacy")}, } diff --git a/yeto/rl/adapters/miles/sampler_island.py b/yeto/rl/adapters/miles/sampler_island.py index 546725b4..b7952790 100644 --- a/yeto/rl/adapters/miles/sampler_island.py +++ b/yeto/rl/adapters/miles/sampler_island.py @@ -100,7 +100,8 @@ def _buffer_config(miles_args: Any, role: str) -> Any: def release_store(config: Any, environ: Any = None) -> PolicyReleaseStore: from yeto.rl.engine.resume import store_for - return PolicyReleaseStore(store_for(config.store, environ=environ)) + return PolicyReleaseStore(store_for(config.store, environ=environ), + cached_latest=not config.release_legacy) # ---------------------------------------------------------------------- sampler side @@ -544,6 +545,7 @@ def publish(self, state: Any, **kwargs: Any) -> Any: weights=release.weights, sha256=release.sha256, engine_checksum=release.engine_checksum, released_at=release.released_at, publish_s=publish_s, release_timings=self.releases.last_timings, + latest_source=self.releases.last_latest_source, release_bytes=sum((release.files or {}).values())) return result diff --git a/yeto/rl/engine/policy_release.py b/yeto/rl/engine/policy_release.py index cc279b16..5c4fa42e 100644 --- a/yeto/rl/engine/policy_release.py +++ b/yeto/rl/engine/policy_release.py @@ -55,6 +55,20 @@ and the sampler side (:meth:`PolicyFollower.refresh`) time their phases with :class:`PhaseTimer` (``last_timings``, seconds); the islands write them in their release and load events. They only measure: no phase changes its order or work. +Writing with a known LATEST (``cached_latest``, default on): the store has one +writer per release directory (one trainer, or the publisher of synced trainers, +:mod:`.merged_release`; a re-quantizer writes its own directory), so the writer +remembers the release it wrote last and does not reload the store (a Modal +Volume reload takes seconds) to read LATEST again. It still reloads and reads +LATEST for its first release in a process (a resumed trainer continues the +store's ``seq``), after a failed release, and when the LATEST it sees without a +reload is not the one it wrote (another writer). The weights, the record and +LATEST then go to the store in one commit: a reader never loads a release +whose files are not all there with their recorded sizes (``check_complete``), +so LATEST becoming visible with or before the weights only delays the load. +``cached_latest=False`` is the earlier behavior (reload before every release, +one commit for the weights and record, a second one for LATEST). + A load hook that checks the weight files itself while it copies them (``verifies_files``, :func:`copy_with_digest`) reads a LoRA release once instead of twice; the follower then skips its own check. @@ -279,7 +293,8 @@ def __init__(self, store: CheckpointStore, *, keep: int = 2, clock: Callable[[], float] = time.time, cache_dir: str | os.PathLike[str] = DEFAULT_CACHE_DIR, release_dir: str = RELEASE_DIR, - timer: Callable[[], float] = time.perf_counter) -> None: + timer: Callable[[], float] = time.perf_counter, + cached_latest: bool = True) -> None: if int(keep) < 1: raise ValueError("keep must be >= 1") self.store = store @@ -288,6 +303,11 @@ def __init__(self, store: CheckpointStore, *, keep: int = 2, self.release_dir = str(release_dir) self.root = Path(store.root) / self.release_dir self.timer = timer + self.cached_latest = bool(cached_latest) + self._written: PolicyRelease | None = None # the release this process wrote last + # how the last release() found LATEST: "cached", "reloaded" (first release, + # after a failure, or cached_latest off) or "other_writer" + self.last_latest_source: str | None = None # seconds per phase of the last release() that wrote a release (module doc) self.last_timings: dict[str, float] = {} self.cache_dir = Path(cache_dir).expanduser() @@ -365,9 +385,43 @@ def release(self, version: int, policy_hash: str, write: Callable[[Path], None], raise ValueError(f"weights must be one of {WEIGHT_FORMATS}, not {weights!r}") timer = PhaseTimer(self.timer) with timer.phase("read_latest_s"): - latest = self.latest() + latest = self._latest_for_release() if latest is not None and (latest.version, latest.policy_hash) == (version, policy_hash): + self._written = latest # LATEST is known: the next release need not read it again return None + self._written = None # until this release is committed + record = self._write_release(timer, latest, version, policy_hash, write, weights=weights, + buffer_selected=buffer_selected, engine_checksum=engine_checksum, + buffer_shards=buffer_shards, derived_from=derived_from, + origin=origin) + self._written = record + self.last_timings = timer.as_dict() + return record + + def _latest_without_reload(self) -> PolicyRelease | None: + path = self.root / LATEST + if not path.is_file(): + return None + return PolicyRelease.from_record(json.loads(path.read_text(encoding="utf-8"))) + + def _latest_for_release(self) -> PolicyRelease | None: + """LATEST as the writer needs it (module doc, ``cached_latest``).""" + known = self._written + if not self.cached_latest or known is None or self.store.objects is not None: + self.last_latest_source = "reloaded" + return self.latest() + seen = self._latest_without_reload() + if seen is None or (seen.seq, seen.path, seen.policy_hash) != (known.seq, known.path, known.policy_hash): + self.last_latest_source = "other_writer" + return self.latest() + self.last_latest_source = "cached" + return known + + def _write_release(self, timer: PhaseTimer, latest: PolicyRelease | None, version: int, + policy_hash: str, write: Callable[[Path], None], *, weights: str, + buffer_selected: int | None, engine_checksum: str | None, + buffer_shards: dict[str, int] | None, derived_from: dict[str, Any] | None, + origin: dict[str, Any] | None) -> PolicyRelease: seq = 1 if latest is None else latest.seq + 1 name = f"r{seq}-v{version}" directory = self.root / name @@ -387,15 +441,20 @@ def release(self, version: int, policy_hash: str, write: Callable[[Path], None], {str(k): int(v) for k, v in sorted(buffer_shards.items())}, derived_from=None if derived_from is None else dict(derived_from), origin=None if origin is None else dict(origin)) - with timer.phase("commit_weights_s"): - _write_json_atomic(directory / RECORD, asdict(record)) - self.store.commit() # the weights and the record are durable before LATEST names them - with timer.phase("commit_latest_s"): - _write_json_atomic(self.root / LATEST, asdict(record)) - self.store.commit() + if self.cached_latest: # one commit (module doc) + with timer.phase("commit_s"): + _write_json_atomic(directory / RECORD, asdict(record)) + _write_json_atomic(self.root / LATEST, asdict(record)) + self.store.commit() + else: + with timer.phase("commit_weights_s"): + _write_json_atomic(directory / RECORD, asdict(record)) + self.store.commit() # the weights and the record are durable before LATEST names them + with timer.phase("commit_latest_s"): + _write_json_atomic(self.root / LATEST, asdict(record)) + self.store.commit() with timer.phase("prune_s"): self._prune() - self.last_timings = timer.as_dict() return record def _prune(self) -> list[str]: diff --git a/yeto/rl/engine/run_flags.py b/yeto/rl/engine/run_flags.py index 7711e3d5..a3605681 100644 --- a/yeto/rl/engine/run_flags.py +++ b/yeto/rl/engine/run_flags.py @@ -589,6 +589,10 @@ def rl_forward_pairs(*dests: str | tuple[str, str]) -> tuple[tuple[str, str], .. help="--rl-sample-role sampler with --tuning lora: check each release's digest " "while copying the adapter to local disk (one read of the store instead of " "two); default off"), + shared_flag("--rl-release-reload-latest", action="store_true", default=False, + help="--rl-sample-role trainer: reload the store and read LATEST before every " + "policy release and commit LATEST in a second commit (the behavior before " + "S26 E5); default: the trainer keeps the LATEST it wrote and commits once"), shared_flag("--rl-sample-wait-s", type=float, default=None, metavar="S", help="--rl-sample-role trainer: max seconds to wait for one full round from " "the buffer (default 3600)"), diff --git a/yeto/rl/engine/sample_roles.py b/yeto/rl/engine/sample_roles.py index d9fcf673..bd9c08fa 100644 --- a/yeto/rl/engine/sample_roles.py +++ b/yeto/rl/engine/sample_roles.py @@ -60,7 +60,10 @@ class SampleBufferConfig: ``min_round_s``: sampler, min seconds from one round start to the next (a throttle, so a test can make the sampler slower than the trainer); ``fused_stage``: LoRA sampler, check the release digest while copying the - adapter to local disk (``--rl-sampler-fused-stage``, S26 E5).""" + adapter to local disk (``--rl-sampler-fused-stage``, S26 E5); + ``release_legacy``: trainer or convertible node, reload the store before every + policy release and commit LATEST separately (``--rl-release-reload-latest``; + ``PolicyReleaseStore.cached_latest``).""" role: str store: str @@ -68,6 +71,7 @@ class SampleBufferConfig: wait_s: float | None = None min_round_s: float | None = None fused_stage: bool = False + release_legacy: bool = False def sample_role(args: Any) -> str: @@ -234,6 +238,8 @@ def island_args(args: Any, island: int, num_islands: int) -> Any: out.rl_sampler_min_round_seconds = None if out.rl_sample_role != SAMPLER: out.rl_sampler_fused_stage = False + if out.rl_sample_role not in (TRAINER, CONVERTIBLE): + out.rl_release_reload_latest = False return out @@ -328,6 +334,8 @@ def check_launch(args: Any) -> None: raise ValueError(f"{flag} needs --rl-sample-role sampler or trainer") if getattr(args, "rl_sampler_fused_stage", False): raise ValueError("--rl-sampler-fused-stage needs --rl-sample-role sampler") + if getattr(args, "rl_release_reload_latest", False): + raise ValueError("--rl-release-reload-latest needs --rl-sample-role trainer") return if not getattr(args, "rl_single_island_no_sync", False) and role != TRAINER: raise ValueError(f"--rl-sample-role {role} needs --rl-single-island-no-sync: a sampler " @@ -360,6 +368,9 @@ def check_launch(args: Any) -> None: role != SAMPLER or (getattr(args, "tuning", None) or "lora") != "lora"): raise ValueError("--rl-sampler-fused-stage needs --rl-sample-role sampler and " "--tuning lora (it fuses the LoRA adapter copy with its digest)") + if getattr(args, "rl_release_reload_latest", False) and role not in (TRAINER, CONVERTIBLE): + raise ValueError("--rl-release-reload-latest needs --rl-sample-role trainer " + "(or a convertible pool): only a trainer writes policy releases") wait_s = getattr(args, "rl_sample_wait_s", None) if wait_s is not None and (role not in (TRAINER, CONVERTIBLE) or not float(wait_s) > 0): raise ValueError("--rl-sample-wait-s needs --rl-sample-role trainer and a value > 0") @@ -402,7 +413,9 @@ def resolve(args: Any, *, groups_per_round: int) -> SampleBufferConfig | None: wait_s=None if wait_s is None else float(wait_s), min_round_s=(float(min_round_s) if role == SAMPLER and min_round_s is not None else None), - fused_stage=role == SAMPLER and bool(getattr(args, "rl_sampler_fused_stage", False))) + fused_stage=role == SAMPLER and bool(getattr(args, "rl_sampler_fused_stage", False)), + release_legacy=role in (TRAINER, CONVERTIBLE) + and bool(getattr(args, "rl_release_reload_latest", False))) def island_env(args: Any) -> dict[str, str]: @@ -434,6 +447,8 @@ def island_flags(args: Any) -> str: flags += f" --rl-sample-wait-s {float(args.rl_sample_wait_s)!r}" if role == SAMPLER and getattr(args, "rl_sampler_fused_stage", False): flags += " --rl-sampler-fused-stage" + if role in (TRAINER, CONVERTIBLE) and getattr(args, "rl_release_reload_latest", False): + flags += " --rl-release-reload-latest" if getattr(args, "rl_sampler_node", None): flags += f" --rl-sampler-node {args.rl_sampler_node}" if getattr(args, "rl_buffer_shards", None):