From e81cf39edcd1a856fdabe349250a1509d582157a Mon Sep 17 00:00:00 2001 From: michael <78113554+chowhsu@users.noreply.github.com> Date: Tue, 20 Oct 2026 02:23:27 +0000 Subject: [PATCH 01/89] Read verl TransferQueue provenance from the extra_fields record verl's TQ agent loop stores AgentLoopOutput.as_dict(), so provenance sits inside each row's extra_fields, not under its own name; the per-source metrics of the TQ path were all 'unknown'. Also resolve the train handles of a bare TrainingWorker (worker.engine). Co-Authored-By: Claude Opus 5.5 --- tests/test_rl_n2_trainer_wiring.py | 29 +++++++++++++++++++++++++++ yeto/rl/adapters/verl/ports_impl.py | 29 ++++++++++++++++++++++----- yeto/rl/adapters/verl/resume_check.py | 5 +++-- 3 files changed, 56 insertions(+), 7 deletions(-) diff --git a/tests/test_rl_n2_trainer_wiring.py b/tests/test_rl_n2_trainer_wiring.py index be3b8a907..b7a7af90d 100644 --- a/tests/test_rl_n2_trainer_wiring.py +++ b/tests/test_rl_n2_trainer_wiring.py @@ -116,3 +116,32 @@ def test_masked_source_samples_checks_row_counts(): masked_source_samples([[0.0]], [[0.0]], [[1], [1]]) with pytest.raises(ValueError, match="provenances"): masked_source_samples([[0.0]], [[0.0]], [[1]], provenances=[None, None]) + + +def test_verl_tq_provenance_is_read_from_the_extra_fields_record(monkeypatch): + """verl's TQ agent loop stores ``AgentLoopOutput.as_dict()``: provenance sits in + ``extra_fields``, not under its own name (S25 T-G4).""" + import sys + import types + + from yeto.rl.adapters.verl import ports_impl + + class Wrapped: # a TransferQueue read can wrap each record (NonTensorData) + def __init__(self, data): + self.data = data + + seen = {} + + def kv_batch_get(*, keys, partition_id, select_fields): + seen["fields"] = select_fields + return {"extra_fields": [Wrapped({"provenance": SOURCE}), {"other": 1}]} + + monkeypatch.setitem(sys.modules, "transfer_queue", types.SimpleNamespace(kv_batch_get=kv_batch_get)) + batch = types.SimpleNamespace(keys=["a", "b"], partition_id="train") + assert ports_impl._row_field(batch, "provenance") == ([SOURCE, None], "ok") + assert seen["fields"] == ["extra_fields"] + + monkeypatch.setitem(sys.modules, "transfer_queue", types.SimpleNamespace( + kv_batch_get=lambda **_: {"extra_fields": [{}, {}]})) + rows, read = ports_impl._row_field(batch, "provenance") + assert rows is None and read.startswith("missing") diff --git a/yeto/rl/adapters/verl/ports_impl.py b/yeto/rl/adapters/verl/ports_impl.py index 6bfaf3b0d..2a44c1719 100644 --- a/yeto/rl/adapters/verl/ports_impl.py +++ b/yeto/rl/adapters/verl/ports_impl.py @@ -522,16 +522,35 @@ def criteria_from_trainer_batch(isl: VerlIsland, kv_batch) -> None: def _row_field(kv_batch, name: str) -> tuple[list | None, str]: - """Per-row values of a non-tensor field (``HarborAgentLoop`` puts ``provenance`` - in ``extra_fields``); ``(None, reason)`` when the batch has none. That verl's - TransferQueue keeps ``extra_fields`` under their own names is not verified.""" + """Per-row values of one ``extra_fields`` entry (``HarborAgentLoop`` puts + ``provenance`` there); ``(None, reason)`` when the batch has none. + + verl's TransferQueue agent loop (``agent_loop_tq._agent_loop_postprocess``) + stores ``AgentLoopOutput.as_dict()``, so the entry sits inside the + ``extra_fields`` record of each row, not under its own name (the classic + agent loop flattens it into ``non_tensor_batch``; the fully async runner + reads that). The batch is read by the record name.""" import transfer_queue as tq try: - data = tq.kv_batch_get(keys=kv_batch.keys, partition_id=kv_batch.partition_id, select_fields=[name]) - return list(data[name]), "ok" + data = tq.kv_batch_get(keys=kv_batch.keys, partition_id=kv_batch.partition_id, + select_fields=["extra_fields"]) + rows = extra_field_values(list(data["extra_fields"]), name) except Exception as exc: # noqa: BLE001 - the rows then group under ``unknown`` return None, f"missing: {type(exc).__name__}: {str(exc)[:200]}" + if all(v is None for v in rows): + return None, f"missing: no {name!r} in extra_fields" + return rows, "ok" + + +def extra_field_values(records: list, name: str) -> list: + """``record[name]`` of each row's ``extra_fields`` (``None`` where absent). A + TransferQueue read can return each record wrapped in a ``NonTensorData``.""" + values = [] + for record in records: + record = getattr(record, "data", record) + values.append(record.get(name) if isinstance(record, dict) else None) + return values def source_event(train_lp, infer_lp, mask, provenances, provenance_read: str, *, rollout_id: int, diff --git a/yeto/rl/adapters/verl/resume_check.py b/yeto/rl/adapters/verl/resume_check.py index f4f008f05..69b521701 100644 --- a/yeto/rl/adapters/verl/resume_check.py +++ b/yeto/rl/adapters/verl/resume_check.py @@ -4,7 +4,7 @@ a verl FSDP actor worker onto them (:func:`train_handles`). Two worker shapes: * the engine worker of the pinned verl (``engine_workers.ActorRolloutRefWorker``): - ``worker.actor.engine`` holds ``module`` / ``optimizer`` / ``lr_scheduler`` + ``worker.actor.engine`` (a bare ``TrainingWorker``: ``worker.engine``) holds ``module`` / ``optimizer`` / ``lr_scheduler`` (``workers/engine/fsdp/transformer_impl.py``); * the legacy fsdp worker: ``actor_module_fsdp`` / ``actor_optimizer`` / ``actor_lr_scheduler``. @@ -55,7 +55,8 @@ def _engine_of(worker: Any) -> Any: inner = getattr(worker, "worker_dict", None) # ray fused worker: {role: worker} if inner: worker = next((w for w in inner.values() if hasattr(w, "actor")), worker) - return getattr(getattr(worker, "actor", None), "engine", None) or worker + actor = getattr(worker, "actor", None) or worker # ActorRolloutRefWorker.actor is a TrainingWorker + return getattr(actor, "engine", None) or worker def train_handles(worker: Any) -> TrainHandles: From 480f96fb6762aa983b00184d46b8ce0caecd3f4f Mon Sep 17 00:00:00 2001 From: michael <78113554+chowhsu@users.noreply.github.com> Date: Tue, 20 Oct 2026 02:38:50 +0000 Subject: [PATCH 02/89] Map Miles convert_to_hf names to checkpoint names in quant hf_slices; add the T-G3 GPU harness convert_to_hf emits SGLang weight-update names (model.layers.*), while the GLM-5.3 checkpoint and the quantization layout keep the text model under model.language_model.; hf_slices now renames with checkpoint_name (same rule as Miles Glm5NextBridge._nest_language_model). tools/probes/g3_miles: Modal H100!:2 harness for N7.5 (TE vs local forward/backward), N8.2 (in-trainer sharded requantization vs single process) and N7.1d (distributed optimizer + CPU offload cut round trip). Co-Authored-By: Claude Opus 5.5 --- tests/test_rl_quant_shards.py | 14 +- tools/probes/g3_miles/compare.py | 135 ++ tools/probes/g3_miles/glm53flash_config.json | 1810 ++++++++++++++++++ tools/probes/g3_miles/harness.py | 716 +++++++ tools/probes/g3_miles/modal_g3.py | 111 ++ yeto/rl/adapters/miles/quant_shards.py | 23 +- 6 files changed, 2803 insertions(+), 6 deletions(-) create mode 100644 tools/probes/g3_miles/compare.py create mode 100644 tools/probes/g3_miles/glm53flash_config.json create mode 100644 tools/probes/g3_miles/harness.py create mode 100644 tools/probes/g3_miles/modal_g3.py diff --git a/tests/test_rl_quant_shards.py b/tests/test_rl_quant_shards.py index 0350ab94f..8868595dc 100644 --- a/tests/test_rl_quant_shards.py +++ b/tests/test_rl_quant_shards.py @@ -23,7 +23,7 @@ from tests import test_rl_quant_mxfp4 as mx from tests import test_rl_quant_requant as w4 -from yeto.rl.adapters.miles.quant_shards import hf_slices, megatron_expert_slices +from yeto.rl.adapters.miles.quant_shards import checkpoint_name, hf_slices, megatron_expert_slices from yeto.rl.adapters.verl.quant_shards import fsdp2_slices, shard0_rows from yeto.rl.engine.policy_release import ( RELEASE_DIR, @@ -225,7 +225,9 @@ def _megatron_rank(src: dict[str, torch.Tensor], ep_rank: int, num_local: int, s def test_megatron_expert_parallel_ranks_merge_to_the_same_bytes(world): src, out = world["src"], world["tmp"] / "megatron" parts = out / PART_DIR - non_expert = [(n, t) for n, t in src.items() if not _PER_EXPERT.match(n)] + # non-expert weights as Miles convert_to_hf emits them (SGLang names, no model.language_model.) + non_expert = [(n.replace("model.language_model.", "model.", 1), t) for n, t in src.items() if not _PER_EXPERT.match(n)] + assert any(n.startswith("model.layers.") for n, _ in non_expert) for ep_rank in (0, 1): # E = 2: one local expert per rank; rank 1 uses the sequential-MLP names slices = list(megatron_expert_slices(_megatron_rank(src, ep_rank, 1, sequential=ep_rank == 1), ep_rank=ep_rank, num_local_experts=1)) @@ -380,3 +382,11 @@ def test_follower_refuses_a_release_of_another_scheme_before_reading_it(tmp_path bf16.release(4, "h4", lambda d: (d / "model.safetensors").write_text("y")) with pytest.raises(SchemeMismatch, match="'bf16'"): PolicyFollower(bf16, lambda r, d: None, scheme="w4a8-int4-channel").refresh() + + +def test_miles_convert_to_hf_names_map_to_checkpoint_names(): + assert checkpoint_name("model.layers.3.mlp.gate.weight") == "model.language_model.layers.3.mlp.gate.weight" + assert checkpoint_name("model.embed_tokens.weight") == "model.language_model.embed_tokens.weight" + assert checkpoint_name("model.norm.weight") == "model.language_model.norm.weight" + for unchanged in ("lm_head.weight", "model.language_model.norm.weight", "model.visual.blocks.0.attn.qkv.weight"): + assert checkpoint_name(unchanged) == unchanged diff --git a/tools/probes/g3_miles/compare.py b/tools/probes/g3_miles/compare.py new file mode 100644 index 000000000..85a024805 --- /dev/null +++ b/tools/probes/g3_miles/compare.py @@ -0,0 +1,135 @@ +"""Judge the T-G3 harness outputs (CPU only). usage: python compare.py + +Writes ``/verdict.json`` and prints it. Thresholds are pre-registered in +SESSION25/gpu-review/T-G3.md; the noise floor (TE vs TE repeat, cont5 vs cont5b) +is reported next to every TE-vs-local and resume-vs-continuous difference. +""" + +from __future__ import annotations + +import json +import math +import sys +from pathlib import Path + +import torch + +sys.path.insert(0, str(Path(__file__).resolve().parent)) +from harness import map_local_to_te # noqa: E402 + +LP_MEAN_MAX, LP_ABS_MAX = 5e-3, 0.1 # label log-prob |delta|, nats +GRAD_REL_MAX, GRAD_GLOBAL_REL_MAX = 0.05, 0.02 # per-tensor / all-tensor relative L2 +LOSS_ABS_MAX = 1e-3 + + +def lp_stats(a: dict, b: dict) -> dict: + d = (a["token_loss"][:-1] - b["token_loss"][:-1]).abs().double() + return {"mean": float(d.mean()), "p99": float(d.quantile(0.99)), "max": float(d.max()), + "loss_a": a["loss"], "loss_b": b["loss"], "loss_abs": abs(a["loss"] - b["loss"]), + "logits_max_abs": float((a["logits"] - b["logits"]).abs().max())} + + +def grad_stats(out: Path, tag_a: str, tag_b: str, local_b: bool) -> dict: + per, worst, num, den, ranks = [], None, 0.0, 0.0, 0 + for r in range(8): + fa, fb = out / f"grads-{tag_a}.rank{r}.pt", out / f"grads-{tag_b}.rank{r}.pt" + if not fa.exists(): + break + ranks += 1 + ga = torch.load(fa, weights_only=True) + gb = torch.load(fb, weights_only=True) + if local_b: + mapping, missing, unused = map_local_to_te(gb, set(ga)) + if missing or unused: + return {"error": f"rank{r}: names do not map: missing={missing[:5]} unused={unused[:5]}"} + gb = {mapping[n]: t for n, t in gb.items()} + if set(ga) != set(gb): + return {"error": f"rank{r}: grad name sets differ: {sorted(set(ga) ^ set(gb))[:6]}"} + for n in sorted(ga): + a, b = ga[n].double(), gb[n].double() + na = float(a.norm()) + diff = float((a - b).norm()) + rel = diff / na if na > 0 else (0.0 if diff == 0 else math.inf) + cos = float((a * b).sum() / (a.norm() * b.norm())) if na > 0 and float(b.norm()) > 0 else None + num += diff ** 2 + den += na ** 2 + item = {"rank": r, "name": n, "rel": rel, "cos": cos, "norm": na} + per.append(item) + if worst is None or rel > worst["rel"]: + worst = item + per.sort(key=lambda x: -x["rel"]) + return {"ranks": ranks, "tensors": len(per), "global_rel": math.sqrt(num / den) if den else None, + "worst": per[:8], "over_threshold": [p for p in per if p["rel"] > GRAD_REL_MAX and p["norm"] > 0]} + + +def n75(out: Path) -> dict: + v: dict = {} + load = lambda t: torch.load(out / f"logits-{t}.pt", weights_only=True) # noqa: E731 + have = {t for t in ("te", "te2", "local") if (out / f"logits-{t}.pt").exists()} + if {"te", "te2"} <= have: + v["noise_te_vs_te2"] = {"logprob": lp_stats(load("te"), load("te2")), "grads": grad_stats(out, "te", "te2", False)} + if {"te", "local"} <= have: + lp = lp_stats(load("te"), load("local")) + g = grad_stats(out, "te", "local", True) + ok = (lp["mean"] <= LP_MEAN_MAX and lp["max"] <= LP_ABS_MAX and lp["loss_abs"] <= LOSS_ABS_MAX + and "error" not in g and not g["over_threshold"] and (g["global_rel"] or 0) <= GRAD_GLOBAL_REL_MAX) + v["te_vs_local"] = {"logprob": lp, "grads": g, "pass": ok} + for t in sorted(have): + info = json.loads((out / f"info-{t}.json").read_text()) + v[f"info_{t}"] = [{k: i.get(k) for k in ("impl", "transformer_impl", "dsa_core_attention", "loss", "timing", + "name_map", "num_params_local", "gpu")} for i in info] + v[f"n82_{t}"] = [i.get("n82") for i in info[:1]] + return v + + +def n71d(out: Path) -> dict: + load = lambda m: json.loads((out / f"n71d-{m}.json").read_text()) if (out / f"n71d-{m}.json").exists() else None # noqa: E731 + c5, c5b, r3, rs = load("cont5"), load("cont5b"), load("run3"), load("resume2") + v: dict = {"present": [m for m, x in (("cont5", c5), ("cont5b", c5b), ("run3", r3), ("resume2", rs)) if x]} + if c5: + v["flags"] = c5[0]["flags"] + v["hdo_cont5_final"] = c5[0]["hdo_final"] + if c5 and c5b: + v["noise_cont5_vs_cont5b"] = { + "params_bitwise_equal": all(a["final_params"]["sha256"] == b["final_params"]["sha256"] for a, b in zip(c5, c5b)), + "loss_abs": [abs(a["loss_mean_dp"] - b["loss_mean_dp"]) for a, b in zip(c5[0]["steps"], c5b[0]["steps"])], + "grad_norm_abs": [abs(a["grad_norm"] - b["grad_norm"]) for a, b in zip(c5[0]["steps"], c5b[0]["steps"])]} + if c5 and r3: + v["run3_vs_cont5_steps1to3_loss_abs"] = [abs(a["loss_mean_dp"] - b["loss_mean_dp"]) + for a, b in zip(r3[0]["steps"], c5[0]["steps"][:3])] + if rs: + r0 = rs + v["restore"] = { + "digest_equal_per_rank": [x["restore_digest_equal"] for x in r0], + "restore_problems": r0[0]["restore_problems"], "first_step_problems": r0[0]["first_step_problems"], + "hdo_steps_after_restore": [[l.get("steps") for l in x["hdo_after_restore"]["leaves"]] for x in r0], + "hdo_steps_after_first_step": [[l.get("steps") for l in x["hdo_after_first_step"]["leaves"]] for x in r0], + "diff": [x["restore_summary"].get("diff") for x in r0]} + if c5: + steps = {s["step"]: s for s in c5[0]["steps"]} + v["resume_vs_cont5"] = { + "params_bitwise_equal": all(a["final_params"]["sha256"] == b["final_params"]["sha256"] + for a, b in zip(rs, c5)), + "loss_abs": {s["step"]: abs(s["loss_mean_dp"] - steps[s["step"]]["loss_mean_dp"]) for s in rs[0]["steps"]}, + "grad_norm_abs": {s["step"]: abs(s["grad_norm"] - steps[s["step"]]["grad_norm"]) for s in rs[0]["steps"]}, + "lr": {s["step"]: (s["lr"], steps[s["step"]]["lr"]) for s in rs[0]["steps"]}, + "param_norm_max_rel": max( + abs(a["final_params"]["norms"][n] - b["final_params"]["norms"][n]) / max(b["final_params"]["norms"][n], 1e-30) + for a, b in zip(rs, c5) for n in b["final_params"]["norms"])} + return v + + +def main() -> None: + out = Path(sys.argv[1]) + verdict = {} + for name, fn, sub in (("n75", n75, "n75"), ("n71d", n71d, "n71d")): + try: + verdict[name] = fn(out / sub) + except Exception as exc: # noqa: BLE001 + verdict[name] = {"error": f"{type(exc).__name__}: {exc}"} + (out / "verdict.json").write_text(json.dumps(verdict, indent=1, default=str)) + print(json.dumps(verdict, indent=1, default=str)[:20000]) + + +if __name__ == "__main__": + main() diff --git a/tools/probes/g3_miles/glm53flash_config.json b/tools/probes/g3_miles/glm53flash_config.json new file mode 100644 index 000000000..f93128cf7 --- /dev/null +++ b/tools/probes/g3_miles/glm53flash_config.json @@ -0,0 +1,1810 @@ +{ + "architectures": [ + "Glm5NextForConditionalGeneration" + ], + "text_config": { + "attention_bias": false, + "attention_dropout": 0.0, + "dtype": "bfloat16", + "eos_token_id": [ + 154820, + 154827, + 154829 + ], + "first_k_dense_replace": 3, + "hc_eps": 1e-06, + "hc_mult": 4, + "hc_sinkhorn_iters": 20, + "head_dim": 0, + "hidden_act": "silu", + "hidden_size": 4096, + "index_head_dim": 128, + "index_kpool": 4, + "index_kpool_always_select_tail": true, + "index_kpool_compress": true, + "index_n_heads": 32, + "index_topk": 2048, + "index_share_for_mtp_iteration": true, + "indexer_rope_interleave": true, + "indexer_types": [ + "full", + "full", + "full", + "full", + "full", + "full", + "full", + "full", + "full", + "full", + "full", + "full", + "full", + "full", + "full", + "full", + "full", + "full", + "full", + "full", + "full", + "full", + "full", + "full", + "full", + "full", + "full", + "full", + "full", + "full", + "full", + "full", + "full", + "full", + "full", + "full", + "full", + "full", + "full", + "full", + "full", + "full", + "full", + "full", + "full" + ], + "initializer_range": 0.02, + "intermediate_size": 12288, + "kv_lora_rank": 512, + "layer_types": [ + "linear_attention", + "linear_attention", + "linear_attention", + "deepseek_sparse_attention", + "linear_attention", + "linear_attention", + "linear_attention", + "deepseek_sparse_attention", + "linear_attention", + "linear_attention", + "linear_attention", + "deepseek_sparse_attention", + "linear_attention", + "linear_attention", + "linear_attention", + "deepseek_sparse_attention", + "linear_attention", + "linear_attention", + "linear_attention", + "deepseek_sparse_attention", + "linear_attention", + "linear_attention", + "linear_attention", + "deepseek_sparse_attention", + "linear_attention", + "linear_attention", + "linear_attention", + "deepseek_sparse_attention", + "linear_attention", + "linear_attention", + "linear_attention", + "deepseek_sparse_attention", + "linear_attention", + "linear_attention", + "linear_attention", + "deepseek_sparse_attention", + "linear_attention", + "linear_attention", + "linear_attention", + "deepseek_sparse_attention", + "linear_attention", + "linear_attention", + "linear_attention", + "deepseek_sparse_attention", + "linear_attention" + ], + "linear_attn_config": { + "num_heads": 64, + "gate_lower_bound": -5.0, + "head_dim": 128, + "short_conv_kernel_size": 4, + "kda_layers": [ + 0, + 1, + 2, + 4, + 5, + 6, + 8, + 9, + 10, + 12, + 13, + 14, + 16, + 17, + 18, + 20, + 21, + 22, + 24, + 25, + 26, + 28, + 29, + 30, + 32, + 33, + 34, + 36, + 37, + 38, + 40, + 41, + 42, + 44 + ], + "full_attn_layers": [ + 3, + 7, + 11, + 15, + 19, + 23, + 27, + 31, + 35, + 39, + 43 + ] + }, + "max_position_embeddings": 1048576, + "mhc": true, + "mla_use_nope": true, + "mlp_layer_types": [ + "dense", + "dense", + "dense", + "sparse", + "sparse", + "sparse", + "sparse", + "sparse", + "sparse", + "sparse", + "sparse", + "sparse", + "sparse", + "sparse", + "sparse", + "sparse", + "sparse", + "sparse", + "sparse", + "sparse", + "sparse", + "sparse", + "sparse", + "sparse", + "sparse", + "sparse", + "sparse", + "sparse", + "sparse", + "sparse", + "sparse", + "sparse", + "sparse", + "sparse", + "sparse", + "sparse", + "sparse", + "sparse", + "sparse", + "sparse", + "sparse", + "sparse", + "sparse", + "sparse", + "sparse" + ], + "model_type": "glm5_next_text", + "moe_intermediate_size": 2048, + "moe_router_dtype": "float32", + "n_group": 1, + "n_routed_experts": 288, + "n_shared_experts": 1, + "norm_topk_prob": true, + "num_attention_heads": 64, + "num_experts_per_tok": 8, + "num_hidden_layers": 45, + "num_key_value_heads": 64, + "num_nextn_predict_layers": 1, + "output_router_logits": false, + "pad_token_id": 154820, + "q_lora_rank": 1536, + "qk_head_dim": 256, + "qk_nope_head_dim": 256, + "qk_rope_head_dim": 0, + "rms_norm_eps": 1e-05, + "routed_scaling_factor": 2.5, + "router_aux_loss_coef": 0.001, + "scoring_func": "sigmoid", + "swiglu_limit": 10.0, + "tie_word_embeddings": false, + "topk_group": 1, + "topk_method": "noaux_tc", + "use_cache": true, + "v_head_dim": 256, + "vocab_size": 154880 + }, + "vision_config": { + "attention_bias": true, + "attention_dropout": 0.0, + "depth": 24, + "hidden_act": "silu", + "hidden_size": 1024, + "image_size": 448, + "in_channels": 3, + "initializer_range": 0.02, + "intermediate_size": 4096, + "model_type": "glm5_next_vision", + "num_heads": 16, + "out_hidden_size": 4096, + "patch_size": 14, + "projection_intermediate_size": 10240, + "rms_norm_eps": 1e-05, + "spatial_merge_size": 2, + "swiglu_limit": 10.0, + "temporal_patch_size": 2 + }, + "image_token_id": 154854, + "video_token_id": 154855, + "image_start_token_id": 154830, + "image_end_token_id": 154831, + "video_start_token_id": 154832, + "video_end_token_id": 154833, + "tie_word_embeddings": false, + "model_type": "glm5_next", + "transformers_version": "5.16.0", + "quantization_config": { + "activation_scheme": "dynamic", + "fmt": "e4m3", + "modules_to_not_convert": [ + "attn_mha", + "attn_mqa", + "dt_bias", + "hyper_connection", + "lm_head", + "mapping_proj", + "model.embed_tokens", + "model.layers.0.hc_attn_base", + "model.layers.0.hc_attn_fn", + "model.layers.0.hc_attn_scale", + "model.layers.0.hc_ffn_base", + "model.layers.0.hc_ffn_fn", + "model.layers.0.hc_ffn_scale", + "model.layers.0.input_layernorm", + "model.layers.0.post_attention_layernorm", + "model.layers.0.self_attn.A_log", + "model.layers.0.self_attn.b_proj", + "model.layers.0.self_attn.dt_bias", + "model.layers.0.self_attn.f_a_proj", + "model.layers.0.self_attn.f_b_proj", + "model.layers.0.self_attn.fused_qkvbfg_a_proj", + "model.layers.0.self_attn.g_a_proj", + "model.layers.0.self_attn.g_b_proj", + "model.layers.0.self_attn.k_conv1d", + "model.layers.0.self_attn.k_proj", + "model.layers.0.self_attn.o_norm", + "model.layers.0.self_attn.o_proj", + "model.layers.0.self_attn.q_conv1d", + "model.layers.0.self_attn.q_proj", + "model.layers.0.self_attn.qkv_proj", + "model.layers.0.self_attn.v_conv1d", + "model.layers.0.self_attn.v_proj", + "model.layers.1.hc_attn_base", + "model.layers.1.hc_attn_fn", + "model.layers.1.hc_attn_scale", + "model.layers.1.hc_ffn_base", + "model.layers.1.hc_ffn_fn", + "model.layers.1.hc_ffn_scale", + "model.layers.1.input_layernorm", + "model.layers.1.post_attention_layernorm", + "model.layers.1.self_attn.A_log", + "model.layers.1.self_attn.b_proj", + "model.layers.1.self_attn.dt_bias", + "model.layers.1.self_attn.f_a_proj", + "model.layers.1.self_attn.f_b_proj", + "model.layers.1.self_attn.fused_qkvbfg_a_proj", + "model.layers.1.self_attn.g_a_proj", + "model.layers.1.self_attn.g_b_proj", + "model.layers.1.self_attn.k_conv1d", + "model.layers.1.self_attn.k_proj", + "model.layers.1.self_attn.o_norm", + "model.layers.1.self_attn.o_proj", + "model.layers.1.self_attn.q_conv1d", + "model.layers.1.self_attn.q_proj", + "model.layers.1.self_attn.qkv_proj", + "model.layers.1.self_attn.v_conv1d", + "model.layers.1.self_attn.v_proj", + "model.layers.10.hc_attn_base", + "model.layers.10.hc_attn_fn", + "model.layers.10.hc_attn_scale", + "model.layers.10.hc_ffn_base", + "model.layers.10.hc_ffn_fn", + "model.layers.10.hc_ffn_scale", + "model.layers.10.input_layernorm", + "model.layers.10.mlp.gate", + "model.layers.10.mlp.gate.e_score_correction_bias", + "model.layers.10.post_attention_layernorm", + "model.layers.10.self_attn.A_log", + "model.layers.10.self_attn.b_proj", + "model.layers.10.self_attn.dt_bias", + "model.layers.10.self_attn.f_a_proj", + "model.layers.10.self_attn.f_b_proj", + "model.layers.10.self_attn.fused_qkvbfg_a_proj", + "model.layers.10.self_attn.g_a_proj", + "model.layers.10.self_attn.g_b_proj", + "model.layers.10.self_attn.k_conv1d", + "model.layers.10.self_attn.k_proj", + "model.layers.10.self_attn.o_norm", + "model.layers.10.self_attn.o_proj", + "model.layers.10.self_attn.q_conv1d", + "model.layers.10.self_attn.q_proj", + "model.layers.10.self_attn.qkv_proj", + "model.layers.10.self_attn.v_conv1d", + "model.layers.10.self_attn.v_proj", + "model.layers.11.hc_attn_base", + "model.layers.11.hc_attn_fn", + "model.layers.11.hc_attn_scale", + "model.layers.11.hc_ffn_base", + "model.layers.11.hc_ffn_fn", + "model.layers.11.hc_ffn_scale", + "model.layers.11.input_layernorm", + "model.layers.11.mlp.gate", + "model.layers.11.mlp.gate.e_score_correction_bias", + "model.layers.11.post_attention_layernorm", + "model.layers.11.self_attn.indexer.index_kpool_compress_ape", + "model.layers.11.self_attn.indexer.index_kpool_compress_gate", + "model.layers.11.self_attn.indexer.k_norm", + "model.layers.11.self_attn.indexer.k_norm.bias", + "model.layers.11.self_attn.indexer.weights_proj", + "model.layers.11.self_attn.indexer.wk", + "model.layers.11.self_attn.indexer.wq_b", + "model.layers.11.self_attn.kv_a_layernorm", + "model.layers.11.self_attn.kv_b_proj", + "model.layers.11.self_attn.q_a_layernorm", + "model.layers.12.hc_attn_base", + "model.layers.12.hc_attn_fn", + "model.layers.12.hc_attn_scale", + "model.layers.12.hc_ffn_base", + "model.layers.12.hc_ffn_fn", + "model.layers.12.hc_ffn_scale", + "model.layers.12.input_layernorm", + "model.layers.12.mlp.gate", + "model.layers.12.mlp.gate.e_score_correction_bias", + "model.layers.12.post_attention_layernorm", + "model.layers.12.self_attn.A_log", + "model.layers.12.self_attn.b_proj", + "model.layers.12.self_attn.dt_bias", + "model.layers.12.self_attn.f_a_proj", + "model.layers.12.self_attn.f_b_proj", + "model.layers.12.self_attn.fused_qkvbfg_a_proj", + "model.layers.12.self_attn.g_a_proj", + "model.layers.12.self_attn.g_b_proj", + "model.layers.12.self_attn.k_conv1d", + "model.layers.12.self_attn.k_proj", + "model.layers.12.self_attn.o_norm", + "model.layers.12.self_attn.o_proj", + "model.layers.12.self_attn.q_conv1d", + "model.layers.12.self_attn.q_proj", + "model.layers.12.self_attn.qkv_proj", + "model.layers.12.self_attn.v_conv1d", + "model.layers.12.self_attn.v_proj", + "model.layers.13.hc_attn_base", + "model.layers.13.hc_attn_fn", + "model.layers.13.hc_attn_scale", + "model.layers.13.hc_ffn_base", + "model.layers.13.hc_ffn_fn", + "model.layers.13.hc_ffn_scale", + "model.layers.13.input_layernorm", + "model.layers.13.mlp.gate", + "model.layers.13.mlp.gate.e_score_correction_bias", + "model.layers.13.post_attention_layernorm", + "model.layers.13.self_attn.A_log", + "model.layers.13.self_attn.b_proj", + "model.layers.13.self_attn.dt_bias", + "model.layers.13.self_attn.f_a_proj", + "model.layers.13.self_attn.f_b_proj", + "model.layers.13.self_attn.fused_qkvbfg_a_proj", + "model.layers.13.self_attn.g_a_proj", + "model.layers.13.self_attn.g_b_proj", + "model.layers.13.self_attn.k_conv1d", + "model.layers.13.self_attn.k_proj", + "model.layers.13.self_attn.o_norm", + "model.layers.13.self_attn.o_proj", + "model.layers.13.self_attn.q_conv1d", + "model.layers.13.self_attn.q_proj", + "model.layers.13.self_attn.qkv_proj", + "model.layers.13.self_attn.v_conv1d", + "model.layers.13.self_attn.v_proj", + "model.layers.14.hc_attn_base", + "model.layers.14.hc_attn_fn", + "model.layers.14.hc_attn_scale", + "model.layers.14.hc_ffn_base", + "model.layers.14.hc_ffn_fn", + "model.layers.14.hc_ffn_scale", + "model.layers.14.input_layernorm", + "model.layers.14.mlp.gate", + "model.layers.14.mlp.gate.e_score_correction_bias", + "model.layers.14.post_attention_layernorm", + "model.layers.14.self_attn.A_log", + "model.layers.14.self_attn.b_proj", + "model.layers.14.self_attn.dt_bias", + "model.layers.14.self_attn.f_a_proj", + "model.layers.14.self_attn.f_b_proj", + "model.layers.14.self_attn.fused_qkvbfg_a_proj", + "model.layers.14.self_attn.g_a_proj", + "model.layers.14.self_attn.g_b_proj", + "model.layers.14.self_attn.k_conv1d", + "model.layers.14.self_attn.k_proj", + "model.layers.14.self_attn.o_norm", + "model.layers.14.self_attn.o_proj", + "model.layers.14.self_attn.q_conv1d", + "model.layers.14.self_attn.q_proj", + "model.layers.14.self_attn.qkv_proj", + "model.layers.14.self_attn.v_conv1d", + "model.layers.14.self_attn.v_proj", + "model.layers.15.hc_attn_base", + "model.layers.15.hc_attn_fn", + "model.layers.15.hc_attn_scale", + "model.layers.15.hc_ffn_base", + "model.layers.15.hc_ffn_fn", + "model.layers.15.hc_ffn_scale", + "model.layers.15.input_layernorm", + "model.layers.15.mlp.gate", + "model.layers.15.mlp.gate.e_score_correction_bias", + "model.layers.15.post_attention_layernorm", + "model.layers.15.self_attn.indexer.index_kpool_compress_ape", + "model.layers.15.self_attn.indexer.index_kpool_compress_gate", + "model.layers.15.self_attn.indexer.k_norm", + "model.layers.15.self_attn.indexer.k_norm.bias", + "model.layers.15.self_attn.indexer.weights_proj", + "model.layers.15.self_attn.indexer.wk", + "model.layers.15.self_attn.indexer.wq_b", + "model.layers.15.self_attn.kv_a_layernorm", + "model.layers.15.self_attn.kv_b_proj", + "model.layers.15.self_attn.q_a_layernorm", + "model.layers.16.hc_attn_base", + "model.layers.16.hc_attn_fn", + "model.layers.16.hc_attn_scale", + "model.layers.16.hc_ffn_base", + "model.layers.16.hc_ffn_fn", + "model.layers.16.hc_ffn_scale", + "model.layers.16.input_layernorm", + "model.layers.16.mlp.gate", + "model.layers.16.mlp.gate.e_score_correction_bias", + "model.layers.16.post_attention_layernorm", + "model.layers.16.self_attn.A_log", + "model.layers.16.self_attn.b_proj", + "model.layers.16.self_attn.dt_bias", + "model.layers.16.self_attn.f_a_proj", + "model.layers.16.self_attn.f_b_proj", + "model.layers.16.self_attn.fused_qkvbfg_a_proj", + "model.layers.16.self_attn.g_a_proj", + "model.layers.16.self_attn.g_b_proj", + "model.layers.16.self_attn.k_conv1d", + "model.layers.16.self_attn.k_proj", + "model.layers.16.self_attn.o_norm", + "model.layers.16.self_attn.o_proj", + "model.layers.16.self_attn.q_conv1d", + "model.layers.16.self_attn.q_proj", + "model.layers.16.self_attn.qkv_proj", + "model.layers.16.self_attn.v_conv1d", + "model.layers.16.self_attn.v_proj", + "model.layers.17.hc_attn_base", + "model.layers.17.hc_attn_fn", + "model.layers.17.hc_attn_scale", + "model.layers.17.hc_ffn_base", + "model.layers.17.hc_ffn_fn", + "model.layers.17.hc_ffn_scale", + "model.layers.17.input_layernorm", + "model.layers.17.mlp.gate", + "model.layers.17.mlp.gate.e_score_correction_bias", + "model.layers.17.post_attention_layernorm", + "model.layers.17.self_attn.A_log", + "model.layers.17.self_attn.b_proj", + "model.layers.17.self_attn.dt_bias", + "model.layers.17.self_attn.f_a_proj", + "model.layers.17.self_attn.f_b_proj", + "model.layers.17.self_attn.fused_qkvbfg_a_proj", + "model.layers.17.self_attn.g_a_proj", + "model.layers.17.self_attn.g_b_proj", + "model.layers.17.self_attn.k_conv1d", + "model.layers.17.self_attn.k_proj", + "model.layers.17.self_attn.o_norm", + "model.layers.17.self_attn.o_proj", + "model.layers.17.self_attn.q_conv1d", + "model.layers.17.self_attn.q_proj", + "model.layers.17.self_attn.qkv_proj", + "model.layers.17.self_attn.v_conv1d", + "model.layers.17.self_attn.v_proj", + "model.layers.18.hc_attn_base", + "model.layers.18.hc_attn_fn", + "model.layers.18.hc_attn_scale", + "model.layers.18.hc_ffn_base", + "model.layers.18.hc_ffn_fn", + "model.layers.18.hc_ffn_scale", + "model.layers.18.input_layernorm", + "model.layers.18.mlp.gate", + "model.layers.18.mlp.gate.e_score_correction_bias", + "model.layers.18.post_attention_layernorm", + "model.layers.18.self_attn.A_log", + "model.layers.18.self_attn.b_proj", + "model.layers.18.self_attn.dt_bias", + "model.layers.18.self_attn.f_a_proj", + "model.layers.18.self_attn.f_b_proj", + "model.layers.18.self_attn.fused_qkvbfg_a_proj", + "model.layers.18.self_attn.g_a_proj", + "model.layers.18.self_attn.g_b_proj", + "model.layers.18.self_attn.k_conv1d", + "model.layers.18.self_attn.k_proj", + "model.layers.18.self_attn.o_norm", + "model.layers.18.self_attn.o_proj", + "model.layers.18.self_attn.q_conv1d", + "model.layers.18.self_attn.q_proj", + "model.layers.18.self_attn.qkv_proj", + "model.layers.18.self_attn.v_conv1d", + "model.layers.18.self_attn.v_proj", + "model.layers.19.hc_attn_base", + "model.layers.19.hc_attn_fn", + "model.layers.19.hc_attn_scale", + "model.layers.19.hc_ffn_base", + "model.layers.19.hc_ffn_fn", + "model.layers.19.hc_ffn_scale", + "model.layers.19.input_layernorm", + "model.layers.19.mlp.gate", + "model.layers.19.mlp.gate.e_score_correction_bias", + "model.layers.19.post_attention_layernorm", + "model.layers.19.self_attn.indexer.index_kpool_compress_ape", + "model.layers.19.self_attn.indexer.index_kpool_compress_gate", + "model.layers.19.self_attn.indexer.k_norm", + "model.layers.19.self_attn.indexer.k_norm.bias", + "model.layers.19.self_attn.indexer.weights_proj", + "model.layers.19.self_attn.indexer.wk", + "model.layers.19.self_attn.indexer.wq_b", + "model.layers.19.self_attn.kv_a_layernorm", + "model.layers.19.self_attn.kv_b_proj", + "model.layers.19.self_attn.q_a_layernorm", + "model.layers.2.hc_attn_base", + "model.layers.2.hc_attn_fn", + "model.layers.2.hc_attn_scale", + "model.layers.2.hc_ffn_base", + "model.layers.2.hc_ffn_fn", + "model.layers.2.hc_ffn_scale", + "model.layers.2.input_layernorm", + "model.layers.2.post_attention_layernorm", + "model.layers.2.self_attn.A_log", + "model.layers.2.self_attn.b_proj", + "model.layers.2.self_attn.dt_bias", + "model.layers.2.self_attn.f_a_proj", + "model.layers.2.self_attn.f_b_proj", + "model.layers.2.self_attn.fused_qkvbfg_a_proj", + "model.layers.2.self_attn.g_a_proj", + "model.layers.2.self_attn.g_b_proj", + "model.layers.2.self_attn.k_conv1d", + "model.layers.2.self_attn.k_proj", + "model.layers.2.self_attn.o_norm", + "model.layers.2.self_attn.o_proj", + "model.layers.2.self_attn.q_conv1d", + "model.layers.2.self_attn.q_proj", + "model.layers.2.self_attn.qkv_proj", + "model.layers.2.self_attn.v_conv1d", + "model.layers.2.self_attn.v_proj", + "model.layers.20.hc_attn_base", + "model.layers.20.hc_attn_fn", + "model.layers.20.hc_attn_scale", + "model.layers.20.hc_ffn_base", + "model.layers.20.hc_ffn_fn", + "model.layers.20.hc_ffn_scale", + "model.layers.20.input_layernorm", + "model.layers.20.mlp.gate", + "model.layers.20.mlp.gate.e_score_correction_bias", + "model.layers.20.post_attention_layernorm", + "model.layers.20.self_attn.A_log", + "model.layers.20.self_attn.b_proj", + "model.layers.20.self_attn.dt_bias", + "model.layers.20.self_attn.f_a_proj", + "model.layers.20.self_attn.f_b_proj", + "model.layers.20.self_attn.fused_qkvbfg_a_proj", + "model.layers.20.self_attn.g_a_proj", + "model.layers.20.self_attn.g_b_proj", + "model.layers.20.self_attn.k_conv1d", + "model.layers.20.self_attn.k_proj", + "model.layers.20.self_attn.o_norm", + "model.layers.20.self_attn.o_proj", + "model.layers.20.self_attn.q_conv1d", + "model.layers.20.self_attn.q_proj", + "model.layers.20.self_attn.qkv_proj", + "model.layers.20.self_attn.v_conv1d", + "model.layers.20.self_attn.v_proj", + "model.layers.21.hc_attn_base", + "model.layers.21.hc_attn_fn", + "model.layers.21.hc_attn_scale", + "model.layers.21.hc_ffn_base", + "model.layers.21.hc_ffn_fn", + "model.layers.21.hc_ffn_scale", + "model.layers.21.input_layernorm", + "model.layers.21.mlp.gate", + "model.layers.21.mlp.gate.e_score_correction_bias", + "model.layers.21.post_attention_layernorm", + "model.layers.21.self_attn.A_log", + "model.layers.21.self_attn.b_proj", + "model.layers.21.self_attn.dt_bias", + "model.layers.21.self_attn.f_a_proj", + "model.layers.21.self_attn.f_b_proj", + "model.layers.21.self_attn.fused_qkvbfg_a_proj", + "model.layers.21.self_attn.g_a_proj", + "model.layers.21.self_attn.g_b_proj", + "model.layers.21.self_attn.k_conv1d", + "model.layers.21.self_attn.k_proj", + "model.layers.21.self_attn.o_norm", + "model.layers.21.self_attn.o_proj", + "model.layers.21.self_attn.q_conv1d", + "model.layers.21.self_attn.q_proj", + "model.layers.21.self_attn.qkv_proj", + "model.layers.21.self_attn.v_conv1d", + "model.layers.21.self_attn.v_proj", + "model.layers.22.hc_attn_base", + "model.layers.22.hc_attn_fn", + "model.layers.22.hc_attn_scale", + "model.layers.22.hc_ffn_base", + "model.layers.22.hc_ffn_fn", + "model.layers.22.hc_ffn_scale", + "model.layers.22.input_layernorm", + "model.layers.22.mlp.gate", + "model.layers.22.mlp.gate.e_score_correction_bias", + "model.layers.22.post_attention_layernorm", + "model.layers.22.self_attn.A_log", + "model.layers.22.self_attn.b_proj", + "model.layers.22.self_attn.dt_bias", + "model.layers.22.self_attn.f_a_proj", + "model.layers.22.self_attn.f_b_proj", + "model.layers.22.self_attn.fused_qkvbfg_a_proj", + "model.layers.22.self_attn.g_a_proj", + "model.layers.22.self_attn.g_b_proj", + "model.layers.22.self_attn.k_conv1d", + "model.layers.22.self_attn.k_proj", + "model.layers.22.self_attn.o_norm", + "model.layers.22.self_attn.o_proj", + "model.layers.22.self_attn.q_conv1d", + "model.layers.22.self_attn.q_proj", + "model.layers.22.self_attn.qkv_proj", + "model.layers.22.self_attn.v_conv1d", + "model.layers.22.self_attn.v_proj", + "model.layers.23.hc_attn_base", + "model.layers.23.hc_attn_fn", + "model.layers.23.hc_attn_scale", + "model.layers.23.hc_ffn_base", + "model.layers.23.hc_ffn_fn", + "model.layers.23.hc_ffn_scale", + "model.layers.23.input_layernorm", + "model.layers.23.mlp.gate", + "model.layers.23.mlp.gate.e_score_correction_bias", + "model.layers.23.post_attention_layernorm", + "model.layers.23.self_attn.indexer.index_kpool_compress_ape", + "model.layers.23.self_attn.indexer.index_kpool_compress_gate", + "model.layers.23.self_attn.indexer.k_norm", + "model.layers.23.self_attn.indexer.k_norm.bias", + "model.layers.23.self_attn.indexer.weights_proj", + "model.layers.23.self_attn.indexer.wk", + "model.layers.23.self_attn.indexer.wq_b", + "model.layers.23.self_attn.kv_a_layernorm", + "model.layers.23.self_attn.kv_b_proj", + "model.layers.23.self_attn.q_a_layernorm", + "model.layers.24.hc_attn_base", + "model.layers.24.hc_attn_fn", + "model.layers.24.hc_attn_scale", + "model.layers.24.hc_ffn_base", + "model.layers.24.hc_ffn_fn", + "model.layers.24.hc_ffn_scale", + "model.layers.24.input_layernorm", + "model.layers.24.mlp.gate", + "model.layers.24.mlp.gate.e_score_correction_bias", + "model.layers.24.post_attention_layernorm", + "model.layers.24.self_attn.A_log", + "model.layers.24.self_attn.b_proj", + "model.layers.24.self_attn.dt_bias", + "model.layers.24.self_attn.f_a_proj", + "model.layers.24.self_attn.f_b_proj", + "model.layers.24.self_attn.fused_qkvbfg_a_proj", + "model.layers.24.self_attn.g_a_proj", + "model.layers.24.self_attn.g_b_proj", + "model.layers.24.self_attn.k_conv1d", + "model.layers.24.self_attn.k_proj", + "model.layers.24.self_attn.o_norm", + "model.layers.24.self_attn.o_proj", + "model.layers.24.self_attn.q_conv1d", + "model.layers.24.self_attn.q_proj", + "model.layers.24.self_attn.qkv_proj", + "model.layers.24.self_attn.v_conv1d", + "model.layers.24.self_attn.v_proj", + "model.layers.25.hc_attn_base", + "model.layers.25.hc_attn_fn", + "model.layers.25.hc_attn_scale", + "model.layers.25.hc_ffn_base", + "model.layers.25.hc_ffn_fn", + "model.layers.25.hc_ffn_scale", + "model.layers.25.input_layernorm", + "model.layers.25.mlp.gate", + "model.layers.25.mlp.gate.e_score_correction_bias", + "model.layers.25.post_attention_layernorm", + "model.layers.25.self_attn.A_log", + "model.layers.25.self_attn.b_proj", + "model.layers.25.self_attn.dt_bias", + "model.layers.25.self_attn.f_a_proj", + "model.layers.25.self_attn.f_b_proj", + "model.layers.25.self_attn.fused_qkvbfg_a_proj", + "model.layers.25.self_attn.g_a_proj", + "model.layers.25.self_attn.g_b_proj", + "model.layers.25.self_attn.k_conv1d", + "model.layers.25.self_attn.k_proj", + "model.layers.25.self_attn.o_norm", + "model.layers.25.self_attn.o_proj", + "model.layers.25.self_attn.q_conv1d", + "model.layers.25.self_attn.q_proj", + "model.layers.25.self_attn.qkv_proj", + "model.layers.25.self_attn.v_conv1d", + "model.layers.25.self_attn.v_proj", + "model.layers.26.hc_attn_base", + "model.layers.26.hc_attn_fn", + "model.layers.26.hc_attn_scale", + "model.layers.26.hc_ffn_base", + "model.layers.26.hc_ffn_fn", + "model.layers.26.hc_ffn_scale", + "model.layers.26.input_layernorm", + "model.layers.26.mlp.gate", + "model.layers.26.mlp.gate.e_score_correction_bias", + "model.layers.26.post_attention_layernorm", + "model.layers.26.self_attn.A_log", + "model.layers.26.self_attn.b_proj", + "model.layers.26.self_attn.dt_bias", + "model.layers.26.self_attn.f_a_proj", + "model.layers.26.self_attn.f_b_proj", + "model.layers.26.self_attn.fused_qkvbfg_a_proj", + "model.layers.26.self_attn.g_a_proj", + "model.layers.26.self_attn.g_b_proj", + "model.layers.26.self_attn.k_conv1d", + "model.layers.26.self_attn.k_proj", + "model.layers.26.self_attn.o_norm", + "model.layers.26.self_attn.o_proj", + "model.layers.26.self_attn.q_conv1d", + "model.layers.26.self_attn.q_proj", + "model.layers.26.self_attn.qkv_proj", + "model.layers.26.self_attn.v_conv1d", + "model.layers.26.self_attn.v_proj", + "model.layers.27.hc_attn_base", + "model.layers.27.hc_attn_fn", + "model.layers.27.hc_attn_scale", + "model.layers.27.hc_ffn_base", + "model.layers.27.hc_ffn_fn", + "model.layers.27.hc_ffn_scale", + "model.layers.27.input_layernorm", + "model.layers.27.mlp.gate", + "model.layers.27.mlp.gate.e_score_correction_bias", + "model.layers.27.post_attention_layernorm", + "model.layers.27.self_attn.indexer.index_kpool_compress_ape", + "model.layers.27.self_attn.indexer.index_kpool_compress_gate", + "model.layers.27.self_attn.indexer.k_norm", + "model.layers.27.self_attn.indexer.k_norm.bias", + "model.layers.27.self_attn.indexer.weights_proj", + "model.layers.27.self_attn.indexer.wk", + "model.layers.27.self_attn.indexer.wq_b", + "model.layers.27.self_attn.kv_a_layernorm", + "model.layers.27.self_attn.kv_b_proj", + "model.layers.27.self_attn.q_a_layernorm", + "model.layers.28.hc_attn_base", + "model.layers.28.hc_attn_fn", + "model.layers.28.hc_attn_scale", + "model.layers.28.hc_ffn_base", + "model.layers.28.hc_ffn_fn", + "model.layers.28.hc_ffn_scale", + "model.layers.28.input_layernorm", + "model.layers.28.mlp.gate", + "model.layers.28.mlp.gate.e_score_correction_bias", + "model.layers.28.post_attention_layernorm", + "model.layers.28.self_attn.A_log", + "model.layers.28.self_attn.b_proj", + "model.layers.28.self_attn.dt_bias", + "model.layers.28.self_attn.f_a_proj", + "model.layers.28.self_attn.f_b_proj", + "model.layers.28.self_attn.fused_qkvbfg_a_proj", + "model.layers.28.self_attn.g_a_proj", + "model.layers.28.self_attn.g_b_proj", + "model.layers.28.self_attn.k_conv1d", + "model.layers.28.self_attn.k_proj", + "model.layers.28.self_attn.o_norm", + "model.layers.28.self_attn.o_proj", + "model.layers.28.self_attn.q_conv1d", + "model.layers.28.self_attn.q_proj", + "model.layers.28.self_attn.qkv_proj", + "model.layers.28.self_attn.v_conv1d", + "model.layers.28.self_attn.v_proj", + "model.layers.29.hc_attn_base", + "model.layers.29.hc_attn_fn", + "model.layers.29.hc_attn_scale", + "model.layers.29.hc_ffn_base", + "model.layers.29.hc_ffn_fn", + "model.layers.29.hc_ffn_scale", + "model.layers.29.input_layernorm", + "model.layers.29.mlp.gate", + "model.layers.29.mlp.gate.e_score_correction_bias", + "model.layers.29.post_attention_layernorm", + "model.layers.29.self_attn.A_log", + "model.layers.29.self_attn.b_proj", + "model.layers.29.self_attn.dt_bias", + "model.layers.29.self_attn.f_a_proj", + "model.layers.29.self_attn.f_b_proj", + "model.layers.29.self_attn.fused_qkvbfg_a_proj", + "model.layers.29.self_attn.g_a_proj", + "model.layers.29.self_attn.g_b_proj", + "model.layers.29.self_attn.k_conv1d", + "model.layers.29.self_attn.k_proj", + "model.layers.29.self_attn.o_norm", + "model.layers.29.self_attn.o_proj", + "model.layers.29.self_attn.q_conv1d", + "model.layers.29.self_attn.q_proj", + "model.layers.29.self_attn.qkv_proj", + "model.layers.29.self_attn.v_conv1d", + "model.layers.29.self_attn.v_proj", + "model.layers.3.hc_attn_base", + "model.layers.3.hc_attn_fn", + "model.layers.3.hc_attn_scale", + "model.layers.3.hc_ffn_base", + "model.layers.3.hc_ffn_fn", + "model.layers.3.hc_ffn_scale", + "model.layers.3.input_layernorm", + "model.layers.3.mlp.gate", + "model.layers.3.mlp.gate.e_score_correction_bias", + "model.layers.3.post_attention_layernorm", + "model.layers.3.self_attn.indexer.index_kpool_compress_ape", + "model.layers.3.self_attn.indexer.index_kpool_compress_gate", + "model.layers.3.self_attn.indexer.k_norm", + "model.layers.3.self_attn.indexer.k_norm.bias", + "model.layers.3.self_attn.indexer.weights_proj", + "model.layers.3.self_attn.indexer.wk", + "model.layers.3.self_attn.indexer.wq_b", + "model.layers.3.self_attn.kv_a_layernorm", + "model.layers.3.self_attn.kv_b_proj", + "model.layers.3.self_attn.q_a_layernorm", + "model.layers.30.hc_attn_base", + "model.layers.30.hc_attn_fn", + "model.layers.30.hc_attn_scale", + "model.layers.30.hc_ffn_base", + "model.layers.30.hc_ffn_fn", + "model.layers.30.hc_ffn_scale", + "model.layers.30.input_layernorm", + "model.layers.30.mlp.gate", + "model.layers.30.mlp.gate.e_score_correction_bias", + "model.layers.30.post_attention_layernorm", + "model.layers.30.self_attn.A_log", + "model.layers.30.self_attn.b_proj", + "model.layers.30.self_attn.dt_bias", + "model.layers.30.self_attn.f_a_proj", + "model.layers.30.self_attn.f_b_proj", + "model.layers.30.self_attn.fused_qkvbfg_a_proj", + "model.layers.30.self_attn.g_a_proj", + "model.layers.30.self_attn.g_b_proj", + "model.layers.30.self_attn.k_conv1d", + "model.layers.30.self_attn.k_proj", + "model.layers.30.self_attn.o_norm", + "model.layers.30.self_attn.o_proj", + "model.layers.30.self_attn.q_conv1d", + "model.layers.30.self_attn.q_proj", + "model.layers.30.self_attn.qkv_proj", + "model.layers.30.self_attn.v_conv1d", + "model.layers.30.self_attn.v_proj", + "model.layers.31.hc_attn_base", + "model.layers.31.hc_attn_fn", + "model.layers.31.hc_attn_scale", + "model.layers.31.hc_ffn_base", + "model.layers.31.hc_ffn_fn", + "model.layers.31.hc_ffn_scale", + "model.layers.31.input_layernorm", + "model.layers.31.mlp.gate", + "model.layers.31.mlp.gate.e_score_correction_bias", + "model.layers.31.post_attention_layernorm", + "model.layers.31.self_attn.indexer.index_kpool_compress_ape", + "model.layers.31.self_attn.indexer.index_kpool_compress_gate", + "model.layers.31.self_attn.indexer.k_norm", + "model.layers.31.self_attn.indexer.k_norm.bias", + "model.layers.31.self_attn.indexer.weights_proj", + "model.layers.31.self_attn.indexer.wk", + "model.layers.31.self_attn.indexer.wq_b", + "model.layers.31.self_attn.kv_a_layernorm", + "model.layers.31.self_attn.kv_b_proj", + "model.layers.31.self_attn.q_a_layernorm", + "model.layers.32.hc_attn_base", + "model.layers.32.hc_attn_fn", + "model.layers.32.hc_attn_scale", + "model.layers.32.hc_ffn_base", + "model.layers.32.hc_ffn_fn", + "model.layers.32.hc_ffn_scale", + "model.layers.32.input_layernorm", + "model.layers.32.mlp.gate", + "model.layers.32.mlp.gate.e_score_correction_bias", + "model.layers.32.post_attention_layernorm", + "model.layers.32.self_attn.A_log", + "model.layers.32.self_attn.b_proj", + "model.layers.32.self_attn.dt_bias", + "model.layers.32.self_attn.f_a_proj", + "model.layers.32.self_attn.f_b_proj", + "model.layers.32.self_attn.fused_qkvbfg_a_proj", + "model.layers.32.self_attn.g_a_proj", + "model.layers.32.self_attn.g_b_proj", + "model.layers.32.self_attn.k_conv1d", + "model.layers.32.self_attn.k_proj", + "model.layers.32.self_attn.o_norm", + "model.layers.32.self_attn.o_proj", + "model.layers.32.self_attn.q_conv1d", + "model.layers.32.self_attn.q_proj", + "model.layers.32.self_attn.qkv_proj", + "model.layers.32.self_attn.v_conv1d", + "model.layers.32.self_attn.v_proj", + "model.layers.33.hc_attn_base", + "model.layers.33.hc_attn_fn", + "model.layers.33.hc_attn_scale", + "model.layers.33.hc_ffn_base", + "model.layers.33.hc_ffn_fn", + "model.layers.33.hc_ffn_scale", + "model.layers.33.input_layernorm", + "model.layers.33.mlp.gate", + "model.layers.33.mlp.gate.e_score_correction_bias", + "model.layers.33.post_attention_layernorm", + "model.layers.33.self_attn.A_log", + "model.layers.33.self_attn.b_proj", + "model.layers.33.self_attn.dt_bias", + "model.layers.33.self_attn.f_a_proj", + "model.layers.33.self_attn.f_b_proj", + "model.layers.33.self_attn.fused_qkvbfg_a_proj", + "model.layers.33.self_attn.g_a_proj", + "model.layers.33.self_attn.g_b_proj", + "model.layers.33.self_attn.k_conv1d", + "model.layers.33.self_attn.k_proj", + "model.layers.33.self_attn.o_norm", + "model.layers.33.self_attn.o_proj", + "model.layers.33.self_attn.q_conv1d", + "model.layers.33.self_attn.q_proj", + "model.layers.33.self_attn.qkv_proj", + "model.layers.33.self_attn.v_conv1d", + "model.layers.33.self_attn.v_proj", + "model.layers.34.hc_attn_base", + "model.layers.34.hc_attn_fn", + "model.layers.34.hc_attn_scale", + "model.layers.34.hc_ffn_base", + "model.layers.34.hc_ffn_fn", + "model.layers.34.hc_ffn_scale", + "model.layers.34.input_layernorm", + "model.layers.34.mlp.gate", + "model.layers.34.mlp.gate.e_score_correction_bias", + "model.layers.34.post_attention_layernorm", + "model.layers.34.self_attn.A_log", + "model.layers.34.self_attn.b_proj", + "model.layers.34.self_attn.dt_bias", + "model.layers.34.self_attn.f_a_proj", + "model.layers.34.self_attn.f_b_proj", + "model.layers.34.self_attn.fused_qkvbfg_a_proj", + "model.layers.34.self_attn.g_a_proj", + "model.layers.34.self_attn.g_b_proj", + "model.layers.34.self_attn.k_conv1d", + "model.layers.34.self_attn.k_proj", + "model.layers.34.self_attn.o_norm", + "model.layers.34.self_attn.o_proj", + "model.layers.34.self_attn.q_conv1d", + "model.layers.34.self_attn.q_proj", + "model.layers.34.self_attn.qkv_proj", + "model.layers.34.self_attn.v_conv1d", + "model.layers.34.self_attn.v_proj", + "model.layers.35.hc_attn_base", + "model.layers.35.hc_attn_fn", + "model.layers.35.hc_attn_scale", + "model.layers.35.hc_ffn_base", + "model.layers.35.hc_ffn_fn", + "model.layers.35.hc_ffn_scale", + "model.layers.35.input_layernorm", + "model.layers.35.mlp.gate", + "model.layers.35.mlp.gate.e_score_correction_bias", + "model.layers.35.post_attention_layernorm", + "model.layers.35.self_attn.indexer.index_kpool_compress_ape", + "model.layers.35.self_attn.indexer.index_kpool_compress_gate", + "model.layers.35.self_attn.indexer.k_norm", + "model.layers.35.self_attn.indexer.k_norm.bias", + "model.layers.35.self_attn.indexer.weights_proj", + "model.layers.35.self_attn.indexer.wk", + "model.layers.35.self_attn.indexer.wq_b", + "model.layers.35.self_attn.kv_a_layernorm", + "model.layers.35.self_attn.kv_b_proj", + "model.layers.35.self_attn.q_a_layernorm", + "model.layers.36.hc_attn_base", + "model.layers.36.hc_attn_fn", + "model.layers.36.hc_attn_scale", + "model.layers.36.hc_ffn_base", + "model.layers.36.hc_ffn_fn", + "model.layers.36.hc_ffn_scale", + "model.layers.36.input_layernorm", + "model.layers.36.mlp.gate", + "model.layers.36.mlp.gate.e_score_correction_bias", + "model.layers.36.post_attention_layernorm", + "model.layers.36.self_attn.A_log", + "model.layers.36.self_attn.b_proj", + "model.layers.36.self_attn.dt_bias", + "model.layers.36.self_attn.f_a_proj", + "model.layers.36.self_attn.f_b_proj", + "model.layers.36.self_attn.fused_qkvbfg_a_proj", + "model.layers.36.self_attn.g_a_proj", + "model.layers.36.self_attn.g_b_proj", + "model.layers.36.self_attn.k_conv1d", + "model.layers.36.self_attn.k_proj", + "model.layers.36.self_attn.o_norm", + "model.layers.36.self_attn.o_proj", + "model.layers.36.self_attn.q_conv1d", + "model.layers.36.self_attn.q_proj", + "model.layers.36.self_attn.qkv_proj", + "model.layers.36.self_attn.v_conv1d", + "model.layers.36.self_attn.v_proj", + "model.layers.37.hc_attn_base", + "model.layers.37.hc_attn_fn", + "model.layers.37.hc_attn_scale", + "model.layers.37.hc_ffn_base", + "model.layers.37.hc_ffn_fn", + "model.layers.37.hc_ffn_scale", + "model.layers.37.input_layernorm", + "model.layers.37.mlp.gate", + "model.layers.37.mlp.gate.e_score_correction_bias", + "model.layers.37.post_attention_layernorm", + "model.layers.37.self_attn.A_log", + "model.layers.37.self_attn.b_proj", + "model.layers.37.self_attn.dt_bias", + "model.layers.37.self_attn.f_a_proj", + "model.layers.37.self_attn.f_b_proj", + "model.layers.37.self_attn.fused_qkvbfg_a_proj", + "model.layers.37.self_attn.g_a_proj", + "model.layers.37.self_attn.g_b_proj", + "model.layers.37.self_attn.k_conv1d", + "model.layers.37.self_attn.k_proj", + "model.layers.37.self_attn.o_norm", + "model.layers.37.self_attn.o_proj", + "model.layers.37.self_attn.q_conv1d", + "model.layers.37.self_attn.q_proj", + "model.layers.37.self_attn.qkv_proj", + "model.layers.37.self_attn.v_conv1d", + "model.layers.37.self_attn.v_proj", + "model.layers.38.hc_attn_base", + "model.layers.38.hc_attn_fn", + "model.layers.38.hc_attn_scale", + "model.layers.38.hc_ffn_base", + "model.layers.38.hc_ffn_fn", + "model.layers.38.hc_ffn_scale", + "model.layers.38.input_layernorm", + "model.layers.38.mlp.gate", + "model.layers.38.mlp.gate.e_score_correction_bias", + "model.layers.38.post_attention_layernorm", + "model.layers.38.self_attn.A_log", + "model.layers.38.self_attn.b_proj", + "model.layers.38.self_attn.dt_bias", + "model.layers.38.self_attn.f_a_proj", + "model.layers.38.self_attn.f_b_proj", + "model.layers.38.self_attn.fused_qkvbfg_a_proj", + "model.layers.38.self_attn.g_a_proj", + "model.layers.38.self_attn.g_b_proj", + "model.layers.38.self_attn.k_conv1d", + "model.layers.38.self_attn.k_proj", + "model.layers.38.self_attn.o_norm", + "model.layers.38.self_attn.o_proj", + "model.layers.38.self_attn.q_conv1d", + "model.layers.38.self_attn.q_proj", + "model.layers.38.self_attn.qkv_proj", + "model.layers.38.self_attn.v_conv1d", + "model.layers.38.self_attn.v_proj", + "model.layers.39.hc_attn_base", + "model.layers.39.hc_attn_fn", + "model.layers.39.hc_attn_scale", + "model.layers.39.hc_ffn_base", + "model.layers.39.hc_ffn_fn", + "model.layers.39.hc_ffn_scale", + "model.layers.39.input_layernorm", + "model.layers.39.mlp.gate", + "model.layers.39.mlp.gate.e_score_correction_bias", + "model.layers.39.post_attention_layernorm", + "model.layers.39.self_attn.indexer.index_kpool_compress_ape", + "model.layers.39.self_attn.indexer.index_kpool_compress_gate", + "model.layers.39.self_attn.indexer.k_norm", + "model.layers.39.self_attn.indexer.k_norm.bias", + "model.layers.39.self_attn.indexer.weights_proj", + "model.layers.39.self_attn.indexer.wk", + "model.layers.39.self_attn.indexer.wq_b", + "model.layers.39.self_attn.kv_a_layernorm", + "model.layers.39.self_attn.kv_b_proj", + "model.layers.39.self_attn.q_a_layernorm", + "model.layers.4.hc_attn_base", + "model.layers.4.hc_attn_fn", + "model.layers.4.hc_attn_scale", + "model.layers.4.hc_ffn_base", + "model.layers.4.hc_ffn_fn", + "model.layers.4.hc_ffn_scale", + "model.layers.4.input_layernorm", + "model.layers.4.mlp.gate", + "model.layers.4.mlp.gate.e_score_correction_bias", + "model.layers.4.post_attention_layernorm", + "model.layers.4.self_attn.A_log", + "model.layers.4.self_attn.b_proj", + "model.layers.4.self_attn.dt_bias", + "model.layers.4.self_attn.f_a_proj", + "model.layers.4.self_attn.f_b_proj", + "model.layers.4.self_attn.fused_qkvbfg_a_proj", + "model.layers.4.self_attn.g_a_proj", + "model.layers.4.self_attn.g_b_proj", + "model.layers.4.self_attn.k_conv1d", + "model.layers.4.self_attn.k_proj", + "model.layers.4.self_attn.o_norm", + "model.layers.4.self_attn.o_proj", + "model.layers.4.self_attn.q_conv1d", + "model.layers.4.self_attn.q_proj", + "model.layers.4.self_attn.qkv_proj", + "model.layers.4.self_attn.v_conv1d", + "model.layers.4.self_attn.v_proj", + "model.layers.40.hc_attn_base", + "model.layers.40.hc_attn_fn", + "model.layers.40.hc_attn_scale", + "model.layers.40.hc_ffn_base", + "model.layers.40.hc_ffn_fn", + "model.layers.40.hc_ffn_scale", + "model.layers.40.input_layernorm", + "model.layers.40.mlp.gate", + "model.layers.40.mlp.gate.e_score_correction_bias", + "model.layers.40.post_attention_layernorm", + "model.layers.40.self_attn.A_log", + "model.layers.40.self_attn.b_proj", + "model.layers.40.self_attn.dt_bias", + "model.layers.40.self_attn.f_a_proj", + "model.layers.40.self_attn.f_b_proj", + "model.layers.40.self_attn.fused_qkvbfg_a_proj", + "model.layers.40.self_attn.g_a_proj", + "model.layers.40.self_attn.g_b_proj", + "model.layers.40.self_attn.k_conv1d", + "model.layers.40.self_attn.k_proj", + "model.layers.40.self_attn.o_norm", + "model.layers.40.self_attn.o_proj", + "model.layers.40.self_attn.q_conv1d", + "model.layers.40.self_attn.q_proj", + "model.layers.40.self_attn.qkv_proj", + "model.layers.40.self_attn.v_conv1d", + "model.layers.40.self_attn.v_proj", + "model.layers.41.hc_attn_base", + "model.layers.41.hc_attn_fn", + "model.layers.41.hc_attn_scale", + "model.layers.41.hc_ffn_base", + "model.layers.41.hc_ffn_fn", + "model.layers.41.hc_ffn_scale", + "model.layers.41.input_layernorm", + "model.layers.41.mlp.gate", + "model.layers.41.mlp.gate.e_score_correction_bias", + "model.layers.41.post_attention_layernorm", + "model.layers.41.self_attn.A_log", + "model.layers.41.self_attn.b_proj", + "model.layers.41.self_attn.dt_bias", + "model.layers.41.self_attn.f_a_proj", + "model.layers.41.self_attn.f_b_proj", + "model.layers.41.self_attn.fused_qkvbfg_a_proj", + "model.layers.41.self_attn.g_a_proj", + "model.layers.41.self_attn.g_b_proj", + "model.layers.41.self_attn.k_conv1d", + "model.layers.41.self_attn.k_proj", + "model.layers.41.self_attn.o_norm", + "model.layers.41.self_attn.o_proj", + "model.layers.41.self_attn.q_conv1d", + "model.layers.41.self_attn.q_proj", + "model.layers.41.self_attn.qkv_proj", + "model.layers.41.self_attn.v_conv1d", + "model.layers.41.self_attn.v_proj", + "model.layers.42.hc_attn_base", + "model.layers.42.hc_attn_fn", + "model.layers.42.hc_attn_scale", + "model.layers.42.hc_ffn_base", + "model.layers.42.hc_ffn_fn", + "model.layers.42.hc_ffn_scale", + "model.layers.42.input_layernorm", + "model.layers.42.mlp.gate", + "model.layers.42.mlp.gate.e_score_correction_bias", + "model.layers.42.post_attention_layernorm", + "model.layers.42.self_attn.A_log", + "model.layers.42.self_attn.b_proj", + "model.layers.42.self_attn.dt_bias", + "model.layers.42.self_attn.f_a_proj", + "model.layers.42.self_attn.f_b_proj", + "model.layers.42.self_attn.fused_qkvbfg_a_proj", + "model.layers.42.self_attn.g_a_proj", + "model.layers.42.self_attn.g_b_proj", + "model.layers.42.self_attn.k_conv1d", + "model.layers.42.self_attn.k_proj", + "model.layers.42.self_attn.o_norm", + "model.layers.42.self_attn.o_proj", + "model.layers.42.self_attn.q_conv1d", + "model.layers.42.self_attn.q_proj", + "model.layers.42.self_attn.qkv_proj", + "model.layers.42.self_attn.v_conv1d", + "model.layers.42.self_attn.v_proj", + "model.layers.43.hc_attn_base", + "model.layers.43.hc_attn_fn", + "model.layers.43.hc_attn_scale", + "model.layers.43.hc_ffn_base", + "model.layers.43.hc_ffn_fn", + "model.layers.43.hc_ffn_scale", + "model.layers.43.input_layernorm", + "model.layers.43.mlp.gate", + "model.layers.43.mlp.gate.e_score_correction_bias", + "model.layers.43.post_attention_layernorm", + "model.layers.43.self_attn.indexer.index_kpool_compress_ape", + "model.layers.43.self_attn.indexer.index_kpool_compress_gate", + "model.layers.43.self_attn.indexer.k_norm", + "model.layers.43.self_attn.indexer.k_norm.bias", + "model.layers.43.self_attn.indexer.weights_proj", + "model.layers.43.self_attn.indexer.wk", + "model.layers.43.self_attn.indexer.wq_b", + "model.layers.43.self_attn.kv_a_layernorm", + "model.layers.43.self_attn.kv_b_proj", + "model.layers.43.self_attn.q_a_layernorm", + "model.layers.44.hc_attn_base", + "model.layers.44.hc_attn_fn", + "model.layers.44.hc_attn_scale", + "model.layers.44.hc_ffn_base", + "model.layers.44.hc_ffn_fn", + "model.layers.44.hc_ffn_scale", + "model.layers.44.input_layernorm", + "model.layers.44.mlp.gate", + "model.layers.44.mlp.gate.e_score_correction_bias", + "model.layers.44.post_attention_layernorm", + "model.layers.44.self_attn.A_log", + "model.layers.44.self_attn.b_proj", + "model.layers.44.self_attn.dt_bias", + "model.layers.44.self_attn.f_a_proj", + "model.layers.44.self_attn.f_b_proj", + "model.layers.44.self_attn.fused_qkvbfg_a_proj", + "model.layers.44.self_attn.g_a_proj", + "model.layers.44.self_attn.g_b_proj", + "model.layers.44.self_attn.k_conv1d", + "model.layers.44.self_attn.k_proj", + "model.layers.44.self_attn.o_norm", + "model.layers.44.self_attn.o_proj", + "model.layers.44.self_attn.q_conv1d", + "model.layers.44.self_attn.q_proj", + "model.layers.44.self_attn.qkv_proj", + "model.layers.44.self_attn.v_conv1d", + "model.layers.44.self_attn.v_proj", + "model.layers.45.eh_proj", + "model.layers.45.enorm", + "model.layers.45.hnorm", + "model.layers.45.input_layernorm", + "model.layers.45.mlp.gate", + "model.layers.45.mlp.gate.e_score_correction_bias", + "model.layers.45.post_attention_layernorm", + "model.layers.45.self_attn.indexer.index_kpool_compress_ape", + "model.layers.45.self_attn.indexer.index_kpool_compress_gate", + "model.layers.45.self_attn.indexer.k_norm", + "model.layers.45.self_attn.indexer.k_norm.bias", + "model.layers.45.self_attn.indexer.weights_proj", + "model.layers.45.self_attn.indexer.wk", + "model.layers.45.self_attn.indexer.wq_b", + "model.layers.45.self_attn.kv_a_layernorm", + "model.layers.45.self_attn.kv_b_proj", + "model.layers.45.self_attn.q_a_layernorm", + "model.layers.45.shared_head.norm", + "model.layers.5.hc_attn_base", + "model.layers.5.hc_attn_fn", + "model.layers.5.hc_attn_scale", + "model.layers.5.hc_ffn_base", + "model.layers.5.hc_ffn_fn", + "model.layers.5.hc_ffn_scale", + "model.layers.5.input_layernorm", + "model.layers.5.mlp.gate", + "model.layers.5.mlp.gate.e_score_correction_bias", + "model.layers.5.post_attention_layernorm", + "model.layers.5.self_attn.A_log", + "model.layers.5.self_attn.b_proj", + "model.layers.5.self_attn.dt_bias", + "model.layers.5.self_attn.f_a_proj", + "model.layers.5.self_attn.f_b_proj", + "model.layers.5.self_attn.fused_qkvbfg_a_proj", + "model.layers.5.self_attn.g_a_proj", + "model.layers.5.self_attn.g_b_proj", + "model.layers.5.self_attn.k_conv1d", + "model.layers.5.self_attn.k_proj", + "model.layers.5.self_attn.o_norm", + "model.layers.5.self_attn.o_proj", + "model.layers.5.self_attn.q_conv1d", + "model.layers.5.self_attn.q_proj", + "model.layers.5.self_attn.qkv_proj", + "model.layers.5.self_attn.v_conv1d", + "model.layers.5.self_attn.v_proj", + "model.layers.6.hc_attn_base", + "model.layers.6.hc_attn_fn", + "model.layers.6.hc_attn_scale", + "model.layers.6.hc_ffn_base", + "model.layers.6.hc_ffn_fn", + "model.layers.6.hc_ffn_scale", + "model.layers.6.input_layernorm", + "model.layers.6.mlp.gate", + "model.layers.6.mlp.gate.e_score_correction_bias", + "model.layers.6.post_attention_layernorm", + "model.layers.6.self_attn.A_log", + "model.layers.6.self_attn.b_proj", + "model.layers.6.self_attn.dt_bias", + "model.layers.6.self_attn.f_a_proj", + "model.layers.6.self_attn.f_b_proj", + "model.layers.6.self_attn.fused_qkvbfg_a_proj", + "model.layers.6.self_attn.g_a_proj", + "model.layers.6.self_attn.g_b_proj", + "model.layers.6.self_attn.k_conv1d", + "model.layers.6.self_attn.k_proj", + "model.layers.6.self_attn.o_norm", + "model.layers.6.self_attn.o_proj", + "model.layers.6.self_attn.q_conv1d", + "model.layers.6.self_attn.q_proj", + "model.layers.6.self_attn.qkv_proj", + "model.layers.6.self_attn.v_conv1d", + "model.layers.6.self_attn.v_proj", + "model.layers.7.hc_attn_base", + "model.layers.7.hc_attn_fn", + "model.layers.7.hc_attn_scale", + "model.layers.7.hc_ffn_base", + "model.layers.7.hc_ffn_fn", + "model.layers.7.hc_ffn_scale", + "model.layers.7.input_layernorm", + "model.layers.7.mlp.gate", + "model.layers.7.mlp.gate.e_score_correction_bias", + "model.layers.7.post_attention_layernorm", + "model.layers.7.self_attn.indexer.index_kpool_compress_ape", + "model.layers.7.self_attn.indexer.index_kpool_compress_gate", + "model.layers.7.self_attn.indexer.k_norm", + "model.layers.7.self_attn.indexer.k_norm.bias", + "model.layers.7.self_attn.indexer.weights_proj", + "model.layers.7.self_attn.indexer.wk", + "model.layers.7.self_attn.indexer.wq_b", + "model.layers.7.self_attn.kv_a_layernorm", + "model.layers.7.self_attn.kv_b_proj", + "model.layers.7.self_attn.q_a_layernorm", + "model.layers.8.hc_attn_base", + "model.layers.8.hc_attn_fn", + "model.layers.8.hc_attn_scale", + "model.layers.8.hc_ffn_base", + "model.layers.8.hc_ffn_fn", + "model.layers.8.hc_ffn_scale", + "model.layers.8.input_layernorm", + "model.layers.8.mlp.gate", + "model.layers.8.mlp.gate.e_score_correction_bias", + "model.layers.8.post_attention_layernorm", + "model.layers.8.self_attn.A_log", + "model.layers.8.self_attn.b_proj", + "model.layers.8.self_attn.dt_bias", + "model.layers.8.self_attn.f_a_proj", + "model.layers.8.self_attn.f_b_proj", + "model.layers.8.self_attn.fused_qkvbfg_a_proj", + "model.layers.8.self_attn.g_a_proj", + "model.layers.8.self_attn.g_b_proj", + "model.layers.8.self_attn.k_conv1d", + "model.layers.8.self_attn.k_proj", + "model.layers.8.self_attn.o_norm", + "model.layers.8.self_attn.o_proj", + "model.layers.8.self_attn.q_conv1d", + "model.layers.8.self_attn.q_proj", + "model.layers.8.self_attn.qkv_proj", + "model.layers.8.self_attn.v_conv1d", + "model.layers.8.self_attn.v_proj", + "model.layers.9.hc_attn_base", + "model.layers.9.hc_attn_fn", + "model.layers.9.hc_attn_scale", + "model.layers.9.hc_ffn_base", + "model.layers.9.hc_ffn_fn", + "model.layers.9.hc_ffn_scale", + "model.layers.9.input_layernorm", + "model.layers.9.mlp.gate", + "model.layers.9.mlp.gate.e_score_correction_bias", + "model.layers.9.post_attention_layernorm", + "model.layers.9.self_attn.A_log", + "model.layers.9.self_attn.b_proj", + "model.layers.9.self_attn.dt_bias", + "model.layers.9.self_attn.f_a_proj", + "model.layers.9.self_attn.f_b_proj", + "model.layers.9.self_attn.fused_qkvbfg_a_proj", + "model.layers.9.self_attn.g_a_proj", + "model.layers.9.self_attn.g_b_proj", + "model.layers.9.self_attn.k_conv1d", + "model.layers.9.self_attn.k_proj", + "model.layers.9.self_attn.o_norm", + "model.layers.9.self_attn.o_proj", + "model.layers.9.self_attn.q_conv1d", + "model.layers.9.self_attn.q_proj", + "model.layers.9.self_attn.qkv_proj", + "model.layers.9.self_attn.v_conv1d", + "model.layers.9.self_attn.v_proj", + "model.norm", + "model.visual", + "router", + "visual", + "visual.blocks.0.attn.k_norm", + "visual.blocks.0.attn.proj", + "visual.blocks.0.attn.proj.bias", + "visual.blocks.0.attn.q_norm", + "visual.blocks.0.attn.qkv", + "visual.blocks.0.attn.qkv.bias", + "visual.blocks.0.mlp.down_proj", + "visual.blocks.0.mlp.down_proj.bias", + "visual.blocks.0.mlp.gate_proj", + "visual.blocks.0.mlp.gate_proj.bias", + "visual.blocks.0.mlp.up_proj", + "visual.blocks.0.mlp.up_proj.bias", + "visual.blocks.0.norm1", + "visual.blocks.0.norm2", + "visual.blocks.1.attn.k_norm", + "visual.blocks.1.attn.proj", + "visual.blocks.1.attn.proj.bias", + "visual.blocks.1.attn.q_norm", + "visual.blocks.1.attn.qkv", + "visual.blocks.1.attn.qkv.bias", + "visual.blocks.1.mlp.down_proj", + "visual.blocks.1.mlp.down_proj.bias", + "visual.blocks.1.mlp.gate_proj", + "visual.blocks.1.mlp.gate_proj.bias", + "visual.blocks.1.mlp.up_proj", + "visual.blocks.1.mlp.up_proj.bias", + "visual.blocks.1.norm1", + "visual.blocks.1.norm2", + "visual.blocks.10.attn.k_norm", + "visual.blocks.10.attn.proj", + "visual.blocks.10.attn.proj.bias", + "visual.blocks.10.attn.q_norm", + "visual.blocks.10.attn.qkv", + "visual.blocks.10.attn.qkv.bias", + "visual.blocks.10.mlp.down_proj", + "visual.blocks.10.mlp.down_proj.bias", + "visual.blocks.10.mlp.gate_proj", + "visual.blocks.10.mlp.gate_proj.bias", + "visual.blocks.10.mlp.up_proj", + "visual.blocks.10.mlp.up_proj.bias", + "visual.blocks.10.norm1", + "visual.blocks.10.norm2", + "visual.blocks.11.attn.k_norm", + "visual.blocks.11.attn.proj", + "visual.blocks.11.attn.proj.bias", + "visual.blocks.11.attn.q_norm", + "visual.blocks.11.attn.qkv", + "visual.blocks.11.attn.qkv.bias", + "visual.blocks.11.mlp.down_proj", + "visual.blocks.11.mlp.down_proj.bias", + "visual.blocks.11.mlp.gate_proj", + "visual.blocks.11.mlp.gate_proj.bias", + "visual.blocks.11.mlp.up_proj", + "visual.blocks.11.mlp.up_proj.bias", + "visual.blocks.11.norm1", + "visual.blocks.11.norm2", + "visual.blocks.12.attn.k_norm", + "visual.blocks.12.attn.proj", + "visual.blocks.12.attn.proj.bias", + "visual.blocks.12.attn.q_norm", + "visual.blocks.12.attn.qkv", + "visual.blocks.12.attn.qkv.bias", + "visual.blocks.12.mlp.down_proj", + "visual.blocks.12.mlp.down_proj.bias", + "visual.blocks.12.mlp.gate_proj", + "visual.blocks.12.mlp.gate_proj.bias", + "visual.blocks.12.mlp.up_proj", + "visual.blocks.12.mlp.up_proj.bias", + "visual.blocks.12.norm1", + "visual.blocks.12.norm2", + "visual.blocks.13.attn.k_norm", + "visual.blocks.13.attn.proj", + "visual.blocks.13.attn.proj.bias", + "visual.blocks.13.attn.q_norm", + "visual.blocks.13.attn.qkv", + "visual.blocks.13.attn.qkv.bias", + "visual.blocks.13.mlp.down_proj", + "visual.blocks.13.mlp.down_proj.bias", + "visual.blocks.13.mlp.gate_proj", + "visual.blocks.13.mlp.gate_proj.bias", + "visual.blocks.13.mlp.up_proj", + "visual.blocks.13.mlp.up_proj.bias", + "visual.blocks.13.norm1", + "visual.blocks.13.norm2", + "visual.blocks.14.attn.k_norm", + "visual.blocks.14.attn.proj", + "visual.blocks.14.attn.proj.bias", + "visual.blocks.14.attn.q_norm", + "visual.blocks.14.attn.qkv", + "visual.blocks.14.attn.qkv.bias", + "visual.blocks.14.mlp.down_proj", + "visual.blocks.14.mlp.down_proj.bias", + "visual.blocks.14.mlp.gate_proj", + "visual.blocks.14.mlp.gate_proj.bias", + "visual.blocks.14.mlp.up_proj", + "visual.blocks.14.mlp.up_proj.bias", + "visual.blocks.14.norm1", + "visual.blocks.14.norm2", + "visual.blocks.15.attn.k_norm", + "visual.blocks.15.attn.proj", + "visual.blocks.15.attn.proj.bias", + "visual.blocks.15.attn.q_norm", + "visual.blocks.15.attn.qkv", + "visual.blocks.15.attn.qkv.bias", + "visual.blocks.15.mlp.down_proj", + "visual.blocks.15.mlp.down_proj.bias", + "visual.blocks.15.mlp.gate_proj", + "visual.blocks.15.mlp.gate_proj.bias", + "visual.blocks.15.mlp.up_proj", + "visual.blocks.15.mlp.up_proj.bias", + "visual.blocks.15.norm1", + "visual.blocks.15.norm2", + "visual.blocks.16.attn.k_norm", + "visual.blocks.16.attn.proj", + "visual.blocks.16.attn.proj.bias", + "visual.blocks.16.attn.q_norm", + "visual.blocks.16.attn.qkv", + "visual.blocks.16.attn.qkv.bias", + "visual.blocks.16.mlp.down_proj", + "visual.blocks.16.mlp.down_proj.bias", + "visual.blocks.16.mlp.gate_proj", + "visual.blocks.16.mlp.gate_proj.bias", + "visual.blocks.16.mlp.up_proj", + "visual.blocks.16.mlp.up_proj.bias", + "visual.blocks.16.norm1", + "visual.blocks.16.norm2", + "visual.blocks.17.attn.k_norm", + "visual.blocks.17.attn.proj", + "visual.blocks.17.attn.proj.bias", + "visual.blocks.17.attn.q_norm", + "visual.blocks.17.attn.qkv", + "visual.blocks.17.attn.qkv.bias", + "visual.blocks.17.mlp.down_proj", + "visual.blocks.17.mlp.down_proj.bias", + "visual.blocks.17.mlp.gate_proj", + "visual.blocks.17.mlp.gate_proj.bias", + "visual.blocks.17.mlp.up_proj", + "visual.blocks.17.mlp.up_proj.bias", + "visual.blocks.17.norm1", + "visual.blocks.17.norm2", + "visual.blocks.18.attn.k_norm", + "visual.blocks.18.attn.proj", + "visual.blocks.18.attn.proj.bias", + "visual.blocks.18.attn.q_norm", + "visual.blocks.18.attn.qkv", + "visual.blocks.18.attn.qkv.bias", + "visual.blocks.18.mlp.down_proj", + "visual.blocks.18.mlp.down_proj.bias", + "visual.blocks.18.mlp.gate_proj", + "visual.blocks.18.mlp.gate_proj.bias", + "visual.blocks.18.mlp.up_proj", + "visual.blocks.18.mlp.up_proj.bias", + "visual.blocks.18.norm1", + "visual.blocks.18.norm2", + "visual.blocks.19.attn.k_norm", + "visual.blocks.19.attn.proj", + "visual.blocks.19.attn.proj.bias", + "visual.blocks.19.attn.q_norm", + "visual.blocks.19.attn.qkv", + "visual.blocks.19.attn.qkv.bias", + "visual.blocks.19.mlp.down_proj", + "visual.blocks.19.mlp.down_proj.bias", + "visual.blocks.19.mlp.gate_proj", + "visual.blocks.19.mlp.gate_proj.bias", + "visual.blocks.19.mlp.up_proj", + "visual.blocks.19.mlp.up_proj.bias", + "visual.blocks.19.norm1", + "visual.blocks.19.norm2", + "visual.blocks.2.attn.k_norm", + "visual.blocks.2.attn.proj", + "visual.blocks.2.attn.proj.bias", + "visual.blocks.2.attn.q_norm", + "visual.blocks.2.attn.qkv", + "visual.blocks.2.attn.qkv.bias", + "visual.blocks.2.mlp.down_proj", + "visual.blocks.2.mlp.down_proj.bias", + "visual.blocks.2.mlp.gate_proj", + "visual.blocks.2.mlp.gate_proj.bias", + "visual.blocks.2.mlp.up_proj", + "visual.blocks.2.mlp.up_proj.bias", + "visual.blocks.2.norm1", + "visual.blocks.2.norm2", + "visual.blocks.20.attn.k_norm", + "visual.blocks.20.attn.proj", + "visual.blocks.20.attn.proj.bias", + "visual.blocks.20.attn.q_norm", + "visual.blocks.20.attn.qkv", + "visual.blocks.20.attn.qkv.bias", + "visual.blocks.20.mlp.down_proj", + "visual.blocks.20.mlp.down_proj.bias", + "visual.blocks.20.mlp.gate_proj", + "visual.blocks.20.mlp.gate_proj.bias", + "visual.blocks.20.mlp.up_proj", + "visual.blocks.20.mlp.up_proj.bias", + "visual.blocks.20.norm1", + "visual.blocks.20.norm2", + "visual.blocks.21.attn.k_norm", + "visual.blocks.21.attn.proj", + "visual.blocks.21.attn.proj.bias", + "visual.blocks.21.attn.q_norm", + "visual.blocks.21.attn.qkv", + "visual.blocks.21.attn.qkv.bias", + "visual.blocks.21.mlp.down_proj", + "visual.blocks.21.mlp.down_proj.bias", + "visual.blocks.21.mlp.gate_proj", + "visual.blocks.21.mlp.gate_proj.bias", + "visual.blocks.21.mlp.up_proj", + "visual.blocks.21.mlp.up_proj.bias", + "visual.blocks.21.norm1", + "visual.blocks.21.norm2", + "visual.blocks.22.attn.k_norm", + "visual.blocks.22.attn.proj", + "visual.blocks.22.attn.proj.bias", + "visual.blocks.22.attn.q_norm", + "visual.blocks.22.attn.qkv", + "visual.blocks.22.attn.qkv.bias", + "visual.blocks.22.mlp.down_proj", + "visual.blocks.22.mlp.down_proj.bias", + "visual.blocks.22.mlp.gate_proj", + "visual.blocks.22.mlp.gate_proj.bias", + "visual.blocks.22.mlp.up_proj", + "visual.blocks.22.mlp.up_proj.bias", + "visual.blocks.22.norm1", + "visual.blocks.22.norm2", + "visual.blocks.23.attn.k_norm", + "visual.blocks.23.attn.proj", + "visual.blocks.23.attn.proj.bias", + "visual.blocks.23.attn.q_norm", + "visual.blocks.23.attn.qkv", + "visual.blocks.23.attn.qkv.bias", + "visual.blocks.23.mlp.down_proj", + "visual.blocks.23.mlp.down_proj.bias", + "visual.blocks.23.mlp.gate_proj", + "visual.blocks.23.mlp.gate_proj.bias", + "visual.blocks.23.mlp.up_proj", + "visual.blocks.23.mlp.up_proj.bias", + "visual.blocks.23.norm1", + "visual.blocks.23.norm2", + "visual.blocks.3.attn.k_norm", + "visual.blocks.3.attn.proj", + "visual.blocks.3.attn.proj.bias", + "visual.blocks.3.attn.q_norm", + "visual.blocks.3.attn.qkv", + "visual.blocks.3.attn.qkv.bias", + "visual.blocks.3.mlp.down_proj", + "visual.blocks.3.mlp.down_proj.bias", + "visual.blocks.3.mlp.gate_proj", + "visual.blocks.3.mlp.gate_proj.bias", + "visual.blocks.3.mlp.up_proj", + "visual.blocks.3.mlp.up_proj.bias", + "visual.blocks.3.norm1", + "visual.blocks.3.norm2", + "visual.blocks.4.attn.k_norm", + "visual.blocks.4.attn.proj", + "visual.blocks.4.attn.proj.bias", + "visual.blocks.4.attn.q_norm", + "visual.blocks.4.attn.qkv", + "visual.blocks.4.attn.qkv.bias", + "visual.blocks.4.mlp.down_proj", + "visual.blocks.4.mlp.down_proj.bias", + "visual.blocks.4.mlp.gate_proj", + "visual.blocks.4.mlp.gate_proj.bias", + "visual.blocks.4.mlp.up_proj", + "visual.blocks.4.mlp.up_proj.bias", + "visual.blocks.4.norm1", + "visual.blocks.4.norm2", + "visual.blocks.5.attn.k_norm", + "visual.blocks.5.attn.proj", + "visual.blocks.5.attn.proj.bias", + "visual.blocks.5.attn.q_norm", + "visual.blocks.5.attn.qkv", + "visual.blocks.5.attn.qkv.bias", + "visual.blocks.5.mlp.down_proj", + "visual.blocks.5.mlp.down_proj.bias", + "visual.blocks.5.mlp.gate_proj", + "visual.blocks.5.mlp.gate_proj.bias", + "visual.blocks.5.mlp.up_proj", + "visual.blocks.5.mlp.up_proj.bias", + "visual.blocks.5.norm1", + "visual.blocks.5.norm2", + "visual.blocks.6.attn.k_norm", + "visual.blocks.6.attn.proj", + "visual.blocks.6.attn.proj.bias", + "visual.blocks.6.attn.q_norm", + "visual.blocks.6.attn.qkv", + "visual.blocks.6.attn.qkv.bias", + "visual.blocks.6.mlp.down_proj", + "visual.blocks.6.mlp.down_proj.bias", + "visual.blocks.6.mlp.gate_proj", + "visual.blocks.6.mlp.gate_proj.bias", + "visual.blocks.6.mlp.up_proj", + "visual.blocks.6.mlp.up_proj.bias", + "visual.blocks.6.norm1", + "visual.blocks.6.norm2", + "visual.blocks.7.attn.k_norm", + "visual.blocks.7.attn.proj", + "visual.blocks.7.attn.proj.bias", + "visual.blocks.7.attn.q_norm", + "visual.blocks.7.attn.qkv", + "visual.blocks.7.attn.qkv.bias", + "visual.blocks.7.mlp.down_proj", + "visual.blocks.7.mlp.down_proj.bias", + "visual.blocks.7.mlp.gate_proj", + "visual.blocks.7.mlp.gate_proj.bias", + "visual.blocks.7.mlp.up_proj", + "visual.blocks.7.mlp.up_proj.bias", + "visual.blocks.7.norm1", + "visual.blocks.7.norm2", + "visual.blocks.8.attn.k_norm", + "visual.blocks.8.attn.proj", + "visual.blocks.8.attn.proj.bias", + "visual.blocks.8.attn.q_norm", + "visual.blocks.8.attn.qkv", + "visual.blocks.8.attn.qkv.bias", + "visual.blocks.8.mlp.down_proj", + "visual.blocks.8.mlp.down_proj.bias", + "visual.blocks.8.mlp.gate_proj", + "visual.blocks.8.mlp.gate_proj.bias", + "visual.blocks.8.mlp.up_proj", + "visual.blocks.8.mlp.up_proj.bias", + "visual.blocks.8.norm1", + "visual.blocks.8.norm2", + "visual.blocks.9.attn.k_norm", + "visual.blocks.9.attn.proj", + "visual.blocks.9.attn.proj.bias", + "visual.blocks.9.attn.q_norm", + "visual.blocks.9.attn.qkv", + "visual.blocks.9.attn.qkv.bias", + "visual.blocks.9.mlp.down_proj", + "visual.blocks.9.mlp.down_proj.bias", + "visual.blocks.9.mlp.gate_proj", + "visual.blocks.9.mlp.gate_proj.bias", + "visual.blocks.9.mlp.up_proj", + "visual.blocks.9.mlp.up_proj.bias", + "visual.blocks.9.norm1", + "visual.blocks.9.norm2", + "visual.downsample", + "visual.downsample.bias", + "visual.merger.down_proj", + "visual.merger.gate_proj", + "visual.merger.post_projection_norm", + "visual.merger.post_projection_norm.bias", + "visual.merger.proj", + "visual.merger.up_proj", + "visual.patch_embed.proj", + "visual.patch_embed.proj.bias", + "visual.post_layernorm", + "weights_proj" + ], + "quant_method": "fp8", + "weight_block_size": [ + 128, + 128 + ] + } +} diff --git a/tools/probes/g3_miles/harness.py b/tools/probes/g3_miles/harness.py new file mode 100644 index 000000000..59cf090fe --- /dev/null +++ b/tools/probes/g3_miles/harness.py @@ -0,0 +1,716 @@ +"""S25 T-G3 GPU harness: real Miles/Megatron glm5_next (4-layer mini) on 2 GPUs, one torchrun per phase. + +Phases (``--phase``): + +* ``n75`` -- N7.5: TP2+SP, EP2. Builds the model with the TE spec (``--impl te``) or the TE-free + fallback (``--impl local``), loads the weights saved by the first TE run (names mapped), + runs one forward/backward on a fixed packed sequence, saves the gathered logits, the + label log-probs and every parameter gradient, and times forward+backward (whole step and + the routed-expert module). Then runs N8.2 on the same weights (``n82`` below). +* ``n71d`` -- N7.1d: TP1, DP2, distributed optimizer + precision-aware CPU offload (real + HybridDeviceOptimizer: CUDA streams, pinned host memory, d2h/h2d overlap). Sub-modes + ``cont5`` (5 steps), ``run3`` (3 steps, then the yeto cut plugin saves a cut) and + ``resume2`` (fresh process: the cut plugin restores, N7.6 assertions, 2 steps). + +N8.2 (inside ``n75``): the reference gathers every weight to HF names through Miles' own weight-update +path (``named_params_and_buffers`` -> ``all_gather_params_async`` -> ``convert_to_hf``) and quantizes it +in one process; the distributed path gives each rank's experts through the yeto Megatron adapter +(``megatron_expert_slices``) and the non-expert weights once (rank 0, ``hf_slices``) to +``requantize_distributed``. Both outputs are compared byte for byte, for W4A8 (R2 AWQ) and MXFP4. + +Writes everything under ``--out``. Not a test: run by ``modal_g3.py`` on Modal ``H100!:2``. +""" + +from __future__ import annotations + +import argparse +import hashlib +import json +import os +import re +import sys +import time +import traceback +from pathlib import Path + +import torch +import torch.distributed as dist + +HERE = Path(__file__).resolve().parent +MILES = Path(os.environ.get("G3_MILES", "/g3/miles")) +sys.path.insert(0, str(MILES / "scripts" / "models")) + +SEQ = 1024 # tokens of the packed sequence (one sample) +TIMING_SEQ = 4096 +VOCAB = 16384 +HIDDEN = 1024 +EXPERTS = 8 + + +# ---------------------------------------------------------------------------------------------- +# Mini HF config (the 45-layer fixture shrunk: 4 layers = 3 KDA + 1 DSA, 1 dense + 3 MoE) +# ---------------------------------------------------------------------------------------------- + +def write_hf_config(src: Path, dst: Path) -> None: + cfg = json.loads(src.read_text()) + cfg.pop("quantization_config", None) + t = cfg["text_config"] + layer_types = ["linear_attention", "linear_attention", "linear_attention", "deepseek_sparse_attention"] + t.update( + hidden_size=HIDDEN, intermediate_size=3 * HIDDEN, moe_intermediate_size=512, num_hidden_layers=4, + first_k_dense_replace=1, n_routed_experts=EXPERTS, num_experts_per_tok=2, vocab_size=VOCAB, + index_topk=256, layer_types=layer_types, indexer_types=["full"] * 4, + mlp_layer_types=["dense", "sparse", "sparse", "sparse"], num_nextn_predict_layers=0, pad_token_id=0, + ) + t["linear_attn_config"] = dict(t["linear_attn_config"], num_heads=16, kda_layers=[0, 1, 2], full_attn_layers=[3]) + dst.mkdir(parents=True, exist_ok=True) + (dst / "config.json").write_text(json.dumps(cfg, indent=1)) + + +def megatron_argv(impl: str, *, tp: int, ep: int, hf: Path, offload: bool, lr: float) -> list[str]: + import importlib.util + + spec = importlib.util.spec_from_file_location("glm53_flash_args", MILES / "scripts" / "models" / "glm5.3-flash.py") + glm = importlib.util.module_from_spec(spec) + spec.loader.exec_module(glm) + model = glm.model_args(nlayers=4, first_k_dense_replace=1, + transformer_impl="transformer_engine" if impl == "te" else "local").split() + overrides = [ + "--hidden-size", str(HIDDEN), "--ffn-hidden-size", str(3 * HIDDEN), "--num-experts", str(EXPERTS), + "--moe-router-topk", "2", "--moe-ffn-hidden-size", "512", "--moe-shared-expert-intermediate-size", "512", + "--vocab-size", str(VOCAB), + ] + par = ["--tensor-model-parallel-size", str(tp), "--pipeline-model-parallel-size", "1", + "--context-parallel-size", "1", "--expert-model-parallel-size", str(ep), + "--expert-tensor-parallel-size", "1"] + if tp > 1: + par.append("--sequence-parallel") + misc = [ + "--qkv-format", "thd", "--micro-batch-size", "1", "--seq-length", str(TIMING_SEQ), + "--attention-dropout", "0.0", "--hidden-dropout", "0.0", "--tokenizer-type", "NullTokenizer", + "--hf-checkpoint", str(hf), "--load", "/nonexistent-g3", "--seed", "1234", + "--optimizer", "adam", "--lr", str(lr), "--lr-decay-style", "constant", "--min-lr", "0", + "--weight-decay", "0.1", "--adam-beta1", "0.9", "--adam-beta2", "0.98", "--clip-grad", "1.0", + "--attention-backend", "auto", "--no-check-for-nan-in-loss-and-grad", "--distributed-timeout-minutes", "10", + ] + if offload: + misc += ["--optimizer-cpu-offload", "--use-precision-aware-optimizer", "--overlap-cpu-optimizer-d2h-h2d", + "--optimizer-offload-fraction", "1.0"] + return model + overrides + par + misc + + +def init_miles(argv: list[str], *, global_batch: int, rollouts: int): + from megatron.training.arguments import parse_args, validate_args + from miles.backends.megatron_utils.arguments import set_default_megatron_args + from miles.backends.megatron_utils.initialize import init + from miles.utils.arguments import get_miles_extra_args_provider, hf_validate_args + from miles.utils.hf_utils.config import load_hf_config + + sys.argv = ["g3"] + argv + ["--global-batch-size", str(global_batch)] + args = parse_args(extra_args_provider=get_miles_extra_args_provider()) + hf_config = load_hf_config(args.hf_checkpoint) + hf_validate_args(args, hf_config) + text = getattr(hf_config, "text_config", None) or hf_config + args.indexer_rope_interleave = bool(getattr(text, "indexer_rope_interleave", False)) + args.compress_ratios = None + # Miles' scheduler sizing (get_optimizer_param_scheduler): train_iters = rollouts*rbs*n // gbs + args.num_rollout, args.rollout_batch_size, args.n_samples_per_prompt = rollouts, 1, global_batch + args.__dict__.setdefault("decrease_batch_size_if_needed", False) + set_default_megatron_args(args) + dist.init_process_group(backend="nccl") + torch.cuda.set_device(int(os.environ.get("LOCAL_RANK", 0))) + validate_args(args) + init(args) + return args + + +def build(args): + from miles.backends.megatron_utils.model import setup_model_and_optimizer, setup_train_iteration_config + + model, optimizer, scheduler = setup_model_and_optimizer(args, "actor") + setup_train_iteration_config(args, model, optimizer, False) + return model, optimizer, scheduler + + +# ---------------------------------------------------------------------------------------------- +# Data, forward/backward +# ---------------------------------------------------------------------------------------------- + +def tokens_for(step: int, dp_rank: int, n: int) -> torch.Tensor: + g = torch.Generator().manual_seed(1000 * step + dp_rank) + return torch.randint(0, VOCAB, (n,), generator=g) + + +def packed(tokens: torch.Tensor): + from megatron.core.packed_seq_params import PackedSeqParams + + n = tokens.numel() + cu = torch.tensor([0, n], dtype=torch.int32, device="cuda") + psp = PackedSeqParams(cu_seqlens_q=cu, cu_seqlens_kv=cu, max_seqlen_q=n, max_seqlen_kv=n, qkv_format="thd") + return tokens.view(1, -1).cuda(), psp + + +def run_fwd_bwd(model, tokens: torch.Tensor, *, capture: list | None = None, forward_only=False): + """One micro-batch through Megatron's schedule; loss = mean next-token CE over the sequence.""" + from megatron.core import tensor_parallel + from megatron.core.pipeline_parallel import get_forward_backward_func + + ids, psp = packed(tokens) + labels = torch.roll(ids, -1, dims=1) + mask = torch.ones_like(ids, dtype=torch.float32) + mask[:, -1] = 0 + + def loss_func(output): + token_loss = tensor_parallel.vocab_parallel_cross_entropy(output.float(), labels) + loss = (token_loss * mask).sum() / mask.sum() + if capture is not None: + capture.append((output.detach(), token_loss.detach())) + return loss, {"loss": loss.detach()} + + def step(data_iter, chunk): + next(data_iter) + out = chunk(input_ids=ids, position_ids=None, attention_mask=None, labels=None, packed_seq_params=psp) + return out, loss_func + + fb = get_forward_backward_func() + res = fb(forward_step_func=step, data_iterator=iter([None]), model=model, num_microbatches=1, + seq_length=ids.shape[1], micro_batch_size=1, forward_only=forward_only) + return float(res[0]["loss"]) if res else float("nan") + + +def zero(model, optimizer): + for chunk in model: + chunk.zero_grad_buffer() + if optimizer is not None: + optimizer.zero_grad() + + +# ---------------------------------------------------------------------------------------------- +# Names: TE spec <-> local spec +# ---------------------------------------------------------------------------------------------- + +def named(model) -> dict[str, torch.Tensor]: + out = {} + for n, p in model[0].named_parameters(): + out[n] = p + for n, b in model[0].named_buffers(): + if "expert_bias" in n: + out[n] = b + return out + + +def local_to_te(name: str) -> str: + from miles.backends.megatron_utils.named_weights import local_expert_name_to_grouped + + name = local_expert_name_to_grouped(name) + # dense MLP: TE fuses the pre-MLP RMSNorm into linear_fc1 (layer_norm_weight) + name = re.sub(r"(decoder\.layers\.\d+)\.pre_mlp_layernorm\.weight$", r"\1.PRE_MLP_NORM", name) + return name + + +def te_canonical(name: str, te_names: set[str]) -> str: + return re.sub(r"(decoder\.layers\.\d+)\.mlp\.linear_fc1\.layer_norm_weight$", r"\1.PRE_MLP_NORM", name) + + +def map_local_to_te(local_names, te_names) -> tuple[dict[str, str], list[str], list[str]]: + canon_te = {te_canonical(n, te_names): n for n in te_names} + mapping, missing = {}, [] + for n in local_names: + c = local_to_te(n) + if c in canon_te: + mapping[n] = canon_te[c] + elif c.endswith("PRE_MLP_NORM") and c.replace("PRE_MLP_NORM", "pre_mlp_layernorm.weight") in canon_te: + mapping[n] = canon_te[c.replace("PRE_MLP_NORM", "pre_mlp_layernorm.weight")] # MoE layers keep the name + else: + missing.append(n) + unused = sorted(set(te_names) - set(mapping.values())) + return mapping, missing, unused + + +# ---------------------------------------------------------------------------------------------- +# N7.5 phase +# ---------------------------------------------------------------------------------------------- + +def rank_file(out: Path, stem: str) -> Path: + return out / f"{stem}.rank{dist.get_rank()}.pt" + + +def log(out: Path, msg: str) -> None: + line = f"[{time.strftime('%H:%M:%S')}] rank{dist.get_rank() if dist.is_initialized() else '-'} {msg}" + print(line, flush=True) + if not dist.is_initialized() or dist.get_rank() == 0: + with open(out / "progress.log", "a") as fh: + fh.write(line + "\n") + + +def phase_n75(ns) -> None: + out = Path(ns.out) + out.mkdir(parents=True, exist_ok=True) + args = init_miles(megatron_argv(ns.impl, tp=2, ep=2, hf=Path(ns.hf), offload=False, lr=1e-5), + global_batch=1, rollouts=1) + log(out, f"n75 impl={ns.impl} tag={ns.tag} transformer_impl={args.transformer_impl}") + model, optimizer, _ = build(args) + params = named(model) + info = {"impl": ns.impl, "tag": ns.tag, "transformer_impl": args.transformer_impl, + "module_classes": sorted({type(m).__name__ for m in model[0].modules()}), + "num_params_local": sum(p.numel() for p in params.values()), + "gpu": torch.cuda.get_device_name(), "names": sorted(params)} + core = [type(m).__name__ for n, m in model[0].named_modules() if n.endswith("self_attention.core_attention")] + info["dsa_core_attention"] = core + weights_file = Path(ns.weights_from) / f"weights.rank{dist.get_rank()}.pt" if ns.weights_from else None + if weights_file is None: + torch.save({n: p.detach().cpu() for n, p in params.items()}, rank_file(out, "weights")) + info["weights"] = "initialized here (seed 1234)" + else: + saved = torch.load(weights_file, map_location="cpu", weights_only=True) + mapping, missing, unused = map_local_to_te(params, set(saved)) + info["name_map"] = {"mapped": len(mapping), "missing": missing, "unused_te": unused, + "renamed": {k: v for k, v in mapping.items() if k != v}} + if missing or unused: + raise RuntimeError(f"weight names do not map: missing={missing[:8]} unused={unused[:8]}") + with torch.no_grad(): + for n, p in params.items(): + src = saved[mapping[n]] + if tuple(src.shape) != tuple(p.shape): + raise RuntimeError(f"{n}: shape {tuple(p.shape)} != saved {tuple(src.shape)}") + p.copy_(src.to(p.device, p.dtype)) + # the DDP/fp32 main copies are rebuilt from the model weights by the optimizer + if optimizer is not None and hasattr(optimizer, "reload_model_params"): + optimizer.reload_model_params() + info["weights"] = f"loaded from {ns.weights_from}" + + # one forward/backward (no optimizer step) + zero(model, optimizer) + cap: list = [] + loss = run_fwd_bwd(model, tokens_for(0, 0, SEQ), capture=cap) + torch.cuda.synchronize() + logits_local, token_loss = cap[0] + from megatron.core import parallel_state as mps + from megatron.core.tensor_parallel.mappings import gather_from_tensor_model_parallel_region + + logits = gather_from_tensor_model_parallel_region(logits_local.float()) # [1, T, V] + grads = {} + for n, p in params.items(): + g = getattr(p, "main_grad", None) + if g is None: + g = p.grad + if g is not None: + grads[n] = g.detach().float().cpu() + torch.save(grads, rank_file(out, f"grads-{ns.tag}")) + if dist.get_rank() == 0: + torch.save({"logits": logits[0].cpu(), "token_loss": token_loss.float().cpu().view(-1), "loss": loss}, + out / f"logits-{ns.tag}.pt") + info["loss"] = loss + info["grad_tensors"] = len(grads) + info["tp_rank"], info["ep_rank"] = mps.get_tensor_model_parallel_rank(), mps.get_expert_model_parallel_rank() + log(out, f"n75 {ns.tag}: loss={loss:.6f} grads={len(grads)}") + + # timing: forward+backward on TIMING_SEQ tokens; routed experts timed with CUDA events + events: list = [] + + def pre(mod, inp): + e = torch.cuda.Event(enable_timing=True) + e.record() + mod._g3_t0 = e + + def post(mod, inp, outp): + e = torch.cuda.Event(enable_timing=True) + e.record() + events.append((mod._g3_t0, e)) + + hooks = [] + for n, m in model[0].named_modules(): + if n.endswith("mlp.experts"): + hooks += [m.register_forward_pre_hook(pre), m.register_forward_hook(post)] + steps = [] + for i in range(7): + zero(model, optimizer) + events.clear() + torch.cuda.synchronize() + t0 = time.perf_counter() + run_fwd_bwd(model, tokens_for(100 + i, 0, TIMING_SEQ)) + torch.cuda.synchronize() + steps.append({"step_s": time.perf_counter() - t0, + "experts_fwd_ms": sum(a.elapsed_time(b) for a, b in events)}) + for h in hooks: + h.remove() + info["timing"] = {"tokens": TIMING_SEQ, "warmup": 2, "steps": steps, + "peak_mem_gib": torch.cuda.max_memory_allocated() / 2**30} + zero(model, optimizer) + + # N8.2 on the same (unchanged) weights + try: + info["n82"] = n82(args, model, out / f"n82-{ns.tag}") + except Exception as exc: # noqa: BLE001 -- recorded; the phase result says what failed + info["n82"] = {"error": f"{type(exc).__name__}: {exc}", "traceback": traceback.format_exc()[-4000:]} + log(out, f"n82 {ns.tag} FAILED: {exc}") + all_info = [None] * dist.get_world_size() + dist.all_gather_object(all_info, info) + if dist.get_rank() == 0: + (out / f"info-{ns.tag}.json").write_text(json.dumps(all_info, indent=1, default=str)) + dist.barrier() + + +# ---------------------------------------------------------------------------------------------- +# N8.2 +# ---------------------------------------------------------------------------------------------- + +def nest_language_model(name: str) -> str: + """GLM-5.3 checkpoint names keep the text model under ``model.language_model.`` (mbridge glm5_next).""" + if name.startswith("model.") and not name.startswith(("model.language_model.", "model.visual.")): + return "model.language_model." + name[len("model."):] + return name + + +EXPERT_HF = re.compile(r"\.mlp\.experts\.\d+\.(gate|up|down)_proj\.weight$") + + +def miles_hf_weights(args, model) -> tuple[dict[str, torch.Tensor], dict[str, torch.Tensor]]: + """Miles' weight-update path: global names, TP all-gather, convert_to_hf. Returns (non-expert on every + rank, this EP rank's experts), HF names as Miles emits them (no nesting).""" + from miles.backends.megatron_utils.megatron_to_hf import convert_to_hf + from miles.backends.megatron_utils.named_weights import named_params_and_buffers + from miles.backends.megatron_utils.update_weight.hf_weight_iterator_direct import ( + all_gather_params_async, + is_routed_expert_param, + ) + from miles.utils.types import ParamInfo + + dense, experts = {}, {} + for name, param in named_params_and_buffers(args, model): + pinfo = ParamInfo(name=name, dtype=param.dtype, shape=param.shape, attrs={}, size=0, src_rank=dist.get_rank()) + full = all_gather_params_async(args, [(pinfo, param)])[0] + for hf_name, t in convert_to_hf(args, "glm5_next", name, full): + (experts if is_routed_expert_param(name) else dense)[hf_name] = t.detach().cpu().clone() + return dense, experts + + +def fake_original(source: dict[str, torch.Tensor], scheme: str, out: Path, smooth: str | None) -> Path: + """An "original" quantized checkpoint (headers matter, values do not) with the source's names.""" + from safetensors.torch import save_file + + out.mkdir(parents=True, exist_ok=True) + tensors = {} + for n, t in source.items(): + if EXPERT_HF.search(n): + mod = n[: -len(".weight")] + rows, cols = t.shape + if scheme == "w4a8": + tensors[n] = torch.zeros(rows, cols, dtype=torch.int8) + tensors[mod + ".weight_scale"] = torch.ones(rows, 1, dtype=torch.bfloat16) + else: + tensors[mod + ".weight_packed"] = torch.zeros(rows, cols // 2, dtype=torch.uint8) + tensors[mod + ".weight_scale"] = torch.zeros(rows, cols // 32, dtype=torch.uint8) + elif n.endswith("mlp.gate.weight"): + tensors[n] = t.float() # the router is stored in F32 + else: + tensors[n] = t + shard = lambda n: ("model-00002-of-00002.safetensors" if re.search(r"layers\.[23]\.mlp\.experts", n) # noqa: E731 + else "model-00001-of-00002.safetensors") + files: dict[str, dict] = {} + for n, t in tensors.items(): + files.setdefault(shard(n), {})[n] = t.contiguous() + for f, ts in files.items(): + save_file(ts, str(out / f), metadata={"format": "pt"}) + (out / "model.safetensors.index.json").write_text(json.dumps({"metadata": {}, "weight_map": {n: shard(n) for n in tensors}})) + if scheme == "w4a8": + q = {"quant_method": "compressed-tensors", "format": "int-quantized", "config_groups": {"group_0": { + "targets": ["re:.*mlp\\.experts\\..*(gate|up|down)_proj$"], + "weights": {"num_bits": 4, "type": "int", "symmetric": True, "strategy": "channel", "group_size": None, + "dynamic": False, "observer": "memoryless_minmax"}, + "input_activations": {"num_bits": 8, "strategy": "token", "dynamic": True}}}} + else: + q = {"quant_method": "compressed-tensors", "format": "mxfp4-pack-quantized", "config_groups": {"MXFP4A16": { + "format": "mxfp4-pack-quantized", "targets": ["re:.*mlp\\.experts\\.\\d+\\.(gate|up|down)_proj$"], + "weights": {"num_bits": 4, "type": "float", "symmetric": True, "strategy": "group", "group_size": 32, + "dynamic": False, "scale_dtype": "torch.uint8", "observer": None}, + "input_activations": None}}, "ignore": ["lm_head", "re:^(?!.*mlp\\.experts\\.\\d+\\.).*$"]} + (out / "config.json").write_text(json.dumps({"model_type": "glm5_next", "quantization_config": q})) + if scheme == "w4a8" and smooth: + layer = smooth.rsplit(".post_attention_layernorm", 1)[0] + esc = re.escape(layer) + (out / "recipe.yaml").write_text( + "default_stage:\n default_modifiers:\n RecordedAWQ:\n mappings:\n" + f" - smooth_layer: {smooth}\n" + f" balance_layers: [{layer}.mlp.gate, 're:{esc}\\.mlp\\.(experts\\.\\d+|shared_experts)\\.(gate|up)_proj$']\n" + " duo_scaling: true\n QuantizationModifier:\n ignore: [lm_head]\n") + return out + + +def n82(args, model, out: Path) -> dict: + from megatron.core import parallel_state as mps + from safetensors.torch import save_file + + from yeto.rl.adapters.miles.quant_shards import hf_slices, megatron_expert_slices + from yeto.rl.quant.awq import AwqSmoothing + from yeto.rl.quant.requant import CheckpointLayout, requantize_checkpoint + from yeto.rl.quant.shards import SourceSlice, requantize_distributed, torch_barrier + + rank, world = dist.get_rank(), dist.get_world_size() + ep_rank, ep_size = mps.get_expert_model_parallel_rank(), mps.get_expert_model_parallel_world_size() + if mps.get_data_parallel_world_size() != 1: + raise RuntimeError("n82 harness expects DP=1 (every rank holds a distinct EP slice)") + out.mkdir(parents=True, exist_ok=True) + dense, experts = miles_hf_weights(args, model) + gathered = [None] * world + dist.all_gather_object(gathered, experts) + res: dict = {"miles_hf_names_sample": sorted(dense)[:6] + sorted(experts)[:3], + "dense_tensors": len(dense), "expert_tensors_this_rank": len(experts)} + # reference source: nested names (the checkpoint layout); every rank builds it (same bytes) + source = {nest_language_model(n): t for n, t in dense.items()} + for part in gathered: + source.update({nest_language_model(n): t for n, t in part.items()}) + smooth = next((n[: -len(".weight")] for n in source if n.endswith("layers.1.post_attention_layernorm.weight")), None) + res["awq_smooth_layer"] = smooth + local = {n: p for n, p in model[0].named_parameters()} + num_local = args.num_experts // ep_size + for scheme in ("w4a8", "mxfp4"): + d = out / scheme + r: dict = {} + try: + if rank == 0: + orig = fake_original(source, scheme, d / "original", smooth) + CheckpointLayout.from_checkpoint(orig).save_template(orig, d / "template") + (d / "bf16").mkdir(parents=True, exist_ok=True) + save_file({n: t.contiguous() for n, t in source.items()}, str(d / "bf16" / "model.safetensors"), + metadata={"format": "pt"}) + if scheme == "w4a8" and smooth: + (d / "awq-scales").mkdir(exist_ok=True) + s = (0.8 + 0.4 * torch.rand(source[smooth + ".weight"].numel(), + generator=torch.Generator().manual_seed(7))).double() + torch.save(s, d / "awq-scales" / f"{smooth}.pt") + dist.barrier() + smoothing = (AwqSmoothing.from_artefact(d / "original" / "recipe.yaml", d / "awq-scales") + if scheme == "w4a8" and smooth else AwqSmoothing.none()) + r["awq"] = bool(scheme == "w4a8" and smooth) + if rank == 0: + t0 = time.perf_counter() + requantize_checkpoint(d / "bf16", d / "template", d / "reference", smoothing) + r["single_process_s"] = time.perf_counter() - t0 + dist.barrier() + # distributed: experts from the yeto adapter on Megatron-local names; non-experts once (rank 0) + mine = list(megatron_expert_slices(((n.removeprefix("module.").removeprefix("module."), p) + for n, p in local.items()), + ep_rank=ep_rank, num_local_experts=num_local)) + mine_cpu = [SourceSlice(x.name, x.tensor.cpu(), x.row_start) for x in mine] + r["expert_slices_this_rank"] = len(mine) + r["expert_slice_fields"] = sorted(getattr(mine[0], "__dataclass_fields__", {})) if mine else [] + # raw: convert_to_hf names as they are; yeto: the adapter's hf_slices (renames to checkpoint + # names), experts first as the GPU tensors a trainer holds, then copied to the host + for variant in ("raw_convert_to_hf_names", "yeto_hf_slices_gpu_experts", "yeto_hf_slices_cpu_experts"): + dst = d / f"dist-{variant}" + slices = list(mine if variant == "yeto_hf_slices_gpu_experts" else mine_cpu) + if rank == 0: + slices += (list(hf_slices(dense.items())) if variant.startswith("yeto_hf_slices") + else [SourceSlice(n, t) for n, t in dense.items()]) + try: + t0 = time.perf_counter() + stats = requantize_distributed(slices, d / "template", dst, smoothing, rank=rank, + world_size=world, barrier=torch_barrier) + r[variant] = {"ok": True, "s": time.perf_counter() - t0, "part": stats.get("part")} + except Exception as exc: # noqa: BLE001 + r[variant] = {"ok": False, "error": f"{type(exc).__name__}: {str(exc)[:600]}"} + # a rank that failed early leaves the others in a barrier: resync before the next variant + dist.barrier() + if rank == 0 and r[variant]["ok"]: + r[variant]["byte_equal"] = compare_dirs(d / "reference", dst) + except Exception as exc: # noqa: BLE001 + r["error"] = f"{type(exc).__name__}: {exc}" + r["traceback"] = traceback.format_exc()[-3000:] + res[scheme] = r + return res + + +def compare_dirs(a: Path, b: Path) -> dict: + names_a = sorted(p.name for p in a.iterdir() if p.is_file()) + names_b = sorted(p.name for p in b.iterdir() if p.is_file()) + diff = [n for n in names_a if n in names_b and (a / n).read_bytes() != (b / n).read_bytes()] + return {"files": names_a, "same_file_set": names_a == names_b, "differing": diff, + "equal": names_a == names_b and not diff, + "sha256": {n: hashlib.sha256((a / n).read_bytes()).hexdigest()[:16] for n in names_a}} + + +# ---------------------------------------------------------------------------------------------- +# N7.1d phase +# ---------------------------------------------------------------------------------------------- + +class Actor: + """The attributes of a Miles MegatronTrainRayActor that the yeto cut plugin reads.""" + + def __init__(self, args, model, optimizer, scheduler): + self.args, self.model, self.optimizer, self.opt_param_scheduler = args, model, optimizer, scheduler + + +def hdo_report(optimizer) -> dict: + try: + return _hdo_report(optimizer) + except Exception as exc: # noqa: BLE001 -- diagnostics only + return {"error": f"{type(exc).__name__}: {exc}"} + + +def _hdo_report(optimizer) -> dict: + from megatron.core.optimizer.cpu_offloading import HybridDeviceOptimizer + + leaves = getattr(optimizer, "chained_optimizers", [optimizer]) + rep = [] + for leaf in leaves: + inner = getattr(leaf, "optimizer", None) + item = {"leaf": type(leaf).__name__, "inner": type(inner).__name__, + "is_hdo": isinstance(inner, HybridDeviceOptimizer)} + if isinstance(inner, HybridDeviceOptimizer): + fp32 = list(inner.param_to_fp32_param.values()) if getattr(inner, "param_to_fp32_param", None) else [] + item.update( + fp32_params=len(fp32), fp32_devices=sorted({str(p.device) for p in fp32}), + fp32_pinned=sum(bool(p.is_pinned()) for p in fp32 if p.device.type == "cpu"), + fp32_dtypes=sorted({str(p.dtype) for p in fp32}), + cpu_optimizers=[type(o).__name__ for o in inner.cpu_optimizers], + overlap=bool(inner.defaults.get("overlap_cpu_optimizer_d2h_h2d")), + state_dtypes=sorted({str(v.dtype) for o in inner.cpu_optimizers for s in o.state.values() + for v in s.values() if torch.is_tensor(v) and v.numel() > 1}), + steps=sorted({float(s["step"]) for o in inner.cpu_optimizers for s in o.state.values() if "step" in s}), + ) + rep.append(item) + return {"leaves": rep} + + +def param_digest(model) -> dict: + h = hashlib.sha256() + norms = {} + for n, p in sorted(model[0].named_parameters()): + t = p.detach().float().cpu().contiguous() + h.update(n.encode()) + h.update(p.detach().cpu().contiguous().view(torch.uint8).numpy().tobytes()) + norms[n] = float(t.norm()) + return {"sha256": h.hexdigest(), "norms": norms} + + +def train_steps(args, model, optimizer, scheduler, first: int, last: int, out: Path) -> list[dict]: + from megatron.core import parallel_state as mps + + dp = mps.get_data_parallel_rank() + rec = [] + for step in range(first, last + 1): + zero(model, optimizer) + t0 = time.perf_counter() + loss = run_fwd_bwd(model, tokens_for(step, dp, SEQ)) + ok, grad_norm, _ = optimizer.step() + assert ok + scheduler.step(increment=args.global_batch_size) + torch.cuda.synchronize() + lt = torch.tensor([loss], device="cuda") + dist.all_reduce(lt) + rec.append({"step": step, "loss_mean_dp": float(lt) / dist.get_world_size(), "loss_rank": loss, + "grad_norm": float(grad_norm), "lr": [g["lr"] for g in optimizer.param_groups], + "scheduler_num_steps": int(scheduler.num_steps), "s": time.perf_counter() - t0}) + log(out, f"step {step} loss={loss:.6f} grad_norm={float(grad_norm):.6f}") + zero(model, optimizer) + return rec + + +def phase_n71d(ns) -> None: + from yeto.rl.adapters.miles import cut_plugin + from yeto.rl.engine.resume_manifest import ( + OptimizerResumeState, + ResumeManifest, + first_step_problems, + require_saveable, + restore_problems, + ) + + out = Path(ns.out) + out.mkdir(parents=True, exist_ok=True) + args = init_miles(megatron_argv("te", tp=1, ep=1, hf=Path(ns.hf), offload=True, lr=1e-4), + global_batch=2, rollouts=10) + flags = {k: getattr(args, k, None) for k in ( + "use_distributed_optimizer", "optimizer_cpu_offload", "use_precision_aware_optimizer", + "overlap_cpu_optimizer_d2h_h2d", "optimizer_offload_fraction", "main_params_dtype", "main_grads_dtype", + "exp_avg_dtype", "exp_avg_sq_dtype", "data_parallel_size", "bf16", "transformer_impl")} + log(out, f"n71d {ns.mode} flags={flags}") + model, optimizer, scheduler = build(args) + actor = Actor(args, model, optimizer, scheduler) + res: dict = {"mode": ns.mode, "flags": {k: str(v) for k, v in flags.items()}, "gpu": torch.cuda.get_device_name()} + res["hdo_before"] = hdo_report(optimizer) + cut_dir = Path(ns.cut_dir) + + def gather(obj): + lst = [None] * dist.get_world_size() + dist.all_gather_object(lst, obj) + return lst + + def manifest(summaries, global_step): + return ResumeManifest(global_step=global_step, data_watermark=global_step * args.global_batch_size, + scheduler_increment=args.global_batch_size, + ranks={str(s["path"]): OptimizerResumeState.from_dict(s["resume"]) for s in summaries}) + + if ns.mode in ("cont5", "cont5b"): + res["steps"] = train_steps(args, model, optimizer, scheduler, 1, 5, out) + elif ns.mode == "run3": + res["steps"] = train_steps(args, model, optimizer, scheduler, 1, 3, out) + summary = cut_plugin.save_cut_shard(actor, directory=str(cut_dir), cut_id="g3-cut-3") + if cut_plugin.REFUSED in summary: + raise RuntimeError(f"save_cut refused: {summary}") + summaries = gather(summary) + res["cut_summary"] = {k: summary[k] for k in ("path", "bytes", "scheduler_samples", "tuning", + "adapter_tensors", "resume")} + if dist.get_rank() == 0: + m = manifest(summaries, 3) + require_saveable(m) + (cut_dir / "files.json").write_text(json.dumps(summaries, default=str)) + (cut_dir / "resume_manifest.json").write_text(json.dumps(m.to_dict())) + res["hdo_after_save"] = hdo_report(optimizer) + elif ns.mode == "resume2": + files = json.loads((cut_dir / "files.json").read_text()) + saved = ResumeManifest.from_dict(json.loads((cut_dir / "resume_manifest.json").read_text())) + summary = cut_plugin.restore_cut_shard(actor, directory=str(cut_dir), files=files, cut_id="g3-cut-3") + if cut_plugin.REFUSED in summary: + raise RuntimeError(f"restore_cut refused: {summary}") + res["restore_summary"] = {k: summary.get(k) for k in ("path", "scheduler_samples", "resume", "diff", + "optimizer_cut_vs_reexport")} + res["restore_digest_equal"] = summary.get("state_digest") == next( + f for f in files if f["path"] == summary["path"])["state_digest"] + res["hdo_after_restore"] = hdo_report(optimizer) + now = manifest(gather(summary), 3) + res["restore_problems"] = restore_problems(saved, now, check_data_watermark=True) + res["steps"] = train_steps(args, model, optimizer, scheduler, 4, 4, out) + st = cut_plugin.resume_state(actor) + after1 = manifest(gather(st), 4) + res["first_step_problems"] = first_step_problems(saved, after1, steps=1) + res["hdo_after_first_step"] = hdo_report(optimizer) + res["steps"] += train_steps(args, model, optimizer, scheduler, 5, 5, out) + else: + raise ValueError(ns.mode) + res["final_params"] = param_digest(model) + res["hdo_final"] = hdo_report(optimizer) + res["peak_mem_gib"] = torch.cuda.max_memory_allocated() / 2**30 + allres = gather(res) + if dist.get_rank() == 0: + (out / f"n71d-{ns.mode}.json").write_text(json.dumps(allres, indent=1, default=str)) + dist.barrier() + + +def main() -> None: + p = argparse.ArgumentParser() + p.add_argument("--phase", required=True, choices=["hfconfig", "n75", "n71d"]) + p.add_argument("--out", required=True) + p.add_argument("--hf", default="/tmp/g3-hf") + p.add_argument("--impl", default="te", choices=["te", "local"]) + p.add_argument("--tag", default="te") + p.add_argument("--weights-from", default="") + p.add_argument("--mode", default="cont5") + p.add_argument("--cut-dir", default="") + p.add_argument("--fixture", default=str(HERE / "glm53flash_config.json")) + ns = p.parse_args() + if ns.phase == "hfconfig": + write_hf_config(Path(ns.fixture), Path(ns.hf)) + return + try: + (phase_n75 if ns.phase == "n75" else phase_n71d)(ns) + except Exception: + traceback.print_exc() + sys.stdout.flush() + os._exit(1) + dist.destroy_process_group() + + +if __name__ == "__main__": + main() diff --git a/tools/probes/g3_miles/modal_g3.py b/tools/probes/g3_miles/modal_g3.py new file mode 100644 index 000000000..fb7e4500b --- /dev/null +++ b/tools/probes/g3_miles/modal_g3.py @@ -0,0 +1,111 @@ +"""Modal launcher for the T-G3 chain (N7.5 / N8.2 / N7.1d) on ``H100!:2``. + +usage (from this directory, after the pre-registration in gpu-spend.md): + G3_RUN= G3_MILES= G3_YETO= \\ + modal run --detach modal_g3.py [--phases all|n75|n71d] + +The entrypoint only ``spawn()``s the chain and returns (killing a detached client would +cancel a blocking ``.remote()``). Results land in the Modal Volume ``t-g3-s25`` under +``//``; pull them with ``modal volume get t-g3-s25 / ``. +""" + +from __future__ import annotations + +import os + +import modal + +RUN = os.environ.get("G3_RUN", "unset") +MILES = os.environ.get("G3_MILES", "/home/michael/work/miles-s25-g3") +YETO = os.environ.get("G3_YETO", "/home/michael/work/s25-g3-miles") +HERE = os.path.dirname(os.path.abspath(__file__)) +IMAGE_REF = "radixark/miles:v0.1.1" # Miles 2806267d + Megatron-LM f148a32b4 (mHC output contraction), TE 2.17 +GPU = os.environ.get("G3_GPU", "H100!:2") # preflight: G3_GPU=L4:2 G3_EXPECT="NVIDIA L4" G3_HARD=900 +EXPECT_GPU = os.environ.get("G3_EXPECT", "NVIDIA H100 80GB HBM3") +HARD_S = int(os.environ.get("G3_HARD", "3600")) + +image = ( + modal.Image.from_registry(IMAGE_REF).entrypoint([]) + .env({"PYTHONPATH": "/g3/miles:/g3/yeto", "CUDA_DEVICE_MAX_CONNECTIONS": "1", "G3_MILES": "/g3/miles", + "PYTHONUNBUFFERED": "1"}) + .add_local_dir(MILES, "/g3/miles", copy=False, ignore=[".git", "**/__pycache__", "tests", "docs", "examples"]) + .add_local_dir(os.path.join(YETO, "yeto"), "/g3/yeto/yeto", copy=False, ignore=["**/__pycache__"]) + .add_local_dir(HERE, "/g3/harness", copy=False, ignore=["**/__pycache__"]) +) +vol = modal.Volume.from_name("t-g3-s25", create_if_missing=True) +app = modal.App(f"t-g3-{RUN}") + +CHAIN = r""" +set -uo pipefail +exec > >(tee -a "$OUT/chain.log") 2>&1 +H=/g3/harness/harness.py +run() { # name timeout cmd... + local name=$1 t=$2; shift 2 + echo "=== PHASE $name start $(date -u +%T)" + timeout "$t" "$@" > "$OUT/logs/$name.log" 2>&1; local rc=$? + echo "=== PHASE $name rc=$rc end $(date -u +%T)"; tail -5 "$OUT/logs/$name.log" + echo "$name $rc" >> "$OUT/phases.txt" + return $rc +} +mkdir -p "$OUT/logs" "$OUT/n75" "$OUT/n71d" +nvidia-smi --query-gpu=index,name,driver_version,memory.total --format=csv | tee "$OUT/nvidia-smi.csv" +n=$(nvidia-smi --query-gpu=name --format=csv,noheader | grep -c "^$EXPECT_GPU\$") +[ "$n" = 2 ] || { echo "GPU ASSERT FAILED: want 2x $EXPECT_GPU"; exit 3; } +pip list 2>/dev/null | grep -iE '^(torch|transformer.engine|megatron|flash.linear|tilelang|mbridge) ' > "$OUT/versions.txt" +git -C /root/Megatron-LM rev-parse HEAD >> "$OUT/versions.txt" +python $H --phase hfconfig --out "$OUT" --hf /tmp/g3-hf && cp /tmp/g3-hf/config.json "$OUT/hf-config.json" +TR="torchrun --nproc-per-node 2 --master-port 29511" +if [[ "$PHASES" == pre ]]; then + run pre-n75-te 600 $TR $H --phase n75 --impl te --tag te --out "$OUT/n75" + run pre-n71d 400 $TR $H --phase n71d --mode cont5 --out "$OUT/n71d" --cut-dir "$OUT/n71d/cut" +fi +if [[ "$PHASES" == all || "$PHASES" == n75 ]]; then + run n75-te 1200 $TR $H --phase n75 --impl te --tag te --out "$OUT/n75" + run n75-te2 900 $TR $H --phase n75 --impl te --tag te2 --weights-from "$OUT/n75" --out "$OUT/n75" + run n75-local 1200 $TR $H --phase n75 --impl local --tag local --weights-from "$OUT/n75" --out "$OUT/n75" +fi +if [[ "$PHASES" == all || "$PHASES" == n71d ]]; then + export NCCL_ALGO=Ring CUBLAS_WORKSPACE_CONFIG=:4096:8 NVIDIA_TF32_OVERRIDE=0 + run n71d-cont5 900 $TR $H --phase n71d --mode cont5 --out "$OUT/n71d" --cut-dir "$OUT/n71d/cut" + run n71d-cont5b 900 $TR $H --phase n71d --mode cont5b --out "$OUT/n71d" --cut-dir "$OUT/n71d/cut" + run n71d-run3 900 $TR $H --phase n71d --mode run3 --out "$OUT/n71d" --cut-dir "$OUT/n71d/cut" + run n71d-resume2 900 $TR $H --phase n71d --mode resume2 --out "$OUT/n71d" --cut-dir "$OUT/n71d/cut" +fi +python /g3/harness/compare.py "$OUT" > "$OUT/logs/compare.log" 2>&1; echo "compare rc=$?" +du -sh "$OUT"; echo CHAIN DONE +""" + + +@app.function(image=image, gpu=GPU, cpu=16.0, memory=65536, timeout=HARD_S, volumes={"/vol": vol}) +def chain(run: str, phases: str) -> int: + import subprocess + import threading + import time + + out = f"/vol/{run}" + os.makedirs(out, exist_ok=True) + stop = threading.Event() + + def committer(): # partial evidence survives a hard timeout + while not stop.wait(60): + try: + vol.commit() + except Exception as exc: # noqa: BLE001 + print("commit failed", exc, flush=True) + + threading.Thread(target=committer, daemon=True).start() + t0 = time.time() + rc = subprocess.call(["bash", "-c", CHAIN], env={**os.environ, "OUT": out, "PHASES": phases, + "EXPECT_GPU": EXPECT_GPU}) + with open(f"{out}/chain.rc", "w") as fh: + fh.write(f"{rc} {time.time() - t0:.0f}s\n") + stop.set() + vol.commit() + return rc + + +@app.local_entrypoint() +def main(phases: str = "all"): + assert RUN != "unset", "set G3_RUN" + call = chain.spawn(RUN, phases) + print("spawned", RUN, call.object_id, flush=True) diff --git a/yeto/rl/adapters/miles/quant_shards.py b/yeto/rl/adapters/miles/quant_shards.py index b9c7c6be7..9b8f564d6 100644 --- a/yeto/rl/adapters/miles/quant_shards.py +++ b/yeto/rl/adapters/miles/quant_shards.py @@ -14,8 +14,9 @@ quantization layout; the neutral code (:mod:`yeto.rl.quant.shards`) does the rest. Non-expert weights are not mapped here: Miles already converts them to HF names -for the sampler weight update (``convert_to_hf``); exactly one rank per -replica group passes those with :func:`hf_slices` (module doc of +for the sampler weight update (``convert_to_hf``, SGLang names ``model.layers.*``); +exactly one rank per replica group passes those with :func:`hf_slices`, which +renames them to the checkpoint names (``model.language_model.layers.*``; module doc of :mod:`yeto.rl.quant.shards`, "Ownership"). Limits (explicit errors): expert tensor parallelism must be 1 (a column split of @@ -68,7 +69,21 @@ def megatron_expert_slices(named_params: Iterable[tuple[str, torch.Tensor]], *, yield SourceSlice(expert_weight_name(prefix, expert, DOWN), weight) +def checkpoint_name(name: str) -> str: + """Miles ``convert_to_hf`` name -> GLM-5.3 checkpoint name. + + ``convert_to_hf`` emits the SGLang weight-update names (``model.layers....``, + ``model.embed_tokens``, ``model.norm``); the GLM-5.3 checkpoint, and so the + quantization layout, keeps the text model under ``model.language_model.`` + (the same rule as Miles ``Glm5NextBridge._nest_language_model``). Names that + already carry it, ``model.visual.*`` and top-level names (``lm_head.weight``) + are unchanged.""" + if name.startswith("model.") and not name.startswith(("model.language_model.", "model.visual.")): + return "model.language_model." + name[len("model."):] + return name + + def hf_slices(named_tensors: Iterable[tuple[str, torch.Tensor]]) -> Iterator[SourceSlice]: - """Whole tensors already in HF names (Miles ``convert_to_hf`` output).""" + """Whole tensors from Miles ``convert_to_hf``, renamed to checkpoint names (:func:`checkpoint_name`).""" for name, tensor in named_tensors: - yield SourceSlice(name, tensor.detach()) + yield SourceSlice(checkpoint_name(name), tensor.detach()) From d62e55478f245c760927ed56489f9ee3ccf8ff7e Mon Sep 17 00:00:00 2001 From: michael <78113554+chowhsu@users.noreply.github.com> Date: Tue, 20 Oct 2026 02:39:51 +0000 Subject: [PATCH 03/89] Serve Harbor's codex agent through the records proxy with Qwen3.8 Codex 0.145.0 only speaks the Responses API (wire_api = "chat" is refused), so the records proxy is its only path to a chat-template engine. A local run of the real Codex binary against the proxy found three gaps: - Codex sends its sub-agent tools as a namespace and a hosted web_search tool; the proxy and gateway now leave those out of the prompt and still refuse other non-function tool types. - Harbor's codex agent asks for reasoning summaries by default; the yeto trial builders now pass reasoning_summary=none and web_search=disabled. - Codex sends its instructions and a developer message as two system messages; the Qwen3.8 template accepts one, so the new qwen38 template merges leading system messages. A template refusal is now a TemplateError (HTTP 400) instead of an unhandled error (HTTP 500 that Codex retries). Register the qwen38 template (Qwen3.8 XML tool calls, thinking kept in every turn) and add scripts/modal_records_endpoint.py, which serves the SGLang adapter plus the records proxy on one pinned Modal GPU behind a secret path prefix. Co-Authored-By: Claude Opus 5.5 --- scripts/modal_records_endpoint.py | 212 +++++++++++++++++++++++++ tests/test_agentic_harbor.py | 2 + tests/test_harness_gateway.py | 2 +- tests/test_records_proxy.py | 53 ++++++- yeto/rl/adapters/miles/harbor_agent.py | 4 +- yeto/rl/harness/chat_templates.py | 65 +++++++- yeto/rl/harness/gateway/proxy.py | 7 +- yeto/rl/harness/gateway/translate.py | 10 ++ yeto/rl/harness/harbor/trial.py | 6 + 9 files changed, 350 insertions(+), 11 deletions(-) create mode 100644 scripts/modal_records_endpoint.py diff --git a/scripts/modal_records_endpoint.py b/scripts/modal_records_endpoint.py new file mode 100644 index 000000000..08e56345a --- /dev/null +++ b/scripts/modal_records_endpoint.py @@ -0,0 +1,212 @@ +#!/usr/bin/env python +"""A records-proxy model endpoint on Modal for Harbor agent trials (S25 N10.7 smoke). + +One GPU container runs the SGLang server adapter (``SGLangServerEngine``, TP1) +on a checkpoint from a Modal Volume and the N10.5 records proxy in front of it +(``RecordsProxy``: Responses -> Chat, token / loss mask / log-prob records). +The proxy is reachable from outside through a Modal tunnel, under a secret path +prefix only (``/k//...``; the proxy itself has no authentication). The +tunnel URL is put into a Modal Dict; the container serves until the Dict's +``stop`` key is set or ``max_serve_s`` passes, then writes every recorded +trajectory to the results Volume. + + python scripts/modal_records_endpoint.py launch --run-name g2-a --secret-file /tmp/secret + python scripts/modal_records_endpoint.py stop --run-name g2-a + +``launch`` spawns the call and returns (killing a client that blocks on +``.remote()`` cancels the input). The GPU type is pinned (``H100!``) and +asserted inside the container. +""" + +from __future__ import annotations + +import argparse +import json +import os +import subprocess +import sys +import threading +import time +from pathlib import Path +from typing import Any + +APP_NAME = "yeto-records-endpoint" +MODELS_VOLUME = "yeto-models-qwen38" +MODELS_MOUNT = "/mnt/models" +RESULTS_VOLUME = "yeto-s25-records-endpoint" +RESULTS_MOUNT = "/mnt/results" +CONTROL_DICT = "yeto-records-endpoint-control" +PROXY_PORT = 8000 +SGLANG_PORT = 30000 + +DEFAULTS: dict[str, Any] = { + "model_subdir": "hf/Qwen3.8-27B/1d4bf0f2", + "served_model_name": "Qwen3.8-27B", + "template": "qwen38", + "template_kwargs": {"reasoning_effort": "medium"}, + "max_model_len": 32768, + "mem_fraction_static": 0.88, + "signed_sampling": {"temperature": 1.0, "top_p": 0.95, "top_k": 20, "max_tokens": 8192}, + "assert_gpu_name": "H100", + "max_serve_s": 3600, +} + + +def _gpu_names() -> list[str]: + out = subprocess.run(["nvidia-smi", "--query-gpu=name", "--format=csv,noheader"], + capture_output=True, text=True, timeout=30).stdout + return [n.strip() for n in out.splitlines() if n.strip()] + + +def serve(config: dict[str, Any]) -> dict[str, Any]: + """Inside the container (module doc).""" + import asyncio + import dataclasses + + import modal + + from yeto.cloud.modal_eval_island import _image_paths + + _image_paths() + from aiohttp import web + + import sglang + + from yeto.rl.adapters.sampler_engines.server import ServerConfig + from yeto.rl.adapters.sampler_engines.sglang_server import SGLangOptions, SGLangServerEngine + from yeto.rl.harness.chat_templates import TEMPLATES, ChatTemplate + from yeto.rl.harness.gateway.proxy import RecordsProxy + + names = _gpu_names() + if not names or not all(config["assert_gpu_name"] in n for n in names): + raise RuntimeError(f"asked for {config['assert_gpu_name']}, got {names}") + run_dir = Path(RESULTS_MOUNT) / config["run_name"] + run_dir.mkdir(parents=True, exist_ok=True) + results = modal.Volume.from_name(RESULTS_VOLUME) + control = modal.Dict.from_name(CONTROL_DICT, create_if_missing=True) + began = time.time() + meta: dict[str, Any] = {"gpu_names": names, "sglang": sglang.__version__, "config": { + k: v for k, v in config.items() if k != "secret"}, "container": os.environ.get("MODAL_TASK_ID")} + + model_dir = Path(MODELS_MOUNT) / config["model_subdir"] + engine = SGLangServerEngine( + ServerConfig(base_dir=model_dir, card_model=names[0], device_family="cuda", + engine_version=sglang.__version__, precision="bf16", quant_scheme="bf16", + topology="tp8", max_model_len=int(config["max_model_len"]), port=SGLANG_PORT, + python=sys.executable, start_timeout_s=1800.0, + extra_args=("--mem-fraction-static", str(config["mem_fraction_static"]))), + SGLangOptions(tp=1, served_model_name=config["served_model_name"])) + report = engine.start(None, 0) + meta["engine_start_s"] = round(report.pause_s, 1) + spec = TEMPLATES[config["template"]] + spec = dataclasses.replace(spec, template_kwargs={**spec.template_kwargs, **config["template_kwargs"]}) + template = ChatTemplate(spec, ChatTemplate.load(spec.name, tokenizer_path=str(model_dir)).tokenizer) + proxy = RecordsProxy(engine, template, model=config["served_model_name"], + signed_sampling=config["signed_sampling"]) + trajectories: set[str] = set() + + @web.middleware + async def note(request: web.Request, handler: Any) -> web.StreamResponse: + tid = request.match_info.get("tid") + if tid: + trajectories.add(tid) + return await handler(request) + + inner = proxy.app() + inner.middlewares.append(note) + outer = web.Application() + outer.add_subapp(f"/k/{config['secret']}/", inner) + + loop = asyncio.new_event_loop() + + async def start_server() -> web.AppRunner: + runner = web.AppRunner(outer) + await runner.setup() + await web.TCPSite(runner, "0.0.0.0", PROXY_PORT).start() + return runner + + async def snapshot() -> str: # runs on the server loop: no concurrent mutation + records = {} + for tid in sorted(trajectories): + try: + records[tid] = proxy.records(tid) + except Exception as exc: # noqa: BLE001 - keep the others + records[tid] = {"error": repr(exc)} + return json.dumps(records) + + def dump() -> None: + text = asyncio.run_coroutine_threadsafe(snapshot(), loop).result(timeout=120) + (run_dir / "records.json").write_text(text) + (run_dir / "meta.json").write_text(json.dumps({**meta, "served_s": round(time.time() - began, 1)})) + results.commit() + + threading.Thread(target=loop.run_forever, daemon=True).start() + runner = asyncio.run_coroutine_threadsafe(start_server(), loop).result(timeout=60) + url_key, stop_key = f"{config['run_name']}/url", f"{config['run_name']}/stop" + try: + with modal.forward(PROXY_PORT) as tunnel: + control[url_key] = tunnel.url + print(f"[records-endpoint] serving, engine start {meta['engine_start_s']} s", flush=True) + deadline = began + float(config["max_serve_s"]) + while time.time() < deadline and not control.get(stop_key, False): + time.sleep(15) + dump() + finally: + control.pop(url_key, None) + dump() + asyncio.run_coroutine_threadsafe(runner.cleanup(), loop).result(timeout=60) + engine.stop() + return {"run_name": config["run_name"], "trajectories": len(trajectories)} + + +def build_app(workdir: str, gpu: str, timeout_s: int) -> tuple[Any, Any]: + import modal + + sys.path.insert(0, workdir) + from yeto.cloud.modal_eval_island import WORKDIR_MOUNT + from yeto.rl.adapters.miles.pins import MILES_NEXT_IMAGE + + image = (modal.Image.from_registry(MILES_NEXT_IMAGE.removeprefix("docker:")) + .env({"HOME": "/root", "PYTHONUNBUFFERED": "1"}) + .add_local_dir(workdir, WORKDIR_MOUNT, copy=False, + ignore=[".git", ".venv", ".claude", "**/__pycache__", "**/*.pyc", "s1-runs", + "syncer/target"])) + app = modal.App(APP_NAME) + models = modal.Volume.from_name(MODELS_VOLUME) + results = modal.Volume.from_name(RESULTS_VOLUME, create_if_missing=True) + fn = app.function(image=image, gpu=gpu, cpu=8.0, memory=65536, timeout=timeout_s, retries=0, + volumes={MODELS_MOUNT: models.read_only(), RESULTS_MOUNT: results}, + name="serve", serialized=True)(serve) + return app, fn + + +def main() -> None: + parser = argparse.ArgumentParser(description=__doc__.splitlines()[0]) + parser.add_argument("action", choices=("launch", "stop", "url")) + parser.add_argument("--run-name", required=True) + parser.add_argument("--secret-file", type=Path, help="launch: file holding the path secret") + parser.add_argument("--gpu", default="H100!") + parser.add_argument("--max-serve-s", type=int, default=DEFAULTS["max_serve_s"]) + parser.add_argument("--workdir", default=str(Path(__file__).resolve().parents[1])) + args = parser.parse_args() + import modal + + control = modal.Dict.from_name(CONTROL_DICT, create_if_missing=True) + if args.action == "stop": + control[f"{args.run_name}/stop"] = True + return + if args.action == "url": + print(control.get(f"{args.run_name}/url", "")) + return + secret = args.secret_file.read_text().strip() + if len(secret) < 24: + raise SystemExit("the path secret must have at least 24 characters") + config = {**DEFAULTS, "run_name": args.run_name, "secret": secret, "max_serve_s": args.max_serve_s} + app, fn = build_app(args.workdir, args.gpu, timeout_s=args.max_serve_s + 1800) + with modal.enable_output(), app.run(detach=True): + call = fn.spawn(config) + print(json.dumps({"function_call_id": call.object_id, "run_name": args.run_name})) + + +if __name__ == "__main__": + main() diff --git a/tests/test_agentic_harbor.py b/tests/test_agentic_harbor.py index e3039cf71..adcba5c46 100644 --- a/tests/test_agentic_harbor.py +++ b/tests/test_agentic_harbor.py @@ -286,6 +286,7 @@ class Trial(_FakeTrial): @classmethod async def create(cls, config): seen["version"] = config.agent.kwargs["version"] + seen["summary"] = config.agent.kwargs["reasoning_summary"] return cls(config) Trial.result = _result({"reward": 1.0}) @@ -295,6 +296,7 @@ async def create(cls, config): out = _run(mod.run("http://ss:30000/sessions/s1", metadata={"instance_id": "fix-git"})) assert example.calls == ["key", ("codex", "http://ss:30000/sessions/s1/v1")] assert seen["version"] == "0.145.0" # Harbor's default would be @latest + assert seen["summary"] == "none" # the gateway refuses reasoning summaries assert (out[OUTCOME_KEY], out["reward"]) == ("ok", 1.0) Trial.result = _result(None, "EnvironmentStartTimeoutError") diff --git a/tests/test_harness_gateway.py b/tests/test_harness_gateway.py index 5d90913db..7b72cb4ac 100644 --- a/tests/test_harness_gateway.py +++ b/tests/test_harness_gateway.py @@ -115,7 +115,7 @@ def test_signed_sampling_cannot_be_overridden_and_unsupported_shapes_fail_closed run(gw.handle("chat", {**base, "model": "other"}, trajectory_id="x3")) resp = {"store": False, "input": [{"type": "message", "role": "user", "content": [{"type": "input_text", "text": "u"}]}]} for bad in ({**resp, "store": True}, {**resp, "parallel_tool_calls": True}, {**resp, "reasoning": {"summary": "auto"}}, - {**resp, "tools": [{"type": "web_search"}]}, {**resp, "input": resp["input"] + [{"type": "compaction"}]}): + {**resp, "tools": [{"type": "local_shell"}]}, {**resp, "input": resp["input"] + [{"type": "compaction"}]}): with pytest.raises(UnsupportedShape): run(gw.handle("responses", bad, trajectory_id="x4")) with pytest.raises(UnsupportedShape): diff --git a/tests/test_records_proxy.py b/tests/test_records_proxy.py index e9a77ace3..108daaf78 100644 --- a/tests/test_records_proxy.py +++ b/tests/test_records_proxy.py @@ -12,6 +12,7 @@ import asyncio import copy import json +import os from typing import Any import pytest @@ -19,10 +20,11 @@ from yeto.rl.adapters.sampler_engines.fake import FakeEngine from yeto.rl.compaction import CompactionConfig from yeto.rl.engine.sampler_engine import SamplingParams, TokenSamplerEngine -from yeto.rl.harness.chat_templates import ChatTemplate, TemplateError, check_alignment, parse_glm, parse_qwen3 +from yeto.rl.harness.chat_templates import ChatTemplate, TemplateError, check_alignment, parse_glm, parse_qwen3, parse_qwen38 from yeto.rl.harness.gateway import Gateway, GatewayConfig from yeto.rl.harness.gateway.proxy import RecordsProxy, chat_tools from yeto.rl.harness.gateway.records import RecordingBackend, RecordsMismatch +from yeto.rl.harness.gateway.translate import UnsupportedShape, responses_to_messages SIGNED = {"temperature": 1.0, "top_p": 1.0, "max_tokens": 256} TOOLS = [{"type": "function", "name": "shell", "description": "run a command", @@ -36,6 +38,16 @@ def _template(name: str) -> ChatTemplate: pytest.skip(f"{name} tokenizer not available: {exc}") +# The Qwen3.8-27B tokenizer files (not on the HF hub cache by default). +QWEN38_TOKENIZER = os.environ.get("YETO_QWEN38_TOKENIZER", "") + + +def _template_at(name: str, path: str) -> ChatTemplate: + if not path: + pytest.skip(f"set YETO_QWEN38_TOKENIZER to run the {name} template tests") + return ChatTemplate.load(name, tokenizer_path=path) + + @pytest.fixture(scope="module") def qwen() -> ChatTemplate: return _template("qwen3") @@ -339,6 +351,18 @@ def test_compaction_on_a_template_that_drops_history_reasoning_fails_closed(qwen ] +def test_codex_namespace_and_hosted_tools_are_left_out_other_types_refused(): + # Codex 0.145.0 sends its sub-agent tools as a namespace and a hosted web_search tool. + namespace = {"type": "namespace", "name": "multi_agent_v1", "tools": [TOOLS[0]]} + tools = chat_tools("responses", {"tools": [*TOOLS, namespace, {"type": "web_search"}]}) + assert [t["function"]["name"] for t in tools] == ["shell"] + with pytest.raises(UnsupportedShape): + chat_tools("responses", {"tools": [{"type": "custom", "name": "apply_patch"}]}) + body = {"store": False, "parallel_tool_calls": False, "tools": [*TOOLS, namespace, {"type": "web_search"}], + "input": [{"type": "message", "role": "user", "content": [{"type": "input_text", "text": "hi"}]}]} + assert responses_to_messages(body) == [{"role": "user", "content": "hi"}] + + @pytest.mark.parametrize("name", ["qwen3", "glm5"]) def test_template_rendering_matches_tokenizer_ids_token_by_token(name, request): template = request.getfixturevalue("qwen" if name == "qwen3" else "glm") @@ -366,6 +390,33 @@ def test_parsers_fail_closed_on_malformed_tool_calls(): parse_qwen3("\nnot json\n") with pytest.raises(TemplateError): parse_glm("ra1") + with pytest.raises(TemplateError): + parse_qwen38('r\n\n\n\n{"name": "f", "arguments": {}}\n') + + +def test_qwen38_merges_codex_system_messages_and_template_refusals_fail_closed(): + template = _template_at("qwen38", QWEN38_TOKENIZER) + tools = chat_tools("responses", {"tools": TOOLS}) + # Codex: instructions + a developer message, both "system" in the gateway form. + messages = [{"role": "system", "content": "a"}, {"role": "system", "content": "b"}, {"role": "user", "content": "u"}] + text = template.render(messages, tools=tools, add_generation_prompt=True) + assert "a\n\nb" in text and text.count("<|im_start|>system") == 1 + assert check_alignment(template, messages, tools=tools).aligned + with pytest.raises(TemplateError, match="refused"): + template.render([*messages, {"role": "system", "content": "late"}], add_generation_prompt=True) + + +def test_qwen38_parser_reads_the_xml_tool_call_format(): + # The Qwen3.8 template's own rendering of one assistant turn (after "\n"). + text = ("look first\n\n\nChecking.\n\n\n\n" + "\nls -la\nwc -l x\n\n\n3\n\n" + "\n\n\n\n\n") + parsed = parse_qwen38(text) + assert parsed["reasoning_content"] == "look first" + assert parsed["content"] == "Checking." + assert [c["function"]["name"] for c in parsed["tool_calls"]] == ["exec_command", "noop"] + assert json.loads(parsed["tool_calls"][0]["function"]["arguments"]) == {"cmd": "ls -la\nwc -l x", "n": 3} + assert parse_qwen38("still thinking") == {"role": "assistant", "content": "", "reasoning_content": "still thinking"} # ---------------------------------------------------------------------- adapters diff --git a/yeto/rl/adapters/miles/harbor_agent.py b/yeto/rl/adapters/miles/harbor_agent.py index 7506cf571..52caae11c 100644 --- a/yeto/rl/adapters/miles/harbor_agent.py +++ b/yeto/rl/adapters/miles/harbor_agent.py @@ -44,7 +44,7 @@ from yeto.rl.agentic.loop_local import LoopLocal from yeto.rl.agentic.trust import check_agent_env from yeto.rl.harness.harbor import CODEX_VERSION, HarborHarness, HarborSettings, build_trial_config -from yeto.rl.harness.harbor.trial import CUSTOM_PROVIDERS, TrialConfigBuilder +from yeto.rl.harness.harbor.trial import CODEX_KWARGS, CUSTOM_PROVIDERS, TrialConfigBuilder EXAMPLE_MODULE = "harbor_agent_function" AGENT_FUNCTION = "yeto.rl.adapters.miles.harbor_agent.run" @@ -105,6 +105,8 @@ def build(task: AgentTask, session_url: str) -> Any: # The example has no codex binding; its default (OPENAI_BASE_URL) is # what Harbor's codex agent reads. Pin the build: Harbor's default is @latest. config.agent.kwargs.setdefault("version", task.agent_version or CODEX_VERSION) + for key, value in CODEX_KWARGS.items(): + config.agent.kwargs.setdefault(key, value) return config return build diff --git a/yeto/rl/harness/chat_templates.py b/yeto/rl/harness/chat_templates.py index 13ea1883d..030daf39a 100644 --- a/yeto/rl/harness/chat_templates.py +++ b/yeto/rl/harness/chat_templates.py @@ -16,7 +16,13 @@ one (same template family). Load the real checkpoint with ``ChatTemplate.load("glm5", tokenizer_path=...)``; whether GLM5.3 uses the same template is not verified; -* ``qwen3``: the Qwen3 hermes tool-call format; the small CPU stand-in. +* ``qwen3``: the Qwen3 hermes tool-call format; the small CPU stand-in; +* ``qwen38``: Qwen3.8 (Qwen3.5 family) XML tool calls (`` + v``), thinking kept in every + turn (``preserve_thinking``). Load with the checkpoint's tokenizer + (``tokenizer_path``); a deployment may set ``reasoning_effort`` by + replacing ``template_kwargs`` (xhigh / medium / low; the template's default + is xhigh). :func:`check_alignment` is the N10.6 check: the template rendering, tokenized the way the proxy sends it, must equal the tokenizer's own @@ -117,6 +123,23 @@ def parse_qwen3(text: str) -> Message: return _assistant(content, reasoning, calls) +_QWEN38_CALL = re.compile(r"\s*\n]+)>(.*?)\s*", re.S) +_QWEN38_PARAM = re.compile(r"\n]+)>\n?(.*?)\n?", re.S) + + +def parse_qwen38(text: str) -> Message: + """Qwen3.8 output after the ``\\n`` generation prompt.""" + reasoning, rest = _split_reasoning(text, opened=True) + calls = [] + for index, match in enumerate(_QWEN38_CALL.finditer(rest)): + arguments = {k.strip(): _glm_value(v) for k, v in _QWEN38_PARAM.findall(match.group(2))} + calls.append(_tool_call(_call_id(text, index), match.group(1).strip(), arguments)) + if rest.count("") != len(calls): + raise TemplateError(f"Qwen3.8 tool call is not XML: {rest[rest.find(''):][:80]!r}") + content = rest.split("", 1)[0].strip() + return _assistant(content, reasoning, calls) + + # ---------------------------------------------------------------------- registry @dataclass(frozen=True) class TemplateSpec: @@ -124,7 +147,10 @@ class TemplateSpec: an assistant turn (sent to the engine as stop token ids). ``keeps_history_reasoning``: with ``template_kwargs`` the template renders the reasoning of every earlier turn (gateway ``ChainRegistry`` input). - ``verified``: what the token alignment check ran on (free text).""" + ``verified``: what the token alignment check ran on (free text). + ``merge_leading_system``: join the leading system messages into one (the + template accepts a single system message; Codex sends its instructions + and a developer message, both system in the gateway form).""" name: str tokenizer: str @@ -133,6 +159,7 @@ class TemplateSpec: template_kwargs: Mapping[str, Any] = field(default_factory=dict) keeps_history_reasoning: bool = True verified: str = "not verified" + merge_leading_system: bool = False TEMPLATES: dict[str, TemplateSpec] = {} @@ -157,9 +184,26 @@ def register_template(spec: TemplateSpec) -> TemplateSpec: # The Qwen3 template drops the reasoning of turns before the last user turn. keeps_history_reasoning=False, verified="Qwen3-0.6B tokenizer (CPU stand-in)")) +register_template(TemplateSpec( + name="qwen38", tokenizer="Qwen/Qwen3.8-27B", parser=parse_qwen38, + stop_tokens=("<|im_end|>", "<|endoftext|>"), + template_kwargs={"enable_thinking": True, "preserve_thinking": True}, + verified="Qwen3.8-27B@1d4bf0f2 tokenizer (S25 T-G2)", merge_leading_system=True)) # ---------------------------------------------------------------------- template +def _merge_leading_system(messages: Sequence[Message]) -> list[Message]: + head = 0 + while head < len(messages) and messages[head].get("role") == "system": + head += 1 + if head < 2: + return list(messages) + if any(not isinstance(m.get("content"), str) for m in messages[:head]): + raise TemplateError("cannot merge system messages with non-text content") + merged = {"role": "system", "content": "\n\n".join(m["content"] for m in messages[:head])} + return [merged, *messages[head:]] + + def _render_ready(messages: Sequence[Message]) -> list[Message]: """Canonical messages -> template input: tool-call arguments as objects.""" out = [] @@ -208,9 +252,18 @@ def load(cls, name: str, *, tokenizer_path: str | None = None, local_files_only: def render(self, messages: Sequence[Message], *, tools: Sequence[Message] | None = None, add_generation_prompt: bool) -> str: - return self.tokenizer.apply_chat_template( - _render_ready(messages), tools=list(tools) if tools else None, tokenize=False, - add_generation_prompt=add_generation_prompt, **dict(self.spec.template_kwargs)) + import jinja2 # comes with transformers + + try: + return self.tokenizer.apply_chat_template( + self._ready(messages), tools=list(tools) if tools else None, tokenize=False, + add_generation_prompt=add_generation_prompt, **dict(self.spec.template_kwargs)) + except jinja2.TemplateError as exc: # raise_exception(...): the template refuses the conversation + raise TemplateError(f"{self.spec.name} template refused the conversation: {exc}") from exc + + def _ready(self, messages: Sequence[Message]) -> list[Message]: + ready = _render_ready(messages) + return _merge_leading_system(ready) if self.spec.merge_leading_system else ready def encode(self, text: str) -> list[int]: """Token ids of rendered text; the template writes its own BOS tokens.""" @@ -252,7 +305,7 @@ def check_alignment(template: ChatTemplate, messages: Sequence[Message], *, text = template.render(messages, tools=tools, add_generation_prompt=add_generation_prompt) ours = template.encode(text) reference = template.tokenizer.apply_chat_template( - _render_ready(messages), tools=list(tools) if tools else None, tokenize=True, + template._ready(messages), tools=list(tools) if tools else None, tokenize=True, add_generation_prompt=add_generation_prompt, **dict(template.spec.template_kwargs)) if isinstance(reference, Mapping): reference = reference["input_ids"] diff --git a/yeto/rl/harness/gateway/proxy.py b/yeto/rl/harness/gateway/proxy.py index f414fa3ae..b5c0e4851 100644 --- a/yeto/rl/harness/gateway/proxy.py +++ b/yeto/rl/harness/gateway/proxy.py @@ -27,18 +27,21 @@ from .core import AdmissionClosed, Gateway, GatewayConfig, GatewayError, SamplingOverrideError, TrajectoryInvalid from .records import RecordingBackend -from .translate import UnsupportedShape +from .translate import OMITTED_TOOL_TYPES, UnsupportedShape Message = dict[str, Any] def chat_tools(entry: str, body: dict[str, Any]) -> list[Message]: - """The request's tool schemas in Chat form (Responses tools are flat).""" + """The request's tool schemas in Chat form (Responses tools are flat). + Types in :data:`OMITTED_TOOL_TYPES` are left out of the prompt.""" tools = body.get("tools") or [] if entry != "responses": return list(tools) out = [] for tool in tools: + if isinstance(tool, dict) and tool.get("type") in OMITTED_TOOL_TYPES: + continue if not isinstance(tool, dict) or tool.get("type") != "function": raise UnsupportedShape(f"unsupported tool {tool!r}") function = {k: tool[k] for k in ("name", "description", "parameters") if k in tool} diff --git a/yeto/rl/harness/gateway/translate.py b/yeto/rl/harness/gateway/translate.py index a533c4d21..0da3e5e90 100644 --- a/yeto/rl/harness/gateway/translate.py +++ b/yeto/rl/harness/gateway/translate.py @@ -73,6 +73,14 @@ def completion_to_chat(completion: dict[str, Any]) -> dict[str, Any]: # ------------------------------------------------------------------ /v1/responses (Codex) +# Responses tool types the policy is not offered: ``namespace`` (Codex 0.145 +# sends its sub-agent tools so; a sub-agent would open a second token stream) +# and OpenAI hosted tools (the served model cannot call them). Other +# non-function tool types are refused. +OMITTED_TOOL_TYPES = frozenset({"namespace", "web_search", "web_search_preview", "image_generation", + "file_search", "code_interpreter"}) + + def responses_to_messages(body: dict[str, Any]) -> list[dict[str, Any]]: if body.get("store") is not False: raise UnsupportedShape("Responses requests must use store:false (full replay)") @@ -82,6 +90,8 @@ def responses_to_messages(body: dict[str, Any]) -> list[dict[str, Any]]: if summary not in ("none", None): raise UnsupportedShape("reasoning.summary must be none") for tool in body.get("tools") or []: + if isinstance(tool, dict) and tool.get("type") in OMITTED_TOOL_TYPES: + continue if not isinstance(tool, dict) or tool.get("type") != "function": raise UnsupportedShape(f"unsupported tool type {tool.get('type') if isinstance(tool, dict) else tool!r}") messages: list[dict[str, Any]] = [] diff --git a/yeto/rl/harness/harbor/trial.py b/yeto/rl/harness/harbor/trial.py index a213c1d38..f506d57e3 100644 --- a/yeto/rl/harness/harbor/trial.py +++ b/yeto/rl/harness/harbor/trial.py @@ -42,6 +42,10 @@ # server, which does not authenticate the agent. AGENT_API_KEY = "dummy" +# Harbor codex CLI flags: the gateway serves no reasoning summaries (it refuses +# ``reasoning.summary`` other than none) and no hosted web search. +CODEX_KWARGS: Mapping[str, str] = {"reasoning_summary": "none", "web_search": "disabled"} + @dataclass(frozen=True) class HarborSettings: @@ -112,6 +116,8 @@ def build_trial_config(task: AgentTask, endpoint: str, settings: HarborSettings) version = task.agent_version or (CODEX_VERSION if task.agent == "codex" else None) if version is not None: kwargs["version"] = version + if task.agent == "codex": + kwargs.update(CODEX_KWARGS) if task.budget.max_seq_len is not None and task.agent.startswith("terminus"): kwargs["max_seq_len"] = task.budget.max_seq_len return TrialConfig( From 3e6604aa21b840c3837962dbd9178547638f6bde Mon Sep 17 00:00:00 2001 From: michael <78113554+chowhsu@users.noreply.github.com> Date: Tue, 20 Oct 2026 02:40:40 +0000 Subject: [PATCH 04/89] Stage-1 mini rehearsal on Modal: plan, SGLang expert load report, requant fixes - glm5_next recipe reads n_routed_experts and first_k_dense_replace from the checkpoint instead of hard-coding GLM5.3-Flash (288 experts, 3 dense layers). - CheckpointLayout.load falls back to the index and shard headers when the template directory is a reference checkpoint without layout.json; stage 1 passes the round-0 MXFP4 export as --template, which failed before. - Re-quantizer --frozen-base/--frozen-pattern: untrained tensors the trainer does not export (e.g. the vision tower) come from the bf16 source; any other missing tensor still fails. - sglang_load_report: a hook for SGLang's GLM5.3 model file that renames routed-expert weight_packed tensors and reports, per rank, the expert matrices that reached a weight loader; the sglang-server adapter reads it for post-load check 1 (it reported no count before, so the check failed). - PostLoadChecks.require_fixed_answers (default True) for a cut-down model that cannot answer; the other fixed-prompt problems still fail. - Island flag --miles-commit (full SHA) for an image that carries another Miles commit than the pin; stage 1 passes the plan's commit when full. - Stage 1: store_backend modal-volume (store hook env on every node), digest-pinned images, tp topologies, sampler python/extra_env/checks, model num_experts/first_k_dense_replace. - examples/stage1/glm5next-mini-modal.yaml + modal-mini-inventory.yaml and the Modal driver tools/probes/s25_g1 (prep, smoke, launch). Co-Authored-By: Claude Opus 5.5 --- .gitignore | 1 + examples/stage1/glm5next-mini-modal.yaml | 131 +++++++++ examples/stage1/modal-mini-inventory.yaml | 31 +++ tests/test_rl_quant_mxfp4.py | 24 ++ tests/test_s25_g1_stage1mini.py | 155 +++++++++++ tools/probes/s25_g1/image_probe.py | 59 ++++ tools/probes/s25_g1/make_mini.py | 148 ++++++++++ tools/probes/s25_g1/modal_app.py | 263 ++++++++++++++++++ yeto/launch/ssh_trainer_island.py | 10 +- yeto/rl/adapters/miles/config.py | 8 +- yeto/rl/adapters/miles/island_entry.py | 15 +- yeto/rl/adapters/miles/models/glm5_next.py | 24 +- yeto/rl/adapters/miles/run_flags.py | 7 + .../sampler_engines/sglang_load_report.py | 141 ++++++++++ .../adapters/sampler_engines/sglang_server.py | 30 +- yeto/rl/engine/sampler_checks.py | 13 +- yeto/rl/quant/release.py | 14 +- yeto/rl/quant/requant.py | 31 ++- yeto/rl/stage1.py | 62 ++++- 19 files changed, 1138 insertions(+), 29 deletions(-) create mode 100644 examples/stage1/glm5next-mini-modal.yaml create mode 100644 examples/stage1/modal-mini-inventory.yaml create mode 100644 tests/test_s25_g1_stage1mini.py create mode 100644 tools/probes/s25_g1/image_probe.py create mode 100644 tools/probes/s25_g1/make_mini.py create mode 100644 tools/probes/s25_g1/modal_app.py create mode 100644 yeto/rl/adapters/sampler_engines/sglang_load_report.py diff --git a/.gitignore b/.gitignore index 2d80f62e8..ece43ff8e 100644 --- a/.gitignore +++ b/.gitignore @@ -95,3 +95,4 @@ openspec/changes/*/evidence/**/*.pyc openspec/**/*.jsonl openspec/**/*.out openspec/**/*.err +.g1-build/ diff --git a/examples/stage1/glm5next-mini-modal.yaml b/examples/stage1/glm5next-mini-modal.yaml new file mode 100644 index 000000000..16ac368b0 --- /dev/null +++ b/examples/stage1/glm5next-mini-modal.yaml @@ -0,0 +1,131 @@ +# Stage-1 mini rehearsal (S25 T-G1) on Modal H100: the same plan schema and +# the same generated commands as glm53-flash-ppu.yaml; only the nodes +# (Modal containers) and the store (a Modal Volume through the yeto store +# hook, not OSS) differ. Model: Miles' 4-layer GLM5.3-Flash CI checkpoint cut +# to 32 routed experts per MoE layer (tools/probes/s25_g1/make_mini.py). +# +# yeto-stage1 plan examples/stage1/glm5next-mini-modal.yaml +# modal run tools/probes/s25_g1/modal_app.py::... (the driver reads this plan) + +run_id: s25-g1-mini +inventory: modal-mini-inventory.yaml +rounds: 3 +seed: 17 + +rendezvous: + modal-trainer: 127.0.0.1 + modal-sampler: 127.0.0.1 + +roles: + trainer: [modal-trainer] + sampler: [modal-sampler] + +model: + name: GLM5.3-Flash-4layer-e32 + architecture: glm5_next + num_layers: 4 # KDA, DSA, KDA, DSA + first_k_dense_replace: 1 + num_experts: 32 # cut from 288 (make_mini.py) + params_b: 5.3 # 推算: 10.5 GB bf16 / 2 + expert_params_b: 2.4 # 3 MoE layers x 32 x 3 x 4096 x 2048 + bf16_source: + revision: 7fb7883207fe73ea04430ac392f6ffe2b9f936c4 # CharyZeng/GLM-5.3-Flash-4layer, then cut + origin: hf://CharyZeng/GLM-5.3-Flash-4layer@7fb7883207fe73ea04430ac392f6ffe2b9f936c4 + staged: modal-volume://yeto-s25-g1-models/glm5next-mini-bf16-7fb7883 + local_dir: /data/models/glm5next-mini-bf16-7fb7883 + sampler_initial: # round 0: our N8.5 converter on the bf16 mini model + repo: local/glm5next-mini-mxfp4-r0-7fb7883 + revision: 7fb7883 + staged: modal-volume://yeto-s25-g1-models/glm5next-mini-mxfp4-r0-7fb7883 + local_dir: /data/models/glm5next-mini-mxfp4-r0-7fb7883 + +quant: + scheme: mxfp4-g32-rtn + format: mxfp4-pack-quantized + release_dir: policy-releases-mxfp4-g32-rtn + expected_expert_matrices: 288 # 3 MoE layers x 32 experts x (gate, up, down) + frozen_pattern: '^model\.visual\.' # Miles trains the text model only; the vision tower stays the base + +store: /data/store/s25-g1 +store_backend: modal-volume # the only difference from OSS: the store hook and its Volume +store_volume: yeto-s25-g1-store +store_access: + head: direct + +images: + sampler: ghcr.io/michaellchung/yeto-miles-ports@sha256:5b9dc2a9e6e315d73dabfbc6faae428b59cc01aabec6189645a5c8c80d6e3721 + trainer: ghcr.io/michaellchung/yeto-miles-ports@sha256:5b9dc2a9e6e315d73dabfbc6faae428b59cc01aabec6189645a5c8c80d6e3721 + docker_device_args: [--gpus, all] # unused on Modal; the containers get the GPUs + +trainer: + backend: miles + miles: + commit: 1618b5c38db80e17ec928cfa47fe832cb18565fc # s25-distopt, checked out in the image + megatron: radixark/Megatron-LM@a84b105 + mbridge: image + workdir: /root/miles + tp: 1 + sequence_parallel: false + pp: 1 + pp_layers: [4] + cp: 1 + ep: 2 + etp: 1 + optimizer_offload_fraction: 1.0 + memory_estimate: # 推算; `plan` checks them + weights_grads_gib_per_card: 22.9 + host_optimizer_gib_per_node: 59.2 + transformer_impl: transformer_engine + recompute: full + max_tokens_per_gpu: 8192 + lr: 1.0e-6 + verl: + commit: null + strategy: fsdp2 + fsdp_size: 2 + ulysses: 1 + param_offload: false + optimizer_offload: true + lr: 1.0e-6 + +sampler: + engine: sglang + engine_version: 0.5.21.dev70+g2fa8801 + python: python3 + topology: tp1 + tp: 1 + precision: mxfp4 + max_model_len: 4096 + port: 30000 + moe_runner_backend: triton + hierarchical_cache: false + nextn: false + renames_weight_packed: true # the yeto hook in the image renames (sglang_load_report) + load_report: true # the hook reports loaded expert matrices (post-load check 1) + in_place_reload: false + extra_env: + PYTHONPATH: /root/sky_workdir + checks: + require_fixed_answers: false # a cut 4-layer model cannot answer GSM8K + fixed_max_tokens: 64 + +sampling: + max_tokens: 512 + max_prompt_len: 512 + temperature: 1.0 + top_p: 1.0 + top_k: -1 + repetition_penalty: 1.0 + frequency_penalty: 0.0 + presence_penalty: 0.0 + +task: + data: zhuzilin/gsm8k + data_revision: 0cbd9f31d91ac21a7613dcbc7fef992adac459ae + reward: yeto.rl.gsm8k_reward:score + rollout_batch_size: 4 + n_samples_per_prompt: 4 + prompts_jsonl: /data/models/datasets/gsm8k-0cbd9f3/train.jsonl + input_key: messages + label_key: label + prompt_format: chat diff --git a/examples/stage1/modal-mini-inventory.yaml b/examples/stage1/modal-mini-inventory.yaml new file mode 100644 index 000000000..fd434c18d --- /dev/null +++ b/examples/stage1/modal-mini-inventory.yaml @@ -0,0 +1,31 @@ +# Stage-1 mini rehearsal on Modal (S25 T-G1): one trainer container with +# 2 x H100 and one sampler container with 1 x H100. "host" is a label only: +# the Modal driver (tools/probes/s25_g1/modal_app.py) starts the containers, +# not SSH. /data/models and /data/store are Modal Volumes in each container. +nodes: + - name: modal-trainer + host: modal-trainer + address: 127.0.0.1 + device_family: cuda + cpu_arch: x86_64 + card: H100 + cards_per_node: 2 + hbm_gib: 80 + host_memory_gib: 400 + mounts: {/data: /data} + image: ghcr.io/michaellchung/yeto-miles-ports@sha256:5b9dc2a9e6e315d73dabfbc6faae428b59cc01aabec6189645a5c8c80d6e3721 + runtime: docker + roles: [trainer] + - name: modal-sampler + host: modal-sampler + address: 127.0.0.1 + device_family: cuda + cpu_arch: x86_64 + card: H100 + cards_per_node: 1 + hbm_gib: 80 + host_memory_gib: 100 + mounts: {/data: /data} + image: ghcr.io/michaellchung/yeto-miles-ports@sha256:5b9dc2a9e6e315d73dabfbc6faae428b59cc01aabec6189645a5c8c80d6e3721 + runtime: docker + roles: [sampler] diff --git a/tests/test_rl_quant_mxfp4.py b/tests/test_rl_quant_mxfp4.py index 34310e09d..b2a896313 100644 --- a/tests/test_rl_quant_mxfp4.py +++ b/tests/test_rl_quant_mxfp4.py @@ -337,3 +337,27 @@ def test_cli_template_convert_compare(world, capsys): assert main(["compare", str(out), str(world["original"]), "--match", r"experts\.\d+\."]) == 0 rep = json.loads(capsys.readouterr().out) assert rep["fp4_match_share"] == 1.0 and rep["e8m0_match_share"] == 1.0 and rep["tensors"] == 3 * 3 * E # packed + scale, plus the BF16 MTP experts + + +def test_reference_checkpoint_serves_as_template(world): + # stage 1 passes the round-0 MXFP4 export (no layout.json) as --template + out = world["tmp"] / "out-ckpt-template" + requantize_checkpoint(world["source"], world["original"], out, AwqSmoothing.none()) + rep = compare_checkpoints(out, world["original"]) + assert rep["unequal_tensors"] == 0 and rep["fp4_match_share"] == 1.0 + + +def test_frozen_tensors_come_from_the_base_only_when_named(world): + partial = world["tmp"] / "release-without-lm-head" + _write_ckpt(partial, {n: t for n, t in world["src"].items() if n != "lm_head.weight"}, + lambda n: "model-00001-of-00001.safetensors") + with pytest.raises(LayoutError, match="1 missing"): + requantize_checkpoint(partial, world["template"], world["tmp"] / "a", AwqSmoothing.none()) + with pytest.raises(LayoutError, match="1 missing"): + requantize_checkpoint(partial, world["template"], world["tmp"] / "b", AwqSmoothing.none(), + frozen=(world["source"], r"^model\.visual\.")) + out = world["tmp"] / "c" + stats = requantize_checkpoint(partial, world["template"], out, AwqSmoothing.none(), + frozen=(world["source"], r"^lm_head\.")) + assert stats["frozen_from_base"] == 1 + assert compare_checkpoints(out, world["original"])["unequal_tensors"] == 0 diff --git a/tests/test_s25_g1_stage1mini.py b/tests/test_s25_g1_stage1mini.py new file mode 100644 index 000000000..d3742ea5e --- /dev/null +++ b/tests/test_s25_g1_stage1mini.py @@ -0,0 +1,155 @@ +"""S25 T-G1: the pieces the stage-1 mini rehearsal added (CPU).""" +from __future__ import annotations + +import json +from pathlib import Path + +import pytest + +from yeto.rl.adapters.miles.models import glm5_next +from yeto.rl.adapters.sampler_engines import sglang_load_report as slr +from yeto.rl.adapters.sampler_engines.server import ServerConfig +from yeto.rl.adapters.sampler_engines.sglang_server import SGLangOptions, SGLangServerEngine +from yeto.rl.engine.sampler_checks import FixedPrompt, PostLoadChecks +from yeto.rl.engine.sampler_engine import Completion + +REPO = Path(__file__).resolve().parents[1] + + +def test_recipe_takes_expert_and_dense_counts_from_the_checkpoint(): + shape = glm5_next.checkpoint_geometry({"text_config": {"num_hidden_layers": 4, "first_k_dense_replace": 1, + "n_routed_experts": 32}}) + assert shape == {"num_layers": 4, "first_k_dense_replace": 1, "num_experts": 32} + argv = glm5_next.model_argv(**shape) + assert argv[argv.index("--num-experts") + 1] == "32" + assert argv[argv.index("--moe-layer-freq") + 1] == "[0,1,1,1]" + full = glm5_next.model_argv() + assert full[full.index("--num-experts") + 1] == "288" and full.count("--num-experts") == 1 + with pytest.raises(ValueError): + glm5_next.model_argv(num_experts=0) + + +class _Param: + def __init__(self, calls: list[str]) -> None: + self.weight_loader = lambda *a, **k: calls.append(a[2] if len(a) > 2 else "?") + + +class _FakeModel: + """Maps names like SGLang's GLM5.3 loader: an expert data tensor reaches a loader + only when its fused name is a parameter.""" + + def __init__(self) -> None: + self.calls: list[str] = [] + self.params = {"model.layers.1.mlp.experts.w13_weight": _Param(self.calls), + "model.layers.1.mlp.experts.w2_weight": _Param(self.calls), + "model.layers.1.mlp.experts.w13_weight_scale": _Param(self.calls)} + + def named_parameters(self): + return iter(self.params.items()) + + def load_weights(self, weights): + for name, tensor in weights: + fused = (name.replace("experts.0.gate_proj.", "experts.w13_").replace("experts.0.up_proj.", "experts.w13_") + .replace("experts.0.down_proj.", "experts.w2_").replace("language_model.", "")) + param = self.params.get(fused) + if param is not None: + param.weight_loader(param, tensor, fused) + + +def _weights(suffix: str): + base = "model.language_model.layers.1.mlp.experts.0." + return [(f"{base}{m}_proj.{suffix}", 0) for m in ("gate", "up", "down")] + \ + [(f"{base}{m}_proj.weight_scale", 0) for m in ("gate", "up", "down")] + + +def test_load_report_counts_only_loaded_expert_matrices(tmp_path, monkeypatch): + cls = type("M", (_FakeModel,), {}) + slr.install(cls) + slr.install(cls) # idempotent + monkeypatch.setenv(slr.REPORT_DIR_ENV, str(tmp_path)) + monkeypatch.delenv(slr.RENAME_ENV, raising=False) + cls().load_weights(_weights("weight_packed")) # no rename: SGLang skips the packed data silently + assert slr.read_reports(tmp_path)["expert_matrices"] == 0 + slr.clear_reports(tmp_path) + assert slr.read_reports(tmp_path) is None + monkeypatch.setenv(slr.RENAME_ENV, "1") + model = cls() + model.load_weights(_weights("weight_packed")) + report = slr.read_reports(tmp_path) + assert report["expert_matrices"] == 3 and report["ranks"] == 1 + assert sorted(set(model.calls)) == ["model.layers.1.mlp.experts.w13_weight", "model.layers.1.mlp.experts.w13_weight_scale", + "model.layers.1.mlp.experts.w2_weight"] + assert list(slr.rename_packed([("a.mlp.experts.3.up_proj.weight_packed", 1), ("x.weight_packed", 2)])) == \ + [("a.mlp.experts.3.up_proj.weight", 1), ("x.weight_packed", 2)] + + +def test_patch_model_file_is_idempotent(tmp_path): + f = tmp_path / "glm5_next.py" + f.write_text("class A: ...\nEntryClass = [A]\n") + assert slr.patch_model_file(f) is True and slr.patch_model_file(f) is False + assert f.read_text().count(slr.MARKER) == 1 + compile(f.read_text(), str(f), "exec") + (tmp_path / "other.py").write_text("x = 1\n") + with pytest.raises(ValueError): + slr.patch_model_file(tmp_path / "other.py") + + +def _engine(tmp_path: Path, tp: int = 1) -> SGLangServerEngine: + config = ServerConfig(base_dir=tmp_path, card_model="H100", device_family="cuda", engine_version="x", + precision="mxfp4", quant_scheme="mxfp4-g32-rtn", topology=f"tp{tp}", max_model_len=64) + return SGLangServerEngine(config, SGLangOptions(tp=tp, renames_weight_packed=True, + load_report_dir=str(tmp_path / "reports"))) + + +def test_adapter_reads_the_report_and_exports_the_hook_env(tmp_path): + engine = _engine(tmp_path, tp=2) + env = engine.env() + assert env[slr.RENAME_ENV] == "1" and env[slr.REPORT_DIR_ENV] == str(tmp_path / "reports") + assert engine.loaded_expert_matrices() is None + reports = tmp_path / "reports" + reports.mkdir() + (reports / "rank0-1.json").write_text(json.dumps({"expert_matrices": 288})) + assert engine.loaded_expert_matrices() == 0 # rank 1 did not report + (reports / "rank1-2.json").write_text(json.dumps({"expert_matrices": 287})) + assert engine.loaded_expert_matrices() == 287 + with pytest.raises(Exception): + SGLangServerEngine(ServerConfig(base_dir=tmp_path, card_model="H100", device_family="cuda", engine_version="x", + precision="mxfp4", quant_scheme="mxfp4-g32-rtn", topology="tp2", + max_model_len=64), SGLangOptions(tp=1)) + + +class _Greedy: + def generate(self, prompts, sampling): + return [[Completion(text="the the the", token_ids=(1, 2, 3), logprobs=(-1.0,) * 3, finish_reason="length")] + for _ in prompts] + + +def test_fixed_answers_may_be_recorded_without_failing(): + prompts = (FixedPrompt("Q: 2+3?", "5"),) + strict = PostLoadChecks(fixed_prompts=prompts).fixed_prompt_answers(_Greedy()) + assert strict["passed"] is False and strict["answers"][0]["problems"] == ["answer missing"] + loose = PostLoadChecks(fixed_prompts=prompts, require_fixed_answers=False).fixed_prompt_answers(_Greedy()) + assert loose["passed"] is True and loose["answers"][0]["answered"] is False + assert loose["answers_required"] is False + + +def test_mini_modal_plan_has_no_errors_and_uses_the_volume_store(): + from yeto.rl.stage1 import expand, load_plan + + out = expand(load_plan(REPO / "examples/stage1/glm5next-mini-modal.yaml")) + assert out.errors == [] and out.gaps == [] + env = out.trainer["env"] + assert env["YETO_RL_STORE_IMPL"] == "yeto.cloud.modal_ckpt_store:store_from_env" + assert env["YETO_RL_STORE_MODAL_VOLUME"] == "yeto-s25-g1-store" + argv = out.trainer["miles"]["island_argv"] + assert argv[argv.index("--miles-commit") + 1] == "1618b5c38db80e17ec928cfa47fe832cb18565fc" + assert out.trainer["miles"]["model_args"][out.trainer["miles"]["model_args"].index("--num-experts") + 1] == "32" + spec = out.sampler["engines"][0]["spec"] + assert spec["checks"] == {"require_fixed_answers": False, "fixed_max_tokens": 64, + "expected_expert_matrices": 288} + assert spec["engine"]["options"]["load_report_dir"].endswith("/load-report") + assert "--frozen-base" in out.requant["argv"] + # the PPU plan keeps OSS and passes no --miles-commit for its short commit + ppu = expand(load_plan(REPO / "examples/stage1/glm53-flash-ppu.yaml")) + assert ppu.errors == [] and "--miles-commit" not in ppu.trainer["miles"]["island_argv"] + assert "--frozen-base" not in ppu.requant["argv"] diff --git a/tools/probes/s25_g1/image_probe.py b/tools/probes/s25_g1/image_probe.py new file mode 100644 index 000000000..24701de94 --- /dev/null +++ b/tools/probes/s25_g1/image_probe.py @@ -0,0 +1,59 @@ +"""CPU-only probe of the Miles ports image for the S25 G1 stage-1 mini rehearsal. + +modal run tools/probes/s25_g1/image_probe.py (no GPU; prints one JSON report) +""" +import modal + +# Same digest as yeto.rl.adapters.miles.pins.MILES_NEXT_IMAGE (the container has no yeto). +MILES_NEXT_IMAGE = ("docker:ghcr.io/michaellchung/yeto-miles-ports@sha256:" + "5b9dc2a9e6e315d73dabfbc6faae428b59cc01aabec6189645a5c8c80d6e3721") + +app = modal.App("yeto-s25-g1-probe") +image = modal.Image.from_registry(MILES_NEXT_IMAGE.removeprefix("docker:")) + +PROBE = r''' +import importlib, json, subprocess, sys +out = {} +def tryit(key, fn): + try: out[key] = fn() + except Exception as e: out[key] = f"ERR {type(e).__name__}: {str(e)[:300]}" +def ver(m): return getattr(importlib.import_module(m), "__version__", "?") +for m in ("torch", "transformer_engine", "megatron.core", "sglang", "fla", "tilelang", "triton", + "triton_kernels", "mbridge", "compressed_tensors", "flash_attn", "ray", "transformers"): + tryit("ver:" + m, lambda m=m: ver(m)) +def tcfg(): + from megatron.core.transformer import TransformerConfig + import dataclasses + names = {f.name for f in dataclasses.fields(TransformerConfig)} + return {k: k in names for k in ("enable_hyper_connections", "num_residual_streams", + "mhc_sinkhorn_iterations", "use_fused_mhc", "experimental_attention_variant")} +tryit("megatron_mhc_fields", tcfg) +tryit("sglang_glm5_next", lambda: bool(importlib.import_module("sglang.srt.models.glm5_next"))) +tryit("sglang_mxfp4", lambda: bool(importlib.import_module("sglang.srt.layers.quantization.mxfp4").Mxfp4MoEMethod)) +tryit("fla_kda", lambda: bool(importlib.import_module("fla.ops.kda"))) +def git(d): + return subprocess.run(["git", "-C", d, "log", "-1", "--format=%H %cd %s"], capture_output=True, text=True).stdout.strip() +for d in ("/root/miles", "/root/Megatron-LM", "/sgl-workspace/sglang"): + tryit("git:" + d, lambda d=d: git(d)) +tryit("manifest", lambda: open("/opt/yeto/image-manifest.json").read()[:1500]) +tryit("python", lambda: sys.version) +def sh(c): return subprocess.run(["bash", "-lc", c], capture_output=True, text=True).stdout[-1500:] +tryit("mcore_mhc_grep", lambda: sh("grep -rln 'enable_hyper_connections' /root/Megatron-LM/megatron | head")) +tryit("mcore_remote", lambda: sh("git -C /root/Megatron-LM remote -v | head -2; git -C /root/Megatron-LM log --oneline -3")) +tryit("te_pip", lambda: sh("pip show transformer_engine transformer-engine-cu12 transformer_engine_torch 2>/dev/null | grep -E 'Name|Version'")) +tryit("triton_kernels", lambda: sh("python3 -c 'import triton_kernels,os;print(os.path.dirname(triton_kernels.__file__))'; pip show triton_kernels 2>/dev/null|head -3")) +tryit("mem_cpu", lambda: sh("nproc; free -g | head -2; df -h /root | tail -1")) +print(json.dumps(out, indent=1)) +''' + + +@app.function(image=image, cpu=2, memory=8192, timeout=900) +def probe(): + import subprocess + r = subprocess.run(["python3", "-c", PROBE], capture_output=True, text=True) + return r.stdout + "\nSTDERR-TAIL:\n" + r.stderr[-3000:] + + +@app.local_entrypoint() +def main(): + print(probe.remote()) diff --git a/tools/probes/s25_g1/make_mini.py b/tools/probes/s25_g1/make_mini.py new file mode 100644 index 000000000..8e49f9d7a --- /dev/null +++ b/tools/probes/s25_g1/make_mini.py @@ -0,0 +1,148 @@ +"""Build the S25 G1 rehearsal checkpoints from Miles' 4-layer GLM5.3-Flash CI model. + +Runs inside the prep container (python3 make_mini.py ROOT). Writes, under ROOT: + +* ``src-4layer/``: CharyZeng/GLM-5.3-Flash-4layer at REVISION (288 experts, bf16); +* ``glm5next-mini-bf16/``: the same model with the first KEEP routed experts per + MoE layer (router rows and expert bias cut to match; n_routed_experts=KEEP); +* ``glm5next-mini-mxfp4-template/``: layout.json + metadata for the MXFP4 layout + of yetalabs/apex-flash-1-mxfp4-experts@70e86a0 (routed experts + ``weight_packed`` U8 + E8M0 ``weight_scale``, everything else as the bf16 source), + with that export's quantization_config; +* ``glm5next-mini-mxfp4-r0/``: round 0 for the sampler, converted from the bf16 + mini model by the N8.5 converter (``requantize_checkpoint``); +* ``datasets/gsm8k-0cbd9f3/train.jsonl`` and the HF dataset cache under ``hf/``. + +Prints one JSON summary line (``MINI-SUMMARY {...}``). +""" +from __future__ import annotations + +import json +import os +import re +import shutil +import sys +from pathlib import Path + +REPO = "CharyZeng/GLM-5.3-Flash-4layer" +REVISION = "7fb7883207fe73ea04430ac392f6ffe2b9f936c4" +MXFP4_REPO = "yetalabs/apex-flash-1-mxfp4-experts" +MXFP4_REVISION = "70e86a0" +DATA_REPO = "zhuzilin/gsm8k" +DATA_REVISION = "0cbd9f31d91ac21a7613dcbc7fef992adac459ae" +KEEP = int(os.environ.get("MINI_EXPERTS", "32")) +EXPERT = re.compile(r"\.mlp\.experts\.(\d+)\.") +ROUTED = re.compile(r"^(.*\.mlp\.experts\.\d+\.(?:gate|up|down)_proj)\.weight$") + + +def truncate(src: Path, dst: Path) -> dict: + import torch + from safetensors import safe_open + from safetensors.torch import save_file + + dst.mkdir(parents=True, exist_ok=True) + index = json.loads((src / "model.safetensors.index.json").read_text())["weight_map"] + weight_map, total, cut = {}, 0, 0 + for shard in sorted(set(index.values())): + out = {} + with safe_open(str(src / shard), framework="pt") as f: + for name in f.keys(): + m = EXPERT.search(name) + if m and int(m.group(1)) >= KEEP: + cut += 1 + continue + t = f.get_tensor(name) + if name.endswith(".mlp.gate.weight") or name.endswith(".mlp.gate.e_score_correction_bias"): + t = t[:KEEP].clone() + out[name] = t.contiguous() + if out: + save_file(out, str(dst / shard), metadata={"format": "pt"}) + for name, t in out.items(): + weight_map[name] = shard + total += t.numel() * t.element_size() + print(f"[mini] {shard}: kept {len(out)}", flush=True) + (dst / "model.safetensors.index.json").write_text( + json.dumps({"metadata": {"total_size": total}, "weight_map": weight_map}, indent=2, sort_keys=True)) + for f in src.iterdir(): + if f.is_file() and f.suffix != ".safetensors" and f.name != "model.safetensors.index.json" \ + and not f.name.startswith("."): + shutil.copy2(f, dst / f.name) + config = json.loads((dst / "config.json").read_text()) + config["text_config"]["n_routed_experts"] = KEEP + (dst / "config.json").write_text(json.dumps(config, indent=2)) + return {"tensors": len(weight_map), "bytes": total, "cut_expert_tensors": cut} + + +def make_template(bf16: Path, template: Path, quant_config: dict) -> dict: + from yeto.rl.quant.requant import (LAYOUT_FILE, LAYOUT_SCHEMA, CheckpointLayout, TensorSpec, + metadata_files, read_header) + from yeto.rl.quant.mxfp4 import GROUP_SIZE + + weight_map = json.loads((bf16 / "model.safetensors.index.json").read_text())["weight_map"] + tensors: dict[str, TensorSpec] = {} + headers = {s: read_header(bf16 / s) for s in sorted(set(weight_map.values()))} + for name, shard in weight_map.items(): + info = headers[shard][name] + m = ROUTED.match(name) + if m: + rows, cols = info["shape"] + tensors[m.group(1) + ".weight_packed"] = TensorSpec("U8", (rows, cols // 2), shard) + tensors[m.group(1) + ".weight_scale"] = TensorSpec("U8", (rows, cols // GROUP_SIZE), shard) + else: + tensors[name] = TensorSpec(info["dtype"], tuple(info["shape"]), shard) + layout = CheckpointLayout(tensors, quant_config) + layout.validate() + template.mkdir(parents=True, exist_ok=True) + for f in metadata_files(bf16): + shutil.copy2(f, template / f.name) + config = json.loads((template / "config.json").read_text()) + config["quantization_config"] = quant_config + (template / "config.json").write_text(json.dumps(config, indent=2)) + payload = {"schema": LAYOUT_SCHEMA, "quantization_config": quant_config, + "tensors": {n: {"dtype": t.dtype, "shape": list(t.shape), "shard": t.shard} + for n, t in sorted(tensors.items())}} + (template / LAYOUT_FILE).write_text(json.dumps(payload, sort_keys=True)) + return {"quantized_modules": len(layout.quantized_modules), "tensors": len(tensors)} + + +def main(root: str) -> None: + from huggingface_hub import hf_hub_download, snapshot_download + + root_dir = Path(root) + os.environ.setdefault("HF_HOME", str(root_dir / "hf")) + summary: dict = {"keep_experts": KEEP} + src = Path(snapshot_download(REPO, revision=REVISION, local_dir=str(root_dir / "src-4layer"))) + bf16 = root_dir / "glm5next-mini-bf16-7fb7883" + if not (bf16 / "model.safetensors.index.json").is_file(): + summary["truncate"] = truncate(src, bf16) + mx_config = json.loads(Path(hf_hub_download(MXFP4_REPO, "config.json", revision=MXFP4_REVISION)).read_text()) + quant_config = mx_config["quantization_config"] + template = root_dir / "glm5next-mini-mxfp4-template" + summary["template"] = make_template(bf16, template, quant_config) + from yeto.rl.quant.awq import AwqSmoothing + from yeto.rl.quant.requant import CheckpointLayout, requantize_checkpoint + + r0 = root_dir / "glm5next-mini-mxfp4-r0-7fb7883" + if r0.exists(): + shutil.rmtree(r0) + summary["r0"] = requantize_checkpoint(bf16, template, r0, AwqSmoothing.none()) + check = CheckpointLayout.from_checkpoint(r0) # the stage-1 template path (no layout.json) + summary["r0_quantized_modules"] = len(check.quantized_modules) + data = Path(snapshot_download(DATA_REPO, repo_type="dataset", revision=DATA_REVISION)) + import pandas as pd + + rows = pd.read_parquet(data / "train.parquet") + out = root_dir / "datasets/gsm8k-0cbd9f3" + out.mkdir(parents=True, exist_ok=True) + with open(out / "train.jsonl", "w", encoding="utf-8") as f: + for row in rows.to_dict(orient="records"): + f.write(json.dumps({k: (v.tolist() if hasattr(v, "tolist") else v) for k, v in row.items()}, + default=str) + "\n") + summary["gsm8k_rows"] = len(rows) + summary["gsm8k_columns"] = list(rows.columns) + summary["gsm8k_row0"] = str(rows.iloc[0].to_dict())[:400] + print("MINI-SUMMARY " + json.dumps(summary), flush=True) + + +if __name__ == "__main__": + main(sys.argv[1]) diff --git a/tools/probes/s25_g1/modal_app.py b/tools/probes/s25_g1/modal_app.py new file mode 100644 index 000000000..ce842c249 --- /dev/null +++ b/tools/probes/s25_g1/modal_app.py @@ -0,0 +1,263 @@ +"""S25 T-G1: stage-1 mini rehearsal on Modal (examples/stage1/glm5next-mini-modal.yaml). + +The containers run the commands ``yeto-stage1 plan`` generates for that plan +(computed inside each container, so ``--source-sha256`` matches the mounted +tree): the Miles trainer-only island with the MXFP4 re-quantizer side process, +and the engine sampler island. Only the nodes (Modal containers instead of SSH +hosts) and the store (Modal Volume through the yeto store hook instead of OSS) +differ from the PPU plan. + + modal run tools/probes/s25_g1/modal_app.py::prep # CPU: checkpoints + data + modal run --detach tools/probes/s25_g1/modal_app.py::smoke # 1 x H100: SGLang + post-load checks + modal run --detach tools/probes/s25_g1/modal_app.py::launch # 3 rounds, cut after 2, resume + +``smoke`` and ``launch`` spawn and return; the local client may exit +(never kill a ``.remote()`` caller: memory modal-detach-kill-cancels). +Raw logs go to the store Volume under ``logs//``. +""" +from __future__ import annotations + +import json +import os +import shlex +import shutil +import subprocess +import sys +import time +from pathlib import Path + +import modal + +HERE = Path(__file__).resolve() +REPO = HERE.parents[3] +WORKDIR = "/root/sky_workdir" +PLAN = "examples/stage1/glm5next-mini-modal.yaml" +IMAGE = ("ghcr.io/michaellchung/yeto-miles-ports@sha256:" + "5b9dc2a9e6e315d73dabfbc6faae428b59cc01aabec6189645a5c8c80d6e3721") +MILES_COMMIT = "1618b5c38db80e17ec928cfa47fe832cb18565fc" # s25-distopt +MODELS_VOLUME, STORE_VOLUME = "yeto-s25-g1-models", "yeto-s25-g1-store" +GPU_NAME = "H100" + +app = modal.App("yeto-s25-g1-stage1mini") +models = modal.Volume.from_name(MODELS_VOLUME, create_if_missing=True, version=2) +store = modal.Volume.from_name(STORE_VOLUME, create_if_missing=True, version=2) +VOLUMES = {"/data/models": models, "/data/store": store} + +_PATCH_SGLANG = ( + "import importlib.util, sys; sys.path.insert(0, '/tmp/yeto-slr');" + "import sglang_load_report as r;" + "root = importlib.util.find_spec('sglang').submodule_search_locations[0];" + "f = root + '/srt/models/glm5_next.py'; print('patched', f, r.patch_model_file(f))" +) +image = ( + modal.Image.from_registry(IMAGE) + .add_local_file(str(REPO / ".g1-build/miles-s25-distopt.bundle"), "/tmp/miles.bundle", copy=True) + .add_local_file(str(REPO / "yeto/rl/adapters/sampler_engines/sglang_load_report.py"), + "/tmp/yeto-slr/sglang_load_report.py", copy=True) + .run_commands( + "git -C /root/miles fetch /tmp/miles.bundle s25-distopt", + f"git -C /root/miles checkout --detach {MILES_COMMIT}", + "git -C /root/miles status --porcelain --untracked-files=all | head -5", + f"python3 -c \"{_PATCH_SGLANG}\"", + ) + # not PYTHONPATH: Modal's own runtime path must stay intact; subprocesses get it (_env) + .env({"HOME": "/root", "PYTHONUNBUFFERED": "1", "HF_HOME": "/data/models/hf"}) + .add_local_dir(str(REPO), WORKDIR, copy=False, + ignore=[".git", ".g1-build", ".venv", ".claude", "syncer/target", "**/__pycache__", + "**/*.pyc", "evidence*", "openspec", "data"]) +) + + +def _hf_secret() -> modal.Secret: + token = Path("~/.cache/huggingface/token").expanduser() + return modal.Secret.from_dict({"HF_TOKEN": token.read_text().strip()} if token.is_file() else {}) + + +def _log_dir(tag: str) -> Path: + path = Path("/data/store/logs") / tag + path.mkdir(parents=True, exist_ok=True) + return path + + +def _commit() -> None: + try: + store.commit() + except Exception as error: # noqa: BLE001 - evidence copy only + print(f"[g1] store commit failed: {error}", flush=True) + + +def _sh(cmd: str, log: Path | None = None, env: dict | None = None, check: bool = True) -> str: + print(f"[g1] $ {cmd}", flush=True) + r = subprocess.run(["bash", "-lc", cmd], capture_output=True, text=True, env=env or _env({}), cwd=WORKDIR) + out = r.stdout + r.stderr + if log is not None: + with open(log, "a", encoding="utf-8") as f: + f.write(f"$ {cmd}\n{out}\n") + if check and r.returncode: + raise RuntimeError(f"{cmd} -> {r.returncode}: {out[-2000:]}") + return out + + +def _gpu_facts(log: Path, want: int) -> dict: + out = _sh("nvidia-smi --query-gpu=index,name,memory.total,driver_version --format=csv,noheader", log) + names = [line.split(",")[1].strip() for line in out.strip().splitlines()] + if len(names) != want or not all(GPU_NAME in n for n in names): + raise RuntimeError(f"GPU assert failed: want {want} x {GPU_NAME}, got {names}") + _sh("nvidia-smi; free -g; nproc; df -h /data/models /data/store /root", log, check=False) + return {"gpus": names} + + +def _plan(log: Path) -> dict: + out = _sh(f"python3 -m yeto.rl.stage1 plan {PLAN} --json 2>/dev/null; true", None, check=False) + start = out.index("{") + plan = json.loads(out[start: out.rindex("}") + 1]) + (log / "plan.json").write_text(json.dumps(plan, indent=1)) + if plan["errors"] or plan["gaps"]: + raise RuntimeError(f"plan errors {plan['errors']} gaps {plan['gaps']}") + return plan + + +def _nvml_sampler(path: Path) -> subprocess.Popen: + return subprocess.Popen( + ["bash", "-c", "nvidia-smi --query-gpu=timestamp,index,utilization.gpu,memory.used,power.draw " + f"--format=csv,noheader -l 5 >> {shlex.quote(str(path))}"]) + + +def _env(extra: dict) -> dict: + path = ":".join(p for p in (WORKDIR, os.environ.get("PYTHONPATH", "")) if p) + return {**os.environ, "PYTHONPATH": path, **{k: str(v) for k, v in extra.items()}} + + +# --------------------------------------------------------------------- prep (CPU) +@app.function(image=image, volumes={"/data/models": models}, cpu=8, memory=65536, timeout=7200, + secrets=[_hf_secret()]) +def prep() -> str: + out = subprocess.run(["python3", f"{WORKDIR}/tools/probes/s25_g1/make_mini.py", "/data/models"], + capture_output=True, text=True, cwd=WORKDIR, env=_env({})) + models.commit() + return out.stdout[-6000:] + "\nSTDERR:\n" + out.stderr[-4000:] + + +# --------------------------------------------------------------------- sampler smoke +@app.function(image=image, volumes=VOLUMES, gpu=f"{GPU_NAME}!:1", cpu=8, memory=65536, timeout=2400) +def smoke_run(tag: str) -> None: + """Start SGLang on round 0 through the adapter and run the post-load checks once.""" + log = _log_dir(tag) + facts = _gpu_facts(log, 1) + nvml = _nvml_sampler(log / "nvml.csv") + try: + plan = _plan(log) + run = plan["sampler"]["engines"][0] + os.environ.update({k: str(v) for k, v in run["env"].items()}) + sys.path.insert(0, WORKDIR) + from yeto.rl.adapters.sampler_engines.from_spec import build + from yeto.rl.engine.sampler_checks import PostLoadChecks + + engine = build(run["spec"]["engine"]) + checks = PostLoadChecks(**run["spec"]["checks"]) + began = time.time() + report = engine.start(None, 0) + started = time.time() - began + results = checks.run(engine, expected_experts=int(run["spec"]["checks"]["expected_expert_matrices"])) + from yeto.rl.engine.sampler_engine import SamplingParams + + sample = engine.generate(["Question: What is 2 + 3?\nAnswer:"], + SamplingParams(max_tokens=32, temperature=1.0, n=2, seed=1).check()) + record = {**facts, "start_s": round(started, 1), "load_report": report.detail, "checks": results, + "sample": [[c.text for c in g] for g in sample]} + engine.stop() + except Exception as error: # noqa: BLE001 - the record is the evidence + record = {**facts, "error": f"{type(error).__name__}: {error}"} + finally: + nvml.terminate() + (log / "smoke.json").write_text(json.dumps(record, indent=1, default=str)) + print(json.dumps(record, default=str)[:4000], flush=True) + _commit() + + +@app.local_entrypoint() +def smoke(tag: str = "smoke-1") -> None: + call = smoke_run.spawn(tag) + print(f"spawned smoke_run {call.object_id}") + + +# --------------------------------------------------------------------- the islands +@app.function(image=image, volumes=VOLUMES, gpu=f"{GPU_NAME}!:1", cpu=8, memory=65536, timeout=7200) +def sampler_island(tag: str) -> int: + log = _log_dir(tag) + _gpu_facts(log, 1) + nvml = _nvml_sampler(log / "nvml.csv") + plan = _plan(log) + run = plan["sampler"]["engines"][0] + spec_path = Path(run["spec_path"]) + spec_path.parent.mkdir(parents=True, exist_ok=True) + spec_path.write_text(json.dumps(run["spec"], indent=1)) + began = time.time() + with open(log / "sampler.log", "w", encoding="utf-8") as f: + rc = subprocess.run(run["argv"], env=_env(run["env"]), cwd=WORKDIR, stdout=f, stderr=subprocess.STDOUT).returncode + nvml.terminate() + (log / "rc.json").write_text(json.dumps({"rc": rc, "wall_s": round(time.time() - began, 1)})) + state = Path(run["spec"]["event_tape"]).parent + if state.is_dir(): + shutil.copytree(state, log / "state", dirs_exist_ok=True) + _commit() + return rc + + +@app.function(image=image, volumes=VOLUMES, gpu=f"{GPU_NAME}!:2", cpu=12, memory=163840, timeout=5400) +def trainer_island(tag: str, stop_after_rounds: int | None = None) -> int: + log = _log_dir(tag) + _gpu_facts(log, 2) + nvml = _nvml_sampler(log / "nvml.csv") + plan = _plan(log) + trainer = plan["trainer"] + env = _env(trainer["env"]) + island = list(trainer["miles"]["island_argv"]) + if stop_after_rounds is not None: + island += ["--rl-stop-after-rounds", str(int(stop_after_rounds))] + (log / "argv.json").write_text(json.dumps({"island": island, "requant": plan["requant"]["loop_argv"], + "env": trainer["env"]}, indent=1)) + Path("/data/yeto", plan["run_id"]).mkdir(parents=True, exist_ok=True) + _sh("ray start --head --port=6379 --node-ip-address=127.0.0.1 --num-gpus=2 --disable-usage-stats", + log / "ray.log", env=env) + began = time.time() + with open(log / "requant.log", "w", encoding="utf-8") as rq, open(log / "trainer.log", "w", encoding="utf-8") as f: + requant = subprocess.Popen(plan["requant"]["loop_argv"], env=env, cwd=WORKDIR, stdout=rq, + stderr=subprocess.STDOUT) + rc = subprocess.run(island, env=env, cwd=WORKDIR, stdout=f, stderr=subprocess.STDOUT).returncode + requant.terminate() + try: + requant.wait(timeout=60) + except subprocess.TimeoutExpired: + requant.kill() + final = subprocess.run(plan["requant"]["argv"], env=env, cwd=WORKDIR, stdout=rq, + stderr=subprocess.STDOUT).returncode + nvml.terminate() + (log / "rc.json").write_text(json.dumps({"rc": rc, "final_requant_rc": final, + "wall_s": round(time.time() - began, 1)})) + shutil.copytree(f"/data/yeto/{plan['run_id']}", log / "state", dirs_exist_ok=True, + ignore=shutil.ignore_patterns("*.pt", "*.safetensors")) + _sh("cp -r /tmp/ray/session_latest/logs " + shlex.quote(str(log / "ray-logs")) + " || true", check=False) + _commit() + return rc + + +@app.function(image=image, cpu=1, memory=1024, timeout=3 * 3600) +def orchestrate(tag: str, cut_after: int = 2) -> dict: + sampler = sampler_island.spawn(f"{tag}/sampler") + first = trainer_island.remote(f"{tag}/trainer-a", cut_after) + second = trainer_island.remote(f"{tag}/trainer-b", None) if first == 0 else None + if first != 0 or second != 0: # no trainer will release again: do not leave the sampler card waiting + sampler.cancel() + sampler_rc = "cancelled" + else: + sampler_rc = sampler.get() + result = {"trainer_a": first, "trainer_b": second, "sampler": sampler_rc} + print(json.dumps(result), flush=True) + return result + + +@app.local_entrypoint() +def launch(tag: str = "run-1", cut_after: int = 2) -> None: + call = orchestrate.spawn(tag, cut_after) + print(f"spawned orchestrate {call.object_id}") diff --git a/yeto/launch/ssh_trainer_island.py b/yeto/launch/ssh_trainer_island.py index 56b3de31c..f47318100 100644 --- a/yeto/launch/ssh_trainer_island.py +++ b/yeto/launch/ssh_trainer_island.py @@ -480,10 +480,16 @@ def ssh_binding(args: Any, island: TrainerIsland, *, journal_dir: Path, files: A def requant_argv(store: str, template: str, *, recipe: str | None = None, awq_scales: str | None = None, - python: str = "python") -> tuple[str, ...]: + python: str = "python", frozen_base: str | None = None, + frozen_pattern: str | None = None) -> tuple[str, ...]: """The one-pass argv of :mod:`yeto.rl.quant.release` (R2 with ``recipe`` and - ``awq_scales``; R1 without). The start script adds ``--loop`` for the side process.""" + ``awq_scales``; R1 without; ``frozen_base``/``frozen_pattern``: untrained + tensors the release omits). The start script adds ``--loop`` for the side process.""" argv = [python, "-m", "yeto.rl.quant.release", "--store", store, "--template", template] + if bool(frozen_base) != bool(frozen_pattern): + raise SshIslandError("give both frozen_base and frozen_pattern, or neither") + if frozen_base: + argv += ["--frozen-base", frozen_base, "--frozen-pattern", frozen_pattern] if recipe or awq_scales: if not (recipe and awq_scales): raise SshIslandError("give both recipe and awq_scales (R2), or neither (R1)") diff --git a/yeto/rl/adapters/miles/config.py b/yeto/rl/adapters/miles/config.py index b6ead7169..449503da3 100644 --- a/yeto/rl/adapters/miles/config.py +++ b/yeto/rl/adapters/miles/config.py @@ -25,6 +25,7 @@ import json import os import sys +from pathlib import Path from collections.abc import Callable, Iterator, Mapping, Sequence from dataclasses import dataclass, field from typing import Any @@ -1150,8 +1151,13 @@ def translate_run_config( values.extend((flag, engine)) if recipe.name == RECIPE_GLM5_NEXT: # the plugin spec replaces the Bridge-geometry model flags + # experts and dense-layer count from the checkpoint (a cut-down model differs from GLM5.3-Flash) + hf_config = Path(config.hf_checkpoint) / "config.json" + shape = (glm5_next.checkpoint_geometry(json.loads(hf_config.read_text())) + if hf_config.is_file() else {}) + shape["num_layers"] = geometry.num_layers values = glm5_next.apply_recipe(values, glm5_next.model_argv( - num_layers=geometry.num_layers, transformer_impl=recipe.transformer_impl)) + **shape, transformer_impl=recipe.transformer_impl)) values.extend(extra_argv) reject_fault_tolerance_flags(values[1:]) # generated argv never carries FT flags _refuse_decoupled_ppo_breakers(values) # also when the spec emitted the flag diff --git a/yeto/rl/adapters/miles/island_entry.py b/yeto/rl/adapters/miles/island_entry.py index 9a031e335..9d4cc0f0a 100644 --- a/yeto/rl/adapters/miles/island_entry.py +++ b/yeto/rl/adapters/miles/island_entry.py @@ -1092,6 +1092,19 @@ def _preflight_codex_openenv_adapter(args, profile_name: str) -> None: forward_legacy_openenv_preflight(args, profile_name) +def miles_revision_pins(args: Any) -> Any: + """The Miles commit the island checks: ``--miles-commit`` (full SHA) or the pin.""" + from yeto.provenance import is_immutable_commit + from yeto.rl.adapters.miles.pins import MILES_NEXT_PINS, MILES_NEXT_REPOSITORY, MilesRevisionPins + + commit = getattr(args, "miles_commit", None) + if not commit: + return MILES_NEXT_PINS + if not is_immutable_commit(commit) or len(str(commit)) != 40: + raise ValueError(f"--miles-commit must be a full 40-hex commit, got {commit!r}") + return MilesRevisionPins(MILES_NEXT_REPOSITORY, str(commit).lower()) + + def _preflight_codex_harness(args) -> None: """Fail before model allocation if the signed Codex surface has drifted.""" @@ -2340,7 +2353,7 @@ def main(argv=None) -> None: miles_root = str(Path(args.miles_root).expanduser().resolve()) if miles_root not in sys.path: sys.path.insert(0, miles_root) - verify_miles_revision(miles_root) + verify_miles_revision(miles_root, expected=miles_revision_pins(args)) _preflight_codex_harness(args) probe_syncers(args) diff --git a/yeto/rl/adapters/miles/models/glm5_next.py b/yeto/rl/adapters/miles/models/glm5_next.py index 5ab9b8c5e..449747b60 100644 --- a/yeto/rl/adapters/miles/models/glm5_next.py +++ b/yeto/rl/adapters/miles/models/glm5_next.py @@ -50,7 +50,6 @@ "--v-head-dim", "256", "--kv-channels", "256", "--qk-layernorm", - "--num-experts", str(NUM_EXPERTS), "--moe-router-topk", "8", "--moe-router-score-function", "sigmoid", "--moe-router-pre-softmax", @@ -93,6 +92,20 @@ def moe_layer_freq(num_layers: int = NUM_LAYERS, first_k_dense_replace: int = FI return "[" + ",".join(["0"] * first_k_dense_replace + ["1"] * (num_layers - first_k_dense_replace)) + "]" +def checkpoint_geometry(config: dict) -> dict[str, int]: + """``num_layers``, ``first_k_dense_replace`` and ``num_experts`` of an HF + ``config.json`` (the text config of ``Glm5NextForConditionalGeneration``), + as :func:`model_argv` keywords; a key the config lacks keeps the default.""" + text = config.get("text_config") or config + out = {} + for key, hf_key in (("num_layers", "num_hidden_layers"), + ("first_k_dense_replace", "first_k_dense_replace"), + ("num_experts", "n_routed_experts")): + if text.get(hf_key) is not None: + out[key] = int(text[hf_key]) + return out + + def check_transformer_impl(transformer_impl: str) -> str: if transformer_impl not in TRANSFORMER_IMPLS: raise ValueError(f"transformer_impl must be one of {TRANSFORMER_IMPLS}, got {transformer_impl!r}") @@ -100,11 +113,16 @@ def check_transformer_impl(transformer_impl: str) -> str: def model_argv(*, num_layers: int = NUM_LAYERS, first_k_dense_replace: int = FIRST_K_DENSE_REPLACE, - transformer_impl: str = TRANSFORMER_ENGINE) -> list[str]: + num_experts: int = NUM_EXPERTS, transformer_impl: str = TRANSFORMER_ENGINE) -> list[str]: """The model arguments of the recipe (module doc) and the run arguments the - model needs (:data:`RUN_ARGV`).""" + model needs (:data:`RUN_ARGV`). ``num_layers``, ``first_k_dense_replace`` and + ``num_experts`` default to GLM5.3-Flash; a cut-down checkpoint (e.g. Miles' + 4-layer CI model) passes its own (:func:`checkpoint_geometry`).""" impl = check_transformer_impl(transformer_impl) + if int(num_experts) < 1: + raise ValueError(f"num_experts must be >= 1, got {num_experts}") return ["--spec", *SPEC, "--num-layers", str(int(num_layers)), *_MODEL_ARGV, + "--num-experts", str(int(num_experts)), "--moe-layer-freq", moe_layer_freq(num_layers, first_k_dense_replace), *(TE_ONLY_ARGV if impl == TRANSFORMER_ENGINE else LOCAL_ARGV), *RUN_ARGV] diff --git a/yeto/rl/adapters/miles/run_flags.py b/yeto/rl/adapters/miles/run_flags.py index a6a9c6dad..28cb07751 100644 --- a/yeto/rl/adapters/miles/run_flags.py +++ b/yeto/rl/adapters/miles/run_flags.py @@ -388,6 +388,13 @@ island_flag("--seed", type=int, required=True), island_flag("--wan-streams", type=int, default=4), island_flag("--miles-root", required=True), + island_flag( + "--miles-commit", + default=None, + metavar="SHA", + help="Full commit the --miles-root checkout must be at (default: the pinned " + "MILES_NEXT_COMMIT); a run whose image carries another Miles commit names it here", + ), island_flag( "--rl-expected-algorithm-sha256", default=None, diff --git a/yeto/rl/adapters/sampler_engines/sglang_load_report.py b/yeto/rl/adapters/sampler_engines/sglang_load_report.py new file mode 100644 index 000000000..b228e163f --- /dev/null +++ b/yeto/rl/adapters/sampler_engines/sglang_load_report.py @@ -0,0 +1,141 @@ +"""Expert load report and ``weight_packed`` rename inside an SGLang server process +(S25 N2.3 post-load check 1; A-INF section 2.3 option A, item 1). + +SGLang's GLM5.3 ``load_weights`` maps ``experts..gate_proj.weight`` to the +fused ``experts.w13_weight`` parameter. A compressed-tensors MXFP4 release +stores that matrix as ``...gate_proj.weight_packed``; the mapped name +``experts.w13_weight_packed`` is not a parameter, and SGLang skips the tensor +without an error. :func:`install` wraps the model class's ``load_weights``: + +* with ``YETO_SGLANG_RENAME_WEIGHT_PACKED=1`` it renames routed-expert + ``*.weight_packed`` tensors to ``*.weight`` before SGLang maps them (the + MXFP4 MoE method stores the packed bytes in ``w13_weight`` / ``w2_weight``); +* with ``YETO_SGLANG_LOAD_REPORT_DIR`` set it counts the calls of the routed + experts' data-weight loaders (one call = one gate, up or down matrix of one + expert on this rank) and writes ``rank-.json`` into that directory + after each load. A skipped tensor never reaches a loader, so the count is + what the engine holds, not what the checkpoint names. + +The image applies :data:`MODEL_FILE_HOOK` to SGLang's model file +(:func:`patch_model_file`); the server process imports this module from the +yeto tree (``PYTHONPATH``). :func:`read_reports` is the adapter's side +(``SGLangServerEngine.loaded_expert_matrices``). + +This module must not import SGLang at import time (CPU tests use a fake model). +""" + +from __future__ import annotations + +import json +import os +import re +import time +from collections.abc import Iterable, Iterator +from pathlib import Path +from typing import Any + +REPORT_DIR_ENV = "YETO_SGLANG_LOAD_REPORT_DIR" +RENAME_ENV = "YETO_SGLANG_RENAME_WEIGHT_PACKED" +ROUTED_PACKED = re.compile(r"(\.mlp\.experts\.\d+\.(?:gate|up|down)_proj)\.weight_packed$") +# FusedMoE data parameters (not the scales); the routed experts' loaders. +EXPERT_DATA_PARAMS = re.compile(r"\.experts\.(?:w13_weight|w2_weight)$") +MARKER = "# yeto: sglang_load_report hook" +MODEL_FILE_HOOK = f""" + +{MARKER} +import os as _yeto_os +if _yeto_os.environ.get("{REPORT_DIR_ENV}") or _yeto_os.environ.get("{RENAME_ENV}") == "1": + from yeto.rl.adapters.sampler_engines.sglang_load_report import install as _yeto_install + for _yeto_cls in EntryClass: + _yeto_install(_yeto_cls) +""" + + +def rename_packed(weights: Iterable[tuple[str, Any]]) -> Iterator[tuple[str, Any]]: + """Routed-expert ``*.weight_packed`` -> ``*.weight``; every other name unchanged.""" + for name, tensor in weights: + yield ROUTED_PACKED.sub(r"\1.weight", name), tensor + + +def _tp_rank() -> int: + try: + from sglang.srt.distributed import get_tensor_model_parallel_rank # noqa: PLC0415 + return int(get_tensor_model_parallel_rank()) + except Exception: # noqa: BLE001 - outside a parallel group (tests) + return 0 + + +def install(model_cls: type) -> type: + """Wrap ``model_cls.load_weights`` once (module doc).""" + original = model_cls.load_weights + if getattr(original, "_yeto_load_report", False): + return model_cls + + def load_weights(self: Any, weights: Iterable[tuple[str, Any]], *args: Any, **kwargs: Any) -> Any: + if os.environ.get(RENAME_ENV) == "1": + weights = rename_packed(weights) + report_dir = os.environ.get(REPORT_DIR_ENV) + if not report_dir: + return original(self, weights, *args, **kwargs) + calls = {"n": 0} + wrapped = [] + for name, param in self.named_parameters(): + loader = getattr(param, "weight_loader", None) + if EXPERT_DATA_PARAMS.search(name) and loader is not None: + def counting(*a: Any, _loader: Any = loader, **kw: Any) -> Any: + calls["n"] += 1 + return _loader(*a, **kw) + param.weight_loader = counting + wrapped.append((param, loader)) + try: + result = original(self, weights, *args, **kwargs) + finally: + for param, loader in wrapped: + param.weight_loader = loader + rank = _tp_rank() + path = Path(report_dir) / f"rank{rank}-{os.getpid()}.json" + path.parent.mkdir(parents=True, exist_ok=True) + tmp = path.with_suffix(".tmp") + tmp.write_text(json.dumps({"tp_rank": rank, "pid": os.getpid(), "expert_matrices": calls["n"], + "expert_params": len(wrapped), "time": time.time()})) + os.replace(tmp, path) + return result + + load_weights._yeto_load_report = True # type: ignore[attr-defined] + model_cls.load_weights = load_weights + return model_cls + + +def patch_model_file(path: str | os.PathLike[str]) -> bool: + """Append :data:`MODEL_FILE_HOOK` to an SGLang model file that defines + ``EntryClass`` (idempotent; False when already patched).""" + file = Path(path) + text = file.read_text(encoding="utf-8") + if MARKER in text: + return False + if "EntryClass" not in text: + raise ValueError(f"{file} defines no EntryClass") + file.write_text(text + MODEL_FILE_HOOK, encoding="utf-8") + return True + + +def clear_reports(report_dir: str | os.PathLike[str]) -> None: + directory = Path(report_dir) + directory.mkdir(parents=True, exist_ok=True) + for f in directory.glob("rank*.json"): + f.unlink() + + +def read_reports(report_dir: str | os.PathLike[str]) -> dict[str, Any] | None: + """The reports of the current server (None: no rank reported). ``expert_matrices`` + is the smallest rank count, so one rank that skipped experts fails the check.""" + reports = [] + for f in sorted(Path(report_dir).glob("rank*.json")): + try: + reports.append(json.loads(f.read_text(encoding="utf-8"))) + except (OSError, ValueError): + continue + if not reports: + return None + counts = [int(r["expert_matrices"]) for r in reports] + return {"expert_matrices": min(counts), "per_rank": counts, "ranks": len(reports)} diff --git a/yeto/rl/adapters/sampler_engines/sglang_server.py b/yeto/rl/adapters/sampler_engines/sglang_server.py index 1be20274e..7d60479f9 100644 --- a/yeto/rl/adapters/sampler_engines/sglang_server.py +++ b/yeto/rl/adapters/sampler_engines/sglang_server.py @@ -6,7 +6,10 @@ * item 1: SGLang's ``load_weights`` does not map ``weight_packed`` to ``weight``; expert weights stored as ``weight_packed`` are skipped without an error. A release whose index names ``*.weight_packed`` is refused unless the - image carries the rename patch (``renames_weight_packed=True``); + image carries the rename patch (``renames_weight_packed=True``). The yeto + patch (:mod:`.sglang_load_report`, applied to the image's model file) renames + when ``YETO_SGLANG_RENAME_WEIGHT_PACKED=1`` and, with ``load_report_dir``, + reports the expert matrices each rank loaded (post-load check 1); * item 5: an MXFP4 export with ``quantization_config.format == "dense"`` (llmcompressor 0.14) is read wrongly; ``dense`` is refused for every quantized scheme; @@ -45,6 +48,7 @@ SamplingParams, ) +from . import sglang_load_report from .server import ServerConfig, ServerEngine TOPOLOGY_TP8 = "tp8" @@ -59,6 +63,8 @@ class SGLangOptions: renames_weight_packed: bool = False # the image patches the weight_packed -> weight map warmup_batch_sizes: tuple[int, ...] = () served_model_name: str = "yeto-policy" + # the yeto load-report patch writes each rank's loaded expert count here (None: no report) + load_report_dir: str | None = None def _finish(meta: dict[str, Any]) -> str: @@ -80,7 +86,8 @@ class SGLangServerEngine(ServerEngine): def __init__(self, config: ServerConfig, options: SGLangOptions | None = None, **kw: Any) -> None: super().__init__(config, **kw) self.options = options or SGLangOptions() - if config.topology not in SGLANG_TOPOLOGIES: + # tp8 on a PPU node; a smaller TP engine (tp1, tp2, ...) names its size the same way + if config.topology not in SGLANG_TOPOLOGIES and config.topology != f"tp{int(self.options.tp)}": raise SamplerEngineError(f"sglang topology {config.topology!r} is not supported " f"(known: {SGLANG_TOPOLOGIES}; PD disaggregation is not used)") if config.topology == TOPOLOGY_DP8 and self.options.nextn: @@ -103,8 +110,27 @@ def env(self) -> dict[str, str]: env = super().env() if self.config.precision == "mxfp4": env["MX_PATCH"] = "1" # item 7 + if self.options.renames_weight_packed: + env[sglang_load_report.RENAME_ENV] = "1" + if self.options.load_report_dir: + env[sglang_load_report.REPORT_DIR_ENV] = str(self.options.load_report_dir) return env + def start(self, release_dir: Path | None, version: int): + if self.options.load_report_dir: # reports of an earlier server must not count for this one + sglang_load_report.clear_reports(self.options.load_report_dir) + return super().start(release_dir, version) + + def loaded_expert_matrices(self) -> int | None: + if not self.options.load_report_dir: + return None + report = sglang_load_report.read_reports(self.options.load_report_dir) + if report is None: + return None + if report["ranks"] < int(self.options.tp): + return 0 # a rank that did not report loaded nothing we can vouch for + return int(report["expert_matrices"]) + def preflight(self, model_dir: Path) -> None: super().preflight(model_dir) index = Path(model_dir) / INDEX_FILE diff --git a/yeto/rl/engine/sampler_checks.py b/yeto/rl/engine/sampler_checks.py index 36d702748..c0439886a 100644 --- a/yeto/rl/engine/sampler_checks.py +++ b/yeto/rl/engine/sampler_checks.py @@ -149,6 +149,9 @@ class PostLoadChecks: allow_unverified_count: bool = False fixed_prompts: Sequence[FixedPrompt] = GSM8K_FIXED_PROMPTS fixed_max_tokens: int = 2048 + # False only for a model that cannot answer (a cut-down rehearsal model): a missing + # answer is recorded, the other fixed-prompt problems still fail the check + require_fixed_answers: bool = True forbidden_markers: Sequence[str] = () max_unreadable_share: float = 0.05 probe_pairs: Sequence[tuple[str, str]] | None = None # None: fixed prompts + their answers @@ -193,15 +196,17 @@ def fixed_prompt_answers(self, engine: SamplerEngine) -> dict[str, Any]: problems = [] if not text.strip(): problems.append("empty") - if not _contains_answer(text, item.answer): + answered = _contains_answer(text, item.answer) + if not answered and self.require_fixed_answers: problems.append("answer missing") if _unreadable_share(text) > self.max_unreadable_share: problems.append("unreadable characters") problems += [f"marker {m!r}" for m in self.forbidden_markers if m in text] - answers.append({"answer": item.answer, "ok": not problems, "problems": problems, - "text": text[:160]}) + answers.append({"answer": item.answer, "answered": answered, "ok": not problems, + "problems": problems, "text": text[:160]}) return {"name": "fixed_prompts", "passed": all(a["ok"] for a in answers) - and len(groups) == len(self.fixed_prompts), "answers": answers} + and len(groups) == len(self.fixed_prompts), "answers": answers, + "answers_required": bool(self.require_fixed_answers)} def logprob_probe(self, engine: SamplerEngine) -> dict[str, Any]: means = [] diff --git a/yeto/rl/quant/release.py b/yeto/rl/quant/release.py index 0eae27629..82f686c55 100644 --- a/yeto/rl/quant/release.py +++ b/yeto/rl/quant/release.py @@ -132,12 +132,13 @@ def run(self, stop: threading.Event, poll_s: float = 10.0) -> None: stop.wait(poll_s) -def requantizer(template_dir: Path, smoothing: Any) -> Requantize: - """The converter of :mod:`.requant` bound to one template and AWQ smoothing.""" +def requantizer(template_dir: Path, smoothing: Any, frozen: tuple[Path, str] | None = None) -> Requantize: + """The converter of :mod:`.requant` bound to one template and AWQ smoothing + (``frozen``: base checkpoint and name pattern of untrained tensors the release omits).""" from .requant import requantize_checkpoint def run(source: Path, out: Path) -> Any: - return requantize_checkpoint(Path(source), Path(template_dir), Path(out), smoothing) + return requantize_checkpoint(Path(source), Path(template_dir), Path(out), smoothing, frozen=frozen) return run @@ -156,7 +157,12 @@ def main(argv: list[str] | None = None) -> int: p.add_argument("--keep", type=int, default=2) p.add_argument("--loop", action="store_true") p.add_argument("--poll-seconds", type=float, default=10.0) + p.add_argument("--frozen-base", help="BF16 checkpoint that supplies untrained tensors the release omits") + p.add_argument("--frozen-pattern", help="regular expression of those tensor names (with --frozen-base)") a = p.parse_args(argv) + if bool(a.frozen_base) != bool(a.frozen_pattern): + p.error("give --frozen-base and --frozen-pattern together") + frozen = (Path(a.frozen_base), a.frozen_pattern) if a.frozen_base else None if a.no_awq: smoothing, awq = AwqSmoothing.none(), AWQ_NONE elif a.recipe and a.awq_scales: @@ -172,7 +178,7 @@ def main(argv: list[str] | None = None) -> int: store = store_for(a.store) releaser = RequantReleaser(PolicyReleaseStore(store, keep=a.keep, release_dir=release_dir_for(BF16)), PolicyReleaseStore(store, keep=a.keep, release_dir=release_dir_for(scheme)), - requantizer(Path(a.template), smoothing), scheme=scheme, awq=awq, + requantizer(Path(a.template), smoothing, frozen), scheme=scheme, awq=awq, awq_source=smoothing.source if smoothing.enabled else None, expert_matrices=len(layout.quantized_modules)) if a.loop: diff --git a/yeto/rl/quant/requant.py b/yeto/rl/quant/requant.py index 44fced0e2..053f51737 100644 --- a/yeto/rl/quant/requant.py +++ b/yeto/rl/quant/requant.py @@ -203,6 +203,11 @@ def from_checkpoint(cls, ckpt_dir: Path) -> "CheckpointLayout": @classmethod def load(cls, template_dir: Path) -> "CheckpointLayout": + """The layout of a template directory (``layout.json``) or, when it has + none, of a reference checkpoint itself (index + shard headers; stage 1 + passes the round-0 MXFP4 export as the template).""" + if not (Path(template_dir) / LAYOUT_FILE).is_file() and (Path(template_dir) / INDEX_FILE).is_file(): + return cls.from_checkpoint(Path(template_dir)) raw = json.loads((Path(template_dir) / LAYOUT_FILE).read_text()) if raw.get("schema") != LAYOUT_SCHEMA: raise LayoutError(f"{template_dir}: unknown layout schema {raw.get('schema')!r}") @@ -383,9 +388,16 @@ def requantize_checkpoint( out_dir: Path, smoothing: AwqSmoothing, shards: Iterable[str] | None = None, + frozen: tuple[Path, str] | None = None, ) -> dict[str, Any]: """Write the target checkpoint (or the given output shards of it) into ``out_dir``. + ``frozen``: ``(base_dir, pattern)``. A layout tensor that the source lacks and + whose source name matches the regular expression ``pattern`` is read from the + BF16 checkpoint ``base_dir`` (weights the trainer does not train and does not + export, e.g. the vision tower of a text-only Miles run). Every other missing + tensor still fails; ``stats["frozen_from_base"]`` counts the filled ones. + ``template_dir`` holds ``layout.json`` and the metadata files (:meth:`CheckpointLayout.save_template`). Without ``shards`` the whole checkpoint is written, then the index and the metadata files, so a @@ -398,6 +410,15 @@ def requantize_checkpoint( source = _SourceReader(source_dir) expected = layout.source_names() missing, extra = expected - set(source.weight_map), set(source.weight_map) - expected + base, filled = None, set() + if frozen is not None and missing: + matcher = re.compile(frozen[1]) + filled = {n for n in missing if matcher.search(n)} + if filled: + base = _SourceReader(Path(frozen[0])) + if absent := filled - set(base.weight_map): + raise LayoutError(f"frozen base {frozen[0]} lacks {len(absent)} tensors, e.g. {sorted(absent)[:3]}") + missing = missing - filled if missing or extra: raise LayoutError( f"source tensors differ from the target layout: {len(missing)} missing (e.g. {sorted(missing)[:3]}), " @@ -410,7 +431,11 @@ def requantize_checkpoint( quantized = set(layout.quantized_weights()) out_dir = Path(out_dir) out_dir.mkdir(parents=True, exist_ok=True) - stats = {"shards": 0, "tensors": 0, "quantized": 0} + stats = {"shards": 0, "tensors": 0, "quantized": 0, "frozen_from_base": len(filled)} + + def read(name: str) -> torch.Tensor: + return base.get(name) if name in filled else source.get(name) + try: for shard in selected: produced: dict[str, torch.Tensor] = {} @@ -419,7 +444,7 @@ def requantize_checkpoint( for name in all_shards[shard]: if name not in produced: src = layout.source_of(name) - produced.update(transform_tensor(src, source.get(src), layout, smoothing)) + produced.update(transform_tensor(src, read(src), layout, smoothing)) tensors = {n: produced[n] for n in all_shards[shard]} write_output_shard(out_dir, shard, tensors, layout) stats["shards"] += 1 @@ -427,6 +452,8 @@ def requantize_checkpoint( stats["tensors"] += len(tensors) finally: source.close() + if base is not None: + base.close() if shards is None: finalize_checkpoint(template_dir, out_dir) return stats diff --git a/yeto/rl/stage1.py b/yeto/rl/stage1.py index 4825e9caf..358035b88 100644 --- a/yeto/rl/stage1.py +++ b/yeto/rl/stage1.py @@ -179,6 +179,27 @@ def record(self) -> dict[str, Any]: } +# Where the run store lives: Alibaba OSS (stage 1 on PPU), or a Modal Volume +# mounted at the store path in every container (the Modal rehearsal, S25 T-G1), +# reached through the yeto store hook (resume.STORE_IMPL_ENV). +STORE_BACKENDS = {"oss": None, "modal-volume": "yeto.cloud.modal_ckpt_store:store_from_env"} + + +def store_env(plan: Stage1Plan) -> dict[str, str]: + """Environment every node needs for the store backend (none for OSS here: the + nodes' ossutil/OSS configuration is set up on the node, runbook S4).""" + backend = str(plan.doc.get("store_backend") or "oss") + if STORE_BACKENDS.get(backend) is None: + return {} + from yeto.cloud.modal_ckpt_store import MODAL_VOLUME_ENV # noqa: PLC0415 + from yeto.rl.engine.resume import STORE_IMPL_ENV # noqa: PLC0415 + + volume = str(plan.doc.get("store_volume") or "") + if not volume: + raise Stage1Error("plan: store_backend modal-volume needs store_volume (the Volume name)") + return {STORE_IMPL_ENV: STORE_BACKENDS[backend], MODAL_VOLUME_ENV: volume} + + def run_store(plan: Stage1Plan) -> str: """The run directory on the store: ``/``.""" return f"{str(_get(plan.doc, 'store', 'plan')).rstrip('/')}/{plan.run_id}" @@ -236,12 +257,13 @@ def megatron_memory(model: Mapping[str, Any], miles: Mapping[str, Any], layout: if hf_config is not None: hf = shape_from_hf_config(json.loads(Path(hf_config).read_text())) shape = dataclasses.replace(hf, params=params, expert_params=experts, - num_experts=hf.num_experts or glm5_next.NUM_EXPERTS) + num_experts=hf.num_experts or int(model.get("num_experts", glm5_next.NUM_EXPERTS))) else: shape = ModelShape(params=params, largest_layer_params=params // int(model["num_layers"]), num_layers=int(model["num_layers"]), hidden_size=glm5_next.HIDDEN_SIZE, vocab_size=glm5_next.VOCAB_SIZE, expert_params=experts, - num_experts=glm5_next.NUM_EXPERTS, source="stage-1 plan + glm5_next recipe") + num_experts=int(model.get("num_experts", glm5_next.NUM_EXPERTS)), + source="stage-1 plan + glm5_next recipe") node = trainers[0] base = TrainPlan(framework="megatron", cards=sum(n.cards_per_node for n in trainers), cards_per_node=node.cards_per_node, hbm_gib=node.hbm_gib, @@ -326,9 +348,12 @@ def _sampler_engine(plan: Stage1Plan, node: Node): quant_scheme=str(q["scheme"]), topology=str(s["topology"]), max_model_len=int(s["max_model_len"]), host="0.0.0.0", port=int(s["port"]), extra_args=tuple(extra), expected_format=str(q["format"]), - in_place_reload=bool(s.get("in_place_reload"))) + in_place_reload=bool(s.get("in_place_reload")), python=str(s.get("python") or "python"), + extra_env={str(k): str(v) for k, v in (s.get("extra_env") or {}).items()}) options = SGLangOptions(tp=int(s["tp"]), nextn=bool(s.get("nextn")), - renames_weight_packed=bool(s.get("renames_weight_packed"))) + renames_weight_packed=bool(s.get("renames_weight_packed")), + load_report_dir=(f"{state_dir_of(plan)}/sampler-{node.name}/load-report" + if s.get("load_report") else None)) return SGLangServerEngine(config, options) @@ -415,7 +440,7 @@ def expand(plan: Stage1Plan, *, hf_config: str | Path | None = None) -> Expanded miles = _section(trainer_cfg, "miles") layout = megatron_layout(miles, cards, per_node) common_env = {DEVICE_FAMILY_ENV: trainers[0].device_family, "YETO_RUN_ID": plan.run_id, - "YETO_RL_STORE": store, "HF_HUB_OFFLINE": "1"} + "YETO_RL_STORE": store, "HF_HUB_OFFLINE": "1", **store_env(plan)} trainer_env = {**common_env, "TORCHINDUCTOR_COMPILE_THREADS": "1"} model_dir = str(model["bf16_source"]["local_dir"]) state_dir = state_dir_of(plan) @@ -435,6 +460,7 @@ def expand(plan: Stage1Plan, *, hf_config: str | Path | None = None) -> Expanded "--tensor-parallel", str(layout["tp"]), "--pipeline-parallel", str(layout["pp"]), "--expert-parallel", str(layout["ep"]), "--miles-root", str(miles["workdir"]), + *(("--miles-commit", str(miles["commit"])) if len(str(miles["commit"])) == 40 else ()), "--trust-remote-code", "--learner-id", "0", "--num-learners", "1", "--rl-single-island-no-sync", "--global-rounds", str(int(plan.doc["rounds"])), @@ -449,7 +475,10 @@ def expand(plan: Stage1Plan, *, hf_config: str | Path | None = None) -> Expanded "--event-tape", f"{state_dir}/events.jsonl", "--audit-dir", f"{state_dir}/audit", *_provenance_args(str(task["reward"]))] try: - model_args = glm5_next.model_argv(num_layers=int(model["num_layers"]), transformer_impl=impl) + model_args = glm5_next.model_argv( + num_layers=int(model["num_layers"]), + first_k_dense_replace=int(model.get("first_k_dense_replace", glm5_next.FIRST_K_DENSE_REPLACE)), + num_experts=int(model.get("num_experts", glm5_next.NUM_EXPERTS)), transformer_impl=impl) except ValueError as error: raise Stage1Error(f"plan: trainer.miles.transformer_impl: {error}") from None out.trainer = { @@ -490,7 +519,11 @@ def expand(plan: Stage1Plan, *, hf_config: str | Path | None = None) -> Expanded # The release directory follows the template's scheme (quant.schemes); the # plan's release_dir is checked against it, not passed to the re-quantizer. - argv = requant_argv(store, str(model["sampler_initial"]["local_dir"]), python="python3") + # quant.frozen_pattern: tensors the trainer does not export (read from the bf16 source) + frozen = quant.get("frozen_pattern") + argv = requant_argv(store, str(model["sampler_initial"]["local_dir"]), python="python3", + frozen_base=str(model["bf16_source"]["local_dir"]) if frozen else None, + frozen_pattern=str(frozen) if frozen else None) out.requant = {"node": trainers[0].name if trainers else None, "scheme": quant["scheme"], "format": quant["format"], "argv": argv, "loop_argv": [*argv, "--loop"]} @@ -537,7 +570,10 @@ def sampler_spec(plan: Stage1Plan, node: Node, engine: Any) -> dict[str, Any]: "groups_per_round": int(task["rollout_batch_size"]), "format": str(task.get("prompt_format", "chat")), "trust_remote_code": True}, "reward": str(task["reward"]), "reward_kind": "miles", - "checks": {"expected_expert_matrices": int(quant["expected_expert_matrices"])}, + # sampler.checks: further PostLoadChecks overrides (e.g. require_fixed_answers for a + # cut-down rehearsal model); the expected expert count always comes from quant + "checks": {**dict(plan.section("sampler").get("checks") or {}), + "expected_expert_matrices": int(quant["expected_expert_matrices"])}, "event_tape": f"{state_dir}/sampler-{node.name}/events.jsonl", } @@ -729,7 +765,8 @@ def check(plan: Stage1Plan, out: Expanded) -> list[str]: if f"tp{int(sampler_cfg['tp'])}" != sampler_cfg["topology"]: errors.append(f"sampler topology {sampler_cfg['topology']} does not match TP {sampler_cfg['tp']}") image = str(plan.section("images").get("sampler") or "") - if not image.endswith(f":v{sampler_cfg['engine_version']}"): + # a tag must name the engine version; a digest pins the image itself (the version is recorded in the plan) + if not image.endswith(f":v{sampler_cfg['engine_version']}") and "@sha256:" not in image: errors.append(f"sampler image {image!r} does not carry engine version v{sampler_cfg['engine_version']}") if int(sampler_cfg["max_model_len"]) < int(sampling["max_prompt_len"]) + int(sampling["max_tokens"]): errors.append("sampler max_model_len is shorter than max_prompt_len + max_tokens") @@ -749,8 +786,13 @@ def check(plan: Stage1Plan, out: Expanded) -> list[str]: # Store. store = out.store["run"] - if not store.startswith("oss://"): + backend = str(plan.doc.get("store_backend") or "oss") + if backend not in STORE_BACKENDS: + errors.append(f"store_backend {backend!r} is not one of {sorted(STORE_BACKENDS)}") + elif backend == "oss" and not store.startswith("oss://"): errors.append(f"store {store!r} is not an oss:// URI (stage 1: Alibaba OSS)") + elif backend == "modal-volume" and not store.startswith("/"): + errors.append(f"store {store!r}: a modal-volume store is the absolute mount path in the containers") if not store.endswith("/" + plan.run_id): errors.append("store run directory does not end with the run id") bucket = store.split("/")[2] if store.startswith("oss://") else None From 5df375d333e13fda25f4a4f5a58132b742ae18a0 Mon Sep 17 00:00:00 2001 From: michael <78113554+chowhsu@users.noreply.github.com> Date: Tue, 20 Oct 2026 02:45:03 +0000 Subject: [PATCH 05/89] s25_g1 driver: print and keep the prep summary Co-Authored-By: Claude Opus 5.5 --- tools/probes/s25_g1/modal_app.py | 11 +++++++++-- 1 file changed, 9 insertions(+), 2 deletions(-) diff --git a/tools/probes/s25_g1/modal_app.py b/tools/probes/s25_g1/modal_app.py index ce842c249..077e93ee9 100644 --- a/tools/probes/s25_g1/modal_app.py +++ b/tools/probes/s25_g1/modal_app.py @@ -29,7 +29,7 @@ import modal HERE = Path(__file__).resolve() -REPO = HERE.parents[3] +REPO = HERE.parents[3] if len(HERE.parents) > 3 else Path("/") # in the container the module sits at /root WORKDIR = "/root/sky_workdir" PLAN = "examples/stage1/glm5next-mini-modal.yaml" IMAGE = ("ghcr.io/michaellchung/yeto-miles-ports@sha256:" @@ -134,8 +134,15 @@ def _env(extra: dict) -> dict: def prep() -> str: out = subprocess.run(["python3", f"{WORKDIR}/tools/probes/s25_g1/make_mini.py", "/data/models"], capture_output=True, text=True, cwd=WORKDIR, env=_env({})) + text = out.stdout[-6000:] + "\nSTDERR:\n" + out.stderr[-4000:] + Path("/data/models/prep.log").write_text(text) models.commit() - return out.stdout[-6000:] + "\nSTDERR:\n" + out.stderr[-4000:] + return text + + +@app.local_entrypoint() +def prep_main() -> None: + print(prep.remote()) # --------------------------------------------------------------------- sampler smoke From f0bbfd911f5be0bda9ccaf30553aaced4706bfad Mon Sep 17 00:00:00 2001 From: michael <78113554+chowhsu@users.noreply.github.com> Date: Tue, 20 Oct 2026 02:49:22 +0000 Subject: [PATCH 06/89] Stage-1 plan: sampler.extra_args; mini plan disables SGLang shared-expert fusion under MXFP4 SGLang fuses the BF16 shared expert into the routed FusedMoE as one more expert; with the MXFP4 MoE method that slot is packed U8 and the BF16 shared weights do not fit (shape error at load, Modal smoke-2). Co-Authored-By: Claude Opus 5.5 --- examples/stage1/glm5next-mini-modal.yaml | 3 +++ tools/probes/s25_g1/modal_app.py | 6 +++++- yeto/rl/stage1.py | 1 + 3 files changed, 9 insertions(+), 1 deletion(-) diff --git a/examples/stage1/glm5next-mini-modal.yaml b/examples/stage1/glm5next-mini-modal.yaml index 16ac368b0..b4dd7a262 100644 --- a/examples/stage1/glm5next-mini-modal.yaml +++ b/examples/stage1/glm5next-mini-modal.yaml @@ -103,6 +103,9 @@ sampler: renames_weight_packed: true # the yeto hook in the image renames (sglang_load_report) load_report: true # the hook reports loaded expert matrices (post-load check 1) in_place_reload: false + # SGLang fuses the BF16 shared expert into the routed FusedMoE as expert 32; under the + # MXFP4 MoE method that slot is packed U8, so the BF16 shared weights do not fit (smoke-2) + extra_args: [--disable-shared-experts-fusion] extra_env: PYTHONPATH: /root/sky_workdir checks: diff --git a/tools/probes/s25_g1/modal_app.py b/tools/probes/s25_g1/modal_app.py index 077e93ee9..932c884c9 100644 --- a/tools/probes/s25_g1/modal_app.py +++ b/tools/probes/s25_g1/modal_app.py @@ -98,7 +98,8 @@ def _sh(cmd: str, log: Path | None = None, env: dict | None = None, check: bool return out -def _gpu_facts(log: Path, want: int) -> dict: +def _gpu_facts(log_dir: Path, want: int) -> dict: + log = log_dir / "gpu.txt" out = _sh("nvidia-smi --query-gpu=index,name,memory.total,driver_version --format=csv,noheader", log) names = [line.split(",")[1].strip() for line in out.strip().splitlines()] if len(names) != want or not all(GPU_NAME in n for n in names): @@ -170,7 +171,10 @@ def smoke_run(tag: str) -> None: sample = engine.generate(["Question: What is 2 + 3?\nAnswer:"], SamplingParams(max_tokens=32, temperature=1.0, n=2, seed=1).check()) + from yeto.rl.adapters.sampler_engines.sglang_load_report import read_reports + record = {**facts, "start_s": round(started, 1), "load_report": report.detail, "checks": results, + "hook_report": read_reports(run["spec"]["engine"]["options"]["load_report_dir"]), "sample": [[c.text for c in g] for g in sample]} engine.stop() except Exception as error: # noqa: BLE001 - the record is the evidence diff --git a/yeto/rl/stage1.py b/yeto/rl/stage1.py index 358035b88..edd7f83f4 100644 --- a/yeto/rl/stage1.py +++ b/yeto/rl/stage1.py @@ -341,6 +341,7 @@ def _sampler_engine(plan: Stage1Plan, node: Node): extra = ["--moe-runner-backend", str(s["moe_runner_backend"])] if s.get("hierarchical_cache"): extra.append("--enable-hierarchical-cache") + extra += [str(a) for a in s.get("extra_args") or ()] # further engine flags of the plan config = ServerConfig( base_dir=Path(plan.section("model")["sampler_initial"]["local_dir"]), card_model=node.card, device_family=node.device_family, From a875e239e2d761223a82c8fd4c7ebc85bce30903 Mon Sep 17 00:00:00 2001 From: michael <78113554+chowhsu@users.noreply.github.com> Date: Tue, 20 Oct 2026 02:53:36 +0000 Subject: [PATCH 07/89] T-G3 harness: Megatron-LM on PYTHONPATH, model_args_utils path, GPU name passed to the container Co-Authored-By: Claude Opus 5.5 --- tools/probes/g3_miles/harness.py | 1 + tools/probes/g3_miles/modal_g3.py | 8 ++++---- 2 files changed, 5 insertions(+), 4 deletions(-) diff --git a/tools/probes/g3_miles/harness.py b/tools/probes/g3_miles/harness.py index 59cf090fe..4b43483d6 100644 --- a/tools/probes/g3_miles/harness.py +++ b/tools/probes/g3_miles/harness.py @@ -39,6 +39,7 @@ HERE = Path(__file__).resolve().parent MILES = Path(os.environ.get("G3_MILES", "/g3/miles")) sys.path.insert(0, str(MILES / "scripts" / "models")) +sys.path.insert(0, str(MILES / "miles" / "utils" / "external_utils")) # model_args_utils SEQ = 1024 # tokens of the packed sequence (one sample) TIMING_SEQ = 4096 diff --git a/tools/probes/g3_miles/modal_g3.py b/tools/probes/g3_miles/modal_g3.py index fb7e4500b..537b2090b 100644 --- a/tools/probes/g3_miles/modal_g3.py +++ b/tools/probes/g3_miles/modal_g3.py @@ -26,7 +26,7 @@ image = ( modal.Image.from_registry(IMAGE_REF).entrypoint([]) - .env({"PYTHONPATH": "/g3/miles:/g3/yeto", "CUDA_DEVICE_MAX_CONNECTIONS": "1", "G3_MILES": "/g3/miles", + .env({"PYTHONPATH": "/g3/miles:/g3/yeto:/root/Megatron-LM", "CUDA_DEVICE_MAX_CONNECTIONS": "1", "G3_MILES": "/g3/miles", "PYTHONUNBUFFERED": "1"}) .add_local_dir(MILES, "/g3/miles", copy=False, ignore=[".git", "**/__pycache__", "tests", "docs", "examples"]) .add_local_dir(os.path.join(YETO, "yeto"), "/g3/yeto/yeto", copy=False, ignore=["**/__pycache__"]) @@ -77,7 +77,7 @@ @app.function(image=image, gpu=GPU, cpu=16.0, memory=65536, timeout=HARD_S, volumes={"/vol": vol}) -def chain(run: str, phases: str) -> int: +def chain(run: str, phases: str, expect_gpu: str) -> int: import subprocess import threading import time @@ -96,7 +96,7 @@ def committer(): # partial evidence survives a hard timeout threading.Thread(target=committer, daemon=True).start() t0 = time.time() rc = subprocess.call(["bash", "-c", CHAIN], env={**os.environ, "OUT": out, "PHASES": phases, - "EXPECT_GPU": EXPECT_GPU}) + "EXPECT_GPU": expect_gpu}) with open(f"{out}/chain.rc", "w") as fh: fh.write(f"{rc} {time.time() - t0:.0f}s\n") stop.set() @@ -107,5 +107,5 @@ def committer(): # partial evidence survives a hard timeout @app.local_entrypoint() def main(phases: str = "all"): assert RUN != "unset", "set G3_RUN" - call = chain.spawn(RUN, phases) + call = chain.spawn(RUN, phases, EXPECT_GPU) # module globals are re-read in the container print("spawned", RUN, call.object_id, flush=True) From e80d9feb16ec622c7b792637630412cf9fd6bc6b Mon Sep 17 00:00:00 2001 From: michael <78113554+chowhsu@users.noreply.github.com> Date: Tue, 20 Oct 2026 02:55:45 +0000 Subject: [PATCH 08/89] s25_g1 driver: a failed sampler cancels the running trainer, and the reverse Co-Authored-By: Claude Opus 5.5 --- tools/probes/s25_g1/modal_app.py | 52 ++++++++++++++++++++++++++------ 1 file changed, 43 insertions(+), 9 deletions(-) diff --git a/tools/probes/s25_g1/modal_app.py b/tools/probes/s25_g1/modal_app.py index 932c884c9..66bac395c 100644 --- a/tools/probes/s25_g1/modal_app.py +++ b/tools/probes/s25_g1/modal_app.py @@ -215,7 +215,7 @@ def sampler_island(tag: str) -> int: return rc -@app.function(image=image, volumes=VOLUMES, gpu=f"{GPU_NAME}!:2", cpu=12, memory=163840, timeout=5400) +@app.function(image=image, volumes=VOLUMES, gpu=f"{GPU_NAME}!:2", cpu=12, memory=163840, timeout=3600) def trainer_island(tag: str, stop_after_rounds: int | None = None) -> int: log = _log_dir(tag) _gpu_facts(log, 2) @@ -253,18 +253,52 @@ def trainer_island(tag: str, stop_after_rounds: int | None = None) -> int: return rc -@app.function(image=image, cpu=1, memory=1024, timeout=3 * 3600) +@app.function(image=image, volumes=VOLUMES, cpu=1, memory=1024, timeout=3 * 3600) def orchestrate(tag: str, cut_after: int = 2) -> dict: + """Sampler for the whole run; trainer A stops after ``cut_after`` rounds (cut), + trainer B resumes in a new container. A trainer that fails cancels the sampler; + a sampler that fails cancels the running trainer (no card waits for nothing).""" sampler = sampler_island.spawn(f"{tag}/sampler") - first = trainer_island.remote(f"{tag}/trainer-a", cut_after) - second = trainer_island.remote(f"{tag}/trainer-b", None) if first == 0 else None - if first != 0 or second != 0: # no trainer will release again: do not leave the sampler card waiting - sampler.cancel() - sampler_rc = "cancelled" + sampler_rc: object = None + + def sampler_done() -> bool: + nonlocal sampler_rc + try: + sampler_rc = sampler.get(timeout=0) + return True + except TimeoutError: + return False + except Exception as error: # noqa: BLE001 - the sampler call failed + sampler_rc = f"error: {error}" + return True + + def run_trainer(sub: str, stop: int | None) -> object: + call = trainer_island.spawn(f"{tag}/{sub}", stop) + while True: + try: + return call.get(timeout=30) + except TimeoutError: + pass + except Exception as error: # noqa: BLE001 + return f"error: {error}" + if sampler_done() and sampler_rc != 0: + call.cancel() + return "cancelled: sampler ended with " + str(sampler_rc) + + first = run_trainer("trainer-a", cut_after) + second = run_trainer("trainer-b", None) if first == 0 else None + if first != 0 or second != 0: + if not sampler_done(): + sampler.cancel() + sampler_rc = "cancelled" else: - sampler_rc = sampler.get() + while not sampler_done(): + time.sleep(30) result = {"trainer_a": first, "trainer_b": second, "sampler": sampler_rc} - print(json.dumps(result), flush=True) + print(json.dumps(result, default=str), flush=True) + _log_dir(tag) + Path(f"/data/store/logs/{tag}/orchestrate.json").write_text(json.dumps(result, default=str)) + _commit() return result From e5e89197f5fdbf182028273e3b10a73618758674 Mon Sep 17 00:00:00 2001 From: michael <78113554+chowhsu@users.noreply.github.com> Date: Tue, 20 Oct 2026 03:01:19 +0000 Subject: [PATCH 09/89] Report model calls that hit their token cap on the verl Harbor path TokenTurn carries the engine's finish reason (the records recorder fills it) and HarborAgentLoop reports length_capped_turns in extra_fields. A Codex trajectory whose reply was cut at max_tokens otherwise looked like a normal episode ("truncated" only covers the trainer's response width). Record the first model-driven Codex runs in the Harbor README. Co-Authored-By: Claude Opus 5.5 --- tests/test_agentic_harbor.py | 4 ++- tests/test_verl_records_recorder.py | 1 + yeto/rl/adapters/verl/harbor_agent_loop.py | 3 +++ yeto/rl/agentic/types.py | 4 ++- yeto/rl/harness/gateway/recorder.py | 3 ++- yeto/rl/harness/harbor/README.md | 29 ++++++++++++++++++++-- 6 files changed, 39 insertions(+), 5 deletions(-) diff --git a/tests/test_agentic_harbor.py b/tests/test_agentic_harbor.py index adcba5c46..e53420c32 100644 --- a/tests/test_agentic_harbor.py +++ b/tests/test_agentic_harbor.py @@ -10,6 +10,7 @@ import importlib import sys import types +from dataclasses import replace from datetime import datetime, timedelta from pathlib import Path from types import SimpleNamespace @@ -424,7 +425,7 @@ def test_verl_loop_token_mask_alignment_with_fake_trial(verl_loop): loop_mod, _ = verl_loop from yeto.rl.harness.harbor import HarborHarness - turns = [_turn([5, 6, 7], [40, 41, 42], [-0.1, -0.2, -0.3]), + turns = [replace(_turn([5, 6, 7], [40, 41, 42], [-0.1, -0.2, -0.3]), finish_reason="length"), _turn([5, 6, 7, 40, 41, 42, 90, 91], [43, 2], [-0.4, -0.5])] recorder = _Recorder(turns) _FakeTrial.result, _FakeTrial.configs = _result({"reward": 1.0}), [] @@ -445,6 +446,7 @@ def test_verl_loop_token_mask_alignment_with_fake_trial(verl_loop): assert out.prompt_ids + out.response_ids == list(turns[-1].prompt_ids + turns[-1].completion_ids) assert out.reward_score == 1.0 and out.num_turns == 4 assert out.extra_fields[OUTCOME_KEY] == "ok" and out.extra_fields["truncated"] is False + assert out.extra_fields["length_capped_turns"] == 1 # the first call hit its token cap def test_verl_loop_aborts_on_misaligned_tokens_and_infra_errors(verl_loop): diff --git a/tests/test_verl_records_recorder.py b/tests/test_verl_records_recorder.py index 0d45d96ac..a8add5c5f 100644 --- a/tests/test_verl_records_recorder.py +++ b/tests/test_verl_records_recorder.py @@ -102,6 +102,7 @@ async def go(): assert sum(out.response_mask) > 0 and out.response_mask[-1] == 1 assert all(lp == 0.0 for lp, m in zip(out.response_logprobs, out.response_mask) if m == 0) assert out.extra_fields["provenance"]["engine"] == "fake" + assert out.extra_fields["length_capped_turns"] == 0 assert proxy.backend.sessions == {} and proxy.gateway.registries == {} # forgotten after close diff --git a/yeto/rl/adapters/verl/harbor_agent_loop.py b/yeto/rl/adapters/verl/harbor_agent_loop.py index 9093839da..21919bd6f 100644 --- a/yeto/rl/adapters/verl/harbor_agent_loop.py +++ b/yeto/rl/adapters/verl/harbor_agent_loop.py @@ -81,6 +81,7 @@ ) from yeto.rl.agentic.loop_local import LoopLocal from yeto.rl.agentic.types import TokenTurn +from yeto.rl.engine.sampler_engine import FINISH_LENGTH from yeto.rl.harness.harbor import HarborHarness, HarborSettings, build_trial_config from yeto.rl.harness.harbor.boat_client import check_in_flight_cap @@ -232,6 +233,8 @@ def _output(self, trajectory: AgentTrajectory, segments: TokenSegments | None, "outcome_reason": trajectory.reason, "policy_version": trajectory.policy_version, "truncated": segments.truncated, + # Model calls that hit their token cap (the agent then saw a cut reply). + "length_capped_turns": sum(t.finish_reason == FINISH_LENGTH for t in turns), "agent_metrics": metadata["agent_metrics"], "provenance": turn_provenance(turns), }, diff --git a/yeto/rl/agentic/types.py b/yeto/rl/agentic/types.py index 2bd677890..39a6adf93 100644 --- a/yeto/rl/agentic/types.py +++ b/yeto/rl/agentic/types.py @@ -122,7 +122,8 @@ class TokenTurn: ``policy_version`` and ``provenance`` (the sampler source record, ``Provenance.to_record()``) are what the recorder knows about the call; - None when it does not record them. + None when it does not record them. ``finish_reason``: the engine's reason + (``stop`` / ``length``; ``length`` = the call hit its token cap). """ prompt_ids: tuple[int, ...] @@ -130,6 +131,7 @@ class TokenTurn: completion_logprobs: tuple[float, ...] policy_version: int | None = None provenance: Mapping[str, str] | None = None + finish_reason: str | None = None def __post_init__(self) -> None: if len(self.completion_ids) != len(self.completion_logprobs): diff --git a/yeto/rl/harness/gateway/recorder.py b/yeto/rl/harness/gateway/recorder.py index 2aadeabcf..f28edad60 100644 --- a/yeto/rl/harness/gateway/recorder.py +++ b/yeto/rl/harness/gateway/recorder.py @@ -74,7 +74,8 @@ def turns_from_records(records: Mapping[str, Any]) -> list[TokenTurn]: turns.append(TokenTurn(prompt_ids=tuple(ids[:cursor]), completion_ids=completion, completion_logprobs=tuple(float(x) for x in turn["logprobs"]), policy_version=turn.get("policy_version"), - provenance=turn.get("provenance"))) + provenance=turn.get("provenance"), + finish_reason=turn.get("finish_reason"))) cursor, pending_stop = end, len(turn.get("stop_added") or ()) if cursor + pending_stop != len(ids): raise TokenAlignmentError(f"segment {sample['segment']} has {len(ids) - cursor - pending_stop} " diff --git a/yeto/rl/harness/harbor/README.md b/yeto/rl/harness/harbor/README.md index 8295646a2..8c7195e97 100644 --- a/yeto/rl/harness/harbor/README.md +++ b/yeto/rl/harness/harbor/README.md @@ -1,8 +1,9 @@ # Harbor harness for agentic RL (S25 N10.1-N10.4) Status: Harbor trials with the oracle/nop agents ran on boat sandboxes (see -"boat sandbox provider"). **No model-driven (codex) trajectory has run yet -(not verified).** +"boat sandbox provider"). Codex trajectories with a real model ran on boat +on 2026-10-20 (see "Codex + records proxy"). A training run on either +backend: not verified. ## Who does what @@ -168,3 +169,27 @@ Miles path, 2026-10-19: the Miles agent function rollout worker or session server) plus `reward_func`, `HARBOR_ENV_TYPE=boat`: fix-git and regex-log, oracle reward 1 on both, nop reward 0 on both (outcome ok); 30 billable s. A full Miles training run on boat: not verified. + +## Codex + records proxy (N10.7 smoke, 2026-10-20) + +Codex 0.145.0 refuses `wire_api = "chat"` ("no longer supported"): it only +speaks the Responses API, so the records proxy (Responses -> Chat, N10.5) is +its path to a chat-template engine. Codex first tries a WebSocket on +`/v1/responses` (5 attempts, the proxy answers 405), then uses HTTP POST. + +What the yeto side sets for Codex: `reasoning_summary=none` and +`web_search=disabled` (`trial.CODEX_KWARGS`); the proxy leaves Codex's +`namespace` (sub-agent) tools out of the prompt; the `qwen38` template merges +Codex's two leading system messages (instructions + developer message). + +Run: `scripts/modal_records_endpoint.py` (Qwen3.8-27B bf16, SGLang TP1 on one +H100, records proxy behind a Modal tunnel and a secret path prefix); +Harbor `codex` on boat small VMs; driver and raw data in +`infra-drafts/SESSION25/logs/T-G2/`. Miles path (agent function + reward +hook, proxy in place of the session server): fix-git and +log-summary-date-ranges, reward 1 on both. verl path (`HarborAgentLoop.run`, +verl module stubbed, HTTP recorder over the proxy): openssl-selfsigned-cert 1, +fix-git 1, regex-log 0 (twice), sqlite-db-truncate 0. All 8 outcome ok; every +recorded turn re-tokenizes to the same ids as the stream (render drift 0); +`boat list` empty afterwards. The regex-log runs ended on one reply that hit +the 8192-token cap (`length_capped_turns`). From 4c05531093a620c337fd68bf311eec48f97255f4 Mon Sep 17 00:00:00 2001 From: michael <78113554+chowhsu@users.noreply.github.com> Date: Tue, 20 Oct 2026 03:01:25 +0000 Subject: [PATCH 10/89] Add a CPU check to the Modal records endpoint and fix its import path The container imports yeto from the mounted workdir, which is not on sys.path before _image_paths() runs. The check action loads the template from the models Volume and imports the engine and proxy modules on a CPU container, so a broken image or path fails before a GPU starts. Co-Authored-By: Claude Opus 5.5 --- scripts/modal_records_endpoint.py | 42 ++++++++++++++++++++++++++++--- 1 file changed, 38 insertions(+), 4 deletions(-) diff --git a/scripts/modal_records_endpoint.py b/scripts/modal_records_endpoint.py index 08e56345a..0c75ca0ff 100644 --- a/scripts/modal_records_endpoint.py +++ b/scripts/modal_records_endpoint.py @@ -36,6 +36,7 @@ RESULTS_VOLUME = "yeto-s25-records-endpoint" RESULTS_MOUNT = "/mnt/results" CONTROL_DICT = "yeto-records-endpoint-control" +CONTAINER_WORKDIR = "/root/sky_workdir" # yeto source in the container (modal_eval_island.WORKDIR_MOUNT) PROXY_PORT = 8000 SGLANG_PORT = 30000 @@ -65,6 +66,7 @@ def serve(config: dict[str, Any]) -> dict[str, Any]: import modal + sys.path.insert(0, CONTAINER_WORKDIR) from yeto.cloud.modal_eval_island import _image_paths _image_paths() @@ -159,13 +161,38 @@ def dump() -> None: return {"run_name": config["run_name"], "trajectories": len(trajectories)} -def build_app(workdir: str, gpu: str, timeout_s: int) -> tuple[Any, Any]: +def check(config: dict[str, Any]) -> dict[str, Any]: + """CPU container: image imports, checkpoint files and the template, no engine.""" + sys.path.insert(0, CONTAINER_WORKDIR) + from yeto.cloud.modal_eval_island import _image_paths + + _image_paths() + import aiohttp + import sglang + + from yeto.rl.adapters.sampler_engines.sglang_server import SGLangServerEngine # noqa: F401 + from yeto.rl.harness.chat_templates import ChatTemplate, check_alignment + from yeto.rl.harness.gateway.proxy import RecordsProxy # noqa: F401 + + model_dir = Path(MODELS_MOUNT) / config["model_subdir"] + template = ChatTemplate.load(config["template"], tokenizer_path=str(model_dir)) + messages = [{"role": "system", "content": "a"}, {"role": "system", "content": "b"}, + {"role": "user", "content": "hi"}] + return {"sglang": sglang.__version__, "aiohttp": aiohttp.__version__, + "shards": len(list(model_dir.glob("*.safetensors"))), + "stop_token_ids": template.stop_token_ids, + "aligned": check_alignment(template, messages).aligned} + + +def build_app(workdir: str, gpu: str, timeout_s: int) -> tuple[Any, Any, Any]: import modal sys.path.insert(0, workdir) from yeto.cloud.modal_eval_island import WORKDIR_MOUNT from yeto.rl.adapters.miles.pins import MILES_NEXT_IMAGE + assert WORKDIR_MOUNT == CONTAINER_WORKDIR + image = (modal.Image.from_registry(MILES_NEXT_IMAGE.removeprefix("docker:")) .env({"HOME": "/root", "PYTHONUNBUFFERED": "1"}) .add_local_dir(workdir, WORKDIR_MOUNT, copy=False, @@ -177,12 +204,14 @@ def build_app(workdir: str, gpu: str, timeout_s: int) -> tuple[Any, Any]: fn = app.function(image=image, gpu=gpu, cpu=8.0, memory=65536, timeout=timeout_s, retries=0, volumes={MODELS_MOUNT: models.read_only(), RESULTS_MOUNT: results}, name="serve", serialized=True)(serve) - return app, fn + checker = app.function(image=image, cpu=2.0, memory=8192, timeout=600, retries=0, + volumes={MODELS_MOUNT: models.read_only()}, name="check")(check) + return app, fn, checker def main() -> None: parser = argparse.ArgumentParser(description=__doc__.splitlines()[0]) - parser.add_argument("action", choices=("launch", "stop", "url")) + parser.add_argument("action", choices=("check", "launch", "stop", "url")) parser.add_argument("--run-name", required=True) parser.add_argument("--secret-file", type=Path, help="launch: file holding the path secret") parser.add_argument("--gpu", default="H100!") @@ -198,11 +227,16 @@ def main() -> None: if args.action == "url": print(control.get(f"{args.run_name}/url", "")) return + if args.action == "check": + app, _, checker = build_app(args.workdir, args.gpu, timeout_s=600) + with modal.enable_output(), app.run(): + print(json.dumps(checker.remote(dict(DEFAULTS)))) + return secret = args.secret_file.read_text().strip() if len(secret) < 24: raise SystemExit("the path secret must have at least 24 characters") config = {**DEFAULTS, "run_name": args.run_name, "secret": secret, "max_serve_s": args.max_serve_s} - app, fn = build_app(args.workdir, args.gpu, timeout_s=args.max_serve_s + 1800) + app, fn, _ = build_app(args.workdir, args.gpu, timeout_s=args.max_serve_s + 1800) with modal.enable_output(), app.run(detach=True): call = fn.spawn(config) print(json.dumps({"function_call_id": call.object_id, "run_name": args.run_name})) From 0ff31dc1b19d90f5e8e4c6b029183ffc3af8901c Mon Sep 17 00:00:00 2001 From: michael <78113554+chowhsu@users.noreply.github.com> Date: Tue, 20 Oct 2026 03:03:32 +0000 Subject: [PATCH 11/89] Stage 1: no --rollout-num-gpus on the trainer-only island; check the argv with parse_args The ports placement check refuses --rollout-num-gpus (0 included) without --rl-placement fixed-partition, and fixed-partition needs >= 1 rollout GPU, so the trainer-only island argv the plan generated never started (Modal run-1). The plan's argv check now runs the island's parse_args with its combination rules, not only the bare parser. Co-Authored-By: Claude Opus 5.5 --- tests/test_s25_g1_stage1mini.py | 9 +++++++++ yeto/rl/stage1.py | 31 +++++++++++++++++++------------ 2 files changed, 28 insertions(+), 12 deletions(-) diff --git a/tests/test_s25_g1_stage1mini.py b/tests/test_s25_g1_stage1mini.py index d3742ea5e..87329c6c5 100644 --- a/tests/test_s25_g1_stage1mini.py +++ b/tests/test_s25_g1_stage1mini.py @@ -153,3 +153,12 @@ def test_mini_modal_plan_has_no_errors_and_uses_the_volume_store(): ppu = expand(load_plan(REPO / "examples/stage1/glm53-flash-ppu.yaml")) assert ppu.errors == [] and "--miles-commit" not in ppu.trainer["miles"]["island_argv"] assert "--frozen-base" not in ppu.requant["argv"] + + +def test_island_argv_check_runs_the_island_combination_rules(): + from yeto.rl.stage1 import _check_island_argv, expand, load_plan + + argv = expand(load_plan(REPO / "examples/stage1/glm5next-mini-modal.yaml")).trainer["miles"]["island_argv"] + assert "--rollout-num-gpus" not in argv and _check_island_argv(argv) is None + problem = _check_island_argv([*argv, "--rollout-num-gpus", "0"]) + assert problem and "fixed-partition" in problem diff --git a/yeto/rl/stage1.py b/yeto/rl/stage1.py index edd7f83f4..584729e6d 100644 --- a/yeto/rl/stage1.py +++ b/yeto/rl/stage1.py @@ -404,20 +404,26 @@ class _ParseError(Exception): def _check_island_argv(argv: Sequence[str]) -> str | None: - """Parse the Miles island argv with the island's own parser.""" - from yeto.rl.adapters.miles.island_entry import build_parser # noqa: PLC0415 + """Parse the Miles island argv with the island's own ``parse_args``: the parser + plus its combination checks (e.g. the ports placement rules; Modal run-1 failed + on one that the bare parser accepted).""" + import argparse # noqa: PLC0415 + from unittest import mock # noqa: PLC0415 - parser = build_parser() + from yeto.rl.adapters.miles.island_entry import build_parser, parse_args # noqa: PLC0415 - def fail(message: str) -> None: + def fail(_parser: Any, message: str) -> None: raise _ParseError(message) - parser.error = fail # type: ignore[method-assign] - no usage dump, no exit - try: - _, unknown = parser.parse_known_args(list(argv[3:])) - except _ParseError as exc: - return f"Miles island argv does not parse: {exc}" - return f"Miles island argv: flags unknown to island_entry: {unknown}" if unknown else None + with mock.patch.object(argparse.ArgumentParser, "error", fail): # no usage dump, no exit + try: + _, unknown = build_parser().parse_known_args(list(argv[3:])) + if unknown: + return f"Miles island argv: flags unknown to island_entry: {unknown}" + parse_args(list(argv[3:])) + except _ParseError as exc: + return f"Miles island argv does not parse: {exc}" + return None def expand(plan: Stage1Plan, *, hf_config: str | Path | None = None) -> Expanded: @@ -468,8 +474,9 @@ def expand(plan: Stage1Plan, *, hf_config: str | Path | None = None) -> Expanded "--groups-per-round", str(int(task["rollout_batch_size"])), "--over-sampling-batch-size", str(int(task["rollout_batch_size"])), "--optimizer-steps", "1", + # trainer-only island: no --rollout-num-gpus (the ports check refuses any value, + # 0 included, without --rl-placement fixed-partition; Modal run-1) "--actor-num-nodes", str(len(trainers)), "--actor-num-gpus-per-node", str(per_node), - "--rollout-num-gpus", "0", "--rl-activation-recompute", "--rl-max-tokens-per-gpu", str(int(miles["max_tokens_per_gpu"])), *(("--rl-optimizer-offload-fraction", repr(offload)) if offload > 0 else ()), "--completed-groups-path", f"{state_dir}/island-checkpoint.pt", @@ -901,7 +908,7 @@ def _full_parameter_layout_gap(plan: Stage1Plan, out: Expanded) -> str | None: problem = full_parameter_layout_problem( data_parallel=out.trainer["miles"]["layout"]["dp"], actor_num_nodes=int(argv[argv.index("--actor-num-nodes") + 1]), - rollout_gpus=int(argv[argv.index("--rollout-num-gpus") + 1]), + rollout_gpus=(int(argv[argv.index("--rollout-num-gpus") + 1]) if "--rollout-num-gpus" in argv else None), rollout_gpus_per_engine=1, # the island's default; it gives no engine sample_role=argv[argv.index("--rl-sample-role") + 1]) if problem is None: From aca8dab6f9232eabdcdf4aaa6d0436f0e63fb02d Mon Sep 17 00:00:00 2001 From: michael <78113554+chowhsu@users.noreply.github.com> Date: Tue, 20 Oct 2026 03:09:14 +0000 Subject: [PATCH 12/89] Stage 1: task.data_local for offline nodes; s25_g1 driver exports the Modal runner interpreter With HF_HUB_OFFLINE=1 the island cannot load a Hub dataset id (Modal run-2); a plan may name the dataset staged on the node instead (no revision flag). The Modal Volume store commits and reloads through the Modal runner's interpreter (YETO_MODAL_PYTHON/YETO_MODAL_SYSPATH, as yeto.modal_runner sets); without them every reload failed and was only logged. Co-Authored-By: Claude Opus 5.5 --- examples/stage1/glm5next-mini-modal.yaml | 1 + tools/probes/s25_g1/modal_app.py | 6 +++++- yeto/rl/stage1.py | 5 ++++- 3 files changed, 10 insertions(+), 2 deletions(-) diff --git a/examples/stage1/glm5next-mini-modal.yaml b/examples/stage1/glm5next-mini-modal.yaml index b4dd7a262..bef79b8c3 100644 --- a/examples/stage1/glm5next-mini-modal.yaml +++ b/examples/stage1/glm5next-mini-modal.yaml @@ -129,6 +129,7 @@ task: rollout_batch_size: 4 n_samples_per_prompt: 4 prompts_jsonl: /data/models/datasets/gsm8k-0cbd9f3/train.jsonl + data_local: /data/models/datasets/gsm8k-0cbd9f3/train.jsonl # the same rows for the trainer (offline nodes) input_key: messages label_key: label prompt_format: chat diff --git a/tools/probes/s25_g1/modal_app.py b/tools/probes/s25_g1/modal_app.py index 66bac395c..cfba2d623 100644 --- a/tools/probes/s25_g1/modal_app.py +++ b/tools/probes/s25_g1/modal_app.py @@ -126,7 +126,11 @@ def _nvml_sampler(path: Path) -> subprocess.Popen: def _env(extra: dict) -> dict: path = ":".join(p for p in (WORKDIR, os.environ.get("PYTHONPATH", "")) if p) - return {**os.environ, "PYTHONPATH": path, **{k: str(v) for k, v in extra.items()}} + # The island's Python cannot import the Modal client; the Volume store commits and + # reloads through this interpreter (yeto.cloud.modal_ckpt_store, as yeto.modal_runner does). + runner = {"YETO_MODAL_PYTHON": sys.executable, + "YETO_MODAL_SYSPATH": os.pathsep.join(p for p in sys.path if p)} + return {**os.environ, "PYTHONPATH": path, **runner, **{k: str(v) for k, v in extra.items()}} # --------------------------------------------------------------------- prep (CPU) diff --git a/yeto/rl/stage1.py b/yeto/rl/stage1.py index 584729e6d..e2e4c3487 100644 --- a/yeto/rl/stage1.py +++ b/yeto/rl/stage1.py @@ -455,7 +455,10 @@ def expand(plan: Stage1Plan, *, hf_config: str | Path | None = None) -> Expanded offload = float(miles.get("optimizer_offload_fraction") or 0.0) island_argv = ["python3", "-m", "yeto.rl.adapters.miles.island_entry", "--model", model_dir, "--model-revision", str(model["bf16_source"]["revision"]), - "--data", str(task["data"]), "--data-revision", str(task["data_revision"]), + # task.data_local: the dataset staged on the node (the nodes run with HF_HUB_OFFLINE=1, + # where a Hub id does not load; Modal run-2); a local path takes no revision + *(("--data", str(task["data_local"])) if task.get("data_local") else + ("--data", str(task["data"]), "--data-revision", str(task["data_revision"]))), "--reward-function", str(task["reward"]), "--samples-per-group", str(int(task["n_samples_per_prompt"])), "--rollout-max-response-len", str(int(sampling["max_tokens"])), From 7314e2bccec2317da7d686a467a87a7bc07e74f3 Mon Sep 17 00:00:00 2001 From: michael <78113554+chowhsu@users.noreply.github.com> Date: Tue, 20 Oct 2026 03:13:14 +0000 Subject: [PATCH 13/89] T-G3 harness: symmetric failure handling in N8.2, logits captured before the in-place CE, TE-to-TE name map, host-copy pin report Co-Authored-By: Claude Opus 5.5 --- tools/probes/g3_miles/compare.py | 13 +++++ tools/probes/g3_miles/harness.py | 88 +++++++++++++++++++++++--------- 2 files changed, 78 insertions(+), 23 deletions(-) diff --git a/tools/probes/g3_miles/compare.py b/tools/probes/g3_miles/compare.py index 85a024805..96dbf180a 100644 --- a/tools/probes/g3_miles/compare.py +++ b/tools/probes/g3_miles/compare.py @@ -82,6 +82,18 @@ def n75(out: Path) -> dict: return v +def _beyond_noise(c5, c5b, rs) -> list: + """Parameters whose final norm after resume differs from cont5 by more than 5x the cont5/cont5b noise.""" + a, b, c = (x[0]["final_params"]["norms"] for x in (c5, c5b, rs)) + out = [] + for n in a: + noise = abs(b[n] - a[n]) / max(a[n], 1e-30) + res = abs(c[n] - a[n]) / max(a[n], 1e-30) + if res > 5 * noise + 1e-6: + out.append({"name": n, "resume_rel": res, "noise_rel": noise, "norm": a[n]}) + return sorted(out, key=lambda x: -x["resume_rel"]) + + def n71d(out: Path) -> dict: load = lambda m: json.loads((out / f"n71d-{m}.json").read_text()) if (out / f"n71d-{m}.json").exists() else None # noqa: E731 c5, c5b, r3, rs = load("cont5"), load("cont5b"), load("run3"), load("resume2") @@ -113,6 +125,7 @@ def n71d(out: Path) -> dict: "loss_abs": {s["step"]: abs(s["loss_mean_dp"] - steps[s["step"]]["loss_mean_dp"]) for s in rs[0]["steps"]}, "grad_norm_abs": {s["step"]: abs(s["grad_norm"] - steps[s["step"]]["grad_norm"]) for s in rs[0]["steps"]}, "lr": {s["step"]: (s["lr"], steps[s["step"]]["lr"]) for s in rs[0]["steps"]}, + "beyond_5x_noise": _beyond_noise(c5, c5b, rs) if c5b else None, "param_norm_max_rel": max( abs(a["final_params"]["norms"][n] - b["final_params"]["norms"][n]) / max(b["final_params"]["norms"][n], 1e-30) for a, b in zip(rs, c5) for n in b["final_params"]["norms"])} diff --git a/tools/probes/g3_miles/harness.py b/tools/probes/g3_miles/harness.py index 4b43483d6..81012782c 100644 --- a/tools/probes/g3_miles/harness.py +++ b/tools/probes/g3_miles/harness.py @@ -162,10 +162,12 @@ def run_fwd_bwd(model, tokens: torch.Tensor, *, capture: list | None = None, for mask[:, -1] = 0 def loss_func(output): + if capture is not None: + logits = output.detach().float().clone() # vocab_parallel_cross_entropy edits its input in place token_loss = tensor_parallel.vocab_parallel_cross_entropy(output.float(), labels) loss = (token_loss * mask).sum() / mask.sum() if capture is not None: - capture.append((output.detach(), token_loss.detach())) + capture.append((logits, token_loss.detach())) return loss, {"loss": loss.detach()} def step(data_iter, chunk): @@ -203,7 +205,7 @@ def named(model) -> dict[str, torch.Tensor]: def local_to_te(name: str) -> str: from miles.backends.megatron_utils.named_weights import local_expert_name_to_grouped - name = local_expert_name_to_grouped(name) + name = te_canonical(local_expert_name_to_grouped(name), set()) # dense MLP: TE fuses the pre-MLP RMSNorm into linear_fc1 (layer_norm_weight) name = re.sub(r"(decoder\.layers\.\d+)\.pre_mlp_layernorm\.weight$", r"\1.PRE_MLP_NORM", name) return name @@ -339,6 +341,10 @@ def post(mod, inp, outp): "peak_mem_gib": torch.cuda.max_memory_allocated() / 2**30} zero(model, optimizer) + pre = [None] * dist.get_world_size() + dist.all_gather_object(pre, info) + if dist.get_rank() == 0: + (out / f"info-{ns.tag}.json").write_text(json.dumps(pre, indent=1, default=str)) # N8.2 on the same (unchanged) weights try: info["n82"] = n82(args, model, out / f"n82-{ns.tag}") @@ -438,14 +444,36 @@ def fake_original(source: dict[str, torch.Tensor], scheme: str, out: Path, smoot return out +def prepare_reference(source: dict[str, torch.Tensor], scheme: str, d: Path, smooth: str | None) -> dict: + """Rank 0: template, BF16 source, AWQ factors and the single-process reference output under ``d``.""" + from safetensors.torch import save_file + + from yeto.rl.quant.awq import AwqSmoothing + from yeto.rl.quant.requant import CheckpointLayout, requantize_checkpoint + + orig = fake_original(source, scheme, d / "original", smooth) + CheckpointLayout.from_checkpoint(orig).save_template(orig, d / "template") + (d / "bf16").mkdir(parents=True, exist_ok=True) + save_file({n: t.contiguous() for n, t in source.items()}, str(d / "bf16" / "model.safetensors"), + metadata={"format": "pt"}) + if scheme == "w4a8" and smooth: + (d / "awq-scales").mkdir(exist_ok=True) + factors = (0.8 + 0.4 * torch.rand(source[smooth + ".weight"].numel(), + generator=torch.Generator().manual_seed(7))).double() + torch.save(factors, d / "awq-scales" / f"{smooth}.pt") + smoothing = (AwqSmoothing.from_artefact(d / "original" / "recipe.yaml", d / "awq-scales") + if scheme == "w4a8" and smooth else AwqSmoothing.none()) + t0 = time.perf_counter() + requantize_checkpoint(d / "bf16", d / "template", d / "reference", smoothing) + return {"single_process_s": time.perf_counter() - t0} + + def n82(args, model, out: Path) -> dict: from megatron.core import parallel_state as mps - from safetensors.torch import save_file from yeto.rl.adapters.miles.quant_shards import hf_slices, megatron_expert_slices from yeto.rl.quant.awq import AwqSmoothing - from yeto.rl.quant.requant import CheckpointLayout, requantize_checkpoint - from yeto.rl.quant.shards import SourceSlice, requantize_distributed, torch_barrier + from yeto.rl.quant.shards import SourceSlice, requantize_distributed rank, world = dist.get_rank(), dist.get_world_size() ep_rank, ep_size = mps.get_expert_model_parallel_rank(), mps.get_expert_model_parallel_world_size() @@ -469,26 +497,19 @@ def n82(args, model, out: Path) -> dict: d = out / scheme r: dict = {} try: + err = None if rank == 0: - orig = fake_original(source, scheme, d / "original", smooth) - CheckpointLayout.from_checkpoint(orig).save_template(orig, d / "template") - (d / "bf16").mkdir(parents=True, exist_ok=True) - save_file({n: t.contiguous() for n, t in source.items()}, str(d / "bf16" / "model.safetensors"), - metadata={"format": "pt"}) - if scheme == "w4a8" and smooth: - (d / "awq-scales").mkdir(exist_ok=True) - s = (0.8 + 0.4 * torch.rand(source[smooth + ".weight"].numel(), - generator=torch.Generator().manual_seed(7))).double() - torch.save(s, d / "awq-scales" / f"{smooth}.pt") - dist.barrier() + try: + r.update(prepare_reference(source, scheme, d, smooth)) + except Exception as exc: # noqa: BLE001 -- re-raised below after telling the peer + err = exc + if err is not None: + signal_failure() + raise err + checked_barrier() # raises PeerFailed on rank 1 when rank 0 failed above smoothing = (AwqSmoothing.from_artefact(d / "original" / "recipe.yaml", d / "awq-scales") if scheme == "w4a8" and smooth else AwqSmoothing.none()) r["awq"] = bool(scheme == "w4a8" and smooth) - if rank == 0: - t0 = time.perf_counter() - requantize_checkpoint(d / "bf16", d / "template", d / "reference", smoothing) - r["single_process_s"] = time.perf_counter() - t0 - dist.barrier() # distributed: experts from the yeto adapter on Megatron-local names; non-experts once (rank 0) mine = list(megatron_expert_slices(((n.removeprefix("module.").removeprefix("module."), p) for n, p in local.items()), @@ -507,11 +528,13 @@ def n82(args, model, out: Path) -> dict: try: t0 = time.perf_counter() stats = requantize_distributed(slices, d / "template", dst, smoothing, rank=rank, - world_size=world, barrier=torch_barrier) + world_size=world, barrier=checked_barrier) r[variant] = {"ok": True, "s": time.perf_counter() - t0, "part": stats.get("part")} + except PeerFailed as exc: + r[variant] = {"ok": False, "error": f"PeerFailed: {exc}"} except Exception as exc: # noqa: BLE001 r[variant] = {"ok": False, "error": f"{type(exc).__name__}: {str(exc)[:600]}"} - # a rank that failed early leaves the others in a barrier: resync before the next variant + signal_failure() # pairs with the peer's next checked_barrier dist.barrier() if rank == 0 and r[variant]["ok"]: r[variant]["byte_equal"] = compare_dirs(d / "reference", dst) @@ -522,6 +545,22 @@ def n82(args, model, out: Path) -> dict: return res +class PeerFailed(RuntimeError): + pass + + +def checked_barrier() -> None: + """Barrier that also carries "a rank failed": every rank either reaches it or calls signal_failure().""" + flag = torch.zeros(1, device="cuda") + dist.all_reduce(flag, op=dist.ReduceOp.MAX) + if flag.item() > 0: + raise PeerFailed("another rank failed in requantize_distributed") + + +def signal_failure() -> None: + dist.all_reduce(torch.ones(1, device="cuda"), op=dist.ReduceOp.MAX) + + def compare_dirs(a: Path, b: Path) -> dict: names_a = sorted(p.name for p in a.iterdir() if p.is_file()) names_b = sorted(p.name for p in b.iterdir() if p.is_file()) @@ -560,7 +599,10 @@ def _hdo_report(optimizer) -> dict: "is_hdo": isinstance(inner, HybridDeviceOptimizer)} if isinstance(inner, HybridDeviceOptimizer): fp32 = list(inner.param_to_fp32_param.values()) if getattr(inner, "param_to_fp32_param", None) else [] + copies = list(getattr(inner, "gpu_params_map_cpu_copy", {}).values()) item.update( + host_copies=len(copies), host_copies_pinned=sum(bool(c.is_pinned()) for c in copies), + host_copy_dtypes=sorted({str(c.dtype) for c in copies}), fp32_params=len(fp32), fp32_devices=sorted({str(p.device) for p in fp32}), fp32_pinned=sum(bool(p.is_pinned()) for p in fp32 if p.device.type == "cpu"), fp32_dtypes=sorted({str(p.dtype) for p in fp32}), From 6f4165684828e54c7762a3d9ec42d33d50134757 Mon Sep 17 00:00:00 2001 From: michael <78113554+chowhsu@users.noreply.github.com> Date: Tue, 20 Oct 2026 03:18:34 +0000 Subject: [PATCH 14/89] Miles island: glm5_next geometry from config.json; stage 1 trainer reads one round behind - AutoBridge cannot read glm5_next (transformers 5.12.1 has no such model type; Modal run-3). Miles trains it through its plugin spec, so the island takes a read-only provider view of config.json (models/glm5_next), as it does for Flash-Next. - A trainer that reads the sample buffer needs --rl-max-policy-age >= 1 and an algorithm spec that tolerates it; stage 1 passes age 1 and trainer.miles.algorithm_spec (default decoupled-ppo-c2: GRPO + token TIS). - CPU test of the island path from the plan argv to the Miles argv. Co-Authored-By: Claude Opus 5.5 --- tests/test_s25_g1_stage1mini.py | 36 +++++++ yeto/rl/adapters/miles/island_entry.py | 12 +++ yeto/rl/adapters/miles/models/glm5_next.py | 103 +++++++++++++++++++++ yeto/rl/stage1.py | 8 ++ 4 files changed, 159 insertions(+) diff --git a/tests/test_s25_g1_stage1mini.py b/tests/test_s25_g1_stage1mini.py index 87329c6c5..a479f7576 100644 --- a/tests/test_s25_g1_stage1mini.py +++ b/tests/test_s25_g1_stage1mini.py @@ -162,3 +162,39 @@ def test_island_argv_check_runs_the_island_combination_rules(): assert "--rollout-num-gpus" not in argv and _check_island_argv(argv) is None problem = _check_island_argv([*argv, "--rollout-num-gpus", "0"]) assert problem and "fixed-partition" in problem + + +_MINI_TEXT_CONFIG = { + "model_type": "glm5_next_text", "num_hidden_layers": 4, "first_k_dense_replace": 1, "hidden_size": 4096, + "intermediate_size": 12288, "num_attention_heads": 64, "num_key_value_heads": 64, "n_routed_experts": 32, + "n_shared_experts": 1, "moe_intermediate_size": 2048, "num_experts_per_tok": 8, "vocab_size": 154880, + "max_position_embeddings": 1048576, "rms_norm_eps": 1e-5, "q_lora_rank": 1536, "kv_lora_rank": 512, + "qk_head_dim": 256, "qk_rope_head_dim": 0, "v_head_dim": 256, "hidden_act": "silu", "attention_bias": False, +} + + +def test_mini_island_argv_translates_to_a_miles_argv_on_cpu(tmp_path): + """The island path up to the Miles argv (Modal runs 1-3 each failed one step of it).""" + from yeto.rl.adapters.miles.island_entry import build_ports_launch, parse_args + from yeto.rl.engine.run_config import resolve_rl_run_config + from yeto.rl.stage1 import expand, load_plan + + model = tmp_path / "model" + model.mkdir() + (model / "config.json").write_text(json.dumps({"model_type": "glm5_next", "text_config": _MINI_TEXT_CONFIG})) + prompts = tmp_path / "p.jsonl" + prompts.write_text('{"messages": [{"role": "user", "content": "hi"}], "label": "1"}\n') + argv = expand(load_plan(REPO / "examples/stage1/glm5next-mini-modal.yaml")).trainer["miles"]["island_argv"][3:] + argv = [str(model) if a == "/data/models/glm5next-mini-bf16-7fb7883" else a for a in argv] + argv = [str(REPO / a) if a.startswith("examples/rl_algorithms/") else a for a in argv] + args = parse_args(argv) + provider = glm5_next.glm5_next_provider(model) + config = resolve_rl_run_config(args, model_path=model, prompt_path=prompts, provider=provider, + target_modules=[], yeto_policy_sync=False, verify_ref_load=False) + assert config.geometry.num_layers == 4 and config.geometry.moe.num_experts == 32 + miles = list(build_ports_launch(args, config, []).argv) + for flag in ("--debug-train-only", "--use-distributed-optimizer", "--optimizer-cpu-offload", + "--use-precision-aware-optimizer", "--use-tis"): + assert flag in miles + assert miles[miles.index("--num-experts") + 1] == "32" and miles.count("--num-experts") == 1 + assert miles[miles.index("--moe-layer-freq") + 1] == "[0,1,1,1]" diff --git a/yeto/rl/adapters/miles/island_entry.py b/yeto/rl/adapters/miles/island_entry.py index 9d4cc0f0a..5511b9bd6 100644 --- a/yeto/rl/adapters/miles/island_entry.py +++ b/yeto/rl/adapters/miles/island_entry.py @@ -1800,6 +1800,18 @@ def run_miles( ) attention_specs = specs = () canonical_targets = flash_next.target_modules() + elif getattr(args, "rl_model_recipe", "generic") == "glm5_next": + # AutoBridge cannot read glm5_next (transformers has no such model type); + # Miles trains it through its plugin spec, so the island reads the geometry + # from config.json (models/glm5_next.provider_view). Full-parameter only. + from yeto.rl.adapters.miles.models import glm5_next as glm5_next_recipe + + if not full_parameter: + raise ValueError("the glm5_next recipe is full-parameter only") + model_bridge = None + provider = glm5_next_recipe.glm5_next_provider(model_path) + attention_specs = specs = () + canonical_targets = [] else: from megatron.bridge import AutoBridge diff --git a/yeto/rl/adapters/miles/models/glm5_next.py b/yeto/rl/adapters/miles/models/glm5_next.py index 449747b60..980dc719a 100644 --- a/yeto/rl/adapters/miles/models/glm5_next.py +++ b/yeto/rl/adapters/miles/models/glm5_next.py @@ -21,7 +21,10 @@ from __future__ import annotations +import dataclasses +import json from collections.abc import Iterator, Sequence +from pathlib import Path MODEL_NAME = "glm5_next" TRANSFORMER_ENGINE = "transformer_engine" @@ -148,3 +151,103 @@ def apply_recipe(argv: Sequence[str], recipe: Sequence[str]) -> list[str]: override.add("--no-gradient-accumulation-fusion") kept = [token for group in _flag_groups(argv[1:]) if group[0] not in override for token in group] return [*argv[:1], *kept, *recipe] + + +# ---------------------------------------------------------------- provider view +# Megatron-Bridge's AutoBridge cannot read a glm5_next checkpoint: transformers +# 5.12.1 has no such model type (Modal run-3), and Miles trains it through its +# plugin spec, not through the Bridge. The island only needs the geometry that +# ``run_config.resolve_rl_run_config`` reads, so it takes this read-only view of +# the checkpoint's config.json (the same approach as the Flash-Next view). +MODEL_TYPES = frozenset({"glm5_next", "glm5_next_text"}) + + +@dataclasses.dataclass(frozen=True) +class Glm5NextModelProvider: + """Read-only stand-in for a Megatron-Bridge provider (the fields run_config reads).""" + + num_layers: int + hidden_size: int + num_attention_heads: int + num_query_groups: int + kv_channels: int + ffn_hidden_size: int + num_moe_experts: int + moe_ffn_hidden_size: int + moe_router_topk: int + moe_layer_freq: tuple[int, ...] + moe_shared_expert_intermediate_size: int + vocab_size: int + seq_length: int + max_position_embeddings: int + layernorm_epsilon: float + q_lora_rank: int + kv_lora_rank: int + qk_head_dim: int + qk_pos_emb_head_dim: int + v_head_dim: int + rotary_base: int = 10000 # the recipe's --rotary-base (the checkpoint has no RoPE dims: qk_rope_head_dim 0) + rotary_percent: float = 1.0 + share_embeddings_and_output_weights: bool = False + multi_latent_attention: bool = True + qk_layernorm: bool = True + gated_linear_unit: bool = True + add_bias_linear: bool = False + add_qkv_bias: bool = False + normalization: str = "RMSNorm" + position_embedding_type: str = "rope" + + def finalize(self) -> None: # provider API parity; nothing to derive + return None + + +def is_glm5_next(model_path: str | Path) -> bool: + path = Path(model_path) / "config.json" + if not path.is_file(): + return False + try: + return json.loads(path.read_text(encoding="utf-8")).get("model_type") in MODEL_TYPES + except (OSError, ValueError): + return False + + +def provider_view(config: dict) -> Glm5NextModelProvider: + """The provider view of a glm5_next ``config.json`` (module section above).""" + if config.get("model_type") not in MODEL_TYPES: + raise ValueError(f"not a glm5_next config: model_type={config.get('model_type')!r}") + text = config.get("text_config") or config + layers = int(text["num_hidden_layers"]) + dense = int(text.get("first_k_dense_replace", FIRST_K_DENSE_REPLACE)) + if text.get("hidden_act", "silu") != "silu": + raise ValueError(f"glm5_next expects SwiGLU (silu), got {text.get('hidden_act')!r}") + if bool(text.get("attention_bias", False)): + raise ValueError("glm5_next expects bias-free linears") + max_positions = int(text["max_position_embeddings"]) + return Glm5NextModelProvider( + num_layers=layers, + hidden_size=int(text["hidden_size"]), + num_attention_heads=int(text["num_attention_heads"]), + num_query_groups=int(text.get("num_key_value_heads", text["num_attention_heads"])), + kv_channels=int(text["v_head_dim"]), + ffn_hidden_size=int(text["intermediate_size"]), + num_moe_experts=int(text["n_routed_experts"]), + moe_ffn_hidden_size=int(text["moe_intermediate_size"]), + moe_router_topk=int(text["num_experts_per_tok"]), + moe_layer_freq=tuple([0] * dense + [1] * (layers - dense)), + moe_shared_expert_intermediate_size=int(text["moe_intermediate_size"]) * int(text.get("n_shared_experts", 1)), + vocab_size=int(text["vocab_size"]), + seq_length=max_positions, + max_position_embeddings=max_positions, + layernorm_epsilon=float(text["rms_norm_eps"]), + q_lora_rank=int(text["q_lora_rank"]), + kv_lora_rank=int(text["kv_lora_rank"]), + qk_head_dim=int(text["qk_head_dim"]), + qk_pos_emb_head_dim=int(text.get("qk_rope_head_dim", 0)), + v_head_dim=int(text["v_head_dim"]), + share_embeddings_and_output_weights=bool(config.get("tie_word_embeddings", False) + or text.get("tie_word_embeddings", False)), + ) + + +def glm5_next_provider(model_path: str | Path) -> Glm5NextModelProvider: + return provider_view(json.loads((Path(model_path) / "config.json").read_text(encoding="utf-8"))) diff --git a/yeto/rl/stage1.py b/yeto/rl/stage1.py index e2e4c3487..408209169 100644 --- a/yeto/rl/stage1.py +++ b/yeto/rl/stage1.py @@ -185,6 +185,10 @@ def record(self) -> dict[str, Any]: STORE_BACKENDS = {"oss": None, "modal-volume": "yeto.cloud.modal_ckpt_store:store_from_env"} +# GRPO + token-level TIS, execution.max_policy_staleness 2 (relative to the yeto workdir). +DEFAULT_ALGORITHM_SPEC = "examples/rl_algorithms/decoupled-ppo-c2.json" + + def store_env(plan: Stage1Plan) -> dict[str, str]: """Environment every node needs for the store backend (none for OSS here: the nodes' ossutil/OSS configuration is set up on the node, runbook S4).""" @@ -473,6 +477,10 @@ def expand(plan: Stage1Plan, *, hf_config: str | Path | None = None) -> Expanded *(("--miles-commit", str(miles["commit"])) if len(str(miles["commit"])) == 40 else ()), "--trust-remote-code", "--learner-id", "0", "--num-learners", "1", "--rl-single-island-no-sync", + # samples come from the buffer, one round ahead at most (sampler max_backlog_groups); + # a policy age >= 1 needs an algorithm spec that tolerates it (TIS correction) + "--rl-max-policy-age", "1", + "--rl-algorithm-spec", str(miles.get("algorithm_spec") or DEFAULT_ALGORITHM_SPEC), "--global-rounds", str(int(plan.doc["rounds"])), "--groups-per-round", str(int(task["rollout_batch_size"])), "--over-sampling-batch-size", str(int(task["rollout_batch_size"])), From 80da1a1b23a8488fe4fc3ba14671a60006c4d861 Mon Sep 17 00:00:00 2001 From: michael <78113554+chowhsu@users.noreply.github.com> Date: Tue, 20 Oct 2026 03:19:50 +0000 Subject: [PATCH 15/89] s25_g1 driver: logs on container disk, mirrored to the Volume An open file on a Modal Volume blocks every reload in that container (run-3: the re-quantizer could not reload while its own log was open there). Co-Authored-By: Claude Opus 5.5 --- tools/probes/s25_g1/modal_app.py | 37 +++++++++++++++++++++++++------- 1 file changed, 29 insertions(+), 8 deletions(-) diff --git a/tools/probes/s25_g1/modal_app.py b/tools/probes/s25_g1/modal_app.py index cfba2d623..d05a4e120 100644 --- a/tools/probes/s25_g1/modal_app.py +++ b/tools/probes/s25_g1/modal_app.py @@ -73,17 +73,35 @@ def _hf_secret() -> modal.Secret: return modal.Secret.from_dict({"HF_TOKEN": token.read_text().strip()} if token.is_file() else {}) +LOCAL_LOGS = Path("/tmp/g1-logs") # container disk: an open file on the Volume blocks every reload + + def _log_dir(tag: str) -> Path: - path = Path("/data/store/logs") / tag + path = LOCAL_LOGS / tag path.mkdir(parents=True, exist_ok=True) return path -def _commit() -> None: +def _commit(tag: str | None = None) -> None: + """Copy the local logs of ``tag`` to the store Volume (closed files only), then commit.""" try: + if tag is not None: + shutil.copytree(LOCAL_LOGS / tag, Path("/data/store/logs") / tag, dirs_exist_ok=True, + ignore=shutil.ignore_patterns("*.sock", "*.pid")) store.commit() except Exception as error: # noqa: BLE001 - evidence copy only - print(f"[g1] store commit failed: {error}", flush=True) + print(f"[g1] log mirror/commit failed: {error}", flush=True) + + +def _mirror_every(tag: str, seconds: float = 180.0) -> None: + import threading + + def loop() -> None: + while True: + time.sleep(seconds) + _commit(tag) + + threading.Thread(target=loop, daemon=True).start() def _sh(cmd: str, log: Path | None = None, env: dict | None = None, check: bool = True) -> str: @@ -187,7 +205,7 @@ def smoke_run(tag: str) -> None: nvml.terminate() (log / "smoke.json").write_text(json.dumps(record, indent=1, default=str)) print(json.dumps(record, default=str)[:4000], flush=True) - _commit() + _commit(tag) @app.local_entrypoint() @@ -202,6 +220,7 @@ def sampler_island(tag: str) -> int: log = _log_dir(tag) _gpu_facts(log, 1) nvml = _nvml_sampler(log / "nvml.csv") + _mirror_every(tag) plan = _plan(log) run = plan["sampler"]["engines"][0] spec_path = Path(run["spec_path"]) @@ -215,7 +234,7 @@ def sampler_island(tag: str) -> int: state = Path(run["spec"]["event_tape"]).parent if state.is_dir(): shutil.copytree(state, log / "state", dirs_exist_ok=True) - _commit() + _commit(tag) return rc @@ -224,6 +243,7 @@ def trainer_island(tag: str, stop_after_rounds: int | None = None) -> int: log = _log_dir(tag) _gpu_facts(log, 2) nvml = _nvml_sampler(log / "nvml.csv") + _mirror_every(tag) plan = _plan(log) trainer = plan["trainer"] env = _env(trainer["env"]) @@ -253,7 +273,7 @@ def trainer_island(tag: str, stop_after_rounds: int | None = None) -> int: shutil.copytree(f"/data/yeto/{plan['run_id']}", log / "state", dirs_exist_ok=True, ignore=shutil.ignore_patterns("*.pt", "*.safetensors")) _sh("cp -r /tmp/ray/session_latest/logs " + shlex.quote(str(log / "ray-logs")) + " || true", check=False) - _commit() + _commit(tag) return rc @@ -300,8 +320,9 @@ def run_trainer(sub: str, stop: int | None) -> object: time.sleep(30) result = {"trainer_a": first, "trainer_b": second, "sampler": sampler_rc} print(json.dumps(result, default=str), flush=True) - _log_dir(tag) - Path(f"/data/store/logs/{tag}/orchestrate.json").write_text(json.dumps(result, default=str)) + out = Path(f"/data/store/logs/{tag}") + out.mkdir(parents=True, exist_ok=True) + (out / "orchestrate.json").write_text(json.dumps(result, default=str)) _commit() return result From f483de5f72a25e03c2733d76747f31c698b9e795 Mon Sep 17 00:00:00 2001 From: michael <78113554+chowhsu@users.noreply.github.com> Date: Tue, 20 Oct 2026 03:26:32 +0000 Subject: [PATCH 16/89] Stage 1: pin the island's AlgorithmSpec hash like the launcher The trainer-only island (partitioned-serial) refuses to start without --rl-expected-algorithm-sha256 (Modal run-5). The plan computes it from the island argv with the launcher's resolve_ports_algorithm. Co-Authored-By: Claude Opus 5.5 --- yeto/rl/stage1.py | 18 ++++++++++++++++++ 1 file changed, 18 insertions(+) diff --git a/yeto/rl/stage1.py b/yeto/rl/stage1.py index 408209169..88c968f7b 100644 --- a/yeto/rl/stage1.py +++ b/yeto/rl/stage1.py @@ -189,6 +189,21 @@ def record(self) -> dict[str, Any]: DEFAULT_ALGORITHM_SPEC = "examples/rl_algorithms/decoupled-ppo-c2.json" +def expected_algorithm_sha256(island_argv: Sequence[str]) -> str: + """The AlgorithmSpec hash the island computes from this argv (the launcher's + ``resolve_ports_algorithm(args).sha256()``). The spec path is resolved against + the yeto tree, where the island runs.""" + from yeto.launch import REPO_ROOT # noqa: PLC0415 + from yeto.rl.adapters.miles.island_entry import parse_args # noqa: PLC0415 + from yeto.rl.engine.algorithm import resolve_ports_algorithm # noqa: PLC0415 + + argv = list(island_argv[3:]) + i = argv.index("--rl-algorithm-spec") + 1 + if not Path(argv[i]).is_absolute(): + argv[i] = str(Path(REPO_ROOT) / argv[i]) + return resolve_ports_algorithm(parse_args(argv)).sha256() + + def store_env(plan: Stage1Plan) -> dict[str, str]: """Environment every node needs for the store backend (none for OSS here: the nodes' ossutil/OSS configuration is set up on the node, runbook S4).""" @@ -493,6 +508,9 @@ def expand(plan: Stage1Plan, *, hf_config: str | Path | None = None) -> Expanded "--completed-groups-path", f"{state_dir}/island-checkpoint.pt", "--event-tape", f"{state_dir}/events.jsonl", "--audit-dir", f"{state_dir}/audit", *_provenance_args(str(task["reward"]))] + # The launcher pins the island's AlgorithmSpec hash (--rl-expected-algorithm-sha256); + # the trainer-only island refuses to start without it (Modal run-5). + island_argv += ["--rl-expected-algorithm-sha256", expected_algorithm_sha256(island_argv)] try: model_args = glm5_next.model_argv( num_layers=int(model["num_layers"]), From ab3e17269047a95e2e4a4f8b9ad268c8b3f25151 Mon Sep 17 00:00:00 2001 From: michael <78113554+chowhsu@users.noreply.github.com> Date: Tue, 20 Oct 2026 03:33:10 +0000 Subject: [PATCH 17/89] Stage 1: export RAY_ADDRESS for the island so its env reaches the Ray actors connect_island_ray forwards the sample-buffer and store variables to the Ray actors only when RAY_ADDRESS is set; without it the buffer rollout function in the RolloutExecutor found no YETO_SAMPLE_BUFFER_MAX_POLICY_AGE (Modal run-6). Co-Authored-By: Claude Opus 5.5 --- tools/probes/s25_g1/modal_app.py | 4 +++- yeto/rl/stage1.py | 3 +++ 2 files changed, 6 insertions(+), 1 deletion(-) diff --git a/tools/probes/s25_g1/modal_app.py b/tools/probes/s25_g1/modal_app.py index d05a4e120..9b79abb39 100644 --- a/tools/probes/s25_g1/modal_app.py +++ b/tools/probes/s25_g1/modal_app.py @@ -259,7 +259,9 @@ def trainer_island(tag: str, stop_after_rounds: int | None = None) -> int: with open(log / "requant.log", "w", encoding="utf-8") as rq, open(log / "trainer.log", "w", encoding="utf-8") as f: requant = subprocess.Popen(plan["requant"]["loop_argv"], env=env, cwd=WORKDIR, stdout=rq, stderr=subprocess.STDOUT) - rc = subprocess.run(island, env=env, cwd=WORKDIR, stdout=f, stderr=subprocess.STDOUT).returncode + # as the stage-1 node script does: the island forwards its env to the Ray actors + rc = subprocess.run(island, env={**env, "RAY_ADDRESS": "127.0.0.1:6379"}, cwd=WORKDIR, stdout=f, + stderr=subprocess.STDOUT).returncode requant.terminate() try: requant.wait(timeout=60) diff --git a/yeto/rl/stage1.py b/yeto/rl/stage1.py index 88c968f7b..a7bac39cb 100644 --- a/yeto/rl/stage1.py +++ b/yeto/rl/stage1.py @@ -655,6 +655,9 @@ def trainer_island_spec(plan: Stage1Plan, out: Expanded, trainers: Sequence[Node "if [ \"$SKYPILOT_NODE_RANK\" = 0 ]; then\n" f" ray start --head --port={RAY_PORT} --node-ip-address=\"$me\" --num-gpus=\"$SKYPILOT_NUM_GPUS_PER_NODE\"\n" f" {join}\n" + # the island forwards its own environment (sample buffer, store) to the Ray + # actors only when RAY_ADDRESS names the cluster (entry.connect_island_ray) + f" export RAY_ADDRESS=\"$me:{RAY_PORT}\"\n" f" exec {island}\n" "fi\n" f"exec ray start --address=\"$head:{RAY_PORT}\" --node-ip-address=\"$me\" " From e875afd0ce31088aab28f8f67c4031cf911e51ca Mon Sep 17 00:00:00 2001 From: michael <78113554+chowhsu@users.noreply.github.com> Date: Tue, 20 Oct 2026 03:35:05 +0000 Subject: [PATCH 18/89] T-G3 harness: extra continuous and resume samples for the N7.1d spread Co-Authored-By: Claude Opus 5.5 --- tools/probes/g3_miles/harness.py | 6 +++--- tools/probes/g3_miles/modal_g3.py | 8 ++++++++ 2 files changed, 11 insertions(+), 3 deletions(-) diff --git a/tools/probes/g3_miles/harness.py b/tools/probes/g3_miles/harness.py index 81012782c..c2bd4c94d 100644 --- a/tools/probes/g3_miles/harness.py +++ b/tools/probes/g3_miles/harness.py @@ -685,9 +685,9 @@ def manifest(summaries, global_step): scheduler_increment=args.global_batch_size, ranks={str(s["path"]): OptimizerResumeState.from_dict(s["resume"]) for s in summaries}) - if ns.mode in ("cont5", "cont5b"): + if ns.mode.startswith("cont5"): res["steps"] = train_steps(args, model, optimizer, scheduler, 1, 5, out) - elif ns.mode == "run3": + elif ns.mode.startswith("run3"): res["steps"] = train_steps(args, model, optimizer, scheduler, 1, 3, out) summary = cut_plugin.save_cut_shard(actor, directory=str(cut_dir), cut_id="g3-cut-3") if cut_plugin.REFUSED in summary: @@ -701,7 +701,7 @@ def manifest(summaries, global_step): (cut_dir / "files.json").write_text(json.dumps(summaries, default=str)) (cut_dir / "resume_manifest.json").write_text(json.dumps(m.to_dict())) res["hdo_after_save"] = hdo_report(optimizer) - elif ns.mode == "resume2": + elif ns.mode.startswith("resume2"): files = json.loads((cut_dir / "files.json").read_text()) saved = ResumeManifest.from_dict(json.loads((cut_dir / "resume_manifest.json").read_text())) summary = cut_plugin.restore_cut_shard(actor, directory=str(cut_dir), files=files, cut_id="g3-cut-3") diff --git a/tools/probes/g3_miles/modal_g3.py b/tools/probes/g3_miles/modal_g3.py index 537b2090b..c8670e776 100644 --- a/tools/probes/g3_miles/modal_g3.py +++ b/tools/probes/g3_miles/modal_g3.py @@ -59,6 +59,14 @@ run pre-n75-te 600 $TR $H --phase n75 --impl te --tag te --out "$OUT/n75" run pre-n71d 400 $TR $H --phase n71d --mode cont5 --out "$OUT/n71d" --cut-dir "$OUT/n71d/cut" fi +if [[ "$PHASES" == n71d-noise ]]; then # more samples of the run-to-run spread (cut in a second directory) + export NCCL_ALGO=Ring CUBLAS_WORKSPACE_CONFIG=:4096:8 NVIDIA_TF32_OVERRIDE=0 + for m in cont5c cont5d; do + run n71d-$m 600 $TR $H --phase n71d --mode $m --out "$OUT/n71d" --cut-dir "$OUT/n71d/cut" + done + run n71d-run3b 600 $TR $H --phase n71d --mode run3b --out "$OUT/n71d" --cut-dir "$OUT/n71d/cutb" + run n71d-resume2b 600 $TR $H --phase n71d --mode resume2b --out "$OUT/n71d" --cut-dir "$OUT/n71d/cutb" +fi if [[ "$PHASES" == all || "$PHASES" == n75 ]]; then run n75-te 1200 $TR $H --phase n75 --impl te --tag te --out "$OUT/n75" run n75-te2 900 $TR $H --phase n75 --impl te --tag te2 --weights-from "$OUT/n75" --out "$OUT/n75" From 55cd48d1ce1db4229d98bbb8f2ace37bfaf754db Mon Sep 17 00:00:00 2001 From: michael <78113554+chowhsu@users.noreply.github.com> Date: Tue, 20 Oct 2026 03:35:52 +0000 Subject: [PATCH 19/89] Hand host copies to the requantizer from the Miles adapter; record the in-trainer N8.2 result On 2xH100 (T-G3) a GPU expert slice with W4A8 AWQ failed (the AWQ factors live on the host), while the host-copy path was byte-equal to the single-process output for W4A8 (R2 AWQ) and MXFP4. Co-Authored-By: Claude Opus 5.5 --- yeto/rl/adapters/miles/quant_shards.py | 13 ++++++++++--- yeto/rl/quant/shards.py | 4 +++- 2 files changed, 13 insertions(+), 4 deletions(-) diff --git a/yeto/rl/adapters/miles/quant_shards.py b/yeto/rl/adapters/miles/quant_shards.py index 9b8f564d6..fc7d5f9c5 100644 --- a/yeto/rl/adapters/miles/quant_shards.py +++ b/yeto/rl/adapters/miles/quant_shards.py @@ -19,9 +19,16 @@ renames them to the checkpoint names (``model.language_model.layers.*``; module doc of :mod:`yeto.rl.quant.shards`, "Ownership"). +Both functions hand host (CPU) copies to the quantizer: the CPU path is the one +checked byte for byte against the single-process converter (GPU T-G3, 2026-10-20); +a GPU expert slice with W4A8 AWQ failed there (factors on the host), and GPU +quantization bytes were not compared for W4A8. + Limits (explicit errors): expert tensor parallelism must be 1 (a column split of ``linear_fc2`` would cut the W4A8 per-row scale). Status: CPU tests with -synthetic Megatron names; not run inside Miles. +synthetic Megatron names; on 2xH100 inside a real Miles/Megatron glm5_next +(TP2, EP2; TE grouped and local sequential experts) the merged W4A8 (R2 AWQ) +and MXFP4 outputs equal the single-process output byte for byte (T-G3). """ from __future__ import annotations @@ -61,7 +68,7 @@ def megatron_expert_slices(named_params: Iterable[tuple[str, torch.Tensor]], *, raise ValueError(f"{name}: local expert {local} >= num_local_experts {num_local_experts}") expert = int(ep_rank) * int(num_local_experts) + local prefix = f"{hf_layer_prefix}{int(layer_offset) + int(match['layer'])}.mlp.experts" - weight = tensor.detach() + weight = tensor.detach().cpu() if (match["fc"] or match["fc2"]) == "1": for hf_name, part in split_gate_up(prefix, expert, weight): yield SourceSlice(hf_name, part) @@ -86,4 +93,4 @@ def checkpoint_name(name: str) -> str: def hf_slices(named_tensors: Iterable[tuple[str, torch.Tensor]]) -> Iterator[SourceSlice]: """Whole tensors from Miles ``convert_to_hf``, renamed to checkpoint names (:func:`checkpoint_name`).""" for name, tensor in named_tensors: - yield SourceSlice(checkpoint_name(name), tensor.detach()) + yield SourceSlice(checkpoint_name(name), tensor.detach().cpu()) diff --git a/yeto/rl/quant/shards.py b/yeto/rl/quant/shards.py index 145441e07..10c6eb28f 100644 --- a/yeto/rl/quant/shards.py +++ b/yeto/rl/quant/shards.py @@ -25,7 +25,9 @@ ``@``. Status: CPU tests (2 gloo ranks, both schemes, byte-equal to the single-process -output; tests/test_rl_quant_shards.py). Not run inside a real trainer. +output; tests/test_rl_quant_shards.py). Inside a real Miles trainer (2xH100, +glm5_next 4-layer mini, TP2/EP2, through the Miles adapter): byte-equal for both +schemes (T-G3, 2026-10-20). Not run inside verl. """ from __future__ import annotations From d44a44b44d1c43ba326c005b07f57505397ac096 Mon Sep 17 00:00:00 2001 From: michael <78113554+chowhsu@users.noreply.github.com> Date: Tue, 20 Oct 2026 03:41:01 +0000 Subject: [PATCH 20/89] glm5_next trains from a Megatron torch_dist checkpoint in raw mode Miles loads an HF checkpoint only through Megatron-Bridge, which has no glm5_next bridge (Modal run-7). The recipe now uses --megatron-to-hf-mode raw and requires --megatron-ref-load; stage 1 takes model.bf16_source.torch_dist, adds a convert step (Miles tools/convert_hf_to_torch_dist.py on trainer node 0) and reports a gap when the path is unset. The Modal driver converts once into the models Volume. Co-Authored-By: Claude Opus 5.5 --- examples/stage1/glm53-flash-ppu.yaml | 1 + examples/stage1/glm5next-mini-modal.yaml | 1 + tests/test_s25_g1_stage1mini.py | 3 +++ tools/probes/s25_g1/modal_app.py | 11 +++++++++++ yeto/rl/adapters/miles/config.py | 4 +++- yeto/rl/engine/run_config.py | 7 +++++++ yeto/rl/stage1.py | 22 ++++++++++++++++++++++ 7 files changed, 48 insertions(+), 1 deletion(-) diff --git a/examples/stage1/glm53-flash-ppu.yaml b/examples/stage1/glm53-flash-ppu.yaml index 5f23b3a1d..c212ec8dd 100644 --- a/examples/stage1/glm53-flash-ppu.yaml +++ b/examples/stage1/glm53-flash-ppu.yaml @@ -37,6 +37,7 @@ model: origin: modal-volume://apex-w4a8-full/source/cecb5eeb9c6b32404a0dd930df81de2c239bdd84 staged: oss://oss-pai-ircyeaqtnfjw11wby2-cn-wulanchabu/yeto/models/apex-flash-1-bf16-cecb5ee/ local_dir: /data/models/apex-flash-1-bf16-cecb5ee + torch_dist: /data/models/apex-flash-1-bf16-cecb5ee_torch_dist # Miles converter output (convert step) sampler_initial: # round 0: the HF MXFP4 export, not our re-quantizer repo: yetalabs/apex-flash-1-mxfp4-experts revision: 70e86a0 diff --git a/examples/stage1/glm5next-mini-modal.yaml b/examples/stage1/glm5next-mini-modal.yaml index bef79b8c3..acb9c690a 100644 --- a/examples/stage1/glm5next-mini-modal.yaml +++ b/examples/stage1/glm5next-mini-modal.yaml @@ -33,6 +33,7 @@ model: origin: hf://CharyZeng/GLM-5.3-Flash-4layer@7fb7883207fe73ea04430ac392f6ffe2b9f936c4 staged: modal-volume://yeto-s25-g1-models/glm5next-mini-bf16-7fb7883 local_dir: /data/models/glm5next-mini-bf16-7fb7883 + torch_dist: /data/models/glm5next-mini-bf16-7fb7883_torch_dist # Miles converter output (convert step) sampler_initial: # round 0: our N8.5 converter on the bf16 mini model repo: local/glm5next-mini-mxfp4-r0-7fb7883 revision: 7fb7883 diff --git a/tests/test_s25_g1_stage1mini.py b/tests/test_s25_g1_stage1mini.py index a479f7576..c3cc14158 100644 --- a/tests/test_s25_g1_stage1mini.py +++ b/tests/test_s25_g1_stage1mini.py @@ -149,6 +149,8 @@ def test_mini_modal_plan_has_no_errors_and_uses_the_volume_store(): "expected_expert_matrices": 288} assert spec["engine"]["options"]["load_report_dir"].endswith("/load-report") assert "--frozen-base" in out.requant["argv"] + assert argv[argv.index("--megatron-ref-load") + 1] == out.trainer["torch_dist"] + assert out.trainer["convert_argv"][out.trainer["convert_argv"].index("--save") + 1] == out.trainer["torch_dist"] # the PPU plan keeps OSS and passes no --miles-commit for its short commit ppu = expand(load_plan(REPO / "examples/stage1/glm53-flash-ppu.yaml")) assert ppu.errors == [] and "--miles-commit" not in ppu.trainer["miles"]["island_argv"] @@ -198,3 +200,4 @@ def test_mini_island_argv_translates_to_a_miles_argv_on_cpu(tmp_path): assert flag in miles assert miles[miles.index("--num-experts") + 1] == "32" and miles.count("--num-experts") == 1 assert miles[miles.index("--moe-layer-freq") + 1] == "[0,1,1,1]" + assert miles[miles.index("--megatron-to-hf-mode") + 1] == "raw" diff --git a/tools/probes/s25_g1/modal_app.py b/tools/probes/s25_g1/modal_app.py index 9b79abb39..ea53ed6e0 100644 --- a/tools/probes/s25_g1/modal_app.py +++ b/tools/probes/s25_g1/modal_app.py @@ -253,6 +253,17 @@ def trainer_island(tag: str, stop_after_rounds: int | None = None) -> int: (log / "argv.json").write_text(json.dumps({"island": island, "requant": plan["requant"]["loop_argv"], "env": trainer["env"]}, indent=1)) Path("/data/yeto", plan["run_id"]).mkdir(parents=True, exist_ok=True) + marker = Path(str(trainer["torch_dist"])) / "latest_checkpointed_iteration.txt" + if trainer.get("convert_argv") and not (marker.is_file() and marker.read_text().strip() == "release"): + began_convert = time.time() + with open(log / "convert.log", "w", encoding="utf-8") as f: + crc = subprocess.run(trainer["convert_argv"], env={**env, "PYTHONPATH": f"/root/miles:{env['PYTHONPATH']}"}, + cwd=WORKDIR, stdout=f, stderr=subprocess.STDOUT).returncode + (log / "convert.json").write_text(json.dumps({"rc": crc, "wall_s": round(time.time() - began_convert, 1)})) + models.commit() # trainer B (a new container) loads the same torch_dist + if crc != 0: + _commit(tag) + return 90 _sh("ray start --head --port=6379 --node-ip-address=127.0.0.1 --num-gpus=2 --disable-usage-stats", log / "ray.log", env=env) began = time.time() diff --git a/yeto/rl/adapters/miles/config.py b/yeto/rl/adapters/miles/config.py index 449503da3..dadbb341d 100644 --- a/yeto/rl/adapters/miles/config.py +++ b/yeto/rl/adapters/miles/config.py @@ -870,7 +870,9 @@ def translate_run_config( "--train-backend", "megatron", "--hf-checkpoint", config.hf_checkpoint, "--ref-load", config.ref_load, - "--megatron-to-hf-mode", "bridge", + # glm5_next: Megatron-Bridge has no glm5_next bridge; Miles loads the torch_dist + # checkpoint and exports through its own converters (raw) + "--megatron-to-hf-mode", "raw" if recipe.name == RECIPE_GLM5_NEXT else "bridge", "--model-name", model_name, *model_recipe_values, "--num-layers", str(geometry.num_layers), diff --git a/yeto/rl/engine/run_config.py b/yeto/rl/engine/run_config.py index b2c82c4ea..616e588a5 100644 --- a/yeto/rl/engine/run_config.py +++ b/yeto/rl/engine/run_config.py @@ -778,6 +778,13 @@ def resolve_rl_run_config( elif recipe == RECIPE_GLM5_NEXT: if parameter_mode != "full": raise ValueError("the glm5_next recipe is full-parameter only") + if getattr(args, "megatron_ref_load", None) is None: + # Miles loads an HF checkpoint only through Megatron-Bridge, which has no + # glm5_next bridge (Modal run-7): train from the converted torch_dist checkpoint + raise ValueError( + "the glm5_next recipe needs --megatron-ref-load (Miles " + "tools/convert_hf_to_torch_dist.py output; Megatron-Bridge cannot load glm5_next)" + ) # Megatron picks the backend (Miles run_glm5_3_flash.py passes none) recipe_name, attention_backend = RECIPE_GLM5_NEXT, "auto" elif fn_variant is not None: diff --git a/yeto/rl/stage1.py b/yeto/rl/stage1.py index a7bac39cb..8fdcf8cc6 100644 --- a/yeto/rl/stage1.py +++ b/yeto/rl/stage1.py @@ -489,6 +489,8 @@ def expand(plan: Stage1Plan, *, hf_config: str | Path | None = None) -> Expanded "--tensor-parallel", str(layout["tp"]), "--pipeline-parallel", str(layout["pp"]), "--expert-parallel", str(layout["ep"]), "--miles-root", str(miles["workdir"]), + *(("--megatron-ref-load", str(model["bf16_source"]["torch_dist"])) + if model["bf16_source"].get("torch_dist") else ()), *(("--miles-commit", str(miles["commit"])) if len(str(miles["commit"])) == 40 else ()), "--trust-remote-code", "--learner-id", "0", "--num-learners", "1", "--rl-single-island-no-sync", @@ -518,8 +520,17 @@ def expand(plan: Stage1Plan, *, hf_config: str | Path | None = None) -> Expanded num_experts=int(model.get("num_experts", glm5_next.NUM_EXPERTS)), transformer_impl=impl) except ValueError as error: raise Stage1Error(f"plan: trainer.miles.transformer_impl: {error}") from None + torch_dist = model["bf16_source"].get("torch_dist") + convert_argv = None if not torch_dist else [ + "torchrun", "--nproc-per-node", str(per_node), f"{miles['workdir']}/tools/convert_hf_to_torch_dist.py", + *model_args, "--hf-checkpoint", model_dir, "--save", str(torch_dist), + "--tensor-model-parallel-size", str(layout["tp"]), "--pipeline-model-parallel-size", str(layout["pp"]), + "--expert-model-parallel-size", str(min(layout["ep"], per_node))] out.trainer = { "backend": backend, "nodes": [n.name for n in trainers], "cards": cards, + # Miles trains glm5_next from a Megatron torch_dist checkpoint (raw mode); converted once + # on trainer node 0 (one node's cards; not run for the 45-layer model) + "torch_dist": torch_dist, "convert_argv": convert_argv, "image": images.get("trainer"), "env": trainer_env, "miles": {"commit": miles["commit"], "megatron": miles["megatron"], "mbridge": miles["mbridge"], "layout": layout, "memory_estimate": megatron_memory(model, miles, layout, trainers, sampling, hf_config, @@ -721,6 +732,14 @@ def _steps(plan: Stage1Plan, out: Expanded, trainers: Sequence[Node], samplers: f"ossutil cp -r -u {shlex.quote(src['staged'])} {shlex.quote(src['local_dir'])}/\n" f"test -f {shlex.quote(src['local_dir'])}/model.safetensors.index.json\n", f"{key} weights on the node's /data (revision {src['revision']})")) + if out.trainer.get("convert_argv") and trainers: + node, dst = trainers[0], str(out.trainer["torch_dist"]) + convert = _docker(plan, node, "convert", out.trainer["image"], out.trainer["convert_argv"], + out.trainer["env"]) + steps.append(Step(f"convert-torch-dist-{node.name}", node.name, + f"set -eu\nif [ \"$(cat {shlex.quote(dst)}/latest_checkpointed_iteration.txt 2>/dev/null)\" " + f"!= release ]; then\n{convert}\nfi\n", + "bf16 source -> Megatron torch_dist (Miles converter; the island's --megatron-ref-load)")) for run in out.sampler["engines"]: node = next(n for n in samplers if n.name == run["node"]) spec_json = json.dumps(run["spec"], indent=1, sort_keys=True) @@ -912,6 +931,9 @@ def gaps(plan: Stage1Plan, out: Expanded) -> list[str]: if gap := _module_gap(module, what): found.append(gap) if out.trainer["backend"] == "miles": + if not out.trainer.get("torch_dist"): + found.append("model.bf16_source.torch_dist is unset: Miles trains glm5_next only from a " + "Megatron torch_dist checkpoint (Megatron-Bridge cannot load glm5_next)") if gap := _check_island_argv(out.trainer["miles"]["island_argv"]): found.append(gap) recipes = _miles_recipes() From d1fc221bd9e62ab1a9f43ba11bf703db4dc66b71 Mon Sep 17 00:00:00 2001 From: michael <78113554+chowhsu@users.noreply.github.com> Date: Tue, 20 Oct 2026 03:46:57 +0000 Subject: [PATCH 21/89] Stage 1: the torch_dist converter takes the model arguments only Miles' convert_hf_to_torch_dist.py rejects the run flag --qkv-format (Modal run-8). Co-Authored-By: Claude Opus 5.5 --- tests/test_s25_g1_stage1mini.py | 1 + yeto/rl/adapters/miles/models/glm5_next.py | 7 +++++++ yeto/rl/stage1.py | 3 ++- 3 files changed, 10 insertions(+), 1 deletion(-) diff --git a/tests/test_s25_g1_stage1mini.py b/tests/test_s25_g1_stage1mini.py index c3cc14158..14f8a5d0a 100644 --- a/tests/test_s25_g1_stage1mini.py +++ b/tests/test_s25_g1_stage1mini.py @@ -151,6 +151,7 @@ def test_mini_modal_plan_has_no_errors_and_uses_the_volume_store(): assert "--frozen-base" in out.requant["argv"] assert argv[argv.index("--megatron-ref-load") + 1] == out.trainer["torch_dist"] assert out.trainer["convert_argv"][out.trainer["convert_argv"].index("--save") + 1] == out.trainer["torch_dist"] + assert "--qkv-format" not in out.trainer["convert_argv"] and "--spec" in out.trainer["convert_argv"] # the PPU plan keeps OSS and passes no --miles-commit for its short commit ppu = expand(load_plan(REPO / "examples/stage1/glm53-flash-ppu.yaml")) assert ppu.errors == [] and "--miles-commit" not in ppu.trainer["miles"]["island_argv"] diff --git a/yeto/rl/adapters/miles/models/glm5_next.py b/yeto/rl/adapters/miles/models/glm5_next.py index 980dc719a..0cb3c4d6f 100644 --- a/yeto/rl/adapters/miles/models/glm5_next.py +++ b/yeto/rl/adapters/miles/models/glm5_next.py @@ -130,6 +130,13 @@ def model_argv(*, num_layers: int = NUM_LAYERS, first_k_dense_replace: int = FIR *(TE_ONLY_ARGV if impl == TRANSFORMER_ENGINE else LOCAL_ARGV), *RUN_ARGV] +def model_argv_without_run_flags(argv: Sequence[str]) -> list[str]: + """:func:`model_argv` without :data:`RUN_ARGV` (Miles' checkpoint converter + parses the model arguments only).""" + run = {group[0] for group in _flag_groups(RUN_ARGV)} + return [token for group in _flag_groups(argv) if group[0] not in run for token in group] + + def _flag_groups(argv: Sequence[str]) -> Iterator[list[str]]: group: list[str] = [] for token in argv: diff --git a/yeto/rl/stage1.py b/yeto/rl/stage1.py index 8fdcf8cc6..aebf8e496 100644 --- a/yeto/rl/stage1.py +++ b/yeto/rl/stage1.py @@ -523,7 +523,8 @@ def expand(plan: Stage1Plan, *, hf_config: str | Path | None = None) -> Expanded torch_dist = model["bf16_source"].get("torch_dist") convert_argv = None if not torch_dist else [ "torchrun", "--nproc-per-node", str(per_node), f"{miles['workdir']}/tools/convert_hf_to_torch_dist.py", - *model_args, "--hf-checkpoint", model_dir, "--save", str(torch_dist), + # model arguments only: the converter's parser has no run flags (--qkv-format; Modal run-8) + *glm5_next.model_argv_without_run_flags(model_args), "--hf-checkpoint", model_dir, "--save", str(torch_dist), "--tensor-model-parallel-size", str(layout["tp"]), "--pipeline-model-parallel-size", str(layout["pp"]), "--expert-model-parallel-size", str(min(layout["ep"], per_node))] out.trainer = { From db07dc1716149fd3d586f5cb36c858e49ed35240 Mon Sep 17 00:00:00 2001 From: michael <78113554+chowhsu@users.noreply.github.com> Date: Tue, 20 Oct 2026 03:54:45 +0000 Subject: [PATCH 22/89] Stage 1: run the torch_dist converter with the Megatron checkout on PYTHONPATH The editable megatron_core install does not expose megatron.post_training, which Megatron's get_model imports in the converter (Modal run-9). Co-Authored-By: Claude Opus 5.5 --- tools/probes/s25_g1/modal_app.py | 15 ++++++++++++++- yeto/rl/stage1.py | 5 ++++- 2 files changed, 18 insertions(+), 2 deletions(-) diff --git a/tools/probes/s25_g1/modal_app.py b/tools/probes/s25_g1/modal_app.py index ea53ed6e0..ae944106f 100644 --- a/tools/probes/s25_g1/modal_app.py +++ b/tools/probes/s25_g1/modal_app.py @@ -257,7 +257,8 @@ def trainer_island(tag: str, stop_after_rounds: int | None = None) -> int: if trainer.get("convert_argv") and not (marker.is_file() and marker.read_text().strip() == "release"): began_convert = time.time() with open(log / "convert.log", "w", encoding="utf-8") as f: - crc = subprocess.run(trainer["convert_argv"], env={**env, "PYTHONPATH": f"/root/miles:{env['PYTHONPATH']}"}, + path = f"{trainer['convert_env']['PYTHONPATH']}:{env['PYTHONPATH']}" + crc = subprocess.run(trainer["convert_argv"], env={**env, "PYTHONPATH": path}, cwd=WORKDIR, stdout=f, stderr=subprocess.STDOUT).returncode (log / "convert.json").write_text(json.dumps({"rc": crc, "wall_s": round(time.time() - began_convert, 1)})) models.commit() # trainer B (a new container) loads the same torch_dist @@ -344,3 +345,15 @@ def run_trainer(sub: str, stop: int | None) -> object: def launch(tag: str = "run-1", cut_after: int = 2) -> None: call = orchestrate.spawn(tag, cut_after) print(f"spawned orchestrate {call.object_id}") + + +@app.function(image=image, volumes={"/data/models": models}, cpu=2, memory=8192, timeout=900) +def cpu_shell(cmd: str) -> str: + """CPU-only inspection of the image (no GPU).""" + r = subprocess.run(["bash", "-lc", cmd], capture_output=True, text=True, cwd=WORKDIR, env=_env({})) + return r.stdout[-20000:] + "\nSTDERR:\n" + r.stderr[-4000:] + + +@app.local_entrypoint() +def shell(cmd: str) -> None: + print(cpu_shell.remote(cmd)) diff --git a/yeto/rl/stage1.py b/yeto/rl/stage1.py index aebf8e496..2d081d3db 100644 --- a/yeto/rl/stage1.py +++ b/yeto/rl/stage1.py @@ -532,6 +532,9 @@ def expand(plan: Stage1Plan, *, hf_config: str | Path | None = None) -> Expanded # Miles trains glm5_next from a Megatron torch_dist checkpoint (raw mode); converted once # on trainer node 0 (one node's cards; not run for the 45-layer model) "torch_dist": torch_dist, "convert_argv": convert_argv, + # the converter imports megatron.post_training, which the editable megatron_core + # install does not expose; Miles' own converter call puts the checkout first (run-9) + "convert_env": {"PYTHONPATH": f"{miles.get('megatron_path', '/root/Megatron-LM')}:{miles['workdir']}"}, "image": images.get("trainer"), "env": trainer_env, "miles": {"commit": miles["commit"], "megatron": miles["megatron"], "mbridge": miles["mbridge"], "layout": layout, "memory_estimate": megatron_memory(model, miles, layout, trainers, sampling, hf_config, @@ -736,7 +739,7 @@ def _steps(plan: Stage1Plan, out: Expanded, trainers: Sequence[Node], samplers: if out.trainer.get("convert_argv") and trainers: node, dst = trainers[0], str(out.trainer["torch_dist"]) convert = _docker(plan, node, "convert", out.trainer["image"], out.trainer["convert_argv"], - out.trainer["env"]) + {**out.trainer["env"], **out.trainer["convert_env"]}) steps.append(Step(f"convert-torch-dist-{node.name}", node.name, f"set -eu\nif [ \"$(cat {shlex.quote(dst)}/latest_checkpointed_iteration.txt 2>/dev/null)\" " f"!= release ]; then\n{convert}\nfi\n", From bd812747e204d264d1e00f4d478061e7131ad156 Mon Sep 17 00:00:00 2001 From: michael <78113554+chowhsu@users.noreply.github.com> Date: Tue, 20 Oct 2026 03:54:56 +0000 Subject: [PATCH 23/89] s25_g1 driver: convert-only entrypoint (2 cards, no sampler) Co-Authored-By: Claude Opus 5.5 --- tools/probes/s25_g1/modal_app.py | 14 +++++++++++++- 1 file changed, 13 insertions(+), 1 deletion(-) diff --git a/tools/probes/s25_g1/modal_app.py b/tools/probes/s25_g1/modal_app.py index ae944106f..5d8ea4cf5 100644 --- a/tools/probes/s25_g1/modal_app.py +++ b/tools/probes/s25_g1/modal_app.py @@ -239,7 +239,7 @@ def sampler_island(tag: str) -> int: @app.function(image=image, volumes=VOLUMES, gpu=f"{GPU_NAME}!:2", cpu=12, memory=163840, timeout=3600) -def trainer_island(tag: str, stop_after_rounds: int | None = None) -> int: +def trainer_island(tag: str, stop_after_rounds: int | None = None, convert_only: bool = False) -> int: log = _log_dir(tag) _gpu_facts(log, 2) nvml = _nvml_sampler(log / "nvml.csv") @@ -265,6 +265,12 @@ def trainer_island(tag: str, stop_after_rounds: int | None = None) -> int: if crc != 0: _commit(tag) return 90 + if convert_only: + _sh(f"ls -la {shlex.quote(str(trainer['torch_dist']))}; du -sh {shlex.quote(str(trainer['torch_dist']))}", + log / "convert-out.txt", check=False) + nvml.terminate() + _commit(tag) + return 0 _sh("ray start --head --port=6379 --node-ip-address=127.0.0.1 --num-gpus=2 --disable-usage-stats", log / "ray.log", env=env) began = time.time() @@ -341,6 +347,12 @@ def run_trainer(sub: str, stop: int | None) -> object: return result +@app.local_entrypoint() +def convert(tag: str = "convert-1") -> None: + call = trainer_island.spawn(tag, None, True) + print(f"spawned convert {call.object_id}") + + @app.local_entrypoint() def launch(tag: str = "run-1", cut_after: int = 2) -> None: call = orchestrate.spawn(tag, cut_after) From 648f47c769ad6ab3daa9bca8e4939d7cb2424e64 Mon Sep 17 00:00:00 2001 From: michael <78113554+chowhsu@users.noreply.github.com> Date: Tue, 20 Oct 2026 04:04:07 +0000 Subject: [PATCH 24/89] Stage 1: trainer.miles.extra_env; mini plan puts the Megatron checkout on PYTHONPATH Megatron's get_model imports megatron.post_training in the trainer actors too, which the ports image's editable megatron_core does not expose (Modal run-10). Co-Authored-By: Claude Opus 5.5 --- examples/stage1/glm5next-mini-modal.yaml | 4 ++++ yeto/rl/stage1.py | 4 +++- 2 files changed, 7 insertions(+), 1 deletion(-) diff --git a/examples/stage1/glm5next-mini-modal.yaml b/examples/stage1/glm5next-mini-modal.yaml index acb9c690a..418f01690 100644 --- a/examples/stage1/glm5next-mini-modal.yaml +++ b/examples/stage1/glm5next-mini-modal.yaml @@ -65,6 +65,10 @@ trainer: megatron: radixark/Megatron-LM@a84b105 mbridge: image workdir: /root/miles + extra_env: + # Megatron's get_model imports megatron.post_training (modelopt check), which the + # image's editable megatron_core install does not expose (Modal run-10) + PYTHONPATH: /root/Megatron-LM:/root/sky_workdir tp: 1 sequence_parallel: false pp: 1 diff --git a/yeto/rl/stage1.py b/yeto/rl/stage1.py index 2d081d3db..67ed28a5e 100644 --- a/yeto/rl/stage1.py +++ b/yeto/rl/stage1.py @@ -467,7 +467,9 @@ def expand(plan: Stage1Plan, *, hf_config: str | Path | None = None) -> Expanded layout = megatron_layout(miles, cards, per_node) common_env = {DEVICE_FAMILY_ENV: trainers[0].device_family, "YETO_RUN_ID": plan.run_id, "YETO_RL_STORE": store, "HF_HUB_OFFLINE": "1", **store_env(plan)} - trainer_env = {**common_env, "TORCHINDUCTOR_COMPILE_THREADS": "1"} + trainer_env = {**common_env, "TORCHINDUCTOR_COMPILE_THREADS": "1", + # trainer.miles.extra_env: image-specific settings (e.g. PYTHONPATH) + **{str(k): str(v) for k, v in (_section(trainer_cfg, "miles").get("extra_env") or {}).items()}} model_dir = str(model["bf16_source"]["local_dir"]) state_dir = state_dir_of(plan) impl = str(miles["transformer_impl"]) From abf0de22bb1e49630260a17b6e27e30d7e387dfa Mon Sep 17 00:00:00 2001 From: michael <78113554+chowhsu@users.noreply.github.com> Date: Tue, 20 Oct 2026 04:11:25 +0000 Subject: [PATCH 25/89] Trainer-only Miles island publishes without engines Under --debug-train-only there is no weight updater and no rollout engine: align_weight_version is a no-op and MilesPublisher skips the engine update and version stamp; the ReleasingPublisher's release is the publication (Modal run-11 stopped at both). Co-Authored-By: Claude Opus 5.5 --- tests/test_rl_miles_adapter_trainer_publish.py | 18 ++++++++++++++++++ tests/test_s25_g1_stage1mini.py | 15 +++++++++++++++ yeto/rl/adapters/miles/publish.py | 7 +++++++ yeto/rl/adapters/miles/state_plugin.py | 5 +++++ 4 files changed, 45 insertions(+) diff --git a/tests/test_rl_miles_adapter_trainer_publish.py b/tests/test_rl_miles_adapter_trainer_publish.py index f60893dd1..6e4ccbe36 100644 --- a/tests/test_rl_miles_adapter_trainer_publish.py +++ b/tests/test_rl_miles_adapter_trainer_publish.py @@ -574,3 +574,21 @@ def test_gspo_and_gmpo_clip_fraction_unchanged_by_unconditional_drain(variant, m t.train_step(handle()) assert t.step_metrics().clip_fraction == 0.25 assert state_plugin._STEP_LOSSES == [] + + +def test_trainer_only_island_publishes_without_engines(monkeypatch): + """--debug-train-only: no engine update, no stamp; the release is the publication.""" + from yeto.rl.adapters.miles import publish as publish_module + + calls = [] + + async def miles_update_weights(args, actor, executor, ctl, **kw): + calls.append("update_weights") + + monkeypatch.setattr(publish_module, "_default_update_weights", lambda: miles_update_weights) + pub = MilesPublisher( + args=SimpleNamespace(debug_train_only=True), actor_model=PluginActor(calls), rollout_executor=object(), + inference_controller=Controller([]), flatten_checksums=flatten, + ) + result = pub.publish(state()) + assert calls == [] and result.members == frozenset() and pub.last_engine_checksums is None diff --git a/tests/test_s25_g1_stage1mini.py b/tests/test_s25_g1_stage1mini.py index 14f8a5d0a..70ea647a6 100644 --- a/tests/test_s25_g1_stage1mini.py +++ b/tests/test_s25_g1_stage1mini.py @@ -202,3 +202,18 @@ def test_mini_island_argv_translates_to_a_miles_argv_on_cpu(tmp_path): assert miles[miles.index("--num-experts") + 1] == "32" and miles.count("--num-experts") == 1 assert miles[miles.index("--moe-layer-freq") + 1] == "[0,1,1,1]" assert miles[miles.index("--megatron-to-hf-mode") + 1] == "raw" + + +def test_align_weight_version_is_a_no_op_on_a_trainer_only_island(): + from types import SimpleNamespace + + from yeto.rl.adapters.miles.state_plugin import align_weight_version + + align_weight_version(SimpleNamespace(weight_updater=None, args=SimpleNamespace(debug_train_only=True)), + next_version=3) + with pytest.raises(RuntimeError, match="no weight_updater"): + align_weight_version(SimpleNamespace(weight_updater=None, args=SimpleNamespace(debug_train_only=False)), + next_version=3) + updater = SimpleNamespace(weight_version=0) + align_weight_version(SimpleNamespace(weight_updater=updater, args=None), next_version=3) + assert updater.weight_version == 2 diff --git a/yeto/rl/adapters/miles/publish.py b/yeto/rl/adapters/miles/publish.py index 29bece01b..990843ef3 100644 --- a/yeto/rl/adapters/miles/publish.py +++ b/yeto/rl/adapters/miles/publish.py @@ -533,6 +533,13 @@ async def _stamp_version(self, version: int, label: str) -> tuple[list[Any], lis async def _publish( self, version: int, token: str ) -> tuple[frozenset[str], dict[str, Any] | None]: + if getattr(self._args, "debug_train_only", False): + # Trainer-only island (Miles --debug-train-only, samples from the store): + # there is no engine to update or stamp. The release written by the + # wrapping ReleasingPublisher is the publication (Modal run-11 stopped at + # the engine stamp: "no rollout engines are eligible"). + self._published_once = True + return frozenset(), None update_weights = self._update_weights_to(version) offload_rollout = bool(getattr(self._args, "offload_rollout", False)) try: diff --git a/yeto/rl/adapters/miles/state_plugin.py b/yeto/rl/adapters/miles/state_plugin.py index ea72b0682..5125287ec 100644 --- a/yeto/rl/adapters/miles/state_plugin.py +++ b/yeto/rl/adapters/miles/state_plugin.py @@ -1828,6 +1828,11 @@ def align_weight_version(actor: Any, *, next_version: int) -> None: updater = getattr(actor, "weight_updater", None) if updater is None: + if getattr(getattr(actor, "args", None), "debug_train_only", False): + # trainer-only island (Miles --debug-train-only): no engine and no weight + # updater; update_weights returns early and the policy release carries + # the version (Modal run-11 failed here) + return raise RuntimeError("the trainer has no weight_updater to carry the policy version") updater.weight_version = int(next_version) - 1 From 0be12c3fc45f429b4d7394a27efcc461a29caa57 Mon Sep 17 00:00:00 2001 From: michael <78113554+chowhsu@users.noreply.github.com> Date: Tue, 20 Oct 2026 04:15:08 +0000 Subject: [PATCH 26/89] Resume a replaced spot sampler island and allow several samplers per trainer A pure sampler island on spot capacity can be reclaimed and started again in a new container (Modal reschedules the input; sky relaunches the cluster). The new incarnation now continues where the buffer says the old one stopped: - every index entry carries a sampler_round record (sampler id, incarnation, rollout id, data cursor after the round, last group of the round); - yeto.rl.engine.sampler_resume reads the resume point: the next rollout id after the newest round with an entry, and that round's data cursor; a round already in the buffer is never written again (RoundGuard), and no prompt is drawn twice (the Miles data source is put at the cursor); - on the reclaim marker (the Modal runner writes it for sampler islands) the island starts no new round, writes no more groups, writes rl_sampler_leave with reason PREEMPTED and a per-incarnation member record in the store, then waits for the cloud to stop the container; a new incarnation writes rl_sampler_join and loads the newest policy release before its first round. --rl-island-sample-role now takes one trainer and one or more samplers. With several samplers on a shared directory (a Modal Volume, where flock does not hold across containers) each sampler writes its own shard file and the trainer merges the shards into index.jsonl (one store commit per refresh); the data ends once every named shard has ended. The trainer scheduler is refused with several samplers (the head does not read the shards). Co-Authored-By: Claude Opus 5.5 --- tests/test_rl_oss_store.py | 2 +- tests/test_rl_sampler_trainer_pair_launch.py | 2 +- tests/test_rl_spot_sampler_replace.py | 376 ++++++++++++++++++ yeto/launch/rl_ports.py | 9 +- yeto/launcher.py | 3 +- yeto/modal_runner.py | 4 +- yeto/rl/adapters/miles/buffer_rollout.py | 115 +++++- .../rl/adapters/miles/buffer_rollout_miles.py | 39 +- yeto/rl/adapters/miles/island_entry.py | 6 +- yeto/rl/adapters/miles/sampler_island.py | 92 ++++- yeto/rl/engine/buffer_payload.py | 10 +- yeto/rl/engine/run_flags.py | 8 + yeto/rl/engine/sample_buffer.py | 118 ++++-- yeto/rl/engine/sample_roles.py | 62 ++- yeto/rl/engine/sampler_resume.py | 185 +++++++++ yeto/rl/engine/trainer_lifecycle.py | 9 +- 16 files changed, 969 insertions(+), 71 deletions(-) create mode 100644 tests/test_rl_spot_sampler_replace.py create mode 100644 yeto/rl/engine/sampler_resume.py diff --git a/tests/test_rl_oss_store.py b/tests/test_rl_oss_store.py index 7aab70d30..7f457ad63 100644 --- a/tests/test_rl_oss_store.py +++ b/tests/test_rl_oss_store.py @@ -156,7 +156,7 @@ def test_each_sampler_writes_its_own_shard_and_the_trainer_merges_them(): assert len(again) == 4 # a reader that does not merge (the head) sees the merged index only assert len(FileSampleBuffer(objects, batch_size=1)) == 4 - with pytest.raises(SampleBufferError, match="only the index of an object store merges"): + with pytest.raises(SampleBufferError, match="a shard writer does not merge"): FileSampleBuffer(objects, batch_size=1, index=shard_index("x"), merge_shards=True) diff --git a/tests/test_rl_sampler_trainer_pair_launch.py b/tests/test_rl_sampler_trainer_pair_launch.py index 4debf1e72..0cbd7e69c 100644 --- a/tests/test_rl_sampler_trainer_pair_launch.py +++ b/tests/test_rl_sampler_trainer_pair_launch.py @@ -134,7 +134,7 @@ def _replace(base, old: tuple, new: tuple) -> list[str]: (("--rl-single-island-no-sync",), (), "no-sync|single-island"), (("--tuning", "full"), ("--tuning", "full", "--rl-sample-role", "trainer"), "replaces --rl-sample-role"), - (("1:trainer",), ("1:sampler",), "one sampler and one trainer"), + (("1:trainer",), ("1:sampler",), "exactly one trainer"), (("--rl-max-policy-age", "2"), ("--rl-max-policy-age", "0"), "max-policy-age"), (("--rl-checkpoint-store", STORE, "--rl-cut-every", "1"), (), "checkpoint-store"), ]) diff --git a/tests/test_rl_spot_sampler_replace.py b/tests/test_rl_spot_sampler_replace.py new file mode 100644 index 000000000..92ff9b416 --- /dev/null +++ b/tests/test_rl_spot_sampler_replace.py @@ -0,0 +1,376 @@ +"""S25 C1: a sampler island on spot capacity is reclaimed and replaced. + +CPU only (no Miles, no Ray): the resume point read from the buffer, the +resumable round writer (no round written twice, no prompt drawn twice), several +sampler islands writing their own shards of one shared-directory buffer that +the trainer merges, the launch composition of one trainer and two samplers, and +the Modal runner's reclaim marker for sampler islands. +""" + +from __future__ import annotations + +import shlex +import signal as _signal +from types import SimpleNamespace + +import pytest + +from test_rl_algorithm_provenance import _fake_sky, _launcher_args, _no_modal_listing # noqa: F401 +from tests.buffer_rollout_procs import sample +from yeto.rl.adapters.miles import buffer_rollout as br +from yeto.rl.engine.policy_release import PolicyFollower, PolicyReleaseStore +from yeto.rl.engine.resume import CheckpointStore +from yeto.rl.engine.sample_buffer import FileSampleBuffer, SampleBufferError +from yeto.rl.engine.sampler_loop import SamplerLoop +from yeto.rl.engine.sampler_resume import ( + RECLAIM_MARKER_ENV, + Member, + SamplerResumeError, + SamplerRound, + reclaim_requested, + resume_point, + write_member, +) + + +def _group(g: int, version: int = 1) -> list[dict]: + return [sample(2 * g, g, version), sample(2 * g + 1, g, version)] + + +def _cursor(offset: int) -> dict[str, int]: + return {"sample_offset": offset, "epoch_id": 0, "sample_group_index": offset, "sample_index": 2 * offset} + + +def _round(rid: int, groups: int) -> list[tuple[str, list[dict]]]: + return [(f"r{rid}-g{g}", _group(g)) for g in range(groups)] + + +def _shard_env(store, sampler=None, shards=("island0", "island1")): + env = {br.BUFFER_ROOT_ENV: str(store), br.SHARDS_ENV: ",".join(shards)} + if sampler is not None: + env[br.SAMPLER_ID_ENV] = sampler + return env + + +# ---------------------------------------------------------------- resume point +def test_resume_point_of_an_empty_buffer_starts_at_round_zero(): + resume = resume_point([], "island0") + assert (resume.next_rollout_id, resume.data_cursor, resume.resumed) == (0, None, False) + + +def test_resume_point_continues_after_the_newest_written_round(tmp_path): + buffer = FileSampleBuffer(CheckpointStore(tmp_path), batch_size=1) + writer = br.SamplerRoundWriter(buffer, "island0", incarnation="a") + assert writer.write_round(0, _round(0, 2), _cursor(2)) == 2 + assert writer.write_round(1, _round(1, 2), _cursor(4)) == 2 + other = br.SamplerRoundWriter(buffer, "island9", incarnation="z") + other.write_round(5, _round(5, 1), _cursor(9)) # another sampler: not ours + resume = resume_point(buffer.entries(), "island0") + assert (resume.next_rollout_id, resume.data_cursor) == (2, _cursor(4)) + assert (resume.written_rounds, resume.partial_rounds, resume.groups) == ((0, 1), (), 4) + assert resume.incarnations == ("a",) + + +def test_a_round_cut_short_by_the_reclaim_keeps_its_groups_and_is_not_written_again(tmp_path): + """The island leaves while it writes round 1: the written group stays, the + rest are dropped; the replacement continues at round 2 from round 1's cursor + (round 1's prompts are not drawn again).""" + store = CheckpointStore(tmp_path) + leaving = {"now": False} + old = br.SamplerRoundWriter(FileSampleBuffer(store, batch_size=1), "island0", incarnation="a", + reclaimed=lambda: leaving["now"]) + old.write_round(0, _round(0, 3), _cursor(3)) + + class Leaving(list): + def __iter__(self): # the reclaim lands after the first group of round 1 + for i, item in enumerate(super().__iter__()): + leaving["now"] = i >= 1 + yield item + + assert old.write_round(1, Leaving(_round(1, 3)), _cursor(6)) == 1 + assert old.dropped_groups == 2 + new = br.SamplerRoundWriter(FileSampleBuffer(store, batch_size=1), "island0", incarnation="b") + assert (new.resume.next_rollout_id, new.resume.data_cursor) == (2, _cursor(6)) + assert new.resume.partial_rounds == (1,) + for rid in (0, 1): + with pytest.raises(SamplerResumeError, match="already in the sample buffer"): + new.write_round(rid, _round(rid, 1), _cursor(0)) + new.write_round(2, _round(2, 3), _cursor(9)) + ids = [e.sample_id for e in FileSampleBuffer(store, batch_size=1).entries()] + assert len(ids) == len(set(ids)) == 3 + 1 + 3 + assert [i.split("-")[1] for i in ids] == ["a"] * 4 + ["b"] * 3 + + +def test_a_round_still_generating_at_the_reclaim_is_generated_again(tmp_path): + store = CheckpointStore(tmp_path) + old = br.SamplerRoundWriter(FileSampleBuffer(store, batch_size=1), "island0", incarnation="a", + reclaimed=lambda: True) + assert old.write_round(0, _round(0, 2), _cursor(2)) == 0 # left before the write + new = br.SamplerRoundWriter(FileSampleBuffer(store, batch_size=1), "island0", incarnation="b") + assert (new.resume.next_rollout_id, new.resume.data_cursor) == (0, None) + assert new.write_round(0, _round(0, 2), _cursor(2)) == 2 + + +def test_resume_point_refuses_a_buffer_with_two_cursors_for_one_round(): + def entry(cursor, last): + return SimpleNamespace(sampler_round=SamplerRound("island0", "a", 0, cursor, last).record()) + + with pytest.raises(SamplerResumeError, match="two data cursors"): + resume_point([entry(_cursor(2), False), entry(_cursor(3), True)], "island0") + with pytest.raises(SamplerResumeError, match="two last groups"): + resume_point([entry(_cursor(2), True), entry(_cursor(2), True)], "island0") + + +def test_seek_puts_the_data_source_where_its_load_would_and_only_once(tmp_path): + class Dataset: + def __init__(self): + self.shuffled = [] + + def shuffle(self, epoch): + self.shuffled.append(epoch) + + source = SimpleNamespace(sample_offset=0, epoch_id=0, sample_group_index=0, sample_index=0, + dataset=Dataset()) + store = CheckpointStore(tmp_path) + br.SamplerRoundWriter(FileSampleBuffer(store, batch_size=1), "island0", incarnation="a" + ).write_round(0, _round(0, 1), {**_cursor(5), "epoch_id": 1}) + writer = br.SamplerRoundWriter(FileSampleBuffer(store, batch_size=1), "island0", incarnation="b") + assert writer.seek(SimpleNamespace(__self__=source), shuffle=True)["sample_offset"] == 5 + assert (source.sample_offset, source.epoch_id, source.sample_index) == (5, 1, 10) + assert source.dataset.shuffled == [1] + assert writer.seek(source, shuffle=True) is None # the second round: no seek + with pytest.raises(br.SampleBufferRolloutError, match="cannot seek"): + br.seek_data_source(SimpleNamespace(sample_offset=0), _cursor(1), shuffle=False) + + +def test_a_writer_without_resume_writes_as_before(tmp_path): + """A convertible node keeps its own rollout numbering: no round check.""" + store = CheckpointStore(tmp_path) + writer = br.SamplerRoundWriter(FileSampleBuffer(store, batch_size=1), "n0", incarnation="a", + resume=False) + writer.write_round(0, _round(0, 1), None) + writer.write_round(0, _round(0, 1), None) + assert writer.seek(object(), shuffle=True) is None + assert len(FileSampleBuffer(store, batch_size=1)) == 2 + + +# ---------------------------------------------------------------- two samplers, one trainer +def test_two_samplers_write_shards_the_trainer_merges_and_a_replacement_continues(tmp_path): + a = br.SamplerRoundWriter(br.open_buffer(_shard_env(tmp_path, "island0"), sampler=True), + "island0", incarnation="a1") + b = br.SamplerRoundWriter(br.open_buffer(_shard_env(tmp_path, "island1"), sampler=True), + "island1", incarnation="b1") + assert a.buffer.shard == "island0" and not (tmp_path / "sample-buffer/index.jsonl").exists() + trainer = br.open_buffer(_shard_env(tmp_path), batch_size=2, reader=True) + reader = br.BatchReader(trainer, max_policy_age=2, wait_s=0.0, poll_s=0.0) + a.write_round(0, _round(0, 2), _cursor(2)) + b.write_round(0, _round(0, 2), _cursor(2)) + taken, _ = reader.take(0, 4, current_version=1) + assert len(taken) == 4 + # island0 is reclaimed after one more round; its replacement continues at round 2 + a.write_round(1, _round(1, 2), _cursor(4)) + a2 = br.SamplerRoundWriter(br.open_buffer(_shard_env(tmp_path, "island0"), sampler=True), + "island0", incarnation="a2") + assert (a2.resume.next_rollout_id, a2.resume.incarnations) == (2, ("a1",)) + a2.write_round(2, _round(2, 2), _cursor(6)) + b.write_round(1, _round(1, 2), _cursor(4)) + taken2, _ = reader.take(1, 6, current_version=2) + ids = [e.sample_id for e in trainer.entries()] + assert len(taken2) == 6 and len(ids) == len(set(ids)) == 10 # none repeated, none lost + assert trainer.shard_positions(trainer.read_cursor) == {"island0": 6, "island1": 4} + # the data ends only when every named shard has ended + a2.buffer.mark_data_end() + trainer.refresh() + assert not trainer.data_ended + b.buffer.mark_data_end() + trainer.refresh() + assert trainer.data_ended + + +def test_the_data_does_not_end_before_a_named_shard_exists(tmp_path): + a = br.open_buffer(_shard_env(tmp_path, "island0"), sampler=True) + a.mark_data_end() + trainer = br.open_buffer(_shard_env(tmp_path), reader=True) + assert not trainer.data_ended # island1 has not written anything yet + + +def test_a_restarted_trainer_merges_no_record_twice(tmp_path): + a = br.SamplerRoundWriter(br.open_buffer(_shard_env(tmp_path, "island0"), sampler=True), + "island0", incarnation="a") + a.write_round(0, _round(0, 3), _cursor(3)) + assert len(br.open_buffer(_shard_env(tmp_path), reader=True)) == 3 + a.write_round(1, _round(1, 1), _cursor(4)) + again = br.open_buffer(_shard_env(tmp_path), reader=True) + assert len(again) == 4 + assert [r.get("shard_pos") for r in map(_json, _lines(tmp_path))] == [0, 1, 2, 3] + + +def _lines(root): + return (root / "sample-buffer/index.jsonl").read_text().splitlines() + + +def _json(line): + import json + + return json.loads(line) + + +def test_the_trainer_merges_one_refresh_with_one_store_commit(tmp_path): + class Counting(CheckpointStore): + commits = 0 + + def commit(self): + Counting.commits += 1 + + writer = br.SamplerRoundWriter(FileSampleBuffer(CheckpointStore(tmp_path), batch_size=1, + index="shards/island0.jsonl"), + "island0", incarnation="a") + writer.write_round(0, _round(0, 5), _cursor(5)) + Counting.commits = 0 + trainer = FileSampleBuffer(Counting(tmp_path), batch_size=1, merge_shards=True, + expected_shards=["island0"]) + assert len(trainer) == 5 and Counting.commits == 1 + + +def test_expected_shards_need_a_merging_buffer(tmp_path): + with pytest.raises(SampleBufferError, match="expected_shards needs merge_shards"): + FileSampleBuffer(CheckpointStore(tmp_path), batch_size=1, expected_shards=["a"]) + + +def test_island_env_names_the_shards_and_refuses_a_sampler_outside_them(): + env = br.island_env("sampler", store_root="/s", sampler_id="island1", shards=["island0", "island1"]) + assert env[br.SHARDS_ENV] == "island0,island1" and env[br.SAMPLER_ID_ENV] == "island1" + with pytest.raises(ValueError, match="not one of the buffer shards"): + br.island_env("sampler", store_root="/s", sampler_id="island2", shards=["island0"]) + assert br.SHARDS_ENV in br.ISLAND_ENVS and br.INCARNATION_ENV in br.ISLAND_ENVS + + +# ---------------------------------------------------------------- the sampler loop leaves +def test_the_sampler_loop_stops_on_the_reclaim_and_marks_no_data_end(tmp_path): + store = CheckpointStore(tmp_path) + marker = tmp_path / "reclaim" + env = {RECLAIM_MARKER_ENV: str(marker)} + buffer = FileSampleBuffer(store, batch_size=1) + writer = br.SamplerRoundWriter(buffer, "island0", incarnation="a", + reclaimed=lambda: reclaim_requested(env)) + events = [] + + def generate(rid): + writer.write_round(rid, _round(rid, 2), _cursor(2 * (rid + 1))) + if rid == 1: + marker.write_text("{}") # the container gets SIGINT during round 1 + + loop = SamplerLoop(follower=PolicyFollower(PolicyReleaseStore(store), lambda r, d: None), + buffer=buffer, generate=generate, num_rollouts=5, max_backlog_groups=100, + emit=lambda e, **f: events.append(e), poll_s=0.0, + should_stop=lambda: reclaim_requested(env)) + assert loop.run() == 2 + assert "rl_sampler_stopped" in events and not buffer.data_ended + resume = resume_point(FileSampleBuffer(store, batch_size=1).entries(), "island0") + assert resume.next_rollout_id == 2 + loop.should_stop = lambda: False + marker.unlink() + assert loop.run(start=resume.next_rollout_id) == 3 and buffer.data_ended + + +def test_member_records_are_one_file_per_incarnation(tmp_path): + store = CheckpointStore(tmp_path) + first = write_member(store, Member("island0", "a", 1.0, {"next_rollout_id": 0})) + member = Member("island0", "b", 2.0, {"next_rollout_id": 3}, leave_reason="PREEMPTED") + second = write_member(store, member) + assert first != second and first.exists() and "PREEMPTED" in second.read_text() + + +# ---------------------------------------------------------------- launch composition +GPU3 = "modal:1xh100,modal:1xh100,modal:1xh100" +STORE = "modal-volume://yeto-c1/run-a" +GROUP = ("--rl-single-island-no-sync", "--tuning", "lora", "--lora-r", "16", + "--rl-island-sample-role", "0:sampler", "--rl-island-sample-role", "1:sampler", + "--rl-island-sample-role", "2:trainer", "--rl-max-policy-age", "2", + "--rl-algorithm-spec", "examples/rl_algorithms/decoupled-ppo-c2.json", + "--rl-checkpoint-store", STORE, "--cluster-prefix", "c1-spot") + + +def _args(extra=()): + from yeto.launcher import _prepare_rl_args + + args = _launcher_args("ports", GROUP + tuple(extra), gpu=GPU3) + args.controller = "local" + _prepare_rl_args(args) + return args + + +def _tokens(command): + return shlex.split(command.split("-m yeto.rl.adapters.miles.island_entry", 1)[1]) + + +def _flag(tokens, flag): + return tokens[tokens.index(flag) + 1] if flag in tokens else None + + +def test_one_trainer_and_two_samplers_get_their_node_and_the_shards(): + from yeto.launcher import dry_run_plan + + plan = dry_run_plan(_args()) + roles = [i["sample_role"] for i in plan["island_requests"]] + assert roles == ["sampler", "sampler", "trainer"] and plan["syncer"] is None + tokens = [_tokens(i["learner_command"]) for i in plan["island_requests"]] + assert [_flag(t, "--rl-sampler-node") for t in tokens] == ["island0", "island1", None] + assert {_flag(t, "--rl-buffer-shards") for t in tokens} == {"island0,island1"} + + +def test_the_island_resolves_its_sampler_id_and_the_shards(): + from yeto.launcher import dry_run_plan + from yeto.rl.adapters.miles import island_entry + + plan = dry_run_plan(_args()) + sampler = island_entry.parse_args([{"$LEARNER_ID": "0"}.get(t, t) for t in + _tokens(plan["island_requests"][1]["learner_command"])]) + env = {} + island_entry.apply_sample_role(sampler, SimpleNamespace(), env) + assert (env[br.SAMPLER_ID_ENV], env[br.SHARDS_ENV]) == ("island1", "island0,island1") + + +def test_only_sampler_islands_get_the_reclaim_marker(): + from yeto.gpu_spec import parse_gpu_spec + from yeto.launch.modal_islands import build_modal_island_config + from yeto.launcher import _island_task, make_miles_island_task + + args = _args() + for m, spec in enumerate(parse_gpu_spec(args.gpu)): + task = _island_task(make_miles_island_task, args, spec, m, 3, "none", {}) + cfg = build_modal_island_config(args, spec, m, task, "none") + assert (RECLAIM_MARKER_ENV in cfg.envs) == (m < 2) + + +def test_two_samplers_refuse_the_trainer_scheduler(): + with pytest.raises(ValueError, match="several sampler islands"): + _args(("--rl-trainer-scheduler", "s1", "--rl-trainer-island", "2")) + + +def test_every_island_needs_a_role_and_one_trainer(): + from yeto.launcher import _prepare_rl_args + + two_trainers = tuple("1:trainer" if t == "1:sampler" else t for t in GROUP) + args = _launcher_args("ports", two_trainers, gpu=GPU3) + args.controller = "local" + with pytest.raises(ValueError, match="exactly one trainer"): + _prepare_rl_args(args) + + +def test_the_modal_runner_turns_the_signal_into_the_marker_on_a_sampler(monkeypatch): + from yeto import modal_runner as mr + + seen = [] + monkeypatch.setattr(mr, "_run_forwarding_to_learner", lambda cmd, env, **kw: seen.append(env) or 0) + + class Sig: + SIGINT, SIGTERM = _signal.SIGINT, _signal.SIGTERM + + @staticmethod + def signal(num, fn): + raise AssertionError("not expected") + + env = {RECLAIM_MARKER_ENV: "/tmp/yeto-reclaim-requested"} + assert mr.run_island_script(["x"], env, popen=lambda c, env, **k: None, signal_mod=Sig) == 0 + assert seen == [env] diff --git a/yeto/launch/rl_ports.py b/yeto/launch/rl_ports.py index c1f2b369c..bdeceda94 100644 --- a/yeto/launch/rl_ports.py +++ b/yeto/launch/rl_ports.py @@ -79,15 +79,16 @@ def _prepare_ports_algorithm(args) -> None: no_sync = bool(getattr(args, "rl_single_island_no_sync", False)) from ..rl.engine import sample_roles - # L5.4: a store-coupled sampler/trainer pair is two no-sync islands; a - # convertible pool (L7.5) is two or more (sample_roles holds their own rules). + # L5.4: a store-coupled pair (one trainer, one or more samplers: S25 C1) and a + # convertible pool (L7.5) are two or more no-sync islands (sample_roles holds + # their own rules). no_sync_islands = (islands if sample_roles.is_convertible_pool(args) - else 2 if sample_roles.is_pair(args) else 1) + or sample_roles.is_pair(args) else 1) if no_sync and (islands != no_sync_islands or getattr(args, "rl_sync_preset", "strict-avg") != "strict-avg" or getattr(args, "rl_initial_adapter", None)): raise AlgorithmSpecError( "--rl-single-island-no-sync needs exactly one island (one --gpu entry, no " - "external learners; two with --rl-island-sample-role), the default sync " + "external learners; one per island with --rl-island-sample-role), the default sync " "preset and no initial adapter" ) args.rl_allow_unverified_mechanism = list( diff --git a/yeto/launcher.py b/yeto/launcher.py index 7955bea99..38b5ac140 100644 --- a/yeto/launcher.py +++ b/yeto/launcher.py @@ -187,7 +187,8 @@ def _island_task(task_factory, args, spec, m, num_learners, syncer_addr, overrid learner_id, island_learners = sample_roles.island_learner(args, m, num_learners) task = task_factory(iargs, spec, learner_id, island_learners, syncer_addr) io.apply_to_task(task, {**io.island_env(args, m, overrides), - **(droppable.role_env(role, spec.cloud, spec.region) if role else {})}) + **(droppable.role_env(role, spec.cloud, spec.region) if role else {}), + **sample_roles.island_env(iargs)}) return task diff --git a/yeto/modal_runner.py b/yeto/modal_runner.py index e2b997379..b25be3c4f 100644 --- a/yeto/modal_runner.py +++ b/yeto/modal_runner.py @@ -925,7 +925,9 @@ def run_island_script(cmd: list[str], env: dict, *, popen=subprocess.Popen, sign sig = signal_mod or _signal if not _is_droppable(env): - if env.get("YETO_SPOT_INFLIGHT_SAVE_DIR"): + # in-flight save (4.1) and sampler islands (S25 C1, the launcher sets the marker + # path) both turn the reclaim signal into the marker file + if env.get("YETO_SPOT_INFLIGHT_SAVE_DIR") or env.get("YETO_SPOT_RECLAIM_MARKER"): return _run_forwarding_to_learner(cmd, env, popen=popen, sig=sig) return subprocess.call(cmd, env=env) if popen is subprocess.Popen else popen(cmd, env=env).wait() killpg = killpg or os.killpg diff --git a/yeto/rl/adapters/miles/buffer_rollout.py b/yeto/rl/adapters/miles/buffer_rollout.py index 2207badd9..6cc78c177 100644 --- a/yeto/rl/adapters/miles/buffer_rollout.py +++ b/yeto/rl/adapters/miles/buffer_rollout.py @@ -69,9 +69,14 @@ POLL_ENV = "YETO_SAMPLE_BUFFER_POLL_S" # reader: seconds between index reads while waiting MAX_POLICY_AGE_ENV = "YETO_SAMPLE_BUFFER_MAX_POLICY_AGE" SHARD_ENV = "YETO_SAMPLE_BUFFER_SHARD" # convertible node: its own shard of the index +# S25 C1: several sampler islands on one shared directory (a Modal Volume): each +# sampler writes its own shard, the trainer merges these shards (comma list) +SHARDS_ENV = "YETO_SAMPLE_BUFFER_SHARDS" +INCARNATION_ENV = "YETO_SAMPLER_INCARNATION" # set by the sampler island before Ray starts # the rollout functions run in Ray workers, which inherit the raylet's environment: # connect_island_ray forwards these to every worker -ISLAND_ENVS = (BUFFER_ROOT_ENV, SAMPLER_ID_ENV, WAIT_ENV, POLL_ENV, MAX_POLICY_AGE_ENV, SHARD_ENV) +ISLAND_ENVS = (BUFFER_ROOT_ENV, SAMPLER_ID_ENV, WAIT_ENV, POLL_ENV, MAX_POLICY_AGE_ENV, SHARD_ENV, + SHARDS_ENV, INCARNATION_ENV) DEFAULT_WAIT_S = 3600.0 DEFAULT_POLL_S = 2.0 @@ -111,10 +116,12 @@ def rollout_argv(role: str | None, *, max_policy_age: int, fully_async: bool = F def island_env(role: str, *, store_root: str, sampler_id: str | None = None, max_policy_age: int | None = None, wait_s: float | None = None, - shard: str | None = None) -> dict[str, str]: + shard: str | None = None, shards: Sequence[str] | None = None) -> dict[str, str]: """The environment the launch layer exports on the island (read in the rollout process). A convertible node writes like a sampler, into its own ``shard`` of - the index, and reads like a trainer.""" + the index, and reads like a trainer. ``shards``: the sampler shards of a + launch with several sampler islands (a sampler writes the one named by its + sampler id, the trainer merges them all).""" if role not in ROLES: raise ValueError(f"sample buffer role {role!r}: expected one of {ROLES}") env = {BUFFER_ROOT_ENV: str(store_root)} @@ -132,9 +139,20 @@ def island_env(role: str, *, store_root: str, sampler_id: str | None = None, env[MAX_POLICY_AGE_ENV] = str(int(max_policy_age)) if wait_s is not None: env[WAIT_ENV] = f"{float(wait_s):g}" + if shards: + if role == CONVERTIBLE: + raise ValueError("a convertible pool names its shards per node, not with shards") + if role == SAMPLER and sampler_id not in shards: + raise ValueError(f"sampler {sampler_id!r} is not one of the buffer shards {list(shards)}") + env[SHARDS_ENV] = ",".join(shards) return env +def buffer_shards(environ: Mapping[str, str]) -> tuple[str, ...]: + """The sampler shards of :data:`SHARDS_ENV` (() when unset).""" + return tuple(s.strip() for s in (environ.get(SHARDS_ENV) or "").split(",") if s.strip()) + + def open_buffer(environ: Mapping[str, str] | None = None, *, batch_size: int = 1, shard: str | None = None, create_index: bool = True, sampler: bool = False, reader: bool = False) -> Any: @@ -142,9 +160,11 @@ def open_buffer(environ: Mapping[str, str] | None = None, *, batch_size: int = 1 with ``shard``, the writer of that shard of the index; ``create_index=False`` for a reader whose index another host writes. - On a store that needs one shard per sampler island (``sample_buffer.needs_shards``): + On a store that needs one shard per sampler island (``sample_buffer.needs_shards``), + or with :data:`SHARDS_ENV` (several sampler islands on one shared directory): ``sampler`` opens the shard named by :data:`SAMPLER_ID_ENV`, and ``reader`` - (the trainer) merges the shards into the index.""" + (the trainer) merges the shards into the index (with :data:`SHARDS_ENV`, the + data ends once every named shard has ended).""" from yeto.rl.engine.resume import store_for from yeto.rl.engine.sample_buffer import INDEX, FileSampleBuffer, needs_shards, shard_index @@ -154,7 +174,8 @@ def open_buffer(environ: Mapping[str, str] | None = None, *, batch_size: int = 1 raise SampleBufferRolloutError(f"{BUFFER_ROOT_ENV} is not set") store = store_for(root, environ=env) merge = False - if needs_shards(store) and shard is None: + named = buffer_shards(env) + if (needs_shards(store) or named) and shard is None: if sampler: shard = env.get(SAMPLER_ID_ENV) if not shard: @@ -162,7 +183,8 @@ def open_buffer(environ: Mapping[str, str] | None = None, *, batch_size: int = 1 merge = reader return FileSampleBuffer(store, batch_size=batch_size, index=INDEX if shard is None else shard_index(shard), - create_index=create_index, merge_shards=merge) + create_index=create_index, merge_shards=merge, + expected_shards=named if merge and named else None) def renumber_group(samples: Sequence[dict[str, Any]], *, position: int, @@ -191,6 +213,85 @@ def renumber_group(samples: Sequence[dict[str, Any]], *, position: int, return out +# ---------------------------------------------------------------------- sampler side +_DATA_CURSOR_FIELDS = ("sample_offset", "epoch_id", "sample_group_index", "sample_index") + + +def seek_data_source(data_source: Any, cursor: Mapping[str, int], *, shuffle: bool) -> None: + """Put a Miles ``RolloutDataSource`` at ``cursor`` the way its ``load`` does + (the four position fields, then the epoch's shuffle). Raises when the source + does not have the fields (a source this resume does not know).""" + source = getattr(data_source, "__self__", data_source) + missing = [f for f in _DATA_CURSOR_FIELDS if not isinstance(getattr(source, f, None), int)] + if missing or set(cursor) != set(_DATA_CURSOR_FIELDS): + raise SampleBufferRolloutError( + f"cannot seek data source {type(source).__name__} to {dict(cursor)} " + f"(missing fields {missing})") + for name in _DATA_CURSOR_FIELDS: + setattr(source, name, int(cursor[name])) + if shuffle and getattr(source, "dataset", None) is not None: + source.dataset.shuffle(int(cursor["epoch_id"])) + + +class SamplerRoundWriter: + """The sampler's buffer writes, resumable across incarnations (S25 C1). + + Opened in the rollout process. On open it reads where this sampler's earlier + incarnations got to (:func:`~yeto.rl.engine.sampler_resume.resume_point`); + :meth:`seek` puts the data source at the resume cursor before the first + round; :meth:`write_round` writes one round's finished groups, each with its + :class:`~yeto.rl.engine.sampler_resume.SamplerRound` record. After the + reclaim marker appears (the island is leaving) nothing more is written: an + old incarnation must not write while its replacement writes the same shard.""" + + def __init__(self, buffer: Any, sampler_id: str, *, incarnation: str, resume: bool = True, + reclaimed: Callable[[], bool] = lambda: False) -> None: + """``resume=False``: write as before (a convertible node: no resume + point, no round check, no seek).""" + from yeto.rl.engine.sampler_resume import RoundGuard, resume_point + + self.buffer = buffer + self.sampler_id = str(sampler_id) + self.incarnation = str(incarnation) + self.reclaimed = reclaimed + self.resume = resume_point(buffer.entries() if resume else (), self.sampler_id) + self.guard = RoundGuard(self.resume) if resume else None + self._writer = GroupWriter(buffer, f"{self.sampler_id}-{self.incarnation}") + self._sought = False + self.dropped_groups = 0 # groups not written because the island is leaving + + def seek(self, data_source: Any, *, shuffle: bool) -> dict[str, int] | None: + """Before the first round only: the resume cursor applied (None: none).""" + if self._sought or self.guard is None: + return None + self._sought = True + cursor = self.resume.data_cursor + if cursor is not None: + seek_data_source(data_source, cursor, shuffle=shuffle) + return cursor + + def write_round(self, rollout_id: int, groups: Sequence[tuple[str, list[dict[str, Any]]]], + data_cursor: Mapping[str, int] | None) -> int: + """Write ``groups`` (``(group key, samples)``, aborted groups already left + out) of round ``rollout_id``; returns the groups written.""" + from yeto.rl.engine.sampler_resume import SamplerRound + + if self.guard is not None: + self.guard.check(rollout_id) + written = 0 + for i, (key, samples) in enumerate(groups): + if self.reclaimed(): + self.dropped_groups += len(groups) - i + break + record = SamplerRound(self.sampler_id, self.incarnation, int(rollout_id), + data_cursor=data_cursor, last=i == len(groups) - 1) + self._writer.write(samples, group_key=key, sampler_round=record.record()) + written += 1 + if written and self.guard is not None: + self.guard.written(rollout_id) + return written + + # ---------------------------------------------------------------------- trainer side @dataclass class TakeReport: diff --git a/yeto/rl/adapters/miles/buffer_rollout_miles.py b/yeto/rl/adapters/miles/buffer_rollout_miles.py index e9f6bad87..354023cb5 100644 --- a/yeto/rl/adapters/miles/buffer_rollout_miles.py +++ b/yeto/rl/adapters/miles/buffer_rollout_miles.py @@ -9,7 +9,6 @@ import asyncio import os -import uuid from typing import Any from miles.rollout.base_types import ( @@ -23,11 +22,13 @@ from miles.utils.types import Sample from yeto.rl.engine.buffer_payload import complete_prompt_tokens +from yeto.rl.engine.sampler_resume import new_incarnation, reclaim_requested from yeto.rl.engine.offpolicy_source import flat_metrics, offpolicy_source_event, payload_source_sample from .buffer_rollout import ( DEFAULT_POLL_S, DEFAULT_WAIT_S, + INCARNATION_ENV, MAX_POLICY_AGE_ENV, POLL_ENV, SAMPLER_ID_ENV, @@ -35,13 +36,14 @@ TRAINER, WAIT_ENV, BatchReader, - GroupWriter, SampleBufferRolloutError, + SamplerRoundWriter, convertible_mode, open_buffer, register_reader, renumber_group, ) +from .rollout_meta_hook import data_cursor def _flat(group: Any) -> list[Any]: @@ -52,32 +54,47 @@ def _flat(group: Any) -> list[Any]: class BufferWriterRolloutFn(InferenceRolloutFn): - """Sampler: Miles' rollout; every finished group is also written to the buffer.""" + """Sampler: Miles' rollout; every finished group is also written to the buffer. + + A replacement incarnation of the island continues where the buffer says the + earlier ones stopped (:class:`.buffer_rollout.SamplerRoundWriter`): the data + source is put at the saved cursor before the first round, and a round + already in the buffer is refused.""" + + resumable = True # a convertible node keeps its own rollout numbering (no resume) def __init__(self, input: RolloutFnConstructorInput) -> None: super().__init__(input) sampler_id = os.environ.get(SAMPLER_ID_ENV) if not sampler_id: raise SampleBufferRolloutError(f"{SAMPLER_ID_ENV} is not set") - # a per-process suffix: a restarted sampler repeats Miles group indexes, so the - # sample ids stay unique; the trainer renumbers groups by buffer position - self._writer = GroupWriter(self._open_writer_buffer(), f"{sampler_id}-{uuid.uuid4().hex[:8]}") + # the incarnation suffix keeps the sample ids unique: a restarted sampler repeats + # Miles group indexes; the trainer renumbers groups by buffer position + incarnation = os.environ.get(INCARNATION_ENV) or new_incarnation() + self._rounds = SamplerRoundWriter(self._open_writer_buffer(), sampler_id, + incarnation=incarnation, resume=self.resumable, + reclaimed=lambda: reclaim_requested(os.environ)) + self._shuffle = bool(getattr(input.args, "rollout_shuffle", False)) # the seek re-shuffles def _open_writer_buffer(self) -> Any: return open_buffer(sampler=True) async def _call_train(self, input: Any) -> Any: + self._rounds.seek(self.data_source, shuffle=self._shuffle) output = await super()._call_train(input) - await asyncio.to_thread(self._write, input.rollout_id, output.samples) + cursor, _ = data_cursor(self.data_source) + await asyncio.to_thread(self._write, input.rollout_id, output.samples, cursor) return output - def _write(self, rollout_id: int, groups: list[Any]) -> None: + def _write(self, rollout_id: int, groups: list[Any], cursor: dict[str, int] | None) -> None: + finished = [] for group in groups: samples = _flat(group) if any(s.status == Sample.Status.ABORTED for s in samples): continue # an aborted group is regenerated, never trained - key = f"r{int(rollout_id)}-g{samples[0].group_index}" - self._writer.write([s.to_dict() for s in samples], group_key=key) + finished.append((f"r{int(rollout_id)}-g{samples[0].group_index}", + [s.to_dict() for s in samples])) + self._rounds.write_round(rollout_id, finished, cursor) class _BufferRound: @@ -158,6 +175,8 @@ class ConvertibleRolloutFn(BufferWriterRolloutFn): :class:`BufferReaderRolloutFn`. The node switches the mode in this process (``buffer_rollout.executor_set_mode``) when it converts.""" + resumable = False + def __init__(self, input: RolloutFnConstructorInput) -> None: super().__init__(input) # the head merges the shards into the index: this node never writes it diff --git a/yeto/rl/adapters/miles/island_entry.py b/yeto/rl/adapters/miles/island_entry.py index 9a031e335..540069584 100644 --- a/yeto/rl/adapters/miles/island_entry.py +++ b/yeto/rl/adapters/miles/island_entry.py @@ -258,10 +258,12 @@ def apply_sample_role(args, miles_args, environ=None) -> None: config = sample_roles.resolve(args, groups_per_round=args.groups_per_round) miles_args.yeto_rl_sample_buffer = config node = getattr(args, "rl_convertible_node", None) # L7.5: a convertible node's own shard + # a pair island runs as learner 0: its sampler id is its node name (S25 C1) + sampler = node or getattr(args, "rl_sampler_node", None) or f"island{int(args.learner_id)}" (os.environ if environ is None else environ).update(island_env( - config.role, store_root=config.store, sampler_id=node or f"island{int(args.learner_id)}", + config.role, store_root=config.store, sampler_id=sampler, max_policy_age=int(getattr(args, "rl_max_policy_age", 0) or 0), wait_s=config.wait_s, - shard=node)) + shard=node, shards=getattr(args, "rl_buffer_shards", None))) def apply_ports_infra_switches(args, miles_args, environ=None) -> None: diff --git a/yeto/rl/adapters/miles/sampler_island.py b/yeto/rl/adapters/miles/sampler_island.py index 4c49446a0..4289c91b4 100644 --- a/yeto/rl/adapters/miles/sampler_island.py +++ b/yeto/rl/adapters/miles/sampler_island.py @@ -32,8 +32,10 @@ import asyncio import hashlib +import os import shutil import tempfile +import time from collections.abc import Callable from pathlib import Path from typing import Any @@ -236,16 +238,41 @@ async def _load(self, path: str | None, version: int) -> list[Any]: return engines -def run_sampler_island(miles_args: Any, *, environ: Any = None) -> int: - """Run the sampler island to its data end; returns the rounds generated.""" +LEAVE_HOLD_S = 300.0 # after a reclaim LEAVE: wait this long for the cloud to stop the container + + +def run_sampler_island(miles_args: Any, *, environ: Any = None, + sleep: Callable[[float], None] = time.sleep) -> int: + """Run the sampler island to its data end; returns the rounds generated. + + A replacement incarnation (S25 C1) continues after the newest round its + island has in the buffer (``sampler_resume``), loads the newest policy + release before its first round, and writes ``rl_sampler_join`` / a member + record. On the reclaim marker (Modal SIGINT/SIGTERM, written by the Modal + runner) it starts no new round, writes no more groups, writes + ``rl_sampler_leave`` with reason PREEMPTED and waits for the cloud to stop + the container (at most :data:`LEAVE_HOLD_S`).""" from miles.ray.placement_group import create_rollout_components from miles.utils.async_utils import Disposer from miles.utils.orchestration_utils import init_orchestration_script + from yeto.rl.elastic_client import LeaveReason from yeto.rl.engine.events import append_tape_event + from yeto.rl.engine.resume import store_for from yeto.rl.engine.sampler_loop import SamplerLoop + from yeto.rl.engine.sampler_resume import ( + JOIN_EVENT, + LEAVE_EVENT, + Member, + new_incarnation, + reclaim_marker, + reclaim_requested, + resume_point, + write_member, + ) from . import LoopRunner + from .buffer_rollout import INCARNATION_ENV, SAMPLER_ID_ENV from .entry import connect_island_ray from .rollout import RayMetadataSink, _release @@ -254,10 +281,23 @@ def run_sampler_island(miles_args: Any, *, environ: Any = None) -> int: raise SampleBufferRolloutError("run_sampler_island needs --rl-sample-role sampler") if not getattr(miles_args, "debug_rollout_only", False): raise SampleBufferRolloutError("a sampler island runs Miles with --debug-rollout-only") + env = os.environ if environ is None else environ def emit(event: str, **fields: Any) -> None: append_tape_event(miles_args, {"event": event, **fields}) + # one incarnation id for this process and its rollout workers (forwarded to Ray) + sampler_id = env.get(SAMPLER_ID_ENV) + if not sampler_id: + raise SampleBufferRolloutError(f"{SAMPLER_ID_ENV} is not set") + incarnation = env[INCARNATION_ENV] = new_incarnation() + reclaim_marker(env).unlink(missing_ok=True) # a restarted machine keeps its /tmp + buffer = open_buffer(env, sampler=True) + resume = resume_point(buffer.entries(), sampler_id) + store = store_for(config.store, environ=environ) + member = Member(sampler=sampler_id, incarnation=incarnation, joined_at=time.time(), + resume=resume.fields()) + connect_island_ray(miles_args=miles_args) # The rollout process's metadata hook pushes each round's group records to this # named actor (on a trainer island the driver creates it and consumes them). @@ -280,7 +320,16 @@ async def generate(rollout_id: int) -> None: _release(await executor.get(rollout_id)) # the rollout function wrote the groups metadata.discard() # the groups are in the buffer; nobody here reads the records + def leave(reason: LeaveReason, rounds: int, **detail: Any) -> None: + member.left_at, member.leave_reason, member.rounds = time.time(), reason.name, rounds + member.policy_version = follower.version + member.detail = detail + emit(LEAVE_EVENT, sampler=sampler_id, incarnation=incarnation, reason=reason.name, + rounds=rounds, policy_version=follower.version, buffer_end=len(buffer), **detail) + write_member(store, member) + error: BaseException | None = None + follower = None try: controller, executor = runner.run(init()) lora = lora_mode(miles_args) @@ -289,19 +338,41 @@ async def generate(rollout_id: int) -> None: loader.stamp(BASE_POLICY_VERSION) # until the first release is loaded follower = PolicyFollower(release_store(config, environ), loader, version=BASE_POLICY_VERSION) + reclaimed = lambda: reclaim_requested(env) # noqa: E731 loop = SamplerLoop( - follower=follower, buffer=open_buffer(environ, sampler=True), emit=emit, + follower=follower, buffer=buffer, emit=emit, generate=lambda rollout_id: runner.run(generate(rollout_id)), num_rollouts=int(miles_args.num_rollout), max_backlog_groups=int(config.max_backlog_groups), - min_round_s=config.min_round_s, require_release=lora, + 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}) + latest = follower.releases.latest() + emit(JOIN_EVENT, sampler=sampler_id, incarnation=incarnation, + latest_release=None if latest is None else {"seq": latest.seq, "policy_version": latest.version}, + **resume.fields()) + write_member(store, member) emit("rl_sampler_start", store=config.store, num_rollouts=int(miles_args.num_rollout), - max_backlog_groups=int(config.max_backlog_groups), weights=loader.weights) - return loop.run() + max_backlog_groups=int(config.max_backlog_groups), weights=loader.weights, + start_rollout_id=resume.next_rollout_id) + if buffer.data_ended: # an earlier incarnation finished and marked the data end + rounds = 0 + else: + rounds = loop.run(start=resume.next_rollout_id) + if reclaimed(): + leave(LeaveReason.PREEMPTED, rounds) + _hold_until_stopped(sleep, LEAVE_HOLD_S) + else: + leave(LeaveReason.FINISHED, rounds) + return rounds except BaseException as exc: error = exc + if follower is not None: + try: + leave(LeaveReason.PREEMPTED if reclaim_requested(env) else LeaveReason.INFRA_FAILED, + -1, error=f"{type(exc).__name__}: {exc}"[:300]) + except Exception: # noqa: BLE001 - the original error is the one to report + pass raise finally: try: @@ -311,6 +382,15 @@ async def generate(rollout_id: int) -> None: runner.close() +def _hold_until_stopped(sleep: Callable[[float], None], hold_s: float) -> None: + """After a reclaim LEAVE the process stays up until the cloud stops the + container: an early clean exit could end the cloud's input as finished + instead of letting the cloud run the island again.""" + print(f"[sampler-island] reclaimed: waiting up to {hold_s:g} s for the container to stop", + flush=True) + sleep(float(hold_s)) + + # ---------------------------------------------------------------------- trainer side class ReleasingPublisher: """A publisher that also releases each published policy for the samplers. diff --git a/yeto/rl/engine/buffer_payload.py b/yeto/rl/engine/buffer_payload.py index 23bb1f8c8..460831f37 100644 --- a/yeto/rl/engine/buffer_payload.py +++ b/yeto/rl/engine/buffer_payload.py @@ -142,14 +142,18 @@ def __init__(self, buffer: Any, sampler_id: str) -> None: self.sampler_id = str(sampler_id) def write(self, samples: Sequence[Mapping[str, Any]], *, group_key: str, - provenance: Mapping[str, str] | None = None) -> SampleIndexEntry: - """``provenance``: the sampler's source record for the index entry (S25 N2.4).""" + provenance: Mapping[str, str] | None = None, + sampler_round: Mapping[str, Any] | None = None) -> SampleIndexEntry: + """``provenance``: the sampler's source record for the index entry (S25 N2.4); + ``sampler_round``: the round record a replacement sampler resumes from + (``sampler_resume.SamplerRound.record()``, S25 C1).""" behavior_version, tokens = group_entry_fields(samples) sample_id = f"{self.sampler_id}-{group_key}" ref = self.buffer.write_payload(sample_id, encode_group(samples)) entry = SampleIndexEntry(sample_id=sample_id, behavior_version=behavior_version, payload_ref=ref, tokens=tokens, - provenance=None if provenance is None else dict(provenance)) + provenance=None if provenance is None else dict(provenance), + sampler_round=None if sampler_round is None else dict(sampler_round)) self.buffer.append(entry) return entry diff --git a/yeto/rl/engine/run_flags.py b/yeto/rl/engine/run_flags.py index 9b1314c3f..b4de91487 100644 --- a/yeto/rl/engine/run_flags.py +++ b/yeto/rl/engine/run_flags.py @@ -588,6 +588,14 @@ def rl_forward_pairs(*dests: str | tuple[str, str]) -> tuple[tuple[str, str], .. 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)"), + shared_flag("--rl-sampler-node", default=None, metavar="NAME", + help="--rl-sample-role sampler: this island's node name (set by the launcher " + "per island of --rl-island-sample-role; the sampler id that a replacement " + "incarnation of the island resumes under)"), + shared_flag("--rl-buffer-shards", default=None, metavar="A,B", + type=lambda text: [s for s in text.split(",") if s], + help="--rl-island-sample-role with several samplers: the sampler shards of the " + "buffer (set by the launcher); each sampler writes its own, the trainer merges them"), ) _SCHEDULING = ( diff --git a/yeto/rl/engine/sample_buffer.py b/yeto/rl/engine/sample_buffer.py index e4feb4023..4c945eb14 100644 --- a/yeto/rl/engine/sample_buffer.py +++ b/yeto/rl/engine/sample_buffer.py @@ -96,12 +96,15 @@ def payload_name(sample_id: str) -> str: class JsonlIndex: """The index as ``index.jsonl`` in a shared directory (see the module doc).""" - def __init__(self, store: CheckpointStore, *, index: str = INDEX, create_index: bool = True) -> None: + def __init__(self, store: CheckpointStore, *, index: str = INDEX, create_index: bool = True, + reload: bool = True) -> None: """``index``: the index file under ``sample-buffer/`` (a shard writer passes :func:`shard_index`). ``create_index=False``: a reader that must not write the index, not even an empty one (on a store mount whose commit would - replace the merger's index).""" + replace the merger's index). ``reload=False``: :meth:`read_new` does not + reload the store (a merger's shard reader: the merger reloaded it once).""" self.store = store + self._reload = bool(reload) self.root = Path(store.root) / BUFFER_DIR (self.root / PAYLOADS).mkdir(parents=True, exist_ok=True) self.path = self.root / index @@ -128,22 +131,30 @@ def has_payload(self, ref: str) -> bool: return (self.root / ref).is_file() def read_new(self) -> list[dict[str, Any]]: - self.store.reload() - if not self.path.exists(): # a reader before the first merge + if self._reload: + self.store.reload() + if not self.path.exists(): # a reader before the first merge (or a shard not created yet) return [] with open(self.path, "rb") as fh: return list(self._lines(fh)) def append(self, absorb: Callable[[list[dict[str, Any]]], None], make: Callable[[], dict[str, Any]]) -> None: + self.append_many(absorb, lambda: [make()]) + + def append_many(self, absorb: Callable[[list[dict[str, Any]]], None], + make: Callable[[], list[dict[str, Any]]]) -> None: + """Append ``make()``'s records with one write, one fsync and one store + commit (a merger appends a whole refresh at once: one Modal Volume commit + per record would cost seconds each).""" with open(self.path, "r+b") as fh: fcntl.flock(fh.fileno(), fcntl.LOCK_EX) try: absorb(list(self._lines(fh))) - record = make() + records = make() fh.truncate(self._offset) # drop a torn last line fh.seek(self._offset) - fh.write((json.dumps(record, sort_keys=True) + "\n").encode()) + fh.write(b"".join((json.dumps(r, sort_keys=True) + "\n").encode() for r in records)) fh.flush() os.fsync(fh.fileno()) absorb(list(self._lines(fh))) @@ -151,6 +162,27 @@ def append(self, absorb: Callable[[list[dict[str, Any]]], None], fcntl.flock(fh.fileno(), fcntl.LOCK_UN) self.store.commit() + # ------------------------------------------------------------ shards (merger side) + def shard_names(self) -> list[str]: + """The shard index files under ``sample-buffer/shards/``, sorted.""" + return sorted(p.name[:-len(".jsonl")] for p in (self.root / SHARDS).glob("*.jsonl")) + + def shard_reader(self, shard: str, start: int) -> "JsonlIndex": + """A reader of one shard's index file that skips its first ``start`` + records (merged before).""" + reader = JsonlIndex(self.store, index=shard_index(shard), create_index=False, reload=False) + if start: + skipped = 0 + if reader.path.exists(): + with open(reader.path, "rb") as fh: + for _ in reader._lines(fh): + skipped += 1 + if skipped == start: + break + if skipped < start: + raise SampleBufferError(f"shard {shard}: {skipped} records, but {start} are merged") + return reader + def _lines(self, fh) -> Iterator[dict[str, Any]]: fh.seek(self._offset) data = fh.read() @@ -223,6 +255,13 @@ def append(self, absorb: Callable[[list[dict[str, Any]]], None], return raise SampleBufferError(f"{self.prefix}: lost {_MAX_CREATE_RACES} appends in a row to other writers") + # ------------------------------------------------------------ shards (merger side) + def shard_names(self) -> list[str]: + return shard_names(self.objects) + + def shard_reader(self, shard: str, start: int) -> "ObjectIndex": + return ObjectIndex(self.objects, prefix=shard_prefix(shard), start=start) + def _objects_of(store: CheckpointStore | ObjectStore) -> ObjectStore | None: return store.objects if isinstance(store, CheckpointStore) else store @@ -249,7 +288,8 @@ def _entry(record: dict[str, Any]) -> SampleIndexEntry: return SampleIndexEntry( sample_id=record["sample_id"], behavior_version=int(record["behavior_version"]), payload_ref=record["payload_ref"], tokens=int(record["tokens"]), - appended_at=float(record["appended_at"]), provenance=record.get("provenance")) + appended_at=float(record["appended_at"]), provenance=record.get("provenance"), + sampler_round=record.get("sampler_round")) class UnmergedShards: @@ -321,21 +361,31 @@ class FileSampleBuffer: def __init__(self, store: CheckpointStore | ObjectStore, *, batch_size: int, arrival_window_s: float = 60.0, clock: Callable[[], float] = time.time, - index: str = INDEX, create_index: bool = True, merge_shards: bool = False) -> None: + index: str = INDEX, create_index: bool = True, merge_shards: bool = False, + expected_shards: Iterable[str] | None = None) -> None: """``index``/``create_index``: see :class:`JsonlIndex`. ``merge_shards``: - this buffer (the consumer) merges the sampler shards of an object store - into the index on every refresh (see the module doc); it must be the - only one that does.""" + this buffer (the consumer) merges the sampler shards into the index on + every refresh (see the module doc); it must be the only one that does. + On a shared directory the shards are ``shards/.jsonl`` files (one + sampler island writes each; S25 C1: several sampler islands on one + Modal Volume, where ``flock`` does not hold across containers). + ``expected_shards``: the shards that must all end before the merged index + ends (None: the shards seen so far), so an early-ending sampler cannot + end the data before a slower one has written its first record.""" if batch_size < 1: raise ValueError("batch_size must be positive") self.store = store self.index = index_for(store, index=index, create_index=create_index) self.shard = _shard_of_index(index) if index != INDEX else None # the shard this buffer writes - if merge_shards and (self.shard is not None or not isinstance(self.index, ObjectIndex)): - raise SampleBufferError("only the index of an object store merges shards") + if merge_shards and self.shard is not None: + raise SampleBufferError("a shard writer does not merge shards (only the consumer's index does)") + if expected_shards is not None and not merge_shards: + raise SampleBufferError("expected_shards needs merge_shards") self._merge = merge_shards + self._expected_shards = (None if expected_shards is None + else frozenset(_shard_name(s) for s in expected_shards)) self._shard_of: list[str | None] = [] # the shard of each position (merged records) - self._shard_readers: dict[str, ObjectIndex] = {} + self._shard_readers: dict[str, JsonlIndex | ObjectIndex] = {} self._shards_ended: set[str] = set() self.batch_size = int(batch_size) self.arrival_window_s = float(arrival_window_s) @@ -368,6 +418,8 @@ def append(self, entry: SampleIndexEntry) -> None: entry = dataclasses.replace(entry, appended_at=self.clock()) # ``provenance`` is written only when set, so a line without it reads as before. source = {} if entry.provenance is None else {"provenance": dict(entry.provenance)} + if entry.sampler_round is not None: + source["sampler_round"] = dict(entry.sampler_round) self._write_record(lambda pos: {"schema": SCHEMA, "pos": pos, "sample_id": entry.sample_id, "behavior_version": int(entry.behavior_version), "payload_ref": entry.payload_ref, "tokens": int(entry.tokens), @@ -386,6 +438,22 @@ def record() -> dict: with self._lock: self.index.append(self._absorb, record) + def _write_records(self, records: list[dict[str, Any]]) -> None: + """Append ``records`` (without ``pos``) at the next positions: one write on + a shared directory, one object per record on an object store.""" + def make() -> list[dict]: + if self._ended: + raise SampleBufferError("append after the data end") + start = len(self._entries) + return [{**record, "pos": start + i} for i, record in enumerate(records)] + + with self._lock: + if isinstance(self.index, JsonlIndex): + self.index.append_many(self._absorb, make) + return + for record in records: # an object store: one create-if-absent per record + self._write_record(lambda pos, record=record: {**record, "pos": pos}) + def _absorb(self, records: list[dict[str, Any]]) -> None: for record in records: if record.get("schema") != SCHEMA or int(record["pos"]) != len(self._entries): @@ -407,14 +475,12 @@ def refresh(self) -> int: return len(self._entries) def _merge_shards(self) -> None: - objects = self.index.objects new: list[list[dict[str, Any]]] = [] # per shard, in shard order - for shard in shard_names(objects): + for shard in self.index.shard_names(): reader = self._shard_readers.get(shard) if reader is None: # first seen here: skip what an earlier merger merged done = sum(1 for s in self._shard_of if s == shard) - reader = self._shard_readers[shard] = ObjectIndex(objects, prefix=shard_prefix(shard), - start=done) + reader = self._shard_readers[shard] = self.index.shard_reader(shard, done) records = [] for record in reader.read_new(): if record.get("end"): @@ -430,11 +496,12 @@ def _merge_shards(self) -> None: # records of a shard must stay a prefix of it, because a restarted merger # resumes each shard after its merged count (a clock step back would # otherwise merge a later record first, then skip the earlier one). - merged = heapq.merge(*new, key=lambda r: (float(r.get("appended_at", 0.0)), r["shard"])) - for record in merged: - self._write_record(lambda pos, record=record: {**record, "pos": pos}) - if self._shard_readers and self._shards_ended >= set(self._shard_readers): - self._write_record(lambda pos: {"schema": SCHEMA, "pos": pos, "end": True}) + merged = list(heapq.merge(*new, key=lambda r: (float(r.get("appended_at", 0.0)), r["shard"]))) + expected = set(self._shard_readers) if self._expected_shards is None else self._expected_shards + if expected and self._shards_ended >= expected: + merged.append({"schema": SCHEMA, "end": True}) + if merged: + self._write_records(merged) def shard_positions(self, cursor: int) -> dict[str, int]: """{shard: records of that shard below merged position ``cursor``}: the @@ -449,6 +516,11 @@ def __len__(self) -> int: with self._lock: return len(self._entries) + def entries(self) -> list[SampleIndexEntry]: + """The entries read so far, in buffer order (a copy).""" + with self._lock: + return list(self._entries) + @property def data_ended(self) -> bool: with self._lock: diff --git a/yeto/rl/engine/sample_roles.py b/yeto/rl/engine/sample_roles.py index 10e68c357..e9b76b2fd 100644 --- a/yeto/rl/engine/sample_roles.py +++ b/yeto/rl/engine/sample_roles.py @@ -18,8 +18,9 @@ The sampler and the trainer mount the same checkpoint store (the buffer and the releases live there). They are either separate single-island launches (``--rl-sample-role``), or one launch of two islands, a *store-coupled pair* -(``--rl-island-sample-role 0:sampler --rl-island-sample-role 1:trainer``, L5.4): -there is no syncer, and each island runs as its own one-learner, no-sync island +(``--rl-island-sample-role 0:sampler --rl-island-sample-role 1:trainer``, L5.4), +or one launch of one trainer and several sampler islands (S25 C1: each sampler +writes its own shard of the buffer, :func:`buffer_shards`): there is no syncer, and each island runs as its own one-learner, no-sync island (learner id 0 of 1) whose role comes from its island args (:func:`island_args`). This module is engine-neutral; the Miles mapping is in ``adapters.miles.buffer_rollout``. @@ -113,10 +114,24 @@ def island_roles(args: Any, num_islands: int) -> dict[int, str]: def is_pair(args: Any) -> bool: - """Whether the launch is a store-coupled sampler/trainer pair.""" + """Whether the launch is a store-coupled group: one trainer island and one or + more sampler islands (``--rl-island-sample-role``; "pair" for short).""" return bool(getattr(args, "rl_island_sample_role", None)) +def sampler_islands(args: Any, num_islands: int) -> list[int]: + """The sampler islands of a pair, sorted.""" + return sorted(m for m, role in island_roles(args, num_islands).items() if role == SAMPLER) + + +def buffer_shards(args: Any, num_islands: int) -> list[str] | None: + """The buffer shards of a pair with several sampler islands (S25 C1): one per + sampler island, named by its node name; None with one sampler (it writes the + index itself, as before).""" + samplers = sampler_islands(args, num_islands) + return [node_name(m) for m in samplers] if len(samplers) > 1 else None + + def is_store_coupled(args: Any) -> bool: """A pair or a convertible pool: no syncer, each island is its own one-learner no-sync island over the shared checkpoint store.""" @@ -133,6 +148,9 @@ def island_args(args: Any, island: int, num_islands: int) -> Any: return args out = copy.copy(args) out.rl_sample_role = roles[island] + if out.rl_sample_role in (SAMPLER, TRAINER): + out.rl_sampler_node = node_name(island) if out.rl_sample_role == SAMPLER else None + out.rl_buffer_shards = buffer_shards(args, num_islands) if out.rl_sample_role == CONVERTIBLE: out.rl_convertible_node = node_name(island) out.rl_sampler_min_round_seconds = None # a node keeps its backlog cap @@ -161,17 +179,23 @@ def buffer_full_parameter(args: Any) -> bool: def check_pair(args: Any, num_islands: int) -> None: - """Launch checks of a store-coupled pair (nothing outside a pair): exactly two - islands, one sampler and one trainer, ``--rl-single-island-no-sync`` (no - syncer), no ``--rl-sample-role`` and no external learner; then every island's - own role checks on its island args.""" + """Launch checks of a store-coupled pair (nothing outside a pair): every + island has a role, exactly one trainer and at least one sampler, + ``--rl-single-island-no-sync`` (no syncer), no ``--rl-sample-role`` and no + external learner; several samplers take no trainer scheduler (the head reads + no sampler shard of a shared directory); then every island's own role checks + on its island args.""" roles = island_roles(args, num_islands) if not roles: return - if num_islands != 2 or sorted(roles.values()) != [SAMPLER, TRAINER] or len(roles) != 2: - raise ValueError("--rl-island-sample-role: a pair is two --gpu islands, one " - f"{SAMPLER} and one {TRAINER} (got {num_islands} island(s), roles " - f"{dict(sorted(roles.items()))})") + values = list(roles.values()) + if len(roles) != num_islands or values.count(TRAINER) != 1 or values.count(SAMPLER) < 1: + raise ValueError("--rl-island-sample-role: a pair gives every --gpu island a role, " + f"at least one {SAMPLER} and exactly one {TRAINER} (got {num_islands} " + f"island(s), roles {dict(sorted(roles.items()))})") + if values.count(SAMPLER) > 1 and getattr(args, "rl_trainer_scheduler", None) not in (None, "off"): + raise ValueError("--rl-trainer-scheduler with several sampler islands is not supported " + "(the head does not read the sampler shards); drop the scheduler") if sample_role(args) != BOTH: raise ValueError("--rl-island-sample-role replaces --rl-sample-role; give only one") if int(getattr(args, "external_learners", 0) or 0): @@ -288,6 +312,18 @@ def resolve(args: Any, *, groups_per_round: int) -> SampleBufferConfig | None: else None)) +def island_env(args: Any) -> dict[str, str]: + """The environment a pair's sampler island gets from the launcher ({} for + every other island; ``args`` are the island args): the reclaim marker path, + so the Modal runner turns the reclaim signal into that file and the sampler + leaves without being signalled (S25 C1).""" + if sample_role(args) != SAMPLER or not getattr(args, "rl_sampler_node", None): + return {} + from .sampler_resume import DEFAULT_RECLAIM_MARKER, RECLAIM_MARKER_ENV + + return {RECLAIM_MARKER_ENV: DEFAULT_RECLAIM_MARKER} + + def island_flags(args: Any) -> str: """The learner flags the launcher forwards (empty for ``both``); ``args`` are the island args (:func:`island_args`).""" @@ -303,4 +339,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 getattr(args, "rl_sampler_node", None): + flags += f" --rl-sampler-node {args.rl_sampler_node}" + if getattr(args, "rl_buffer_shards", None): + flags += f" --rl-buffer-shards {','.join(args.rl_buffer_shards)}" return flags diff --git a/yeto/rl/engine/sampler_resume.py b/yeto/rl/engine/sampler_resume.py new file mode 100644 index 000000000..e25eefda9 --- /dev/null +++ b/yeto/rl/engine/sampler_resume.py @@ -0,0 +1,185 @@ +"""Sampler island replacement: where a new incarnation continues (S25 C1). + +A sampler island on spot capacity can be reclaimed at any time; the cloud (or +the launcher) then starts a new incarnation of the same island. The island's +sample-buffer writes must neither repeat nor lose groups, and the trainer keeps +training on what the buffer holds. + +Each index entry a sampler writes carries a :class:`SamplerRound` record +(``SampleIndexEntry.sampler_round``): the sampler id (the island, stable across +incarnations), the incarnation, the rollout id, the data cursor after the round +drew its prompts, and ``last`` on the round's last group. A Miles round writes +its groups after the whole round is generated, so the buffer itself says how +far the island got: + +* a round with any entry in the buffer is never written again + (:class:`RoundGuard`); the new incarnation continues at the next rollout id + from that round's data cursor, so no prompt is drawn twice; +* a round without its ``last`` entry (reclaimed while it wrote) keeps the + groups it wrote (each whole and written once); its unwritten groups are lost + (``partial_rounds`` counts such rounds); +* a round still generating when the island was reclaimed wrote nothing: the new + incarnation generates that rollout id again from the cursor before it. + +:func:`resume_point` reads the entries; the island loop starts at +``next_rollout_id`` and the rollout function seeks the data source to +``data_cursor``. The incarnation's arrival and departure are written as +``rl_sampler_join`` / ``rl_sampler_leave`` (reason: a ``LeaveReason`` name) and as +one member record per incarnation in the store (:func:`write_member`), which +outlives the island's own tape. Engine-neutral. +""" + +from __future__ import annotations + +import json +import uuid +from collections.abc import Iterable, Mapping +from dataclasses import asdict, dataclass, field +from pathlib import Path +from typing import Any + +JOIN_EVENT = "rl_sampler_join" +LEAVE_EVENT = "rl_sampler_leave" +DROPPED_EVENT = "rl_sampler_groups_dropped" # groups not written after the leave +MEMBERS_DIR = "sample-buffer/members" # one record per sampler incarnation, under the store root +RECLAIM_MARKER_ENV = "YETO_SPOT_RECLAIM_MARKER" # written by the Modal runner on the reclaim signal +DEFAULT_RECLAIM_MARKER = "/tmp/yeto-reclaim-requested" + + +class SamplerResumeError(RuntimeError): + pass + + +def new_incarnation() -> str: + """A fresh incarnation id (unique per island process).""" + return uuid.uuid4().hex[:8] + + +@dataclass(frozen=True) +class SamplerRound: + """The ``sampler_round`` record of one index entry.""" + + sampler: str + incarnation: str + rollout_id: int + data_cursor: Mapping[str, int] | None = None # the data source after the round drew its prompts + last: bool = False # the round's last group + + def record(self) -> dict[str, Any]: + out = {"sampler": self.sampler, "incarnation": self.incarnation, + "rollout_id": int(self.rollout_id), "last": bool(self.last)} + if self.data_cursor is not None: + out["data_cursor"] = {k: int(v) for k, v in self.data_cursor.items()} + return out + + +@dataclass(frozen=True) +class SamplerResume: + """Where the next incarnation of a sampler continues.""" + + next_rollout_id: int # after the newest round with an entry in the buffer + data_cursor: dict[str, int] | None # that round's data cursor (None: start of the data) + written_rounds: tuple[int, ...] # rounds with at least one entry: never written again + partial_rounds: tuple[int, ...] # written rounds without their last group + incarnations: tuple[str, ...] # earlier incarnations, in the order they first wrote + groups: int # every group this sampler has in the buffer + + @property + def resumed(self) -> bool: + return bool(self.incarnations) + + def fields(self) -> dict[str, Any]: + out = asdict(self) + out["written_rounds"] = len(self.written_rounds) + out["partial_rounds"] = list(self.partial_rounds) + out["incarnations"] = list(self.incarnations) + return out + + +def resume_point(entries: Iterable[Any], sampler: str) -> SamplerResume: + """Read the ``sampler_round`` records of ``sampler``'s entries (other samplers' + entries and entries without the record are skipped).""" + cursors: dict[int, dict[str, int] | None] = {} + ended: set[int] = set() + incarnations: list[str] = [] + groups = 0 + for entry in entries: + record = getattr(entry, "sampler_round", None) + if not record or record.get("sampler") != sampler: + continue + groups += 1 + rid = int(record["rollout_id"]) + raw = record.get("data_cursor") + cursor = None if raw is None else {k: int(v) for k, v in raw.items()} + if rid in cursors and cursors[rid] != cursor: + raise SamplerResumeError(f"sampler {sampler}: round {rid} has two data cursors " + f"({cursors[rid]} and {cursor})") + cursors[rid] = cursor + if record.get("incarnation") not in incarnations: + incarnations.append(str(record.get("incarnation"))) + if record.get("last"): + if rid in ended: + raise SamplerResumeError(f"sampler {sampler}: round {rid} has two last groups") + ended.add(rid) + newest = max(cursors) if cursors else None + return SamplerResume(next_rollout_id=0 if newest is None else newest + 1, + data_cursor=None if newest is None else cursors[newest], + written_rounds=tuple(sorted(cursors)), + partial_rounds=tuple(sorted(set(cursors) - ended)), + incarnations=tuple(incarnations), groups=groups) + + +class RoundGuard: + """Write-side check: a round with entries in the buffer is never written by + another incarnation, and an incarnation never writes before its first round.""" + + def __init__(self, resume: SamplerResume) -> None: + self._written = set(resume.written_rounds) + self._first = int(resume.next_rollout_id) + + def check(self, rollout_id: int) -> None: + rid = int(rollout_id) + if rid in self._written or rid < self._first: + raise SamplerResumeError( + f"round {rid} is already in the sample buffer (this incarnation starts at " + f"round {self._first}); a replacement sampler must not write it again") + + def written(self, rollout_id: int) -> None: + self._written.add(int(rollout_id)) + + +def reclaim_marker(environ: Mapping[str, str]) -> Path: + return Path(environ.get(RECLAIM_MARKER_ENV) or DEFAULT_RECLAIM_MARKER) + + +def reclaim_requested(environ: Mapping[str, str]) -> bool: + """The Modal runner wrote the reclaim marker (the container got SIGINT/SIGTERM).""" + return reclaim_marker(environ).exists() + + +@dataclass +class Member: + """One sampler incarnation's record in the store (``members//.json``).""" + + sampler: str + incarnation: str + joined_at: float + resume: dict[str, Any] + policy_version: int | None = None + left_at: float | None = None + leave_reason: str | None = None + rounds: int | None = None # rounds this incarnation completed + detail: dict[str, Any] = field(default_factory=dict) + + +def write_member(store: Any, member: Member) -> Path: + """Write (or rewrite) the incarnation's member record and commit the store. + One file per incarnation: an old and a new incarnation never write the same + file (a Modal Volume keeps the last commit of a file).""" + path = Path(store.root) / MEMBERS_DIR / member.sampler / f"{member.incarnation}.json" + path.parent.mkdir(parents=True, exist_ok=True) + tmp = path.with_suffix(".json.tmp") + tmp.write_text(json.dumps(asdict(member), sort_keys=True) + "\n") + tmp.replace(path) + store.commit() + return path diff --git a/yeto/rl/engine/trainer_lifecycle.py b/yeto/rl/engine/trainer_lifecycle.py index 731ecf418..ec3416aee 100644 --- a/yeto/rl/engine/trainer_lifecycle.py +++ b/yeto/rl/engine/trainer_lifecycle.py @@ -139,7 +139,13 @@ class SampleIndexEntry: record: device family, card, engine and version, precision, quantization scheme, topology); None for entries written before it existed or by a sampler that does not set it (read as ``unknown``). Data only: it never - enters a hash or the version gate (rl-hetero-sampler-engines D3).""" + enters a hash or the version gate (rl-hetero-sampler-engines D3). + + ``sampler_round``: which sampler round wrote the entry + (``sampler_resume.SamplerRound`` record: sampler, incarnation, rollout id, + last group of the round, data cursor after the round); None for entries + written before it existed. A replacement sampler island reads it to resume + after the last complete round (S25 C1).""" sample_id: str behavior_version: int @@ -147,6 +153,7 @@ class SampleIndexEntry: tokens: int = 0 appended_at: float = 0.0 provenance: Mapping[str, str] | None = field(default=None, compare=False, hash=False) + sampler_round: Mapping[str, Any] | None = field(default=None, compare=False, hash=False) @dataclass(frozen=True) From 79130ce75f16249803de704abe7de4ecdaaafc5f Mon Sep 17 00:00:00 2001 From: michael <78113554+chowhsu@users.noreply.github.com> Date: Tue, 20 Oct 2026 04:15:16 +0000 Subject: [PATCH 27/89] verl: trainer island over the sample buffer (S25 C2) A verl trainer island now trains from groups a separate sampler island wrote to the durable sample buffer, and releases each policy to the store for the samplers. Only the verl adapter layer changes; verl is not patched. - buffer_feed: the fully_async rollout port's queue calls served from the buffer (versions mapped through the island's VersionMap, so the RoundCollector keeps its max_policy_age judgement); kept groups become padded verl rows; every pulled buffer position is logged once. - buffer_ports: BufferReleasePublisher merges the LoRA into the base and releases full weights (the engine sampler loads full checkpoints); the base policy v0 (B all zero) is not released. - merged_release: W + B@A in float32, stored in the base dtype, same file layout. - buffer_trainer_runner / verl_main.main_buffer_trainer: trainer-only set-up (no rollouter, no MessageQueue) with the fully_async trainer actor plus yeto_add_group / yeto_param_version. - island_entry: --rl-sample-role trainer plans the buffer trainer; other roles are refused (the sampler is the engine sampler island). Co-Authored-By: Claude Opus 5.5 --- tests/test_rl_verl_buffer_trainer.py | 358 ++++++++++++++++++ yeto/rl/adapters/verl/buffer_feed.py | 274 ++++++++++++++ yeto/rl/adapters/verl/buffer_ports.py | 148 ++++++++ .../rl/adapters/verl/buffer_trainer_runner.py | 143 +++++++ yeto/rl/adapters/verl/fully_async_runner.py | 18 +- yeto/rl/adapters/verl/island_entry.py | 27 +- yeto/rl/adapters/verl/merged_release.py | 112 ++++++ yeto/rl/adapters/verl/trainer.py | 22 +- yeto/rl/adapters/verl/verl_main.py | 22 ++ 9 files changed, 1112 insertions(+), 12 deletions(-) create mode 100644 tests/test_rl_verl_buffer_trainer.py create mode 100644 yeto/rl/adapters/verl/buffer_feed.py create mode 100644 yeto/rl/adapters/verl/buffer_ports.py create mode 100644 yeto/rl/adapters/verl/buffer_trainer_runner.py create mode 100644 yeto/rl/adapters/verl/merged_release.py diff --git a/tests/test_rl_verl_buffer_trainer.py b/tests/test_rl_verl_buffer_trainer.py new file mode 100644 index 000000000..37e3c71cc --- /dev/null +++ b/tests/test_rl_verl_buffer_trainer.py @@ -0,0 +1,358 @@ +"""S25 C2: verl trainer island over the sample buffer (CPU, no verl/Ray). + +The trainer actor is a fake reached through the island's ``call``; the buffer, +the release store, the yeto driver, the fully_async rollout port (version map, +policy-age judgement), the feed, the publisher and the merged checkpoint are +the real ones. The end-to-end test runs the real engine sampler island (fake +engine) in a thread against the same store. +""" + +from __future__ import annotations + +import json +import threading +import time +from pathlib import Path + +import pytest +import torch +from safetensors.torch import load_file, save_file + +from yeto.rl.adapters.sampler_engines.fake import FakeEngine +from yeto.rl.adapters.verl import buffer_feed as bf +from yeto.rl.adapters.verl import fully_async_round as far +from yeto.rl.adapters.verl.buffer_ports import (CONSUMED_LOG, RELEASE_LOG, BufferReleasePublisher, + install_buffer_ports, is_base_adapter) +from yeto.rl.adapters.verl.fully_async_translate import VersionMap +from yeto.rl.adapters.verl.merged_release import MergeError, lora_pairs, write_merged_checkpoint +from yeto.rl.engine.algorithm import AlgorithmSpec, CorrectionSpec, ExecutionSpec +from yeto.rl.engine.buffer_payload import GroupWriter +from yeto.rl.engine.engine_sampler import EngineSamplerConfig, PromptItem, build_engine_sampler +from yeto.rl.engine.policy_release import PolicyReleaseStore +from yeto.rl.engine.resume import CheckpointStore +from yeto.rl.engine.sample_buffer import FileSampleBuffer +from yeto.rl.engine.sampler_checks import PostLoadChecks +from yeto.rl.engine.sampler_engine import SamplingParams + +TOLERANT = AlgorithmSpec(execution=ExecutionSpec(max_policy_staleness=1), + correction=CorrectionSpec(method="tis", tis_clip=2.0, tis_clip_low=0.0)) +MODULE = "model.layers.0.q_proj" +NAMES = (f"base_model.model.{MODULE}.lora_A.weight", f"base_model.model.{MODULE}.lora_B.weight") +SHAPE = bf.RowShape(prompt_length=6, response_length=8, pad_token_id=0) + + +def _sample(version, *, prompt=(11, 12), response=(21, 22, 23), reward=1.0, known=True, index=0, + provenance=None): + tokens = (list(prompt) if known else []) + list(response) + return {"index": index, "prompt": "p", "tokens": tokens, "response_length": len(response), + "rollout_log_probs": [-0.5] * len(response), "reward": reward, + "metadata": {"weight_version": str(version), "prompt_tokens_known": known, + "finish_reason": "stop", "provenance": provenance or {"engine": "fake"}}} + + +# ---------------------------------------------------------------------- rows +def test_group_rows_pad_like_the_agent_loop(): + rows = bf.group_rows([_sample(3, index=0), _sample(3, response=(31,), reward=0.0, index=1)], + uid="buf7", param_versions=[2, 2], shape=SHAPE) + t, nt = rows["tensors"], rows["non_tensor"] + assert t["prompts"][0] == [0, 0, 0, 0, 11, 12] + assert t["responses"][0] == [21, 22, 23, 0, 0, 0, 0, 0] + assert t["response_mask"][1] == [1, 0, 0, 0, 0, 0, 0, 0] + assert t["input_ids"][0] == [0, 0, 0, 0, 11, 12, 21, 22, 23, 0, 0, 0, 0, 0] + assert t["attention_mask"][0] == [0] * 4 + [1] * 5 + [0] * 5 + assert t["position_ids"][0] == [0, 0, 0, 0, 0, 1, 2, 3, 4, 4, 4, 4, 4, 4] + assert t["rm_scores"][0][:4] == [0.0, 0.0, 1.0, 0.0] and sum(t["rm_scores"][1]) == 0.0 + assert t["rollout_log_probs"][0][:4] == [-0.5, -0.5, -0.5, 0.0] + assert nt["uid"] == ["buf7", "buf7"] and nt["min_global_steps"] == [2, 2] == nt["max_global_steps"] + assert nt["provenance"][0] == {"engine": "fake"} and nt["finish_reason"] == ["stop", "stop"] + assert len(rows["meta_info"]["metrics"]) == 2 + + +def test_group_rows_rebuild_unknown_prompt_ids_with_the_trainer_tokenizer(): + rows = bf.group_rows([_sample(0, known=False)], uid="u", param_versions=[0], shape=SHAPE, + tokenize=lambda text: [5, 6, 7]) + assert rows["tensors"]["prompts"][0] == [0, 0, 0, 5, 6, 7] + with pytest.raises(Exception, match="no tokenizer"): + bf.group_rows([_sample(0, known=False)], uid="u", param_versions=[0], shape=SHAPE) + + +def test_group_rows_refuse_samples_wider_than_the_rows(): + with pytest.raises(bf.BufferFeedError, match="exceed"): + bf.group_rows([_sample(0, prompt=tuple(range(1, 8)))], uid="u", param_versions=[0], shape=SHAPE) + with pytest.raises(bf.BufferFeedError, match="param versions"): + bf.group_rows([_sample(0)], uid="u", param_versions=[], shape=SHAPE) + + +# ---------------------------------------------------------------------- feed +def _buffer(tmp_path): + return FileSampleBuffer(CheckpointStore(tmp_path), batch_size=1) + + +def test_feed_maps_versions_judges_and_records_each_position_once(tmp_path): + writer = GroupWriter(_buffer(tmp_path), "s0") + for key, version in (("a", 0), ("b", 1), ("c", 5)): + writer.write([_sample(version, index=0), _sample(version, index=1)], group_key=key) + reader, vmap, added, records = _buffer(tmp_path), VersionMap(), [], [] + vmap.record(0, 0) + vmap.record(1, 1) + feed = bf.BufferFeed(reader, vmap=lambda: vmap, shape=SHAPE, add_group=added.append, + record=records.append, sleep=lambda s: None) + collector = far.RoundCollector(vmap, 1, 1, 3) + for _ in range(3): + meta = feed.pull() + keep = collector.offer(far.SampleMeta(**meta)) + feed("yeto_keep_last" if keep else "yeto_drop_last") + assert [r["position"] for r in records] == [0, 1, 2] == feed.pulled_positions + assert [r["kept"] for r in records] == [True, True, False] # v5 was never published here + assert records[2]["param_versions"] == [None, None] + assert [rows["non_tensor"]["uid"][0] for rows in added] == ["buf0", "buf1"] + assert collector.discarded[0].unknown + assert feed.read_cursor() == 3 + + +def test_feed_returns_none_at_the_data_end_and_times_out_without_data(tmp_path): + buffer = _buffer(tmp_path) + clock = iter(range(0, 100, 10)) + feed = bf.BufferFeed(buffer, vmap=VersionMap, shape=SHAPE, add_group=print, record=print, + wait_s=15, clock=lambda: next(clock), sleep=lambda s: None) + with pytest.raises(bf.BufferFeedError, match="no buffer group within 15 s"): + feed.pull() + buffer.mark_data_end() + assert bf.BufferFeed(_buffer(tmp_path), vmap=VersionMap, shape=SHAPE, add_group=print, + record=print).pull() is None + with pytest.raises(bf.BufferFeedError, match="without a pulled group"): + feed.keep() + + +def test_buffer_call_routes_queue_calls_to_the_feed(): + seen = [] + call = bf.BufferCall(lambda name, *a: seen.append(("feed", name)), lambda name, *a, **k: seen.append(("actor", name))) + call("yeto_pull_sample") + call("yeto_train_step") + assert seen == [("feed", "yeto_pull_sample"), ("actor", "yeto_train_step")] + + +# ---------------------------------------------------------------------- merge +def _base(tmp_path, *, dtype=torch.bfloat16): + base = tmp_path / "base" + base.mkdir() + weights = {f"{MODULE}.weight": torch.randn(4, 4).to(dtype), "lm_head.weight": torch.randn(3, 4).to(dtype)} + save_file(weights, str(base / "model.safetensors")) + (base / "config.json").write_text(json.dumps({"model_type": "qwen3"})) + (base / "tokenizer.json").write_text("{}") + return base, weights + + +def test_merge_adds_b_times_a_and_keeps_the_layout(tmp_path): + base, weights = _base(tmp_path, dtype=torch.float32) + a, b = torch.randn(2, 4), torch.randn(4, 2) + summary = write_merged_checkpoint(base, {NAMES[0]: a, NAMES[1]: b}, tmp_path / "out") + out = load_file(str(tmp_path / "out" / "model.safetensors")) + assert torch.allclose(out[f"{MODULE}.weight"], weights[f"{MODULE}.weight"] + b @ a, atol=1e-6) + assert torch.equal(out["lm_head.weight"], weights["lm_head.weight"]) + assert summary["merged_modules"] == 1 and summary["max_relative_change"] > 0 + assert (tmp_path / "out" / "config.json").read_text() == (base / "config.json").read_text() + assert (tmp_path / "out" / "tokenizer.json").is_file() + + +def test_merge_keeps_the_base_dtype_and_refuses_unknown_or_misshaped_modules(tmp_path): + base, _ = _base(tmp_path) + write_merged_checkpoint(base, {NAMES[0]: torch.ones(2, 4), NAMES[1]: torch.ones(4, 2)}, tmp_path / "o1") + assert load_file(str(tmp_path / "o1" / "model.safetensors"))[f"{MODULE}.weight"].dtype == torch.bfloat16 + other = "base_model.model.model.layers.9.q_proj" + with pytest.raises(MergeError, match="no weight for 1 LoRA modules"): + write_merged_checkpoint(base, {f"{other}.lora_A.weight": torch.ones(2, 4), + f"{other}.lora_B.weight": torch.ones(4, 2)}, tmp_path / "o2") + with pytest.raises(MergeError, match="do not fit"): + write_merged_checkpoint(base, {NAMES[0]: torch.ones(2, 3), NAMES[1]: torch.ones(4, 2)}, tmp_path / "o3") + with pytest.raises(MergeError, match="only side"): + lora_pairs({NAMES[0]: torch.ones(2, 4)}) + + +def test_base_adapter_is_version_zero_with_zero_b(): + zero = {NAMES[0]: torch.ones(2, 4), NAMES[1]: torch.zeros(4, 2)} + assert is_base_adapter(0, zero) + assert not is_base_adapter(1, zero) + assert not is_base_adapter(0, {**zero, NAMES[1]: torch.full((4, 2), 1e-3)}) + + +# ---------------------------------------------------------------------- driver end to end +class FakeBufferActor: + """Stands in for YetoBufferTrainer: kept rows, a LoRA pair, a param version.""" + + def __init__(self): + self.tensors = {NAMES[0]: torch.full((2, 4), 0.1), NAMES[1]: torch.zeros(4, 2)} + self.param_version = 0 + self.kept, self.assembled = [], [] + + def __call__(self, name, *args, timeout=None): + return getattr(self, name)(*args) + + def yeto_configure(self, *_): + return None + + def yeto_add_group(self, rows): + self.kept.append(rows) + return len(self.kept) + + def yeto_assemble(self): + self.assembled.append([r["non_tensor"]["uid"][0] for r in self.kept]) + n, self.kept = len(self.kept), [] + return {"samples": n, "queue_size": None} + + def yeto_train_step(self): + self.tensors[NAMES[1]] += 0.01 + self.param_version += 1 + return {"metrics": {"actor/grad_norm": 0.5, "actor/pg_loss": 0.1, "actor/lr": 1e-5}, + "criteria": {}, "param_version": self.param_version} + + def yeto_param_version(self): + return self.param_version + + def yeto_export_lora(self): + return {k: v.clone() for k, v in self.tensors.items()}, {"count": 2} + + def yeto_call_workers(self, fn, tensors, reset): + for k in self.tensors: + self.tensors[k] = tensors[k].clone().float() + return [{"applied": len(tensors)}] + + +def _trainer(tmp_path, store_dir, actor, *, rounds, required, base): + from yeto.rl.adapters.verl.entry import verl_capabilities + from yeto.rl.adapters.verl.fully_async_ports import FullyAsyncIsland, FullyAsyncTrainerGroup + from yeto.rl.adapters.verl.ports_impl import VerlPolicyState + from yeto.rl.core import CanonicalTensorSpec, canonical_layout_hash + from yeto.rl.engine.bridges import LocalOnlySync + from yeto.rl.engine.driver import EventTape, IslandDriver + + specs = (CanonicalTensorSpec(NAMES[0], (2, 4), "float32", 8), CanonicalTensorSpec(NAMES[1], (4, 2), "float32", 8)) + island = FullyAsyncIsland(actor, learner_id=0, limit=1, required=required, out_dir=str(tmp_path / "out"), + readback_dir=str(tmp_path / "rb"), base_model_revision="a" * 40, + lora_config_hash="b" * 64, layout_hash=canonical_layout_hash(specs), + expected_specs=specs) + tape = EventTape(tmp_path / "tape.jsonl", 0) + island.emit = lambda event, **f: tape.append({"event": event, **f}) + store = CheckpointStore(store_dir) + rollout, publisher, placement = install_buffer_ports( + island, buffer=FileSampleBuffer(store, batch_size=1, create_index=False), + releases=PolicyReleaseStore(store), shape=bf.RowShape(64, 16, 0), base_dir=base, + tokenize=lambda text: [7, 8], poll_s=0.01, wait_s=30) + profile = far.execution_profile_for(TOLERANT, 1, groups_per_round=required, samples_per_group=2, sync="none") + driver = IslandDriver( + learner_id=0, rollout=rollout, trainer=FullyAsyncTrainerGroup(island), policy_state=VerlPolicyState(island), + publisher=publisher, placement=placement, + capabilities=verl_capabilities("sha256:" + "0" * 64, fully_async_limit=1), + algorithm=TOLERANT, sync=LocalOnlySync(rounds), events=tape, profile=profile) + return driver, island, publisher + + +def _lines(path): + return [json.loads(line) for line in Path(path).read_text().splitlines()] + + +def test_sampler_island_and_verl_buffer_trainer_run_three_rounds(tmp_path): + store_dir = tmp_path / "store" + store = CheckpointStore(store_dir) + base, weights = _base(tmp_path, dtype=torch.float32) + sampler_events, done = [], threading.Event() + engine = FakeEngine(max_model_len=64) + config = EngineSamplerConfig(sampler_id="island0", num_rollouts=40, max_backlog_groups=2, + sampling=SamplingParams(max_tokens=6, n=2), max_prompt_tokens=32, poll_s=0.01) + sampler = build_engine_sampler( + config, engine=engine, releases=PolicyReleaseStore(store), buffer=FileSampleBuffer(store, batch_size=1), + emit=lambda n, **f: sampler_events.append((n, f)), + prompts=lambda r: [PromptItem(f"p{r}", f"prompt {r}", label="x")], reward=lambda item, c: float(len(c.token_ids) % 2), + checks=PostLoadChecks(expected_expert_matrices=None, fixed_prompts=(), probe_pairs=(("Q", " 7"),)), + should_stop=done.is_set, + sleep=lambda s: time.sleep(0.005)) + thread = threading.Thread(target=sampler.run, daemon=True) + thread.start() + actor = FakeBufferActor() + driver, island, publisher = _trainer(tmp_path, store_dir, actor, rounds=3, required=2, base=base) + try: + state = driver.run() + finally: + done.set() + thread.join(timeout=10) + assert driver.rounds_completed == 3 and state.policy_version == 3 + assert island.vmap.to_dict() == {"0": 0, "1": 1, "2": 2, "3": 3} + # no group consumed twice; trained groups are those handed to the actor + consumed = _lines(tmp_path / "out" / CONSUMED_LOG.format(learner=0)) + positions = [r["position"] for r in consumed] + assert positions == sorted(set(positions)) == list(range(len(positions))) + kept = [r for r in consumed if r["kept"]] + assert len(kept) == 6 and [len(a) for a in actor.assembled] == [2, 2, 2] + assert sorted(u for a in actor.assembled for u in a) == sorted(f"buf{r['position']}" for r in kept) + # versions: v0 (the base) is not released; v1..v3 are, in order + releases = _lines(tmp_path / "out" / RELEASE_LOG.format(learner=0)) + assert [(r["policy_version"], r["released"]) for r in releases] == [(0, False), (1, True), (2, True), (3, True)] + assert [r["seq"] for r in releases if r["released"]] == [1, 2, 3] + latest = PolicyReleaseStore(store).latest() + assert (latest.version, latest.weights, latest.buffer_selected) == (3, "full", releases[-1]["buffer_selected"]) + merged = load_file(str(store_dir / "policy-releases" / latest.path / "model.safetensors")) + a, b = actor.tensors[NAMES[0]], actor.tensors[NAMES[1]] + assert torch.allclose(merged[f"{MODULE}.weight"], weights[f"{MODULE}.weight"] + b @ a, atol=1e-6) + # the sampler loaded the releases in order, and the trained groups respect the age limit + loaded = [f["policy_version"] for n, f in sampler_events if n == "rl_sampler_policy_loaded"] + assert loaded == sorted(set(loaded)) and set(loaded) <= {1, 2, 3} + events = _lines(tmp_path / "tape.jsonl") + carry = [e for e in events if e["event"] == "rl_rollout_carry_over"] + assert [e["rollout_id"] for e in carry] == [0, 1, 2] + rounds = [e for e in events if e["event"] == "rl_verl_fa_round"] + assert sum(e["kept"] for e in rounds) == 6 + assert sum(e["pulled"] for e in rounds) == len(consumed) + + +def test_publisher_releases_the_base_when_its_adapter_is_not_zero(tmp_path): + from yeto.rl.adapters.verl.fully_async_ports import FullyAsyncIsland + + actor = FakeBufferActor() + actor.tensors[NAMES[1]] += 1.0 + island = FullyAsyncIsland(actor, learner_id=0, limit=1, required=1, out_dir=str(tmp_path / "out"), + readback_dir=str(tmp_path / "rb")) + store = CheckpointStore(tmp_path / "store") + feed = bf.BufferFeed(FileSampleBuffer(store, batch_size=1), vmap=lambda: island.vmap, shape=SHAPE, + add_group=print, record=print) + + class State: + policy_version = 0 + + @staticmethod + def policy_tensor_hash(): + return "0" * 64 + + @staticmethod + def to_lora(): + return type("L", (), {"tensors": actor.tensors}) + + calls = [] + publisher = BufferReleasePublisher(island, PolicyReleaseStore(store), feed, base_dir=tmp_path, + merge=lambda base, tensors, d: calls.append(d) or {"merged_modules": 1}) + result = publisher.publish(State()) + assert len(calls) == 1 and island.published == (0, "0" * 64) + assert result.manifest.target_policy_version == 0 and island.vmap.to_dict() == {"0": 0} + + +# ---------------------------------------------------------------------- island entry +def _entry_args(*extra): + from yeto.rl.adapters.verl.island_entry import parse_args + + return parse_args(["--model", "m", "--model-revision", "r", "--data", "d", "--learner-id", "0", + "--reward-function", "gsm8k_reward:score", "--global-rounds", "3", + "--groups-per-round", "2", "--samples-per-group", "2", + "--rollout-max-response-len", "64", "--event-tape", "t", "--inner-lr", "1e-5", + "--seq-len", "512", "--seed", "1", *extra])[0] + + +def test_island_entry_plans_a_buffer_trainer_and_refuses_other_roles(): + from yeto.rl.adapters.verl.island_entry import sample_buffer_section + + assert sample_buffer_section(_entry_args()) is None + section = sample_buffer_section(_entry_args("--rl-sample-role", "trainer", "--rl-resume-store", "/s", + "--rl-max-policy-age", "1", "--rl-sample-wait-s", "60")) + assert section == {"store": "/s", "wait_s": 60.0} + with pytest.raises(SystemExit, match="engine sampler"): + sample_buffer_section(_entry_args("--rl-sample-role", "sampler", "--rl-resume-store", "/s")) + with pytest.raises(SystemExit, match="max-policy-age"): + sample_buffer_section(_entry_args("--rl-sample-role", "trainer", "--rl-resume-store", "/s")) diff --git a/yeto/rl/adapters/verl/buffer_feed.py b/yeto/rl/adapters/verl/buffer_feed.py new file mode 100644 index 000000000..e93e7854a --- /dev/null +++ b/yeto/rl/adapters/verl/buffer_feed.py @@ -0,0 +1,274 @@ +"""verl trainer island over the durable sample buffer: the feed (S25 C2). + +A verl *buffer trainer* trains on groups that a separate sampler island wrote +to the sample buffer (``yeto.rl.engine.sample_buffer``; the engine sampler +island ``yeto.rl.engine.engine_sampler`` or the Miles sampler island). It reuses +the fully_async consumption unchanged: the yeto ``IslandDriver`` and the ports +of :mod:`.fully_async_ports` (``RoundCollector`` with the ``max_policy_age`` +judgement, carry-over and discard reports, the trainer step with the off-policy +correction). Only the source of the queued samples differs: verl's +``MessageQueue`` (filled by its rollouter) is replaced by :class:`BufferFeed`. + +:class:`BufferFeed` answers the three queue calls of the rollout port +(``yeto_pull_sample`` / ``yeto_keep_last`` / ``yeto_drop_last``) in the driver +process, so everything that decides which group is trained runs here, on plain +data, without verl or Ray: + +* ``pull``: the next buffer entry (one GRPO group), waiting up to ``wait_s`` for + the sampler; its :class:`~.fully_async_round.SampleMeta` with every response's + behavior version translated to verl's ``param_version`` through the island's + :class:`~.fully_async_translate.VersionMap` (a version the trainer never + published here maps to None: the collector discards the group as unknown); +* ``keep``: the group's padded rows (:func:`group_rows`) go to the trainer actor + (``add_group``), which wraps them as a verl ``RolloutSample``; +* ``drop``: nothing reaches the actor; the group stays consumed. + +Every pulled group is recorded once (``record``) with its buffer position, so a +run can show that no group was consumed twice. The buffer positions are those of +the trainer's single read cursor (``FileSampleBuffer.select``). + +:func:`group_rows` builds the rows the way verl's agent loop postprocess pads +them (``AgentLoopWorker._postprocess``): prompts left-padded to the prompt +length, responses right-padded to the response length, the reward on the last +response token (``rm_scores``), the behavior log-probs as ``rollout_log_probs`` +and the sample's source record as ``provenance`` (the per-source off-policy +metrics read it). +""" + +from __future__ import annotations + +import time +from collections.abc import Callable, Mapping, Sequence +from dataclasses import dataclass +from typing import Any + +from yeto.rl.engine.buffer_payload import ( + SampleBufferRolloutError, + complete_prompt_tokens, + decode_group, + group_entry_fields, +) + +from .fully_async_translate import VerlTranslationError, VersionMap + +TRAINER_CONSUMER = "trainer" # the buffer's single consumer name (as the Miles reader) +QUEUE_METHODS = ("yeto_pull_sample", "yeto_keep_last", "yeto_drop_last") +UID_PREFIX = "buf" + + +class BufferFeedError(SampleBufferRolloutError): + """A buffer group cannot become verl rows.""" + + +@dataclass(frozen=True) +class RowShape: + """The padded widths of the trainer's rows (``data.max_prompt_length`` and + ``data.max_response_length`` of the verl config) and the pad token id.""" + + prompt_length: int + response_length: int + pad_token_id: int + + def __post_init__(self) -> None: + if self.prompt_length < 1 or self.response_length < 1: + raise ValueError(f"row widths must be positive, got {self.prompt_length}/{self.response_length}") + + +def sample_version(sample: Mapping[str, Any]) -> int: + """The oldest engine weight version of one sample (its behavior version).""" + return group_entry_fields([sample])[0] + + +def group_rows(samples: Sequence[Mapping[str, Any]], *, uid: str, param_versions: Sequence[int], + shape: RowShape, tokenize: Callable[[str], Sequence[int]] | None = None) -> dict[str, Any]: + """The verl rows of one group as plain lists (module doc). + + ``param_versions``: verl's ``param_version`` of each sample's behavior + version (``min_global_steps`` = ``max_global_steps``: a buffer sample comes + from one engine version). ``tokenize`` rebuilds the prompt ids of a sample + whose engine did not return them (``complete_prompt_tokens``). Returns + ``{"tensors": {name: rows}, "non_tensor": {name: values}, "meta_info": {...}}``. + Raises :class:`BufferFeedError` when a sample does not fit the widths.""" + if len(param_versions) != len(samples): + raise BufferFeedError(f"{len(param_versions)} param versions for {len(samples)} samples") + tensors: dict[str, list[list[Any]]] = {k: [] for k in ( + "prompts", "responses", "response_mask", "input_ids", "attention_mask", "position_ids", + "rollout_log_probs", "rm_scores")} + provenance, finish = [], [] + pad = int(shape.pad_token_id) + for sample in samples: + sample = complete_prompt_tokens(sample, tokenize) + tokens = [int(t) for t in sample.get("tokens") or ()] + length = int(sample.get("response_length") or 0) + logprobs = [float(x) for x in sample.get("rollout_log_probs") or ()] + where = f"group {uid} sample {sample.get('index')}" + if length < 1 or len(logprobs) != length or len(tokens) <= length: + raise BufferFeedError(f"{where}: {len(tokens)} tokens, response_length {length}, " + f"{len(logprobs)} log-probs") + prompt, response = tokens[:-length], tokens[-length:] + if len(prompt) > shape.prompt_length or length > shape.response_length: + raise BufferFeedError(f"{where}: prompt {len(prompt)} / response {length} tokens exceed " + f"the rows ({shape.prompt_length} / {shape.response_length})") + left = shape.prompt_length - len(prompt) + right = shape.response_length - length + attention = [0] * left + [1] * (len(prompt) + length) + [0] * right + tensors["prompts"].append([pad] * left + prompt) + tensors["responses"].append(response + [pad] * right) + tensors["response_mask"].append([1] * length + [0] * right) + tensors["input_ids"].append([pad] * left + prompt + response + [pad] * right) + tensors["attention_mask"].append(attention) + tensors["position_ids"].append(_position_ids(attention)) + tensors["rollout_log_probs"].append(logprobs + [0.0] * right) + tensors["rm_scores"].append([0.0] * (length - 1) + [float(sample.get("reward") or 0.0)] + + [0.0] * right) + metadata = sample.get("metadata") or {} + provenance.append(dict(metadata.get("provenance") or {}) or None) + finish.append(metadata.get("finish_reason")) + n = len(samples) + versions = [int(v) for v in param_versions] + return { + "tensors": tensors, + "non_tensor": {"uid": [uid] * n, "min_global_steps": versions, "max_global_steps": list(versions), + "provenance": provenance, "finish_reason": finish, "__num_turns__": [1] * n}, + # verl's assemble step pops one ``metrics`` record per row (addition_process) + "meta_info": {"metrics": [{"generate_sequences": 0.0, "tool_calls": 0.0}] * n, + "reward_extra_keys": []}, + } + + +def _position_ids(attention: Sequence[int]) -> list[int]: + """verl's ``compute_position_id_with_mask``: cumulative sum minus one, floored at 0.""" + out, seen = [], 0 + for bit in attention: + seen += int(bit) + out.append(max(seen - 1, 0)) + return out + + +@dataclass +class _Pulled: + """The group the collector is judging: its rows' inputs and its record.""" + + position: int + samples: list[dict[str, Any]] + param_versions: tuple[int | None, ...] + record: dict[str, Any] + + +class BufferFeed: + """The rollout port's queue calls over the sample buffer (module doc). + + ``buffer``: the trainer's ``FileSampleBuffer``; ``vmap``: returns the + island's live version map; ``add_group(rows)``: hands kept rows to the + trainer actor; ``record(fields)``: one record per pulled group.""" + + def __init__(self, buffer: Any, *, vmap: Callable[[], VersionMap], shape: RowShape, + add_group: Callable[[dict[str, Any]], Any], record: Callable[[dict[str, Any]], None], + tokenize: Callable[[str], Sequence[int]] | None = None, wait_s: float = 3600.0, + poll_s: float = 2.0, consumer: str = TRAINER_CONSUMER, + clock: Callable[[], float] = time.monotonic, + sleep: Callable[[float], None] = time.sleep) -> None: + self.buffer = buffer + self.vmap = vmap + self.shape = shape + self.add_group = add_group + self.record = record + self.tokenize = tokenize + self.wait_s = float(wait_s) + self.poll_s = float(poll_s) + self.consumer = consumer + self.clock = clock + self.sleep = sleep + self._last: _Pulled | None = None + self.pulled_positions: list[int] = [] + + def __call__(self, name: str, *args: Any) -> Any: + """``call(name)`` for the three queue methods (:data:`QUEUE_METHODS`).""" + if name == "yeto_pull_sample": + return self.pull() + if name == "yeto_keep_last": + return self.keep() + if name == "yeto_drop_last": + return self.drop() + raise ValueError(f"{name} is not a queue call of the buffer feed") + + def _next(self) -> Any | None: + """The next selected entry; None after the sampler's data end.""" + started = self.clock() + while True: + selected = self.buffer.select(self.consumer, 1) + if selected: + return selected[0] + if self.buffer.data_ended: + return None + if self.clock() - started > self.wait_s: + raise BufferFeedError(f"no buffer group within {self.wait_s:g} s " + f"(read position {self.buffer.read_cursor})") + self.sleep(self.poll_s) + + def _local(self, version: int) -> int | None: + try: + return self.vmap().local(version) + except VerlTranslationError: + return None + + def pull(self) -> dict[str, Any] | None: + if self._last is not None: + raise BufferFeedError("pull before the previous group was kept or dropped") + item = self._next() + if item is None: + return None + samples = decode_group(self.buffer.read_payload(item.entry.payload_ref)) + if not samples: + raise BufferFeedError(f"buffer position {item.position}: empty group") + versions = [sample_version(s) for s in samples] + local = tuple(self._local(v) for v in versions) + lengths = tuple(int(s.get("response_length") or 0) for s in samples) + position = int(item.position) + self._last = _Pulled(position, samples, local, { + "position": position, "sample_id": item.entry.sample_id, + "behavior_versions": versions, "param_versions": list(local), + "source": (item.entry.provenance or {}).get("sampler_id"), + "samples": len(samples), "response_tokens": sum(lengths)}) + self.pulled_positions.append(position) + return {"uid": f"{UID_PREFIX}{int(item.position)}", "min_steps": local, "max_steps": local, + "calls": (None,) * len(samples), "lengths": lengths, + "rewards": tuple(float(s.get("reward") or 0.0) for s in samples)} + + def keep(self) -> None: + last = self._take_last() + if None in last.param_versions: # the collector keeps only groups whose versions it knows + raise BufferFeedError(f"buffer position {last.position}: kept with an unknown version " + f"{last.param_versions}") + self.add_group(group_rows(last.samples, uid=f"{UID_PREFIX}{last.position}", + param_versions=last.param_versions, shape=self.shape, + tokenize=self.tokenize)) + self.record({**last.record, "kept": True}) + + def drop(self) -> None: + self.record({**self._take_last().record, "kept": False}) + + def _take_last(self) -> _Pulled: + if self._last is None: + raise BufferFeedError("keep/drop without a pulled group") + last, self._last = self._last, None + return last + + def read_cursor(self) -> int: + """The buffer position after the last pulled group (released to the + samplers for their backlog cap).""" + return int(self.buffer.read_cursor) + + +class BufferCall: + """The island's ``call(method, *args)``: the queue calls go to the feed, every + other call to the trainer actor (``actor_call``).""" + + def __init__(self, feed: BufferFeed, actor_call: Callable[..., Any]) -> None: + self.feed = feed + self.actor_call = actor_call + + def __call__(self, name: str, *args: Any, **kwargs: Any) -> Any: + if name in QUEUE_METHODS: + return self.feed(name, *args) + return self.actor_call(name, *args, **kwargs) diff --git a/yeto/rl/adapters/verl/buffer_ports.py b/yeto/rl/adapters/verl/buffer_ports.py new file mode 100644 index 000000000..e354ff119 --- /dev/null +++ b/yeto/rl/adapters/verl/buffer_ports.py @@ -0,0 +1,148 @@ +"""verl buffer trainer: the ports that differ from the fully_async path (S25 C2). + +:func:`install_buffer_ports` turns a :class:`~.fully_async_ports.FullyAsyncIsland` +into a trainer island over the sample buffer (``--rl-sample-role trainer``): + +* the island's ``call`` becomes a :class:`~.buffer_feed.BufferCall`: the queue + calls go to the :class:`~.buffer_feed.BufferFeed`, the rest to the trainer + actor. The rollout port, the trainer port and the policy state are the + fully_async ones, unchanged; +* :class:`BufferReleasePublisher` replaces the checkpoint-engine push: the + trained adapter is merged into the base (:mod:`.merged_release`) and released + into the store for the samplers (``PolicyReleaseStore``, full weights) with + the trainer's buffer read position. The pairing ``param_version <-> outer + version`` is recorded at every publication, as on the fully_async path. + +The base policy (version 0, LoRA B all zero: the merged weights are the base) +is not released: the samplers start on the base weights as version 0. A +version-0 state with any non-zero B is released like any other version. + +The publication is complete when the release record and ``LATEST.json`` are in +the store. Whether a sampler has loaded it is the sampler's report +(``rl_sampler_policy_loaded`` on its tape) and the behavior versions of the +groups it writes next; the trainer does not wait for it (the samplers run +elsewhere and may be stopped). +""" + +from __future__ import annotations + +import hashlib +import json +import time +from collections.abc import Callable, Mapping, Sequence +from pathlib import Path +from typing import Any + +from yeto.rl.contracts import InferencePublicationManifest +from yeto.rl.engine.policy_release import FULL, PolicyReleaseStore +from yeto.rl.engine.ports import PlacementDescription, PublicationResult + +from .buffer_feed import BufferCall, BufferFeed, RowShape +from .fully_async_ports import MEMBER, FullyAsyncIsland, FullyAsyncRolloutPool +from .fully_async_round import FULLY_ASYNC_PLACEMENT +from .merged_release import LORA_B, write_merged_checkpoint + +BASE_POLICY_VERSION = 0 +RELEASE_EVENT = "rl_policy_released" # the Miles trainer island's event name +CONSUMED_LOG = "verl-buffer-consumed-{learner}.jsonl" +RELEASE_LOG = "verl-buffer-release-{learner}.jsonl" + + +def _sha(text: str) -> str: + return hashlib.sha256(text.encode()).hexdigest() + + +def is_base_adapter(version: int, tensors: Mapping[str, Any]) -> bool: + """Version 0 with every LoRA B zero: base + adapter is the base itself.""" + return int(version) == BASE_POLICY_VERSION and all( + not bool(t.detach().abs().max() > 0) for name, t in tensors.items() if name.endswith(LORA_B)) + + +class BufferReleasePublisher: + """yeto publication as a store release of the merged policy (module doc). + + ``feed`` gives the buffer read position; ``base_dir`` is the base HF + checkpoint the adapter trains over.""" + + def __init__(self, island: FullyAsyncIsland, releases: PolicyReleaseStore, feed: BufferFeed, + *, base_dir: str | Path, clock: Callable[[], float] = time.monotonic, + merge: Callable[..., dict[str, Any]] = write_merged_checkpoint) -> None: + self.island = island + self.releases = releases + self.feed = feed + self.base_dir = Path(base_dir) + self.clock = clock + self.merge = merge + + def publish(self, state: Any) -> PublicationResult: + isl = self.island + version = int(state.policy_version) + param_version = int(isl.call("yeto_param_version")) + isl.vmap.record(param_version, version) # raises on contradiction / regression + digest = state.policy_tensor_hash() + tensors = dict(state.to_lora().tensors) + started = self.clock() + merge: dict[str, Any] = {} + release = None + position = self.feed.read_cursor() + if not is_base_adapter(version, tensors): + def write(directory: Path) -> None: + merge.update(self.merge(self.base_dir, tensors, directory)) + + release = self.releases.release(version, digest, write, buffer_selected=position, weights=FULL) + record = {"policy_version": version, "param_version": param_version, "policy_hash": digest, + "released": release is not None, "buffer_selected": position, + "seconds": round(self.clock() - started, 3), "merge": merge or None, + "version_map": isl.vmap.to_dict()} + if release is not None: + record.update(seq=release.seq, path=release.path, + payload_bytes=sum(int(v) for v in (release.files or {}).values())) + isl.emit(RELEASE_EVENT, seq=release.seq, policy_version=version, policy_hash=digest, + buffer_selected=position, weights=FULL, sha256=release.sha256, + engine_checksum=None, merged_modules=merge.get("merged_modules"), + max_relative_change=merge.get("max_relative_change"), + seconds=record["seconds"]) + isl.append_jsonl(RELEASE_LOG.format(learner=isl.learner_id), record) + isl.emit("rl_verl_version_map", param_version=param_version, outer_version=version) + isl.published = (version, digest) + isl.hashes[version] = digest + manifest = InferencePublicationManifest( + publication_mode="full", base_policy_version=None, target_policy_version=version, + target_policy_hash=digest, target_manifest_hash=_sha(f"{version}:{digest}"), + payload_hash=_sha(json.dumps({"release": record.get("path"), "hash": digest})), + # the release's files, or (the base, not released) the adapter in bf16 as on the fully_async path + payload_bytes=int(record.get("payload_bytes") or sum(int(t.numel()) * 2 for t in tensors.values())), + complete=True) + return PublicationResult(manifest, frozenset({MEMBER})) + + +class BufferPlacement: + """The trainer GPUs only: the samplers run on other islands.""" + + def describe(self) -> PlacementDescription: + return PlacementDescription(FULLY_ASYNC_PLACEMENT, ("cuda:trainer-0",), ("store:sample-buffer",), + extra={"engine": "verl-buffer-trainer"}) + + +def install_buffer_ports(island: FullyAsyncIsland, *, buffer: Any, releases: PolicyReleaseStore, + shape: RowShape, base_dir: str | Path, + tokenize: Callable[[str], Sequence[int]] | None = None, + wait_s: float = 3600.0, poll_s: float = 2.0, + **feed_kwargs: Any) -> tuple[FullyAsyncRolloutPool, BufferReleasePublisher, + BufferPlacement]: + """Route the island's queue calls to a buffer feed and return the rollout + port, the publisher and the placement of a buffer trainer (module doc); the + feed is ``publisher.feed``.""" + actor_call = island.call + log = CONSUMED_LOG.format(learner=island.learner_id) + + def record(fields: dict[str, Any]) -> None: + island.append_jsonl(log, fields) + + feed = BufferFeed(buffer, vmap=lambda: island.vmap, shape=shape, + add_group=lambda rows: actor_call("yeto_add_group", rows), record=record, + tokenize=tokenize, wait_s=wait_s, poll_s=poll_s, **feed_kwargs) + island.call = BufferCall(feed, actor_call) + island.start_rollouter = None # no rollouter: the samplers run elsewhere + publisher = BufferReleasePublisher(island, releases, feed, base_dir=base_dir) + return FullyAsyncRolloutPool(island), publisher, BufferPlacement() diff --git a/yeto/rl/adapters/verl/buffer_trainer_runner.py b/yeto/rl/adapters/verl/buffer_trainer_runner.py new file mode 100644 index 000000000..595737b79 --- /dev/null +++ b/yeto/rl/adapters/verl/buffer_trainer_runner.py @@ -0,0 +1,143 @@ +"""verl buffer trainer island, image side (S25 C2; runs inside the verl image). + +The trainer island of a sampler/trainer split (``--rl-sample-role trainer``): +verl's ``FullyAsyncTrainer`` trains, yeto's driver runs the rounds over the +fully_async ports, and the queued samples come from the durable sample buffer +instead of verl's rollouter (:mod:`.buffer_feed`, :mod:`.buffer_ports`). + +* :class:`YetoBufferTaskRunner` builds only the trainer (tokenizer, worker + mapping, the trainer actor and its workers): no rollouter, no MessageQueue, + no checkpoint engine, no initial weight push. Each publication is a store + release instead (``BufferReleasePublisher``). +* :data:`YetoBufferTrainer` is the fully_async trainer actor + (``_YetoTrainerMixin``: assemble, train step with the mismatch criteria and + the per-source off-policy metrics, LoRA export) plus ``yeto_add_group`` (one + kept group's rows as a verl ``RolloutSample``) and ``yeto_param_version``. + +verl itself is not changed. Only exercised on a GPU; the logic it relies on is +in the CPU-tested :mod:`.buffer_feed` and :mod:`.buffer_ports`. +""" + +from __future__ import annotations + +import os + +import numpy as np +import ray + +from .fully_async_runner import YetoFullyAsyncTaskRunner, _TrainerBase, _YetoTrainerMixin + +FLOAT_ROWS = ("rollout_log_probs", "rm_scores") +INT32_FIELDS = ("__num_turns__",) +SAMPLE_BUFFER_PLAN_KEY = "sample_buffer" # plan section of a buffer trainer (island_entry) + + +def _object_array(values) -> np.ndarray: + out = np.empty(len(values), dtype=object) + out[:] = list(values) + return out + + +def rows_to_dataproto(rows: dict): + """The plain rows of :func:`.buffer_feed.group_rows` as a verl ``DataProto``.""" + import torch + from tensordict import TensorDict + from verl import DataProto + + tensors = rows["tensors"] + n = len(tensors["prompts"]) + batch = TensorDict({name: torch.tensor(values, dtype=torch.float32 if name in FLOAT_ROWS else torch.long) + for name, values in tensors.items()}, batch_size=n) + non_tensor = {name: (np.asarray(values, dtype=np.int32) if name in INT32_FIELDS else _object_array(values)) + for name, values in rows["non_tensor"].items()} + meta = {key: (list(value) if isinstance(value, list) else value) for key, value in rows["meta_info"].items()} + return DataProto(batch=batch, non_tensor_batch=non_tensor, meta_info=meta) + + +class _YetoBufferTrainerMixin(_YetoTrainerMixin): + def yeto_add_group(self, rows: dict) -> int: + from verl.experimental.fully_async_policy.detach_utils import RolloutSample + + if getattr(self, "_yeto_kept", None) is None: + self._yeto_kept = [] + proto = rows_to_dataproto(rows) + self._yeto_kept.append(RolloutSample(full_batch=proto, sample_id=str(rows["non_tensor"]["uid"][0]), + epoch=0, rollout_status={})) + return len(self._yeto_kept) + + def yeto_param_version(self) -> int: + return int(self.current_param_version) + + async def _yeto_queue_size(self): + return None # no MessageQueue: the buffer backlog is the samplers' report + + +YetoBufferTrainer = ray.remote(num_cpus=10)( + type("YetoBufferTrainer", (_YetoBufferTrainerMixin, _TrainerBase), {})) + + +class YetoBufferTaskRunner(YetoFullyAsyncTaskRunner): + """The trainer-only set-up, then the yeto driver over the buffer ports.""" + + trainer_actor = YetoBufferTrainer + + def _initialize_trainer(self, config, total_train_steps: int) -> None: + """verl's ``_initialize_components`` (fork acad9875, fully_async_main.py:51) + up to the trainer, without the rollouter and the MessageQueue.""" + from omegaconf import OmegaConf + from verl.experimental.separation.utils import create_role_worker_mapping + from verl.utils import hf_processor, hf_tokenizer + from verl.utils.fs import copy_to_local + + OmegaConf.resolve(config) + local_path = copy_to_local(config.actor_rollout_ref.model.path, + use_shm=config.actor_rollout_ref.model.get("use_shm", False)) + trust_remote_code = config.data.get("trust_remote_code", False) + self.components["tokenizer"] = hf_tokenizer(local_path, trust_remote_code=trust_remote_code) + self.components["processor"] = hf_processor(local_path, trust_remote_code=trust_remote_code, use_fast=True) + self.components["config"] = config + mapping, group_cls = create_role_worker_mapping(config) + self.components["role_worker_mapping"] = mapping + self.components["ray_worker_group_cls"] = group_cls + self._create_trainer(config) + trainer = self.components["trainer"] + ray.get(trainer.set_total_train_steps.remote(int(total_train_steps))) + ray.get(trainer.load_checkpoint.remote()) + + def run(self, config, plan: dict): + from omegaconf import OmegaConf + + from yeto.island_credential_guard import check_island_credentials + from yeto.rl.engine.policy_release import PolicyReleaseStore + from yeto.rl.engine.resume import store_for + from yeto.rl.engine.sample_buffer import FileSampleBuffer + + from .buffer_feed import RowShape + from .buffer_ports import install_buffer_ports + from .trainer import run_fully_async_island + + check_island_credentials() # secret-handling-hardening D4 (buffer trainer Ray actor) + section = plan[SAMPLE_BUFFER_PLAN_KEY] + self._initialize_trainer(config, int(plan["global_rounds"])) + trainer = self.components["trainer"] + tokenizer = self.components["tokenizer"] + + def call(name, *args, timeout=None): + return ray.get(getattr(trainer, name).remote(*args), timeout=timeout) + + environ = dict(os.environ) + store = store_for(section["store"], environ=environ) + buffer = FileSampleBuffer(store, batch_size=1, create_index=False) + releases = PolicyReleaseStore(store_for(section.get("release_store") or section["store"], + environ=environ)) + shape = RowShape(int(config.data.max_prompt_length), int(config.data.max_response_length), + int(tokenizer.pad_token_id)) + + def ports(island): + return install_buffer_ports(island, buffer=buffer, releases=releases, shape=shape, + base_dir=plan["model_path"], tokenize=tokenizer.encode, + wait_s=float(section.get("wait_s") or 3600.0), + poll_s=float(section.get("poll_s") or 2.0)) + + return run_fully_async_island(call, None, plan, OmegaConf.to_container(config, resolve=True), + buffer_ports=ports) diff --git a/yeto/rl/adapters/verl/fully_async_runner.py b/yeto/rl/adapters/verl/fully_async_runner.py index 0d07f881e..900ccf6d3 100644 --- a/yeto/rl/adapters/verl/fully_async_runner.py +++ b/yeto/rl/adapters/verl/fully_async_runner.py @@ -96,7 +96,11 @@ async def yeto_assemble(self): self._collect_metrics_from_samples(batch, self.metrics) batch.meta_info["temperature"] = self.config.actor_rollout_ref.rollout.temperature self._yeto_batch = batch - return {"samples": len(kept), "queue_size": int(await self.message_queue_client.get_queue_size())} + return {"samples": len(kept), "queue_size": await self._yeto_queue_size()} + + async def _yeto_queue_size(self): + """Samples left in verl's MessageQueue (a buffer trainer has none: None).""" + return int(await self.message_queue_client.get_queue_size()) def _yeto_criteria(self, batch) -> dict: from yeto.rl.engine import mismatch_criteria @@ -196,7 +200,11 @@ async def yeto_export_lora(self): class YetoFullyAsyncTaskRunner(_TaskRunnerBase): """verl's component set-up (trainer, rollouter, MessageQueue, initial weight - sync); then the yeto driver instead of ``_run_training_loop``.""" + sync); then the yeto driver instead of ``_run_training_loop``. + + ``trainer_actor``: the Ray actor class :meth:`_create_trainer` builds.""" + + trainer_actor = YetoFullyAsyncTrainer def run(self, config, plan: dict): from yeto.island_credential_guard import check_island_credentials @@ -240,7 +248,7 @@ def call(name, *args, timeout=None): def _create_trainer(self, config) -> None: """verl's ``_create_trainer`` (fork acad9875, fully_async_main.py:138), building - :data:`YetoFullyAsyncTrainer`. Setting ``fa_main.FullyAsyncTrainer`` did not + :attr:`trainer_actor`. Setting ``fa_main.FullyAsyncTrainer`` did not reach verl's method inside the Ray actor (s19-verl64b-async7-20261010a: the trainer was verl's own class, "no attribute yeto_configure"), so the method is overridden here instead of relying on that module global.""" @@ -250,7 +258,7 @@ def _create_trainer(self, config) -> None: print("[ASYNC MAIN] Starting create trainer (yeto)...") mapping = {role: cls for role, cls in self.components["role_worker_mapping"].items() if role != Role.Rollout} - trainer = YetoFullyAsyncTrainer.remote( + trainer = self.trainer_actor.remote( config=config, tokenizer=self.components["tokenizer"], role_worker_mapping=mapping, @@ -261,4 +269,4 @@ def _create_trainer(self, config) -> None: trainer.yeto_configure # noqa: B018 - fail fast (AttributeError) if this is not our class ray.get(trainer.init_workers.remote()) self.components["trainer"] = trainer - print("[ASYNC MAIN] YetoFullyAsyncTrainer created and initialized successfully") + print("[ASYNC MAIN] yeto trainer actor created and initialized successfully") diff --git a/yeto/rl/adapters/verl/island_entry.py b/yeto/rl/adapters/verl/island_entry.py index 0508e441b..010e909cb 100644 --- a/yeto/rl/adapters/verl/island_entry.py +++ b/yeto/rl/adapters/verl/island_entry.py @@ -67,6 +67,11 @@ def parse_args(argv): p.add_argument("--rl-island-lease-s", type=float, default=None) p.add_argument("--rl-echo-events", action="store_true") p.add_argument("--rl-max-policy-age", type=int, default=0) + # S25 C2: a trainer island over the sample buffer (sample_roles; the sampler is + # the engine sampler island, yeto.rl.engine.engine_sampler) + p.add_argument("--rl-sample-role", default=None) + p.add_argument("--rl-resume-store", default=None) + p.add_argument("--rl-sample-wait-s", type=float, default=None) args, ignored = p.parse_known_args(argv) return args, ignored @@ -252,6 +257,7 @@ def ensure_ray_cpus(min_cpus: int = MIN_RAY_CPUS) -> dict: def main(argv=None) -> int: args, ignored = parse_args(sys.argv[1:] if argv is None else argv) run_config_rules.refuse_agentic_argv(ignored) # S25 N10.4: never ignore agentic flags + buffer_section = sample_buffer_section(args) # S25 C2: refused roles stop before any work OUTPUT.mkdir(parents=True, exist_ok=True) lid = args.learner_id t0 = time.time() @@ -342,6 +348,7 @@ def main(argv=None) -> int: "asserted": {k: _override_value(overrides, k) for k in asserted_keys}, "max_policy_age": limit, "fully_async": bool(limit), "versions": manifest["versions"], + **({"sample_buffer": buffer_section} if buffer_section else {}), } plan_path = work / f"plan-{lid}.json" plan_path.write_text(json.dumps(plan, indent=1, default=str)) @@ -357,7 +364,7 @@ def main(argv=None) -> int: env["YETO_VERL_READBACK_DIR"] = str(readback) env["VERL_FILE_LOGGER_PATH"] = str(OUTPUT / f"verl-file-logger-{lid}.jsonl") readback.mkdir(parents=True, exist_ok=True) - entry = "main_fully_async" if limit else "main" + entry = "main_buffer_trainer" if buffer_section else "main_fully_async" if limit else "main" cmd = [sys.executable, "-c", f"from yeto.rl.adapters.verl.verl_main import {entry}; {entry}()", *overrides] print("[yeto-verl] " + " ".join(shlex.quote(c) for c in cmd), flush=True) log = OUTPUT / f"verl-train-{lid}.log" @@ -372,6 +379,24 @@ def main(argv=None) -> int: return code +def sample_buffer_section(args) -> dict | None: + """The plan's ``sample_buffer`` section of a buffer trainer (S25 C2); None for + role ``both``. A verl island is never a sampler: a sampler island runs the + engine sampler (``python -m yeto.rl.engine.engine_sampler``) on any engine.""" + from yeto.rl.engine import sample_roles + + role = sample_roles.sample_role(args) + if role == sample_roles.BOTH: + return None + if role != sample_roles.TRAINER: + raise SystemExit(f"--rl-sample-role {role}: a verl island only trains from the buffer; " + "run the sampler as the engine sampler island (yeto.rl.engine.engine_sampler)") + if int(args.rl_max_policy_age or 0) < 1: + raise SystemExit("--rl-sample-role trainer needs --rl-max-policy-age >= 1 (the fully_async path)") + config = sample_roles.resolve(args, groups_per_round=args.groups_per_round) + return {"store": config.store, "wait_s": config.wait_s} + + def fully_async_plan(limit: int, overrides: list[str], spec_dict: dict | None, args, sync: str): """Limit > 0: fully_async overrides + asserted keys; the algorithm contract is checked here (bounded staleness needs ``execution.max_policy_staleness >= limit`` diff --git a/yeto/rl/adapters/verl/merged_release.py b/yeto/rl/adapters/verl/merged_release.py new file mode 100644 index 000000000..a4d829f6c --- /dev/null +++ b/yeto/rl/adapters/verl/merged_release.py @@ -0,0 +1,112 @@ +"""A verl LoRA policy as a full HF checkpoint for a sampler island (S25 C2). + +The engine sampler island loads full checkpoints only (``SamplerEngine.load``: +the SGLang and vllm-ascend adapters reload weights from a model directory). A +verl buffer trainer trains a LoRA adapter over a frozen base, so each release +merges the adapter into the base: ``W' = W + B @ A`` for every LoRA module +(yeto canonical LoRA has no scaling: ``lora_alpha == rank``, +``config.validate``), summed in float32 and stored in the base tensor's dtype. + +The output has the base's file layout (the same safetensors file names, the same +index, the same non-weight files: config, tokenizer, generation config), so the +engine reads it exactly as it reads the base. A module the adapter names that +the base lacks, or an A/B pair with mismatched shapes, is an error: a silently +skipped module would serve the base weights for it. + +Rounding: a bf16 base keeps about 8 bits of mantissa, so a LoRA delta far below +``|W| / 256`` does not change the stored weight. The sampler then serves the +merged weights, not the trainer's exact ``W + BA``; the trainer's log-prob +recomputation shows the difference as off-policy mismatch. +""" + +from __future__ import annotations + +import json +import shutil +from collections.abc import Mapping +from pathlib import Path +from typing import Any + +from .param_names import CANONICAL_PREFIX, ParamNameError + +INDEX_FILE = "model.safetensors.index.json" +SINGLE_FILE = "model.safetensors" +LORA_A, LORA_B = ".lora_A.weight", ".lora_B.weight" + + +class MergeError(ValueError): + """The adapter does not fit the base checkpoint.""" + + +def lora_pairs(lora: Mapping[str, Any]) -> dict[str, tuple[Any, Any]]: + """{base weight name: (A, B)} from canonical LoRA tensors + (``base_model.model..lora_{A,B}.weight``).""" + pairs: dict[str, dict[str, Any]] = {} + for name, tensor in lora.items(): + if not name.startswith(CANONICAL_PREFIX) or not name.endswith((LORA_A, LORA_B)): + raise ParamNameError(f"not a canonical LoRA name: {name!r}") + side = "A" if name.endswith(LORA_A) else "B" + module = name[len(CANONICAL_PREFIX):-len(LORA_A)] + pairs.setdefault(f"{module}.weight", {})[side] = tensor + out = {} + for weight, sides in sorted(pairs.items()): + if set(sides) != {"A", "B"}: + raise MergeError(f"{weight}: LoRA has only side {sorted(sides)}") + out[weight] = (sides["A"], sides["B"]) + return out + + +def weight_files(base_dir: Path) -> dict[str, list[str]]: + """{safetensors file: [tensor names]} of a HF checkpoint directory.""" + index = base_dir / INDEX_FILE + if index.is_file(): + files: dict[str, list[str]] = {} + for name, file in json.loads(index.read_text(encoding="utf-8"))["weight_map"].items(): + files.setdefault(file, []).append(name) + return files + if (base_dir / SINGLE_FILE).is_file(): + from safetensors import safe_open + + with safe_open(str(base_dir / SINGLE_FILE), framework="pt") as handle: + return {SINGLE_FILE: list(handle.keys())} + raise MergeError(f"{base_dir}: no {INDEX_FILE} and no {SINGLE_FILE}") + + +def write_merged_checkpoint(base_dir: str | Path, lora: Mapping[str, Any], out_dir: str | Path) -> dict[str, Any]: + """Write base + LoRA into ``out_dir`` (module doc); returns a summary + (merged modules, files, the largest relative change of a merged weight).""" + import torch + from safetensors.torch import load_file, save_file + + base, out = Path(base_dir), Path(out_dir) + out.mkdir(parents=True, exist_ok=True) + pairs = lora_pairs(lora) + files = weight_files(base) + known = {name for names in files.values() for name in names} + missing = sorted(set(pairs) - known) + if missing: + raise MergeError(f"the base has no weight for {len(missing)} LoRA modules, e.g. {missing[0]}") + merged, max_rel = 0, 0.0 + for file in sorted(files): + tensors = load_file(str(base / file)) + for name in files[file]: + if name not in pairs: + continue + a, b = (t.detach().to("cpu", torch.float32) for t in pairs[name]) + weight = tensors[name] + if a.dim() != 2 or b.dim() != 2 or b.shape[1] != a.shape[0] or \ + (b.shape[0], a.shape[1]) != tuple(weight.shape): + raise MergeError(f"{name}: A {tuple(a.shape)}, B {tuple(b.shape)} do not fit " + f"the weight {tuple(weight.shape)}") + delta = b @ a + new = (weight.to(torch.float32) + delta).to(weight.dtype) + norm = float(weight.to(torch.float32).norm()) + if norm > 0: + max_rel = max(max_rel, float((new.to(torch.float32) - weight.to(torch.float32)).norm()) / norm) + tensors[name] = new.contiguous() + merged += 1 + save_file(tensors, str(out / file), metadata={"format": "pt"}) + for item in sorted(base.iterdir()): + if item.is_file() and item.name not in files and not item.name.startswith("."): + shutil.copy2(item, out / item.name) + return {"merged_modules": merged, "files": sorted(files), "max_relative_change": max_rel} diff --git a/yeto/rl/adapters/verl/trainer.py b/yeto/rl/adapters/verl/trainer.py index 72e7859e0..5febc9718 100644 --- a/yeto/rl/adapters/verl/trainer.py +++ b/yeto/rl/adapters/verl/trainer.py @@ -275,11 +275,15 @@ def run_island(trainer, agent_loop_manager, plan: dict, cfg: dict) -> dict: return result -def run_fully_async_island(call, start_rollouter, plan: dict, cfg: dict) -> dict: +def run_fully_async_island(call, start_rollouter, plan: dict, cfg: dict, *, buffer_ports=None) -> dict: """agentic-rollout-utilization 6.4b: the yeto driver over the fully_async ports. ``call(method, *args)`` reaches the trainer actor; ``start_rollouter`` starts the - rollouter's generation loop (called once, after the first verified publication).""" + rollouter's generation loop (called once, after the first verified publication). + + ``buffer_ports(island)``: a buffer trainer (S25 C2, ``buffer_ports.install_buffer_ports`` + bound to the run's buffer and release store) returns ``(rollout, publisher, + placement)``; the trainer port and the policy state stay the fully_async ones.""" from yeto.rl.core import canonical_layout_hash, canonical_lora_config_hash from yeto.rl.engine.algorithm import AlgorithmSpec from yeto.rl.engine.driver import EventTape, IslandDriver @@ -317,16 +321,22 @@ def run_fully_async_island(call, start_rollouter, plan: dict, cfg: dict) -> dict tape.append({"event": "rl_verl_island_start", "plan": plan, "backend_identity": identity.to_dict(), "backend_identity_sha256": identity_sha, "layout_hash": layout_hash, "lora_config_hash": lora_hash, "n_specs": len(specs), "asserted": asserted, - "fully_async": True, "max_policy_age": limit, "init_seconds": time.time() - t0}) + "fully_async": True, "max_policy_age": limit, "init_seconds": time.time() - t0, + **({"sample_role": "trainer"} if buffer_ports is not None else {})}) sync = build_sync(plan, specs, lora_hash=lora_hash, layout_hash=layout_hash, identity_sha256=identity_sha, compat_group=identity.compat_group) spec = AlgorithmSpec.from_dict(plan["algorithm_spec"]) if plan.get("algorithm_spec") else AlgorithmSpec() profile = execution_profile_for(spec, limit, groups_per_round=int(plan["groups_per_round"]), samples_per_group=int(plan["samples_per_group"]), sync=plan["sync"]) - publisher = FullyAsyncPublisher(island, strict=bool(plan.get("publish_strict", True))) + if buffer_ports is None: + rollout, publisher, placement = (FullyAsyncRolloutPool(island), + FullyAsyncPublisher(island, strict=bool(plan.get("publish_strict", True))), + FullyAsyncPlacement()) + else: + rollout, publisher, placement = buffer_ports(island) driver = IslandDriver( - learner_id=learner_id, rollout=FullyAsyncRolloutPool(island), trainer=FullyAsyncTrainerGroup(island), - policy_state=VerlPolicyState(island), publisher=publisher, placement=FullyAsyncPlacement(), + learner_id=learner_id, rollout=rollout, trainer=FullyAsyncTrainerGroup(island), + policy_state=VerlPolicyState(island), publisher=publisher, placement=placement, capabilities=verl_capabilities(_fingerprint(cfg), fully_async_limit=limit), algorithm=spec, sync=sync, events=tape, profile=profile) state = driver.run() diff --git a/yeto/rl/adapters/verl/verl_main.py b/yeto/rl/adapters/verl/verl_main.py index ff0989c13..106a0093f 100644 --- a/yeto/rl/adapters/verl/verl_main.py +++ b/yeto/rl/adapters/verl/verl_main.py @@ -140,6 +140,28 @@ def run_ppo(config, task_runner_class=None): fully_async_hydra_entry(fa_main)() +def main_buffer_trainer(argv: list[str] | None = None) -> None: + """S25 C2: the fully_async hydra config and fix-ups as :func:`main_fully_async`, + with :class:`~.buffer_trainer_runner.YetoBufferTaskRunner`: a trainer island + over the sample buffer (the plan's ``sample_buffer`` section), no rollouter.""" + 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() + + def run_ppo(config, task_runner_class=None): + from yeto.rl.adapters.verl.buffer_trainer_runner import YetoBufferTaskRunner + + runner = ray.remote(num_cpus=1)(YetoBufferTaskRunner) + result = ray.get(runner.remote().run.remote(config, plan)) + print("[yeto-verl] island result " + json.dumps(result, sort_keys=True, default=str), flush=True) + + main_ppo.run_ppo = run_ppo + if argv is not None: + sys.argv = [sys.argv[0], *argv] + fully_async_hydra_entry(fa_main)() + + def fully_async_hydra_entry(fa_main): """``fa_main.main`` re-decorated with an absolute config directory. From 39aa81ffcbe991b34ea3b64ad699c945825538f2 Mon Sep 17 00:00:00 2001 From: michael <78113554+chowhsu@users.noreply.github.com> Date: Tue, 20 Oct 2026 04:16:08 +0000 Subject: [PATCH 28/89] Wire the yeto round-cut controller into the verl island; digest the stepped state (S25 N7.6b, N7.6c) N7.6b: the round-cut controller (RoundCutCheckpoint, CutSource) moves from the Miles adapter to yeto/rl/engine/round_cut.py, so both islands run one controller; the Miles modules re-export it and keep their in-flight export check. The data cursor of a cut now comes from the cut's own cursor_fields (verl: sample_offset). adapters/verl/round_cuts.py wires it into the verl sync island the way the Miles island does: --rl-resume-store / --rl-resume-state-dir / --rl-cut-every / --rl-cut-keep / --rl-stop-after-rounds / --rl-resume-allow-config-change, a ResumeController that restores the store before the driver is built (the driver's ledger is the restored one), a preemption marker handler, and the trainer_remote commanded mode whose stop cuts with the SCHEDULED trigger. Synced islands and the fully_async path are refused. N7.6c: the verl rank digest is taken over optimizer.param_groups and optimizer.state (the tensors the optimizer steps), not the module state dict. A regression test injects an optimizer that steps a build-time snapshot after the load and requires the restore to be refused. Co-Authored-By: Claude Opus 5.5 --- tests/test_rl_verl_cut.py | 39 ++ tests/test_rl_verl_round_cuts.py | 240 +++++++++++ yeto/rl/adapters/miles/rebuild_wiring.py | 107 +---- yeto/rl/adapters/miles/round_cut.py | 388 +----------------- yeto/rl/adapters/verl/cut.py | 68 ++-- yeto/rl/adapters/verl/island_entry.py | 14 + yeto/rl/adapters/verl/round_cuts.py | 155 +++++++ yeto/rl/adapters/verl/trainer.py | 9 +- yeto/rl/engine/round_cut.py | 493 +++++++++++++++++++++++ 9 files changed, 1014 insertions(+), 499 deletions(-) create mode 100644 tests/test_rl_verl_round_cuts.py create mode 100644 yeto/rl/adapters/verl/round_cuts.py create mode 100644 yeto/rl/engine/round_cut.py diff --git a/tests/test_rl_verl_cut.py b/tests/test_rl_verl_cut.py index 0736c287a..5b15d0641 100644 --- a/tests/test_rl_verl_cut.py +++ b/tests/test_rl_verl_cut.py @@ -271,3 +271,42 @@ class _Cfg(dict): """Attribute + ``get`` access, like an omegaconf node.""" __getattr__ = dict.__getitem__ + + +# ---------------------------------------------------------------- N7.6c: digest over the stepped state + + +def test_restore_whose_optimizer_steps_a_stale_snapshot_is_refused(tmp_path, monkeypatch): + """Inject the T-G3 bug shape: after the load the module holds the cut's weights, but the + optimizer steps a snapshot taken when the fresh trainer was built. The restore digest + reads the stepped tensors, so the restore is refused (a module-side digest would match).""" + import torch.distributed.checkpoint.state_dict as dcp_sd + + real = dcp_sd.set_state_dict + + def load_then_step_stale(module, optimizer, **kwargs): + stale = {id(p): p.detach().clone() for p in module.parameters()} # build-time values + real(module, optimizer, **kwargs) + for group in optimizer.param_groups: + for i, param in enumerate(group["params"]): + copy = torch.nn.Parameter(stale[id(param)]) + optimizer.state[copy] = optimizer.state.pop(param) + group["params"][i] = copy + + monkeypatch.setattr(dcp_sd, "set_state_dict", load_then_step_stale) + trainer, fresh = _single(tmp_path) + saved = torch.load(tmp_path / "cut-v" / vcut.rank_file(0), weights_only=True)["model"] + with pytest.raises(ResumeAssertionError, match="loaded train state differs from the cut"): + trainer.restore_cut("cut-v", epoch=1, root=str(tmp_path), expect=_expect()) + for name, p in fresh.module.named_parameters(): # the module copy alone looks restored + assert torch.equal(p.detach(), saved[name]), name + + +def test_digest_reads_the_tensors_the_optimizer_steps(): + h = _handles(_model()) + _round(h, 0) + before = vcut._digest(h) + first = h.optimizer.param_groups[0]["params"][0] + with torch.no_grad(): + first.add_(1.0) + assert vcut._digest(h) != before diff --git a/tests/test_rl_verl_round_cuts.py b/tests/test_rl_verl_round_cuts.py new file mode 100644 index 000000000..84a12d548 --- /dev/null +++ b/tests/test_rl_verl_round_cuts.py @@ -0,0 +1,240 @@ +"""S25 N7.6b (CPU): the verl island's round-cut controller. + +The real cut controller (ResumeController + RoundCutCheckpoint on LocalOnlySync, +wired by ``adapters.verl.round_cuts``) and the real verl cut trainer +(``VerlCutTrainer``: rank file, digests, resume assertions) over a torch model, +behind the fake driver ports. The policy the driver publishes is read from the +torch model, so the resume's policy-hash check compares real restored weights. +""" + +from __future__ import annotations + +from types import SimpleNamespace + +import pytest +import torch +from test_rl_verl_cut import K, _handles, _model, _round + +from yeto.rl.adapters.verl import cut as vcut +from yeto.rl.adapters.verl import round_cuts +from yeto.rl.elastic_client import LeaveReason +from yeto.rl.engine import recovery, trainer_remote +from yeto.rl.engine.algorithm import AlgorithmSpec +from yeto.rl.engine.bridges import LocalOnlySync +from yeto.rl.engine.driver import EventTape, IslandDriver +from yeto.rl.engine.fake import FakeEngine, FakeRolloutPool, FakeTrainerGroup, fake_capabilities +from yeto.rl.engine.recovery import CUT_KEY +from yeto.rl.engine.resume_manifest import ResumeAssertionError + +NAME = "base_model.model.layer.lora_A.weight" +PROMPTS = 4 # prompts drawn per round +ROUNDS = 5 + + +class _Pool(FakeRolloutPool): + """The fake pool with the verl data cursor shape (``sample_offset`` only).""" + + def __init__(self, engine): + super().__init__(engine) + self.drawn = 0 + + def generate(self, rollout_id, **kwargs): + self.drawn += PROMPTS + return super().generate(rollout_id, **kwargs) + + def data_cursor(self): + return {"sample_offset": self.drawn} + + def seek_data_cursor(self, cursor): + target = int(cursor["sample_offset"]) + if target < self.drawn or (target - self.drawn) % PROMPTS: + raise ValueError(f"cannot seek {self.drawn} -> {target}") + self.drawn = target + return self.data_cursor() + + +class _TorchTrainer(FakeTrainerGroup): + """Fake trainer port whose train state is a torch model behind ``VerlCutTrainer``.""" + + def __init__(self, engine, handles, *, skip_optimizer_round=None): + super().__init__(engine) + self.handles = handles + worker = SimpleNamespace(module=handles.module, optimizer=handles.optimizer, + lr_scheduler=handles.scheduler) + self.cuts = vcut.VerlCutTrainer(lambda fn, *a: [fn(worker, *a)], optimizer_steps_per_round=K) + self.skip_optimizer_round = skip_optimizer_round + self._sync_policy() + + def _sync_policy(self): + self.engine.tensors[NAME] = self.handles.module[0].weight.detach()[:1, :2].clone().float() + + def actual_layout(self): + return self.cuts.actual_layout() + + def layout(self): + return self.actual_layout() + + def save_cut(self, *, epoch, context): + return self.cuts.save_cut(epoch=epoch, context=context) + + def restore_cut(self, cut_id, *, epoch, root, expect, shared_filesystem=True): + manifest = self.cuts.restore_cut(cut_id, epoch=epoch, root=root, expect=expect, + shared_filesystem=shared_filesystem) + self._sync_policy() + return manifest + + def confirm_resumed_data_cursor(self, cursor): + self.cuts.confirm_resumed_data_cursor(cursor) + + def train_step(self, batch): + receipt = super().train_step(batch) + if batch.rollout_id == self.skip_optimizer_round: + self.handles.scheduler.step() # injected: the round took no optimizer step + else: + _round(self.handles, batch.rollout_id) + self._sync_policy() + self.cuts.check_first_step(1) + return receipt + + +def _plan(store=None, state_dir=None, **resume): + config = None + if store is not None: + config = {"store": str(store), "state_dir": str(state_dir), "every": 1, "keep": 2, + "allow_config_change": False, "stop_after": None, **resume} + return {"sync": "none", "resume": config, "model_revision": "rev", "lora_targets": "all-linear", + "run_config": {"lr": 1e-2, "seed": 1, "total_rounds": ROUNDS, "out_dir": "/x"}} + + +def _island(tmp_path, plan, *, name, rounds=ROUNDS, **trainer_kwargs): + engine = FakeEngine(tensors={NAME: torch.zeros(1, 2)}, step_delta=1.0) + engine.rollout = _Pool(engine) + handles = _handles(_model()) + engine.trainer = _TorchTrainer(engine, handles, **trainer_kwargs) + resume = round_cuts.resume_wiring_for(plan, runtime_fingerprint="verl@test") + tape = tmp_path / f"{name}.jsonl" + driver = IslandDriver( + learner_id=0, rollout=engine.rollout, trainer=engine.trainer, policy_state=engine.policy_state, + publisher=engine.publisher, placement=engine.placement, algorithm=AlgorithmSpec(), + capabilities=fake_capabilities("miles-like"), sync=LocalOnlySync(rounds), + events=EventTape(tape, 0), ledger=None if resume is None else resume.ledger) + checkpoint = round_cuts.wire_round_cuts(driver, resume, plan=plan, algorithm=AlgorithmSpec(), + install_signals=False) + return driver, handles, checkpoint, resume, tape + + +def _events(path, kind): + import json + + return [e for e in map(json.loads, path.read_text().splitlines()) if e.get("event") == kind] + + +def _params(handles): + return {n: p.detach().clone() for n, p in handles.module.named_parameters()} + + +def test_stop_after_then_resume_equals_continuous_run(tmp_path): + store = tmp_path / "store" + driver, _, checkpoint, _, tape = _island(tmp_path, _plan(store, tmp_path / "a", stop_after=3), name="a") + round_cuts.run_driver(driver, None) + pointer = checkpoint.pointer() + assert pointer["next_rollout_id"] == 3 and pointer["final"] is True + assert pointer[CUT_KEY]["trigger"] == recovery.MANUAL + assert [e["rollout_id"] for e in _events(tape, "rl_cut_saved")] == [1, 2, 3] + + # a fresh container: new state dir, restored from the store's LATEST + driver, resumed, checkpoint, _, tape = _island(tmp_path, _plan(store, tmp_path / "b"), name="b") + round_cuts.run_driver(driver, None) + [resume] = _events(tape, "rl_resume") + assert resume["rollout_id"] == 3 and resume["checks"]["policy_hash"] == "same" + assert _events(tape, "rl_data_cursor_restored")[0]["data_cursor"] == {"sample_offset": 3 * PROMPTS} + assert checkpoint.pointer()["next_rollout_id"] == ROUNDS + assert driver.rounds_completed == ROUNDS - 3 + + driver, continuous, _, _, _ = _island(tmp_path, _plan(), name="c") + round_cuts.run_driver(driver, None) + a, b = _params(resumed), _params(continuous) + assert all(torch.equal(a[n], b[n]) for n in a) + assert resumed.scheduler.last_epoch == continuous.scheduler.last_epoch == ROUNDS + + +def test_cut_every_and_keep(tmp_path): + store = tmp_path / "store" + driver, _, _, _, tape = _island(tmp_path, _plan(store, tmp_path / "a", every=2, keep=1), name="a") + round_cuts.run_driver(driver, None) + assert [e["rollout_id"] for e in _events(tape, "rl_cut_saved")] == [2, 4, 5] # every 2, and the end + + +def test_first_resumed_round_without_an_optimizer_step_is_refused(tmp_path): + store = tmp_path / "store" + driver, _, _, _, _ = _island(tmp_path, _plan(store, tmp_path / "a", stop_after=2), name="a") + round_cuts.run_driver(driver, None) + driver, _, _, _, _ = _island(tmp_path, _plan(store, tmp_path / "b"), name="b", skip_optimizer_round=2) + with pytest.raises(Exception) as info: + round_cuts.run_driver(driver, None) + chain, error = [], info.value + while error is not None: + chain.append(error) + error = error.__cause__ or error.__context__ + assert any(isinstance(e, ResumeAssertionError) and "Adam step" in str(e) for e in chain) + + +def test_scheduler_stop_cuts_with_scheduled_and_resumes(tmp_path): + """The head's trainer scheduler: run 2 rounds, stop (STOPPED -> SCHEDULED cut); + the next incarnation restores that cut.""" + store = tmp_path / "store" + driver, _, checkpoint, resume, tape = _island(tmp_path, _plan(store, tmp_path / "a"), name="a") + files = trainer_remote.LocalStoreFiles(resume.controller.store) + writer = trainer_remote.CommandWriter(files, 0) + writer.send(trainer_remote.RUN, rounds=2) + writer.send(trainer_remote.STOP, reason=LeaveReason.STOPPED) + env = {trainer_remote.CONTROL_ENV: "1", trainer_remote.INCARNATION_ENV: "0", + trainer_remote.RESUME_ENV: "", trainer_remote.POLL_ENV: "0.01"} + round_cuts.run_driver(driver, resume, environ=env) + pointer = checkpoint.pointer() + assert pointer["next_rollout_id"] == 2 and pointer[CUT_KEY]["trigger"] == recovery.SCHEDULED + phases = [e["phase"] for e in _events(tape, trainer_remote.EVENT)] + assert phases == ["started", "ready", "trained", "stopped"] + + driver, _, _, _, tape = _island(tmp_path, _plan(store, tmp_path / "b"), name="b") + round_cuts.run_driver(driver, None) + assert _events(tape, "rl_resume")[0]["rollout_id"] == 2 + + +def test_commanded_run_needs_round_cuts(tmp_path): + driver, _, _, _, _ = _island(tmp_path, _plan(), name="a") + with pytest.raises(RuntimeError, match="needs round cuts"): + round_cuts.run_driver(driver, None, environ={trainer_remote.CONTROL_ENV: "1"}) + + +def _args(**kw): + base = dict(rl_resume_store=None, rl_resume_state_dir=None, rl_cut_every=None, rl_cut_keep=None, + rl_stop_after_rounds=None, rl_resume_allow_config_change=False, + rl_single_island_no_sync=True, rl_max_policy_age=0) + return SimpleNamespace(**{**base, **kw}) + + +def test_resume_config_from_flags(): + assert round_cuts.resume_config(_args()) is None + config = round_cuts.resume_config(_args(rl_resume_store="/s", rl_cut_every=2, rl_stop_after_rounds=3)) + assert config == {"store": "/s", "state_dir": round_cuts.DEFAULT_STATE_DIR, "every": 2, "keep": 2, + "allow_config_change": False, "stop_after": 3} + with pytest.raises(ValueError, match="need --rl-resume-store"): + round_cuts.resume_config(_args(rl_cut_every=2)) + with pytest.raises(ValueError, match="--rl-single-island-no-sync"): + round_cuts.resume_config(_args(rl_resume_store="/s", rl_single_island_no_sync=False)) + with pytest.raises(ValueError, match="fully_async"): + round_cuts.resume_config(_args(rl_resume_store="/s", rl_max_policy_age=2)) + with pytest.raises(ValueError, match="must be >= 1"): + round_cuts.resume_config(_args(rl_resume_store="/s", rl_cut_keep=0)) + + +def test_changed_run_config_is_refused(tmp_path): + store = tmp_path / "store" + driver, _, _, _, _ = _island(tmp_path, _plan(store, tmp_path / "a", stop_after=2), name="a") + round_cuts.run_driver(driver, None) + plan = _plan(store, tmp_path / "b") + plan["run_config"] = {**plan["run_config"], "lr": 2e-2, "total_rounds": 9} # rounds may change + driver, _, _, _, _ = _island(tmp_path, plan, name="b") + with pytest.raises(Exception, match="run configuration differs"): + round_cuts.run_driver(driver, None) diff --git a/yeto/rl/adapters/miles/rebuild_wiring.py b/yeto/rl/adapters/miles/rebuild_wiring.py index ca2009f52..36f5e60c1 100644 --- a/yeto/rl/adapters/miles/rebuild_wiring.py +++ b/yeto/rl/adapters/miles/rebuild_wiring.py @@ -29,104 +29,17 @@ from collections.abc import Callable, Mapping from typing import Any +from yeto.rl.engine.round_cut import CutSource as _CoreCutSource -def _algorithm_identity(algorithm: Any, *, ref_model: Mapping[str, Any] | None) -> Any: - from yeto.rl.engine.cut import AlgorithmIdentity - - runtime_attrs = None - to_attrs = getattr(algorithm, "to_legacy_runtime_attrs", None) - if callable(to_attrs): - runtime_attrs = dict(to_attrs()) or None - return AlgorithmIdentity.from_spec(algorithm, runtime_attrs=runtime_attrs, ref_model=ref_model) - - -class CutSource: - """What the driver contributes to a cut (E2 ``CutContext``) and what a - restore must match (``RestoreExpectation``); shared by the 4.4 same-shape - rebuilder and E3's ``MilesTrainerOps`` (4.7).""" - - def __init__(self, *, driver: Any, trainer: Any, rollout: Any, ledger: Any, algorithm: Any, - backend_fingerprint: str, cut_root: str, global_batch_size: int, - ref_model: Mapping[str, Any] | None = None, - shared_filesystem: bool = True) -> None: - self.driver_ref = driver # the driver, or a zero-arg callable returning it - self.trainer = trainer - self.rollout = rollout - self.ledger = ledger - self.identity = _algorithm_identity(algorithm, ref_model=ref_model) - self.backend_fingerprint = backend_fingerprint - self.cut_root = cut_root - self.global_batch_size = int(global_batch_size) - self.shared_filesystem = shared_filesystem - - @property - 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: - from yeto.rl.engine.controller import RebuildRefused - from yeto.rl.engine.cut import CutProgress - from .trainer import CutContext - - driver = self.driver - state = driver.published_state - if state is None or driver.published_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) - if cursor is None: - raise RebuildRefused("rollout data cursor unknown (needs the elastic metadata)") - data = dict(cursor) - summary = dict(self.ledger.cut_summary()) if self.ledger is not None else {} - if summary.get("engine_buffer_length") is not None: - data.setdefault("buffer_length", summary["engine_buffer_length"]) - in_flight: tuple = () - limit = int(getattr(getattr(driver, "profile", None), "max_policy_age", 0) or 0) - if limit > 0: - # agentic-rollout-utilization 3.3: the unfinished work the rollout - # process holds NOW (buffered groups, suspended agentic trajectories). - exporter = getattr(self.rollout, "export_in_flight", None) - if not callable(exporter): - raise RebuildRefused("the rollout pool cannot export its in-flight trajectories") - from .in_flight import check_export - - exported = exporter() - problems = check_export(exported) - if problems: - raise RebuildRefused("; ".join(problems)) - in_flight = tuple(exported["entries"]) - data["buffer_length"] = int(exported["buffer_groups"]) - summary["carried_over"] = len(in_flight) - summary["max_policy_age"] = limit - progress = CutProgress( - 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), - policy_version=int(state.policy_version), - policy_hash=state.policy_tensor_hash(), - ) - return CutContext( - root=self.cut_root, cut_id=cut_id, backend_fingerprint=self.backend_fingerprint, - progress=progress, algorithm=self.identity, data=data, ledger=summary, - # The safe point is after boundary() returned non-stop and every - # member acknowledged the policy: the outer commit is settled. - outer={"settled": bool(driver.at_safe_point), "policy_version": state.policy_version, - "policy_token": driver.expected_token}, - shared_filesystem=self.shared_filesystem, - in_flight=in_flight, - ) - - def expectation(self, layout: Mapping[str, int], *, epoch: int) -> Any: - from yeto.rl.engine.cut import RestoreExpectation - - driver = self.driver - return RestoreExpectation( - algorithm=self.identity, layout=dict(layout), - backend_fingerprint=self.backend_fingerprint, local_step=int(driver.local_step), - policy_version=int(driver.published_state.policy_version), epoch=epoch, - ) + +class CutSource(_CoreCutSource): + """:class:`yeto.rl.engine.round_cut.CutSource` with the Miles in-flight export check.""" + + def __init__(self, **kwargs: Any) -> None: + from .in_flight import check_export + + kwargs.setdefault("check_in_flight", check_export) + super().__init__(**kwargs) def make_trainer_rebuilder( diff --git a/yeto/rl/adapters/miles/round_cut.py b/yeto/rl/adapters/miles/round_cut.py index f43932acc..bfef3fbd8 100644 --- a/yeto/rl/adapters/miles/round_cut.py +++ b/yeto/rl/adapters/miles/round_cut.py @@ -1,380 +1,14 @@ -"""rl-multinode-island M4: policy-weight continuation across a machine replacement. - -Without an outer syncer (``--rl-single-island-no-sync``, :class:`~..bridges.LocalOnlySync`) -a restarted learner used to start again at rollout 0 from the startup weights: the -checkpoint store (Q4, C5) carried the journal/ledger, but no LoRA/optimizer state. - -With ``--rl-elastic-checkpoint-store`` set this module keeps a 4.2 cut of the trainer -(every rank's LoRA + optimizer + scheduler + RNG shard, manifest committed last) at -every round-boundary safe point and resumes from the newest one at start: - -* the cut root is ``/round-cuts`` -- the store is visible on every node (sky - Storage MOUNT of the bucket), so the shard of a rank on another node (PP across - nodes) lands where the driver can check it (``shared_filesystem=True``); -* after the manifest is committed, ``round-cut.json`` (the pointer: cut id, next - rollout id, local step, policy hash, writing incarnation) is written into the state - dir and the state dir is synced to the store -- pointer, ledger and journal travel - in ONE store snapshot (``STORE-MANIFEST.json`` last), so a restored pointer never - names a cut newer than the restored ledger; -* at start (in-place restart or a fresh machine after the store restore) the cut the - pointer names is verified (``verify_cut``: algorithm, layout, backend, step, - policy version, every shard digest) and loaded on the fresh trainer - (``MilesTrainerGroup.restore_cut``); the driver then rebases the ledger and seeks - the data source to the pointer's rollout id. Any failure is fatal (no silent - restart at 0 once a pointer exists). -""" +"""Miles round cuts: the cut controller moved to :mod:`yeto.rl.engine.round_cut` +(S25 N7.6b, shared with the verl island); this module keeps the Miles import path.""" from __future__ import annotations -import json -import os -from pathlib import Path -from typing import Any - -from yeto.rl.engine import recovery -from yeto.rl.engine.recovery import CUT_KEY, FINAL_TRIGGERS, MANUAL, PERIODIC, ConsistentCut -from yeto.rl.engine.resume import ROUND_CUTS # under the checkpoint store (not under the state dir) - -ISLAND = "local" # round cuts run without an outer syncer: one island, the local one -POINTER = "round-cut.json" # under the state dir (synced with the ledger/journal) - - -class RoundCutError(RuntimeError): - pass - - -class RoundCutCheckpoint: - """``save`` at a safe point, ``resume`` at sync start (see the module doc). - - rl-resume-from-checkpoint (S17 C4): ``every`` (``--rl-cut-every``) cuts only every N - rounds (plus the last round and our own stop, ``final=True``); the pointer carries - ``run_fingerprint`` / ``lr_at_next_round`` / ``consumed_prompt_ids_digest``; a resume - checks the fingerprint (refused on a difference unless ``allow_config_change``), the - consumed-group digest and, on the first resumed round, the learning rate.""" - - 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: - if getattr(controller, "checkpoint_store", None) is None: - raise ValueError("round cuts need the checkpoint store") - if int(every) < 1: - raise ValueError("--rl-cut-every must be >= 1") - self.controller = controller - self.source = source # rebuild_wiring.CutSource (cut_root is set here) - self.trainer = trainer # MilesTrainerGroup: save_cut / restore_cut / actual_layout - self.root = str(Path(controller.checkpoint_store) / ROUND_CUTS) - source.cut_root = self.root - self.resumed: dict[str, Any] | None = None - self.every = int(every) - self.fingerprint = dict(fingerprint) if fingerprint is not None else None - self.allow_config_change = bool(allow_config_change) - self.next_lr = next_lr # callable(driver) -> list[float] | None - self.pending_lr_check: dict[str, Any] | None = None - # set by resume(): where the rollout data source continues (the cut's cursor) - self.resume_data_cursor: dict[str, int] | None = None - # rl-trainer-start-stop D4: the durable sample buffer's consumed watermark, - # written into the cut record (None: the run reads no sample buffer). - # buffer_watermark(next_rollout_id) commits and returns the position after the - # last trained round; buffer_resume(watermark) continues there after a restore - # (buffer_rollout.wire_sample_buffer sets both). - self.buffer_watermark: Any = None # callable(int) -> int | None - self.buffer_resume: Any = None # callable(int) -> int - - # ---------------------------------------------------------------- pointer - @property - def pointer_path(self) -> Path: - return Path(self.controller.state_dir) / POINTER - - def pointer(self) -> dict[str, Any] | None: - path = self.pointer_path - if not path.is_file(): - 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") - 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) - - def _epoch(self) -> int: - return int(self.controller.journal.epochs.config_epoch) - - def _consumed(self) -> tuple[int, str | None]: - ledger = getattr(self.source, "ledger", None) - groups = getattr(ledger, "_consumed_groups", None) - if groups is None: - return 0, None - from yeto.rl.engine.resume import consumed_digest - - return len(groups), consumed_digest(groups) - - # ---------------------------------------------------------------- save - def due(self, rollout_id: int, *, final: bool = False) -> bool: - from yeto.rl.engine.resume import should_cut - - return should_cut(int(rollout_id), every=self.every, final=final) - - 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 - version), point the state dir at it, sync the store. Returns the pointer, or None - 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). - - ``trigger`` names a final cut's trigger (default ``MANUAL``; ``SCHEDULED`` for - a stop the trainer scheduler ordered); a periodic cut takes none.""" - import time - - from yeto.rl.engine.resume import directory_bytes - - 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): - 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']}" - 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 - progress = context.progress - consumed_n, consumed_sha = self._consumed() - next_lr = self.next_lr(driver) if callable(self.next_lr) else None - info = { - "cut_id": cut_id, "root": ROUND_CUTS, "epoch": self._epoch(), - "next_rollout_id": int(progress.next_rollout_id), - "local_step": int(progress.local_step), - "policy_version": int(progress.policy_version), - "policy_hash": progress.policy_hash, - "incarnation": ctl.incarnation["id"], - "incarnation_index": ctl.incarnation.get("index"), - "run_fingerprint": self.fingerprint, - "lr_at_next_round": None if next_lr is None else [float(x) for x in next_lr], - "consumed_groups": consumed_n, - "consumed_prompt_ids_digest": consumed_sha, - "outer_version": getattr(getattr(driver, "sync", None), "outer_version", None), - "final": bool(final), - } - cut = ConsistentCut( - checkpoint_id=cut_id, trigger=(trigger or MANUAL) if final else PERIODIC, - outer_step=info["outer_version"], island_steps={ISLAND: info["local_step"]}, - 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)) - info[CUT_KEY] = saved[CUT_KEY] - synced = saved["synced"] - sync = dict(getattr(ctl, "last_store_sync", None) or {}) - cut_bytes = directory_bytes(Path(self.root) / cut_id) - total_s = time.monotonic() - started - info["store_synced"] = bool(synced) - 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["write_bytes_per_s"] = (cut_bytes / save_s) if save_s > 0 else None - info["pruned"] = sync.get("pruned") - info["latest_seq"] = sync.get("seq") - info["trainer_onloaded_for_cut"] = offloaded - if not synced: - info["store_error"] = sync.get("error") - return info - - def _write_and_sync(self, info: dict[str, Any], cut_record: dict[str, Any]) -> dict[str, Any]: - """The round cut's ``recovery.save`` writer: the pointer (with the cut record) - into the state dir, then the state dir into the store (which commits).""" - self._write_pointer({**info, CUT_KEY: cut_record}) - synced = self.controller.sync_checkpoint_store("round_cut", skip_if_recovery=True) - # 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} - - # ---------------------------------------------------------------- resume - def resume(self, driver: Any) -> dict[str, Any] | None: - """Load the pointed cut into the freshly built trainer and set the driver's local - step. Returns the pointer (``next_rollout_id`` = where the run continues), or None - when this island never cut (fresh run). Raises on any inconsistency.""" - import time - - from yeto.rl.engine.cut import RestoreExpectation - from yeto.rl.engine.resume import check_fingerprint - - info = self.pointer() - if info is None: - if callable(self.buffer_resume): - self.buffer_resume(0) - return None - started = time.monotonic() - rid = int(info["next_rollout_id"]) - checks: dict[str, Any] = {} - diff = {} - if self.fingerprint is not None and info.get("run_fingerprint") is None: - checks["run_fingerprint"] = "absent (pointer written before rl-resume-from-checkpoint)" - elif self.fingerprint is not None: - diff = check_fingerprint(info.get("run_fingerprint"), self.fingerprint, - allow_change=self.allow_config_change) - checks["run_fingerprint"] = "same" if not diff else "changed(allowed)" - ledger = getattr(self.source, "ledger", None) - if ledger is not None and rid > 0: - last = ledger.batch(rid - 1) - if last is None or last.get("state") != "outer_recorded": - raise RoundCutError(f"round cut {info['cut_id']} continues at rollout {rid} but the " - f"ledger has no outer-recorded rollout {rid - 1}: {last}") - checks["ledger_last_round"] = "outer_recorded" - if info.get("consumed_prompt_ids_digest") is not None: - n, sha = self._consumed() - if sha != info["consumed_prompt_ids_digest"] or n != int(info.get("consumed_groups", n)): - raise RoundCutError(f"round cut {info['cut_id']}: the restored ledger's consumed groups " - f"({n}, {sha}) differ from the pointer's " - f"({info.get('consumed_groups')}, {info['consumed_prompt_ids_digest']})") - checks["consumed_prompt_ids_digest"] = "same" - backend_fingerprint = self.source.backend_fingerprint - if self.allow_config_change: - # The backend fingerprint hashes the whole Miles argv, so a longer run - # (--num-rollout, --lr-decay-iters) changes it; allowed changes are recorded. - from yeto.rl.engine.cut import load_manifest - - saved = load_manifest(self.root, info["cut_id"]).runtime.get("backend_fingerprint") - if saved != backend_fingerprint: - checks["backend_fingerprint"] = {"result": "changed(allowed)", "cut": saved, - "run": backend_fingerprint} - backend_fingerprint = saved - expect = RestoreExpectation( - algorithm=self.source.identity, layout=dict(self.trainer.actual_layout()), - backend_fingerprint=backend_fingerprint, - local_step=int(info["local_step"]), policy_version=int(info["policy_version"]), - epoch=self._epoch(), - ) - manifest = self.trainer.restore_cut(info["cut_id"], epoch=self._epoch(), root=self.root, - expect=expect, shared_filesystem=True) - checks["cut_shards_sha256"] = "ok" - if manifest.progress.policy_hash != info["policy_hash"] or manifest.progress.next_rollout_id != rid: - raise RoundCutError(f"round cut {info['cut_id']} manifest disagrees with its pointer") - restored = driver.export_local().policy_tensor_hash() - if restored != info["policy_hash"]: - raise RoundCutError(f"restored trainer holds {restored}, round cut holds {info['policy_hash']}") - checks["policy_hash"] = "same" - checks["policy_version"] = int(info["policy_version"]) - in_flight = self._restore_in_flight(driver, manifest, rid) - if in_flight is not None: - checks["in_flight"] = in_flight - self.resume_data_cursor = _cut_data_cursor(manifest) - if callable(self.buffer_resume): - watermark = (info.get(CUT_KEY) or {}).get("buffer_watermark") - if watermark is None: - raise RoundCutError(f"round cut {info['cut_id']} has no sample buffer watermark " - "but this trainer reads the sample buffer") - checks["buffer_watermark"] = int(self.buffer_resume(int(watermark))) - driver.local_step = int(info["local_step"]) - self.resumed = dict(info) - if info.get("lr_at_next_round") is not None: - self.pending_lr_check = {"rollout_id": rid, "expected": list(info["lr_at_next_round"])} - self.controller._record("round_cut", tx_id=None, action="restore", cut_id=info["cut_id"], - next_rollout_id=rid, local_step=int(info["local_step"]), - policy_hash=info["policy_hash"], cut_incarnation=info["incarnation"], - incarnation=self.controller.incarnation["id"]) - driver.emit("rl_round_cut_restored", rollout_id=rid, cut_id=info["cut_id"], - local_step=int(info["local_step"]), policy_hash=info["policy_hash"], - cut_incarnation=info["incarnation"]) - restored_store = getattr(self.controller, "restored", None) or {} - cut_bytes = sum(int(f.bytes) for f in manifest.files) - driver.emit( - "rl_resume", rollout_id=rid, cut_id=info["cut_id"], - incarnation=self.controller.incarnation.get("index"), - incarnation_id=self.controller.incarnation["id"], - cut_incarnation=info["incarnation"], cut_incarnation_index=info.get("incarnation_index"), - policy_version=int(info["policy_version"]), policy_hash=info["policy_hash"], - local_step=int(info["local_step"]), lr_at_next_round=info.get("lr_at_next_round"), - checks=checks, config_diff=diff or None, - discarded_rounds=list(getattr(self, "discarded_rounds", []) or []), - restore_s=round(time.monotonic() - started, 3), cut_bytes=cut_bytes, - state_restore=restored_store or None, - ) - return info - - def _restore_in_flight(self, driver: Any, manifest: Any, rollout_id: int) -> dict[str, Any] | None: - """agentic-rollout-utilization 3.3: the cut's unfinished trajectories, by the - CURRENT run's limit: within it and not already trained -> continued - (buffered groups back into the data buffer); else discarded and reported. - Suspended agentic trajectories never survive a restart (discarded). - None when the cut carries none.""" - entries = tuple(getattr(manifest, "in_flight", ()) or ()) - if not entries: - return None - limit = int(getattr(getattr(driver, "profile", None), "max_policy_age", 0) or 0) - ledger = getattr(self.source, "ledger", None) - completed = sorted(getattr(ledger, "_consumed_groups", None) or ()) - rollout = getattr(self.source, "rollout", None) - importer = getattr(rollout, "import_in_flight", None) - if not callable(importer): - raise RoundCutError("the cut carries in-flight trajectories but the rollout pool " - "cannot import them") - report = importer(entries, max_policy_age=limit, current_version=int(rollout_id), - completed_group_ids=completed) - driver.emit("rl_in_flight_restored", rollout_id=int(rollout_id), max_policy_age=limit, - entries=len(entries), **report) - return report - - def check_first_round(self, driver: Any, *, rollout_id: int, applied_lrs: Any) -> None: - """Design §4.5, on the boundary of the first resumed round (before it is recorded - or published): the trained learning rate continues bitwise. Raises, unless - ``allow_config_change`` (the event then records the difference).""" - from yeto.rl.engine.resume import check_lr_continues - - pending = self.pending_lr_check - if pending is None or int(rollout_id) != int(pending["rollout_id"]): - return - self.pending_lr_check = None - problem = check_lr_continues(pending["expected"], None if applied_lrs is None else list(applied_lrs)) - driver.emit("rl_resume_check", rollout_id=int(rollout_id), check="lr_at_next_round", - expected=pending["expected"], - applied=None if applied_lrs is None else [float(x) for x in applied_lrs], - ok=problem is None, allowed=problem is not None and self.allow_config_change) - if problem is not None and not self.allow_config_change: - raise RoundCutError(problem) - - -def _cut_data_cursor(manifest: Any) -> dict[str, int] | None: - """The data source position recorded at the cut (None when incomplete). - - Under fully_async the background worker keeps drawing prompts after the - round's trained groups were reported, and the cut puts the untrained ones - back into the data buffer. The position to resume from is therefore the - one read at the cut, not the one the batch ledger recorded for the round.""" - from .rollout_meta_hook import _CURSOR_FIELDS - - data = dict(getattr(manifest, "data", None) or {}) - if not all(isinstance(data.get(f), int) for f in _CURSOR_FIELDS): - return None - return {f: int(data[f]) for f in _CURSOR_FIELDS} - - -def wire_round_cuts(driver: Any, *, controller: Any, source: Any, **options: Any) -> RoundCutCheckpoint | None: - """Give a :class:`~..bridges.LocalOnlySync` driver round cuts when a checkpoint store is - configured (other syncs own the policy: their syncer is authoritative).""" - from yeto.rl.engine.bridges import LocalOnlySync - - if getattr(controller, "checkpoint_store", None) is None or not isinstance(driver.sync, LocalOnlySync): - return None - checkpoint = RoundCutCheckpoint(controller=controller, source=source, trainer=driver.trainer, **options) - driver.sync.round_cuts = checkpoint - return checkpoint +from yeto.rl.engine.round_cut import ( # noqa: F401 - re-exported + ISLAND, + POINTER, + ROUND_CUTS, + RoundCutCheckpoint, + RoundCutError, + _cut_data_cursor, + wire_round_cuts, +) diff --git a/yeto/rl/adapters/verl/cut.py b/yeto/rl/adapters/verl/cut.py index 5a3e2d2de..360268628 100644 --- a/yeto/rl/adapters/verl/cut.py +++ b/yeto/rl/adapters/verl/cut.py @@ -19,7 +19,8 @@ Restore needs the same world size (FSDP2 shards are per rank; a DP change is refused before any rank writes). Each rank checks its file checksum, loads it into the freshly built module/optimizer (DTensors rebuilt with the fresh -module's mesh and placements), and reports a digest of the loaded local state +module's mesh and placements), and reports a digest of the state its optimizer +now steps (N7.6c: read from ``optimizer.param_groups``, not the module copy) (must equal the one recorded at save) plus its resume state. Units (verl differs from Miles here): the run step is one trainer round (one @@ -99,11 +100,14 @@ def _is_dtensor(value: Any) -> bool: return hasattr(value, "to_local") and hasattr(value, "placements") +def _local_tensor(value: Any) -> Any: + return value.to_local() if _is_dtensor(value) else value + + def _to_local(value: Any) -> Any: import torch - if _is_dtensor(value): - value = value.to_local() + value = _local_tensor(value) return value.detach().to("cpu").clone() if isinstance(value, torch.Tensor) else value @@ -122,26 +126,45 @@ def _local_state(handles: TrainHandles) -> tuple[dict[str, Any], dict[str, Any], return model, optim, dtensors -def _digest(model: Mapping[str, Any], optim: Mapping[str, Any]) -> str: - """sha256 over the local trainable state (names, dtypes, shapes, bytes; group hyper-parameters).""" +def _stepped_state(handles: TrainHandles) -> tuple[dict[str, dict[str, Any]], list[dict[str, Any]]]: + """The tensors the optimizer actually steps on this rank, keyed by parameter name. + + Read from ``optimizer.param_groups`` and ``optimizer.state`` (local shards), not + from the module's state dict: a digest over the module copy cannot see an + optimizer that steps another tensor (S25 N7.6c; on Miles the CPU-offload host + master went stale after a restore while the cut digest matched, GPU T-G3). A + stepped tensor that is no module parameter is named by its group position.""" + names = {id(p): n for n, p in handles.module.named_parameters()} + state: dict[str, dict[str, Any]] = {} + groups = [] + for gi, group in enumerate(handles.optimizer.param_groups): + groups.append({k: v for k, v in group.items() if k != "params"}) + for pi, param in enumerate(group["params"]): + name = names.get(id(param), f"") + entry = {"param": param} + entry.update({k: v for k, v in handles.optimizer.state.get(param, {}).items()}) + state[name] = entry + return state, groups + + +def _digest(handles: TrainHandles) -> str: + """sha256 over the stepped state (names, dtypes, shapes, local bytes) and the group + hyper-parameters; see :func:`_stepped_state`.""" import torch h = hashlib.sha256() - - def tensor(name: str, t: Any) -> None: - if not isinstance(t, torch.Tensor): - h.update(f"{name}={t!r};".encode()) - return - flat = t.detach().to("cpu").contiguous().reshape(-1) - h.update(f"{name}|{flat.dtype}|{tuple(t.shape)}|".encode()) - h.update(flat.view(torch.uint8).numpy().tobytes()) - - for name in sorted(model): - tensor(f"model/{name}", model[name]) - for name in sorted(optim["state"]): - for key in sorted(optim["state"][name]): - tensor(f"optim/{name}/{key}", optim["state"][name][key]) - h.update(json.dumps(optim["param_groups"], sort_keys=True, default=repr).encode()) + state, groups = _stepped_state(handles) + for name in sorted(state): + for key in sorted(state[name]): + value = state[name][key] + label = f"{name}/{key}" + if not isinstance(value, torch.Tensor): + h.update(f"{label}={value!r};".encode()) + continue + flat = _local_tensor(value).detach().to("cpu").contiguous().reshape(-1) + h.update(f"{label}|{flat.dtype}|{tuple(value.shape)}|".encode()) + h.update(flat.view(torch.uint8).numpy().tobytes()) + h.update(json.dumps(groups, sort_keys=True, default=repr).encode()) return h.hexdigest() @@ -200,7 +223,7 @@ def save_rank(handles: TrainHandles, *, directory: str, cut_id: str) -> dict[str fsync_file(tmp) os.replace(tmp, path) fsync_dir(path.parent) - summary = _rank_summary(handles, rank=rank, world=world, digest=_digest(model, optim)) + summary = _rank_summary(handles, rank=rank, world=world, digest=_digest(handles)) return {"path": name, "sha256": sha256_file(path), "bytes": path.stat().st_size, **summary} @@ -258,8 +281,7 @@ def value(key: str, local: Any, like: str) -> Any: options=_state_options()) handles.scheduler.load_state_dict(saved["scheduler"]) _set_rng_state(saved["rng"]) - model, optim, _ = _local_state(handles) - return {"path": entry.path, **_rank_summary(handles, rank=rank, world=world, digest=_digest(model, optim))} + return {"path": entry.path, **_rank_summary(handles, rank=rank, world=world, digest=_digest(handles))} def resume_rank(handles: TrainHandles) -> dict[str, Any]: diff --git a/yeto/rl/adapters/verl/island_entry.py b/yeto/rl/adapters/verl/island_entry.py index 0508e441b..74fc8edd3 100644 --- a/yeto/rl/adapters/verl/island_entry.py +++ b/yeto/rl/adapters/verl/island_entry.py @@ -67,6 +67,13 @@ def parse_args(argv): p.add_argument("--rl-island-lease-s", type=float, default=None) p.add_argument("--rl-echo-events", action="store_true") p.add_argument("--rl-max-policy-age", type=int, default=0) + # N7.6b round cuts (same flags as the Miles island) + p.add_argument("--rl-resume-store", default=None) + p.add_argument("--rl-resume-state-dir", default=None) + p.add_argument("--rl-cut-every", type=int, default=None) + p.add_argument("--rl-cut-keep", type=int, default=None) + p.add_argument("--rl-stop-after-rounds", type=int, default=None) + p.add_argument("--rl-resume-allow-config-change", action="store_true") args, ignored = p.parse_known_args(argv) return args, ignored @@ -316,6 +323,12 @@ def main(argv=None) -> int: sync = "strict" if sync != "none" and not args.syncer: raise SystemExit("--syncer is required unless --rl-single-island-no-sync") + from .round_cuts import resume_config + + try: + resume = resume_config(args) + except ValueError as exc: + raise SystemExit(str(exc)) from exc asserted_keys = vconf.ASSERTED_KEYS if limit: # agentic-rollout-utilization 6.4b: fully_async path overrides, asserted_keys = fully_async_plan(limit, overrides, spec_dict, args, sync) @@ -341,6 +354,7 @@ def main(argv=None) -> int: "publish_selftest": os.environ.get("YETO_VERL_PUBLISH_SELFTEST") == "1", "run_config": run.to_dict(), "overrides": overrides, "asserted": {k: _override_value(overrides, k) for k in asserted_keys}, "max_policy_age": limit, "fully_async": bool(limit), + "resume": resume, "versions": manifest["versions"], } plan_path = work / f"plan-{lid}.json" diff --git a/yeto/rl/adapters/verl/round_cuts.py b/yeto/rl/adapters/verl/round_cuts.py new file mode 100644 index 000000000..ecaffbd9e --- /dev/null +++ b/yeto/rl/adapters/verl/round_cuts.py @@ -0,0 +1,155 @@ +"""verl island round cuts: when to cut, the pointer, stop triggers, resume (S25 N7.6b). + +The verl sync island runs the same cut controller as the Miles island +(:mod:`yeto.rl.engine.round_cut`), wired the same way: with ``--rl-resume-store`` +and no outer syncer (``--rl-single-island-no-sync``), + +* a :class:`~yeto.rl.engine.resume.ResumeController` restores the store's + ``LATEST`` snapshot into the state dir at start (pointer ``round-cut.json``, + batch ledger, journal) and syncs the state dir back after every cut; +* a :class:`~yeto.rl.engine.round_cut.RoundCutCheckpoint` on the + :class:`~yeto.rl.engine.bridges.LocalOnlySync` cuts the trainer + (:class:`~.cut.VerlCutTrainer`: one yeto file per FSDP rank) every + ``--rl-cut-every`` rounds, at the last round and at our own stop + (``--rl-stop-after-rounds``), and resumes from the pointed cut at start; +* triggers of the final cut: ``MANUAL`` (last round, our own stop), + ``SCHEDULED`` (the head's trainer scheduler stopped us: + ``trainer_remote`` commands, :func:`run_driver`); a preemption signal only + writes a marker (``install_preempt_handler``): the next launch resumes from + the newest complete cut. + +The resume assertions of :mod:`.cut` run on this path: the restored rank digests +and Adam steps at restore, the landed data cursor against the cut's data +watermark (driver ``_restore_data_cursor`` -> ``confirm_resumed_data_cursor``), +and the Adam/scheduler advance after the first resumed round. +""" + +from __future__ import annotations + +import os +from collections.abc import Mapping +from dataclasses import dataclass +from typing import Any + +RESUME_FLAGS = ("rl_cut_every", "rl_cut_keep", "rl_stop_after_rounds") +DEFAULT_STATE_DIR = "~/yeto-rl/resume-state" +# VerlRunConfig fields that may change between a cut and the launch resuming it: +# paths of this container, run length (extending a run is the normal resume), +# and settings that do not change the trained numbers. +FINGERPRINT_EXCLUDED = ("out_dir", "train_file", "val_file", "total_rounds", "experiment_name", + "gpu_memory_utilization", "reward_path") + + +def _flag(name: str) -> str: + return "--" + name.replace("_", "-") + + +def resume_config(args: Any) -> dict[str, Any] | None: + """``plan["resume"]`` from the island flags (None: no resume store). Raises on + a flag combination the cut controller cannot run.""" + store = getattr(args, "rl_resume_store", None) + given = [_flag(n) for n in RESUME_FLAGS if getattr(args, n, None) is not None] + if getattr(args, "rl_resume_allow_config_change", False): + given.append("--rl-resume-allow-config-change") + for name in RESUME_FLAGS: + value = getattr(args, name, None) + if value is not None and int(value) < 1: + raise ValueError(f"{_flag(name)} must be >= 1") + if not store: + if given: + raise ValueError(", ".join(given) + " need --rl-resume-store") + return None + if not getattr(args, "rl_single_island_no_sync", False): + raise ValueError("--rl-resume-store needs --rl-single-island-no-sync (a synced island " + "resumes from the syncer)") + if int(getattr(args, "rl_max_policy_age", 0) or 0): + raise ValueError("--rl-resume-store is not wired for the verl fully_async path " + "(--rl-max-policy-age > 0)") + return { + "store": str(store), + "state_dir": str(getattr(args, "rl_resume_state_dir", None) or DEFAULT_STATE_DIR), + "every": int(getattr(args, "rl_cut_every", None) or 1), + "keep": int(getattr(args, "rl_cut_keep", None) or 2), + "allow_config_change": bool(getattr(args, "rl_resume_allow_config_change", False)), + "stop_after": (None if getattr(args, "rl_stop_after_rounds", None) is None + else int(args.rl_stop_after_rounds)), + } + + +def run_fingerprint_of(plan: Mapping[str, Any], algorithm: Any) -> dict[str, Any]: + """What must be identical between a cut and the launch resuming it.""" + run = {k: v for k, v in dict(plan.get("run_config") or {}).items() if k not in FINGERPRINT_EXCLUDED} + sha = getattr(algorithm, "sha256", None) + return {**run, "model_revision": plan.get("model_revision"), "lora_targets": plan.get("lora_targets"), + "algorithm_spec_sha256": sha() if callable(sha) else None} + + +@dataclass +class ResumeWiring: + """The round-cut controller and batch ledger of one verl island run.""" + + controller: Any # ResumeController + ledger: Any # BatchLedger (in the state dir: restored and synced with the pointer) + config: dict[str, Any] + + +def resume_wiring_for(plan: Mapping[str, Any], *, runtime_fingerprint: str, + environ: Mapping[str, str] | None = None) -> ResumeWiring | None: + """Restore the store into the state dir (before the driver is built: the + ledger the driver gets is the restored one). None without ``plan["resume"]``.""" + config = plan.get("resume") + if not config: + return None + if plan.get("sync") != "none": + raise ValueError("verl round cuts need sync 'none' (--rl-single-island-no-sync)") + if not runtime_fingerprint: + raise ValueError("verl round cuts need a runtime fingerprint (a cut without it is refused)") + from yeto.rl.engine.ledger import BatchLedger + from yeto.rl.engine.resume import ResumeController, store_for + + store = store_for(config["store"], environ=environ) + controller = ResumeController(state_dir=os.path.expanduser(config["state_dir"]), store=store, + runtime_fingerprint=runtime_fingerprint, keep=int(config.get("keep", 2))) + return ResumeWiring(controller=controller, ledger=BatchLedger(controller.state_dir), config=dict(config)) + + +def wire_round_cuts(driver: Any, resume: ResumeWiring | None, *, plan: Mapping[str, Any], + algorithm: Any, install_signals: bool = True) -> Any: + """Give the driver's LocalOnlySync the cut controller (the Miles wiring, verl + units: one scheduler step per round, so the cut's ``global_batch_size`` is 1). + Returns the :class:`~yeto.rl.engine.round_cut.RoundCutCheckpoint` or None.""" + if resume is None: + return None + from yeto.rl.engine import round_cut + from yeto.rl.engine.resume import install_preempt_handler + + controller = resume.controller + source = round_cut.CutSource( + driver=lambda: driver, trainer=driver.trainer, rollout=driver.rollout, ledger=resume.ledger, + algorithm=algorithm, backend_fingerprint=controller.runtime_fingerprint, cut_root="", + global_batch_size=1) + checkpoint = round_cut.wire_round_cuts( + driver, controller=controller, source=source, every=int(resume.config.get("every", 1)), + fingerprint=run_fingerprint_of(plan, algorithm), + allow_config_change=bool(resume.config.get("allow_config_change", False))) + if checkpoint is None: + raise ValueError("verl round cuts need the driver's sync to be LocalOnlySync") + driver.sync.stop_after = resume.config.get("stop_after") + if install_signals: + install_preempt_handler(controller, emit=driver.emit) + return checkpoint + + +def run_driver(driver: Any, resume: ResumeWiring | None, *, environ: Mapping[str, str] | None = None) -> Any: + """Run the island: as one commanded incarnation when the head's trainer + scheduler starts and stops it (``trainer_remote``: its stop cuts with the + SCHEDULED trigger), else to the end of the run. Returns the final policy state.""" + from yeto.rl.engine import trainer_remote + + if not trainer_remote.commanded(environ): + return driver.run() + if resume is None: + raise RuntimeError("a trainer started by the head's scheduler needs round cuts " + "(--rl-resume-store with --rl-single-island-no-sync)") + trainer_remote.run_commanded_island(driver, store=resume.controller.store, environ=environ) + return driver.published_state diff --git a/yeto/rl/adapters/verl/trainer.py b/yeto/rl/adapters/verl/trainer.py index 72e7859e0..55ed0cfb3 100644 --- a/yeto/rl/adapters/verl/trainer.py +++ b/yeto/rl/adapters/verl/trainer.py @@ -206,6 +206,7 @@ def run_island(trainer, agent_loop_manager, plan: dict, cfg: dict) -> dict: from yeto.rl.engine.driver import EventTape, IslandDriver from yeto.rl.export import adapter_targets, derive_peft_lora_specs + from . import round_cuts from .identity import backend_identity from .ports_impl import (VerlIsland, VerlPlacement, VerlPolicyState, VerlPublisher, VerlRolloutPool, VerlTrainerGroup) @@ -256,16 +257,20 @@ def run_island(trainer, agent_loop_manager, plan: dict, cfg: dict) -> dict: spec = (AlgorithmSpec.from_dict(plan["algorithm_spec"]) if plan.get("algorithm_spec") else AlgorithmSpec()) publisher = VerlPublisher(island, strict=bool(plan.get("publish_strict", True))) + # N7.6b: round cuts (--rl-resume-store): the store is restored before the driver + # is built, so the driver's ledger is the restored one + resume = round_cuts.resume_wiring_for(plan, runtime_fingerprint=_fingerprint(cfg)) driver = IslandDriver( learner_id=learner_id, rollout=VerlRolloutPool(island), trainer=VerlTrainerGroup(island), policy_state=VerlPolicyState(island), publisher=_TestExitPublisher(publisher, plan.get("test_exit_after_version"), tape), placement=VerlPlacement(), capabilities=_capabilities(_fingerprint(cfg)), algorithm=spec, - sync=sync, events=tape) + sync=sync, events=tape, ledger=None if resume is None else resume.ledger) + round_cuts.wire_round_cuts(driver, resume, plan=plan, algorithm=spec) if plan.get("publish_selftest"): tape.append({"event": "rl_verl_publish_selftest", **publish_selftest(island, publisher, VerlPolicyState(island), specs)}) - state = driver.run() + state = round_cuts.run_driver(driver, resume) result = {"learner_id": learner_id, "rounds_completed": driver.rounds_completed, "final_policy_version": state.policy_version, "final_policy_tensor_hash": state.policy_tensor_hash(), diff --git a/yeto/rl/engine/round_cut.py b/yeto/rl/engine/round_cut.py new file mode 100644 index 000000000..32ef0fb18 --- /dev/null +++ b/yeto/rl/engine/round_cut.py @@ -0,0 +1,493 @@ +"""Round cuts of a run without an outer syncer: save at safe points, resume at start. + +Neutral core (moved from the Miles adapter for the verl island, S25 N7.6b): the +backend gives a ``CuttableTrainer`` (``save_cut`` / ``restore_cut`` / +``actual_layout``) and a rollout pool with ``data_cursor`` / ``seek_data_cursor``; +:class:`CutSource` builds the cut context from the driver, and +:class:`RoundCutCheckpoint` is the cut controller wired into +:class:`~.bridges.LocalOnlySync` (:func:`wire_round_cuts`). + +rl-multinode-island M4: policy-weight continuation across a machine replacement. + + +Without an outer syncer (``--rl-single-island-no-sync``, :class:`~..bridges.LocalOnlySync`) +a restarted learner used to start again at rollout 0 from the startup weights: the +checkpoint store (Q4, C5) carried the journal/ledger, but no LoRA/optimizer state. + +With ``--rl-elastic-checkpoint-store`` set this module keeps a 4.2 cut of the trainer +(every rank's LoRA + optimizer + scheduler + RNG shard, manifest committed last) at +every round-boundary safe point and resumes from the newest one at start: + +* the cut root is ``/round-cuts`` -- the store is visible on every node (sky + Storage MOUNT of the bucket), so the shard of a rank on another node (PP across + nodes) lands where the driver can check it (``shared_filesystem=True``); +* after the manifest is committed, ``round-cut.json`` (the pointer: cut id, next + rollout id, local step, policy hash, writing incarnation) is written into the state + dir and the state dir is synced to the store -- pointer, ledger and journal travel + in ONE store snapshot (``STORE-MANIFEST.json`` last), so a restored pointer never + names a cut newer than the restored ledger; +* at start (in-place restart or a fresh machine after the store restore) the cut the + pointer names is verified (``verify_cut``: algorithm, layout, backend, step, + policy version, every shard digest) and loaded on the fresh trainer + (``MilesTrainerGroup.restore_cut``); the driver then rebases the ledger and seeks + the data source to the pointer's rollout id. Any failure is fatal (no silent + restart at 0 once a pointer exists). +""" + +from __future__ import annotations + +import json +import os +from collections.abc import Callable, Mapping +from pathlib import Path +from typing import Any + +from . import recovery +from .recovery import CUT_KEY, FINAL_TRIGGERS, MANUAL, PERIODIC, ConsistentCut +from .resume import ROUND_CUTS # under the checkpoint store (not under the state dir) + +ISLAND = "local" # round cuts run without an outer syncer: one island, the local one +POINTER = "round-cut.json" # under the state dir (synced with the ledger/journal) + + +class RoundCutError(RuntimeError): + pass + + +def _algorithm_identity(algorithm: Any, *, ref_model: Mapping[str, Any] | None) -> Any: + from yeto.rl.engine.cut import AlgorithmIdentity + + runtime_attrs = None + to_attrs = getattr(algorithm, "to_legacy_runtime_attrs", None) + if callable(to_attrs): + runtime_attrs = dict(to_attrs()) or None + return AlgorithmIdentity.from_spec(algorithm, runtime_attrs=runtime_attrs, ref_model=ref_model) + + +class CutSource: + """What the driver contributes to a cut (E2 ``CutContext``) and what a + restore must match (``RestoreExpectation``); shared by the 4.4 same-shape + rebuilder and E3's ``MilesTrainerOps`` (4.7).""" + + def __init__(self, *, driver: Any, trainer: Any, rollout: Any, ledger: Any, algorithm: Any, + backend_fingerprint: str, cut_root: str, global_batch_size: int, + ref_model: Mapping[str, Any] | None = None, + shared_filesystem: bool = True, + check_in_flight: Callable[[Mapping[str, Any]], list[str]] | None = None) -> None: + self.driver_ref = driver # the driver, or a zero-arg callable returning it + self.trainer = trainer + self.rollout = rollout + self.ledger = ledger + self.identity = _algorithm_identity(algorithm, ref_model=ref_model) + self.backend_fingerprint = backend_fingerprint + self.cut_root = cut_root + self.global_batch_size = int(global_batch_size) + self.shared_filesystem = shared_filesystem + # backend check of an in-flight export (problems refuse the cut); None: no check + self.check_in_flight = check_in_flight + + @property + 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: + from yeto.rl.engine.controller import RebuildRefused + from yeto.rl.engine.cut import CutProgress + from yeto.rl.engine.cut import CutContext + + driver = self.driver + state = driver.published_state + if state is None or driver.published_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) + if cursor is None: + raise RebuildRefused("rollout data cursor unknown (needs the elastic metadata)") + data = dict(cursor) + summary = dict(self.ledger.cut_summary()) if self.ledger is not None else {} + if summary.get("engine_buffer_length") is not None: + data.setdefault("buffer_length", summary["engine_buffer_length"]) + in_flight: tuple = () + limit = int(getattr(getattr(driver, "profile", None), "max_policy_age", 0) or 0) + if limit > 0: + # agentic-rollout-utilization 3.3: the unfinished work the rollout + # process holds NOW (buffered groups, suspended agentic trajectories). + exporter = getattr(self.rollout, "export_in_flight", None) + if not callable(exporter): + raise RebuildRefused("the rollout pool cannot export its in-flight trajectories") + exported = exporter() + problems = list(self.check_in_flight(exported)) if self.check_in_flight else [] + if problems: + raise RebuildRefused("; ".join(problems)) + in_flight = tuple(exported["entries"]) + data["buffer_length"] = int(exported["buffer_groups"]) + summary["carried_over"] = len(in_flight) + summary["max_policy_age"] = limit + progress = CutProgress( + 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), + policy_version=int(state.policy_version), + policy_hash=state.policy_tensor_hash(), + ) + return CutContext( + root=self.cut_root, cut_id=cut_id, backend_fingerprint=self.backend_fingerprint, + progress=progress, algorithm=self.identity, data=data, ledger=summary, + # The safe point is after boundary() returned non-stop and every + # member acknowledged the policy: the outer commit is settled. + outer={"settled": bool(driver.at_safe_point), "policy_version": state.policy_version, + "policy_token": driver.expected_token}, + shared_filesystem=self.shared_filesystem, + in_flight=in_flight, + ) + + def expectation(self, layout: Mapping[str, int], *, epoch: int) -> Any: + from yeto.rl.engine.cut import RestoreExpectation + + driver = self.driver + return RestoreExpectation( + algorithm=self.identity, layout=dict(layout), + backend_fingerprint=self.backend_fingerprint, local_step=int(driver.local_step), + policy_version=int(driver.published_state.policy_version), epoch=epoch, + ) + + +class RoundCutCheckpoint: + """``save`` at a safe point, ``resume`` at sync start (see the module doc). + + rl-resume-from-checkpoint (S17 C4): ``every`` (``--rl-cut-every``) cuts only every N + rounds (plus the last round and our own stop, ``final=True``); the pointer carries + ``run_fingerprint`` / ``lr_at_next_round`` / ``consumed_prompt_ids_digest``; a resume + checks the fingerprint (refused on a difference unless ``allow_config_change``), the + consumed-group digest and, on the first resumed round, the learning rate.""" + + 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: + if getattr(controller, "checkpoint_store", None) is None: + raise ValueError("round cuts need the checkpoint store") + if int(every) < 1: + raise ValueError("--rl-cut-every must be >= 1") + self.controller = controller + self.source = source # rebuild_wiring.CutSource (cut_root is set here) + self.trainer = trainer # MilesTrainerGroup: save_cut / restore_cut / actual_layout + self.root = str(Path(controller.checkpoint_store) / ROUND_CUTS) + source.cut_root = self.root + self.resumed: dict[str, Any] | None = None + self.every = int(every) + self.fingerprint = dict(fingerprint) if fingerprint is not None else None + self.allow_config_change = bool(allow_config_change) + self.next_lr = next_lr # callable(driver) -> list[float] | None + self.pending_lr_check: dict[str, Any] | None = None + # set by resume(): where the rollout data source continues (the cut's cursor) + self.resume_data_cursor: dict[str, int] | None = None + # rl-trainer-start-stop D4: the durable sample buffer's consumed watermark, + # written into the cut record (None: the run reads no sample buffer). + # buffer_watermark(next_rollout_id) commits and returns the position after the + # last trained round; buffer_resume(watermark) continues there after a restore + # (buffer_rollout.wire_sample_buffer sets both). + self.buffer_watermark: Any = None # callable(int) -> int | None + self.buffer_resume: Any = None # callable(int) -> int + + # ---------------------------------------------------------------- pointer + @property + def pointer_path(self) -> Path: + return Path(self.controller.state_dir) / POINTER + + def pointer(self) -> dict[str, Any] | None: + path = self.pointer_path + if not path.is_file(): + 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") + 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) + + def _epoch(self) -> int: + return int(self.controller.journal.epochs.config_epoch) + + def _consumed(self) -> tuple[int, str | None]: + ledger = getattr(self.source, "ledger", None) + groups = getattr(ledger, "_consumed_groups", None) + if groups is None: + return 0, None + from yeto.rl.engine.resume import consumed_digest + + return len(groups), consumed_digest(groups) + + # ---------------------------------------------------------------- save + def due(self, rollout_id: int, *, final: bool = False) -> bool: + from yeto.rl.engine.resume import should_cut + + return should_cut(int(rollout_id), every=self.every, final=final) + + 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 + version), point the state dir at it, sync the store. Returns the pointer, or None + 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). + + ``trigger`` names a final cut's trigger (default ``MANUAL``; ``SCHEDULED`` for + a stop the trainer scheduler ordered); a periodic cut takes none.""" + import time + + from yeto.rl.engine.resume import directory_bytes + + 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): + 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']}" + 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 + progress = context.progress + consumed_n, consumed_sha = self._consumed() + next_lr = self.next_lr(driver) if callable(self.next_lr) else None + info = { + "cut_id": cut_id, "root": ROUND_CUTS, "epoch": self._epoch(), + "next_rollout_id": int(progress.next_rollout_id), + "local_step": int(progress.local_step), + "policy_version": int(progress.policy_version), + "policy_hash": progress.policy_hash, + "incarnation": ctl.incarnation["id"], + "incarnation_index": ctl.incarnation.get("index"), + "run_fingerprint": self.fingerprint, + "lr_at_next_round": None if next_lr is None else [float(x) for x in next_lr], + "consumed_groups": consumed_n, + "consumed_prompt_ids_digest": consumed_sha, + "outer_version": getattr(getattr(driver, "sync", None), "outer_version", None), + "final": bool(final), + } + cut = ConsistentCut( + checkpoint_id=cut_id, trigger=(trigger or MANUAL) if final else PERIODIC, + outer_step=info["outer_version"], island_steps={ISLAND: info["local_step"]}, + 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)) + info[CUT_KEY] = saved[CUT_KEY] + synced = saved["synced"] + sync = dict(getattr(ctl, "last_store_sync", None) or {}) + cut_bytes = directory_bytes(Path(self.root) / cut_id) + total_s = time.monotonic() - started + info["store_synced"] = bool(synced) + 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["write_bytes_per_s"] = (cut_bytes / save_s) if save_s > 0 else None + info["pruned"] = sync.get("pruned") + info["latest_seq"] = sync.get("seq") + info["trainer_onloaded_for_cut"] = offloaded + if not synced: + info["store_error"] = sync.get("error") + return info + + def _write_and_sync(self, info: dict[str, Any], cut_record: dict[str, Any]) -> dict[str, Any]: + """The round cut's ``recovery.save`` writer: the pointer (with the cut record) + into the state dir, then the state dir into the store (which commits).""" + self._write_pointer({**info, CUT_KEY: cut_record}) + synced = self.controller.sync_checkpoint_store("round_cut", skip_if_recovery=True) + # 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} + + # ---------------------------------------------------------------- resume + def resume(self, driver: Any) -> dict[str, Any] | None: + """Load the pointed cut into the freshly built trainer and set the driver's local + step. Returns the pointer (``next_rollout_id`` = where the run continues), or None + when this island never cut (fresh run). Raises on any inconsistency.""" + import time + + from yeto.rl.engine.cut import RestoreExpectation + from yeto.rl.engine.resume import check_fingerprint + + info = self.pointer() + if info is None: + if callable(self.buffer_resume): + self.buffer_resume(0) + return None + started = time.monotonic() + rid = int(info["next_rollout_id"]) + checks: dict[str, Any] = {} + diff = {} + if self.fingerprint is not None and info.get("run_fingerprint") is None: + checks["run_fingerprint"] = "absent (pointer written before rl-resume-from-checkpoint)" + elif self.fingerprint is not None: + diff = check_fingerprint(info.get("run_fingerprint"), self.fingerprint, + allow_change=self.allow_config_change) + checks["run_fingerprint"] = "same" if not diff else "changed(allowed)" + ledger = getattr(self.source, "ledger", None) + if ledger is not None and rid > 0: + last = ledger.batch(rid - 1) + if last is None or last.get("state") != "outer_recorded": + raise RoundCutError(f"round cut {info['cut_id']} continues at rollout {rid} but the " + f"ledger has no outer-recorded rollout {rid - 1}: {last}") + checks["ledger_last_round"] = "outer_recorded" + if info.get("consumed_prompt_ids_digest") is not None: + n, sha = self._consumed() + if sha != info["consumed_prompt_ids_digest"] or n != int(info.get("consumed_groups", n)): + raise RoundCutError(f"round cut {info['cut_id']}: the restored ledger's consumed groups " + f"({n}, {sha}) differ from the pointer's " + f"({info.get('consumed_groups')}, {info['consumed_prompt_ids_digest']})") + checks["consumed_prompt_ids_digest"] = "same" + backend_fingerprint = self.source.backend_fingerprint + if self.allow_config_change: + # The backend fingerprint hashes the whole Miles argv, so a longer run + # (--num-rollout, --lr-decay-iters) changes it; allowed changes are recorded. + from yeto.rl.engine.cut import load_manifest + + saved = load_manifest(self.root, info["cut_id"]).runtime.get("backend_fingerprint") + if saved != backend_fingerprint: + checks["backend_fingerprint"] = {"result": "changed(allowed)", "cut": saved, + "run": backend_fingerprint} + backend_fingerprint = saved + expect = RestoreExpectation( + algorithm=self.source.identity, layout=dict(self.trainer.actual_layout()), + backend_fingerprint=backend_fingerprint, + local_step=int(info["local_step"]), policy_version=int(info["policy_version"]), + epoch=self._epoch(), + ) + manifest = self.trainer.restore_cut(info["cut_id"], epoch=self._epoch(), root=self.root, + expect=expect, shared_filesystem=True) + checks["cut_shards_sha256"] = "ok" + if manifest.progress.policy_hash != info["policy_hash"] or manifest.progress.next_rollout_id != rid: + raise RoundCutError(f"round cut {info['cut_id']} manifest disagrees with its pointer") + restored = driver.export_local().policy_tensor_hash() + if restored != info["policy_hash"]: + raise RoundCutError(f"restored trainer holds {restored}, round cut holds {info['policy_hash']}") + checks["policy_hash"] = "same" + checks["policy_version"] = int(info["policy_version"]) + in_flight = self._restore_in_flight(driver, manifest, rid) + if in_flight is not None: + checks["in_flight"] = in_flight + self.resume_data_cursor = _cut_data_cursor(manifest) + if callable(self.buffer_resume): + watermark = (info.get(CUT_KEY) or {}).get("buffer_watermark") + if watermark is None: + raise RoundCutError(f"round cut {info['cut_id']} has no sample buffer watermark " + "but this trainer reads the sample buffer") + checks["buffer_watermark"] = int(self.buffer_resume(int(watermark))) + driver.local_step = int(info["local_step"]) + self.resumed = dict(info) + if info.get("lr_at_next_round") is not None: + self.pending_lr_check = {"rollout_id": rid, "expected": list(info["lr_at_next_round"])} + self.controller._record("round_cut", tx_id=None, action="restore", cut_id=info["cut_id"], + next_rollout_id=rid, local_step=int(info["local_step"]), + policy_hash=info["policy_hash"], cut_incarnation=info["incarnation"], + incarnation=self.controller.incarnation["id"]) + driver.emit("rl_round_cut_restored", rollout_id=rid, cut_id=info["cut_id"], + local_step=int(info["local_step"]), policy_hash=info["policy_hash"], + cut_incarnation=info["incarnation"]) + restored_store = getattr(self.controller, "restored", None) or {} + cut_bytes = sum(int(f.bytes) for f in manifest.files) + driver.emit( + "rl_resume", rollout_id=rid, cut_id=info["cut_id"], + incarnation=self.controller.incarnation.get("index"), + incarnation_id=self.controller.incarnation["id"], + cut_incarnation=info["incarnation"], cut_incarnation_index=info.get("incarnation_index"), + policy_version=int(info["policy_version"]), policy_hash=info["policy_hash"], + local_step=int(info["local_step"]), lr_at_next_round=info.get("lr_at_next_round"), + checks=checks, config_diff=diff or None, + discarded_rounds=list(getattr(self, "discarded_rounds", []) or []), + restore_s=round(time.monotonic() - started, 3), cut_bytes=cut_bytes, + state_restore=restored_store or None, + ) + return info + + def _restore_in_flight(self, driver: Any, manifest: Any, rollout_id: int) -> dict[str, Any] | None: + """agentic-rollout-utilization 3.3: the cut's unfinished trajectories, by the + CURRENT run's limit: within it and not already trained -> continued + (buffered groups back into the data buffer); else discarded and reported. + Suspended agentic trajectories never survive a restart (discarded). + None when the cut carries none.""" + entries = tuple(getattr(manifest, "in_flight", ()) or ()) + if not entries: + return None + limit = int(getattr(getattr(driver, "profile", None), "max_policy_age", 0) or 0) + ledger = getattr(self.source, "ledger", None) + completed = sorted(getattr(ledger, "_consumed_groups", None) or ()) + rollout = getattr(self.source, "rollout", None) + importer = getattr(rollout, "import_in_flight", None) + if not callable(importer): + raise RoundCutError("the cut carries in-flight trajectories but the rollout pool " + "cannot import them") + report = importer(entries, max_policy_age=limit, current_version=int(rollout_id), + completed_group_ids=completed) + driver.emit("rl_in_flight_restored", rollout_id=int(rollout_id), max_policy_age=limit, + entries=len(entries), **report) + return report + + def check_first_round(self, driver: Any, *, rollout_id: int, applied_lrs: Any) -> None: + """Design §4.5, on the boundary of the first resumed round (before it is recorded + or published): the trained learning rate continues bitwise. Raises, unless + ``allow_config_change`` (the event then records the difference).""" + from yeto.rl.engine.resume import check_lr_continues + + pending = self.pending_lr_check + if pending is None or int(rollout_id) != int(pending["rollout_id"]): + return + self.pending_lr_check = None + problem = check_lr_continues(pending["expected"], None if applied_lrs is None else list(applied_lrs)) + driver.emit("rl_resume_check", rollout_id=int(rollout_id), check="lr_at_next_round", + expected=pending["expected"], + applied=None if applied_lrs is None else [float(x) for x in applied_lrs], + ok=problem is None, allowed=problem is not None and self.allow_config_change) + if problem is not None and not self.allow_config_change: + raise RoundCutError(problem) + + +def _cut_data_cursor(manifest: Any) -> dict[str, int] | None: + """The data source position recorded at the cut (None when incomplete). + + Under fully_async the background worker keeps drawing prompts after the + round's trained groups were reported, and the cut puts the untrained ones + back into the data buffer. The position to resume from is therefore the + one read at the cut, not the one the batch ledger recorded for the round.""" + from .cut import DATA_CURSOR_FIELDS + + data = dict(getattr(manifest, "data", None) or {}) + # a backend with fewer position fields (verl: sample_offset) names them in the cut + fields = tuple(data.get("cursor_fields") or DATA_CURSOR_FIELDS) + if not all(isinstance(data.get(f), int) for f in fields): + return None + return {f: int(data[f]) for f in fields} + + +def wire_round_cuts(driver: Any, *, controller: Any, source: Any, **options: Any) -> RoundCutCheckpoint | None: + """Give a :class:`~..bridges.LocalOnlySync` driver round cuts when a checkpoint store is + configured (other syncs own the policy: their syncer is authoritative).""" + from .bridges import LocalOnlySync + + if getattr(controller, "checkpoint_store", None) is None or not isinstance(driver.sync, LocalOnlySync): + return None + checkpoint = RoundCutCheckpoint(controller=controller, source=source, trainer=driver.trainer, **options) + driver.sync.round_cuts = checkpoint + return checkpoint From a9fad17ac76c7a842b748f6a11b57cc9f87e53e4 Mon Sep 17 00:00:00 2001 From: michael <78113554+chowhsu@users.noreply.github.com> Date: Tue, 20 Oct 2026 04:16:19 +0000 Subject: [PATCH 29/89] Split the trainer's rounds evenly over several sampler islands Co-Authored-By: Claude Opus 5.5 --- tests/test_rl_spot_sampler_replace.py | 10 +++++++++- yeto/rl/engine/sample_roles.py | 18 +++++++++++++++--- 2 files changed, 24 insertions(+), 4 deletions(-) diff --git a/tests/test_rl_spot_sampler_replace.py b/tests/test_rl_spot_sampler_replace.py index 92ff9b416..0ae3c3125 100644 --- a/tests/test_rl_spot_sampler_replace.py +++ b/tests/test_rl_spot_sampler_replace.py @@ -288,7 +288,7 @@ def test_member_records_are_one_file_per_incarnation(tmp_path): "--rl-island-sample-role", "0:sampler", "--rl-island-sample-role", "1:sampler", "--rl-island-sample-role", "2:trainer", "--rl-max-policy-age", "2", "--rl-algorithm-spec", "examples/rl_algorithms/decoupled-ppo-c2.json", - "--rl-checkpoint-store", STORE, "--cluster-prefix", "c1-spot") + "--rl-checkpoint-store", STORE, "--cluster-prefix", "c1-spot", "--total-steps", "4") def _args(extra=()): @@ -317,6 +317,14 @@ def test_one_trainer_and_two_samplers_get_their_node_and_the_shards(): tokens = [_tokens(i["learner_command"]) for i in plan["island_requests"]] assert [_flag(t, "--rl-sampler-node") for t in tokens] == ["island0", "island1", None] assert {_flag(t, "--rl-buffer-shards") for t in tokens} == {"island0,island1"} + # the trainer trains --total-steps rounds; each of the two samplers generates half + assert [_flag(t, "--global-rounds") for t in tokens] == ["2", "2", "4"] + assert [_flag(t, "--total-fragment-steps") for t in tokens] == ["2", "2", "4"] + + +def test_the_samplers_split_the_rounds_evenly(): + with pytest.raises(ValueError, match="multiple of the sampler count"): + _args(("--total-steps", "3")) def test_the_island_resolves_its_sampler_id_and_the_shards(): diff --git a/yeto/rl/engine/sample_roles.py b/yeto/rl/engine/sample_roles.py index e9b76b2fd..69f72d7c8 100644 --- a/yeto/rl/engine/sample_roles.py +++ b/yeto/rl/engine/sample_roles.py @@ -151,6 +151,13 @@ def island_args(args: Any, island: int, num_islands: int) -> Any: if out.rl_sample_role in (SAMPLER, TRAINER): out.rl_sampler_node = node_name(island) if out.rl_sample_role == SAMPLER else None out.rl_buffer_shards = buffer_shards(args, num_islands) + samplers = len(sampler_islands(args, num_islands)) + if out.rl_sample_role == SAMPLER and samplers > 1: + # the trainer trains --total-steps rounds; the samplers share them, so each + # one reaches its data end (check_pair needs an even split) + out.total_steps = int(args.total_steps) // samplers + if getattr(args, "rl_total_fragment_steps", None) is not None: + out.rl_total_fragment_steps = out.total_steps * int(getattr(args, "fragments", 1) or 1) if out.rl_sample_role == CONVERTIBLE: out.rl_convertible_node = node_name(island) out.rl_sampler_min_round_seconds = None # a node keeps its backlog cap @@ -182,8 +189,9 @@ def check_pair(args: Any, num_islands: int) -> None: """Launch checks of a store-coupled pair (nothing outside a pair): every island has a role, exactly one trainer and at least one sampler, ``--rl-single-island-no-sync`` (no syncer), no ``--rl-sample-role`` and no - external learner; several samplers take no trainer scheduler (the head reads - no sampler shard of a shared directory); then every island's own role checks + external learner; several samplers split ``--total-steps`` evenly (each + generates its share of the trainer's rounds) and take no trainer scheduler + (the head reads no sampler shard of a shared directory); then every island's own role checks on its island args.""" roles = island_roles(args, num_islands) if not roles: @@ -193,7 +201,11 @@ def check_pair(args: Any, num_islands: int) -> None: raise ValueError("--rl-island-sample-role: a pair gives every --gpu island a role, " f"at least one {SAMPLER} and exactly one {TRAINER} (got {num_islands} " f"island(s), roles {dict(sorted(roles.items()))})") - if values.count(SAMPLER) > 1 and getattr(args, "rl_trainer_scheduler", None) not in (None, "off"): + samplers = values.count(SAMPLER) + if samplers > 1 and int(getattr(args, "total_steps", 0) or 0) % samplers: + raise ValueError(f"--total-steps {args.total_steps}: {samplers} sampler islands share the " + "trainer's rounds; it must be a multiple of the sampler count") + if samplers > 1 and getattr(args, "rl_trainer_scheduler", None) not in (None, "off"): raise ValueError("--rl-trainer-scheduler with several sampler islands is not supported " "(the head does not read the sampler shards); drop the scheduler") if sample_role(args) != BOTH: From b651fe07ae1a06c3cd59b347d10d5e22313db7ad Mon Sep 17 00:00:00 2001 From: michael <78113554+chowhsu@users.noreply.github.com> Date: Tue, 20 Oct 2026 04:16:54 +0000 Subject: [PATCH 30/89] Let synced trainer islands feed sampler islands through merged releases Trainer islands (--rl-sample-role trainer) may now run with a syncer (strict or elastic OuterStep merge) instead of only as no-sync islands. Learner 0 is the one publisher: it releases the merged policy with the syncer's outer version, only above the newest release (gaps are reported). Every trainer gates max_policy_age on the merged version and reads its own lane of the shared sample buffer (position % N == learner id), with a durable lane cursor so a re-joined island never trains a group twice. The launcher forwards the store as --rl-sample-store plus --rl-sample-lanes; round cuts stay refused for synced trainers. Samplers still run as no-sync islands. Co-Authored-By: Claude Opus 5.5 --- tests/test_rl_merged_release.py | 288 ++++++++++++++++++ tests/test_rl_synced_trainers_sampler.py | 203 ++++++++++++ yeto/launch/rl_ports.py | 24 +- yeto/rl/adapters/miles/buffer_rollout.py | 54 +++- .../rl/adapters/miles/buffer_rollout_miles.py | 4 +- yeto/rl/adapters/miles/island_entry.py | 23 +- yeto/rl/adapters/miles/run_flags.py | 5 + yeto/rl/adapters/miles/sampler_island.py | 90 +++++- yeto/rl/engine/merged_release.py | 175 +++++++++++ yeto/rl/engine/policy_release.py | 12 +- yeto/rl/engine/sample_buffer.py | 26 +- yeto/rl/engine/sample_roles.py | 61 +++- 12 files changed, 934 insertions(+), 31 deletions(-) create mode 100644 tests/test_rl_merged_release.py create mode 100644 tests/test_rl_synced_trainers_sampler.py create mode 100644 yeto/rl/engine/merged_release.py diff --git a/tests/test_rl_merged_release.py b/tests/test_rl_merged_release.py new file mode 100644 index 000000000..5446f2b66 --- /dev/null +++ b/tests/test_rl_merged_release.py @@ -0,0 +1,288 @@ +"""S25 C3: synced trainer islands feed sampler islands (yeto.rl.engine.merged_release). + +Unit level: buffer lanes, the lane cursor log, the release guard, the publisher +wrappers, the role checks and the launch/island flags. The end-to-end case with +the real Rust syncer is in test_rl_synced_trainers_sampler.py.""" + +from __future__ import annotations + +import json +import shlex +from types import SimpleNamespace + +import pytest +from test_rl_algorithm_provenance import _fake_sky, _no_modal_listing # noqa: F401 + +from yeto.rl.adapters.miles import buffer_rollout as br +from yeto.rl.adapters.miles.sampler_island import GateVersionPublisher, ReleasingPublisher +from yeto.rl.engine import sample_roles +from yeto.rl.engine.merged_release import ( + BufferLane, + LaneCursorLog, + MergedReleaseGuard, + gate_version, + is_publisher, + merged_version, +) +from yeto.rl.engine.policy_release import PolicyReleaseStore +from yeto.rl.engine.resume import CheckpointStore +from yeto.rl.engine.sample_buffer import FileSampleBuffer, SampleBufferError +from yeto.rl.engine.trainer_lifecycle import SampleIndexEntry + + +def _fill(buffer: FileSampleBuffer, n: int, version: int = 0) -> None: + for i in range(len(buffer), len(buffer) + n): + ref = buffer.write_payload(f"s{i}", br.encode_group([{"tokens": [i]}])) + buffer.append(SampleIndexEntry(f"s{i}", version, ref)) + + +# ---------------------------------------------------------------- lanes +def test_lane_parse_and_ownership(): + assert BufferLane.parse(None) is None and BufferLane.parse("0/1") is None + lane = BufferLane.parse("1/3") + assert lane == BufferLane(1, 3) and lane.env() == "1/3" + assert [p for p in range(9) if lane.owns(p)] == [1, 4, 7] + with pytest.raises(ValueError): + BufferLane(2, 2) + with pytest.raises(ValueError): + BufferLane.parse("x/2") + + +def test_two_lane_readers_split_one_buffer_without_overlap(tmp_path): + writer = FileSampleBuffer(CheckpointStore(tmp_path), batch_size=1) + _fill(writer, 7) + readers = [FileSampleBuffer(CheckpointStore(tmp_path), batch_size=1, lane=BufferLane(m, 2)) + for m in range(2)] + taken = [[s.position for s in r.select("trainer", 10)] for r in readers] + assert taken == [[0, 2, 4, 6], [1, 3, 5]] + assert readers[0].read_cursor == readers[1].read_cursor == 7 # foreign positions are passed + # a limited select stops at the limit and resumes after the last lane position + again = FileSampleBuffer(CheckpointStore(tmp_path), batch_size=1, lane=BufferLane(0, 2)) + assert [s.position for s in again.select("trainer", 2)] == [0, 2] + assert [s.position for s in again.select("trainer", 2)] == [4, 6] + + +def test_lane_reader_refuses_an_object_store_and_a_shard_merge(tmp_path): + with pytest.raises(SampleBufferError, match="shared-directory"): + FileSampleBuffer(CheckpointStore(tmp_path), batch_size=1, lane=BufferLane(0, 2), + merge_shards=True) + + +def test_reader_continues_after_its_lane_cursor_and_never_retrains(tmp_path): + store = CheckpointStore(tmp_path) + writer = FileSampleBuffer(store, batch_size=1) + _fill(writer, 8) + lane = BufferLane(1, 2) + + def reader(): + buffer = FileSampleBuffer(CheckpointStore(tmp_path), batch_size=1, lane=lane) + return br.BatchReader(buffer, max_policy_age=2, cursor_log=br.open_lane_log(buffer)) + + first = reader() + groups, report = first.take(0, 2, current_version=0) + assert [p for p, _ in groups] == [1, 3] and report.end == 4 + # restart (a re-joined island): the next round continues after position 3 + second = reader() + assert second.buffer.read_cursor == 4 + groups, _ = second.take(1, 2, current_version=1) + assert [p for p, _ in groups] == [5, 7] + lines = [json.loads(line) for line in + (tmp_path / "sample-buffer" / "lanes" / "1-of-2.jsonl").read_text().splitlines()] + assert [(r["rollout_id"], r["positions"], r["cursor"]) for r in lines] == \ + [(0, [1, 3], 4), (1, [5, 7], 8)] + + +def test_lane_cursor_log_skips_a_torn_last_line(tmp_path): + log = LaneCursorLog(tmp_path, BufferLane(0, 2)) + log.append(0, [0, 2], 3) + with log.path.open("a") as fh: + fh.write('{"cursor": 9') + assert log.last_cursor() == 3 + + +def test_gate_uses_the_merged_version_once_set(tmp_path): + writer = FileSampleBuffer(CheckpointStore(tmp_path), batch_size=1) + _fill(writer, 2, version=1) + _fill(writer, 2, version=4) + reader = br.BatchReader(FileSampleBuffer(CheckpointStore(tmp_path), batch_size=1), + max_policy_age=1, wait_s=0.0) + reader.gate_version = 5 # merged version 5: version 1 is 4 behind -> passed over + groups, report = reader.take(0, 2, current_version=1) + assert [p for p, _ in groups] == [2, 3] and report.stale_groups == 2 + assert gate_version(3, None) == 3 and gate_version(3, 7) == 7 + + +def test_executor_gate_version_setter_reaches_the_active_reader(tmp_path, monkeypatch): + monkeypatch.setattr(br, "_ACTIVE", {}) + assert br.executor_set_gate_version(None, 3) is None # no reader yet + reader = br.BatchReader(FileSampleBuffer(CheckpointStore(tmp_path), batch_size=1), + max_policy_age=1) + br.register_reader(reader) + assert br.executor_set_gate_version(None, 3) is None + assert br.executor_set_gate_version(None, 4) == 3 and reader.gate_version == 4 + + +def test_island_env_carries_the_lane_only_for_a_trainer(): + env = br.island_env("trainer", store_root="/s", max_policy_age=2, lane=BufferLane(1, 2)) + assert env[br.LANE_ENV] == "1/2" and br.LANE_ENV in br.ISLAND_ENVS + with pytest.raises(ValueError): + br.island_env("sampler", store_root="/s", sampler_id="a", lane=BufferLane(1, 2)) + + +# ---------------------------------------------------------------- versions and the publisher +def test_publisher_and_merged_version_rules(): + assert is_publisher(0, 2) and not is_publisher(1, 2) + with pytest.raises(ValueError): + is_publisher(2, 2) + state = SimpleNamespace(policy_version=4) + assert merged_version(SimpleNamespace(outer_version=None), state) == 4 + assert merged_version(SimpleNamespace(outer_version=6), state) == 6 + assert merged_version(SimpleNamespace(), state) == 4 + + +def test_guard_releases_increasing_versions_and_reports_gaps(): + gaps = [] + guard = MergedReleaseGuard(lambda a, b: gaps.append((a, b))) + assert guard.check(0, None) == "release" + assert guard.check(1, 0) == "release" + assert guard.check(1, 1) == "stale" and guard.check(0, 1) == "stale" + assert guard.check(4, 1) == "release" and gaps == [(1, 4)] + + +class _State: + def __init__(self, version: int, digest: str) -> None: + self.policy_version = version + self._digest = digest + + def policy_tensor_hash(self) -> str: + return self._digest + + +class _Inner: + def __init__(self) -> None: + self.published: list[int] = [] + + def publish(self, state): + self.published.append(state.policy_version) + return "ok" + + +def test_releasing_publisher_releases_merged_versions_in_order(tmp_path): + events: list[tuple[str, dict]] = [] + releases = PolicyReleaseStore(CheckpointStore(tmp_path)) + driver = SimpleNamespace(outer_version=None) + gates: list[int] = [] + inner = GateVersionPublisher(_Inner(), gates.append, lambda s: merged_version(driver, s)) + publisher = ReleasingPublisher( + inner, releases, lambda v, d: (d / "w.bin").write_bytes(b"%d" % v), lambda: 3, + lambda event, **f: events.append((event, f)), + version_of=lambda s: merged_version(driver, s), + guard=MergedReleaseGuard(lambda a, b: events.append(("gap", {"a": a, "b": b}))), + origin=lambda s: {"kind": "merged", "learner_id": 0, "num_learners": 2, + "local_version": s.policy_version}) + for local, outer, digest in ((0, 0, "h0"), (1, 1, "h1"), (2, 1, "h1"), (3, 3, "h3")): + driver.outer_version = outer + assert publisher.publish(_State(local, digest)) == "ok" + assert gates == [0, 1, 1, 3] + assert [e for e, _ in events] == ["rl_policy_released", "rl_policy_released", + "rl_policy_release_skipped", "gap", "rl_policy_released"] + latest = releases.latest() + assert (latest.seq, latest.version, latest.policy_hash) == (3, 3, "h3") + assert latest.origin == {"kind": "merged", "learner_id": 0, "num_learners": 2, + "local_version": 3} + assert latest.buffer_selected == 3 + + +# ---------------------------------------------------------------- role checks +def _role_args(**kw): + base = dict(rl_sample_role="trainer", rl_single_island_no_sync=False, + rl_checkpoint_store="modal-volume://v/p", rl_max_policy_age=2, tuning="lora") + return SimpleNamespace(**{**base, **kw}) + + +def test_a_trainer_may_sync_a_sampler_may_not(): + sample_roles.check_launch(_role_args()) + sample_roles.check_synced_trainers(_role_args(), 2) + with pytest.raises(ValueError, match="only trainers may sync"): + sample_roles.check_launch(_role_args(rl_sample_role="sampler")) + with pytest.raises(ValueError, match="placement reconfiguration"): + sample_roles.check_launch(_role_args(rl_elastic=True)) + + +@pytest.mark.parametrize("kw,islands,message", [ + ({}, 1, "at least two islands"), + ({"tuning": "full"}, 2, "--tuning lora"), + ({"rl_checkpoint_store": "s3://b/p"}, 2, "object store"), +]) +def test_synced_trainer_refusals(kw, islands, message): + with pytest.raises(ValueError, match=message): + sample_roles.check_synced_trainers(_role_args(**kw), islands) + + +def test_buffer_lane_of_island_args(): + args = _role_args(learner_id=1, rl_sample_lanes=2) + assert sample_roles.buffer_lane(args) == BufferLane(1, 2) + assert sample_roles.buffer_lane(_role_args(learner_id=0, rl_single_island_no_sync=True, + rl_sample_lanes=2)) is None + resolved = sample_roles.resolve(_role_args(rl_sample_store="/root/store"), groups_per_round=4) + assert resolved.store == "/root/store" and resolved.role == "trainer" + + +# ---------------------------------------------------------------- launch and island flags +SYNCED = ("--rl-sample-role", "trainer", "--rl-max-policy-age", "2", + "--rl-checkpoint-store", "modal-volume://yeto-c3/run-a", "--rl-sample-wait-s", "900", + "--rl-island-scheduling", "elastic", "--cluster-prefix", "c3-synced", + "--rl-algorithm-spec", "examples/rl_algorithms/decoupled-ppo-c2.json") + + +@pytest.fixture(autouse=True) +def _island_hmac_key(monkeypatch): + monkeypatch.setenv("YETO_ISLAND_HMAC_KEY", "c3-test-key") + + +def _synced_plan(extra=()): + from test_rl_sampler_trainer_pair_launch import _args + + from yeto.launcher import dry_run_plan + + args = _args(tuple(extra), base=SYNCED) + return args, dry_run_plan(args) + + +def _island_tokens(command: str, learner: str) -> list[str]: + tail = command.split("-m yeto.rl.adapters.miles.island_entry", 1)[1] + return [{"$LEARNER_ID": learner}.get(t, t) for t in shlex.split(tail)] + + +def test_synced_trainer_plan_forwards_the_sample_store_and_lanes(): + from yeto.rl.adapters.miles import island_entry + + _, plan = _synced_plan() + assert plan["syncer"] is not None + islands = plan["island_requests"] + assert len(islands) == 2 + for m, island in enumerate(islands): + tokens = _island_tokens(island["learner_command"], str(m)) + assert tokens[tokens.index("--rl-sample-role") + 1] == "trainer" + assert tokens[tokens.index("--rl-sample-lanes") + 1] == "2" + assert "--rl-sample-store" in tokens and "--rl-resume-store" not in tokens + parsed = island_entry.parse_args(tokens) + assert sample_roles.buffer_lane(parsed) == BufferLane(m, 2) + + +def test_synced_trainer_launch_refuses_round_cuts(): + with pytest.raises(ValueError, match="no round cuts"): + _synced_plan(("--rl-cut-every", "1")) + + +def test_island_refuses_a_sample_store_outside_a_synced_trainer(): + from yeto.rl.adapters.miles import island_entry + + with pytest.raises(ValueError, match="synced trainer island needs"): + island_entry._check_sample_store(SimpleNamespace( + rl_sample_role="trainer", rl_single_island_no_sync=False, rl_sample_store=None, + rl_sample_lanes=None, rl_resume_store=None)) + with pytest.raises(ValueError, match="need --rl-sample-role trainer"): + island_entry._check_sample_store(SimpleNamespace( + rl_sample_role="sampler", rl_single_island_no_sync=True, rl_sample_store="/s", + rl_sample_lanes=2, rl_resume_store=None)) diff --git a/tests/test_rl_synced_trainers_sampler.py b/tests/test_rl_synced_trainers_sampler.py new file mode 100644 index 000000000..61f4d0966 --- /dev/null +++ b/tests/test_rl_synced_trainers_sampler.py @@ -0,0 +1,203 @@ +"""S25 C3 end to end on CPU: two synced trainer islands, the real Rust syncer +(elastic OuterStep merge), one sampler that follows the releases. + +The trainers are fake ports islands (``yeto.rl.engine.fake``) wired with the +production publisher wrappers (``sampler_island.wire_policy_releases``) and the +production lane reader (``buffer_rollout.BatchReader`` over a lane of one +``FileSampleBuffer``). The sampler is a thread: ``PolicyFollower`` loads each +release and writes groups stamped with the version it serves. + +Checks (the same as the GPU run's, SESSION25/gpu-review/T-C3.md): +1. both trainers apply the same merged policy hash at every outer version; +2. the sampler loads exactly the merged versions 0..R, in order, each with the + merged hash, all released by learner 0; +3. the two lanes are disjoint and no buffer position is trained twice; +4. the syncer ran R OuterSteps. + +Runs when YETO_TEST_ELASTIC_SYNCER points to a syncer binary (as +test_rl_inter_island_ports.py).""" + +from __future__ import annotations + +import asyncio +import json +import os +import subprocess +import threading +import time +from pathlib import Path +from types import SimpleNamespace + +import pytest +import torch + +from test_rl_engine_driver import NAME, _driver, _engine, _strict_config +from test_rl_inter_island_ports import _free_port +from yeto.rl.adapters.miles import buffer_rollout as br +from yeto.rl.adapters.miles.sampler_island import wire_policy_releases +from yeto.rl.elastic_client import ElasticClientConfig, ElasticIslandClient +from yeto.rl.engine.bridges import ElasticAvgSync +from yeto.rl.engine.merged_release import BufferLane +from yeto.rl.engine.policy_release import PolicyFollower, PolicyReleaseStore +from yeto.rl.engine.resume import CheckpointStore +from yeto.rl.engine.sample_buffer import FileSampleBuffer +from yeto.rl.engine.sample_roles import SampleBufferConfig +from yeto.rl.engine.trainer_lifecycle import SampleIndexEntry + +ROUNDS = 3 +MAX_POLICY_AGE = 2 +_EXECUTOR_LOCK = threading.Lock() + + +class _LaneRollout: + """The fake rollout pool of one trainer: each round first takes its groups + from the island's buffer lane (as ``BufferReaderRolloutFn``), then lets the + fake pool build the batch. ``_in_executor`` runs a function as in the Miles + rollout executor, with this island's reader active.""" + + def __init__(self, pool, reader: br.BatchReader, groups: int) -> None: + self.pool, self.reader, self.groups = pool, reader, groups + self.taken: list[tuple[int, list[int], dict[int, int]]] = [] + + def generate(self, rollout_id, **kwargs): + groups, report = self.reader.take(rollout_id, self.groups, current_version=None) + self.taken.append((rollout_id, [p for p, _ in groups], report.behavior_versions)) + return self.pool.generate(rollout_id, **kwargs) + + def _in_executor(self, fn, *args): + with _EXECUTOR_LOCK: + br._ACTIVE[br.TRAINER] = self.reader + return fn(None, *args) + + def __getattr__(self, name): + return getattr(self.pool, name) + + +class _Actor: + """``actor_model.export_hf``: the full weights the engine serves.""" + + def __init__(self, engine) -> None: + self.engine = engine + + async def export_hf(self, *, rollout_id: int, path: str) -> None: + torch.save(self.engine.published_tensors, Path(path) / "model.pt") + + +def _sampler(store_root: Path, stop: threading.Event, loads: list, groups_ahead: int) -> None: + store = CheckpointStore(store_root) + buffer = FileSampleBuffer(store, batch_size=1) + follower = PolicyFollower( + PolicyReleaseStore(store), + lambda release, directory: loads.append( + (release.version, release.policy_hash, release.origin, + torch.load(directory / "model.pt")[NAME].tolist()))) + n = 0 + while not stop.is_set(): + follower.refresh() + release = follower.release + if release is None or len(buffer) - int(release.buffer_selected or 0) >= groups_ahead: + time.sleep(0.005) + continue + ref = buffer.write_payload(f"g{n}", br.encode_group([{"tokens": [n]}])) + buffer.append(SampleIndexEntry(f"g{n}", follower.version, ref)) + n += 1 + buffer.mark_data_end() + + +def _events(path: Path, kind: str) -> list[dict]: + return [r for r in (json.loads(line) for line in path.read_text().splitlines()) + if r.get("event") == kind] + + +@pytest.mark.skipif(not os.environ.get("YETO_TEST_ELASTIC_SYNCER"), + reason="needs YETO_TEST_ELASTIC_SYNCER (syncer binary with elastic mode)") +def test_two_synced_trainers_feed_one_sampler_through_merged_releases(tmp_path): + port, key = _free_port(), "c3-e2e" + out = Path(os.environ.get("YETO_TEST_ELASTIC_OUT", tmp_path)) + out.mkdir(parents=True, exist_ok=True) + for old in out.glob("*.jsonl"): + old.unlink() + store_root = out / "store" + store_root.mkdir(exist_ok=True) + syncer = subprocess.Popen( + [os.environ["YETO_TEST_ELASTIC_SYNCER"], "--port", str(port), "--learners", "2", + "--total-steps", str(ROUNDS), "--event-tape", str(out / "syncer-tape.jsonl"), + "--island-scheduling-mode", "elastic", "--quorum-theta", "0.75", "--carry-gamma", "0.5", + "--soft-deadline-s", "30", "--q-min", "1", "--max-carry-lag", "2", + "--island-lease-s", "6", "--syncer-epoch", "5"], + env=dict(os.environ, YETO_ISLAND_HMAC_KEY=key), + stdout=open(out / "syncer.log", "w"), stderr=subprocess.STDOUT) + stop, loads = threading.Event(), [] + try: + time.sleep(1.0) + engines = (_engine(torch.tensor([1.0, 3.0])), _engine(torch.tensor([3.0, 5.0]))) + sampler = threading.Thread(target=_sampler, daemon=True, + args=(store_root, stop, loads, 4 * engines[0].groups)) + sampler.start() + drivers, rollouts = [], [] + for i, engine in enumerate(engines): + client = ElasticIslandClient(ElasticClientConfig(("127.0.0.1", port), i, lease_s=6.0, + syncer_epoch=5), key.encode()) + cfg = _strict_config(out, engine, learner_id=i, rounds=ROUNDS, tape=f"island-{i}.jsonl") + driver = _driver(engine, ElasticAvgSync(cfg, client=client, base_wait_s=30.0), out, + name=f"events-{i}.jsonl", learner_id=i) + lane = BufferLane(i, 2) + buffer = FileSampleBuffer(CheckpointStore(store_root), batch_size=engine.groups, lane=lane) + reader = br.BatchReader(buffer, max_policy_age=MAX_POLICY_AGE, wait_s=30.0, poll_s=0.01, + cursor_log=br.open_lane_log(buffer)) + rollout = _LaneRollout(driver.rollout, reader, engine.groups) + driver.rollout = rollout + wire_policy_releases( + driver, actor_model=_Actor(engine), runner=SimpleNamespace(run=asyncio.run), + miles_args=SimpleNamespace( + yeto_rl_sample_buffer=SampleBufferConfig(role="trainer", store=str(store_root)), + yeto_rl_sample_lane=lane, lora_rank=0)) + drivers.append(driver) + rollouts.append(rollout) + results, errors = {}, {} + + def run(i, d): + try: + results[i] = d.run() + except BaseException as error: # noqa: BLE001 + errors[i] = error + + threads = [threading.Thread(target=run, args=(i, d), daemon=True) + for i, d in enumerate(drivers)] + for t in threads: + t.start() + for t in threads: + t.join(timeout=60) + assert not t.is_alive() + assert errors == {} + assert syncer.wait(15) == 0 + finally: + stop.set() + if syncer.poll() is None: + syncer.kill() + + # 1. the same merged policy on both trainers at every outer version + applied = [{e["sync/outer_version"]: e["sync/global_policy_hash"] + for e in _events(out / f"events-{i}.jsonl", "rl_policy_apply")} for i in range(2)] + assert sorted(applied[0]) == list(range(ROUNDS + 1)) and applied[0] == applied[1] + assert torch.equal(results[0].tensors[NAME], results[1].tensors[NAME]) + # 2. the sampler loaded the merged versions in order, from the publisher only + assert [v for v, *_ in loads] == list(range(ROUNDS + 1)) + assert all(h == applied[0][v] for v, h, *_ in loads) + assert {o["learner_id"] for _, _, o, _ in loads} == {0} + assert loads[-1][3] == results[0].tensors[NAME].tolist() + released = _events(out / "events-0.jsonl", "rl_policy_released") + assert [e["policy_version"] for e in released] == list(range(ROUNDS + 1)) + assert _events(out / "events-1.jsonl", "rl_policy_released") == [] + # 3. disjoint lanes, no position trained twice + positions = [[p for _, ps, _ in r.taken for p in ps] for r in rollouts] + assert all(p % 2 == i for i, ps in enumerate(positions) for p in ps) + flat = positions[0] + positions[1] + assert len(flat) == len(set(flat)) == 2 * ROUNDS * engines[0].groups + # every trained group is within max_policy_age of the merged version + for r in rollouts: + for rollout_id, _, versions in r.taken: + assert all(rollout_id - v <= MAX_POLICY_AGE for v in versions) + # 4. R OuterSteps + kinds = [json.loads(line).get("kind") for line in (out / "syncer-tape.jsonl").read_text().splitlines()] + assert kinds.count("outer_step") == ROUNDS diff --git a/yeto/launch/rl_ports.py b/yeto/launch/rl_ports.py index c1f2b369c..ebe45cacd 100644 --- a/yeto/launch/rl_ports.py +++ b/yeto/launch/rl_ports.py @@ -255,6 +255,20 @@ def _check_resume_launch_flags(args) -> bool: raise ValueError(f"{flag} needs --rl-checkpoint-store") if getattr(args, "rl_resume_allow_config_change", False) and not store: raise ValueError("--rl-resume-allow-config-change needs --rl-checkpoint-store") + from ..rl.engine.sample_roles import is_synced_trainer + + if store and not elastic and is_synced_trainer(args): + # S25 C3: the store holds the sample buffer and the releases of synced trainers + # (forwarded as --rl-sample-store); multi-island round cuts do not exist + for name, flag in (("rl_cut_every", "--rl-cut-every"), ("rl_cut_keep", "--rl-cut-keep"), + ("rl_stop_after_rounds", "--rl-stop-after-rounds")): + if getattr(args, name, None) is not None: + raise ValueError(f"{flag}: synced trainers write no round cuts (multi-island " + "resume: rl-resume-from-checkpoint 3.3)") + if getattr(args, "rl_resume_allow_config_change", False): + raise ValueError("--rl-resume-allow-config-change: synced trainers write no round cuts") + rl_checkpoint_store_plan(args) # validates the form + return False if elastic: for name, flag in (("rl_cut_keep", "--rl-cut-keep"), ("rl_stop_after_rounds", "--rl-stop-after-rounds")): if getattr(args, name, None) is not None: @@ -378,6 +392,7 @@ def _check_ports_infra_switches(args) -> None: elif sample_roles.is_pair(args): # L5.4 --rl-island-sample-role: per-island roles sample_roles.check_pair(args, rl_island_count(args)) else: + sample_roles.check_synced_trainers(args, rl_island_count(args)) # S25 C3 sample_roles.check_launch(args) # L4.6 --rl-sample-role (both checks nothing) from .storage import store_object_envs @@ -397,7 +412,8 @@ def _check_ports_infra_switches(args) -> None: if getattr(args, "rl_elastic_accept_rebind", False): given.append("--rl-elastic-accept-rebind") resume_store = _check_resume_launch_flags(args) - if getattr(args, "rl_checkpoint_store", None) and not resume_store: + if getattr(args, "rl_checkpoint_store", None) and not resume_store \ + and not sample_roles.is_synced_trainer(args): given.append("--rl-checkpoint-store") if getattr(args, "rl_elastic_declare_cells", False): given.append("--rl-elastic-declare-cells") @@ -531,6 +547,7 @@ def _ports_infra_flags(args) -> tuple[str, str]: if getattr(args, "rl_max_policy_age", 0): # 0 (default) adds nothing flags += f" --rl-max-policy-age {int(args.rl_max_policy_age)}" from ..rl.engine.sample_roles import island_flags as sample_role_flags + from ..rl.engine.sample_roles import is_synced_trainer as sample_role_is_synced_trainer flags += sample_role_flags(args) # L4.6: role both (default) adds nothing if (getattr(args, "rl_lr_schedule", None) or "auto") != "auto": # auto adds nothing @@ -542,6 +559,11 @@ def _ports_infra_flags(args) -> tuple[str, str]: resume_store = (rl_checkpoint_store_plan(args) if getattr(args, "rl_checkpoint_store", None) and not getattr(args, "rl_elastic", False) else None) + if resume_store is not None and sample_role_is_synced_trainer(args): + # S25 C3: synced trainers mount the store for the buffer and the releases only + flags += (f" --rl-sample-store {shlex.quote(resume_store[0])}" + f" --rl-sample-lanes {rl_island_count(args)}") + resume_store = None if resume_store is not None: # rl-resume-from-checkpoint: non-elastic resume store flags += f" --rl-resume-store {shlex.quote(resume_store[0])}" for name, flag in (("rl_cut_keep", "--rl-cut-keep"), ("rl_stop_after_rounds", "--rl-stop-after-rounds")): diff --git a/yeto/rl/adapters/miles/buffer_rollout.py b/yeto/rl/adapters/miles/buffer_rollout.py index 2207badd9..6b2ae54f1 100644 --- a/yeto/rl/adapters/miles/buffer_rollout.py +++ b/yeto/rl/adapters/miles/buffer_rollout.py @@ -28,7 +28,8 @@ Limits (first implementation): -* one trainer reads a buffer; on a shared directory (``index.jsonl``) one +* one trainer reads a buffer, or synced trainers read one lane each + (:mod:`yeto.rl.engine.merged_release`); on a shared directory (``index.jsonl``) one sampler island writes it (``flock`` holds only on one host); * the buffer lives in the checkpoint store (:func:`~yeto.rl.engine.resume.store_for` at ``BUFFER_ROOT_ENV``): a Modal Volume or a shared directory that both islands @@ -59,6 +60,7 @@ encode_group, group_entry_fields, ) +from yeto.rl.engine.merged_release import LANE_ENV, BufferLane, LaneCursorLog, gate_version SAMPLER, TRAINER, CONVERTIBLE = "sampler", "trainer", "convertible" ROLES = (SAMPLER, TRAINER, CONVERTIBLE) @@ -71,7 +73,8 @@ SHARD_ENV = "YETO_SAMPLE_BUFFER_SHARD" # convertible node: its own shard of the index # the rollout functions run in Ray workers, which inherit the raylet's environment: # connect_island_ray forwards these to every worker -ISLAND_ENVS = (BUFFER_ROOT_ENV, SAMPLER_ID_ENV, WAIT_ENV, POLL_ENV, MAX_POLICY_AGE_ENV, SHARD_ENV) +ISLAND_ENVS = (BUFFER_ROOT_ENV, SAMPLER_ID_ENV, WAIT_ENV, POLL_ENV, MAX_POLICY_AGE_ENV, SHARD_ENV, + LANE_ENV) DEFAULT_WAIT_S = 3600.0 DEFAULT_POLL_S = 2.0 @@ -111,10 +114,10 @@ def rollout_argv(role: str | None, *, max_policy_age: int, fully_async: bool = F def island_env(role: str, *, store_root: str, sampler_id: str | None = None, max_policy_age: int | None = None, wait_s: float | None = None, - shard: str | None = None) -> dict[str, str]: + shard: str | None = None, lane: BufferLane | None = None) -> dict[str, str]: """The environment the launch layer exports on the island (read in the rollout process). A convertible node writes like a sampler, into its own ``shard`` of - the index, and reads like a trainer.""" + the index, and reads like a trainer. ``lane``: a synced trainer's buffer lane.""" if role not in ROLES: raise ValueError(f"sample buffer role {role!r}: expected one of {ROLES}") env = {BUFFER_ROOT_ENV: str(store_root)} @@ -132,6 +135,10 @@ def island_env(role: str, *, store_root: str, sampler_id: str | None = None, env[MAX_POLICY_AGE_ENV] = str(int(max_policy_age)) if wait_s is not None: env[WAIT_ENV] = f"{float(wait_s):g}" + if lane is not None: + if role != TRAINER: + raise ValueError("only a trainer reads a buffer lane") + env[LANE_ENV] = lane.env() return env @@ -160,9 +167,22 @@ def open_buffer(environ: Mapping[str, str] | None = None, *, batch_size: int = 1 if not shard: raise SampleBufferRolloutError(f"{SAMPLER_ID_ENV} is not set (the sampler's buffer shard)") merge = reader + lane = BufferLane.parse(env.get(LANE_ENV)) if reader else None return FileSampleBuffer(store, batch_size=batch_size, index=INDEX if shard is None else shard_index(shard), - create_index=create_index, merge_shards=merge) + create_index=create_index, merge_shards=merge and lane is None, + lane=lane) + + +def open_lane_log(buffer: Any) -> LaneCursorLog | None: + """The lane cursor log of a lane reader (None for a single reader).""" + lane = getattr(buffer, "lane", None) + if lane is None: + return None + from yeto.rl.engine.sample_buffer import BUFFER_DIR + + store = buffer.store + return LaneCursorLog(f"{store.root}/{BUFFER_DIR}", lane, commit=store.commit) def renumber_group(samples: Sequence[dict[str, Any]], *, position: int, @@ -215,7 +235,10 @@ class BatchReader: def __init__(self, buffer: Any, *, max_policy_age: int, wait_s: float = DEFAULT_WAIT_S, poll_s: float = DEFAULT_POLL_S, consumer: str = TRAINER_CONSUMER, clock: Callable[[], float] = time.monotonic, - sleep: Callable[[float], None] = time.sleep) -> None: + sleep: Callable[[float], None] = time.sleep, + cursor_log: LaneCursorLog | None = None) -> None: + """``cursor_log``: a synced trainer's lane cursor (:mod:`yeto.rl.engine.merged_release`): + the reader continues after its last line and appends one per round.""" if int(max_policy_age) < 1: raise ValueError("max_policy_age must be >= 1") self.buffer = buffer @@ -226,6 +249,11 @@ def __init__(self, buffer: Any, *, max_policy_age: int, wait_s: float = DEFAULT_ self.clock = clock self.sleep = sleep self._ends: dict[int, int] = {} # rollout id -> buffer position after its round + # a synced trainer gates on the merged version it last published (set_gate_version) + self.gate_version: int | None = None + self.cursor_log = cursor_log + if cursor_log is not None: + self.buffer.reset_to(cursor_log.last_cursor()) def take(self, rollout_id: int, groups: int, *, current_version: int | None ) -> tuple[list[tuple[int, list[dict[str, Any]]]], TakeReport]: @@ -240,6 +268,7 @@ def take(self, rollout_id: int, groups: int, *, current_version: int | None out: list[tuple[int, list[dict[str, Any]]]] = [] versions: dict[int, int] = {} stale = 0 + current_version = gate_version(current_version, self.gate_version) floor = None if current_version is None else int(current_version) - self.max_policy_age while len(out) < groups: selected = self.buffer.select(self.consumer, groups - len(out)) @@ -264,6 +293,8 @@ def take(self, rollout_id: int, groups: int, *, current_version: int | None raise SampleBufferRolloutError(f"rollout {rollout_id}: the sample buffer ended " f"with no trainable group (position {end})") self._ends[rollout_id] = end + if self.cursor_log is not None: + self.cursor_log.append(rollout_id, [position for position, _ in out], end) return out, TakeReport(rollout_id=rollout_id, groups=len(out), stale_groups=stale, start=start, end=end, behavior_versions=dict(sorted(versions.items())), wait_s=round(self.clock() - started, 3)) @@ -340,6 +371,17 @@ def executor_read_shards(_executor: Any) -> dict[str, int] | None: return reader.buffer.shard_positions(reader.buffer.read_cursor) or None +def executor_set_gate_version(_executor: Any, version: int) -> int | None: + """Runs INSIDE the rollout executor actor: a synced trainer's merged version, + the version its ``max_policy_age`` gate compares against. Returns the previous + one (None before the reader exists: nothing to set).""" + reader = _ACTIVE.get(TRAINER) + if reader is None: + return None + previous, reader.gate_version = reader.gate_version, int(version) + return previous + + def executor_reset(_executor: Any, watermark: int) -> int: """Runs INSIDE the rollout executor actor (``__ray_call__``).""" return _active_reader().reset(int(watermark)) diff --git a/yeto/rl/adapters/miles/buffer_rollout_miles.py b/yeto/rl/adapters/miles/buffer_rollout_miles.py index e9f6bad87..55552c7fa 100644 --- a/yeto/rl/adapters/miles/buffer_rollout_miles.py +++ b/yeto/rl/adapters/miles/buffer_rollout_miles.py @@ -39,6 +39,7 @@ SampleBufferRolloutError, convertible_mode, open_buffer, + open_lane_log, register_reader, renumber_group, ) @@ -94,8 +95,9 @@ def __init__(self, args: Any, *, create_index: bool = True, merge_shards: bool = self.args = args self.groups = int(args.rollout_batch_size) self.samples_per_group = int(args.n_samples_per_prompt) + buffer = open_buffer(batch_size=self.groups, create_index=create_index, reader=merge_shards) self.reader = BatchReader( - open_buffer(batch_size=self.groups, create_index=create_index, reader=merge_shards), + buffer, cursor_log=open_lane_log(buffer), max_policy_age=int(age), wait_s=float(os.environ.get(WAIT_ENV, DEFAULT_WAIT_S)), poll_s=float(os.environ.get(POLL_ENV, DEFAULT_POLL_S))) diff --git a/yeto/rl/adapters/miles/island_entry.py b/yeto/rl/adapters/miles/island_entry.py index 9a031e335..73d27d2cd 100644 --- a/yeto/rl/adapters/miles/island_entry.py +++ b/yeto/rl/adapters/miles/island_entry.py @@ -192,6 +192,7 @@ def _check_ports_infra_switches(args) -> None: if getattr(args, "rl_elastic_declare_cells", False) and not getattr(args, "rl_elastic", False): raise ValueError("--rl-elastic-declare-cells needs --rl-elastic") _check_resume_switches(args) + _check_sample_store(args) if not getattr(args, "rl_elastic", False): if given: raise ValueError(", ".join(given) + " need --rl-elastic") @@ -210,6 +211,25 @@ def _check_ports_infra_switches(args) -> None: raise ValueError(f"--{name.replace('_', '-')} must be positive") +def _check_sample_store(args) -> None: + """S25 C3: ``--rl-sample-store`` / ``--rl-sample-lanes`` belong to a synced + trainer island (``yeto.rl.engine.merged_release``).""" + from yeto.rl.engine.sample_roles import buffer_lane, is_synced_trainer + + store, lanes = getattr(args, "rl_sample_store", None), getattr(args, "rl_sample_lanes", None) + if not is_synced_trainer(args): + if store is not None or lanes is not None: + raise ValueError("--rl-sample-store / --rl-sample-lanes need --rl-sample-role trainer " + "with a syncer (no --rl-single-island-no-sync)") + return + if not store or lanes is None or int(lanes) < 2: + raise ValueError("a synced trainer island needs --rl-sample-store and --rl-sample-lanes " + ">= 2 (the launcher forwards both)") + if getattr(args, "rl_resume_store", None): + raise ValueError("a synced trainer island takes --rl-sample-store, not --rl-resume-store") + buffer_lane(args) # learner id inside 0..lanes-1 + + def _check_resume_switches(args) -> None: """rl-resume-from-checkpoint: ``--rl-resume-store`` is single-island/no-sync; the cut cadence flags need a store (elastic or resume).""" @@ -261,7 +281,8 @@ def apply_sample_role(args, miles_args, environ=None) -> None: (os.environ if environ is None else environ).update(island_env( config.role, store_root=config.store, sampler_id=node or f"island{int(args.learner_id)}", max_policy_age=int(getattr(args, "rl_max_policy_age", 0) or 0), wait_s=config.wait_s, - shard=node)) + shard=node, lane=sample_roles.buffer_lane(args))) + miles_args.yeto_rl_sample_lane = sample_roles.buffer_lane(args) # S25 C3: None off a synced trainer def apply_ports_infra_switches(args, miles_args, environ=None) -> None: diff --git a/yeto/rl/adapters/miles/run_flags.py b/yeto/rl/adapters/miles/run_flags.py index a6a9c6dad..b4157b6fa 100644 --- a/yeto/rl/adapters/miles/run_flags.py +++ b/yeto/rl/adapters/miles/run_flags.py @@ -372,6 +372,11 @@ # bucket mount or a shared path); round cuts + a hashed state snapshot + LATEST go there. island_flag("--rl-resume-store", default=None, metavar="PATH"), island_flag("--rl-resume-state-dir", default="~/yeto-rl/resume-state", metavar="PATH"), + # S25 C3: a synced trainer island (--rl-sample-role trainer with a syncer): the + # store of the sample buffer and the policy releases (no round cuts), and the + # number of synced trainers (its buffer lane is learner id of that count) + island_flag("--rl-sample-store", default=None, metavar="PATH"), + island_flag("--rl-sample-lanes", type=int, default=None, metavar="N"), island_flag("--lora-r", type=int, default=None), island_flag( "--lora-targets", diff --git a/yeto/rl/adapters/miles/sampler_island.py b/yeto/rl/adapters/miles/sampler_island.py index 4c49446a0..dd2982279 100644 --- a/yeto/rl/adapters/miles/sampler_island.py +++ b/yeto/rl/adapters/miles/sampler_island.py @@ -9,7 +9,9 @@ while it trains), every successful publish is also released into the store (:class:`ReleasingPublisher`): a full-parameter trainer exports an HF checkpoint (``actor.export_hf``), a LoRA trainer only its - adapter (:mod:`.release_plugin`), into the release directory. + adapter (:mod:`.release_plugin`), into the release directory. Synced trainers + (several islands and a syncer): only learner 0 releases, with the merged + version (:mod:`yeto.rl.engine.merged_release`). Engine weight load (the order of the publisher, ``publish.MilesPublisher``): under the controller's update lock, every engine loads the release, then the @@ -45,10 +47,17 @@ PolicyRelease, PolicyReleaseStore, ) +from yeto.rl.engine.merged_release import MergedReleaseGuard, is_publisher, merged_version from yeto.rl.engine.policy_versions import engine_weight_version from yeto.rl.engine.sample_roles import CONVERTIBLE, SAMPLER, TRAINER -from .buffer_rollout import SampleBufferRolloutError, executor_read_cursor, executor_read_shards, open_buffer +from .buffer_rollout import ( + SampleBufferRolloutError, + executor_read_cursor, + executor_read_shards, + executor_set_gate_version, + open_buffer, +) from .publish import ( _engine_id, _maybe_await, @@ -326,10 +335,19 @@ def __init__(self, inner: Any, releases: PolicyReleaseStore, export: Callable[[int, Path], None], read_position: Callable[[], int | None], emit: Callable[..., None], *, weights: str = FULL, - read_shards: Callable[[], dict[str, int] | None] | None = None) -> None: + read_shards: Callable[[], dict[str, int] | None] | None = None, + version_of: Callable[[Any], int] | None = None, + guard: MergedReleaseGuard | None = None, + origin: Callable[[Any], dict[str, Any]] | None = None) -> None: """``read_shards``: the trainer's read position per sampler shard (None without merged shards), released with every policy for the samplers' - backlog caps.""" + backlog caps. Synced trainers (:mod:`yeto.rl.engine.merged_release`): + ``version_of(state)`` is the merged version to release (default: the + policy version), ``guard`` keeps the release versions increasing and + ``origin(state)`` is the release record's ``origin``.""" + self.version_of = version_of or (lambda state: int(state.policy_version)) + self.guard = guard + self.origin = origin self.inner = inner self.releases = releases self.export = export @@ -346,7 +364,15 @@ def _engine_checksum(self) -> str | None: def publish(self, state: Any, **kwargs: Any) -> Any: result = self.inner.publish(state, **kwargs) - version = int(state.policy_version) + version = int(self.version_of(state)) + if self.guard is not None: + latest = self.releases.latest() + if self.guard.check(version, None if latest is None else latest.version) == "stale": + self.emit("rl_policy_release_skipped", policy_version=version, + local_version=int(state.policy_version), + latest_version=None if latest is None else latest.version, + reason="not_newer_than_latest") + return result position = self.read_position() shards = self.read_shards() if self.read_shards is not None and position is not None else None if position is None: # nothing read yet in this process: keep the last known positions @@ -356,9 +382,11 @@ def publish(self, state: Any, **kwargs: Any) -> Any: release = self.releases.release(version, state.policy_tensor_hash(), lambda directory: self.export(version, directory), buffer_selected=position, weights=self.weights, - engine_checksum=self._engine_checksum(), buffer_shards=shards) + engine_checksum=self._engine_checksum(), buffer_shards=shards, + origin=None if self.origin is None else self.origin(state)) if release is not None: self.emit("rl_policy_released", seq=release.seq, policy_version=version, + 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) @@ -368,10 +396,35 @@ def __getattr__(self, name: str) -> Any: return getattr(self.inner, name) +class GateVersionPublisher: + """A synced trainer's publisher (:mod:`yeto.rl.engine.merged_release`): after + each publish, the island's buffer reader gates ``max_policy_age`` on the merged + version (``set_gate_version(version)``). Every other attribute is the wrapped + publisher's.""" + + def __init__(self, inner: Any, set_gate_version: Callable[[int], Any], + version_of: Callable[[Any], int]) -> None: + self.inner = inner + self.set_gate_version = set_gate_version + self.version_of = version_of + + def publish(self, state: Any, **kwargs: Any) -> Any: + result = self.inner.publish(state, **kwargs) + self.set_gate_version(int(self.version_of(state))) + return result + + def __getattr__(self, name: str) -> Any: + return getattr(self.inner, name) + + def wire_policy_releases(driver: Any, *, miles_args: Any, actor_model: Any, runner: Any, environ: Any = None) -> ReleasingPublisher | None: """A trainer island (``--rl-sample-role trainer``) releases every publish; so does - the trainer of a convertible node (the other nodes follow the releases).""" + the trainer of a convertible node (the other nodes follow the releases). + + Synced trainers (a buffer lane on ``miles_args.yeto_rl_sample_lane``): every + island gates on the merged version; only the publisher (learner id 0) releases, + with the merged version (:mod:`yeto.rl.engine.merged_release`).""" config = _buffer_config(miles_args, TRAINER) or _buffer_config(miles_args, CONVERTIBLE) if config is None: return None @@ -379,6 +432,27 @@ def wire_policy_releases(driver: Any, *, miles_args: Any, actor_model: Any, runn if not callable(in_executor): raise SampleBufferRolloutError(f"rollout pool {type(driver.rollout).__name__} cannot " "run code in the rollout executor") + lane = getattr(miles_args, "yeto_rl_sample_lane", None) + synced: dict[str, Any] = {} + if lane is not None: + def version_of(state: Any) -> int: + return merged_version(driver, state) + + def origin(state: Any) -> dict[str, Any]: + return {"kind": "merged", "learner_id": lane.index, "num_learners": lane.count, + "local_version": int(state.policy_version)} + + def on_gap(previous: int, version: int) -> None: + driver.emit("rl_policy_release_gap", previous_version=previous, policy_version=version) + + driver.publisher = GateVersionPublisher( + driver.publisher, lambda v: in_executor(executor_set_gate_version, int(v)), version_of) + publisher_island = is_publisher(lane.index, lane.count) + driver.emit("rl_policy_release_role", learner_id=lane.index, num_learners=lane.count, + publisher=publisher_island) + if not publisher_island: + return None + synced = {"version_of": version_of, "guard": MergedReleaseGuard(on_gap), "origin": origin} lora = lora_mode(miles_args) @@ -391,6 +465,6 @@ def export(version: int, directory: Path) -> None: publisher = ReleasingPublisher(driver.publisher, release_store(config, environ), export, read_position=lambda: in_executor(executor_read_cursor), read_shards=lambda: in_executor(executor_read_shards), - emit=driver.emit, weights=LORA if lora else FULL) + emit=driver.emit, weights=LORA if lora else FULL, **synced) driver.publisher = publisher return publisher diff --git a/yeto/rl/engine/merged_release.py b/yeto/rl/engine/merged_release.py new file mode 100644 index 000000000..0f953ff4a --- /dev/null +++ b/yeto/rl/engine/merged_release.py @@ -0,0 +1,175 @@ +"""Synced trainer islands feed sampler islands (S25 C3). + +N trainer islands (``--rl-sample-role trainer`` without +``--rl-single-island-no-sync``) train against one syncer (strict or elastic +OuterStep merge) and read their rounds from one durable sample buffer. One or +more sampler islands (separate no-sync launches over the same store) write the +buffer and follow the policy releases. This module holds the engine-neutral +rules; the Miles wiring is in ``adapters.miles.sampler_island`` and +``adapters.miles.buffer_rollout``. + +Who releases: exactly one island, the *publisher*, the trainer with learner id +:data:`PUBLISHER_LEARNER_ID` (0). Every trainer applies the same merged policy +at a merge boundary (the syncer sends one base to all), so one writer is enough, +and one writer keeps ``policy-releases/`` free of write races between hosts (a +Modal Volume or a bucket has no cross-host lock). The other trainers release +nothing. + +When: after every successful publish of the publisher, the same point a no-sync +trainer releases (``ReleasingPublisher``): the merged policy is applied, +verified by hash and loaded on the trainer's own engines before the sampler sees +it. + +Version rule: a release carries the *merged* version, the syncer's outer +version of the applied base (``IslandDriver.outer_version``; the strict syncer +has none, its committed version is the policy version). The publisher releases +only versions above the newest release (:class:`MergedReleaseGuard`), so the +release versions are strictly increasing. They are consecutive (v, v+1, ...) +when the publisher applies every merge; a publisher that misses merges (it was +late or re-joined) releases the next merged version it applies, and the gap is +reported (``rl_policy_release_gap``). The sampler stamps the release version on +its engines, so every buffered group carries the merged version as its behavior +version, and each trainer gates with ``max_policy_age`` against its own merged +version (:func:`gate_version`). + +Buffer: the trainers read disjoint *lanes* of the one buffer: trainer m of N +trains the positions ``p`` with ``p % N == m`` (:class:`BufferLane`), so no +group is trained by two islands. The store must be a shared directory (local +disk, Modal Volume, a bucket mount that appends): a store read through an object +API (``CheckpointStore.objects``) is refused for now. A trainer keeps its lane +cursor in ``sample-buffer/lanes/-of-.jsonl`` (one line per round, written +when the round is selected), and a restarted or re-joined trainer continues +after it: a group selected by a round that did not finish is lost, never +trained twice (at most once). + +Failures: + +* the publisher island stops: no new release appears. The samplers keep the + last release, up to their backlog cap; the other trainers keep training on + groups within ``max_policy_age`` and pass over older ones (consumed, not + trained). When the publisher re-joins (a replacement island with learner id + 0), it applies the syncer's newest base and releases that version: the gap + is reported, the order stays increasing; +* another trainer stops: its lane is not read until it re-joins; the samplers' + backlog cap follows the publisher's read position, so the buffer grows by at + most one lane share of that cap; +* a sampler stops: the trainers wait up to ``--rl-sample-wait-s`` for groups. +""" + +from __future__ import annotations + +import json +from collections.abc import Callable +from dataclasses import dataclass +from pathlib import Path +from typing import Any + +PUBLISHER_LEARNER_ID = 0 +LANES_DIR = "lanes" +LANE_ENV = "YETO_RL_SAMPLE_BUFFER_LANE" # "M/N" on a synced trainer island + + +@dataclass(frozen=True) +class BufferLane: + """Lane ``index`` of ``count``: the buffer positions ``p % count == index``.""" + + index: int + count: int + + def __post_init__(self) -> None: + if not 0 <= int(self.index) < int(self.count): + raise ValueError(f"buffer lane {self.index}/{self.count}: need 0 <= M < N") + + def owns(self, position: int) -> bool: + return int(position) % self.count == self.index + + def env(self) -> str: + return f"{self.index}/{self.count}" + + @classmethod + def parse(cls, text: str | None) -> "BufferLane | None": + if not text: + return None + index, sep, count = str(text).partition("/") + if not sep or not index.strip().isdigit() or not count.strip().isdigit(): + raise ValueError(f"{LANE_ENV}={text!r}: expected M/N") + lane = cls(int(index), int(count)) + return None if lane.count == 1 else lane + + def file_name(self) -> str: + return f"{self.index}-of-{self.count}.jsonl" + + +def is_publisher(learner_id: int, num_learners: int) -> bool: + """Whether learner ``learner_id`` of ``num_learners`` synced trainers releases.""" + if not 0 <= int(learner_id) < int(num_learners): + raise ValueError(f"learner id {learner_id} outside 0..{int(num_learners) - 1}") + return int(learner_id) == PUBLISHER_LEARNER_ID + + +def merged_version(driver: Any, state: Any) -> int: + """The version a synced trainer releases and gates with: the syncer's outer + version of the applied base when the sync session sets one (elastic), else + the policy version (strict and no-sync: the committed version).""" + outer = getattr(driver, "outer_version", None) + return int(state.policy_version) if outer is None else int(outer) + + +def gate_version(weight_version: int | None, merged: int | None) -> int | None: + """The trainer version the ``max_policy_age`` gate compares against: the + merged version once the island published one, else the engines' version.""" + return merged if merged is not None else weight_version + + +class MergedReleaseGuard: + """Keeps the publisher's release versions strictly increasing. + + ``check(version, latest)`` returns ``"release"``, or ``"stale"`` for a + version at or below the newest release (nothing is released: the samplers + already serve that version or a newer one). A version more than one above + the newest release is released and reported by ``on_gap(previous, version)``.""" + + def __init__(self, on_gap: Callable[[int, int], None]) -> None: + self.on_gap = on_gap + + def check(self, version: int, latest_version: int | None) -> str: + if latest_version is None: + return "release" + if int(version) <= int(latest_version): + return "stale" + if int(version) > int(latest_version) + 1: + self.on_gap(int(latest_version), int(version)) + return "release" + + +class LaneCursorLog: + """The lane cursor of one synced trainer, durable in the store. + + One JSON line per selected round: ``{"rollout_id", "positions", "cursor"}`` + (``cursor``: the buffer read position after the round). :meth:`last_cursor` + is where a restarted trainer continues. ``commit`` makes the line durable + (the store's ``commit``; a Modal Volume commit).""" + + def __init__(self, root: str | Path, lane: BufferLane, + commit: Callable[[], None] = lambda: None) -> None: + self.path = Path(root) / LANES_DIR / lane.file_name() + self.commit = commit + + def last_cursor(self) -> int: + if not self.path.is_file(): + return 0 + cursor = 0 + for line in self.path.read_text(encoding="utf-8").splitlines(): + try: + cursor = int(json.loads(line)["cursor"]) + except (ValueError, KeyError, TypeError): + continue # a torn last line (crash during the write) + return cursor + + def append(self, rollout_id: int, positions: list[int], cursor: int) -> None: + self.path.parent.mkdir(parents=True, exist_ok=True) + line = json.dumps({"rollout_id": int(rollout_id), "positions": [int(p) for p in positions], + "cursor": int(cursor)}, sort_keys=True) + with self.path.open("a", encoding="utf-8") as fh: + fh.write(line + "\n") + self.commit() diff --git a/yeto/rl/engine/policy_release.py b/yeto/rl/engine/policy_release.py index fbc081d9e..45aeef17d 100644 --- a/yeto/rl/engine/policy_release.py +++ b/yeto/rl/engine/policy_release.py @@ -40,6 +40,9 @@ GB). The trainer still writes through the mount (each file is new and written once). +Synced trainers (several islands merged by a syncer) release through one +publisher, with the merged version: :mod:`.merged_release`. + A trainer that publishes the same version and hash again (a restart) releases nothing. A trainer restored from an older cut publishes an older version with other weights (a new lineage): that is a new release with a higher ``seq``. @@ -101,6 +104,9 @@ class PolicyRelease: # A release made from another one (a W4A8 re-quantization of a BF16 release): # {"release_dir", "seq", "version", "policy_hash", "scheme", ...}; None: trained weights. derived_from: dict[str, Any] | None = None + # Who released it when the trainers sync (S25 C3, yeto.rl.engine.merged_release): + # {"kind": "merged", "learner_id", "num_learners", "local_version"}; None: one trainer. + origin: dict[str, Any] | None = None schema: str = SCHEMA @classmethod @@ -276,7 +282,8 @@ def release(self, version: int, policy_hash: str, write: Callable[[Path], None], buffer_selected: int | None = None, weights: str = FULL, engine_checksum: str | None = None, buffer_shards: dict[str, int] | None = None, - derived_from: dict[str, Any] | None = None) -> PolicyRelease | None: + derived_from: dict[str, Any] | None = None, + origin: dict[str, Any] | None = None) -> PolicyRelease | None: """Write the weights with ``write(directory)``, then the record (it lists the weight files and their sizes), then LATEST. @@ -303,7 +310,8 @@ def release(self, version: int, policy_hash: str, write: Callable[[Path], None], engine_checksum=engine_checksum or None, files=file_sizes(directory), 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)) + 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)) diff --git a/yeto/rl/engine/sample_buffer.py b/yeto/rl/engine/sample_buffer.py index e4feb4023..c9f652d78 100644 --- a/yeto/rl/engine/sample_buffer.py +++ b/yeto/rl/engine/sample_buffer.py @@ -27,7 +27,8 @@ Arrival notification is by polling the index (``poll`` / ``start_polling``). -Limits of this first implementation: one consumer; with ``index.jsonl``, +Limits of this first implementation: one consumer (or one per lane: synced +trainer islands read disjoint lanes, :mod:`.merged_release`); with ``index.jsonl``, concurrent appenders must share one host (``flock``); the version gate on selection is the consumer's job (P0-S ``max_policy_age``). @@ -67,6 +68,7 @@ from pathlib import Path from typing import Any +from .merged_release import BufferLane from .object_store import ObjectStore from .resume import CheckpointStore from .trainer_lifecycle import BacklogStats, SampleIndexEntry, Selected @@ -321,11 +323,19 @@ class FileSampleBuffer: def __init__(self, store: CheckpointStore | ObjectStore, *, batch_size: int, arrival_window_s: float = 60.0, clock: Callable[[], float] = time.time, - index: str = INDEX, create_index: bool = True, merge_shards: bool = False) -> None: + index: str = INDEX, create_index: bool = True, merge_shards: bool = False, + lane: BufferLane | None = None) -> None: """``index``/``create_index``: see :class:`JsonlIndex`. ``merge_shards``: this buffer (the consumer) merges the sampler shards of an object store into the index on every refresh (see the module doc); it must be the - only one that does.""" + only one that does. ``lane``: this consumer is one of several synced + trainers and selects only the positions of its lane + (:mod:`.merged_release`); the read cursor and the watermark stay + positions of the whole index.""" + if lane is not None and (merge_shards or _objects_of(store) is not None): + raise SampleBufferError("a lane reader needs a shared-directory store (no object " + "index, no shard merge): see yeto.rl.engine.merged_release") + self.lane = lane if batch_size < 1: raise ValueError("batch_size must be positive") self.store = store @@ -481,9 +491,13 @@ def select(self, consumer: str, limit: int) -> list[Selected]: elif consumer != self._consumer: raise SampleBufferError(f"buffer is read by {self._consumer!r}, not {consumer!r}") self.refresh() - stop = min(len(self._entries), self._read + max(0, int(limit))) - out = [Selected(pos, self._entries[pos]) for pos in range(self._read, stop)] - self._read = stop + out: list[Selected] = [] + pos = self._read + while len(out) < max(0, int(limit)) and pos < len(self._entries): + if self.lane is None or self.lane.owns(pos): + out.append(Selected(pos, self._entries[pos])) + pos += 1 + self._read = pos return out def commit_consumed(self, cursor: int) -> None: diff --git a/yeto/rl/engine/sample_roles.py b/yeto/rl/engine/sample_roles.py index 10e68c357..755a20e1f 100644 --- a/yeto/rl/engine/sample_roles.py +++ b/yeto/rl/engine/sample_roles.py @@ -21,6 +21,15 @@ (``--rl-island-sample-role 0:sampler --rl-island-sample-role 1:trainer``, L5.4): there is no syncer, and each island runs as its own one-learner, no-sync island (learner id 0 of 1) whose role comes from its island args (:func:`island_args`). +Synced trainers (S25 C3, :mod:`.merged_release`): a launch of N >= 2 islands +with ``--rl-sample-role trainer`` and a syncer (strict or elastic sync, no +``--rl-single-island-no-sync``): every island is a trainer, learner m of N, and +reads lane m of the buffer; the island with learner id 0 releases the merged +policy. The store reaches the island as ``--rl-sample-store`` (buffer and +releases only: multi-island round cuts do not exist), the trainer count as +``--rl-sample-lanes``. The samplers are separate +no-sync launches over the same store. + This module is engine-neutral; the Miles mapping is in ``adapters.miles.buffer_rollout``. """ @@ -64,6 +73,42 @@ def sample_role(args: Any) -> str: return role +def is_synced_trainer(args: Any) -> bool: + """A trainer island of a synced multi-island launch (module doc).""" + return sample_role(args) == TRAINER and not getattr(args, "rl_single_island_no_sync", False) + + +def buffer_lane(args: Any) -> Any: + """The buffer lane of a synced trainer island (island args: ``learner_id`` and + ``--rl-sample-lanes``, the trainer count the launcher forwards); None for every + other island.""" + from .merged_release import BufferLane + + if not is_synced_trainer(args): + return None + count = int(getattr(args, "rl_sample_lanes", None) or 1) + return None if count == 1 else BufferLane(int(getattr(args, "learner_id", 0) or 0), count) + + +def check_synced_trainers(args: Any, num_islands: int) -> None: + """Launch checks of synced trainers (module doc), before the per-role checks.""" + if not is_synced_trainer(args): + return + if int(num_islands) < 2: + raise ValueError("--rl-sample-role trainer without --rl-single-island-no-sync is a synced " + f"multi-island launch: it needs at least two islands (got {num_islands})") + if is_pair(args) or is_convertible_pool(args): + raise ValueError("synced trainers do not combine with --rl-island-sample-role or " + "convertible nodes") + if getattr(args, "tuning", "lora") != "lora": + raise ValueError("synced trainers merge LoRA deltas: they need --tuning lora") + store = str(getattr(args, "rl_checkpoint_store", None) or "") + if store.startswith(("s3://", "oss://")): + raise ValueError(f"--rl-checkpoint-store {store}: synced trainers read buffer lanes from " + "a shared directory store (a Modal Volume or a shared path), not an " + "object store") + + def node_name(island: int) -> str: """The node name of island ``island`` (``--rl-node-class NAME=...``, buffer shard).""" return f"island{int(island)}" @@ -221,14 +266,17 @@ def check_launch(args: Any) -> None: if getattr(args, dest, None) is not None: raise ValueError(f"{flag} needs --rl-sample-role sampler or trainer") return - if not getattr(args, "rl_single_island_no_sync", False): + 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 " - "and a trainer are separate no-sync islands over one store") + "is a no-sync island over the store (only trainers may sync)") if not getattr(args, "rl_checkpoint_store", None): raise ValueError(f"--rl-sample-role {role} needs --rl-checkpoint-store (the sample " "buffer and the policy releases live in it)") if getattr(args, "rl_elastic", False): - raise ValueError(f"--rl-sample-role {role} does not support --rl-elastic") + # --rl-elastic is placement reconfiguration (T/R cells), not the island sync: + # synced trainers use --rl-island-scheduling elastic (or strict sync) + raise ValueError(f"--rl-sample-role {role} does not support --rl-elastic " + "(placement reconfiguration)") if role in (TRAINER, CONVERTIBLE) and not int(getattr(args, "rl_max_policy_age", 0) or 0) >= 1: raise ValueError(f"--rl-sample-role {role} trains on samples of earlier versions: " "it needs --rl-max-policy-age >= 1") @@ -269,14 +317,15 @@ def _max_backlog_rounds(args: Any, role: str) -> int | None: def resolve(args: Any, *, groups_per_round: int) -> SampleBufferConfig | None: - """The island's buffer config (island side: the store is ``--rl-resume-store``).""" + """The island's buffer config (island side: the store is ``--rl-resume-store``, + or ``--rl-sample-store`` on a synced trainer).""" role = sample_role(args) if role == BOTH: return None - store = getattr(args, "rl_resume_store", None) + store = getattr(args, "rl_sample_store", None) or getattr(args, "rl_resume_store", None) if not store: raise ValueError(f"--rl-sample-role {role} needs the checkpoint store " - "(--rl-resume-store on the island)") + "(--rl-resume-store, or --rl-sample-store on a synced trainer)") rounds = _max_backlog_rounds(args, role) wait_s = getattr(args, "rl_sample_wait_s", None) min_round_s = getattr(args, "rl_sampler_min_round_seconds", None) From af2303b8ad1cc5dcdcc8199e1bcd50984d20c050 Mon Sep 17 00:00:00 2001 From: michael <78113554+chowhsu@users.noreply.github.com> Date: Tue, 20 Oct 2026 04:17:52 +0000 Subject: [PATCH 31/89] s25_g1 driver: the re-quantizer runs in its own CPU container A Modal Volume reload discards the container's uncommitted writes; the re-quantizer's poll removed a release directory while Miles wrote it in the same container (run-12). Co-Authored-By: Claude Opus 5.5 --- tools/probes/s25_g1/modal_app.py | 33 ++++++++++++++++++++------------ 1 file changed, 21 insertions(+), 12 deletions(-) diff --git a/tools/probes/s25_g1/modal_app.py b/tools/probes/s25_g1/modal_app.py index 5d8ea4cf5..2daadf1b1 100644 --- a/tools/probes/s25_g1/modal_app.py +++ b/tools/probes/s25_g1/modal_app.py @@ -274,22 +274,15 @@ def trainer_island(tag: str, stop_after_rounds: int | None = None, convert_only: _sh("ray start --head --port=6379 --node-ip-address=127.0.0.1 --num-gpus=2 --disable-usage-stats", log / "ray.log", env=env) began = time.time() - with open(log / "requant.log", "w", encoding="utf-8") as rq, open(log / "trainer.log", "w", encoding="utf-8") as f: - requant = subprocess.Popen(plan["requant"]["loop_argv"], env=env, cwd=WORKDIR, stdout=rq, - stderr=subprocess.STDOUT) + # The re-quantizer runs in its own container (requant_island), not as a side process + # here: a Modal Volume reload in this container (its poll) discards this container's + # uncommitted writes, and it removed a release directory while Miles wrote it (run-12). + with open(log / "trainer.log", "w", encoding="utf-8") as f: # as the stage-1 node script does: the island forwards its env to the Ray actors rc = subprocess.run(island, env={**env, "RAY_ADDRESS": "127.0.0.1:6379"}, cwd=WORKDIR, stdout=f, stderr=subprocess.STDOUT).returncode - requant.terminate() - try: - requant.wait(timeout=60) - except subprocess.TimeoutExpired: - requant.kill() - final = subprocess.run(plan["requant"]["argv"], env=env, cwd=WORKDIR, stdout=rq, - stderr=subprocess.STDOUT).returncode nvml.terminate() - (log / "rc.json").write_text(json.dumps({"rc": rc, "final_requant_rc": final, - "wall_s": round(time.time() - began, 1)})) + (log / "rc.json").write_text(json.dumps({"rc": rc, "wall_s": round(time.time() - began, 1)})) shutil.copytree(f"/data/yeto/{plan['run_id']}", log / "state", dirs_exist_ok=True, ignore=shutil.ignore_patterns("*.pt", "*.safetensors")) _sh("cp -r /tmp/ray/session_latest/logs " + shlex.quote(str(log / "ray-logs")) + " || true", check=False) @@ -297,12 +290,26 @@ def trainer_island(tag: str, stop_after_rounds: int | None = None, convert_only: return rc +@app.function(image=image, volumes=VOLUMES, cpu=8, memory=65536, timeout=3 * 3600) +def requant_island(tag: str) -> int: + """The plan's re-quantizer loop (node: its own CPU container; same argv).""" + log = _log_dir(tag) + _mirror_every(tag, 60.0) + plan = _plan(log) + with open(log / "requant.log", "w", encoding="utf-8") as f: + rc = subprocess.run(plan["requant"]["loop_argv"], env=_env(plan["trainer"]["env"]), cwd=WORKDIR, + stdout=f, stderr=subprocess.STDOUT).returncode + _commit(tag) + return rc + + @app.function(image=image, volumes=VOLUMES, cpu=1, memory=1024, timeout=3 * 3600) def orchestrate(tag: str, cut_after: int = 2) -> dict: """Sampler for the whole run; trainer A stops after ``cut_after`` rounds (cut), trainer B resumes in a new container. A trainer that fails cancels the sampler; a sampler that fails cancels the running trainer (no card waits for nothing).""" sampler = sampler_island.spawn(f"{tag}/sampler") + requant = requant_island.spawn(f"{tag}/requant") sampler_rc: object = None def sampler_done() -> bool: @@ -338,6 +345,8 @@ def run_trainer(sub: str, stop: int | None) -> object: else: while not sampler_done(): time.sleep(30) + time.sleep(90) # the re-quantizer's last log mirror + requant.cancel() result = {"trainer_a": first, "trainer_b": second, "sampler": sampler_rc} print(json.dumps(result, default=str), flush=True) out = Path(f"/data/store/logs/{tag}") From bf94fa7c0290d0a32afb50f91704a97bc9bd994a Mon Sep 17 00:00:00 2001 From: michael <78113554+chowhsu@users.noreply.github.com> Date: Tue, 20 Oct 2026 04:18:24 +0000 Subject: [PATCH 32/89] Leave at once on the reclaim marker from a watcher thread The cloud stops the container seconds after its signal, while the sampler loop may still be inside a round or its throttle sleep. A daemon watcher now writes rl_sampler_leave (reason PREEMPTED, rounds done, the rollout id in flight) and the member record as soon as the marker appears; the exit path leaves only when the watcher has not. Co-Authored-By: Claude Opus 5.5 --- tests/test_rl_spot_sampler_replace.py | 26 +++++++++ yeto/rl/adapters/miles/sampler_island.py | 51 ++++++++--------- yeto/rl/engine/sampler_resume.py | 73 +++++++++++++++++++++++- 3 files changed, 122 insertions(+), 28 deletions(-) diff --git a/tests/test_rl_spot_sampler_replace.py b/tests/test_rl_spot_sampler_replace.py index 0ae3c3125..7e8d2c204 100644 --- a/tests/test_rl_spot_sampler_replace.py +++ b/tests/test_rl_spot_sampler_replace.py @@ -382,3 +382,29 @@ def signal(num, fn): env = {RECLAIM_MARKER_ENV: "/tmp/yeto-reclaim-requested"} assert mr.run_island_script(["x"], env, popen=lambda c, env, **k: None, signal_mod=Sig) == 0 assert seen == [env] + + +def test_the_reclaim_watcher_leaves_once_while_the_loop_is_still_in_a_round(tmp_path): + import threading + + from yeto.rl.engine.sampler_resume import ReclaimWatcher, SamplerMembership + + store = CheckpointStore(tmp_path) + events = [] + membership = SamplerMembership(store, "island0", "b", resume_point([], "island0"), + lambda e, **f: events.append((e, f)), clock=lambda: 7.0) + membership.join(latest_release={"seq": 3, "policy_version": 3}) + flag = threading.Event() + watcher = ReclaimWatcher(flag.is_set, lambda: membership.leave("PREEMPTED", rounds=2, + in_flight_rollout_id=4), + poll_s=0.01).start() + flag.set() + assert watcher.fired.wait(2.0) + watcher.stop() + assert membership.leave("FINISHED") is False # the exit path after the watcher: no second LEAVE + names = [e for e, _ in events] + assert names == ["rl_sampler_join", "rl_sampler_leave"] + leave = events[1][1] + assert (leave["reason"], leave["rounds"], leave["in_flight_rollout_id"]) == ("PREEMPTED", 2, 4) + record = (tmp_path / "sample-buffer/members/island0/b.json").read_text() + assert '"leave_reason": "PREEMPTED"' in record and '"joined_at": 7.0' in record diff --git a/yeto/rl/adapters/miles/sampler_island.py b/yeto/rl/adapters/miles/sampler_island.py index 4289c91b4..3fd9b6092 100644 --- a/yeto/rl/adapters/miles/sampler_island.py +++ b/yeto/rl/adapters/miles/sampler_island.py @@ -261,14 +261,12 @@ def run_sampler_island(miles_args: Any, *, environ: Any = None, from yeto.rl.engine.resume import store_for from yeto.rl.engine.sampler_loop import SamplerLoop from yeto.rl.engine.sampler_resume import ( - JOIN_EVENT, - LEAVE_EVENT, - Member, + ReclaimWatcher, + SamplerMembership, new_incarnation, reclaim_marker, reclaim_requested, resume_point, - write_member, ) from . import LoopRunner @@ -294,9 +292,9 @@ def emit(event: str, **fields: Any) -> None: reclaim_marker(env).unlink(missing_ok=True) # a restarted machine keeps its /tmp buffer = open_buffer(env, sampler=True) resume = resume_point(buffer.entries(), sampler_id) - store = store_for(config.store, environ=environ) - member = Member(sampler=sampler_id, incarnation=incarnation, joined_at=time.time(), - resume=resume.fields()) + membership = SamplerMembership(store_for(config.store, environ=environ), sampler_id, + incarnation, resume, emit) + progress = {"rollout_id": None, "rounds": 0} # the round in flight, rounds done here connect_island_ray(miles_args=miles_args) # The rollout process's metadata hook pushes each round's group records to this @@ -316,20 +314,22 @@ async def init() -> tuple[Any, Any]: return controller, executor async def generate(rollout_id: int) -> None: + progress["rollout_id"] = rollout_id await controller.prepare_rollout(rollout_id) _release(await executor.get(rollout_id)) # the rollout function wrote the groups metadata.discard() # the groups are in the buffer; nobody here reads the records + progress["rounds"] += 1 + progress["rollout_id"] = None - def leave(reason: LeaveReason, rounds: int, **detail: Any) -> None: - member.left_at, member.leave_reason, member.rounds = time.time(), reason.name, rounds - member.policy_version = follower.version - member.detail = detail - emit(LEAVE_EVENT, sampler=sampler_id, incarnation=incarnation, reason=reason.name, - rounds=rounds, policy_version=follower.version, buffer_end=len(buffer), **detail) - write_member(store, member) + def leave(reason: LeaveReason, **detail: Any) -> None: + membership.leave(reason.name, rounds=progress["rounds"], + policy_version=None if follower is None else follower.version, + in_flight_rollout_id=progress["rollout_id"], **detail) error: BaseException | None = None follower = None + reclaimed = lambda: reclaim_requested(env) # noqa: E731 + watcher = ReclaimWatcher(reclaimed, lambda: leave(LeaveReason.PREEMPTED)).start() try: controller, executor = runner.run(init()) lora = lora_mode(miles_args) @@ -338,7 +338,6 @@ def leave(reason: LeaveReason, rounds: int, **detail: Any) -> None: loader.stamp(BASE_POLICY_VERSION) # until the first release is loaded follower = PolicyFollower(release_store(config, environ), loader, version=BASE_POLICY_VERSION) - reclaimed = lambda: reclaim_requested(env) # noqa: E731 loop = SamplerLoop( follower=follower, buffer=buffer, emit=emit, generate=lambda rollout_id: runner.run(generate(rollout_id)), @@ -348,10 +347,8 @@ def leave(reason: LeaveReason, rounds: int, **detail: Any) -> None: loaded_fields=lambda: {"served_checksum": loader.last_checksum, "probe": loader.last_probe}) latest = follower.releases.latest() - emit(JOIN_EVENT, sampler=sampler_id, incarnation=incarnation, - latest_release=None if latest is None else {"seq": latest.seq, "policy_version": latest.version}, - **resume.fields()) - write_member(store, member) + membership.join(latest_release=None if latest is None else + {"seq": latest.seq, "policy_version": latest.version}) emit("rl_sampler_start", store=config.store, num_rollouts=int(miles_args.num_rollout), max_backlog_groups=int(config.max_backlog_groups), weights=loader.weights, start_rollout_id=resume.next_rollout_id) @@ -360,21 +357,21 @@ def leave(reason: LeaveReason, rounds: int, **detail: Any) -> None: else: rounds = loop.run(start=resume.next_rollout_id) if reclaimed(): - leave(LeaveReason.PREEMPTED, rounds) + leave(LeaveReason.PREEMPTED) # no-op when the watcher already left _hold_until_stopped(sleep, LEAVE_HOLD_S) else: - leave(LeaveReason.FINISHED, rounds) + leave(LeaveReason.FINISHED) return rounds except BaseException as exc: error = exc - if follower is not None: - try: - leave(LeaveReason.PREEMPTED if reclaim_requested(env) else LeaveReason.INFRA_FAILED, - -1, error=f"{type(exc).__name__}: {exc}"[:300]) - except Exception: # noqa: BLE001 - the original error is the one to report - pass + try: + leave(LeaveReason.PREEMPTED if reclaimed() else LeaveReason.INFRA_FAILED, + error=f"{type(exc).__name__}: {exc}"[:300]) + except Exception: # noqa: BLE001 - the original error is the one to report + pass raise finally: + watcher.stop() try: runner.run(disposer.__aexit__(None if error is None else type(error), error, None if error is None else error.__traceback__)) diff --git a/yeto/rl/engine/sampler_resume.py b/yeto/rl/engine/sampler_resume.py index e25eefda9..203671aed 100644 --- a/yeto/rl/engine/sampler_resume.py +++ b/yeto/rl/engine/sampler_resume.py @@ -32,8 +32,10 @@ from __future__ import annotations import json +import threading +import time import uuid -from collections.abc import Iterable, Mapping +from collections.abc import Callable, Iterable, Mapping from dataclasses import asdict, dataclass, field from pathlib import Path from typing import Any @@ -183,3 +185,72 @@ def write_member(store: Any, member: Member) -> Path: tmp.replace(path) store.commit() return path + + +class SamplerMembership: + """One incarnation's JOIN and LEAVE: the tape events and the member record. + :meth:`leave` acts once (the reclaim watcher and the island's exit path may + both call it); later calls return False.""" + + def __init__(self, store: Any, sampler: str, incarnation: str, resume: SamplerResume, + emit: Callable[..., None], clock: Callable[[], float] = time.time) -> None: + self.store = store + self.emit = emit + self.clock = clock + self.member = Member(sampler=sampler, incarnation=incarnation, joined_at=clock(), + resume=resume.fields()) + self._lock = threading.Lock() + self.left = False + + def join(self, **fields: Any) -> None: + m = self.member + self.emit(JOIN_EVENT, sampler=m.sampler, incarnation=m.incarnation, **m.resume, **fields) + write_member(self.store, m) + + def leave(self, reason: str, *, rounds: int | None = None, policy_version: int | None = None, + **detail: Any) -> bool: + with self._lock: + if self.left: + return False + self.left = True + m = self.member + m.left_at, m.leave_reason, m.rounds = self.clock(), str(reason), rounds + m.policy_version, m.detail = policy_version, dict(detail) + self.emit(LEAVE_EVENT, sampler=m.sampler, incarnation=m.incarnation, reason=str(reason), + rounds=rounds, policy_version=policy_version, **detail) + write_member(self.store, m) + return True + + +class ReclaimWatcher: + """Polls ``reclaimed()`` on a daemon thread and calls ``on_reclaim()`` once, at + once: the cloud stops the container a few seconds after its signal (Modal: + 30 s), while the island's loop may still be inside a round.""" + + def __init__(self, reclaimed: Callable[[], bool], on_reclaim: Callable[[], Any], *, + poll_s: float = 0.5) -> None: + self.reclaimed = reclaimed + self.on_reclaim = on_reclaim + self.poll_s = float(poll_s) + self._stop = threading.Event() + self.fired = threading.Event() + self._thread = threading.Thread(target=self._run, name="sampler-reclaim-watcher", daemon=True) + + def start(self) -> "ReclaimWatcher": + self._thread.start() + return self + + def stop(self) -> None: + self._stop.set() + self._thread.join(timeout=5.0) + + def _run(self) -> None: + while not self._stop.is_set(): + if self.reclaimed(): + self.fired.set() + try: + self.on_reclaim() + except Exception as error: # noqa: BLE001 - the island loop still stops on the marker + print(f"[sampler-island] reclaim handler failed: {error!r}", flush=True) + return + self._stop.wait(self.poll_s) From 7c21f7f2110db92e124fa24c000a29fd7af2029c Mon Sep 17 00:00:00 2001 From: michael <78113554+chowhsu@users.noreply.github.com> Date: Tue, 20 Oct 2026 04:21:53 +0000 Subject: [PATCH 33/89] T-C4 GPU probes: verl SCHEDULED cut + resume chain, Miles stale-master injection phase; record the round-cut wiring Co-Authored-By: Claude Opus 5.5 --- tools/probes/c4_verl/modal_c4.py | 137 ++++++++++++++++++++++++++++ tools/probes/g3_miles/harness.py | 30 ++++++ tools/probes/g3_miles/modal_g3.py | 7 ++ yeto/rl/adapters/verl/round_cuts.py | 8 +- 4 files changed, 180 insertions(+), 2 deletions(-) create mode 100644 tools/probes/c4_verl/modal_c4.py diff --git a/tools/probes/c4_verl/modal_c4.py b/tools/probes/c4_verl/modal_c4.py new file mode 100644 index 000000000..dfe9d5e6c --- /dev/null +++ b/tools/probes/c4_verl/modal_c4.py @@ -0,0 +1,137 @@ +"""T-C4 (S25 N7.6b) GPU check: the verl island's round cuts on a real verl island. + +Three island runs of ``yeto.rl.adapters.verl.island_entry`` (Qwen3-0.6B, LoRA, +GSM8K, ``--rl-single-island-no-sync``) in the pinned verl image, one GPU each: + +* ``a`` -- commanded by a pre-written trainer-control file (``run 2``, then + ``stop STOPPED``): trains rounds 0-1, the stop cuts with the SCHEDULED trigger + into the store (Modal Volume), the island exits; +* ``b`` -- a fresh container: restores the store into a new state dir, resumes + at rollout 2 (data cursor seek, digest + Adam/scheduler checks), trains 2-3; +* ``c`` -- the continuous control: rounds 0-3, no store (runs beside ``a``). + +usage (after the gpu-spend.md pre-registration): + C4_RUN= [C4_GPU=A10G] modal run --detach modal_c4.py +The entrypoint only spawns the chain (a killed detached client would cancel a +blocking ``.remote()``). Evidence: Volume ``t-c4-s25`` under ``//``. +""" + +from __future__ import annotations + +import json +import os +import shlex +import sys + +import modal + +RUN = os.environ.get("C4_RUN", "unset") +YETO = os.environ.get("C4_YETO", "/home/michael/work/s25-c4-cut-ctl") +GPU = os.environ.get("C4_GPU", "A10G") +EXPECT = {"A10G": ("A10G", "nvidia-a10g"), "L4": ("L4", "nvidia-l4"), "H100!": ("H100", "nvidia-h100")} +sys.path.insert(0, YETO) +from yeto.rl.adapters.verl.image import modal_image # noqa: E402 + +vol = modal.Volume.from_name("t-c4-s25", create_if_missing=True) +image = (modal_image(modal) + .env({"PYTHONPATH": "/yeto", "HF_HUB_DISABLE_TELEMETRY": "1", "PYTHONUNBUFFERED": "1"}) + .add_local_dir(f"{YETO}/yeto", "/yeto/yeto", copy=False, ignore=["**/__pycache__"])) +app = modal.App(f"t-c4-{RUN}", image=image) + +ROUNDS = 4 +PY = "/workspace/verl/.venv/bin/python" # the pinned verl venv (image.VENV_BIN) +ISLAND_ARGS = [ + "--model", "Qwen/Qwen3-0.6B", "--model-revision", "c1899de289a04d12100db370d81485cdf75e47ca", + "--data", "zhuzilin/gsm8k", "--data-revision", "0cbd9f31d91ac21a7613dcbc7fef992adac459ae", + "--rl-single-island-no-sync", "--learner-id", "0", "--reward-function", "gsm8k_reward:score", + "--global-rounds", str(ROUNDS), "--groups-per-round", "8", "--samples-per-group", "4", + "--rollout-max-response-len", "256", "--event-tape", "/root/yeto-output/rl-island-0.jsonl", + "--lora-r", "32", "--lora-targets", "all-linear", "--inner-lr", "1e-5", "--seq-len", "1024", + "--seed", "17", "--apply-chat-template-kwargs", '{"enable_thinking": false}', +] + + +def _sh(cmd: str, log) -> int: + import subprocess + + log.write(f"$ {cmd}\n") + log.flush() + return subprocess.run(cmd, shell=True, stdout=log, stderr=subprocess.STDOUT).returncode + + +@app.function(gpu=GPU, cpu=8.0, memory=49152, timeout=40 * 60, volumes={"/vol": vol}) +def island(run: str, phase: str, gpu: str) -> dict: + import shutil + import subprocess + import time + + want_name, group = EXPECT[gpu] + out = f"/vol/{run}/{phase}" + os.makedirs(out, exist_ok=True) + log = open(f"{out}/container.log", "a") + name = subprocess.run(["nvidia-smi", "--query-gpu=name", "--format=csv,noheader"], + capture_output=True, text=True).stdout.strip() + _sh(f"nvidia-smi > {out}/nvidia-smi.txt", log) + if want_name not in name: + log.write(f"GPU ASSERT FAILED: {name!r} is not {want_name}\n") + log.close() + vol.commit() + return {"phase": phase, "rc": 3, "gpu": name} + env = dict(os.environ, YETO_RL_COMPAT_GROUP=group, RAY_ADDRESS="127.0.0.1:6379") + store = f"/vol/{run}/store" + args = list(ISLAND_ARGS) + if phase in ("a", "b"): + args += ["--rl-resume-store", store, "--rl-resume-state-dir", f"/root/state-{phase}"] + if phase == "a": # the head's commands for incarnation 0, written through yeto's own writer + rc = _sh(f"{PY} -c " + shlex.quote( + "import sys; from yeto.rl.engine import trainer_remote as t; " + "from yeto.rl.engine.resume import store_for; from yeto.rl.elastic_client import LeaveReason; " + "w = t.CommandWriter(t.LocalStoreFiles(store_for(sys.argv[1])), 0); " + "w.send(t.RUN, rounds=2); w.send(t.STOP, reason=LeaveReason.STOPPED); " + "print(t.island_env(0, None, poll_s=2.0))") + f" {store}", log) + if rc != 0: + log.close() + vol.commit() + return {"phase": phase, "rc": 4, "gpu": name, "error": "command file not written"} + env.update({"YETO_TRAINER_CONTROL": "1", "YETO_TRAINER_INCARNATION": "0", + "YETO_TRAINER_RESUME": "", "YETO_TRAINER_COMMAND_POLL_S": "2.0"}) + _sh("ray start --head --node-ip-address=127.0.0.1 --port=6379 --num-cpus=32 --num-gpus=1 " + "--include-dashboard=false --temp-dir=/root/miles-ray", log) + t0 = time.time() + with open(f"{out}/island.log", "w") as ilog: + rc = subprocess.run([PY, "-m", "yeto.rl.adapters.verl.island_entry", *args], env=env, + stdout=ilog, stderr=subprocess.STDOUT, cwd="/root").returncode + seconds = time.time() - t0 + _sh("ray stop --force", log) + if os.path.isdir("/root/yeto-output"): + shutil.copytree("/root/yeto-output", f"{out}/yeto-output", dirs_exist_ok=True) + if os.path.isdir(f"/root/state-{phase}"): + shutil.copytree(f"/root/state-{phase}", f"{out}/state-dir", dirs_exist_ok=True) + result = {"phase": phase, "rc": rc, "gpu": name, "island_seconds": round(seconds, 1), "args": args} + with open(f"{out}/result.json", "w") as fh: + json.dump(result, fh, indent=1) + log.close() + vol.commit() + return result + + +@app.function(timeout=90 * 60, volumes={"/vol": vol}) +def chain(run: str, gpu: str) -> dict: + control = island.spawn(run, "c", gpu) + results = {"a": island.remote(run, "a", gpu)} + vol.reload() + if results["a"]["rc"] == 0: + results["b"] = island.remote(run, "b", gpu) + results["c"] = control.get() + vol.reload() + with open(f"/vol/{run}/chain.json", "w") as fh: + json.dump(results, fh, indent=1) + vol.commit() + return results + + +@app.local_entrypoint() +def main(): + assert RUN != "unset", "set C4_RUN" + call = chain.spawn(RUN, GPU) + print("spawned", RUN, call.object_id, flush=True) diff --git a/tools/probes/g3_miles/harness.py b/tools/probes/g3_miles/harness.py index c2bd4c94d..84cf8cbfa 100644 --- a/tools/probes/g3_miles/harness.py +++ b/tools/probes/g3_miles/harness.py @@ -616,6 +616,28 @@ def _hdo_report(optimizer) -> dict: return {"leaves": rep} +def stale_report(optimizer) -> dict: + """Per HDO parameter: does the tensor the sub-optimizer steps equal the HDO's "master_param"? + (the old export read the master; N7.6c reads the stepped tensor)""" + from miles.backends.megatron_utils.lora import dp_invariant_state as dps + + stale, same = [], 0 + for leaf in getattr(optimizer, "chained_optimizers", [optimizer]): + hdo = getattr(leaf, "optimizer", None) + if not hasattr(hdo, "sub_optimizers"): + continue + for param, state in hdo.state.items(): + master = state.get("master_param") + if not torch.is_tensor(master): + continue + inner = dps._hdo_inner_param(hdo, param) + if torch.equal(inner.detach().cpu().reshape(-1), master.detach().cpu().reshape(-1)): + same += 1 + else: + stale.append({"shape": list(param.shape), "dtype": str(param.dtype)}) + return {"stepped_equals_master": same, "stepped_differs_from_master": len(stale), "differs": stale[:40]} + + def param_digest(model) -> dict: h = hashlib.sha256() norms = {} @@ -704,9 +726,17 @@ def manifest(summaries, global_step): elif ns.mode.startswith("resume2"): files = json.loads((cut_dir / "files.json").read_text()) saved = ResumeManifest.from_dict(json.loads((cut_dir / "resume_manifest.json").read_text())) + inject = ns.mode.endswith("inj") + if inject: # T-C4 N7.6c: put the T-G3 bug back (the restore leaves the HDO stepping its stale host copy) + from miles.backends.megatron_utils.lora import dp_invariant_state as dps + + dps._bind_masters_to_stepped_tensors = lambda hdo: None summary = cut_plugin.restore_cut_shard(actor, directory=str(cut_dir), files=files, cut_id="g3-cut-3") if cut_plugin.REFUSED in summary: raise RuntimeError(f"restore_cut refused: {summary}") + if inject: + res["injected"] = "dp_invariant_state._bind_masters_to_stepped_tensors = no-op" + res["stale_vs_master"] = stale_report(optimizer) res["restore_summary"] = {k: summary.get(k) for k in ("path", "scheduler_samples", "resume", "diff", "optimizer_cut_vs_reexport")} res["restore_digest_equal"] = summary.get("state_digest") == next( diff --git a/tools/probes/g3_miles/modal_g3.py b/tools/probes/g3_miles/modal_g3.py index c8670e776..eccadf663 100644 --- a/tools/probes/g3_miles/modal_g3.py +++ b/tools/probes/g3_miles/modal_g3.py @@ -67,6 +67,13 @@ run n71d-run3b 600 $TR $H --phase n71d --mode run3b --out "$OUT/n71d" --cut-dir "$OUT/n71d/cutb" run n71d-resume2b 600 $TR $H --phase n71d --mode resume2b --out "$OUT/n71d" --cut-dir "$OUT/n71d/cutb" fi +if [[ "$PHASES" == c4 ]]; then # T-C4 N7.6c: cut, clean resume, resume with the T-G3 bug injected + export NCCL_ALGO=Ring CUBLAS_WORKSPACE_CONFIG=:4096:8 NVIDIA_TF32_OVERRIDE=0 + run n71d-run3 900 $TR $H --phase n71d --mode run3 --out "$OUT/n71d" --cut-dir "$OUT/n71d/cut" + run n71d-resume2 900 $TR $H --phase n71d --mode resume2 --out "$OUT/n71d" --cut-dir "$OUT/n71d/cut" + run n71d-resume2inj 900 $TR $H --phase n71d --mode resume2inj --out "$OUT/n71d" --cut-dir "$OUT/n71d/cut" + du -sh "$OUT"; echo CHAIN DONE; exit 0 +fi if [[ "$PHASES" == all || "$PHASES" == n75 ]]; then run n75-te 1200 $TR $H --phase n75 --impl te --tag te --out "$OUT/n75" run n75-te2 900 $TR $H --phase n75 --impl te --tag te2 --weights-from "$OUT/n75" --out "$OUT/n75" diff --git a/yeto/rl/adapters/verl/round_cuts.py b/yeto/rl/adapters/verl/round_cuts.py index ecaffbd9e..0e5abe309 100644 --- a/yeto/rl/adapters/verl/round_cuts.py +++ b/yeto/rl/adapters/verl/round_cuts.py @@ -135,8 +135,12 @@ def wire_round_cuts(driver: Any, resume: ResumeWiring | None, *, plan: Mapping[s if checkpoint is None: raise ValueError("verl round cuts need the driver's sync to be LocalOnlySync") driver.sync.stop_after = resume.config.get("stop_after") - if install_signals: - install_preempt_handler(controller, emit=driver.emit) + # the handler needs the main thread (verl: the task runner actor's); [] when it is not + signals = install_preempt_handler(controller, emit=driver.emit) if install_signals else [] + driver.emit("rl_round_cuts_wired", backend="verl", every=checkpoint.every, + keep=int(controller.keep), stop_after=driver.sync.stop_after, + store=str(controller.checkpoint_store), preempt_signals=list(signals), + restored_store=controller.restored is not None) return checkpoint From 07ec4ca3704a0a99b848bbbc5d5dd448623864ca Mon Sep 17 00:00:00 2001 From: michael <78113554+chowhsu@users.noreply.github.com> Date: Tue, 20 Oct 2026 04:24:36 +0000 Subject: [PATCH 34/89] Driver accepts a publication without members on an engineless pool A trainer-only Miles island (--debug-train-only) has no engine, so its rollout pool reports engineless and no members; the driver no longer calls that a partial publication (Modal run-13). Every other pool still needs at least one member and full acknowledgement. Co-Authored-By: Claude Opus 5.5 --- yeto/rl/adapters/miles/rollout.py | 9 +++++++++ yeto/rl/engine/driver.py | 8 ++++++-- 2 files changed, 15 insertions(+), 2 deletions(-) diff --git a/yeto/rl/adapters/miles/rollout.py b/yeto/rl/adapters/miles/rollout.py index f5490f47f..15947270e 100644 --- a/yeto/rl/adapters/miles/rollout.py +++ b/yeto/rl/adapters/miles/rollout.py @@ -1068,7 +1068,16 @@ def _offload_after_rollout(self) -> None: def abort(self) -> None: self._run(self._controller.abort_all()) + @property + def engineless(self) -> bool: + """True on a trainer-only island (Miles --debug-train-only): samples come + from the store and no engine serves a policy, so a publication has no + members to acknowledge it (the driver accepts an empty member set).""" + return bool(getattr(self._args, "debug_train_only", False)) + def members(self) -> frozenset[str]: + if self.engineless: + return frozenset() return self._run(running_members(self._controller)) # -- E1 membership verbs (3.4) --------------------------------------------- diff --git a/yeto/rl/engine/driver.py b/yeto/rl/engine/driver.py index c4ecee2c4..6b020e83f 100644 --- a/yeto/rl/engine/driver.py +++ b/yeto/rl/engine/driver.py @@ -718,7 +718,10 @@ def publish(self, state: TrainableState, *, rollout_id: int) -> None: f"v{manifest.target_policy_version} {manifest.target_policy_hash}" ) members = frozenset(self.rollout.members()) - if not members or result.members != members: + # an engineless pool (trainer-only island) has no member to acknowledge; + # every other pool must have at least one and all must acknowledge + engineless = bool(getattr(self.rollout, "engineless", False)) + if (not members and not engineless) or result.members != members: missing = sorted(members - result.members) raise PublicationError( f"partial publication of policy v{rollout_id}; unacknowledged " @@ -1694,7 +1697,8 @@ def rebuild_trainer(self, rebuild: Callable[[], Any], *, cut_policy_hash: str) - or manifest.target_policy_hash != cut_policy_hash): raise PublicationError("re-publication after rebuild acknowledged another policy") members = frozenset(self.rollout.members()) - if not members or result_pub.members != members: + engineless = bool(getattr(self.rollout, "engineless", False)) + if (not members and not engineless) or result_pub.members != members: raise PublicationError( f"partial re-publication after rebuild; missing {sorted(members - result_pub.members)}" ) From 9e0ec8270c939598cbef3a1bdd9a01634d6f2fd3 Mon Sep 17 00:00:00 2001 From: michael <78113554+chowhsu@users.noreply.github.com> Date: Tue, 20 Oct 2026 04:25:08 +0000 Subject: [PATCH 35/89] Make the in-flight export check a CutSource class attribute; fix the A10G name assert Co-Authored-By: Claude Opus 5.5 --- tools/probes/c4_verl/modal_c4.py | 3 ++- yeto/rl/adapters/miles/rebuild_wiring.py | 6 +++--- yeto/rl/engine/round_cut.py | 9 +++++---- 3 files changed, 10 insertions(+), 8 deletions(-) diff --git a/tools/probes/c4_verl/modal_c4.py b/tools/probes/c4_verl/modal_c4.py index dfe9d5e6c..39a337a19 100644 --- a/tools/probes/c4_verl/modal_c4.py +++ b/tools/probes/c4_verl/modal_c4.py @@ -28,7 +28,8 @@ RUN = os.environ.get("C4_RUN", "unset") YETO = os.environ.get("C4_YETO", "/home/michael/work/s25-c4-cut-ctl") GPU = os.environ.get("C4_GPU", "A10G") -EXPECT = {"A10G": ("A10G", "nvidia-a10g"), "L4": ("L4", "nvidia-l4"), "H100!": ("H100", "nvidia-h100")} +# Modal A10G reports "NVIDIA A10" +EXPECT = {"A10G": ("NVIDIA A10", "nvidia-a10g"), "L4": ("L4", "nvidia-l4"), "H100!": ("H100", "nvidia-h100")} sys.path.insert(0, YETO) from yeto.rl.adapters.verl.image import modal_image # noqa: E402 diff --git a/yeto/rl/adapters/miles/rebuild_wiring.py b/yeto/rl/adapters/miles/rebuild_wiring.py index 36f5e60c1..0f7297547 100644 --- a/yeto/rl/adapters/miles/rebuild_wiring.py +++ b/yeto/rl/adapters/miles/rebuild_wiring.py @@ -35,11 +35,11 @@ class CutSource(_CoreCutSource): """:class:`yeto.rl.engine.round_cut.CutSource` with the Miles in-flight export check.""" - def __init__(self, **kwargs: Any) -> None: + @staticmethod + def check_in_flight(exported: Mapping[str, Any]) -> list[str]: from .in_flight import check_export - kwargs.setdefault("check_in_flight", check_export) - super().__init__(**kwargs) + return check_export(exported) def make_trainer_rebuilder( diff --git a/yeto/rl/engine/round_cut.py b/yeto/rl/engine/round_cut.py index 32ef0fb18..cadc5047e 100644 --- a/yeto/rl/engine/round_cut.py +++ b/yeto/rl/engine/round_cut.py @@ -69,11 +69,14 @@ class CutSource: restore must match (``RestoreExpectation``); shared by the 4.4 same-shape rebuilder and E3's ``MilesTrainerOps`` (4.7).""" + # backend check of an in-flight export: problems refuse the cut (None: no check; + # the Miles adapter overrides it with its export check) + check_in_flight: Callable[[Mapping[str, Any]], list[str]] | None = None + def __init__(self, *, driver: Any, trainer: Any, rollout: Any, ledger: Any, algorithm: Any, backend_fingerprint: str, cut_root: str, global_batch_size: int, ref_model: Mapping[str, Any] | None = None, - shared_filesystem: bool = True, - check_in_flight: Callable[[Mapping[str, Any]], list[str]] | None = None) -> None: + shared_filesystem: bool = True) -> None: self.driver_ref = driver # the driver, or a zero-arg callable returning it self.trainer = trainer self.rollout = rollout @@ -83,8 +86,6 @@ def __init__(self, *, driver: Any, trainer: Any, rollout: Any, ledger: Any, algo self.cut_root = cut_root self.global_batch_size = int(global_batch_size) self.shared_filesystem = shared_filesystem - # backend check of an in-flight export (problems refuse the cut); None: no check - self.check_in_flight = check_in_flight @property def driver(self) -> Any: From 6361010a40ee39a305252d5363abe8c354732b0d Mon Sep 17 00:00:00 2001 From: michael <78113554+chowhsu@users.noreply.github.com> Date: Tue, 20 Oct 2026 04:31:14 +0000 Subject: [PATCH 36/89] Re-quantizer accepts the text-model prefix of Miles' raw glm5_next export Miles' raw export names the text model model.layers.*, the HF composite layout model.language_model.layers.* (Modal run-14: 398 tensors missing and 398 extra). A missing layout tensor is read from the source under the text prefix only when that exact name exists. Driver: requant-only entrypoint. Co-Authored-By: Claude Opus 5.5 --- tests/test_rl_quant_mxfp4.py | 11 +++++++++++ tools/probes/s25_g1/modal_app.py | 7 +++++++ yeto/rl/quant/requant.py | 14 ++++++++++++-- 3 files changed, 30 insertions(+), 2 deletions(-) diff --git a/tests/test_rl_quant_mxfp4.py b/tests/test_rl_quant_mxfp4.py index b2a896313..d7b321505 100644 --- a/tests/test_rl_quant_mxfp4.py +++ b/tests/test_rl_quant_mxfp4.py @@ -361,3 +361,14 @@ def test_frozen_tensors_come_from_the_base_only_when_named(world): frozen=(world["source"], r"^lm_head\.")) assert stats["frozen_from_base"] == 1 assert compare_checkpoints(out, world["original"])["unequal_tensors"] == 0 + + +def test_text_model_prefix_of_a_miles_raw_export_is_accepted(world): + # Miles' raw glm5_next export writes "model.layers..."; the layout has "model.language_model.layers..." + raw = world["tmp"] / "raw-export" + _write_ckpt(raw, {n.replace("model.language_model.", "model.", 1): t for n, t in world["src"].items()}, + lambda n: "model-00001-of-00001.safetensors") + out = world["tmp"] / "out-raw" + stats = requantize_checkpoint(raw, world["template"], out, AwqSmoothing.none()) + assert stats["renamed_from_text_prefix"] == sum(n.startswith("model.language_model.") for n in world["src"]) + assert compare_checkpoints(out, world["original"])["unequal_tensors"] == 0 diff --git a/tools/probes/s25_g1/modal_app.py b/tools/probes/s25_g1/modal_app.py index 2daadf1b1..50c2f8021 100644 --- a/tools/probes/s25_g1/modal_app.py +++ b/tools/probes/s25_g1/modal_app.py @@ -356,6 +356,13 @@ def run_trainer(sub: str, stop: int | None) -> object: return result +@app.local_entrypoint() +def requant(tag: str = "requant-1") -> None: + """Start (or restart) only the re-quantizer loop of a running rehearsal.""" + call = requant_island.spawn(tag) + print(f"spawned requant {call.object_id}") + + @app.local_entrypoint() def convert(tag: str = "convert-1") -> None: call = trainer_island.spawn(tag, None, True) diff --git a/yeto/rl/quant/requant.py b/yeto/rl/quant/requant.py index 053f51737..31b61582d 100644 --- a/yeto/rl/quant/requant.py +++ b/yeto/rl/quant/requant.py @@ -382,6 +382,9 @@ def close(self) -> None: self._open.clear() +_COMPOSITE_PREFIX, _TEXT_PREFIX = "model.language_model.", "model." + + def requantize_checkpoint( source_dir: Path, template_dir: Path, @@ -410,6 +413,12 @@ def requantize_checkpoint( source = _SourceReader(source_dir) expected = layout.source_names() missing, extra = expected - set(source.weight_map), set(source.weight_map) - expected + # Miles' raw glm5_next export names the text model "model.layers..."; the HF + # composite checkpoint (and so the layout) has "model.language_model.layers..." + # (Modal run-14). Only an exact one-to-one prefix match is taken as the same tensor. + alias = {n: n.replace(_COMPOSITE_PREFIX, _TEXT_PREFIX, 1) for n in missing + if n.startswith(_COMPOSITE_PREFIX) and n.replace(_COMPOSITE_PREFIX, _TEXT_PREFIX, 1) in extra} + missing, extra = missing - set(alias), extra - set(alias.values()) base, filled = None, set() if frozen is not None and missing: matcher = re.compile(frozen[1]) @@ -431,10 +440,11 @@ def requantize_checkpoint( quantized = set(layout.quantized_weights()) out_dir = Path(out_dir) out_dir.mkdir(parents=True, exist_ok=True) - stats = {"shards": 0, "tensors": 0, "quantized": 0, "frozen_from_base": len(filled)} + stats = {"shards": 0, "tensors": 0, "quantized": 0, "frozen_from_base": len(filled), + "renamed_from_text_prefix": len(alias)} def read(name: str) -> torch.Tensor: - return base.get(name) if name in filled else source.get(name) + return base.get(name) if name in filled else source.get(alias.get(name, name)) try: for shard in selected: From 58c2e3246cd2670ca0f6c529844706c6a85e8bb5 Mon Sep 17 00:00:00 2001 From: michael <78113554+chowhsu@users.noreply.github.com> Date: Tue, 20 Oct 2026 04:37:44 +0000 Subject: [PATCH 37/89] verl cut: load the trainable state non-strictly (frozen LoRA base weights are not in the cut) GPU T-C4: a relaunched verl island (Qwen3-0.6B LoRA) refused its own cut with 'Missing key(s) in state_dict: ...base_layer.weight': the cut skips frozen parameters, and set_state_dict loads strictly by default. restore_rank already compares the trainable names with the cut before loading, so the load itself is not strict. The probe chain no longer reruns a finished phase after a Modal preemption restart. Co-Authored-By: Claude Opus 5.5 --- tests/test_rl_verl_cut.py | 32 ++++++++++++++++++++++++++++++++ tools/probes/c4_verl/modal_c4.py | 5 +++++ yeto/rl/adapters/verl/cut.py | 9 ++++++--- 3 files changed, 43 insertions(+), 3 deletions(-) diff --git a/tests/test_rl_verl_cut.py b/tests/test_rl_verl_cut.py index 5b15d0641..390d43519 100644 --- a/tests/test_rl_verl_cut.py +++ b/tests/test_rl_verl_cut.py @@ -310,3 +310,35 @@ def test_digest_reads_the_tensors_the_optimizer_steps(): with torch.no_grad(): first.add_(1.0) assert vcut._digest(h) != before + + +def test_frozen_base_weights_are_not_in_the_cut_and_do_not_block_the_restore(tmp_path): + """LoRA shape: frozen base weights are skipped at save; a strict load would report + them missing (GPU T-C4 a, 2026-10-20).""" + + def frozen_model(): + model = _model() + model[0].weight.requires_grad_(False) + return model + + def handles(model): + opt = torch.optim.AdamW([p for p in model.parameters() if p.requires_grad], lr=1e-2) + return TrainHandles(model, opt, torch.optim.lr_scheduler.LambdaLR(opt, lambda s: 1.0 / (1 + s))) + + trained = handles(frozen_model()) + for r in range(CUT): + _round(trained, r) + worker = SimpleNamespace(module=trained.module, optimizer=trained.optimizer, lr_scheduler=trained.scheduler) + vcut.VerlCutTrainer(lambda fn, *a: [fn(worker, *a)], optimizer_steps_per_round=K).save_cut( + epoch=1, context=_context(tmp_path)) + saved = torch.load(tmp_path / "cut-v" / vcut.rank_file(0), weights_only=True)["model"] + assert "0.weight" not in saved and "0.bias" in saved + fresh = handles(frozen_model()) + fresh_worker = SimpleNamespace(module=fresh.module, optimizer=fresh.optimizer, lr_scheduler=fresh.scheduler) + vcut.VerlCutTrainer(lambda fn, *a: [fn(fresh_worker, *a)], optimizer_steps_per_round=K).restore_cut( + "cut-v", epoch=1, root=str(tmp_path), expect=_expect()) + for (name, p), (_, q) in zip(trained.module.named_parameters(), fresh.module.named_parameters()): + if p.requires_grad: # the frozen base keeps the fresh build's values (both from seed 0 here) + assert torch.equal(p, q), name + for key in ("exp_avg", "exp_avg_sq", "step"): + assert torch.equal(trained.optimizer.state[p][key], fresh.optimizer.state[q][key]), key diff --git a/tools/probes/c4_verl/modal_c4.py b/tools/probes/c4_verl/modal_c4.py index 39a337a19..a020bf9c8 100644 --- a/tools/probes/c4_verl/modal_c4.py +++ b/tools/probes/c4_verl/modal_c4.py @@ -69,6 +69,11 @@ def island(run: str, phase: str, gpu: str) -> dict: want_name, group = EXPECT[gpu] out = f"/vol/{run}/{phase}" os.makedirs(out, exist_ok=True) + done = f"{out}/result.json" + if os.path.isfile(done): # Modal restarts an input after a preemption: never rerun a finished phase + previous = json.load(open(done)) + if previous.get("rc") == 0: + return {**previous, "reused_after_restart": True} log = open(f"{out}/container.log", "a") name = subprocess.run(["nvidia-smi", "--query-gpu=name", "--format=csv,noheader"], capture_output=True, text=True).stdout.strip() diff --git a/yeto/rl/adapters/verl/cut.py b/yeto/rl/adapters/verl/cut.py index 360268628..20b6b3ae3 100644 --- a/yeto/rl/adapters/verl/cut.py +++ b/yeto/rl/adapters/verl/cut.py @@ -90,10 +90,13 @@ def layout_for(world: int) -> dict[str, int]: # -------------------------------------------------------------------------- -def _state_options() -> Any: +def _state_options(*, load: bool = False) -> Any: + """Local shards of the trainable state. A load is not strict: the cut holds no + frozen base weights (LoRA), which a strict load reports as missing (GPU T-C4). + ``restore_rank`` checks the trainable names against the cut before loading.""" from torch.distributed.checkpoint.state_dict import StateDictOptions - return StateDictOptions(full_state_dict=False, ignore_frozen_params=True) + return StateDictOptions(full_state_dict=False, ignore_frozen_params=True, strict=not load) def _is_dtensor(value: Any) -> bool: @@ -278,7 +281,7 @@ def value(key: str, local: Any, like: str) -> Any: "param_groups": saved["optim"]["param_groups"], } set_state_dict(handles.module, handles.optimizer, model_state_dict=model_sd, optim_state_dict=optim_sd, - options=_state_options()) + options=_state_options(load=True)) handles.scheduler.load_state_dict(saved["scheduler"]) _set_rng_state(saved["rng"]) return {"path": entry.path, **_rank_summary(handles, rank=rank, world=world, digest=_digest(handles))} From 615ebfc30853ca94b868d4f72e082da8ef044e71 Mon Sep 17 00:00:00 2001 From: michael <78113554+chowhsu@users.noreply.github.com> Date: Tue, 20 Oct 2026 04:44:35 +0000 Subject: [PATCH 38/89] Round cut walks a ChainedOptimizer through its members Megatron's ChainedOptimizer.optimizer asserts a single member; an MoE model chains the dense and the expert optimizer, so the cut's side-effect-free state read failed on the first round (Modal run-14). Driver: mirror the islands' local state (tapes) with the logs. Co-Authored-By: Claude Opus 5.5 --- tests/test_s25_g1_stage1mini.py | 23 +++++++++++++++++++++++ tools/probes/s25_g1/modal_app.py | 18 +++++++++++++++--- yeto/rl/adapters/miles/cut_plugin.py | 9 +++++++-- 3 files changed, 45 insertions(+), 5 deletions(-) diff --git a/tests/test_s25_g1_stage1mini.py b/tests/test_s25_g1_stage1mini.py index 70ea647a6..edf56d5c9 100644 --- a/tests/test_s25_g1_stage1mini.py +++ b/tests/test_s25_g1_stage1mini.py @@ -217,3 +217,26 @@ def test_align_weight_version_is_a_no_op_on_a_trainer_only_island(): updater = SimpleNamespace(weight_version=0) align_weight_version(SimpleNamespace(weight_updater=updater, args=None), next_version=3) assert updater.weight_version == 2 + + +def test_cut_state_walk_does_not_touch_the_chained_optimizer_property(): + from types import SimpleNamespace + + from yeto.rl.adapters.miles.cut_plugin import _state_dicts, side_effect_free_state + + class Chained: + def __init__(self, members): + self.chained_optimizers = members + self.state = {} + + @property + def optimizer(self): # Megatron asserts len(chained_optimizers) == 1 + raise AssertionError("ChainedOptimizer has more than one optimizer when accessing self.optimizer") + + dense = SimpleNamespace(optimizer=SimpleNamespace(state={"a": {"step": 1}}), state=None) + expert = SimpleNamespace(optimizer=SimpleNamespace(state={}), state=None) + chained = Chained([dense, expert]) + assert len(_state_dicts(chained)) == 3 + with side_effect_free_state(chained): + expert.optimizer.state["new"] = {} + assert "new" not in expert.optimizer.state and dense.optimizer.state == {"a": {"step": 1}} diff --git a/tools/probes/s25_g1/modal_app.py b/tools/probes/s25_g1/modal_app.py index 50c2f8021..ee4ce90b2 100644 --- a/tools/probes/s25_g1/modal_app.py +++ b/tools/probes/s25_g1/modal_app.py @@ -93,12 +93,20 @@ def _commit(tag: str | None = None) -> None: print(f"[g1] log mirror/commit failed: {error}", flush=True) -def _mirror_every(tag: str, seconds: float = 180.0) -> None: +def _mirror_every(tag: str, seconds: float = 180.0, state: list | None = None) -> None: + """Mirror the logs (and, once known, ``state[0]``: the island's local state dir with its + tape) every ``seconds``; a cancelled container keeps what was mirrored last.""" import threading def loop() -> None: while True: time.sleep(seconds) + if state: + try: + shutil.copytree(state[0], LOCAL_LOGS / tag / "state", dirs_exist_ok=True, + ignore=shutil.ignore_patterns("*.pt", "*.safetensors")) + except Exception as error: # noqa: BLE001 + print(f"[g1] state copy failed: {error}", flush=True) _commit(tag) threading.Thread(target=loop, daemon=True).start() @@ -220,9 +228,11 @@ def sampler_island(tag: str) -> int: log = _log_dir(tag) _gpu_facts(log, 1) nvml = _nvml_sampler(log / "nvml.csv") - _mirror_every(tag) + state_dir: list = [] + _mirror_every(tag, 120.0, state_dir) plan = _plan(log) run = plan["sampler"]["engines"][0] + state_dir.append(Path(run["spec"]["event_tape"]).parent) spec_path = Path(run["spec_path"]) spec_path.parent.mkdir(parents=True, exist_ok=True) spec_path.write_text(json.dumps(run["spec"], indent=1)) @@ -243,8 +253,10 @@ def trainer_island(tag: str, stop_after_rounds: int | None = None, convert_only: log = _log_dir(tag) _gpu_facts(log, 2) nvml = _nvml_sampler(log / "nvml.csv") - _mirror_every(tag) + state_dir: list = [] + _mirror_every(tag, 120.0, state_dir) plan = _plan(log) + state_dir.append(Path("/data/yeto", plan["run_id"])) trainer = plan["trainer"] env = _env(trainer["env"]) island = list(trainer["miles"]["island_argv"]) diff --git a/yeto/rl/adapters/miles/cut_plugin.py b/yeto/rl/adapters/miles/cut_plugin.py index bb22edb07..50fe3afde 100644 --- a/yeto/rl/adapters/miles/cut_plugin.py +++ b/yeto/rl/adapters/miles/cut_plugin.py @@ -165,8 +165,13 @@ def _state_dicts(optimizer: Any) -> list[Any]: if obj is None or id(obj) in seen: continue seen.add(id(obj)) - stack.extend(getattr(obj, "chained_optimizers", None) or ()) - stack.append(getattr(obj, "optimizer", None)) + chained = getattr(obj, "chained_optimizers", None) + if chained: + # Megatron's ChainedOptimizer.optimizer asserts a single member (an MoE model + # chains the dense and the expert optimizer; Modal run-14 failed the cut here) + stack.extend(chained) + else: + stack.append(getattr(obj, "optimizer", None)) state = getattr(obj, "state", None) if isinstance(state, Mapping): out.append(state) From 6f2a05639ef1cee06674e69001ada3488a6efc40 Mon Sep 17 00:00:00 2001 From: michael <78113554+chowhsu@users.noreply.github.com> Date: Tue, 20 Oct 2026 04:44:49 +0000 Subject: [PATCH 39/89] Mini plan: capture SGLang CUDA graphs up to batch 16 only Co-Authored-By: Claude Opus 5.5 --- examples/stage1/glm5next-mini-modal.yaml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/examples/stage1/glm5next-mini-modal.yaml b/examples/stage1/glm5next-mini-modal.yaml index 418f01690..fcf318289 100644 --- a/examples/stage1/glm5next-mini-modal.yaml +++ b/examples/stage1/glm5next-mini-modal.yaml @@ -110,7 +110,7 @@ sampler: in_place_reload: false # SGLang fuses the BF16 shared expert into the routed FusedMoE as expert 32; under the # MXFP4 MoE method that slot is packed U8, so the BF16 shared weights do not fit (smoke-2) - extra_args: [--disable-shared-experts-fusion] + extra_args: [--disable-shared-experts-fusion, --cuda-graph-max-bs, "16"] # 16 = 4 prompts x 4 samples; capture took 118 s at the default sizes extra_env: PYTHONPATH: /root/sky_workdir checks: From 7a25a049df1647e2563c2e8af547fb8cf24abb11 Mon Sep 17 00:00:00 2001 From: michael <78113554+chowhsu@users.noreply.github.com> Date: Tue, 20 Oct 2026 04:45:24 +0000 Subject: [PATCH 40/89] Turn the reclaim signal into the marker inside the sampler island Modal signals the whole container: in T-C1 run a the sampler process got SIGINT next to the runner, raised KeyboardInterrupt inside its round throttle and left as INFRA_FAILED. A pair sampler island (the launcher sets the marker path) now handles SIGINT/SIGTERM on its main thread by writing the reclaim marker, so the watcher leaves as PREEMPTED and the island waits for the container to stop. Co-Authored-By: Claude Opus 5.5 --- tests/test_rl_spot_sampler_replace.py | 19 ++++++++++++++++ yeto/rl/adapters/miles/sampler_island.py | 4 ++++ yeto/rl/engine/sampler_resume.py | 29 ++++++++++++++++++++++++ 3 files changed, 52 insertions(+) diff --git a/tests/test_rl_spot_sampler_replace.py b/tests/test_rl_spot_sampler_replace.py index 7e8d2c204..184c969a5 100644 --- a/tests/test_rl_spot_sampler_replace.py +++ b/tests/test_rl_spot_sampler_replace.py @@ -408,3 +408,22 @@ def test_the_reclaim_watcher_leaves_once_while_the_loop_is_still_in_a_round(tmp_ assert (leave["reason"], leave["rounds"], leave["in_flight_rollout_id"]) == ("PREEMPTED", 2, 4) record = (tmp_path / "sample-buffer/members/island0/b.json").read_text() assert '"leave_reason": "PREEMPTED"' in record and '"joined_at": 7.0' in record + + +def test_the_reclaim_signal_writes_the_marker_instead_of_raising(tmp_path): + from yeto.rl.engine.sampler_resume import install_marker_signals + + handlers = {} + + class Sig: + SIGINT, SIGTERM = _signal.SIGINT, _signal.SIGTERM + + @staticmethod + def signal(num, fn): + handlers[num] = fn + + env = {RECLAIM_MARKER_ENV: str(tmp_path / "marker")} + assert install_marker_signals(env, signal_mod=Sig) == ["SIGINT", "SIGTERM"] + assert not reclaim_requested(env) + handlers[_signal.SIGINT](_signal.SIGINT, None) # Modal's SIGINT reaches the island process + assert reclaim_requested(env) and '"signal": 2' in (tmp_path / "marker").read_text() diff --git a/yeto/rl/adapters/miles/sampler_island.py b/yeto/rl/adapters/miles/sampler_island.py index 3fd9b6092..0f347e510 100644 --- a/yeto/rl/adapters/miles/sampler_island.py +++ b/yeto/rl/adapters/miles/sampler_island.py @@ -261,8 +261,10 @@ def run_sampler_island(miles_args: Any, *, environ: Any = None, from yeto.rl.engine.resume import store_for from yeto.rl.engine.sampler_loop import SamplerLoop from yeto.rl.engine.sampler_resume import ( + RECLAIM_MARKER_ENV, ReclaimWatcher, SamplerMembership, + install_marker_signals, new_incarnation, reclaim_marker, reclaim_requested, @@ -290,6 +292,8 @@ def emit(event: str, **fields: Any) -> None: raise SampleBufferRolloutError(f"{SAMPLER_ID_ENV} is not set") incarnation = env[INCARNATION_ENV] = new_incarnation() reclaim_marker(env).unlink(missing_ok=True) # a restarted machine keeps its /tmp + if env.get(RECLAIM_MARKER_ENV): # set by the launcher on a pair's sampler island + install_marker_signals(env) buffer = open_buffer(env, sampler=True) resume = resume_point(buffer.entries(), sampler_id) membership = SamplerMembership(store_for(config.store, environ=environ), sampler_id, diff --git a/yeto/rl/engine/sampler_resume.py b/yeto/rl/engine/sampler_resume.py index 203671aed..36ea72635 100644 --- a/yeto/rl/engine/sampler_resume.py +++ b/yeto/rl/engine/sampler_resume.py @@ -154,6 +154,35 @@ def reclaim_marker(environ: Mapping[str, str]) -> Path: return Path(environ.get(RECLAIM_MARKER_ENV) or DEFAULT_RECLAIM_MARKER) +def write_reclaim_marker(environ: Mapping[str, str], signum: int) -> Path: + """Write the reclaim marker (what the Modal runner writes on its signal).""" + path = reclaim_marker(environ) + tmp = path.with_name(path.name + ".tmp") + tmp.write_text(json.dumps({"signal": int(signum), "time": time.time()}) + "\n") + tmp.replace(path) + return path + + +def install_marker_signals(environ: Mapping[str, str], signal_mod: Any = None) -> list[str]: + """On the island's main thread: SIGINT and SIGTERM write the reclaim marker + instead of raising ``KeyboardInterrupt``. Modal signals the whole container, + so the island process gets the reclaim signal next to the runner; without + this it dies inside its round and leaves as INFRA_FAILED (T-C1 run a). + Returns the signals handled ([] off the main thread).""" + import signal as _signal + + sig = signal_mod or _signal + if threading.current_thread() is not threading.main_thread(): + return [] + done = [] + for name in ("SIGINT", "SIGTERM"): + num = getattr(sig, name, None) + if num is not None: + sig.signal(num, lambda signum, _frame: write_reclaim_marker(environ, signum)) + done.append(name) + return done + + def reclaim_requested(environ: Mapping[str, str]) -> bool: """The Modal runner wrote the reclaim marker (the container got SIGINT/SIGTERM).""" return reclaim_marker(environ).exists() From 8e8253e929c9324a726f7e970e669ab1f40a3870 Mon Sep 17 00:00:00 2001 From: michael <78113554+chowhsu@users.noreply.github.com> Date: Tue, 20 Oct 2026 04:49:15 +0000 Subject: [PATCH 41/89] T-C4 verl probe: fresh container per phase and a Volume reload before it Co-Authored-By: Claude Opus 5.5 --- tools/probes/c4_verl/modal_c4.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/tools/probes/c4_verl/modal_c4.py b/tools/probes/c4_verl/modal_c4.py index a020bf9c8..5b0bbfd21 100644 --- a/tools/probes/c4_verl/modal_c4.py +++ b/tools/probes/c4_verl/modal_c4.py @@ -60,12 +60,15 @@ def _sh(cmd: str, log) -> int: return subprocess.run(cmd, shell=True, stdout=log, stderr=subprocess.STDOUT).returncode -@app.function(gpu=GPU, cpu=8.0, memory=49152, timeout=40 * 60, volumes={"/vol": vol}) +# max_inputs=1: every phase gets a fresh container (a reused warm container kept the +# previous phase's ~/yeto-output and its stale Volume view: run c4-verl-20261020c b). +@app.function(gpu=GPU, cpu=8.0, memory=49152, timeout=40 * 60, volumes={"/vol": vol}, max_inputs=1) def island(run: str, phase: str, gpu: str) -> dict: import shutil import subprocess import time + vol.reload() # see the store the previous phase committed want_name, group = EXPECT[gpu] out = f"/vol/{run}/{phase}" os.makedirs(out, exist_ok=True) From 94b0a6026c0696091aa086ee5ae4e1ed9313c2c2 Mon Sep 17 00:00:00 2001 From: michael <78113554+chowhsu@users.noreply.github.com> Date: Tue, 20 Oct 2026 04:49:41 +0000 Subject: [PATCH 42/89] Add --rl-sampler-spot: spot sampler islands, on-demand trainer A pair's sampler islands can now use spot capacity while the trainer stays on demand. The sampler island copy gets spot billing without the legacy --spot volume (its durable state is the sample buffer); the flag is refused outside --rl-island-sample-role and together with --spot. Co-Authored-By: Claude Opus 5.5 --- tests/test_rl_spot_sampler_replace.py | 35 +++++++++++++++++++++++++++ yeto/cloud/droppable.py | 6 +++-- yeto/launch/rl_ports.py | 3 +++ yeto/launch_preflight.py | 11 ++++++--- yeto/rl/engine/sample_roles.py | 7 ++++++ 5 files changed, 57 insertions(+), 5 deletions(-) diff --git a/tests/test_rl_spot_sampler_replace.py b/tests/test_rl_spot_sampler_replace.py index 184c969a5..8e8b4f4f1 100644 --- a/tests/test_rl_spot_sampler_replace.py +++ b/tests/test_rl_spot_sampler_replace.py @@ -427,3 +427,38 @@ def signal(num, fn): assert not reclaim_requested(env) handlers[_signal.SIGINT](_signal.SIGINT, None) # Modal's SIGINT reaches the island process assert reclaim_requested(env) and '"signal": 2' in (tmp_path / "marker").read_text() + + +def test_sampler_spot_bills_only_the_samplers_and_adds_no_spot_volume(monkeypatch): + from yeto.cloud import s3_store + from yeto.gpu_spec import parse_gpu_spec + from yeto.launch.rl_island_task import make_miles_island_task + from yeto.launcher import _island_task, _prepare_rl_args + + monkeypatch.setenv(s3_store.KEY_ID_ENV, "AKIDSTORE") + monkeypatch.setenv(s3_store.SECRET_ENV, "store-secret") + gpu = "nebius:1xh100@eu-north1,nebius:1xh100@eu-north1" + base = ("--rl-single-island-no-sync", "--tuning", "lora", "--rl-island-sample-role", "0:sampler", + "--rl-island-sample-role", "1:trainer", "--rl-max-policy-age", "2", + "--rl-algorithm-spec", "examples/rl_algorithms/decoupled-ppo-c2.json", + "--rl-checkpoint-store", "s3://yeto-ckpt/run-a", "--on-demand", "--rl-sampler-spot") + args = _launcher_args("ports", base, gpu=gpu) + args.controller = "local" + _prepare_rl_args(args) + spots, volumes = [], [] + for m, spec in enumerate(parse_gpu_spec(args.gpu)): + task = _island_task(make_miles_island_task, args, spec, m, 2, "none", {}) + spots.append(task.resources.use_spot) + volumes.append(sorted(getattr(task, "storage_mounts", None) or {})) + # the sampler gets the store mount only, like the trainer: no legacy --spot volume + assert spots == [True, False] and volumes[0] == volumes[1] and len(volumes[0]) == 1 + + +def test_sampler_spot_needs_a_pair(): + from yeto.launcher import _prepare_rl_args + + args = _launcher_args("ports", ("--rl-single-island-no-sync", "--rl-sampler-spot"), + gpu="modal:1xh100") + args.controller = "local" + with pytest.raises(ValueError, match="--rl-sampler-spot needs --rl-island-sample-role"): + _prepare_rl_args(args) diff --git a/yeto/cloud/droppable.py b/yeto/cloud/droppable.py index abc6cb22c..abd1e3759 100644 --- a/yeto/cloud/droppable.py +++ b/yeto/cloud/droppable.py @@ -126,8 +126,10 @@ def role_args(args, role: str | None): def island_keeps_spot_volume(args) -> bool: """The legacy ``--spot`` per-island volume (completed groups, re-synced on a - rebuild) is a resume from an archive; a droppable island must not get it.""" - return bool(getattr(args, "spot", False)) and getattr(args, "_island_role", None) != DROPPABLE + rebuild) is a resume from an archive; a droppable island must not get it, nor + a spot sampler island (``spot_without_volume``: its state is the buffer).""" + return (bool(getattr(args, "spot", False)) and getattr(args, "_island_role", None) != DROPPABLE + and not getattr(args, "spot_without_volume", False)) def role_env(role: str | None, cloud: str, region: str | None) -> dict[str, str]: diff --git a/yeto/launch/rl_ports.py b/yeto/launch/rl_ports.py index bdeceda94..22470ce98 100644 --- a/yeto/launch/rl_ports.py +++ b/yeto/launch/rl_ports.py @@ -380,6 +380,9 @@ def _check_ports_infra_switches(args) -> None: sample_roles.check_pair(args, rl_island_count(args)) else: sample_roles.check_launch(args) # L4.6 --rl-sample-role (both checks nothing) + if getattr(args, "rl_sampler_spot", False) and not sample_roles.is_pair(args): + raise ValueError("--rl-sampler-spot needs --rl-island-sample-role (it bills the sampler " + "islands of a pair as spot)") from .storage import store_object_envs for cloud in {s.cloud for s in parse_gpu_spec(args.gpu)} if getattr(args, "gpu", None) else (): diff --git a/yeto/launch_preflight.py b/yeto/launch_preflight.py index 419432c2e..539b07a07 100644 --- a/yeto/launch_preflight.py +++ b/yeto/launch_preflight.py @@ -363,12 +363,17 @@ def add_cli_args(p) -> None: "base, never from a checkpoint store); repeatable; an island without a role is " "an anchor; replaces --spot") g.add_argument("--rl-island-sample-role", action="append", default=None, metavar="ISLAND:ROLE", - help="rl-trainer-start-stop L5.4: one launch of a sampler island and a trainer " + help="rl-trainer-start-stop L5.4: one launch of sampler islands and a trainer " "island over one --rl-checkpoint-store (ROLE sampler or trainer, one entry per " - "--gpu island; exactly two islands). Needs --rl-single-island-no-sync: no " + "--gpu island; exactly one trainer, one or more samplers; several samplers " + "split --total-steps). Needs --rl-single-island-no-sync: no " "syncer, each island runs as a one-learner no-sync island with that " - "--rl-sample-role. Combine with --rl-trainer-scheduler s1 " + "--rl-sample-role. With one sampler, combine with --rl-trainer-scheduler s1 " "--rl-trainer-island M to start and stop the trainer island") + g.add_argument("--rl-sampler-spot", action="store_true", + help="with --rl-island-sample-role: the sampler islands use spot (preemptible) " + "capacity, the trainer island stays on demand; a reclaimed sampler island is " + "started again and continues after its last round in the buffer") g.add_argument("--rl-negative-test-run", action="store_true", help="mark this run as a negative-test run (required by --rl-island-override); " "a normal run refuses to resume its checkpoint store and `yeto merge` refuses " diff --git a/yeto/rl/engine/sample_roles.py b/yeto/rl/engine/sample_roles.py index 69f72d7c8..072a3e3ab 100644 --- a/yeto/rl/engine/sample_roles.py +++ b/yeto/rl/engine/sample_roles.py @@ -152,6 +152,11 @@ def island_args(args: Any, island: int, num_islands: int) -> Any: out.rl_sampler_node = node_name(island) if out.rl_sample_role == SAMPLER else None out.rl_buffer_shards = buffer_shards(args, num_islands) samplers = len(sampler_islands(args, num_islands)) + if out.rl_sample_role == SAMPLER and getattr(args, "rl_sampler_spot", False): + # spot billing without the legacy --spot volume: the buffer is the sampler's + # durable state (S25 C1) + out.spot = True + out.spot_without_volume = True if out.rl_sample_role == SAMPLER and samplers > 1: # the trainer trains --total-steps rounds; the samplers share them, so each # one reaches its data end (check_pair needs an even split) @@ -205,6 +210,8 @@ def check_pair(args: Any, num_islands: int) -> None: if samplers > 1 and int(getattr(args, "total_steps", 0) or 0) % samplers: raise ValueError(f"--total-steps {args.total_steps}: {samplers} sampler islands share the " "trainer's rounds; it must be a multiple of the sampler count") + if getattr(args, "rl_sampler_spot", False) and getattr(args, "spot", False): + raise ValueError("--rl-sampler-spot bills only the samplers as spot; it does not take --spot") if samplers > 1 and getattr(args, "rl_trainer_scheduler", None) not in (None, "off"): raise ValueError("--rl-trainer-scheduler with several sampler islands is not supported " "(the head does not read the sampler shards); drop the scheduler") From 8105fc4398ed4340f0b5c45cb4d016a61f7d5b36 Mon Sep 17 00:00:00 2001 From: michael <78113554+chowhsu@users.noreply.github.com> Date: Tue, 20 Oct 2026 04:50:53 +0000 Subject: [PATCH 43/89] Mini plan: the SGLang flag is --cuda-graph-max-bs-decode Co-Authored-By: Claude Opus 5.5 --- examples/stage1/glm5next-mini-modal.yaml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/examples/stage1/glm5next-mini-modal.yaml b/examples/stage1/glm5next-mini-modal.yaml index fcf318289..71e3b4d2a 100644 --- a/examples/stage1/glm5next-mini-modal.yaml +++ b/examples/stage1/glm5next-mini-modal.yaml @@ -110,7 +110,7 @@ sampler: in_place_reload: false # SGLang fuses the BF16 shared expert into the routed FusedMoE as expert 32; under the # MXFP4 MoE method that slot is packed U8, so the BF16 shared weights do not fit (smoke-2) - extra_args: [--disable-shared-experts-fusion, --cuda-graph-max-bs, "16"] # 16 = 4 prompts x 4 samples; capture took 118 s at the default sizes + extra_args: [--disable-shared-experts-fusion, --cuda-graph-max-bs-decode, "16"] # 16 = 4 prompts x 4 samples; capture took 118 s at the default sizes extra_env: PYTHONPATH: /root/sky_workdir checks: From 45006c7809eabc463e048af576ec74155b34a846 Mon Sep 17 00:00:00 2001 From: michael <78113554+chowhsu@users.noreply.github.com> Date: Tue, 20 Oct 2026 05:01:41 +0000 Subject: [PATCH 44/89] verl island: order-independent prompt digest per round; record the first-step resume check Co-Authored-By: Claude Opus 5.5 --- yeto/rl/adapters/verl/ports_impl.py | 13 +++++++++++-- 1 file changed, 11 insertions(+), 2 deletions(-) diff --git a/yeto/rl/adapters/verl/ports_impl.py b/yeto/rl/adapters/verl/ports_impl.py index 2a44c1719..4952a73c2 100644 --- a/yeto/rl/adapters/verl/ports_impl.py +++ b/yeto/rl/adapters/verl/ports_impl.py @@ -216,6 +216,10 @@ def generate(self, rollout_id: int, *, expected_policy_version: str | None = Non padded = data.to_padded_tensor() uids = [str(u) for u in padded.pop("uid").tolist()] prompt_sha = hashlib.sha256(padded["prompts"].cpu().numpy().tobytes()).hexdigest() + # the same digest independent of arrival order (rows sorted): compares the prompts two + # runs drew for one round (T-C4: a resumed island vs a continuous one) + prompt_set_sha = hashlib.sha256(b"".join(sorted( + row.tobytes() for row in padded["prompts"].cpu().numpy()))).hexdigest() gen_steps = sorted({int(tag.get("global_steps", -1)) for tag in batch.tags}) rewards = padded["rm_scores"].float().sum(-1).tolist() lengths = padded["response_mask"].float().sum(-1).tolist() @@ -246,7 +250,8 @@ def generate(self, rollout_id: int, *, expected_policy_version: str | None = Non "rollout_id": rollout_id, "policy_token": token, "groups": len(metas), "samples": len(uids), "reward_mean": float(np.mean(rewards)) if rewards else None, "gen_seconds": gen_seconds, "expected_policy_version": expected_policy_version, - "prompts_sha256": prompt_sha, "sample_submit_global_steps": gen_steps, + "prompts_sha256": prompt_sha, "prompt_set_sha256": prompt_set_sha, + "sample_submit_global_steps": gen_steps, "trainer_global_step": t.global_steps, "rewards": rewards, "uids": uids, "data_cursor": self.data_cursor()}) return RolloutBatchHandle( @@ -352,7 +357,11 @@ def train_step(self, batch: RolloutBatchHandle) -> LocalStepReceipt: applied_lrs=None if lr is None else (lr,), ) if grad_norm is not None and self._cut_trainer is not None: - self._cut_trainer.check_first_step(1) # no-op unless a restore is pending + pending = self._cut_trainer.resume_reference is not None + self._cut_trainer.check_first_step(1) # no-op unless a restore is pending; raises on a problem + if pending: + isl.emit("rl_resume_check", rollout_id=batch.rollout_id, check="first_step", ok=True, + optimizer_steps=self._cut_trainer.optimizer_steps_per_round) ids = tuple(s for g in batch.groups for s in g.sample_ids) return LocalStepReceipt( algorithm="grpo", learner_id=isl.learner_id, learner_generation=0, From 51a6ebf0fee2905774c8651d1796df2430bd3306 Mon Sep 17 00:00:00 2001 From: michael <78113554+chowhsu@users.noreply.github.com> Date: Tue, 20 Oct 2026 05:07:55 +0000 Subject: [PATCH 45/89] Cut coverage: dense tensors are covered over the EP ranks of a (tp, pp) With EP>1 the dense DistOpt shards span the EP ranks of one (tp, pp); the cut checked them per (tp, pp, ep) and refused every cut (Modal run-16). Expert tensors stay per (tp, pp, ep). Driver: one input per GPU container (trainer B got trainer A's warm container with Ray still running), ray stop first. Co-Authored-By: Claude Opus 5.5 --- tests/test_s25_g1_stage1mini.py | 20 ++++++++++++ tools/probes/s25_g1/modal_app.py | 32 ++++++++++++------ yeto/rl/adapters/miles/trainer.py | 54 ++++++++++++++++++++++--------- 3 files changed, 82 insertions(+), 24 deletions(-) diff --git a/tests/test_s25_g1_stage1mini.py b/tests/test_s25_g1_stage1mini.py index edf56d5c9..e6d2d3113 100644 --- a/tests/test_s25_g1_stage1mini.py +++ b/tests/test_s25_g1_stage1mini.py @@ -240,3 +240,23 @@ def optimizer(self): # Megatron asserts len(chained_optimizers) == 1 with side_effect_free_state(chained): expert.optimizer.state["new"] = {} assert "new" not in expert.optimizer.state and dense.optimizer.state == {"a": {"step": 1}} + + +def test_cut_coverage_groups_dense_tensors_over_the_ep_ranks(): + from yeto.rl.adapters.miles.trainer import optimizer_coverage_problems + + dense, expert = "decoder.layers.0.mlp.linear_fc1.weight", "decoder.layers.1.mlp.experts.linear_fc1.weight0" + shared = "decoder.layers.1.mlp.shared_experts.linear_fc1.weight" + + def rank(ep, optimized): + return {"coord": {"tp": 0, "pp": 0, "ep": ep}, "adapter_names": [dense, shared, expert], + "optimizer_names": optimized} + + # EP2 x DP2: the dense DistOpt shards split over both EP ranks; each rank owns its experts + ok = [rank(0, [dense, expert]), rank(1, [shared, expert])] + assert optimizer_coverage_problems(ok) == [] + bad = [rank(0, [dense, expert]), rank(1, [shared])] + problems = optimizer_coverage_problems(bad) + assert problems == [f"tp0/pp0/ep1: no optimizer state for ['{expert}']"] + assert bad[1]["has_optimizer_state"] is False and "has_optimizer_state" not in bad[0] + assert optimizer_coverage_problems([]) == ["no rank reported any adapter tensor"] diff --git a/tools/probes/s25_g1/modal_app.py b/tools/probes/s25_g1/modal_app.py index ee4ce90b2..dbc244b87 100644 --- a/tools/probes/s25_g1/modal_app.py +++ b/tools/probes/s25_g1/modal_app.py @@ -13,7 +13,7 @@ ``smoke`` and ``launch`` spawn and return; the local client may exit (never kill a ``.remote()`` caller: memory modal-detach-kill-cancels). -Raw logs go to the store Volume under ``logs//``. +Raw logs go to the Volume ``yeto-s25-g1-logs`` under ``/``. """ from __future__ import annotations @@ -41,7 +41,10 @@ app = modal.App("yeto-s25-g1-stage1mini") models = modal.Volume.from_name(MODELS_VOLUME, create_if_missing=True, version=2) store = modal.Volume.from_name(STORE_VOLUME, create_if_missing=True, version=2) -VOLUMES = {"/data/models": models, "/data/store": store} +# Logs go to their own Volume: the islands reload the store Volume, and a reload +# drops this container's uncommitted copies there (run-16 lost its log mirror). +logs = modal.Volume.from_name("yeto-s25-g1-logs", create_if_missing=True, version=2) +VOLUMES = {"/data/models": models, "/data/store": store, "/data/logs": logs} _PATCH_SGLANG = ( "import importlib.util, sys; sys.path.insert(0, '/tmp/yeto-slr');" @@ -82,13 +85,21 @@ def _log_dir(tag: str) -> Path: return path +def _copy_files(src: Path, dst: Path) -> None: + """Copy file contents only (no copystat: a Volume refuses some metadata calls).""" + for path in src.rglob("*"): + if path.is_file() and path.suffix not in (".sock", ".pid"): + target = dst / path.relative_to(src) + target.parent.mkdir(parents=True, exist_ok=True) + shutil.copyfile(path, target) + + def _commit(tag: str | None = None) -> None: """Copy the local logs of ``tag`` to the store Volume (closed files only), then commit.""" try: if tag is not None: - shutil.copytree(LOCAL_LOGS / tag, Path("/data/store/logs") / tag, dirs_exist_ok=True, - ignore=shutil.ignore_patterns("*.sock", "*.pid")) - store.commit() + _copy_files(LOCAL_LOGS / tag, Path("/data/logs") / tag) + logs.commit() except Exception as error: # noqa: BLE001 - evidence copy only print(f"[g1] log mirror/commit failed: {error}", flush=True) @@ -223,7 +234,8 @@ def smoke(tag: str = "smoke-1") -> None: # --------------------------------------------------------------------- the islands -@app.function(image=image, volumes=VOLUMES, gpu=f"{GPU_NAME}!:1", cpu=8, memory=65536, timeout=7200) +@app.function(image=image, volumes=VOLUMES, gpu=f"{GPU_NAME}!:1", cpu=8, memory=65536, timeout=7200, + max_inputs=1) # a fresh container per call: trainer B must not inherit A's disk or Ray def sampler_island(tag: str) -> int: log = _log_dir(tag) _gpu_facts(log, 1) @@ -248,7 +260,8 @@ def sampler_island(tag: str) -> int: return rc -@app.function(image=image, volumes=VOLUMES, gpu=f"{GPU_NAME}!:2", cpu=12, memory=163840, timeout=3600) +@app.function(image=image, volumes=VOLUMES, gpu=f"{GPU_NAME}!:2", cpu=12, memory=163840, timeout=3600, + max_inputs=1) # a fresh container per call: trainer B must not inherit A's disk or Ray def trainer_island(tag: str, stop_after_rounds: int | None = None, convert_only: bool = False) -> int: log = _log_dir(tag) _gpu_facts(log, 2) @@ -283,6 +296,7 @@ def trainer_island(tag: str, stop_after_rounds: int | None = None, convert_only: nvml.terminate() _commit(tag) return 0 + _sh("ray stop --force || true", log / "ray.log", env=env, check=False) _sh("ray start --head --port=6379 --node-ip-address=127.0.0.1 --num-gpus=2 --disable-usage-stats", log / "ray.log", env=env) began = time.time() @@ -361,10 +375,10 @@ def run_trainer(sub: str, stop: int | None) -> object: requant.cancel() result = {"trainer_a": first, "trainer_b": second, "sampler": sampler_rc} print(json.dumps(result, default=str), flush=True) - out = Path(f"/data/store/logs/{tag}") + out = Path(f"/data/logs/{tag}") out.mkdir(parents=True, exist_ok=True) (out / "orchestrate.json").write_text(json.dumps(result, default=str)) - _commit() + logs.commit() return result diff --git a/yeto/rl/adapters/miles/trainer.py b/yeto/rl/adapters/miles/trainer.py index ff548c45a..54c65ed90 100644 --- a/yeto/rl/adapters/miles/trainer.py +++ b/yeto/rl/adapters/miles/trainer.py @@ -44,6 +44,44 @@ class TrainStepError(RuntimeError): pass + +def _is_expert_tensor(name: str) -> bool: + """A routed-expert tensor (Megatron ``mlp.experts.*``; not ``shared_experts``).""" + return ".mlp.experts." in name + +def optimizer_coverage_problems(summaries: list[dict[str, Any]]) -> list[str]: + """Check that every adapter tensor has optimizer state on at least one rank of + the group that shards it (DistOpt ranges); marks ``has_optimizer_state`` False on + the ranks of a group that fails. Expert tensors are sharded over the expert-DP + ranks of one (tp, pp, ep); every other tensor over all DP ranks of its (tp, pp), + which with EP>1 spans the EP ranks too (Modal run-16: the dense tensors were + checked per EP rank and found missing).""" + problems: list[str] = [] + covered: dict[tuple[int, ...], tuple[set[str], set[str]]] = {} + for s in summaries: + tp, pp, ep = s["coord"]["tp"], s["coord"]["pp"], s["coord"].get("ep", 0) + optimized = set(s.get("optimizer_names") or ()) + for name in s.get("adapter_names") or (): + key = (tp, pp, ep) if _is_expert_tensor(name) else (tp, pp) + names, owned = covered.setdefault(key, (set(), set())) + names.add(name) + if name in optimized: + owned.add(name) + if not covered: + problems.append("no rank reported any adapter tensor") + for key, (adapters, optimized) in sorted(covered.items()): + if not adapters <= optimized: + missing = sorted(adapters - optimized)[:4] + for s in summaries: + if (s["coord"]["tp"], s["coord"]["pp"]) == key[:2] and ( + len(key) == 2 or s["coord"].get("ep", 0) == key[2]): + s["has_optimizer_state"] = False + where = f"tp{key[0]}/pp{key[1]}" + (f"/ep{key[2]}" if len(key) == 3 else " (dense)") + problems.append(f"{where}: no optimizer state for {missing}") + return problems + + + def batch_hash(batch: RolloutBatchHandle) -> str: """Content identity of the trained batch, computed from metadata only.""" @@ -801,21 +839,7 @@ def save_cut(self, *, epoch: int, context: "CutContext") -> str: expected = trainer_workers(self._args) if len(summaries) != expected: raise TrainStepError(f"expected {expected} cut shards (one per rank), got {len(summaries)}") - # has_optimizer_state means: every adapter tensor of each (tp, pp) has - # optimizer state on at least one of its DP ranks (DistOpt ranges). - covered: dict[tuple[int, int], tuple[set[str], set[str]]] = {} - for s in summaries: - key = (s["coord"]["tp"], s["coord"]["pp"], s["coord"].get("ep", 0)) - adapters, optimized = covered.setdefault(key, (set(), set())) - adapters.update(s.get("adapter_names") or ()) - optimized.update(s.get("optimizer_names") or ()) - for key, (adapters, optimized) in sorted(covered.items()): - if not adapters or not adapters <= optimized: - missing = sorted(adapters - optimized)[:4] - for s in summaries: - if (s["coord"]["tp"], s["coord"]["pp"], s["coord"].get("ep", 0)) == key: - s["has_optimizer_state"] = False - problems.append(f"tp{key[0]}/pp{key[1]}/ep{key[2]}: no optimizer state for {missing or 'any adapter'}") + problems += optimizer_coverage_problems(summaries) if problems: raise CutError("refusing an incomplete cut: " + "; ".join(problems)) resume = self._resume_manifest(context, summaries) From cfc2fb49ee3107e8c0cda849db3c290213d2f809 Mon Sep 17 00:00:00 2001 From: michael <78113554+chowhsu@users.noreply.github.com> Date: Tue, 20 Oct 2026 05:08:13 +0000 Subject: [PATCH 46/89] s25_g1 driver: every launch starts from an empty run store Co-Authored-By: Claude Opus 5.5 --- tools/probes/s25_g1/modal_app.py | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/tools/probes/s25_g1/modal_app.py b/tools/probes/s25_g1/modal_app.py index dbc244b87..6eb1992cf 100644 --- a/tools/probes/s25_g1/modal_app.py +++ b/tools/probes/s25_g1/modal_app.py @@ -334,6 +334,10 @@ def orchestrate(tag: str, cut_after: int = 2) -> dict: """Sampler for the whole run; trainer A stops after ``cut_after`` rounds (cut), trainer B resumes in a new container. A trainer that fails cancels the sampler; a sampler that fails cancels the running trainer (no card waits for nothing).""" + run_store = Path("/data/store/s25-g1") # the plan's store: every run starts from an empty one + if run_store.exists(): + shutil.rmtree(run_store) + store.commit() sampler = sampler_island.spawn(f"{tag}/sampler") requant = requant_island.spawn(f"{tag}/requant") sampler_rc: object = None From 2c482b64ce5a073c28bfe79a110b58e9f83f2f5e Mon Sep 17 00:00:00 2001 From: michael <78113554+chowhsu@users.noreply.github.com> Date: Tue, 20 Oct 2026 05:08:50 +0000 Subject: [PATCH 47/89] Never raise from the sampler's reclaim signal handler In T-C1 run b the Modal runner and the sampler's handler wrote the marker through the same temp name at the same moment; the rename failed and the handler raised inside the island's main thread. The temp name is now per process and a failed write is printed, not raised. Co-Authored-By: Claude Opus 5.5 --- tests/test_rl_spot_sampler_replace.py | 8 ++++++++ yeto/rl/engine/sampler_resume.py | 15 +++++++++++---- 2 files changed, 19 insertions(+), 4 deletions(-) diff --git a/tests/test_rl_spot_sampler_replace.py b/tests/test_rl_spot_sampler_replace.py index 8e8b4f4f1..18e99c2d8 100644 --- a/tests/test_rl_spot_sampler_replace.py +++ b/tests/test_rl_spot_sampler_replace.py @@ -462,3 +462,11 @@ def test_sampler_spot_needs_a_pair(): args.controller = "local" with pytest.raises(ValueError, match="--rl-sampler-spot needs --rl-island-sample-role"): _prepare_rl_args(args) + + +def test_a_failed_marker_write_in_the_signal_handler_does_not_raise(tmp_path, capsys): + from yeto.rl.engine.sampler_resume import write_reclaim_marker + + env = {RECLAIM_MARKER_ENV: str(tmp_path / "missing-dir" / "marker")} + write_reclaim_marker(env, _signal.SIGINT) # T-C1 run b: the runner raced the same temp name + assert "reclaim marker not written" in capsys.readouterr().out diff --git a/yeto/rl/engine/sampler_resume.py b/yeto/rl/engine/sampler_resume.py index 36ea72635..50b947f0d 100644 --- a/yeto/rl/engine/sampler_resume.py +++ b/yeto/rl/engine/sampler_resume.py @@ -32,6 +32,7 @@ from __future__ import annotations import json +import os import threading import time import uuid @@ -155,11 +156,17 @@ def reclaim_marker(environ: Mapping[str, str]) -> Path: def write_reclaim_marker(environ: Mapping[str, str], signum: int) -> Path: - """Write the reclaim marker (what the Modal runner writes on its signal).""" + """Write the reclaim marker (what the Modal runner writes on its signal). The + runner may write it at the same moment: the temp file is per process, and a + failed write is reported, never raised (it runs in a signal handler; the + runner's marker serves as well).""" path = reclaim_marker(environ) - tmp = path.with_name(path.name + ".tmp") - tmp.write_text(json.dumps({"signal": int(signum), "time": time.time()}) + "\n") - tmp.replace(path) + tmp = path.with_name(f"{path.name}.{os.getpid()}.tmp") + try: + tmp.write_text(json.dumps({"signal": int(signum), "time": time.time()}) + "\n") + tmp.replace(path) + except OSError as error: + print(f"[sampler-island] reclaim marker not written ({error})", flush=True) return path From 3eb1780f5796872a08ee17c654a49adc1ec8aa5d Mon Sep 17 00:00:00 2001 From: michael <78113554+chowhsu@users.noreply.github.com> Date: Tue, 20 Oct 2026 05:26:08 +0000 Subject: [PATCH 48/89] Synced verl trainer islands feed sampler islands; merged releases over object stores - verl buffer trainer: read one buffer lane per synced island with a durable lane cursor (written when a round's collection ends, at most once); only the publisher (learner 0) releases, once per new merged version, with the syncer's outer version, kept increasing by MergedReleaseGuard; merged behavior versions map to the first local version that applied them. - verl island entry accepts --rl-sample-store / --rl-sample-lanes (shared check in sample_roles.check_sample_store). - Object stores (S3, OSS): lane cursors are one object per round; on a store whose samplers write shards, only the publisher's lane merges the index. The s3/oss refusal for synced trainers is removed. - N2-fix: PostLoadChecks no longer divides by zero without probe pairs; the SGLang adapter labels and checks the real tp/dp topology. Co-Authored-By: Claude Opus 5.5 --- tests/test_rl_hetero_sampler_engines.py | 28 +++ tests/test_rl_merged_release.py | 73 ++++++- tests/test_rl_verl_buffer_trainer.py | 115 ++++++++++- tests/test_rl_verl_synced_trainers.py | 186 ++++++++++++++++++ yeto/rl/adapters/miles/buffer_rollout.py | 24 +-- yeto/rl/adapters/miles/island_entry.py | 20 +- .../adapters/sampler_engines/sglang_server.py | 30 ++- yeto/rl/adapters/verl/buffer_feed.py | 80 +++++++- yeto/rl/adapters/verl/buffer_ports.py | 93 +++++++-- .../rl/adapters/verl/buffer_trainer_runner.py | 10 +- yeto/rl/adapters/verl/island_entry.py | 11 +- yeto/rl/adapters/verl/trainer.py | 2 + yeto/rl/engine/merged_release.py | 103 ++++++++-- yeto/rl/engine/sample_buffer.py | 21 +- yeto/rl/engine/sample_roles.py | 23 ++- yeto/rl/engine/sampler_checks.py | 2 + 16 files changed, 739 insertions(+), 82 deletions(-) create mode 100644 tests/test_rl_verl_synced_trainers.py diff --git a/tests/test_rl_hetero_sampler_engines.py b/tests/test_rl_hetero_sampler_engines.py index 7f2a5a442..1427f7d1c 100644 --- a/tests/test_rl_hetero_sampler_engines.py +++ b/tests/test_rl_hetero_sampler_engines.py @@ -624,3 +624,31 @@ def test_stop_sends_sigterm(tmp_path): engine.stop() assert process.signals == [signal.SIGTERM] and engine.last_exit["normal"] engine.stop() # idempotent + + +def test_sglang_topology_names_the_real_layout(tmp_path): + """S25 N2-fix: the source label is tp/dp with N = --tp, not tp8 for every layout.""" + spawned = [] + http = _Http({("GET", "/health"): {}}) + one = SGLangServerEngine(_config(tmp_path, topology="tp1"), SGLangOptions(tp=1), + spawn=lambda a, e: spawned.append(a) or _Process(), http=http, sleep=lambda s: None) + assert one.identity.topology == "tp1" + one.start(None, 0) + argv = spawned[0] + assert argv[argv.index("--tp") + 1] == "1" and "--dp" not in argv + dp = SGLangServerEngine(_config(tmp_path, topology="dp4"), SGLangOptions(tp=4), + spawn=lambda a, e: spawned.append(a) or _Process(), http=http, sleep=lambda s: None) + dp.start(None, 0) + assert spawned[1][spawned[1].index("--dp") + 1] == "4" + for topology, tp in (("tp8", 1), ("tp1", 8), ("pd8", 8), ("tp", 8)): + with pytest.raises(SamplerEngineError, match="does not match TP"): + SGLangServerEngine(_config(tmp_path, topology=topology), SGLangOptions(tp=tp)) + + +def test_probe_without_pairs_passes_every_load(): + """S25 N2-fix: no fixed prompts and no probe pairs: the second load divided by zero.""" + checks = PostLoadChecks(expected_expert_matrices=None, fixed_prompts=()) + for _ in range(2): + result = checks.logprob_probe(engine=None) + checks.accept() + assert result == {"name": "logprob_probe", "passed": True, "detail": "no probe pairs"} diff --git a/tests/test_rl_merged_release.py b/tests/test_rl_merged_release.py index 5446f2b66..c0b10287c 100644 --- a/tests/test_rl_merged_release.py +++ b/tests/test_rl_merged_release.py @@ -20,13 +20,16 @@ BufferLane, LaneCursorLog, MergedReleaseGuard, + ObjectLaneCursorLog, gate_version, is_publisher, + lane_cursor_log, merged_version, + merges_shards, ) from yeto.rl.engine.policy_release import PolicyReleaseStore from yeto.rl.engine.resume import CheckpointStore -from yeto.rl.engine.sample_buffer import FileSampleBuffer, SampleBufferError +from yeto.rl.engine.sample_buffer import FileSampleBuffer, SampleBufferError, open_trainer_buffer from yeto.rl.engine.trainer_lifecycle import SampleIndexEntry @@ -62,12 +65,68 @@ def test_two_lane_readers_split_one_buffer_without_overlap(tmp_path): assert [s.position for s in again.select("trainer", 2)] == [4, 6] -def test_lane_reader_refuses_an_object_store_and_a_shard_merge(tmp_path): - with pytest.raises(SampleBufferError, match="shared-directory"): - FileSampleBuffer(CheckpointStore(tmp_path), batch_size=1, lane=BufferLane(0, 2), +def test_only_the_publishers_lane_merges_shards(tmp_path): + assert merges_shards(None) and merges_shards(BufferLane(0, 2)) and not merges_shards(BufferLane(1, 2)) + with pytest.raises(SampleBufferError, match="only the publisher's lane merges"): + FileSampleBuffer(CheckpointStore(tmp_path), batch_size=1, lane=BufferLane(1, 2), merge_shards=True) +# ---------------------------------------------------------------- object stores (S3, OSS) +def _s3_store(tmp_path, s3): + from yeto.cloud import s3_store + from yeto.rl.engine.resume import STORE_IMPL_ENV, store_for + + env = {STORE_IMPL_ENV: s3_store.STORE_IMPL, s3_store.URI_ENV: "s3://bkt/runs/a"} + return store_for(tmp_path / "mount", environ=env, client=s3) + + +def test_lane_readers_and_cursors_on_an_s3_store_use_the_object_api(tmp_path): + from test_rl_object_store import FakeS3 + + s3 = FakeS3() + _fill(FileSampleBuffer(_s3_store(tmp_path, s3), batch_size=1), 6) + + def reader(m): + buffer = open_trainer_buffer(_s3_store(tmp_path, s3), batch_size=1, lane=BufferLane(m, 2)) + return br.BatchReader(buffer, max_policy_age=2, cursor_log=lane_cursor_log(buffer)) + + first = [reader(m) for m in range(2)] + assert all(isinstance(r.cursor_log, ObjectLaneCursorLog) for r in first) + taken = [[p for p, _ in r.take(0, 2, current_version=0)[0]] for r in first] + assert taken == [[0, 2], [1, 3]] + keys = sorted(k for _, k in s3.objects if "/lanes/" in k) + assert keys == ["runs/a/sample-buffer/lanes/0-of-2/0000000000.json", + "runs/a/sample-buffer/lanes/1-of-2/0000000000.json"] + assert not (tmp_path / "mount").exists() # nothing through the mount + # a restarted lane reader continues after its last round (at most once) + again = reader(1) + assert again.buffer.read_cursor == 4 + assert [p for p, _ in again.take(1, 1, current_version=0)[0]] == [5] + assert again.cursor_log.last_cursor() == 6 + assert len([k for _, k in s3.objects if "/lanes/1-of-2/" in k]) == 2 + + +def test_on_an_oss_store_the_publisher_merges_and_the_other_lane_reads_the_merged_index(): + from test_rl_oss_store import _entry, _objects + + from yeto.rl.engine.sample_buffer import shard_index + + objects = _objects() + sampler = FileSampleBuffer(objects, batch_size=1, index=shard_index("s0")) + for i in range(4): + ref = sampler.write_payload(f"s{i}", br.encode_group([{"tokens": [i]}])) + sampler.append(SampleIndexEntry(f"s{i}", 0, ref, appended_at=float(i))) + other = open_trainer_buffer(objects, batch_size=1, lane=BufferLane(1, 2)) + assert other.select("trainer", 10) == [] # nothing merged yet: lane 1 does not merge + publisher = open_trainer_buffer(objects, batch_size=1, lane=BufferLane(0, 2)) + assert [s.position for s in publisher.select("trainer", 10)] == [0, 2] + assert [s.entry.sample_id for s in other.select("trainer", 10)] == ["s1", "s3"] + log = lane_cursor_log(other) + log.append(0, [1, 3], other.read_cursor) + assert lane_cursor_log(open_trainer_buffer(objects, batch_size=1, lane=BufferLane(1, 2))).last_cursor() == 4 + + def test_reader_continues_after_its_lane_cursor_and_never_retrains(tmp_path): store = CheckpointStore(tmp_path) writer = FileSampleBuffer(store, batch_size=1) @@ -212,13 +271,17 @@ def test_a_trainer_may_sync_a_sampler_may_not(): @pytest.mark.parametrize("kw,islands,message", [ ({}, 1, "at least two islands"), ({"tuning": "full"}, 2, "--tuning lora"), - ({"rl_checkpoint_store": "s3://b/p"}, 2, "object store"), ]) def test_synced_trainer_refusals(kw, islands, message): with pytest.raises(ValueError, match=message): sample_roles.check_synced_trainers(_role_args(**kw), islands) +@pytest.mark.parametrize("store", ["s3://b/p", "oss://b/p", "modal-volume://v/p"]) +def test_synced_trainers_take_an_object_store(store): + sample_roles.check_synced_trainers(_role_args(rl_checkpoint_store=store), 2) + + def test_buffer_lane_of_island_args(): args = _role_args(learner_id=1, rl_sample_lanes=2) assert sample_roles.buffer_lane(args) == BufferLane(1, 2) diff --git a/tests/test_rl_verl_buffer_trainer.py b/tests/test_rl_verl_buffer_trainer.py index 37e3c71cc..58dea8247 100644 --- a/tests/test_rl_verl_buffer_trainer.py +++ b/tests/test_rl_verl_buffer_trainer.py @@ -350,9 +350,122 @@ def test_island_entry_plans_a_buffer_trainer_and_refuses_other_roles(): assert sample_buffer_section(_entry_args()) is None section = sample_buffer_section(_entry_args("--rl-sample-role", "trainer", "--rl-resume-store", "/s", + "--rl-single-island-no-sync", "--rl-max-policy-age", "1", "--rl-sample-wait-s", "60")) - assert section == {"store": "/s", "wait_s": 60.0} + assert section == {"store": "/s", "wait_s": 60.0, "lane": None} with pytest.raises(SystemExit, match="engine sampler"): sample_buffer_section(_entry_args("--rl-sample-role", "sampler", "--rl-resume-store", "/s")) with pytest.raises(SystemExit, match="max-policy-age"): sample_buffer_section(_entry_args("--rl-sample-role", "trainer", "--rl-resume-store", "/s")) + + +def test_island_entry_plans_a_synced_trainer_lane_and_refuses_partial_flags(): + from yeto.rl.adapters.verl.island_entry import sample_buffer_section + + synced = ("--rl-sample-role", "trainer", "--syncer", "h:1", "--rl-max-policy-age", "2") + section = sample_buffer_section(_entry_args(*synced, "--rl-sample-store", "/s", "--rl-sample-lanes", "2")) + assert section == {"store": "/s", "wait_s": None, "lane": "0/2"} + for extra, message in ((("--rl-sample-store", "/s"), "needs --rl-sample-store and --rl-sample-lanes"), + (("--rl-sample-store", "/s", "--rl-sample-lanes", "2", "--rl-resume-store", "/s"), + "not --rl-resume-store")): + with pytest.raises(SystemExit, match=message): + sample_buffer_section(_entry_args(*synced, *extra)) + with pytest.raises(SystemExit, match="need --rl-sample-role trainer with a syncer"): + sample_buffer_section(_entry_args("--rl-sample-role", "trainer", "--rl-single-island-no-sync", + "--rl-max-policy-age", "1", "--rl-resume-store", "/s", + "--rl-sample-lanes", "2")) + + +# ---------------------------------------------------------------------- synced trainers (S25 C3) +class _State: + """A published policy state: local version, hash, the LoRA tensors.""" + + def __init__(self, version, tensors, label): + import hashlib + + self.policy_version, self._tensors = version, tensors + self.label, self._digest = label, hashlib.sha256(label.encode()).hexdigest() + + def policy_tensor_hash(self): + return self._digest + + def to_lora(self): + return type("L", (), {"tensors": self._tensors}) + + +def _synced_publisher(tmp_path, store, lane, events): + from yeto.rl.adapters.verl.fully_async_ports import FullyAsyncIsland + from yeto.rl.engine.merged_release import BufferLane, lane_cursor_log + + actor = FakeBufferActor() + island = FullyAsyncIsland(actor, learner_id=lane, limit=2, required=1, + out_dir=str(tmp_path / f"out{lane}"), readback_dir=str(tmp_path / "rb")) + (tmp_path / f"out{lane}").mkdir(exist_ok=True) + island.emit = lambda event, **f: events.append((event, f)) + buffer = FileSampleBuffer(store, batch_size=1, lane=BufferLane(lane, 2)) + _, publisher, _ = install_buffer_ports( + island, buffer=buffer, releases=PolicyReleaseStore(store), shape=SHAPE, base_dir=tmp_path, + cursor_log=lane_cursor_log(buffer)) + def merge(base, tensors, directory): + (directory / "model.safetensors").write_bytes(b"w") + return {} + + publisher.merge = merge + driver = type("D", (), {"outer_version": None})() + publisher.bind(driver) + return actor, island, publisher, driver + + +def test_only_the_publisher_releases_each_new_merged_version_once(tmp_path): + store = CheckpointStore(tmp_path / "store") + nonzero = {NAMES[0]: torch.ones(2, 4), NAMES[1]: torch.ones(4, 2)} + zero = {NAMES[0]: torch.ones(2, 4), NAMES[1]: torch.zeros(4, 2)} + events = {0: [], 1: []} + islands = {m: _synced_publisher(tmp_path, store, m, events[m]) for m in (0, 1)} + # (local version, outer version of the applied base): base, a local round without a + # merge, merges 1 and 2, then merge 4 (merge 3 missed: a gap) + steps = [(0, 0, zero), (1, 0, nonzero), (2, 1, nonzero), (3, 2, nonzero), (4, 4, nonzero)] + for local, outer, tensors in steps: + for m, (actor, island, publisher, driver) in islands.items(): + actor.param_version = local + driver.outer_version = outer + publisher.publish(_State(local, tensors, f"h{outer}-{local}")) + released = [f for e, f in events[0] if e == "rl_policy_released"] + assert [(f["policy_version"], f["local_version"]) for f in released] == [(1, 2), (2, 3), (4, 4)] + assert all(f["origin"] == {"kind": "merged", "learner_id": 0, "num_learners": 2, + "local_version": f["local_version"]} for f in released) + assert [f for e, f in events[0] if e == "rl_policy_release_gap"] == [{"previous_version": 2, "policy_version": 4}] + assert not [e for e, _ in events[1] if e == "rl_policy_released"] + assert ("rl_policy_release_role", {"learner_id": 1, "num_learners": 2, "publisher": False}) in events[1] + latest = PolicyReleaseStore(store).latest() + assert (latest.seq, latest.version, latest.policy_hash) == (3, 4, _State(4, {}, "h4-4").policy_tensor_hash()) + # both islands pair each merged version with the first local version that applied it + for _, island, publisher, _ in islands.values(): + assert publisher.feed.merged.to_dict() == {"0": 0, "1": 2, "2": 3, "4": 4} + + +def test_a_synced_feed_reads_its_lane_translates_merged_versions_and_logs_each_round(tmp_path): + store = CheckpointStore(tmp_path / "store") + writer = GroupWriter(FileSampleBuffer(store, batch_size=1), "s0") + for i, version in enumerate((0, 0, 1, 1, 3, 3)): # merged behavior versions + writer.write([_sample(version, index=0)], group_key=f"g{i}") + events = [] + actor, island, publisher, driver = _synced_publisher(tmp_path, store, 1, events) + for local, outer in ((0, 0), (1, 1), (2, 1)): # merge 1 applied at local 1; local 2 trains on + actor.param_version = local + driver.outer_version = outer + publisher.publish(_State(local, {NAMES[1]: torch.zeros(4, 2)}, f"h{local}")) + feed = publisher.feed + pulled = [] + for _ in range(3): + meta = feed.pull() + pulled.append((meta["uid"], meta["min_steps"])) + feed.keep() if None not in meta["min_steps"] else feed.drop() + # lane 1 of 2: positions 1, 3, 5; merged v0 -> local 0, v1 -> local 1, v3 never applied here + assert pulled == [("buf1", (0,)), ("buf3", (1,)), ("buf5", (None,))] + island.call("yeto_assemble") # the round's collection ended: one lane cursor line + lane_file = tmp_path / "store" / "sample-buffer" / "lanes" / "1-of-2.jsonl" + assert _lines(lane_file) == [{"rollout_id": 2, "positions": [1, 3], "cursor": 6}] + # a restarted island continues after the recorded cursor + again = _synced_publisher(tmp_path, store, 1, [])[2].feed + assert again.buffer.read_cursor == 6 diff --git a/tests/test_rl_verl_synced_trainers.py b/tests/test_rl_verl_synced_trainers.py new file mode 100644 index 000000000..85547c6db --- /dev/null +++ b/tests/test_rl_verl_synced_trainers.py @@ -0,0 +1,186 @@ +"""S25 C3V end to end on CPU: two synced verl buffer trainers, the real Rust +syncer (elastic OuterStep merge), one engine sampler island that follows the +merged releases. + +The trainers are the verl buffer trainer of test_rl_verl_buffer_trainer.py +(fake trainer actor; real driver, fully_async rollout port, feed, publisher, +merged checkpoint) over one lane each of one buffer, with ``ElasticAvgSync``. +The sampler is the real engine sampler island (fake engine) in a thread. + +Checks (the same as the GPU run's, SESSION25/gpu-review/T-C3V.md): +1. both trainers apply the same merged policy hash at every outer version; +2. the sampler loads only versions the publisher released, in order, and each + released hash is the merged hash of that outer version; learner 1 releases + nothing; +3. the two lanes are disjoint and no buffer position is trained twice; +4. the syncer ran R OuterSteps. + +Runs when YETO_TEST_ELASTIC_SYNCER points to a syncer binary (as +test_rl_synced_trainers_sampler.py).""" + +from __future__ import annotations + +import json +import os +import subprocess +import threading +import time +from pathlib import Path + +import pytest +import torch + +from test_rl_inter_island_ports import _free_port +from test_rl_verl_buffer_trainer import NAMES, FakeBufferActor, _base +from yeto.rl.adapters.sampler_engines.fake import FakeEngine +from yeto.rl.adapters.verl import buffer_feed as bf +from yeto.rl.adapters.verl import fully_async_round as far +from yeto.rl.adapters.verl.buffer_ports import RELEASE_LOG, install_buffer_ports +from yeto.rl.engine.algorithm import AlgorithmSpec, CorrectionSpec, ExecutionSpec +from yeto.rl.engine.engine_sampler import EngineSamplerConfig, PromptItem, build_engine_sampler +from yeto.rl.engine.merged_release import BufferLane, lane_cursor_log +from yeto.rl.engine.policy_release import PolicyReleaseStore +from yeto.rl.engine.resume import CheckpointStore +from yeto.rl.engine.sample_buffer import FileSampleBuffer, open_trainer_buffer +from yeto.rl.engine.sampler_checks import PostLoadChecks +from yeto.rl.engine.sampler_engine import SamplingParams + +ROUNDS = 3 +GROUPS = 2 +MAX_POLICY_AGE = 2 +TOLERANT = AlgorithmSpec(execution=ExecutionSpec(max_policy_staleness=MAX_POLICY_AGE), + correction=CorrectionSpec(method="tis", tis_clip=2.0, tis_clip_low=0.0)) + + +def _lines(path: Path) -> list[dict]: + return [json.loads(line) for line in path.read_text().splitlines()] if path.is_file() else [] + + +def _island(out: Path, store_root: Path, base: Path, learner_id: int, port: int, key: str): + from yeto.rl.adapters.verl.entry import verl_capabilities + from yeto.rl.adapters.verl.fully_async_ports import FullyAsyncIsland, FullyAsyncTrainerGroup + from yeto.rl.adapters.verl.ports_impl import VerlPolicyState + from yeto.rl.bridge import BridgeConfig + from yeto.rl.core import CanonicalTensorSpec, canonical_layout_hash + from yeto.rl.elastic_client import ElasticClientConfig, ElasticIslandClient + from yeto.rl.engine.bridges import ElasticAvgSync + from yeto.rl.engine.driver import EventTape, IslandDriver + + specs = (CanonicalTensorSpec(NAMES[0], (2, 4), "float32", 8), + CanonicalTensorSpec(NAMES[1], (4, 2), "float32", 8)) + actor = FakeBufferActor() + island = FullyAsyncIsland(actor, learner_id=learner_id, limit=MAX_POLICY_AGE, required=GROUPS, + out_dir=str(out / f"out{learner_id}"), readback_dir=str(out / f"rb{learner_id}"), + base_model_revision="a" * 40, lora_config_hash="b" * 64, + layout_hash=canonical_layout_hash(specs), expected_specs=specs) + (out / f"out{learner_id}").mkdir(exist_ok=True) + tape = EventTape(out / f"tape-{learner_id}.jsonl", learner_id) + island.emit = lambda event, **f: tape.append({"event": event, **f}) + store = CheckpointStore(store_root) + buffer = open_trainer_buffer(store, batch_size=1, create_index=False, lane=BufferLane(learner_id, 2)) + rollout, publisher, placement = install_buffer_ports( + island, buffer=buffer, releases=PolicyReleaseStore(store), shape=bf.RowShape(64, 16, 0), + base_dir=base, tokenize=lambda text: [7, 8], poll_s=0.01, wait_s=30, + cursor_log=lane_cursor_log(buffer)) + bridge = BridgeConfig(syncer_addr=("127.0.0.1", port), learner_id=learner_id, global_rounds=ROUNDS, + groups_per_round=GROUPS, samples_per_group=2, local_optimizer_steps=1, + expected_specs=specs, base_model_revision="a" * 40, lora_config_hash="b" * 64, + layout_hash=canonical_layout_hash(specs), event_tape=str(out / f"bridge-{learner_id}.jsonl"), + wan_streams=0) + client = ElasticIslandClient(ElasticClientConfig(("127.0.0.1", port), learner_id, lease_s=6.0, + syncer_epoch=5), key.encode()) + profile = far.execution_profile_for(TOLERANT, MAX_POLICY_AGE, groups_per_round=GROUPS, samples_per_group=2, + sync="elastic") + driver = IslandDriver( + learner_id=learner_id, rollout=rollout, trainer=FullyAsyncTrainerGroup(island), + policy_state=VerlPolicyState(island), publisher=publisher, placement=placement, + capabilities=verl_capabilities("sha256:" + "0" * 64, fully_async_limit=MAX_POLICY_AGE), + algorithm=TOLERANT, sync=ElasticAvgSync(bridge, client=client, base_wait_s=30.0, syncer_epoch=5), + events=tape, profile=profile) + publisher.bind(driver) + return driver, island + + +@pytest.mark.skipif(not os.environ.get("YETO_TEST_ELASTIC_SYNCER"), + reason="needs YETO_TEST_ELASTIC_SYNCER (syncer binary with elastic mode)") +def test_two_synced_verl_trainers_feed_one_engine_sampler(tmp_path): + port, key = _free_port(), "c3v-e2e" + out = tmp_path + store_root = out / "store" + store = CheckpointStore(store_root) + base, _ = _base(tmp_path, dtype=torch.float32) + syncer = subprocess.Popen( + [os.environ["YETO_TEST_ELASTIC_SYNCER"], "--port", str(port), "--learners", "2", + "--total-steps", str(ROUNDS), "--event-tape", str(out / "syncer-tape.jsonl"), + "--island-scheduling-mode", "elastic", "--quorum-theta", "0.75", "--carry-gamma", "0.5", + "--soft-deadline-s", "30", "--q-min", "1", "--max-carry-lag", "2", + "--island-lease-s", "6", "--syncer-epoch", "5"], + env=dict(os.environ, YETO_ISLAND_HMAC_KEY=key), + stdout=open(out / "syncer.log", "w"), stderr=subprocess.STDOUT) + done, sampler_events = threading.Event(), [] + try: + time.sleep(1.0) + config = EngineSamplerConfig(sampler_id="s0", num_rollouts=200, max_backlog_groups=4 * GROUPS, + sampling=SamplingParams(max_tokens=6, n=2), max_prompt_tokens=32, + poll_s=0.01) + sampler = build_engine_sampler( + config, engine=FakeEngine(max_model_len=64), releases=PolicyReleaseStore(store), + buffer=FileSampleBuffer(store, batch_size=1), emit=lambda n, **f: sampler_events.append((n, f)), + prompts=lambda r: [PromptItem(f"p{r}", f"prompt {r}", label="x")], + reward=lambda item, c: float(len(c.token_ids) % 2), + checks=PostLoadChecks(expected_expert_matrices=None, fixed_prompts=()), # N2-fix: no probe pairs + should_stop=done.is_set, sleep=lambda s: time.sleep(0.005)) + sampler_thread = threading.Thread(target=sampler.run, daemon=True) + sampler_thread.start() + islands = [_island(out, store_root, base, i, port, key) for i in range(2)] + results, errors = {}, {} + + def run(i, driver): + try: + results[i] = driver.run() + except BaseException as error: # noqa: BLE001 + errors[i] = error + + threads = [threading.Thread(target=run, args=(i, d), daemon=True) for i, (d, _) in enumerate(islands)] + for t in threads: + t.start() + for t in threads: + t.join(timeout=90) + assert not t.is_alive() + assert errors == {} + assert syncer.wait(15) == 0 + finally: + done.set() + if syncer.poll() is None: + syncer.kill() + + tapes = [_lines(out / f"tape-{i}.jsonl") for i in range(2)] + # 1. the same merged policy on both trainers at every outer version + applied = [{e["sync/outer_version"]: e["sync/global_policy_hash"] for e in t if e["event"] == "rl_policy_apply"} + for t in tapes] + assert sorted(applied[0]) == list(range(ROUNDS + 1)) and applied[0] == applied[1] + # 2. releases: the publisher only, merged versions increasing, each with the merged hash + released = [e for e in tapes[0] if e["event"] == "rl_policy_released"] + assert [e["policy_version"] for e in released] == list(range(1, ROUNDS + 1)) # v0 is the base + assert all(e["policy_hash"] == applied[0][e["policy_version"]] for e in released) + assert all(e["origin"]["learner_id"] == 0 for e in released) + assert not [e for e in tapes[1] if e["event"] == "rl_policy_released"] + assert not [e for t in tapes for e in t if e["event"] == "rl_policy_release_gap"] + loaded = [f["policy_version"] for n, f in sampler_events if n == "rl_sampler_policy_loaded"] + assert loaded == sorted(set(loaded)) and set(loaded) <= {e["policy_version"] for e in released} + latest = PolicyReleaseStore(store).latest() + assert latest.version == ROUNDS and latest.origin["learner_id"] == 0 + # 3. disjoint lanes, no position trained twice, lane cursor lines per round + lanes = [_lines(store_root / "sample-buffer" / "lanes" / f"{i}-of-2.jsonl") for i in range(2)] + positions = [[p for line in lane for p in line["positions"]] for lane in lanes] + assert all(p % 2 == i for i, ps in enumerate(positions) for p in ps) + flat = positions[0] + positions[1] + assert len(flat) == len(set(flat)) == 2 * ROUNDS * GROUPS + assert [len(lane) for lane in lanes] == [ROUNDS, ROUNDS] + # the islands paired merged versions with local ones identically (one merge per round) + for i in range(2): + records = _lines(out / f"out{i}" / RELEASE_LOG.format(learner=i)) + assert records[-1]["merged_versions"] == {str(v): v for v in range(ROUNDS + 1)} + # 4. R OuterSteps + kinds = [json.loads(line).get("kind") for line in (out / "syncer-tape.jsonl").read_text().splitlines()] + assert kinds.count("outer_step") == ROUNDS diff --git a/yeto/rl/adapters/miles/buffer_rollout.py b/yeto/rl/adapters/miles/buffer_rollout.py index 6b2ae54f1..794837fcd 100644 --- a/yeto/rl/adapters/miles/buffer_rollout.py +++ b/yeto/rl/adapters/miles/buffer_rollout.py @@ -60,7 +60,15 @@ encode_group, group_entry_fields, ) -from yeto.rl.engine.merged_release import LANE_ENV, BufferLane, LaneCursorLog, gate_version +from yeto.rl.engine.merged_release import ( + LANE_ENV, + BufferLane, + LaneCursorLog, + ObjectLaneCursorLog, + gate_version, + lane_cursor_log, + merges_shards, +) SAMPLER, TRAINER, CONVERTIBLE = "sampler", "trainer", "convertible" ROLES = (SAMPLER, TRAINER, CONVERTIBLE) @@ -170,19 +178,13 @@ def open_buffer(environ: Mapping[str, str] | None = None, *, batch_size: int = 1 lane = BufferLane.parse(env.get(LANE_ENV)) if reader else None return FileSampleBuffer(store, batch_size=batch_size, index=INDEX if shard is None else shard_index(shard), - create_index=create_index, merge_shards=merge and lane is None, + create_index=create_index, merge_shards=merge and merges_shards(lane), lane=lane) -def open_lane_log(buffer: Any) -> LaneCursorLog | None: +def open_lane_log(buffer: Any) -> LaneCursorLog | ObjectLaneCursorLog | None: """The lane cursor log of a lane reader (None for a single reader).""" - lane = getattr(buffer, "lane", None) - if lane is None: - return None - from yeto.rl.engine.sample_buffer import BUFFER_DIR - - store = buffer.store - return LaneCursorLog(f"{store.root}/{BUFFER_DIR}", lane, commit=store.commit) + return lane_cursor_log(buffer) def renumber_group(samples: Sequence[dict[str, Any]], *, position: int, @@ -236,7 +238,7 @@ def __init__(self, buffer: Any, *, max_policy_age: int, wait_s: float = DEFAULT_ poll_s: float = DEFAULT_POLL_S, consumer: str = TRAINER_CONSUMER, clock: Callable[[], float] = time.monotonic, sleep: Callable[[float], None] = time.sleep, - cursor_log: LaneCursorLog | None = None) -> None: + cursor_log: LaneCursorLog | ObjectLaneCursorLog | None = None) -> None: """``cursor_log``: a synced trainer's lane cursor (:mod:`yeto.rl.engine.merged_release`): the reader continues after its last line and appends one per round.""" if int(max_policy_age) < 1: diff --git a/yeto/rl/adapters/miles/island_entry.py b/yeto/rl/adapters/miles/island_entry.py index 73d27d2cd..ed13592a6 100644 --- a/yeto/rl/adapters/miles/island_entry.py +++ b/yeto/rl/adapters/miles/island_entry.py @@ -212,22 +212,10 @@ def _check_ports_infra_switches(args) -> None: def _check_sample_store(args) -> None: - """S25 C3: ``--rl-sample-store`` / ``--rl-sample-lanes`` belong to a synced - trainer island (``yeto.rl.engine.merged_release``).""" - from yeto.rl.engine.sample_roles import buffer_lane, is_synced_trainer - - store, lanes = getattr(args, "rl_sample_store", None), getattr(args, "rl_sample_lanes", None) - if not is_synced_trainer(args): - if store is not None or lanes is not None: - raise ValueError("--rl-sample-store / --rl-sample-lanes need --rl-sample-role trainer " - "with a syncer (no --rl-single-island-no-sync)") - return - if not store or lanes is None or int(lanes) < 2: - raise ValueError("a synced trainer island needs --rl-sample-store and --rl-sample-lanes " - ">= 2 (the launcher forwards both)") - if getattr(args, "rl_resume_store", None): - raise ValueError("a synced trainer island takes --rl-sample-store, not --rl-resume-store") - buffer_lane(args) # learner id inside 0..lanes-1 + """S25 C3: ``--rl-sample-store`` / ``--rl-sample-lanes`` (engine-neutral rule).""" + from yeto.rl.engine.sample_roles import check_sample_store + + check_sample_store(args) def _check_resume_switches(args) -> None: diff --git a/yeto/rl/adapters/sampler_engines/sglang_server.py b/yeto/rl/adapters/sampler_engines/sglang_server.py index 1be20274e..36028b84c 100644 --- a/yeto/rl/adapters/sampler_engines/sglang_server.py +++ b/yeto/rl/adapters/sampler_engines/sglang_server.py @@ -31,6 +31,7 @@ from __future__ import annotations import json +import re from collections.abc import Sequence from dataclasses import dataclass from pathlib import Path @@ -49,7 +50,17 @@ TOPOLOGY_TP8 = "tp8" TOPOLOGY_DP8 = "dp8" -SGLANG_TOPOLOGIES = (TOPOLOGY_TP8, TOPOLOGY_DP8) +_TOPOLOGY = re.compile(r"(tp|dp)([1-9][0-9]*)") + + +def sglang_topology(tp: int, *, dp_attention: bool = False) -> str: + """The topology label of one SGLang instance over ``tp`` cards: ``tp``, or + ``dp`` with DP attention (``--dp N --enable-dp-attention``).""" + return f"{'dp' if dp_attention else 'tp'}{int(tp)}" + + +def is_dp_attention(topology: str) -> bool: + return topology.startswith("dp") @dataclass @@ -80,11 +91,16 @@ class SGLangServerEngine(ServerEngine): def __init__(self, config: ServerConfig, options: SGLangOptions | None = None, **kw: Any) -> None: super().__init__(config, **kw) self.options = options or SGLangOptions() - if config.topology not in SGLANG_TOPOLOGIES: - raise SamplerEngineError(f"sglang topology {config.topology!r} is not supported " - f"(known: {SGLANG_TOPOLOGIES}; PD disaggregation is not used)") - if config.topology == TOPOLOGY_DP8 and self.options.nextn: - raise SamplerEngineError("sglang dp8 + DeepEP cannot run NEXTN (A-INF item 8); " + # The topology is the source label of every sample: it must name the real + # layout, tp or dp with N = the instance's card count (options.tp). + match = _TOPOLOGY.fullmatch(config.topology) + if match is None or int(match.group(2)) != int(self.options.tp): + raise SamplerEngineError( + f"sglang topology {config.topology!r} does not match TP {self.options.tp}: expected " + f"{sglang_topology(self.options.tp)!r} or {sglang_topology(self.options.tp, dp_attention=True)!r} " + "(PD disaggregation is not used)") + if is_dp_attention(config.topology) and self.options.nextn: + raise SamplerEngineError("sglang DP attention + DeepEP cannot run NEXTN (A-INF item 8); " "turn one of them off") def argv(self, model_dir: Path) -> list[str]: @@ -93,7 +109,7 @@ def argv(self, model_dir: Path) -> list[str]: "--host", c.host, "--port", str(c.port), "--tp", str(o.tp), "--context-length", str(c.max_model_len), "--served-model-name", o.served_model_name, "--trust-remote-code"] - if c.topology == TOPOLOGY_DP8: + if is_dp_attention(c.topology): argv += ["--dp", str(o.tp), "--enable-dp-attention", "--moe-a2a-backend", "deepep"] if o.nextn: argv += ["--speculative-algorithm", "NEXTN"] diff --git a/yeto/rl/adapters/verl/buffer_feed.py b/yeto/rl/adapters/verl/buffer_feed.py index e93e7854a..cac0808a8 100644 --- a/yeto/rl/adapters/verl/buffer_feed.py +++ b/yeto/rl/adapters/verl/buffer_feed.py @@ -25,7 +25,20 @@ Every pulled group is recorded once (``record``) with its buffer position, so a run can show that no group was consumed twice. The buffer positions are those of -the trainer's single read cursor (``FileSampleBuffer.select``). +the trainer's read cursor (``FileSampleBuffer.select``). + +Synced trainers (S25 C3, ``yeto.rl.engine.merged_release``): the buffer is a +lane reader (positions ``p % N == m`` only) and the feed keeps the lane cursor +(``cursor_log``): it continues after the last recorded round, and +:meth:`BufferFeed.close_round` records the round's trained positions when its +collection ends, before the round trains (at most once, as the Miles reader). +The groups carry the *merged* version as behavior version; +:class:`MergedVersions` turns it into the island's local policy version (the +first local version that applied that merge, whose weights are the merged +policy), and the version map turns that into verl's ``param_version``. So the +``max_policy_age`` judgement counts local versions: equal to merged versions +when every round merges, stricter when an island trains several rounds per +merge (elastic sync). :func:`group_rows` builds the rows the way verl's agent loop postprocess pads them (``AgentLoopWorker._postprocess``): prompts left-padded to the prompt @@ -60,6 +73,32 @@ class BufferFeedError(SampleBufferRolloutError): """A buffer group cannot become verl rows.""" +class MergedVersions: + """Synced trainers: merged (syncer outer) version -> the first local policy + version of this island that applied it (module doc). Merged versions only + increase; a local version pairs with one merged version.""" + + def __init__(self) -> None: + self._local: dict[int, int] = {} + + def record(self, merged: int, local: int) -> bool: + """Record one publication; True when ``merged`` is new on this island.""" + merged, local = int(merged), int(local) + if merged in self._local: + return False + if self._local and (merged < max(self._local) or local <= self._local[max(self._local)]): + raise BufferFeedError(f"merged v{merged} at local v{local} after merged v{max(self._local)} " + f"at local v{self._local[max(self._local)]}: versions must increase") + self._local[merged] = local + return True + + def local(self, merged: int) -> int | None: + return self._local.get(int(merged)) + + def to_dict(self) -> dict[str, int]: + return {str(k): v for k, v in sorted(self._local.items())} + + @dataclass(frozen=True) class RowShape: """The padded widths of the trainer's rows (``data.max_prompt_length`` and @@ -167,7 +206,11 @@ def __init__(self, buffer: Any, *, vmap: Callable[[], VersionMap], shape: RowSha tokenize: Callable[[str], Sequence[int]] | None = None, wait_s: float = 3600.0, poll_s: float = 2.0, consumer: str = TRAINER_CONSUMER, clock: Callable[[], float] = time.monotonic, - sleep: Callable[[float], None] = time.sleep) -> None: + sleep: Callable[[float], None] = time.sleep, + cursor_log: Any = None, merged: MergedVersions | None = None) -> None: + """``cursor_log``: a synced trainer's lane cursor log + (``merged_release.lane_cursor_log``); ``merged``: a synced trainer's + merged version pairing (module doc).""" self.buffer = buffer self.vmap = vmap self.shape = shape @@ -181,6 +224,11 @@ def __init__(self, buffer: Any, *, vmap: Callable[[], VersionMap], shape: RowSha self.sleep = sleep self._last: _Pulled | None = None self.pulled_positions: list[int] = [] + self.cursor_log = cursor_log + self.merged = merged + self._round_kept: list[int] = [] # trained positions of the round being collected + if cursor_log is not None: + self.buffer.reset_to(cursor_log.last_cursor()) def __call__(self, name: str, *args: Any) -> Any: """``call(name)`` for the three queue methods (:data:`QUEUE_METHODS`).""" @@ -207,6 +255,11 @@ def _next(self) -> Any | None: self.sleep(self.poll_s) def _local(self, version: int) -> int | None: + if self.merged is not None: + local = self.merged.local(version) + if local is None: # a merge this island never applied: unknown + return None + version = local try: return self.vmap().local(version) except VerlTranslationError: @@ -243,6 +296,7 @@ def keep(self) -> None: self.add_group(group_rows(last.samples, uid=f"{UID_PREFIX}{last.position}", param_versions=last.param_versions, shape=self.shape, tokenize=self.tokenize)) + self._round_kept.append(last.position) self.record({**last.record, "kept": True}) def drop(self) -> None: @@ -254,21 +308,39 @@ def _take_last(self) -> _Pulled: last, self._last = self._last, None return last + def close_round(self, rollout_id: int) -> None: + """The round's collection ended (before it trains): record its trained + positions and the read position in the lane cursor log (module doc).""" + if self._last is not None: + raise BufferFeedError("round closed with a pulled group not kept or dropped") + kept, self._round_kept = self._round_kept, [] + if self.cursor_log is not None: + self.cursor_log.append(int(rollout_id), kept, int(self.buffer.read_cursor)) + def read_cursor(self) -> int: """The buffer position after the last pulled group (released to the samplers for their backlog cap).""" return int(self.buffer.read_cursor) +ASSEMBLE_METHOD = "yeto_assemble" # the rollout port's call after a round's collection + + class BufferCall: """The island's ``call(method, *args)``: the queue calls go to the feed, every - other call to the trainer actor (``actor_call``).""" + other call to the trainer actor (``actor_call``). ``round_id()``: the round + being collected; its collection ends at the assemble call, where the feed + closes the round (:meth:`BufferFeed.close_round`).""" - def __init__(self, feed: BufferFeed, actor_call: Callable[..., Any]) -> None: + def __init__(self, feed: BufferFeed, actor_call: Callable[..., Any], + round_id: Callable[[], int] | None = None) -> None: self.feed = feed self.actor_call = actor_call + self.round_id = round_id def __call__(self, name: str, *args: Any, **kwargs: Any) -> Any: if name in QUEUE_METHODS: return self.feed(name, *args) + if name == ASSEMBLE_METHOD and self.round_id is not None: + self.feed.close_round(self.round_id()) return self.actor_call(name, *args, **kwargs) diff --git a/yeto/rl/adapters/verl/buffer_ports.py b/yeto/rl/adapters/verl/buffer_ports.py index e354ff119..e8a7ba10f 100644 --- a/yeto/rl/adapters/verl/buffer_ports.py +++ b/yeto/rl/adapters/verl/buffer_ports.py @@ -17,6 +17,18 @@ is not released: the samplers start on the base weights as version 0. A version-0 state with any non-zero B is released like any other version. +Synced trainers (S25 C3, ``yeto.rl.engine.merged_release``; the same rules as +the Miles trainer island): every island records, at each publication, the +merged version (``merged_version``: the syncer's outer version of the applied +base) it first applied (:class:`~.buffer_feed.MergedVersions`; the feed turns a +group's merged behavior version into a local one with it). Only the publisher +(learner id 0) releases, once per new merged version, with the merged version +as release version, kept increasing by ``MergedReleaseGuard`` (a skipped +merge is reported as ``rl_policy_release_gap``). A publication without a new +merge (elastic: several local rounds per merge) releases nothing: its weights +are not a merged policy. The driver is bound after it is built +(:meth:`BufferReleasePublisher.bind`). + The publication is complete when the release record and ``LATEST.json`` are in the store. Whether a sampler has loaded it is the sampler's report (``rl_sampler_policy_loaded`` on its tape) and the behavior versions of the @@ -34,10 +46,11 @@ from typing import Any from yeto.rl.contracts import InferencePublicationManifest +from yeto.rl.engine.merged_release import BufferLane, MergedReleaseGuard, is_publisher, merged_version from yeto.rl.engine.policy_release import FULL, PolicyReleaseStore from yeto.rl.engine.ports import PlacementDescription, PublicationResult -from .buffer_feed import BufferCall, BufferFeed, RowShape +from .buffer_feed import BufferCall, BufferFeed, MergedVersions, RowShape from .fully_async_ports import MEMBER, FullyAsyncIsland, FullyAsyncRolloutPool from .fully_async_round import FULLY_ASYNC_PLACEMENT from .merged_release import LORA_B, write_merged_checkpoint @@ -62,17 +75,59 @@ class BufferReleasePublisher: """yeto publication as a store release of the merged policy (module doc). ``feed`` gives the buffer read position; ``base_dir`` is the base HF - checkpoint the adapter trains over.""" + checkpoint the adapter trains over; ``lane``: a synced trainer's buffer lane + (module doc; None for a single trainer).""" def __init__(self, island: FullyAsyncIsland, releases: PolicyReleaseStore, feed: BufferFeed, *, base_dir: str | Path, clock: Callable[[], float] = time.monotonic, - merge: Callable[..., dict[str, Any]] = write_merged_checkpoint) -> None: + merge: Callable[..., dict[str, Any]] = write_merged_checkpoint, + lane: BufferLane | None = None) -> None: self.island = island self.releases = releases self.feed = feed self.base_dir = Path(base_dir) self.clock = clock self.merge = merge + self.lane = lane + self.driver: Any = None + self.guard = MergedReleaseGuard(self._on_gap) + self.releasing = lane is None or is_publisher(lane.index, lane.count) + + def bind(self, driver: Any) -> None: + """The island's driver (synced trainers read its outer version).""" + self.driver = driver + if self.lane is not None: + self.island.emit("rl_policy_release_role", learner_id=self.lane.index, + num_learners=self.lane.count, publisher=self.releasing) + + def _on_gap(self, previous: int, version: int) -> None: + self.island.emit("rl_policy_release_gap", previous_version=previous, policy_version=version) + + def _release_version(self, state: Any) -> int | None: + """The version to release, or None (nothing to release). Single trainer: + the policy version. Synced trainer: the merged version, on the publisher, + the first time this island applies it (module doc).""" + local = int(state.policy_version) + if self.lane is None: + return local + if self.driver is None: + raise RuntimeError("synced buffer trainer: publish before bind(driver)") + merged = merged_version(self.driver, state) + new = self.feed.merged.record(merged, local) + if not (new and self.releasing): + return None + latest = self.releases.latest() + if self.guard.check(merged, None if latest is None else latest.version) == "stale": + self.island.emit("rl_policy_release_skipped", policy_version=merged, local_version=local, + latest_version=latest.version, reason="not_newer_than_latest") + return None + return merged + + def _origin(self, state: Any) -> dict[str, Any] | None: + if self.lane is None: + return None + return {"kind": "merged", "learner_id": self.lane.index, "num_learners": self.lane.count, + "local_version": int(state.policy_version)} def publish(self, state: Any) -> PublicationResult: isl = self.island @@ -85,19 +140,25 @@ def publish(self, state: Any) -> PublicationResult: merge: dict[str, Any] = {} release = None position = self.feed.read_cursor() - if not is_base_adapter(version, tensors): + released = self._release_version(state) + if released is not None and not is_base_adapter(released, tensors): def write(directory: Path) -> None: merge.update(self.merge(self.base_dir, tensors, directory)) - release = self.releases.release(version, digest, write, buffer_selected=position, weights=FULL) + release = self.releases.release(released, digest, write, buffer_selected=position, + weights=FULL, origin=self._origin(state)) record = {"policy_version": version, "param_version": param_version, "policy_hash": digest, "released": release is not None, "buffer_selected": position, "seconds": round(self.clock() - started, 3), "merge": merge or None, "version_map": isl.vmap.to_dict()} + if self.lane is not None: + record.update(merged_version=merged_version(self.driver, state), + merged_versions=self.feed.merged.to_dict()) if release is not None: - record.update(seq=release.seq, path=release.path, + record.update(seq=release.seq, path=release.path, release_version=release.version, payload_bytes=sum(int(v) for v in (release.files or {}).values())) - isl.emit(RELEASE_EVENT, seq=release.seq, policy_version=version, policy_hash=digest, + isl.emit(RELEASE_EVENT, seq=release.seq, policy_version=release.version, + local_version=version, origin=release.origin, policy_hash=digest, buffer_selected=position, weights=FULL, sha256=release.sha256, engine_checksum=None, merged_modules=merge.get("merged_modules"), max_relative_change=merge.get("max_relative_change"), @@ -110,7 +171,7 @@ def write(directory: Path) -> None: publication_mode="full", base_policy_version=None, target_policy_version=version, target_policy_hash=digest, target_manifest_hash=_sha(f"{version}:{digest}"), payload_hash=_sha(json.dumps({"release": record.get("path"), "hash": digest})), - # the release's files, or (the base, not released) the adapter in bf16 as on the fully_async path + # the release's files, or (nothing released) the adapter in bf16 as on the fully_async path payload_bytes=int(record.get("payload_bytes") or sum(int(t.numel()) * 2 for t in tensors.values())), complete=True) return PublicationResult(manifest, frozenset({MEMBER})) @@ -127,12 +188,17 @@ def describe(self) -> PlacementDescription: def install_buffer_ports(island: FullyAsyncIsland, *, buffer: Any, releases: PolicyReleaseStore, shape: RowShape, base_dir: str | Path, tokenize: Callable[[str], Sequence[int]] | None = None, - wait_s: float = 3600.0, poll_s: float = 2.0, + wait_s: float = 3600.0, poll_s: float = 2.0, cursor_log: Any = None, **feed_kwargs: Any) -> tuple[FullyAsyncRolloutPool, BufferReleasePublisher, BufferPlacement]: """Route the island's queue calls to a buffer feed and return the rollout port, the publisher and the placement of a buffer trainer (module doc); the - feed is ``publisher.feed``.""" + feed is ``publisher.feed``. A lane reader ``buffer`` (``buffer.lane``) makes + a synced trainer: ``cursor_log`` is its lane cursor log + (``merged_release.lane_cursor_log``).""" + lane = getattr(buffer, "lane", None) + if (lane is None) != (cursor_log is None): + raise ValueError("a synced buffer trainer needs both a lane reader and its cursor log") actor_call = island.call log = CONSUMED_LOG.format(learner=island.learner_id) @@ -141,8 +207,9 @@ def record(fields: dict[str, Any]) -> None: feed = BufferFeed(buffer, vmap=lambda: island.vmap, shape=shape, add_group=lambda rows: actor_call("yeto_add_group", rows), record=record, - tokenize=tokenize, wait_s=wait_s, poll_s=poll_s, **feed_kwargs) - island.call = BufferCall(feed, actor_call) + tokenize=tokenize, wait_s=wait_s, poll_s=poll_s, cursor_log=cursor_log, + merged=None if lane is None else MergedVersions(), **feed_kwargs) + island.call = BufferCall(feed, actor_call, round_id=lambda: int(island.published[0])) island.start_rollouter = None # no rollouter: the samplers run elsewhere - publisher = BufferReleasePublisher(island, releases, feed, base_dir=base_dir) + publisher = BufferReleasePublisher(island, releases, feed, base_dir=base_dir, lane=lane) return FullyAsyncRolloutPool(island), publisher, BufferPlacement() diff --git a/yeto/rl/adapters/verl/buffer_trainer_runner.py b/yeto/rl/adapters/verl/buffer_trainer_runner.py index 595737b79..5601eee24 100644 --- a/yeto/rl/adapters/verl/buffer_trainer_runner.py +++ b/yeto/rl/adapters/verl/buffer_trainer_runner.py @@ -110,7 +110,8 @@ def run(self, config, plan: dict): from yeto.island_credential_guard import check_island_credentials from yeto.rl.engine.policy_release import PolicyReleaseStore from yeto.rl.engine.resume import store_for - from yeto.rl.engine.sample_buffer import FileSampleBuffer + from yeto.rl.engine.merged_release import BufferLane, lane_cursor_log + from yeto.rl.engine.sample_buffer import open_trainer_buffer from .buffer_feed import RowShape from .buffer_ports import install_buffer_ports @@ -127,7 +128,9 @@ def call(name, *args, timeout=None): environ = dict(os.environ) store = store_for(section["store"], environ=environ) - buffer = FileSampleBuffer(store, batch_size=1, create_index=False) + # S25 C3: a synced trainer reads its lane of the buffer (section "lane": "M/N") + buffer = open_trainer_buffer(store, batch_size=1, create_index=False, + lane=BufferLane.parse(section.get("lane"))) releases = PolicyReleaseStore(store_for(section.get("release_store") or section["store"], environ=environ)) shape = RowShape(int(config.data.max_prompt_length), int(config.data.max_response_length), @@ -137,7 +140,8 @@ def ports(island): return install_buffer_ports(island, buffer=buffer, releases=releases, shape=shape, base_dir=plan["model_path"], tokenize=tokenizer.encode, wait_s=float(section.get("wait_s") or 3600.0), - poll_s=float(section.get("poll_s") or 2.0)) + poll_s=float(section.get("poll_s") or 2.0), + cursor_log=lane_cursor_log(buffer)) return run_fully_async_island(call, None, plan, OmegaConf.to_container(config, resolve=True), buffer_ports=ports) diff --git a/yeto/rl/adapters/verl/island_entry.py b/yeto/rl/adapters/verl/island_entry.py index 010e909cb..e49c8c497 100644 --- a/yeto/rl/adapters/verl/island_entry.py +++ b/yeto/rl/adapters/verl/island_entry.py @@ -72,6 +72,10 @@ def parse_args(argv): p.add_argument("--rl-sample-role", default=None) p.add_argument("--rl-resume-store", default=None) p.add_argument("--rl-sample-wait-s", type=float, default=None) + # S25 C3: a synced trainer island (syncer, learner m of N) reads lane m of the + # buffer in this store; the launcher forwards both (yeto.rl.engine.merged_release) + p.add_argument("--rl-sample-store", default=None) + p.add_argument("--rl-sample-lanes", type=int, default=None) args, ignored = p.parse_known_args(argv) return args, ignored @@ -393,8 +397,13 @@ def sample_buffer_section(args) -> dict | None: "run the sampler as the engine sampler island (yeto.rl.engine.engine_sampler)") if int(args.rl_max_policy_age or 0) < 1: raise SystemExit("--rl-sample-role trainer needs --rl-max-policy-age >= 1 (the fully_async path)") + try: + sample_roles.check_sample_store(args) + lane = sample_roles.buffer_lane(args) + except ValueError as error: + raise SystemExit(str(error)) from None config = sample_roles.resolve(args, groups_per_round=args.groups_per_round) - return {"store": config.store, "wait_s": config.wait_s} + return {"store": config.store, "wait_s": config.wait_s, "lane": None if lane is None else lane.env()} def fully_async_plan(limit: int, overrides: list[str], spec_dict: dict | None, args, sync: str): diff --git a/yeto/rl/adapters/verl/trainer.py b/yeto/rl/adapters/verl/trainer.py index 5febc9718..82e4ddf59 100644 --- a/yeto/rl/adapters/verl/trainer.py +++ b/yeto/rl/adapters/verl/trainer.py @@ -339,6 +339,8 @@ def run_fully_async_island(call, start_rollouter, plan: dict, cfg: dict, *, buff policy_state=VerlPolicyState(island), publisher=publisher, placement=placement, capabilities=verl_capabilities(_fingerprint(cfg), fully_async_limit=limit), algorithm=spec, sync=sync, events=tape, profile=profile) + if buffer_ports is not None: + publisher.bind(driver) # synced buffer trainers release the driver's merged version state = driver.run() result = {"learner_id": learner_id, "rounds_completed": driver.rounds_completed, "final_policy_version": state.policy_version, diff --git a/yeto/rl/engine/merged_release.py b/yeto/rl/engine/merged_release.py index 0f953ff4a..a6376b758 100644 --- a/yeto/rl/engine/merged_release.py +++ b/yeto/rl/engine/merged_release.py @@ -34,13 +34,28 @@ Buffer: the trainers read disjoint *lanes* of the one buffer: trainer m of N trains the positions ``p`` with ``p % N == m`` (:class:`BufferLane`), so no -group is trained by two islands. The store must be a shared directory (local -disk, Modal Volume, a bucket mount that appends): a store read through an object -API (``CheckpointStore.objects``) is refused for now. A trainer keeps its lane -cursor in ``sample-buffer/lanes/-of-.jsonl`` (one line per round, written -when the round is selected), and a restarted or re-joined trainer continues -after it: a group selected by a round that did not finish is lost, never -trained twice (at most once). +group is trained by two islands. A trainer keeps its lane cursor, one record per +round, written when the round is selected (:func:`lane_cursor_log`), and a +restarted or re-joined trainer continues after it: a group selected by a round +that did not finish is lost, never trained twice (at most once). + +Stores: a shared directory (local disk, Modal Volume) or an object store read +through its object API (``CheckpointStore.objects``: S3, OSS). Each lane has one +writer, so every layout below has one writer per file or key: + +* lane cursor: ``sample-buffer/lanes/-of-.jsonl`` in a directory + (:class:`LaneCursorLog`, appended); ``sample-buffer/lanes/-of-/.json`` + in an object store (:class:`ObjectLaneCursorLog`, one new object per round: + an object store cannot append); +* buffer index: on a store whose samplers write their own shards + (``objects.buffer_shards``, OSS), the index has one writer, the merger + (``sample_buffer`` module doc). With synced trainers the publisher merges + (:func:`merges_shards`) and the other trainers only read the merged index; + so a lane whose publisher is stopped waits for the publisher's return; +* releases: written by the publisher only, through the store mount, in the + order of ``PolicyReleaseStore.release`` (weights, record, commit, then + ``LATEST.json``), so a sampler that reads ``LATEST.json`` by key sees a + complete release. Nothing here needs a cross-host lock or create-if-absent. Failures: @@ -96,8 +111,18 @@ def parse(cls, text: str | None) -> "BufferLane | None": lane = cls(int(index), int(count)) return None if lane.count == 1 else lane + def name(self) -> str: + return f"{self.index}-of-{self.count}" + def file_name(self) -> str: - return f"{self.index}-of-{self.count}.jsonl" + return f"{self.name()}.jsonl" + + +def merges_shards(lane: BufferLane | None) -> bool: + """Whether this trainer merges the sampler shards into the buffer index (on a + store that shards, ``sample_buffer.needs_shards``): a single trainer does, and + of synced trainers only the publisher (module doc: one index writer).""" + return lane is None or lane.index == PUBLISHER_LEARNER_ID def is_publisher(learner_id: int, num_learners: int) -> bool: @@ -143,7 +168,7 @@ def check(self, version: int, latest_version: int | None) -> str: class LaneCursorLog: - """The lane cursor of one synced trainer, durable in the store. + """The lane cursor of one synced trainer, durable in a shared-directory store. One JSON line per selected round: ``{"rollout_id", "positions", "cursor"}`` (``cursor``: the buffer read position after the round). :meth:`last_cursor` @@ -168,8 +193,64 @@ def last_cursor(self) -> int: def append(self, rollout_id: int, positions: list[int], cursor: int) -> None: self.path.parent.mkdir(parents=True, exist_ok=True) - line = json.dumps({"rollout_id": int(rollout_id), "positions": [int(p) for p in positions], - "cursor": int(cursor)}, sort_keys=True) + line = _record(rollout_id, positions, cursor) with self.path.open("a", encoding="utf-8") as fh: fh.write(line + "\n") self.commit() + + +def _record(rollout_id: int, positions: list[int], cursor: int) -> str: + return json.dumps({"rollout_id": int(rollout_id), "positions": [int(p) for p in positions], + "cursor": int(cursor)}, sort_keys=True) + + +class ObjectLaneCursorLog: + """The lane cursor of one synced trainer in an object store (module doc). + + Round ``k`` of this lane (counted over every start of the lane) is the object + ``.json`` with the record of :class:`LaneCursorLog`. One writer per + lane, so the next sequence number is the count of the listed objects; the + write is a plain put (no create-if-absent between hosts needed).""" + + def __init__(self, objects: Any, prefix: str, lane: BufferLane) -> None: + self.objects = objects + self.prefix = f"{prefix.rstrip('/')}/{LANES_DIR}/{lane.name()}/" + self._keys: list[str] | None = None + + def _listed(self) -> list[str]: + if self._keys is None: + self._keys = sorted(self.objects.list(self.prefix)) + return self._keys + + def last_cursor(self) -> int: + for key in reversed(self._listed()): + raw = self.objects.read(key) + try: + return int(json.loads(raw)["cursor"]) + except (TypeError, ValueError, KeyError): + continue # an unreadable object (should not happen: a put is atomic) + return 0 + + def append(self, rollout_id: int, positions: list[int], cursor: int) -> None: + keys = self._listed() + key = f"{self.prefix}{len(keys):010d}.json" + self.objects.write(key, _record(rollout_id, positions, cursor).encode()) + keys.append(key) + + +def lane_cursor_log(buffer: Any) -> LaneCursorLog | ObjectLaneCursorLog | None: + """The lane cursor log of a lane reader (``FileSampleBuffer`` with a lane; None + for a single reader): in the store's object API when it has one, else in the + store directory (committed after each line).""" + lane = getattr(buffer, "lane", None) + if lane is None: + return None + from .resume import CheckpointStore # noqa: PLC0415 + from .sample_buffer import BUFFER_DIR # noqa: PLC0415 - sample_buffer imports this module + + store = buffer.store # a CheckpointStore, or an ObjectStore at the store root + objects = store.objects if isinstance(store, CheckpointStore) else store + if objects is not None: + return ObjectLaneCursorLog(objects, BUFFER_DIR, lane) + return LaneCursorLog(f"{store.root}/{BUFFER_DIR}", lane, commit=store.commit) + diff --git a/yeto/rl/engine/sample_buffer.py b/yeto/rl/engine/sample_buffer.py index c9f652d78..11cf98728 100644 --- a/yeto/rl/engine/sample_buffer.py +++ b/yeto/rl/engine/sample_buffer.py @@ -68,7 +68,7 @@ from pathlib import Path from typing import Any -from .merged_release import BufferLane +from .merged_release import BufferLane, merges_shards from .object_store import ObjectStore from .resume import CheckpointStore from .trainer_lifecycle import BacklogStats, SampleIndexEntry, Selected @@ -314,6 +314,16 @@ def index_for(store: CheckpointStore | ObjectStore, *, index: str = INDEX, return ObjectIndex(objects, prefix=shard_prefix(shard)) +def open_trainer_buffer(store: CheckpointStore | ObjectStore, *, batch_size: int, + lane: BufferLane | None = None, create_index: bool = True, + **kwargs: Any) -> "FileSampleBuffer": + """The buffer of a trainer (the reader): it merges the sampler shards of a + store that shards (:func:`needs_shards`) unless it is a synced trainer other + than the publisher (``merged_release.merges_shards``); ``lane``: its lane.""" + return FileSampleBuffer(store, batch_size=batch_size, create_index=create_index, lane=lane, + merge_shards=needs_shards(store) and merges_shards(lane), **kwargs) + + class FileSampleBuffer: """:class:`~yeto.rl.engine.trainer_lifecycle.DurableSampleBuffer` over a store (a :class:`CheckpointStore`, or an :class:`ObjectStore` at the store root). @@ -331,10 +341,11 @@ def __init__(self, store: CheckpointStore | ObjectStore, *, batch_size: int, only one that does. ``lane``: this consumer is one of several synced trainers and selects only the positions of its lane (:mod:`.merged_release`); the read cursor and the watermark stay - positions of the whole index.""" - if lane is not None and (merge_shards or _objects_of(store) is not None): - raise SampleBufferError("a lane reader needs a shared-directory store (no object " - "index, no shard merge): see yeto.rl.engine.merged_release") + positions of the whole index. Of the lane readers only the publisher's + lane merges shards (``merged_release.merges_shards``).""" + if lane is not None and merge_shards and not merges_shards(lane): + raise SampleBufferError(f"buffer lane {lane.env()}: only the publisher's lane merges the " + "sampler shards (yeto.rl.engine.merged_release)") self.lane = lane if batch_size < 1: raise ValueError("batch_size must be positive") diff --git a/yeto/rl/engine/sample_roles.py b/yeto/rl/engine/sample_roles.py index 755a20e1f..8317437f0 100644 --- a/yeto/rl/engine/sample_roles.py +++ b/yeto/rl/engine/sample_roles.py @@ -102,11 +102,24 @@ def check_synced_trainers(args: Any, num_islands: int) -> None: "convertible nodes") if getattr(args, "tuning", "lora") != "lora": raise ValueError("synced trainers merge LoRA deltas: they need --tuning lora") - store = str(getattr(args, "rl_checkpoint_store", None) or "") - if store.startswith(("s3://", "oss://")): - raise ValueError(f"--rl-checkpoint-store {store}: synced trainers read buffer lanes from " - "a shared directory store (a Modal Volume or a shared path), not an " - "object store") + + +def check_sample_store(args: Any) -> None: + """Island-side check: ``--rl-sample-store`` / ``--rl-sample-lanes`` belong to a + synced trainer island, which needs both (the launcher forwards them) and no + ``--rl-resume-store``; its learner id must name one of the lanes.""" + store, lanes = getattr(args, "rl_sample_store", None), getattr(args, "rl_sample_lanes", None) + if not is_synced_trainer(args): + if store is not None or lanes is not None: + raise ValueError("--rl-sample-store / --rl-sample-lanes need --rl-sample-role trainer " + "with a syncer (no --rl-single-island-no-sync)") + return + if not store or lanes is None or int(lanes) < 2: + raise ValueError("a synced trainer island needs --rl-sample-store and --rl-sample-lanes " + ">= 2 (the launcher forwards both)") + if getattr(args, "rl_resume_store", None): + raise ValueError("a synced trainer island takes --rl-sample-store, not --rl-resume-store") + buffer_lane(args) # learner id inside 0..lanes-1 def node_name(island: int) -> str: diff --git a/yeto/rl/engine/sampler_checks.py b/yeto/rl/engine/sampler_checks.py index 36d702748..404dee244 100644 --- a/yeto/rl/engine/sampler_checks.py +++ b/yeto/rl/engine/sampler_checks.py @@ -204,6 +204,8 @@ def fixed_prompt_answers(self, engine: SamplerEngine) -> dict[str, Any]: and len(groups) == len(self.fixed_prompts), "answers": answers} def logprob_probe(self, engine: SamplerEngine) -> dict[str, Any]: + if not self.pairs(): # no fixed prompts and no probe pairs: nothing to score or compare + return {"name": "logprob_probe", "passed": True, "detail": "no probe pairs"} means = [] for prompt, continuation in self.pairs(): logprobs = engine.score(prompt, continuation) From e44933443a8e77cad2951b9a503ee63cc3a11ebd Mon Sep 17 00:00:00 2001 From: michael <78113554+chowhsu@users.noreply.github.com> Date: Tue, 20 Oct 2026 05:43:31 +0000 Subject: [PATCH 49/89] Register the qwen3_2507 chat template (Qwen3-4B-Instruct-2507) with token alignment tests Co-Authored-By: Claude Opus 5.5 --- tests/test_records_proxy.py | 20 +++++++++++++++----- yeto/rl/harness/chat_templates.py | 6 ++++++ 2 files changed, 21 insertions(+), 5 deletions(-) diff --git a/tests/test_records_proxy.py b/tests/test_records_proxy.py index 108daaf78..918793a75 100644 --- a/tests/test_records_proxy.py +++ b/tests/test_records_proxy.py @@ -58,6 +58,14 @@ def glm() -> ChatTemplate: return _template("glm5") +@pytest.fixture(scope="module") +def qwen3_2507() -> ChatTemplate: + return _template("qwen3_2507") + + +FIXTURES = {"qwen3": "qwen", "glm5": "glm", "qwen3_2507": "qwen3_2507"} + + def call(call_id: str, cmd: str) -> dict[str, Any]: return {"id": call_id, "type": "function", "function": {"name": "shell", "arguments": json.dumps({"cmd": cmd})}} @@ -154,13 +162,15 @@ def run_codex(proxy: Any, trajectory_id: str, *, turns: int = 3) -> list[dict[st # ---------------------------------------------------------------------- N10.5 records -@pytest.mark.parametrize("name", ["qwen3", "glm5"]) +@pytest.mark.parametrize("name", ["qwen3", "glm5", "qwen3_2507"]) def test_codex_responses_loop_records_one_aligned_stream(name, request): - template = request.getfixturevalue("qwen" if name == "qwen3" else "glm") + template = request.getfixturevalue(FIXTURES[name]) engine = engine_for(ScriptedModel(template, SCRIPT)) proxy = RecordsProxy(engine, template, model="policy", signed_sampling=SIGNED, max_chains=1) outputs = run_codex(proxy, "t1") - assert [o["type"] for o in outputs[0]["output"]] == ["reasoning", "function_call"] + # Qwen3-4B-Instruct-2507 has no reasoning block: its template renders no reasoning. + first = ["function_call"] if name == "qwen3_2507" else ["reasoning", "function_call"] + assert [o["type"] for o in outputs[0]["output"]] == first assert outputs[-1]["output"][-1]["content"][0]["text"] == "The file says hello." records = proxy.records("t1") @@ -363,9 +373,9 @@ def test_codex_namespace_and_hosted_tools_are_left_out_other_types_refused(): assert responses_to_messages(body) == [{"role": "user", "content": "hi"}] -@pytest.mark.parametrize("name", ["qwen3", "glm5"]) +@pytest.mark.parametrize("name", ["qwen3", "glm5", "qwen3_2507"]) def test_template_rendering_matches_tokenizer_ids_token_by_token(name, request): - template = request.getfixturevalue("qwen" if name == "qwen3" else "glm") + template = request.getfixturevalue(FIXTURES[name]) tools = chat_tools("responses", {"tools": TOOLS}) for end in range(2, len(CONVERSATION) + 1): report = check_alignment(template, CONVERSATION[:end], tools=tools) diff --git a/yeto/rl/harness/chat_templates.py b/yeto/rl/harness/chat_templates.py index 030daf39a..9662b10c7 100644 --- a/yeto/rl/harness/chat_templates.py +++ b/yeto/rl/harness/chat_templates.py @@ -17,6 +17,8 @@ ``ChatTemplate.load("glm5", tokenizer_path=...)``; whether GLM5.3 uses the same template is not verified; * ``qwen3``: the Qwen3 hermes tool-call format; the small CPU stand-in; +* ``qwen3_2507``: Qwen3-4B-Instruct-2507 (non-thinking Qwen3, hermes JSON tool + calls, no reasoning block); the Codex system messages are merged into one; * ``qwen38``: Qwen3.8 (Qwen3.5 family) XML tool calls (`` v``), thinking kept in every turn (``preserve_thinking``). Load with the checkpoint's tokenizer @@ -184,6 +186,10 @@ def register_template(spec: TemplateSpec) -> TemplateSpec: # The Qwen3 template drops the reasoning of turns before the last user turn. keeps_history_reasoning=False, verified="Qwen3-0.6B tokenizer (CPU stand-in)")) +register_template(TemplateSpec( + name="qwen3_2507", tokenizer="Qwen/Qwen3-4B-Instruct-2507", parser=parse_qwen3, + stop_tokens=("<|im_end|>", "<|endoftext|>"), + verified="Qwen3-4B-Instruct-2507@cdbee75f tokenizer (S25 T-N109)", merge_leading_system=True)) register_template(TemplateSpec( name="qwen38", tokenizer="Qwen/Qwen3.8-27B", parser=parse_qwen38, stop_tokens=("<|im_end|>", "<|endoftext|>"), From 4de5c78403cb9e370e9ec862073ccabfb4aa6c73 Mon Sep 17 00:00:00 2001 From: michael <78113554+chowhsu@users.noreply.github.com> Date: Tue, 20 Oct 2026 05:51:51 +0000 Subject: [PATCH 50/89] verl image: add yeto's object-store client (boto3) for s3/oss stores The verl venv comes from the fork's uv.lock, which has no boto3, so a verl island over an s3:// store failed on import (found by the S25 C3V S3 run). The client is installed in its own layer at yeto's pin. Co-Authored-By: Claude Opus 5.5 --- tests/test_rl_verl_adapter.py | 13 +++++++++++++ yeto/rl/adapters/verl/image.py | 14 +++++++++++--- yeto/rl/adapters/verl/pins.py | 4 ++++ 3 files changed, 28 insertions(+), 3 deletions(-) diff --git a/tests/test_rl_verl_adapter.py b/tests/test_rl_verl_adapter.py index 44f53542f..96c588451 100644 --- a/tests/test_rl_verl_adapter.py +++ b/tests/test_rl_verl_adapter.py @@ -355,6 +355,19 @@ def test_image_recipe_is_pinned(): image.build_commands("c" * 40) +def test_image_adds_the_store_client_at_the_yeto_pin(): + import tomllib + from pathlib import Path + + from yeto.rl.adapters.verl import image + from yeto.rl.adapters.verl.pins import STORE_CLIENT_REQUIREMENTS + + project = tomllib.loads((Path(__file__).parents[1] / "pyproject.toml").read_text())["project"] + pins = set(project["dependencies"]).union(*project["optional-dependencies"].values()) + assert set(STORE_CLIENT_REQUIREMENTS) <= pins + assert image.store_client_commands()[0].endswith(" ".join(STORE_CLIENT_REQUIREMENTS)) + + def test_verl_capabilities_accept_default_and_tis_specs_only(): from yeto.rl.engine.algorithm import AlgorithmSpec from yeto.rl.engine.capabilities import CapabilityMismatch diff --git a/yeto/rl/adapters/verl/image.py b/yeto/rl/adapters/verl/image.py index e7e9ca14e..b096d1b18 100644 --- a/yeto/rl/adapters/verl/image.py +++ b/yeto/rl/adapters/verl/image.py @@ -5,7 +5,8 @@ dependencies from the fork's own ``uv.lock`` (``uv sync --frozen --all-packages --extra vllm --extra fsdp``), plus yeto's one-hook local patch (``patch_verl``). ``--rl-image verl-build:`` selects it; Modal caches -the layers, so only the first run builds. yeto's own image default +the layers, so only the first run builds. yeto's object-store client (boto3, +``STORE_CLIENT_REQUIREMENTS``) is added to the same venv. yeto's own image default (ghcr.io/michaellchung/yeto-miles-ports) is unchanged. """ @@ -13,8 +14,8 @@ from pathlib import Path -from .pins import (VERL_BASE_IMAGE, VERL_COMMIT, VERL_PYTHON, VERL_REPOSITORY, VERL_ROOT, - VERL_UV_VERSION) +from .pins import (STORE_CLIENT_REQUIREMENTS, VERL_BASE_IMAGE, VERL_COMMIT, VERL_PYTHON, VERL_REPOSITORY, + VERL_ROOT, VERL_UV_VERSION) VENV_BIN = f"{VERL_ROOT}/.venv/bin" @@ -30,6 +31,12 @@ def build_commands(commit: str = VERL_COMMIT) -> list[str]: ] +def store_client_commands() -> list[str]: + """An island over an s3:// or oss:// store reads it with boto3 (S25 C3V); its own + layer, so the cached verl build above stays valid.""" + return [f"uv pip install --python {VENV_BIN}/python " + " ".join(STORE_CLIENT_REQUIREMENTS)] + + def image_env() -> dict[str, str]: return { "PATH": f"{VENV_BIN}:/usr/local/nvidia/bin:/usr/local/cuda/bin:/usr/local/sbin:/usr/local/bin:" @@ -46,6 +53,7 @@ def modal_image(modal, commit: str = VERL_COMMIT): .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()) diff --git a/yeto/rl/adapters/verl/pins.py b/yeto/rl/adapters/verl/pins.py index a7922c408..9736ae1e8 100644 --- a/yeto/rl/adapters/verl/pins.py +++ b/yeto/rl/adapters/verl/pins.py @@ -20,6 +20,10 @@ "peft": "0.19.1", } +# yeto's object-store client (s3:// and oss:// checkpoint stores: yeto.cloud.s3_store), +# not in the fork's uv.lock; the same pin as yeto's pyproject (a test keeps them equal). +STORE_CLIENT_REQUIREMENTS = ("boto3==1.43.104",) + # --- Ascend NPU stack (rl-verl-backend 3.11/3.14) ---------------------------- # # The NPU stack is a second, separate image: vLLM 0.23 + torch 2.10 + CANN From e4a77406a6f98a12e0ccdc2c277ebfc2b56a5e7e Mon Sep 17 00:00:00 2001 From: michael <78113554+chowhsu@users.noreply.github.com> Date: Tue, 20 Oct 2026 05:53:24 +0000 Subject: [PATCH 51/89] Records endpoint: model presets (Qwen3-4B-Instruct-2507 for N10.9) Co-Authored-By: Claude Opus 5.5 --- scripts/modal_records_endpoint.py | 31 +++++++++++++++++++++++++------ 1 file changed, 25 insertions(+), 6 deletions(-) diff --git a/scripts/modal_records_endpoint.py b/scripts/modal_records_endpoint.py index 0c75ca0ff..7749929f2 100644 --- a/scripts/modal_records_endpoint.py +++ b/scripts/modal_records_endpoint.py @@ -52,6 +52,21 @@ "max_serve_s": 3600, } +# Model presets: the Modal Volume holding the checkpoint and the DEFAULTS keys that differ. +MODELS: dict[str, dict[str, Any]] = { + "qwen38-27b": {"volume": MODELS_VOLUME, "config": {}}, + # S25 N10.9: full-parameter agentic RL model. Sampling is the trainer's + # (temperature 1, no top-p / top-k truncation), so the pass rate measured + # here is the one the RL rollouts see. + "qwen3-4b-2507": {"volume": "yeto-fn-models", "config": { + "model_subdir": "hf/Qwen3-4B-Instruct-2507/cdbee75f", + "served_model_name": "Qwen3-4B-Instruct-2507", + "template": "qwen3_2507", + "template_kwargs": {}, + "signed_sampling": {"temperature": 1.0, "top_p": 1.0, "top_k": -1, "max_tokens": 4096}, + }}, +} + def _gpu_names() -> list[str]: out = subprocess.run(["nvidia-smi", "--query-gpu=name", "--format=csv,noheader"], @@ -184,7 +199,7 @@ def check(config: dict[str, Any]) -> dict[str, Any]: "aligned": check_alignment(template, messages).aligned} -def build_app(workdir: str, gpu: str, timeout_s: int) -> tuple[Any, Any, Any]: +def build_app(workdir: str, gpu: str, timeout_s: int, models_volume: str = MODELS_VOLUME) -> tuple[Any, Any, Any]: import modal sys.path.insert(0, workdir) @@ -199,7 +214,7 @@ def build_app(workdir: str, gpu: str, timeout_s: int) -> tuple[Any, Any, Any]: ignore=[".git", ".venv", ".claude", "**/__pycache__", "**/*.pyc", "s1-runs", "syncer/target"])) app = modal.App(APP_NAME) - models = modal.Volume.from_name(MODELS_VOLUME) + models = modal.Volume.from_name(models_volume) results = modal.Volume.from_name(RESULTS_VOLUME, create_if_missing=True) fn = app.function(image=image, gpu=gpu, cpu=8.0, memory=65536, timeout=timeout_s, retries=0, volumes={MODELS_MOUNT: models.read_only(), RESULTS_MOUNT: results}, @@ -215,6 +230,7 @@ def main() -> None: parser.add_argument("--run-name", required=True) parser.add_argument("--secret-file", type=Path, help="launch: file holding the path secret") parser.add_argument("--gpu", default="H100!") + parser.add_argument("--model", choices=sorted(MODELS), default="qwen38-27b") parser.add_argument("--max-serve-s", type=int, default=DEFAULTS["max_serve_s"]) parser.add_argument("--workdir", default=str(Path(__file__).resolve().parents[1])) args = parser.parse_args() @@ -227,16 +243,19 @@ def main() -> None: if args.action == "url": print(control.get(f"{args.run_name}/url", "")) return + preset = MODELS[args.model] + defaults = {**DEFAULTS, **preset["config"]} if args.action == "check": - app, _, checker = build_app(args.workdir, args.gpu, timeout_s=600) + app, _, checker = build_app(args.workdir, args.gpu, timeout_s=600, models_volume=preset["volume"]) with modal.enable_output(), app.run(): - print(json.dumps(checker.remote(dict(DEFAULTS)))) + print(json.dumps(checker.remote(defaults))) return secret = args.secret_file.read_text().strip() if len(secret) < 24: raise SystemExit("the path secret must have at least 24 characters") - config = {**DEFAULTS, "run_name": args.run_name, "secret": secret, "max_serve_s": args.max_serve_s} - app, fn, _ = build_app(args.workdir, args.gpu, timeout_s=args.max_serve_s + 1800) + config = {**defaults, "run_name": args.run_name, "secret": secret, "max_serve_s": args.max_serve_s} + app, fn, _ = build_app(args.workdir, args.gpu, timeout_s=args.max_serve_s + 1800, + models_volume=preset["volume"]) with modal.enable_output(), app.run(detach=True): call = fn.spawn(config) print(json.dumps({"function_call_id": call.object_id, "run_name": args.run_name})) From ef6400c92be5646d5d2f79c64a63184fa646bd98 Mon Sep 17 00:00:00 2001 From: michael <78113554+chowhsu@users.noreply.github.com> Date: Tue, 20 Oct 2026 06:04:24 +0000 Subject: [PATCH 52/89] Island keeps the default verify_miles_revision call without --miles-commit; allowlist the requant frozen flags Co-Authored-By: Claude Opus 5.5 --- tests/test_run_flags.py | 2 +- yeto/rl/adapters/miles/island_entry.py | 13 +++++++++---- 2 files changed, 10 insertions(+), 5 deletions(-) diff --git a/tests/test_run_flags.py b/tests/test_run_flags.py index afc9b9ea6..9ef5c8eab 100644 --- a/tests/test_run_flags.py +++ b/tests/test_run_flags.py @@ -269,7 +269,7 @@ def test_forward_tables_name_flags_of_both_layers(): "--temp-dir", "--ulimit", "--untracked-files", "--wandb-group", "--yes", # torchrun, tail and the W4A8 re-quantizer CLI in the SSH trainer island (yeto/launch/ssh_trainer_island.py) "--awq-scales", "--loop", "--master-addr", "--master-port", "--no-awq", "--node-rank", "--nproc-per-node", - "--pid", "--poll-seconds", "--recipe", "--store", "--template", + "--pid", "--poll-seconds", "--recipe", "--store", "--template", "--frozen-base", "--frozen-pattern", # docker run of a runtime: docker SSH trainer node (ssh_trainer_island.py, node_inventory.py example) "--device", "--network", }) diff --git a/yeto/rl/adapters/miles/island_entry.py b/yeto/rl/adapters/miles/island_entry.py index 5511b9bd6..d45980835 100644 --- a/yeto/rl/adapters/miles/island_entry.py +++ b/yeto/rl/adapters/miles/island_entry.py @@ -1093,13 +1093,14 @@ def _preflight_codex_openenv_adapter(args, profile_name: str) -> None: def miles_revision_pins(args: Any) -> Any: - """The Miles commit the island checks: ``--miles-commit`` (full SHA) or the pin.""" + """The Miles commit the island checks for ``--miles-commit`` (full SHA); None + without it (``verify_miles_revision`` then checks the pinned MILES_NEXT_PINS).""" from yeto.provenance import is_immutable_commit - from yeto.rl.adapters.miles.pins import MILES_NEXT_PINS, MILES_NEXT_REPOSITORY, MilesRevisionPins + from yeto.rl.adapters.miles.pins import MILES_NEXT_REPOSITORY, MilesRevisionPins commit = getattr(args, "miles_commit", None) if not commit: - return MILES_NEXT_PINS + return None if not is_immutable_commit(commit) or len(str(commit)) != 40: raise ValueError(f"--miles-commit must be a full 40-hex commit, got {commit!r}") return MilesRevisionPins(MILES_NEXT_REPOSITORY, str(commit).lower()) @@ -2365,7 +2366,11 @@ def main(argv=None) -> None: miles_root = str(Path(args.miles_root).expanduser().resolve()) if miles_root not in sys.path: sys.path.insert(0, miles_root) - verify_miles_revision(miles_root, expected=miles_revision_pins(args)) + pins = miles_revision_pins(args) + if pins is None: + verify_miles_revision(miles_root) # the pinned MILES_NEXT_PINS + else: + verify_miles_revision(miles_root, expected=pins) _preflight_codex_harness(args) probe_syncers(args) From ac4f8526b09b4465c6c0d9dcb6973db1835f8d12 Mon Sep 17 00:00:00 2001 From: michael <78113554+chowhsu@users.noreply.github.com> Date: Tue, 20 Oct 2026 06:12:35 +0000 Subject: [PATCH 53/89] verl agentic path: records proxy over verl's rollout client, path-secret recorder, agentic entry VerlRolloutEngine forwards the records proxy's token-id prompts to verl's LLMServerClient on the agent-loop worker's event loop; HarborAgentLoop binds it before the first call. ProxyTokenRecorder can serve under a /k/ prefix and advertise a tunnel URL. agentic_main starts Ray with verl's runtime env, the forwarded agentic variables and a setup hook that registers yeto_grpo and the yeto_harbor agent loop. Co-Authored-By: Claude Opus 5.5 --- tests/test_verl_rollout_engine.py | 135 ++++++++++++++ yeto/rl/adapters/verl/agentic_main.py | 63 +++++++ yeto/rl/adapters/verl/harbor_agent_loop.py | 4 + yeto/rl/adapters/verl/rollout_engine.py | 206 +++++++++++++++++++++ yeto/rl/harness/gateway/recorder.py | 31 +++- 5 files changed, 435 insertions(+), 4 deletions(-) create mode 100644 tests/test_verl_rollout_engine.py create mode 100644 yeto/rl/adapters/verl/agentic_main.py create mode 100644 yeto/rl/adapters/verl/rollout_engine.py diff --git a/tests/test_verl_rollout_engine.py b/tests/test_verl_rollout_engine.py new file mode 100644 index 000000000..1f59e710b --- /dev/null +++ b/tests/test_verl_rollout_engine.py @@ -0,0 +1,135 @@ +"""S25 N10.9: the records proxy over verl's own rollout client, behind a path secret (CPU). + +Real pieces: ``HarborAgentLoop.run``, ``ProxyTokenRecorder`` serving a +``RecordsProxy`` over HTTP, the gateway and recording backend, and +``VerlRolloutEngine``. Fakes: verl's agent-loop module (stub), verl's +``LLMServerClient`` (an async ``generate`` that replies with the scripted +model's tokens), the agent (Codex-like Responses loop). Tokenizer: Qwen3-0.6B +from the local HF cache (skip if absent). +""" + +from __future__ import annotations + +import asyncio +import copy +from types import SimpleNamespace +from typing import Any + +import pytest +import test_agentic_harbor as agentic +from test_agentic_harbor import _make_loop +from test_records_proxy import SCRIPT, SIGNED, ScriptedModel, _template +from test_verl_records_recorder import ROW, VERL_SAMPLING, CodexLikeAgent + +from yeto.rl.agentic import OUTCOME_KEY, HarnessConfigError +from yeto.rl.engine.sampler_engine import FINISH_LENGTH, FINISH_STOP, EngineIdentity, SamplerEngineError +from yeto.rl.harness.gateway.proxy import RecordsProxy +from yeto.rl.harness.gateway.recorder import ProxyTokenRecorder + +pytest.importorskip("aiohttp") + +verl_loop = agentic.verl_loop # the verl agent-loop stub fixture +SECRET = "s" * 32 +IDENTITY = EngineIdentity(engine="verl-rollout", engine_version="test", device_family="cuda", + card_model="H100", precision="bf16", quant_scheme="bf16", topology="unknown") + + +@pytest.fixture(scope="module") +def qwen(): + return _template("qwen3") + + +class FakeVerlClient: + """verl ``LLMServerClient`` stand-in: replies with the scripted model's tokens.""" + + def __init__(self, model: ScriptedModel, *, aborted: bool = False) -> None: + self.model = model + self.aborted = aborted + self.calls: list[dict[str, Any]] = [] + self.loops: set[int] = set() + + async def generate(self, request_id: str, *, prompt_ids: list[int], sampling_params: dict[str, Any]) -> Any: + self.loops.add(id(asyncio.get_running_loop())) + self.calls.append({"prompt_ids": list(prompt_ids), "sampling_params": dict(sampling_params)}) + if self.aborted: + return SimpleNamespace(token_ids=[], log_probs=None, stop_reason="aborted") + ids = self.model(tuple(prompt_ids), 0) + return SimpleNamespace(token_ids=ids, log_probs=[-0.25] * len(ids), stop_reason="completed") + + +def _verl_proxy(template) -> tuple[RecordsProxy, Any]: + from yeto.rl.adapters.verl.rollout_engine import VerlRolloutEngine + + engine = VerlRolloutEngine(max_model_len=8192, identity=IDENTITY) + return RecordsProxy(engine, template, model="policy", signed_sampling=SIGNED, max_chains=1), engine + + +def test_agent_loop_samples_through_verls_client_behind_the_path_secret(verl_loop, qwen): + loop_mod, _ = verl_loop + proxy, engine = _verl_proxy(qwen) + recorder = ProxyTokenRecorder(proxy, path_secret=SECRET) + agent = CodexLikeAgent() + loop = _make_loop(loop_mod, recorder, agent, prompt_length=4096, response_length=4096) + client = FakeVerlClient(ScriptedModel(qwen, copy.deepcopy(SCRIPT))) + loop.server_manager = client # verl's AgentLoopBase sets it; the stub does not + + async def go(): + try: + out = await loop.run(VERL_SAMPLING, **ROW) + import aiohttp + + base = recorder._base_url.split("/k/")[0] + async with aiohttp.ClientSession() as http: # the root is not served, only the secret prefix + async with http.get(f"{base}/t/traj-1/records") as reply: + root_status = reply.status + finally: + await recorder.aclose() + return out, root_status + + out, root_status = asyncio.run(go()) + assert out.extra_fields[OUTCOME_KEY] == "ok" and out.reward_score == 1.0 and out.num_turns == 6 + (endpoint,) = agent.endpoints + assert f"/k/{SECRET}/t/traj-1-" in endpoint and root_status == 404 + assert engine.bound and len(client.calls) == 3 and len(client.loops) == 1 + params = client.calls[0]["sampling_params"] + assert params["logprobs"] is True and params["temperature"] == 1.0 and params["top_k"] == -1 + assert params["stop_token_ids"] == sorted(qwen.stop_token_ids) and params["max_tokens"] == 256 + # trained == sampled: the last prompt verl got is a prefix of prompt + response + last_prompt = client.calls[-1]["prompt_ids"] + assert out.prompt_ids + out.response_ids[:len(last_prompt) - len(out.prompt_ids)] == last_prompt + trained = [lp for lp, m in zip(out.response_logprobs, out.response_mask) if m == 1] + assert trained and all(lp == -0.25 for lp in trained) + assert out.extra_fields["provenance"]["engine"] == "verl-rollout" + assert out.extra_fields["length_capped_turns"] == 0 + + +def test_an_aborted_verl_request_aborts_the_sample(verl_loop, qwen): + loop_mod, _ = verl_loop + proxy, _ = _verl_proxy(qwen) + recorder = ProxyTokenRecorder(proxy) + agent = CodexLikeAgent() + loop = _make_loop(loop_mod, recorder, agent, prompt_length=4096, response_length=4096) + loop.server_manager = FakeVerlClient(ScriptedModel(qwen, copy.deepcopy(SCRIPT)), aborted=True) + + async def go(): + try: + return await loop.run(VERL_SAMPLING, **ROW) + finally: + await recorder.aclose() + + out = asyncio.run(go()) + assert out.extra_fields[OUTCOME_KEY] == "aborted" and sum(out.response_mask) == 0 + + +def test_unbound_engine_and_short_secret_fail_closed(qwen): + from yeto.rl.adapters.verl.rollout_engine import VerlRolloutEngine, finish_reason + from yeto.rl.engine.sampler_engine import SamplingParams + + engine = VerlRolloutEngine(max_model_len=64, identity=IDENTITY) + with pytest.raises(SamplerEngineError, match="not bound"): + engine.generate_tokens([[1, 2]], SamplingParams(max_tokens=4)) + with pytest.raises(HarnessConfigError, match="at least"): + ProxyTokenRecorder(_verl_proxy(qwen)[0], path_secret="short") + assert finish_reason((5, 7), {7}, 2) == FINISH_STOP + assert finish_reason((5, 6), {7}, 2) == FINISH_LENGTH + assert finish_reason((5,), {7}, 2) == FINISH_STOP diff --git a/yeto/rl/adapters/verl/agentic_main.py b/yeto/rl/adapters/verl/agentic_main.py new file mode 100644 index 000000000..1fdaba314 --- /dev/null +++ b/yeto/rl/adapters/verl/agentic_main.py @@ -0,0 +1,63 @@ +"""verl's own synchronous trainer on the Harbor agentic path (S25 N10.9; runs inside the verl image). + +``python -m yeto.rl.adapters.verl.agentic_main `` is verl's +``main_ppo.main`` unchanged, with Ray started here so that every Ray worker: + +* imports the yeto pieces verl selects by name -- the ``yeto_grpo`` advantage + estimator (trainer process) and the ``yeto_harbor`` agent loop (agent-loop + workers) -- through Ray's ``worker_process_setup_hook``; +* carries the agentic environment: ``YETO_*`` (recorder, records proxy, boat + caps), ``HARBOR_*`` (task directory, provider) and ``BOAT_API_KEY`` (the + sandbox account, read by ``boat_client``), on top of verl's own runtime env. + +The round loop, weight updates and rollout servers are verl's. The records +proxy runs in each agent-loop worker over verl's rollout client +(:mod:`.rollout_engine`); select it with +``YETO_RECORDS_PROXY=yeto.rl.adapters.verl.rollout_engine:records_proxy_from_env``. +""" + +from __future__ import annotations + +import os +import sys + +FORWARD_ENV_PREFIXES = ("YETO_", "HARBOR_", "HF_", "PYTHONPATH") +FORWARD_ENV_NAMES = ("BOAT_API_KEY",) +SETUP_HOOK = "yeto.rl.adapters.verl.agentic_main.setup_worker" + + +def setup_worker() -> None: + """Ray ``worker_process_setup_hook``: register the yeto estimator and agent loop.""" + import yeto.rl.adapters.verl.advantage # noqa: F401 (registers yeto_grpo) + import yeto.rl.adapters.verl.harbor_agent_loop # noqa: F401 (registers yeto_harbor) + + +def runtime_env(environ: dict[str, str] | None = None) -> dict: + """verl's PPO runtime env plus the forwarded agentic variables and the setup hook.""" + from verl.trainer.constants_ppo import get_ppo_ray_runtime_env + + environ = dict(os.environ if environ is None else environ) + env = get_ppo_ray_runtime_env(None) + env.setdefault("env_vars", {}) + env["env_vars"]["TRANSFER_QUEUE_ENABLE"] = "1" # verl's V1 trainer runs on TransferQueue + for key, value in environ.items(): + if key.startswith(FORWARD_ENV_PREFIXES) or key in FORWARD_ENV_NAMES: + env["env_vars"][key] = value + env["worker_process_setup_hook"] = SETUP_HOOK + return env + + +def main(argv: list[str] | None = None) -> None: + import ray + import verl.trainer.main_ppo as main_ppo + + setup_worker() # the driver validates algorithm.adv_estimator against the registry too + if not ray.is_initialized(): + ray.init(runtime_env=runtime_env(), namespace="yeto-verl-agentic") + if argv is not None: + sys.argv = [sys.argv[0], *argv] + main_ppo.main() + + +if __name__ == "__main__": + main() diff --git a/yeto/rl/adapters/verl/harbor_agent_loop.py b/yeto/rl/adapters/verl/harbor_agent_loop.py index 21919bd6f..c50b16c2e 100644 --- a/yeto/rl/adapters/verl/harbor_agent_loop.py +++ b/yeto/rl/adapters/verl/harbor_agent_loop.py @@ -85,6 +85,7 @@ from yeto.rl.harness.harbor import HarborHarness, HarborSettings, build_trial_config from yeto.rl.harness.harbor.boat_client import check_in_flight_cap +from .rollout_engine import bind_rollout_client from .run_config_rules import AGENTIC_ADV_ESTIMATOR as YETO_GRPO logger = logging.getLogger(__name__) @@ -191,6 +192,9 @@ def __init__(self, *args: Any, **kwargs: Any) -> None: async def run(self, sampling_params: dict[str, Any], **kwargs: Any) -> AgentLoopOutput: task = task_from_row(kwargs) recorder = self.recorder_factory() + client = getattr(self, "server_manager", None) + if client is not None: # a recorder over verl's own rollout needs verl's client (N10.9) + bind_rollout_client(recorder, client) session = recorder_session_id(task) endpoint = await recorder.open(session, sampling_params) try: diff --git a/yeto/rl/adapters/verl/rollout_engine.py b/yeto/rl/adapters/verl/rollout_engine.py new file mode 100644 index 000000000..f78590b3a --- /dev/null +++ b/yeto/rl/adapters/verl/rollout_engine.py @@ -0,0 +1,206 @@ +"""verl's own rollout servers behind the records proxy (S25 N10.9). + +In a verl training run the policy is served by verl's rollout replicas (vLLM +or SGLang, started and updated by verl). ``HarborAgentLoop`` needs a token +recorder in front of that policy: the sandboxed agent talks OpenAI to the +records proxy, the proxy renders the chat template and asks an engine for +tokens. :class:`VerlRolloutEngine` is that engine. It owns no server: it +forwards token-id prompts to verl's ``LLMServerClient`` (the agent loop's +``server_manager``) on the agent-loop worker's event loop, so verl keeps +routing, sticky sessions and weight updates. + +Wiring (all in the agent-loop worker process): + +* ``YETO_AGENTIC_RECORDER=yeto.rl.harness.gateway.recorder:recorder_from_env``; +* ``YETO_RECORDS_PROXY=yeto.rl.adapters.verl.rollout_engine:records_proxy_from_env`` + (this module) builds the proxy over a :class:`VerlRolloutEngine`; +* ``HarborAgentLoop.run`` binds the engine to its ``server_manager`` and the + running event loop before the first call (:func:`bind_rollout_client`). + +Environment of :func:`records_proxy_from_env`: + YETO_RECORDS_TEMPLATE chat template name (``chat_templates.TEMPLATES``) + YETO_RECORDS_TOKENIZER tokenizer directory (the checkpoint) + YETO_RECORDS_MODEL model name the agent sends (``AgentTask.model_name``) + YETO_RECORDS_SAMPLING JSON signed sampling (temperature, top_p, top_k, max_tokens) + YETO_RECORDS_MAX_MODEL_LEN context length of verl's rollout (``rollout.max_model_len``) + YETO_RECORDS_CARD GPU name (provenance; default ``unknown``) + +Policy version: verl's synchronous trainer updates the rollout weights only +between steps, never while a batch is generated, so a turn cannot straddle a +reload. The rollout reply carries verl's ``global_steps`` only in the fully +async mode; :meth:`VerlRolloutEngine.served_version` therefore returns None +(unknown) and the trainer's step log is the version record. + +Finish reason: verl reports ``completed`` for both an end token and the token +cap. A reply that ends without a stop token and fills ``max_tokens`` is +reported as ``length``; everything else as ``stop``. +""" + +from __future__ import annotations + +import asyncio +import json +import os +import uuid +from collections.abc import Sequence +from pathlib import Path +from typing import Any + +from yeto.rl.engine.sampler_engine import ( + FINISH_LENGTH, + FINISH_STOP, + UNKNOWN, + Completion, + EngineIdentity, + LoadReport, + SamplerEngineError, + SamplingParams, +) + +TEMPLATE_ENV = "YETO_RECORDS_TEMPLATE" +TOKENIZER_ENV = "YETO_RECORDS_TOKENIZER" +MODEL_ENV = "YETO_RECORDS_MODEL" +SAMPLING_ENV = "YETO_RECORDS_SAMPLING" +MAX_MODEL_LEN_ENV = "YETO_RECORDS_MAX_MODEL_LEN" +CARD_ENV = "YETO_RECORDS_CARD" + +CALL_TIMEOUT_S = 1800.0 + + +class VerlRolloutEngine: + """``TokenSamplerEngine`` over verl's ``LLMServerClient`` (module doc).""" + + def __init__(self, *, max_model_len: int, identity: EngineIdentity) -> None: + self._max_model_len = int(max_model_len) + self._identity = identity + self._client: Any = None + self._loop: asyncio.AbstractEventLoop | None = None + + # ------------------------------------------------------------ binding + def bind(self, client: Any, loop: asyncio.AbstractEventLoop) -> None: + """Use ``client`` (verl ``server_manager``) on ``loop``; the same pair for the worker's life.""" + if self._client is not None and (self._client is not client or self._loop is not loop): + raise SamplerEngineError("VerlRolloutEngine is already bound to another rollout client or loop") + self._client, self._loop = client, loop + + @property + def bound(self) -> bool: + return self._client is not None + + # ------------------------------------------------------------ engine port + @property + def identity(self) -> EngineIdentity: + return self._identity + + @property + def max_model_len(self) -> int: + return self._max_model_len + + def served_version(self) -> int | None: + return None + + def generate_tokens(self, prompts: Sequence[Sequence[int]], sampling: SamplingParams, *, + stop_token_ids: Sequence[int] = ()) -> list[list[Completion]]: + """Called from a worker thread (the records backend uses ``asyncio.to_thread``).""" + if self._client is None or self._loop is None: + raise SamplerEngineError("VerlRolloutEngine is not bound to verl's rollout client") + stops = {int(t) for t in stop_token_ids} + out: list[list[Completion]] = [] + for prompt in prompts: + group = [] + for _ in range(int(sampling.n)): + future = asyncio.run_coroutine_threadsafe( + self._generate(list(prompt), sampling, sorted(stops)), self._loop) + group.append(future.result(timeout=CALL_TIMEOUT_S)) + out.append(group) + return out + + async def _generate(self, prompt: list[int], sampling: SamplingParams, stops: list[int]) -> Completion: + params: dict[str, Any] = { + "temperature": float(sampling.temperature), "top_p": float(sampling.top_p), + "top_k": int(sampling.top_k), "repetition_penalty": float(sampling.repetition_penalty), + "max_tokens": int(sampling.max_tokens), "logprobs": True, + } + if stops: + params["stop_token_ids"] = stops + output = await self._client.generate(request_id=uuid.uuid4().hex, prompt_ids=prompt, + sampling_params=params) + if getattr(output, "stop_reason", None) == "aborted": + raise SamplerEngineError("verl rollout aborted the request") + tokens = tuple(int(t) for t in output.token_ids) + logprobs = output.log_probs + if logprobs is None or len(logprobs) != len(tokens): + raise SamplerEngineError(f"verl rollout returned {0 if logprobs is None else len(logprobs)} " + f"log-probs for {len(tokens)} tokens") + return Completion(token_ids=tokens, logprobs=tuple(float(x) for x in logprobs), text="", + finish_reason=finish_reason(tokens, stops, int(sampling.max_tokens)), + prompt_token_ids=tuple(prompt)) + + # verl starts, updates and stops its own servers: the proxy never calls these. + def start(self, release_dir: Path | None, version: int) -> LoadReport: + raise SamplerEngineError("verl owns its rollout servers") + + def load(self, release_dir: Path, version: int) -> LoadReport: + raise SamplerEngineError("verl owns its rollout servers") + + def generate(self, prompts: Sequence[str], sampling: SamplingParams) -> list[list[Completion]]: + raise SamplerEngineError("text prompts are not used; the records proxy sends token ids") + + def score(self, prompt: str, continuation: str) -> list[float]: + raise SamplerEngineError("scoring is not used on the verl rollout") + + def loaded_expert_matrices(self) -> int | None: + return None + + def weight_gib_per_card(self) -> float | None: + return None + + def stop(self) -> None: + return None + + +def finish_reason(tokens: Sequence[int], stops: set[int] | Sequence[int], max_tokens: int) -> str: + """``length`` when the reply filled ``max_tokens`` without a stop token (module doc).""" + if tokens and int(tokens[-1]) in set(stops): + return FINISH_STOP + return FINISH_LENGTH if len(tokens) >= max_tokens else FINISH_STOP + + +def bind_rollout_client(recorder: Any, client: Any) -> None: + """Bind the recorder's :class:`VerlRolloutEngine` (if it has one) to verl's client on this loop.""" + proxy = getattr(recorder, "proxy", None) + engine = getattr(getattr(proxy, "backend", None), "engine", None) + if isinstance(engine, VerlRolloutEngine): + engine.bind(client, asyncio.get_running_loop()) + + +def records_proxy_from_env() -> Any: + """``YETO_RECORDS_PROXY`` callable: a RecordsProxy over a :class:`VerlRolloutEngine` (module doc).""" + from yeto.rl.agentic import HarnessConfigError + from yeto.rl.harness.chat_templates import ChatTemplate + from yeto.rl.harness.gateway.proxy import RecordsProxy + + try: + template_name = os.environ[TEMPLATE_ENV] + tokenizer = os.environ[TOKENIZER_ENV] + model = os.environ[MODEL_ENV] + signed = json.loads(os.environ[SAMPLING_ENV]) + max_model_len = int(os.environ[MAX_MODEL_LEN_ENV]) + except (KeyError, ValueError) as exc: + raise HarnessConfigError(f"records proxy over verl rollout: bad or missing setting {exc}") from exc + identity = EngineIdentity(engine="verl-rollout", engine_version=_verl_version(), device_family="cuda", + card_model=os.environ.get(CARD_ENV, UNKNOWN), precision="bf16", + quant_scheme="bf16", topology=UNKNOWN) + engine = VerlRolloutEngine(max_model_len=max_model_len, identity=identity) + template = ChatTemplate.load(template_name, tokenizer_path=tokenizer) + return RecordsProxy(engine, template, model=model, signed_sampling=signed, + sampler_id="verl-rollout-records") + + +def _verl_version() -> str: + try: + import verl + + return str(getattr(verl, "__version__", UNKNOWN)) + except ImportError: + return UNKNOWN diff --git a/yeto/rl/harness/gateway/recorder.py b/yeto/rl/harness/gateway/recorder.py index f28edad60..66f9ff995 100644 --- a/yeto/rl/harness/gateway/recorder.py +++ b/yeto/rl/harness/gateway/recorder.py @@ -28,6 +28,11 @@ YETO_RECORDS_PROXY_HOST bind address (default 127.0.0.1) YETO_RECORDS_PROXY_PORT bind port (default 0: a free port) YETO_RECORDS_PROXY_PUBLIC_HOST host the sandbox dials (default: the bind address) + YETO_RECORDS_PROXY_PUBLIC_BASE full URL the sandbox dials instead (a tunnel, e.g. + ``https://.modal.host``); overrides PUBLIC_HOST and the port + YETO_RECORDS_PROXY_SECRET_FILE file holding a path secret: the proxy is served only under + ``/k//`` (a public tunnel never exposes the + unauthenticated root); the secret is never logged """ from __future__ import annotations @@ -35,6 +40,7 @@ import asyncio import importlib import os +from pathlib import Path from collections.abc import Mapping, Sequence from typing import Any @@ -47,6 +53,9 @@ HOST_ENV = "YETO_RECORDS_PROXY_HOST" PORT_ENV = "YETO_RECORDS_PROXY_PORT" PUBLIC_HOST_ENV = "YETO_RECORDS_PROXY_PUBLIC_HOST" +PUBLIC_BASE_ENV = "YETO_RECORDS_PROXY_PUBLIC_BASE" +SECRET_FILE_ENV = "YETO_RECORDS_PROXY_SECRET_FILE" +MIN_SECRET_CHARS = 24 # verl sampling keys -> SamplingParams fields; other keys (``logprobs``) do not change sampling. CHECKED_SAMPLING = ("temperature", "top_p", "top_k", "repetition_penalty") @@ -87,11 +96,16 @@ class ProxyTokenRecorder: """:class:`~yeto.rl.agentic.TokenRecorder` over one in-process :class:`RecordsProxy` (module doc).""" def __init__(self, proxy: RecordsProxy, *, host: str = "127.0.0.1", port: int = 0, - public_host: str | None = None) -> None: + public_host: str | None = None, public_base: str | None = None, + path_secret: str | None = None) -> None: + if path_secret is not None and len(path_secret) < MIN_SECRET_CHARS: + raise HarnessConfigError(f"the path secret must have at least {MIN_SECRET_CHARS} characters") self.proxy = proxy self.host = host self.port = int(port) self.public_host = public_host or host + self.public_base = public_base.rstrip("/") if public_base else None + self._prefix = f"/k/{path_secret}" if path_secret else "" self._signed = sampling_params(proxy.gateway.config.signed_sampling) self._runner: Any = None self._base_url: str | None = None @@ -103,13 +117,19 @@ async def _base(self) -> str: if self._base_url is None: from aiohttp import web - runner = web.AppRunner(self.proxy.app()) + app = self.proxy.app() + if self._prefix: + outer = web.Application() + outer.add_subapp(f"{self._prefix}/", app) + app = outer + runner = web.AppRunner(app) await runner.setup() site = web.TCPSite(runner, self.host, self.port) await site.start() port = runner.addresses[0][1] self._runner = runner - self._base_url = f"http://{self.public_host}:{port}" + base = self.public_base or f"http://{self.public_host}:{port}" + self._base_url = f"{base}{self._prefix}" return self._base_url def check_sampling(self, sampling: Mapping[str, Any]) -> None: @@ -160,5 +180,8 @@ def recorder_from_env() -> ProxyTokenRecorder: if not isinstance(proxy, RecordsProxy): raise HarnessConfigError(f"{spec} returned {type(proxy).__name__}, not a RecordsProxy") host = os.environ.get(HOST_ENV, "127.0.0.1") + secret_file = os.environ.get(SECRET_FILE_ENV, "").strip() + secret = Path(secret_file).read_text().strip() if secret_file else None return ProxyTokenRecorder(proxy, host=host, port=int(os.environ.get(PORT_ENV, "0")), - public_host=os.environ.get(PUBLIC_HOST_ENV) or host) + public_host=os.environ.get(PUBLIC_HOST_ENV) or host, + public_base=os.environ.get(PUBLIC_BASE_ENV) or None, path_secret=secret) From 5f7d529f3c86ff270dab71172267ab87aef1a956 Mon Sep 17 00:00:00 2001 From: michael <78113554+chowhsu@users.noreply.github.com> Date: Tue, 20 Oct 2026 06:16:19 +0000 Subject: [PATCH 54/89] Agentic outcome: a malformed model reply that ends the agent is a trained 0, not ABORTED The records backend marks a completion the chat template cannot parse (parse_error, policy_output_errors counter); the sampled ids stay in the stream. When Harbor reports NonZeroAgentExitCodeError and the recorder saw such a reply, policy_output_failure re-classifies the trajectory as POLICY_LIMIT with reward 0. Every other ABORTED stays ABORTED. Co-Authored-By: Claude Opus 5.5 --- tests/test_verl_rollout_engine.py | 68 ++++++++++++++++++++++ yeto/rl/adapters/verl/harbor_agent_loop.py | 11 ++++ yeto/rl/agentic/outcome.py | 26 ++++++++- yeto/rl/agentic/types.py | 4 ++ yeto/rl/harness/gateway/recorder.py | 3 +- yeto/rl/harness/gateway/records.py | 14 ++++- 6 files changed, 121 insertions(+), 5 deletions(-) diff --git a/tests/test_verl_rollout_engine.py b/tests/test_verl_rollout_engine.py index 1f59e710b..c2e97bbf4 100644 --- a/tests/test_verl_rollout_engine.py +++ b/tests/test_verl_rollout_engine.py @@ -133,3 +133,71 @@ def test_unbound_engine_and_short_secret_fail_closed(qwen): assert finish_reason((5, 7), {7}, 2) == FINISH_STOP assert finish_reason((5, 6), {7}, 2) == FINISH_LENGTH assert finish_reason((5,), {7}, 2) == FINISH_STOP + + +class ExitOnErrorAgent(CodexLikeAgent): + """Codex on a 400: the turn fails and the agent process exits non-zero (Harbor's verdict).""" + + def __init__(self, exception_type: str = "NonZeroAgentExitCodeError") -> None: + super().__init__() + self.exception_type = exception_type + self.statuses: list[int] = [] + + async def run(self, task: Any, endpoint: str) -> Any: + import aiohttp + + from test_records_proxy import echo, responses_body, tool_output, user_item + + from yeto.rl.agentic import TrialVerdict + + self.endpoints.append(endpoint) + items = [user_item("read a.txt")] + async with aiohttp.ClientSession() as http: + for turn in range(3): + async with http.post(f"{endpoint}/responses", json=responses_body(copy.deepcopy(items))) as reply: + self.statuses.append(reply.status) + if reply.status != 200: + return TrialVerdict(rewards=None, exception_type=self.exception_type) + response = await reply.json() + items += echo(response["output"]) + items += [tool_output(o["call_id"], f"result {turn}") + for o in response["output"] if o["type"] == "function_call"] + return TrialVerdict(rewards={"reward": 1.0}) + + +class MalformedSecondReply(ScriptedModel): + """The second reply is a tool call that is not JSON (as Qwen3-4B wrote with heredocs).""" + + def __call__(self, prompt: tuple[int, ...], index: int) -> list[int]: + if len(self.replies) == len(SCRIPT) - 1: + self.replies.pop(0) + return self.template.encode('\n{"name": "shell", "arguments": {"cmd": "cat > a\nb"}\n') + \ + [self.template.encode("<|im_end|>")[0]] + return super().__call__(prompt, index) + + +@pytest.mark.parametrize("exception_type,outcome", [("NonZeroAgentExitCodeError", "policy_limit"), + ("EnvironmentStartTimeoutError", "aborted")]) +def test_malformed_tool_call_is_a_trained_zero_not_an_abort(verl_loop, qwen, exception_type, outcome): + loop_mod, _ = verl_loop + proxy, _ = _verl_proxy(qwen) + recorder = ProxyTokenRecorder(proxy) + agent = ExitOnErrorAgent(exception_type) + loop = _make_loop(loop_mod, recorder, agent, prompt_length=4096, response_length=4096) + loop.server_manager = FakeVerlClient(MalformedSecondReply(qwen, copy.deepcopy(SCRIPT))) + + async def go(): + try: + return await loop.run(VERL_SAMPLING, **ROW) + finally: + await recorder.aclose() + + out = asyncio.run(go()) + assert agent.statuses == [200, 400] + assert out.extra_fields[OUTCOME_KEY] == outcome and out.reward_score == 0.0 + if outcome == "policy_limit": + # the malformed reply was sampled and is trained (mask 1) as part of a 0-reward sample + assert out.extra_fields["outcome_reason"].startswith("policy_output_error") + assert sum(out.response_mask) > 0 + else: + assert sum(out.response_mask) == 0 # another exit cause stays an infrastructure abort diff --git a/yeto/rl/adapters/verl/harbor_agent_loop.py b/yeto/rl/adapters/verl/harbor_agent_loop.py index c50b16c2e..6f966c635 100644 --- a/yeto/rl/adapters/verl/harbor_agent_loop.py +++ b/yeto/rl/adapters/verl/harbor_agent_loop.py @@ -80,6 +80,7 @@ run_agent_task, ) from yeto.rl.agentic.loop_local import LoopLocal +from yeto.rl.agentic.outcome import policy_output_failure from yeto.rl.agentic.types import TokenTurn from yeto.rl.engine.sampler_engine import FINISH_LENGTH from yeto.rl.harness.harbor import HarborHarness, HarborSettings, build_trial_config @@ -208,6 +209,8 @@ async def run(self, sampling_params: dict[str, Any], **kwargs: Any) -> AgentLoop except TokenAlignmentError as exc: if trajectory.outcome is not Outcome.ABORTED: trajectory = _abort(trajectory, f"token_alignment: {exc}") + else: + trajectory = _policy_output_failure(trajectory, turns) segments: TokenSegments | None = None if trajectory.outcome is not Outcome.ABORTED: @@ -250,6 +253,14 @@ def _placeholder(self) -> TokenSegments: return TokenSegments(prompt_ids=(pad,), response_ids=(pad,), response_mask=(0,), response_logprobs=(0.0,)) +def _policy_output_failure(trajectory: AgentTrajectory, turns: Sequence[TokenTurn]) -> AgentTrajectory: + """A malformed model reply that ended the agent is a 0 reward, not ABORTED (N10.9).""" + updated = policy_output_failure(trajectory, [t.policy_error for t in turns if t.policy_error]) + if updated.outcome is not trajectory.outcome: + _gate().metrics.reclassify(trajectory.outcome, updated.outcome) + return updated + + def _abort(trajectory: AgentTrajectory, reason: str) -> AgentTrajectory: _gate().metrics.reclassify(trajectory.outcome, Outcome.ABORTED) return AgentTrajectory(task=trajectory.task, outcome=Outcome.ABORTED, reward=0.0, reason=reason, diff --git a/yeto/rl/agentic/outcome.py b/yeto/rl/agentic/outcome.py index 2b3c34ac0..ea6b2c24d 100644 --- a/yeto/rl/agentic/outcome.py +++ b/yeto/rl/agentic/outcome.py @@ -13,8 +13,11 @@ from __future__ import annotations +from collections.abc import Sequence +from dataclasses import replace + from .trust import verified_reward -from .types import Outcome, TrialVerdict +from .types import AgentTrajectory, Outcome, TrialVerdict # The policy ran out of its own budget: time, or sequence/context length. POLICY_LIMIT_EXCEPTIONS = frozenset({ @@ -30,6 +33,27 @@ # Reason given when the outer wall-clock cap fires (the harness never returned). WALL_CLOCK_REASON = "wall_clock_cap" +# The agent process exited non-zero. On its own this is ABORTED (could be the +# sandbox); with a recorded malformed model reply it is the policy's failure. +AGENT_EXIT_EXCEPTIONS = frozenset({"NonZeroAgentExitCodeError"}) +POLICY_OUTPUT_ERROR_REASON = "policy_output_error" + + +def policy_output_failure(trajectory: AgentTrajectory, policy_errors: Sequence[str]) -> AgentTrajectory: + """ABORTED by the agent's exit after the recorder refused a malformed model reply -> POLICY_LIMIT, reward 0. + + S25 N10.9 (main agent decision): a tool call that is not valid JSON is the + policy's own failure, trained as a 0 reward, not an infrastructure abort. + Fail closed: only the agent-exit exception, and only with a recorded + refusal, is re-classified; every other ABORTED stays ABORTED. + """ + verdict = trajectory.verdict + if (trajectory.outcome is not Outcome.ABORTED or not policy_errors or verdict is None + or verdict.exception_type not in AGENT_EXIT_EXCEPTIONS): + return trajectory + return replace(trajectory, outcome=Outcome.POLICY_LIMIT, reward=0.0, + reason=f"{POLICY_OUTPUT_ERROR_REASON}: {policy_errors[0][:200]}") + def classify(verdict: TrialVerdict, *, reward_range: tuple[float, float] = (0.0, 1.0)) -> tuple[Outcome, float, str]: """Return ``(outcome, reward, reason)`` for one verdict. diff --git a/yeto/rl/agentic/types.py b/yeto/rl/agentic/types.py index 39a6adf93..2b5bd50f7 100644 --- a/yeto/rl/agentic/types.py +++ b/yeto/rl/agentic/types.py @@ -124,6 +124,9 @@ class TokenTurn: ``Provenance.to_record()``) are what the recorder knows about the call; None when it does not record them. ``finish_reason``: the engine's reason (``stop`` / ``length``; ``length`` = the call hit its token cap). + ``policy_error``: why the recorder refused the reply as malformed model + output (e.g. a tool call that is not valid JSON); the tokens were sampled + and are kept, the agent got an error for that call. """ prompt_ids: tuple[int, ...] @@ -132,6 +135,7 @@ class TokenTurn: policy_version: int | None = None provenance: Mapping[str, str] | None = None finish_reason: str | None = None + policy_error: str | None = None def __post_init__(self) -> None: if len(self.completion_ids) != len(self.completion_logprobs): diff --git a/yeto/rl/harness/gateway/recorder.py b/yeto/rl/harness/gateway/recorder.py index 66f9ff995..4e6b85ef2 100644 --- a/yeto/rl/harness/gateway/recorder.py +++ b/yeto/rl/harness/gateway/recorder.py @@ -84,7 +84,8 @@ def turns_from_records(records: Mapping[str, Any]) -> list[TokenTurn]: completion_logprobs=tuple(float(x) for x in turn["logprobs"]), policy_version=turn.get("policy_version"), provenance=turn.get("provenance"), - finish_reason=turn.get("finish_reason"))) + finish_reason=turn.get("finish_reason"), + policy_error=turn.get("parse_error"))) cursor, pending_stop = end, len(turn.get("stop_added") or ()) if cursor + pending_stop != len(ids): raise TokenAlignmentError(f"segment {sample['segment']} has {len(ids) - cursor - pending_stop} " diff --git a/yeto/rl/harness/gateway/records.py b/yeto/rl/harness/gateway/records.py index c9b9f95da..8ceae491e 100644 --- a/yeto/rl/harness/gateway/records.py +++ b/yeto/rl/harness/gateway/records.py @@ -41,7 +41,7 @@ SamplingParams, TokenSamplerEngine, ) -from yeto.rl.harness.chat_templates import ChatTemplate, first_difference +from yeto.rl.harness.chat_templates import ChatTemplate, TemplateError, first_difference from .core import GatewayError @@ -84,7 +84,9 @@ class TurnRecord: ``completion``: the sampled ids (mask 1, stop token included when the engine returned it). ``stop_added``: a stop token the engine did not return, added with mask 0. ``render_diff``: first position where the - stream differs from a fresh tokenization (None: aligned).""" + stream differs from a fresh tokenization (None: aligned). ``parse_error``: + the completion is not valid model output for the template (the agent got + an HTTP 400 for this call; the sampled ids stay in the stream).""" turn: int kind: str @@ -96,6 +98,7 @@ class TurnRecord: provenance: dict[str, str] stop_added: tuple[int, ...] = () render_diff: int | None = None + parse_error: str | None = None @dataclass @@ -200,7 +203,12 @@ async def chat(self, session_id: str, messages: list[Message], sampling: dict[st completion = await self._sample(session, segment, pending, params, kind=KIND_ACTION) segment.n_context = len(context) + 1 prompt_tokens = len(segment.token_ids) - len(completion.token_ids) - message = self._parse(completion) + try: + message = self._parse(completion) + except TemplateError as exc: # malformed model output: the policy's failure (N10.9) + segment.turns[-1].parse_error = str(exc)[:500] + self.counters[session.trajectory_id]["policy_output_errors"] += 1 + raise return self._chat_response(session_id, message, completion, prompt_tokens=prompt_tokens) # -------------------------------------------------------------- records From e44e4ed23e69390e6a75ce4142947885976ccaa2 Mon Sep 17 00:00:00 2001 From: michael <78113554+chowhsu@users.noreply.github.com> Date: Tue, 20 Oct 2026 06:18:03 +0000 Subject: [PATCH 55/89] G1B: dense word-diversity reward plan and per-release tensor digests The cut 4-layer rehearsal model gets a constant GSM8K reward, so GRPO advantages and grad_norm were 0 and no release changed. The dense plan swaps in a word-diversity test reward and lr 3e-4 (1e-6 Adam steps fall below one bf16 ulp). A CPU hasher digests every release while the store still keeps it. Co-Authored-By: Claude Opus 5.5 --- .../stage1/glm5next-mini-modal-dense.yaml | 148 ++++++++++++++++++ tests/test_word_diversity_reward.py | 17 ++ tools/probes/s25_g1/modal_app.py | 99 +++++++++++- yeto/rl/word_diversity_reward.py | 28 ++++ 4 files changed, 290 insertions(+), 2 deletions(-) create mode 100644 examples/stage1/glm5next-mini-modal-dense.yaml create mode 100644 tests/test_word_diversity_reward.py create mode 100644 yeto/rl/word_diversity_reward.py diff --git a/examples/stage1/glm5next-mini-modal-dense.yaml b/examples/stage1/glm5next-mini-modal-dense.yaml new file mode 100644 index 000000000..28230adfa --- /dev/null +++ b/examples/stage1/glm5next-mini-modal-dense.yaml @@ -0,0 +1,148 @@ +# S25 T-G1B: glm5next-mini-modal.yaml with two changes, so the weights move and the +# chain "trained weights -> re-quantized MXFP4 experts -> sampler" is checked by value: +# task.reward yeto.rl.word_diversity_reward:score (GSM8K is constant 0 on this model, +# so GRPO advantages and grad_norm were 0 in T-G1 run-17) +# trainer.miles.lr 3.0e-4 (1e-6 Adam steps are below one bf16 ulp of most weights, +# so a release hash could stay the same even with a non-zero gradient) +# Everything else (nodes, store, images, sampler, sampling) is the T-G1 plan. +# +# Stage-1 mini rehearsal (S25 T-G1) on Modal H100: the same plan schema and +# the same generated commands as glm53-flash-ppu.yaml; only the nodes +# (Modal containers) and the store (a Modal Volume through the yeto store +# hook, not OSS) differ. Model: Miles' 4-layer GLM5.3-Flash CI checkpoint cut +# to 32 routed experts per MoE layer (tools/probes/s25_g1/make_mini.py). +# +# yeto-stage1 plan examples/stage1/glm5next-mini-modal.yaml +# modal run tools/probes/s25_g1/modal_app.py::... (the driver reads this plan) + +run_id: s25-g1-mini +inventory: modal-mini-inventory.yaml +rounds: 3 +seed: 17 + +rendezvous: + modal-trainer: 127.0.0.1 + modal-sampler: 127.0.0.1 + +roles: + trainer: [modal-trainer] + sampler: [modal-sampler] + +model: + name: GLM5.3-Flash-4layer-e32 + architecture: glm5_next + num_layers: 4 # KDA, DSA, KDA, DSA + first_k_dense_replace: 1 + num_experts: 32 # cut from 288 (make_mini.py) + params_b: 5.3 # 推算: 10.5 GB bf16 / 2 + expert_params_b: 2.4 # 3 MoE layers x 32 x 3 x 4096 x 2048 + bf16_source: + revision: 7fb7883207fe73ea04430ac392f6ffe2b9f936c4 # CharyZeng/GLM-5.3-Flash-4layer, then cut + origin: hf://CharyZeng/GLM-5.3-Flash-4layer@7fb7883207fe73ea04430ac392f6ffe2b9f936c4 + staged: modal-volume://yeto-s25-g1-models/glm5next-mini-bf16-7fb7883 + local_dir: /data/models/glm5next-mini-bf16-7fb7883 + torch_dist: /data/models/glm5next-mini-bf16-7fb7883_torch_dist # Miles converter output (convert step) + sampler_initial: # round 0: our N8.5 converter on the bf16 mini model + repo: local/glm5next-mini-mxfp4-r0-7fb7883 + revision: 7fb7883 + staged: modal-volume://yeto-s25-g1-models/glm5next-mini-mxfp4-r0-7fb7883 + local_dir: /data/models/glm5next-mini-mxfp4-r0-7fb7883 + +quant: + scheme: mxfp4-g32-rtn + format: mxfp4-pack-quantized + release_dir: policy-releases-mxfp4-g32-rtn + expected_expert_matrices: 288 # 3 MoE layers x 32 experts x (gate, up, down) + frozen_pattern: '^model\.visual\.' # Miles trains the text model only; the vision tower stays the base + +store: /data/store/s25-g1 +store_backend: modal-volume # the only difference from OSS: the store hook and its Volume +store_volume: yeto-s25-g1-store +store_access: + head: direct + +images: + sampler: ghcr.io/michaellchung/yeto-miles-ports@sha256:5b9dc2a9e6e315d73dabfbc6faae428b59cc01aabec6189645a5c8c80d6e3721 + trainer: ghcr.io/michaellchung/yeto-miles-ports@sha256:5b9dc2a9e6e315d73dabfbc6faae428b59cc01aabec6189645a5c8c80d6e3721 + docker_device_args: [--gpus, all] # unused on Modal; the containers get the GPUs + +trainer: + backend: miles + miles: + commit: 1618b5c38db80e17ec928cfa47fe832cb18565fc # s25-distopt, checked out in the image + megatron: radixark/Megatron-LM@a84b105 + mbridge: image + workdir: /root/miles + extra_env: + # Megatron's get_model imports megatron.post_training (modelopt check), which the + # image's editable megatron_core install does not expose (Modal run-10) + PYTHONPATH: /root/Megatron-LM:/root/sky_workdir + tp: 1 + sequence_parallel: false + pp: 1 + pp_layers: [4] + cp: 1 + ep: 2 + etp: 1 + optimizer_offload_fraction: 1.0 + memory_estimate: # 推算; `plan` checks them + weights_grads_gib_per_card: 22.9 + host_optimizer_gib_per_node: 59.2 + transformer_impl: transformer_engine + recompute: full + max_tokens_per_gpu: 8192 + lr: 3.0e-4 # T-G1B: see the header + verl: + commit: null + strategy: fsdp2 + fsdp_size: 2 + ulysses: 1 + param_offload: false + optimizer_offload: true + lr: 1.0e-6 + +sampler: + engine: sglang + engine_version: 0.5.21.dev70+g2fa8801 + python: python3 + topology: tp1 + tp: 1 + precision: mxfp4 + max_model_len: 4096 + port: 30000 + moe_runner_backend: triton + hierarchical_cache: false + nextn: false + renames_weight_packed: true # the yeto hook in the image renames (sglang_load_report) + load_report: true # the hook reports loaded expert matrices (post-load check 1) + in_place_reload: false + # SGLang fuses the BF16 shared expert into the routed FusedMoE as expert 32; under the + # MXFP4 MoE method that slot is packed U8, so the BF16 shared weights do not fit (smoke-2) + extra_args: [--disable-shared-experts-fusion, --cuda-graph-max-bs-decode, "16"] # 16 = 4 prompts x 4 samples; capture took 118 s at the default sizes + extra_env: + PYTHONPATH: /root/sky_workdir + checks: + require_fixed_answers: false # a cut 4-layer model cannot answer GSM8K + fixed_max_tokens: 64 + +sampling: + max_tokens: 512 + max_prompt_len: 512 + temperature: 1.0 + top_p: 1.0 + top_k: -1 + repetition_penalty: 1.0 + frequency_penalty: 0.0 + presence_penalty: 0.0 + +task: + data: zhuzilin/gsm8k + data_revision: 0cbd9f31d91ac21a7613dcbc7fef992adac459ae + reward: yeto.rl.word_diversity_reward:score # T-G1B: see the header + rollout_batch_size: 4 + n_samples_per_prompt: 4 + prompts_jsonl: /data/models/datasets/gsm8k-0cbd9f3/train.jsonl + data_local: /data/models/datasets/gsm8k-0cbd9f3/train.jsonl # the same rows for the trainer (offline nodes) + input_key: messages + label_key: label + prompt_format: chat diff --git a/tests/test_word_diversity_reward.py b/tests/test_word_diversity_reward.py new file mode 100644 index 000000000..a0c1de2e4 --- /dev/null +++ b/tests/test_word_diversity_reward.py @@ -0,0 +1,17 @@ +import asyncio +import types + +from yeto.rl.word_diversity_reward import diversity, score + + +def test_diversity_values(): + assert diversity("") == 0.0 + assert diversity("unload unload unload unload") == 0.25 + assert diversity("a b c d") == 1.0 + + +def test_score_has_variance_on_truncated_degenerate_samples(): + texts = ["Ref unload unload unload", "adam unload unloadimers unloadimers"] + values = [asyncio.run(score(None, types.SimpleNamespace(response=t))) for t in texts] + assert values[0] != values[1] + assert asyncio.run(score(None, types.SimpleNamespace(response=None))) == 0.0 diff --git a/tools/probes/s25_g1/modal_app.py b/tools/probes/s25_g1/modal_app.py index 6eb1992cf..dee8bc043 100644 --- a/tools/probes/s25_g1/modal_app.py +++ b/tools/probes/s25_g1/modal_app.py @@ -31,7 +31,9 @@ HERE = Path(__file__).resolve() REPO = HERE.parents[3] if len(HERE.parents) > 3 else Path("/") # in the container the module sits at /root WORKDIR = "/root/sky_workdir" -PLAN = "examples/stage1/glm5next-mini-modal.yaml" +# T-G1B: YETO_G1_PLAN=examples/stage1/glm5next-mini-modal-dense.yaml (set at deploy time; +# the image env carries it into every container, so all of them plan the same file) +PLAN = os.environ.get("YETO_G1_PLAN", "examples/stage1/glm5next-mini-modal.yaml") IMAGE = ("ghcr.io/michaellchung/yeto-miles-ports@sha256:" "5b9dc2a9e6e315d73dabfbc6faae428b59cc01aabec6189645a5c8c80d6e3721") MILES_COMMIT = "1618b5c38db80e17ec928cfa47fe832cb18565fc" # s25-distopt @@ -64,7 +66,7 @@ f"python3 -c \"{_PATCH_SGLANG}\"", ) # not PYTHONPATH: Modal's own runtime path must stay intact; subprocesses get it (_env) - .env({"HOME": "/root", "PYTHONUNBUFFERED": "1", "HF_HOME": "/data/models/hf"}) + .env({"HOME": "/root", "PYTHONUNBUFFERED": "1", "HF_HOME": "/data/models/hf", "YETO_G1_PLAN": PLAN}) .add_local_dir(str(REPO), WORKDIR, copy=False, ignore=[".git", ".g1-build", ".venv", ".claude", "syncer/target", "**/__pycache__", "**/*.pyc", "evidence*", "openspec", "data"]) @@ -329,6 +331,94 @@ def requant_island(tag: str) -> int: return rc +# --------------------------------------------------------------------- release digests (T-G1B) +def _tensor_digests(directory: Path) -> dict[str, str]: + """sha256 of every tensor's bytes in the safetensors files of ``directory``.""" + import hashlib + + out: dict[str, str] = {} + for file in sorted(directory.glob("*.safetensors")): + with open(file, "rb") as f: + size = int.from_bytes(f.read(8), "little") + header = json.loads(f.read(size)) + base = 8 + size + for name, meta in sorted(header.items(), key=lambda kv: kv[1]["data_offsets"][0] + if kv[0] != "__metadata__" else -1): + if name == "__metadata__": + continue + begin, end = meta["data_offsets"] + f.seek(base + begin) + digest, left = hashlib.sha256(), end - begin + while left: + chunk = f.read(min(left, 1 << 24)) + digest.update(chunk) + left -= len(chunk) + out[name] = digest.hexdigest() + return out + + +def _summarize(digests: dict[str, str]) -> dict: + import hashlib + + def combined(names: list[str]) -> str: + h = hashlib.sha256() + for name in names: + h.update(f"{name}={digests[name]}\n".encode()) + return h.hexdigest() + + experts = sorted(n for n in digests if ".mlp.experts." in n) + others = sorted(n for n in digests if ".mlp.experts." not in n) + return {"tensors": len(digests), "expert_tensors": len(experts), "experts_sha256": combined(experts), + "others_sha256": combined(others)} + + +@app.function(image=image, volumes=VOLUMES, cpu=4, memory=8192, timeout=3 * 3600) +def release_hasher(tag: str, base_dir: str, store_root: str, release_dirs: list[str], poll_s: float = 20.0) -> int: + """Digest each release as soon as it is complete (the store keeps only the newest 2): + per tensor sha256, plus combined digests of the routed-expert tensors and of the rest. + Writes ``/release-digests.jsonl`` and ``/tensors//.json``.""" + log = _log_dir(tag) + out = log / "release-digests.jsonl" + seen: set[str] = set() + + def record(kind: str, name: str, directory: Path, extra: dict) -> None: + began = time.time() + try: + digests = _tensor_digests(directory) + except Exception as error: # noqa: BLE001 - pruned or still open: record, go on + row = {"kind": kind, "release": name, "error": repr(error)} + else: + target = log / "tensors" / kind / f"{name}.json" + target.parent.mkdir(parents=True, exist_ok=True) + target.write_text(json.dumps(digests, indent=0, sort_keys=True)) + row = {"kind": kind, "release": name, **extra, **_summarize(digests), + "hash_s": round(time.time() - began, 1)} + row["time_unix"] = time.time() + with open(out, "a", encoding="utf-8") as f: + f.write(json.dumps(row) + "\n") + print(f"[g1b] {json.dumps(row)}", flush=True) + _commit(tag) + + record("base", Path(base_dir).name, Path(base_dir), {}) + while True: + try: + store.reload() + except Exception as error: # noqa: BLE001 - an open file blocks the reload; retry + print(f"[g1b] reload failed: {error}", flush=True) + for kind in release_dirs: + root = Path(store_root) / kind + # RELEASE.json is written after the weights (weights -> record -> commit -> LATEST) + for marker in sorted(root.glob("*/RELEASE.json")) if root.is_dir() else (): + name = marker.parent.name + if f"{kind}/{name}" in seen: + continue + seen.add(f"{kind}/{name}") + meta = json.loads(marker.read_text()) + record(kind, name, root / name, {"version": meta.get("version"), + "policy_hash": meta.get("policy_hash")}) + time.sleep(poll_s) + + @app.function(image=image, volumes=VOLUMES, cpu=1, memory=1024, timeout=3 * 3600) def orchestrate(tag: str, cut_after: int = 2) -> dict: """Sampler for the whole run; trainer A stops after ``cut_after`` rounds (cut), @@ -340,6 +430,10 @@ def orchestrate(tag: str, cut_after: int = 2) -> dict: store.commit() sampler = sampler_island.spawn(f"{tag}/sampler") requant = requant_island.spawn(f"{tag}/requant") + plan = _plan(_log_dir(f"{tag}/orchestrate")) + spec = plan["sampler"]["engines"][0]["spec"] + hasher = release_hasher.spawn(f"{tag}/hashes", str(spec["base_dir"]), str(spec["store"]), + ["policy-releases", str(spec["release_dir"])]) sampler_rc: object = None def sampler_done() -> bool: @@ -377,6 +471,7 @@ def run_trainer(sub: str, stop: int | None) -> object: time.sleep(30) time.sleep(90) # the re-quantizer's last log mirror requant.cancel() + hasher.cancel() result = {"trainer_a": first, "trainer_b": second, "sampler": sampler_rc} print(json.dumps(result, default=str), flush=True) out = Path(f"/data/logs/{tag}") diff --git a/yeto/rl/word_diversity_reward.py b/yeto/rl/word_diversity_reward.py new file mode 100644 index 000000000..9eac4b51e --- /dev/null +++ b/yeto/rl/word_diversity_reward.py @@ -0,0 +1,28 @@ +"""Dense word-diversity reward for rehearsals on models that cannot solve the task. + +A test reward, not a product reward. The cut 4-layer GLM5.3-Flash rehearsal model +(S25 T-G1) never answers GSM8K and every response runs to ``max_tokens``: the +GSM8K reward and the length reward are constant within a group, so GRPO +advantages, the loss gradient and grad_norm are all 0 and the weights never move. +This reward depends on the text itself, so samples of one prompt differ even when +all of them are truncated: + + score = distinct whitespace words / words (0.0 for an empty response) + +Range [0, 1]; a response that repeats one word scores low. Launch with + + --reward-function yeto.rl.word_diversity_reward:score +""" + +from __future__ import annotations + +__all__ = ["diversity", "score"] + + +def diversity(text: str) -> float: + words = text.split() + return len(set(words)) / len(words) if words else 0.0 + + +async def score(args, sample, **kwargs) -> float: + return diversity(sample.response or "") From 12b7cd7821e2791ef708c4ae1be0eb92947e65ed Mon Sep 17 00:00:00 2001 From: michael <78113554+chowhsu@users.noreply.github.com> Date: Tue, 20 Oct 2026 06:18:47 +0000 Subject: [PATCH 56/89] Miles Harbor path with the records proxy as recorder (N10.9) records_generate.generate runs one Harbor trial per sample under the yeto outer layer; the agent calls the records proxy, which samples token ids from Miles' SGLang router (SGLangRouterEngine). The recorded calls are joined into the Miles sample (tokens, loss mask, rollout log-probs) and the verifier reward is set there. reward_post_process normalises each GRPO group without its ABORTED members and marks them remove_sample, so groups are kept and an ABORTED sample never starts another sandbox. Co-Authored-By: Claude Opus 5.5 --- tests/test_miles_records_generate.py | 102 +++++++ yeto/rl/adapters/miles/records_generate.py | 323 +++++++++++++++++++++ 2 files changed, 425 insertions(+) create mode 100644 tests/test_miles_records_generate.py create mode 100644 yeto/rl/adapters/miles/records_generate.py diff --git a/tests/test_miles_records_generate.py b/tests/test_miles_records_generate.py new file mode 100644 index 000000000..22a1a86e2 --- /dev/null +++ b/tests/test_miles_records_generate.py @@ -0,0 +1,102 @@ +"""S25 N10.9: the Miles Harbor path with the records proxy as recorder (CPU). + +Real pieces: ``records_generate.run_sample`` (outer layer, outcome rules, +``join_turns``, sample filling), ``ProxyTokenRecorder`` over HTTP, the gateway +and recording backend, ``SGLangRouterEngine`` (its HTTP call replaced by the +scripted model). Fakes: the agent (Codex-like Responses loop), Miles' Sample. +""" + +from __future__ import annotations + +import asyncio +import copy +from types import SimpleNamespace + +import pytest +from test_records_proxy import SCRIPT, SIGNED, ScriptedModel, _template +from test_verl_rollout_engine import IDENTITY, ExitOnErrorAgent, MalformedSecondReply +from test_verl_records_recorder import CodexLikeAgent + +from yeto.rl.adapters.miles import records_generate as rg +from yeto.rl.agentic import OUTCOME_KEY +from yeto.rl.harness.gateway.proxy import RecordsProxy +from yeto.rl.harness.gateway.recorder import ProxyTokenRecorder + +pytest.importorskip("aiohttp") + +MILES_SAMPLING = {"temperature": 1.0, "top_p": 1.0, "top_k": -1, "max_new_tokens": 256} +ARGS = SimpleNamespace(rollout_max_response_len=4096, n_samples_per_prompt=4, grpo_std_normalization=True) + + +@pytest.fixture(scope="module") +def qwen(): + return _template("qwen3") + + +def _router_http(model: ScriptedModel, calls: list): + def http(url: str, body: dict) -> dict: + calls.append(body) + ids = model(tuple(body["input_ids"]), 0) + return {"text": "", "meta_info": {"output_token_logprobs": [[-0.5, t, None] for t in ids], + "finish_reason": {"type": "stop"}}} + return http + + +def _run(qwen, agent, model, monkeypatch): + monkeypatch.setenv("HARBOR_ENV_TYPE", "boat") + calls: list = [] + engine = rg.SGLangRouterEngine("http://router", max_model_len=8192, identity=IDENTITY, + http=_router_http(model, calls)) + recorder = ProxyTokenRecorder(RecordsProxy(engine, qwen, model="policy", signed_sampling=SIGNED, max_chains=1)) + sample = SimpleNamespace(metadata={"instance_id": "fix-git", "agent_name": "codex"}, tokens=[], + response_length=0, loss_mask=None, rollout_log_probs=None, status=None, + reward=None, remove_sample=False) + + async def go(): + try: + return await rg.run_sample(sample, MILES_SAMPLING, ARGS, recorder=recorder, harness=agent) + finally: + await recorder.aclose() + + return asyncio.run(go()), calls + + +def test_sample_carries_the_tokens_the_router_sampled(qwen, monkeypatch): + out, calls = _run(qwen, CodexLikeAgent(), ScriptedModel(qwen, copy.deepcopy(SCRIPT)), monkeypatch) + assert out.metadata[OUTCOME_KEY] == "ok" and out.reward == 1.0 and out.status.name == "COMPLETED" + assert len(calls) == 3 and calls[0]["return_logprob"] is True + assert calls[0]["sampling_params"]["stop_token_ids"] == list(qwen.stop_token_ids) + assert len(out.loss_mask) == len(out.rollout_log_probs) == out.response_length + assert out.tokens[:len(calls[-1]["input_ids"])] == calls[-1]["input_ids"] # trained == sampled + assert all(lp == -0.5 for lp, m in zip(out.rollout_log_probs, out.loss_mask) if m) + assert out.metadata["model_calls"] == 3 and out.metadata["truncated"] is False + + +def test_malformed_reply_is_a_trained_zero(qwen, monkeypatch): + out, _ = _run(qwen, ExitOnErrorAgent(), MalformedSecondReply(qwen, copy.deepcopy(SCRIPT)), monkeypatch) + assert out.metadata[OUTCOME_KEY] == "policy_limit" and out.reward == 0.0 + assert sum(out.loss_mask) > 0 and not out.remove_sample + + +def test_infrastructure_exit_is_aborted_and_never_trained(qwen, monkeypatch): + out, _ = _run(qwen, ExitOnErrorAgent("EnvironmentStartTimeoutError"), + MalformedSecondReply(qwen, copy.deepcopy(SCRIPT)), monkeypatch) + assert out.metadata[OUTCOME_KEY] == "aborted" and out.status.name == "ABORTED" + assert out.remove_sample and out.loss_mask == [0] and out.reward == 0.0 + + +def _s(reward, outcome="ok"): + return SimpleNamespace(reward=reward, metadata={OUTCOME_KEY: outcome}, remove_sample=False) + + +def test_reward_post_process_leaves_aborted_out_of_the_group_baseline(): + samples = [_s(1.0), _s(0.0), _s(0.0, "aborted"), _s(0.0, "policy_limit"), # group 1 + _s(0.0), _s(0.0), _s(0.0), _s(0.0)] # group 2: no variance + raw, adv = rg.reward_post_process(ARGS, samples) + assert raw == [1.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0] + mean, std = 1 / 3, (((2 / 3) ** 2 + 2 * (1 / 3) ** 2) / 2) ** 0.5 + assert adv[0] == pytest.approx((1 - mean) / (std + 1e-6)) and adv[1] == pytest.approx(-mean / (std + 1e-6)) + assert adv[2] == 0.0 and samples[2].remove_sample and not samples[0].remove_sample + assert adv[4:] == [0.0] * 4 # all-zero group: no gradient (expected for most TB2.1 groups) + with pytest.raises(ValueError, match="whole groups"): + rg.reward_post_process(ARGS, samples[:3]) diff --git a/yeto/rl/adapters/miles/records_generate.py b/yeto/rl/adapters/miles/records_generate.py new file mode 100644 index 000000000..561d3a30b --- /dev/null +++ b/yeto/rl/adapters/miles/records_generate.py @@ -0,0 +1,323 @@ +"""Miles Harbor path with the records proxy as the token recorder (S25 N10.9). + +Miles' own Harbor example records tokens through Miles' session server, which +speaks Chat Completions only; Harbor's codex agent (0.145.0) speaks Responses +only (T-G2). On this path the yeto records proxy (N10.5) is the recorder, as +on verl: + +* ``--custom-generate-function-path yeto.rl.adapters.miles.records_generate.generate``: + one Harbor trial per sample under the yeto outer layer; the agent calls the + records proxy, the proxy sends token ids to Miles' SGLang router + (:class:`SGLangRouterEngine`), and the recorded calls become the Miles + sample (``join_turns``: tokens, loss mask, rollout log-probs). The reward + is the Harbor verifier's and is set here, so Miles' reward hook is not called. +* ``--custom-reward-post-process-path yeto.rl.adapters.miles.records_generate.reward_post_process``: + GRPO normalisation over each prompt group without its ABORTED members; an + ABORTED sample has advantage 0 and ``remove_sample`` (loss mask 0). Groups + are kept (no dynamic filter), so an ABORTED sample never triggers a new + sandbox. + +A malformed model reply that ends the agent is a trained 0 (POLICY_LIMIT), +not ABORTED (:func:`~yeto.rl.agentic.outcome.policy_output_failure`). + +Rollout-process environment (besides the Harbor/boat variables of +``harbor_agent``): + YETO_RECORDS_PROXY ``yeto.rl.adapters.miles.records_generate:records_proxy_from_env`` + YETO_RECORDS_PROXY_HOST/PORT/PUBLIC_BASE/SECRET_FILE (``gateway.recorder``) + YETO_RECORDS_TEMPLATE, YETO_RECORDS_TOKENIZER, YETO_RECORDS_MODEL, + YETO_RECORDS_SAMPLING, YETO_RECORDS_MAX_MODEL_LEN, YETO_RECORDS_CARD + (same meaning as ``adapters.verl.rollout_engine``) + YETO_SGLANG_ROUTER_URL ``http://:`` of Miles' router (the launcher sets it) + +Policy version: Miles' synchronous loop updates the engines between +rollouts, never during one; the proxy records version None and the trainer's +rollout id is the version record (as on verl). +""" + +from __future__ import annotations + +import contextlib +import json +import os +from collections.abc import Sequence +from functools import partial +from typing import Any + +from yeto.rl.agentic import ( + OUTCOME_KEY, + AgentTrajectory, + HarnessConfigError, + Outcome, + TokenAlignmentError, + join_turns, + run_agent_task, +) +from yeto.rl.agentic.loop_local import LoopLocal +from yeto.rl.agentic.outcome import policy_output_failure +from yeto.rl.engine.sampler_engine import ( + FINISH_LENGTH, + FINISH_STOP, + UNKNOWN, + Completion, + EngineIdentity, + LoadReport, + SamplerEngineError, + SamplingParams, +) +from yeto.rl.harness.harbor import HarborHarness, HarborSettings, build_trial_config +from yeto.rl.harness.harbor.trial import CUSTOM_PROVIDERS + +from .harbor_agent import gate, task_from_metadata + +ROUTER_URL_ENV = "YETO_SGLANG_ROUTER_URL" +RECORDER_ENV = "YETO_AGENTIC_RECORDER" +DEFAULT_RECORDER = "yeto.rl.harness.gateway.recorder:recorder_from_env" +GENERATE_FUNCTION = "yeto.rl.adapters.miles.records_generate.generate" +REWARD_POST_PROCESS = "yeto.rl.adapters.miles.records_generate.reward_post_process" +HTTP_TIMEOUT_S = 1800.0 + + +# ---------------------------------------------------------------------- engine +class SGLangRouterEngine: + """``TokenSamplerEngine`` over an SGLang router (or server) someone else runs (Miles).""" + + def __init__(self, url: str, *, max_model_len: int, identity: EngineIdentity, http: Any = None) -> None: + self.url = url.rstrip("/") + self._max_model_len = int(max_model_len) + self._identity = identity + self._http = http + + @property + def identity(self) -> EngineIdentity: + return self._identity + + @property + def max_model_len(self) -> int: + return self._max_model_len + + def served_version(self) -> int | None: + return None + + def _post(self, body: dict[str, Any]) -> Any: + if self._http is not None: + return self._http(f"{self.url}/generate", body) + import httpx + + reply = httpx.post(f"{self.url}/generate", json=body, timeout=HTTP_TIMEOUT_S) + reply.raise_for_status() + return reply.json() + + def generate_tokens(self, prompts: Sequence[Sequence[int]], sampling: SamplingParams, *, + stop_token_ids: Sequence[int] = ()) -> list[list[Completion]]: + """Called from a worker thread (the records backend uses ``asyncio.to_thread``).""" + params: dict[str, Any] = { + "temperature": float(sampling.temperature), "top_p": float(sampling.top_p), + "top_k": int(sampling.top_k), "repetition_penalty": float(sampling.repetition_penalty), + "max_new_tokens": int(sampling.max_tokens), "skip_special_tokens": False, + "no_stop_trim": True, + } + if stop_token_ids: + params["stop_token_ids"] = [int(t) for t in stop_token_ids] + out: list[list[Completion]] = [] + for prompt in prompts: + group = [] + for _ in range(int(sampling.n)): + body = self._post({"input_ids": [int(t) for t in prompt], "sampling_params": params, + "return_logprob": True}) + body = body[0] if isinstance(body, list) else body + meta = body.get("meta_info") or {} + pairs = meta.get("output_token_logprobs") or [] + tokens = tuple(int(p[1]) for p in pairs) + finish = (meta.get("finish_reason") or {}).get("type") + if finish == "abort": + raise SamplerEngineError(f"sglang aborted the request: {meta.get('finish_reason')}") + group.append(Completion(token_ids=tokens, logprobs=tuple(float(p[0]) for p in pairs), + text=str(body.get("text", "")), + finish_reason=FINISH_LENGTH if finish == "length" else FINISH_STOP, + prompt_token_ids=tuple(int(t) for t in prompt))) + out.append(group) + return out + + # Miles starts, updates and stops its engines: the proxy never calls these. + def start(self, release_dir: Any, version: int) -> LoadReport: + raise SamplerEngineError("Miles owns its rollout engines") + + def load(self, release_dir: Any, version: int) -> LoadReport: + raise SamplerEngineError("Miles owns its rollout engines") + + def generate(self, prompts: Sequence[str], sampling: SamplingParams) -> list[list[Completion]]: + raise SamplerEngineError("text prompts are not used; the records proxy sends token ids") + + def score(self, prompt: str, continuation: str) -> list[float]: + raise SamplerEngineError("scoring is not used on the Miles rollout") + + def loaded_expert_matrices(self) -> int | None: + return None + + def weight_gib_per_card(self) -> float | None: + return None + + def stop(self) -> None: + return None + + +def records_proxy_from_env() -> Any: + """``YETO_RECORDS_PROXY`` callable: a RecordsProxy over Miles' router (module doc).""" + from yeto.rl.adapters.verl.rollout_engine import ( + CARD_ENV, + MAX_MODEL_LEN_ENV, + MODEL_ENV, + SAMPLING_ENV, + TEMPLATE_ENV, + TOKENIZER_ENV, + ) + from yeto.rl.harness.chat_templates import ChatTemplate + from yeto.rl.harness.gateway.proxy import RecordsProxy + + try: + url = os.environ[ROUTER_URL_ENV] + template = ChatTemplate.load(os.environ[TEMPLATE_ENV], tokenizer_path=os.environ[TOKENIZER_ENV]) + signed = json.loads(os.environ[SAMPLING_ENV]) + engine = SGLangRouterEngine(url, max_model_len=int(os.environ[MAX_MODEL_LEN_ENV]), identity=EngineIdentity( + engine="miles-sglang-router", engine_version=UNKNOWN, device_family="cuda", + card_model=os.environ.get(CARD_ENV, UNKNOWN), precision="bf16", quant_scheme="bf16", + topology=UNKNOWN)) + model = os.environ[MODEL_ENV] + except (KeyError, ValueError) as exc: + raise HarnessConfigError(f"records proxy over Miles' router: bad or missing setting {exc}") from exc + return RecordsProxy(engine, template, model=model, signed_sampling=signed, sampler_id="miles-records") + + +# ---------------------------------------------------------------------- recorder +def _recorder_from_env() -> Any: + import importlib + + spec = os.environ.get(RECORDER_ENV, DEFAULT_RECORDER) + module, _, attr = spec.partition(":") + return getattr(importlib.import_module(module), attr)() + + +_recorders: LoopLocal[Any] = LoopLocal(_recorder_from_env) + + +def _harness(task: Any) -> HarborHarness: + if task.provider not in CUSTOM_PROVIDERS: + raise HarnessConfigError(f"records_generate runs yeto's TrialConfig builder; provider {task.provider!r} " + f"is not one of {sorted(CUSTOM_PROVIDERS)} (use Miles' example path)") + return HarborHarness(partial(build_trial_config, settings=HarborSettings.from_env())) + + +# ---------------------------------------------------------------------- generate +async def generate(input: Any) -> Any: + """Miles ``--custom-generate-function-path`` (module doc).""" + from miles.rollout.base_types import GenerateFnOutput + + sample = await run_sample(input.sample, input.sampling_params, input.args, recorder=_recorders.get()) + return GenerateFnOutput(samples=sample) + + +async def run_sample(sample: Any, sampling_params: dict[str, Any], args: Any, *, recorder: Any, + harness: Any = None) -> Any: + """One Harbor trial -> the filled Miles sample (separate from ``generate`` for CPU tests).""" + task = task_from_metadata(dict(sample.metadata or {})) + session = f"miles-{task.task_id}-{os.urandom(6).hex()}" + endpoint = await recorder.open(session, _verl_style_sampling(sampling_params)) + try: + trajectory = await run_agent_task(task, harness or _harness(task), endpoint, gate()) + except BaseException: + with contextlib.suppress(Exception): # the trial's error is the one to report + await recorder.close(session) + raise + turns: Sequence[Any] = () + try: + turns = await recorder.close(session) + except TokenAlignmentError as exc: + trajectory = _reclassify(trajectory, Outcome.ABORTED, f"token_alignment: {exc}") + else: + updated = policy_output_failure(trajectory, [t.policy_error for t in turns if t.policy_error]) + if updated.outcome is not trajectory.outcome: + gate().metrics.reclassify(trajectory.outcome, updated.outcome) + trajectory = updated + segments = None + if trajectory.outcome is not Outcome.ABORTED: + try: + segments = join_turns(turns, max_response_len=int(args.rollout_max_response_len)) + except TokenAlignmentError as exc: + trajectory = _reclassify(trajectory, Outcome.ABORTED, f"token_alignment: {exc}") + return fill_sample(sample, trajectory, segments, turns) + + +def _reclassify(trajectory: AgentTrajectory, outcome: Outcome, reason: str) -> AgentTrajectory: + if trajectory.outcome is not outcome: + gate().metrics.reclassify(trajectory.outcome, outcome) + return AgentTrajectory(task=trajectory.task, outcome=outcome, reward=0.0 if outcome is Outcome.ABORTED + else trajectory.reward, reason=reason, verdict=trajectory.verdict, + policy_version=trajectory.policy_version) + + +def _verl_style_sampling(sampling_params: dict[str, Any]) -> dict[str, Any]: + """Miles' SGLang sampling keys -> the recorder's check keys (temperature, top_p, top_k).""" + return {k: sampling_params[k] for k in ("temperature", "top_p", "top_k", "repetition_penalty") + if k in sampling_params} + + +def fill_sample(sample: Any, trajectory: AgentTrajectory, segments: Any, turns: Sequence[Any]) -> Any: + status = _status() + metadata = {**(sample.metadata or {}), **trajectory.to_metadata(), + "length_capped_turns": sum(t.finish_reason == FINISH_LENGTH for t in turns), + "model_calls": len(turns)} + if segments is None: # ABORTED: one masked placeholder token, never trained + sample.tokens, sample.response_length = [0, 0], 1 + sample.loss_mask, sample.rollout_log_probs = [0], [0.0] + sample.remove_sample = True + sample.status = status.ABORTED + sample.reward = 0.0 + else: + sample.tokens = [*segments.prompt_ids, *segments.response_ids] + sample.response_length = len(segments.response_ids) + sample.loss_mask = list(segments.response_mask) + sample.rollout_log_probs = list(segments.response_logprobs) + sample.status = status.TRUNCATED if segments.truncated else status.COMPLETED + sample.reward = float(trajectory.reward) + metadata["truncated"] = bool(segments.truncated) + sample.metadata = metadata + return sample + + +def _status() -> Any: + try: + from miles.utils.types import Sample + + return Sample.Status + except ImportError: # CPU tests without Miles + from enum import Enum + + return Enum("Status", "PENDING COMPLETED TRUNCATED ABORTED FAILED") + + +# ---------------------------------------------------------------------- rewards +def is_aborted(sample: Any) -> bool: + return (sample.metadata or {}).get(OUTCOME_KEY) == Outcome.ABORTED.value + + +def reward_post_process(args: Any, samples: Sequence[Any]) -> tuple[list[float], list[float]]: + """GRPO over each group of ``n_samples_per_prompt`` without its ABORTED members (module doc).""" + n = int(args.n_samples_per_prompt) + if len(samples) % n: + raise ValueError(f"{len(samples)} samples are not whole groups of {n}") + by_std = bool(getattr(args, "grpo_std_normalization", True)) + raw = [0.0 if is_aborted(s) else float(s.reward) for s in samples] + out = [0.0] * len(samples) + for start in range(0, len(samples), n): + members = [i for i in range(start, start + n) if not is_aborted(samples[i])] + if not members: + continue + values = [raw[i] for i in members] + mean = sum(values) / len(values) + std = (sum((v - mean) ** 2 for v in values) / (len(values) - 1)) ** 0.5 if len(values) > 1 else 0.0 + for i in members: + out[i] = (raw[i] - mean) / (std + 1e-6) if by_std else raw[i] - mean + for sample in samples: + if is_aborted(sample): + sample.remove_sample = True + return raw, out From f1584009ec81944aa35535c060f9c0521d8c1cc3 Mon Sep 17 00:00:00 2001 From: michael <78113554+chowhsu@users.noreply.github.com> Date: Tue, 20 Oct 2026 06:29:29 +0000 Subject: [PATCH 57/89] Add the RL capability evaluation suite (tools/eval) One repeatable Modal evaluation per capability (E1 island rollout+trainer, E2 sampler island reclaim/replace, E3 synced trainers -> sampler island, E4 async rollout + off-policy correction + over-age drop, E5 SCHEDULED stop + resume) and backend (Miles, verl), with one CLI (python -m tools.eval list|run|judge|report), one verdict format and a Markdown/HTML report. verl cells run a direct Modal harness derived from the S25 C2/C3V/C4 harnesses; Miles cells run the yeto launcher through scripts derived from the S25 C1/C3 and S24 L5.4 run scripts. The judges replay the pre-registered criteria of those tasks and pass on their recorded GPU data. Co-Authored-By: Claude Opus 5.5 --- tests/test_tools_eval.py | 91 +++++++ tools/eval/__init__.py | 22 ++ tools/eval/__main__.py | 152 +++++++++++ tools/eval/cells.py | 86 ++++++ tools/eval/evidence.py | 197 ++++++++++++++ tools/eval/judges.py | 304 +++++++++++++++++++++ tools/eval/miles.py | 119 ++++++++ tools/eval/miles_head.sh | 116 ++++++++ tools/eval/miles_local.sh | 101 +++++++ tools/eval/modal_verl.py | 552 ++++++++++++++++++++++++++++++++++++++ tools/eval/report.py | 93 +++++++ tools/eval/verdict.py | 109 ++++++++ 12 files changed, 1942 insertions(+) create mode 100644 tests/test_tools_eval.py create mode 100644 tools/eval/__init__.py create mode 100644 tools/eval/__main__.py create mode 100644 tools/eval/cells.py create mode 100644 tools/eval/evidence.py create mode 100644 tools/eval/judges.py create mode 100644 tools/eval/miles.py create mode 100755 tools/eval/miles_head.sh create mode 100644 tools/eval/miles_local.sh create mode 100644 tools/eval/modal_verl.py create mode 100644 tools/eval/report.py create mode 100644 tools/eval/verdict.py diff --git a/tests/test_tools_eval.py b/tests/test_tools_eval.py new file mode 100644 index 000000000..f5685860a --- /dev/null +++ b/tests/test_tools_eval.py @@ -0,0 +1,91 @@ +"""CPU tests of the capability evaluation suite (tools/eval): matrix, verdict +status, evidence loading, a judge on synthetic data, Miles argument building, +and the report.""" + +from __future__ import annotations + +import json +from pathlib import Path + +import pytest + +from tools.eval import __main__ as cli +from tools.eval import cells, miles, report +from tools.eval.verdict import Verdict + + +def test_matrix_covers_every_capability_and_backend(): + assert {(c.capability, c.backend) for c in cells.CELLS.values()} == { + (cap, b) for cap in cells.CAPABILITIES for b in cells.BACKENDS} + assert [c.key for c in cells.CELLS.values() if not c.supported] == ["E2-verl"] + with pytest.raises(SystemExit): + cells.get("E9", "verl") + + +@pytest.mark.parametrize("supported,passed,status", [ + (True, True, "pass"), (True, False, "fail"), (False, False, "expected_fail"), (False, True, "unexpected_pass")]) +def test_verdict_status(supported, passed, status): + v = Verdict("E1", "verl", "r", supported) + v.check("a", "a", True) + v.check("b", "b", passed) + assert v.status == status + assert Verdict("E1", "verl", "r", True).status == "error" # no criteria + + +def _island_run(tmp: Path, rounds: int, loss: float = 0.1) -> Path: + run = tmp / "run" + comp = run / "raw" / "run" / "island-c" + (comp / "yeto-output").mkdir(parents=True) + cell = cells.get("E1", "verl") + (run / "cell.json").write_text(json.dumps({"capability": "E1", "backend": "verl", "run_id": "run", + "config": cell.config, "code_sha": "x"})) + (comp / "meta.json").write_text(json.dumps({"rc": 0, "gpu": ["NVIDIA A10"], "seconds": 100.0})) + rows = [] + for r in range(rounds): + rows.append({"event": "rl_round_trained", "rollout_id": r, "time_unix": 10.0 + r}) + rows.append({"event": "rl_local_round", "loss": loss, "grad_norm": 0.2, "reward_mean": 0.5, + "base_policy_version": r, "time_unix": 10.5 + r}) + rows.append({"event": "rl_verl_island_done", "rounds_completed": rounds, "time_unix": 99.0}) + (comp / "yeto-output" / "rl-island-0.jsonl").write_text("\n".join(json.dumps(x) for x in rows) + "\n") + return run + + +def test_e1_judge_passes_and_reports(tmp_path): + run = _island_run(tmp_path, rounds=4) + verdict = cli.judge(run) + assert verdict.status == "pass", [c for c in verdict.criteria if not c.passed] + assert verdict.run["card_seconds"] == 100.0 + verdict.write(run) + md = report.markdown([run], title="t") + assert "| E1 island rollout + trainer | not run | PASS" in md + assert "" in report.html([run], title="t") + + +def test_e1_judge_fails_on_missing_round_and_nan(tmp_path): + run = _island_run(tmp_path, rounds=3, loss=float("nan")) + verdict = cli.judge(run) + failed = {c.id for c in verdict.criteria if not c.passed} + assert verdict.status == "fail" and {"rounds", "finite"} <= failed + + +def test_miles_launcher_args(): + run = "eval-x" + e2 = miles.launcher_args(cells.get("E2", "miles"), run) + assert "--rl-island-sample-role 0:sampler --rl-island-sample-role 1:sampler" in e2["ARGS"] + assert "--rl-island-sample-role 2:trainer" in e2["ARGS"] and "modal:1xa10g,modal:1xa10g,modal:1xa10g" in e2["ARGS"] + assert f"modal-volume://{miles.STORE_VOLUME}/{run}" in e2["ARGS"] + assert (e2["RECLAIM_LABEL"], e2["RECLAIM_ROUND"]) == ("l0", "1") + e4 = miles.launcher_args(cells.get("E4", "miles"), run) + assert "--rl-sample-role trainer" in e4["ARGS"] and "--rl-sample-role sampler" in e4["SARGS"] + assert f"--cluster-prefix {run}-s" in e4["SARGS"] and f"/{run} " in e4["SARGS"] # same store + assert "--rl-trainer-scheduler s1" in miles.launcher_args(cells.get("E5", "miles"), run)["ARGS"] + with pytest.raises(ValueError): + miles.launcher_args(cells.get("E3", "miles"), run) + + +def test_run_refuses_unregistered_run_id(tmp_path, monkeypatch): + ledger = tmp_path / "gpu-spend.md" + ledger.write_text("| other |\n") + monkeypatch.setattr(cli, "LEDGER", ledger) + with pytest.raises(SystemExit, match="not pre-registered"): + cli._preflight("eval-unregistered") diff --git a/tools/eval/__init__.py b/tools/eval/__init__.py new file mode 100644 index 000000000..abfbd6fea --- /dev/null +++ b/tools/eval/__init__.py @@ -0,0 +1,22 @@ +"""yeto RL capability evaluation suite (S25 T-EVAL). + +One repeatable GPU evaluation per capability and backend, one CLI, one verdict +format and one report: + + python -m tools.eval list + python -m tools.eval run --backend miles|verl --run-id [--dry-run] + python -m tools.eval judge + python -m tools.eval report ... --out + +* :mod:`.cells` -- the capability x backend matrix (what each cell runs, its + configuration, its judge, and whether the backend supports it); +* :mod:`.verdict` -- the verdict written by every judge (``verdict.json``); +* :mod:`.judges` -- the pre-registered criteria of each capability; +* :mod:`.modal_verl` -- verl cells: a direct Modal harness (syncer, engine sampler + island, verl islands), derived from the S25 C2/C3V/C4 harnesses; +* ``miles_launch.sh`` -- Miles cells: the yeto launcher on Modal, derived from + the S25 C1/C3 and S24 L5.4 run scripts; +* :mod:`.report` -- Markdown and HTML report over a set of verdicts. + +Every run writes its raw data under ``//raw/`` before judging. +""" diff --git a/tools/eval/__main__.py b/tools/eval/__main__.py new file mode 100644 index 000000000..eee7b4dcb --- /dev/null +++ b/tools/eval/__main__.py @@ -0,0 +1,152 @@ +"""CLI of the evaluation suite (see the package doc). + + python -m tools.eval list + python -m tools.eval run E3 --backend verl --run-id eval-e3-verl-a [--out DIR] [--dry-run] + python -m tools.eval judge DIR/eval-e3-verl-a + python -m tools.eval report DIR/* --out report.md [--html report.html] + +``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. +""" + +from __future__ import annotations + +import argparse +import json +import subprocess +import sys +import time +from pathlib import Path + +from . import cells, judges, report +from .verdict import Verdict + +DEFAULT_OUT = Path("/home/michael/work/infra-drafts/SESSION25/logs/T-EVAL") +LEDGER = Path("/home/michael/work/infra-drafts/gpu-spend.md") +THREAD_LIMIT = 3000 + + +def _preflight(run_id: str) -> None: + if not LEDGER.is_file() or run_id not in LEDGER.read_text(): + raise SystemExit(f"run id {run_id} is not pre-registered in {LEDGER}") + threads = subprocess.run("ps -u michael -L --no-headers | wc -l", shell=True, capture_output=True, + text=True).stdout.strip() + if int(threads or 0) >= THREAD_LIMIT: + raise SystemExit(f"{threads} local threads (>= {THREAD_LIMIT}): do not launch") + if subprocess.run(["pgrep", "-x", "raylet"], capture_output=True).returncode == 0: + raise SystemExit("a local Ray is running: do not launch") + + +def cmd_list(_args) -> int: + for key, cell in sorted(cells.CELLS.items()): + flag = "" if cell.supported else " [unsupported: expected to fail]" + print(f"{key:10s} {cell.harness:12s} {cell.cards} card(s) {cell.config.get('gpu')} " + f"{cells.CAPABILITIES[cell.capability]}{flag}") + return 0 + + +def cmd_run(args) -> int: + cell = cells.get(args.capability, args.backend) + run_dir = Path(args.out) / args.run_id + if (run_dir / "cell.json").exists() and not args.dry_run: + raise SystemExit(f"{run_dir} already used") + run_dir.mkdir(parents=True, exist_ok=True) + sha = subprocess.run(["git", "-C", str(Path(__file__).resolve().parents[2]), "rev-parse", "HEAD"], + capture_output=True, text=True).stdout.strip() + (run_dir / "cell.json").write_text(json.dumps({ + "capability": cell.capability, "backend": cell.backend, "run_id": args.run_id, "harness": cell.harness, + "cards": cell.cards, "supported": cell.supported, "config": cell.config, "code_sha": sha}, indent=1)) + began = time.time() + if cell.harness == "modal_verl": + if args.dry_run: + print(json.dumps(cell.config, indent=1)) + return 0 + _preflight(args.run_id) + from . import modal_verl + + summary = modal_verl.run(cell, args.run_id, run_dir) + else: + from . import miles + + if not args.dry_run: + _preflight(args.run_id) + summary = miles.run(cell, args.run_id, run_dir, dry_run=args.dry_run) + if args.dry_run: + print((run_dir / "dry-run.log").read_text()) + return summary["rc"] + summary["wall_s"] = round(time.time() - began, 1) + (run_dir / "run.json").write_text(json.dumps(summary, indent=1, default=str)) + return cmd_judge(argparse.Namespace(run_dir=str(run_dir))) + + +def judge(run_dir: Path) -> Verdict: + from . import evidence, miles, modal_verl + + cell_doc = json.loads((run_dir / "cell.json").read_text()) + cell = cells.get(cell_doc["capability"], cell_doc["backend"]) + cfg = cell_doc["config"] + verdict = Verdict(capability=cell.capability, backend=cell.backend, run_id=cell_doc["run_id"], + supported=cell.supported) + run_doc = json.loads((run_dir / "run.json").read_text()) if (run_dir / "run.json").is_file() else {} + raw = run_dir / "raw" + seconds = (modal_verl.card_seconds(raw) if cell.backend == "verl" else miles.card_seconds(raw)) + rate = cells.A10G_USD_PER_HOUR if cfg.get("gpu") == "A10G" else 5.3 + verdict.run.update(code_sha=cell_doc.get("code_sha"), wall_s=run_doc.get("wall_s"), + card_seconds=round(seconds, 1), usd_estimate=round(seconds / 3600 * rate, 3), + usd_rate=f"{rate}/card-hour (list price estimate, not a bill)") + try: + ev = evidence.load(run_dir) + judges.JUDGES[cell.capability](ev, verdict, cfg) + except Exception as exc: # noqa: BLE001 - a judge crash is an error verdict, with the reason + verdict.error = f"{type(exc).__name__}: {exc}" + return verdict + + +def cmd_judge(args) -> int: + run_dir = Path(args.run_dir) + verdict = judge(run_dir) + verdict.write(run_dir) + doc = verdict.to_dict() + print(json.dumps({"run": doc["run_id"], "cell": f"{doc['capability']}-{doc['backend']}", + "status": doc["status"], "failed": [c["id"] for c in doc["criteria"] if not c["passed"]], + "error": doc["error"]})) + return 0 if doc["status"] in ("pass", "expected_fail") else 1 + + +def cmd_report(args) -> int: + dirs = [Path(d) for d in args.run_dirs if (Path(d) / "verdict.json").is_file()] + md = report.markdown(dirs, title=args.title) + Path(args.out).write_text(md) + if args.html: + Path(args.html).write_text(report.html(dirs, title=args.title)) + print(f"{len(dirs)} verdict(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) + sub.add_parser("list").set_defaults(fn=cmd_list) + r = sub.add_parser("run") + r.add_argument("capability", choices=sorted(cells.CAPABILITIES)) + r.add_argument("--backend", choices=cells.BACKENDS, required=True) + r.add_argument("--run-id", required=True) + r.add_argument("--out", default=str(DEFAULT_OUT)) + r.add_argument("--dry-run", action="store_true") + r.set_defaults(fn=cmd_run) + j = sub.add_parser("judge") + j.add_argument("run_dir") + j.set_defaults(fn=cmd_judge) + rep = sub.add_parser("report") + rep.add_argument("run_dirs", nargs="+") + rep.add_argument("--out", required=True) + rep.add_argument("--html", default=None) + rep.add_argument("--title", default="yeto RL capability evaluation") + rep.set_defaults(fn=cmd_report) + args = p.parse_args(argv) + return args.fn(args) + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/tools/eval/cells.py b/tools/eval/cells.py new file mode 100644 index 000000000..56b73ceb9 --- /dev/null +++ b/tools/eval/cells.py @@ -0,0 +1,86 @@ +"""The capability x backend matrix of the evaluation suite. + +Model Qwen3-0.6B (pinned revision), GSM8K, LoRA, Modal A10G, the fewest cards +the capability needs. Each cell names its harness, its configuration (what the +judge reads too) and whether the backend supports the capability: an +unsupported cell still runs, and its failure is the expected baseline. +""" + +from __future__ import annotations + +from dataclasses import dataclass, field +from typing import Any + +CAPABILITIES = { + "E1": "island rollout + trainer", + "E2": "pure sampler island reclaimed and replaced", + "E3": "synced trainer islands -> sampler island", + "E4": "async rollout + off-policy correction + over-age drop", + "E5": "round cut: SCHEDULED stop + resume assertions", +} +BACKENDS = ("miles", "verl") +A10G_USD_PER_HOUR = 1.10 # Modal list price used by the S25 C2/C3V estimates (not a bill) + + +@dataclass(frozen=True) +class Cell: + capability: str + backend: str + harness: str # "modal_verl" | "miles_local" | "miles_head" + cards: int + config: dict[str, Any] = field(default_factory=dict) + supported: bool = True + note: str = "" + + @property + def key(self) -> str: + return f"{self.capability}-{self.backend}" + + +# ------------------------------------------------------------------ verl (direct Modal harness) +_VERL = {"gpu": "A10G", "rounds": 4, "groups_per_round": 4, "samples_per_group": 4, "max_tokens": 256, + "lr": 1e-4, "lora_r": 32, "seq_len": 768, "seed": 1, "sglang_mem": 0.6, "max_model_len": 1024, + "timeout_s": 2700, "learners": 1, "max_policy_age": 1, "sampler_rounds": 24, "max_backlog_groups": 8, + "reclaim_after_round": None} +# the plain island cells reuse the S25 C4 configuration (8 groups x 4, lr 1e-5) +_VERL_ISLAND = {**_VERL, "kind": "island", "groups_per_round": 8, "lr": 1e-5, "seq_len": 1024, "seed": 17, + "stop_after_rounds": 2} + +# ------------------------------------------------------------------ Miles (yeto launcher on Modal) +_MILES = {"gpu": "A10G", "hard_s": 3000, "rollout_batch_size": 16, "n_samples": 4, "max_response_len": 512, + "seq_len": 1024, "lora_r": 16, "lr": 1e-5, "seed": 17} + +CELLS: dict[str, Cell] = {c.key: c for c in ( + Cell("E1", "verl", "modal_verl", 1, {**_VERL_ISLAND, "phases": ["c"]}), + Cell("E1", "miles", "miles_local", 1, {**_MILES, "rounds": 4, "islands": 1}), + # E2: the S25 C1 shape (two samplers, one reclaimed after it wrote round 1) on Miles; the + # neutral engine sampler (verl's sampler island) has no rejoin path: expected to fail + Cell("E2", "verl", "modal_verl", 2, + {**_VERL, "kind": "buffer", "max_policy_age": 2, "max_backlog_groups": 12, "reclaim_after_round": 1}, + supported=False, note="engine sampler island: no LEAVE/JOIN, no resume after a reclaim (S25 C1 is Miles only)"), + Cell("E2", "miles", "miles_local", 3, + {**_MILES, "rounds": 8, "islands": 3, "samplers": 2, "max_policy_age": 4, "reclaim_label": "l0", + "reclaim_round": 1, "reclaim_component": "l0"}), + Cell("E3", "verl", "modal_verl", 3, + {**_VERL, "kind": "buffer", "learners": 2, "max_policy_age": 2, "max_backlog_groups": 16, + "sampler_rounds": 20}), + Cell("E3", "miles", "miles_head", 3, {**_MILES, "rounds": 4, "learners": 2, "hard_s": 2700}), + # E4: a sampler island far ahead of the trainer (backlog 3 rounds, short responses) and + # max_policy_age 1: round 2 meets v0 groups, which must be dropped + Cell("E4", "verl", "modal_verl", 2, + {**_VERL, "kind": "buffer", "max_policy_age": 1, "max_backlog_groups": 12, "max_tokens": 128}), + Cell("E4", "miles", "miles_local", 2, + {**_MILES, "rounds": 4, "islands": 2, "max_policy_age": 1, "sampler_rounds": 10, "backlog_rounds": 3, + "max_response_len": 256}), + Cell("E5", "verl", "modal_verl", 1, + {**_VERL_ISLAND, "phases": ["a", "b", "c"], "compare_prompt_sets": True}), + Cell("E5", "miles", "miles_local", 2, + {**_MILES, "rounds": 6, "islands": 2, "max_policy_age": 2, "min_round_s": 120}), +)} + + +def get(capability: str, backend: str) -> Cell: + key = f"{capability.upper()}-{backend}" + if key not in CELLS: + raise SystemExit(f"no cell {key}; known: {', '.join(sorted(CELLS))}") + return CELLS[key] diff --git a/tools/eval/evidence.py b/tools/eval/evidence.py new file mode 100644 index 000000000..be8fda9d0 --- /dev/null +++ b/tools/eval/evidence.py @@ -0,0 +1,197 @@ +"""Raw run data, read the same way for both backends. + +A run directory holds ``cell.json`` (the cell and run id), ``driver.json`` and +``raw/`` (everything pulled). :class:`Evidence` turns the raw data into: + +* ``islands``: event lists keyed by component (verl: ``trainer0``, ``sampler-0``, + ``island-a``...; Miles: the launcher's island labels ``l0``, ``l1``, and + ``s:l0`` for a separate sampler launch), each sorted by time and de-duplicated + across sources (launch log echo, per-container tapes, the tape Volume); +* ``index``: the sample buffer index (shared-directory JSONL or object files); +* ``lanes``: per-lane read-position logs of synced trainers; +* return codes, GPU names and card seconds. +""" + +from __future__ import annotations + +import csv +import json +import re +from collections import defaultdict +from collections.abc import Iterable, Mapping +from dataclasses import dataclass, field +from pathlib import Path +from typing import Any + +from .verdict import events, find, find_all, jsonl, objects + +TRAINER_EVENTS = ("rl_round_trained", "rl_verl_island_done") +SAMPLER_EVENTS = ("rl_sampler_round", "rl_sampler_start") + + +def _merge(target: list[dict], rows: Iterable[dict], seen: set) -> None: + for row in rows: + key = (row.get("event"), row.get("time_unix"), row.get("seq"), row.get("rollout_id")) + if row.get("event") and key not in seen: + seen.add(key) + target.append(row) + + +@dataclass +class Evidence: + run_dir: Path + backend: str + islands: dict[str, list[dict[str, Any]]] = field(default_factory=dict) + index: list[dict[str, Any]] = field(default_factory=list) + lanes: dict[int, list[dict[str, Any]]] = field(default_factory=dict) + shards: dict[str, list[dict[str, Any]]] = field(default_factory=dict) # sampler shards of the index + rc: dict[str, Any] = field(default_factory=dict) + gpu_names: dict[str, int] = field(default_factory=dict) + store: Path | None = None + files: dict[str, Path] = field(default_factory=dict) + + # ------------------------------------------------------------- queries + def of(self, component: str, name: str) -> list[dict[str, Any]]: + return list(events(self.islands.get(component, []), name)) + + def all(self, name: str) -> list[dict[str, Any]]: + return [e for tape in self.islands.values() for e in events(tape, name)] + + def trainers(self) -> list[str]: + return sorted(k for k, tape in self.islands.items() if any(e.get("event") in TRAINER_EVENTS for e in tape)) + + def samplers(self) -> list[str]: + return sorted(k for k, tape in self.islands.items() if any(e.get("event") in SAMPLER_EVENTS for e in tape)) + + def entries(self) -> list[dict[str, Any]]: + """Index records that are groups (not end markers).""" + return [e for e in self.index if not e.get("end") and ("behavior_version" in e or "sample_id" in e)] + + +# ------------------------------------------------------------------ loaders +def load(run_dir: Path) -> Evidence: + run_dir = Path(run_dir) + cell = json.loads((run_dir / "cell.json").read_text()) + ev = Evidence(run_dir=run_dir, backend=cell["backend"]) + raw = run_dir / "raw" + (_load_verl if ev.backend == "verl" else _load_miles)(ev, raw, cell["run_id"]) + for tape in ev.islands.values(): + tape.sort(key=lambda e: float(e.get("time_unix") or e.get("t") or 0.0)) + if ev.store is not None: + buffer = ev.store / "sample-buffer" + if not buffer.is_dir(): # miles_head.sh pulls index.jsonl and lanes/ into the store dir itself + buffer = ev.store + ev.index = jsonl(buffer / "index.jsonl") or objects(buffer / "index") + for path in sorted((buffer / "shards").glob("*.jsonl")) if (buffer / "shards").is_dir() else []: + ev.shards[path.stem] = jsonl(path) + lanes = buffer / "lanes" + for path in sorted(lanes.glob("*-of-*")) if lanes.is_dir() else []: + m = re.match(r"(\d+)-of-\d+", path.name) + if m: + ev.lanes[int(m.group(1))] = jsonl(path) if path.is_file() else objects(path) + return ev + + +def _load_verl(ev: Evidence, raw: Path, run_id: str) -> None: + root = raw / run_id if (raw / run_id).is_dir() else raw + for comp in sorted(p for p in root.iterdir() if p.is_dir() and p.name != "store") if root.is_dir() else []: + rows: list[dict] = [] + seen: set = set() + for name in ("trainer-tape.jsonl", "sampler-tape.jsonl", "syncer-tape.jsonl"): + _merge(rows, jsonl(comp / name), seen) + for tape in find_all(comp / "yeto-output", "rl-island-*.jsonl"): + _merge(rows, jsonl(tape), seen) + if comp.name == "syncer": + rows = [dict(r, event=r.get("event") or r.get("kind")) for r in jsonl(comp / "syncer-tape.jsonl")] + ev.islands[comp.name] = rows + meta = comp / "meta.json" + if meta.is_file(): + data = json.loads(meta.read_text()) + # a sampler or a phase returns without "rc": its error (if any) is the outcome + ev.rc[comp.name] = data["rc"] if "rc" in data else (data["error"] if data.get("error") else 0) + for name in data.get("gpu") or []: + ev.gpu_names[name] = ev.gpu_names.get(name, 0) + 1 + ev.files[comp.name] = comp + ev.store = root / "store" + + +_EVENT_LINE = r"^(?:\S+ )?\[{prefix}-l(\d+)-modal\] YETO_RL_EVENT (\{{.*\}})\s*$" + + +def _launch_events(path: Path, prefix: str) -> dict[int, list[dict]]: + out: dict[int, list[dict]] = defaultdict(list) + if not path.is_file(): + return out + pattern = re.compile(_EVENT_LINE.format(prefix=re.escape(prefix))) + for line in path.read_text(errors="replace").splitlines(): + m = pattern.match(line) + if m: + try: + out[int(m.group(1))].append(json.loads(m.group(2))) + except ValueError: + pass + return out + + +def _load_miles(ev: Evidence, raw: Path, run_id: str) -> None: + seen: dict[str, set] = defaultdict(set) + for app, tag, log in ((run_id, "", "launch.ts.log"), (f"{run_id}-s", "s:", "slaunch.ts.log")): + for island, rows in _launch_events(raw / log, app).items(): + _merge(ev.islands.setdefault(f"{tag}l{island}", []), rows, seen[f"{tag}l{island}"]) + # the tape Volume holds every incarnation's tape (rl-island-0.inc.jsonl after a + # replacement container); the launch log only echoes the first container + for tape in find_all(raw / "tape-direct" / f"yeto-{app}", "rl-island-*.jsonl"): + label = next((part for part in tape.parts if re.fullmatch(r"l\d+", part)), None) + if label is None: + continue + key = f"{tag}{label}" + _merge(ev.islands.setdefault(key, []), jsonl(tape), seen[key]) + for name, key in (("rc", "launch"), ("src", "sampler_launch")): + path = raw / name + if path.is_file(): + ev.rc[key] = int(path.read_text().strip() or -1) + head_rc = raw / "rc.txt" + if head_rc.is_file(): + m = re.search(r"rc=(\d+)", head_rc.read_text()) + ev.rc["launch"] = int(m.group(1)) if m else -1 + nvml = raw / "nvml.csv" + if nvml.is_file(): + for row in csv.reader(nvml.read_text(errors="replace").splitlines()): + if len(row) > 2: + name = row[2].strip() + ev.gpu_names[name] = ev.gpu_names.get(name, 0) + 1 + store = raw / "store" + # `modal volume get

/sample-buffer ` lands as /sample-buffer + ev.store = store + ev.files["raw"] = raw + + +def gpu_ok(ev: Evidence, want: str) -> bool: + return bool(ev.gpu_names) and all(want in name for name in ev.gpu_names) + + +def finite(values: Iterable[Any]) -> bool: + import math + + vals = list(values) + return bool(vals) and all(isinstance(v, (int, float)) and math.isfinite(v) for v in vals) + + +def numeric_sum(rows: Iterable[Mapping[str, Any]], needle: str) -> float: + """Sum of every numeric field whose (flattened) name contains ``needle``.""" + total = 0.0 + + def walk(value: Any, name: str) -> None: + nonlocal total + if isinstance(value, Mapping): + for k, v in value.items(): + walk(v, f"{name}/{k}") + elif isinstance(value, (int, float)) and not isinstance(value, bool) and needle in name: + total += float(value) + + for row in rows: + walk(row, "") + return total + + +__all__ = ["Evidence", "load", "gpu_ok", "finite", "numeric_sum", "find"] diff --git a/tools/eval/judges.py b/tools/eval/judges.py new file mode 100644 index 000000000..f544cac03 --- /dev/null +++ b/tools/eval/judges.py @@ -0,0 +1,304 @@ +"""Pre-registered criteria of each capability (S25 T-EVAL; review SESSION25/gpu-review/T-EVAL.md). + +One judge per capability, the same criteria for both backends; only where the +evidence lives differs (:mod:`.evidence`). Each judge fills a +:class:`~.verdict.Verdict` with criteria (pass/fail + evidence) and metrics. +The criteria come from the GPU acceptance of the capability's S25 task: E2 from +T-C1 (C1-C6), E3 from T-C3/T-C3V (1-5), E4 from T-C2 (off-policy) plus a forced +over-age drop, E5 from T-C4 / S24 L5.4 (SCHEDULED stop, resume checks). +""" + +from __future__ import annotations + +import re +import statistics +from collections import defaultdict +from collections.abc import Callable +from typing import Any + +from .evidence import Evidence, finite, gpu_ok, numeric_sum +from .verdict import Verdict, find_all, jsonl + +GPU_WANT = {"A10G": "A10", "H100!": "H100", "L4": "L4"} +CONTROL = "island-c" # verl E5: the continuous control island (not part of the resumed run) + + +# ------------------------------------------------------------------ shared criteria +def _rc_ok(ev: Evidence, v: Verdict, ignore: tuple[str, ...] = ()) -> None: + rcs = {k: r for k, r in ev.rc.items() if k not in ignore} + v.check("rc", "every launch / container exited 0", bool(rcs) and all(r == 0 for r in rcs.values()), rcs) + + +def _gpu(ev: Evidence, v: Verdict, cfg: dict) -> None: + want = GPU_WANT[cfg["gpu"]] + v.check("gpu", f"every container ran on the requested card ({want})", gpu_ok(ev, want), ev.gpu_names) + + +def _round_metrics(ev: Evidence, trainer: str | None = None) -> dict[str, Any]: + rounds = ev.of(trainer, "rl_local_round") if trainer else ev.all("rl_local_round") + walls = ev.of(trainer, "rl_round_wall") if trainer else ev.all("rl_round_wall") + return { + "rounds": len(rounds), + "reward_mean": [round(r.get("reward_mean", float("nan")), 4) for r in rounds], + "loss": [r.get("loss") for r in rounds], + "grad_norm": [r.get("grad_norm") for r in rounds], + "rollout_s": [r.get("rollout_seconds") for r in rounds], + "train_s": [r.get("train_seconds") for r in rounds], + "round_wall_s_median": (statistics.median(w["seconds"] for w in walls if "seconds" in w) + if any("seconds" in w for w in walls) else None), + } + + +def _trained_rounds(ev: Evidence, trainer: str) -> list[int]: + return sorted({int(e["rollout_id"]) for e in ev.of(trainer, "rl_round_trained") if "rollout_id" in e}) + + +def _released(ev: Evidence, trainer: str) -> list[int]: + return [int(e["policy_version"]) for e in ev.of(trainer, "rl_policy_released")] + + +# ------------------------------------------------------------------ E1 +def judge_e1(ev: Evidence, v: Verdict, cfg: dict) -> None: + """Island rollout + trainer: R rounds trained in one island, finite numbers.""" + trainers = ev.trainers() + v.check("island", "one island trained (rollout and trainer in the same island)", len(trainers) == 1, trainers) + t = trainers[0] if trainers else "" + _rc_ok(ev, v) + rounds = _trained_rounds(ev, t) + v.check("rounds", f"rounds 0..{cfg['rounds'] - 1} trained once each", rounds == list(range(cfg["rounds"])), rounds) + local = ev.of(t, "rl_local_round") + v.check("finite", "loss, grad norm and reward finite every round", + len(local) == cfg["rounds"] and finite(r.get("loss") for r in local) + and finite(r.get("grad_norm") for r in local) and finite(r.get("reward_mean") for r in local), + {"loss": [r.get("loss") for r in local], "reward_mean": [r.get("reward_mean") for r in local]}) + bases = [r.get("base_policy_version") for r in local] + v.check("versions", "the policy version advances by one per round", bases == list(range(len(bases))) and bases, + bases) + _gpu(ev, v, cfg) + v.metrics.update(_round_metrics(ev, t)) + + +# ------------------------------------------------------------------ E2 +def _reclaim(ev: Evidence) -> tuple[int, float | None, str]: + log = ev.run_dir / ("raw/reclaim.log" if ev.backend == "miles" else "reclaim.log") + text = log.read_text() if log.is_file() else "" + m = re.search(r"stop rc=(\d+) \S+ (\d+\.\d+)", text) # Miles script + if m: + return int(m.group(1)), float(m.group(2)), text + m = re.search(r"rc=(\d+)", text) # verl driver: the stop time is the line's ISO time + if m: + import datetime + + stamp = datetime.datetime.strptime(text.split()[0], "%Y-%m-%dT%H:%M:%SZ") + return int(m.group(1)), stamp.replace(tzinfo=datetime.timezone.utc).timestamp(), text + return -1, None, text + + +def judge_e2(ev: Evidence, v: Verdict, cfg: dict) -> None: + """A pure sampler island is reclaimed once and replaced; the trainer keeps training.""" + rc, stop_t, text = _reclaim(ev) + v.check("reclaim", "one reclaim of the sampler island's container, stop rc 0", + rc == 0 and text.count("stop") >= 1, text.strip().splitlines()[-3:]) + samplers = ev.samplers() + target = cfg.get("reclaim_component") or (samplers[0] if samplers else "") + if ev.backend == "verl": # each sampler container writes its own component dir + tapes = [ev.islands.get(s, []) for s in samplers] + merged = [e for t in tapes for e in t] + else: + merged = ev.islands.get(target, []) + leaves = [e for e in merged if e.get("event") == "rl_sampler_leave"] + joins = [e for e in merged if e.get("event") == "rl_sampler_join"] + pre = [e for e in leaves if e.get("reason") == "PREEMPTED"] + v.check("leave", "the reclaimed incarnation wrote rl_sampler_leave reason=PREEMPTED within 30 s of the stop", + bool(pre) and stop_t is not None and float(pre[0]["time_unix"]) - stop_t <= 30.0, + {"leaves": [(e.get("incarnation"), e.get("reason")) for e in leaves], "stop_t": stop_t}) + rounds = [e for e in merged if e.get("event") == "rl_sampler_round"] + ok = False + detail: dict[str, Any] = {"joins": len(joins)} + if len(joins) >= 2: + old, new = joins[0], joins[-1] + written = [int(e["rollout_id"]) for e in rounds if e["time_unix"] < new["time_unix"] and e.get("groups", 0)] + ok = (new.get("incarnations") == [old.get("incarnation")] + and written and new.get("next_rollout_id") == max(written) + 1) + detail.update(old_written=written, next_rollout_id=new.get("next_rollout_id"), + incarnations=new.get("incarnations")) + if stop_t: + v.metrics["replacement_join_s"] = round(float(new["time_unix"]) - stop_t, 1) + v.check("join", "a new incarnation joined, named the reclaimed one and resumed at the next round", ok, detail) + groups = ev.entries() + ids = [e.get("sample_id") or e.get("group_id") or e.get("payload") for e in groups] + per_round: dict[Any, set] = defaultdict(set) + for e in groups + [r for rows in ev.shards.values() for r in rows if not r.get("end")]: + sr = e.get("sampler_round") or {} + if sr: + per_round[(sr.get("sampler"), sr.get("rollout_id"))].add(sr.get("incarnation")) + v.check("buffer", "buffer groups unique; each sampler round written by one incarnation only", + bool(ids) and len(ids) == len(set(ids)) and bool(per_round) and all(len(s) == 1 for s in per_round.values()), + {"groups": len(ids), "unique": len(set(ids)), "rounds_with_incarnation": len(per_round)}) + trainers = ev.trainers() + t = trainers[0] if trainers else "" + trained = _trained_rounds(ev, t) + during = [e for e in ev.of(t, "rl_round_trained") + if stop_t and joins and len(joins) >= 2 and stop_t < float(e["time_unix"]) < float(joins[-1]["time_unix"])] + v.check("trainer", f"the trainer trained all {cfg['rounds']} rounds", len(trained) == cfg["rounds"], trained) + v.metrics.update(trainer_rounds_during_replacement=len(during), sampler_rounds=len(rounds), + buffer_groups=len(ids)) + _rc_ok(ev, v) + _gpu(ev, v, cfg) + + +# ------------------------------------------------------------------ E3 +def judge_e3(ev: Evidence, v: Verdict, cfg: dict) -> None: + """N synced trainer islands; learner 0 releases the merged version to a sampler island.""" + trainers = ev.trainers() + n = cfg["learners"] + v.check("trainers", f"{n} synced trainer islands trained", len(trainers) == n, trainers) + applied = [{e["sync/outer_version"]: e["sync/global_policy_hash"] for e in ev.of(t, "rl_policy_apply") + if "sync/outer_version" in e} for t in trainers] + v.check("merged_hash", "every outer version has the same merged hash on every trainer (>= rounds versions)", + len(applied) == n and applied[0] and all(a == applied[0] for a in applied) + and len(applied[0]) >= cfg["rounds"], {"outer_versions": sorted(applied[0]) if applied else []}) + publisher = next((t for t in trainers if any(e.get("publisher") for e in ev.of(t, "rl_policy_release_role"))), None) + released = ev.of(publisher, "rl_policy_released") if publisher else [] + others = [t for t in trainers if t != publisher] + sampler_tapes = [e for s in ev.samplers() for e in ev.islands[s]] + loaded = [e for e in sampler_tapes if e.get("event") == "rl_sampler_policy_loaded"] + merged = applied[0] if applied else {} + rel_versions = [int(e["policy_version"]) for e in released] + load_versions = [int(e["policy_version"]) for e in loaded] + v.check("sampler_hash", "the sampler loads released versions in order; loaded hash = merged hash; " + "releases come from learner 0", + bool(loaded) and load_versions == sorted(set(load_versions)) and all(x in rel_versions for x in load_versions) + and all(e.get("policy_hash") == merged.get(int(e["policy_version"])) for e in loaded + released) + and all((e.get("origin") or {"learner_id": 0}).get("learner_id") == 0 for e in released), + {"loaded": load_versions, "released": rel_versions}) + gaps = [e for t in trainers for e in ev.of(t, "rl_policy_release_gap")] + contiguous = rel_versions == list(range(rel_versions[0], rel_versions[0] + len(rel_versions))) if rel_versions else False + v.check("one_publisher", "one publisher, contiguous release versions (>= 3), no gap; others release nothing", + publisher is not None and contiguous and len(rel_versions) >= 3 and not gaps + and not any(ev.of(t, "rl_policy_released") for t in others), + {"publisher": publisher, "gaps": len(gaps)}) + positions = {m: [p for row in rows for p in row.get("positions", [])] for m, rows in ev.lanes.items()} + flat = [p for ps in positions.values() for p in ps] + v.check("lanes", "lane read positions disjoint, no duplicate, p % N == m", + len(positions) == n and bool(flat) and len(flat) == len(set(flat)) + and all(p % n == m for m, ps in positions.items() for p in ps), + {m: len(ps) for m, ps in positions.items()}) + steps = [e for e in ev.all("outer_step")] or _head_outer_steps(ev) + v.check("outer_steps", "the syncer ran >= 3 OuterSteps", len(steps) >= 3, len(steps)) + _rc_ok(ev, v, ignore=("sampler_launch", "sampler")) + _gpu(ev, v, cfg) + v.metrics.update(release_seconds=[e.get("seconds") for e in released], buffer_groups=len(ev.entries()), + outer_steps=len(steps)) + + +def _head_outer_steps(ev: Evidence) -> list[dict]: + rows = [r for p in find_all(ev.run_dir / "raw" / "head", "*tape*.jsonl") for r in jsonl(p)] + return [r for r in rows if r.get("kind") == "outer_step"] + + +# ------------------------------------------------------------------ E4 +def judge_e4(ev: Evidence, v: Verdict, cfg: dict) -> None: + """Asynchronous rollout (a sampler island ahead of the trainer) with off-policy + correction; the trainer drops groups older than max_policy_age at least once.""" + trainers = ev.trainers() + t = trainers[0] if trainers else "" + age = int(cfg["max_policy_age"]) + trained = _trained_rounds(ev, t) + v.check("rounds", f"the trainer trained all {cfg['rounds']} rounds", len(trained) == cfg["rounds"], trained) + rt = ev.of(t, "rl_round_trained") + tis = [((e.get("mismatch") or {}).get("tis")) for e in rt] + # verl: per-source importance ratio (offpolicy_source); Miles: the round's TIS mean + src = ev.of(t, "offpolicy_source") + ratio = [s.get("ratio_p50") for e in src for s in e.get("sources") or []] or [x for x in tis if x is not None] + v.check("correction", "off-policy correction metrics present and finite every round", + len(rt) == cfg["rounds"] and finite(ratio) and finite( + (e.get("mismatch") or {}).get("abs_diff", (e.get("mismatch") or {}).get("tis_abs")) for e in rt), + {"ratio_or_tis": ratio}) + dropped = (numeric_sum(ev.of(t, "rl_rollout_carry_over"), "over_age_discarded_groups") + + numeric_sum(ev.of(t, "rl_round_labels") + rt + ev.of(t, "rl_local_round"), "stale_groups") + + numeric_sum(ev.of(t, "offpolicy_gate"), "group_dropped")) + if ev.backend == "verl": + files = find_all(ev.files.get(t, ev.run_dir), "verl-fa-rollout-*.jsonl") + dropped += numeric_sum([r.get("carry_over") or {} for p in files for r in jsonl(p)], "over_age_discarded_groups") + v.check("over_age_drop", "at least one group older than max_policy_age was dropped, not trained", dropped >= 1, + dropped) + over = _over_age_trained(ev, t, age) + v.check("no_over_age_trained", "no trained group is older than max_policy_age", over == 0, over) + loads = [e for s in ev.samplers() for e in ev.of(s, "rl_sampler_policy_loaded")] + v.check("sampler_follows", "the sampler island loaded >= 2 trainer releases while it ran ahead", + len({e.get("policy_version") for e in loads}) >= 2, sorted({e.get("policy_version") for e in loads})) + _rc_ok(ev, v, ignore=("sampler_launch", "sampler-0", "sampler-1")) + _gpu(ev, v, cfg) + v.metrics.update(over_age_dropped_groups=dropped, ratio_p50=ratio, **_round_metrics(ev, t)) + + +def _over_age_trained(ev: Evidence, trainer: str, age: int) -> int: + over = 0 + if ev.backend == "verl": + for path in find_all(ev.files.get(trainer, ev.run_dir), "verl-fa-rollout-*.jsonl"): + for r in jsonl(path): + for versions in (r.get("versions") or {}).values(): + over += min(versions) < r["rollout_id"] - age + return over + for e in ev.of(trainer, "rl_group_policies"): + current = e.get("policy_version") + for g in e.get("groups") or []: + oldest = g.get("oldest_version", g.get("min_version")) if isinstance(g, dict) else None + if current is not None and oldest is not None: + over += oldest < current - age + return over + + +# ------------------------------------------------------------------ E5 +def judge_e5(ev: Evidence, v: Verdict, cfg: dict) -> None: + """A SCHEDULED stop cuts the trainer; a new incarnation resumes from the cut + and the resume assertions pass.""" + life = ev.all("rl_trainer_lifecycle") + stopped = [e for e in life if e.get("phase") == "stopped" and (e.get("cut") or {}).get("trigger") == "scheduled"] + if not stopped: # verl: the cut record carries the trigger + stopped = [e for e in ev.all("rl_round_cut") + ev.all("rl_cut_saved") + if (e.get("consistent_cut") or e).get("trigger") == "scheduled"] + v.check("scheduled_stop", "the trainer stopped on a scheduler command with a SCHEDULED cut", bool(stopped), + [(e.get("incarnation"), (e.get("cut") or {}).get("cut_id", e.get("cut_id"))) for e in stopped]) + resumes = ev.all("rl_resume") + checks_ok = bool(resumes) and all( + all(str(val) in ("same", "ok", "outer_recorded") for key, val in (r.get("checks") or {}).items() + if key in ("run_fingerprint", "consumed_prompt_ids_digest", "cut_shards_sha256", "policy_hash", + "ledger_last_round")) + for r in resumes) + v.check("resume", "a new incarnation resumed from the cut; fingerprint, ledger, prompts digest, shards and " + "policy hash all match", checks_ok, + [{"rollout_id": r.get("rollout_id"), "checks": r.get("checks")} for r in resumes]) + first = [e for e in ev.all("rl_resume_check")] + v.check("first_step", "the resume checks after the first resumed round pass", bool(first) + and all(e.get("ok") for e in first), [(e.get("check"), e.get("ok")) for e in first]) + records = [e for t in ev.trainers() if t != CONTROL for e in ev.of(t, "rl_round_trained")] + trained = sorted({int(e["rollout_id"]) for e in records}) + v.check("rounds", f"rounds 0..{cfg['rounds'] - 1} trained exactly once across incarnations", + trained == list(range(cfg["rounds"])) and len(records) == cfg["rounds"], + {"rounds": trained, "records": len(records)}) + if cfg.get("compare_prompt_sets"): # verl: the resumed rounds read the same prompts as the control + v.check("data_cursor", "stopped + resumed rounds read the same prompt sets as the continuous control", + _prompt_sets_match(ev), None) + _rc_ok(ev, v) + _gpu(ev, v, cfg) + restore = [r.get("restore_s") for r in resumes] + v.metrics.update(restore_s=restore, scheduled_stops=len(stopped), resumes=len(resumes)) + + +def _prompt_sets_match(ev: Evidence) -> bool: + """verl: each round's prompt set (``verl-rollout-0.jsonl``, sorted-row hash; S25 C4).""" + def sets(component: str) -> dict[int, str]: + rows = [r for p in find_all(ev.files.get(component, ev.run_dir / "none"), "verl-rollout-*.jsonl") + for r in jsonl(p)] + return {i: r["prompt_set_sha256"] for i, r in enumerate(rows) if r.get("prompt_set_sha256")} + + control = sets(CONTROL) + a, b = sets("island-a"), sets("island-b") + resumed = {**a, **{len(a) + i: h for i, h in b.items()}} # b continues where a stopped + return bool(control) and len(resumed) == len(control) and resumed == control + + +JUDGES: dict[str, Callable[[Evidence, Verdict, dict], None]] = { + "E1": judge_e1, "E2": judge_e2, "E3": judge_e3, "E4": judge_e4, "E5": judge_e5} diff --git a/tools/eval/miles.py b/tools/eval/miles.py new file mode 100644 index 000000000..b15d29f98 --- /dev/null +++ b/tools/eval/miles.py @@ -0,0 +1,119 @@ +"""Miles cells: the yeto launcher on Modal through ``miles_local.sh`` / ``miles_head.sh``. + +:func:`launcher_args` builds the launcher argument strings of a cell (the store +is ``modal-volume:///``); :func:`run` writes them into the +script environment and runs the script, which pulls everything into ``raw/``. +""" + +from __future__ import annotations + +import os +import subprocess +from pathlib import Path +from typing import Any + +HERE = Path(__file__).resolve().parent +REPO = HERE.parents[1] +PY = "/home/michael/work/gpu-head/venv/bin/python" # modal client + launcher deps (S25 C1/C3) +SKY = "/home/michael/work/gpu-head/venv/bin/sky" +STORE_VOLUME = "yeto-eval-store" +TAPE_VOLUME = "yeto-event-tapes" +MODEL = "--model Qwen/Qwen3-0.6B --model-revision c1899de289a04d12100db370d81485cdf75e47ca" +DATA = "--data zhuzilin/gsm8k --data-revision 0cbd9f31d91ac21a7613dcbc7fef992adac459ae" +SPEC = "--rl-algorithm-spec examples/rl_algorithms/decoupled-ppo-c2.json" # max_policy_staleness 2 +SPEC_AGE4 = "--rl-algorithm-spec examples/rl_algorithms/decoupled-ppo-c2-agentic.json" # staleness 4 (S25 C1) + + +def _gpu(cfg: dict, n: int) -> str: + card = {"A10G": "a10g", "H100!": "h100"}[cfg["gpu"]] + return ",".join([f"modal:1x{card}"] * n) + + +def _common(cfg: dict, run_id: str, prefix: str) -> str: + return (f"launch --training-mode rl --rl-engine ports --modal-gpu-exact --modal-retries 0 --modal-cpu 8 " + f"--modal-memory 64 --modal-tape-volume {TAPE_VOLUME} --controller local --cluster-prefix {prefix} " + f"--tuning lora --lora-r {cfg['lora_r']} --rl-observe-timeline {MODEL} {DATA} " + f"--rl-checkpoint-store modal-volume://{STORE_VOLUME}/{run_id} " + "--reward-function gsm8k_reward:score --fragments 1 --pipeline 1 " + f"--rollout-batch-size {cfg['rollout_batch_size']} --n-samples-per-prompt {cfg['n_samples']} " + f"--rollout-max-response-len {cfg['max_response_len']} --seq-len {cfg['seq_len']} " + f"--inner-lr {cfg['lr']} --seed {cfg['seed']} " + "--apply-chat-template-kwargs '{\"enable_thinking\": false}' --trust-remote-code") + + +def launcher_args(cell, run_id: str) -> dict[str, str]: + """The script environment of a local-controller Miles cell (ARGS, SARGS, reclaim).""" + cfg, cap = cell.config, cell.capability + if cell.harness != "miles_local": + raise ValueError(f"{cell.key} is not a local-controller Miles cell") + common = _common(cfg, run_id, run_id) + env: dict[str, str] = {"ISLANDS": str(cfg["islands"])} + if cap == "E1": + env["ARGS"] = f"{common} --gpu {_gpu(cfg, 1)} --rl-single-island-no-sync --total-steps {cfg['rounds']}" + elif cap == "E2": # S25 C1: islands 0..k-1 sample, island k trains + k = cfg["samplers"] + roles = " ".join(f"--rl-island-sample-role {i}:sampler" for i in range(k)) + env["ARGS"] = (f"{common} --gpu {_gpu(cfg, k + 1)} --rl-single-island-no-sync {roles} " + f"--rl-island-sample-role {k}:trainer {SPEC_AGE4} --rl-max-policy-age {cfg['max_policy_age']} " + "--rl-sampler-max-backlog-rounds 2 --rl-sampler-min-round-seconds 60 --rl-sample-wait-s 900 " + f"--total-steps {cfg['rounds']}") + env.update(RECLAIM_LABEL=cfg["reclaim_label"], RECLAIM_ROUND=str(cfg["reclaim_round"])) + elif cap == "E4": # two launches on one store: the sampler runs ahead of the trainer + age = cfg["max_policy_age"] + env["ARGS"] = (f"{common} --gpu {_gpu(cfg, 1)} --rl-single-island-no-sync --rl-sample-role trainer " + f"{SPEC} --rl-max-policy-age {age} --rl-sample-wait-s 900 --total-steps {cfg['rounds']}") + env["SARGS"] = (f"{_common(cfg, run_id, run_id + '-s')} --gpu {_gpu(cfg, 1)} --rl-single-island-no-sync " + f"--rl-sample-role sampler {SPEC} --rl-max-policy-age {age} " + f"--rl-sampler-max-backlog-rounds {cfg['backlog_rounds']} " + f"--total-steps {cfg['sampler_rounds']}") + elif cap == "E5": # S24 L5.4: a pair whose trainer the head's S1 scheduler stops and restarts + env["ARGS"] = (f"{common} --gpu {_gpu(cfg, 2)} --rl-single-island-no-sync " + "--rl-island-sample-role 0:sampler --rl-island-sample-role 1:trainer " + f"{SPEC} --rl-max-policy-age {cfg['max_policy_age']} --rl-cut-every 1 " + f"--rl-sampler-max-backlog-rounds 3 --rl-sampler-min-round-seconds {cfg['min_round_s']} " + "--rl-sample-wait-s 900 --rl-trainer-scheduler s1 --rl-trainer-island 1 " + "--rl-trainer-scheduler-interval 15 --rl-trainer-start-seconds 30 " + "--rl-trainer-start-seconds-fixed --rl-trainer-ready-seconds 2700 " + f"--rl-trainer-drain-seconds 2400 --total-steps {cfg['rounds']}") + return env + + +def run(cell, run_id: str, run_dir: Path, *, dry_run: bool = False) -> dict[str, Any]: + raw = run_dir / "raw" + cfg = cell.config + env = dict(os.environ, P=run_id, R=str(raw), REPO=str(REPO), PY=PY, SKY=SKY, STOREVOL=STORE_VOLUME, + TAPEVOL=TAPE_VOLUME, HARD=str(cfg["hard_s"]), HOME="/home/michael") + if cell.harness == "miles_head": + script = HERE / "miles_head.sh" + env.update(GPU={"A10G": "a10g", "H100!": "h100"}[cfg["gpu"]], STEPS=str(cfg["rounds"])) + if dry_run: + env["PLAN_ONLY"] = "1" + else: + script = HERE / "miles_local.sh" + env.update(launcher_args(cell, run_id)) + 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(script)], env=env, stdout=fh, stderr=subprocess.STDOUT, cwd="/tmp").returncode + return {"script": script.name, "rc": rc, "log": str(log)} + + +def card_seconds(raw: Path) -> float: + """Card seconds from the container list the script polled every 20 s: per + container, last seen - first seen + 40 s (the S25 C1 estimate).""" + import datetime + + seen: dict[str, list[float]] = {} + path = Path(raw) / "containers.log" + if not path.is_file(): + return 0.0 + for line in path.read_text().splitlines(): + parts = line.split() + if not parts: + continue + stamp = datetime.datetime.strptime(parts[0], "%Y-%m-%dT%H:%M:%SZ").replace( + tzinfo=datetime.timezone.utc).timestamp() + for container in parts[1:]: + seen.setdefault(container, []).append(stamp) + return sum(max(t) - min(t) + 40.0 for t in seen.values()) diff --git a/tools/eval/miles_head.sh b/tools/eval/miles_head.sh new file mode 100755 index 000000000..2c27480d5 --- /dev/null +++ b/tools/eval/miles_head.sh @@ -0,0 +1,116 @@ +#!/bin/bash +# Miles eval cell E3: two synced trainer islands (elastic OuterStep; head + syncer on a +# Nebius CPU VM) feed one sampler island (a separate no-sync launch) through merged +# releases in one Modal Volume store. Copied from s1-runs/s25-c3/run.sh (S25 C3, passed in +# run s25-c3-20261020b); parameterised by the eval driver (tools/eval/miles.py). +# env: P (run id), R (raw dir), REPO, GPU (modal gpu, e.g. a10g), HARD, STEPS, STOREVOL, +# PY, SKY, PLAN_ONLY=1 (dry-run only) +set -u +: "${P:?}" "${R:?}" "${REPO:?}" "${PY:?}" "${SKY:?}" +HARD=${HARD:-2700}; WD=$(( HARD + 300 )); STEPS=${STEPS:-4}; SROUNDS=$(( 2 * STEPS + 6 )); GPU=${GPU:-a10g} +SHA=${SHA:-HEAD}; MODAL="$PY -m modal" +APP=yeto-$P; SP=$P-s; SAPP=yeto-$SP; TAPEVOL=yeto-event-tapes; STOREVOL=${STOREVOL:-yeto-eval-store}; HEAD_REGION=nebius/eu-north1 +export CARGO_HOME=/home/michael/work/gpu-default-modal/cargo RUSTUP_HOME=/home/michael/work/gpu-default-modal/rustup +export PATH=$CARGO_HOME/bin:$PATH +T=$(ps -u michael -L --no-headers | wc -l) +[ "${PLAN_ONLY:-0}" = 1 ] || [ "$T" -lt 3000 ] || { echo "abort: $T user threads (>= 3000)"; exit 3; } +pgrep -x raylet >/dev/null && { echo "abort: local Ray running"; exit 3; } +HEAD_CL=$(cd $REPO && PYTHONPATH=$REPO /usr/bin/python3 -c "from yeto.launch.sky_ops import sky_cluster_name; print(sky_cluster_name('$P-head'))") +[ -e $R/start_utc.txt ] && { echo "abort: $R already used"; exit 67; } +mkdir -p $R/runs $R/sruns $R/head $R/tapes $R/store +# a real clone: the code pack needs a git checkout. HOME stays the real one: the remote sky API +# server cannot unpack a file mount reached through a symlinked HOME dir (run a, FileExistsError). +git clone -q --no-local $REPO $R/yeto && git -C $R/yeto checkout -q $SHA +touch $R/yeto/yeto-rl-echo-events +cp $REPO/yeto/rl/gsm8k_reward.py $R/yeto/gsm8k_reward.py +(cd $R/yeto && git add -f gsm8k_reward.py yeto-rl-echo-events && git -c user.email=eval@yeto -c user.name=t-eval commit -qm 'run inputs') +MODEL="--model Qwen/Qwen3-0.6B --model-revision c1899de289a04d12100db370d81485cdf75e47ca" +DATA="--data zhuzilin/gsm8k --data-revision 0cbd9f31d91ac21a7613dcbc7fef992adac459ae" +LORA="--tuning lora --lora-r 16" +MODALX="--modal-gpu-exact --modal-retries 0 --modal-cpu 8 --modal-memory 64 --modal-tape-volume $TAPEVOL" +STORE="--rl-checkpoint-store modal-volume://$STOREVOL/$P" +COMMON="--training-mode rl --rl-engine ports $MODALX $MODEL $DATA $LORA $STORE --rl-algorithm-spec examples/rl_algorithms/decoupled-ppo-c2.json --rl-max-policy-age 2 --rl-observe-timeline --reward-function gsm8k_reward:score --fragments 1 --pipeline 1 --rollout-batch-size 32 --n-samples-per-prompt 4 --rollout-max-response-len 1024 --seq-len 2048 --inner-lr 1e-5 --seed 17 --apply-chat-template-kwargs '{\"enable_thinking\": false}' --trust-remote-code" +TARGS="launch --controller head --syncer-region $HEAD_REGION --preflight-threads wait --rl-stall-timeout 2400 --on-demand --no-island-relaunch --modal-timeout-s $HARD --cluster-prefix $P --gpu modal:1x$GPU,modal:1x$GPU --rl-island-scheduling elastic --rl-sample-role trainer --rl-sample-wait-s 900 $COMMON --total-steps $STEPS" +SARGS="launch --controller local --cluster-prefix $SP --gpu modal:1x$GPU --rl-single-island-no-sync --rl-sample-role sampler --rl-sampler-max-backlog-rounds 4 $COMMON --total-steps $SROUNDS" +git -C $REPO rev-parse $SHA > $R/yeto_sha.txt; echo "$TARGS" > $R/args.txt; echo "$SARGS" > $R/sargs.txt +echo "$APP $SAPP $HEAD_CL" > $R/resources.txt +eval "$(/usr/bin/python3 - <<'PYC' +import shlex, tomllib +t = tomllib.load(open("/home/michael/.modal.toml", "rb")) +prof = next((v for v in t.values() if isinstance(v, dict) and v.get("active")), None) or next(v for v in t.values() if isinstance(v, dict)) +print(f"export MODAL_TOKEN_ID={shlex.quote(prof['token_id'])} MODAL_TOKEN_SECRET={shlex.quote(prof['token_secret'])}") +PYC +)" +export YETO_ISLAND_HMAC_KEY=${YETO_ISLAND_HMAC_KEY:-$(openssl rand -hex 32)} +[ -n "${MODAL_TOKEN_ID:-}" ] || { echo "abort: no modal token"; exit 65; } +export PYTHONPATH=$R/yeto +if [ "${PLAN_ONLY:-0}" = 1 ]; then + echo "PLAN_ONLY run=$P apps=$APP,$SAPP head=$HEAD_CL steps=$STEPS sampler_rounds=$SROUNDS hard=${HARD}s code=$(cut -c1-8 $R/yeto_sha.txt) threads=$T" + (cd $R/yeto && eval "YETO_RUNS_DIR=$R/runs $PY -m yeto.cli $TARGS --dry-run") > $R/plan.json 2> $R/plan.err; echo "trainer dry-run rc=$?"; tail -3 $R/plan.err + (cd $R/yeto && eval "YETO_RUNS_DIR=$R/sruns $PY -m yeto.cli $SARGS --dry-run") > $R/splan.json 2> $R/splan.err; echo "sampler dry-run rc=$?"; tail -3 $R/splan.err + exit 0 +fi +date -u +%FT%TZ > $R/start_utc.txt +cat > $R/teardown.sh </dev/null) +if timeout 60 $SKY status $HEAD_CL 2>/dev/null | grep -q ' UP '; then + timeout 600 $SKY logs $HEAD_CL \$JOB --no-follow > $R/\${1:-final}-launch.log 2>&1 + timeout 300 $SKY queue $HEAD_CL > $R/head/queue.txt 2>&1 + timeout 600 \$S $HEAD_CL 'cd ~ && tar czf - yeto-output yeto-syncer.log 2>/dev/null' > $R/head/head-\${1:-final}.tgz 2>/dev/null + (cd $R/head && tar xzf head-\${1:-final}.tgz 2>/dev/null) +fi +for a in $APP $SAPP; do case \$a in yeto-$P|yeto-$P-s) timeout 300 $MODAL app stop --yes \$a;; *) echo REFUSE-STOP \$a;; esac; done > $R/\${1:-final}-stop.out 2>&1 +timeout 900 $SKY down -y $HEAD_CL > $R/\${1:-final}-down.out 2>&1; echo "down rc=\$?" >> $R/\${1:-final}-down.out +sleep 15; timeout 120 $MODAL app list --json > $R/\${1:-final}-applist.json 2>&1 +timeout 120 $SKY status $HEAD_CL > $R/\${1:-final}-skystatus.txt 2>&1 +EOF +echo "ps_threads_at_start=$T" > $R/threads-ps.txt +# container tape + NVML puller (both apps) +( while sleep 20; do c=$(cd /tmp && HOME=/home/michael timeout 60 $MODAL container list --json 2>/dev/null | python3 -c "import json,sys; print(' '.join(x['container_id'] for x in json.load(sys.stdin) if x.get('app_name') in ('$APP','$SAPP')))" 2>/dev/null) + echo "$(date -u +%FT%TZ) $c" >> $R/containers.log + for x in $c; do + (cd /tmp && HOME=/home/michael timeout 60 $MODAL container exec $x -- sh -c "cat /root/yeto-output/rl-island-*.jsonl 2>/dev/null") > $R/tapes/.$x 2>/dev/null && [ -s $R/tapes/.$x ] && grep -q '"event"' $R/tapes/.$x && mv $R/tapes/.$x $R/tapes/$x.jsonl + (cd /tmp && HOME=/home/michael timeout 30 $MODAL container exec $x -- sh -c "nvidia-smi --query-gpu=timestamp,name,utilization.gpu,memory.used,power.draw --format=csv,noheader") 2>/dev/null | sed "s/^/$x,/" >> $R/nvml.csv; done; done ) 9>&- & +PULLER=$! +cd $R/yeto +exec 9>/home/michael/work/infra-drafts/GPU-LAUNCH.lock; flock 9 +date -u +%FT%TZ > $R/lock_utc.txt +systemd-run --user --unit=$P-wd --on-active=${WD}s /bin/bash -c "bash $R/teardown.sh watchdog; touch $R/WATCHDOG_FIRED" > $R/watchdog.out 2>&1 +( eval "YETO_RUNS_DIR=$R/runs timeout $HARD $PY -m yeto.cli $TARGS" 2>&1 | tee $R/launch.stream.log | awk '{ print strftime("%FT%TZ", systime(), 1) " " $0; fflush() }' > $R/launch.ts.log; echo "rc=${PIPESTATUS[0]}" > $R/submit_rc.txt ) 9>&- & +SUB=$! +JOB="" +for i in $(seq 1 360); do + JOB=$(grep -oP "submitted: job \K\d+" $R/launch.stream.log 2>/dev/null | head -1); [ -n "$JOB" ] && break + kill -0 $SUB 2>/dev/null || break; sleep 10 +done +echo "head_job=$JOB" > $R/head_job.txt +for i in $(seq 1 90); do [ "$(grep -c 'modal-island [01]\] requested' $R/launch.stream.log 2>/dev/null)" -ge 2 ] && break; kill -0 $SUB 2>/dev/null || break; sleep 10; done +# the sampler starts once both trainer islands are requested (it waits for the first release) +( eval "YETO_RUNS_DIR=$R/sruns timeout $HARD $PY -m yeto.cli $SARGS" 2>&1 | awk '{ print strftime("%FT%TZ", systime(), 1) " " $0; fflush() }' > $R/slaunch.ts.log; echo "rc=${PIPESTATUS[0]}" > $R/ssubmit_rc.txt ) 9>&- & +SSUB=$! +sleep 60; exec 9>&- +( cd /tmp; HOME=/home/michael timeout $HARD $MODAL app logs --timestamps $APP > $R/modal-live-logs.txt 2>&1 ) 9>&- & +wait $SUB +if [ -n "$JOB" ]; then + END=$(( $(date +%s) + HARD )) + while [ $(date +%s) -lt $END ]; do + st=$(cd /tmp && HOME=/home/michael timeout 120 $SKY queue $HEAD_CL 2>/dev/null | awk -v j=$JOB '$1==j{print}') + echo "$st" | grep -qE "RUNNING|PENDING|SETTING_UP|INIT" || break; sleep 30 + done +fi +date -u +%FT%TZ > $R/end_utc.txt +kill $PULLER 2>/dev/null +export HOME=/home/michael +bash $R/teardown.sh final +kill $SSUB 2>/dev/null +grep -E "^ *$JOB " $R/head/queue.txt 2>/dev/null | grep -q SUCCEEDED && echo "rc=0" > $R/rc.txt || echo "rc=1 ($(grep -E "^ *$JOB " $R/head/queue.txt 2>/dev/null | tr -s ' ' | cut -c1-120))" > $R/rc.txt +for a in $APP $SAPP; do + AID=$(/usr/bin/python3 -c "import json,sys;print(next((a['app_id'] for a in json.load(open(sys.argv[1])) if a.get('description')==sys.argv[2]),''))" $R/final-applist.json $a 2>/dev/null) + [ -n "$AID" ] && (cd /tmp && timeout 180 $MODAL app logs --timestamps $AID > $R/modal-app-logs-$a.txt 2>&1) + mkdir -p $R/tape-direct/$a; (cd /tmp && timeout 600 $MODAL volume get --force $TAPEVOL $a $R/tape-direct/$a) >> $R/tape-direct.out 2>&1 +done +(cd /tmp && timeout 900 $MODAL volume get --force $STOREVOL $P/policy-releases $R/store; timeout 900 $MODAL volume get --force $STOREVOL $P/sample-buffer/index.jsonl $R/store; timeout 300 $MODAL volume get --force $STOREVOL $P/sample-buffer/lanes $R/store; timeout 300 $MODAL volume ls --json $STOREVOL $P) > $R/store.out 2>&1 +systemctl --user stop $P-wd.timer 2>/dev/null +echo "done $P $(cat $R/rc.txt)" diff --git a/tools/eval/miles_local.sh b/tools/eval/miles_local.sh new file mode 100644 index 000000000..6a24cba14 --- /dev/null +++ b/tools/eval/miles_local.sh @@ -0,0 +1,101 @@ +#!/bin/bash +# Miles eval cell, local controller: one yeto launcher run on Modal (`ARGS`), and optionally +# a second, sampler-only launcher run against the same store (`SARGS`), and optionally one +# reclaim of an island container (`modal container stop`) after it wrote a sampler round. +# Derived from s1-runs/s25-c1/run.sh (S25 C1) and s1-runs/s24-l5.4/run.sh (S24 L5.4). +# Called by `python -m tools.eval run` (tools/eval/miles.py), which writes the env below. +# +# env: P (run id; app yeto-$P, store prefix $P), R (raw dir), REPO (yeto checkout), +# 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) +set -u +: "${P:?}" "${R:?}" "${REPO:?}" "${ARGS:?}" "${STOREVOL:?}" "${PY:?}" +HARD=${HARD:-3000}; ISLANDS=${ISLANDS:-1}; TAPEVOL=${TAPEVOL:-yeto-event-tapes} +MODAL="$PY -m modal"; APP=yeto-$P; SAPP=yeto-$P-s; LOCK=/home/michael/work/infra-drafts/GPU-LAUNCH.lock +CLONE=$R/yeto +T=$(ps -u michael -L --no-headers | wc -l); [ -n "${DRY:-}" ] || [ "$T" -lt 3000 ] || { echo "abort: $T threads"; exit 3; } +pgrep -x raylet >/dev/null && { echo "abort: local Ray"; exit 3; } +[ -e $R/start_utc.txt ] && { echo "abort: $R used"; exit 67; } +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') +eval "$(/usr/bin/python3 - <<'PYC' +import shlex, tomllib +t = tomllib.load(open("/home/michael/.modal.toml", "rb")) +prof = next((v for v in t.values() if isinstance(v, dict) and v.get("active")), None) or next(v for v in t.values() if isinstance(v, dict)) +print(f"export MODAL_TOKEN_ID={shlex.quote(prof['token_id'])} MODAL_TOKEN_SECRET={shlex.quote(prof['token_secret'])}") +PYC +)" +echo "$ARGS" > $R/args.txt; echo "${SARGS:-}" > $R/sargs.txt +if [ -n "${DRY:-}" ]; then + (cd $CLONE && eval "PYTHONPATH=$CLONE YETO_RUNS_DIR=$R/runs $PY -m yeto.cli $ARGS --dry-run") > $R/plan.json 2> $R/plan.err + echo "dry-run rc=$?"; tail -3 $R/plan.err + if [ -n "${SARGS:-}" ]; then + (cd $CLONE && eval "PYTHONPATH=$CLONE YETO_RUNS_DIR=$R/sruns $PY -m yeto.cli $SARGS --dry-run") > $R/splan.json 2> $R/splan.err + echo "sampler dry-run rc=$?"; tail -3 $R/splan.err + fi + exit 0 +fi +stop_apps(){ for a in $APP $SAPP; do case $a in yeto-$P|yeto-$P-s) (cd /tmp && timeout 120 $MODAL app stop -y $a);; esac; done > $R/final-stop.out 2>&1; } +trap stop_apps EXIT +setsid nohup bash -c "sleep $(( HARD + 300 )); cd /tmp; $MODAL app stop -y $APP; $MODAL app stop -y $SAPP; touch $R/WATCHDOG_FIRED" >/dev/null 2>&1 < /dev/null & +echo $! > $R/watchdog.pid +# container tape + NVML puller (both apps) +( while sleep 20; do c=$(cd /tmp && timeout 60 $MODAL container list --json 2>/dev/null | python3 -c "import json,sys; print(' '.join(x['container_id'] for x in json.load(sys.stdin) if x.get('app_name') in ('$APP','$SAPP')))" 2>/dev/null) + echo "$(date -u +%FT%TZ) $c" >> $R/containers.log + for x in $c; do + (cd /tmp && timeout 60 $MODAL container exec $x -- sh -c "cat /root/yeto-output/rl-island-*.jsonl 2>/dev/null") > $R/tapes/.$x 2>/dev/null && [ -s $R/tapes/.$x ] && grep -q '"event"' $R/tapes/.$x && mv $R/tapes/.$x $R/tapes/$x.jsonl + (cd /tmp && timeout 30 $MODAL container exec $x -- sh -c "nvidia-smi --query-gpu=timestamp,name,utilization.gpu,memory.used,power.draw --format=csv,noheader") 2>/dev/null | sed "s/^/$x,/" >> $R/nvml.csv; done + done ) 9>&- & +PULLER=$! +exec 9>$LOCK; flock 9 +date -u +%FT%TZ > $R/start_utc.txt +( cd $CLONE && eval "PYTHONPATH=$CLONE YETO_RUNS_DIR=$R/runs timeout $HARD $PY -m yeto.cli $ARGS" 2>&1 | awk '{ print strftime("%FT%TZ", systime(), 1) " " $0; fflush() }' > $R/launch.ts.log; echo ${PIPESTATUS[0]} > $R/rc ) 9>&- & +SUB=$! +SSUB="" +if [ -n "${SARGS:-}" ]; then + ( cd $CLONE && eval "PYTHONPATH=$CLONE YETO_RUNS_DIR=$R/sruns timeout $HARD $PY -m yeto.cli $SARGS" 2>&1 | awk '{ print strftime("%FT%TZ", systime(), 1) " " $0; fflush() }' > $R/slaunch.ts.log; echo ${PIPESTATUS[0]} > $R/src ) 9>&- & + SSUB=$! +fi +# release the launch lock once every island has a container (or the launch ended) +for i in $(seq 1 120); do + n=$(cat $R/launch.ts.log $R/slaunch.ts.log 2>/dev/null | grep -ac 'rank 0 container') + [ "$n" -ge "$ISLANDS" ] && break + kill -0 $SUB 2>/dev/null || break; sleep 10 +done +exec 9>&- +# one reclaim: island $RECLAIM_LABEL wrote rollout $RECLAIM_ROUND -> stop its container +if [ -n "${RECLAIM_LABEL:-}" ]; then +( cd /tmp + for i in $(seq 1 600); do + kill -0 $SUB 2>/dev/null || { echo "$(date -u +%FT%TZ) launch ended before the reclaim" >> $R/reclaim.log; exit 0; } + if grep -a -- "-$RECLAIM_LABEL-modal" $R/launch.ts.log | grep -a '"event":"rl_sampler_round"' | grep -aq "\"rollout_id\":$RECLAIM_ROUND[,}]"; then + CID=$(grep -a -oP "\[$P-$RECLAIM_LABEL-modal\] \[modal-island \d+\] rank 0 container \K\S+" $R/launch.ts.log | tail -1) + echo "$(date -u +%FT%TZ) trigger container=$CID" >> $R/reclaim.log + [ -n "$CID" ] && timeout 120 $MODAL container stop --yes $CID >> $R/reclaim.log 2>&1; echo "stop rc=$? $(date -u +%FT%TZ) $(date +%s.%N)" >> $R/reclaim.log + exit 0 + fi + sleep 5 + done ) 9>&- & +fi +wait $SUB +# a sampler-only launch ends at its backlog cap once the trainer is gone: stop it +if [ -n "$SSUB" ]; then sleep 30; (cd /tmp && timeout 120 $MODAL app stop -y $SAPP) >> $R/final-stop.out 2>&1; wait $SSUB; fi +date -u +%FT%TZ > $R/end_utc.txt +sed 's/^[^ ]* //' $R/launch.ts.log > $R/launch.log +kill $PULLER 2>/dev/null; kill $(cat $R/watchdog.pid) 2>/dev/null +stop_apps; trap - EXIT; sleep 10 +(cd /tmp && timeout 120 $MODAL app list --json) > $R/final-applist.json 2>&1 +mkdir -p $R/tape-direct +for a in $APP $SAPP; do + (cd /tmp && timeout 600 $MODAL volume get --force $TAPEVOL $a $R/tape-direct) >> $R/tape-direct.out 2>&1 + AID=$(/usr/bin/python3 -c "import json,sys;print(next((a['app_id'] for a in json.load(open(sys.argv[1])) if a.get('description')==sys.argv[2]),''))" $R/final-applist.json $a 2>/dev/null) + [ -n "$AID" ] && (cd /tmp && timeout 180 $MODAL app logs --timestamps $AID > $R/modal-app-logs-$a.txt 2>&1) +done +(cd /tmp && for d in sample-buffer policy-releases round-cuts trainer-control; do timeout 900 $MODAL volume get --force $STOREVOL $P/$d $R/store; done + timeout 300 $MODAL volume ls --json $STOREVOL $P) > $R/store.out 2>&1 +echo "$P rc=$(cat $R/rc) src=$(cat $R/src 2>/dev/null)" diff --git a/tools/eval/modal_verl.py b/tools/eval/modal_verl.py new file mode 100644 index 000000000..81067c3a3 --- /dev/null +++ b/tools/eval/modal_verl.py @@ -0,0 +1,552 @@ +"""verl cells: a direct Modal harness (no yeto launcher; the S25 C2/C3V/C4 path). + +Container functions (one GPU each unless noted), all writing their raw data to +the results Volume under ``//`` before they return: + +* ``syncer`` (CPU): the elastic yeto-syncer built from this checkout's + ``syncer/``, reachable through a Modal TCP tunnel (address in the control Dict); +* ``sampler``: the neutral engine sampler island (SGLang server adapter, TP1) in + the Miles ports image; optionally reclaimed once by the local driver + (``modal container stop``) after it wrote round ``reclaim_after_round``; +* ``trainer``: ``yeto.rl.adapters.verl.island_entry --rl-sample-role trainer`` + (a buffer trainer; with ``learners`` >= 2 synced through the syncer, lane m); +* ``island``: one plain verl island (rollout + trainer, ``--rl-single-island-no-sync``); + phase ``a`` is commanded (``run 2`` then ``stop STOPPED``: SCHEDULED cut into + the store), ``b`` resumes from the store, ``c`` is the continuous control; +* ``collect`` (CPU): copies the store's records (no weights) into the results. + +The local side (:func:`run`) checks the preflight, spawns the calls in a +detached app (a killed client never cancels a spawned call), triggers the +reclaim, waits, collects, pulls ``/`` from the results Volume into +``/raw/`` and stops only its own app (``yeto-eval-``). +""" + +from __future__ import annotations + +import json +import os +import secrets +import shlex +import shutil +import subprocess +import sys +import time +from pathlib import Path +from typing import Any + +REPO = Path(__file__).resolve().parents[2] +APP_PREFIX = "yeto-eval-" +STORE_VOLUME = "yeto-eval-store" +RESULTS_VOLUME = "yeto-eval-results" +CONTROL = "yeto-eval-control" +STORE = "/mnt/store" +RESULTS = "/mnt/results" +WORKDIR = "/root/sky_workdir" +VERL_PY = "/workspace/verl/.venv/bin/python" +MODEL = "Qwen/Qwen3-0.6B" +MODEL_REV = "c1899de289a04d12100db370d81485cdf75e47ca" +DATA = "zhuzilin/gsm8k" +DATA_REV = "0cbd9f31d91ac21a7613dcbc7fef992adac459ae" +SYNCER_PORT = 29400 +GPU_NAMES = {"A10G": "NVIDIA A10", "L4": "L4", "H100!": "H100"} # Modal A10G reports "NVIDIA A10" +WEIGHT_SUFFIXES = (".safetensors", ".bin", ".pt", ".pth", ".distcp") +MAX_COLLECT_BYTES = 32 << 20 +# runs in the verl venv: the head's commands for trainer incarnation 0 (run N rounds, +# then stop STOPPED), printed island environment as JSON (S25 C4 modal_c4.py) +COMMAND_WRITER = ( + "import json, sys\n" + "from yeto.rl.elastic_client import LeaveReason\n" + "from yeto.rl.engine import trainer_remote as t\n" + "from yeto.rl.engine.resume import store_for\n" + "w = t.CommandWriter(t.LocalStoreFiles(store_for(sys.argv[1])), 0)\n" + "w.send(t.RUN, rounds=int(sys.argv[2]))\n" + "w.send(t.STOP, reason=LeaveReason.STOPPED)\n" + "print(json.dumps(t.island_env(0, None, poll_s=2.0)))\n") + + +# ====================================================================== in-container helpers +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): + raise RuntimeError(f"asked for {want}, got {names}") + return names + + +def _nvml(path: Path) -> subprocess.Popen: + return subprocess.Popen(["bash", "-c", "nvidia-smi --query-gpu=timestamp,name,utilization.gpu,memory.used," + f"power.draw --format=csv -l 5 >> {path}"]) + + +def _commit_results() -> None: + import modal + + modal.Volume.from_name(RESULTS_VOLUME).commit() + + +def _store_env() -> dict[str, str]: + from yeto.cloud.modal_ckpt_store import MODAL_VOLUME_ENV, STORE_IMPL + from yeto.rl.engine.resume import STORE_IMPL_ENV + + return {STORE_IMPL_ENV: STORE_IMPL, MODAL_VOLUME_ENV: STORE_VOLUME, + # the islands' interpreters lack the Modal client; commit/reload run here + "YETO_MODAL_PYTHON": sys.executable, "YETO_MODAL_SYSPATH": ":".join(p for p in sys.path if p)} + + +def _finish(out: Path, meta: dict, began: float, copies: dict[str, Path]) -> None: + meta["seconds"] = round(time.time() - began, 1) + for name, src in copies.items(): + if src.is_dir(): + shutil.copytree(src, out / name, dirs_exist_ok=True) + elif src.is_file(): + shutil.copy2(src, out / name) + (out / "meta.json").write_text(json.dumps(meta, indent=1, default=str)) + _commit_results() + + +def _incarnation(control, key: str) -> int: + """0 for the first container of an input, 1 for the next (Modal re-runs an + input in a new container after a reclaim).""" + n = int(control.get(key, -1)) + 1 + control[key] = n + return n + + +# ---------------------------------------------------------------------- syncer +def syncer(run: str, cfg: dict) -> dict: + import modal + + out = Path(RESULTS) / run / "syncer" + out.mkdir(parents=True, exist_ok=True) + control = modal.Dict.from_name(CONTROL, create_if_missing=True) + began = time.time() + argv = ["/opt/yeto-syncer", "--port", str(SYNCER_PORT), "--learners", str(cfg["learners"]), + "--total-steps", str(cfg["rounds"]), "--event-tape", "/tmp/syncer-tape.jsonl", + "--island-scheduling-mode", "elastic", "--quorum-theta", "0.75", "--carry-gamma", "0.5", + "--soft-deadline-s", "600", "--q-min", "1", "--max-carry-lag", "2", "--syncer-epoch", "0", + "--island-lease-s", "150.0", "--final-grace-s", "300"] + meta: dict[str, Any] = {"argv": argv, "container": os.environ.get("MODAL_TASK_ID")} + try: + with modal.forward(SYNCER_PORT, unencrypted=True) as tunnel, (out / "syncer.log").open("w") as log: + host, port = tunnel.tcp_socket + proc = subprocess.Popen(argv, stdout=log, stderr=subprocess.STDOUT) + time.sleep(2.0) + control[f"{run}/syncer"] = meta["address"] = f"{host}:{port}" + meta["rc"] = proc.wait(timeout=cfg["timeout_s"] - 300) + except BaseException as exc: # noqa: BLE001 - recorded, then re-raised + meta["error"] = f"{type(exc).__name__}: {exc}" + raise + finally: + _finish(out, meta, began, {"syncer-tape.jsonl": Path("/tmp/syncer-tape.jsonl")}) + return meta + + +# ---------------------------------------------------------------------- sampler +def sampler(run: str, cfg: dict) -> dict: + import modal + + sys.path.insert(0, WORKDIR) + from yeto.cloud.modal_eval_island import _image_paths + + _image_paths() + control = modal.Dict.from_name(CONTROL, create_if_missing=True) + incarnation = _incarnation(control, f"{run}/sampler/incarnation") + out = Path(RESULTS) / run / f"sampler-{incarnation}" + out.mkdir(parents=True, exist_ok=True) + control[f"{run}/sampler/container"] = os.environ.get("MODAL_TASK_ID") + names = _gpu(out, cfg) + nvml = _nvml(out / "nvml.csv") + os.environ.update(_store_env()) + began = time.time() + meta: dict[str, Any] = {"gpu": names, "container": os.environ.get("MODAL_TASK_ID"), "cfg": cfg, + "incarnation": incarnation} + tape = Path("/tmp/sampler-tape.jsonl") + try: + import pandas as pd + import sglang + from huggingface_hub import hf_hub_download, snapshot_download + from transformers import AutoTokenizer + + base = snapshot_download(MODEL, revision=MODEL_REV, local_dir="/tmp/qwen3-0.6b") + tok = AutoTokenizer.from_pretrained(base) + frame = pd.read_parquet(hf_hub_download(DATA, "train.parquet", repo_type="dataset", revision=DATA_REV)) + prompts = Path("/tmp/prompts.jsonl") + with prompts.open("w") as handle: + for _, row in frame.head(512).iterrows(): + messages = [{"role": m["role"], "content": m["content"]} for m in row["messages"]] + text = tok.apply_chat_template(messages, tokenize=False, add_generation_prompt=True, + enable_thinking=False) + handle.write(json.dumps({"text": text, "label": str(row["label"])}) + "\n") + from yeto.rl.engine import engine_sampler_entry as entry + from yeto.rl.engine.engine_sampler import build_engine_sampler + from yeto.rl.engine.sampler_checks import PostLoadChecks + + spec = { + "sampler_id": "eval-sglang-0", + "engine": {"factory": "yeto.rl.adapters.sampler_engines.from_spec:build", "adapter": "sglang-server", + "server": {"base_dir": base, "card_model": names[0], "device_family": "cuda", + "engine_version": sglang.__version__, "precision": "bf16", "quant_scheme": "bf16", + "topology": "tp1", "max_model_len": cfg["max_model_len"], "port": 30000, + "python": sys.executable, "start_timeout_s": 1200.0, "in_place_reload": True, + "extra_args": ["--mem-fraction-static", str(cfg["sglang_mem"])]}, + "options": {"tp": 1, "served_model_name": "qwen3-0.6b"}}, + "sampling": {"max_tokens": cfg["max_tokens"], "temperature": 1.0, "top_p": 1.0, "top_k": -1, + "n": cfg["samples_per_group"]}, + "num_rollouts": cfg["sampler_rounds"], "max_backlog_groups": cfg["max_backlog_groups"], + "max_prompt_tokens": 512, "base_version": 0, "base_dir": base, "poll_s": 2.0, + "store": f"{STORE}/{run}", + "prompts": {"path": str(prompts), "input_key": "text", "label_key": "label", + "groups_per_round": cfg["groups_per_round"], "format": "text"}, + "reward": "yeto.rl.gsm8k_reward:score", "reward_kind": "miles", + "checks": {"expected_expert_matrices": None, "fixed_prompts": []}, + "event_tape": str(tape), + } + (out / "spec.json").write_text(json.dumps(spec, indent=1)) + config = entry.island_config(spec) + write = entry.tape_writer(spec["event_tape"], config.sampler_id) + + def emit(event: str, **fields: Any) -> None: + write(event, **fields) + if event == "rl_sampler_round": # the local driver's reclaim trigger + control[f"{run}/sampler/round"] = {"incarnation": incarnation, **{ + k: fields.get(k) for k in ("rollout_id", "groups")}} + + releases, buffer = entry.build_stores(spec) + engine = entry.build_engine(spec) + done_keys = [f"{run}/done/{i}" for i in range(cfg["learners"])] + deadline = began + cfg["timeout_s"] - 300 + + def should_stop() -> bool: + return all(control.get(k, False) for k in done_keys) or time.time() > deadline + + try: + island = build_engine_sampler(config, engine=engine, releases=releases, buffer=buffer, emit=emit, + prompts=entry.build_prompts(spec), reward=entry.build_reward(spec), + checks=PostLoadChecks(**spec["checks"]), should_stop=should_stop) + meta["rounds"] = island.run() + finally: + engine.stop() + except BaseException as exc: # noqa: BLE001 - recorded, then re-raised + meta["error"] = f"{type(exc).__name__}: {exc}" + raise + finally: + nvml.terminate() + _finish(out, meta, began, {"sampler-tape.jsonl": tape}) + return meta + + +# ---------------------------------------------------------------------- trainer +def _ray_start(env: dict, log) -> None: + subprocess.run(["/workspace/verl/.venv/bin/ray", "start", "--head", "--node-ip-address=127.0.0.1", + "--port=6379", "--num-cpus=32", "--num-gpus=1", "--include-dashboard=false", + "--temp-dir=/root/eval-ray"], env=env, check=True, timeout=300, stdout=log, stderr=log) + + +def _island_env(cfg: dict, learner_id: int) -> dict[str, str]: + from yeto.hw.catalog import COMPAT_GROUP_ENV, compat_group + + # the launcher's island environment (rl_island_task): compat group and learner id + return dict(os.environ, PYTHONPATH=WORKDIR, HF_HUB_DISABLE_TELEMETRY="1", NCCL_DEBUG="WARN", + RAY_ADDRESS="127.0.0.1:6379", LEARNER_ID=str(learner_id), + **{COMPAT_GROUP_ENV: compat_group(cfg["gpu"].rstrip("!"))}) + + +def _common_island_argv(cfg: dict, learner_id: int, tape: Path) -> list[str]: + return [VERL_PY, "-m", "yeto.rl.adapters.verl.island_entry", "--model", MODEL, "--model-revision", MODEL_REV, + "--data", DATA, "--data-revision", DATA_REV, "--learner-id", str(learner_id), + "--reward-function", "gsm8k_reward:score", "--global-rounds", str(cfg["rounds"]), + "--groups-per-round", str(cfg["groups_per_round"]), "--samples-per-group", str(cfg["samples_per_group"]), + "--rollout-max-response-len", str(cfg["max_tokens"]), "--event-tape", str(tape), + "--inner-lr", str(cfg["lr"]), "--lora-r", str(cfg["lora_r"]), "--seq-len", str(cfg["seq_len"]), + "--seed", str(cfg["seed"]), "--apply-chat-template-kwargs", '{"enable_thinking": false}'] + + +def trainer(run: str, cfg: dict, learner_id: int) -> dict: + import modal + + sys.path.insert(0, WORKDIR) + out = Path(RESULTS) / run / f"trainer{learner_id}" + out.mkdir(parents=True, exist_ok=True) + names = _gpu(out, cfg) + nvml = _nvml(out / "nvml.csv") + control = modal.Dict.from_name(CONTROL, create_if_missing=True) + began = time.time() + meta: dict[str, Any] = {"gpu": names, "container": os.environ.get("MODAL_TASK_ID"), "cfg": cfg, + "learner_id": learner_id} + from yeto.rl.engine.algorithm import AlgorithmSpec, CorrectionSpec, ExecutionSpec + + spec = AlgorithmSpec(execution=ExecutionSpec(max_policy_staleness=cfg["max_policy_age"]), + correction=CorrectionSpec(method="tis", tis_clip=2.0, tis_clip_low=0.0)) + Path("/tmp/algo.json").write_text(json.dumps(spec.to_dict())) + env = {**_island_env(cfg, learner_id), **_store_env()} + tape = out / "trainer-tape.jsonl" + try: + argv = _common_island_argv(cfg, learner_id, tape) + [ + "--rl-max-policy-age", str(cfg["max_policy_age"]), "--rl-algorithm-spec", "/tmp/algo.json", + "--rl-sample-role", "trainer", "--rl-sample-wait-s", "1200"] + if cfg["learners"] > 1: + address = None + for _ in range(600): # the syncer container publishes its tunnel address + address = control.get(f"{run}/syncer") + if address: + break + time.sleep(2.0) + if not address: + raise RuntimeError("no syncer address within 1200 s") + argv += ["--syncer", address, "--rl-island-scheduling", "elastic", "--rl-syncer-epoch", "0", + "--rl-island-lease-s", "150.0", "--rl-sample-store", f"{STORE}/{run}", + "--rl-sample-lanes", str(cfg["learners"])] + else: + argv += ["--rl-single-island-no-sync", "--rl-resume-store", f"{STORE}/{run}"] + meta["argv"] = argv + with (out / "island.log").open("w") as log: + _ray_start(env, log) + meta["rc"] = subprocess.run(argv, env=env, stdout=log, stderr=subprocess.STDOUT, + cwd="/workspace/verl", timeout=cfg["timeout_s"] - 300).returncode + except BaseException as exc: # noqa: BLE001 - recorded, then re-raised + meta["error"] = f"{type(exc).__name__}: {exc}" + raise + finally: + control[f"{run}/done/{learner_id}"] = True + nvml.terminate() + _finish(out, meta, began, {"yeto-output": Path("~/yeto-output").expanduser()}) + return meta + + +# ---------------------------------------------------------------------- plain island (E1, E5) +def island(run: str, cfg: dict, phase: str) -> dict: + """One verl island (rollout + trainer); phases as in the module doc.""" + import modal + + sys.path.insert(0, WORKDIR) + modal.Volume.from_name(STORE_VOLUME).reload() # see what the previous phase committed + out = Path(RESULTS) / run / f"island-{phase}" + out.mkdir(parents=True, exist_ok=True) + done = out / "meta.json" + if done.is_file() and json.loads(done.read_text()).get("rc") == 0: + return {**json.loads(done.read_text()), "reused_after_restart": True} # never rerun a finished phase + names = _gpu(out, cfg) + nvml = _nvml(out / "nvml.csv") + began = time.time() + meta: dict[str, Any] = {"gpu": names, "container": os.environ.get("MODAL_TASK_ID"), "cfg": cfg, "phase": phase} + env = {**_island_env(cfg, 0)} + store = f"{STORE}/{run}/cuts" + state = Path(f"/root/state-{phase}") + try: + argv = _common_island_argv(cfg, 0, Path("/root/yeto-output/rl-island-0.jsonl")) + [ + "--rl-single-island-no-sync", "--lora-targets", "all-linear"] + if phase in ("a", "b"): + argv += ["--rl-resume-store", store, "--rl-resume-state-dir", str(state)] + if phase == "a": # the head's commands for incarnation 0, through yeto's own writer + write = subprocess.run([VERL_PY, "-c", COMMAND_WRITER, store, str(cfg["stop_after_rounds"])], + env=env, capture_output=True, text=True, timeout=300) + if write.returncode != 0: + raise RuntimeError(f"command file not written: {write.stderr[-2000:]}") + modal.Volume.from_name(STORE_VOLUME).commit() + env.update(json.loads(write.stdout.strip().splitlines()[-1])) + meta["argv"] = argv + with (out / "island.log").open("w") as log: + _ray_start(env, log) + meta["rc"] = subprocess.run(argv, env=env, stdout=log, stderr=subprocess.STDOUT, cwd="/root", + timeout=cfg["timeout_s"] - 300).returncode + subprocess.run(["/workspace/verl/.venv/bin/ray", "stop", "--force"], stdout=log, stderr=log) + modal.Volume.from_name(STORE_VOLUME).commit() + except BaseException as exc: # noqa: BLE001 - recorded, then re-raised + meta["error"] = f"{type(exc).__name__}: {exc}" + raise + finally: + nvml.terminate() + _finish(out, meta, began, {"yeto-output": Path("/root/yeto-output"), "state-dir": state}) + return meta + + +# ---------------------------------------------------------------------- collect +def collect(run: str) -> dict: + """Copy the store's records under ``/`` (no weights, no large files) + into the results Volume ``/store/``.""" + import modal + + modal.Volume.from_name(STORE_VOLUME).reload() + src, dst = Path(STORE) / run, Path(RESULTS) / run / "store" + copied, skipped = 0, [] + for path in sorted(src.rglob("*")) if src.is_dir() else []: + if not path.is_file(): + continue + rel = path.relative_to(src) + if path.suffix in WEIGHT_SUFFIXES or path.stat().st_size > MAX_COLLECT_BYTES: + skipped.append([str(rel), path.stat().st_size]) + continue + (dst / rel).parent.mkdir(parents=True, exist_ok=True) + shutil.copy2(path, dst / rel) + copied += 1 + dst.mkdir(parents=True, exist_ok=True) + (dst / "_skipped.json").write_text(json.dumps(skipped, indent=0)) + _commit_results() + return {"copied": copied, "skipped": len(skipped)} + + +# ====================================================================== app +def build_app(run: str, cfg: dict): + import modal + + sys.path.insert(0, str(REPO)) + from yeto.rl.adapters.miles.pins import MILES_NEXT_IMAGE + from yeto.rl.adapters.verl.image import modal_image + + ignore = [".git", ".venv", ".claude", "**/__pycache__", "**/*.pyc", "s1-runs", "syncer/target", + "evidence*", "openspec", "docs", "tutorial", "data"] + sampler_image = (modal.Image.from_registry(MILES_NEXT_IMAGE.removeprefix("docker:")) + .env({"HOME": "/root", "PYTHONUNBUFFERED": "1"}) + .add_local_dir(str(REPO), WORKDIR, copy=False, ignore=ignore)) + verl_image = (modal_image(modal).env({"PYTHONUNBUFFERED": "1", "PYTHONPATH": WORKDIR}) + .add_local_dir(str(REPO), WORKDIR, copy=False, ignore=ignore)) + cpu_image = modal.Image.debian_slim(python_version="3.12").add_local_dir( + str(REPO / "yeto"), f"{WORKDIR}/yeto", copy=False, ignore=["**/__pycache__"]) + app = modal.App(APP_PREFIX + run) + store = modal.Volume.from_name(STORE_VOLUME, create_if_missing=True, version=2) + results = modal.Volume.from_name(RESULTS_VOLUME, create_if_missing=True, version=2) + vols = {STORE: store, RESULTS: results} + hmac = modal.Secret.from_dict({"YETO_ISLAND_HMAC_KEY": secrets.token_hex(32)}) + gpu, timeout = cfg["gpu"], cfg["timeout_s"] + fns = { + "sampler": app.function(image=sampler_image, gpu=gpu, cpu=8.0, memory=32768, timeout=timeout, retries=0, + volumes=vols, name="sampler", serialized=True)(sampler), + "trainer": app.function(image=verl_image, gpu=gpu, cpu=16.0, memory=65536, timeout=timeout, retries=0, + volumes=vols, secrets=[hmac], name="trainer", serialized=True)(trainer), + # max_inputs=1: every phase in a fresh container (a warm one keeps ~/yeto-output: C4 run c) + "island": app.function(image=verl_image, gpu=gpu, cpu=8.0, memory=49152, timeout=timeout, retries=0, + volumes=vols, max_inputs=1, name="island", serialized=True)(island), + "collect": app.function(image=cpu_image, cpu=1.0, memory=2048, timeout=900, volumes=vols, + name="collect", serialized=True)(collect), + } + if cfg.get("learners", 1) > 1: + syncer_image = (modal.Image.debian_slim(python_version="3.12") + .apt_install("curl", "build-essential", "pkg-config") + .run_commands("curl -sSf https://sh.rustup.rs | sh -s -- -y --profile minimal") + .add_local_dir(str(REPO / "syncer"), "/build/syncer", copy=True, ignore=["target"]) + .run_commands("cd /build/syncer && /root/.cargo/bin/cargo build --release --quiet" + " && cp target/release/yeto-syncer /opt/yeto-syncer")) + fns["syncer"] = app.function(image=syncer_image, cpu=2.0, memory=4096, timeout=timeout, retries=0, + volumes={RESULTS: results}, secrets=[hmac], name="syncer", + serialized=True)(syncer) + return app, fns + + +# ====================================================================== local driver +def _modal_cli(*args: str, timeout: float = 900) -> subprocess.CompletedProcess: + return subprocess.run([sys.executable, "-m", "modal", *args], capture_output=True, text=True, + timeout=timeout, cwd="/tmp") + + +def _log(path: Path, message: str) -> None: + with path.open("a") as fh: + fh.write(f"{time.strftime('%Y-%m-%dT%H:%M:%SZ', time.gmtime())} {message}\n") + + +def _wait_calls(calls: dict, deadline: float, log: Path, on_poll=None) -> dict[str, Any]: + import modal + + results: dict[str, Any] = {} + while len(results) < len(calls) and time.time() < deadline: + for name, call in calls.items(): + if name in results: + continue + try: + results[name] = call.get(timeout=0) + _log(log, f"{name} returned {json.dumps(results[name], default=str)[:300]}") + except TimeoutError: + pass + except modal.exception.OutputExpiredError: + results[name] = {"error": "output expired"} + except Exception as exc: # noqa: BLE001 - a failed call is a result + results[name] = {"error": f"{type(exc).__name__}: {exc}"} + _log(log, f"{name} failed: {results[name]['error']}") + if on_poll is not None: + on_poll() + time.sleep(10) + for name in calls: + results.setdefault(name, {"error": "timeout waiting for the call"}) + return results + + +def run(cell, run_id: str, run_dir: Path) -> dict[str, Any]: + """Run one verl cell end to end (spawn, reclaim, wait, collect, pull).""" + import modal + + cfg = dict(cell.config) + raw = run_dir / "raw" + raw.mkdir(parents=True, exist_ok=True) + log = run_dir / "driver.log" + control = modal.Dict.from_name(CONTROL, create_if_missing=True) + if control.get(f"{run_id}/started", None) is not None: + raise SystemExit(f"run id {run_id} already used (control Dict {CONTROL})") + control[f"{run_id}/started"] = time.time() + app, fns = build_app(run_id, cfg) + began = time.time() + with modal.enable_output(), app.run(detach=True): + calls: dict[str, Any] = {} + if cfg["kind"] == "island": + calls["chain"] = None # phases run in order from here (a, then b; c beside them) + if "c" in cfg["phases"]: + calls["island-c"] = fns["island"].spawn(run_id, cfg, "c") + first = [p for p in cfg["phases"] if p != "c"] + if first: + calls["island-" + first[0]] = fns["island"].spawn(run_id, cfg, first[0]) + calls.pop("chain") + else: + if cfg["learners"] > 1: + calls["syncer"] = fns["syncer"].spawn(run_id, cfg) + calls["sampler"] = fns["sampler"].spawn(run_id, cfg) + for i in range(cfg["learners"]): + calls[f"trainer{i}"] = fns["trainer"].spawn(run_id, cfg, i) + _log(log, f"spawned {sorted(calls)} app {APP_PREFIX + run_id}") + deadline = began + cfg["timeout_s"] + 600 + reclaim = {"done": cfg.get("reclaim_after_round") is None} + + def maybe_reclaim() -> None: + if reclaim["done"]: + return + progress = control.get(f"{run_id}/sampler/round", None) or {} + if progress.get("incarnation") == 0 and int(progress.get("rollout_id", -1)) >= cfg["reclaim_after_round"]: + container = control.get(f"{run_id}/sampler/container", None) + result = _modal_cli("container", "stop", "--yes", str(container), timeout=120) + reclaim.update(done=True, container=container, rc=result.returncode, at=time.time(), + rollout_id=progress.get("rollout_id")) + _log(run_dir / "reclaim.log", f"stop container={container} rc={result.returncode} " + f"after rollout {progress.get('rollout_id')} {result.stdout.strip()} {result.stderr.strip()}") + + results = _wait_calls(calls, deadline, log, maybe_reclaim) + if cfg["kind"] == "island" and "b" in cfg["phases"]: + if (results.get("island-a") or {}).get("rc") == 0: + with modal.enable_output(), app.run(detach=True): + call_b = fns["island"].spawn(run_id, cfg, "b") + results.update(_wait_calls({"island-b": call_b}, time.time() + cfg["timeout_s"] + 600, log)) + else: + results["island-b"] = {"error": "phase a failed; b not run"} + with modal.enable_output(), app.run(detach=True): + collect_call = fns["collect"].spawn(run_id) + results["collect"] = _wait_calls({"collect": collect_call}, time.time() + 1200, log)["collect"] + wall = time.time() - began + pulled = _modal_cli("volume", "get", "--force", RESULTS_VOLUME, run_id, str(raw), timeout=1800) + stopped = _modal_cli("app", "stop", "-y", APP_PREFIX + run_id, timeout=120) # only our own app + summary = {"calls": results, "reclaim": reclaim, "wall_s": round(wall, 1), + "pull_rc": pulled.returncode, "stop_rc": stopped.returncode} + (run_dir / "driver.json").write_text(json.dumps(summary, indent=1, default=str)) + return summary + + +def card_seconds(raw: Path) -> float: + """GPU container seconds from the components' meta.json (one card each).""" + total = 0.0 + for meta in sorted(Path(raw).rglob("meta.json")): + data = json.loads(meta.read_text()) + if data.get("gpu"): + total += float(data.get("seconds") or 0.0) + return total + + +def describe(cell) -> str: + return shlex.join(["modal", "app", APP_PREFIX + "", json.dumps(cell.config, sort_keys=True)]) diff --git a/tools/eval/report.py b/tools/eval/report.py new file mode 100644 index 000000000..bc848332f --- /dev/null +++ b/tools/eval/report.py @@ -0,0 +1,93 @@ +"""Markdown and HTML reports over a set of run verdicts (one table, then one +section per run with every criterion, its evidence, and the metrics).""" + +from __future__ import annotations + +import html as _html +import json +from pathlib import Path +from typing import Any + +from .cells import BACKENDS, CAPABILITIES +from .verdict import load + +MARK = {"pass": "PASS", "fail": "FAIL", "expected_fail": "EXPECTED FAIL (unsupported)", + "unexpected_pass": "UNEXPECTED PASS", "error": "ERROR"} + + +def _latest(dirs: list[Path]) -> dict[tuple[str, str], dict[str, Any]]: + """The last verdict per cell (dirs in the given order).""" + out: dict[tuple[str, str], dict[str, Any]] = {} + for d in dirs: + doc = load(d) + out[(doc["capability"], doc["backend"])] = doc + return out + + +def _short(value: Any, limit: int = 160) -> str: + text = json.dumps(value, default=str) if not isinstance(value, str) else value + return text if len(text) <= limit else text[: limit - 3] + "..." + + +def markdown(dirs: list[Path], *, title: str) -> str: + cells = _latest(dirs) + lines = [f"# {title}", "", "| Capability | " + " | ".join(BACKENDS) + " |", "|---|" + "---|" * len(BACKENDS)] + for cap, text in CAPABILITIES.items(): + row = [f"{cap} {text}"] + for backend in BACKENDS: + doc = cells.get((cap, backend)) + if doc is None: + row.append("not run") + continue + r = doc.get("run") or {} + row.append(f"{MARK[doc['status']]}
`{doc['run_id']}` {r.get('card_seconds', '?')} card-s, " + f"~${r.get('usd_estimate', '?')}, wall {r.get('wall_s', '?')} s") + lines.append("| " + " | ".join(row) + " |") + total = sum(float((d.get("run") or {}).get("usd_estimate") or 0) for d in cells.values()) + lines += ["", f"Estimated GPU cost of the cells above: ${total:.2f} (list-price estimate, not a bill).", ""] + for (cap, backend), doc in sorted(cells.items()): + lines += [f"## {cap} {backend}: {MARK[doc['status']]}", "", + f"run `{doc['run_id']}`, code `{(doc.get('run') or {}).get('code_sha', '')[:10]}`" + + (f", error: {doc['error']}" if doc.get("error") else ""), "", + "| Criterion | Result | Evidence |", "|---|---|---|"] + for c in doc["criteria"]: + lines.append(f"| {c['id']}: {c['text']} | {'pass' if c['passed'] else '**FAIL**'} | " + f"`{_short(c['evidence'])}` |") + if doc.get("metrics"): + lines += ["", "Metrics: `" + _short(doc["metrics"], 900) + "`"] + lines.append("") + return "\n".join(lines) + "\n" + + +def html(dirs: list[Path], *, title: str) -> str: + cells = _latest(dirs) + esc = _html.escape + color = {"pass": "#1a7f37", "fail": "#cf222e", "expected_fail": "#9a6700", "unexpected_pass": "#0969da", + "error": "#cf222e"} + head = "".join(f"{esc(b)}" for b in BACKENDS) + rows = [] + for cap, text in CAPABILITIES.items(): + tds = [] + for backend in BACKENDS: + doc = cells.get((cap, backend)) + if doc is None: + tds.append("not run") + continue + r = doc.get("run") or {} + tds.append(f"{esc(MARK[doc['status']])}
" + f"{esc(doc['run_id'])}
{r.get('card_seconds', '?')} card-s, " + f"~${r.get('usd_estimate', '?')}") + rows.append(f"{esc(cap)} {esc(text)}{''.join(tds)}") + sections = [] + for (cap, backend), doc in sorted(cells.items()): + crit = "".join(f"{esc(c['id'])}: {esc(c['text'])}{'pass' if c['passed'] else 'FAIL'}" + f"{esc(_short(c['evidence'], 400))}" for c in doc["criteria"]) + sections.append(f"

{esc(cap)} {esc(backend)}: {esc(MARK[doc['status']])}

" + f"{crit}
CriterionResultEvidence
" + f"

{esc(_short(doc.get('metrics') or {}, 1500))}

") + style = ("body{font-family:sans-serif;margin:16px;background:#fff;color:#1f2328}" + "table{border-collapse:collapse;margin:8px 0}td,th{border:1px solid #d0d7de;padding:4px 8px;" + "vertical-align:top}code{font-size:12px;word-break:break-all}") + return (f"{esc(title)}" + f"

{esc(title)}

{head}{''.join(rows)}
Capability
" + f"{''.join(sections)}\n") diff --git a/tools/eval/verdict.py b/tools/eval/verdict.py new file mode 100644 index 000000000..2af2da038 --- /dev/null +++ b/tools/eval/verdict.py @@ -0,0 +1,109 @@ +"""The verdict every capability judge writes (``/verdict.json``). + +Status values: + +* ``pass`` -- every criterion passed; +* ``fail`` -- a criterion failed on a backend that supports the capability; +* ``expected_fail`` -- a criterion failed and the backend does not support the + capability (:attr:`cells.Cell.supported` is False): the run is kept as the + regression baseline of the gap; +* ``unexpected_pass`` -- an unsupported cell passed (update the matrix); +* ``error`` -- the run produced no judgeable data (harness or launch failure). +""" + +from __future__ import annotations + +import json +from collections.abc import Iterable, Mapping +from dataclasses import asdict, dataclass, field +from pathlib import Path +from typing import Any + +PASS, FAIL, EXPECTED_FAIL, UNEXPECTED_PASS, ERROR = "pass", "fail", "expected_fail", "unexpected_pass", "error" +VERDICT_FILE = "verdict.json" + + +@dataclass +class Criterion: + id: str + text: str + passed: bool + evidence: Any = None + + +@dataclass +class Verdict: + capability: str + backend: str + run_id: str + supported: bool + criteria: list[Criterion] = field(default_factory=list) + metrics: dict[str, Any] = field(default_factory=dict) + run: dict[str, Any] = field(default_factory=dict) # wall seconds, card seconds, cost estimate, code sha + error: str | None = None + + def check(self, cid: str, text: str, passed: bool, evidence: Any = None) -> bool: + self.criteria.append(Criterion(cid, text, bool(passed), evidence)) + return bool(passed) + + @property + def status(self) -> str: + if self.error is not None or not self.criteria: + return ERROR + ok = all(c.passed for c in self.criteria) + if self.supported: + return PASS if ok else FAIL + return UNEXPECTED_PASS if ok else EXPECTED_FAIL + + def to_dict(self) -> dict[str, Any]: + return {**asdict(self), "status": self.status} + + def write(self, run_dir: Path) -> Path: + path = Path(run_dir) / VERDICT_FILE + path.write_text(json.dumps(self.to_dict(), indent=1, default=str) + "\n") + return path + + +def load(run_dir: Path) -> dict[str, Any]: + return json.loads((Path(run_dir) / VERDICT_FILE).read_text()) + + +# ---------------------------------------------------------------- raw data readers +def jsonl(path: Path) -> list[dict[str, Any]]: + """The JSON objects of a JSONL file (missing file: []; bad lines skipped).""" + path = Path(path) + if not path.is_file(): + return [] + out = [] + for line in path.read_text(errors="replace").splitlines(): + line = line.strip() + start = line.find("{") + if start < 0: + continue + try: + value = json.loads(line[start:]) + except ValueError: + continue + if isinstance(value, dict): + out.append(value) + return out + + +def find(root: Path, name: str) -> Path: + """The first file called ``name`` under ``root`` (sorted), else ``root/name``.""" + hits = sorted(Path(root).rglob(name)) if Path(root).is_dir() else [] + return hits[0] if hits else Path(root) / name + + +def find_all(root: Path, pattern: str) -> list[Path]: + return sorted(Path(root).rglob(pattern)) if Path(root).is_dir() else [] + + +def events(tape: Iterable[Mapping[str, Any]], name: str) -> list[Mapping[str, Any]]: + return [e for e in tape if e.get("event") == name] + + +def objects(directory: Path) -> list[dict[str, Any]]: + """One JSON object per ``*.json`` file of a directory (object-store layouts).""" + directory = Path(directory) + return [json.loads(p.read_text()) for p in sorted(directory.glob("*.json"))] if directory.is_dir() else [] From c05b6cab20b21f217c88f3a0892334ea143ae802 Mon Sep 17 00:00:00 2001 From: michael <78113554+chowhsu@users.noreply.github.com> Date: Tue, 20 Oct 2026 06:29:47 +0000 Subject: [PATCH 58/89] tools/eval: price A10G containers with their CPU and memory requests Co-Authored-By: Claude Opus 5.5 --- tools/eval/__main__.py | 2 +- tools/eval/cells.py | 4 +++- 2 files changed, 4 insertions(+), 2 deletions(-) diff --git a/tools/eval/__main__.py b/tools/eval/__main__.py index eee7b4dcb..aaa1a3191 100644 --- a/tools/eval/__main__.py +++ b/tools/eval/__main__.py @@ -94,7 +94,7 @@ def judge(run_dir: Path) -> Verdict: rate = cells.A10G_USD_PER_HOUR if cfg.get("gpu") == "A10G" else 5.3 verdict.run.update(code_sha=cell_doc.get("code_sha"), wall_s=run_doc.get("wall_s"), card_seconds=round(seconds, 1), usd_estimate=round(seconds / 3600 * rate, 3), - usd_rate=f"{rate}/card-hour (list price estimate, not a bill)") + usd_rate=f"{rate}/container-hour incl. CPU+memory (list-price estimate, not a bill)") try: ev = evidence.load(run_dir) judges.JUDGES[cell.capability](ev, verdict, cfg) diff --git a/tools/eval/cells.py b/tools/eval/cells.py index 56b73ceb9..038b032e1 100644 --- a/tools/eval/cells.py +++ b/tools/eval/cells.py @@ -19,7 +19,9 @@ "E5": "round cut: SCHEDULED stop + resume assertions", } BACKENDS = ("miles", "verl") -A10G_USD_PER_HOUR = 1.10 # Modal list price used by the S25 C2/C3V estimates (not a bill) +# all-in estimate per A10G container hour: GPU $1.10 + the CPU (8-16 cores) and memory +# (32-64 GiB) the harness requests, which Modal bills too (list prices; not a bill) +A10G_USD_PER_HOUR = 2.0 @dataclass(frozen=True) From 8a9a6597ed14f7a308b4e33140831d98f25462fc Mon Sep 17 00:00:00 2001 From: michael <78113554+chowhsu@users.noreply.github.com> Date: Tue, 20 Oct 2026 06:31:27 +0000 Subject: [PATCH 59/89] tools/eval: one detached app per verl run; Modal 1.6 call timeouts Co-Authored-By: Claude Opus 5.5 --- tools/eval/modal_verl.py | 62 ++++++++++++++++------------------------ 1 file changed, 25 insertions(+), 37 deletions(-) diff --git a/tools/eval/modal_verl.py b/tools/eval/modal_verl.py index 81067c3a3..31e78fb1e 100644 --- a/tools/eval/modal_verl.py +++ b/tools/eval/modal_verl.py @@ -26,7 +26,6 @@ import json import os import secrets -import shlex import shutil import subprocess import sys @@ -403,8 +402,7 @@ def build_app(run: str, cfg: dict): .add_local_dir(str(REPO), WORKDIR, copy=False, ignore=ignore)) verl_image = (modal_image(modal).env({"PYTHONUNBUFFERED": "1", "PYTHONPATH": WORKDIR}) .add_local_dir(str(REPO), WORKDIR, copy=False, ignore=ignore)) - cpu_image = modal.Image.debian_slim(python_version="3.12").add_local_dir( - str(REPO / "yeto"), f"{WORKDIR}/yeto", copy=False, ignore=["**/__pycache__"]) + cpu_image = modal.Image.debian_slim(python_version="3.12") app = modal.App(APP_PREFIX + run) store = modal.Volume.from_name(STORE_VOLUME, create_if_missing=True, version=2) results = modal.Volume.from_name(RESULTS_VOLUME, create_if_missing=True, version=2) @@ -457,8 +455,8 @@ def _wait_calls(calls: dict, deadline: float, log: Path, on_poll=None) -> dict[s try: results[name] = call.get(timeout=0) _log(log, f"{name} returned {json.dumps(results[name], default=str)[:300]}") - except TimeoutError: - pass + except (TimeoutError, modal.exception.TimeoutError): + pass # still running except modal.exception.OutputExpiredError: results[name] = {"error": "output expired"} except Exception as exc: # noqa: BLE001 - a failed call is a result @@ -486,25 +484,7 @@ def run(cell, run_id: str, run_dir: Path) -> dict[str, Any]: control[f"{run_id}/started"] = time.time() app, fns = build_app(run_id, cfg) began = time.time() - with modal.enable_output(), app.run(detach=True): - calls: dict[str, Any] = {} - if cfg["kind"] == "island": - calls["chain"] = None # phases run in order from here (a, then b; c beside them) - if "c" in cfg["phases"]: - calls["island-c"] = fns["island"].spawn(run_id, cfg, "c") - first = [p for p in cfg["phases"] if p != "c"] - if first: - calls["island-" + first[0]] = fns["island"].spawn(run_id, cfg, first[0]) - calls.pop("chain") - else: - if cfg["learners"] > 1: - calls["syncer"] = fns["syncer"].spawn(run_id, cfg) - calls["sampler"] = fns["sampler"].spawn(run_id, cfg) - for i in range(cfg["learners"]): - calls[f"trainer{i}"] = fns["trainer"].spawn(run_id, cfg, i) - _log(log, f"spawned {sorted(calls)} app {APP_PREFIX + run_id}") - deadline = began + cfg["timeout_s"] + 600 - reclaim = {"done": cfg.get("reclaim_after_round") is None} + reclaim: dict[str, Any] = {"done": cfg.get("reclaim_after_round") is None} def maybe_reclaim() -> None: if reclaim["done"]: @@ -518,17 +498,28 @@ def maybe_reclaim() -> None: _log(run_dir / "reclaim.log", f"stop container={container} rc={result.returncode} " f"after rollout {progress.get('rollout_id')} {result.stdout.strip()} {result.stderr.strip()}") - results = _wait_calls(calls, deadline, log, maybe_reclaim) - if cfg["kind"] == "island" and "b" in cfg["phases"]: - if (results.get("island-a") or {}).get("rc") == 0: - with modal.enable_output(), app.run(detach=True): - call_b = fns["island"].spawn(run_id, cfg, "b") - results.update(_wait_calls({"island-b": call_b}, time.time() + cfg["timeout_s"] + 600, log)) - else: - results["island-b"] = {"error": "phase a failed; b not run"} + # one detached app for the whole run: a killed driver never cancels a spawned call with modal.enable_output(), app.run(detach=True): - collect_call = fns["collect"].spawn(run_id) - results["collect"] = _wait_calls({"collect": collect_call}, time.time() + 1200, log)["collect"] + calls: dict[str, Any] = {} + if cfg["kind"] == "island": # c (control) beside a; b after a, from a's store + for phase in [p for p in cfg["phases"] if p != "b"]: + calls[f"island-{phase}"] = fns["island"].spawn(run_id, cfg, phase) + else: + if cfg["learners"] > 1: + calls["syncer"] = fns["syncer"].spawn(run_id, cfg) + calls["sampler"] = fns["sampler"].spawn(run_id, cfg) + for i in range(cfg["learners"]): + calls[f"trainer{i}"] = fns["trainer"].spawn(run_id, cfg, i) + _log(log, f"spawned {sorted(calls)} app {APP_PREFIX + run_id}") + results = _wait_calls(calls, began + cfg["timeout_s"] + 600, log, maybe_reclaim) + if "b" in cfg.get("phases", ()): + if (results.get("island-a") or {}).get("rc") == 0: + call_b = fns["island"].spawn(run_id, cfg, "b") + results.update(_wait_calls({"island-b": call_b}, time.time() + cfg["timeout_s"] + 600, log)) + else: + results["island-b"] = {"error": "phase a failed; b not run"} + results["collect"] = _wait_calls({"collect": fns["collect"].spawn(run_id)}, time.time() + 1200, + log)["collect"] wall = time.time() - began pulled = _modal_cli("volume", "get", "--force", RESULTS_VOLUME, run_id, str(raw), timeout=1800) stopped = _modal_cli("app", "stop", "-y", APP_PREFIX + run_id, timeout=120) # only our own app @@ -547,6 +538,3 @@ def card_seconds(raw: Path) -> float: total += float(data.get("seconds") or 0.0) return total - -def describe(cell) -> str: - return shlex.join(["modal", "app", APP_PREFIX + "", json.dumps(cell.config, sort_keys=True)]) From 78d63277f85dbf2dc4b07645725dfe768dd93788 Mon Sep 17 00:00:00 2001 From: michael <78113554+chowhsu@users.noreply.github.com> Date: Tue, 20 Oct 2026 06:32:44 +0000 Subject: [PATCH 60/89] tools/eval: no image-level PYTHONPATH on the verl image; stop the run's app by id Co-Authored-By: Claude Opus 5.5 --- tools/eval/modal_verl.py | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/tools/eval/modal_verl.py b/tools/eval/modal_verl.py index 31e78fb1e..10ae0930a 100644 --- a/tools/eval/modal_verl.py +++ b/tools/eval/modal_verl.py @@ -400,7 +400,7 @@ def build_app(run: str, cfg: dict): sampler_image = (modal.Image.from_registry(MILES_NEXT_IMAGE.removeprefix("docker:")) .env({"HOME": "/root", "PYTHONUNBUFFERED": "1"}) .add_local_dir(str(REPO), WORKDIR, copy=False, ignore=ignore)) - verl_image = (modal_image(modal).env({"PYTHONUNBUFFERED": "1", "PYTHONPATH": WORKDIR}) + verl_image = (modal_image(modal).env({"PYTHONUNBUFFERED": "1"}) .add_local_dir(str(REPO), WORKDIR, copy=False, ignore=ignore)) cpu_image = modal.Image.debian_slim(python_version="3.12") app = modal.App(APP_PREFIX + run) @@ -500,6 +500,8 @@ def maybe_reclaim() -> None: # one detached app for the whole run: a killed driver never cancels a spawned call with modal.enable_output(), app.run(detach=True): + app_id = app.app_id # an ephemeral app is stopped by id (`modal app stop ` misses it) + (run_dir / "app_id.txt").write_text(f"{app_id}\n") calls: dict[str, Any] = {} if cfg["kind"] == "island": # c (control) beside a; b after a, from a's store for phase in [p for p in cfg["phases"] if p != "b"]: @@ -522,7 +524,7 @@ def maybe_reclaim() -> None: log)["collect"] wall = time.time() - began pulled = _modal_cli("volume", "get", "--force", RESULTS_VOLUME, run_id, str(raw), timeout=1800) - stopped = _modal_cli("app", "stop", "-y", APP_PREFIX + run_id, timeout=120) # only our own app + stopped = _modal_cli("app", "stop", "-y", app_id, timeout=120) # only our own app summary = {"calls": results, "reclaim": reclaim, "wall_s": round(wall, 1), "pull_rc": pulled.returncode, "stop_rc": stopped.returncode} (run_dir / "driver.json").write_text(json.dumps(summary, indent=1, default=str)) From 934e5055849c416978f2738fb59e83fddd7e059c Mon Sep 17 00:00:00 2001 From: michael <78113554+chowhsu@users.noreply.github.com> Date: Tue, 20 Oct 2026 06:33:40 +0000 Subject: [PATCH 61/89] tools/eval: check the preflight before claiming the run dir Co-Authored-By: Claude Opus 5.5 --- tools/eval/__main__.py | 9 ++++----- 1 file changed, 4 insertions(+), 5 deletions(-) diff --git a/tools/eval/__main__.py b/tools/eval/__main__.py index aaa1a3191..1d047988e 100644 --- a/tools/eval/__main__.py +++ b/tools/eval/__main__.py @@ -49,8 +49,10 @@ def cmd_list(_args) -> int: def cmd_run(args) -> int: cell = cells.get(args.capability, args.backend) run_dir = Path(args.out) / args.run_id - if (run_dir / "cell.json").exists() and not args.dry_run: - raise SystemExit(f"{run_dir} already used") + if not args.dry_run: + if (run_dir / "cell.json").exists(): + raise SystemExit(f"{run_dir} already used") + _preflight(args.run_id) run_dir.mkdir(parents=True, exist_ok=True) sha = subprocess.run(["git", "-C", str(Path(__file__).resolve().parents[2]), "rev-parse", "HEAD"], capture_output=True, text=True).stdout.strip() @@ -62,15 +64,12 @@ def cmd_run(args) -> int: if args.dry_run: print(json.dumps(cell.config, indent=1)) return 0 - _preflight(args.run_id) from . import modal_verl summary = modal_verl.run(cell, args.run_id, run_dir) else: from . import miles - if not args.dry_run: - _preflight(args.run_id) summary = miles.run(cell, args.run_id, run_dir, dry_run=args.dry_run) if args.dry_run: print((run_dir / "dry-run.log").read_text()) From 8dd269f4c9cf5506cb36cd3036ded74bf95f6edc Mon Sep 17 00:00:00 2001 From: michael <78113554+chowhsu@users.noreply.github.com> Date: Tue, 20 Oct 2026 06:35:03 +0000 Subject: [PATCH 62/89] tools/eval: ship the verl harness functions by value (tools is not in the images) Co-Authored-By: Claude Opus 5.5 --- tools/eval/modal_verl.py | 8 ++++++-- 1 file changed, 6 insertions(+), 2 deletions(-) diff --git a/tools/eval/modal_verl.py b/tools/eval/modal_verl.py index 10ae0930a..a0839afff 100644 --- a/tools/eval/modal_verl.py +++ b/tools/eval/modal_verl.py @@ -392,6 +392,10 @@ def build_app(run: str, cfg: dict): import modal sys.path.insert(0, str(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__]) from yeto.rl.adapters.miles.pins import MILES_NEXT_IMAGE from yeto.rl.adapters.verl.image import modal_image @@ -414,9 +418,9 @@ def build_app(run: str, cfg: dict): volumes=vols, name="sampler", serialized=True)(sampler), "trainer": app.function(image=verl_image, gpu=gpu, cpu=16.0, memory=65536, timeout=timeout, retries=0, volumes=vols, secrets=[hmac], name="trainer", serialized=True)(trainer), - # max_inputs=1: every phase in a fresh container (a warm one keeps ~/yeto-output: C4 run c) + # single_use_containers: every phase in a fresh container (a warm one keeps ~/yeto-output: C4 run c) "island": app.function(image=verl_image, gpu=gpu, cpu=8.0, memory=49152, timeout=timeout, retries=0, - volumes=vols, max_inputs=1, name="island", serialized=True)(island), + volumes=vols, single_use_containers=True, name="island", serialized=True)(island), "collect": app.function(image=cpu_image, cpu=1.0, memory=2048, timeout=900, volumes=vols, name="collect", serialized=True)(collect), } From 2d017361886a8944b1623525ed18b648614d81a6 Mon Sep 17 00:00:00 2001 From: michael <78113554+chowhsu@users.noreply.github.com> Date: Tue, 20 Oct 2026 06:46:53 +0000 Subject: [PATCH 63/89] verl agentic entry: register yeto_grpo lazily so the Ray setup hook imports no CUDA code The worker_process_setup_hook runs before Ray sets an actor's CUDA_VISIBLE_DEVICES; importing verl there initialised CUDA early and put both FSDP ranks on GPU 0 (NCCL Duplicate GPU detected, T-N109 run a). The hook now installs an import hook that loads the estimator right after verl's core_algos executes. Co-Authored-By: Claude Opus 5.5 --- tests/test_verl_agentic_main.py | 31 ++++++++++++++ yeto/rl/adapters/verl/agentic_main.py | 59 +++++++++++++++++++++++---- 2 files changed, 83 insertions(+), 7 deletions(-) create mode 100644 tests/test_verl_agentic_main.py diff --git a/tests/test_verl_agentic_main.py b/tests/test_verl_agentic_main.py new file mode 100644 index 000000000..e988ac178 --- /dev/null +++ b/tests/test_verl_agentic_main.py @@ -0,0 +1,31 @@ +"""S25 N10.9: the verl agentic entry's Ray setup hook imports nothing CUDA-related (CPU).""" + +from __future__ import annotations + +import subprocess +import sys +import textwrap + + +def test_setup_hook_defers_the_estimator_until_core_algos_loads(tmp_path): + # Stand-ins for verl's core_algos and yeto's estimator module, in a fresh interpreter. + (tmp_path / "fakeverl").mkdir() + (tmp_path / "fakeverl" / "__init__.py").write_text("") + (tmp_path / "fakeverl" / "core_algos.py").write_text("REG = []\n") + (tmp_path / "fakeest.py").write_text("import fakeverl.core_algos as c\nc.REG.append('yeto_grpo')\n") + code = textwrap.dedent(""" + import sys + from yeto.rl.adapters.verl import agentic_main as m + m.ESTIMATOR_HOST, m.ESTIMATOR_MODULE = "fakeverl.core_algos", "fakeest" + m.setup_worker() + assert "fakeverl.core_algos" not in sys.modules and "torch" not in sys.modules + assert not any(k == "verl" or k.startswith("verl.") for k in sys.modules) + import fakeverl.core_algos as c + assert c.REG == ["yeto_grpo"], c.REG + m.setup_worker() # idempotent once loaded + assert c.REG == ["yeto_grpo"] + print("ok") + """) + out = subprocess.run([sys.executable, "-c", code], cwd=tmp_path, capture_output=True, text=True, + env={"PYTHONPATH": f"{tmp_path}:{__import__('os').getcwd()}"}) + assert out.returncode == 0 and out.stdout.strip() == "ok", out.stderr diff --git a/yeto/rl/adapters/verl/agentic_main.py b/yeto/rl/adapters/verl/agentic_main.py index 1fdaba314..ea7d69e31 100644 --- a/yeto/rl/adapters/verl/agentic_main.py +++ b/yeto/rl/adapters/verl/agentic_main.py @@ -3,9 +3,13 @@ ``python -m yeto.rl.adapters.verl.agentic_main `` is verl's ``main_ppo.main`` unchanged, with Ray started here so that every Ray worker: -* imports the yeto pieces verl selects by name -- the ``yeto_grpo`` advantage - estimator (trainer process) and the ``yeto_harbor`` agent loop (agent-loop - workers) -- through Ray's ``worker_process_setup_hook``; +* registers the ``yeto_grpo`` advantage estimator right after verl's + ``core_algos`` is imported (Ray ``worker_process_setup_hook`` installs an + import hook). The hook itself must not import verl or torch: it runs + before Ray sets the actor's ``CUDA_VISIBLE_DEVICES``, and an early CUDA + init puts every rank on GPU 0 (NCCL "Duplicate GPU detected"; S24 T-GLM, + S25 T-N109 run a). The ``yeto_harbor`` agent loop is imported by verl + through ``agent_loop_config_path``; * carries the agentic environment: ``YETO_*`` (recorder, records proxy, boat caps), ``HARBOR_*`` (task directory, provider) and ``BOAT_API_KEY`` (the sandbox account, read by ``boat_client``), on top of verl's own runtime env. @@ -18,18 +22,59 @@ from __future__ import annotations +import importlib +import importlib.abc +import importlib.util import os import sys +from typing import Any FORWARD_ENV_PREFIXES = ("YETO_", "HARBOR_", "HF_", "PYTHONPATH") FORWARD_ENV_NAMES = ("BOAT_API_KEY",) SETUP_HOOK = "yeto.rl.adapters.verl.agentic_main.setup_worker" +ESTIMATOR_HOST = "verl.trainer.ppo.core_algos" +ESTIMATOR_MODULE = "yeto.rl.adapters.verl.advantage" + + +class _AfterImport(importlib.abc.MetaPathFinder): + """Import ``then`` right after ``target`` has been executed (no import of either before that).""" + + def __init__(self, target: str, then: str) -> None: + self.target, self.then = target, then + + def find_spec(self, name: str, path: Any, target: Any = None) -> Any: + if name != self.target: + return None + sys.meta_path.remove(self) # find the real spec without this finder + try: + spec = importlib.util.find_spec(name) + finally: + sys.meta_path.insert(0, self) + if spec is None or spec.loader is None: + return spec + loader, then = spec.loader, self.then + + class _Loader(importlib.abc.Loader): + def create_module(self, spec: Any) -> Any: + return loader.create_module(spec) + + def exec_module(self, module: Any) -> None: + loader.exec_module(module) + importlib.import_module(then) + + spec.loader = _Loader() + return spec + + def setup_worker() -> None: - """Ray ``worker_process_setup_hook``: register the yeto estimator and agent loop.""" - import yeto.rl.adapters.verl.advantage # noqa: F401 (registers yeto_grpo) - import yeto.rl.adapters.verl.harbor_agent_loop # noqa: F401 (registers yeto_harbor) + """Ray ``worker_process_setup_hook``: register yeto_grpo once verl's core_algos loads (module doc).""" + if ESTIMATOR_HOST in sys.modules: + importlib.import_module(ESTIMATOR_MODULE) + return + if not any(isinstance(f, _AfterImport) for f in sys.meta_path): + sys.meta_path.insert(0, _AfterImport(ESTIMATOR_HOST, ESTIMATOR_MODULE)) def runtime_env(environ: dict[str, str] | None = None) -> dict: @@ -51,7 +96,7 @@ def main(argv: list[str] | None = None) -> None: import ray import verl.trainer.main_ppo as main_ppo - setup_worker() # the driver validates algorithm.adv_estimator against the registry too + importlib.import_module(ESTIMATOR_MODULE) # the driver validates algorithm.adv_estimator too if not ray.is_initialized(): ray.init(runtime_env=runtime_env(), namespace="yeto-verl-agentic") if argv is not None: From 4895f1acfbf24c83428e46d4120de08416277e9d Mon Sep 17 00:00:00 2001 From: michael <78113554+chowhsu@users.noreply.github.com> Date: Tue, 20 Oct 2026 06:54:05 +0000 Subject: [PATCH 64/89] tools/eval: Miles launches skip the launcher thread preflight (the scripts check first); count each over-age drop once Co-Authored-By: Claude Opus 5.5 --- tools/eval/judges.py | 15 ++++++++------- tools/eval/miles.py | 5 ++++- tools/eval/miles_head.sh | 4 ++-- 3 files changed, 14 insertions(+), 10 deletions(-) diff --git a/tools/eval/judges.py b/tools/eval/judges.py index f544cac03..b42baf312 100644 --- a/tools/eval/judges.py +++ b/tools/eval/judges.py @@ -215,12 +215,12 @@ def judge_e4(ev: Evidence, v: Verdict, cfg: dict) -> None: len(rt) == cfg["rounds"] and finite(ratio) and finite( (e.get("mismatch") or {}).get("abs_diff", (e.get("mismatch") or {}).get("tis_abs")) for e in rt), {"ratio_or_tis": ratio}) - dropped = (numeric_sum(ev.of(t, "rl_rollout_carry_over"), "over_age_discarded_groups") - + numeric_sum(ev.of(t, "rl_round_labels") + rt + ev.of(t, "rl_local_round"), "stale_groups") - + numeric_sum(ev.of(t, "offpolicy_gate"), "group_dropped")) - if ev.backend == "verl": - files = find_all(ev.files.get(t, ev.run_dir), "verl-fa-rollout-*.jsonl") - dropped += numeric_sum([r.get("carry_over") or {} for p in files for r in jsonl(p)], "over_age_discarded_groups") + # two distinct drop points, each counted once: the round collector's over-age discard + # (rl_rollout_carry_over, both backends) and the Miles buffer reader passing over stale + # groups (sample_buffer/stale_groups in the round labels) + carried = numeric_sum(ev.of(t, "rl_rollout_carry_over"), "over_age_discarded_groups") + passed_over = numeric_sum(ev.of(t, "rl_round_labels"), "stale_groups") + dropped = carried + passed_over v.check("over_age_drop", "at least one group older than max_policy_age was dropped, not trained", dropped >= 1, dropped) over = _over_age_trained(ev, t, age) @@ -230,7 +230,8 @@ def judge_e4(ev: Evidence, v: Verdict, cfg: dict) -> None: len({e.get("policy_version") for e in loads}) >= 2, sorted({e.get("policy_version") for e in loads})) _rc_ok(ev, v, ignore=("sampler_launch", "sampler-0", "sampler-1")) _gpu(ev, v, cfg) - v.metrics.update(over_age_dropped_groups=dropped, ratio_p50=ratio, **_round_metrics(ev, t)) + v.metrics.update(over_age_discarded_groups=carried, stale_passed_over_groups=passed_over, ratio_p50=ratio, + **_round_metrics(ev, t)) def _over_age_trained(ev: Evidence, trainer: str, age: int) -> int: diff --git a/tools/eval/miles.py b/tools/eval/miles.py index b15d29f98..3f9538e69 100644 --- a/tools/eval/miles.py +++ b/tools/eval/miles.py @@ -30,7 +30,10 @@ def _gpu(cfg: dict, n: int) -> str: def _common(cfg: dict, run_id: str, prefix: str) -> str: - return (f"launch --training-mode rl --rl-engine ports --modal-gpu-exact --modal-retries 0 --modal-cpu 8 " + # --preflight-threads off: Modal islands put no SkyPilot threads on this host; the scripts + # check the host thread count (< 3000) before the launch instead + return (f"launch --training-mode rl --rl-engine ports --preflight-threads off " + f"--modal-gpu-exact --modal-retries 0 --modal-cpu 8 " f"--modal-memory 64 --modal-tape-volume {TAPE_VOLUME} --controller local --cluster-prefix {prefix} " f"--tuning lora --lora-r {cfg['lora_r']} --rl-observe-timeline {MODEL} {DATA} " f"--rl-checkpoint-store modal-volume://{STORE_VOLUME}/{run_id} " diff --git a/tools/eval/miles_head.sh b/tools/eval/miles_head.sh index 2c27480d5..dac653735 100755 --- a/tools/eval/miles_head.sh +++ b/tools/eval/miles_head.sh @@ -30,8 +30,8 @@ LORA="--tuning lora --lora-r 16" MODALX="--modal-gpu-exact --modal-retries 0 --modal-cpu 8 --modal-memory 64 --modal-tape-volume $TAPEVOL" STORE="--rl-checkpoint-store modal-volume://$STOREVOL/$P" COMMON="--training-mode rl --rl-engine ports $MODALX $MODEL $DATA $LORA $STORE --rl-algorithm-spec examples/rl_algorithms/decoupled-ppo-c2.json --rl-max-policy-age 2 --rl-observe-timeline --reward-function gsm8k_reward:score --fragments 1 --pipeline 1 --rollout-batch-size 32 --n-samples-per-prompt 4 --rollout-max-response-len 1024 --seq-len 2048 --inner-lr 1e-5 --seed 17 --apply-chat-template-kwargs '{\"enable_thinking\": false}' --trust-remote-code" -TARGS="launch --controller head --syncer-region $HEAD_REGION --preflight-threads wait --rl-stall-timeout 2400 --on-demand --no-island-relaunch --modal-timeout-s $HARD --cluster-prefix $P --gpu modal:1x$GPU,modal:1x$GPU --rl-island-scheduling elastic --rl-sample-role trainer --rl-sample-wait-s 900 $COMMON --total-steps $STEPS" -SARGS="launch --controller local --cluster-prefix $SP --gpu modal:1x$GPU --rl-single-island-no-sync --rl-sample-role sampler --rl-sampler-max-backlog-rounds 4 $COMMON --total-steps $SROUNDS" +TARGS="launch --controller head --syncer-region $HEAD_REGION --preflight-threads off --rl-stall-timeout 2400 --on-demand --no-island-relaunch --modal-timeout-s $HARD --cluster-prefix $P --gpu modal:1x$GPU,modal:1x$GPU --rl-island-scheduling elastic --rl-sample-role trainer --rl-sample-wait-s 900 $COMMON --total-steps $STEPS" +SARGS="launch --preflight-threads off --controller local --cluster-prefix $SP --gpu modal:1x$GPU --rl-single-island-no-sync --rl-sample-role sampler --rl-sampler-max-backlog-rounds 4 $COMMON --total-steps $SROUNDS" git -C $REPO rev-parse $SHA > $R/yeto_sha.txt; echo "$TARGS" > $R/args.txt; echo "$SARGS" > $R/sargs.txt echo "$APP $SAPP $HEAD_CL" > $R/resources.txt eval "$(/usr/bin/python3 - <<'PYC' From 25b418af138ca600d3386b5703dc2ec09a1e72c5 Mon Sep 17 00:00:00 2001 From: michael <78113554+chowhsu@users.noreply.github.com> Date: Tue, 20 Oct 2026 06:56:43 +0000 Subject: [PATCH 65/89] verl agentic path: report each sample's weight step range (min/max_global_steps) verl's V1 trainer reads min_global_steps / max_global_steps from every sample's extra fields (T-N109 run b failed on None). The records proxy exposes the trajectory being served in a context variable; VerlRolloutEngine keeps the global_steps of each reply per trajectory and HarborAgentLoop reports their range. Co-Authored-By: Claude Opus 5.5 --- tests/test_verl_rollout_engine.py | 7 +++- yeto/rl/adapters/verl/harbor_agent_loop.py | 8 ++-- yeto/rl/adapters/verl/rollout_engine.py | 48 +++++++++++++++++++--- yeto/rl/harness/gateway/proxy.py | 12 +++++- 4 files changed, 65 insertions(+), 10 deletions(-) diff --git a/tests/test_verl_rollout_engine.py b/tests/test_verl_rollout_engine.py index c2e97bbf4..89a60a94d 100644 --- a/tests/test_verl_rollout_engine.py +++ b/tests/test_verl_rollout_engine.py @@ -54,7 +54,8 @@ async def generate(self, request_id: str, *, prompt_ids: list[int], sampling_par if self.aborted: return SimpleNamespace(token_ids=[], log_probs=None, stop_reason="aborted") ids = self.model(tuple(prompt_ids), 0) - return SimpleNamespace(token_ids=ids, log_probs=[-0.25] * len(ids), stop_reason="completed") + return SimpleNamespace(token_ids=ids, log_probs=[-0.25] * len(ids), stop_reason="completed", + extra_fields={"global_steps": 3}) def _verl_proxy(template) -> tuple[RecordsProxy, Any]: @@ -101,6 +102,9 @@ async def go(): assert trained and all(lp == -0.25 for lp in trained) assert out.extra_fields["provenance"]["engine"] == "verl-rollout" assert out.extra_fields["length_capped_turns"] == 0 + # verl V1 needs the weights' step range per sample + assert out.extra_fields["min_global_steps"] == out.extra_fields["max_global_steps"] == 3 + assert engine.pop_steps(endpoint.split("/t/")[1].split("/")[0]) == [] # taken by the loop def test_an_aborted_verl_request_aborts_the_sample(verl_loop, qwen): @@ -119,6 +123,7 @@ async def go(): out = asyncio.run(go()) assert out.extra_fields[OUTCOME_KEY] == "aborted" and sum(out.response_mask) == 0 + assert out.extra_fields["min_global_steps"] == out.extra_fields["max_global_steps"] == 0 # no reply seen def test_unbound_engine_and_short_secret_fail_closed(qwen): diff --git a/yeto/rl/adapters/verl/harbor_agent_loop.py b/yeto/rl/adapters/verl/harbor_agent_loop.py index 6f966c635..07ce7898f 100644 --- a/yeto/rl/adapters/verl/harbor_agent_loop.py +++ b/yeto/rl/adapters/verl/harbor_agent_loop.py @@ -86,7 +86,7 @@ from yeto.rl.harness.harbor import HarborHarness, HarborSettings, build_trial_config from yeto.rl.harness.harbor.boat_client import check_in_flight_cap -from .rollout_engine import bind_rollout_client +from .rollout_engine import bind_rollout_client, weight_steps from .run_config_rules import AGENTIC_ADV_ESTIMATOR as YETO_GRPO logger = logging.getLogger(__name__) @@ -220,10 +220,10 @@ async def run(self, sampling_params: dict[str, Any], **kwargs: Any) -> AgentLoop trajectory = _abort(trajectory, f"token_alignment: {exc}") if segments is not None and len(segments.prompt_ids) > self.prompt_length: trajectory, segments = _abort(trajectory, "prompt_exceeds_prompt_length"), None - return self._output(trajectory, segments, turns) + return self._output(trajectory, segments, turns, steps=weight_steps(recorder, session)) def _output(self, trajectory: AgentTrajectory, segments: TokenSegments | None, - turns: Sequence[TokenTurn]) -> AgentLoopOutput: + turns: Sequence[TokenTurn], *, steps: tuple[int, int] | None = None) -> AgentLoopOutput: if segments is None: segments = self._placeholder() metadata = trajectory.to_metadata() @@ -244,6 +244,8 @@ def _output(self, trajectory: AgentTrajectory, segments: TokenSegments | None, "length_capped_turns": sum(t.finish_reason == FINISH_LENGTH for t in turns), "agent_metrics": metadata["agent_metrics"], "provenance": turn_provenance(turns), + # verl V1 reads both for its staleness metrics; recorders without verl's rollout leave them out + **({"min_global_steps": steps[0], "max_global_steps": steps[1]} if steps else {}), }, ) diff --git a/yeto/rl/adapters/verl/rollout_engine.py b/yeto/rl/adapters/verl/rollout_engine.py index f78590b3a..68bdfbf1b 100644 --- a/yeto/rl/adapters/verl/rollout_engine.py +++ b/yeto/rl/adapters/verl/rollout_engine.py @@ -27,9 +27,12 @@ Policy version: verl's synchronous trainer updates the rollout weights only between steps, never while a batch is generated, so a turn cannot straddle a -reload. The rollout reply carries verl's ``global_steps`` only in the fully -async mode; :meth:`VerlRolloutEngine.served_version` therefore returns None -(unknown) and the trainer's step log is the version record. +reload. Each rollout reply carries the weights' ``global_steps``; the engine +keeps them per trajectory (the records proxy's ``CURRENT_TRAJECTORY``) and +``HarborAgentLoop`` reports their min / max as the sample's +``min_global_steps`` / ``max_global_steps`` (verl's V1 trainer requires both). +:meth:`VerlRolloutEngine.served_version` returns None: the version is known +only after a reply, not before it. Finish reason: verl reports ``completed`` for both an end token and the token cap. A reply that ends without a stop token and fills ``max_tokens`` is @@ -75,6 +78,8 @@ def __init__(self, *, max_model_len: int, identity: EngineIdentity) -> None: self._identity = identity self._client: Any = None self._loop: asyncio.AbstractEventLoop | None = None + self._steps: dict[str, list[int]] = {} + self.last_global_steps: int | None = None # ------------------------------------------------------------ binding def bind(self, client: Any, loop: asyncio.AbstractEventLoop) -> None: @@ -87,6 +92,10 @@ def bind(self, client: Any, loop: asyncio.AbstractEventLoop) -> None: def bound(self) -> bool: return self._client is not None + def pop_steps(self, trajectory: str) -> list[int]: + """The ``global_steps`` of every reply sampled for ``trajectory`` (and forget them).""" + return self._steps.pop(trajectory, []) + # ------------------------------------------------------------ engine port @property def identity(self) -> EngineIdentity: @@ -104,18 +113,22 @@ def generate_tokens(self, prompts: Sequence[Sequence[int]], sampling: SamplingPa """Called from a worker thread (the records backend uses ``asyncio.to_thread``).""" if self._client is None or self._loop is None: raise SamplerEngineError("VerlRolloutEngine is not bound to verl's rollout client") + from yeto.rl.harness.gateway.proxy import CURRENT_TRAJECTORY + + trajectory = CURRENT_TRAJECTORY.get() # copied into this thread by asyncio.to_thread stops = {int(t) for t in stop_token_ids} out: list[list[Completion]] = [] for prompt in prompts: group = [] for _ in range(int(sampling.n)): future = asyncio.run_coroutine_threadsafe( - self._generate(list(prompt), sampling, sorted(stops)), self._loop) + self._generate(list(prompt), sampling, sorted(stops), trajectory), self._loop) group.append(future.result(timeout=CALL_TIMEOUT_S)) out.append(group) return out - async def _generate(self, prompt: list[int], sampling: SamplingParams, stops: list[int]) -> Completion: + async def _generate(self, prompt: list[int], sampling: SamplingParams, stops: list[int], + trajectory: str | None = None) -> Completion: params: dict[str, Any] = { "temperature": float(sampling.temperature), "top_p": float(sampling.top_p), "top_k": int(sampling.top_k), "repetition_penalty": float(sampling.repetition_penalty), @@ -127,6 +140,11 @@ async def _generate(self, prompt: list[int], sampling: SamplingParams, stops: li sampling_params=params) if getattr(output, "stop_reason", None) == "aborted": raise SamplerEngineError("verl rollout aborted the request") + steps = (getattr(output, "extra_fields", None) or {}).get("global_steps") + if steps is not None: + self.last_global_steps = int(steps) + if trajectory is not None: + self._steps.setdefault(trajectory, []).append(int(steps)) tokens = tuple(int(t) for t in output.token_ids) logprobs = output.log_probs if logprobs is None or len(logprobs) != len(tokens): @@ -166,6 +184,26 @@ def finish_reason(tokens: Sequence[int], stops: set[int] | Sequence[int], max_to return FINISH_LENGTH if len(tokens) >= max_tokens else FINISH_STOP +def _engine_of(recorder: Any) -> Any: + engine = getattr(getattr(getattr(recorder, "proxy", None), "backend", None), "engine", None) + return engine if isinstance(engine, VerlRolloutEngine) else None + + +def weight_steps(recorder: Any, trajectory: str) -> tuple[int, int] | None: + """``(min, max)`` verl ``global_steps`` of the trajectory's replies; None without a verl engine. + + A trajectory with no reply (aborted before its first call) gets the + latest step the engine saw: in the synchronous trainer that is the batch's. + Before the engine saw any reply it gets ``(0, 0)``: verl's V1 trainer + needs an integer, and such a row is an ABORTED placeholder (no loss tokens). + """ + engine = _engine_of(recorder) + if engine is None: + return None + steps = engine.pop_steps(trajectory) or ([engine.last_global_steps] if engine.last_global_steps is not None else []) + return (min(steps), max(steps)) if steps else (0, 0) + + def bind_rollout_client(recorder: Any, client: Any) -> None: """Bind the recorder's :class:`VerlRolloutEngine` (if it has one) to verl's client on this loop.""" proxy = getattr(recorder, "proxy", None) diff --git a/yeto/rl/harness/gateway/proxy.py b/yeto/rl/harness/gateway/proxy.py index b5c0e4851..ac6e7a342 100644 --- a/yeto/rl/harness/gateway/proxy.py +++ b/yeto/rl/harness/gateway/proxy.py @@ -19,6 +19,7 @@ import json import secrets +from contextvars import ContextVar from typing import Any from yeto.rl.compaction import CompactionConfig @@ -29,6 +30,11 @@ from .records import RecordingBackend from .translate import OMITTED_TOOL_TYPES, UnsupportedShape +# The trajectory whose request is being served. Engine calls run in a thread +# started by ``asyncio.to_thread``, which copies this context: an engine that +# owns no version of its own (verl's rollout) can tag what it saw per trajectory. +CURRENT_TRAJECTORY: ContextVar[str | None] = ContextVar("records_proxy_trajectory", default=None) + Message = dict[str, Any] @@ -68,7 +74,11 @@ def open(self) -> str: async def handle(self, entry: str, body: dict[str, Any], *, trajectory_id: str) -> dict[str, Any]: self.backend.declare_tools(trajectory_id, chat_tools(entry, body)) - return await self.gateway.handle(entry, body, trajectory_id=trajectory_id) + token = CURRENT_TRAJECTORY.set(trajectory_id) + try: + return await self.gateway.handle(entry, body, trajectory_id=trajectory_id) + finally: + CURRENT_TRAJECTORY.reset(token) def records(self, trajectory_id: str) -> dict[str, Any]: out = self.backend.records(trajectory_id) From c8efa159710ccceba92d09f626db6add39827d29 Mon Sep 17 00:00:00 2001 From: michael <78113554+chowhsu@users.noreply.github.com> Date: Tue, 20 Oct 2026 07:02:20 +0000 Subject: [PATCH 66/89] Miles records generate: take the router address from Miles' arguments Co-Authored-By: Claude Opus 5.5 --- yeto/rl/adapters/miles/records_generate.py | 8 ++++++-- 1 file changed, 6 insertions(+), 2 deletions(-) diff --git a/yeto/rl/adapters/miles/records_generate.py b/yeto/rl/adapters/miles/records_generate.py index 561d3a30b..9b157f989 100644 --- a/yeto/rl/adapters/miles/records_generate.py +++ b/yeto/rl/adapters/miles/records_generate.py @@ -27,7 +27,8 @@ YETO_RECORDS_TEMPLATE, YETO_RECORDS_TOKENIZER, YETO_RECORDS_MODEL, YETO_RECORDS_SAMPLING, YETO_RECORDS_MAX_MODEL_LEN, YETO_RECORDS_CARD (same meaning as ``adapters.verl.rollout_engine``) - YETO_SGLANG_ROUTER_URL ``http://:`` of Miles' router (the launcher sets it) + YETO_SGLANG_ROUTER_URL ``http://:`` of Miles' router (default: from + ``--sglang-router-ip`` / ``--sglang-router-port``) Policy version: Miles' synchronous loop updates the engines between rollouts, never during one; the proxy records version None and the trainer's @@ -212,7 +213,10 @@ async def generate(input: Any) -> Any: """Miles ``--custom-generate-function-path`` (module doc).""" from miles.rollout.base_types import GenerateFnOutput - sample = await run_sample(input.sample, input.sampling_params, input.args, recorder=_recorders.get()) + args = input.args + # Miles' router address is an argument, not an env variable; the proxy is built from env once per loop. + os.environ.setdefault(ROUTER_URL_ENV, f"http://{args.sglang_router_ip}:{args.sglang_router_port}") + sample = await run_sample(input.sample, input.sampling_params, args, recorder=_recorders.get()) return GenerateFnOutput(samples=sample) From 8731a0f191b9cdffc43a3848bc857a8373ffce6c Mon Sep 17 00:00:00 2001 From: michael <78113554+chowhsu@users.noreply.github.com> Date: Tue, 20 Oct 2026 07:05:02 +0000 Subject: [PATCH 67/89] codepack: reuse a tree another launch unpacked first instead of failing Two launches from one checkout (a trainer and a sampler launch on one store, S25 T-EVAL E4) unpack the same digest at the same time; the later rename hit ENOTEMPTY and the trainer launch exited 1. Keep the existing tree when it verifies, drop the staging copy. Co-Authored-By: Claude Opus 5.5 --- tests/test_codepack.py | 21 +++++++++++++++++++++ yeto/codepack.py | 10 +++++++++- 2 files changed, 30 insertions(+), 1 deletion(-) diff --git a/tests/test_codepack.py b/tests/test_codepack.py index 3d09fb3fe..24a9b35cd 100644 --- a/tests/test_codepack.py +++ b/tests/test_codepack.py @@ -77,6 +77,27 @@ def test_verify_tree_rejects_changed_tree(repo: Path) -> None: codepack.verify_tree(tree, pack.digest) +def test_unpack_keeps_a_tree_another_launch_placed_first(repo: Path, monkeypatch) -> None: + # S25 T-EVAL: two launches from one checkout unpack the same digest at once; the + # later rename finds the earlier tree in place (ENOTEMPTY) and must reuse it + pack, data = codepack.build_pack(repo) + cache = repo.parent / "cache" + first = codepack.unpack(data, pack.digest, cache) + real_is_dir = Path.is_dir + calls = {"n": 0} + + def racing_is_dir(self: Path) -> bool: # the first existence check misses the tree + if self == first and calls["n"] == 0: + calls["n"] += 1 + return False + return real_is_dir(self) + + monkeypatch.setattr(Path, "is_dir", racing_is_dir) + assert codepack.unpack(data, pack.digest, cache) == first + assert codepack.verify_tree(first, pack.digest) == pack.digest + assert [p.name for p in cache.iterdir()] == [first.name] # no staging dir left behind + + def test_unpack_rejects_wrong_bytes(repo: Path) -> None: pack, data = codepack.build_pack(repo) with pytest.raises(CodePackError, match="does not have digest"): diff --git a/yeto/codepack.py b/yeto/codepack.py index 73fbc4ac5..0e09abb25 100644 --- a/yeto/codepack.py +++ b/yeto/codepack.py @@ -263,7 +263,15 @@ def unpack(data: bytes, digest: str, cache_root: str | os.PathLike[str]) -> Path _safe_member(member) tar.extractall(staging, members=members, filter="tar") verify_tree(staging, digest) - staging.rename(dest) + try: + staging.rename(dest) + except OSError: + # another process (a second launch from the same checkout) unpacked the + # same digest first: keep its tree when it is intact + if not dest.is_dir(): + raise + verify_tree(dest, digest) + shutil.rmtree(staging, ignore_errors=True) except BaseException: shutil.rmtree(staging, ignore_errors=True) raise From 7e031aa69301664b306a4f5f21d20227502cd801 Mon Sep 17 00:00:00 2001 From: michael <78113554+chowhsu@users.noreply.github.com> Date: Tue, 20 Oct 2026 07:26:46 +0000 Subject: [PATCH 68/89] tools/eval: E4 counts the Miles buffer reader's passed-over stale groups from consumed minus trained; over-age check reads rl_group_policies per round Co-Authored-By: Claude Opus 5.5 --- tools/eval/judges.py | 25 ++++++++++++++++--------- 1 file changed, 16 insertions(+), 9 deletions(-) diff --git a/tools/eval/judges.py b/tools/eval/judges.py index b42baf312..aa6650fde 100644 --- a/tools/eval/judges.py +++ b/tools/eval/judges.py @@ -216,10 +216,15 @@ def judge_e4(ev: Evidence, v: Verdict, cfg: dict) -> None: (e.get("mismatch") or {}).get("abs_diff", (e.get("mismatch") or {}).get("tis_abs")) for e in rt), {"ratio_or_tis": ratio}) # two distinct drop points, each counted once: the round collector's over-age discard - # (rl_rollout_carry_over, both backends) and the Miles buffer reader passing over stale - # groups (sample_buffer/stale_groups in the round labels) + # (rl_rollout_carry_over, both backends) and the Miles buffer reader passing over + # stale groups, which no yeto event reports: buffer positions consumed (the last + # release's buffer_selected) minus groups trained (round labels; the S25 C1 C6 count) carried = numeric_sum(ev.of(t, "rl_rollout_carry_over"), "over_age_discarded_groups") - passed_over = numeric_sum(ev.of(t, "rl_round_labels"), "stale_groups") + passed_over = 0.0 + releases = ev.of(t, "rl_policy_released") + if ev.backend == "miles" and releases and releases[-1].get("buffer_selected") is not None: + trained_groups = numeric_sum(ev.of(t, "rl_round_labels"), "rl/groups") + passed_over = float(releases[-1]["buffer_selected"]) - trained_groups dropped = carried + passed_over v.check("over_age_drop", "at least one group older than max_policy_age was dropped, not trained", dropped >= 1, dropped) @@ -242,12 +247,14 @@ def _over_age_trained(ev: Evidence, trainer: str, age: int) -> int: for versions in (r.get("versions") or {}).values(): over += min(versions) < r["rollout_id"] - age return over - for e in ev.of(trainer, "rl_group_policies"): - current = e.get("policy_version") - for g in e.get("groups") or []: - oldest = g.get("oldest_version", g.get("min_version")) if isinstance(g, dict) else None - if current is not None and oldest is not None: - over += oldest < current - age + # Miles: one rl_group_policies record per trained round, in round order: [group, token, + # [versions]]; the round trains at its base policy version (rl_local_round) + bases = [r.get("base_policy_version") for r in ev.of(trainer, "rl_local_round")] + for record, base in zip(ev.of(trainer, "rl_group_policies"), bases): + for group in record.get("groups") or []: + versions = group[2] if isinstance(group, (list, tuple)) and len(group) > 2 else [] + if base is not None and versions: + over += min(versions) < int(base) - age return over From 8b2a03193f8e2889d1ed66a342dbf70f4bd7c3fd Mon Sep 17 00:00:00 2001 From: michael <78113554+chowhsu@users.noreply.github.com> Date: Tue, 20 Oct 2026 04:15:08 +0000 Subject: [PATCH 69/89] Resume a replaced spot sampler island and allow several samplers per trainer A pure sampler island on spot capacity can be reclaimed and started again in a new container (Modal reschedules the input; sky relaunches the cluster). The new incarnation now continues where the buffer says the old one stopped: - every index entry carries a sampler_round record (sampler id, incarnation, rollout id, data cursor after the round, last group of the round); - yeto.rl.engine.sampler_resume reads the resume point: the next rollout id after the newest round with an entry, and that round's data cursor; a round already in the buffer is never written again (RoundGuard), and no prompt is drawn twice (the Miles data source is put at the cursor); - on the reclaim marker (the Modal runner writes it for sampler islands) the island starts no new round, writes no more groups, writes rl_sampler_leave with reason PREEMPTED and a per-incarnation member record in the store, then waits for the cloud to stop the container; a new incarnation writes rl_sampler_join and loads the newest policy release before its first round. --rl-island-sample-role now takes one trainer and one or more samplers. With several samplers on a shared directory (a Modal Volume, where flock does not hold across containers) each sampler writes its own shard file and the trainer merges the shards into index.jsonl (one store commit per refresh); the data ends once every named shard has ended. The trainer scheduler is refused with several samplers (the head does not read the shards). Co-Authored-By: Claude Opus 5.5 --- tests/test_rl_oss_store.py | 2 +- tests/test_rl_sampler_trainer_pair_launch.py | 2 +- tests/test_rl_spot_sampler_replace.py | 376 ++++++++++++++++++ yeto/launch/rl_ports.py | 9 +- yeto/launcher.py | 3 +- yeto/modal_runner.py | 4 +- yeto/rl/adapters/miles/buffer_rollout.py | 115 +++++- .../rl/adapters/miles/buffer_rollout_miles.py | 39 +- yeto/rl/adapters/miles/island_entry.py | 6 +- yeto/rl/adapters/miles/sampler_island.py | 92 ++++- yeto/rl/engine/buffer_payload.py | 10 +- yeto/rl/engine/run_flags.py | 8 + yeto/rl/engine/sample_buffer.py | 118 ++++-- yeto/rl/engine/sample_roles.py | 62 ++- yeto/rl/engine/sampler_resume.py | 185 +++++++++ yeto/rl/engine/trainer_lifecycle.py | 9 +- 16 files changed, 969 insertions(+), 71 deletions(-) create mode 100644 tests/test_rl_spot_sampler_replace.py create mode 100644 yeto/rl/engine/sampler_resume.py diff --git a/tests/test_rl_oss_store.py b/tests/test_rl_oss_store.py index 7aab70d30..7f457ad63 100644 --- a/tests/test_rl_oss_store.py +++ b/tests/test_rl_oss_store.py @@ -156,7 +156,7 @@ def test_each_sampler_writes_its_own_shard_and_the_trainer_merges_them(): assert len(again) == 4 # a reader that does not merge (the head) sees the merged index only assert len(FileSampleBuffer(objects, batch_size=1)) == 4 - with pytest.raises(SampleBufferError, match="only the index of an object store merges"): + with pytest.raises(SampleBufferError, match="a shard writer does not merge"): FileSampleBuffer(objects, batch_size=1, index=shard_index("x"), merge_shards=True) diff --git a/tests/test_rl_sampler_trainer_pair_launch.py b/tests/test_rl_sampler_trainer_pair_launch.py index 4debf1e72..0cbd7e69c 100644 --- a/tests/test_rl_sampler_trainer_pair_launch.py +++ b/tests/test_rl_sampler_trainer_pair_launch.py @@ -134,7 +134,7 @@ def _replace(base, old: tuple, new: tuple) -> list[str]: (("--rl-single-island-no-sync",), (), "no-sync|single-island"), (("--tuning", "full"), ("--tuning", "full", "--rl-sample-role", "trainer"), "replaces --rl-sample-role"), - (("1:trainer",), ("1:sampler",), "one sampler and one trainer"), + (("1:trainer",), ("1:sampler",), "exactly one trainer"), (("--rl-max-policy-age", "2"), ("--rl-max-policy-age", "0"), "max-policy-age"), (("--rl-checkpoint-store", STORE, "--rl-cut-every", "1"), (), "checkpoint-store"), ]) diff --git a/tests/test_rl_spot_sampler_replace.py b/tests/test_rl_spot_sampler_replace.py new file mode 100644 index 000000000..92ff9b416 --- /dev/null +++ b/tests/test_rl_spot_sampler_replace.py @@ -0,0 +1,376 @@ +"""S25 C1: a sampler island on spot capacity is reclaimed and replaced. + +CPU only (no Miles, no Ray): the resume point read from the buffer, the +resumable round writer (no round written twice, no prompt drawn twice), several +sampler islands writing their own shards of one shared-directory buffer that +the trainer merges, the launch composition of one trainer and two samplers, and +the Modal runner's reclaim marker for sampler islands. +""" + +from __future__ import annotations + +import shlex +import signal as _signal +from types import SimpleNamespace + +import pytest + +from test_rl_algorithm_provenance import _fake_sky, _launcher_args, _no_modal_listing # noqa: F401 +from tests.buffer_rollout_procs import sample +from yeto.rl.adapters.miles import buffer_rollout as br +from yeto.rl.engine.policy_release import PolicyFollower, PolicyReleaseStore +from yeto.rl.engine.resume import CheckpointStore +from yeto.rl.engine.sample_buffer import FileSampleBuffer, SampleBufferError +from yeto.rl.engine.sampler_loop import SamplerLoop +from yeto.rl.engine.sampler_resume import ( + RECLAIM_MARKER_ENV, + Member, + SamplerResumeError, + SamplerRound, + reclaim_requested, + resume_point, + write_member, +) + + +def _group(g: int, version: int = 1) -> list[dict]: + return [sample(2 * g, g, version), sample(2 * g + 1, g, version)] + + +def _cursor(offset: int) -> dict[str, int]: + return {"sample_offset": offset, "epoch_id": 0, "sample_group_index": offset, "sample_index": 2 * offset} + + +def _round(rid: int, groups: int) -> list[tuple[str, list[dict]]]: + return [(f"r{rid}-g{g}", _group(g)) for g in range(groups)] + + +def _shard_env(store, sampler=None, shards=("island0", "island1")): + env = {br.BUFFER_ROOT_ENV: str(store), br.SHARDS_ENV: ",".join(shards)} + if sampler is not None: + env[br.SAMPLER_ID_ENV] = sampler + return env + + +# ---------------------------------------------------------------- resume point +def test_resume_point_of_an_empty_buffer_starts_at_round_zero(): + resume = resume_point([], "island0") + assert (resume.next_rollout_id, resume.data_cursor, resume.resumed) == (0, None, False) + + +def test_resume_point_continues_after_the_newest_written_round(tmp_path): + buffer = FileSampleBuffer(CheckpointStore(tmp_path), batch_size=1) + writer = br.SamplerRoundWriter(buffer, "island0", incarnation="a") + assert writer.write_round(0, _round(0, 2), _cursor(2)) == 2 + assert writer.write_round(1, _round(1, 2), _cursor(4)) == 2 + other = br.SamplerRoundWriter(buffer, "island9", incarnation="z") + other.write_round(5, _round(5, 1), _cursor(9)) # another sampler: not ours + resume = resume_point(buffer.entries(), "island0") + assert (resume.next_rollout_id, resume.data_cursor) == (2, _cursor(4)) + assert (resume.written_rounds, resume.partial_rounds, resume.groups) == ((0, 1), (), 4) + assert resume.incarnations == ("a",) + + +def test_a_round_cut_short_by_the_reclaim_keeps_its_groups_and_is_not_written_again(tmp_path): + """The island leaves while it writes round 1: the written group stays, the + rest are dropped; the replacement continues at round 2 from round 1's cursor + (round 1's prompts are not drawn again).""" + store = CheckpointStore(tmp_path) + leaving = {"now": False} + old = br.SamplerRoundWriter(FileSampleBuffer(store, batch_size=1), "island0", incarnation="a", + reclaimed=lambda: leaving["now"]) + old.write_round(0, _round(0, 3), _cursor(3)) + + class Leaving(list): + def __iter__(self): # the reclaim lands after the first group of round 1 + for i, item in enumerate(super().__iter__()): + leaving["now"] = i >= 1 + yield item + + assert old.write_round(1, Leaving(_round(1, 3)), _cursor(6)) == 1 + assert old.dropped_groups == 2 + new = br.SamplerRoundWriter(FileSampleBuffer(store, batch_size=1), "island0", incarnation="b") + assert (new.resume.next_rollout_id, new.resume.data_cursor) == (2, _cursor(6)) + assert new.resume.partial_rounds == (1,) + for rid in (0, 1): + with pytest.raises(SamplerResumeError, match="already in the sample buffer"): + new.write_round(rid, _round(rid, 1), _cursor(0)) + new.write_round(2, _round(2, 3), _cursor(9)) + ids = [e.sample_id for e in FileSampleBuffer(store, batch_size=1).entries()] + assert len(ids) == len(set(ids)) == 3 + 1 + 3 + assert [i.split("-")[1] for i in ids] == ["a"] * 4 + ["b"] * 3 + + +def test_a_round_still_generating_at_the_reclaim_is_generated_again(tmp_path): + store = CheckpointStore(tmp_path) + old = br.SamplerRoundWriter(FileSampleBuffer(store, batch_size=1), "island0", incarnation="a", + reclaimed=lambda: True) + assert old.write_round(0, _round(0, 2), _cursor(2)) == 0 # left before the write + new = br.SamplerRoundWriter(FileSampleBuffer(store, batch_size=1), "island0", incarnation="b") + assert (new.resume.next_rollout_id, new.resume.data_cursor) == (0, None) + assert new.write_round(0, _round(0, 2), _cursor(2)) == 2 + + +def test_resume_point_refuses_a_buffer_with_two_cursors_for_one_round(): + def entry(cursor, last): + return SimpleNamespace(sampler_round=SamplerRound("island0", "a", 0, cursor, last).record()) + + with pytest.raises(SamplerResumeError, match="two data cursors"): + resume_point([entry(_cursor(2), False), entry(_cursor(3), True)], "island0") + with pytest.raises(SamplerResumeError, match="two last groups"): + resume_point([entry(_cursor(2), True), entry(_cursor(2), True)], "island0") + + +def test_seek_puts_the_data_source_where_its_load_would_and_only_once(tmp_path): + class Dataset: + def __init__(self): + self.shuffled = [] + + def shuffle(self, epoch): + self.shuffled.append(epoch) + + source = SimpleNamespace(sample_offset=0, epoch_id=0, sample_group_index=0, sample_index=0, + dataset=Dataset()) + store = CheckpointStore(tmp_path) + br.SamplerRoundWriter(FileSampleBuffer(store, batch_size=1), "island0", incarnation="a" + ).write_round(0, _round(0, 1), {**_cursor(5), "epoch_id": 1}) + writer = br.SamplerRoundWriter(FileSampleBuffer(store, batch_size=1), "island0", incarnation="b") + assert writer.seek(SimpleNamespace(__self__=source), shuffle=True)["sample_offset"] == 5 + assert (source.sample_offset, source.epoch_id, source.sample_index) == (5, 1, 10) + assert source.dataset.shuffled == [1] + assert writer.seek(source, shuffle=True) is None # the second round: no seek + with pytest.raises(br.SampleBufferRolloutError, match="cannot seek"): + br.seek_data_source(SimpleNamespace(sample_offset=0), _cursor(1), shuffle=False) + + +def test_a_writer_without_resume_writes_as_before(tmp_path): + """A convertible node keeps its own rollout numbering: no round check.""" + store = CheckpointStore(tmp_path) + writer = br.SamplerRoundWriter(FileSampleBuffer(store, batch_size=1), "n0", incarnation="a", + resume=False) + writer.write_round(0, _round(0, 1), None) + writer.write_round(0, _round(0, 1), None) + assert writer.seek(object(), shuffle=True) is None + assert len(FileSampleBuffer(store, batch_size=1)) == 2 + + +# ---------------------------------------------------------------- two samplers, one trainer +def test_two_samplers_write_shards_the_trainer_merges_and_a_replacement_continues(tmp_path): + a = br.SamplerRoundWriter(br.open_buffer(_shard_env(tmp_path, "island0"), sampler=True), + "island0", incarnation="a1") + b = br.SamplerRoundWriter(br.open_buffer(_shard_env(tmp_path, "island1"), sampler=True), + "island1", incarnation="b1") + assert a.buffer.shard == "island0" and not (tmp_path / "sample-buffer/index.jsonl").exists() + trainer = br.open_buffer(_shard_env(tmp_path), batch_size=2, reader=True) + reader = br.BatchReader(trainer, max_policy_age=2, wait_s=0.0, poll_s=0.0) + a.write_round(0, _round(0, 2), _cursor(2)) + b.write_round(0, _round(0, 2), _cursor(2)) + taken, _ = reader.take(0, 4, current_version=1) + assert len(taken) == 4 + # island0 is reclaimed after one more round; its replacement continues at round 2 + a.write_round(1, _round(1, 2), _cursor(4)) + a2 = br.SamplerRoundWriter(br.open_buffer(_shard_env(tmp_path, "island0"), sampler=True), + "island0", incarnation="a2") + assert (a2.resume.next_rollout_id, a2.resume.incarnations) == (2, ("a1",)) + a2.write_round(2, _round(2, 2), _cursor(6)) + b.write_round(1, _round(1, 2), _cursor(4)) + taken2, _ = reader.take(1, 6, current_version=2) + ids = [e.sample_id for e in trainer.entries()] + assert len(taken2) == 6 and len(ids) == len(set(ids)) == 10 # none repeated, none lost + assert trainer.shard_positions(trainer.read_cursor) == {"island0": 6, "island1": 4} + # the data ends only when every named shard has ended + a2.buffer.mark_data_end() + trainer.refresh() + assert not trainer.data_ended + b.buffer.mark_data_end() + trainer.refresh() + assert trainer.data_ended + + +def test_the_data_does_not_end_before_a_named_shard_exists(tmp_path): + a = br.open_buffer(_shard_env(tmp_path, "island0"), sampler=True) + a.mark_data_end() + trainer = br.open_buffer(_shard_env(tmp_path), reader=True) + assert not trainer.data_ended # island1 has not written anything yet + + +def test_a_restarted_trainer_merges_no_record_twice(tmp_path): + a = br.SamplerRoundWriter(br.open_buffer(_shard_env(tmp_path, "island0"), sampler=True), + "island0", incarnation="a") + a.write_round(0, _round(0, 3), _cursor(3)) + assert len(br.open_buffer(_shard_env(tmp_path), reader=True)) == 3 + a.write_round(1, _round(1, 1), _cursor(4)) + again = br.open_buffer(_shard_env(tmp_path), reader=True) + assert len(again) == 4 + assert [r.get("shard_pos") for r in map(_json, _lines(tmp_path))] == [0, 1, 2, 3] + + +def _lines(root): + return (root / "sample-buffer/index.jsonl").read_text().splitlines() + + +def _json(line): + import json + + return json.loads(line) + + +def test_the_trainer_merges_one_refresh_with_one_store_commit(tmp_path): + class Counting(CheckpointStore): + commits = 0 + + def commit(self): + Counting.commits += 1 + + writer = br.SamplerRoundWriter(FileSampleBuffer(CheckpointStore(tmp_path), batch_size=1, + index="shards/island0.jsonl"), + "island0", incarnation="a") + writer.write_round(0, _round(0, 5), _cursor(5)) + Counting.commits = 0 + trainer = FileSampleBuffer(Counting(tmp_path), batch_size=1, merge_shards=True, + expected_shards=["island0"]) + assert len(trainer) == 5 and Counting.commits == 1 + + +def test_expected_shards_need_a_merging_buffer(tmp_path): + with pytest.raises(SampleBufferError, match="expected_shards needs merge_shards"): + FileSampleBuffer(CheckpointStore(tmp_path), batch_size=1, expected_shards=["a"]) + + +def test_island_env_names_the_shards_and_refuses_a_sampler_outside_them(): + env = br.island_env("sampler", store_root="/s", sampler_id="island1", shards=["island0", "island1"]) + assert env[br.SHARDS_ENV] == "island0,island1" and env[br.SAMPLER_ID_ENV] == "island1" + with pytest.raises(ValueError, match="not one of the buffer shards"): + br.island_env("sampler", store_root="/s", sampler_id="island2", shards=["island0"]) + assert br.SHARDS_ENV in br.ISLAND_ENVS and br.INCARNATION_ENV in br.ISLAND_ENVS + + +# ---------------------------------------------------------------- the sampler loop leaves +def test_the_sampler_loop_stops_on_the_reclaim_and_marks_no_data_end(tmp_path): + store = CheckpointStore(tmp_path) + marker = tmp_path / "reclaim" + env = {RECLAIM_MARKER_ENV: str(marker)} + buffer = FileSampleBuffer(store, batch_size=1) + writer = br.SamplerRoundWriter(buffer, "island0", incarnation="a", + reclaimed=lambda: reclaim_requested(env)) + events = [] + + def generate(rid): + writer.write_round(rid, _round(rid, 2), _cursor(2 * (rid + 1))) + if rid == 1: + marker.write_text("{}") # the container gets SIGINT during round 1 + + loop = SamplerLoop(follower=PolicyFollower(PolicyReleaseStore(store), lambda r, d: None), + buffer=buffer, generate=generate, num_rollouts=5, max_backlog_groups=100, + emit=lambda e, **f: events.append(e), poll_s=0.0, + should_stop=lambda: reclaim_requested(env)) + assert loop.run() == 2 + assert "rl_sampler_stopped" in events and not buffer.data_ended + resume = resume_point(FileSampleBuffer(store, batch_size=1).entries(), "island0") + assert resume.next_rollout_id == 2 + loop.should_stop = lambda: False + marker.unlink() + assert loop.run(start=resume.next_rollout_id) == 3 and buffer.data_ended + + +def test_member_records_are_one_file_per_incarnation(tmp_path): + store = CheckpointStore(tmp_path) + first = write_member(store, Member("island0", "a", 1.0, {"next_rollout_id": 0})) + member = Member("island0", "b", 2.0, {"next_rollout_id": 3}, leave_reason="PREEMPTED") + second = write_member(store, member) + assert first != second and first.exists() and "PREEMPTED" in second.read_text() + + +# ---------------------------------------------------------------- launch composition +GPU3 = "modal:1xh100,modal:1xh100,modal:1xh100" +STORE = "modal-volume://yeto-c1/run-a" +GROUP = ("--rl-single-island-no-sync", "--tuning", "lora", "--lora-r", "16", + "--rl-island-sample-role", "0:sampler", "--rl-island-sample-role", "1:sampler", + "--rl-island-sample-role", "2:trainer", "--rl-max-policy-age", "2", + "--rl-algorithm-spec", "examples/rl_algorithms/decoupled-ppo-c2.json", + "--rl-checkpoint-store", STORE, "--cluster-prefix", "c1-spot") + + +def _args(extra=()): + from yeto.launcher import _prepare_rl_args + + args = _launcher_args("ports", GROUP + tuple(extra), gpu=GPU3) + args.controller = "local" + _prepare_rl_args(args) + return args + + +def _tokens(command): + return shlex.split(command.split("-m yeto.rl.adapters.miles.island_entry", 1)[1]) + + +def _flag(tokens, flag): + return tokens[tokens.index(flag) + 1] if flag in tokens else None + + +def test_one_trainer_and_two_samplers_get_their_node_and_the_shards(): + from yeto.launcher import dry_run_plan + + plan = dry_run_plan(_args()) + roles = [i["sample_role"] for i in plan["island_requests"]] + assert roles == ["sampler", "sampler", "trainer"] and plan["syncer"] is None + tokens = [_tokens(i["learner_command"]) for i in plan["island_requests"]] + assert [_flag(t, "--rl-sampler-node") for t in tokens] == ["island0", "island1", None] + assert {_flag(t, "--rl-buffer-shards") for t in tokens} == {"island0,island1"} + + +def test_the_island_resolves_its_sampler_id_and_the_shards(): + from yeto.launcher import dry_run_plan + from yeto.rl.adapters.miles import island_entry + + plan = dry_run_plan(_args()) + sampler = island_entry.parse_args([{"$LEARNER_ID": "0"}.get(t, t) for t in + _tokens(plan["island_requests"][1]["learner_command"])]) + env = {} + island_entry.apply_sample_role(sampler, SimpleNamespace(), env) + assert (env[br.SAMPLER_ID_ENV], env[br.SHARDS_ENV]) == ("island1", "island0,island1") + + +def test_only_sampler_islands_get_the_reclaim_marker(): + from yeto.gpu_spec import parse_gpu_spec + from yeto.launch.modal_islands import build_modal_island_config + from yeto.launcher import _island_task, make_miles_island_task + + args = _args() + for m, spec in enumerate(parse_gpu_spec(args.gpu)): + task = _island_task(make_miles_island_task, args, spec, m, 3, "none", {}) + cfg = build_modal_island_config(args, spec, m, task, "none") + assert (RECLAIM_MARKER_ENV in cfg.envs) == (m < 2) + + +def test_two_samplers_refuse_the_trainer_scheduler(): + with pytest.raises(ValueError, match="several sampler islands"): + _args(("--rl-trainer-scheduler", "s1", "--rl-trainer-island", "2")) + + +def test_every_island_needs_a_role_and_one_trainer(): + from yeto.launcher import _prepare_rl_args + + two_trainers = tuple("1:trainer" if t == "1:sampler" else t for t in GROUP) + args = _launcher_args("ports", two_trainers, gpu=GPU3) + args.controller = "local" + with pytest.raises(ValueError, match="exactly one trainer"): + _prepare_rl_args(args) + + +def test_the_modal_runner_turns_the_signal_into_the_marker_on_a_sampler(monkeypatch): + from yeto import modal_runner as mr + + seen = [] + monkeypatch.setattr(mr, "_run_forwarding_to_learner", lambda cmd, env, **kw: seen.append(env) or 0) + + class Sig: + SIGINT, SIGTERM = _signal.SIGINT, _signal.SIGTERM + + @staticmethod + def signal(num, fn): + raise AssertionError("not expected") + + env = {RECLAIM_MARKER_ENV: "/tmp/yeto-reclaim-requested"} + assert mr.run_island_script(["x"], env, popen=lambda c, env, **k: None, signal_mod=Sig) == 0 + assert seen == [env] diff --git a/yeto/launch/rl_ports.py b/yeto/launch/rl_ports.py index c1f2b369c..bdeceda94 100644 --- a/yeto/launch/rl_ports.py +++ b/yeto/launch/rl_ports.py @@ -79,15 +79,16 @@ def _prepare_ports_algorithm(args) -> None: no_sync = bool(getattr(args, "rl_single_island_no_sync", False)) from ..rl.engine import sample_roles - # L5.4: a store-coupled sampler/trainer pair is two no-sync islands; a - # convertible pool (L7.5) is two or more (sample_roles holds their own rules). + # L5.4: a store-coupled pair (one trainer, one or more samplers: S25 C1) and a + # convertible pool (L7.5) are two or more no-sync islands (sample_roles holds + # their own rules). no_sync_islands = (islands if sample_roles.is_convertible_pool(args) - else 2 if sample_roles.is_pair(args) else 1) + or sample_roles.is_pair(args) else 1) if no_sync and (islands != no_sync_islands or getattr(args, "rl_sync_preset", "strict-avg") != "strict-avg" or getattr(args, "rl_initial_adapter", None)): raise AlgorithmSpecError( "--rl-single-island-no-sync needs exactly one island (one --gpu entry, no " - "external learners; two with --rl-island-sample-role), the default sync " + "external learners; one per island with --rl-island-sample-role), the default sync " "preset and no initial adapter" ) args.rl_allow_unverified_mechanism = list( diff --git a/yeto/launcher.py b/yeto/launcher.py index 7955bea99..38b5ac140 100644 --- a/yeto/launcher.py +++ b/yeto/launcher.py @@ -187,7 +187,8 @@ def _island_task(task_factory, args, spec, m, num_learners, syncer_addr, overrid learner_id, island_learners = sample_roles.island_learner(args, m, num_learners) task = task_factory(iargs, spec, learner_id, island_learners, syncer_addr) io.apply_to_task(task, {**io.island_env(args, m, overrides), - **(droppable.role_env(role, spec.cloud, spec.region) if role else {})}) + **(droppable.role_env(role, spec.cloud, spec.region) if role else {}), + **sample_roles.island_env(iargs)}) return task diff --git a/yeto/modal_runner.py b/yeto/modal_runner.py index e2b997379..b25be3c4f 100644 --- a/yeto/modal_runner.py +++ b/yeto/modal_runner.py @@ -925,7 +925,9 @@ def run_island_script(cmd: list[str], env: dict, *, popen=subprocess.Popen, sign sig = signal_mod or _signal if not _is_droppable(env): - if env.get("YETO_SPOT_INFLIGHT_SAVE_DIR"): + # in-flight save (4.1) and sampler islands (S25 C1, the launcher sets the marker + # path) both turn the reclaim signal into the marker file + if env.get("YETO_SPOT_INFLIGHT_SAVE_DIR") or env.get("YETO_SPOT_RECLAIM_MARKER"): return _run_forwarding_to_learner(cmd, env, popen=popen, sig=sig) return subprocess.call(cmd, env=env) if popen is subprocess.Popen else popen(cmd, env=env).wait() killpg = killpg or os.killpg diff --git a/yeto/rl/adapters/miles/buffer_rollout.py b/yeto/rl/adapters/miles/buffer_rollout.py index 2207badd9..6cc78c177 100644 --- a/yeto/rl/adapters/miles/buffer_rollout.py +++ b/yeto/rl/adapters/miles/buffer_rollout.py @@ -69,9 +69,14 @@ POLL_ENV = "YETO_SAMPLE_BUFFER_POLL_S" # reader: seconds between index reads while waiting MAX_POLICY_AGE_ENV = "YETO_SAMPLE_BUFFER_MAX_POLICY_AGE" SHARD_ENV = "YETO_SAMPLE_BUFFER_SHARD" # convertible node: its own shard of the index +# S25 C1: several sampler islands on one shared directory (a Modal Volume): each +# sampler writes its own shard, the trainer merges these shards (comma list) +SHARDS_ENV = "YETO_SAMPLE_BUFFER_SHARDS" +INCARNATION_ENV = "YETO_SAMPLER_INCARNATION" # set by the sampler island before Ray starts # the rollout functions run in Ray workers, which inherit the raylet's environment: # connect_island_ray forwards these to every worker -ISLAND_ENVS = (BUFFER_ROOT_ENV, SAMPLER_ID_ENV, WAIT_ENV, POLL_ENV, MAX_POLICY_AGE_ENV, SHARD_ENV) +ISLAND_ENVS = (BUFFER_ROOT_ENV, SAMPLER_ID_ENV, WAIT_ENV, POLL_ENV, MAX_POLICY_AGE_ENV, SHARD_ENV, + SHARDS_ENV, INCARNATION_ENV) DEFAULT_WAIT_S = 3600.0 DEFAULT_POLL_S = 2.0 @@ -111,10 +116,12 @@ def rollout_argv(role: str | None, *, max_policy_age: int, fully_async: bool = F def island_env(role: str, *, store_root: str, sampler_id: str | None = None, max_policy_age: int | None = None, wait_s: float | None = None, - shard: str | None = None) -> dict[str, str]: + shard: str | None = None, shards: Sequence[str] | None = None) -> dict[str, str]: """The environment the launch layer exports on the island (read in the rollout process). A convertible node writes like a sampler, into its own ``shard`` of - the index, and reads like a trainer.""" + the index, and reads like a trainer. ``shards``: the sampler shards of a + launch with several sampler islands (a sampler writes the one named by its + sampler id, the trainer merges them all).""" if role not in ROLES: raise ValueError(f"sample buffer role {role!r}: expected one of {ROLES}") env = {BUFFER_ROOT_ENV: str(store_root)} @@ -132,9 +139,20 @@ def island_env(role: str, *, store_root: str, sampler_id: str | None = None, env[MAX_POLICY_AGE_ENV] = str(int(max_policy_age)) if wait_s is not None: env[WAIT_ENV] = f"{float(wait_s):g}" + if shards: + if role == CONVERTIBLE: + raise ValueError("a convertible pool names its shards per node, not with shards") + if role == SAMPLER and sampler_id not in shards: + raise ValueError(f"sampler {sampler_id!r} is not one of the buffer shards {list(shards)}") + env[SHARDS_ENV] = ",".join(shards) return env +def buffer_shards(environ: Mapping[str, str]) -> tuple[str, ...]: + """The sampler shards of :data:`SHARDS_ENV` (() when unset).""" + return tuple(s.strip() for s in (environ.get(SHARDS_ENV) or "").split(",") if s.strip()) + + def open_buffer(environ: Mapping[str, str] | None = None, *, batch_size: int = 1, shard: str | None = None, create_index: bool = True, sampler: bool = False, reader: bool = False) -> Any: @@ -142,9 +160,11 @@ def open_buffer(environ: Mapping[str, str] | None = None, *, batch_size: int = 1 with ``shard``, the writer of that shard of the index; ``create_index=False`` for a reader whose index another host writes. - On a store that needs one shard per sampler island (``sample_buffer.needs_shards``): + On a store that needs one shard per sampler island (``sample_buffer.needs_shards``), + or with :data:`SHARDS_ENV` (several sampler islands on one shared directory): ``sampler`` opens the shard named by :data:`SAMPLER_ID_ENV`, and ``reader`` - (the trainer) merges the shards into the index.""" + (the trainer) merges the shards into the index (with :data:`SHARDS_ENV`, the + data ends once every named shard has ended).""" from yeto.rl.engine.resume import store_for from yeto.rl.engine.sample_buffer import INDEX, FileSampleBuffer, needs_shards, shard_index @@ -154,7 +174,8 @@ def open_buffer(environ: Mapping[str, str] | None = None, *, batch_size: int = 1 raise SampleBufferRolloutError(f"{BUFFER_ROOT_ENV} is not set") store = store_for(root, environ=env) merge = False - if needs_shards(store) and shard is None: + named = buffer_shards(env) + if (needs_shards(store) or named) and shard is None: if sampler: shard = env.get(SAMPLER_ID_ENV) if not shard: @@ -162,7 +183,8 @@ def open_buffer(environ: Mapping[str, str] | None = None, *, batch_size: int = 1 merge = reader return FileSampleBuffer(store, batch_size=batch_size, index=INDEX if shard is None else shard_index(shard), - create_index=create_index, merge_shards=merge) + create_index=create_index, merge_shards=merge, + expected_shards=named if merge and named else None) def renumber_group(samples: Sequence[dict[str, Any]], *, position: int, @@ -191,6 +213,85 @@ def renumber_group(samples: Sequence[dict[str, Any]], *, position: int, return out +# ---------------------------------------------------------------------- sampler side +_DATA_CURSOR_FIELDS = ("sample_offset", "epoch_id", "sample_group_index", "sample_index") + + +def seek_data_source(data_source: Any, cursor: Mapping[str, int], *, shuffle: bool) -> None: + """Put a Miles ``RolloutDataSource`` at ``cursor`` the way its ``load`` does + (the four position fields, then the epoch's shuffle). Raises when the source + does not have the fields (a source this resume does not know).""" + source = getattr(data_source, "__self__", data_source) + missing = [f for f in _DATA_CURSOR_FIELDS if not isinstance(getattr(source, f, None), int)] + if missing or set(cursor) != set(_DATA_CURSOR_FIELDS): + raise SampleBufferRolloutError( + f"cannot seek data source {type(source).__name__} to {dict(cursor)} " + f"(missing fields {missing})") + for name in _DATA_CURSOR_FIELDS: + setattr(source, name, int(cursor[name])) + if shuffle and getattr(source, "dataset", None) is not None: + source.dataset.shuffle(int(cursor["epoch_id"])) + + +class SamplerRoundWriter: + """The sampler's buffer writes, resumable across incarnations (S25 C1). + + Opened in the rollout process. On open it reads where this sampler's earlier + incarnations got to (:func:`~yeto.rl.engine.sampler_resume.resume_point`); + :meth:`seek` puts the data source at the resume cursor before the first + round; :meth:`write_round` writes one round's finished groups, each with its + :class:`~yeto.rl.engine.sampler_resume.SamplerRound` record. After the + reclaim marker appears (the island is leaving) nothing more is written: an + old incarnation must not write while its replacement writes the same shard.""" + + def __init__(self, buffer: Any, sampler_id: str, *, incarnation: str, resume: bool = True, + reclaimed: Callable[[], bool] = lambda: False) -> None: + """``resume=False``: write as before (a convertible node: no resume + point, no round check, no seek).""" + from yeto.rl.engine.sampler_resume import RoundGuard, resume_point + + self.buffer = buffer + self.sampler_id = str(sampler_id) + self.incarnation = str(incarnation) + self.reclaimed = reclaimed + self.resume = resume_point(buffer.entries() if resume else (), self.sampler_id) + self.guard = RoundGuard(self.resume) if resume else None + self._writer = GroupWriter(buffer, f"{self.sampler_id}-{self.incarnation}") + self._sought = False + self.dropped_groups = 0 # groups not written because the island is leaving + + def seek(self, data_source: Any, *, shuffle: bool) -> dict[str, int] | None: + """Before the first round only: the resume cursor applied (None: none).""" + if self._sought or self.guard is None: + return None + self._sought = True + cursor = self.resume.data_cursor + if cursor is not None: + seek_data_source(data_source, cursor, shuffle=shuffle) + return cursor + + def write_round(self, rollout_id: int, groups: Sequence[tuple[str, list[dict[str, Any]]]], + data_cursor: Mapping[str, int] | None) -> int: + """Write ``groups`` (``(group key, samples)``, aborted groups already left + out) of round ``rollout_id``; returns the groups written.""" + from yeto.rl.engine.sampler_resume import SamplerRound + + if self.guard is not None: + self.guard.check(rollout_id) + written = 0 + for i, (key, samples) in enumerate(groups): + if self.reclaimed(): + self.dropped_groups += len(groups) - i + break + record = SamplerRound(self.sampler_id, self.incarnation, int(rollout_id), + data_cursor=data_cursor, last=i == len(groups) - 1) + self._writer.write(samples, group_key=key, sampler_round=record.record()) + written += 1 + if written and self.guard is not None: + self.guard.written(rollout_id) + return written + + # ---------------------------------------------------------------------- trainer side @dataclass class TakeReport: diff --git a/yeto/rl/adapters/miles/buffer_rollout_miles.py b/yeto/rl/adapters/miles/buffer_rollout_miles.py index e9f6bad87..354023cb5 100644 --- a/yeto/rl/adapters/miles/buffer_rollout_miles.py +++ b/yeto/rl/adapters/miles/buffer_rollout_miles.py @@ -9,7 +9,6 @@ import asyncio import os -import uuid from typing import Any from miles.rollout.base_types import ( @@ -23,11 +22,13 @@ from miles.utils.types import Sample from yeto.rl.engine.buffer_payload import complete_prompt_tokens +from yeto.rl.engine.sampler_resume import new_incarnation, reclaim_requested from yeto.rl.engine.offpolicy_source import flat_metrics, offpolicy_source_event, payload_source_sample from .buffer_rollout import ( DEFAULT_POLL_S, DEFAULT_WAIT_S, + INCARNATION_ENV, MAX_POLICY_AGE_ENV, POLL_ENV, SAMPLER_ID_ENV, @@ -35,13 +36,14 @@ TRAINER, WAIT_ENV, BatchReader, - GroupWriter, SampleBufferRolloutError, + SamplerRoundWriter, convertible_mode, open_buffer, register_reader, renumber_group, ) +from .rollout_meta_hook import data_cursor def _flat(group: Any) -> list[Any]: @@ -52,32 +54,47 @@ def _flat(group: Any) -> list[Any]: class BufferWriterRolloutFn(InferenceRolloutFn): - """Sampler: Miles' rollout; every finished group is also written to the buffer.""" + """Sampler: Miles' rollout; every finished group is also written to the buffer. + + A replacement incarnation of the island continues where the buffer says the + earlier ones stopped (:class:`.buffer_rollout.SamplerRoundWriter`): the data + source is put at the saved cursor before the first round, and a round + already in the buffer is refused.""" + + resumable = True # a convertible node keeps its own rollout numbering (no resume) def __init__(self, input: RolloutFnConstructorInput) -> None: super().__init__(input) sampler_id = os.environ.get(SAMPLER_ID_ENV) if not sampler_id: raise SampleBufferRolloutError(f"{SAMPLER_ID_ENV} is not set") - # a per-process suffix: a restarted sampler repeats Miles group indexes, so the - # sample ids stay unique; the trainer renumbers groups by buffer position - self._writer = GroupWriter(self._open_writer_buffer(), f"{sampler_id}-{uuid.uuid4().hex[:8]}") + # the incarnation suffix keeps the sample ids unique: a restarted sampler repeats + # Miles group indexes; the trainer renumbers groups by buffer position + incarnation = os.environ.get(INCARNATION_ENV) or new_incarnation() + self._rounds = SamplerRoundWriter(self._open_writer_buffer(), sampler_id, + incarnation=incarnation, resume=self.resumable, + reclaimed=lambda: reclaim_requested(os.environ)) + self._shuffle = bool(getattr(input.args, "rollout_shuffle", False)) # the seek re-shuffles def _open_writer_buffer(self) -> Any: return open_buffer(sampler=True) async def _call_train(self, input: Any) -> Any: + self._rounds.seek(self.data_source, shuffle=self._shuffle) output = await super()._call_train(input) - await asyncio.to_thread(self._write, input.rollout_id, output.samples) + cursor, _ = data_cursor(self.data_source) + await asyncio.to_thread(self._write, input.rollout_id, output.samples, cursor) return output - def _write(self, rollout_id: int, groups: list[Any]) -> None: + def _write(self, rollout_id: int, groups: list[Any], cursor: dict[str, int] | None) -> None: + finished = [] for group in groups: samples = _flat(group) if any(s.status == Sample.Status.ABORTED for s in samples): continue # an aborted group is regenerated, never trained - key = f"r{int(rollout_id)}-g{samples[0].group_index}" - self._writer.write([s.to_dict() for s in samples], group_key=key) + finished.append((f"r{int(rollout_id)}-g{samples[0].group_index}", + [s.to_dict() for s in samples])) + self._rounds.write_round(rollout_id, finished, cursor) class _BufferRound: @@ -158,6 +175,8 @@ class ConvertibleRolloutFn(BufferWriterRolloutFn): :class:`BufferReaderRolloutFn`. The node switches the mode in this process (``buffer_rollout.executor_set_mode``) when it converts.""" + resumable = False + def __init__(self, input: RolloutFnConstructorInput) -> None: super().__init__(input) # the head merges the shards into the index: this node never writes it diff --git a/yeto/rl/adapters/miles/island_entry.py b/yeto/rl/adapters/miles/island_entry.py index d45980835..6186a1482 100644 --- a/yeto/rl/adapters/miles/island_entry.py +++ b/yeto/rl/adapters/miles/island_entry.py @@ -258,10 +258,12 @@ def apply_sample_role(args, miles_args, environ=None) -> None: config = sample_roles.resolve(args, groups_per_round=args.groups_per_round) miles_args.yeto_rl_sample_buffer = config node = getattr(args, "rl_convertible_node", None) # L7.5: a convertible node's own shard + # a pair island runs as learner 0: its sampler id is its node name (S25 C1) + sampler = node or getattr(args, "rl_sampler_node", None) or f"island{int(args.learner_id)}" (os.environ if environ is None else environ).update(island_env( - config.role, store_root=config.store, sampler_id=node or f"island{int(args.learner_id)}", + config.role, store_root=config.store, sampler_id=sampler, max_policy_age=int(getattr(args, "rl_max_policy_age", 0) or 0), wait_s=config.wait_s, - shard=node)) + shard=node, shards=getattr(args, "rl_buffer_shards", None))) def apply_ports_infra_switches(args, miles_args, environ=None) -> None: diff --git a/yeto/rl/adapters/miles/sampler_island.py b/yeto/rl/adapters/miles/sampler_island.py index 4c49446a0..4289c91b4 100644 --- a/yeto/rl/adapters/miles/sampler_island.py +++ b/yeto/rl/adapters/miles/sampler_island.py @@ -32,8 +32,10 @@ import asyncio import hashlib +import os import shutil import tempfile +import time from collections.abc import Callable from pathlib import Path from typing import Any @@ -236,16 +238,41 @@ async def _load(self, path: str | None, version: int) -> list[Any]: return engines -def run_sampler_island(miles_args: Any, *, environ: Any = None) -> int: - """Run the sampler island to its data end; returns the rounds generated.""" +LEAVE_HOLD_S = 300.0 # after a reclaim LEAVE: wait this long for the cloud to stop the container + + +def run_sampler_island(miles_args: Any, *, environ: Any = None, + sleep: Callable[[float], None] = time.sleep) -> int: + """Run the sampler island to its data end; returns the rounds generated. + + A replacement incarnation (S25 C1) continues after the newest round its + island has in the buffer (``sampler_resume``), loads the newest policy + release before its first round, and writes ``rl_sampler_join`` / a member + record. On the reclaim marker (Modal SIGINT/SIGTERM, written by the Modal + runner) it starts no new round, writes no more groups, writes + ``rl_sampler_leave`` with reason PREEMPTED and waits for the cloud to stop + the container (at most :data:`LEAVE_HOLD_S`).""" from miles.ray.placement_group import create_rollout_components from miles.utils.async_utils import Disposer from miles.utils.orchestration_utils import init_orchestration_script + from yeto.rl.elastic_client import LeaveReason from yeto.rl.engine.events import append_tape_event + from yeto.rl.engine.resume import store_for from yeto.rl.engine.sampler_loop import SamplerLoop + from yeto.rl.engine.sampler_resume import ( + JOIN_EVENT, + LEAVE_EVENT, + Member, + new_incarnation, + reclaim_marker, + reclaim_requested, + resume_point, + write_member, + ) from . import LoopRunner + from .buffer_rollout import INCARNATION_ENV, SAMPLER_ID_ENV from .entry import connect_island_ray from .rollout import RayMetadataSink, _release @@ -254,10 +281,23 @@ def run_sampler_island(miles_args: Any, *, environ: Any = None) -> int: raise SampleBufferRolloutError("run_sampler_island needs --rl-sample-role sampler") if not getattr(miles_args, "debug_rollout_only", False): raise SampleBufferRolloutError("a sampler island runs Miles with --debug-rollout-only") + env = os.environ if environ is None else environ def emit(event: str, **fields: Any) -> None: append_tape_event(miles_args, {"event": event, **fields}) + # one incarnation id for this process and its rollout workers (forwarded to Ray) + sampler_id = env.get(SAMPLER_ID_ENV) + if not sampler_id: + raise SampleBufferRolloutError(f"{SAMPLER_ID_ENV} is not set") + incarnation = env[INCARNATION_ENV] = new_incarnation() + reclaim_marker(env).unlink(missing_ok=True) # a restarted machine keeps its /tmp + buffer = open_buffer(env, sampler=True) + resume = resume_point(buffer.entries(), sampler_id) + store = store_for(config.store, environ=environ) + member = Member(sampler=sampler_id, incarnation=incarnation, joined_at=time.time(), + resume=resume.fields()) + connect_island_ray(miles_args=miles_args) # The rollout process's metadata hook pushes each round's group records to this # named actor (on a trainer island the driver creates it and consumes them). @@ -280,7 +320,16 @@ async def generate(rollout_id: int) -> None: _release(await executor.get(rollout_id)) # the rollout function wrote the groups metadata.discard() # the groups are in the buffer; nobody here reads the records + def leave(reason: LeaveReason, rounds: int, **detail: Any) -> None: + member.left_at, member.leave_reason, member.rounds = time.time(), reason.name, rounds + member.policy_version = follower.version + member.detail = detail + emit(LEAVE_EVENT, sampler=sampler_id, incarnation=incarnation, reason=reason.name, + rounds=rounds, policy_version=follower.version, buffer_end=len(buffer), **detail) + write_member(store, member) + error: BaseException | None = None + follower = None try: controller, executor = runner.run(init()) lora = lora_mode(miles_args) @@ -289,19 +338,41 @@ async def generate(rollout_id: int) -> None: loader.stamp(BASE_POLICY_VERSION) # until the first release is loaded follower = PolicyFollower(release_store(config, environ), loader, version=BASE_POLICY_VERSION) + reclaimed = lambda: reclaim_requested(env) # noqa: E731 loop = SamplerLoop( - follower=follower, buffer=open_buffer(environ, sampler=True), emit=emit, + follower=follower, buffer=buffer, emit=emit, generate=lambda rollout_id: runner.run(generate(rollout_id)), num_rollouts=int(miles_args.num_rollout), max_backlog_groups=int(config.max_backlog_groups), - min_round_s=config.min_round_s, require_release=lora, + 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}) + latest = follower.releases.latest() + emit(JOIN_EVENT, sampler=sampler_id, incarnation=incarnation, + latest_release=None if latest is None else {"seq": latest.seq, "policy_version": latest.version}, + **resume.fields()) + write_member(store, member) emit("rl_sampler_start", store=config.store, num_rollouts=int(miles_args.num_rollout), - max_backlog_groups=int(config.max_backlog_groups), weights=loader.weights) - return loop.run() + max_backlog_groups=int(config.max_backlog_groups), weights=loader.weights, + start_rollout_id=resume.next_rollout_id) + if buffer.data_ended: # an earlier incarnation finished and marked the data end + rounds = 0 + else: + rounds = loop.run(start=resume.next_rollout_id) + if reclaimed(): + leave(LeaveReason.PREEMPTED, rounds) + _hold_until_stopped(sleep, LEAVE_HOLD_S) + else: + leave(LeaveReason.FINISHED, rounds) + return rounds except BaseException as exc: error = exc + if follower is not None: + try: + leave(LeaveReason.PREEMPTED if reclaim_requested(env) else LeaveReason.INFRA_FAILED, + -1, error=f"{type(exc).__name__}: {exc}"[:300]) + except Exception: # noqa: BLE001 - the original error is the one to report + pass raise finally: try: @@ -311,6 +382,15 @@ async def generate(rollout_id: int) -> None: runner.close() +def _hold_until_stopped(sleep: Callable[[float], None], hold_s: float) -> None: + """After a reclaim LEAVE the process stays up until the cloud stops the + container: an early clean exit could end the cloud's input as finished + instead of letting the cloud run the island again.""" + print(f"[sampler-island] reclaimed: waiting up to {hold_s:g} s for the container to stop", + flush=True) + sleep(float(hold_s)) + + # ---------------------------------------------------------------------- trainer side class ReleasingPublisher: """A publisher that also releases each published policy for the samplers. diff --git a/yeto/rl/engine/buffer_payload.py b/yeto/rl/engine/buffer_payload.py index 23bb1f8c8..460831f37 100644 --- a/yeto/rl/engine/buffer_payload.py +++ b/yeto/rl/engine/buffer_payload.py @@ -142,14 +142,18 @@ def __init__(self, buffer: Any, sampler_id: str) -> None: self.sampler_id = str(sampler_id) def write(self, samples: Sequence[Mapping[str, Any]], *, group_key: str, - provenance: Mapping[str, str] | None = None) -> SampleIndexEntry: - """``provenance``: the sampler's source record for the index entry (S25 N2.4).""" + provenance: Mapping[str, str] | None = None, + sampler_round: Mapping[str, Any] | None = None) -> SampleIndexEntry: + """``provenance``: the sampler's source record for the index entry (S25 N2.4); + ``sampler_round``: the round record a replacement sampler resumes from + (``sampler_resume.SamplerRound.record()``, S25 C1).""" behavior_version, tokens = group_entry_fields(samples) sample_id = f"{self.sampler_id}-{group_key}" ref = self.buffer.write_payload(sample_id, encode_group(samples)) entry = SampleIndexEntry(sample_id=sample_id, behavior_version=behavior_version, payload_ref=ref, tokens=tokens, - provenance=None if provenance is None else dict(provenance)) + provenance=None if provenance is None else dict(provenance), + sampler_round=None if sampler_round is None else dict(sampler_round)) self.buffer.append(entry) return entry diff --git a/yeto/rl/engine/run_flags.py b/yeto/rl/engine/run_flags.py index 9b1314c3f..b4de91487 100644 --- a/yeto/rl/engine/run_flags.py +++ b/yeto/rl/engine/run_flags.py @@ -588,6 +588,14 @@ def rl_forward_pairs(*dests: str | tuple[str, str]) -> tuple[tuple[str, str], .. 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)"), + shared_flag("--rl-sampler-node", default=None, metavar="NAME", + help="--rl-sample-role sampler: this island's node name (set by the launcher " + "per island of --rl-island-sample-role; the sampler id that a replacement " + "incarnation of the island resumes under)"), + shared_flag("--rl-buffer-shards", default=None, metavar="A,B", + type=lambda text: [s for s in text.split(",") if s], + help="--rl-island-sample-role with several samplers: the sampler shards of the " + "buffer (set by the launcher); each sampler writes its own, the trainer merges them"), ) _SCHEDULING = ( diff --git a/yeto/rl/engine/sample_buffer.py b/yeto/rl/engine/sample_buffer.py index e4feb4023..4c945eb14 100644 --- a/yeto/rl/engine/sample_buffer.py +++ b/yeto/rl/engine/sample_buffer.py @@ -96,12 +96,15 @@ def payload_name(sample_id: str) -> str: class JsonlIndex: """The index as ``index.jsonl`` in a shared directory (see the module doc).""" - def __init__(self, store: CheckpointStore, *, index: str = INDEX, create_index: bool = True) -> None: + def __init__(self, store: CheckpointStore, *, index: str = INDEX, create_index: bool = True, + reload: bool = True) -> None: """``index``: the index file under ``sample-buffer/`` (a shard writer passes :func:`shard_index`). ``create_index=False``: a reader that must not write the index, not even an empty one (on a store mount whose commit would - replace the merger's index).""" + replace the merger's index). ``reload=False``: :meth:`read_new` does not + reload the store (a merger's shard reader: the merger reloaded it once).""" self.store = store + self._reload = bool(reload) self.root = Path(store.root) / BUFFER_DIR (self.root / PAYLOADS).mkdir(parents=True, exist_ok=True) self.path = self.root / index @@ -128,22 +131,30 @@ def has_payload(self, ref: str) -> bool: return (self.root / ref).is_file() def read_new(self) -> list[dict[str, Any]]: - self.store.reload() - if not self.path.exists(): # a reader before the first merge + if self._reload: + self.store.reload() + if not self.path.exists(): # a reader before the first merge (or a shard not created yet) return [] with open(self.path, "rb") as fh: return list(self._lines(fh)) def append(self, absorb: Callable[[list[dict[str, Any]]], None], make: Callable[[], dict[str, Any]]) -> None: + self.append_many(absorb, lambda: [make()]) + + def append_many(self, absorb: Callable[[list[dict[str, Any]]], None], + make: Callable[[], list[dict[str, Any]]]) -> None: + """Append ``make()``'s records with one write, one fsync and one store + commit (a merger appends a whole refresh at once: one Modal Volume commit + per record would cost seconds each).""" with open(self.path, "r+b") as fh: fcntl.flock(fh.fileno(), fcntl.LOCK_EX) try: absorb(list(self._lines(fh))) - record = make() + records = make() fh.truncate(self._offset) # drop a torn last line fh.seek(self._offset) - fh.write((json.dumps(record, sort_keys=True) + "\n").encode()) + fh.write(b"".join((json.dumps(r, sort_keys=True) + "\n").encode() for r in records)) fh.flush() os.fsync(fh.fileno()) absorb(list(self._lines(fh))) @@ -151,6 +162,27 @@ def append(self, absorb: Callable[[list[dict[str, Any]]], None], fcntl.flock(fh.fileno(), fcntl.LOCK_UN) self.store.commit() + # ------------------------------------------------------------ shards (merger side) + def shard_names(self) -> list[str]: + """The shard index files under ``sample-buffer/shards/``, sorted.""" + return sorted(p.name[:-len(".jsonl")] for p in (self.root / SHARDS).glob("*.jsonl")) + + def shard_reader(self, shard: str, start: int) -> "JsonlIndex": + """A reader of one shard's index file that skips its first ``start`` + records (merged before).""" + reader = JsonlIndex(self.store, index=shard_index(shard), create_index=False, reload=False) + if start: + skipped = 0 + if reader.path.exists(): + with open(reader.path, "rb") as fh: + for _ in reader._lines(fh): + skipped += 1 + if skipped == start: + break + if skipped < start: + raise SampleBufferError(f"shard {shard}: {skipped} records, but {start} are merged") + return reader + def _lines(self, fh) -> Iterator[dict[str, Any]]: fh.seek(self._offset) data = fh.read() @@ -223,6 +255,13 @@ def append(self, absorb: Callable[[list[dict[str, Any]]], None], return raise SampleBufferError(f"{self.prefix}: lost {_MAX_CREATE_RACES} appends in a row to other writers") + # ------------------------------------------------------------ shards (merger side) + def shard_names(self) -> list[str]: + return shard_names(self.objects) + + def shard_reader(self, shard: str, start: int) -> "ObjectIndex": + return ObjectIndex(self.objects, prefix=shard_prefix(shard), start=start) + def _objects_of(store: CheckpointStore | ObjectStore) -> ObjectStore | None: return store.objects if isinstance(store, CheckpointStore) else store @@ -249,7 +288,8 @@ def _entry(record: dict[str, Any]) -> SampleIndexEntry: return SampleIndexEntry( sample_id=record["sample_id"], behavior_version=int(record["behavior_version"]), payload_ref=record["payload_ref"], tokens=int(record["tokens"]), - appended_at=float(record["appended_at"]), provenance=record.get("provenance")) + appended_at=float(record["appended_at"]), provenance=record.get("provenance"), + sampler_round=record.get("sampler_round")) class UnmergedShards: @@ -321,21 +361,31 @@ class FileSampleBuffer: def __init__(self, store: CheckpointStore | ObjectStore, *, batch_size: int, arrival_window_s: float = 60.0, clock: Callable[[], float] = time.time, - index: str = INDEX, create_index: bool = True, merge_shards: bool = False) -> None: + index: str = INDEX, create_index: bool = True, merge_shards: bool = False, + expected_shards: Iterable[str] | None = None) -> None: """``index``/``create_index``: see :class:`JsonlIndex`. ``merge_shards``: - this buffer (the consumer) merges the sampler shards of an object store - into the index on every refresh (see the module doc); it must be the - only one that does.""" + this buffer (the consumer) merges the sampler shards into the index on + every refresh (see the module doc); it must be the only one that does. + On a shared directory the shards are ``shards/.jsonl`` files (one + sampler island writes each; S25 C1: several sampler islands on one + Modal Volume, where ``flock`` does not hold across containers). + ``expected_shards``: the shards that must all end before the merged index + ends (None: the shards seen so far), so an early-ending sampler cannot + end the data before a slower one has written its first record.""" if batch_size < 1: raise ValueError("batch_size must be positive") self.store = store self.index = index_for(store, index=index, create_index=create_index) self.shard = _shard_of_index(index) if index != INDEX else None # the shard this buffer writes - if merge_shards and (self.shard is not None or not isinstance(self.index, ObjectIndex)): - raise SampleBufferError("only the index of an object store merges shards") + if merge_shards and self.shard is not None: + raise SampleBufferError("a shard writer does not merge shards (only the consumer's index does)") + if expected_shards is not None and not merge_shards: + raise SampleBufferError("expected_shards needs merge_shards") self._merge = merge_shards + self._expected_shards = (None if expected_shards is None + else frozenset(_shard_name(s) for s in expected_shards)) self._shard_of: list[str | None] = [] # the shard of each position (merged records) - self._shard_readers: dict[str, ObjectIndex] = {} + self._shard_readers: dict[str, JsonlIndex | ObjectIndex] = {} self._shards_ended: set[str] = set() self.batch_size = int(batch_size) self.arrival_window_s = float(arrival_window_s) @@ -368,6 +418,8 @@ def append(self, entry: SampleIndexEntry) -> None: entry = dataclasses.replace(entry, appended_at=self.clock()) # ``provenance`` is written only when set, so a line without it reads as before. source = {} if entry.provenance is None else {"provenance": dict(entry.provenance)} + if entry.sampler_round is not None: + source["sampler_round"] = dict(entry.sampler_round) self._write_record(lambda pos: {"schema": SCHEMA, "pos": pos, "sample_id": entry.sample_id, "behavior_version": int(entry.behavior_version), "payload_ref": entry.payload_ref, "tokens": int(entry.tokens), @@ -386,6 +438,22 @@ def record() -> dict: with self._lock: self.index.append(self._absorb, record) + def _write_records(self, records: list[dict[str, Any]]) -> None: + """Append ``records`` (without ``pos``) at the next positions: one write on + a shared directory, one object per record on an object store.""" + def make() -> list[dict]: + if self._ended: + raise SampleBufferError("append after the data end") + start = len(self._entries) + return [{**record, "pos": start + i} for i, record in enumerate(records)] + + with self._lock: + if isinstance(self.index, JsonlIndex): + self.index.append_many(self._absorb, make) + return + for record in records: # an object store: one create-if-absent per record + self._write_record(lambda pos, record=record: {**record, "pos": pos}) + def _absorb(self, records: list[dict[str, Any]]) -> None: for record in records: if record.get("schema") != SCHEMA or int(record["pos"]) != len(self._entries): @@ -407,14 +475,12 @@ def refresh(self) -> int: return len(self._entries) def _merge_shards(self) -> None: - objects = self.index.objects new: list[list[dict[str, Any]]] = [] # per shard, in shard order - for shard in shard_names(objects): + for shard in self.index.shard_names(): reader = self._shard_readers.get(shard) if reader is None: # first seen here: skip what an earlier merger merged done = sum(1 for s in self._shard_of if s == shard) - reader = self._shard_readers[shard] = ObjectIndex(objects, prefix=shard_prefix(shard), - start=done) + reader = self._shard_readers[shard] = self.index.shard_reader(shard, done) records = [] for record in reader.read_new(): if record.get("end"): @@ -430,11 +496,12 @@ def _merge_shards(self) -> None: # records of a shard must stay a prefix of it, because a restarted merger # resumes each shard after its merged count (a clock step back would # otherwise merge a later record first, then skip the earlier one). - merged = heapq.merge(*new, key=lambda r: (float(r.get("appended_at", 0.0)), r["shard"])) - for record in merged: - self._write_record(lambda pos, record=record: {**record, "pos": pos}) - if self._shard_readers and self._shards_ended >= set(self._shard_readers): - self._write_record(lambda pos: {"schema": SCHEMA, "pos": pos, "end": True}) + merged = list(heapq.merge(*new, key=lambda r: (float(r.get("appended_at", 0.0)), r["shard"]))) + expected = set(self._shard_readers) if self._expected_shards is None else self._expected_shards + if expected and self._shards_ended >= expected: + merged.append({"schema": SCHEMA, "end": True}) + if merged: + self._write_records(merged) def shard_positions(self, cursor: int) -> dict[str, int]: """{shard: records of that shard below merged position ``cursor``}: the @@ -449,6 +516,11 @@ def __len__(self) -> int: with self._lock: return len(self._entries) + def entries(self) -> list[SampleIndexEntry]: + """The entries read so far, in buffer order (a copy).""" + with self._lock: + return list(self._entries) + @property def data_ended(self) -> bool: with self._lock: diff --git a/yeto/rl/engine/sample_roles.py b/yeto/rl/engine/sample_roles.py index 10e68c357..e9b76b2fd 100644 --- a/yeto/rl/engine/sample_roles.py +++ b/yeto/rl/engine/sample_roles.py @@ -18,8 +18,9 @@ The sampler and the trainer mount the same checkpoint store (the buffer and the releases live there). They are either separate single-island launches (``--rl-sample-role``), or one launch of two islands, a *store-coupled pair* -(``--rl-island-sample-role 0:sampler --rl-island-sample-role 1:trainer``, L5.4): -there is no syncer, and each island runs as its own one-learner, no-sync island +(``--rl-island-sample-role 0:sampler --rl-island-sample-role 1:trainer``, L5.4), +or one launch of one trainer and several sampler islands (S25 C1: each sampler +writes its own shard of the buffer, :func:`buffer_shards`): there is no syncer, and each island runs as its own one-learner, no-sync island (learner id 0 of 1) whose role comes from its island args (:func:`island_args`). This module is engine-neutral; the Miles mapping is in ``adapters.miles.buffer_rollout``. @@ -113,10 +114,24 @@ def island_roles(args: Any, num_islands: int) -> dict[int, str]: def is_pair(args: Any) -> bool: - """Whether the launch is a store-coupled sampler/trainer pair.""" + """Whether the launch is a store-coupled group: one trainer island and one or + more sampler islands (``--rl-island-sample-role``; "pair" for short).""" return bool(getattr(args, "rl_island_sample_role", None)) +def sampler_islands(args: Any, num_islands: int) -> list[int]: + """The sampler islands of a pair, sorted.""" + return sorted(m for m, role in island_roles(args, num_islands).items() if role == SAMPLER) + + +def buffer_shards(args: Any, num_islands: int) -> list[str] | None: + """The buffer shards of a pair with several sampler islands (S25 C1): one per + sampler island, named by its node name; None with one sampler (it writes the + index itself, as before).""" + samplers = sampler_islands(args, num_islands) + return [node_name(m) for m in samplers] if len(samplers) > 1 else None + + def is_store_coupled(args: Any) -> bool: """A pair or a convertible pool: no syncer, each island is its own one-learner no-sync island over the shared checkpoint store.""" @@ -133,6 +148,9 @@ def island_args(args: Any, island: int, num_islands: int) -> Any: return args out = copy.copy(args) out.rl_sample_role = roles[island] + if out.rl_sample_role in (SAMPLER, TRAINER): + out.rl_sampler_node = node_name(island) if out.rl_sample_role == SAMPLER else None + out.rl_buffer_shards = buffer_shards(args, num_islands) if out.rl_sample_role == CONVERTIBLE: out.rl_convertible_node = node_name(island) out.rl_sampler_min_round_seconds = None # a node keeps its backlog cap @@ -161,17 +179,23 @@ def buffer_full_parameter(args: Any) -> bool: def check_pair(args: Any, num_islands: int) -> None: - """Launch checks of a store-coupled pair (nothing outside a pair): exactly two - islands, one sampler and one trainer, ``--rl-single-island-no-sync`` (no - syncer), no ``--rl-sample-role`` and no external learner; then every island's - own role checks on its island args.""" + """Launch checks of a store-coupled pair (nothing outside a pair): every + island has a role, exactly one trainer and at least one sampler, + ``--rl-single-island-no-sync`` (no syncer), no ``--rl-sample-role`` and no + external learner; several samplers take no trainer scheduler (the head reads + no sampler shard of a shared directory); then every island's own role checks + on its island args.""" roles = island_roles(args, num_islands) if not roles: return - if num_islands != 2 or sorted(roles.values()) != [SAMPLER, TRAINER] or len(roles) != 2: - raise ValueError("--rl-island-sample-role: a pair is two --gpu islands, one " - f"{SAMPLER} and one {TRAINER} (got {num_islands} island(s), roles " - f"{dict(sorted(roles.items()))})") + values = list(roles.values()) + if len(roles) != num_islands or values.count(TRAINER) != 1 or values.count(SAMPLER) < 1: + raise ValueError("--rl-island-sample-role: a pair gives every --gpu island a role, " + f"at least one {SAMPLER} and exactly one {TRAINER} (got {num_islands} " + f"island(s), roles {dict(sorted(roles.items()))})") + if values.count(SAMPLER) > 1 and getattr(args, "rl_trainer_scheduler", None) not in (None, "off"): + raise ValueError("--rl-trainer-scheduler with several sampler islands is not supported " + "(the head does not read the sampler shards); drop the scheduler") if sample_role(args) != BOTH: raise ValueError("--rl-island-sample-role replaces --rl-sample-role; give only one") if int(getattr(args, "external_learners", 0) or 0): @@ -288,6 +312,18 @@ def resolve(args: Any, *, groups_per_round: int) -> SampleBufferConfig | None: else None)) +def island_env(args: Any) -> dict[str, str]: + """The environment a pair's sampler island gets from the launcher ({} for + every other island; ``args`` are the island args): the reclaim marker path, + so the Modal runner turns the reclaim signal into that file and the sampler + leaves without being signalled (S25 C1).""" + if sample_role(args) != SAMPLER or not getattr(args, "rl_sampler_node", None): + return {} + from .sampler_resume import DEFAULT_RECLAIM_MARKER, RECLAIM_MARKER_ENV + + return {RECLAIM_MARKER_ENV: DEFAULT_RECLAIM_MARKER} + + def island_flags(args: Any) -> str: """The learner flags the launcher forwards (empty for ``both``); ``args`` are the island args (:func:`island_args`).""" @@ -303,4 +339,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 getattr(args, "rl_sampler_node", None): + flags += f" --rl-sampler-node {args.rl_sampler_node}" + if getattr(args, "rl_buffer_shards", None): + flags += f" --rl-buffer-shards {','.join(args.rl_buffer_shards)}" return flags diff --git a/yeto/rl/engine/sampler_resume.py b/yeto/rl/engine/sampler_resume.py new file mode 100644 index 000000000..e25eefda9 --- /dev/null +++ b/yeto/rl/engine/sampler_resume.py @@ -0,0 +1,185 @@ +"""Sampler island replacement: where a new incarnation continues (S25 C1). + +A sampler island on spot capacity can be reclaimed at any time; the cloud (or +the launcher) then starts a new incarnation of the same island. The island's +sample-buffer writes must neither repeat nor lose groups, and the trainer keeps +training on what the buffer holds. + +Each index entry a sampler writes carries a :class:`SamplerRound` record +(``SampleIndexEntry.sampler_round``): the sampler id (the island, stable across +incarnations), the incarnation, the rollout id, the data cursor after the round +drew its prompts, and ``last`` on the round's last group. A Miles round writes +its groups after the whole round is generated, so the buffer itself says how +far the island got: + +* a round with any entry in the buffer is never written again + (:class:`RoundGuard`); the new incarnation continues at the next rollout id + from that round's data cursor, so no prompt is drawn twice; +* a round without its ``last`` entry (reclaimed while it wrote) keeps the + groups it wrote (each whole and written once); its unwritten groups are lost + (``partial_rounds`` counts such rounds); +* a round still generating when the island was reclaimed wrote nothing: the new + incarnation generates that rollout id again from the cursor before it. + +:func:`resume_point` reads the entries; the island loop starts at +``next_rollout_id`` and the rollout function seeks the data source to +``data_cursor``. The incarnation's arrival and departure are written as +``rl_sampler_join`` / ``rl_sampler_leave`` (reason: a ``LeaveReason`` name) and as +one member record per incarnation in the store (:func:`write_member`), which +outlives the island's own tape. Engine-neutral. +""" + +from __future__ import annotations + +import json +import uuid +from collections.abc import Iterable, Mapping +from dataclasses import asdict, dataclass, field +from pathlib import Path +from typing import Any + +JOIN_EVENT = "rl_sampler_join" +LEAVE_EVENT = "rl_sampler_leave" +DROPPED_EVENT = "rl_sampler_groups_dropped" # groups not written after the leave +MEMBERS_DIR = "sample-buffer/members" # one record per sampler incarnation, under the store root +RECLAIM_MARKER_ENV = "YETO_SPOT_RECLAIM_MARKER" # written by the Modal runner on the reclaim signal +DEFAULT_RECLAIM_MARKER = "/tmp/yeto-reclaim-requested" + + +class SamplerResumeError(RuntimeError): + pass + + +def new_incarnation() -> str: + """A fresh incarnation id (unique per island process).""" + return uuid.uuid4().hex[:8] + + +@dataclass(frozen=True) +class SamplerRound: + """The ``sampler_round`` record of one index entry.""" + + sampler: str + incarnation: str + rollout_id: int + data_cursor: Mapping[str, int] | None = None # the data source after the round drew its prompts + last: bool = False # the round's last group + + def record(self) -> dict[str, Any]: + out = {"sampler": self.sampler, "incarnation": self.incarnation, + "rollout_id": int(self.rollout_id), "last": bool(self.last)} + if self.data_cursor is not None: + out["data_cursor"] = {k: int(v) for k, v in self.data_cursor.items()} + return out + + +@dataclass(frozen=True) +class SamplerResume: + """Where the next incarnation of a sampler continues.""" + + next_rollout_id: int # after the newest round with an entry in the buffer + data_cursor: dict[str, int] | None # that round's data cursor (None: start of the data) + written_rounds: tuple[int, ...] # rounds with at least one entry: never written again + partial_rounds: tuple[int, ...] # written rounds without their last group + incarnations: tuple[str, ...] # earlier incarnations, in the order they first wrote + groups: int # every group this sampler has in the buffer + + @property + def resumed(self) -> bool: + return bool(self.incarnations) + + def fields(self) -> dict[str, Any]: + out = asdict(self) + out["written_rounds"] = len(self.written_rounds) + out["partial_rounds"] = list(self.partial_rounds) + out["incarnations"] = list(self.incarnations) + return out + + +def resume_point(entries: Iterable[Any], sampler: str) -> SamplerResume: + """Read the ``sampler_round`` records of ``sampler``'s entries (other samplers' + entries and entries without the record are skipped).""" + cursors: dict[int, dict[str, int] | None] = {} + ended: set[int] = set() + incarnations: list[str] = [] + groups = 0 + for entry in entries: + record = getattr(entry, "sampler_round", None) + if not record or record.get("sampler") != sampler: + continue + groups += 1 + rid = int(record["rollout_id"]) + raw = record.get("data_cursor") + cursor = None if raw is None else {k: int(v) for k, v in raw.items()} + if rid in cursors and cursors[rid] != cursor: + raise SamplerResumeError(f"sampler {sampler}: round {rid} has two data cursors " + f"({cursors[rid]} and {cursor})") + cursors[rid] = cursor + if record.get("incarnation") not in incarnations: + incarnations.append(str(record.get("incarnation"))) + if record.get("last"): + if rid in ended: + raise SamplerResumeError(f"sampler {sampler}: round {rid} has two last groups") + ended.add(rid) + newest = max(cursors) if cursors else None + return SamplerResume(next_rollout_id=0 if newest is None else newest + 1, + data_cursor=None if newest is None else cursors[newest], + written_rounds=tuple(sorted(cursors)), + partial_rounds=tuple(sorted(set(cursors) - ended)), + incarnations=tuple(incarnations), groups=groups) + + +class RoundGuard: + """Write-side check: a round with entries in the buffer is never written by + another incarnation, and an incarnation never writes before its first round.""" + + def __init__(self, resume: SamplerResume) -> None: + self._written = set(resume.written_rounds) + self._first = int(resume.next_rollout_id) + + def check(self, rollout_id: int) -> None: + rid = int(rollout_id) + if rid in self._written or rid < self._first: + raise SamplerResumeError( + f"round {rid} is already in the sample buffer (this incarnation starts at " + f"round {self._first}); a replacement sampler must not write it again") + + def written(self, rollout_id: int) -> None: + self._written.add(int(rollout_id)) + + +def reclaim_marker(environ: Mapping[str, str]) -> Path: + return Path(environ.get(RECLAIM_MARKER_ENV) or DEFAULT_RECLAIM_MARKER) + + +def reclaim_requested(environ: Mapping[str, str]) -> bool: + """The Modal runner wrote the reclaim marker (the container got SIGINT/SIGTERM).""" + return reclaim_marker(environ).exists() + + +@dataclass +class Member: + """One sampler incarnation's record in the store (``members//.json``).""" + + sampler: str + incarnation: str + joined_at: float + resume: dict[str, Any] + policy_version: int | None = None + left_at: float | None = None + leave_reason: str | None = None + rounds: int | None = None # rounds this incarnation completed + detail: dict[str, Any] = field(default_factory=dict) + + +def write_member(store: Any, member: Member) -> Path: + """Write (or rewrite) the incarnation's member record and commit the store. + One file per incarnation: an old and a new incarnation never write the same + file (a Modal Volume keeps the last commit of a file).""" + path = Path(store.root) / MEMBERS_DIR / member.sampler / f"{member.incarnation}.json" + path.parent.mkdir(parents=True, exist_ok=True) + tmp = path.with_suffix(".json.tmp") + tmp.write_text(json.dumps(asdict(member), sort_keys=True) + "\n") + tmp.replace(path) + store.commit() + return path diff --git a/yeto/rl/engine/trainer_lifecycle.py b/yeto/rl/engine/trainer_lifecycle.py index 731ecf418..ec3416aee 100644 --- a/yeto/rl/engine/trainer_lifecycle.py +++ b/yeto/rl/engine/trainer_lifecycle.py @@ -139,7 +139,13 @@ class SampleIndexEntry: record: device family, card, engine and version, precision, quantization scheme, topology); None for entries written before it existed or by a sampler that does not set it (read as ``unknown``). Data only: it never - enters a hash or the version gate (rl-hetero-sampler-engines D3).""" + enters a hash or the version gate (rl-hetero-sampler-engines D3). + + ``sampler_round``: which sampler round wrote the entry + (``sampler_resume.SamplerRound`` record: sampler, incarnation, rollout id, + last group of the round, data cursor after the round); None for entries + written before it existed. A replacement sampler island reads it to resume + after the last complete round (S25 C1).""" sample_id: str behavior_version: int @@ -147,6 +153,7 @@ class SampleIndexEntry: tokens: int = 0 appended_at: float = 0.0 provenance: Mapping[str, str] | None = field(default=None, compare=False, hash=False) + sampler_round: Mapping[str, Any] | None = field(default=None, compare=False, hash=False) @dataclass(frozen=True) From c7812a808c206e77c5298f490cc4a44e5e44df34 Mon Sep 17 00:00:00 2001 From: michael <78113554+chowhsu@users.noreply.github.com> Date: Tue, 20 Oct 2026 04:16:19 +0000 Subject: [PATCH 70/89] Split the trainer's rounds evenly over several sampler islands Co-Authored-By: Claude Opus 5.5 --- tests/test_rl_spot_sampler_replace.py | 10 +++++++++- yeto/rl/engine/sample_roles.py | 18 +++++++++++++++--- 2 files changed, 24 insertions(+), 4 deletions(-) diff --git a/tests/test_rl_spot_sampler_replace.py b/tests/test_rl_spot_sampler_replace.py index 92ff9b416..0ae3c3125 100644 --- a/tests/test_rl_spot_sampler_replace.py +++ b/tests/test_rl_spot_sampler_replace.py @@ -288,7 +288,7 @@ def test_member_records_are_one_file_per_incarnation(tmp_path): "--rl-island-sample-role", "0:sampler", "--rl-island-sample-role", "1:sampler", "--rl-island-sample-role", "2:trainer", "--rl-max-policy-age", "2", "--rl-algorithm-spec", "examples/rl_algorithms/decoupled-ppo-c2.json", - "--rl-checkpoint-store", STORE, "--cluster-prefix", "c1-spot") + "--rl-checkpoint-store", STORE, "--cluster-prefix", "c1-spot", "--total-steps", "4") def _args(extra=()): @@ -317,6 +317,14 @@ def test_one_trainer_and_two_samplers_get_their_node_and_the_shards(): tokens = [_tokens(i["learner_command"]) for i in plan["island_requests"]] assert [_flag(t, "--rl-sampler-node") for t in tokens] == ["island0", "island1", None] assert {_flag(t, "--rl-buffer-shards") for t in tokens} == {"island0,island1"} + # the trainer trains --total-steps rounds; each of the two samplers generates half + assert [_flag(t, "--global-rounds") for t in tokens] == ["2", "2", "4"] + assert [_flag(t, "--total-fragment-steps") for t in tokens] == ["2", "2", "4"] + + +def test_the_samplers_split_the_rounds_evenly(): + with pytest.raises(ValueError, match="multiple of the sampler count"): + _args(("--total-steps", "3")) def test_the_island_resolves_its_sampler_id_and_the_shards(): diff --git a/yeto/rl/engine/sample_roles.py b/yeto/rl/engine/sample_roles.py index e9b76b2fd..69f72d7c8 100644 --- a/yeto/rl/engine/sample_roles.py +++ b/yeto/rl/engine/sample_roles.py @@ -151,6 +151,13 @@ def island_args(args: Any, island: int, num_islands: int) -> Any: if out.rl_sample_role in (SAMPLER, TRAINER): out.rl_sampler_node = node_name(island) if out.rl_sample_role == SAMPLER else None out.rl_buffer_shards = buffer_shards(args, num_islands) + samplers = len(sampler_islands(args, num_islands)) + if out.rl_sample_role == SAMPLER and samplers > 1: + # the trainer trains --total-steps rounds; the samplers share them, so each + # one reaches its data end (check_pair needs an even split) + out.total_steps = int(args.total_steps) // samplers + if getattr(args, "rl_total_fragment_steps", None) is not None: + out.rl_total_fragment_steps = out.total_steps * int(getattr(args, "fragments", 1) or 1) if out.rl_sample_role == CONVERTIBLE: out.rl_convertible_node = node_name(island) out.rl_sampler_min_round_seconds = None # a node keeps its backlog cap @@ -182,8 +189,9 @@ def check_pair(args: Any, num_islands: int) -> None: """Launch checks of a store-coupled pair (nothing outside a pair): every island has a role, exactly one trainer and at least one sampler, ``--rl-single-island-no-sync`` (no syncer), no ``--rl-sample-role`` and no - external learner; several samplers take no trainer scheduler (the head reads - no sampler shard of a shared directory); then every island's own role checks + external learner; several samplers split ``--total-steps`` evenly (each + generates its share of the trainer's rounds) and take no trainer scheduler + (the head reads no sampler shard of a shared directory); then every island's own role checks on its island args.""" roles = island_roles(args, num_islands) if not roles: @@ -193,7 +201,11 @@ def check_pair(args: Any, num_islands: int) -> None: raise ValueError("--rl-island-sample-role: a pair gives every --gpu island a role, " f"at least one {SAMPLER} and exactly one {TRAINER} (got {num_islands} " f"island(s), roles {dict(sorted(roles.items()))})") - if values.count(SAMPLER) > 1 and getattr(args, "rl_trainer_scheduler", None) not in (None, "off"): + samplers = values.count(SAMPLER) + if samplers > 1 and int(getattr(args, "total_steps", 0) or 0) % samplers: + raise ValueError(f"--total-steps {args.total_steps}: {samplers} sampler islands share the " + "trainer's rounds; it must be a multiple of the sampler count") + if samplers > 1 and getattr(args, "rl_trainer_scheduler", None) not in (None, "off"): raise ValueError("--rl-trainer-scheduler with several sampler islands is not supported " "(the head does not read the sampler shards); drop the scheduler") if sample_role(args) != BOTH: From 8ffc3816af18451ad55147eea5619f9ffa625487 Mon Sep 17 00:00:00 2001 From: michael <78113554+chowhsu@users.noreply.github.com> Date: Tue, 20 Oct 2026 04:18:24 +0000 Subject: [PATCH 71/89] Leave at once on the reclaim marker from a watcher thread The cloud stops the container seconds after its signal, while the sampler loop may still be inside a round or its throttle sleep. A daemon watcher now writes rl_sampler_leave (reason PREEMPTED, rounds done, the rollout id in flight) and the member record as soon as the marker appears; the exit path leaves only when the watcher has not. Co-Authored-By: Claude Opus 5.5 --- tests/test_rl_spot_sampler_replace.py | 26 +++++++++ yeto/rl/adapters/miles/sampler_island.py | 51 ++++++++--------- yeto/rl/engine/sampler_resume.py | 73 +++++++++++++++++++++++- 3 files changed, 122 insertions(+), 28 deletions(-) diff --git a/tests/test_rl_spot_sampler_replace.py b/tests/test_rl_spot_sampler_replace.py index 0ae3c3125..7e8d2c204 100644 --- a/tests/test_rl_spot_sampler_replace.py +++ b/tests/test_rl_spot_sampler_replace.py @@ -382,3 +382,29 @@ def signal(num, fn): env = {RECLAIM_MARKER_ENV: "/tmp/yeto-reclaim-requested"} assert mr.run_island_script(["x"], env, popen=lambda c, env, **k: None, signal_mod=Sig) == 0 assert seen == [env] + + +def test_the_reclaim_watcher_leaves_once_while_the_loop_is_still_in_a_round(tmp_path): + import threading + + from yeto.rl.engine.sampler_resume import ReclaimWatcher, SamplerMembership + + store = CheckpointStore(tmp_path) + events = [] + membership = SamplerMembership(store, "island0", "b", resume_point([], "island0"), + lambda e, **f: events.append((e, f)), clock=lambda: 7.0) + membership.join(latest_release={"seq": 3, "policy_version": 3}) + flag = threading.Event() + watcher = ReclaimWatcher(flag.is_set, lambda: membership.leave("PREEMPTED", rounds=2, + in_flight_rollout_id=4), + poll_s=0.01).start() + flag.set() + assert watcher.fired.wait(2.0) + watcher.stop() + assert membership.leave("FINISHED") is False # the exit path after the watcher: no second LEAVE + names = [e for e, _ in events] + assert names == ["rl_sampler_join", "rl_sampler_leave"] + leave = events[1][1] + assert (leave["reason"], leave["rounds"], leave["in_flight_rollout_id"]) == ("PREEMPTED", 2, 4) + record = (tmp_path / "sample-buffer/members/island0/b.json").read_text() + assert '"leave_reason": "PREEMPTED"' in record and '"joined_at": 7.0' in record diff --git a/yeto/rl/adapters/miles/sampler_island.py b/yeto/rl/adapters/miles/sampler_island.py index 4289c91b4..3fd9b6092 100644 --- a/yeto/rl/adapters/miles/sampler_island.py +++ b/yeto/rl/adapters/miles/sampler_island.py @@ -261,14 +261,12 @@ def run_sampler_island(miles_args: Any, *, environ: Any = None, from yeto.rl.engine.resume import store_for from yeto.rl.engine.sampler_loop import SamplerLoop from yeto.rl.engine.sampler_resume import ( - JOIN_EVENT, - LEAVE_EVENT, - Member, + ReclaimWatcher, + SamplerMembership, new_incarnation, reclaim_marker, reclaim_requested, resume_point, - write_member, ) from . import LoopRunner @@ -294,9 +292,9 @@ def emit(event: str, **fields: Any) -> None: reclaim_marker(env).unlink(missing_ok=True) # a restarted machine keeps its /tmp buffer = open_buffer(env, sampler=True) resume = resume_point(buffer.entries(), sampler_id) - store = store_for(config.store, environ=environ) - member = Member(sampler=sampler_id, incarnation=incarnation, joined_at=time.time(), - resume=resume.fields()) + membership = SamplerMembership(store_for(config.store, environ=environ), sampler_id, + incarnation, resume, emit) + progress = {"rollout_id": None, "rounds": 0} # the round in flight, rounds done here connect_island_ray(miles_args=miles_args) # The rollout process's metadata hook pushes each round's group records to this @@ -316,20 +314,22 @@ async def init() -> tuple[Any, Any]: return controller, executor async def generate(rollout_id: int) -> None: + progress["rollout_id"] = rollout_id await controller.prepare_rollout(rollout_id) _release(await executor.get(rollout_id)) # the rollout function wrote the groups metadata.discard() # the groups are in the buffer; nobody here reads the records + progress["rounds"] += 1 + progress["rollout_id"] = None - def leave(reason: LeaveReason, rounds: int, **detail: Any) -> None: - member.left_at, member.leave_reason, member.rounds = time.time(), reason.name, rounds - member.policy_version = follower.version - member.detail = detail - emit(LEAVE_EVENT, sampler=sampler_id, incarnation=incarnation, reason=reason.name, - rounds=rounds, policy_version=follower.version, buffer_end=len(buffer), **detail) - write_member(store, member) + def leave(reason: LeaveReason, **detail: Any) -> None: + membership.leave(reason.name, rounds=progress["rounds"], + policy_version=None if follower is None else follower.version, + in_flight_rollout_id=progress["rollout_id"], **detail) error: BaseException | None = None follower = None + reclaimed = lambda: reclaim_requested(env) # noqa: E731 + watcher = ReclaimWatcher(reclaimed, lambda: leave(LeaveReason.PREEMPTED)).start() try: controller, executor = runner.run(init()) lora = lora_mode(miles_args) @@ -338,7 +338,6 @@ def leave(reason: LeaveReason, rounds: int, **detail: Any) -> None: loader.stamp(BASE_POLICY_VERSION) # until the first release is loaded follower = PolicyFollower(release_store(config, environ), loader, version=BASE_POLICY_VERSION) - reclaimed = lambda: reclaim_requested(env) # noqa: E731 loop = SamplerLoop( follower=follower, buffer=buffer, emit=emit, generate=lambda rollout_id: runner.run(generate(rollout_id)), @@ -348,10 +347,8 @@ def leave(reason: LeaveReason, rounds: int, **detail: Any) -> None: loaded_fields=lambda: {"served_checksum": loader.last_checksum, "probe": loader.last_probe}) latest = follower.releases.latest() - emit(JOIN_EVENT, sampler=sampler_id, incarnation=incarnation, - latest_release=None if latest is None else {"seq": latest.seq, "policy_version": latest.version}, - **resume.fields()) - write_member(store, member) + membership.join(latest_release=None if latest is None else + {"seq": latest.seq, "policy_version": latest.version}) emit("rl_sampler_start", store=config.store, num_rollouts=int(miles_args.num_rollout), max_backlog_groups=int(config.max_backlog_groups), weights=loader.weights, start_rollout_id=resume.next_rollout_id) @@ -360,21 +357,21 @@ def leave(reason: LeaveReason, rounds: int, **detail: Any) -> None: else: rounds = loop.run(start=resume.next_rollout_id) if reclaimed(): - leave(LeaveReason.PREEMPTED, rounds) + leave(LeaveReason.PREEMPTED) # no-op when the watcher already left _hold_until_stopped(sleep, LEAVE_HOLD_S) else: - leave(LeaveReason.FINISHED, rounds) + leave(LeaveReason.FINISHED) return rounds except BaseException as exc: error = exc - if follower is not None: - try: - leave(LeaveReason.PREEMPTED if reclaim_requested(env) else LeaveReason.INFRA_FAILED, - -1, error=f"{type(exc).__name__}: {exc}"[:300]) - except Exception: # noqa: BLE001 - the original error is the one to report - pass + try: + leave(LeaveReason.PREEMPTED if reclaimed() else LeaveReason.INFRA_FAILED, + error=f"{type(exc).__name__}: {exc}"[:300]) + except Exception: # noqa: BLE001 - the original error is the one to report + pass raise finally: + watcher.stop() try: runner.run(disposer.__aexit__(None if error is None else type(error), error, None if error is None else error.__traceback__)) diff --git a/yeto/rl/engine/sampler_resume.py b/yeto/rl/engine/sampler_resume.py index e25eefda9..203671aed 100644 --- a/yeto/rl/engine/sampler_resume.py +++ b/yeto/rl/engine/sampler_resume.py @@ -32,8 +32,10 @@ from __future__ import annotations import json +import threading +import time import uuid -from collections.abc import Iterable, Mapping +from collections.abc import Callable, Iterable, Mapping from dataclasses import asdict, dataclass, field from pathlib import Path from typing import Any @@ -183,3 +185,72 @@ def write_member(store: Any, member: Member) -> Path: tmp.replace(path) store.commit() return path + + +class SamplerMembership: + """One incarnation's JOIN and LEAVE: the tape events and the member record. + :meth:`leave` acts once (the reclaim watcher and the island's exit path may + both call it); later calls return False.""" + + def __init__(self, store: Any, sampler: str, incarnation: str, resume: SamplerResume, + emit: Callable[..., None], clock: Callable[[], float] = time.time) -> None: + self.store = store + self.emit = emit + self.clock = clock + self.member = Member(sampler=sampler, incarnation=incarnation, joined_at=clock(), + resume=resume.fields()) + self._lock = threading.Lock() + self.left = False + + def join(self, **fields: Any) -> None: + m = self.member + self.emit(JOIN_EVENT, sampler=m.sampler, incarnation=m.incarnation, **m.resume, **fields) + write_member(self.store, m) + + def leave(self, reason: str, *, rounds: int | None = None, policy_version: int | None = None, + **detail: Any) -> bool: + with self._lock: + if self.left: + return False + self.left = True + m = self.member + m.left_at, m.leave_reason, m.rounds = self.clock(), str(reason), rounds + m.policy_version, m.detail = policy_version, dict(detail) + self.emit(LEAVE_EVENT, sampler=m.sampler, incarnation=m.incarnation, reason=str(reason), + rounds=rounds, policy_version=policy_version, **detail) + write_member(self.store, m) + return True + + +class ReclaimWatcher: + """Polls ``reclaimed()`` on a daemon thread and calls ``on_reclaim()`` once, at + once: the cloud stops the container a few seconds after its signal (Modal: + 30 s), while the island's loop may still be inside a round.""" + + def __init__(self, reclaimed: Callable[[], bool], on_reclaim: Callable[[], Any], *, + poll_s: float = 0.5) -> None: + self.reclaimed = reclaimed + self.on_reclaim = on_reclaim + self.poll_s = float(poll_s) + self._stop = threading.Event() + self.fired = threading.Event() + self._thread = threading.Thread(target=self._run, name="sampler-reclaim-watcher", daemon=True) + + def start(self) -> "ReclaimWatcher": + self._thread.start() + return self + + def stop(self) -> None: + self._stop.set() + self._thread.join(timeout=5.0) + + def _run(self) -> None: + while not self._stop.is_set(): + if self.reclaimed(): + self.fired.set() + try: + self.on_reclaim() + except Exception as error: # noqa: BLE001 - the island loop still stops on the marker + print(f"[sampler-island] reclaim handler failed: {error!r}", flush=True) + return + self._stop.wait(self.poll_s) From e47429a2a9d3ef9d5aac96bfce0067a07927cdaa Mon Sep 17 00:00:00 2001 From: michael <78113554+chowhsu@users.noreply.github.com> Date: Tue, 20 Oct 2026 04:45:24 +0000 Subject: [PATCH 72/89] Turn the reclaim signal into the marker inside the sampler island Modal signals the whole container: in T-C1 run a the sampler process got SIGINT next to the runner, raised KeyboardInterrupt inside its round throttle and left as INFRA_FAILED. A pair sampler island (the launcher sets the marker path) now handles SIGINT/SIGTERM on its main thread by writing the reclaim marker, so the watcher leaves as PREEMPTED and the island waits for the container to stop. Co-Authored-By: Claude Opus 5.5 --- tests/test_rl_spot_sampler_replace.py | 19 ++++++++++++++++ yeto/rl/adapters/miles/sampler_island.py | 4 ++++ yeto/rl/engine/sampler_resume.py | 29 ++++++++++++++++++++++++ 3 files changed, 52 insertions(+) diff --git a/tests/test_rl_spot_sampler_replace.py b/tests/test_rl_spot_sampler_replace.py index 7e8d2c204..184c969a5 100644 --- a/tests/test_rl_spot_sampler_replace.py +++ b/tests/test_rl_spot_sampler_replace.py @@ -408,3 +408,22 @@ def test_the_reclaim_watcher_leaves_once_while_the_loop_is_still_in_a_round(tmp_ assert (leave["reason"], leave["rounds"], leave["in_flight_rollout_id"]) == ("PREEMPTED", 2, 4) record = (tmp_path / "sample-buffer/members/island0/b.json").read_text() assert '"leave_reason": "PREEMPTED"' in record and '"joined_at": 7.0' in record + + +def test_the_reclaim_signal_writes_the_marker_instead_of_raising(tmp_path): + from yeto.rl.engine.sampler_resume import install_marker_signals + + handlers = {} + + class Sig: + SIGINT, SIGTERM = _signal.SIGINT, _signal.SIGTERM + + @staticmethod + def signal(num, fn): + handlers[num] = fn + + env = {RECLAIM_MARKER_ENV: str(tmp_path / "marker")} + assert install_marker_signals(env, signal_mod=Sig) == ["SIGINT", "SIGTERM"] + assert not reclaim_requested(env) + handlers[_signal.SIGINT](_signal.SIGINT, None) # Modal's SIGINT reaches the island process + assert reclaim_requested(env) and '"signal": 2' in (tmp_path / "marker").read_text() diff --git a/yeto/rl/adapters/miles/sampler_island.py b/yeto/rl/adapters/miles/sampler_island.py index 3fd9b6092..0f347e510 100644 --- a/yeto/rl/adapters/miles/sampler_island.py +++ b/yeto/rl/adapters/miles/sampler_island.py @@ -261,8 +261,10 @@ def run_sampler_island(miles_args: Any, *, environ: Any = None, from yeto.rl.engine.resume import store_for from yeto.rl.engine.sampler_loop import SamplerLoop from yeto.rl.engine.sampler_resume import ( + RECLAIM_MARKER_ENV, ReclaimWatcher, SamplerMembership, + install_marker_signals, new_incarnation, reclaim_marker, reclaim_requested, @@ -290,6 +292,8 @@ def emit(event: str, **fields: Any) -> None: raise SampleBufferRolloutError(f"{SAMPLER_ID_ENV} is not set") incarnation = env[INCARNATION_ENV] = new_incarnation() reclaim_marker(env).unlink(missing_ok=True) # a restarted machine keeps its /tmp + if env.get(RECLAIM_MARKER_ENV): # set by the launcher on a pair's sampler island + install_marker_signals(env) buffer = open_buffer(env, sampler=True) resume = resume_point(buffer.entries(), sampler_id) membership = SamplerMembership(store_for(config.store, environ=environ), sampler_id, diff --git a/yeto/rl/engine/sampler_resume.py b/yeto/rl/engine/sampler_resume.py index 203671aed..36ea72635 100644 --- a/yeto/rl/engine/sampler_resume.py +++ b/yeto/rl/engine/sampler_resume.py @@ -154,6 +154,35 @@ def reclaim_marker(environ: Mapping[str, str]) -> Path: return Path(environ.get(RECLAIM_MARKER_ENV) or DEFAULT_RECLAIM_MARKER) +def write_reclaim_marker(environ: Mapping[str, str], signum: int) -> Path: + """Write the reclaim marker (what the Modal runner writes on its signal).""" + path = reclaim_marker(environ) + tmp = path.with_name(path.name + ".tmp") + tmp.write_text(json.dumps({"signal": int(signum), "time": time.time()}) + "\n") + tmp.replace(path) + return path + + +def install_marker_signals(environ: Mapping[str, str], signal_mod: Any = None) -> list[str]: + """On the island's main thread: SIGINT and SIGTERM write the reclaim marker + instead of raising ``KeyboardInterrupt``. Modal signals the whole container, + so the island process gets the reclaim signal next to the runner; without + this it dies inside its round and leaves as INFRA_FAILED (T-C1 run a). + Returns the signals handled ([] off the main thread).""" + import signal as _signal + + sig = signal_mod or _signal + if threading.current_thread() is not threading.main_thread(): + return [] + done = [] + for name in ("SIGINT", "SIGTERM"): + num = getattr(sig, name, None) + if num is not None: + sig.signal(num, lambda signum, _frame: write_reclaim_marker(environ, signum)) + done.append(name) + return done + + def reclaim_requested(environ: Mapping[str, str]) -> bool: """The Modal runner wrote the reclaim marker (the container got SIGINT/SIGTERM).""" return reclaim_marker(environ).exists() From 944c6ee127fc8e98e08bfcebb73184f07185e364 Mon Sep 17 00:00:00 2001 From: michael <78113554+chowhsu@users.noreply.github.com> Date: Tue, 20 Oct 2026 04:49:41 +0000 Subject: [PATCH 73/89] Add --rl-sampler-spot: spot sampler islands, on-demand trainer A pair's sampler islands can now use spot capacity while the trainer stays on demand. The sampler island copy gets spot billing without the legacy --spot volume (its durable state is the sample buffer); the flag is refused outside --rl-island-sample-role and together with --spot. Co-Authored-By: Claude Opus 5.5 --- tests/test_rl_spot_sampler_replace.py | 35 +++++++++++++++++++++++++++ yeto/cloud/droppable.py | 6 +++-- yeto/launch/rl_ports.py | 3 +++ yeto/launch_preflight.py | 11 ++++++--- yeto/rl/engine/sample_roles.py | 7 ++++++ 5 files changed, 57 insertions(+), 5 deletions(-) diff --git a/tests/test_rl_spot_sampler_replace.py b/tests/test_rl_spot_sampler_replace.py index 184c969a5..8e8b4f4f1 100644 --- a/tests/test_rl_spot_sampler_replace.py +++ b/tests/test_rl_spot_sampler_replace.py @@ -427,3 +427,38 @@ def signal(num, fn): assert not reclaim_requested(env) handlers[_signal.SIGINT](_signal.SIGINT, None) # Modal's SIGINT reaches the island process assert reclaim_requested(env) and '"signal": 2' in (tmp_path / "marker").read_text() + + +def test_sampler_spot_bills_only_the_samplers_and_adds_no_spot_volume(monkeypatch): + from yeto.cloud import s3_store + from yeto.gpu_spec import parse_gpu_spec + from yeto.launch.rl_island_task import make_miles_island_task + from yeto.launcher import _island_task, _prepare_rl_args + + monkeypatch.setenv(s3_store.KEY_ID_ENV, "AKIDSTORE") + monkeypatch.setenv(s3_store.SECRET_ENV, "store-secret") + gpu = "nebius:1xh100@eu-north1,nebius:1xh100@eu-north1" + base = ("--rl-single-island-no-sync", "--tuning", "lora", "--rl-island-sample-role", "0:sampler", + "--rl-island-sample-role", "1:trainer", "--rl-max-policy-age", "2", + "--rl-algorithm-spec", "examples/rl_algorithms/decoupled-ppo-c2.json", + "--rl-checkpoint-store", "s3://yeto-ckpt/run-a", "--on-demand", "--rl-sampler-spot") + args = _launcher_args("ports", base, gpu=gpu) + args.controller = "local" + _prepare_rl_args(args) + spots, volumes = [], [] + for m, spec in enumerate(parse_gpu_spec(args.gpu)): + task = _island_task(make_miles_island_task, args, spec, m, 2, "none", {}) + spots.append(task.resources.use_spot) + volumes.append(sorted(getattr(task, "storage_mounts", None) or {})) + # the sampler gets the store mount only, like the trainer: no legacy --spot volume + assert spots == [True, False] and volumes[0] == volumes[1] and len(volumes[0]) == 1 + + +def test_sampler_spot_needs_a_pair(): + from yeto.launcher import _prepare_rl_args + + args = _launcher_args("ports", ("--rl-single-island-no-sync", "--rl-sampler-spot"), + gpu="modal:1xh100") + args.controller = "local" + with pytest.raises(ValueError, match="--rl-sampler-spot needs --rl-island-sample-role"): + _prepare_rl_args(args) diff --git a/yeto/cloud/droppable.py b/yeto/cloud/droppable.py index abc6cb22c..abd1e3759 100644 --- a/yeto/cloud/droppable.py +++ b/yeto/cloud/droppable.py @@ -126,8 +126,10 @@ def role_args(args, role: str | None): def island_keeps_spot_volume(args) -> bool: """The legacy ``--spot`` per-island volume (completed groups, re-synced on a - rebuild) is a resume from an archive; a droppable island must not get it.""" - return bool(getattr(args, "spot", False)) and getattr(args, "_island_role", None) != DROPPABLE + rebuild) is a resume from an archive; a droppable island must not get it, nor + a spot sampler island (``spot_without_volume``: its state is the buffer).""" + return (bool(getattr(args, "spot", False)) and getattr(args, "_island_role", None) != DROPPABLE + and not getattr(args, "spot_without_volume", False)) def role_env(role: str | None, cloud: str, region: str | None) -> dict[str, str]: diff --git a/yeto/launch/rl_ports.py b/yeto/launch/rl_ports.py index bdeceda94..22470ce98 100644 --- a/yeto/launch/rl_ports.py +++ b/yeto/launch/rl_ports.py @@ -380,6 +380,9 @@ def _check_ports_infra_switches(args) -> None: sample_roles.check_pair(args, rl_island_count(args)) else: sample_roles.check_launch(args) # L4.6 --rl-sample-role (both checks nothing) + if getattr(args, "rl_sampler_spot", False) and not sample_roles.is_pair(args): + raise ValueError("--rl-sampler-spot needs --rl-island-sample-role (it bills the sampler " + "islands of a pair as spot)") from .storage import store_object_envs for cloud in {s.cloud for s in parse_gpu_spec(args.gpu)} if getattr(args, "gpu", None) else (): diff --git a/yeto/launch_preflight.py b/yeto/launch_preflight.py index 419432c2e..539b07a07 100644 --- a/yeto/launch_preflight.py +++ b/yeto/launch_preflight.py @@ -363,12 +363,17 @@ def add_cli_args(p) -> None: "base, never from a checkpoint store); repeatable; an island without a role is " "an anchor; replaces --spot") g.add_argument("--rl-island-sample-role", action="append", default=None, metavar="ISLAND:ROLE", - help="rl-trainer-start-stop L5.4: one launch of a sampler island and a trainer " + help="rl-trainer-start-stop L5.4: one launch of sampler islands and a trainer " "island over one --rl-checkpoint-store (ROLE sampler or trainer, one entry per " - "--gpu island; exactly two islands). Needs --rl-single-island-no-sync: no " + "--gpu island; exactly one trainer, one or more samplers; several samplers " + "split --total-steps). Needs --rl-single-island-no-sync: no " "syncer, each island runs as a one-learner no-sync island with that " - "--rl-sample-role. Combine with --rl-trainer-scheduler s1 " + "--rl-sample-role. With one sampler, combine with --rl-trainer-scheduler s1 " "--rl-trainer-island M to start and stop the trainer island") + g.add_argument("--rl-sampler-spot", action="store_true", + help="with --rl-island-sample-role: the sampler islands use spot (preemptible) " + "capacity, the trainer island stays on demand; a reclaimed sampler island is " + "started again and continues after its last round in the buffer") g.add_argument("--rl-negative-test-run", action="store_true", help="mark this run as a negative-test run (required by --rl-island-override); " "a normal run refuses to resume its checkpoint store and `yeto merge` refuses " diff --git a/yeto/rl/engine/sample_roles.py b/yeto/rl/engine/sample_roles.py index 69f72d7c8..072a3e3ab 100644 --- a/yeto/rl/engine/sample_roles.py +++ b/yeto/rl/engine/sample_roles.py @@ -152,6 +152,11 @@ def island_args(args: Any, island: int, num_islands: int) -> Any: out.rl_sampler_node = node_name(island) if out.rl_sample_role == SAMPLER else None out.rl_buffer_shards = buffer_shards(args, num_islands) samplers = len(sampler_islands(args, num_islands)) + if out.rl_sample_role == SAMPLER and getattr(args, "rl_sampler_spot", False): + # spot billing without the legacy --spot volume: the buffer is the sampler's + # durable state (S25 C1) + out.spot = True + out.spot_without_volume = True if out.rl_sample_role == SAMPLER and samplers > 1: # the trainer trains --total-steps rounds; the samplers share them, so each # one reaches its data end (check_pair needs an even split) @@ -205,6 +210,8 @@ def check_pair(args: Any, num_islands: int) -> None: if samplers > 1 and int(getattr(args, "total_steps", 0) or 0) % samplers: raise ValueError(f"--total-steps {args.total_steps}: {samplers} sampler islands share the " "trainer's rounds; it must be a multiple of the sampler count") + if getattr(args, "rl_sampler_spot", False) and getattr(args, "spot", False): + raise ValueError("--rl-sampler-spot bills only the samplers as spot; it does not take --spot") if samplers > 1 and getattr(args, "rl_trainer_scheduler", None) not in (None, "off"): raise ValueError("--rl-trainer-scheduler with several sampler islands is not supported " "(the head does not read the sampler shards); drop the scheduler") From 31af37f10f87cc5d47a3abb28730500d178d74b4 Mon Sep 17 00:00:00 2001 From: michael <78113554+chowhsu@users.noreply.github.com> Date: Tue, 20 Oct 2026 05:08:50 +0000 Subject: [PATCH 74/89] Never raise from the sampler's reclaim signal handler In T-C1 run b the Modal runner and the sampler's handler wrote the marker through the same temp name at the same moment; the rename failed and the handler raised inside the island's main thread. The temp name is now per process and a failed write is printed, not raised. Co-Authored-By: Claude Opus 5.5 --- tests/test_rl_spot_sampler_replace.py | 8 ++++++++ yeto/rl/engine/sampler_resume.py | 15 +++++++++++---- 2 files changed, 19 insertions(+), 4 deletions(-) diff --git a/tests/test_rl_spot_sampler_replace.py b/tests/test_rl_spot_sampler_replace.py index 8e8b4f4f1..18e99c2d8 100644 --- a/tests/test_rl_spot_sampler_replace.py +++ b/tests/test_rl_spot_sampler_replace.py @@ -462,3 +462,11 @@ def test_sampler_spot_needs_a_pair(): args.controller = "local" with pytest.raises(ValueError, match="--rl-sampler-spot needs --rl-island-sample-role"): _prepare_rl_args(args) + + +def test_a_failed_marker_write_in_the_signal_handler_does_not_raise(tmp_path, capsys): + from yeto.rl.engine.sampler_resume import write_reclaim_marker + + env = {RECLAIM_MARKER_ENV: str(tmp_path / "missing-dir" / "marker")} + write_reclaim_marker(env, _signal.SIGINT) # T-C1 run b: the runner raced the same temp name + assert "reclaim marker not written" in capsys.readouterr().out diff --git a/yeto/rl/engine/sampler_resume.py b/yeto/rl/engine/sampler_resume.py index 36ea72635..50b947f0d 100644 --- a/yeto/rl/engine/sampler_resume.py +++ b/yeto/rl/engine/sampler_resume.py @@ -32,6 +32,7 @@ from __future__ import annotations import json +import os import threading import time import uuid @@ -155,11 +156,17 @@ def reclaim_marker(environ: Mapping[str, str]) -> Path: def write_reclaim_marker(environ: Mapping[str, str], signum: int) -> Path: - """Write the reclaim marker (what the Modal runner writes on its signal).""" + """Write the reclaim marker (what the Modal runner writes on its signal). The + runner may write it at the same moment: the temp file is per process, and a + failed write is reported, never raised (it runs in a signal handler; the + runner's marker serves as well).""" path = reclaim_marker(environ) - tmp = path.with_name(path.name + ".tmp") - tmp.write_text(json.dumps({"signal": int(signum), "time": time.time()}) + "\n") - tmp.replace(path) + tmp = path.with_name(f"{path.name}.{os.getpid()}.tmp") + try: + tmp.write_text(json.dumps({"signal": int(signum), "time": time.time()}) + "\n") + tmp.replace(path) + except OSError as error: + print(f"[sampler-island] reclaim marker not written ({error})", flush=True) return path From 00064e40ef16f6a79680b1902d52842ecef9379d Mon Sep 17 00:00:00 2001 From: michael <78113554+chowhsu@users.noreply.github.com> Date: Tue, 20 Oct 2026 11:04:00 +0000 Subject: [PATCH 75/89] Agentic outcome: a conversation that outgrows the context is a trained 0; shared records env names The records backend marks the last sampled turn when the next prompt no longer fits max_model_len (context_full); with the agent's non-zero exit the outer layer re-classifies the trajectory as POLICY_LIMIT, reward 0, as for a malformed tool call (main agent decision). The records-proxy env names move to gateway.records_env so the Miles adapter no longer imports the verl one (import boundary test). Co-Authored-By: Claude Opus 5.5 --- tests/test_verl_rollout_engine.py | 33 ++++++++++++++++++++++ yeto/rl/adapters/miles/records_generate.py | 4 +-- yeto/rl/adapters/verl/rollout_engine.py | 14 +++++---- yeto/rl/agentic/outcome.py | 5 ++-- yeto/rl/agentic/types.py | 7 +++-- yeto/rl/harness/gateway/records.py | 14 +++++++-- yeto/rl/harness/gateway/records_env.py | 12 ++++++++ 7 files changed, 73 insertions(+), 16 deletions(-) create mode 100644 yeto/rl/harness/gateway/records_env.py diff --git a/tests/test_verl_rollout_engine.py b/tests/test_verl_rollout_engine.py index 89a60a94d..45c64911e 100644 --- a/tests/test_verl_rollout_engine.py +++ b/tests/test_verl_rollout_engine.py @@ -206,3 +206,36 @@ async def go(): assert sum(out.response_mask) > 0 else: assert sum(out.response_mask) == 0 # another exit cause stays an infrastructure abort + + +def test_context_full_is_a_trained_zero_not_an_abort(verl_loop, qwen): + """The conversation outgrows the context: the proxy answers 409, Codex exits; the policy's budget (N10.9).""" + from yeto.rl.adapters.verl.rollout_engine import VerlRolloutEngine + + loop_mod, _ = verl_loop + + def run(max_model_len: int, client: FakeVerlClient): + engine = VerlRolloutEngine(max_model_len=max_model_len, identity=IDENTITY) + recorder = ProxyTokenRecorder(RecordsProxy(engine, qwen, model="policy", signed_sampling=SIGNED, + max_chains=1)) + agent = ExitOnErrorAgent() + loop = _make_loop(loop_mod, recorder, agent, prompt_length=4096, response_length=4096) + loop.server_manager = client + + async def go(): + try: + return await loop.run(VERL_SAMPLING, **ROW) + finally: + await recorder.aclose() + + return asyncio.run(go()), agent + + probe = FakeVerlClient(ScriptedModel(qwen, copy.deepcopy(SCRIPT))) + run(8192, probe) + second = probe.calls[1]["prompt_ids"] + # room for the second turn, not for the third prompt (which adds the tool output) + limit = len(second) + 40 + out, agent = run(limit, FakeVerlClient(ScriptedModel(qwen, copy.deepcopy(SCRIPT)))) + assert 409 in agent.statuses + assert out.extra_fields[OUTCOME_KEY] == "policy_limit" and out.reward_score == 0.0 + assert "context_full" in out.extra_fields["outcome_reason"] and sum(out.response_mask) > 0 diff --git a/yeto/rl/adapters/miles/records_generate.py b/yeto/rl/adapters/miles/records_generate.py index 9b157f989..354024cab 100644 --- a/yeto/rl/adapters/miles/records_generate.py +++ b/yeto/rl/adapters/miles/records_generate.py @@ -26,7 +26,7 @@ YETO_RECORDS_PROXY_HOST/PORT/PUBLIC_BASE/SECRET_FILE (``gateway.recorder``) YETO_RECORDS_TEMPLATE, YETO_RECORDS_TOKENIZER, YETO_RECORDS_MODEL, YETO_RECORDS_SAMPLING, YETO_RECORDS_MAX_MODEL_LEN, YETO_RECORDS_CARD - (same meaning as ``adapters.verl.rollout_engine``) + (``gateway.records_env``) YETO_SGLANG_ROUTER_URL ``http://:`` of Miles' router (default: from ``--sglang-router-ip`` / ``--sglang-router-port``) @@ -164,7 +164,7 @@ def stop(self) -> None: def records_proxy_from_env() -> Any: """``YETO_RECORDS_PROXY`` callable: a RecordsProxy over Miles' router (module doc).""" - from yeto.rl.adapters.verl.rollout_engine import ( + from yeto.rl.harness.gateway.records_env import ( CARD_ENV, MAX_MODEL_LEN_ENV, MODEL_ENV, diff --git a/yeto/rl/adapters/verl/rollout_engine.py b/yeto/rl/adapters/verl/rollout_engine.py index 68bdfbf1b..93db3586d 100644 --- a/yeto/rl/adapters/verl/rollout_engine.py +++ b/yeto/rl/adapters/verl/rollout_engine.py @@ -60,12 +60,14 @@ SamplingParams, ) -TEMPLATE_ENV = "YETO_RECORDS_TEMPLATE" -TOKENIZER_ENV = "YETO_RECORDS_TOKENIZER" -MODEL_ENV = "YETO_RECORDS_MODEL" -SAMPLING_ENV = "YETO_RECORDS_SAMPLING" -MAX_MODEL_LEN_ENV = "YETO_RECORDS_MAX_MODEL_LEN" -CARD_ENV = "YETO_RECORDS_CARD" +from yeto.rl.harness.gateway.records_env import ( + CARD_ENV, + MAX_MODEL_LEN_ENV, + MODEL_ENV, + SAMPLING_ENV, + TEMPLATE_ENV, + TOKENIZER_ENV, +) CALL_TIMEOUT_S = 1800.0 diff --git a/yeto/rl/agentic/outcome.py b/yeto/rl/agentic/outcome.py index ea6b2c24d..0f9d87ca3 100644 --- a/yeto/rl/agentic/outcome.py +++ b/yeto/rl/agentic/outcome.py @@ -42,8 +42,9 @@ def policy_output_failure(trajectory: AgentTrajectory, policy_errors: Sequence[str]) -> AgentTrajectory: """ABORTED by the agent's exit after the recorder refused a malformed model reply -> POLICY_LIMIT, reward 0. - S25 N10.9 (main agent decision): a tool call that is not valid JSON is the - policy's own failure, trained as a 0 reward, not an infrastructure abort. + S25 N10.9 (main agent decisions): a tool call that is not valid JSON, and + a conversation that outgrew the context (the policy's own length budget), + are the policy's failures, trained as a 0 reward, not infrastructure aborts. Fail closed: only the agent-exit exception, and only with a recorded refusal, is re-classified; every other ABORTED stays ABORTED. """ diff --git a/yeto/rl/agentic/types.py b/yeto/rl/agentic/types.py index 2b5bd50f7..e6d425f44 100644 --- a/yeto/rl/agentic/types.py +++ b/yeto/rl/agentic/types.py @@ -124,9 +124,10 @@ class TokenTurn: ``Provenance.to_record()``) are what the recorder knows about the call; None when it does not record them. ``finish_reason``: the engine's reason (``stop`` / ``length``; ``length`` = the call hit its token cap). - ``policy_error``: why the recorder refused the reply as malformed model - output (e.g. a tool call that is not valid JSON); the tokens were sampled - and are kept, the agent got an error for that call. + ``policy_error``: the policy's own failure the recorder saw after this + call: a malformed reply (e.g. a tool call that is not valid JSON) or a + conversation that no longer fits the context (``context_full: ...``). + The tokens were sampled and are kept; the agent got an error. """ prompt_ids: tuple[int, ...] diff --git a/yeto/rl/harness/gateway/records.py b/yeto/rl/harness/gateway/records.py index 8ceae491e..0a42f789c 100644 --- a/yeto/rl/harness/gateway/records.py +++ b/yeto/rl/harness/gateway/records.py @@ -85,8 +85,10 @@ class TurnRecord: engine returned it). ``stop_added``: a stop token the engine did not return, added with mask 0. ``render_diff``: first position where the stream differs from a fresh tokenization (None: aligned). ``parse_error``: - the completion is not valid model output for the template (the agent got - an HTTP 400 for this call; the sampled ids stay in the stream).""" + the policy failed after this completion: the completion is not valid + model output for the template (HTTP 400), or the next prompt no longer + fits the context (``context_full: ...``, HTTP 409). The sampled ids stay + in the stream.""" turn: int kind: str @@ -303,7 +305,13 @@ async def _sample(self, session: _Session, segment: Segment, pending: _Prompt, p prompt = [*segment.token_ids, *pending.stop, *pending.tail] room = int(self.engine.max_model_len) - len(prompt) if room < 1: - raise ContextFull(f"prompt has {len(prompt)} tokens; max_model_len is {self.engine.max_model_len}") + message = f"prompt has {len(prompt)} tokens; max_model_len is {self.engine.max_model_len}" + # The policy filled its own context (a budget, like a token cap): mark it on the last + # sampled turn so the outer layer trains it as a 0 instead of an infrastructure abort. + if segment.turns: + segment.turns[-1].parse_error = f"context_full: {message}" + self.counters[session.trajectory_id]["context_full"] += 1 + raise ContextFull(message) params = params.replace(max_tokens=min(params.max_tokens, room)) version = self.engine.served_version() (group,) = await asyncio.to_thread(self.engine.generate_tokens, [prompt], params, diff --git a/yeto/rl/harness/gateway/records_env.py b/yeto/rl/harness/gateway/records_env.py new file mode 100644 index 000000000..6f434e545 --- /dev/null +++ b/yeto/rl/harness/gateway/records_env.py @@ -0,0 +1,12 @@ +"""Environment names of a records proxy built inside a trainer's rollout process (S25 N10.9). + +Shared by the backend adapters (verl ``rollout_engine``, Miles +``records_generate``); kept here so adapters do not import each other. +""" + +TEMPLATE_ENV = "YETO_RECORDS_TEMPLATE" # chat template name (``chat_templates.TEMPLATES``) +TOKENIZER_ENV = "YETO_RECORDS_TOKENIZER" # tokenizer directory (the checkpoint) +MODEL_ENV = "YETO_RECORDS_MODEL" # model name the agent sends +SAMPLING_ENV = "YETO_RECORDS_SAMPLING" # JSON signed sampling +MAX_MODEL_LEN_ENV = "YETO_RECORDS_MAX_MODEL_LEN" # context length of the rollout engine +CARD_ENV = "YETO_RECORDS_CARD" # GPU name (provenance) From d9330ea41a08f3d474c6ca425aac9c63f25893cb Mon Sep 17 00:00:00 2001 From: michael <78113554+chowhsu@users.noreply.github.com> Date: Tue, 20 Oct 2026 11:11:09 +0000 Subject: [PATCH 76/89] Nebius spot sampler island: host unit forwards the VM stop, member records by object A spot sampler island on Nebius gets the yeto-reclaim-notice host unit: its notice file is the sampler's reclaim marker, and the sampler writes the ack file once its LEAVE is stored, so the VM stops only after the LEAVE. Member records go through the store's object API when it has one (the staged s3/OSS mirror does not carry sample-buffer/, so the record stayed on the island disk). Co-Authored-By: Claude Opus 5.5 --- tests/test_rl_spot_sampler_replace.py | 45 +++++++++++++++++++++++- yeto/cloud/nebius_reclaim.py | 19 ++++++++-- yeto/launcher.py | 15 +++++++- yeto/rl/adapters/miles/sampler_island.py | 14 +++++--- yeto/rl/engine/sampler_resume.py | 43 ++++++++++++++++++++-- 5 files changed, 125 insertions(+), 11 deletions(-) diff --git a/tests/test_rl_spot_sampler_replace.py b/tests/test_rl_spot_sampler_replace.py index 18e99c2d8..23ddfabf8 100644 --- a/tests/test_rl_spot_sampler_replace.py +++ b/tests/test_rl_spot_sampler_replace.py @@ -445,13 +445,56 @@ def test_sampler_spot_bills_only_the_samplers_and_adds_no_spot_volume(monkeypatc args = _launcher_args("ports", base, gpu=gpu) args.controller = "local" _prepare_rl_args(args) - spots, volumes = [], [] + from yeto.cloud import nebius_reclaim + + spots, volumes, hooks, markers = [], [], [], [] for m, spec in enumerate(parse_gpu_spec(args.gpu)): task = _island_task(make_miles_island_task, args, spec, m, 2, "none", {}) spots.append(task.resources.use_spot) volumes.append(sorted(getattr(task, "storage_mounts", None) or {})) + hooks.append(nebius_reclaim.wants_host_hook(task.envs)) + markers.append(task.envs.get(RECLAIM_MARKER_ENV)) # the sampler gets the store mount only, like the trainer: no legacy --spot volume assert spots == [True, False] and volumes[0] == volumes[1] and len(volumes[0]) == 1 + # Nebius stops the VM without signalling the container: the host unit writes the + # sampler's marker (S25 C1b); the on-demand trainer gets no unit + assert hooks == [True, False] and markers == [nebius_reclaim.NOTICE_FILE, None] + + +def test_an_on_demand_sampler_on_nebius_gets_no_host_unit(): + from yeto.launcher import _spot_sampler_cloud_env + + sampler = SimpleNamespace(rl_sample_role="sampler", rl_sampler_node="island0", spot=False) + assert _spot_sampler_cloud_env(sampler, "nebius") == {} + sampler.spot = True + assert _spot_sampler_cloud_env(sampler, "modal") == {} # Modal signals the container + assert _spot_sampler_cloud_env(sampler, "nebius") + + +def test_the_reclaim_ack_is_written_only_where_a_host_hook_waits(tmp_path): + from yeto.rl.engine.sampler_resume import RECLAIM_ACK_ENV, acknowledge_reclaim, clear_reclaim_files + + assert acknowledge_reclaim({}) is None + env = {RECLAIM_MARKER_ENV: str(tmp_path / "notice"), RECLAIM_ACK_ENV: str(tmp_path / "ack")} + assert acknowledge_reclaim(env) == tmp_path / "ack" and (tmp_path / "ack").exists() + (tmp_path / "notice").write_text("1\n") + clear_reclaim_files(env) # a restarted VM keeps /tmp: the next incarnation starts clean + assert not (tmp_path / "notice").exists() and not (tmp_path / "ack").exists() + + +def test_member_records_go_by_object_on_an_object_store(tmp_path): + from yeto.rl.engine.object_store import DirectoryObjectStore + + class StagedStore(CheckpointStore): # an s3/OSS store: its stage does not carry sample-buffer/ + def __init__(self, root, objects): + super().__init__(root) + self.objects = objects + + objects = DirectoryObjectStore(tmp_path / "bucket") + store = StagedStore(tmp_path / "stage", objects) + write_member(store, Member("island0", "a", 1.0, {"next_rollout_id": 0}, leave_reason="PREEMPTED")) + assert '"leave_reason": "PREEMPTED"' in objects.read("sample-buffer/members/island0/a.json").decode() + assert not (tmp_path / "stage" / "sample-buffer").exists() def test_sampler_spot_needs_a_pair(): diff --git a/yeto/cloud/nebius_reclaim.py b/yeto/cloud/nebius_reclaim.py index 31d8b6b96..56e5f8cd5 100644 --- a/yeto/cloud/nebius_reclaim.py +++ b/yeto/cloud/nebius_reclaim.py @@ -106,11 +106,26 @@ def install_command(container: str = CONTAINER, notice: str = NOTICE_FILE, ack: ) +def spot_sampler_env() -> dict[str, str]: + """The environment of a spot sampler island on Nebius (S25 C1b): the sampler's + reclaim marker is the notice file the host unit writes, and the sampler + writes the ack file once its LEAVE is stored. The two env names of this + module also tell :func:`wants_host_hook` to install the unit.""" + from yeto.rl.engine.sampler_resume import RECLAIM_ACK_ENV, RECLAIM_MARKER_ENV + + assert RECLAIM_ACK_ENV == ACK_ENV + return {RECLAIM_MARKER_ENV: NOTICE_FILE, NOTICE_ENV: NOTICE_FILE, ACK_ENV: ACK_FILE} + + def wants_host_hook(task_envs) -> bool: - """True for a droppable island on Nebius (from the task's YETO_ISLAND_ROLE).""" + """True for a droppable island on Nebius (from the task's YETO_ISLAND_ROLE) and + for a Nebius spot sampler island (:func:`spot_sampler_env`).""" from yeto.cloud import droppable - doc = droppable.env_role(dict(task_envs or {})) + envs = dict(task_envs or {}) + if envs.get(NOTICE_ENV) == NOTICE_FILE and envs.get(ACK_ENV) == ACK_FILE: + return True + doc = droppable.env_role(envs) return bool(doc) and doc.get("role") == droppable.DROPPABLE and doc.get("cloud") == "nebius" diff --git a/yeto/launcher.py b/yeto/launcher.py index 38b5ac140..89d37836f 100644 --- a/yeto/launcher.py +++ b/yeto/launcher.py @@ -188,10 +188,23 @@ def _island_task(task_factory, args, spec, m, num_learners, syncer_addr, overrid task = task_factory(iargs, spec, learner_id, island_learners, syncer_addr) io.apply_to_task(task, {**io.island_env(args, m, overrides), **(droppable.role_env(role, spec.cloud, spec.region) if role else {}), - **sample_roles.island_env(iargs)}) + **sample_roles.island_env(iargs), + **_spot_sampler_cloud_env(iargs, spec.cloud)}) return task +def _spot_sampler_cloud_env(iargs, cloud: str) -> dict[str, str]: + """A spot sampler island on Nebius: the VM stop reaches the sampler through the + host unit of ``yeto.cloud.nebius_reclaim`` (Modal signals the container itself).""" + from .rl.engine import sample_roles + + if cloud != "nebius" or not getattr(iargs, "spot", False) or not sample_roles.island_env(iargs): + return {} + from .cloud import nebius_reclaim + + return nebius_reclaim.spot_sampler_env() + + def _ssh_trainer_spec(args, spec: ClusterSpec, task): """The SSH spec of a held-back ``ssh:`` trainer island (None on other clouds).""" if spec.cloud != "ssh": diff --git a/yeto/rl/adapters/miles/sampler_island.py b/yeto/rl/adapters/miles/sampler_island.py index d91d1732e..22da4c7d9 100644 --- a/yeto/rl/adapters/miles/sampler_island.py +++ b/yeto/rl/adapters/miles/sampler_island.py @@ -273,9 +273,10 @@ def run_sampler_island(miles_args: Any, *, environ: Any = None, RECLAIM_MARKER_ENV, ReclaimWatcher, SamplerMembership, + acknowledge_reclaim, + clear_reclaim_files, install_marker_signals, new_incarnation, - reclaim_marker, reclaim_requested, resume_point, ) @@ -300,7 +301,7 @@ def emit(event: str, **fields: Any) -> None: if not sampler_id: raise SampleBufferRolloutError(f"{SAMPLER_ID_ENV} is not set") incarnation = env[INCARNATION_ENV] = new_incarnation() - reclaim_marker(env).unlink(missing_ok=True) # a restarted machine keeps its /tmp + clear_reclaim_files(env) # a restarted machine keeps its /tmp if env.get(RECLAIM_MARKER_ENV): # set by the launcher on a pair's sampler island install_marker_signals(env) buffer = open_buffer(env, sampler=True) @@ -342,7 +343,12 @@ def leave(reason: LeaveReason, **detail: Any) -> None: error: BaseException | None = None follower = None reclaimed = lambda: reclaim_requested(env) # noqa: E731 - watcher = ReclaimWatcher(reclaimed, lambda: leave(LeaveReason.PREEMPTED)).start() + + def leave_reclaimed() -> None: + leave(LeaveReason.PREEMPTED) + acknowledge_reclaim(env) # a host hook (Nebius) lets the VM stop once the LEAVE is stored + + watcher = ReclaimWatcher(reclaimed, leave_reclaimed).start() try: controller, executor = runner.run(init()) lora = lora_mode(miles_args) @@ -370,7 +376,7 @@ def leave(reason: LeaveReason, **detail: Any) -> None: else: rounds = loop.run(start=resume.next_rollout_id) if reclaimed(): - leave(LeaveReason.PREEMPTED) # no-op when the watcher already left + leave_reclaimed() # the LEAVE is a no-op when the watcher already left _hold_until_stopped(sleep, LEAVE_HOLD_S) else: leave(LeaveReason.FINISHED) diff --git a/yeto/rl/engine/sampler_resume.py b/yeto/rl/engine/sampler_resume.py index 50b947f0d..234b06fc3 100644 --- a/yeto/rl/engine/sampler_resume.py +++ b/yeto/rl/engine/sampler_resume.py @@ -47,6 +47,9 @@ MEMBERS_DIR = "sample-buffer/members" # one record per sampler incarnation, under the store root RECLAIM_MARKER_ENV = "YETO_SPOT_RECLAIM_MARKER" # written by the Modal runner on the reclaim signal DEFAULT_RECLAIM_MARKER = "/tmp/yeto-reclaim-requested" +# Set where a host hook waits for the sampler's answer before the VM stops +# (Nebius: yeto.cloud.nebius_reclaim writes the marker and waits for this file). +RECLAIM_ACK_ENV = "YETO_RECLAIM_ACK_FILE" class SamplerResumeError(RuntimeError): @@ -190,6 +193,29 @@ def install_marker_signals(environ: Mapping[str, str], signal_mod: Any = None) - return done +def clear_reclaim_files(environ: Mapping[str, str]) -> None: + """Remove a reclaim marker and ack left from an earlier boot (a restarted VM keeps /tmp).""" + reclaim_marker(environ).unlink(missing_ok=True) + ack = environ.get(RECLAIM_ACK_ENV) + if ack: + Path(ack).unlink(missing_ok=True) + + +def acknowledge_reclaim(environ: Mapping[str, str]) -> Path | None: + """Write the ack file a host hook waits for (``RECLAIM_ACK_ENV``) once the + LEAVE is durable, so the host lets the VM stop; None when no hook waits.""" + ack = environ.get(RECLAIM_ACK_ENV) + if not ack: + return None + path = Path(ack) + try: + path.write_text(f"{time.time():.3f}\n") + except OSError as error: + print(f"[sampler-island] reclaim ack not written ({error})", flush=True) + return None + return path + + def reclaim_requested(environ: Mapping[str, str]) -> bool: """The Modal runner wrote the reclaim marker (the container got SIGINT/SIGTERM).""" return reclaim_marker(environ).exists() @@ -210,14 +236,25 @@ class Member: detail: dict[str, Any] = field(default_factory=dict) +def member_key(member: Member) -> str: + return f"{MEMBERS_DIR}/{member.sampler}/{member.incarnation}.json" + + def write_member(store: Any, member: Member) -> Path: """Write (or rewrite) the incarnation's member record and commit the store. One file per incarnation: an old and a new incarnation never write the same - file (a Modal Volume keeps the last commit of a file).""" - path = Path(store.root) / MEMBERS_DIR / member.sampler / f"{member.incarnation}.json" + file (a Modal Volume keeps the last commit of a file). A store with an object + API (s3, OSS) gets the record by object: its staged mirror does not carry + ``sample-buffer/``, so a file would stay on the island's disk.""" + data = (json.dumps(asdict(member), sort_keys=True) + "\n").encode() + objects = getattr(store, "objects", None) + if objects is not None: + objects.write(member_key(member), data) + return Path(member_key(member)) + path = Path(store.root) / member_key(member) path.parent.mkdir(parents=True, exist_ok=True) tmp = path.with_suffix(".json.tmp") - tmp.write_text(json.dumps(asdict(member), sort_keys=True) + "\n") + tmp.write_bytes(data) tmp.replace(path) store.commit() return path From a4b9962dd3e5e39cbb34dc92a2ef514ac3ee93d9 Mon Sep 17 00:00:00 2001 From: michael <78113554+chowhsu@users.noreply.github.com> Date: Tue, 20 Oct 2026 11:23:35 +0000 Subject: [PATCH 77/89] Resume a replaced engine sampler island; never overwrite a payload T-G1B g1b-1 lost the engine sampler's container; Modal ran it again and the new container started at round 0, wrote payloads under the same names (r0-row0, ...) and overwrote the first container's buffer payloads and tape. The verl buffer path's engine sampler island has the same gap. - The engine sampler island reuses the neutral sampler_resume layer (C1): each process is one incarnation, reads the resume point from the buffer, starts after the newest written round, refuses written rounds (RoundGuard), writes rl_sampler_join / rl_sampler_leave and a member record per incarnation, and stops writing on the reclaim marker. - Sample ids are --r-, and the buffer refuses a payload name that is already stored. - Each incarnation writes its own tape (events..jsonl). Co-Authored-By: Claude Opus 5.5 --- tests/test_rl_spot_sampler_replace.py | 2 +- tests/test_s25_g1c_sampler_resume.py | 152 +++++++++++++++++++++++++ tests/test_stage1_gaps.py | 5 +- yeto/rl/engine/engine_sampler.py | 76 +++++++++++-- yeto/rl/engine/engine_sampler_entry.py | 92 +++++++++++++-- yeto/rl/engine/sample_buffer.py | 10 +- 6 files changed, 312 insertions(+), 25 deletions(-) create mode 100644 tests/test_s25_g1c_sampler_resume.py diff --git a/tests/test_rl_spot_sampler_replace.py b/tests/test_rl_spot_sampler_replace.py index 18e99c2d8..5826db51b 100644 --- a/tests/test_rl_spot_sampler_replace.py +++ b/tests/test_rl_spot_sampler_replace.py @@ -149,7 +149,7 @@ def test_a_writer_without_resume_writes_as_before(tmp_path): writer = br.SamplerRoundWriter(FileSampleBuffer(store, batch_size=1), "n0", incarnation="a", resume=False) writer.write_round(0, _round(0, 1), None) - writer.write_round(0, _round(0, 1), None) + writer.write_round(0, [("r0-again", _round(0, 1)[0][1])], None) # payload names stay unique assert writer.seek(object(), shuffle=True) is None assert len(FileSampleBuffer(store, batch_size=1)) == 2 diff --git a/tests/test_s25_g1c_sampler_resume.py b/tests/test_s25_g1c_sampler_resume.py new file mode 100644 index 000000000..c8995bcd2 --- /dev/null +++ b/tests/test_s25_g1c_sampler_resume.py @@ -0,0 +1,152 @@ +"""S25 T-G1C: a replaced engine sampler island continues where the buffer says +and never overwrites a payload or an earlier incarnation's tape (CPU).""" + +from __future__ import annotations + +import json +import os +import signal +from pathlib import Path + +import pytest + +from tests.test_stage1_gaps import _sampler_spec +from yeto.rl.engine import engine_sampler_entry as entry +from yeto.rl.engine.resume import CheckpointStore +from yeto.rl.engine.sample_buffer import FileSampleBuffer, SampleBufferError +from yeto.rl.engine.sampler_resume import MEMBERS_DIR, RECLAIM_MARKER_ENV, SamplerResumeError + +STATE_ENV = "YETO_TEST_G1C_REWARD_STATE" # the reward is imported by name: share state via a file + + +def counting_reward(item, completion) -> float: + """Fails (a lost container) or writes the reclaim marker after the limits in the state file.""" + path = Path(os.environ[STATE_ENV]) + state = json.loads(path.read_text()) + state["n"] += 1 + path.write_text(json.dumps(state)) + if state.get("limit") is not None and state["n"] > state["limit"]: + raise RuntimeError("container lost") + if state.get("mark_at") is not None and state["n"] == state["mark_at"]: + Path(state["marker"]).write_text("{}") + return 0.5 + + +class _Calls: + def __init__(self, path: Path) -> None: + self.path = path + self.update(n=0) + + def update(self, **fields) -> None: + state = json.loads(self.path.read_text()) if self.path.exists() else {} + self.path.write_text(json.dumps({**state, **fields})) + + +@pytest.fixture(autouse=True) +def _reset(monkeypatch, tmp_path): + global CALLS + monkeypatch.setenv(STATE_ENV, str(tmp_path / "reward-state.json")) + CALLS = _Calls(tmp_path / "reward-state.json") + monkeypatch.setenv(RECLAIM_MARKER_ENV, str(tmp_path / "reclaim-marker")) + handlers = {s: signal.getsignal(s) for s in (signal.SIGINT, signal.SIGTERM)} + yield + for sig, handler in handlers.items(): + signal.signal(sig, handler) + + +def _spec(tmp_path: Path, rounds: int = 4) -> dict: + return _sampler_spec(tmp_path, num_rollouts=rounds, reward="tests.test_s25_g1c_sampler_resume:counting_reward") + + +def _buffer(tmp_path: Path) -> FileSampleBuffer: + buffer = FileSampleBuffer(CheckpointStore(tmp_path / "store"), batch_size=1) + buffer.refresh() + return buffer + + +def _payloads(tmp_path: Path) -> dict[str, bytes]: + return {p.name: p.read_bytes() for p in (tmp_path / "store" / "sample-buffer" / "payloads").iterdir()} + + +def _tape(path: Path) -> list[dict]: + return [json.loads(line) for line in path.read_text().splitlines()] + + +def test_replacement_continues_after_the_written_rounds_and_overwrites_nothing(tmp_path): + spec = _spec(tmp_path) + CALLS.update(limit=8) # 2 prompts x n=2 per round: the third round fails before it writes + with pytest.raises(RuntimeError, match="container lost"): + entry.run(spec, resource_sampler=lambda emit: None, incarnation="aaaa0001") + first = _payloads(tmp_path) + assert len(first) == 4 and all(name.startswith("ppu4-aaaa0001-r") for name in first) + + CALLS.update(n=0, limit=None) + assert entry.run(spec, resource_sampler=lambda emit: None, incarnation="bbbb0002") == 2 + after = _payloads(tmp_path) + assert {k: after[k] for k in first} == first # the first incarnation's payloads are untouched + assert len(after) == 8 + buffer = _buffer(tmp_path) + entries = buffer.entries() + assert buffer.data_ended and len({e.sample_id for e in entries}) == len(entries) == 8 + rounds = [(e.sampler_round["incarnation"], e.sampler_round["rollout_id"]) for e in entries] + assert rounds == [("aaaa0001", 0)] * 2 + [("aaaa0001", 1)] * 2 + [("bbbb0002", 2)] * 2 \ + + [("bbbb0002", 3)] * 2 + assert all(buffer.read_payload(e.payload_ref) for e in entries) + + tapes = {p.name: _tape(p) for p in (tmp_path / "tape").glob("events.*.jsonl")} + assert set(tapes) == {"events.aaaa0001.jsonl", "events.bbbb0002.jsonl"} + old = {e["event"]: e for e in tapes["events.aaaa0001.jsonl"]} + assert old["rl_sampler_leave"]["reason"] == "INFRA_FAILED" + new = {e["event"]: e for e in tapes["events.bbbb0002.jsonl"]} + assert new["rl_sampler_join"]["next_rollout_id"] == 2 + assert new["rl_sampler_join"]["incarnations"] == ["aaaa0001"] + assert new["rl_sampler_leave"]["reason"] == "FINISHED" + members = sorted(p.name for p in (tmp_path / "store" / MEMBERS_DIR / "ppu4").iterdir()) + assert members == ["aaaa0001.json", "bbbb0002.json"] + + +def test_a_finished_sampler_restarted_generates_nothing(tmp_path): + spec = _spec(tmp_path, rounds=2) + assert entry.run(spec, resource_sampler=lambda emit: None, incarnation="aaaa0001") == 2 + before = _payloads(tmp_path) + assert entry.run(spec, resource_sampler=lambda emit: None, incarnation="bbbb0002") == 0 + assert _payloads(tmp_path) == before + + +def test_a_written_round_is_refused_to_a_new_incarnation(tmp_path): + spec = _spec(tmp_path) + CALLS.update(limit=4) + with pytest.raises(RuntimeError): + entry.run(spec, resource_sampler=lambda emit: None, incarnation="aaaa0001") + from yeto.rl.engine.engine_sampler import EngineRound + from yeto.rl.engine.sampler_resume import resume_point + + resume = resume_point(_buffer(tmp_path).entries(), "ppu4") + assert resume.next_rollout_id == 1 + round_ = EngineRound.__new__(EngineRound) + from yeto.rl.engine.sampler_resume import RoundGuard + + round_.guard = RoundGuard(resume) + with pytest.raises(SamplerResumeError, match="round 0 is already"): + EngineRound.__call__(round_, 0) + + +def test_a_payload_is_never_overwritten(tmp_path): + buffer = FileSampleBuffer(CheckpointStore(tmp_path / "store"), batch_size=1) + buffer.write_payload("ppu4-a-r0-row0", b"one") + with pytest.raises(SampleBufferError, match="never overwritten"): + buffer.write_payload("ppu4-a-r0-row0", b"two") + assert (tmp_path / "store" / "sample-buffer" / "payloads" / "ppu4-a-r0-row0.bin").read_bytes() == b"one" + + +def test_reclaim_marker_stops_writes_and_leaves_preempted(tmp_path): + spec = _spec(tmp_path) + CALLS.update(marker=str(tmp_path / "reclaim-marker"), mark_at=6) # inside round 1 + # round 1 was generated (2 rounds), but none of its groups is written + assert entry.run(spec, resource_sampler=lambda emit: None, incarnation="aaaa0001") == 2 + buffer = _buffer(tmp_path) + assert not buffer.data_ended + assert [e.sampler_round["rollout_id"] for e in buffer.entries()] == [0, 0] + [tape] = (tmp_path / "tape").glob("events.*.jsonl") + leave = [e for e in _tape(tape) if e["event"] == "rl_sampler_leave"] + assert len(leave) == 1 and leave[0]["reason"] == "PREEMPTED" diff --git a/tests/test_stage1_gaps.py b/tests/test_stage1_gaps.py index 9ff34c178..0e4a33896 100644 --- a/tests/test_stage1_gaps.py +++ b/tests/test_stage1_gaps.py @@ -66,9 +66,10 @@ def test_engine_sampler_entry_runs_the_island_from_a_spec(tmp_path): reader = FileSampleBuffer(CheckpointStore(tmp_path / "store"), batch_size=1) reader.refresh() assert len(reader.select("t", 10)) == 4 # 2 rounds x 2 prompts - events = [json.loads(line) for line in (tmp_path / "tape" / "events.jsonl").read_text().splitlines()] + [tape] = (tmp_path / "tape").glob("events.*.jsonl") # one tape per incarnation + events = [json.loads(line) for line in tape.read_text().splitlines()] names = [e["event"] for e in events] - assert names[0] == "rl_sampler_start" and names.count("rl_sampler_engine_round") == 2 + assert names[0] == "rl_sampler_start" and "rl_sampler_join" in names and names.count("rl_sampler_engine_round") == 2 assert all(e["sampler_id"] == "ppu4" for e in events) diff --git a/yeto/rl/engine/engine_sampler.py b/yeto/rl/engine/engine_sampler.py index 7ef136bea..5ce580e3f 100644 --- a/yeto/rl/engine/engine_sampler.py +++ b/yeto/rl/engine/engine_sampler.py @@ -16,6 +16,13 @@ (:class:`~.sampler_engine.Provenance`) in the index entry; 6. write ``rl_sampler_engine_round`` (finish reasons, tokens, provenance key). +Replacement (S25 G1C, the engine-sampler side of C1's :mod:`.sampler_resume`): +each group's sample id carries the incarnation, the round and the prompt row, +so a new incarnation never writes an earlier one's payload name; each index +entry carries a :class:`~.sampler_resume.SamplerRound` record; a round already +in the buffer is never generated or written again (:class:`~.sampler_resume.RoundGuard`); +after the reclaim marker no group is written. + :func:`build_engine_sampler` wires the island: the start checks, the engine start on the base weights with the post-load checks (:mod:`.sampler_checks`; a failure there stops the island), the follower and the loop. The caller @@ -47,6 +54,7 @@ SamplingParams, ) from .sampler_loop import SamplerLoop +from .sampler_resume import RoundGuard, SamplerResume, SamplerRound ROUND_EVENT = "rl_sampler_engine_round" START_EVENT = "rl_sampler_start" @@ -91,7 +99,15 @@ class EngineRound: def __init__(self, engine: SamplerEngine, writer: GroupWriter, *, prompts: Callable[[int], Sequence[PromptItem]], reward: Reward, sampling: SamplingParams, provenance: Provenance, emit: Callable[..., None], - loop_guard: LoopGuard | None = None) -> None: + loop_guard: LoopGuard | None = None, incarnation: str | None = None, + resume: SamplerResume | None = None, + data_cursor: Callable[[int], Mapping[str, int]] | None = None, + reclaimed: Callable[[], bool] = lambda: False) -> None: + """``incarnation``: this process's id (``sampler_resume.new_incarnation``); + with it, sample ids are ``--r-`` and + every entry carries its round record. ``resume``: where the earlier + incarnations got to (its rounds are refused). ``data_cursor(rollout_id)``: + the prompt source's position after the round drew its prompts.""" self.engine = engine self.writer = writer self.prompts = prompts @@ -100,9 +116,24 @@ def __init__(self, engine: SamplerEngine, writer: GroupWriter, *, self.provenance = provenance self.emit = emit self.loop_guard = loop_guard + self.incarnation = None if incarnation is None else str(incarnation) + self.guard = None if resume is None else RoundGuard(resume) + self.data_cursor = data_cursor + self.reclaimed = reclaimed + self.dropped_groups = 0 # groups not written because the island is leaving + + def _round_record(self, rollout_id: int, last: bool) -> dict[str, Any] | None: + if self.incarnation is None: + return None + cursor = None if self.data_cursor is None else self.data_cursor(rollout_id) + return SamplerRound(self.writer.sampler_id, self.incarnation, rollout_id, + data_cursor=cursor, last=last).record() def __call__(self, rollout_id: int) -> None: - items = list(self.prompts(int(rollout_id))) + rollout_id = int(rollout_id) + if self.guard is not None: + self.guard.check(rollout_id) + items = list(self.prompts(rollout_id)) version = self.engine.served_version() if version is None: raise SamplerEngineError("the engine serves no version; it cannot generate") @@ -114,6 +145,7 @@ def __call__(self, rollout_id: int) -> None: record = self.provenance.to_record() finish: Counter[str] = Counter() tokens = 0 + scored = [] for position, (item, group) in enumerate(zip(items, groups)): if len(group) != self.sampling.n: raise SamplerEngineError(f"prompt {item.prompt_id}: {len(group)} completions, " @@ -128,11 +160,23 @@ def __call__(self, rollout_id: int) -> None: item, completion, reward=self.reward(item, completion), version=version, provenance=record, group_index=position, index=position * self.sampling.n + j)) - self.writer.write(samples, group_key=f"r{int(rollout_id)}-{item.prompt_id}", - provenance=record) - self.emit(ROUND_EVENT, rollout_id=int(rollout_id), policy_version=version, - groups=len(items), samples=sum(finish.values()), response_tokens=tokens, - finish=dict(sorted(finish.items())), source=self.provenance.key()) + scored.append((item, samples)) + prefix = "" if self.incarnation is None else f"{self.incarnation}-" + written = 0 + for position, (item, samples) in enumerate(scored): + if self.reclaimed(): # leaving: the replacement owns the shard from now on + self.dropped_groups += len(scored) - position + break + self.writer.write(samples, group_key=f"{prefix}r{rollout_id}-{item.prompt_id}", + provenance=record, + sampler_round=self._round_record(rollout_id, position == len(scored) - 1)) + written += 1 + if written and self.guard is not None: + self.guard.written(rollout_id) + self.emit(ROUND_EVENT, rollout_id=rollout_id, policy_version=version, + groups=written, samples=sum(finish.values()), response_tokens=tokens, + finish=dict(sorted(finish.items())), source=self.provenance.key(), + dropped_groups=len(scored) - written) @dataclass @@ -159,7 +203,12 @@ class EngineSamplerIsland: follower: PolicyFollower round: EngineRound - def run(self, start: int = 0) -> int: + resume: SamplerResume | None = None + + def run(self, start: int | None = None) -> int: + """Rounds from ``start`` (default: the resume point, else 0).""" + if start is None: + start = 0 if self.resume is None else self.resume.next_rollout_id return self.loop.run(start) @@ -168,7 +217,10 @@ def build_engine_sampler(config: EngineSamplerConfig, *, engine: SamplerEngine, prompts: Callable[[int], Sequence[PromptItem]], reward: Reward, checks: PostLoadChecks | None = None, should_stop: Callable[[], bool] = lambda: False, - sleep: Callable[[float], None] | None = None) -> EngineSamplerIsland: + sleep: Callable[[float], None] | None = None, + incarnation: str | None = None, resume: SamplerResume | None = None, + data_cursor: Callable[[int], Mapping[str, int]] | None = None, + ) -> EngineSamplerIsland: """Check the sampling against the engine, start it on the base weights, run the post-load checks, and return the wired island (module doc).""" sampling = config.sampling.check() @@ -188,7 +240,8 @@ def build_engine_sampler(config: EngineSamplerConfig, *, engine: SamplerEngine, follower = PolicyFollower(releases, loader, version=int(config.base_version)) engine_round = EngineRound(engine, GroupWriter(buffer, config.sampler_id), prompts=prompts, reward=reward, sampling=sampling, provenance=provenance, emit=emit, - loop_guard=config.loop_guard) + loop_guard=config.loop_guard, incarnation=incarnation, + resume=resume, data_cursor=data_cursor, reclaimed=should_stop) extra = {} if sleep is None else {"sleep": sleep} loop = SamplerLoop(follower=follower, buffer=buffer, generate=engine_round, num_rollouts=int(config.num_rollouts), @@ -196,7 +249,8 @@ def build_engine_sampler(config: EngineSamplerConfig, *, engine: SamplerEngine, poll_s=config.poll_s, should_stop=should_stop, min_round_s=config.min_round_s, loaded_fields=loader.loaded_fields, ready=lambda: loader.ready, **extra) - return EngineSamplerIsland(loop=loop, loader=loader, follower=follower, round=engine_round) + return EngineSamplerIsland(loop=loop, loader=loader, follower=follower, round=engine_round, + resume=resume) def main(argv: Sequence[str] | None = None) -> int: diff --git a/yeto/rl/engine/engine_sampler_entry.py b/yeto/rl/engine/engine_sampler_entry.py index 1eda00f9b..f842074ba 100644 --- a/yeto/rl/engine/engine_sampler_entry.py +++ b/yeto/rl/engine/engine_sampler_entry.py @@ -31,7 +31,17 @@ ``reward_kind`` is ``"miles"``; * ``checks``: :class:`~.sampler_checks.PostLoadChecks` keyword overrides (optional; e.g. ``expected_expert_matrices``); -* ``event_tape``: the JSONL tape file of the island's events. +* ``event_tape``: the JSONL tape file of the island's events; each + incarnation writes its own file next to it (``events..jsonl``), + so a replacement container never overwrites an earlier one's tape. + +Replacement (S25 G1C): the process is one incarnation of the sampler +(``sampler_resume.new_incarnation``). It reads where the earlier incarnations +got to from the buffer (``sampler_resume.resume_point``), starts at the next +round, and writes ``rl_sampler_join`` / ``rl_sampler_leave`` plus one member +record per incarnation in the buffer store. With ``YETO_SPOT_RECLAIM_MARKER`` +set, SIGINT/SIGTERM write the reclaim marker; the island then starts no new +round, writes no more groups and leaves as PREEMPTED. Nothing here was run on a PPU or NPU node. """ @@ -107,6 +117,11 @@ def __init__(self, path: str | Path, *, input_key: str, label_key: str | None, g self.groups_per_round = int(groups_per_round) self.render = render + def cursor(self, rollout_id: int) -> dict[str, int]: + """The data position after round ``rollout_id`` drew its prompts (the + ``sampler_round`` record; rounds draw by rollout id, so it only documents).""" + return {"row_offset": (int(rollout_id) + 1) * self.groups_per_round} + def __call__(self, rollout_id: int) -> list[PromptItem]: start = int(rollout_id) * self.groups_per_round items = [] @@ -191,6 +206,12 @@ def build_stores(spec: Mapping[str, Any], environ: Mapping[str, str] | None = No return releases, FileSampleBuffer(buffer_store, batch_size=1, index=index) # as open_buffer(sampler=True) +def incarnation_tape(path: str | Path, incarnation: str) -> Path: + """``events.jsonl`` -> ``events..jsonl`` (one tape per incarnation).""" + tape = Path(path).expanduser() + return tape.with_name(f"{tape.stem}.{incarnation}{tape.suffix or '.jsonl'}") + + def tape_writer(path: str | Path, sampler_id: str) -> Callable[..., None]: from yeto.rl.event_echo import append_record # noqa: PLC0415 @@ -219,16 +240,52 @@ def island_config(spec: Mapping[str, Any]) -> EngineSamplerConfig: def run(spec: Mapping[str, Any], *, environ: Mapping[str, str] | None = None, - resource_sampler: Callable[[Callable[..., None]], Any] | None = None) -> int: - """Build every part of the island from ``spec``, run it, stop the engine; - returns the rounds the loop ran. ``resource_sampler(emit)`` returns a - started telemetry thread (default: ``telemetry.ResourceSampler`` of the - island's device family; None from the factory = no telemetry).""" + resource_sampler: Callable[[Callable[..., None]], Any] | None = None, + incarnation: str | None = None) -> int: + """Build every part of the island from ``spec``, run it from the resume + point, stop the engine; returns the rounds this incarnation ran. + ``resource_sampler(emit)`` returns a started telemetry thread (default: + ``telemetry.ResourceSampler`` of the island's device family; None from the + factory = no telemetry).""" + import os # noqa: PLC0415 + + from yeto.rl.elastic_client import LeaveReason # noqa: PLC0415 + + from .sampler_resume import ( # noqa: PLC0415 + RECLAIM_MARKER_ENV, + ReclaimWatcher, + SamplerMembership, + install_marker_signals, + new_incarnation, + reclaim_marker, + reclaim_requested, + resume_point, + ) + + env = os.environ if environ is None else environ config = island_config(spec) - emit = tape_writer(_require(spec, "event_tape"), config.sampler_id) + incarnation = incarnation or new_incarnation() + tape = incarnation_tape(_require(spec, "event_tape"), incarnation) + emit = tape_writer(tape, config.sampler_id) releases, buffer = build_stores(spec, environ) prompts, reward = build_prompts(spec), build_reward(spec) checks = PostLoadChecks(**dict(spec.get("checks") or {})) + reclaim_marker(env).unlink(missing_ok=True) # a restarted machine keeps its /tmp + if env.get(RECLAIM_MARKER_ENV): + install_marker_signals(env) + buffer.refresh() + resume = resume_point(buffer.entries(), config.sampler_id) + from .resume import store_for # noqa: PLC0415 + + member_store = store_for(str(spec.get("buffer_root") or spec["store"]), environ=environ) + membership = SamplerMembership(member_store, config.sampler_id, incarnation, resume, emit) + reclaimed = lambda: reclaim_requested(env) # noqa: E731 + island = None + + def leave(reason: LeaveReason, **detail: Any) -> None: + version = None if island is None else island.follower.version + membership.leave(reason.name, policy_version=version, buffer_end=len(buffer), **detail) + engine = build_engine(spec) if resource_sampler is None: from .telemetry import ResourceSampler # noqa: PLC0415 @@ -237,11 +294,28 @@ def resource_sampler(sink: Callable[..., None]) -> Any: return ResourceSampler(lambda event, **fields: sink(event, **fields)).start() telemetry = resource_sampler(emit) + watcher = ReclaimWatcher(reclaimed, lambda: leave(LeaveReason.PREEMPTED)).start() try: island = build_engine_sampler(config, engine=engine, releases=releases, buffer=buffer, emit=emit, - prompts=prompts, reward=reward, checks=checks) - return island.run() + prompts=prompts, reward=reward, checks=checks, + should_stop=reclaimed, incarnation=incarnation, resume=resume, + data_cursor=getattr(prompts, "cursor", None)) + latest = releases.latest() + membership.join(tape=str(tape), latest_release=None if latest is None else + {"seq": latest.seq, "policy_version": latest.version}) + rounds = 0 if buffer.data_ended else island.run() + leave(LeaveReason.PREEMPTED if reclaimed() else LeaveReason.FINISHED, rounds=rounds, + dropped_groups=island.round.dropped_groups) + return rounds + except BaseException as error: + try: + leave(LeaveReason.PREEMPTED if reclaimed() else LeaveReason.INFRA_FAILED, + error=f"{type(error).__name__}: {error}"[:300]) + except Exception: # noqa: BLE001 - the original error is the one to report + pass + raise finally: + watcher.stop() try: engine.stop() finally: diff --git a/yeto/rl/engine/sample_buffer.py b/yeto/rl/engine/sample_buffer.py index 4c945eb14..e2f64906c 100644 --- a/yeto/rl/engine/sample_buffer.py +++ b/yeto/rl/engine/sample_buffer.py @@ -405,8 +405,14 @@ def __init__(self, store: CheckpointStore | ObjectStore, *, batch_size: int, # ---------------------------------------------------------------- writing def write_payload(self, sample_id: str, data: bytes) -> str: - """Store one payload; returns the ``payload_ref`` for its index entry.""" - return self.index.write_payload(payload_name(sample_id), data) + """Store one payload; returns the ``payload_ref`` for its index entry. + A payload is written once: a name already in the store is refused (a + replacement sampler must never overwrite a group an index entry names).""" + name = payload_name(sample_id) + if self.index.has_payload(f"{PAYLOADS}/{name}"): + raise SampleBufferError(f"payload {name} is already in the sample buffer; " + "a payload is never overwritten") + return self.index.write_payload(name, data) def read_payload(self, payload_ref: str) -> bytes: return self.index.read_payload(payload_ref) From 7346fc212a799557f26696f3d52500f05b81f5f6 Mon Sep 17 00:00:00 2001 From: michael <78113554+chowhsu@users.noreply.github.com> Date: Tue, 20 Oct 2026 11:23:35 +0000 Subject: [PATCH 78/89] Token-diversity test reward; engine sampler rewards see token ids The word-diversity rehearsal reward was gamed in T-G1B g1b-1: one 512-token "word" without spaces scored 1.0 for every sample, so round 2 had no gradient. yeto.rl.token_diversity_reward scores distinct token ids over response tokens. The engine sampler's Miles reward sample now carries tokens and response_length as a Miles Sample does. The G1C plan uses the new reward with lr 1e-4. Co-Authored-By: Claude Opus 5.5 --- examples/stage1/glm5next-mini-modal-g1c.yaml | 154 +++++++++++++++++++ tests/test_s25_g1c_reward_limits.py | 67 ++++++++ yeto/rl/engine/engine_sampler_entry.py | 9 +- yeto/rl/token_diversity_reward.py | 41 +++++ 4 files changed, 269 insertions(+), 2 deletions(-) create mode 100644 examples/stage1/glm5next-mini-modal-g1c.yaml create mode 100644 tests/test_s25_g1c_reward_limits.py create mode 100644 yeto/rl/token_diversity_reward.py diff --git a/examples/stage1/glm5next-mini-modal-g1c.yaml b/examples/stage1/glm5next-mini-modal-g1c.yaml new file mode 100644 index 000000000..ad1d75335 --- /dev/null +++ b/examples/stage1/glm5next-mini-modal-g1c.yaml @@ -0,0 +1,154 @@ +# S25 T-G1C: glm5next-mini-modal-dense.yaml (T-G1B) with the reward and lr changed again: +# task.reward yeto.rl.token_diversity_reward:score (the word reward was gamed in g1b-1: +# one 512-token word scored 1.0 for every sample, round 2 had grad_norm 0) +# trainer.miles.lr 1.0e-4 (3e-4 drove the collapse in two steps; 1e-4 Adam steps are +# still about one bf16 ulp of a 0.02 weight, so releases still change) +# +# T-G1B: glm5next-mini-modal.yaml with two changes, so the weights move and the +# chain "trained weights -> re-quantized MXFP4 experts -> sampler" is checked by value: +# task.reward yeto.rl.word_diversity_reward:score (GSM8K is constant 0 on this model, +# so GRPO advantages and grad_norm were 0 in T-G1 run-17) +# trainer.miles.lr 3.0e-4 (1e-6 Adam steps are below one bf16 ulp of most weights, +# so a release hash could stay the same even with a non-zero gradient) +# Everything else (nodes, store, images, sampler, sampling) is the T-G1 plan. +# +# Stage-1 mini rehearsal (S25 T-G1) on Modal H100: the same plan schema and +# the same generated commands as glm53-flash-ppu.yaml; only the nodes +# (Modal containers) and the store (a Modal Volume through the yeto store +# hook, not OSS) differ. Model: Miles' 4-layer GLM5.3-Flash CI checkpoint cut +# to 32 routed experts per MoE layer (tools/probes/s25_g1/make_mini.py). +# +# yeto-stage1 plan examples/stage1/glm5next-mini-modal.yaml +# modal run tools/probes/s25_g1/modal_app.py::... (the driver reads this plan) + +run_id: s25-g1-mini +inventory: modal-mini-inventory.yaml +rounds: 3 +seed: 17 + +rendezvous: + modal-trainer: 127.0.0.1 + modal-sampler: 127.0.0.1 + +roles: + trainer: [modal-trainer] + sampler: [modal-sampler] + +model: + name: GLM5.3-Flash-4layer-e32 + architecture: glm5_next + num_layers: 4 # KDA, DSA, KDA, DSA + first_k_dense_replace: 1 + num_experts: 32 # cut from 288 (make_mini.py) + params_b: 5.3 # 推算: 10.5 GB bf16 / 2 + expert_params_b: 2.4 # 3 MoE layers x 32 x 3 x 4096 x 2048 + bf16_source: + revision: 7fb7883207fe73ea04430ac392f6ffe2b9f936c4 # CharyZeng/GLM-5.3-Flash-4layer, then cut + origin: hf://CharyZeng/GLM-5.3-Flash-4layer@7fb7883207fe73ea04430ac392f6ffe2b9f936c4 + staged: modal-volume://yeto-s25-g1-models/glm5next-mini-bf16-7fb7883 + local_dir: /data/models/glm5next-mini-bf16-7fb7883 + torch_dist: /data/models/glm5next-mini-bf16-7fb7883_torch_dist # Miles converter output (convert step) + sampler_initial: # round 0: our N8.5 converter on the bf16 mini model + repo: local/glm5next-mini-mxfp4-r0-7fb7883 + revision: 7fb7883 + staged: modal-volume://yeto-s25-g1-models/glm5next-mini-mxfp4-r0-7fb7883 + local_dir: /data/models/glm5next-mini-mxfp4-r0-7fb7883 + +quant: + scheme: mxfp4-g32-rtn + format: mxfp4-pack-quantized + release_dir: policy-releases-mxfp4-g32-rtn + expected_expert_matrices: 288 # 3 MoE layers x 32 experts x (gate, up, down) + frozen_pattern: '^model\.visual\.' # Miles trains the text model only; the vision tower stays the base + +store: /data/store/s25-g1 +store_backend: modal-volume # the only difference from OSS: the store hook and its Volume +store_volume: yeto-s25-g1-store +store_access: + head: direct + +images: + sampler: ghcr.io/michaellchung/yeto-miles-ports@sha256:5b9dc2a9e6e315d73dabfbc6faae428b59cc01aabec6189645a5c8c80d6e3721 + trainer: ghcr.io/michaellchung/yeto-miles-ports@sha256:5b9dc2a9e6e315d73dabfbc6faae428b59cc01aabec6189645a5c8c80d6e3721 + docker_device_args: [--gpus, all] # unused on Modal; the containers get the GPUs + +trainer: + backend: miles + miles: + commit: 1618b5c38db80e17ec928cfa47fe832cb18565fc # s25-distopt, checked out in the image + megatron: radixark/Megatron-LM@a84b105 + mbridge: image + workdir: /root/miles + extra_env: + # Megatron's get_model imports megatron.post_training (modelopt check), which the + # image's editable megatron_core install does not expose (Modal run-10) + PYTHONPATH: /root/Megatron-LM:/root/sky_workdir + tp: 1 + sequence_parallel: false + pp: 1 + pp_layers: [4] + cp: 1 + ep: 2 + etp: 1 + optimizer_offload_fraction: 1.0 + memory_estimate: # 推算; `plan` checks them + weights_grads_gib_per_card: 22.9 + host_optimizer_gib_per_node: 59.2 + transformer_impl: transformer_engine + recompute: full + max_tokens_per_gpu: 8192 + lr: 1.0e-4 # T-G1C: see the header + verl: + commit: null + strategy: fsdp2 + fsdp_size: 2 + ulysses: 1 + param_offload: false + optimizer_offload: true + lr: 1.0e-6 + +sampler: + engine: sglang + engine_version: 0.5.21.dev70+g2fa8801 + python: python3 + topology: tp1 + tp: 1 + precision: mxfp4 + max_model_len: 4096 + port: 30000 + moe_runner_backend: triton + hierarchical_cache: false + nextn: false + renames_weight_packed: true # the yeto hook in the image renames (sglang_load_report) + load_report: true # the hook reports loaded expert matrices (post-load check 1) + in_place_reload: false + # SGLang fuses the BF16 shared expert into the routed FusedMoE as expert 32; under the + # MXFP4 MoE method that slot is packed U8, so the BF16 shared weights do not fit (smoke-2) + extra_args: [--disable-shared-experts-fusion, --cuda-graph-max-bs-decode, "16"] # 16 = 4 prompts x 4 samples; capture took 118 s at the default sizes + extra_env: + PYTHONPATH: /root/sky_workdir + checks: + require_fixed_answers: false # a cut 4-layer model cannot answer GSM8K + fixed_max_tokens: 64 + +sampling: + max_tokens: 512 + max_prompt_len: 512 + temperature: 1.0 + top_p: 1.0 + top_k: -1 + repetition_penalty: 1.0 + frequency_penalty: 0.0 + presence_penalty: 0.0 + +task: + data: zhuzilin/gsm8k + data_revision: 0cbd9f31d91ac21a7613dcbc7fef992adac459ae + reward: yeto.rl.token_diversity_reward:score # T-G1C: see the header + rollout_batch_size: 4 + n_samples_per_prompt: 4 + prompts_jsonl: /data/models/datasets/gsm8k-0cbd9f3/train.jsonl + data_local: /data/models/datasets/gsm8k-0cbd9f3/train.jsonl # the same rows for the trainer (offline nodes) + input_key: messages + label_key: label + prompt_format: chat diff --git a/tests/test_s25_g1c_reward_limits.py b/tests/test_s25_g1c_reward_limits.py new file mode 100644 index 000000000..ca14a0260 --- /dev/null +++ b/tests/test_s25_g1c_reward_limits.py @@ -0,0 +1,67 @@ +"""S25 T-G1C: the token-diversity test reward and the G1 driver's wall-clock limits (CPU).""" + +from __future__ import annotations + +import sys +from pathlib import Path + +import pytest + +from yeto.rl import token_diversity_reward as tdr +from yeto.rl.engine import engine_sampler_entry as entry + +sys.path.insert(0, str(Path(__file__).resolve().parents[1] / "tools" / "probes" / "s25_g1")) +from run_limits import Deadline, wait_until # noqa: E402 + + +def test_token_diversity_reward_scores_token_ids_not_words(): + import asyncio + from types import SimpleNamespace + + one_word = SimpleNamespace(tokens=[1, 2] + [7] * 8, response_length=8) # "PrototypePrototype..." + varied = SimpleNamespace(tokens=[1, 2, 3, 4, 3, 5], response_length=4) + assert asyncio.run(tdr.score(None, one_word)) == pytest.approx(1 / 8) + assert asyncio.run(tdr.score(None, varied)) == pytest.approx(3 / 4) + assert tdr.diversity([]) == 0.0 + with pytest.raises(ValueError): + tdr.response_ids(SimpleNamespace(tokens=[1], response_length=3)) + + +def test_miles_reward_wrapper_passes_token_ids(): + from yeto.rl.engine.engine_sampler import PromptItem + from yeto.rl.engine.sampler_engine import Completion + + reward = entry.miles_reward(tdr.score) + done = Completion(text="x", token_ids=(5, 5, 6, 7), logprobs=(0.0,) * 4, finish_reason="length", + prompt_token_ids=(1, 2)) + assert reward(PromptItem("p", "q"), done) == pytest.approx(3 / 4) + + +class _Clock: + def __init__(self) -> None: + self.now = 0.0 + + def __call__(self) -> float: + return self.now + + def sleep(self, s: float) -> None: + self.now += s + + +def test_sampler_wait_after_the_trainers_has_a_limit(): + clock = _Clock() + assert wait_until(lambda: False, 300, poll_s=15, clock=clock, sleep=clock.sleep) is False + assert 300 <= clock.now < 315 + clock.now = 0.0 + assert wait_until(lambda: clock.now >= 60, 300, poll_s=15, clock=clock, sleep=clock.sleep) is True + + +def test_run_deadline_ends_a_wait_early(): + clock = _Clock() + deadline = Deadline(100, clock=clock) + assert not deadline.passed() + assert wait_until(lambda: False, 1000, poll_s=10, clock=clock, sleep=clock.sleep, + deadline=deadline) is False + assert deadline.passed() and clock.now == 100 + with pytest.raises(ValueError): + Deadline(0) diff --git a/yeto/rl/engine/engine_sampler_entry.py b/yeto/rl/engine/engine_sampler_entry.py index f842074ba..3dfae23b8 100644 --- a/yeto/rl/engine/engine_sampler_entry.py +++ b/yeto/rl/engine/engine_sampler_entry.py @@ -170,11 +170,16 @@ def build_prompts(spec: Mapping[str, Any]) -> JsonlPrompts: def miles_reward(score: Callable[..., Any]) -> Reward: """A Miles reward (``async score(args, sample)``) as an engine-sampler reward: - the sample carries ``response``, ``label`` and a ``metadata`` dict.""" + the sample carries ``response``, ``label``, a ``metadata`` dict, ``tokens`` + (prompt ids when the engine gave them, then the response ids) and + ``response_length``, as a Miles ``Sample`` does.""" def reward(item: PromptItem, completion: Completion) -> float: + response = list(getattr(completion, "token_ids", ())) + prompt_ids = list(getattr(completion, "prompt_token_ids", ())) sample = types.SimpleNamespace(prompt=item.prompt, response=completion.text, label=item.label, - metadata=dict(item.metadata)) + metadata=dict(item.metadata), tokens=prompt_ids + response, + response_length=len(response)) value = score(None, sample) return float(asyncio.run(value) if inspect.isawaitable(value) else value) diff --git a/yeto/rl/token_diversity_reward.py b/yeto/rl/token_diversity_reward.py new file mode 100644 index 000000000..8547b677c --- /dev/null +++ b/yeto/rl/token_diversity_reward.py @@ -0,0 +1,41 @@ +"""Dense token-diversity reward for rehearsals on models that cannot solve the task. + +A test reward, not a product reward (S25 T-G1C). The word-diversity reward +(:mod:`yeto.rl.word_diversity_reward`) was gamed in T-G1B: after two steps the +cut model wrote one 512-token "word" without spaces, every sample scored 1.0, +the group variance was 0 and round 2 had no gradient. This reward counts token +ids, so one repeated token scores low however it is spelled: + + score = distinct response token ids / response tokens (0.0 when empty) + +The response ids are ``sample.tokens[-response_length:]`` (Miles ``Sample`` +and the engine sampler's reward sample both carry ``tokens`` and +``response_length``). Range (0, 1]; the maximum needs every response token to +differ, which a sampled 512-token response at temperature 1 does not reach. +Launch with + + --reward-function yeto.rl.token_diversity_reward:score +""" + +from __future__ import annotations + +from collections.abc import Sequence +from typing import Any + +__all__ = ["diversity", "response_ids", "score"] + + +def diversity(ids: Sequence[int]) -> float: + return len(set(ids)) / len(ids) if ids else 0.0 + + +def response_ids(sample: Any) -> list[int]: + tokens = list(getattr(sample, "tokens", None) or ()) + length = int(getattr(sample, "response_length", 0) or 0) + if length > len(tokens): + raise ValueError(f"response_length {length} > {len(tokens)} tokens") + return tokens[len(tokens) - length:] if length else [] + + +async def score(args, sample, **kwargs) -> float: + return diversity(response_ids(sample)) From e7125c82f3fefae835c57cc0fa8e9eafd5a5c13d Mon Sep 17 00:00:00 2001 From: michael <78113554+chowhsu@users.noreply.github.com> Date: Tue, 20 Oct 2026 11:23:35 +0000 Subject: [PATCH 79/89] G1 driver: wall-clock limits and one log directory per sampler container g1b-1 held a sampler H100 about 30 min after both trainers had finished. orchestrate now cancels every call past a whole-run limit, and the sampler 5 min after both trainers returned 0. A sampler container Modal runs again writes its own log directory. Co-Authored-By: Claude Opus 5.5 --- tools/probes/s25_g1/modal_app.py | 39 +++++++++++++++++++++---- tools/probes/s25_g1/run_limits.py | 47 +++++++++++++++++++++++++++++++ 2 files changed, 80 insertions(+), 6 deletions(-) create mode 100644 tools/probes/s25_g1/run_limits.py diff --git a/tools/probes/s25_g1/modal_app.py b/tools/probes/s25_g1/modal_app.py index dee8bc043..73f0cf205 100644 --- a/tools/probes/s25_g1/modal_app.py +++ b/tools/probes/s25_g1/modal_app.py @@ -28,9 +28,14 @@ import modal + HERE = Path(__file__).resolve() REPO = HERE.parents[3] if len(HERE.parents) > 3 else Path("/") # in the container the module sits at /root WORKDIR = "/root/sky_workdir" +# run_limits.py sits next to this file locally and under the mounted tree in a container +sys.path.insert(0, str(HERE.parent) if (HERE.parent / "run_limits.py").is_file() + else f"{WORKDIR}/tools/probes/s25_g1") +from run_limits import Deadline, wait_until # noqa: E402 # T-G1B: YETO_G1_PLAN=examples/stage1/glm5next-mini-modal-dense.yaml (set at deploy time; # the image env carries it into every container, so all of them plan the same file) PLAN = os.environ.get("YETO_G1_PLAN", "examples/stage1/glm5next-mini-modal.yaml") @@ -239,6 +244,10 @@ def smoke(tag: str = "smoke-1") -> None: @app.function(image=image, volumes=VOLUMES, gpu=f"{GPU_NAME}!:1", cpu=8, memory=65536, timeout=7200, max_inputs=1) # a fresh container per call: trainer B must not inherit A's disk or Ray def sampler_island(tag: str) -> int: + # Modal may run this input again in a new container (a lost worker): each container + # keeps its own log directory, so a replacement never overwrites the first one's logs + container = os.environ.get("MODAL_TASK_ID") or f"t{int(time.time())}" + tag = f"{tag}/{container}" log = _log_dir(tag) _gpu_facts(log, 1) nvml = _nvml_sampler(log / "nvml.csv") @@ -419,11 +428,21 @@ def record(kind: str, name: str, directory: Path, extra: dict) -> None: time.sleep(poll_s) +RUN_LIMIT_S = 100 * 60 # whole run (g1b-1: trainers about 40 min, sampler first start about 6 min) +SAMPLER_DRAIN_S = 5 * 60 # after both trainers returned 0: the sampler only ends its last round + + @app.function(image=image, volumes=VOLUMES, cpu=1, memory=1024, timeout=3 * 3600) -def orchestrate(tag: str, cut_after: int = 2) -> dict: +def orchestrate(tag: str, cut_after: int = 2, run_limit_s: float = RUN_LIMIT_S, + sampler_drain_s: float = SAMPLER_DRAIN_S) -> dict: """Sampler for the whole run; trainer A stops after ``cut_after`` rounds (cut), trainer B resumes in a new container. A trainer that fails cancels the sampler; - a sampler that fails cancels the running trainer (no card waits for nothing).""" + a sampler that fails cancels the running trainer (no card waits for nothing). + Two wall-clock limits (``run_limits``): the whole run (``run_limit_s``: a lost + container that Modal runs again, or a hang, cancels every call), and the sampler + after both trainers returned 0 (``sampler_drain_s``).""" + deadline = Deadline(run_limit_s) + limits: list[str] = [] run_store = Path("/data/store/s25-g1") # the plan's store: every run starts from an empty one if run_store.exists(): shutil.rmtree(run_store) @@ -438,6 +457,8 @@ def orchestrate(tag: str, cut_after: int = 2) -> dict: def sampler_done() -> bool: nonlocal sampler_rc + if sampler_rc is not None: + return True try: sampler_rc = sampler.get(timeout=0) return True @@ -456,6 +477,10 @@ def run_trainer(sub: str, stop: int | None) -> object: pass except Exception as error: # noqa: BLE001 return f"error: {error}" + if deadline.passed(): + call.cancel() + limits.append(f"{sub}: run limit {run_limit_s:g} s passed") + return "cancelled: run limit" if sampler_done() and sampler_rc != 0: call.cancel() return "cancelled: sampler ended with " + str(sampler_rc) @@ -466,13 +491,15 @@ def run_trainer(sub: str, stop: int | None) -> object: if not sampler_done(): sampler.cancel() sampler_rc = "cancelled" - else: - while not sampler_done(): - time.sleep(30) + elif not wait_until(sampler_done, sampler_drain_s, deadline=deadline): + sampler.cancel() + limits.append(f"sampler: not done {sampler_drain_s:g} s after both trainers returned 0") + sampler_rc = "cancelled: drain limit" time.sleep(90) # the re-quantizer's last log mirror requant.cancel() hasher.cancel() - result = {"trainer_a": first, "trainer_b": second, "sampler": sampler_rc} + result = {"trainer_a": first, "trainer_b": second, "sampler": sampler_rc, "limits": limits, + "wall_s": round(run_limit_s - deadline.left(), 1)} print(json.dumps(result, default=str), flush=True) out = Path(f"/data/logs/{tag}") out.mkdir(parents=True, exist_ok=True) diff --git a/tools/probes/s25_g1/run_limits.py b/tools/probes/s25_g1/run_limits.py new file mode 100644 index 000000000..6433c8688 --- /dev/null +++ b/tools/probes/s25_g1/run_limits.py @@ -0,0 +1,47 @@ +"""Wall-clock limits of the S25 G1 rehearsal driver (``modal_app.orchestrate``). + +No card may wait for nothing (T-G1B g1b-1: the sampler held 1 x H100 for about +30 min after both trainers had finished): + +* :class:`Deadline`: the whole run's wall-clock limit. A lost container that + Modal runs again (or one that hangs) can not keep the cards past it. +* :func:`wait_until`: wait for a call to finish, at most ``limit_s``; the driver + uses it for the sampler after every trainer returned rc 0 (the sampler then + has nothing to wait for but its backlog). + +Pure functions over ``done()``, ``clock()`` and ``sleep()``: CPU-tested in +``tests/test_s25_g1c_sampler_resume.py``. +""" +from __future__ import annotations + +import time +from collections.abc import Callable + + +class Deadline: + def __init__(self, limit_s: float, clock: Callable[[], float] = time.monotonic) -> None: + if float(limit_s) <= 0: + raise ValueError("limit_s must be > 0") + self.clock = clock + self.limit_s = float(limit_s) + self.started = clock() + + def left(self) -> float: + return self.limit_s - (self.clock() - self.started) + + def passed(self) -> bool: + return self.left() <= 0 + + +def wait_until(done: Callable[[], bool], limit_s: float, *, poll_s: float = 15.0, + clock: Callable[[], float] = time.monotonic, + sleep: Callable[[float], None] = time.sleep, + deadline: Deadline | None = None) -> bool: + """True once ``done()``; False when ``limit_s`` (or ``deadline``) passed first.""" + started = clock() + while True: + if done(): + return True + if clock() - started >= float(limit_s) or (deadline is not None and deadline.passed()): + return False + sleep(poll_s) From bd58b188cce771a11887c5405a924fbf767d2d13 Mon Sep 17 00:00:00 2001 From: michael <78113554+chowhsu@users.noreply.github.com> Date: Tue, 20 Oct 2026 11:23:35 +0000 Subject: [PATCH 80/89] Resume a replaced engine sampler island; never overwrite a payload T-G1B g1b-1 lost the engine sampler's container; Modal ran it again and the new container started at round 0, wrote payloads under the same names (r0-row0, ...) and overwrote the first container's buffer payloads and tape. The verl buffer path's engine sampler island has the same gap. - The engine sampler island reuses the neutral sampler_resume layer (C1): each process is one incarnation, reads the resume point from the buffer, starts after the newest written round, refuses written rounds (RoundGuard), writes rl_sampler_join / rl_sampler_leave and a member record per incarnation, and stops writing on the reclaim marker. - Sample ids are --r-, and the buffer refuses a payload name that is already stored. - Each incarnation writes its own tape (events..jsonl). Co-Authored-By: Claude Opus 5.5 (cherry picked from commit a4b9962dd3e5e39cbb34dc92a2ef514ac3ee93d9) --- tests/test_rl_spot_sampler_replace.py | 2 +- tests/test_s25_g1c_sampler_resume.py | 152 +++++++++++++++++++++++++ tests/test_stage1_gaps.py | 5 +- yeto/rl/engine/engine_sampler.py | 76 +++++++++++-- yeto/rl/engine/engine_sampler_entry.py | 92 +++++++++++++-- yeto/rl/engine/sample_buffer.py | 10 +- 6 files changed, 312 insertions(+), 25 deletions(-) create mode 100644 tests/test_s25_g1c_sampler_resume.py diff --git a/tests/test_rl_spot_sampler_replace.py b/tests/test_rl_spot_sampler_replace.py index 18e99c2d8..5826db51b 100644 --- a/tests/test_rl_spot_sampler_replace.py +++ b/tests/test_rl_spot_sampler_replace.py @@ -149,7 +149,7 @@ def test_a_writer_without_resume_writes_as_before(tmp_path): writer = br.SamplerRoundWriter(FileSampleBuffer(store, batch_size=1), "n0", incarnation="a", resume=False) writer.write_round(0, _round(0, 1), None) - writer.write_round(0, _round(0, 1), None) + writer.write_round(0, [("r0-again", _round(0, 1)[0][1])], None) # payload names stay unique assert writer.seek(object(), shuffle=True) is None assert len(FileSampleBuffer(store, batch_size=1)) == 2 diff --git a/tests/test_s25_g1c_sampler_resume.py b/tests/test_s25_g1c_sampler_resume.py new file mode 100644 index 000000000..c8995bcd2 --- /dev/null +++ b/tests/test_s25_g1c_sampler_resume.py @@ -0,0 +1,152 @@ +"""S25 T-G1C: a replaced engine sampler island continues where the buffer says +and never overwrites a payload or an earlier incarnation's tape (CPU).""" + +from __future__ import annotations + +import json +import os +import signal +from pathlib import Path + +import pytest + +from tests.test_stage1_gaps import _sampler_spec +from yeto.rl.engine import engine_sampler_entry as entry +from yeto.rl.engine.resume import CheckpointStore +from yeto.rl.engine.sample_buffer import FileSampleBuffer, SampleBufferError +from yeto.rl.engine.sampler_resume import MEMBERS_DIR, RECLAIM_MARKER_ENV, SamplerResumeError + +STATE_ENV = "YETO_TEST_G1C_REWARD_STATE" # the reward is imported by name: share state via a file + + +def counting_reward(item, completion) -> float: + """Fails (a lost container) or writes the reclaim marker after the limits in the state file.""" + path = Path(os.environ[STATE_ENV]) + state = json.loads(path.read_text()) + state["n"] += 1 + path.write_text(json.dumps(state)) + if state.get("limit") is not None and state["n"] > state["limit"]: + raise RuntimeError("container lost") + if state.get("mark_at") is not None and state["n"] == state["mark_at"]: + Path(state["marker"]).write_text("{}") + return 0.5 + + +class _Calls: + def __init__(self, path: Path) -> None: + self.path = path + self.update(n=0) + + def update(self, **fields) -> None: + state = json.loads(self.path.read_text()) if self.path.exists() else {} + self.path.write_text(json.dumps({**state, **fields})) + + +@pytest.fixture(autouse=True) +def _reset(monkeypatch, tmp_path): + global CALLS + monkeypatch.setenv(STATE_ENV, str(tmp_path / "reward-state.json")) + CALLS = _Calls(tmp_path / "reward-state.json") + monkeypatch.setenv(RECLAIM_MARKER_ENV, str(tmp_path / "reclaim-marker")) + handlers = {s: signal.getsignal(s) for s in (signal.SIGINT, signal.SIGTERM)} + yield + for sig, handler in handlers.items(): + signal.signal(sig, handler) + + +def _spec(tmp_path: Path, rounds: int = 4) -> dict: + return _sampler_spec(tmp_path, num_rollouts=rounds, reward="tests.test_s25_g1c_sampler_resume:counting_reward") + + +def _buffer(tmp_path: Path) -> FileSampleBuffer: + buffer = FileSampleBuffer(CheckpointStore(tmp_path / "store"), batch_size=1) + buffer.refresh() + return buffer + + +def _payloads(tmp_path: Path) -> dict[str, bytes]: + return {p.name: p.read_bytes() for p in (tmp_path / "store" / "sample-buffer" / "payloads").iterdir()} + + +def _tape(path: Path) -> list[dict]: + return [json.loads(line) for line in path.read_text().splitlines()] + + +def test_replacement_continues_after_the_written_rounds_and_overwrites_nothing(tmp_path): + spec = _spec(tmp_path) + CALLS.update(limit=8) # 2 prompts x n=2 per round: the third round fails before it writes + with pytest.raises(RuntimeError, match="container lost"): + entry.run(spec, resource_sampler=lambda emit: None, incarnation="aaaa0001") + first = _payloads(tmp_path) + assert len(first) == 4 and all(name.startswith("ppu4-aaaa0001-r") for name in first) + + CALLS.update(n=0, limit=None) + assert entry.run(spec, resource_sampler=lambda emit: None, incarnation="bbbb0002") == 2 + after = _payloads(tmp_path) + assert {k: after[k] for k in first} == first # the first incarnation's payloads are untouched + assert len(after) == 8 + buffer = _buffer(tmp_path) + entries = buffer.entries() + assert buffer.data_ended and len({e.sample_id for e in entries}) == len(entries) == 8 + rounds = [(e.sampler_round["incarnation"], e.sampler_round["rollout_id"]) for e in entries] + assert rounds == [("aaaa0001", 0)] * 2 + [("aaaa0001", 1)] * 2 + [("bbbb0002", 2)] * 2 \ + + [("bbbb0002", 3)] * 2 + assert all(buffer.read_payload(e.payload_ref) for e in entries) + + tapes = {p.name: _tape(p) for p in (tmp_path / "tape").glob("events.*.jsonl")} + assert set(tapes) == {"events.aaaa0001.jsonl", "events.bbbb0002.jsonl"} + old = {e["event"]: e for e in tapes["events.aaaa0001.jsonl"]} + assert old["rl_sampler_leave"]["reason"] == "INFRA_FAILED" + new = {e["event"]: e for e in tapes["events.bbbb0002.jsonl"]} + assert new["rl_sampler_join"]["next_rollout_id"] == 2 + assert new["rl_sampler_join"]["incarnations"] == ["aaaa0001"] + assert new["rl_sampler_leave"]["reason"] == "FINISHED" + members = sorted(p.name for p in (tmp_path / "store" / MEMBERS_DIR / "ppu4").iterdir()) + assert members == ["aaaa0001.json", "bbbb0002.json"] + + +def test_a_finished_sampler_restarted_generates_nothing(tmp_path): + spec = _spec(tmp_path, rounds=2) + assert entry.run(spec, resource_sampler=lambda emit: None, incarnation="aaaa0001") == 2 + before = _payloads(tmp_path) + assert entry.run(spec, resource_sampler=lambda emit: None, incarnation="bbbb0002") == 0 + assert _payloads(tmp_path) == before + + +def test_a_written_round_is_refused_to_a_new_incarnation(tmp_path): + spec = _spec(tmp_path) + CALLS.update(limit=4) + with pytest.raises(RuntimeError): + entry.run(spec, resource_sampler=lambda emit: None, incarnation="aaaa0001") + from yeto.rl.engine.engine_sampler import EngineRound + from yeto.rl.engine.sampler_resume import resume_point + + resume = resume_point(_buffer(tmp_path).entries(), "ppu4") + assert resume.next_rollout_id == 1 + round_ = EngineRound.__new__(EngineRound) + from yeto.rl.engine.sampler_resume import RoundGuard + + round_.guard = RoundGuard(resume) + with pytest.raises(SamplerResumeError, match="round 0 is already"): + EngineRound.__call__(round_, 0) + + +def test_a_payload_is_never_overwritten(tmp_path): + buffer = FileSampleBuffer(CheckpointStore(tmp_path / "store"), batch_size=1) + buffer.write_payload("ppu4-a-r0-row0", b"one") + with pytest.raises(SampleBufferError, match="never overwritten"): + buffer.write_payload("ppu4-a-r0-row0", b"two") + assert (tmp_path / "store" / "sample-buffer" / "payloads" / "ppu4-a-r0-row0.bin").read_bytes() == b"one" + + +def test_reclaim_marker_stops_writes_and_leaves_preempted(tmp_path): + spec = _spec(tmp_path) + CALLS.update(marker=str(tmp_path / "reclaim-marker"), mark_at=6) # inside round 1 + # round 1 was generated (2 rounds), but none of its groups is written + assert entry.run(spec, resource_sampler=lambda emit: None, incarnation="aaaa0001") == 2 + buffer = _buffer(tmp_path) + assert not buffer.data_ended + assert [e.sampler_round["rollout_id"] for e in buffer.entries()] == [0, 0] + [tape] = (tmp_path / "tape").glob("events.*.jsonl") + leave = [e for e in _tape(tape) if e["event"] == "rl_sampler_leave"] + assert len(leave) == 1 and leave[0]["reason"] == "PREEMPTED" diff --git a/tests/test_stage1_gaps.py b/tests/test_stage1_gaps.py index 9ff34c178..0e4a33896 100644 --- a/tests/test_stage1_gaps.py +++ b/tests/test_stage1_gaps.py @@ -66,9 +66,10 @@ def test_engine_sampler_entry_runs_the_island_from_a_spec(tmp_path): reader = FileSampleBuffer(CheckpointStore(tmp_path / "store"), batch_size=1) reader.refresh() assert len(reader.select("t", 10)) == 4 # 2 rounds x 2 prompts - events = [json.loads(line) for line in (tmp_path / "tape" / "events.jsonl").read_text().splitlines()] + [tape] = (tmp_path / "tape").glob("events.*.jsonl") # one tape per incarnation + events = [json.loads(line) for line in tape.read_text().splitlines()] names = [e["event"] for e in events] - assert names[0] == "rl_sampler_start" and names.count("rl_sampler_engine_round") == 2 + assert names[0] == "rl_sampler_start" and "rl_sampler_join" in names and names.count("rl_sampler_engine_round") == 2 assert all(e["sampler_id"] == "ppu4" for e in events) diff --git a/yeto/rl/engine/engine_sampler.py b/yeto/rl/engine/engine_sampler.py index 7ef136bea..5ce580e3f 100644 --- a/yeto/rl/engine/engine_sampler.py +++ b/yeto/rl/engine/engine_sampler.py @@ -16,6 +16,13 @@ (:class:`~.sampler_engine.Provenance`) in the index entry; 6. write ``rl_sampler_engine_round`` (finish reasons, tokens, provenance key). +Replacement (S25 G1C, the engine-sampler side of C1's :mod:`.sampler_resume`): +each group's sample id carries the incarnation, the round and the prompt row, +so a new incarnation never writes an earlier one's payload name; each index +entry carries a :class:`~.sampler_resume.SamplerRound` record; a round already +in the buffer is never generated or written again (:class:`~.sampler_resume.RoundGuard`); +after the reclaim marker no group is written. + :func:`build_engine_sampler` wires the island: the start checks, the engine start on the base weights with the post-load checks (:mod:`.sampler_checks`; a failure there stops the island), the follower and the loop. The caller @@ -47,6 +54,7 @@ SamplingParams, ) from .sampler_loop import SamplerLoop +from .sampler_resume import RoundGuard, SamplerResume, SamplerRound ROUND_EVENT = "rl_sampler_engine_round" START_EVENT = "rl_sampler_start" @@ -91,7 +99,15 @@ class EngineRound: def __init__(self, engine: SamplerEngine, writer: GroupWriter, *, prompts: Callable[[int], Sequence[PromptItem]], reward: Reward, sampling: SamplingParams, provenance: Provenance, emit: Callable[..., None], - loop_guard: LoopGuard | None = None) -> None: + loop_guard: LoopGuard | None = None, incarnation: str | None = None, + resume: SamplerResume | None = None, + data_cursor: Callable[[int], Mapping[str, int]] | None = None, + reclaimed: Callable[[], bool] = lambda: False) -> None: + """``incarnation``: this process's id (``sampler_resume.new_incarnation``); + with it, sample ids are ``--r-`` and + every entry carries its round record. ``resume``: where the earlier + incarnations got to (its rounds are refused). ``data_cursor(rollout_id)``: + the prompt source's position after the round drew its prompts.""" self.engine = engine self.writer = writer self.prompts = prompts @@ -100,9 +116,24 @@ def __init__(self, engine: SamplerEngine, writer: GroupWriter, *, self.provenance = provenance self.emit = emit self.loop_guard = loop_guard + self.incarnation = None if incarnation is None else str(incarnation) + self.guard = None if resume is None else RoundGuard(resume) + self.data_cursor = data_cursor + self.reclaimed = reclaimed + self.dropped_groups = 0 # groups not written because the island is leaving + + def _round_record(self, rollout_id: int, last: bool) -> dict[str, Any] | None: + if self.incarnation is None: + return None + cursor = None if self.data_cursor is None else self.data_cursor(rollout_id) + return SamplerRound(self.writer.sampler_id, self.incarnation, rollout_id, + data_cursor=cursor, last=last).record() def __call__(self, rollout_id: int) -> None: - items = list(self.prompts(int(rollout_id))) + rollout_id = int(rollout_id) + if self.guard is not None: + self.guard.check(rollout_id) + items = list(self.prompts(rollout_id)) version = self.engine.served_version() if version is None: raise SamplerEngineError("the engine serves no version; it cannot generate") @@ -114,6 +145,7 @@ def __call__(self, rollout_id: int) -> None: record = self.provenance.to_record() finish: Counter[str] = Counter() tokens = 0 + scored = [] for position, (item, group) in enumerate(zip(items, groups)): if len(group) != self.sampling.n: raise SamplerEngineError(f"prompt {item.prompt_id}: {len(group)} completions, " @@ -128,11 +160,23 @@ def __call__(self, rollout_id: int) -> None: item, completion, reward=self.reward(item, completion), version=version, provenance=record, group_index=position, index=position * self.sampling.n + j)) - self.writer.write(samples, group_key=f"r{int(rollout_id)}-{item.prompt_id}", - provenance=record) - self.emit(ROUND_EVENT, rollout_id=int(rollout_id), policy_version=version, - groups=len(items), samples=sum(finish.values()), response_tokens=tokens, - finish=dict(sorted(finish.items())), source=self.provenance.key()) + scored.append((item, samples)) + prefix = "" if self.incarnation is None else f"{self.incarnation}-" + written = 0 + for position, (item, samples) in enumerate(scored): + if self.reclaimed(): # leaving: the replacement owns the shard from now on + self.dropped_groups += len(scored) - position + break + self.writer.write(samples, group_key=f"{prefix}r{rollout_id}-{item.prompt_id}", + provenance=record, + sampler_round=self._round_record(rollout_id, position == len(scored) - 1)) + written += 1 + if written and self.guard is not None: + self.guard.written(rollout_id) + self.emit(ROUND_EVENT, rollout_id=rollout_id, policy_version=version, + groups=written, samples=sum(finish.values()), response_tokens=tokens, + finish=dict(sorted(finish.items())), source=self.provenance.key(), + dropped_groups=len(scored) - written) @dataclass @@ -159,7 +203,12 @@ class EngineSamplerIsland: follower: PolicyFollower round: EngineRound - def run(self, start: int = 0) -> int: + resume: SamplerResume | None = None + + def run(self, start: int | None = None) -> int: + """Rounds from ``start`` (default: the resume point, else 0).""" + if start is None: + start = 0 if self.resume is None else self.resume.next_rollout_id return self.loop.run(start) @@ -168,7 +217,10 @@ def build_engine_sampler(config: EngineSamplerConfig, *, engine: SamplerEngine, prompts: Callable[[int], Sequence[PromptItem]], reward: Reward, checks: PostLoadChecks | None = None, should_stop: Callable[[], bool] = lambda: False, - sleep: Callable[[float], None] | None = None) -> EngineSamplerIsland: + sleep: Callable[[float], None] | None = None, + incarnation: str | None = None, resume: SamplerResume | None = None, + data_cursor: Callable[[int], Mapping[str, int]] | None = None, + ) -> EngineSamplerIsland: """Check the sampling against the engine, start it on the base weights, run the post-load checks, and return the wired island (module doc).""" sampling = config.sampling.check() @@ -188,7 +240,8 @@ def build_engine_sampler(config: EngineSamplerConfig, *, engine: SamplerEngine, follower = PolicyFollower(releases, loader, version=int(config.base_version)) engine_round = EngineRound(engine, GroupWriter(buffer, config.sampler_id), prompts=prompts, reward=reward, sampling=sampling, provenance=provenance, emit=emit, - loop_guard=config.loop_guard) + loop_guard=config.loop_guard, incarnation=incarnation, + resume=resume, data_cursor=data_cursor, reclaimed=should_stop) extra = {} if sleep is None else {"sleep": sleep} loop = SamplerLoop(follower=follower, buffer=buffer, generate=engine_round, num_rollouts=int(config.num_rollouts), @@ -196,7 +249,8 @@ def build_engine_sampler(config: EngineSamplerConfig, *, engine: SamplerEngine, poll_s=config.poll_s, should_stop=should_stop, min_round_s=config.min_round_s, loaded_fields=loader.loaded_fields, ready=lambda: loader.ready, **extra) - return EngineSamplerIsland(loop=loop, loader=loader, follower=follower, round=engine_round) + return EngineSamplerIsland(loop=loop, loader=loader, follower=follower, round=engine_round, + resume=resume) def main(argv: Sequence[str] | None = None) -> int: diff --git a/yeto/rl/engine/engine_sampler_entry.py b/yeto/rl/engine/engine_sampler_entry.py index 1eda00f9b..f842074ba 100644 --- a/yeto/rl/engine/engine_sampler_entry.py +++ b/yeto/rl/engine/engine_sampler_entry.py @@ -31,7 +31,17 @@ ``reward_kind`` is ``"miles"``; * ``checks``: :class:`~.sampler_checks.PostLoadChecks` keyword overrides (optional; e.g. ``expected_expert_matrices``); -* ``event_tape``: the JSONL tape file of the island's events. +* ``event_tape``: the JSONL tape file of the island's events; each + incarnation writes its own file next to it (``events..jsonl``), + so a replacement container never overwrites an earlier one's tape. + +Replacement (S25 G1C): the process is one incarnation of the sampler +(``sampler_resume.new_incarnation``). It reads where the earlier incarnations +got to from the buffer (``sampler_resume.resume_point``), starts at the next +round, and writes ``rl_sampler_join`` / ``rl_sampler_leave`` plus one member +record per incarnation in the buffer store. With ``YETO_SPOT_RECLAIM_MARKER`` +set, SIGINT/SIGTERM write the reclaim marker; the island then starts no new +round, writes no more groups and leaves as PREEMPTED. Nothing here was run on a PPU or NPU node. """ @@ -107,6 +117,11 @@ def __init__(self, path: str | Path, *, input_key: str, label_key: str | None, g self.groups_per_round = int(groups_per_round) self.render = render + def cursor(self, rollout_id: int) -> dict[str, int]: + """The data position after round ``rollout_id`` drew its prompts (the + ``sampler_round`` record; rounds draw by rollout id, so it only documents).""" + return {"row_offset": (int(rollout_id) + 1) * self.groups_per_round} + def __call__(self, rollout_id: int) -> list[PromptItem]: start = int(rollout_id) * self.groups_per_round items = [] @@ -191,6 +206,12 @@ def build_stores(spec: Mapping[str, Any], environ: Mapping[str, str] | None = No return releases, FileSampleBuffer(buffer_store, batch_size=1, index=index) # as open_buffer(sampler=True) +def incarnation_tape(path: str | Path, incarnation: str) -> Path: + """``events.jsonl`` -> ``events..jsonl`` (one tape per incarnation).""" + tape = Path(path).expanduser() + return tape.with_name(f"{tape.stem}.{incarnation}{tape.suffix or '.jsonl'}") + + def tape_writer(path: str | Path, sampler_id: str) -> Callable[..., None]: from yeto.rl.event_echo import append_record # noqa: PLC0415 @@ -219,16 +240,52 @@ def island_config(spec: Mapping[str, Any]) -> EngineSamplerConfig: def run(spec: Mapping[str, Any], *, environ: Mapping[str, str] | None = None, - resource_sampler: Callable[[Callable[..., None]], Any] | None = None) -> int: - """Build every part of the island from ``spec``, run it, stop the engine; - returns the rounds the loop ran. ``resource_sampler(emit)`` returns a - started telemetry thread (default: ``telemetry.ResourceSampler`` of the - island's device family; None from the factory = no telemetry).""" + resource_sampler: Callable[[Callable[..., None]], Any] | None = None, + incarnation: str | None = None) -> int: + """Build every part of the island from ``spec``, run it from the resume + point, stop the engine; returns the rounds this incarnation ran. + ``resource_sampler(emit)`` returns a started telemetry thread (default: + ``telemetry.ResourceSampler`` of the island's device family; None from the + factory = no telemetry).""" + import os # noqa: PLC0415 + + from yeto.rl.elastic_client import LeaveReason # noqa: PLC0415 + + from .sampler_resume import ( # noqa: PLC0415 + RECLAIM_MARKER_ENV, + ReclaimWatcher, + SamplerMembership, + install_marker_signals, + new_incarnation, + reclaim_marker, + reclaim_requested, + resume_point, + ) + + env = os.environ if environ is None else environ config = island_config(spec) - emit = tape_writer(_require(spec, "event_tape"), config.sampler_id) + incarnation = incarnation or new_incarnation() + tape = incarnation_tape(_require(spec, "event_tape"), incarnation) + emit = tape_writer(tape, config.sampler_id) releases, buffer = build_stores(spec, environ) prompts, reward = build_prompts(spec), build_reward(spec) checks = PostLoadChecks(**dict(spec.get("checks") or {})) + reclaim_marker(env).unlink(missing_ok=True) # a restarted machine keeps its /tmp + if env.get(RECLAIM_MARKER_ENV): + install_marker_signals(env) + buffer.refresh() + resume = resume_point(buffer.entries(), config.sampler_id) + from .resume import store_for # noqa: PLC0415 + + member_store = store_for(str(spec.get("buffer_root") or spec["store"]), environ=environ) + membership = SamplerMembership(member_store, config.sampler_id, incarnation, resume, emit) + reclaimed = lambda: reclaim_requested(env) # noqa: E731 + island = None + + def leave(reason: LeaveReason, **detail: Any) -> None: + version = None if island is None else island.follower.version + membership.leave(reason.name, policy_version=version, buffer_end=len(buffer), **detail) + engine = build_engine(spec) if resource_sampler is None: from .telemetry import ResourceSampler # noqa: PLC0415 @@ -237,11 +294,28 @@ def resource_sampler(sink: Callable[..., None]) -> Any: return ResourceSampler(lambda event, **fields: sink(event, **fields)).start() telemetry = resource_sampler(emit) + watcher = ReclaimWatcher(reclaimed, lambda: leave(LeaveReason.PREEMPTED)).start() try: island = build_engine_sampler(config, engine=engine, releases=releases, buffer=buffer, emit=emit, - prompts=prompts, reward=reward, checks=checks) - return island.run() + prompts=prompts, reward=reward, checks=checks, + should_stop=reclaimed, incarnation=incarnation, resume=resume, + data_cursor=getattr(prompts, "cursor", None)) + latest = releases.latest() + membership.join(tape=str(tape), latest_release=None if latest is None else + {"seq": latest.seq, "policy_version": latest.version}) + rounds = 0 if buffer.data_ended else island.run() + leave(LeaveReason.PREEMPTED if reclaimed() else LeaveReason.FINISHED, rounds=rounds, + dropped_groups=island.round.dropped_groups) + return rounds + except BaseException as error: + try: + leave(LeaveReason.PREEMPTED if reclaimed() else LeaveReason.INFRA_FAILED, + error=f"{type(error).__name__}: {error}"[:300]) + except Exception: # noqa: BLE001 - the original error is the one to report + pass + raise finally: + watcher.stop() try: engine.stop() finally: diff --git a/yeto/rl/engine/sample_buffer.py b/yeto/rl/engine/sample_buffer.py index b3fa1b36d..aa2715883 100644 --- a/yeto/rl/engine/sample_buffer.py +++ b/yeto/rl/engine/sample_buffer.py @@ -426,8 +426,14 @@ def __init__(self, store: CheckpointStore | ObjectStore, *, batch_size: int, # ---------------------------------------------------------------- writing def write_payload(self, sample_id: str, data: bytes) -> str: - """Store one payload; returns the ``payload_ref`` for its index entry.""" - return self.index.write_payload(payload_name(sample_id), data) + """Store one payload; returns the ``payload_ref`` for its index entry. + A payload is written once: a name already in the store is refused (a + replacement sampler must never overwrite a group an index entry names).""" + name = payload_name(sample_id) + if self.index.has_payload(f"{PAYLOADS}/{name}"): + raise SampleBufferError(f"payload {name} is already in the sample buffer; " + "a payload is never overwritten") + return self.index.write_payload(name, data) def read_payload(self, payload_ref: str) -> bytes: return self.index.read_payload(payload_ref) From 0e29001bb43b6440e4ce0c2ca346ca0ebc77d4ae Mon Sep 17 00:00:00 2001 From: michael <78113554+chowhsu@users.noreply.github.com> Date: Tue, 20 Oct 2026 11:28:02 +0000 Subject: [PATCH 81/89] Miles buffer reader tapes passed-over stale groups (rl_buffer_stale_skipped); eval E4 reads it and cross-checks consumed minus trained Co-Authored-By: Claude Opus 5.5 --- tests/test_rl_buffer_rollout.py | 26 +++++++++++ tests/test_rl_miles_adapter_rollout.py | 6 ++- tests/test_rl_policy_age.py | 24 ++++++++++ tests/test_tools_eval.py | 44 +++++++++++++++++++ tools/eval/judges.py | 21 ++++++--- .../rl/adapters/miles/buffer_rollout_miles.py | 4 +- yeto/rl/adapters/miles/rollout.py | 3 ++ yeto/rl/adapters/miles/rollout_meta_hook.py | 4 +- yeto/rl/engine/driver.py | 6 +++ yeto/rl/engine/ports.py | 3 ++ 10 files changed, 131 insertions(+), 10 deletions(-) diff --git a/tests/test_rl_buffer_rollout.py b/tests/test_rl_buffer_rollout.py index 1079914c9..d8febda6a 100644 --- a/tests/test_rl_buffer_rollout.py +++ b/tests/test_rl_buffer_rollout.py @@ -370,3 +370,29 @@ def test_full_parameter_trainer_stops_at_a_scheduled_cut_and_continues_without_r assert trained == list(range(2 * rounds)) # every group once, in order: none trained twice for name, value in params(reference).items(): assert torch.equal(value, params(rank_b)[name]), name + + +def test_reader_round_reports_its_stale_groups_as_round_metadata(monkeypatch): + """T-EVAL E4: the groups the reader passed over this round reach the driver as + per-round metadata (driver event rl_buffer_stale_skipped), before the + end-of-round hooks put the round's main record.""" + from types import SimpleNamespace + + calls = [] + registry = {"all": lambda args, samples, source: calls.append(("all",))} + plugin = _import_plugin_with_stub_miles(monkeypatch, registry) + monkeypatch.setattr(plugin, "record_round_metadata", + lambda args, rid, **counts: calls.append(("meta", rid, counts))) + monkeypatch.setattr(plugin, "Sample", SimpleNamespace(from_dict=lambda d: d)) + monkeypatch.setattr(plugin, "RolloutFnTrainOutput", lambda **kw: kw) + monkeypatch.setattr(plugin, "offpolicy_source_event", lambda *a, **kw: {}) + monkeypatch.setattr(plugin, "flat_metrics", lambda source: {}) + group = [sample(0, 0, 4), sample(1, 0, 4)] + report = br.TakeReport(rollout_id=2, groups=1, stale_groups=3, start=0, end=4) + reader = plugin._BufferRound.__new__(plugin._BufferRound) + reader.args = SimpleNamespace(rollout_sample_filter_path=None, rollout_all_samples_process_path="all") + reader.groups, reader.samples_per_group = 1, 2 + reader.reader = SimpleNamespace(take=lambda rid, n, current_version: ([(3, group)], report)) + out = reader(SimpleNamespace(evaluation=False, rollout_id=2, weight_version=4)) + assert calls == [("meta", 2, {"buffer_stale_groups": 3}), ("all",)] + assert out["metrics"]["sample_buffer/stale_groups"] == 3.0 diff --git a/tests/test_rl_miles_adapter_rollout.py b/tests/test_rl_miles_adapter_rollout.py index 0ecdf8b99..366138820 100644 --- a/tests/test_rl_miles_adapter_rollout.py +++ b/tests/test_rl_miles_adapter_rollout.py @@ -262,7 +262,7 @@ async def get(self, rollout_id): result = await super().get(rollout_id) nonzero = counts[rollout_id] # computed during reward post-processing hook.record_round_metadata(self.args, self.kept, sink=self.sink, - nonzero_advantages=nonzero) + nonzero_advantages=nonzero, buffer_stale_groups=nonzero % 5) return result kept = [group(0, [1.0, 0.0]), group(1, [0.0, 1.0])] @@ -278,7 +278,9 @@ def expected(): seen = {} for rid in (3, 4, 5): current["v"] = next(versions) - seen[rid] = p.generate(rid).nonzero_advantages + batch = p.generate(rid) + seen[rid] = batch.nonzero_advantages + assert batch.buffer_stale_groups == counts[rid] % 5 assert seen == counts diff --git a/tests/test_rl_policy_age.py b/tests/test_rl_policy_age.py index 3f82926f6..bc9b5b1e0 100644 --- a/tests/test_rl_policy_age.py +++ b/tests/test_rl_policy_age.py @@ -227,3 +227,27 @@ def test_miles_switches_are_derived_from_the_limit(): assert row.translate(AlgorithmSpec(execution=ExecutionSpec(max_policy_staleness=1))) == [] with pytest.raises((AlgorithmSpecError, ValueError), match="rl-max-policy-age"): af.absorb_extra_argv(AlgorithmSpec(), ["--partial-rollout"]) + + +def test_driver_tapes_the_buffer_readers_stale_groups_per_round(tmp_path): + """T-EVAL E4: groups a sample-buffer reader passed over as too old reach the + tape as rl_buffer_stale_skipped (one per round, 0 included); an engine that + does not report them tapes nothing.""" + import dataclasses + + stale = {0: 0, 1: 3, 2: 1} + driver = _driver(tmp_path, 1) + original = driver.rollout.generate + + def gen(rid, **kw): + return dataclasses.replace(original(rid, **kw), buffer_stale_groups=stale[rid]) + + driver.rollout.generate = gen + driver.run() + events = [json.loads(l) for l in (tmp_path / "e.jsonl").read_text().splitlines()] + skipped = [e for e in events if e["event"] == "rl_buffer_stale_skipped"] + assert [(e["rollout_id"], e["stale_groups"], e["max_policy_age"]) for e in skipped] == [ + (0, 0, 1), (1, 3, 1), (2, 1, 1)] + (tmp_path / "x").mkdir() + _driver(tmp_path / "x", 1).run() + assert '"rl_buffer_stale_skipped"' not in (tmp_path / "x" / "e.jsonl").read_text() diff --git a/tests/test_tools_eval.py b/tests/test_tools_eval.py index f5685860a..c15de19ed 100644 --- a/tests/test_tools_eval.py +++ b/tests/test_tools_eval.py @@ -89,3 +89,47 @@ def test_run_refuses_unregistered_run_id(tmp_path, monkeypatch): monkeypatch.setattr(cli, "LEDGER", ledger) with pytest.raises(SystemExit, match="not pre-registered"): cli._preflight("eval-unregistered") + + +class _E4Evidence: + """Just the event reads judge_e4 makes (one Miles trainer, one sampler).""" + backend = "miles" + + def __init__(self, trainer_events, sampler_events): + self.events = {"trainer": trainer_events, "sampler-0": sampler_events} + + def trainers(self): + return ["trainer"] + + def samplers(self): + return ["sampler-0"] + + def of(self, component, name): + return [e for e in self.events.get(component, []) if e["event"] == name] + + +@pytest.mark.parametrize("stale,consumed,checks", [ + ((0, 16, 0, 0), 80, {"stale_event": True, "stale_cross_check": True, "over_age_drop": True}), + ((0, 15, 0, 0), 80, {"stale_event": True, "stale_cross_check": False, "over_age_drop": True}), + ((), 80, {"stale_event": False, "stale_cross_check": False, "over_age_drop": False}), +]) +def test_e4_miles_reads_the_stale_event_and_cross_checks_consumed_minus_trained(monkeypatch, stale, consumed, + checks): + from tools.eval import judges + + for name in ("_rc_ok", "_gpu"): + monkeypatch.setattr(judges, name, lambda *a, **kw: None) + monkeypatch.setattr(judges, "_round_metrics", lambda ev, t: {}) + rounds = 4 + trainer = [{"event": "rl_round_trained", "rollout_id": r, "mismatch": {"tis": 1.0, "tis_abs": 0.01}} + for r in range(rounds)] + trainer += [{"event": "rl_local_round", "base_policy_version": r} for r in range(rounds)] + trainer += [{"event": "rl_round_labels", "rollout_id": r, "rl/groups": 16} for r in range(rounds)] + trainer += [{"event": "rl_buffer_stale_skipped", "rollout_id": r, "stale_groups": n, "max_policy_age": 1} + for r, n in enumerate(stale)] + trainer += [{"event": "rl_policy_released", "policy_version": rounds, "buffer_selected": consumed}] + monkeypatch.setattr(judges, "_trained_rounds", lambda ev, t: list(range(rounds))) + sampler = [{"event": "rl_sampler_policy_loaded", "policy_version": v} for v in (0, 1)] + v = Verdict("E4", "miles", "r", True) + judges.judge_e4(_E4Evidence(trainer, sampler), v, {"rounds": rounds, "max_policy_age": 1}) + assert {c.id: c.passed for c in v.criteria if c.id in checks} == checks diff --git a/tools/eval/judges.py b/tools/eval/judges.py index aa6650fde..38a0b7c28 100644 --- a/tools/eval/judges.py +++ b/tools/eval/judges.py @@ -217,14 +217,22 @@ def judge_e4(ev: Evidence, v: Verdict, cfg: dict) -> None: {"ratio_or_tis": ratio}) # two distinct drop points, each counted once: the round collector's over-age discard # (rl_rollout_carry_over, both backends) and the Miles buffer reader passing over - # stale groups, which no yeto event reports: buffer positions consumed (the last - # release's buffer_selected) minus groups trained (round labels; the S25 C1 C6 count) + # stale groups (rl_buffer_stale_skipped, one per round) carried = numeric_sum(ev.of(t, "rl_rollout_carry_over"), "over_age_discarded_groups") - passed_over = 0.0 + skipped = ev.of(t, "rl_buffer_stale_skipped") + passed_over = numeric_sum(skipped, "stale_groups") + # cross-check (the reading before the event existed): buffer positions consumed (the + # last release's buffer_selected) minus groups trained (round labels; S25 C1 C6 count) + legacy = None releases = ev.of(t, "rl_policy_released") if ev.backend == "miles" and releases and releases[-1].get("buffer_selected") is not None: - trained_groups = numeric_sum(ev.of(t, "rl_round_labels"), "rl/groups") - passed_over = float(releases[-1]["buffer_selected"]) - trained_groups + legacy = float(releases[-1]["buffer_selected"]) - numeric_sum(ev.of(t, "rl_round_labels"), "rl/groups") + if ev.backend == "miles": + v.check("stale_event", "the Miles buffer reader taped rl_buffer_stale_skipped every round", + len(skipped) == cfg["rounds"], len(skipped)) + if legacy is not None: + v.check("stale_cross_check", "taped passed-over groups = consumed positions - trained groups", + passed_over == legacy, {"event": passed_over, "consumed_minus_trained": legacy}) dropped = carried + passed_over v.check("over_age_drop", "at least one group older than max_policy_age was dropped, not trained", dropped >= 1, dropped) @@ -235,7 +243,8 @@ def judge_e4(ev: Evidence, v: Verdict, cfg: dict) -> None: len({e.get("policy_version") for e in loads}) >= 2, sorted({e.get("policy_version") for e in loads})) _rc_ok(ev, v, ignore=("sampler_launch", "sampler-0", "sampler-1")) _gpu(ev, v, cfg) - v.metrics.update(over_age_discarded_groups=carried, stale_passed_over_groups=passed_over, ratio_p50=ratio, + v.metrics.update(over_age_discarded_groups=carried, stale_passed_over_groups=passed_over, + stale_consumed_minus_trained=legacy, ratio_p50=ratio, **_round_metrics(ev, t)) diff --git a/yeto/rl/adapters/miles/buffer_rollout_miles.py b/yeto/rl/adapters/miles/buffer_rollout_miles.py index 60f43e841..2f285e2b8 100644 --- a/yeto/rl/adapters/miles/buffer_rollout_miles.py +++ b/yeto/rl/adapters/miles/buffer_rollout_miles.py @@ -44,7 +44,7 @@ register_reader, renumber_group, ) -from .rollout_meta_hook import data_cursor +from .rollout_meta_hook import data_cursor, record_round_metadata def _flat(group: Any) -> list[Any]: @@ -136,6 +136,8 @@ def __call__(self, input: RolloutFnInput) -> RolloutFnTrainOutput: samples = [[Sample.from_dict(d) for d in renumber_group( group, position=position, samples_per_group=self.samples_per_group)] for position, group in groups] + record_round_metadata(self.args, int(input.rollout_id), + buffer_stale_groups=int(report.stale_groups)) self._run_round_hooks(samples) return RolloutFnTrainOutput(samples=samples, metrics={**report.metrics(), **flat_metrics(source)}) diff --git a/yeto/rl/adapters/miles/rollout.py b/yeto/rl/adapters/miles/rollout.py index f5490f47f..f86d1ed6d 100644 --- a/yeto/rl/adapters/miles/rollout.py +++ b/yeto/rl/adapters/miles/rollout.py @@ -476,6 +476,9 @@ def handle_from_metadata( nonzero_advantages=( int(payload["nonzero_advantages"]) if "nonzero_advantages" in payload else None ), + buffer_stale_groups=( + int(payload["buffer_stale_groups"]) if "buffer_stale_groups" in payload else None + ), policy_age_violation=( int(payload["policy_age_violation"]) if "policy_age_violation" in payload else None ), diff --git a/yeto/rl/adapters/miles/rollout_meta_hook.py b/yeto/rl/adapters/miles/rollout_meta_hook.py index bcb3440df..d34e761d9 100644 --- a/yeto/rl/adapters/miles/rollout_meta_hook.py +++ b/yeto/rl/adapters/miles/rollout_meta_hook.py @@ -54,7 +54,9 @@ # rollout's metadata by the driver-side reader, so they are never attributed # to the next round. ROUND_META_SCHEMA = "yeto-rl-round-metadata-v1" -ROUND_METADATA_KEYS = frozenset({"nonzero_advantages"}) +# 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"}) def current_round_id(samples: Sequence[Any] = (), sink: str | None = None) -> int | None: diff --git a/yeto/rl/engine/driver.py b/yeto/rl/engine/driver.py index c4ecee2c4..a05ee2f16 100644 --- a/yeto/rl/engine/driver.py +++ b/yeto/rl/engine/driver.py @@ -897,6 +897,12 @@ def _generate(self, rollout_id: int) -> RolloutBatchHandle: f"(policy v{batch.policy_version}) for rollout {rollout_id}" ) batch = self._with_policy_labels(batch) + stale = getattr(batch, "buffer_stale_groups", None) + if stale is not None: + # The sample-buffer reader consumed these groups without training + # them (older than max_policy_age); one event per round, 0 included. + self.emit("rl_buffer_stale_skipped", rollout_id=rollout_id, + stale_groups=int(stale), max_policy_age=self._max_policy_age()) violations = getattr(batch, "policy_age_violation", None) if violations: # IR-3 (age 0): a sample generated by a policy other than the diff --git a/yeto/rl/engine/ports.py b/yeto/rl/engine/ports.py index 9399d2bcf..c5a514603 100644 --- a/yeto/rl/engine/ports.py +++ b/yeto/rl/engine/ports.py @@ -103,6 +103,9 @@ class RolloutBatchHandle: # (rl-algo-seq-and-adv R2); None = not reported. Read by per-algorithm # gradient expectations (``gradient_expectation(batch, metrics)``). nonzero_advantages: int | None = None + # Groups a sample-buffer reader passed over this round as older than + # max_policy_age (consumed, never trained); None = not a buffer reader. + buffer_stale_groups: int | None = None # rl-infra-spec 1.7: summed non-generation (tool) time of the rollout's # samples; None = not reported. tool_wait_seconds: float | None = None From 58329f910febc53de6f6ecf5e47a28c17378d72f Mon Sep 17 00:00:00 2001 From: michael <78113554+chowhsu@users.noreply.github.com> Date: Tue, 20 Oct 2026 11:30:52 +0000 Subject: [PATCH 82/89] Eval E2-verl runs the production engine sampler entry: resume after a reclaim, hold until the container stops; E2 checks continuous trainer versions - engine_sampler_entry.run: optional caller stop (leaves FINISHED) and on_event hook; build_engine_sampler separates the loop stop from the round reclaim. - engine_sampler_entry.hold_after_reclaim: a reclaimed process waits for the cloud to stop the container instead of exiting cleanly (main() calls it). - tools/eval: E2-verl is supported; the sampler function sets the reclaim marker and calls entry.run; judge_e2 adds the versions check. Co-Authored-By: Claude Opus 5.5 --- tests/test_s25_g1c_sampler_resume.py | 26 ++++++++++++++++ tests/test_tools_eval.py | 2 +- tools/eval/cells.py | 7 ++--- tools/eval/judges.py | 3 ++ tools/eval/modal_verl.py | 42 ++++++++++++++------------ yeto/rl/engine/engine_sampler.py | 8 +++-- yeto/rl/engine/engine_sampler_entry.py | 39 +++++++++++++++++++++--- 7 files changed, 96 insertions(+), 31 deletions(-) diff --git a/tests/test_s25_g1c_sampler_resume.py b/tests/test_s25_g1c_sampler_resume.py index c8995bcd2..88ad0fa05 100644 --- a/tests/test_s25_g1c_sampler_resume.py +++ b/tests/test_s25_g1c_sampler_resume.py @@ -150,3 +150,29 @@ def test_reclaim_marker_stops_writes_and_leaves_preempted(tmp_path): [tape] = (tmp_path / "tape").glob("events.*.jsonl") leave = [e for e in _tape(tape) if e["event"] == "rl_sampler_leave"] assert len(leave) == 1 and leave[0]["reason"] == "PREEMPTED" + + +def test_caller_stop_ends_the_island_finished_and_on_event_sees_the_tape(tmp_path): + """T-E2V: the eval harness's end condition (trainers done) stops the loop + without a reclaim: the round in progress is still written and the island + leaves FINISHED; on_event sees every record written to the tape.""" + spec = _spec(tmp_path) + seen = [] + rounds = entry.run(spec, resource_sampler=lambda emit: None, incarnation="aaaa0001", + stop=lambda: any(e == "rl_sampler_round" for e, _ in seen), + on_event=lambda event, **fields: seen.append((event, fields))) + assert rounds == 1 + assert [e.sampler_round["rollout_id"] for e in _buffer(tmp_path).entries()] == [0, 0] + [tape] = (tmp_path / "tape").glob("events.*.jsonl") + assert [e["event"] for e in _tape(tape)] == [e for e, _ in seen] + assert [f["reason"] for e, f in seen if e == "rl_sampler_leave"] == ["FINISHED"] + + +def test_hold_after_reclaim_only_when_the_marker_is_there(tmp_path): + """A reclaimed engine sampler process does not exit cleanly at once (the cloud + could end its input as finished instead of running it again).""" + env = {RECLAIM_MARKER_ENV: str(tmp_path / "marker")} + slept = [] + assert entry.hold_after_reclaim(env, hold_s=5.0, sleep=slept.append) is False and slept == [] + (tmp_path / "marker").write_text("{}") + assert entry.hold_after_reclaim(env, hold_s=5.0, sleep=slept.append) is True and slept == [5.0] diff --git a/tests/test_tools_eval.py b/tests/test_tools_eval.py index c15de19ed..a0a56d2a5 100644 --- a/tests/test_tools_eval.py +++ b/tests/test_tools_eval.py @@ -17,7 +17,7 @@ def test_matrix_covers_every_capability_and_backend(): assert {(c.capability, c.backend) for c in cells.CELLS.values()} == { (cap, b) for cap in cells.CAPABILITIES for b in cells.BACKENDS} - assert [c.key for c in cells.CELLS.values() if not c.supported] == ["E2-verl"] + assert [c.key for c in cells.CELLS.values() if not c.supported] == [] with pytest.raises(SystemExit): cells.get("E9", "verl") diff --git a/tools/eval/cells.py b/tools/eval/cells.py index 038b032e1..ec6cb8fec 100644 --- a/tools/eval/cells.py +++ b/tools/eval/cells.py @@ -55,11 +55,10 @@ def key(self) -> str: CELLS: dict[str, Cell] = {c.key: c for c in ( Cell("E1", "verl", "modal_verl", 1, {**_VERL_ISLAND, "phases": ["c"]}), Cell("E1", "miles", "miles_local", 1, {**_MILES, "rounds": 4, "islands": 1}), - # E2: the S25 C1 shape (two samplers, one reclaimed after it wrote round 1) on Miles; the - # neutral engine sampler (verl's sampler island) has no rejoin path: expected to fail + # E2: the S25 C1 shape (two samplers, one reclaimed after it wrote round 1) on Miles; verl's + # sampler island is the neutral engine sampler, which resumes after a reclaim since S25 G1C Cell("E2", "verl", "modal_verl", 2, - {**_VERL, "kind": "buffer", "max_policy_age": 2, "max_backlog_groups": 12, "reclaim_after_round": 1}, - supported=False, note="engine sampler island: no LEAVE/JOIN, no resume after a reclaim (S25 C1 is Miles only)"), + {**_VERL, "kind": "buffer", "max_policy_age": 2, "max_backlog_groups": 12, "reclaim_after_round": 1}), Cell("E2", "miles", "miles_local", 3, {**_MILES, "rounds": 8, "islands": 3, "samplers": 2, "max_policy_age": 4, "reclaim_label": "l0", "reclaim_round": 1, "reclaim_component": "l0"}), diff --git a/tools/eval/judges.py b/tools/eval/judges.py index 38a0b7c28..7c583d18d 100644 --- a/tools/eval/judges.py +++ b/tools/eval/judges.py @@ -141,6 +141,9 @@ def judge_e2(ev: Evidence, v: Verdict, cfg: dict) -> None: during = [e for e in ev.of(t, "rl_round_trained") if stop_t and joins and len(joins) >= 2 and stop_t < float(e["time_unix"]) < float(joins[-1]["time_unix"])] v.check("trainer", f"the trainer trained all {cfg['rounds']} rounds", len(trained) == cfg["rounds"], trained) + bases = [r.get("base_policy_version") for r in ev.of(t, "rl_local_round")] + v.check("versions", "the trainer's policy version advanced by one every round (0..R-1, no restart)", + bases == list(range(cfg["rounds"])), bases) v.metrics.update(trainer_rounds_during_replacement=len(during), sampler_rounds=len(rounds), buffer_groups=len(ids)) _rc_ok(ev, v) diff --git a/tools/eval/modal_verl.py b/tools/eval/modal_verl.py index a0839afff..67afe1072 100644 --- a/tools/eval/modal_verl.py +++ b/tools/eval/modal_verl.py @@ -162,7 +162,18 @@ def sampler(run: str, cfg: dict) -> dict: began = time.time() meta: dict[str, Any] = {"gpu": names, "container": os.environ.get("MODAL_TASK_ID"), "cfg": cfg, "incarnation": incarnation} - tape = Path("/tmp/sampler-tape.jsonl") + from yeto.rl.engine import engine_sampler_entry as entry + from yeto.rl.engine.sampler_resume import ( + DEFAULT_RECLAIM_MARKER, + RECLAIM_MARKER_ENV, + new_incarnation, + reclaim_requested, + ) + + os.environ[RECLAIM_MARKER_ENV] = DEFAULT_RECLAIM_MARKER # the reclaim signal writes the marker + island_incarnation = new_incarnation() + tape = Path("/tmp/sampler-tape.jsonl") # the entry writes sampler-tape..jsonl + tape = entry.incarnation_tape(tape, island_incarnation) try: import pandas as pd import sglang @@ -179,10 +190,6 @@ def sampler(run: str, cfg: dict) -> dict: text = tok.apply_chat_template(messages, tokenize=False, add_generation_prompt=True, enable_thinking=False) handle.write(json.dumps({"text": text, "label": str(row["label"])}) + "\n") - from yeto.rl.engine import engine_sampler_entry as entry - from yeto.rl.engine.engine_sampler import build_engine_sampler - from yeto.rl.engine.sampler_checks import PostLoadChecks - spec = { "sampler_id": "eval-sglang-0", "engine": {"factory": "yeto.rl.adapters.sampler_engines.from_spec:build", "adapter": "sglang-server", @@ -201,39 +208,34 @@ def sampler(run: str, cfg: dict) -> dict: "groups_per_round": cfg["groups_per_round"], "format": "text"}, "reward": "yeto.rl.gsm8k_reward:score", "reward_kind": "miles", "checks": {"expected_expert_matrices": None, "fixed_prompts": []}, - "event_tape": str(tape), + "event_tape": "/tmp/sampler-tape.jsonl", } (out / "spec.json").write_text(json.dumps(spec, indent=1)) - config = entry.island_config(spec) - write = entry.tape_writer(spec["event_tape"], config.sampler_id) + meta["island_incarnation"] = island_incarnation - def emit(event: str, **fields: Any) -> None: - write(event, **fields) + def on_event(event: str, **fields: Any) -> None: if event == "rl_sampler_round": # the local driver's reclaim trigger control[f"{run}/sampler/round"] = {"incarnation": incarnation, **{ k: fields.get(k) for k in ("rollout_id", "groups")}} - releases, buffer = entry.build_stores(spec) - engine = entry.build_engine(spec) done_keys = [f"{run}/done/{i}" for i in range(cfg["learners"])] deadline = began + cfg["timeout_s"] - 300 - def should_stop() -> bool: + def trainers_done() -> bool: return all(control.get(k, False) for k in done_keys) or time.time() > deadline - try: - island = build_engine_sampler(config, engine=engine, releases=releases, buffer=buffer, emit=emit, - prompts=entry.build_prompts(spec), reward=entry.build_reward(spec), - checks=PostLoadChecks(**spec["checks"]), should_stop=should_stop) - meta["rounds"] = island.run() - finally: - engine.stop() + # the production entry: resumes after the rounds earlier incarnations wrote, + # turns the reclaim signal into the marker and leaves PREEMPTED (S25 G1C) + meta["rounds"] = entry.run(spec, incarnation=island_incarnation, stop=trainers_done, on_event=on_event) + meta["reclaimed"] = reclaim_requested(os.environ) except BaseException as exc: # noqa: BLE001 - recorded, then re-raised meta["error"] = f"{type(exc).__name__}: {exc}" raise finally: nvml.terminate() _finish(out, meta, began, {"sampler-tape.jsonl": tape}) + # reclaimed: stay up until Modal stops the container, so Modal runs the input again + entry.hold_after_reclaim(hold_s=900.0) return meta diff --git a/yeto/rl/engine/engine_sampler.py b/yeto/rl/engine/engine_sampler.py index 5ce580e3f..1d803e2fd 100644 --- a/yeto/rl/engine/engine_sampler.py +++ b/yeto/rl/engine/engine_sampler.py @@ -220,9 +220,12 @@ def build_engine_sampler(config: EngineSamplerConfig, *, engine: SamplerEngine, sleep: Callable[[float], None] | None = None, incarnation: str | None = None, resume: SamplerResume | None = None, data_cursor: Callable[[int], Mapping[str, int]] | None = None, + reclaimed: Callable[[], bool] | None = None, ) -> EngineSamplerIsland: """Check the sampling against the engine, start it on the base weights, - run the post-load checks, and return the wired island (module doc).""" + run the post-load checks, and return the wired island (module doc). + ``should_stop``: the loop starts no new round. ``reclaimed``: also stop + writing the current round's groups (default: ``should_stop``).""" sampling = config.sampling.check() check_sampling_fits(sampling, engine, max_prompt_tokens=config.max_prompt_tokens) provenance = Provenance.of(engine.identity, sampler_id=config.sampler_id, sampling=sampling) @@ -241,7 +244,8 @@ def build_engine_sampler(config: EngineSamplerConfig, *, engine: SamplerEngine, engine_round = EngineRound(engine, GroupWriter(buffer, config.sampler_id), prompts=prompts, reward=reward, sampling=sampling, provenance=provenance, emit=emit, loop_guard=config.loop_guard, incarnation=incarnation, - resume=resume, data_cursor=data_cursor, reclaimed=should_stop) + resume=resume, data_cursor=data_cursor, + reclaimed=should_stop if reclaimed is None else reclaimed) extra = {} if sleep is None else {"sleep": sleep} loop = SamplerLoop(follower=follower, buffer=buffer, generate=engine_round, num_rollouts=int(config.num_rollouts), diff --git a/yeto/rl/engine/engine_sampler_entry.py b/yeto/rl/engine/engine_sampler_entry.py index f842074ba..aed7075c7 100644 --- a/yeto/rl/engine/engine_sampler_entry.py +++ b/yeto/rl/engine/engine_sampler_entry.py @@ -241,12 +241,15 @@ def island_config(spec: Mapping[str, Any]) -> EngineSamplerConfig: def run(spec: Mapping[str, Any], *, environ: Mapping[str, str] | None = None, resource_sampler: Callable[[Callable[..., None]], Any] | None = None, - incarnation: str | None = None) -> int: + incarnation: str | None = None, stop: Callable[[], bool] | None = None, + on_event: Callable[..., None] | None = None) -> int: """Build every part of the island from ``spec``, run it from the resume point, stop the engine; returns the rounds this incarnation ran. ``resource_sampler(emit)`` returns a started telemetry thread (default: ``telemetry.ResourceSampler`` of the island's device family; None from the - factory = no telemetry).""" + factory = no telemetry). ``stop()``: the caller's own end condition (no new + round starts; the island leaves FINISHED, not PREEMPTED). ``on_event(event, + **fields)`` sees every tape record after it is written.""" import os # noqa: PLC0415 from yeto.rl.elastic_client import LeaveReason # noqa: PLC0415 @@ -266,7 +269,13 @@ def run(spec: Mapping[str, Any], *, environ: Mapping[str, str] | None = None, config = island_config(spec) incarnation = incarnation or new_incarnation() tape = incarnation_tape(_require(spec, "event_tape"), incarnation) - emit = tape_writer(tape, config.sampler_id) + write = tape_writer(tape, config.sampler_id) + + def emit(event: str, **fields: Any) -> None: + write(event, **fields) + if on_event is not None: + on_event(event, **fields) + releases, buffer = build_stores(spec, environ) prompts, reward = build_prompts(spec), build_reward(spec) checks = PostLoadChecks(**dict(spec.get("checks") or {})) @@ -298,7 +307,8 @@ def resource_sampler(sink: Callable[..., None]) -> Any: try: island = build_engine_sampler(config, engine=engine, releases=releases, buffer=buffer, emit=emit, prompts=prompts, reward=reward, checks=checks, - should_stop=reclaimed, incarnation=incarnation, resume=resume, + should_stop=lambda: reclaimed() or (stop is not None and stop()), + reclaimed=reclaimed, incarnation=incarnation, resume=resume, data_cursor=getattr(prompts, "cursor", None)) latest = releases.latest() membership.join(tape=str(tape), latest_release=None if latest is None else @@ -333,6 +343,26 @@ def load_spec(path: str | Path) -> dict[str, Any]: return doc +RECLAIM_HOLD_S = 300.0 # after a reclaim LEAVE: wait this long for the cloud to stop the container + + +def hold_after_reclaim(environ: Mapping[str, str] | None = None, *, hold_s: float = RECLAIM_HOLD_S, + sleep: Callable[[float], None] = time.sleep) -> bool: + """After a reclaim the process stays up until the cloud stops the container: + a clean early exit could end the cloud's input as finished instead of letting + the cloud run the island again (as the Miles sampler island does). Returns + whether a reclaim was requested (False: returns at once).""" + import os # noqa: PLC0415 + + from .sampler_resume import reclaim_requested # noqa: PLC0415 + + if not reclaim_requested(os.environ if environ is None else environ): + return False + print(f"[engine-sampler] reclaimed: waiting up to {hold_s:g} s for the container to stop", flush=True) + sleep(hold_s) + return True + + def main(argv: Sequence[str] | None = None) -> int: parser = argparse.ArgumentParser(prog="python -m yeto.rl.engine.engine_sampler", description="run an engine sampler island from a spec file") @@ -345,4 +375,5 @@ def main(argv: Sequence[str] | None = None) -> int: print(f"[engine-sampler] spec rejected: {error}", file=sys.stderr) return 2 print(f"[engine-sampler] done: {rounds} round(s)", flush=True) + hold_after_reclaim() return 0 From c34612b1064e4db7bda354e0a9055866ccf7f4fc Mon Sep 17 00:00:00 2001 From: michael <78113554+chowhsu@users.noreply.github.com> Date: Tue, 20 Oct 2026 11:31:38 +0000 Subject: [PATCH 83/89] tests: the Miles reader stub reports its stale groups Co-Authored-By: Claude Opus 5.5 --- tests/test_rl_n2_trainer_wiring.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/tests/test_rl_n2_trainer_wiring.py b/tests/test_rl_n2_trainer_wiring.py index b7a7af90d..4e23caefe 100644 --- a/tests/test_rl_n2_trainer_wiring.py +++ b/tests/test_rl_n2_trainer_wiring.py @@ -68,7 +68,8 @@ def test_miles_reader_rebuilds_prompts_and_reports_per_source_metrics(monkeypatc reader.args = types.SimpleNamespace(rollout_sample_filter_path=None, rollout_all_samples_process_path=None) reader.samples_per_group = 2 reader._tokenizer = lambda text, add_special_tokens: {"input_ids": [7, 8]} - report = types.SimpleNamespace(metrics=lambda: {"sample_buffer/groups": 1.0}) + report = types.SimpleNamespace(metrics=lambda: {"sample_buffer/groups": 1.0}, stale_groups=0) + monkeypatch.setattr(plugin, "record_round_metadata", lambda args, rid, **counts: None) group = [_record(index=0), _record(finish=FINISH_LENGTH, index=1)] reader.reader = types.SimpleNamespace(take=lambda rollout_id, groups, current_version: ([(5, group)], report)) reader.groups = 1 From 434435bcd4de6eaf119c83da690d2a7621833572 Mon Sep 17 00:00:00 2001 From: michael <78113554+chowhsu@users.noreply.github.com> Date: Tue, 20 Oct 2026 11:32:31 +0000 Subject: [PATCH 84/89] eval E2-verl: save the sampler tape when the island leaves (a reclaimed container may die before the function ends) Co-Authored-By: Claude Opus 5.5 --- tools/eval/modal_verl.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/tools/eval/modal_verl.py b/tools/eval/modal_verl.py index 67afe1072..eadc46995 100644 --- a/tools/eval/modal_verl.py +++ b/tools/eval/modal_verl.py @@ -217,6 +217,8 @@ def on_event(event: str, **fields: Any) -> None: if event == "rl_sampler_round": # the local driver's reclaim trigger control[f"{run}/sampler/round"] = {"incarnation": incarnation, **{ k: fields.get(k) for k in ("rollout_id", "groups")}} + elif event == "rl_sampler_leave": # save the tape now: a reclaimed container may die soon + _finish(out, {**meta, "leave": fields.get("reason")}, began, {"sampler-tape.jsonl": tape}) done_keys = [f"{run}/done/{i}" for i in range(cfg["learners"])] deadline = began + cfg["timeout_s"] - 300 From 3e1c833bd4405923d24a65a1b9d21b11935b9e7f Mon Sep 17 00:00:00 2001 From: michael <78113554+chowhsu@users.noreply.github.com> Date: Tue, 20 Oct 2026 11:41:40 +0000 Subject: [PATCH 85/89] codepack: unpacked trees get a 2000-01-01 mtime so sky can zip them The pack's entries carry mtime 0. With a remote sky API server, sky zips the workdir it uploads and ZIP refuses timestamps before 1980, so every launch with a code pack failed before provisioning. The digest covers names and bytes only; a tree unpacked earlier is stamped when it is reused. Co-Authored-By: Claude Opus 5.5 --- tests/test_codepack.py | 19 +++++++++++++++++++ yeto/codepack.py | 18 +++++++++++++++++- 2 files changed, 36 insertions(+), 1 deletion(-) diff --git a/tests/test_codepack.py b/tests/test_codepack.py index 24a9b35cd..af8d6e87f 100644 --- a/tests/test_codepack.py +++ b/tests/test_codepack.py @@ -98,6 +98,25 @@ def racing_is_dir(self: Path) -> bool: # the first existence check misses the t assert [p.name for p in cache.iterdir()] == [first.name] # no staging dir left behind +def test_an_unpacked_tree_can_be_zipped_for_a_remote_api_server(repo: Path) -> None: + # T-C1b: the pack's entries carry mtime 0; sky zips the workdir it uploads to a + # remote API server and ZIP refuses times before 1980 + import os + import zipfile + + pack, data = codepack.build_pack(repo) + cache = repo.parent / "cache" + tree = codepack.unpack(data, pack.digest, cache) + files = [p for p in tree.rglob("*") if p.is_file()] + assert files and all(p.stat().st_mtime == codepack.UNPACKED_MTIME for p in files) + with zipfile.ZipFile(repo.parent / "upload.zip", "w") as zf: + for p in files: + zf.write(p, p.relative_to(tree)) + os.utime(files[0], (0, 0)) # a tree unpacked before the fix is stamped on reuse + assert codepack.unpack(data, pack.digest, cache) == tree + assert files[0].stat().st_mtime == codepack.UNPACKED_MTIME + + def test_unpack_rejects_wrong_bytes(repo: Path) -> None: pack, data = codepack.build_pack(repo) with pytest.raises(CodePackError, match="does not have digest"): diff --git a/yeto/codepack.py b/yeto/codepack.py index 0e09abb25..d18c74b87 100644 --- a/yeto/codepack.py +++ b/yeto/codepack.py @@ -241,9 +241,23 @@ def _safe_member(member: tarfile.TarInfo) -> None: raise CodePackError(f"code pack symlink {member.name!r} leaves the tree") +# The pack's entries carry mtime 0 (a deterministic tar); an unpacked tree gets this +# time instead. sky zips a workdir it uploads to a remote API server, and ZIP refuses +# timestamps before 1980 ("ZIP does not support timestamps before 1980", T-C1b). +UNPACKED_MTIME = 946684800 # 2000-01-01T00:00:00Z + + +def _stamp_times(tree: Path) -> None: + """Give every entry of ``tree`` :data:`UNPACKED_MTIME` (links themselves, not + their targets). The digest covers names and bytes only, so it is unchanged.""" + for path in [tree, *tree.rglob("*")]: + os.utime(path, (UNPACKED_MTIME, UNPACKED_MTIME), follow_symlinks=False) + + def unpack(data: bytes, digest: str, cache_root: str | os.PathLike[str]) -> Path: """Check ``data`` against ``digest`` and unpack it to ``/`` - (reused when already there and intact). Returns the tree.""" + (reused when already there and intact; a tree unpacked before + :data:`UNPACKED_MTIME` existed gets its times now). Returns the tree.""" if hashlib.sha256(data).hexdigest() != digest.lower(): raise CodePackError(f"downloaded code pack does not have digest {digest}") @@ -251,6 +265,7 @@ def unpack(data: bytes, digest: str, cache_root: str | os.PathLike[str]) -> Path if dest.is_dir(): try: verify_tree(dest, digest) + _stamp_times(dest) return dest except CodePackError: shutil.rmtree(dest) @@ -262,6 +277,7 @@ def unpack(data: bytes, digest: str, cache_root: str | os.PathLike[str]) -> Path for member in members: _safe_member(member) tar.extractall(staging, members=members, filter="tar") + _stamp_times(staging) verify_tree(staging, digest) try: staging.rename(dest) From 3bf65245ca4c4793065a971bdcf49d89ec956c01 Mon Sep 17 00:00:00 2001 From: michael <78113554+chowhsu@users.noreply.github.com> Date: Tue, 20 Oct 2026 11:41:53 +0000 Subject: [PATCH 86/89] Engine server adapters start the server with SIGINT ignored: a cloud reclaim signals every process, and the island must finish its round and leave PREEMPTED Found by eval-e2-verl-e2v-a: the reclaimed SGLang server died under round 2 and the reclaimed incarnation ended with an engine error instead of a clean leave. Co-Authored-By: Claude Opus 5.5 --- tests/test_rl_hetero_sampler_engines.py | 12 ++++++++++++ yeto/rl/adapters/sampler_engines/server.py | 12 +++++++++++- 2 files changed, 23 insertions(+), 1 deletion(-) diff --git a/tests/test_rl_hetero_sampler_engines.py b/tests/test_rl_hetero_sampler_engines.py index 1427f7d1c..e6c41ded6 100644 --- a/tests/test_rl_hetero_sampler_engines.py +++ b/tests/test_rl_hetero_sampler_engines.py @@ -652,3 +652,15 @@ def test_probe_without_pairs_passes_every_load(): result = checks.logprob_probe(engine=None) checks.accept() assert result == {"name": "logprob_probe", "passed": True, "detail": "no probe pairs"} + + +def test_server_process_ignores_the_reclaim_sigint(): + """T-E2V: a cloud reclaim sends SIGINT to every process of the container; the + engine server must survive it so the island can finish its round and leave.""" + import sys + + from yeto.rl.adapters.sampler_engines.server import popen_spawn + + probe = popen_spawn([sys.executable, "-c", "import signal, sys; " + "sys.exit(7 if signal.getsignal(signal.SIGINT) is signal.SIG_IGN else 3)"], {}) + assert probe.wait(timeout=60) == 7 diff --git a/yeto/rl/adapters/sampler_engines/server.py b/yeto/rl/adapters/sampler_engines/server.py index 5d868c46e..ad33b0ace 100644 --- a/yeto/rl/adapters/sampler_engines/server.py +++ b/yeto/rl/adapters/sampler_engines/server.py @@ -58,8 +58,18 @@ def wait(self, timeout: float | None = None) -> int: ... Http = Callable[[str, str, Any, float], tuple[int, Any]] +def _ignore_sigint() -> None: + signal.signal(signal.SIGINT, signal.SIG_IGN) + + def popen_spawn(argv: Sequence[str], env: Mapping[str, str]) -> Process: - return subprocess.Popen(list(argv), env={**os.environ, **env}) # noqa: S603 + """Start the server with SIGINT ignored (inherited by its workers). A cloud + reclaim (Modal ``container stop``) sends SIGINT to every process of the + container: the island turns it into the reclaim marker and leaves after + the round in progress, so the server must not die under that round + (T-E2V eval-e2-verl-e2v-a). ``stop()`` still ends it with SIGTERM.""" + return subprocess.Popen(list(argv), env={**os.environ, **env}, # noqa: S603 + preexec_fn=_ignore_sigint) # noqa: PLW1509 - only resets one signal def httpx_call(method: str, url: str, body: Any, timeout_s: float) -> tuple[int, Any]: From 37c00f6a9853c70a67b0b55325cc4100fcd95ddf Mon Sep 17 00:00:00 2001 From: michael <78113554+chowhsu@users.noreply.github.com> Date: Tue, 20 Oct 2026 11:54:54 +0000 Subject: [PATCH 87/89] nebius_reclaim: host unit install works with a remote sky API server The cluster yaml lives on the API server, so reading the SSH credentials from it failed ('Cluster yaml ... not found') and the reclaim unit was never installed. Without a local yaml, ask the server for the cluster's credentials (sky status with credentials, as the sky CLI does) and write the key locally. Checked against a live Nebius cluster. Co-Authored-By: Claude Opus 5.5 --- tests/test_nebius_reclaim.py | 27 +++++++++++++++++++++++++++ yeto/cloud/nebius_reclaim.py | 24 ++++++++++++++++++++++-- 2 files changed, 49 insertions(+), 2 deletions(-) diff --git a/tests/test_nebius_reclaim.py b/tests/test_nebius_reclaim.py index 43f1a358e..9cf6fb407 100644 --- a/tests/test_nebius_reclaim.py +++ b/tests/test_nebius_reclaim.py @@ -8,6 +8,8 @@ import threading from types import SimpleNamespace +import pytest + from yeto import launcher from yeto.cloud import droppable as dr from yeto.cloud import nebius_reclaim as nr @@ -156,3 +158,28 @@ def test_launcher_installs_hook_only_for_nebius_droppable(monkeypatch): assert launcher.maybe_install_nebius_reclaim_hook(SimpleNamespace(envs=envs), (7, "H"), "c") is None assert launcher.maybe_install_nebius_reclaim_hook(neb, (7, None), "c2") is False assert len(seen) == 1 + + +def test_host_runner_takes_credentials_from_a_remote_api_server(monkeypatch, tmp_path): + # T-C1b: with the sky API server on another machine the cluster yaml is not here + pytest.importorskip("sky") + from sky.utils import command_runner + + from yeto.cloud import nebius_reclaim as nr + + asked = [] + key = tmp_path / "c.key" + key.write_text("key\n") + + def server_credentials(cluster): + asked.append(cluster) + return {"ssh_user": "ubuntu", "ssh_private_key": str(key), + "ssh_control_name": "c", "ssh_proxy_command": None, "docker_user": "root"} + + monkeypatch.setattr(nr, "_server_credentials", server_credentials) + handle = SimpleNamespace(cluster_yaml=str(tmp_path / "missing.yml"), cluster_name="c-1", + ssh_user="ubuntu", head_ip="10.0.0.5") + runner = nr._host_runner(handle) + assert asked == ["c-1"] and isinstance(runner, command_runner.SSHCommandRunner) + assert runner.ip == "10.0.0.5" and runner.ssh_user == "ubuntu" + assert getattr(runner, "docker_user", None) is None # the HOST, not the container diff --git a/yeto/cloud/nebius_reclaim.py b/yeto/cloud/nebius_reclaim.py index 56e5f8cd5..6903dd73a 100644 --- a/yeto/cloud/nebius_reclaim.py +++ b/yeto/cloud/nebius_reclaim.py @@ -129,12 +129,32 @@ def wants_host_hook(task_envs) -> bool: return bool(doc) and doc.get("role") == droppable.DROPPABLE and doc.get("cloud") == "nebius" +def _server_credentials(cluster: str) -> dict[str, Any]: + """The cluster's SSH credentials from a remote sky API server (the cluster yaml + stays on the server), as ``sky status`` gets them; the private key is written + to the client's per-cluster key file, as the sky CLI does.""" + from sky.client import sdk + from sky.utils import cluster_utils + + records = sdk.stream_and_get(sdk.status([cluster], _include_credentials=True)) + creds = dict(next((r.get("credentials") for r in records if r.get("name") == cluster), None) or {}) + if not creds: + raise RuntimeError(f"the sky API server returned no SSH credentials for {cluster}") + cluster_utils.SSHConfigHelper.generate_local_key_file(cluster, creds) + return creds + + def _host_runner(handle: Any) -> Any: - """A sky SSH runner on the VM HOST (docker_user=None), not in the container.""" + """A sky SSH runner on the VM HOST (docker_user=None), not in the container. + The credentials come from the local cluster yaml (sky API server on this + machine), else from the remote API server (T-C1b).""" from sky.backends import backend_utils from sky.utils import command_runner - creds = backend_utils.ssh_credential_from_yaml(handle.cluster_yaml, None, handle.ssh_user) + if os.path.exists(os.path.expanduser(handle.cluster_yaml)): + creds = backend_utils.ssh_credential_from_yaml(handle.cluster_yaml, None, handle.ssh_user) + else: + creds = _server_credentials(handle.cluster_name) creds.pop("docker_user", None) ip = handle.head_ip return command_runner.SSHCommandRunner((ip, 22), **creds) From 89985b8d891e12bcd21023b1a3575a7ea6601b71 Mon Sep 17 00:00:00 2001 From: michael <78113554+chowhsu@users.noreply.github.com> Date: Tue, 20 Oct 2026 11:58:38 +0000 Subject: [PATCH 88/89] codepack: the pack lists its executables; verify restores lost executable bits sky uploads the workdir to a remote API server as a zip that the server unpacks without file modes, so every executable file reached the island as 0644 and the island refused the tree (code pack digest mismatch). The pack info now lists the executable files; verify_tree sets their bit again before it rebuilds the digest. Packs without the list are left alone. Co-Authored-By: Claude Opus 5.5 --- tests/test_codepack.py | 17 +++++++++++++++++ yeto/codepack.py | 30 ++++++++++++++++++++++++++++-- 2 files changed, 45 insertions(+), 2 deletions(-) diff --git a/tests/test_codepack.py b/tests/test_codepack.py index af8d6e87f..351cbfea8 100644 --- a/tests/test_codepack.py +++ b/tests/test_codepack.py @@ -117,6 +117,23 @@ def test_an_unpacked_tree_can_be_zipped_for_a_remote_api_server(repo: Path) -> N assert files[0].stat().st_mtime == codepack.UNPACKED_MTIME +def test_a_tree_that_lost_its_executable_bits_still_verifies(repo: Path) -> None: + # T-C1b: sky's upload to a remote API server unpacks the workdir zip without file + # modes; the island restores the bits the pack lists before it checks the digest + import os + + script = repo / "run.sh" + script.write_text("#!/bin/sh\n") + script.chmod(0o755) + subprocess.run(["git", "-C", str(repo), "add", "run.sh"], check=True) + pack, data = codepack.build_pack(repo) + tree = codepack.unpack(data, pack.digest, repo.parent / "cache") + assert "run.sh" in codepack.read_pack_info(tree)["executables"] + os.chmod(tree / "run.sh", 0o644) + assert codepack.verify_tree(tree, pack.digest) == pack.digest + assert os.stat(tree / "run.sh").st_mode & 0o111 + + def test_unpack_rejects_wrong_bytes(repo: Path) -> None: pack, data = codepack.build_pack(repo) with pytest.raises(CodePackError, match="does not have digest"): diff --git a/yeto/codepack.py b/yeto/codepack.py index d18c74b87..23ea2a4c2 100644 --- a/yeto/codepack.py +++ b/yeto/codepack.py @@ -190,7 +190,8 @@ def build_pack(repo: str | os.PathLike[str]) -> tuple[CodePack, bytes]: files = pack_file_list(root) patch_sha = hashlib.sha256(patch).hexdigest() info = {"format": PACK_FORMAT, "mode": CODE_PACK_MODE_GIT_PATCH, "commit": commit, - "dirty": bool(patch), "patch_sha": patch_sha, "files": files} + "dirty": bool(patch), "patch_sha": patch_sha, "files": files, + "executables": _executables(root, files)} data = _tar_bytes(root, files, info) pack = CodePack(commit=commit, dirty=bool(patch), patch_sha=patch_sha, digest=hashlib.sha256(data).hexdigest()) @@ -221,9 +222,34 @@ def tree_digest(tree: str | os.PathLike[str]) -> str: return hashlib.sha256(_tar_bytes(root, files, info)).hexdigest() +def _executables(root: Path, files: Iterable[str]) -> list[str]: + """The listed regular files with an executable bit (the pack records them as 0o755).""" + return [n for n in files + if not (root / n).is_symlink() and os.stat(root / n).st_mode & 0o111] + + +def restore_executables(tree: str | os.PathLike[str]) -> list[str]: + """Set the executable bit again on the files the pack info lists as executable + and returns the ones it changed. A transport may drop the bit: sky's workdir + upload to a remote API server goes through a zip that the server unpacks + without file modes (T-C1b). A pack without the list (older packs) is left alone.""" + root = Path(tree) + fixed = [] + for name in read_pack_info(root).get("executables") or []: + path = root / name + if path.is_file() and not path.is_symlink(): + mode = os.stat(path).st_mode + if not mode & 0o111: + os.chmod(path, mode | 0o755) + fixed.append(name) + return fixed + + def verify_tree(tree: str | os.PathLike[str], expected: str) -> str: - """Refuse a tree whose rebuilt tar does not have the digest ``expected``.""" + """Refuse a tree whose rebuilt tar does not have the digest ``expected`` + (after :func:`restore_executables`: modes are part of the digest).""" + restore_executables(tree) actual = tree_digest(tree) if actual != expected.lower(): raise CodePackError(f"code pack digest mismatch: run manifest has {expected.lower()}, " From 58e10d4e1b1e5343c9cd6ada3986800cd25ca8b0 Mon Sep 17 00:00:00 2001 From: michael <78113554+chowhsu@users.noreply.github.com> Date: Tue, 20 Oct 2026 12:56:20 +0000 Subject: [PATCH 89/89] launcher: the spot sampler env tolerates an island task built without a spec Co-Authored-By: Claude Opus 5.5 --- yeto/launcher.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/yeto/launcher.py b/yeto/launcher.py index 89d37836f..0abd2f4b5 100644 --- a/yeto/launcher.py +++ b/yeto/launcher.py @@ -189,11 +189,11 @@ def _island_task(task_factory, args, spec, m, num_learners, syncer_addr, overrid io.apply_to_task(task, {**io.island_env(args, m, overrides), **(droppable.role_env(role, spec.cloud, spec.region) if role else {}), **sample_roles.island_env(iargs), - **_spot_sampler_cloud_env(iargs, spec.cloud)}) + **_spot_sampler_cloud_env(iargs, getattr(spec, "cloud", None))}) return task -def _spot_sampler_cloud_env(iargs, cloud: str) -> dict[str, str]: +def _spot_sampler_cloud_env(iargs, cloud: str | None) -> dict[str, str]: """A spot sampler island on Nebius: the VM stop reaches the sampler through the host unit of ``yeto.cloud.nebius_reclaim`` (Modal signals the container itself).""" from .rl.engine import sample_roles