From ae98f7ec98655b16145c2c1ea2978685c2210e96 Mon Sep 17 00:00:00 2001 From: michael <78113554+chowhsu@users.noreply.github.com> Date: Mon, 19 Oct 2026 07:58:45 +0000 Subject: [PATCH 01/34] Add the engine-neutral SamplerEngine port and a fake engine (S25 N2.3) The port (yeto/rl/engine/sampler_engine.py) has start/load/served_version, generate with per-token log-probs, teacher-forced score, and the expert-count and weight-memory introspection the post-load checks need. Provenance and the run-wide SamplingParams (repetition penalties refused unless allowed) live next to it. Adapters register in yeto/rl/adapters/sampler_engines. Co-Authored-By: Claude Opus 5.5 --- tests/test_import_boundaries.py | 3 +- yeto/rl/adapters/sampler_engines/__init__.py | 34 +++ yeto/rl/adapters/sampler_engines/fake.py | 126 +++++++++ yeto/rl/engine/sampler_engine.py | 267 +++++++++++++++++++ 4 files changed, 429 insertions(+), 1 deletion(-) create mode 100644 yeto/rl/adapters/sampler_engines/__init__.py create mode 100644 yeto/rl/adapters/sampler_engines/fake.py create mode 100644 yeto/rl/engine/sampler_engine.py diff --git a/tests/test_import_boundaries.py b/tests/test_import_boundaries.py index 80a8025f8..7994a6840 100644 --- a/tests/test_import_boundaries.py +++ b/tests/test_import_boundaries.py @@ -62,7 +62,8 @@ ) # Backend adapters: name -> directory. Each must not import the others. ADAPTER_DIRS = {"miles": ("yeto/rl/adapters/miles",), - "verl": ("yeto/rl/adapters/verl",)} + "verl": ("yeto/rl/adapters/verl",), + "sampler_engines": ("yeto/rl/adapters/sampler_engines",)} # Cloud / launch layer (audit A8, A9, C1-C3). CLOUD_FILES = ("yeto/launcher.py", "yeto/modal_runner.py", "yeto/cli.py", "yeto/rl/stage_w_entry.py") diff --git a/yeto/rl/adapters/sampler_engines/__init__.py b/yeto/rl/adapters/sampler_engines/__init__.py new file mode 100644 index 000000000..4d453a843 --- /dev/null +++ b/yeto/rl/adapters/sampler_engines/__init__.py @@ -0,0 +1,34 @@ +"""Sampler engine adapters: implementations of the neutral +:class:`~yeto.rl.engine.sampler_engine.SamplerEngine` port (rl-hetero-sampler-engines D1). + +* ``sglang-server`` (:mod:`.sglang_server`): an SGLang HTTP server, PPU stage 1; +* ``vllm-server`` (:mod:`.vllm_server`): a vLLM OpenAI-compatible server + (vllm-ascend on 910B), stage 2; +* ``fake`` (:mod:`.fake`): deterministic in-memory engine for CPU tests. + +Vendor facts (patches, environment variables, launch flags, weight-name maps) +live only here. This package does not import another backend adapter. +""" + +from __future__ import annotations + +import importlib +from typing import Any + +# adapter name -> "module:class"; imported on use so that a host without httpx +# or the engine package can still import this package. +ENGINE_ADAPTERS = { + "sglang-server": "yeto.rl.adapters.sampler_engines.sglang_server:SGLangServerEngine", + "vllm-server": "yeto.rl.adapters.sampler_engines.vllm_server:VllmServerEngine", + "fake": "yeto.rl.adapters.sampler_engines.fake:FakeEngine", +} + + +def engine_class(name: str) -> Any: + """The adapter class registered under ``name``; unknown -> ValueError.""" + try: + target = ENGINE_ADAPTERS[name] + except KeyError: + raise ValueError(f"unknown sampler engine {name!r} (known: {sorted(ENGINE_ADAPTERS)})") from None + module, _, attr = target.partition(":") + return getattr(importlib.import_module(module), attr) diff --git a/yeto/rl/adapters/sampler_engines/fake.py b/yeto/rl/adapters/sampler_engines/fake.py new file mode 100644 index 000000000..927d0fb6e --- /dev/null +++ b/yeto/rl/adapters/sampler_engines/fake.py @@ -0,0 +1,126 @@ +"""A fake :class:`~yeto.rl.engine.sampler_engine.SamplerEngine` for CPU tests. + +Deterministic: the token ids and log-probs of a completion depend only on the +prompt, the sample index and the served version. Faults are injected by +attribute (``expert_matrices``, ``answers``, ``probe_shift``, ``fail_load``). +""" + +from __future__ import annotations + +import hashlib +from collections.abc import Callable, Mapping, Sequence +from pathlib import Path + +from yeto.rl.engine.sampler_engine import ( + FINISH_LENGTH, + FINISH_STOP, + RELOAD_IN_PLACE, + RELOAD_RESTART, + Completion, + EngineIdentity, + LoadReport, + SamplerEngineError, + SamplingParams, +) + +FAKE_IDENTITY = EngineIdentity(engine="fake", engine_version="0", device_family="cpu", + card_model="none", precision="bf16", quant_scheme="bf16", + topology="tp1") + + +def _unit(*parts: object) -> float: + digest = hashlib.sha256("|".join(map(str, parts)).encode()).digest() + return int.from_bytes(digest[:8], "big") / 2**64 + + +class FakeEngine: + """In-memory engine. ``answers``: prompt -> greedy answer text (fixed-prompt + check); other prompts get ``"fake answer"``. ``tokens_for``: response length + per (prompt, sample index); a completion that reaches ``max_tokens`` ends + with ``length``.""" + + def __init__(self, *, identity: EngineIdentity = FAKE_IDENTITY, max_model_len: int = 4096, + in_place: bool = True, expert_matrices: int | None = 36_288, + weight_gib: float | None = None, answers: Mapping[str, str] | None = None, + tokens_for: Callable[[str, int], int] = lambda prompt, i: 4 + i, + token_ids_for: Callable[[str, int, int], list[int]] | None = None) -> None: + self._identity = identity + self._max_model_len = int(max_model_len) + self.in_place = bool(in_place) + self.expert_matrices = expert_matrices + self.weight_gib = weight_gib + self.answers = dict(answers or {}) + self.tokens_for = tokens_for + self.token_ids_for = token_ids_for + self.probe_shift = 0.0 # added to every score log-prob (a "silently wrong" load) + self.fail_load: str | None = None # raise this message on the next load + self._version: int | None = None + self.release_dir: Path | None = None + self.loads: list[tuple[str | None, int]] = [] + self.requests: list[tuple[tuple[str, ...], SamplingParams]] = [] + self.stopped = False + + @property + def identity(self) -> EngineIdentity: + return self._identity + + @property + def max_model_len(self) -> int: + return self._max_model_len + + def start(self, release_dir: Path | None, version: int) -> LoadReport: + self.stopped = False + return self._serve(release_dir, version, RELOAD_RESTART) + + def load(self, release_dir: Path, version: int) -> LoadReport: + return self._serve(release_dir, version, RELOAD_IN_PLACE if self.in_place else RELOAD_RESTART) + + def _serve(self, release_dir: Path | None, version: int, mode: str) -> LoadReport: + if self.fail_load is not None: + message, self.fail_load = self.fail_load, None + raise SamplerEngineError(message) + self.release_dir = None if release_dir is None else Path(release_dir) + self._version = int(version) + self.loads.append((None if release_dir is None else str(release_dir), int(version))) + return LoadReport(reload_mode=mode, pause_s=0.0) + + def served_version(self) -> int | None: + return self._version + + def generate(self, prompts: Sequence[str], sampling: SamplingParams) -> list[list[Completion]]: + if self._version is None or self.stopped: + raise SamplerEngineError("fake engine is not serving") + self.requests.append((tuple(prompts), sampling)) + out = [] + for prompt in prompts: + group = [] + for i in range(int(sampling.n)): + if prompt in self.answers: + text = self.answers[prompt] + ids = [ord(c) for c in text][: sampling.max_tokens] + else: + length = int(self.tokens_for(prompt, i)) + ids = (self.token_ids_for(prompt, i, length) if self.token_ids_for + else [int(_unit(prompt, i, k) * 1000) for k in range(length)]) + ids = ids[: sampling.max_tokens] + text = "fake answer" + finish = FINISH_LENGTH if len(ids) >= sampling.max_tokens else FINISH_STOP + logprobs = tuple(-_unit(prompt, i, k, self._version) for k in range(len(ids))) + group.append(Completion(token_ids=tuple(ids), logprobs=logprobs, text=text, + finish_reason=finish)) + out.append(group) + return out + + def score(self, prompt: str, continuation: str) -> list[float]: + if self._version is None: + raise SamplerEngineError("fake engine is not serving") + return [-_unit(prompt, continuation, k) - self.probe_shift for k in range(len(continuation))] + + def loaded_expert_matrices(self) -> int | None: + return self.expert_matrices + + def weight_gib_per_card(self) -> float | None: + return self.weight_gib + + def stop(self) -> None: + self.stopped = True diff --git a/yeto/rl/engine/sampler_engine.py b/yeto/rl/engine/sampler_engine.py new file mode 100644 index 000000000..431a2fc25 --- /dev/null +++ b/yeto/rl/engine/sampler_engine.py @@ -0,0 +1,267 @@ +"""The ``SamplerEngine`` port: one interface for every inference engine a +sampler island drives (rl-hetero-sampler-engines D1, S25 N2.3). + +The sampler island's loop, buffer writer, reward call and release follower are +engine-neutral. Everything that depends on the engine (process start, HTTP API, +weight reload, vendor patches and environment) lives in an adapter that +implements :class:`SamplerEngine`. Adapters: ``sglang-server`` (PPU, stage 1) +and ``vllm-server`` (vllm-ascend on 910B, stage 2), in +:mod:`yeto.rl.adapters.sampler_engines`. Tests use +:class:`yeto.rl.adapters.sampler_engines.fake.FakeEngine`. + +The port has four kinds of operation: + +* ``load(release_dir, version)``: serve the weights of a release. The adapter + reloads in place or restarts the engine; :class:`LoadReport` says which and + how long the engine did not serve (``reload_mode``, ``pause_s``); +* ``generate(prompts, sampling)``: ``sampling.n`` completions per prompt, each + with the per-token log-probabilities the engine sampled from; +* ``score(prompt, continuation)``: teacher-forced log-probabilities of a fixed + continuation (the post-load log-prob probe, :mod:`.sampler_checks`); +* introspection for the post-load checks: ``served_version()``, + ``loaded_expert_matrices()``, ``weight_gib_per_card()``. + +Who uses it: :class:`yeto.rl.engine.sampler_checks.PortReleaseLoader` (the +``PolicyFollower`` load hook) and :mod:`yeto.rl.engine.engine_sampler` +(the island's round). A records proxy (S25 N10.5) can wrap a ``SamplerEngine`` +and forward every call; the port has no hidden state the proxy must copy. + +Sampling parameters are the same for every source of a run (A-INF 5, guard +D11): :class:`SamplingParams` is built once from the run config, and +:meth:`SamplingParams.check` refuses repetition penalties, which change the +behavior distribution away from the model's own distribution. +""" + +from __future__ import annotations + +import hashlib +import json +from collections.abc import Mapping, Sequence +from dataclasses import asdict, dataclass, field, fields +from pathlib import Path +from typing import Any, Protocol, runtime_checkable + +# finish_reason of a completion. ``loop_abort``: the island cut a repeating +# completion (:mod:`.loop_guard`); the kept tokens are unchanged. +FINISH_STOP = "stop" +FINISH_LENGTH = "length" +FINISH_LOOP_ABORT = "loop_abort" +FINISH_REASONS = (FINISH_STOP, FINISH_LENGTH, FINISH_LOOP_ABORT) + +# reload_mode of a LoadReport. +RELOAD_IN_PLACE = "in_place" +RELOAD_RESTART = "restart" +RELOAD_MODES = (RELOAD_IN_PLACE, RELOAD_RESTART) + +# Neutral values of the repetition penalties (no change to the distribution). +NEUTRAL_PENALTIES = {"repetition_penalty": 1.0, "frequency_penalty": 0.0, "presence_penalty": 0.0} + + +class SamplerEngineError(RuntimeError): + """An engine call failed (start, load, generate, score).""" + + +class SamplingConfigError(ValueError): + """The sampling parameters are not allowed for a rollout (guard D11).""" + + +@dataclass(frozen=True) +class SamplingParams: + """Sampling parameters of one run, identical for every sampler source. + + ``max_tokens`` comes from the task config, never from an adapter default: + a source with a lower limit would truncate more and bias the per-source + reward (A-INF 5.1). ``allow_penalties`` must be set explicitly (and is + written to the tape) to use a non-neutral penalty.""" + + max_tokens: int + temperature: float = 1.0 + top_p: float = 1.0 + top_k: int = -1 # -1: no top-k limit + n: int = 1 + seed: int | None = None + stop: tuple[str, ...] = () + repetition_penalty: float = 1.0 + frequency_penalty: float = 0.0 + presence_penalty: float = 0.0 + allow_penalties: bool = False + + def check(self) -> "SamplingParams": + if int(self.max_tokens) < 1: + raise SamplingConfigError(f"max_tokens must be >= 1, got {self.max_tokens}") + if int(self.n) < 1: + raise SamplingConfigError(f"n must be >= 1, got {self.n}") + if not self.allow_penalties: + changed = {name: getattr(self, name) for name, neutral in NEUTRAL_PENALTIES.items() + if float(getattr(self, name)) != neutral} + if changed: + raise SamplingConfigError( + f"rollout sampling sets {changed}; repetition penalties change the behavior " + "distribution and engines differ in whether the returned log-probs include " + "them. Use the loop guard (finish_reason loop_abort) instead, or set " + "allow_penalties explicitly (recorded in the tape)") + return self + + def to_dict(self) -> dict[str, Any]: + out = asdict(self) + out["stop"] = list(self.stop) + return out + + def digest(self) -> str: + """Short digest of the parameters; equal on every source of a run.""" + text = json.dumps(self.to_dict(), sort_keys=True, separators=(",", ":")) + return hashlib.sha256(text.encode("utf-8")).hexdigest()[:16] + + def replace(self, **changes: Any) -> "SamplingParams": + return SamplingParams(**{**{f.name: getattr(self, f.name) for f in fields(self)}, **changes}) + + +@dataclass(frozen=True) +class Completion: + """One completion: response tokens, the log-probability of each sampled + token (same length), the response text and why it ended.""" + + token_ids: tuple[int, ...] + logprobs: tuple[float, ...] + text: str + finish_reason: str + + def __post_init__(self) -> None: + if len(self.token_ids) != len(self.logprobs): + raise SamplerEngineError(f"completion has {len(self.token_ids)} tokens but " + f"{len(self.logprobs)} log-probs") + if self.finish_reason not in FINISH_REASONS: + raise SamplerEngineError(f"unknown finish_reason {self.finish_reason!r} " + f"(known: {FINISH_REASONS})") + + +@dataclass(frozen=True) +class LoadReport: + """What one ``load`` did. ``pause_s``: seconds the engine served no request.""" + + reload_mode: str + pause_s: float + detail: Mapping[str, Any] = field(default_factory=dict) + + def __post_init__(self) -> None: + if self.reload_mode not in RELOAD_MODES: + raise SamplerEngineError(f"unknown reload_mode {self.reload_mode!r}") + + +@dataclass(frozen=True) +class EngineIdentity: + """The source facts of an engine (the provenance of its samples). + + ``precision``: the weight precision class (``bf16``, ``w4a8``, ``mxfp4``...); + ``quant_scheme``: the release scheme the engine serves (``bf16`` for trained + weights; else the ``derived_from.scheme`` of a derived release); + ``topology``: the parallel layout of one engine instance (``tp8``, ``dp8``, + ``tp4``...).""" + + engine: str + engine_version: str + device_family: str + card_model: str + precision: str + quant_scheme: str + topology: str + + +@runtime_checkable +class SamplerEngine(Protocol): + """The port (module doc). Every method may raise :class:`SamplerEngineError`.""" + + @property + def identity(self) -> EngineIdentity: ... + + @property + def max_model_len(self) -> int: + """Prompt plus response token limit of the served model.""" + ... + + def start(self, release_dir: Path | None, version: int) -> LoadReport: + """Start the engine on a release (None: the base weights it was configured with).""" + ... + + def load(self, release_dir: Path, version: int) -> LoadReport: + """Serve the weights of ``release_dir`` as policy ``version``.""" + ... + + def served_version(self) -> int | None: + """The policy version the engine serves (None: not started).""" + ... + + def generate(self, prompts: Sequence[str], sampling: SamplingParams) -> list[list[Completion]]: + """``sampling.n`` completions for each prompt, in prompt order.""" + ... + + def score(self, prompt: str, continuation: str) -> list[float]: + """Teacher-forced log-probabilities of the continuation tokens.""" + ... + + def loaded_expert_matrices(self) -> int | None: + """Quantized routed-expert matrices the engine loaded (None: it cannot tell).""" + ... + + def weight_gib_per_card(self) -> float | None: + """Weight memory per card after the load (None: it cannot tell).""" + ... + + def stop(self) -> None: + """Stop the engine (idempotent).""" + ... + + +# ---------------------------------------------------------------------- provenance +UNKNOWN = "unknown" + + +@dataclass(frozen=True) +class Provenance: + """The source of a buffer entry (rl-hetero-sampler-engines D3, S25 N2.4). + + Data only: it never enters ``algorithm_sha256`` or the session contract + hash. The behavior policy version stays in the entry's ``behavior_version``. + An entry written before this field exists reads as :meth:`unknown`.""" + + device_family: str + card_model: str + engine: str + engine_version: str + precision: str + quant_scheme: str + topology: str = UNKNOWN + sampler_id: str = UNKNOWN + sampling_digest: str = UNKNOWN + + # The per-source metrics group by these fields (D4, A-INF 4.1). + KEY_FIELDS = ("device_family", "card_model", "engine", "engine_version", "precision", + "topology") + + @classmethod + def unknown(cls) -> "Provenance": + return cls(*([UNKNOWN] * 6)) + + @classmethod + def of(cls, identity: EngineIdentity, *, sampler_id: str, sampling: SamplingParams) -> "Provenance": + return cls(device_family=identity.device_family, card_model=identity.card_model, + engine=identity.engine, engine_version=identity.engine_version, + precision=identity.precision, quant_scheme=identity.quant_scheme, + topology=identity.topology, sampler_id=str(sampler_id), + sampling_digest=sampling.digest()) + + @classmethod + def from_record(cls, record: Mapping[str, Any] | None) -> "Provenance": + """The provenance of an index record field; missing or partial -> ``unknown``.""" + if not isinstance(record, Mapping): + return cls.unknown() + names = [f.name for f in fields(cls)] + return cls(**{name: str(record.get(name, UNKNOWN)) for name in names}) + + def to_record(self) -> dict[str, str]: + return {f.name: getattr(self, f.name) for f in fields(self)} + + def key(self) -> str: + """The grouping key as one string, ``family/card/engine@version/precision/topology``.""" + return (f"{self.device_family}/{self.card_model}/{self.engine}@{self.engine_version}/" + f"{self.precision}/{self.topology}") From a1e00f07f025c15cb81a5f6f18546aa9f6918fe0 Mon Sep 17 00:00:00 2001 From: michael <78113554+chowhsu@users.noreply.github.com> Date: Mon, 19 Oct 2026 08:11:23 +0000 Subject: [PATCH 02/34] Add resume assertions so a restored optimizer must continue the cut (S25 N7.6) A-TRN report #8: restoring optimizer state on PPU can silently reset the Adam step. A cut now records a backend-neutral resume manifest (Adam step per group, LR scheduler progress and applied LRs, exp_avg/exp_avg_sq norm and sha256 of fixed sampled parameters, global step, data watermark). After loading, every rank must match it; after the first step every counter must advance by exactly the steps taken. Any mismatch raises ResumeAssertionError (a CutError); a cut without the manifest is refused. - yeto/rl/engine/resume_manifest.py: neutral types, checks, file form, torch.optim reader (FSDP2 DTensor local shards). - Miles: cut shards report the rank state; save_cut embeds the manifest, restore_cut checks it before the digest comparison, the first train_step after a restore checks the advance. - verl/FSDP2: worker binding (actor_module_fsdp/actor_optimizer/ actor_lr_scheduler). - CPU tests: train 3 + save + resume 2 equals train 5 (torch and FSDP2 over gloo, 2 ranks); reset step, partial reset, fresh moments, reset scheduler and a non-advancing first step are all refused. Co-Authored-By: Claude Opus 5.5 --- tests/test_rl_resume_manifest.py | 398 ++++++++++++++++++++++++++ yeto/rl/adapters/miles/cut_plugin.py | 51 ++++ yeto/rl/adapters/miles/trainer.py | 75 ++++- yeto/rl/adapters/verl/resume_check.py | 75 +++++ yeto/rl/engine/cut.py | 6 + yeto/rl/engine/resume_manifest.py | 374 ++++++++++++++++++++++++ 6 files changed, 978 insertions(+), 1 deletion(-) create mode 100644 tests/test_rl_resume_manifest.py create mode 100644 yeto/rl/adapters/verl/resume_check.py create mode 100644 yeto/rl/engine/resume_manifest.py diff --git a/tests/test_rl_resume_manifest.py b/tests/test_rl_resume_manifest.py new file mode 100644 index 000000000..21219e77d --- /dev/null +++ b/tests/test_rl_resume_manifest.py @@ -0,0 +1,398 @@ +"""S25 N7.6: resume assertions (A-TRN report #8, a silently reset Adam step). CPU only. + +* "train 3, save, resume, train 2" equals "train 5" bit for bit (plain torch and + FSDP2 over gloo with two CPU ranks); +* a restored optimizer whose Adam step, scheduler or moments differ is refused, + and a first step that does not advance the step by exactly one is refused; +* the Miles cut path and the verl worker binding raise on the same faults. +""" + +from __future__ import annotations + +import copy +import os +import socket +from types import SimpleNamespace + +import pytest +import torch + +from yeto.rl.engine.cut import CutError +from yeto.rl.engine.resume_manifest import ( + RESUME_MANIFEST_NAME, + ResumeAssertionError, + ResumeManifest, + read_resume_manifest, + require_first_step, + require_restored, + sample_names, + torch_optimizer_state, + write_resume_manifest, +) + +TOTAL, CUT = 5, 3 + + +# ---------------------------------------------------------------- tiny trainer + + +def _model() -> torch.nn.Module: + torch.manual_seed(0) + return torch.nn.Sequential(torch.nn.Linear(8, 16), torch.nn.Tanh(), torch.nn.Linear(16, 8), + torch.nn.Tanh(), torch.nn.Linear(8, 1)) + + +def _optim(model): + opt = torch.optim.AdamW(model.parameters(), lr=1e-2, weight_decay=0.01) + sched = torch.optim.lr_scheduler.LambdaLR(opt, lambda s: 1.0 / (1 + s)) + return opt, sched + + +def _batch(step: int, rank: int = 0) -> tuple[torch.Tensor, torch.Tensor]: + g = torch.Generator().manual_seed(1000 * rank + step) # data watermark -> batch + return torch.randn(4, 8, generator=g), torch.randn(4, 1, generator=g) + + +def _step(model, opt, sched, step: int, rank: int = 0) -> None: + x, y = _batch(step, rank) + opt.zero_grad(set_to_none=True) + torch.nn.functional.mse_loss(model(x), y).backward() + opt.step() + sched.step() + + +def _manifest(model, opt, sched, step: int, rank: str = "rank0") -> ResumeManifest: + return ResumeManifest(global_step=step, data_watermark=4 * step, scheduler_increment=1, + ranks={rank: torch_optimizer_state(opt, model.named_parameters(), sched)}) + + +def _save(directory, model, opt, sched, step: int) -> None: + torch.save({"model": model.state_dict(), "optim": opt.state_dict(), "sched": sched.state_dict(), + "step": step}, os.path.join(directory, "ckpt.pt")) + write_resume_manifest(directory, _manifest(model, opt, sched, step)) # after the checkpoint + + +def _load(directory, *, tamper=None): + model = _model() + opt, sched = _optim(model) + ckpt = torch.load(os.path.join(directory, "ckpt.pt"), weights_only=True) + model.load_state_dict(ckpt["model"]) + opt.load_state_dict(ckpt["optim"]) + sched.load_state_dict(ckpt["sched"]) + if tamper is not None: + tamper(model, opt, sched) + return model, opt, sched, int(ckpt["step"]) + + +def _equal_state(a_model, a_opt, b_model, b_opt) -> None: + for (name, p), (_, q) in zip(a_model.named_parameters(), b_model.named_parameters()): + assert torch.equal(p, q), name + for p, q in zip(a_model.parameters(), b_model.parameters()): + sa, sb = a_opt.state[p], b_opt.state[q] + assert set(sa) == set(sb) + for key in sa: + assert torch.equal(torch.as_tensor(sa[key]), torch.as_tensor(sb[key])), key + + +def _continuous(): + model = _model() + opt, sched = _optim(model) + for step in range(TOTAL): + _step(model, opt, sched, step) + return model, opt, sched + + +def _resumed(tmp_path, *, tamper=None): + model = _model() + opt, sched = _optim(model) + for step in range(CUT): + _step(model, opt, sched, step) + _save(tmp_path, model, opt, sched, CUT) + saved = read_resume_manifest(tmp_path) + model, opt, sched, start = _load(tmp_path, tamper=tamper) + require_restored(saved, _manifest(model, opt, sched, start)) + _step(model, opt, sched, start) + require_first_step(saved, _manifest(model, opt, sched, start + 1)) + for step in range(start + 1, TOTAL): + _step(model, opt, sched, step) + return model, opt, sched + + +# ---------------------------------------------------------------- equivalence + + +def test_train3_save_resume2_equals_train5(tmp_path): + a = _continuous() + b = _resumed(tmp_path) + _equal_state(a[0], a[1], b[0], b[1]) + assert a[2].last_epoch == b[2].last_epoch == TOTAL + assert a[2].get_last_lr() == b[2].get_last_lr() + assert all(int(s["step"]) == TOTAL for s in b[1].state.values()) + + +def test_manifest_round_trips_and_samples_are_fixed(tmp_path): + model, opt, sched = _continuous() + manifest = _manifest(model, opt, sched, TOTAL) + write_resume_manifest(tmp_path, manifest) + assert read_resume_manifest(tmp_path) == manifest + names = [n for n, _ in model.named_parameters()] + assert sample_names(names, 4) == sample_names(reversed(names), 4) + assert len(manifest.ranks["rank0"].moments) == 4 + assert manifest.ranks["rank0"].adam_steps["group0"].max == TOTAL + + +def test_missing_manifest_is_refused(tmp_path): + with pytest.raises(ResumeAssertionError, match=RESUME_MANIFEST_NAME): + read_resume_manifest(tmp_path) + + +# ---------------------------------------------------------------- simulated faults + + +def _reset_steps(model, opt, sched): # what A-TRN report #8 describes on PPU + for state in opt.state.values(): + state["step"] = torch.zeros_like(state["step"]) + + +def _reset_one_step(model, opt, sched): + first = next(iter(model.parameters())) + opt.state[first]["step"] = torch.zeros_like(opt.state[first]["step"]) + + +def _fresh_moments(model, opt, sched): + for state in opt.state.values(): + state["exp_avg"].zero_() + + +def _reset_scheduler(model, opt, sched): + sched.last_epoch = 0 + + +@pytest.mark.parametrize("tamper, message", [ + (_reset_steps, "Adam step was reset"), + (_reset_one_step, "Adam step differs across parameters"), + (_fresh_moments, "exp_avg norm"), + (_reset_scheduler, "scheduler was reset"), +]) +def test_restore_faults_are_refused(tmp_path, tamper, message): + with pytest.raises(ResumeAssertionError, match=message): + _resumed(tmp_path, tamper=tamper) + + +def test_resume_assertion_is_a_cut_error(): + assert issubclass(ResumeAssertionError, CutError) + + +def test_first_step_that_does_not_advance_is_refused(tmp_path): + model = _model() + opt, sched = _optim(model) + for step in range(CUT): + _step(model, opt, sched, step) + saved = _manifest(model, opt, sched, CUT) + _step(model, opt, sched, CUT) + _reset_steps(model, opt, sched) # e.g. a lazy re-init on the first step after the load + with pytest.raises(ResumeAssertionError, match="Adam step 0 != expected 4"): + require_first_step(saved, _manifest(model, opt, sched, CUT + 1)) + + +def test_first_step_lr_follows_the_schedule(tmp_path): + model = _model() + opt, sched = _optim(model) + for step in range(CUT): + _step(model, opt, sched, step) + saved = _manifest(model, opt, sched, CUT) + _step(model, opt, sched, CUT) + now = _manifest(model, opt, sched, CUT + 1) + require_first_step(saved, now, expected_lrs={"rank0": [1e-2 / (CUT + 2)]}) + with pytest.raises(ResumeAssertionError, match="scheduled"): + require_first_step(saved, now, expected_lrs={"rank0": [1e-2]}) + + +def test_data_watermark_and_global_step_are_checked(): + model, opt, sched = _continuous() + saved = _manifest(model, opt, sched, TOTAL) + moved = ResumeManifest(TOTAL, saved.data_watermark + 4, 1, saved.ranks) + with pytest.raises(ResumeAssertionError, match="data watermark"): + require_restored(saved, moved) + with pytest.raises(ResumeAssertionError, match="global step"): + require_restored(saved, ResumeManifest(0, saved.data_watermark, 1, saved.ranks)) + + +# ---------------------------------------------------------------- verl / FSDP2 binding + + +def test_verl_worker_binding_refuses_a_reset_step(tmp_path): + from yeto.rl.adapters.verl.resume_check import ( + require_worker_first_step, + require_worker_restored, + worker_manifest, + ) + + model = _model() + opt, sched = _optim(model) + for step in range(CUT): + _step(model, opt, sched, step) + worker = SimpleNamespace(actor_module_fsdp=model, actor_optimizer=opt, actor_lr_scheduler=sched) + saved = worker_manifest(worker, rank=0, global_step=CUT, data_watermark=12) + require_worker_restored(saved, worker, rank=0, global_step=CUT, data_watermark=12) + _step(model, opt, sched, CUT) + require_worker_first_step(saved, worker, rank=0, global_step=CUT + 1, data_watermark=16) + _reset_steps(model, opt, sched) + with pytest.raises(ResumeAssertionError, match="Adam step was reset"): + require_worker_restored(saved, worker, rank=0, global_step=CUT, data_watermark=12) + with pytest.raises(ResumeAssertionError, match="lacks"): + worker_manifest(SimpleNamespace(), rank=0, global_step=0, data_watermark=0) + + +def _free_port() -> int: + with socket.socket() as s: + s.bind(("127.0.0.1", 0)) + return s.getsockname()[1] + + +def _fsdp2_worker(rank: int, world: int, port: int, out_dir: str) -> None: + """Two CPU ranks, FSDP2 over gloo: continuous 5 steps vs 3 + save/resume + 2.""" + import torch.distributed as dist + from torch.distributed.checkpoint.state_dict import get_state_dict, set_state_dict + from torch.distributed.fsdp import fully_shard + + os.environ.update(MASTER_ADDR="127.0.0.1", MASTER_PORT=str(port)) + dist.init_process_group("gloo", rank=rank, world_size=world) + try: + def build(): + model = _model() + for layer in model: + if isinstance(layer, torch.nn.Linear): + fully_shard(layer) + fully_shard(model) + return (model, *_optim(model)) + + def run(model, opt, sched, steps): + for step in steps: + _step(model, opt, sched, step, rank) + + def label(): + return f"rank{rank}" + + a = build() + run(*a, range(TOTAL)) + + b = build() + run(*b, range(CUT)) + model_sd, optim_sd = get_state_dict(b[0], b[1]) + saved = _manifest(b[0], b[1], b[2], CUT, label()) + snapshot = (copy.deepcopy(model_sd), copy.deepcopy(optim_sd), b[2].state_dict()) + + c = build() + set_state_dict(c[0], c[1], model_state_dict=snapshot[0], optim_state_dict=snapshot[1]) + c[2].load_state_dict(snapshot[2]) + require_restored(saved, _manifest(c[0], c[1], c[2], CUT, label())) + run(*c, [CUT]) + require_first_step(saved, _manifest(c[0], c[1], c[2], CUT + 1, label())) + run(*c, range(CUT + 1, TOTAL)) + for (name, p), (_, q) in zip(a[0].named_parameters(), c[0].named_parameters()): + assert torch.equal(p.to_local(), q.to_local()), name + for p, q in zip(a[0].parameters(), c[0].parameters()): + for key in ("exp_avg", "exp_avg_sq", "step"): + x, y = a[1].state[p][key], c[1].state[q][key] + x = x.to_local() if hasattr(x, "to_local") else x + y = y.to_local() if hasattr(y, "to_local") else y + assert torch.equal(x, y), key + + _reset_steps(c[0], c[1], c[2]) # the PPU fault on this rank's shard + try: + require_restored(saved, _manifest(c[0], c[1], c[2], CUT, label())) + except ResumeAssertionError as error: + assert "Adam step was reset" in str(error) + else: + raise AssertionError("a reset Adam step was not caught") + open(os.path.join(out_dir, f"ok{rank}"), "w").close() + finally: + dist.destroy_process_group() + + +def test_fsdp2_gloo_two_ranks_resume_equals_continuous(tmp_path): + import torch.multiprocessing as mp + + pytest.importorskip("torch.distributed.fsdp") + if not torch.distributed.is_gloo_available(): + pytest.skip("gloo backend not built") + world = 2 + mp.spawn(_fsdp2_worker, args=(world, _free_port(), str(tmp_path)), nprocs=world, join=True) + assert sorted(os.listdir(tmp_path)) == [f"ok{r}" for r in range(world)] + + +# ---------------------------------------------------------------- Miles cut path + + +def _miles(): + from tests.test_rl_trainer_cut import RankGroup, _context, _expect, _trained_rank, _trainer + + return RankGroup, _context, _expect, _trained_rank, _trainer + + +def test_miles_cut_records_and_checks_the_resume_manifest(tmp_path): + from tests.rl_cut_fakes import GBS, make_rank, train_step + + RankGroup, _context, _expect, _trained_rank, _trainer = _miles() + _trainer(RankGroup([_trained_rank()])).save_cut(epoch=1, context=_context(tmp_path)) + fresh = make_rank(9) + trainer = _trainer(RankGroup([fresh])) + manifest = trainer.restore_cut("cut-a", epoch=1, root=str(tmp_path), expect=_expect()) + resume = ResumeManifest.from_dict(manifest.resume) + assert (resume.global_step, resume.data_watermark, resume.scheduler_increment) == (2, 4, GBS) + state = next(iter(resume.ranks.values())) + assert {s.max for s in state.adam_steps.values()} == {2} and state.moments + train_step(fresh, torch.randn(3, 6)) + trainer._check_first_step_after_resume(1) # step 2 -> 3, scheduler 2*GBS -> 3*GBS + assert trainer._resume_reference is None + + +def test_miles_restore_refuses_a_reset_adam_step(tmp_path): + from tests.rl_cut_fakes import make_rank + + RankGroup, _context, _expect, _trained_rank, _trainer = _miles() + _trainer(RankGroup([_trained_rank()])).save_cut(epoch=1, context=_context(tmp_path)) + fresh = make_rank(9) + real_load = fresh._yeto_cut_backend.load_optimizer + + def resetting(optimizer, named, merged): + real_load(optimizer, named, merged) + for state in optimizer.state.values(): + state["step"] = torch.zeros_like(state["step"]) + + fresh._yeto_cut_backend.load_optimizer = resetting + with pytest.raises(CutError, match="Adam step was reset"): + _trainer(RankGroup([fresh])).restore_cut("cut-a", epoch=1, root=str(tmp_path), expect=_expect()) + + +def test_miles_first_step_without_advance_is_refused(tmp_path): + from tests.rl_cut_fakes import make_rank + + RankGroup, _context, _expect, _trained_rank, _trainer = _miles() + _trainer(RankGroup([_trained_rank()])).save_cut(epoch=1, context=_context(tmp_path)) + fresh = make_rank(9) + trainer = _trainer(RankGroup([fresh])) + trainer.restore_cut("cut-a", epoch=1, root=str(tmp_path), expect=_expect()) + with pytest.raises(ResumeAssertionError, match="Adam step 2 != expected 3"): + trainer._check_first_step_after_resume(1) # no optimizer step was taken + + +def test_miles_cut_without_resume_section_is_refused(tmp_path): + """A cut written before N7.6 (no resume section, hash still valid) is not resumed.""" + import hashlib + import json + + from tests.rl_cut_fakes import make_rank + from yeto.rl.engine.cut import MANIFEST_NAME, canonical_json + + RankGroup, _context, _expect, _trained_rank, _trainer = _miles() + _trainer(RankGroup([_trained_rank()])).save_cut(epoch=1, context=_context(tmp_path)) + path = tmp_path / "cut-a" / MANIFEST_NAME + raw = json.loads(path.read_text()) + raw.pop("resume") + body = {k: v for k, v in raw.items() if k != "manifest_sha256"} + raw["manifest_sha256"] = hashlib.sha256(canonical_json(body).encode()).hexdigest() + path.write_text(json.dumps(raw)) + with pytest.raises(CutError, match="no resume manifest"): + _trainer(RankGroup([make_rank(9)])).restore_cut("cut-a", epoch=1, root=str(tmp_path), expect=_expect()) diff --git a/yeto/rl/adapters/miles/cut_plugin.py b/yeto/rl/adapters/miles/cut_plugin.py index 266bc2ead..2f17dac02 100644 --- a/yeto/rl/adapters/miles/cut_plugin.py +++ b/yeto/rl/adapters/miles/cut_plugin.py @@ -499,6 +499,55 @@ def train_state_digest(snap: Mapping[str, Any]) -> str: return state_digest({k: v for k, v in snap.items() if k != "miles_counters"}) +def _entry_moments(entry: Mapping[str, Any]) -> Mapping[str, Any]: + """Moment tensors of one named-export entry (fork-M5 ``tensors`` or a plain ``state``).""" + return entry.get("tensors") or entry.get("state") or {} + + +def _entry_step(entry: Mapping[str, Any]) -> Any: + """Adam step of one named-export entry: per-parameter state (torch Adam) or the + group hyper-parameters (TE/apex fused Adam); None before the first step.""" + for source in (entry.get("scalars"), entry.get("state"), entry.get("hyper")): + if isinstance(source, Mapping) and source.get("step") is not None: + return source["step"] + return None + + +def resume_rank_state(actor: Any, optimizer_named: Any) -> dict[str, Any]: + """This rank's :class:`~yeto.rl.engine.resume_manifest.OptimizerResumeState` (as a dict), + read from the named optimizer export and the Megatron LR scheduler (S25 N7.6).""" + from yeto.rl.engine.resume_manifest import build_rank_state, scheduler_progress_of + + entries = (optimizer_named or {}).get("entries") or {} + steps = [s for s in (_entry_step(e) for e in entries.values()) if s is not None] + lrs = [float(g["lr"]) for g in getattr(actor.optimizer, "param_groups", ()) if "lr" in g] + return build_rank_state( + steps_by_group={"named": steps}, + scheduler_progress=scheduler_progress_of(actor.opt_param_scheduler), + lrs=lrs, + moments_by_name={n: _entry_moments(e) for n, e in entries.items()}, + ).to_dict() + + +def resume_state(actor: Any) -> dict[str, Any]: + """Plugin: ``{"path": shard name, "resume": rank state}`` of the live trainer (first-step check).""" + import torch + + try: + backend = _backend(actor) + coord = _require(actor, backend) + with trainer_resident(actor): + named = _adapters(actor, backend) + with torch.no_grad(): + exported = backend.export_optimizer(actor.optimizer, named) + return {"path": shard_name(coord), "resume": resume_rank_state(actor, exported)} + except Exception as error: # noqa: BLE001 -- returned, not raised (see save_cut_shard) + return {REFUSED: f"{type(error).__name__}: {error}", REFUSAL_KIND: "failed_before_write"} + + +RESUME_STATE = f"{_MODULE}.resume_state" + + def _snapshot(actor: Any, backend: Any, named: list) -> dict[str, Any]: return { "adapter": {n: p.detach().to("cpu").clone() for n, p in named}, @@ -590,6 +639,7 @@ def _save(actor: Any, *, directory: str, cut_id: str) -> dict[str, Any]: "adapter_tensors": len(named), "adapter_names": sorted(n for n, _ in named), "optimizer_names": sorted((snap["optimizer_named"] or {}).get("entries", {})), + "resume": resume_rank_state(actor, snap["optimizer_named"]), } @@ -829,6 +879,7 @@ def _restore(actor: Any, *, directory: str, files: list[Mapping[str, Any]], cut_ "state_digest": state_digest(snap), "rng_digest": state_digest(rng_now), "components": component_digests({"state": snap, "rng": rng_now}), + "resume": resume_rank_state(actor, snap["optimizer_named"]), } saved_snap = {k: shard.get(k) for k in snap} if state_digest(saved_snap) != summary["state_digest"]: diff --git a/yeto/rl/adapters/miles/trainer.py b/yeto/rl/adapters/miles/trainer.py index f8a62240c..3d7e5ebe2 100644 --- a/yeto/rl/adapters/miles/trainer.py +++ b/yeto/rl/adapters/miles/trainer.py @@ -291,6 +291,8 @@ def __init__( # (``last_step_losses`` stays gated on the mechanisms that consume it). self.last_round_step_losses: list[dict[str, Any]] | None = None self.optimizer_steps_total = 0 + # S25 N7.6: the restored cut's resume manifest until the first step after it is checked + self._resume_reference: Any = None # A LoRA actor installs its step recorders through the state plugins # the publisher calls before the first step; a full-parameter actor # calls none of them, so the trainer installs the recorders itself. @@ -380,6 +382,8 @@ def train_step(self, batch: RolloutBatchHandle) -> LocalStepReceipt: self._release(self._args, batch.payload) steps = int(self._args.num_steps_per_rollout) if succeeded else 0 self.optimizer_steps_total += steps + if succeeded and self._resume_reference is not None: + self._check_first_step_after_resume(steps) return LocalStepReceipt( algorithm=self._algorithm, learner_id=self._learner_id, @@ -394,6 +398,31 @@ def train_step(self, batch: RolloutBatchHandle) -> LocalStepReceipt: parameter_layout_hash=self._layout_hash(), ) + def _rank_resume_states(self, results: list[Mapping[str, Any]], what: str) -> dict[str, Any]: + from yeto.rl.engine.cut import CutError + from yeto.rl.engine.resume_manifest import OptimizerResumeState + + refused = sorted({str(r["refused"]) for r in results if "refused" in r}) + if refused: + raise CutError(f"{what}: rank could not report its optimizer state: " + "; ".join(refused)) + return {str(r["path"]): OptimizerResumeState.from_dict(r["resume"]) for r in results} + + def _check_first_step_after_resume(self, steps: int) -> None: + """The first round after ``restore_cut`` advanced every rank's Adam step and the + LR scheduler by exactly ``steps`` optimizer steps (S25 N7.6; A-TRN report #8).""" + from yeto.rl.engine.resume_manifest import ResumeManifest, require_first_step + from .cut_plugin import RESUME_STATE + + saved, self._resume_reference = self._resume_reference, None + results = [dict(r) for r in self._run(self._actor.run_plugin(RESUME_STATE, {}))] + now = ResumeManifest( + global_step=saved.global_step + steps, # the driver owns the global counter + data_watermark=saved.data_watermark, + scheduler_increment=saved.scheduler_increment, + ranks=self._rank_resume_states(results, "first step after resume"), + ) + require_first_step(saved, now, steps=steps) + def _install_actor_recorders(self) -> None: """Record every optimizer step's applied LR and grad norm in the actor processes, once per trainer; without them ``_applied_lrs`` finds no @@ -769,6 +798,7 @@ def save_cut(self, *, epoch: int, context: "CutContext") -> str: problems.append(f"tp{key[0]}/pp{key[1]}/ep{key[2]}: no optimizer state for {missing or 'any adapter'}") if problems: raise CutError("refusing an incomplete cut: " + "; ".join(problems)) + resume = self._resume_manifest(context, summaries) critic = self._save_critic_cut(directory / "critic", context.progress.policy_version) if critic is not None: # 4.3; absent for every critic-free trainer (manifest unchanged) runtime = {**runtime, "critic": critic} @@ -782,6 +812,7 @@ def save_cut(self, *, epoch: int, context: "CutContext") -> str: ledger=context.ledger, outer=context.outer, in_flight=tuple(getattr(context, "in_flight", ()) or ()), + resume=resume.to_dict(), files=tuple( CutFile(s["path"], s["sha256"], int(s["bytes"]), s["coord"]) for s in summaries ), @@ -797,6 +828,19 @@ def save_cut(self, *, epoch: int, context: "CutContext") -> str: commit_manifest(context.root, manifest, check_files=context.shared_filesystem) return cut_id + def _resume_manifest(self, context: "CutContext", summaries: list[Mapping[str, Any]]) -> Any: + """Resume manifest of this cut: run counters plus every rank's optimizer state (S25 N7.6).""" + from yeto.rl.engine.resume_manifest import ResumeManifest, require_saveable + + manifest = ResumeManifest( + global_step=int(context.progress.local_step), + data_watermark=int(context.data["sample_offset"]), + scheduler_increment=int(context.progress.global_batch_size), + ranks=self._rank_resume_states(list(summaries), "save_cut"), + ) + require_saveable(manifest) + return manifest + def _runtime(self, context: "CutContext") -> dict[str, Any]: return { "backend_fingerprint": context.backend_fingerprint, @@ -846,6 +890,7 @@ def restore_cut(self, cut_id: str, *, epoch: int, root: str, expect: Any, it as RECOVERY_REQUIRED and never resumes training on it. """ from yeto.rl.engine.cut import CutError, cut_dir, verify_cut + from yeto.rl.engine.resume_manifest import require from .cut_plugin import RESTORE_CUT_SHARD self._reshard_plan = None # exact same-shape restore: no DP change to guard (review L-1) @@ -876,21 +921,46 @@ def restore_cut(self, cut_id: str, *, epoch: int, root: str, expect: Any, saved = {s["path"]: s for s in manifest.rank_summaries} if len(results) != len(saved) or {r["path"] for r in results} != set(saved): raise CutError(f"restored shards {sorted(r['path'] for r in results)} != cut {sorted(saved)}") + reference, resume_problems = self._restored_resume(manifest, expect, results) + # The resume assertions name a reset Adam step / LR scheduler; the digest + # diagnostics below locate any other difference. Both go into one error. + prefix = ("resume assertions failed: " + "; ".join(resume_problems) + "; ") if resume_problems else "" for r in results: s = saved[r["path"]] for key in ("scheduler_samples", "state_digest", "rng_digest"): if r[key] != s[key]: saved_c, now_c = s.get("components") or {}, r.get("components") or {} differ = sorted(k for k in set(saved_c) | set(now_c) if saved_c.get(k) != now_c.get(k)) - raise CutError(f"{r['path']}: restored {key} differs from the cut; " + raise CutError(f"{prefix}{r['path']}: restored {key} differs from the cut; " f"cut->reexport {r.get('optimizer_cut_vs_reexport')}; " f"cut->after_load {r.get('optimizer_cut_vs_after_load')}; " f"after_load->reexport {r.get('optimizer_after_load_vs_reexport')}; " f"rank diff {r.get('diff')}; " f"differing components ({len(differ)}): {differ[:40]}") + require(resume_problems, "restored optimizer does not continue the cut; refusing to resume") self._restore_critic_cut(cut_dir(root, cut_id) / "critic", manifest) + self._resume_reference = reference return manifest + def _restored_resume(self, manifest: Any, expect: Any, + results: list[Mapping[str, Any]]) -> tuple[Any, list[str]]: + """The cut's resume manifest and every resume assertion the restored ranks fail + (S25 N7.6; A-TRN report #8: a silently reset Adam step).""" + from yeto.rl.engine.cut import CutError + from yeto.rl.engine.resume_manifest import ResumeManifest, restore_problems + + if not manifest.resume: + raise CutError(f"cut {manifest.cut_id!r} has no resume manifest (written before S25 N7.6); " + "refusing to resume without the optimizer-continuity checks") + saved = ResumeManifest.from_dict(manifest.resume) + now = ResumeManifest( + global_step=int(expect.local_step), + data_watermark=int(manifest.data["sample_offset"]), # the cursor the driver restores + scheduler_increment=saved.scheduler_increment, + ranks=self._rank_resume_states(results, "restore_cut"), + ) + return saved, restore_problems(saved, now) + def restore_cut_resharded(self, cut_id: str, *, epoch: int, root: str, expect: Any, plan: Any, certified: Any = None, shared_filesystem: bool = True) -> dict[str, Any]: """Load a cut written at ``plan.source`` DP into this freshly built trainer at ``plan.target`` DP (4.6). @@ -955,6 +1025,9 @@ def restore_cut_resharded(self, cut_id: str, *, epoch: int, root: str, expect: A if expected.get((c["tp"], c["pp"], c["dp"])) != r["rng"]: raise CutError(f"rank {c}: RNG {r['rng']} does not follow the recorded mapping") self._reshard_plan = plan # every later batch is guarded (batch_guard_problems) + # DP change: shard names and optimizer ranges change, so the per-rank resume + # manifest has no counterpart; the resharded restore keeps its own slice checks. + self._resume_reference = None return {"manifest": manifest, "plan": plan.to_dict(), "rng_mapping": mapping, "full_state_digests": {f"tp{k[0]}_pp{k[1]}": g[0]["full_state_digest"] for k, g in sorted(by_group.items())}, diff --git a/yeto/rl/adapters/verl/resume_check.py b/yeto/rl/adapters/verl/resume_check.py new file mode 100644 index 000000000..8de6cf605 --- /dev/null +++ b/yeto/rl/adapters/verl/resume_check.py @@ -0,0 +1,75 @@ +"""verl / FSDP2 binding of the resume assertions (S25 N7.6; A-TRN report #8). + +The checks live in :mod:`yeto.rl.engine.resume_manifest`; this module only maps +a verl FSDP actor worker onto them. verl's ``ActorRolloutRefWorker`` holds the +FSDP2-wrapped actor, its ``torch.optim`` optimizer and LR scheduler as +``actor_module_fsdp``, ``actor_optimizer`` and ``actor_lr_scheduler`` (attribute +names from verl's fsdp worker; not yet checked against a live verl build). + +Each rank reads its own FSDP2 shard (``DTensor.to_local``), so the per-rank +state is keyed by rank. Use: + +* after the checkpoint is complete: ``write_resume_manifest(dir, worker_manifest(...))``; +* after loading it, before the first step: :func:`require_worker_restored`; +* after the first step: :func:`require_worker_first_step`. + +A multi-rank run gathers the per-rank states (e.g. ``all_gather_object``) into +one :class:`ResumeManifest` before writing; every rank then checks its own entry. +""" + +from __future__ import annotations + +from typing import Any + +from yeto.rl.engine.resume_manifest import ( + DEFAULT_SAMPLES, + OptimizerResumeState, + ResumeAssertionError, + ResumeManifest, + require_first_step, + require_restored, + torch_optimizer_state, +) + +WORKER_ATTRS = ("actor_module_fsdp", "actor_optimizer", "actor_lr_scheduler") + + +def rank_label(rank: int) -> str: + return f"rank{int(rank)}" + + +def worker_rank_state(worker: Any, *, samples: int = DEFAULT_SAMPLES) -> OptimizerResumeState: + """This rank's optimizer/scheduler state of a verl FSDP actor worker.""" + missing = [a for a in WORKER_ATTRS if getattr(worker, a, None) is None] + if missing: + raise ResumeAssertionError(f"verl worker {type(worker).__name__} lacks {missing}; " + "cannot check the resumed optimizer") + return torch_optimizer_state(worker.actor_optimizer, worker.actor_module_fsdp.named_parameters(), + worker.actor_lr_scheduler, samples=samples) + + +def worker_manifest(worker: Any, *, rank: int, global_step: int, data_watermark: int, + samples: int = DEFAULT_SAMPLES) -> ResumeManifest: + """Single-rank manifest (``scheduler_increment`` is 1: verl steps its torch scheduler once per step).""" + return ResumeManifest(global_step=int(global_step), data_watermark=int(data_watermark), scheduler_increment=1, + ranks={rank_label(rank): worker_rank_state(worker, samples=samples)}) + + +def _own(saved: ResumeManifest, rank: int) -> ResumeManifest: + label = rank_label(rank) + if label not in saved.ranks: + raise ResumeAssertionError(f"resume manifest has no state for {label} (ranks: {sorted(saved.ranks)})") + return ResumeManifest(saved.global_step, saved.data_watermark, saved.scheduler_increment, + {label: saved.ranks[label]}) + + +def require_worker_restored(saved: ResumeManifest, worker: Any, *, rank: int, global_step: int, + data_watermark: int) -> None: + require_restored(_own(saved, rank), worker_manifest(worker, rank=rank, global_step=global_step, + data_watermark=data_watermark)) + + +def require_worker_first_step(saved: ResumeManifest, worker: Any, *, rank: int, global_step: int, + data_watermark: int, steps: int = 1) -> None: + require_first_step(_own(saved, rank), worker_manifest(worker, rank=rank, global_step=global_step, + data_watermark=data_watermark), steps=steps) diff --git a/yeto/rl/engine/cut.py b/yeto/rl/engine/cut.py index 983da061e..0d6a645b9 100644 --- a/yeto/rl/engine/cut.py +++ b/yeto/rl/engine/cut.py @@ -260,6 +260,10 @@ class CutManifest: # cut (InFlightTrajectory.to_dict()); written only under a policy-age limit # > 0 and serialized only when non-empty (default manifests unchanged). in_flight: tuple[Mapping[str, Any], ...] = () + # S25 N7.6: ResumeManifest.to_dict() (yeto.rl.engine.resume_manifest); + # serialized only when non-empty. Restoring a cut without it is refused + # by the backends that write it. + resume: Mapping[str, Any] = field(default_factory=dict) def body(self) -> dict[str, Any]: return { @@ -276,6 +280,7 @@ def body(self) -> dict[str, Any]: "rank_summaries": [_jsonable(s) for s in self.rank_summaries], **({"backend_state": _jsonable(self.backend_state)} if self.backend_state else {}), **({"in_flight": [_jsonable(t) for t in self.in_flight]} if self.in_flight else {}), + **({"resume": _jsonable(self.resume)} if self.resume else {}), } def digest(self) -> str: @@ -309,6 +314,7 @@ def from_json(cls, text: str) -> "CutManifest": rank_summaries=tuple(raw.get("rank_summaries") or ()), backend_state=raw.get("backend_state") or {}, in_flight=tuple(raw.get("in_flight") or ()), + resume=raw.get("resume") or {}, ) except (KeyError, TypeError, ValueError) as exc: raise CutError(f"malformed cut manifest: {exc}") from exc diff --git a/yeto/rl/engine/resume_manifest.py b/yeto/rl/engine/resume_manifest.py new file mode 100644 index 000000000..bd6785139 --- /dev/null +++ b/yeto/rl/engine/resume_manifest.py @@ -0,0 +1,374 @@ +"""Resume manifest: prove that a restored optimizer continues where the cut stopped. + +A-TRN report #8: restoring optimizer state on some stacks (seen on PPU) silently +resets the Adam ``step``. Bias correction and the LR schedule then restart while +every other check passes. This module is the backend-neutral guard (S25 N7.6): + +1. At save time each rank records an :class:`OptimizerResumeState`: the Adam + ``step`` of every parameter group (min and max over its parameters), the LR + scheduler progress and the applied learning rates, and the ``exp_avg`` / + ``exp_avg_sq`` norm and sha256 of a few fixed sampled parameters. The run + adds the global optimizer step and the data watermark (:class:`ResumeManifest`). +2. After loading, and before the first training step, the restored state of + every rank must equal the saved one (:func:`restore_problems`). A step that + came back as 0 is named explicitly. +3. After the first training step the Adam step and the scheduler progress must + have advanced by exactly the steps taken (:func:`first_step_problems`). + +Any failure raises :class:`ResumeAssertionError` (a :class:`~yeto.rl.engine.cut.CutError`); +the caller refuses to resume and never falls back to a fresh optimizer. + +Backends only produce the per-rank state: :func:`build_rank_state` takes plain +values; :func:`torch_optimizer_state` reads any ``torch.optim`` optimizer +(FSDP2 ``DTensor`` parameters are read as their local shard). torch is imported lazily. +""" + +from __future__ import annotations + +import hashlib +import json +import math +import os +from collections.abc import Iterable, Mapping, Sequence +from dataclasses import dataclass, field +from pathlib import Path +from typing import Any + +from yeto.rl.engine.cut import CutError, fsync_dir + +RESUME_SCHEMA = "yeto.resume_manifest/v1" +RESUME_MANIFEST_NAME = "resume_manifest.json" +MOMENT_KEYS = ("exp_avg", "exp_avg_sq") +DEFAULT_SAMPLES = 4 +LR_REL_TOL = 1e-12 + + +class ResumeAssertionError(CutError): + """The restored (or first-stepped) optimizer does not continue the saved run.""" + + +# -------------------------------------------------------------------------- +# Value types +# -------------------------------------------------------------------------- + + +@dataclass(frozen=True) +class TensorFingerprint: + norm: float # L2 norm in float64 + sha256: str # over the raw bytes, dtype and shape + + def to_dict(self) -> dict[str, Any]: + return {"norm": self.norm, "sha256": self.sha256} + + @classmethod + def from_dict(cls, raw: Mapping[str, Any]) -> "TensorFingerprint": + return cls(float(raw["norm"]), str(raw["sha256"])) + + +def fingerprint(tensor: Any) -> TensorFingerprint: + """Norm and checksum of a tensor (an FSDP2 ``DTensor`` contributes its local shard).""" + import torch + + local = tensor.to_local() if hasattr(tensor, "to_local") else tensor + flat = local.detach().to("cpu").contiguous().reshape(-1) + digest = hashlib.sha256(f"{flat.dtype}|{tuple(local.shape)}|".encode()) + digest.update(flat.view(torch.uint8).numpy().tobytes()) + return TensorFingerprint(float(flat.to(torch.float64).norm()), digest.hexdigest()) + + +@dataclass(frozen=True) +class GroupSteps: + """Adam ``step`` over the parameters of one group; a healthy group has min == max.""" + + min: int + max: int + + def to_dict(self) -> dict[str, int]: + return {"min": self.min, "max": self.max} + + +@dataclass(frozen=True) +class OptimizerResumeState: + """What one rank records about its optimizer and LR scheduler.""" + + adam_steps: Mapping[str, GroupSteps] # group label -> steps + scheduler_progress: int # torch last_epoch, Megatron num_steps (samples) + lrs: tuple[float, ...] # applied LR per param group + moments: Mapping[str, Mapping[str, TensorFingerprint]] # sampled param -> exp_avg/exp_avg_sq + + def to_dict(self) -> dict[str, Any]: + return { + "adam_steps": {k: v.to_dict() for k, v in sorted(self.adam_steps.items())}, + "scheduler_progress": self.scheduler_progress, + "lrs": list(self.lrs), + "moments": {n: {k: f.to_dict() for k, f in sorted(m.items())} for n, m in sorted(self.moments.items())}, + } + + @classmethod + def from_dict(cls, raw: Mapping[str, Any]) -> "OptimizerResumeState": + return cls( + adam_steps={k: GroupSteps(int(v["min"]), int(v["max"])) for k, v in raw["adam_steps"].items()}, + scheduler_progress=int(raw["scheduler_progress"]), + lrs=tuple(float(x) for x in raw["lrs"]), + moments={n: {k: TensorFingerprint.from_dict(f) for k, f in m.items()} for n, m in raw["moments"].items()}, + ) + + +@dataclass(frozen=True) +class ResumeManifest: + """Run-level counters plus one :class:`OptimizerResumeState` per rank (keyed by shard name).""" + + global_step: int # optimizer steps applied by this run + data_watermark: int # samples consumed (data cursor position) + scheduler_increment: int # scheduler progress per optimizer step (1 for torch, GBS for Megatron) + ranks: Mapping[str, OptimizerResumeState] = field(default_factory=dict) + schema: str = RESUME_SCHEMA + + def to_dict(self) -> dict[str, Any]: + return { + "schema": self.schema, + "global_step": self.global_step, + "data_watermark": self.data_watermark, + "scheduler_increment": self.scheduler_increment, + "ranks": {k: v.to_dict() for k, v in sorted(self.ranks.items())}, + } + + @classmethod + def from_dict(cls, raw: Mapping[str, Any]) -> "ResumeManifest": + if raw.get("schema") != RESUME_SCHEMA: + raise ResumeAssertionError(f"unknown resume manifest schema {raw.get('schema')!r}") + try: + return cls( + global_step=int(raw["global_step"]), + data_watermark=int(raw["data_watermark"]), + scheduler_increment=int(raw["scheduler_increment"]), + ranks={k: OptimizerResumeState.from_dict(v) for k, v in raw["ranks"].items()}, + ) + except (KeyError, TypeError, ValueError) as exc: + raise ResumeAssertionError(f"malformed resume manifest: {exc}") from exc + + +# -------------------------------------------------------------------------- +# Building the state +# -------------------------------------------------------------------------- + + +def sample_names(names: Iterable[str], count: int = DEFAULT_SAMPLES) -> list[str]: + """A fixed, evenly spread subset of the sorted names (first and last always included).""" + ordered = sorted(set(names)) + if count <= 0 or not ordered: + return [] + if len(ordered) <= count: + return ordered + if count == 1: + return [ordered[0]] + picks = {round(i * (len(ordered) - 1) / (count - 1)) for i in range(count)} + return [ordered[i] for i in sorted(picks)] + + +def _as_int_step(value: Any) -> int: + step = float(value.item() if hasattr(value, "item") else value) + if not math.isfinite(step) or step != int(step): + raise ResumeAssertionError(f"Adam step {step!r} is not a whole number") + return int(step) + + +def build_rank_state( + *, + steps_by_group: Mapping[str, Iterable[Any]], + scheduler_progress: int, + lrs: Sequence[float], + moments_by_name: Mapping[str, Mapping[str, Any]], + samples: int = DEFAULT_SAMPLES, +) -> OptimizerResumeState: + """Plain values -> rank state. ``moments_by_name`` maps a parameter name to its + moment tensors; only sampled names that carry every key in :data:`MOMENT_KEYS` + are fingerprinted. A group without any step value counts as step 0 (no state yet).""" + adam_steps = {} + for label, values in steps_by_group.items(): + steps = [_as_int_step(v) for v in values] or [0] + adam_steps[str(label)] = GroupSteps(min(steps), max(steps)) + eligible = [n for n, m in moments_by_name.items() if all(k in m for k in MOMENT_KEYS)] + moments = {n: {k: fingerprint(moments_by_name[n][k]) for k in MOMENT_KEYS} for n in sample_names(eligible, samples)} + return OptimizerResumeState(adam_steps, int(scheduler_progress), tuple(float(x) for x in lrs), moments) + + +def scheduler_progress_of(scheduler: Any) -> int: + """torch ``LRScheduler.last_epoch`` or Megatron ``OptimizerParamScheduler.num_steps``.""" + for attr in ("last_epoch", "num_steps"): + value = getattr(scheduler, attr, None) + if value is not None: + return int(value) + raise ResumeAssertionError(f"cannot read the progress of LR scheduler {type(scheduler).__name__}") + + +def torch_optimizer_state( + optimizer: Any, + named_parameters: Iterable[tuple[str, Any]], + scheduler: Any, + *, + samples: int = DEFAULT_SAMPLES, +) -> OptimizerResumeState: + """Read any ``torch.optim`` Adam-family optimizer (plain module or FSDP2).""" + names = {id(p): n for n, p in named_parameters} + steps: dict[str, list[Any]] = {} + moments: dict[str, Mapping[str, Any]] = {} + for index, group in enumerate(optimizer.param_groups): + values = steps.setdefault(f"group{index}", []) + if "step" in group: # apex/TE fused Adam keeps one step per group + values.append(group["step"]) + for p in group["params"]: + state = optimizer.state.get(p) or {} + if "step" in state: + values.append(state["step"]) + if id(p) in names: + moments[names[id(p)]] = state + lrs = [float(g["lr"]) for g in optimizer.param_groups] + return build_rank_state(steps_by_group=steps, scheduler_progress=scheduler_progress_of(scheduler), + lrs=lrs, moments_by_name=moments, samples=samples) + + +# -------------------------------------------------------------------------- +# Checks +# -------------------------------------------------------------------------- + + +def _lr_close(a: float, b: float, rel: float = LR_REL_TOL) -> bool: + return a == b or abs(a - b) <= rel * max(abs(a), abs(b)) + + +def _step_problems(rank: str, saved: Mapping[str, GroupSteps], now: Mapping[str, GroupSteps], + advance: int) -> list[str]: + out = [] + if set(saved) != set(now): + return [f"{rank}: param groups {sorted(now)} != saved {sorted(saved)}"] + for label in sorted(saved): + want = saved[label].max + advance + got = now[label] + if got.min != got.max: + out.append(f"{rank}/{label}: Adam step differs across parameters ({got.min}..{got.max}); " + "part of the optimizer state was reset") + elif got.max != want: + reset = " (Adam step was reset)" if got.max == 0 and want > 0 else "" + out.append(f"{rank}/{label}: Adam step {got.max} != expected {want}{reset}") + return out + + +def _rank_pairs(saved: ResumeManifest, now: ResumeManifest) -> tuple[list[str], list[tuple[str, Any, Any]]]: + if set(saved.ranks) != set(now.ranks): + return [f"ranks {sorted(now.ranks)} != saved {sorted(saved.ranks)}"], [] + return [], [(r, saved.ranks[r], now.ranks[r]) for r in sorted(saved.ranks)] + + +def save_problems(manifest: ResumeManifest) -> list[str]: + """A manifest that cannot serve as a resume reference.""" + out = [] + if manifest.global_step < 0 or manifest.data_watermark < 0 or manifest.scheduler_increment <= 0: + out.append("global_step/data_watermark must be >= 0 and scheduler_increment > 0") + if not manifest.ranks: + out.append("no rank state") + for rank, state in sorted(manifest.ranks.items()): + for label, steps in sorted(state.adam_steps.items()): + if steps.min != steps.max: + out.append(f"{rank}/{label}: Adam step differs across parameters ({steps.min}..{steps.max})") + if manifest.global_step > 0 and manifest.ranks and not any(s.moments for s in manifest.ranks.values()): + # a DistOpt rank may own no range; the run as a whole must have sampled moments + out.append(f"no rank has a parameter with exp_avg/exp_avg_sq after {manifest.global_step} steps") + return out + + +def restore_problems(saved: ResumeManifest, now: ResumeManifest) -> list[str]: + """Restored state vs the saved one; every item must match exactly (LR within 1e-12 relative).""" + out = [] + if now.global_step != saved.global_step: + out.append(f"global step {now.global_step} != saved {saved.global_step}") + if now.data_watermark != saved.data_watermark: + out.append(f"data watermark {now.data_watermark} != saved {saved.data_watermark}") + if now.scheduler_increment != saved.scheduler_increment: + out.append(f"scheduler increment {now.scheduler_increment} != saved {saved.scheduler_increment}") + problems, pairs = _rank_pairs(saved, now) + out += problems + for rank, s, n in pairs: + out += _step_problems(rank, s.adam_steps, n.adam_steps, 0) + if n.scheduler_progress != s.scheduler_progress: + reset = " (scheduler was reset)" if n.scheduler_progress == 0 and s.scheduler_progress > 0 else "" + out.append(f"{rank}: scheduler progress {n.scheduler_progress} != saved {s.scheduler_progress}{reset}") + if len(n.lrs) != len(s.lrs) or not all(_lr_close(a, b) for a, b in zip(n.lrs, s.lrs)): + out.append(f"{rank}: learning rates {list(n.lrs)} != saved {list(s.lrs)}") + for name in sorted(s.moments): + if name not in n.moments: + out.append(f"{rank}: sampled parameter {name!r} has no exp_avg/exp_avg_sq after restore") + continue + for key in MOMENT_KEYS: + if n.moments[name][key] != s.moments[name][key]: + out.append(f"{rank}: {name}.{key} norm {n.moments[name][key].norm!r} / checksum differs " + f"from saved (norm {s.moments[name][key].norm!r})") + return out + + +def first_step_problems(saved: ResumeManifest, now: ResumeManifest, *, steps: int = 1, + expected_lrs: Mapping[str, Sequence[float]] | None = None) -> list[str]: + """After ``steps`` optimizer steps on the restored state: every counter advanced by exactly ``steps``.""" + out = [] + if now.global_step != saved.global_step + steps: + out.append(f"global step {now.global_step} != saved {saved.global_step} + {steps}") + problems, pairs = _rank_pairs(saved, now) + out += problems + for rank, s, n in pairs: + out += _step_problems(rank, s.adam_steps, n.adam_steps, steps) + want = s.scheduler_progress + steps * saved.scheduler_increment + if n.scheduler_progress != want: + out.append(f"{rank}: scheduler progress {n.scheduler_progress} != expected {want}") + lrs = (expected_lrs or {}).get(rank) + if lrs is not None and (len(lrs) != len(n.lrs) or not all(_lr_close(a, b) for a, b in zip(n.lrs, lrs))): + out.append(f"{rank}: learning rates {list(n.lrs)} != scheduled {list(lrs)}") + return out + + +def require(problems: Sequence[str], what: str) -> None: + if problems: + raise ResumeAssertionError(f"{what}: " + "; ".join(problems)) + + +def require_saveable(manifest: ResumeManifest) -> None: + require(save_problems(manifest), "refusing to record a resume manifest") + + +def require_restored(saved: ResumeManifest, now: ResumeManifest) -> None: + require(restore_problems(saved, now), "restored optimizer does not continue the cut; refusing to resume") + + +def require_first_step(saved: ResumeManifest, now: ResumeManifest, *, steps: int = 1, + expected_lrs: Mapping[str, Sequence[float]] | None = None) -> None: + require(first_step_problems(saved, now, steps=steps, expected_lrs=expected_lrs), + "first step after resume did not advance the optimizer by exactly the steps taken") + + +# -------------------------------------------------------------------------- +# File form (backends whose checkpoint has no cut manifest, e.g. FSDP2) +# -------------------------------------------------------------------------- + + +def write_resume_manifest(directory: str | os.PathLike[str], manifest: ResumeManifest) -> Path: + """Write ``resume_manifest.json`` atomically; call it after the checkpoint itself is complete.""" + require_saveable(manifest) + final = Path(directory) / RESUME_MANIFEST_NAME + tmp = final.with_name(RESUME_MANIFEST_NAME + ".tmp") + with open(tmp, "w", encoding="utf-8") as fh: + json.dump(manifest.to_dict(), fh, indent=1, sort_keys=True) + fh.flush() + os.fsync(fh.fileno()) + os.replace(tmp, final) + fsync_dir(final.parent) + return final + + +def read_resume_manifest(directory: str | os.PathLike[str]) -> ResumeManifest: + path = Path(directory) / RESUME_MANIFEST_NAME + if not path.is_file(): + raise ResumeAssertionError(f"checkpoint {directory} has no {RESUME_MANIFEST_NAME}; refusing to resume") + try: + raw = json.loads(path.read_text(encoding="utf-8")) + except json.JSONDecodeError as exc: + raise ResumeAssertionError(f"{path} is not JSON: {exc}") from exc + return ResumeManifest.from_dict(raw) From a7a646e73bfb6519393bf0f165590dc00a7c886a Mon Sep 17 00:00:00 2001 From: michael <78113554+chowhsu@users.noreply.github.com> Date: Mon, 19 Oct 2026 08:16:48 +0000 Subject: [PATCH 03/34] Add MXFP4 group-32 RTN re-quantization for routed experts (S25 N8.5) New yeto/rl/quant/mxfp4.py implements the compressed-tensors mxfp4-pack-quantized rules: E8M0 shared scale per 32 inputs (BF16 amax rounded to a power of two at mantissa 1.75, minus 2), FP4 E2M1 ties to even with saturation at 6, sign bit for w < 0, even column in the low nibble. requant.py now selects a QuantScheme from the template's quantization_config.format: w4a8-int4-channel (int-quantized, unchanged behaviour) or mxfp4-g32-rtn (weight_packed + weight_scale, refuses AWQ smoothing). compare_checkpoints reports FP4 and E8M0 code match shares. The release CLI takes its scheme tag from the template. Checked on CPU against yetalabs/apex-flash-1-mxfp4-experts@70e86a0: layer-10 routed experts (296 modules) equal in every FP4 and E8M0 code. Co-Authored-By: Claude Opus 5.5 --- tests/test_rl_quant_mxfp4.py | 339 +++++++++++++++++++++++++++++++++++ yeto/rl/quant/__init__.py | 12 +- yeto/rl/quant/mxfp4.py | 139 ++++++++++++++ yeto/rl/quant/release.py | 6 +- yeto/rl/quant/requant.py | 260 ++++++++++++++++++++------- 5 files changed, 685 insertions(+), 71 deletions(-) create mode 100644 tests/test_rl_quant_mxfp4.py create mode 100644 yeto/rl/quant/mxfp4.py diff --git a/tests/test_rl_quant_mxfp4.py b/tests/test_rl_quant_mxfp4.py new file mode 100644 index 000000000..34310e09d --- /dev/null +++ b/tests/test_rl_quant_mxfp4.py @@ -0,0 +1,339 @@ +"""Re-quantization BF16 -> MXFP4 experts, group 32, RTN (S25 N8.5). + +Codec rules (E8M0 rounding, FP4 ties, nibble order, sign of zero) are checked +against a literal float re-statement of compressed-tensors 0.18, written +without the package code; the layout tests build a tiny reference artefact in +the format of ``yetalabs/apex-flash-1-mxfp4-experts@70e86a0`` (only routed +experts as ``weight_packed`` + ``weight_scale``, everything else BF16/F32). +""" + +from __future__ import annotations + +import json +import math +from pathlib import Path + +import pytest +import torch +from safetensors.torch import save_file + +from yeto.rl.quant.awq import AwqSmoothing +from yeto.rl.quant.mxfp4 import ( + FP4_VALUES, + GROUP_SIZE, + dequantize_mxfp4, + e8m0_codes, + e8m0_to_scale, + fp4_decode, + fp4_encode, + pack_fp4, + quantize_mxfp4, + unpack_fp4, +) +from yeto.rl.quant.requant import ( + INDEX_FILE, + MXFP4_G32_RTN, + CheckpointLayout, + LayoutError, + compare_checkpoints, + main, + requantize_checkpoint, + transform_tensor, +) + + +def _bf16(*xs: float) -> torch.Tensor: + return torch.tensor(xs, dtype=torch.bfloat16) + + +def _reference_quantize(w: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: + """compressed-tensors 0.18 MXFP4 path, re-stated in plain float math (no package code).""" + rows, cols = w.shape + out_codes, out_scale = [], [] + for r in range(rows): + row_codes, row_scale = [], [] + for g in range(0, cols, GROUP_SIZE): + grp = [float(v) for v in w[r, g:g + GROUP_SIZE].tolist()] + amax = max(abs(v) for v in grp) + if amax == 0: + exp = 0 + else: + m, e = math.frexp(amax) # amax = m * 2**e, m in [0.5, 1) + e -= 1 # amax = (2m) * 2**e, 2m in [1, 2) + exp = max(0, 127 + e + (1 if 2 * m >= 1.75 else 0) - 2) + scale = 2.0 ** (exp - 127) + row_scale.append(exp) + for v in grp: + x = v / scale + mag = min(abs(x), 6.0) + # nearest E2M1 value, ties to the even mantissa (0, 1, 2, 4 are "even") + best = min(range(8), key=lambda k: (abs(FP4_VALUES[k] - mag), k % 2)) + row_codes.append(best | (8 if v < 0 else 0)) + out_codes.append(row_codes) + out_scale.append(row_scale) + codes = torch.tensor(out_codes, dtype=torch.uint8) + return codes[:, 0::2] | (codes[:, 1::2] << 4), torch.tensor(out_scale, dtype=torch.uint8) + + +# -- codec -------------------------------------------------------------------- + + +def test_e8m0_rounds_up_from_mantissa_1_75(): + # amax = 1.0 -> 2**0 -> code 127 - 2; 1.7421875 (largest BF16 below 1.75) stays; 1.75 carries. + codes = e8m0_codes(_bf16(1.0, 1.7421875, 1.75, 6.0, 0.5, 2.0**-130, 0.0)) + assert codes.tolist() == [125, 125, 126, 127, 124, 0, 0] + + +def test_e8m0_extremes(): + tiny = torch.tensor([2.0**-126, 2.0**-125, 2.0**-124], dtype=torch.bfloat16) + assert e8m0_codes(tiny).tolist() == [0, 0, 1] # 127 - 126 - 2 < 0 clamps to 0 + assert e8m0_codes(_bf16(2.0**127)).item() == 252 + with pytest.raises(ValueError, match="too large"): + e8m0_codes(torch.tensor([torch.finfo(torch.bfloat16).max], dtype=torch.bfloat16)) + with pytest.raises(ValueError, match="non-finite"): + e8m0_codes(_bf16(float("inf"))) + with pytest.raises(TypeError): + e8m0_codes(torch.ones(2)) + assert e8m0_to_scale(torch.tensor([127, 125, 0], dtype=torch.uint8)).tolist() == [1.0, 0.25, 2.0**-127] + + +def test_fp4_ties_to_even_saturation_and_sign_of_zero(): + x = torch.tensor([0.25, 0.26, 0.75, 1.25, 1.75, 2.5, 3.5, 5.0, 5.01, 100.0, -0.1, -0.0, 0.0, -6.0]) + assert fp4_encode(x).tolist() == [0, 1, 2, 2, 4, 4, 6, 6, 7, 7, 8, 0, 0, 15] + codes = torch.arange(16, dtype=torch.uint8) + assert fp4_decode(codes).tolist() == list(FP4_VALUES) + [-v for v in FP4_VALUES] + + +def test_pack_low_nibble_is_the_even_column_and_round_trips(): + assert pack_fp4(torch.tensor([[1, 2, 15, 0]], dtype=torch.uint8)).tolist() == [[0x21, 0x0F]] + codes = torch.randint(0, 16, (5, 64), dtype=torch.uint8, generator=torch.Generator().manual_seed(1)) + assert torch.equal(unpack_fp4(pack_fp4(codes)), codes) + with pytest.raises(ValueError): + pack_fp4(torch.zeros(2, 3, dtype=torch.uint8)) + with pytest.raises(ValueError): + pack_fp4(torch.full((1, 2), 16, dtype=torch.uint8)) + + +def test_quantize_matches_literal_reference_on_mixed_groups(): + g = torch.Generator().manual_seed(3) + w = (torch.randn(6, 4 * GROUP_SIZE, generator=g) * 0.02).to(torch.bfloat16) + w[1, :GROUP_SIZE] = 0 # all-zero group + w[2, GROUP_SIZE:2 * GROUP_SIZE] = -0.0 # all negative zeros + w[2, 2 * GROUP_SIZE] = -1e-6 # negative, rounds to zero: stored as -0 + w[3, :GROUP_SIZE] = 1e-4 + w[3, 5] = 300.0 # one extreme value: the others round to zero + w[4, 2 * GROUP_SIZE:3 * GROUP_SIZE] = torch.tensor(FP4_VALUES * 4, dtype=torch.bfloat16) * 2.0**-3 + packed, scale = quantize_mxfp4(w) + ref_packed, ref_scale = _reference_quantize(w) + assert torch.equal(scale, ref_scale) and torch.equal(packed, ref_packed) + assert scale.shape == (6, 4) and packed.shape == (6, 2 * GROUP_SIZE) + assert scale[1, 0] == 0 and (unpack_fp4(packed)[1, :GROUP_SIZE] == 0).all() + assert (unpack_fp4(packed)[2, GROUP_SIZE:2 * GROUP_SIZE] == 0).all() # input -0.0 is stored as +0 + assert int((unpack_fp4(packed)[3, :GROUP_SIZE] & 7).count_nonzero()) == 1 + assert unpack_fp4(packed)[2, 2 * GROUP_SIZE] == 8 + + +def test_representable_values_round_trip_exactly(): + g = torch.Generator().manual_seed(4) + codes = torch.randint(0, 16, (8, 4 * GROUP_SIZE), dtype=torch.uint8, generator=g) + codes[:, ::GROUP_SIZE] = 7 # every group reaches 6 * scale, so the scale is recovered + exps = torch.randint(110, 135, (8, 4), dtype=torch.uint8, generator=g) + w = dequantize_mxfp4(pack_fp4(codes), exps, torch.bfloat16) + packed, scale = quantize_mxfp4(w) + assert torch.equal(scale, exps) + assert torch.equal(dequantize_mxfp4(packed, scale), w.float()) + + +def test_round_trip_error_is_bounded_by_the_group_step(): + g = torch.Generator().manual_seed(5) + w = (torch.randn(16, 8 * GROUP_SIZE, generator=g) * 0.03).to(torch.bfloat16) + packed, scale = quantize_mxfp4(w) + err = (dequantize_mxfp4(packed, scale) - w.float()).abs().reshape(16, 8, GROUP_SIZE) + step = e8m0_to_scale(scale).unsqueeze(-1) + # Largest gap between E2M1 values is 2 (4 -> 6); saturation above 6 loses < 1 step (amax < 7 * scale). + assert (err <= step).all() + + +def test_rows_are_independent(): + w = (torch.randn(8, 2 * GROUP_SIZE, generator=torch.Generator().manual_seed(6))).to(torch.bfloat16) + packed, scale = quantize_mxfp4(w) + p2, s2 = quantize_mxfp4(w[3:6]) + assert torch.equal(p2, packed[3:6]) and torch.equal(s2, scale[3:6]) + + +def test_quantize_rejects_bad_input(): + with pytest.raises(ValueError, match="multiple of the group size"): + quantize_mxfp4(torch.zeros(2, 48, dtype=torch.bfloat16)) + with pytest.raises(TypeError, match="BF16"): + quantize_mxfp4(torch.zeros(2, 32)) + with pytest.raises(ValueError, match="2-D"): + quantize_mxfp4(torch.zeros(32, dtype=torch.bfloat16)) + with pytest.raises(ValueError, match="does not fit"): + dequantize_mxfp4(torch.zeros(2, 16, dtype=torch.uint8), torch.zeros(2, 2, dtype=torch.uint8)) + + +def test_matches_compressed_tensors_if_installed(): + """Cross-check with the library that wrote the reference artefact (skipped where it is not installed).""" + ct = pytest.importorskip("compressed_tensors.compressors.mxfp4.base") + from compressed_tensors.quantization import QuantizationArgs, QuantizationScheme + from compressed_tensors.quantization.utils.helpers import calculate_qparams + + g = torch.Generator().manual_seed(8) + w = (torch.randn(4, 4 * GROUP_SIZE, generator=g) * 0.05).to(torch.bfloat16) + w[0, :GROUP_SIZE] = 0 + w[1, 3] = -0.0 + w[2, 7] = 0.4 + w[3, 9] = -1e-6 + args = QuantizationArgs(num_bits=4, type="float", strategy="group", group_size=GROUP_SIZE, + symmetric=True, scale_dtype=torch.uint8) + groups = w.reshape(4, -1, GROUP_SIZE) + scale, _ = calculate_qparams(groups.amin(-1), groups.amax(-1), args) + out = ct.MXFP4PackedCompressor.compress({"weight": w, "weight_scale": scale}, + QuantizationScheme(targets=["Linear"], weights=args)) + packed, codes = quantize_mxfp4(w) + assert torch.equal(out["weight_scale"], codes) and torch.equal(out["weight_packed"], packed) + + +# -- checkpoint layout ---------------------------------------------------------- + +H, I, E = 64, 32, 2 # hidden, expert width, routed experts +L = "model.language_model.layers.3." +MTP = "model.language_model.layers.4." # stands for the ignored MTP layer 45 + +QCONFIG = { + "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:.*layers\\.4\\..*", "re:^(?!.*mlp\\.experts\\.\\d+\\.).*$"], +} + + +def _source_tensors() -> dict[str, torch.Tensor]: + g = torch.Generator().manual_seed(11) + + def rnd(*shape: int) -> torch.Tensor: + return (torch.randn(*shape, generator=g) * 0.05).to(torch.bfloat16) + + t = {L + "post_attention_layernorm.weight": rnd(H), L + "mlp.gate.weight": rnd(E, H), + L + "mlp.gate.e_score_correction_bias": torch.randn(E, generator=g), + L + "mlp.shared_experts.gate_proj.weight": rnd(I, H), "lm_head.weight": rnd(8, H)} + for layer in (L, MTP): + for e in range(E): + t[f"{layer}mlp.experts.{e}.gate_proj.weight"] = rnd(I, H) + t[f"{layer}mlp.experts.{e}.up_proj.weight"] = rnd(I, H) + t[f"{layer}mlp.experts.{e}.down_proj.weight"] = rnd(H, I) + t[L + "mlp.experts.0.down_proj.weight"][2, :GROUP_SIZE] = 0 # an all-zero group + return t + + +def _reference_artefact(src: dict[str, torch.Tensor]) -> dict[str, torch.Tensor]: + out: dict[str, torch.Tensor] = {} + for name, w in src.items(): + if name.startswith(L + "mlp.experts."): + module = name.removesuffix(".weight") + out[module + ".weight_packed"], out[module + ".weight_scale"] = _reference_quantize(w) + else: + out[name] = w + return out + + +def _shard_of(name: str) -> str: + return "model-00002-of-00002.safetensors" if "experts.1." in name else "model-00001-of-00002.safetensors" + + +def _write_ckpt(d: Path, tensors: dict[str, torch.Tensor], shard_of) -> None: + d.mkdir(parents=True, exist_ok=True) + shards: dict[str, dict[str, torch.Tensor]] = {} + for n, t in tensors.items(): + shards.setdefault(shard_of(n), {})[n] = t.contiguous() + for f, ts in shards.items(): + save_file(ts, str(d / f), metadata={"format": "pt"}) + (d / INDEX_FILE).write_text(json.dumps({"metadata": {}, "weight_map": {n: shard_of(n) for n in tensors}})) + + +@pytest.fixture +def world(tmp_path: Path): + src = _source_tensors() + original = tmp_path / "mxfp4-experts" + _write_ckpt(original, _reference_artefact(src), _shard_of) + (original / "config.json").write_text(json.dumps({"model_type": "glm5_next", "quantization_config": QCONFIG})) + source = tmp_path / "bf16" + _write_ckpt(source, src, lambda n: "model-00001-of-00001.safetensors") + template = tmp_path / "template" + CheckpointLayout.from_checkpoint(original).save_template(original, template) + return {"original": original, "source": source, "template": template, "src": src, "tmp": tmp_path} + + +def test_layout_selects_mxfp4_scheme_and_source_names(world): + layout = CheckpointLayout.load(world["template"]) + assert layout.scheme is MXFP4_G32_RTN and layout.scheme.name == "mxfp4-g32-rtn" + assert len(layout.quantized_weights()) == 3 * E and all(n.endswith(".weight_packed") for n in layout.quantized_weights()) + assert layout.source_names() == set(world["src"]) + + +def test_untrained_source_reproduces_reference_bit_for_bit(world): + out = world["tmp"] / "out" + stats = requantize_checkpoint(world["source"], world["template"], out, AwqSmoothing.none()) + assert stats["quantized"] == 3 * E + rep = compare_checkpoints(out, world["original"]) + assert rep["fp4_match_share"] == 1.0 and rep["e8m0_match_share"] == 1.0 and rep["unequal_tensors"] == 0 + assert rep["tensors"] == len(_reference_artefact(world["src"])) + assert (out / "config.json").is_file() + + +def test_awq_smoothing_is_refused(world, tmp_path): + torch.save(torch.ones(H, dtype=torch.float64), tmp_path / f"{L}post_attention_layernorm.pt") + (tmp_path / "recipe.yaml").write_text( + f"default_stage:\n default_modifiers:\n RecordedAWQ:\n mappings:\n" + f" - smooth_layer: {L}post_attention_layernorm\n balance_layers: ['{L}mlp.gate']\n") + smoothing = AwqSmoothing.from_artefact(tmp_path / "recipe.yaml", tmp_path) + with pytest.raises(LayoutError, match="round-to-nearest only"): + requantize_checkpoint(world["source"], world["template"], world["tmp"] / "x", smoothing) + + +def test_transform_row_slice_gives_the_same_rows(world): + layout = CheckpointLayout.load(world["template"]) + name = L + "mlp.experts.1.down_proj.weight" + full = transform_tensor(name, world["src"][name], layout, AwqSmoothing.none()) + part = transform_tensor(name, world["src"][name][10:20], layout, AwqSmoothing.none()) + for k in full: + assert torch.equal(part[k], full[k][10:20]) + with pytest.raises(LayoutError, match="not a source tensor"): + transform_tensor(L + "mlp.experts.1.down_proj.weight_scale", full[L + "mlp.experts.1.down_proj.weight_scale"], + layout, AwqSmoothing.none()) + + +@pytest.mark.parametrize("field,value", [("group_size", 16), ("scale_dtype", "torch.bfloat16"), ("type", "int")]) +def test_unsupported_mxfp4_config_is_rejected(world, field, value): + cfg = json.loads((world["original"] / "config.json").read_text()) + cfg["quantization_config"]["config_groups"]["MXFP4A16"]["weights"][field] = value + (world["original"] / "config.json").write_text(json.dumps(cfg)) + with pytest.raises(LayoutError, match="not supported"): + CheckpointLayout.from_checkpoint(world["original"]) + + +def test_dense_format_is_rejected(world): + # Pit 2 (A-QNT): llmcompressor 0.14 wrote format=dense for MXFP4 and readers mis-read the weights. + cfg = json.loads((world["original"] / "config.json").read_text()) + cfg["quantization_config"]["format"] = "dense" + (world["original"] / "config.json").write_text(json.dumps(cfg)) + with pytest.raises(LayoutError, match="unsupported quantization format"): + CheckpointLayout.from_checkpoint(world["original"]) + + +def test_cli_template_convert_compare(world, capsys): + tpl, out = world["tmp"] / "tpl-cli", world["tmp"] / "out-cli" + assert main(["template", str(world["original"]), str(tpl)]) == 0 + assert main(["convert", "--source", str(world["source"]), "--template", str(tpl), "--out", str(out), "--no-awq"]) == 0 + capsys.readouterr() + 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 diff --git a/yeto/rl/quant/__init__.py b/yeto/rl/quant/__init__.py index 147a247ec..61357cde1 100644 --- a/yeto/rl/quant/__init__.py +++ b/yeto/rl/quant/__init__.py @@ -14,10 +14,18 @@ (``awq-scales/.pt``), so option R2 reads them instead of deriving them. +A second target (N8.5) is the MXFP4 experts artefact +``yetalabs/apex-flash-1-mxfp4-experts@70e86a0``: compressed-tensors +``mxfp4-pack-quantized``, routed experts only, ``weight_packed`` (U8, two FP4 +E2M1 codes per byte) + ``weight_scale`` (U8 E8M0, one per 32 inputs), RTN +without AWQ; all other tensors BF16/F32 as in the source. The converter picks +the scheme from the template's ``quantization_config.format``. + Modules: * :mod:`.int4`: per-channel symmetric int4 quantize / dequantize. +* :mod:`.mxfp4`: MXFP4 group-32 quantize / dequantize, E8M0 scales, FP4 packing. * :mod:`.awq`: the fixed AWQ smoothing (R2) and the no-smoothing fallback (R1). -* :mod:`.requant`: the checkpoint layout, the per-tensor transform, the - shard-by-shard converter, and the tensor-by-tensor comparison. +* :mod:`.requant`: the checkpoint layout and its scheme, the per-tensor + transform, the shard-by-shard converter, and the tensor-by-tensor comparison. """ diff --git a/yeto/rl/quant/mxfp4.py b/yeto/rl/quant/mxfp4.py new file mode 100644 index 000000000..371650a37 --- /dev/null +++ b/yeto/rl/quant/mxfp4.py @@ -0,0 +1,139 @@ +"""MXFP4 weight quantization, group 32, RTN (compressed-tensors ``mxfp4-pack-quantized``). + +The rules reproduce compressed-tensors 0.18 (the reference artefact +``yetalabs/apex-flash-1-mxfp4-experts@70e86a0`` declares 0.19.0) for +``num_bits=4, type=float, strategy=group, group_size=32, scale_dtype=uint8``: + +* Groups are 32 consecutive input columns of one output row. +* Shared scale (E8M0, uint8, bias 127): take the group ``amax`` in BF16, round + it to a power of two with the mantissa threshold 1.75 (``round_to_power_2``: + add ``1 << 5`` to the BF16 bits, keep sign and exponent), then subtract 2 + (``floor(log2(6))``, the FP4 E2M1 maximum) from the exponent. The code is + clamped below at 0, so codes lie in ``[0, 252]``; an all-zero group gets + code 0. A group whose amax rounds to infinity (``>= 1.75 * 2**127``) is + refused; compressed-tensors would store code 255 (E8M0 NaN) there. +* Element (FP4 E2M1): ``w / 2**(code - 127)`` rounded to the nearest value of + ``{0, 0.5, 1, 1.5, 2, 3, 4, 6}``, ties to even mantissa, saturating at 6. The + sign bit is set for ``w < 0``, so a negative value that rounds to zero is + stored as code ``0b1000`` (``-0``) while an input ``-0.0`` is stored as code 0, + as ``cast_to_fp4`` + ``pack_fp4_to_uint8`` do. +* Packing: two codes per byte along the input dimension, the even column in the + low nibble. ``weight_packed`` is ``uint8 [out, in / 2]``, ``weight_scale`` is + ``uint8 [out, in / 32]``. + +No calibration data and no smoothing: round to nearest only (RTN). Each output +row is quantized on its own, so any row slice of a weight gives the same rows +as the whole weight (the basis for per-rank quantization, D7). + +Evidence (CPU, 2026-10-19): the BF16 layer-10 routed experts of +``apex-flash-1`` (288 ``down_proj``, 4 ``gate_proj`` / ``up_proj``; native +layout copy on Modal) quantized with this module equal the reference artefact +in every FP4 code (2.48e9, 1.2e8 of them ``-0``) and every E8M0 code (7.8e7). +``tests/test_rl_quant_mxfp4.py`` also cross-checks compressed-tensors itself +where it is installed. +""" + +from __future__ import annotations + +import torch + +GROUP_SIZE = 32 +E8M0_BIAS = 127 +# floor(log2(FP4_E2M1_MAX)) with FP4_E2M1_MAX = 6.0. +_ELEM_EXP_OFFSET = 2 +# BF16: 7 mantissa bits; FP4 E2M1: 1. Adding 1 << (7 - 1 - 1) carries every +# mantissa >= 1.75 into the next exponent before the mantissa is dropped. +_BF16_ROUND_ADD = 1 << 5 +_BF16_EXP_MASK = 0x7F80 + +# Magnitude code k (0..7) decodes to FP4_VALUES[k]; bit 3 is the sign. +FP4_VALUES = (0.0, 0.5, 1.0, 1.5, 2.0, 3.0, 4.0, 6.0) +# Rounding boundaries between neighbouring magnitudes and whether the boundary +# value itself goes up (True) or down (False): ties to the even mantissa. +_BOUNDARIES = ((0.25, False), (0.75, True), (1.25, False), (1.75, True), (2.5, False), (3.5, True), (5.0, False)) +_SIGN_BIT = 0x8 + + +def e8m0_codes(amax: torch.Tensor) -> torch.Tensor: + """Group ``amax`` (non-negative BF16) -> E8M0 scale code (uint8), as described in the module doc.""" + if amax.dtype != torch.bfloat16: + raise TypeError(f"amax must be bfloat16, got {amax.dtype}") + if not torch.isfinite(amax).all(): + raise ValueError("weight has non-finite values; refusing to quantize") + bits = amax.view(torch.int16).to(torch.int32) & 0xFFFF + biased_exp = ((bits + _BF16_ROUND_ADD) & _BF16_EXP_MASK) >> 7 + if (biased_exp == 0xFF).any(): + # amax >= 1.75 * 2**127 rounds to infinity; compressed-tensors would store code 255 (E8M0 NaN). + raise ValueError("group amax too large for an E8M0 scale; refusing to quantize") + # biased_exp == 0 means amax rounded to zero (all-zero or tiny subnormal group): + # compressed-tensors computes log2(0) = -inf there and clamps the code to 0. + codes = torch.where(biased_exp == 0, 0, biased_exp - _ELEM_EXP_OFFSET) + return codes.clamp_(min=0).to(torch.uint8) + + +def e8m0_to_scale(codes: torch.Tensor, dtype: torch.dtype = torch.float32) -> torch.Tensor: + """E8M0 code -> power-of-two scale ``2 ** (code - 127)``.""" + if codes.dtype != torch.uint8: + raise TypeError(f"E8M0 codes must be uint8, got {codes.dtype}") + return torch.exp2(codes.to(torch.float32) - E8M0_BIAS).to(dtype) + + +def fp4_encode(x: torch.Tensor) -> torch.Tensor: + """Round already-scaled values to FP4 E2M1 codes (uint8, 0..15), as described in the module doc.""" + mag = x.abs() + k = torch.zeros(x.shape, dtype=torch.uint8, device=x.device) + for bound, tie_up in _BOUNDARIES: + k += (mag >= bound) if tie_up else (mag > bound) + return k | ((x < 0).to(torch.uint8) << 3) + + +def fp4_decode(codes: torch.Tensor, dtype: torch.dtype = torch.float32) -> torch.Tensor: + """FP4 E2M1 codes (uint8, 0..15) -> values.""" + table = torch.tensor(FP4_VALUES, dtype=torch.float32, device=codes.device) + mag = table[(codes & 0x7).long()] + return torch.where((codes & _SIGN_BIT).bool(), -mag, mag).to(dtype) + + +def pack_fp4(codes: torch.Tensor) -> torch.Tensor: + """``uint8 [out, in]`` codes -> ``uint8 [out, in / 2]``; the even column goes to the low nibble.""" + if codes.dim() != 2 or codes.shape[1] % 2: + raise ValueError(f"need a 2-D code tensor with an even column count, got {tuple(codes.shape)}") + if codes.dtype != torch.uint8 or (codes > 0xF).any(): + raise ValueError("FP4 codes must be uint8 in [0, 15]") + return codes[:, 0::2] | (codes[:, 1::2] << 4) + + +def unpack_fp4(packed: torch.Tensor) -> torch.Tensor: + """Inverse of :func:`pack_fp4`.""" + if packed.dtype != torch.uint8 or packed.dim() != 2: + raise ValueError(f"expected uint8 [out, in / 2], got {packed.dtype} {tuple(packed.shape)}") + return torch.stack((packed & 0xF, packed >> 4), dim=-1).reshape(packed.shape[0], -1) + + +def quantize_mxfp4(weight: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: + """Quantize a BF16 weight ``[out, in]``; returns ``(weight_packed [out, in/2], weight_scale [out, in/32])``. + + The division by a power of two is exact in FP32 for every BF16 input, so + the result does not depend on the compute dtype. + """ + if weight.dim() != 2: + raise ValueError(f"MXFP4 needs a 2-D weight, got shape {tuple(weight.shape)}") + if weight.dtype != torch.bfloat16: + raise TypeError(f"MXFP4 RTN reads the BF16 trainer weight, got {weight.dtype}") + rows, cols = weight.shape + if cols % GROUP_SIZE: + raise ValueError(f"input dimension {cols} is not a multiple of the group size {GROUP_SIZE}") + groups = weight.reshape(rows, cols // GROUP_SIZE, GROUP_SIZE) + codes = e8m0_codes(groups.abs().amax(dim=-1)) + scaled = groups.float() / e8m0_to_scale(codes).unsqueeze(-1) + return pack_fp4(fp4_encode(scaled).reshape(rows, cols)), codes + + +def dequantize_mxfp4(packed: torch.Tensor, scale: torch.Tensor, dtype: torch.dtype = torch.float32) -> torch.Tensor: + """Inverse of :func:`quantize_mxfp4`: FP4 value times its group scale, in ``dtype``.""" + values = fp4_decode(unpack_fp4(packed)) + rows, cols = values.shape + if scale.dtype != torch.uint8 or scale.shape != (rows, cols // GROUP_SIZE) or cols % GROUP_SIZE: + raise ValueError(f"scale {scale.dtype} {tuple(scale.shape)} does not fit packed {tuple(packed.shape)}") + groups = values.reshape(rows, -1, GROUP_SIZE) * e8m0_to_scale(scale).unsqueeze(-1) + return groups.reshape(rows, cols).to(dtype) diff --git a/yeto/rl/quant/release.py b/yeto/rl/quant/release.py index 9805b4d68..2bdfd52f8 100644 --- a/yeto/rl/quant/release.py +++ b/yeto/rl/quant/release.py @@ -147,10 +147,14 @@ def main(argv: list[str] | None = None) -> int: smoothing, awq = AwqSmoothing.from_artefact(Path(a.recipe), Path(a.awq_scales)), "fixed-artefact" else: p.error("give --recipe and --awq-scales (option R2) or --no-awq (option R1)") + from .requant import CheckpointLayout + + # The scheme tag follows the template format (W4A8 int4 or MXFP4), so readers can check it (N8.3). + scheme = CheckpointLayout.load(Path(a.template)).scheme.name store = store_for(a.store) releaser = RequantReleaser(PolicyReleaseStore(store, keep=a.keep, release_dir=a.source_dir), PolicyReleaseStore(store, keep=a.keep, release_dir=a.target_dir), - requantizer(Path(a.template), smoothing), awq=awq) + requantizer(Path(a.template), smoothing), scheme=scheme, awq=awq) if a.loop: releaser.run(threading.Event(), a.poll_seconds) else: diff --git a/yeto/rl/quant/requant.py b/yeto/rl/quant/requant.py index 0bfb02184..eaba2d783 100644 --- a/yeto/rl/quant/requant.py +++ b/yeto/rl/quant/requant.py @@ -1,13 +1,22 @@ -"""BF16 checkpoint -> W4A8 checkpoint in the layout of the original artefact (S25 N8.1). +"""BF16 checkpoint -> quantized checkpoint in the layout of a reference artefact (S25 N8.1, N8.5). -The target layout is taken from the original W4A8 checkpoint, not hard-coded: +The target layout is taken from the reference checkpoint, not hard-coded: :class:`CheckpointLayout` records every tensor's name, dtype, shape and output -shard. A tensor is quantized exactly when the layout has it as ``I8`` with a -``.weight_scale`` next to it; every other tensor is cast to its layout -dtype (the original artefact changes some dtypes, e.g. routers BF16 -> F32 and -``e_score_correction_bias`` F32 -> BF16). AWQ smoothing (:mod:`.awq`) is -applied first. The result therefore has the same file names, tensor names, -dtypes and shapes as the original, and the same metadata files. +shard. The ``format`` of its ``quantization_config`` selects the scheme: + +* ``int-quantized`` -> :data:`W4A8_INT4_CHANNEL` (N8.1, :mod:`.int4`): a module + is quantized exactly when the layout has ``.weight`` as ``I8`` with a + ``.weight_scale`` next to it. +* ``mxfp4-pack-quantized`` -> :data:`MXFP4_G32_RTN` (N8.5, :mod:`.mxfp4`): a + module is quantized exactly when the layout has ``.weight_packed`` as + ``U8`` with a ``.weight_scale``; the BF16 source provides + ``.weight``. Use it without AWQ smoothing (RTN). + +Every other tensor is cast to its layout dtype (the W4A8 artefact changes some +dtypes, e.g. routers BF16 -> F32 and ``e_score_correction_bias`` F32 -> BF16; +the MXFP4 artefact keeps them). AWQ smoothing (:mod:`.awq`) is applied first. +The result therefore has the same file names, tensor names, dtypes and shapes +as the reference, and the same metadata files. A layout can be saved as a small template directory (``layout.json`` plus the non-weight files such as ``config.json`` and ``recipe.yaml``), so later rounds @@ -18,9 +27,12 @@ works on any row slice of a 2-D weight with the same result rows, which is the basis for quantizing on the trainer ranks (D7, task N8.2; not done here). -Status: transform rules checked bit for bit against the original artefact on -layer 30 (router, shared experts, norm, one routed expert); the full-model -comparison (:func:`compare_checkpoints` over all 74169 tensors) is not run. +Status: W4A8 transform rules checked bit for bit against the original artefact +on layer 30 (router, shared experts, norm, one routed expert); MXFP4 rules +checked bit for bit against ``yetalabs/apex-flash-1-mxfp4-experts@70e86a0`` on +layer 10 routed experts and five non-expert tensors of layers 9-10 (CPU, +2026-10-19). The full-model comparisons (:func:`compare_checkpoints` over every +tensor) are not run. """ from __future__ import annotations @@ -35,6 +47,7 @@ from collections import defaultdict from collections.abc import Iterable, Mapping from dataclasses import dataclass +from functools import cached_property from pathlib import Path from typing import Any @@ -44,12 +57,14 @@ from .awq import AwqSmoothing from .int4 import QMAX, QMIN, quantize_per_channel +from .mxfp4 import GROUP_SIZE, quantize_mxfp4, unpack_fp4 INDEX_FILE = "model.safetensors.index.json" LAYOUT_FILE = "layout.json" LAYOUT_SCHEMA = "yeto.rl.quant.layout/v1" SCALE_SUFFIX = ".weight_scale" WEIGHT_SUFFIX = ".weight" +PACKED_SUFFIX = ".weight_packed" _DTYPES = { "BF16": torch.bfloat16, @@ -76,6 +91,85 @@ class TensorSpec: shard: str +class QuantScheme: + """How one compressed-tensors format stores a quantized linear ````. + + The quantized data lives in ```` (dtype ``data_dtype``) + and its scale in ``.weight_scale``; the BF16 source tensor is + ``.weight``. ``name`` is the scheme tag written into release records + (N8.3 readers compare it). + """ + + name: str + format: str + rtn_only: bool # True: the scheme refuses AWQ smoothing + data_suffix: str + data_dtype: str + weight_scheme: dict[str, Any] + + def check_weights(self, group: str, weights: Mapping[str, Any]) -> None: + got = {k: weights.get(k) for k in self.weight_scheme} + if got != self.weight_scheme: + raise LayoutError(f"config group {group}: weight scheme {got} is not supported by {self.name}") + + def check_pair(self, module: str, data: TensorSpec, scale: TensorSpec) -> None: + raise NotImplementedError + + def quantize(self, weight: torch.Tensor, scale: TensorSpec) -> tuple[torch.Tensor, torch.Tensor]: + raise NotImplementedError + + def check_values(self, name: str, data: torch.Tensor) -> None: + """Reject impossible stored values (default: none).""" + + +class _Int4Channel(QuantScheme): + name = "w4a8-int4-channel" + rtn_only = False + format = "int-quantized" + data_suffix = WEIGHT_SUFFIX + data_dtype = "I8" + weight_scheme = {"num_bits": 4, "type": "int", "symmetric": True, "strategy": "channel", "group_size": None, "dynamic": False} + + def check_weights(self, group: str, weights: Mapping[str, Any]) -> None: + super().check_weights(group, weights) + if weights.get("observer") not in ("memoryless_minmax", "minmax"): + raise LayoutError(f"config group {group}: observer {weights.get('observer')!r} is not supported by {self.name}") + + def check_pair(self, module: str, data: TensorSpec, scale: TensorSpec) -> None: + if scale.shape != (data.shape[0], 1): + raise LayoutError(f"{module}: scale shape {scale.shape} is not per-channel [{data.shape[0]}, 1]") + + def quantize(self, weight: torch.Tensor, scale: TensorSpec) -> tuple[torch.Tensor, torch.Tensor]: + return quantize_per_channel(weight, scale_dtype=_DTYPES[scale.dtype]) + + def check_values(self, name: str, data: torch.Tensor) -> None: + if data.min() < QMIN or data.max() > QMAX: + raise LayoutError(f"{name}: int4 values outside [{QMIN}, {QMAX}]") + + +class _Mxfp4Group32(QuantScheme): + name = "mxfp4-g32-rtn" + rtn_only = True + format = "mxfp4-pack-quantized" + data_suffix = PACKED_SUFFIX + data_dtype = "U8" + weight_scheme = {"num_bits": 4, "type": "float", "symmetric": True, "strategy": "group", + "group_size": GROUP_SIZE, "dynamic": False, "scale_dtype": "torch.uint8"} + + def check_pair(self, module: str, data: TensorSpec, scale: TensorSpec) -> None: + rows, cols = data.shape[0], 2 * data.shape[1] + if scale.dtype != "U8" or scale.shape != (rows, cols // GROUP_SIZE) or cols % GROUP_SIZE: + raise LayoutError(f"{module}: scale {scale.dtype}{scale.shape} is not E8M0 per {GROUP_SIZE} of [{rows}, {cols}]") + + def quantize(self, weight: torch.Tensor, scale: TensorSpec) -> tuple[torch.Tensor, torch.Tensor]: + return quantize_mxfp4(weight) + + +W4A8_INT4_CHANNEL = _Int4Channel() +MXFP4_G32_RTN = _Mxfp4Group32() +SCHEMES = {s.format: s for s in (W4A8_INT4_CHANNEL, MXFP4_G32_RTN)} + + @dataclass class CheckpointLayout: """Tensor specs of the target checkpoint, plus the quantization config it declares.""" @@ -126,36 +220,53 @@ def save_template(self, ckpt_dir: Path, template_dir: Path) -> None: _write_text_atomic(template_dir / LAYOUT_FILE, json.dumps(payload, sort_keys=True)) # -- queries -------------------------------------------------------------- - def validate(self) -> None: - """Check that the quantization config is the one :mod:`.int4` implements, and the I8/scale pairs.""" + @property + def scheme(self) -> QuantScheme: q = self.quantization_config - if q.get("quant_method") != "compressed-tensors" or q.get("format") != "int-quantized": + scheme = SCHEMES.get(q.get("format")) if q.get("quant_method") == "compressed-tensors" else None + if scheme is None: raise LayoutError(f"unsupported quantization format: {q.get('quant_method')!r} / {q.get('format')!r}") - for group, cfg in (q.get("config_groups") or {}).items(): - w = cfg.get("weights") or {} - want = {"num_bits": 4, "type": "int", "symmetric": True, "strategy": "channel", "group_size": None, "dynamic": False} - got = {k: w.get(k) for k in want} - if got != want or w.get("observer") not in ("memoryless_minmax", "minmax"): - raise LayoutError(f"config group {group}: weight scheme {got} (observer {w.get('observer')!r}) is not supported") - quantized = set(self.quantized_weights()) - for name in quantized: - scale = self.tensors[name[: -len(WEIGHT_SUFFIX)] + SCALE_SUFFIX] - rows = self.tensors[name].shape[0] - if scale.shape != (rows, 1): - raise LayoutError(f"{name}: scale shape {scale.shape} is not per-channel [{rows}, 1]") - bare = [n for n, t in self.tensors.items() if t.dtype == "I8" and n.endswith(WEIGHT_SUFFIX) and n not in quantized] + return scheme + + def validate(self) -> None: + """Check that the quantization config is one the schemes implement, and the data/scale pairs.""" + scheme = self.scheme + for group, cfg in (self.quantization_config.get("config_groups") or {}).items(): + if cfg.get("format", scheme.format) != scheme.format: + raise LayoutError(f"config group {group}: format {cfg.get('format')!r} differs from {scheme.format!r}") + scheme.check_weights(group, cfg.get("weights") or {}) + modules = self.quantized_modules + for module in modules: + scheme.check_pair(module, self.tensors[module + scheme.data_suffix], self.tensors[module + SCALE_SUFFIX]) + bare = [n for n, t in self.tensors.items() + if t.dtype == scheme.data_dtype and n.endswith(scheme.data_suffix) and n[: -len(scheme.data_suffix)] not in modules] if bare: - raise LayoutError(f"{len(bare)} I8 weights without a weight_scale, e.g. {bare[:3]}") + raise LayoutError(f"{len(bare)} {scheme.data_dtype} weights without a weight_scale, e.g. {bare[:3]}") + if clash := [m for m in modules if m + WEIGHT_SUFFIX in self.tensors and scheme.data_suffix != WEIGHT_SUFFIX]: + raise LayoutError(f"{len(clash)} quantized modules also keep a plain weight, e.g. {clash[:3]}") + + @cached_property + def quantized_modules(self) -> frozenset[str]: + """Modules ``m`` stored as ``m`` (scheme data dtype) plus ``m.weight_scale``.""" + scheme = self.scheme + return frozenset( + n[: -len(scheme.data_suffix)] for n, t in self.tensors.items() + if t.dtype == scheme.data_dtype and n.endswith(scheme.data_suffix) + and n[: -len(scheme.data_suffix)] + SCALE_SUFFIX in self.tensors + ) def quantized_weights(self) -> list[str]: - return sorted( - n for n, t in self.tensors.items() - if t.dtype == "I8" and n.endswith(WEIGHT_SUFFIX) and n[: -len(WEIGHT_SUFFIX)] + SCALE_SUFFIX in self.tensors - ) + """Names of the quantized data tensors (``.weight`` for W4A8, ``.weight_packed`` for MXFP4).""" + return sorted(m + self.scheme.data_suffix for m in self.quantized_modules) + + def source_of(self, name: str) -> str: + """The source tensor a target tensor is computed from (``.weight`` for data and scale).""" + module = _quantized_module_of(name, self, (self.scheme.data_suffix, SCALE_SUFFIX)) + return name if module is None else module + WEIGHT_SUFFIX def source_names(self) -> set[str]: - """Tensor names the BF16 source must provide: every target tensor except the scales.""" - return {n for n in self.tensors if not _is_quant_scale(n, self)} + """Tensor names the BF16 source must provide.""" + return {self.source_of(n) for n in self.tensors} def shards(self) -> dict[str, list[str]]: out: dict[str, list[str]] = defaultdict(list) @@ -186,22 +297,22 @@ def transform_tensor( ) -> dict[str, torch.Tensor]: """Map one source tensor (or a row slice of it) to its target tensors. - A quantized weight yields ``{name: int8, .weight_scale: scale}``; - any other tensor yields ``{name: tensor in the layout dtype}``. Shapes are - not checked here, so row slices pass; :func:`requantize_checkpoint` checks - whole tensors. + A quantized weight yields ``{: data, .weight_scale: + scale}``; any other tensor yields ``{name: tensor in the layout dtype}``. + Shapes are not checked here, so row slices pass; :func:`requantize_checkpoint` + checks whole tensors. """ - spec = layout.tensors.get(name) - if spec is None or _is_quant_scale(name, layout): - raise LayoutError(f"{name}: not a source tensor of the target layout") - if spec.dtype != "I8": + scheme = layout.scheme + module = _quantized_module_of(name, layout, (WEIGHT_SUFFIX,)) + if module is None: + spec = layout.tensors.get(name) + if spec is None or layout.source_of(name) != name: + raise LayoutError(f"{name}: not a source tensor of the target layout") return {name: smoothing.apply(name, tensor, _DTYPES[spec.dtype])} - module = name[: -len(WEIGHT_SUFFIX)] - scale_spec = layout.tensors[module + SCALE_SUFFIX] - # The original smoothed the expert in FP64 and stored it in the source dtype before quantizing. + # The W4A8 original smoothed the expert in FP64 and stored it in the source dtype before quantizing. smoothed = smoothing.apply(name, tensor, tensor.dtype) - q, scale = quantize_per_channel(smoothed, scale_dtype=_DTYPES[scale_spec.dtype]) - return {name: q, module + SCALE_SUFFIX: scale} + data, scale = scheme.quantize(smoothed, layout.tensors[module + SCALE_SUFFIX]) + return {module + scheme.data_suffix: data, module + SCALE_SUFFIX: scale} class _SourceReader: @@ -243,6 +354,8 @@ def requantize_checkpoint( The source must provide exactly the layout's source tensors. """ layout = CheckpointLayout.load(template_dir) + if layout.scheme.rtn_only and smoothing.enabled: + raise LayoutError(f"{layout.scheme.name} is round-to-nearest only; use no AWQ smoothing") source = _SourceReader(source_dir) expected = layout.source_names() missing, extra = expected - set(source.weight_map), set(source.weight_map) - expected @@ -255,28 +368,26 @@ def requantize_checkpoint( selected = list(all_shards) if shards is None else list(shards) if unknown := set(selected) - set(all_shards): raise LayoutError(f"unknown output shards {sorted(unknown)}") + 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} try: for shard in selected: produced: dict[str, torch.Tensor] = {} - for name in all_shards[shard]: - if name in expected: - for out_name, t in transform_tensor(name, source.get(name), layout, smoothing).items(): - produced[out_name] = t - stats["quantized"] += layout.tensors[name].dtype == "I8" - # Scales may live in another shard than their weight (seen in the original); recompute them there. + # A scale may live in another shard than its data (seen in the W4A8 original); + # its module is then quantized again in that shard. for name in all_shards[shard]: if name not in produced: - weight = name[: -len(SCALE_SUFFIX)] + WEIGHT_SUFFIX - produced[name] = transform_tensor(weight, source.get(weight), layout, smoothing)[name] + src = layout.source_of(name) + produced.update(transform_tensor(src, source.get(src), layout, smoothing)) tensors = {n: produced[n].contiguous() for n in all_shards[shard]} _check_against_layout(tensors, layout) tmp = out_dir / (shard + ".tmp") save_file(tensors, str(tmp), metadata={"format": "pt"}) os.replace(tmp, out_dir / shard) stats["shards"] += 1 + stats["quantized"] += sum(n in quantized for n in tensors) stats["tensors"] += len(tensors) finally: source.close() @@ -300,17 +411,18 @@ def finalize_checkpoint(template_dir: Path, out_dir: Path) -> None: def compare_checkpoints(a_dir: Path, b_dir: Path, names: Iterable[str] | None = None) -> dict[str, Any]: - """Compare two checkpoints of one layout tensor by tensor (N8.1 acceptance). + """Compare two checkpoints of one layout tensor by tensor (N8.1 and N8.5 acceptance). Reports the share of equal int values over all I8 tensors, the largest - relative error of the weight scales, and the high-precision tensors that - are not bit-equal. No threshold is applied here. + relative error of the float weight scales, the share of equal FP4 codes + (``U8 .weight_packed``) and of equal E8M0 codes (``U8 .weight_scale``), and + the other tensors that are not bit-equal. No threshold is applied here. """ ra, rb = _SourceReader(a_dir), _SourceReader(b_dir) todo = sorted(names) if names is not None else sorted(ra.weight_map) if missing := [n for n in todo if n not in rb.weight_map]: raise LayoutError(f"{len(missing)} tensors missing in {b_dir}, e.g. {missing[:3]}") - int_total = int_equal = 0 + int_total = int_equal = fp4_total = fp4_equal = e8m0_total = e8m0_equal = 0 scale_rel = 0.0 unequal: list[str] = [] try: @@ -321,6 +433,12 @@ def compare_checkpoints(a_dir: Path, b_dir: Path, names: Iterable[str] | None = elif a.dtype == torch.int8: int_total += a.numel() int_equal += int((a == b).sum()) + elif a.dtype == torch.uint8 and name.endswith(PACKED_SUFFIX): + fp4_total += 2 * a.numel() + fp4_equal += int((unpack_fp4(a) == unpack_fp4(b)).sum()) + elif a.dtype == torch.uint8 and name.endswith(SCALE_SUFFIX): + e8m0_total += a.numel() + e8m0_equal += int((a == b).sum()) elif name.endswith(SCALE_SUFFIX): rel = ((a.double() - b.double()).abs() / b.double().abs().clamp_min(1e-30)).max().item() scale_rel = max(scale_rel, rel) @@ -334,25 +452,31 @@ def compare_checkpoints(a_dir: Path, b_dir: Path, names: Iterable[str] | None = "int_values": int_total, "int_match_share": int_equal / int_total if int_total else None, "scale_max_rel_err": scale_rel, + "fp4_values": fp4_total, + "fp4_match_share": fp4_equal / fp4_total if fp4_total else None, + "e8m0_match_share": e8m0_equal / e8m0_total if e8m0_total else None, "unequal_tensors": len(unequal), "unequal_examples": unequal[:20], } -def _is_quant_scale(name: str, layout: CheckpointLayout) -> bool: - if not name.endswith(SCALE_SUFFIX): - return False - weight = layout.tensors.get(name[: -len(SCALE_SUFFIX)] + WEIGHT_SUFFIX) - return weight is not None and weight.dtype == "I8" +def _quantized_module_of(name: str, layout: CheckpointLayout, suffixes: Iterable[str]) -> str | None: + """```` if ``name`` is ```` for a quantized module, else ``None``.""" + modules = layout.quantized_modules + for suffix in suffixes: + if name.endswith(suffix) and name[: -len(suffix)] in modules: + return name[: -len(suffix)] + return None def _check_against_layout(tensors: Mapping[str, torch.Tensor], layout: CheckpointLayout) -> None: + scheme = layout.scheme for name, t in tensors.items(): spec = layout.tensors[name] if _DTYPE_NAMES.get(t.dtype) != spec.dtype or tuple(t.shape) != spec.shape: raise LayoutError(f"{name}: produced {t.dtype}{tuple(t.shape)}, layout wants {spec.dtype}{spec.shape}") - if t.dtype == torch.int8 and name.endswith(WEIGHT_SUFFIX) and (t.min() < QMIN or t.max() > QMAX): - raise LayoutError(f"{name}: int4 values outside [{QMIN}, {QMAX}]") + if name.endswith(scheme.data_suffix) and spec.dtype == scheme.data_dtype: + scheme.check_values(name, t) def _spec_bytes(spec: TensorSpec) -> int: @@ -379,10 +503,10 @@ def _smoothing_from_args(args: argparse.Namespace) -> AwqSmoothing: def main(argv: list[str] | None = None) -> int: p = argparse.ArgumentParser(prog="python -m yeto.rl.quant.requant", description=__doc__.splitlines()[0]) sub = p.add_subparsers(dest="cmd", required=True) - t = sub.add_parser("template", help="record the layout of an original W4A8 checkpoint") + t = sub.add_parser("template", help="record the layout of a reference quantized checkpoint (W4A8 or MXFP4)") t.add_argument("checkpoint") t.add_argument("template_dir") - c = sub.add_parser("convert", help="BF16 checkpoint -> W4A8 checkpoint in the template layout") + c = sub.add_parser("convert", help="BF16 checkpoint -> quantized checkpoint in the template layout") c.add_argument("--source", required=True) c.add_argument("--template", required=True) c.add_argument("--out", required=True) From 24656d0ed6e3229a18fb4d276a54d7e8f2e2c8d5 Mon Sep 17 00:00:00 2001 From: michael <78113554+chowhsu@users.noreply.github.com> Date: Mon, 19 Oct 2026 08:13:17 +0000 Subject: [PATCH 04/34] Put device family and HBM size in the compat group (S25 N2.1) compat_group(card, hbm_gib) gives e.g. ascend-910b4-64g, so 32 GiB and 64 GiB 910B4 islands never merge; without HBM the group is unchanged. Inventory nodes derive and record their group and refuse an unknown HBM variant at parse time. M890 is 144 GiB as read on the PPU machine. Co-Authored-By: Claude Opus 5.5 --- yeto/hw/catalog.py | 59 +++++++++++++++++++++++++++++++++-- yeto/launch/node_inventory.py | 15 +++++++-- 2 files changed, 69 insertions(+), 5 deletions(-) diff --git a/yeto/hw/catalog.py b/yeto/hw/catalog.py index 9a8e521e3..32059743f 100644 --- a/yeto/hw/catalog.py +++ b/yeto/hw/catalog.py @@ -2,6 +2,12 @@ (yeto-framework-decoupling 7.7a, design D9a-1). ``compat_group`` is ``"-"`` in lower case, e.g. ``nvidia-h100``. +With the HBM size (S25 N2.1) it is ``"--g"``, e.g. +``ascend-910b4-64g``: the 910B4 ships with 32 GiB and with 64 GiB (the user's +cluster changed from one to the other on 2026-10-16), and the two must not +merge. The vendor names the device family one to one +(:data:`DEVICE_FAMILY_BY_VENDOR`; a PPU is ``thead``), so the group always +carries the device family. Phase 1 rule: two islands may merge only when their groups are identical (H100 and H200 are different groups). The value comes from this catalog, never from a cloud name; an unknown card is an error before launch (no @@ -35,6 +41,17 @@ "810E": "thead", "M890": "thead", } +# Nominal HBM sizes (GiB) a card ships with. Read from our machines: 910B4 +# 64 GiB (inventory, 2026-10-16) and 32 GiB (ModelArts), M890 144 GiB +# (ppu-smi 147456 MiB, SESSION25/N0.1-ppu-survey.md). The others are public +# product figures, not read on our nodes. 810E: not known, so any reported +# size is used as is (rounded). +CARD_HBM_GIB: Mapping[str, tuple[int, ...]] = { + "910B": (32, 64), "910B4": (32, 64), "910B2": (64,), "910B3": (64,), "910C": (128,), + "810E": (), "M890": (144,), + "H100": (80,), "H200": (141,), "A100": (40, 80), "A100-80GB": (80,), "B200": (180,), +} + # torch device type -> device family in the backend identity. A PPU is absent # on purpose: its torch device type is "cuda", so only the card name tells it # apart from an NVIDIA GPU (see VENDOR_BY_DEVICE_FAMILY). @@ -43,6 +60,7 @@ # Launcher device family (yeto.gpu_spec.DEVICE_FAMILIES) -> vendor. VENDOR_BY_DEVICE_FAMILY: Mapping[str, str] = {"cuda": "nvidia", "npu": "ascend", "ppu": "thead"} +DEVICE_FAMILY_BY_VENDOR: Mapping[str, str] = {v: k for k, v in VENDOR_BY_DEVICE_FAMILY.items()} def device_family(device_type: str | None = None) -> str: @@ -92,10 +110,45 @@ def _canonical(card: str) -> str: f"(known: {sorted(CARD_VENDOR)}); cannot derive compat_group") -def compat_group(card: str) -> str: - """``"H100"`` -> ``"nvidia-h100"``; unknown card -> :class:`UnknownCardType`.""" +def nominal_hbm_gib(card: str, reported_gib: float) -> int: + """The nominal HBM size of ``card`` that a reported (usable) size belongs to. + + Card tools report the size minus the driver reservation, so the largest + nominal size within ``card_tools.HBM_TOLERANCE`` below the report wins. + A card without known sizes returns the report rounded; a report that fits + no known size raises :class:`UnknownCardType`.""" + from .card_tools import hbm_matches # noqa: PLC0415 - card_tools imports gpu_spec + name = _canonical(card) - return f"{CARD_VENDOR[name]}-{name.lower()}" + if not reported_gib or reported_gib <= 0: + raise UnknownCardType(f"{name}: HBM size must be positive, got {reported_gib!r}") + sizes = CARD_HBM_GIB.get(name, ()) + if not sizes: + return int(round(reported_gib)) + fitting = [size for size in sizes if hbm_matches(size, reported_gib) and reported_gib <= size * 1.02] + if not fitting: + raise UnknownCardType(f"{name} with {reported_gib:g} GiB HBM is not a known variant " + f"(known sizes: {list(sizes)} GiB); cannot derive compat_group") + return max(fitting) + + +def compat_group(card: str, hbm_gib: float | None = None) -> str: + """``"H100"`` -> ``"nvidia-h100"``; with ``hbm_gib``: ``("910B4", 61.9)`` -> + ``"ascend-910b4-64g"``. Unknown card or HBM variant -> :class:`UnknownCardType`.""" + name = _canonical(card) + group = f"{CARD_VENDOR[name]}-{name.lower()}" + if hbm_gib is None: + return group + return f"{group}-{nominal_hbm_gib(name, float(hbm_gib))}g" + + +def compat_group_family(group: str) -> str: + """The launcher device family (``cuda``/``npu``/``ppu``) of a compat group.""" + vendor = validate_compat_group(group).split("-", 1)[0] + try: + return DEVICE_FAMILY_BY_VENDOR[vendor] + except KeyError: + raise ValueError(f"compat_group {group!r}: unknown vendor {vendor!r}") from None def validate_compat_group(value: str) -> str: diff --git a/yeto/launch/node_inventory.py b/yeto/launch/node_inventory.py index 6a31652ac..104c7d9e5 100644 --- a/yeto/launch/node_inventory.py +++ b/yeto/launch/node_inventory.py @@ -33,6 +33,7 @@ from ..gpu_spec import DEVICE_FAMILIES, UnknownCard, canonical_card, device_family_of from ..hw import card_tools +from ..hw.catalog import UnknownCardType, compat_group # Fabrics an inventory may declare. "tcp" is Ethernet without RDMA; "roce" and # "ib" are RDMA over Ethernet and InfiniBand; "hccs" is the Ascend on-node @@ -75,6 +76,11 @@ class Node: def hbm_gib_per_node(self) -> float: return self.hbm_gib * self.cards_per_node + @property + def compat_group(self) -> str: + """Compatibility group with device family and HBM, e.g. ``ascend-910b4-64g``.""" + return compat_group(self.card, hbm_gib=self.hbm_gib) + def ssh_options(self) -> list[str]: """``ssh`` options that reach this node, ``ProxyJump`` included. @@ -96,7 +102,7 @@ def record(self) -> dict[str, Any]: """Manifest / event-tape row of this node (no SSH user, no mounts).""" return {"name": self.name, "device_family": self.device_family, "card": self.card, "cards_per_node": self.cards_per_node, "hbm_gib": self.hbm_gib, - "host_memory_gib": self.host_memory_gib, "fabric": self.fabric, + "compat_group": self.compat_group, "host_memory_gib": self.host_memory_gib, "fabric": self.fabric, "runtime": self.runtime, "behind_jump_host": self.proxy_jump is not None} @@ -164,7 +170,7 @@ def parse_node(raw: Any, *, index: int) -> Node: roles = raw.get("roles") or () if not isinstance(roles, (list, tuple)) or not all(isinstance(r, str) for r in roles): raise InventoryError(f"node {name}: roles must be a list of strings, got {roles!r}") - return Node( + node = Node( name=name, host=_require_text(raw["host"], name, "host"), device_family=family, @@ -184,6 +190,11 @@ def parse_node(raw: Any, *, index: int) -> Node: price_per_card_hour=(float(price) if price is not None else None), roles=tuple(roles), ) + try: + node.compat_group # noqa: B018 - an unknown HBM variant fails here, not at launch + except UnknownCardType as exc: + raise InventoryError(f"node {name}: {exc}") from None + return node def parse_inventory(document: Any) -> list[Node]: From 8b92f21470707fc4ea33f1e72fc7fd9fa2994ab8 Mon Sep 17 00:00:00 2001 From: michael <78113554+chowhsu@users.noreply.github.com> Date: Mon, 19 Oct 2026 08:13:17 +0000 Subject: [PATCH 05/34] Add npu-smi and ppu-smi card collectors for resource samples (S25 N2.2) ResourceSampler reads a CardCollector of the island's device family (YETO_RL_DEVICE_FAMILY, NVML when unset). npu-smi samples carry utilization, HBM used, power, temperature and health per chip; an unusable tool writes one available:false record with the reason. Co-Authored-By: Claude Opus 5.5 --- yeto/hw/card_telemetry.py | 240 ++++++++++++++++++++++++++++++++++++ yeto/rl/engine/telemetry.py | 74 +++++------ 2 files changed, 272 insertions(+), 42 deletions(-) create mode 100644 yeto/hw/card_telemetry.py diff --git a/yeto/hw/card_telemetry.py b/yeto/hw/card_telemetry.py new file mode 100644 index 000000000..4763091a8 --- /dev/null +++ b/yeto/hw/card_telemetry.py @@ -0,0 +1,240 @@ +"""Per-card telemetry collectors for ``rl_resource_sample`` (S25 N2.2, +rl-hetero-sampler-engines "Telemetry by device family"). + +One collector per device family, one card-sample shape: + +* ``cuda``: NVML (``pynvml``), unchanged from fleet-dashboard 2.2; +* ``npu``: ``npu-smi info`` (Ascend); +* ``ppu``: ``ppu-smi --query-gpu=... --format=csv,noheader,nounits``. + +A card sample has the keys of :data:`CARD_KEYS`; a value the tool does not +print is None. NVML samples keep their historical six keys +(``telemetry.GPU_SAMPLE_KEYS``); ``npu-smi`` and ``ppu-smi`` samples add +``temp_c`` and ``health``. + +Verification state: the ``npu-smi info`` table shape is the one already +recorded in ``yeto/gpu_spec.py`` and ``card_tools`` (910B4); the power and +temperature columns are read from the same rows and are NOT checked against a +driver 25.x machine. ``ppu-smi`` is assumed to accept the nvidia-smi query +form (T-PPU-gaps.md 2.2 says the stack is NVML compatible); NOT verified on a +PPU. Both parsers raise :class:`CardTelemetryError` on output they cannot read, +so a format change shows up as the island's single ``available: false`` record +with the reason, never as silent zeros. +""" + +from __future__ import annotations + +import re +import subprocess +from collections.abc import Callable, Sequence +from typing import Any, Protocol + +CARD_KEYS = ("index", "uuid", "util_pct", "mem_used_mb", "mem_total_mb", "power_w", "temp_c", + "health") + +Run = Callable[[Sequence[str]], str] # argv -> stdout; raises on a failed command + + +class CardTelemetryError(RuntimeError): + """The card tool is missing or its output cannot be read.""" + + +class CardCollector(Protocol): + """``source``: the tool name in the record; ``open()`` raises when the tool is + not usable; ``sample()`` returns one dict per card.""" + + source: str + + def open(self) -> None: ... + + def sample(self) -> list[dict[str, Any]]: ... + + +def run_command(argv: Sequence[str], timeout_s: float = 20.0) -> str: + try: + result = subprocess.run(list(argv), capture_output=True, text=True, timeout=timeout_s) + except (OSError, subprocess.SubprocessError) as exc: + raise CardTelemetryError(f"{argv[0]}: {type(exc).__name__}: {exc}") from None + if result.returncode != 0: + raise CardTelemetryError(f"{' '.join(argv)} exited {result.returncode}: " + f"{result.stderr.strip()[:200]}") + return result.stdout + + +def _num(text: str) -> float | None: + try: + return float(text) + except ValueError: + return None + + +# --- NVML --------------------------------------------------------------------- + +def _safe(fn: Callable[[], Any]) -> Any: + try: + return fn() + except Exception: # noqa: BLE001 - one unreadable field is None, not a failed sample + return None + + +class NvmlCollector: + """NVML through ``pynvml``; ``loader`` returns the initialised module.""" + + source = "nvml" + + def __init__(self, loader: Callable[[], Any]) -> None: + self._loader = loader + self._nvml: Any = None + + def open(self) -> None: + self._nvml = self._loader() + + def sample(self) -> list[dict[str, Any]]: + nv = self._nvml + cards = [] + for index in range(int(nv.nvmlDeviceGetCount())): + handle = nv.nvmlDeviceGetHandleByIndex(index) + mem = _safe(lambda: nv.nvmlDeviceGetMemoryInfo(handle)) + util = _safe(lambda: nv.nvmlDeviceGetUtilizationRates(handle)) + power = _safe(lambda: nv.nvmlDeviceGetPowerUsage(handle)) + uuid = _safe(lambda: nv.nvmlDeviceGetUUID(handle)) + cards.append({ + "index": index, + "uuid": uuid.decode() if isinstance(uuid, bytes) else uuid, + "util_pct": None if util is None else int(util.gpu), + "mem_used_mb": None if mem is None else int(mem.used) / 2**20, + "mem_total_mb": None if mem is None else int(mem.total) / 2**20, + "power_w": None if power is None else float(power) / 1000.0, + }) + return cards + + +# --- npu-smi ------------------------------------------------------------------ +# +# ``npu-smi info`` prints two rows per chip (card_tools has the same sample): +# +# | NPU Name | Health | Power(W) Temp(C) Hugepages-Usage(page) | +# | Chip | Bus-Id | AICore(%) Memory-Usage(MB) HBM-Usage(MB) | +# | 0 910B4 | OK | 92.5 47 0 / 0 | +# | 0 | 0000:C1:00.0 | 0 0 / 0 3161 / 65536 | +# +# Name row: NPU id, card, health, then power and temperature. Usage row: chip +# id, bus id, then AICore %, memory used / total, HBM used / total (the last +# pair). A card with two chips (910C) has two usage rows; each chip is one sample. + +_CELL_SPLIT = re.compile(r"\s*\|\s*") +_NUMBERS = re.compile(r"-?\d+(?:\.\d+)?") + + +def _cells(line: str) -> list[str]: + return [c for c in _CELL_SPLIT.split(line.strip().strip("|")) if c != ""] if line.strip().startswith("|") else [] + + +def parse_npu_smi_usage(text: str) -> list[dict[str, Any]]: + """Card samples from ``npu-smi info``; raises when no chip row is found.""" + samples: list[dict[str, Any]] = [] + card: dict[str, Any] | None = None + chips = 0 + for line in text.splitlines(): + cells = _cells(line) + if len(cells) < 3: + continue + head = cells[0].split() + if not head or not head[0].isdigit(): + continue # header rows ("NPU Name", "Chip") and process tables + if len(head) >= 2 and re.fullmatch(r"[A-Za-z]+", cells[1]): # name row + numbers = [_num(n) for n in _NUMBERS.findall(cells[2])] + card = {"index": int(head[0]), "card": head[1], "health": cells[1], + "power_w": numbers[0] if numbers else None, + "temp_c": numbers[1] if len(numbers) > 1 else None} + chips = 0 + continue + if card is None or len(head) != 1 or ":" not in cells[1]: # usage row has a bus id + continue + numbers = [_num(n) for n in _NUMBERS.findall(cells[2])] + if len(numbers) < 3: + raise CardTelemetryError(f"npu-smi usage row has no HBM column: {line.strip()[:160]!r}") + samples.append({ + "index": card["index"] if chips == 0 else f"{card['index']}.{chips}", + "uuid": cells[1], # bus id: stable per chip, npu-smi prints no uuid here + "util_pct": numbers[0], + "mem_used_mb": numbers[-2], + "mem_total_mb": numbers[-1], + "power_w": card["power_w"], + "temp_c": card["temp_c"], + "health": card["health"], + }) + chips += 1 + if not samples: + raise CardTelemetryError(f"npu-smi info printed no chip usage row: {text.strip()[:200]!r}") + return samples + + +class NpuSmiCollector: + source = "npu-smi" + ARGV = ("npu-smi", "info") + + def __init__(self, run: Run = run_command) -> None: + self._run = run + + def open(self) -> None: + parse_npu_smi_usage(self._run(self.ARGV)) + + def sample(self) -> list[dict[str, Any]]: + return parse_npu_smi_usage(self._run(self.ARGV)) + + +# --- ppu-smi (nvidia-smi query form) ------------------------------------------ + +PPU_QUERY = ("index", "uuid", "utilization.gpu", "memory.used", "memory.total", "power.draw", + "temperature.gpu") + + +def parse_smi_query_csv(text: str, tool: str = "ppu-smi") -> list[dict[str, Any]]: + """Rows of ``--query-gpu= --format=csv,noheader,nounits``.""" + samples = [] + for line in text.splitlines(): + parts = [p.strip() for p in line.split(",")] + if len(parts) != len(PPU_QUERY) or not parts[0].isdigit(): + continue + values = [None if p in ("", "[N/A]", "N/A", "[Not Supported]") else p for p in parts] + samples.append({ + "index": int(parts[0]), "uuid": values[1], + "util_pct": None if values[2] is None else _num(values[2]), + "mem_used_mb": None if values[3] is None else _num(values[3]), + "mem_total_mb": None if values[4] is None else _num(values[4]), + "power_w": None if values[5] is None else _num(values[5]), + "temp_c": None if values[6] is None else _num(values[6]), + "health": None, # the query form has no health field + }) + if not samples: + raise CardTelemetryError(f"{tool} printed no card row: {text.strip()[:200]!r}") + return samples + + +class PpuSmiCollector: + source = "ppu-smi" + ARGV = ("ppu-smi", f"--query-gpu={','.join(PPU_QUERY)}", "--format=csv,noheader,nounits") + + def __init__(self, run: Run = run_command) -> None: + self._run = run + + def open(self) -> None: + parse_smi_query_csv(self._run(self.ARGV)) + + def sample(self) -> list[dict[str, Any]]: + return parse_smi_query_csv(self._run(self.ARGV)) + + +def collector_for(family: str, *, nvml_loader: Callable[[], Any] | None = None, + run: Run = run_command) -> CardCollector: + """The collector of a launcher device family (``cuda``, ``npu``, ``ppu``).""" + if family == "cuda": + if nvml_loader is None: + raise ValueError("cuda collector needs an NVML loader") + return NvmlCollector(nvml_loader) + if family == "npu": + return NpuSmiCollector(run) + if family == "ppu": + return PpuSmiCollector(run) + raise ValueError(f"unknown device family {family!r} (known: cuda, npu, ppu)") diff --git a/yeto/rl/engine/telemetry.py b/yeto/rl/engine/telemetry.py index 8eedbb85d..db6ba30f5 100644 --- a/yeto/rl/engine/telemetry.py +++ b/yeto/rl/engine/telemetry.py @@ -2,8 +2,8 @@ ``HeartbeatThread`` writes ``rl_heartbeat`` (phase, rollout_id, policy_version, pid, uptime_s) every ``interval_s``; ``ResourceSampler`` writes -``rl_resource_sample`` (per-GPU NVML index/uuid/util/memory/power plus the -process RSS). Both run on daemon threads, only call ``emit`` (the tape writer +``rl_resource_sample`` (per-card index/uuid/util/memory/power from the +device family's card tool, plus the process RSS). Both run on daemon threads, only call ``emit`` (the tape writer is lock-protected), never block the training loop, and stop cleanly on ``stop()``. When NVML cannot be loaded the sampler writes exactly one ``available: false`` record and stops. Both are opt-in (the driver starts @@ -22,12 +22,16 @@ import time from typing import Any, Callable, Mapping +from yeto.hw.card_telemetry import CardCollector, collector_for + HEARTBEAT_EVENT = "rl_heartbeat" RESOURCE_EVENT = "rl_resource_sample" DEFAULT_HEARTBEAT_INTERVAL_S = 30.0 DEFAULT_RESOURCE_INTERVAL_S = 60.0 # Per-GPU keys of an available ``rl_resource_sample`` (None = not readable). GPU_SAMPLE_KEYS = ("index", "uuid", "util_pct", "mem_used_mb", "mem_total_mb", "power_w") +# The island's launcher device family (cuda / npu / ppu): selects the card collector. +DEVICE_FAMILY_ENV = "YETO_RL_DEVICE_FAMILY" Emit = Callable[..., None] @@ -121,88 +125,74 @@ def _cpu_rss_bytes() -> int | None: return None -def _safe(fn: Callable[[], Any]) -> Any: - try: - return fn() - except Exception: # noqa: BLE001 - one unreadable field is None, not a failed sample - return None - - class ResourceSampler(_Periodic): - """2.2: ``rl_resource_sample`` every ``interval_s`` (NVML), single fallback record.""" + """2.2: ``rl_resource_sample`` every ``interval_s``, single fallback record. + + The card source is a :class:`yeto.hw.card_telemetry.CardCollector` of the + island's device family (S25 N2.2): ``collector`` when given, else the one of + ``$YETO_RL_DEVICE_FAMILY`` (``cuda`` when unset: NVML through + ``nvml_loader``, as before). A collector that cannot open writes one + ``available: false`` record that names the reason, and no thread starts.""" name = "yeto-resource-sampler" def __init__(self, emit: Emit, *, interval_s: float = DEFAULT_RESOURCE_INTERVAL_S, nvml_loader: Callable[[], Any] = load_nvml, + collector: CardCollector | None = None, rss: Callable[[], int | None] = _cpu_rss_bytes, clock: Callable[[], float] = time.monotonic, labels: Callable[[], Mapping[str, Any]] | None = None, **kw: Any): super().__init__(interval_s, **kw) self._emit = emit - self._loader = nvml_loader + family = os.environ.get(DEVICE_FAMILY_ENV) or "cuda" + self._collector = collector or collector_for(family, nvml_loader=nvml_loader) self._rss = rss self._clock = clock self._labels = labels or (lambda: {}) - self._nvml: Any = None self.available: bool | None = None self._peak_lock = threading.Lock() self._peak_gpu: int | None = None self._peak_rss: int | None = None - def _ensure_nvml(self) -> bool: + @property + def source(self) -> str: + return self._collector.source + + def _ensure_open(self) -> bool: if self.available is None: try: - self._nvml = self._loader() + self._collector.open() self.available = True - except Exception as error: # noqa: BLE001 - ImportError, NVMLError, ... + except Exception as error: # noqa: BLE001 - ImportError, NVMLError, tool missing ... self.available = False - self._emit(RESOURCE_EVENT, available=False, t=self._clock(), + self._emit(RESOURCE_EVENT, available=False, t=self._clock(), source=self.source, reason=f"{type(error).__name__}: {error}"[:300], **self._labels()) return bool(self.available) def start(self) -> "ResourceSampler": - # Probe NVML once up front: an unavailable source is recorded at start - # (single ``available: false`` record) and no thread is started. - if self._ensure_nvml(): + # Probe the source once up front: an unavailable source is recorded at + # start (single ``available: false`` record) and no thread is started. + if self._ensure_open(): super().start() return self def sample_gpus(self) -> list[dict[str, Any]]: - nv = self._nvml - gpus = [] - self._used_bytes: list[int] = [] - for index in range(int(nv.nvmlDeviceGetCount())): - handle = nv.nvmlDeviceGetHandleByIndex(index) - mem = _safe(lambda: nv.nvmlDeviceGetMemoryInfo(handle)) - util = _safe(lambda: nv.nvmlDeviceGetUtilizationRates(handle)) - power = _safe(lambda: nv.nvmlDeviceGetPowerUsage(handle)) - uuid = _safe(lambda: nv.nvmlDeviceGetUUID(handle)) - if mem is not None: - self._used_bytes.append(int(mem.used)) - gpus.append({ - "index": index, - "uuid": uuid.decode() if isinstance(uuid, bytes) else uuid, - "util_pct": None if util is None else int(util.gpu), - "mem_used_mb": None if mem is None else int(mem.used) / 2**20, - "mem_total_mb": None if mem is None else int(mem.total) / 2**20, - "power_w": None if power is None else float(power) / 1000.0, - }) - return gpus + return self._collector.sample() def tick(self) -> bool: - if not self._ensure_nvml(): + if not self._ensure_open(): return False # one ``available: false`` record, then no more sampling gpus = self.sample_gpus() rss = self._rss() - used = self._used_bytes + used = [int(card["mem_used_mb"] * 2**20) for card in gpus + if card.get("mem_used_mb") is not None] with self._peak_lock: if used: self._peak_gpu = max(self._peak_gpu or 0, max(used)) if rss is not None: self._peak_rss = max(self._peak_rss or 0, int(rss)) self._emit(RESOURCE_EVENT, available=True, t=self._clock(), pid=os.getpid(), - gpus=gpus, cpu_rss_bytes=rss, **self._labels()) + source=self.source, gpus=gpus, cpu_rss_bytes=rss, **self._labels()) return True def take_peaks(self) -> dict[str, int]: From 79faef10454a00075509a1cf369ac95604f2a992 Mon Sep 17 00:00:00 2001 From: michael <78113554+chowhsu@users.noreply.github.com> Date: Mon, 19 Oct 2026 08:13:17 +0000 Subject: [PATCH 06/34] Record sample provenance in buffer index lines (S25 N2.4) SampleIndexEntry gets an optional provenance record that index lines carry only when set; old lines read as unknown. The payload format and GroupWriter move to the neutral yeto.rl.engine.buffer_payload (re-exported by the Miles adapter) so the engine sampler island writes the same payloads. Co-Authored-By: Claude Opus 5.5 --- yeto/rl/adapters/miles/buffer_rollout.py | 110 ++------------------- yeto/rl/engine/buffer_payload.py | 121 +++++++++++++++++++++++ yeto/rl/engine/sample_buffer.py | 10 +- yeto/rl/engine/trainer_lifecycle.py | 9 +- 4 files changed, 145 insertions(+), 105 deletions(-) create mode 100644 yeto/rl/engine/buffer_payload.py diff --git a/yeto/rl/adapters/miles/buffer_rollout.py b/yeto/rl/adapters/miles/buffer_rollout.py index 3265b523e..2207badd9 100644 --- a/yeto/rl/adapters/miles/buffer_rollout.py +++ b/yeto/rl/adapters/miles/buffer_rollout.py @@ -42,7 +42,6 @@ from __future__ import annotations -import base64 import json import os import time @@ -50,8 +49,16 @@ from dataclasses import dataclass, field from typing import Any -from yeto.rl.engine.policy_versions import engine_version_number -from yeto.rl.engine.trainer_lifecycle import SampleIndexEntry +# The payload format and the group writer are engine-neutral (S25 N2.3: the +# engine sampler island writes the same payloads); re-exported here. +from yeto.rl.engine.buffer_payload import ( # noqa: F401 + PAYLOAD_SCHEMA, + GroupWriter, + SampleBufferRolloutError, + decode_group, + encode_group, + group_entry_fields, +) SAMPLER, TRAINER, CONVERTIBLE = "sampler", "trainer", "convertible" ROLES = (SAMPLER, TRAINER, CONVERTIBLE) @@ -78,12 +85,6 @@ } # the rollout functions that read the buffer (the trainer keeps a watermark) READER_FUNCTION_PATHS = (ROLLOUT_FUNCTION_PATHS[TRAINER], ROLLOUT_FUNCTION_PATHS[CONVERTIBLE]) -PAYLOAD_SCHEMA = "yeto.rl.sample-buffer-payload/v1" -_NDARRAY = "__ndarray__" - - -class SampleBufferRolloutError(RuntimeError): - pass # ---------------------------------------------------------------------- argv @@ -164,77 +165,6 @@ def open_buffer(environ: Mapping[str, str] | None = None, *, batch_size: int = 1 create_index=create_index, merge_shards=merge) -# ---------------------------------------------------------------------- payload -def _encode_value(value: Any) -> Any: - tolist = getattr(value, "tolist", None) - if callable(tolist) and hasattr(value, "dtype") and hasattr(value, "shape"): - import numpy - - array = numpy.ascontiguousarray(value) - return {_NDARRAY: base64.b64encode(array.tobytes()).decode("ascii"), - "dtype": str(array.dtype), "shape": list(array.shape)} - if hasattr(value, "value") and type(value).__module__ != "builtins": # enum - return value.value - raise TypeError(f"sample field of type {type(value).__name__} is not serializable") - - -def _decode_hook(obj: dict[str, Any]) -> Any: - if _NDARRAY in obj: - import numpy - - data = base64.b64decode(obj[_NDARRAY]) - return numpy.frombuffer(data, dtype=obj["dtype"]).reshape(obj["shape"]).copy() - return obj - - -def encode_group(samples: Sequence[Mapping[str, Any]]) -> bytes: - """One payload: the group's samples as ``Sample.to_dict()`` mappings.""" - record = {"schema": PAYLOAD_SCHEMA, "samples": [dict(s) for s in samples]} - return json.dumps(record, default=_encode_value, sort_keys=True, allow_nan=True).encode() - - -def decode_group(data: bytes) -> list[dict[str, Any]]: - record = json.loads(data, object_hook=_decode_hook) - if record.get("schema") != PAYLOAD_SCHEMA: - raise SampleBufferRolloutError(f"payload schema {record.get('schema')!r} != {PAYLOAD_SCHEMA}") - return list(record["samples"]) - - -def _sample_versions(sample: Mapping[str, Any]) -> list[int]: - versions = [] - for call in sample.get("weight_versions") or (): - for span in call or (): - versions.append(span["version"]) - if not versions and (sample.get("metadata") or {}).get("weight_version") is not None: - versions.append(sample["metadata"]["weight_version"]) - numbers = [engine_version_number(v) for v in versions] - if not numbers or None in numbers: - raise SampleBufferRolloutError( - f"sample {sample.get('index')}: no numeric engine weight version ({versions!r}); " - "the buffer needs the behavior version of every sample") - return [int(n) for n in numbers] - - -def group_entry_fields(samples: Sequence[Mapping[str, Any]]) -> tuple[int, int]: - """``(behavior_version, tokens)`` of a group: the oldest engine version of any - token, and the response tokens. Every sample must carry per-token behavior - log-probs (decoupled PPO needs them).""" - if not samples: - raise SampleBufferRolloutError("empty group") - oldest, tokens = None, 0 - for sample in samples: - length = int(sample.get("response_length") or 0) - logprobs = sample.get("rollout_log_probs") - if logprobs is None or len(logprobs) != length: - raise SampleBufferRolloutError( - f"sample {sample.get('index')}: rollout_log_probs " - f"{'missing' if logprobs is None else len(logprobs)} != response_length {length}") - version = min(_sample_versions(sample)) - oldest = version if oldest is None else min(oldest, version) - tokens += length - return int(oldest), tokens - - def renumber_group(samples: Sequence[dict[str, Any]], *, position: int, samples_per_group: int) -> list[dict[str, Any]]: """Give a group read from the buffer an identity unique in the buffer. @@ -261,26 +191,6 @@ def renumber_group(samples: Sequence[dict[str, Any]], *, position: int, return out -# ---------------------------------------------------------------------- sampler side -class GroupWriter: - """Writes finished groups to the buffer: payload first, then the index entry.""" - - def __init__(self, buffer: Any, sampler_id: str) -> None: - if not sampler_id: - raise ValueError("sampler_id must be non-empty") - self.buffer = buffer - self.sampler_id = str(sampler_id) - - def write(self, samples: Sequence[Mapping[str, Any]], *, group_key: str) -> SampleIndexEntry: - 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) - self.buffer.append(entry) - return entry - - # ---------------------------------------------------------------------- trainer side @dataclass class TakeReport: diff --git a/yeto/rl/engine/buffer_payload.py b/yeto/rl/engine/buffer_payload.py new file mode 100644 index 000000000..dd6cb3292 --- /dev/null +++ b/yeto/rl/engine/buffer_payload.py @@ -0,0 +1,121 @@ +"""The sample-buffer payload: one group of samples as ``Sample.to_dict()``-shaped +mappings, and the writer that stores a group (payload first, then its index entry). + +Engine-neutral (moved from ``adapters.miles.buffer_rollout``, which re-exports +these names): the Miles sampler and the engine sampler island +(:mod:`.engine_sampler`) write the same payloads, so one trainer reader serves +both. A sample carries ``response_length``, per-token ``rollout_log_probs`` and +its engine weight version (``weight_versions`` spans or +``metadata.weight_version``); the entry's ``behavior_version`` is the oldest. +""" + +from __future__ import annotations + +import base64 +import json +from collections.abc import Mapping, Sequence +from typing import Any + +from .policy_versions import engine_version_number +from .trainer_lifecycle import SampleIndexEntry + +PAYLOAD_SCHEMA = "yeto.rl.sample-buffer-payload/v1" +_NDARRAY = "__ndarray__" + + +class SampleBufferRolloutError(RuntimeError): + pass + + +def _encode_value(value: Any) -> Any: + tolist = getattr(value, "tolist", None) + if callable(tolist) and hasattr(value, "dtype") and hasattr(value, "shape"): + import numpy + + array = numpy.ascontiguousarray(value) + return {_NDARRAY: base64.b64encode(array.tobytes()).decode("ascii"), + "dtype": str(array.dtype), "shape": list(array.shape)} + if hasattr(value, "value") and type(value).__module__ != "builtins": # enum + return value.value + raise TypeError(f"sample field of type {type(value).__name__} is not serializable") + + +def _decode_hook(obj: dict[str, Any]) -> Any: + if _NDARRAY in obj: + import numpy + + data = base64.b64decode(obj[_NDARRAY]) + return numpy.frombuffer(data, dtype=obj["dtype"]).reshape(obj["shape"]).copy() + return obj + + +def encode_group(samples: Sequence[Mapping[str, Any]]) -> bytes: + """One payload: the group's samples as ``Sample.to_dict()`` mappings.""" + record = {"schema": PAYLOAD_SCHEMA, "samples": [dict(s) for s in samples]} + return json.dumps(record, default=_encode_value, sort_keys=True, allow_nan=True).encode() + + +def decode_group(data: bytes) -> list[dict[str, Any]]: + record = json.loads(data, object_hook=_decode_hook) + if record.get("schema") != PAYLOAD_SCHEMA: + raise SampleBufferRolloutError(f"payload schema {record.get('schema')!r} != {PAYLOAD_SCHEMA}") + return list(record["samples"]) + + +def _sample_versions(sample: Mapping[str, Any]) -> list[int]: + versions = [] + for call in sample.get("weight_versions") or (): + for span in call or (): + versions.append(span["version"]) + if not versions and (sample.get("metadata") or {}).get("weight_version") is not None: + versions.append(sample["metadata"]["weight_version"]) + numbers = [engine_version_number(v) for v in versions] + if not numbers or None in numbers: + raise SampleBufferRolloutError( + f"sample {sample.get('index')}: no numeric engine weight version ({versions!r}); " + "the buffer needs the behavior version of every sample") + return [int(n) for n in numbers] + + +def group_entry_fields(samples: Sequence[Mapping[str, Any]]) -> tuple[int, int]: + """``(behavior_version, tokens)`` of a group: the oldest engine version of any + token, and the response tokens. Every sample must carry per-token behavior + log-probs (decoupled PPO needs them).""" + if not samples: + raise SampleBufferRolloutError("empty group") + oldest, tokens = None, 0 + for sample in samples: + length = int(sample.get("response_length") or 0) + logprobs = sample.get("rollout_log_probs") + if logprobs is None or len(logprobs) != length: + raise SampleBufferRolloutError( + f"sample {sample.get('index')}: rollout_log_probs " + f"{'missing' if logprobs is None else len(logprobs)} != response_length {length}") + version = min(_sample_versions(sample)) + oldest = version if oldest is None else min(oldest, version) + tokens += length + return int(oldest), tokens + + +class GroupWriter: + """Writes finished groups to the buffer: payload first, then the index entry.""" + + def __init__(self, buffer: Any, sampler_id: str) -> None: + if not sampler_id: + raise ValueError("sampler_id must be non-empty") + self.buffer = buffer + 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).""" + 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)) + self.buffer.append(entry) + return entry + + diff --git a/yeto/rl/engine/sample_buffer.py b/yeto/rl/engine/sample_buffer.py index 75200f7fc..e4feb4023 100644 --- a/yeto/rl/engine/sample_buffer.py +++ b/yeto/rl/engine/sample_buffer.py @@ -54,6 +54,7 @@ from __future__ import annotations +import dataclasses import fcntl import heapq import json @@ -248,7 +249,7 @@ 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"])) + appended_at=float(record["appended_at"]), provenance=record.get("provenance")) class UnmergedShards: @@ -364,12 +365,13 @@ def append(self, entry: SampleIndexEntry) -> None: if entry.payload_ref and not self.index.has_payload(entry.payload_ref): raise SampleBufferError(f"sample {entry.sample_id}: payload {entry.payload_ref} is missing") if entry.appended_at == 0.0: - entry = SampleIndexEntry(entry.sample_id, entry.behavior_version, entry.payload_ref, - entry.tokens, self.clock()) + 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)} 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), - "appended_at": float(entry.appended_at)}) + "appended_at": float(entry.appended_at), **source}) def mark_data_end(self) -> None: """The samplers will append nothing more.""" diff --git a/yeto/rl/engine/trainer_lifecycle.py b/yeto/rl/engine/trainer_lifecycle.py index 7ed96ae4f..731ecf418 100644 --- a/yeto/rl/engine/trainer_lifecycle.py +++ b/yeto/rl/engine/trainer_lifecycle.py @@ -133,13 +133,20 @@ class RunReceipt: class SampleIndexEntry: """One sample in the durable buffer: its id, the policy version that produced it (behavior version, P0-S schema 4) and where its payload is. One sample per - entry; the cross-island group index is ``sample_pool.SampleIndexEntry``.""" + entry; the cross-island group index is ``sample_pool.SampleIndexEntry``. + + ``provenance``: the source of the sample (``sampler_engine.Provenance`` + 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).""" sample_id: str behavior_version: int payload_ref: str tokens: int = 0 appended_at: float = 0.0 + provenance: Mapping[str, str] | None = field(default=None, compare=False, hash=False) @dataclass(frozen=True) From e77affc84e8479039692e12f92a4a6f3c069b71c Mon Sep 17 00:00:00 2001 From: michael <78113554+chowhsu@users.noreply.github.com> Date: Mon, 19 Oct 2026 08:13:17 +0000 Subject: [PATCH 07/34] Add the per-source offpolicy_source event (S25 N2.5) Groups the backend's recorded train and rollout log-probs by provenance key and reports ratio quantiles, TIS clip share, |dlogp| mean and p99, k3, finish reasons and truncated share, also per length bucket. yeto still computes no importance weight. Co-Authored-By: Claude Opus 5.5 --- yeto/rl/engine/offpolicy_source.py | 162 +++++++++++++++++++++++++++++ 1 file changed, 162 insertions(+) create mode 100644 yeto/rl/engine/offpolicy_source.py diff --git a/yeto/rl/engine/offpolicy_source.py b/yeto/rl/engine/offpolicy_source.py new file mode 100644 index 000000000..dc8b7fb98 --- /dev/null +++ b/yeto/rl/engine/offpolicy_source.py @@ -0,0 +1,162 @@ +"""Per-source off-policy metrics: the ``offpolicy_source`` event +(rl-hetero-sampler-engines D4, S25 N2.5; grouping per A-INF 4.1). + +yeto computes no importance weight for training. The training backend logs, per +consumed sample, its recomputed log-probs next to the rollout log-probs (the +pass-through of ``rl-offpolicy-correction-passthrough``). After each round the +trainer adapter hands those recorded values here, one :class:`SourceSample` per +consumed sample, and this module only groups them by sampler source and +summarises them. The ratio is ``exp(train_logp - rollout_logp)`` per token, as +in :mod:`.mismatch_criteria`. + +One entry per provenance key (``Provenance.key()``: device family, card, +engine@version, precision, topology); samples without provenance group under +``unknown``. Per entry: + +* ``samples``, ``tokens``; +* ``ratio_p50`` / ``ratio_p90`` / ``ratio_p99`` / ``ratio_max`` over tokens; +* ``tis_clipfrac``: share of tokens with ratio above ``tis_upper``, per-sample + mean then mean over samples (the :mod:`.mismatch_criteria` convention); +* ``abs_diff_mean`` (per-sample mean, then mean) and ``abs_diff_p99`` (tokens); +* ``k3`` (same convention as :mod:`.mismatch_criteria`); +* ``finish``: count per ``finish_reason``; ``truncated_share``: share ending + with ``length``; ``loop_abort_share``; +* ``by_length``: the same ``tis_clipfrac`` / ``abs_diff_mean`` / ``samples`` + per response-length bucket (``offpolicy_correction.length_bucket``). + +Counts add up to the round's consumed samples (``consumed`` in the event). +Pure Python: no torch. +""" + +from __future__ import annotations + +import math +from collections import Counter, defaultdict +from collections.abc import Iterable, Mapping, Sequence +from dataclasses import dataclass +from typing import Any + +from .mismatch_criteria import DEFAULT_TIS_UPPER, K3_TOKEN_CLAMP, LOG_RATIO_CLAMP +from .offpolicy_correction import length_bucket +from .sampler_engine import FINISH_LENGTH, FINISH_LOOP_ABORT, Provenance + +OFFPOLICY_SOURCE_EVENT = "offpolicy_source" +SCHEMA = "yeto-offpolicy-source-v1" + + +@dataclass(frozen=True) +class SourceSample: + """One consumed sample: its source record, the backend's recomputed + log-probs and the rollout log-probs (same length; response tokens only), + and the completion's ``finish_reason`` (None: not recorded).""" + + provenance: Mapping[str, str] | None + train_logprobs: Sequence[float] + rollout_logprobs: Sequence[float] + finish_reason: str | None = None + + def __post_init__(self) -> None: + if len(self.train_logprobs) != len(self.rollout_logprobs): + raise ValueError(f"train log-probs ({len(self.train_logprobs)}) and rollout " + f"log-probs ({len(self.rollout_logprobs)}) differ in length") + + +def _quantile(sorted_values: Sequence[float], q: float) -> float | None: + """Nearest-rank quantile of sorted values (None when empty).""" + if not sorted_values: + return None + rank = max(1, math.ceil(q * len(sorted_values))) + return float(sorted_values[min(rank, len(sorted_values)) - 1]) + + +def _diff(train: float, rollout: float) -> float: + d = float(train) - float(rollout) + if not math.isfinite(d): + return 0.0 # as mismatch_criteria: a non-finite difference counts as 0 + return d + + +def _mean(values: Sequence[float]) -> float | None: + return float(sum(values) / len(values)) if values else None + + +class _Group: + def __init__(self) -> None: + self.samples = 0 + self.ratios: list[float] = [] + self.abs_diffs: list[float] = [] + self.sample_clip: list[float] = [] + self.sample_abs: list[float] = [] + self.sample_k3: list[float] = [] + self.finish: Counter[str] = Counter() + + def add(self, sample: SourceSample, tis_upper: float) -> tuple[float | None, float | None]: + self.samples += 1 + self.finish[sample.finish_reason or "unknown"] += 1 + diffs = [_diff(t, r) for t, r in zip(sample.train_logprobs, sample.rollout_logprobs)] + if not diffs: + return None, None + clipped = [max(-LOG_RATIO_CLAMP, min(LOG_RATIO_CLAMP, d)) for d in diffs] + ratios = [math.exp(d) for d in clipped] + k3 = [max(-K3_TOKEN_CLAMP, min(K3_TOKEN_CLAMP, math.exp(d) - 1 - d)) for d in clipped] + self.ratios += ratios + self.abs_diffs += [abs(d) for d in diffs] + clip = sum(1 for r in ratios if r > tis_upper) / len(ratios) + absd = sum(abs(d) for d in diffs) / len(diffs) + self.sample_clip.append(clip) + self.sample_abs.append(absd) + self.sample_k3.append(sum(k3) / len(k3)) + return clip, absd + + def summary(self) -> dict[str, Any]: + ratios = sorted(self.ratios) + abs_diffs = sorted(self.abs_diffs) + return { + "samples": self.samples, + "tokens": len(ratios), + "ratio_p50": _quantile(ratios, 0.50), + "ratio_p90": _quantile(ratios, 0.90), + "ratio_p99": _quantile(ratios, 0.99), + "ratio_max": ratios[-1] if ratios else None, + "tis_clipfrac": _mean(self.sample_clip), + "abs_diff_mean": _mean(self.sample_abs), + "abs_diff_p99": _quantile(abs_diffs, 0.99), + "k3": _mean(self.sample_k3), + "finish": dict(sorted(self.finish.items())), + "truncated_share": self.finish[FINISH_LENGTH] / self.samples if self.samples else None, + "loop_abort_share": (self.finish[FINISH_LOOP_ABORT] / self.samples + if self.samples else None), + } + + +def offpolicy_source_event(samples: Iterable[SourceSample], *, rollout_id: int, + trainer_version: int, backend: str, + tis_upper: float = DEFAULT_TIS_UPPER) -> dict[str, Any]: + """The ``offpolicy_source`` event fields of one training round (module doc).""" + groups: dict[str, _Group] = defaultdict(_Group) + buckets: dict[str, dict[str, dict[str, list[float]]]] = defaultdict( + lambda: defaultdict(lambda: {"clip": [], "abs": [], "n": []})) + records: dict[str, dict[str, str]] = {} + consumed = 0 + for sample in samples: + consumed += 1 + provenance = Provenance.from_record(sample.provenance) + key = provenance.key() + records.setdefault(key, {f: getattr(provenance, f) for f in Provenance.KEY_FIELDS}) + clip, absd = groups[key].add(sample, float(tis_upper)) + bucket = buckets[key][length_bucket(len(sample.rollout_logprobs))] + bucket["n"].append(1.0) + if clip is not None: + bucket["clip"].append(clip) + bucket["abs"].append(absd) + sources = [] + for key in sorted(groups): + entry = {"key": key, **records[key], **groups[key].summary()} + entry["by_length"] = { + name: {"samples": len(b["n"]), "tis_clipfrac": _mean(b["clip"]), + "abs_diff_mean": _mean(b["abs"])} + for name, b in sorted(buckets[key].items())} + sources.append(entry) + return {"schema": SCHEMA, "rollout_id": int(rollout_id), "trainer_version": int(trainer_version), + "backend": str(backend), "tis_upper": float(tis_upper), "consumed": consumed, + "sources": sources} From fef682e37cb714ebfba6f917d8a95349d747c091 Mon Sep 17 00:00:00 2001 From: michael <78113554+chowhsu@users.noreply.github.com> Date: Mon, 19 Oct 2026 08:13:17 +0000 Subject: [PATCH 08/34] Add the engine sampler island, post-load checks and SGLang/vLLM adapters (S25 N2.3) PortReleaseLoader loads releases through the SamplerEngine port and serves them only after the expert-count, fixed-prompt and log-prob probe checks and a served-version read-back pass; failures write sampler_load_check_failed, stop generation and restart on retry. The island refuses repetition penalties and a max_tokens the engine cannot serve, and can cut looping completions as loop_abort. The SGLang (PPU) and vLLM (910B) adapters carry the A-INF vendor guards: weight_packed and dense-format refusal, MX_PATCH for MXFP4, dp8+NEXTN refusal, DeepGEMM warm-up, exit 134 after stop, start timeout, host-memory floor, eager refusal, restart-only reload on vllm-ascend. Co-Authored-By: Claude Opus 5.5 --- tests/test_rl_hetero_sampler_engines.py | 563 ++++++++++++++++++ yeto/rl/adapters/sampler_engines/server.py | 295 +++++++++ .../adapters/sampler_engines/sglang_server.py | 192 ++++++ .../adapters/sampler_engines/vllm_server.py | 146 +++++ yeto/rl/engine/engine_sampler.py | 198 ++++++ yeto/rl/engine/loop_guard.py | 54 ++ yeto/rl/engine/policy_release.py | 12 + yeto/rl/engine/sampler_checks.py | 334 +++++++++++ yeto/rl/engine/sampler_engine.py | 5 +- yeto/rl/engine/sampler_loop.py | 27 +- 10 files changed, 1821 insertions(+), 5 deletions(-) create mode 100644 tests/test_rl_hetero_sampler_engines.py create mode 100644 yeto/rl/adapters/sampler_engines/server.py create mode 100644 yeto/rl/adapters/sampler_engines/sglang_server.py create mode 100644 yeto/rl/adapters/sampler_engines/vllm_server.py create mode 100644 yeto/rl/engine/engine_sampler.py create mode 100644 yeto/rl/engine/loop_guard.py create mode 100644 yeto/rl/engine/sampler_checks.py diff --git a/tests/test_rl_hetero_sampler_engines.py b/tests/test_rl_hetero_sampler_engines.py new file mode 100644 index 000000000..0121c4a80 --- /dev/null +++ b/tests/test_rl_hetero_sampler_engines.py @@ -0,0 +1,563 @@ +"""CPU tests of the engine-neutral sampler (S25 N2.1-N2.5, rl-hetero-sampler-engines). + +No engine, card or network: the island runs on :class:`FakeEngine`; the SGLang +and vLLM adapters run on a fake process and a fake HTTP layer; the card tools +on recorded output. +""" + +from __future__ import annotations + +import json +import math +import signal +from pathlib import Path + +import pytest + +from yeto.hw import card_telemetry as ct +from yeto.hw.catalog import UnknownCardType, compat_group, compat_group_family +from yeto.launch.node_inventory import InventoryError, parse_inventory +from yeto.rl.adapters.sampler_engines import engine_class +from yeto.rl.adapters.sampler_engines.fake import FakeEngine +from yeto.rl.adapters.sampler_engines.server import ServerConfig +from yeto.rl.adapters.sampler_engines.sglang_server import SGLangOptions, SGLangServerEngine +from yeto.rl.adapters.sampler_engines.vllm_server import VllmOptions, VllmServerEngine +from yeto.rl.engine import telemetry +from yeto.rl.engine.buffer_payload import GroupWriter, decode_group +from yeto.rl.engine.engine_sampler import EngineSamplerConfig, PromptItem, build_engine_sampler +from yeto.rl.engine.loop_guard import LoopGuard +from yeto.rl.engine.offpolicy_source import SourceSample, offpolicy_source_event +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 ( + GSM8K_FIXED_PROMPTS, + LOAD_CHECK_FAILED_EVENT, + PortReleaseLoader, + PostLoadChecks, + SamplerLoadFatal, + count_release_experts, +) +from yeto.rl.engine.sampler_engine import ( + FINISH_LENGTH, + FINISH_LOOP_ABORT, + FINISH_STOP, + Completion, + EngineIdentity, + Provenance, + SamplerEngine, + SamplerEngineError, + SamplingConfigError, + SamplingParams, +) +from yeto.rl.engine.trainer_lifecycle import SampleIndexEntry + +ANSWERS = {p.prompt: f"The answer is {p.answer}." for p in GSM8K_FIXED_PROMPTS} +W4A8 = EngineIdentity(engine="fake", engine_version="0", device_family="npu", card_model="910B4", + precision="w4a8", quant_scheme="w4a8-int4-channel", topology="tp4") + + +# ---------------------------------------------------------------------- N2.1 catalog +def test_compat_group_carries_family_and_hbm(): + assert compat_group("H100") == "nvidia-h100" # without HBM: unchanged + assert compat_group("910B4", 61.9) == "ascend-910b4-64g" + assert compat_group("910b4", 32) == "ascend-910b4-32g" + assert compat_group("910B2", 64) == "ascend-910b2-64g" + assert compat_group("910C", 128) == "ascend-910c-128g" + assert compat_group("M890", 144.0) == "thead-m890-144g" + assert compat_group("810E", 63.9) == "thead-810e-64g" # sizes unknown: rounded + assert compat_group_family("ascend-910b4-64g") == "npu" + assert compat_group_family("thead-m890-144g") == "ppu" + with pytest.raises(UnknownCardType, match="not a known variant"): + compat_group("910B3", 32) + + +def _inventory(**node): + base = {"name": "npu1", "host": "u@h", "device_family": "npu", "card": "910B4", + "cards_per_node": 8, "hbm_gib": 64, "host_memory_gib": 1024, "image": "img"} + return parse_inventory({"nodes": [{**base, **node}]}) + + +def test_inventory_node_records_its_compat_group_and_refuses_unknown_hbm(): + (node,) = _inventory() + assert node.compat_group == "ascend-910b4-64g" + assert node.record()["compat_group"] == "ascend-910b4-64g" + with pytest.raises(InventoryError, match="not a known variant"): + _inventory(hbm_gib=48) + + +# ---------------------------------------------------------------------- N2.2 telemetry +NPU_SMI = """+---------------------------+---------------+----------------------------------+ +| NPU Name | Health | Power(W) Temp(C) Hugepages | +| Chip | Bus-Id | AICore(%) Memory-Usage(MB) | ++===========================+===============+==================================+ +| 0 910B4 | OK | 92.5 47 0 / 0 | +| 0 | 0000:C1:00.0 | 12 3161 / 65536 | ++===========================+===============+==================================+ +| 1 910B4 | Warning | 91.3 46 0 / 0 | +| 0 | 0000:C2:00.0 | 0 0 / 0 3160 / 65536 | ++===========================+===============+==================================+ +""" + + +def test_npu_smi_usage_has_util_hbm_power_temp_and_health_per_card(): + samples = ct.parse_npu_smi_usage(NPU_SMI) + assert samples[0] == {"index": 0, "uuid": "0000:C1:00.0", "util_pct": 12.0, + "mem_used_mb": 3161.0, "mem_total_mb": 65536.0, "power_w": 92.5, + "temp_c": 47.0, "health": "OK"} + assert (samples[1]["health"], samples[1]["mem_used_mb"]) == ("Warning", 3160.0) + assert set(samples[0]) == set(ct.CARD_KEYS) + with pytest.raises(ct.CardTelemetryError): + ct.parse_npu_smi_usage("npu-smi: command not found") + + +def test_ppu_smi_query_rows(): + text = "0, PPU-1, 50, 1024, 98304, 300.5, 61\n1, PPU-2, [N/A], 0, 98304, [N/A], 60\n" + samples = ct.parse_smi_query_csv(text) + assert samples[0]["util_pct"] == 50.0 and samples[0]["temp_c"] == 61.0 + assert samples[1]["util_pct"] is None and samples[1]["power_w"] is None + assert "--query-gpu=" + ",".join(ct.PPU_QUERY) in ct.PpuSmiCollector.ARGV + + +def test_resource_sampler_reads_the_family_collector(monkeypatch): + events = [] + collector = ct.NpuSmiCollector(run=lambda argv: NPU_SMI) + sampler = telemetry.ResourceSampler(lambda n, **f: events.append((n, f)), collector=collector, + rss=lambda: 1) + assert sampler.tick() + name, fields = events[0] + assert name == telemetry.RESOURCE_EVENT and fields["source"] == "npu-smi" + assert fields["gpus"][1]["health"] == "Warning" + assert sampler.take_peaks()["peak_gpu_mem_bytes"] == 3161 * 2**20 + # the device family from the environment selects the collector + monkeypatch.setenv(telemetry.DEVICE_FAMILY_ENV, "ppu") + assert telemetry.ResourceSampler(lambda *a, **k: None).source == "ppu-smi" + + +def test_missing_card_tool_writes_one_record_with_the_reason(): + def missing(argv): + raise ct.CardTelemetryError("npu-smi: No such file or directory") + + events = [] + sampler = telemetry.ResourceSampler(lambda n, **f: events.append((n, f)), interval_s=0.001, + collector=ct.NpuSmiCollector(run=missing)).start() + sampler.stop() + assert len(events) == 1 + assert events[0][1]["available"] is False and "No such file" in events[0][1]["reason"] + assert events[0][1]["source"] == "npu-smi" + + +# ---------------------------------------------------------------------- N2.3 port types +def test_sampling_refuses_repetition_penalties_unless_allowed(): + with pytest.raises(SamplingConfigError, match="repetition_penalty"): + SamplingParams(max_tokens=8, repetition_penalty=1.1).check() + with pytest.raises(SamplingConfigError, match="presence_penalty"): + SamplingParams(max_tokens=8, presence_penalty=0.5).check() + allowed = SamplingParams(max_tokens=8, frequency_penalty=0.2, allow_penalties=True).check() + assert allowed.digest() != SamplingParams(max_tokens=8).digest() + assert SamplingParams(max_tokens=8).digest() == SamplingParams(max_tokens=8).digest() + + +def test_fake_engine_implements_the_port(): + assert isinstance(FakeEngine(), SamplerEngine) + assert engine_class("fake") is FakeEngine + with pytest.raises(ValueError, match="unknown sampler engine"): + engine_class("trt") + + +def test_loop_guard_cuts_a_repeating_completion_and_keeps_its_logprobs(): + guard = LoopGuard(max_period=3, repeats=4) + ids = (5, 6, 1, 2, 3, 1, 2, 3, 1, 2, 3, 1, 2, 3, 1, 2, 3) + completion = Completion(token_ids=ids, logprobs=tuple(-0.1 * i for i in range(len(ids))), + text="loop", finish_reason=FINISH_LENGTH) + cut = guard.apply(completion) + assert cut.finish_reason == FINISH_LOOP_ABORT + assert cut.token_ids == ids[:14] and cut.logprobs == completion.logprobs[:14] + plain = Completion(token_ids=(1, 2, 3, 4), logprobs=(0.0,) * 4, text="x", finish_reason=FINISH_STOP) + assert guard.apply(plain) is plain + + +# ---------------------------------------------------------------------- N2.4 provenance +def test_index_lines_carry_provenance_and_old_lines_read_as_unknown(tmp_path): + buffer = FileSampleBuffer(CheckpointStore(tmp_path), batch_size=1) + old = SampleIndexEntry("old", 0, buffer.write_payload("old", b"{}"), tokens=1) + buffer.append(old) + record = Provenance.of(W4A8, sampler_id="npu1", sampling=SamplingParams(max_tokens=8)).to_record() + buffer.append(SampleIndexEntry("new", 1, buffer.write_payload("new", b"{}"), tokens=1, + provenance=record)) + lines = [json.loads(line) for line in + (tmp_path / "sample-buffer" / "index.jsonl").read_text().splitlines()] + assert "provenance" not in lines[0] and lines[1]["provenance"]["card_model"] == "910B4" + reader = FileSampleBuffer(CheckpointStore(tmp_path), batch_size=1) + reader.refresh() + first, second = (s.entry for s in reader.select("t", 2)) + assert Provenance.from_record(first.provenance) == Provenance.unknown() + assert Provenance.from_record(second.provenance).quant_scheme == "w4a8-int4-channel" + assert Provenance.from_record(second.provenance).key() == "npu/910B4/fake@0/w4a8/tp4" + + +# ---------------------------------------------------------------------- N2.5 per-source metrics +def _npu_ppu(record_npu, record_ppu): + small = SourceSample(record_npu, [-1.0, -2.0], [-1.0, -2.01], FINISH_STOP) + shifted = SourceSample(record_ppu, [-1.0, -1.0, -1.0], [-2.0, -1.0, -1.0], FINISH_LENGTH) + return [small, small, shifted, SourceSample(None, [-1.0], [-1.0])] + + +def test_offpolicy_source_groups_recorded_backend_values_by_source(): + npu = Provenance.of(W4A8, sampler_id="a", sampling=SamplingParams(max_tokens=8)).to_record() + ppu_identity = EngineIdentity(engine="sglang", engine_version="0.5.17", device_family="ppu", + card_model="810E", precision="mxfp4", quant_scheme="MXFP4_EXPERTS_RTN", + topology="tp8") + ppu = Provenance.of(ppu_identity, sampler_id="b", sampling=SamplingParams(max_tokens=8)).to_record() + event = offpolicy_source_event(_npu_ppu(npu, ppu), rollout_id=3, trainer_version=7, backend="verl") + assert event["consumed"] == 4 + sources = {s["key"]: s for s in event["sources"]} + assert sum(s["samples"] for s in sources.values()) == event["consumed"] + assert set(sources) == {"npu/910B4/fake@0/w4a8/tp4", "ppu/810E/sglang@0.5.17/mxfp4/tp8", + "unknown/unknown/unknown@unknown/unknown/unknown"} + p = sources["ppu/810E/sglang@0.5.17/mxfp4/tp8"] + assert p["tokens"] == 3 and p["truncated_share"] == 1.0 + assert math.isclose(p["ratio_max"], math.e) and math.isclose(p["tis_clipfrac"], 1 / 3) + assert math.isclose(p["abs_diff_mean"], 1 / 3) + n = sources["npu/910B4/fake@0/w4a8/tp4"] + assert n["tis_clipfrac"] == 0.0 and math.isclose(n["abs_diff_p99"], 0.01, abs_tol=1e-9) + assert n["finish"] == {"stop": 2} and n["by_length"]["0-255"]["samples"] == 2 + + +# ---------------------------------------------------------------------- post-load checks +def _w4a8_release(releases, version, *, experts=4, scheme="w4a8-int4-channel", recorded=None): + def write(directory: Path) -> None: + names = {f"model.layers.0.mlp.experts.{e}.down_proj.weight": "a.safetensors" + for e in range(experts)} + names["model.embed_tokens.weight"] = "a.safetensors" + (directory / "model.safetensors.index.json").write_text(json.dumps({"weight_map": names})) + (directory / "a.safetensors").write_bytes(b"x") + + derived = {"scheme": scheme, "version": version} + if recorded is not None: + derived["expert_matrices"] = recorded + return releases.release(version, f"h{version}", write, derived_from=derived) + + +def _loader(engine, events, **checks): + return PortReleaseLoader(engine, PostLoadChecks(**{"expected_expert_matrices": 4, **checks}), + lambda n, **f: events.append((n, f)), max_consecutive_failures=2) + + +def test_release_expert_count_reads_the_index(tmp_path): + releases = PolicyReleaseStore(CheckpointStore(tmp_path)) + release = _w4a8_release(releases, 1, experts=3) + assert count_release_experts(releases.resolve(release)) == 3 + assert count_release_experts(tmp_path) is None + + +def test_loader_serves_a_release_only_after_all_checks_pass(tmp_path): + releases = PolicyReleaseStore(CheckpointStore(tmp_path)) + engine = FakeEngine(identity=W4A8, answers=ANSWERS, expert_matrices=4) + events = [] + loader = _loader(engine, events, max_probe_jump=0.5) + loader.start(None, 0) + release = _w4a8_release(releases, 1) + loader(release, releases.resolve(release)) + assert loader.ready and engine.served_version() == 1 + checked = [f for n, f in events if n == "rl_sampler_load_checked"] + assert [c["reload_mode"] for c in checked] == ["restart", "in_place"] + assert {c["name"] for c in checked[1]["checks"]} == {"expert_count", "fixed_prompts", + "logprob_probe", "served_version"} + assert checked[1]["checks"][2]["against"] == "previous" + + +@pytest.mark.parametrize("fault, step", [ + ("silent_skip", "expert_count"), # the engine loaded fewer experts (A-INF item 1) + ("wrong_answers", "fixed_prompts"), + ("probe_jump", "logprob_probe"), + ("release_short", "release_expert_count"), +]) +def test_a_failed_check_marks_the_engine_not_ready_and_restarts_next_time(tmp_path, fault, step): + from yeto.rl.engine.policy_release import ReleaseNotLoaded + + releases = PolicyReleaseStore(CheckpointStore(tmp_path)) + engine = FakeEngine(identity=W4A8, answers=ANSWERS, expert_matrices=4) + events = [] + loader = _loader(engine, events, max_probe_jump=0.5) + loader.start(None, 0) + release = _w4a8_release(releases, 1, experts=3 if fault == "release_short" else 4) + if fault == "silent_skip": + engine.expert_matrices = 0 + elif fault == "wrong_answers": + engine.answers = {} + elif fault == "probe_jump": + engine.probe_shift = 2.0 + with pytest.raises(ReleaseNotLoaded, match=step): + loader(release, releases.resolve(release)) + assert not loader.ready + (failed,) = [f for n, f in events if n == LOAD_CHECK_FAILED_EVENT] + assert failed["step"] == step and failed["failures"] == 1 + with pytest.raises(SamplerLoadFatal): # second failure in a row: the island stops + loader(release, releases.resolve(release)) + + +def test_engine_that_reports_no_count_needs_memory_or_an_explicit_waiver(tmp_path): + engine = FakeEngine(identity=W4A8, answers=ANSWERS, expert_matrices=None) + checks = PostLoadChecks(expected_expert_matrices=4) + assert checks.expert_count(engine.__class__(identity=W4A8, expert_matrices=None), 4)["passed"] is False + engine.weight_gib = 20.8 + assert PostLoadChecks(expected_weight_gib_per_card=20.0).expert_count(engine, 4)["passed"] + assert not PostLoadChecks(expected_weight_gib_per_card=15.0).expert_count(engine, 4)["passed"] + waived = PostLoadChecks(allow_unverified_count=True).expert_count( + FakeEngine(expert_matrices=None), 4) + assert waived["passed"] and waived["unverified"] + + +def test_scheme_mismatch_keeps_the_served_version(tmp_path): + from yeto.rl.engine.policy_release import ReleaseNotLoaded + + releases = PolicyReleaseStore(CheckpointStore(tmp_path)) + engine = FakeEngine(identity=W4A8, answers=ANSWERS, expert_matrices=4) + loader = _loader(engine, []) + loader.start(None, 0) + bf16 = releases.release(1, "h1", lambda d: (d / "w").write_bytes(b"x")) + with pytest.raises(ReleaseNotLoaded) as info: + loader(bf16, releases.resolve(bf16)) + assert info.value.reason == "scheme_mismatch" + assert engine.served_version() == 0 and loader.ready and loader.failures == 0 + + +# ---------------------------------------------------------------------- island end to end +def test_engine_sampler_island_writes_groups_with_version_and_provenance(tmp_path): + store = CheckpointStore(tmp_path) + releases = PolicyReleaseStore(store) + buffer = FileSampleBuffer(store, batch_size=1) + engine = FakeEngine(identity=W4A8, answers=ANSWERS, expert_matrices=4, max_model_len=64) + events = [] + config = EngineSamplerConfig(sampler_id="npu1", num_rollouts=2, max_backlog_groups=100, + sampling=SamplingParams(max_tokens=6, n=2), max_prompt_tokens=32) + + def prompts(rollout_id): + if rollout_id == 0: # the trainer releases v1 during round 0 + _w4a8_release(releases, 1) + return [PromptItem(f"p{rollout_id}", f"prompt {rollout_id}", label="x")] + + island = build_engine_sampler(config, engine=engine, releases=releases, buffer=buffer, + emit=lambda n, **f: events.append((n, f)), prompts=prompts, + reward=lambda item, c: 1.0, + checks=PostLoadChecks(expected_expert_matrices=4), + sleep=lambda s: None) + assert island.run() == 2 + reader = FileSampleBuffer(store, batch_size=1) + reader.refresh() + entries = [s.entry for s in reader.select("t", 10)] + assert [e.behavior_version for e in entries] == [0, 1] + assert all(Provenance.from_record(e.provenance).card_model == "910B4" for e in entries) + samples = decode_group(reader.read_payload(entries[0].payload_ref)) + assert len(samples) == 2 and samples[0]["metadata"]["finish_reason"] in (FINISH_STOP, FINISH_LENGTH) + assert len(samples[0]["rollout_log_probs"]) == samples[0]["response_length"] + names = [n for n, _ in events] + assert names[0] == "rl_sampler_start" and "rl_sampler_engine_round" in names + start = events[0][1] + assert start["sampling"]["repetition_penalty"] == 1.0 and start["provenance"]["precision"] == "w4a8" + + +def test_island_refuses_max_tokens_the_engine_cannot_serve(tmp_path): + store = CheckpointStore(tmp_path) + config = EngineSamplerConfig(sampler_id="s", num_rollouts=1, max_backlog_groups=1, + sampling=SamplingParams(max_tokens=16384), max_prompt_tokens=1024) + with pytest.raises(SamplerEngineError, match="max_model_len 8192"): + build_engine_sampler(config, engine=FakeEngine(max_model_len=8192), + releases=PolicyReleaseStore(store), buffer=FileSampleBuffer(store, batch_size=1), + emit=lambda *a, **k: None, prompts=lambda r: [], reward=lambda i, c: 0.0) + + +def test_group_writer_puts_provenance_in_the_entry(tmp_path): + buffer = FileSampleBuffer(CheckpointStore(tmp_path), batch_size=1) + sample = {"response_length": 1, "rollout_log_probs": [-0.5], "metadata": {"weight_version": "3"}} + entry = GroupWriter(buffer, "s").write([sample], group_key="g", provenance={"engine": "vllm"}) + assert entry.behavior_version == 3 and entry.provenance == {"engine": "vllm"} + + +# ---------------------------------------------------------------------- server adapters +class _Process: + def __init__(self, exit_on_term=0): + self.code = None + self.exit_on_term = exit_on_term + self.signals = [] + + def poll(self): + return self.code + + def send_signal(self, sig): + self.signals.append(sig) + self.code = self.exit_on_term + + def kill(self): + self.code = -9 + + def wait(self, timeout=None): + return self.code + + +class _Http: + def __init__(self, routes): + self.routes = routes + self.calls = [] + + def __call__(self, method, url, body, timeout): + path = url.split(":30000", 1)[1] if ":30000" in url else url + self.calls.append((method, path, body)) + handler = self.routes.get((method, path)) + return (404, None) if handler is None else (200, handler(body) if callable(handler) else handler) + + +def _config(tmp_path, **kw): + base = tmp_path / "base" + base.mkdir(exist_ok=True) + values = dict(base_dir=base, card_model="810E", device_family="ppu", engine_version="0.5.17", + precision="mxfp4", quant_scheme="MXFP4_EXPERTS_RTN", topology="tp8", max_model_len=32768) + values.update(kw) + return ServerConfig(**values) + + +def _sglang_generate(body): + texts = body["text"] if isinstance(body["text"], list) else [body["text"]] + if body["sampling_params"]["max_new_tokens"] == 0: # teacher-forced score + return [{"meta_info": {"input_token_logprobs": [[None, 1, None]] + [[-0.5, 2, None]] * (len(t) - 1)}} + for t in texts] + return [{"text": f"out {t}", "meta_info": {"output_token_logprobs": [[-0.1, 7, None], [-0.2, 8, None]], + "finish_reason": {"type": "length"}}} for t in texts] + + +def test_sglang_adapter_argv_env_generate_and_score(tmp_path): + spawned = [] + http = _Http({("GET", "/health"): {}, ("POST", "/generate"): _sglang_generate}) + engine = SGLangServerEngine(_config(tmp_path), spawn=lambda a, e: spawned.append((a, e)) or _Process(134), + http=http, sleep=lambda s: None) + report = engine.start(None, 0) + argv, env = spawned[0] + assert argv[:3] == ["python", "-m", "sglang.launch_server"] and "--tp" in argv + assert env["MX_PATCH"] == "1" and report.detail["env"] == {"MX_PATCH": "1"} + groups = engine.generate(["a", "b"], SamplingParams(max_tokens=2, n=2)) + assert [len(g) for g in groups] == [2, 2] + assert groups[0][0].token_ids == (7, 8) and groups[0][0].finish_reason == FINISH_LENGTH + sent = [b for m, p, b in http.calls if p == "/generate"][0] + assert sent["text"] == ["a", "a", "b", "b"] and sent["sampling_params"]["repetition_penalty"] == 1.0 + assert engine.score("abc", "de") == [-0.5, -0.5] + engine.stop() # TP8 exits 134 (pcclCommAbort) after a stop request: normal + assert engine.last_exit["code"] == 134 and engine.last_exit["normal"] + + +def test_sglang_exit_134_while_serving_is_a_failure(tmp_path): + process = _Process() + engine = SGLangServerEngine(_config(tmp_path), spawn=lambda a, e: process, + http=_Http({("GET", "/health"): {}}), sleep=lambda s: None) + engine.start(None, 0) + process.code = 134 + with pytest.raises(SamplerEngineError, match="exited with code 134 while serving"): + engine.generate(["a"], SamplingParams(max_tokens=2)) + assert engine.served_version() is None + + +def test_sglang_refuses_dp8_with_nextn_packed_names_and_dense_format(tmp_path): + with pytest.raises(SamplerEngineError, match="NEXTN"): + SGLangServerEngine(_config(tmp_path, topology="dp8"), SGLangOptions(nextn=True)) + release = tmp_path / "r1" + release.mkdir() + (release / "model.safetensors.index.json").write_text(json.dumps( + {"weight_map": {"model.layers.0.mlp.experts.0.down_proj.weight_packed": "a"}})) + engine = SGLangServerEngine(_config(tmp_path), spawn=lambda a, e: _Process(), + http=_Http({("GET", "/health"): {}}), sleep=lambda s: None) + with pytest.raises(SamplerEngineError, match="weight_packed"): + engine.start(release, 1) + patched = SGLangServerEngine(_config(tmp_path), SGLangOptions(renames_weight_packed=True), + spawn=lambda a, e: _Process(), http=_Http({("GET", "/health"): {}}), + sleep=lambda s: None) + assert patched.start(release, 1).reload_mode == "restart" + (release / "config.json").write_text(json.dumps({"quantization_config": {"quant_method": "compressed-tensors", + "format": "dense"}})) + with pytest.raises(SamplerEngineError, match="dense"): + patched.start(release, 1) + + +def test_sglang_dp8_warms_up_batch_sizes_before_the_checks(tmp_path): + http = _Http({("GET", "/health"): {}, ("POST", "/generate"): _sglang_generate}) + engine = SGLangServerEngine(_config(tmp_path, topology="dp8"), SGLangOptions(warmup_batch_sizes=(1, 4)), + spawn=lambda a, e: _Process(), http=http, sleep=lambda s: None) + report = engine.start(None, 0) + assert report.detail["warmup_batch_sizes"] == [1, 4] + assert [len(b["text"]) for m, p, b in http.calls if p == "/generate"] == [1, 4] + + +def test_start_timeout_kills_the_server_and_names_the_hang(tmp_path): + clock = iter(range(0, 10_000, 100)) + process = _Process() + engine = VllmServerEngine(_vllm(tmp_path, start_timeout_s=300), + spawn=lambda a, e: process, http=_Http({}), clock=lambda: next(clock), + sleep=lambda s: None) + with pytest.raises(SamplerEngineError, match="HCCL hang"): + engine.start(None, 0) + assert process.code == -9 + + +def _vllm(tmp_path, **kw): + cfg = _config(tmp_path, device_family="npu", card_model="910B4", precision="w4a8", + quant_scheme="w4a8-int4-channel", topology="tp4", max_model_len=8192, **kw) + (cfg.base_dir / "config.json").write_text(json.dumps( + {"quantization_config": {"quant_method": "compressed-tensors", "format": "int-quantized"}})) + return cfg + + +def test_vllm_adapter_restarts_on_each_release_and_parses_completions(tmp_path): + spawned = [] + + def completions(body): + if body.get("echo"): + n = len(body["prompt"]) + return {"choices": [{"index": 0, "logprobs": {"token_logprobs": [None] + [-0.25] * (n - 1) + [-1.0]}}]} + return {"choices": [{"index": i, "text": "t", "finish_reason": "stop", + "logprobs": {"tokens": ["token_id:4", "token_id:5"], + "token_logprobs": [-0.3, -0.4]}} + for i in range(len(body["prompt"]) * body["n"])]} + + http = _Http({("GET", "/health"): {}, ("POST", "/v1/completions"): completions}) + engine = VllmServerEngine(_vllm(tmp_path), VllmOptions(max_num_seqs=32), + spawn=lambda a, e: spawned.append((a, e)) or _Process(), http=http, + sleep=lambda s: None) + engine.start(None, 0) + release = tmp_path / "r1" + release.mkdir() + (release / "config.json").write_text(json.dumps({"quantization_config": {"format": "int-quantized"}})) + assert engine.load(release, 1).reload_mode == "restart" and engine.served_version() == 1 + assert len(spawned) == 2 and spawned[0][1]["VLLM_ASCEND_W4A8_CPU_STAGE"] == "1" + assert "--max-num-seqs" in spawned[0][0] + (group,) = engine.generate(["p"], SamplingParams(max_tokens=4)) + assert group[0].token_ids == (4, 5) and group[0].logprobs == (-0.3, -0.4) + assert engine.score("ab", "cd") == [-0.25, -0.25] + + +def test_vllm_refuses_eager_pack_quantized_in_place_and_low_host_memory(tmp_path): + with pytest.raises(SamplerEngineError, match="enforce-eager"): + VllmServerEngine(_vllm(tmp_path, extra_args=("--enforce-eager",))) + assert VllmServerEngine(_vllm(tmp_path, extra_args=("--enforce-eager",)), VllmOptions(allow_eager=True)) + with pytest.raises(SamplerEngineError, match="in-place reload"): + VllmServerEngine(_vllm(tmp_path, in_place_reload=True)) + release = tmp_path / "r1" + release.mkdir() + (release / "config.json").write_text(json.dumps({"quantization_config": {"format": "pack-quantized"}})) + engine = VllmServerEngine(_vllm(tmp_path), spawn=lambda a, e: _Process(), + http=_Http({("GET", "/health"): {}}), sleep=lambda s: None) + with pytest.raises(SamplerEngineError, match="pack-quantized"): + engine.start(release, 1) + low = VllmServerEngine(_vllm(tmp_path, min_host_free_gib=610), spawn=lambda a, e: _Process(), + http=_Http({("GET", "/health"): {}}), sleep=lambda s: None, + host_free_gib=lambda: 400.0) + with pytest.raises(SamplerEngineError, match="610"): + low.start(None, 0) + + +def test_stop_sends_sigterm(tmp_path): + process = _Process() + engine = VllmServerEngine(_vllm(tmp_path), spawn=lambda a, e: process, + http=_Http({("GET", "/health"): {}}), sleep=lambda s: None) + engine.start(None, 0) + engine.stop() + assert process.signals == [signal.SIGTERM] and engine.last_exit["normal"] + engine.stop() # idempotent diff --git a/yeto/rl/adapters/sampler_engines/server.py b/yeto/rl/adapters/sampler_engines/server.py new file mode 100644 index 000000000..5d868c46e --- /dev/null +++ b/yeto/rl/adapters/sampler_engines/server.py @@ -0,0 +1,295 @@ +"""Shared part of the HTTP-server engine adapters: one engine server process per +sampler instance, started on a release directory, checked for readiness with +a timeout, and stopped with a known exit-code rule. + +Guards that apply to every server engine (A-INF §3): + +* start timeout (item 15): a server that is not ready in ``start_timeout_s`` + is killed and reported with the adapter's ``hang_hint`` (910B: HCCL hang + without error when the CANN libprofapi stub is missing); +* host memory (item 17): before a start, the host must have at least + ``min_host_free_gib`` available (910B W4A8 stages int8 weights in host + memory, peak about 610 GiB per the report); +* exit codes (item 13): an exit after :meth:`ServerEngine.stop` asked the + server to stop is normal when the code is in ``exit_codes_after_stop`` (PPU + SGLang TP8 exits 134 from ``pcclCommAbort``); any exit while serving is a + failure; +* request timeouts on every HTTP call. + +The process and HTTP layers are injected (``spawn``, ``http``) so CPU tests +drive the adapters without a server. Defaults: ``subprocess.Popen`` and httpx. +""" + +from __future__ import annotations + +import os +import signal +import subprocess +import time +from collections.abc import Callable, Mapping, Sequence +from dataclasses import dataclass, field +from pathlib import Path +from typing import Any, Protocol + +from yeto.rl.engine.sampler_checks import release_quant_format +from yeto.rl.engine.sampler_engine import ( + RELOAD_IN_PLACE, + RELOAD_RESTART, + EngineIdentity, + LoadReport, + SamplerEngineError, +) + +GIB = 1024 ** 3 + + +class Process(Protocol): + def poll(self) -> int | None: ... + + def send_signal(self, sig: int) -> None: ... + + def kill(self) -> None: ... + + def wait(self, timeout: float | None = None) -> int: ... + + +Spawn = Callable[[Sequence[str], Mapping[str, str]], Process] +# (method, url, json body or None, timeout seconds) -> (HTTP status, decoded JSON or None) +Http = Callable[[str, str, Any, float], tuple[int, Any]] + + +def popen_spawn(argv: Sequence[str], env: Mapping[str, str]) -> Process: + return subprocess.Popen(list(argv), env={**os.environ, **env}) # noqa: S603 + + +def httpx_call(method: str, url: str, body: Any, timeout_s: float) -> tuple[int, Any]: + import httpx # noqa: PLC0415 - optional on hosts that only run tests + + try: + response = httpx.request(method, url, json=body, timeout=timeout_s) + except httpx.HTTPError as exc: + return 0, {"error": f"{type(exc).__name__}: {exc}"} + try: + return response.status_code, response.json() if response.content else None + except ValueError: + return response.status_code, {"error": response.text[:300]} + + +def host_available_gib(meminfo: str = "/proc/meminfo") -> float | None: + try: + with open(meminfo, encoding="ascii") as handle: + for line in handle: + if line.startswith("MemAvailable:"): + return int(line.split()[1]) * 1024 / GIB + except (OSError, ValueError, IndexError): + return None + return None + + +@dataclass +class ServerConfig: + """Settings every server adapter takes (run config / inventory, not code).""" + + base_dir: Path # the weights the engine starts on before any release + card_model: str + device_family: str + engine_version: str + precision: str # bf16 / w4a8 / mxfp4 ... + quant_scheme: str # the release scheme served (derived_from.scheme; bf16 for trained weights) + topology: str + max_model_len: int + host: str = "127.0.0.1" + port: int = 30000 + python: str = "python" + extra_args: tuple[str, ...] = () + extra_env: Mapping[str, str] = field(default_factory=dict) + start_timeout_s: float = 3600.0 + request_timeout_s: float = 3600.0 + stop_grace_s: float = 120.0 + min_host_free_gib: float | None = None + expected_format: str | None = None # quantization_config.format the release must carry + in_place_reload: bool = False # only after an on-card test passed (N4.2) + + @property + def url(self) -> str: + return f"http://{self.host}:{self.port}" + + +class ServerEngine: + """Base of the server adapters. Subclasses set ``ENGINE``, ``HANG_HINT``, + ``EXIT_CODES_AFTER_STOP`` and implement ``argv``, ``env``, ``generate``, + ``score``; optionally ``preflight`` and ``reload_in_place``.""" + + ENGINE = "server" + HANG_HINT = "" + EXIT_CODES_AFTER_STOP: frozenset[int] = frozenset() + REFUSED_FORMATS: frozenset[str] = frozenset() + + def __init__(self, config: ServerConfig, *, spawn: Spawn = popen_spawn, http: Http = httpx_call, + clock: Callable[[], float] = time.monotonic, + sleep: Callable[[float], None] = time.sleep, + host_free_gib: Callable[[], float | None] = host_available_gib) -> None: + self.config = config + self._spawn = spawn + self._http = http + self._clock = clock + self._sleep = sleep + self._host_free_gib = host_free_gib + self._process: Process | None = None + self._stopping = False + self._version: int | None = None + self.last_exit: dict[str, Any] | None = None + + # ------------------------------------------------------------------ identity + @property + def identity(self) -> EngineIdentity: + c = self.config + return EngineIdentity(engine=self.ENGINE, engine_version=c.engine_version, + device_family=c.device_family, card_model=c.card_model, + precision=c.precision, quant_scheme=c.quant_scheme, topology=c.topology) + + @property + def max_model_len(self) -> int: + return int(self.config.max_model_len) + + # ------------------------------------------------------------------ subclass hooks + def argv(self, model_dir: Path) -> list[str]: + raise NotImplementedError + + def env(self) -> dict[str, str]: + return dict(self.config.extra_env) + + def preflight(self, model_dir: Path) -> None: + """Refuse a release or configuration the engine would serve wrongly.""" + self.check_format(model_dir) + + def reload_in_place(self, model_dir: Path, version: int) -> None: + raise SamplerEngineError(f"{self.ENGINE}: in-place reload is not implemented") + + def after_ready(self) -> dict[str, Any]: + """Work after the server is ready and before it is checked (warm-up).""" + return {} + + # ------------------------------------------------------------------ guards + def check_format(self, model_dir: Path) -> None: + method, fmt = release_quant_format(model_dir) + want = self.config.expected_format + if fmt in self.REFUSED_FORMATS: + raise SamplerEngineError(f"{model_dir}: quantization_config.format {fmt!r} is refused " + f"by {self.ENGINE} (known to load wrong weights without error)") + if want is not None and fmt != want: + raise SamplerEngineError(f"{model_dir}: quantization_config.format is {fmt!r} " + f"(quant_method {method!r}); this engine is configured for {want!r}") + if want is None and self.config.quant_scheme == "bf16" and method is not None: + raise SamplerEngineError(f"{model_dir}: a bf16 engine got a quantized checkpoint " + f"({method}/{fmt})") + + def _check_host_memory(self) -> None: + need = self.config.min_host_free_gib + if need is None: + return + free = self._host_free_gib() + if free is None or free < need: + raise SamplerEngineError(f"{self.ENGINE}: host has {free} GiB available, the load needs " + f">= {need} GiB (weight staging in host memory)") + + # ------------------------------------------------------------------ http + def call(self, method: str, path: str, body: Any = None, timeout_s: float | None = None) -> Any: + self._check_alive() + status, data = self._http(method, f"{self.config.url}{path}", body, + float(timeout_s or self.config.request_timeout_s)) + if status != 200: + raise SamplerEngineError(f"{self.ENGINE} {method} {path}: HTTP {status}: {str(data)[:300]}") + return data + + def _check_alive(self) -> None: + if self._process is None: + raise SamplerEngineError(f"{self.ENGINE} is not running") + code = self._process.poll() + if code is not None and not self._stopping: + self._process = None + self._version = None + self.last_exit = {"code": code, "during": "service", "normal": False} + raise SamplerEngineError(f"{self.ENGINE} exited with code {code} while serving") + + # ------------------------------------------------------------------ lifecycle + def start(self, release_dir: Path | None, version: int) -> LoadReport: + model_dir = Path(release_dir) if release_dir is not None else Path(self.config.base_dir) + self.preflight(model_dir) + began = self._clock() + self.stop() + self._check_host_memory() + self._stopping = False + self._version = None + self._process = self._spawn(self.argv(model_dir), self.env()) + self._wait_ready() + detail = self.after_ready() + self._version = int(version) + return LoadReport(reload_mode=RELOAD_RESTART, pause_s=self._clock() - began, + detail={"model_dir": str(model_dir), **detail}) + + def load(self, release_dir: Path, version: int) -> LoadReport: + if not self.config.in_place_reload or self._process is None: + return self.start(release_dir, version) + self.preflight(Path(release_dir)) + began = self._clock() + self.reload_in_place(Path(release_dir), int(version)) + self._version = int(version) + return LoadReport(reload_mode=RELOAD_IN_PLACE, pause_s=self._clock() - began, + detail={"model_dir": str(release_dir)}) + + def _wait_ready(self) -> None: + deadline = self._clock() + float(self.config.start_timeout_s) + while True: + code = None if self._process is None else self._process.poll() + if code is not None: + self._process = None + self.last_exit = {"code": code, "during": "start", "normal": False} + raise SamplerEngineError(f"{self.ENGINE} exited with code {code} during start") + status, _ = self._http("GET", f"{self.config.url}/health", None, 10.0) + if status == 200: + return + if self._clock() >= deadline: + self._kill() + raise SamplerEngineError(f"{self.ENGINE} not ready after {self.config.start_timeout_s:g} s; " + f"killed. {self.HANG_HINT}".strip()) + self._sleep(5.0) + + def _kill(self) -> None: + if self._process is not None: + self._stopping = True + self._process.kill() + self._process.wait(timeout=30) + self._process = None + self._version = None + + def stop(self) -> None: + """Ask the server to stop; an exit code in ``EXIT_CODES_AFTER_STOP`` is normal.""" + process, self._process = self._process, None + self._version = None + if process is None: + return + self._stopping = True + code = process.poll() + if code is None: + process.send_signal(signal.SIGTERM) + try: + code = process.wait(timeout=self.config.stop_grace_s) + except subprocess.TimeoutExpired: + process.kill() + code = process.wait(timeout=30) + known = code in self.EXIT_CODES_AFTER_STOP + self.last_exit = {"code": code, "during": "stop", + "normal": code in (0, -signal.SIGTERM) or known, + **({"note": "known exit code after stop"} if known else {})} + + def served_version(self) -> int | None: + return self._version if self._process is not None else None + + # The servers report neither count; the post-load check then uses the + # configured weight memory or fails (sampler_checks.PostLoadChecks). + def loaded_expert_matrices(self) -> int | None: + return None + + def weight_gib_per_card(self) -> float | None: + return None diff --git a/yeto/rl/adapters/sampler_engines/sglang_server.py b/yeto/rl/adapters/sampler_engines/sglang_server.py new file mode 100644 index 000000000..ed879083e --- /dev/null +++ b/yeto/rl/adapters/sampler_engines/sglang_server.py @@ -0,0 +1,192 @@ +"""``sglang-server``: an SGLang HTTP server behind the ``SamplerEngine`` port +(PPU, stage 1: TP8 per machine, MXFP4 experts RTN; user decision 2026-10-19). + +Vendor facts handled here (A-INF §3; none verified on a PPU by this code): + +* 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``); +* item 5: an MXFP4 export with ``quantization_config.format == "dense"`` + (llmcompressor 0.14) is read wrongly; ``dense`` is refused for every + quantized scheme; +* item 7: the MXFP4 W4A16 path needs ``MX_PATCH=1``; set when the precision is + ``mxfp4`` and recorded in the load report; +* item 8: dp8 + DeepEP cannot run NEXTN; the pair is refused before start; +* item 9: dp8 + DeepGEMM compiles per new batch size M (75-113 s each); with + ``warmup_batch_sizes`` the adapter sends one batch per size before the + engine is checked and records the seconds; +* item 13: TP8 exits 134 from ``pcclCommAbort`` after a stop request; normal + then, a failure while serving. + +API calls (SGLang native HTTP API; the exact fields of SGLang 0.5.17 on PPU are +NOT verified): ``GET /health``; ``POST /generate`` with ``return_logprob`` +(``meta_info.output_token_logprobs`` = ``[logprob, token_id, ...]`` per +token, ``meta_info.finish_reason.type``); teacher-forced scores from +``meta_info.input_token_logprobs`` with ``max_new_tokens: 0``; +``POST /update_weights_from_disk`` for an in-place reload (off unless +``in_place_reload``: refit of quantized experts is not tested, A-INF 2.1). +""" + +from __future__ import annotations + +import json +from collections.abc import Sequence +from dataclasses import dataclass +from pathlib import Path +from typing import Any + +from yeto.rl.engine.sampler_checks import INDEX_FILE +from yeto.rl.engine.sampler_engine import ( + FINISH_LENGTH, + FINISH_STOP, + Completion, + SamplerEngineError, + SamplingParams, +) + +from .server import ServerConfig, ServerEngine + +TOPOLOGY_TP8 = "tp8" +TOPOLOGY_DP8 = "dp8" +SGLANG_TOPOLOGIES = (TOPOLOGY_TP8, TOPOLOGY_DP8) + + +@dataclass +class SGLangOptions: + tp: int = 8 + nextn: bool = False + renames_weight_packed: bool = False # the image patches the weight_packed -> weight map + warmup_batch_sizes: tuple[int, ...] = () + served_model_name: str = "yeto-policy" + + +def _finish(meta: dict[str, Any]) -> str: + reason = meta.get("finish_reason") + kind = reason.get("type") if isinstance(reason, dict) else reason + if kind == "length": + return FINISH_LENGTH + if kind == "stop": + return FINISH_STOP + raise SamplerEngineError(f"sglang finish_reason {reason!r}: request aborted or unknown") + + +class SGLangServerEngine(ServerEngine): + ENGINE = "sglang" + HANG_HINT = "Check the PCCL/ACCL-P collective setup of the PPU image." + EXIT_CODES_AFTER_STOP = frozenset({134}) # pcclCommAbort at TP8 shutdown (item 13) + REFUSED_FORMATS = frozenset({"dense"}) # item 5 + + 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); " + "turn one of them off") + + def argv(self, model_dir: Path) -> list[str]: + c, o = self.config, self.options + argv = [c.python, "-m", "sglang.launch_server", "--model-path", str(model_dir), + "--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: + argv += ["--dp", str(o.tp), "--enable-dp-attention", "--moe-a2a-backend", "deepep"] + if o.nextn: + argv += ["--speculative-algorithm", "NEXTN"] + return argv + list(c.extra_args) + + def env(self) -> dict[str, str]: + env = super().env() + if self.config.precision == "mxfp4": + env["MX_PATCH"] = "1" # item 7 + return env + + def preflight(self, model_dir: Path) -> None: + super().preflight(model_dir) + index = Path(model_dir) / INDEX_FILE + if index.is_file() and not self.options.renames_weight_packed: + names = json.loads(index.read_text(encoding="utf-8")).get("weight_map") or {} + packed = [n for n in names if n.endswith(".weight_packed")] + if packed: + raise SamplerEngineError( + f"{model_dir}: {len(packed)} tensors are named *.weight_packed (e.g. " + f"{packed[0]}); SGLang skips them without an error (A-INF item 1). Use an " + "image with the rename patch and set renames_weight_packed") + + def after_ready(self) -> dict[str, Any]: + detail: dict[str, Any] = {"env": {k: v for k, v in self.env().items() if k == "MX_PATCH"}} + if self.options.warmup_batch_sizes: + began = self._clock() + for size in self.options.warmup_batch_sizes: + self._generate_raw(["warm-up"] * int(size), + SamplingParams(max_tokens=8, temperature=0.0)) + detail["warmup_batch_sizes"] = list(self.options.warmup_batch_sizes) + detail["warmup_s"] = round(self._clock() - began, 3) + return detail + + def reload_in_place(self, model_dir: Path, version: int) -> None: + body = self.call("POST", "/update_weights_from_disk", + {"model_path": str(model_dir), "weight_version": str(int(version))}) + if isinstance(body, dict) and body.get("success") is False: + raise SamplerEngineError(f"sglang update_weights_from_disk: {body.get('message')}") + + # ------------------------------------------------------------------ generation + @staticmethod + def _params(sampling: SamplingParams, max_new_tokens: int | None = None) -> dict[str, Any]: + params: dict[str, Any] = { + "max_new_tokens": int(sampling.max_tokens if max_new_tokens is None else max_new_tokens), + "temperature": float(sampling.temperature), "top_p": float(sampling.top_p), + "top_k": int(sampling.top_k), + "repetition_penalty": float(sampling.repetition_penalty), + "frequency_penalty": float(sampling.frequency_penalty), + "presence_penalty": float(sampling.presence_penalty), + } + if sampling.stop: + params["stop"] = list(sampling.stop) + if sampling.seed is not None: + params["sampling_seed"] = int(sampling.seed) + return params + + def _generate_raw(self, texts: Sequence[str], sampling: SamplingParams) -> list[dict[str, Any]]: + body = self.call("POST", "/generate", {"text": list(texts), "sampling_params": self._params(sampling), + "return_logprob": True}) + items = body if isinstance(body, list) else [body] + if len(items) != len(texts): + raise SamplerEngineError(f"sglang returned {len(items)} outputs for {len(texts)} inputs") + return items + + def generate(self, prompts: Sequence[str], sampling: SamplingParams) -> list[list[Completion]]: + n = int(sampling.n) + texts = [p for p in prompts for _ in range(n)] # one request per sample: order is explicit + items = self._generate_raw(texts, sampling.replace(n=1)) + out: list[list[Completion]] = [] + for i in range(len(prompts)): + group = [] + for item in items[i * n:(i + 1) * n]: + meta = item.get("meta_info") or {} + tokens = meta.get("output_token_logprobs") or [] + group.append(Completion(token_ids=tuple(int(t[1]) for t in tokens), + logprobs=tuple(float(t[0]) for t in tokens), + text=str(item.get("text", "")), finish_reason=_finish(meta))) + out.append(group) + return out + + def _input_logprobs(self, text: str) -> list[float | None]: + body = self.call("POST", "/generate", { + "text": text, "sampling_params": {"max_new_tokens": 0, "temperature": 0.0}, + "return_logprob": True, "logprob_start_len": 0}) + body = body[0] if isinstance(body, list) else body + return [None if t[0] is None else float(t[0]) + for t in (body.get("meta_info") or {}).get("input_token_logprobs") or []] + + def score(self, prompt: str, continuation: str) -> list[float]: + prompt_len = len(self._input_logprobs(prompt)) + full = self._input_logprobs(prompt + continuation) + tail = full[prompt_len:] + if not tail or any(x is None for x in tail): + raise SamplerEngineError(f"sglang gave no input log-probs for continuation {continuation!r}") + return [float(x) for x in tail] diff --git a/yeto/rl/adapters/sampler_engines/vllm_server.py b/yeto/rl/adapters/sampler_engines/vllm_server.py new file mode 100644 index 000000000..af5f9067e --- /dev/null +++ b/yeto/rl/adapters/sampler_engines/vllm_server.py @@ -0,0 +1,146 @@ +"""``vllm-server``: a vLLM OpenAI-compatible server behind the ``SamplerEngine`` +port (vllm-ascend 0.30.0rc1 on 910B, stage 2: TP4 per instance, W4A8). + +Vendor facts handled here (A-INF §3, report items 15-18; none verified on a +910B by this code): + +* reload: restart only. In-place reload of W4A8_DYNAMIC on vllm-ascend is not + tested (GOAL §8); ``in_place_reload`` stays False until N4.2 measures it, and + this adapter has no in-place path yet; +* item 17: 910B reads ``int-quantized`` only, not ``pack-quantized``; the + default ``expected_format`` for a quantized scheme is ``int-quantized``. + W4A8 stages int8 weights in host memory (``VLLM_ASCEND_W4A8_CPU_STAGE=1``), + so set ``min_host_free_gib`` from the measured peak (report: about 610 GiB); +* item 18: eager mode serves about 2.75 tok/s against about 38 tok/s with + graphs; ``--enforce-eager`` is refused unless ``allow_eager`` (debug only); +* item 15: a missing CANN 9.1 libprofapi stub hangs HCCL without an error; the + start timeout names it; +* items 16 and the PT_INTERP fix belong to the image or rootfs start script + (N0.2), not to this adapter. + +API calls (vLLM OpenAI-compatible server): ``GET /health``; +``POST /v1/completions`` with ``logprobs: 1`` and +``return_tokens_as_token_ids`` (tokens come back as ``token_id:``); +teacher-forced scores with ``echo: true`` and ``max_tokens: 1``. Field +support in vllm-ascend 0.30.0rc1 is NOT verified. +""" + +from __future__ import annotations + +from collections.abc import Sequence +from dataclasses import dataclass +from pathlib import Path +from typing import Any + +from yeto.rl.engine.sampler_engine import ( + FINISH_LENGTH, + FINISH_STOP, + Completion, + SamplerEngineError, + SamplingParams, +) + +from .server import ServerConfig, ServerEngine + +INT_QUANTIZED = "int-quantized" + + +@dataclass +class VllmOptions: + tp: int = 4 + max_num_seqs: int | None = None # N4.1 measures 8 / 32 / 64 + served_model_name: str = "yeto-policy" + allow_eager: bool = False + + +def _token_id(token: str) -> int: + if not isinstance(token, str) or not token.startswith("token_id:"): + raise SamplerEngineError(f"vllm token {token!r} is not 'token_id:'; " + "the server ignored return_tokens_as_token_ids") + return int(token.split(":", 1)[1]) + + +class VllmServerEngine(ServerEngine): + ENGINE = "vllm" + HANG_HINT = ("Suspected HCCL hang: check that the CANN 9.1 libprofapi stub is installed " + "(report item 15).") + REFUSED_FORMATS = frozenset({"pack-quantized", "dense"}) + + def __init__(self, config: ServerConfig, options: VllmOptions | None = None, **kw: Any) -> None: + if config.quant_scheme != "bf16" and config.expected_format is None: + config.expected_format = INT_QUANTIZED # item 17 + if config.in_place_reload: + raise SamplerEngineError("vllm-server: in-place reload is not built; restart is the " + "only reload mode until N4.2 tests it") + super().__init__(config, **kw) + self.options = options or VllmOptions() + if "--enforce-eager" in config.extra_args and not self.options.allow_eager: + raise SamplerEngineError("vllm-server: --enforce-eager serves about 14x slower on 910B " + "(report item 18); set allow_eager only for debugging") + + def argv(self, model_dir: Path) -> list[str]: + c, o = self.config, self.options + argv = [c.python, "-m", "vllm.entrypoints.openai.api_server", "--model", str(model_dir), + "--host", c.host, "--port", str(c.port), "--tensor-parallel-size", str(o.tp), + "--max-model-len", str(c.max_model_len), "--served-model-name", o.served_model_name, + "--trust-remote-code"] + if o.max_num_seqs is not None: + argv += ["--max-num-seqs", str(int(o.max_num_seqs))] + return argv + list(c.extra_args) + + def env(self) -> dict[str, str]: + env = super().env() + if self.config.precision == "w4a8": + env["VLLM_ASCEND_W4A8_CPU_STAGE"] = "1" # int8 staging in host memory + return env + + def _body(self, prompt: Any, sampling: SamplingParams, **extra: Any) -> dict[str, Any]: + body: dict[str, Any] = { + "model": self.options.served_model_name, "prompt": prompt, "n": int(sampling.n), + "max_tokens": int(sampling.max_tokens), "temperature": float(sampling.temperature), + "top_p": float(sampling.top_p), "top_k": int(sampling.top_k), + "repetition_penalty": float(sampling.repetition_penalty), + "frequency_penalty": float(sampling.frequency_penalty), + "presence_penalty": float(sampling.presence_penalty), + "logprobs": 1, "return_tokens_as_token_ids": True, **extra} + if sampling.stop: + body["stop"] = list(sampling.stop) + if sampling.seed is not None: + body["seed"] = int(sampling.seed) + return body + + def generate(self, prompts: Sequence[str], sampling: SamplingParams) -> list[list[Completion]]: + n = int(sampling.n) + body = self.call("POST", "/v1/completions", self._body(list(prompts), sampling)) + choices = sorted(body.get("choices") or [], key=lambda c: int(c["index"])) + if len(choices) != len(prompts) * n: + raise SamplerEngineError(f"vllm returned {len(choices)} choices for {len(prompts)} x {n}") + out: list[list[Completion]] = [] + for i in range(len(prompts)): + group = [] + for choice in choices[i * n:(i + 1) * n]: + logprobs = choice.get("logprobs") or {} + reason = choice.get("finish_reason") + if reason not in ("stop", "length"): + raise SamplerEngineError(f"vllm finish_reason {reason!r}") + group.append(Completion( + token_ids=tuple(_token_id(t) for t in logprobs.get("tokens") or []), + logprobs=tuple(float(x) for x in logprobs.get("token_logprobs") or []), + text=str(choice.get("text", "")), + finish_reason=FINISH_LENGTH if reason == "length" else FINISH_STOP)) + out.append(group) + return out + + def _echo_logprobs(self, text: str) -> list[float | None]: + greedy = SamplingParams(max_tokens=1, temperature=0.0) + body = self.call("POST", "/v1/completions", self._body(text, greedy, echo=True)) + choice = (body.get("choices") or [{}])[0] + values = (choice.get("logprobs") or {}).get("token_logprobs") or [] + return list(values[:-1]) # the last value is the one generated token + + def score(self, prompt: str, continuation: str) -> list[float]: + prompt_len = len(self._echo_logprobs(prompt)) + tail = self._echo_logprobs(prompt + continuation)[prompt_len:] + if not tail or any(x is None for x in tail): + raise SamplerEngineError(f"vllm gave no prompt log-probs for continuation {continuation!r}") + return [float(x) for x in tail] diff --git a/yeto/rl/engine/engine_sampler.py b/yeto/rl/engine/engine_sampler.py new file mode 100644 index 000000000..1ec63a1bf --- /dev/null +++ b/yeto/rl/engine/engine_sampler.py @@ -0,0 +1,198 @@ +"""The engine sampler island: ``SamplerLoop`` over a ``SamplerEngine`` adapter +(rl-hetero-sampler-engines D1, S25 N2.3). + +The Miles sampler island (``adapters.miles.sampler_island``) lets Miles drive +SGLang. This island drives any engine behind the port instead, so the same +loop, buffer writer, reward call and release follower serve SGLang on PPU and +vllm-ascend on 910B. Per round (:class:`EngineRound`): + +1. take the round's prompts (``prompts(rollout_id)``); +2. ``engine.generate`` with the run's :class:`SamplingParams` (one value for + every source of the run, checked at start); +3. cut looping completions (:class:`~.loop_guard.LoopGuard`, off by default); +4. score each completion (``reward(item, completion)``); +5. write one group per prompt to the buffer (:class:`~.buffer_payload.GroupWriter`) + with the behavior version the engine serves and the source record + (:class:`~.sampler_engine.Provenance`) in the index entry; +6. write ``rl_sampler_engine_round`` (finish reasons, tokens, provenance key). + +: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 +starts the resource sampler of its device family +(``telemetry.ResourceSampler`` with ``card_telemetry.collector_for``) and +calls ``island.run()``. +""" + +from __future__ import annotations + +from collections import Counter +from collections.abc import Callable, Mapping, Sequence +from dataclasses import dataclass, field +from pathlib import Path +from typing import Any + +from .buffer_payload import GroupWriter +from .loop_guard import LoopGuard +from .policy_release import PolicyFollower, PolicyReleaseStore, ReleaseNotLoaded +from .policy_versions import engine_weight_version +from .sampler_checks import PortReleaseLoader, PostLoadChecks, SamplerLoadFatal, check_sampling_fits +from .sampler_engine import ( + FINISH_STOP, + Completion, + Provenance, + SamplerEngine, + SamplerEngineError, + SamplingParams, +) +from .sampler_loop import SamplerLoop + +ROUND_EVENT = "rl_sampler_engine_round" +START_EVENT = "rl_sampler_start" + + +@dataclass(frozen=True) +class PromptItem: + """One prompt of a round; ``label`` and ``metadata`` go to the reward and payload.""" + + prompt_id: str + prompt: str + label: Any = None + metadata: Mapping[str, Any] = field(default_factory=dict) + + +Reward = Callable[[PromptItem, Completion], float] + + +def sample_record(item: PromptItem, completion: Completion, *, reward: float, version: int, + provenance: Mapping[str, str], group_index: int, index: int) -> dict[str, Any]: + """One sample in the buffer payload (the ``Sample.to_dict()`` keys the + trainer readers use; ``tokens`` is prompt + response when the engine gave + the prompt tokens, else the response only and ``prompt_tokens_known`` False).""" + response = list(completion.token_ids) + prompt_ids = list(completion.prompt_token_ids) + return { + "group_index": group_index, "index": index, "prompt": item.prompt, + "response": completion.text, "label": item.label, + "tokens": prompt_ids + response, "response_length": len(response), + "rollout_log_probs": list(completion.logprobs), "reward": float(reward), + "status": "completed" if completion.finish_reason == FINISH_STOP else "truncated", + "metadata": {**dict(item.metadata), "prompt_id": item.prompt_id, + "weight_version": engine_weight_version(int(version)), + "finish_reason": completion.finish_reason, + "prompt_tokens_known": bool(prompt_ids), "provenance": dict(provenance)}, + } + + +class EngineRound: + """``generate(rollout_id)`` of the :class:`SamplerLoop` (module doc).""" + + 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: + self.engine = engine + self.writer = writer + self.prompts = prompts + self.reward = reward + self.sampling = sampling.check() + self.provenance = provenance + self.emit = emit + self.loop_guard = loop_guard + + def __call__(self, rollout_id: int) -> None: + items = list(self.prompts(int(rollout_id))) + version = self.engine.served_version() + if version is None: + raise SamplerEngineError("the engine serves no version; it cannot generate") + groups = self.engine.generate([item.prompt for item in items], self.sampling) + if len(groups) != len(items): + raise SamplerEngineError(f"engine returned {len(groups)} groups for {len(items)} prompts") + if self.engine.served_version() != version: # a reload during the round would mix versions + raise SamplerEngineError(f"served version changed from v{version} during round {rollout_id}") + record = self.provenance.to_record() + finish: Counter[str] = Counter() + tokens = 0 + 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, " + f"expected n={self.sampling.n}") + samples = [] + for j, completion in enumerate(group): + if self.loop_guard is not None: + completion = self.loop_guard.apply(completion) + finish[completion.finish_reason] += 1 + tokens += len(completion.token_ids) + samples.append(sample_record( + 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()) + + +@dataclass +class EngineSamplerConfig: + """What the run config gives an engine sampler island.""" + + sampler_id: str + num_rollouts: int + max_backlog_groups: int + sampling: SamplingParams + max_prompt_tokens: int + base_version: int = 0 + base_dir: Path | None = None # None: the engine's configured base weights + min_round_s: float | None = None + poll_s: float = 5.0 + loop_guard: LoopGuard | None = None + max_consecutive_load_failures: int = 3 + + +@dataclass +class EngineSamplerIsland: + loop: SamplerLoop + loader: PortReleaseLoader + follower: PolicyFollower + round: EngineRound + + def run(self, start: int = 0) -> int: + return self.loop.run(start) + + +def build_engine_sampler(config: EngineSamplerConfig, *, engine: SamplerEngine, + releases: PolicyReleaseStore, buffer: Any, emit: Callable[..., None], + prompts: Callable[[int], Sequence[PromptItem]], reward: Reward, + checks: PostLoadChecks | None = None, + should_stop: Callable[[], bool] = lambda: False, + sleep: Callable[[float], None] | 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() + check_sampling_fits(sampling, engine, max_prompt_tokens=config.max_prompt_tokens) + provenance = Provenance.of(engine.identity, sampler_id=config.sampler_id, sampling=sampling) + emit(START_EVENT, sampler_id=config.sampler_id, num_rollouts=int(config.num_rollouts), + max_backlog_groups=int(config.max_backlog_groups), sampling=sampling.to_dict(), + sampling_digest=sampling.digest(), allow_penalties=sampling.allow_penalties, + max_prompt_tokens=int(config.max_prompt_tokens), provenance=provenance.to_record(), + loop_guard=None if config.loop_guard is None else vars(config.loop_guard).copy()) + loader = PortReleaseLoader(engine, checks or PostLoadChecks(), emit, + max_consecutive_failures=config.max_consecutive_load_failures) + try: + loader.start(config.base_dir, int(config.base_version)) + except ReleaseNotLoaded as error: # no served version to fall back to + raise SamplerLoadFatal(f"the engine failed its start checks: {error}") from error + 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) + extra = {} if sleep is None else {"sleep": sleep} + loop = SamplerLoop(follower=follower, buffer=buffer, generate=engine_round, + num_rollouts=int(config.num_rollouts), + max_backlog_groups=int(config.max_backlog_groups), emit=emit, + 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) diff --git a/yeto/rl/engine/loop_guard.py b/yeto/rl/engine/loop_guard.py new file mode 100644 index 000000000..a10e03b66 --- /dev/null +++ b/yeto/rl/engine/loop_guard.py @@ -0,0 +1,54 @@ +"""Cut repeating completions instead of penalising repetition (A-INF 5.5, D11). + +A repetition penalty changes the behavior distribution, so rollouts do not use +one (``SamplingParams.check``). A completion that falls into a loop is cut at +the point the loop is detected instead: the kept tokens and their log-probs are +exactly what the engine sampled, so the behavior log-probs stay consistent; the +completion ends with ``finish_reason: loop_abort``. + +Detection: some period ``p`` in ``1..max_period`` whose last ``p * repeats`` +tokens are one ``p``-gram repeated ``repeats`` times. ``run[p]`` counts the +consecutive positions with ``token[i] == token[i - p]``; the loop is detected +at the first ``i`` with ``run[p] >= p * (repeats - 1)``. O(len * max_period). + +Off by default: the thresholds must first be measured on round-0 samples +(A-INF 5.5). Cutting after the engine finished does not save generation time; +a streaming cut inside the adapter would, and is not built. +""" + +from __future__ import annotations + +from collections.abc import Sequence +from dataclasses import dataclass + +from .sampler_engine import FINISH_LOOP_ABORT, Completion + + +@dataclass(frozen=True) +class LoopGuard: + max_period: int + repeats: int + + def __post_init__(self) -> None: + if self.max_period < 1 or self.repeats < 2: + raise ValueError("LoopGuard needs max_period >= 1 and repeats >= 2") + + def cut_at(self, token_ids: Sequence[int]) -> int | None: + """Number of tokens to keep (through the detection point), or None.""" + run = [0] * (self.max_period + 1) + for i in range(1, len(token_ids)): + for p in range(1, min(self.max_period, i) + 1): + run[p] = run[p] + 1 if token_ids[i] == token_ids[i - p] else 0 + if run[p] >= p * (self.repeats - 1): + return i + 1 + return None + + def apply(self, completion: Completion) -> Completion: + keep = self.cut_at(completion.token_ids) + if keep is None or keep >= len(completion.token_ids): + return completion + # The text of the kept tokens is not known without the tokenizer; the + # sample's text is the engine text marked as cut (training reads tokens). + return Completion(token_ids=completion.token_ids[:keep], logprobs=completion.logprobs[:keep], + text=completion.text, finish_reason=FINISH_LOOP_ABORT, + prompt_token_ids=completion.prompt_token_ids) diff --git a/yeto/rl/engine/policy_release.py b/yeto/rl/engine/policy_release.py index 4b65583d1..afd4540bc 100644 --- a/yeto/rl/engine/policy_release.py +++ b/yeto/rl/engine/policy_release.py @@ -129,6 +129,18 @@ class ReleaseIncomplete(PolicyReleaseError): """A release whose files are not all present with their recorded sizes yet.""" +class ReleaseNotLoaded(PolicyReleaseError): + """The load hook did not take this release, and the sampler may go on: the + release scheme is not the one the engine serves, or a post-load check + failed (``yeto.rl.engine.sampler_checks``). The follower keeps its served + release and tries the newest one again on the next refresh. ``reason`` is + a short machine-readable cause.""" + + def __init__(self, message: str, *, reason: str) -> None: + super().__init__(message) + self.reason = reason + + def file_sizes(directory: Path) -> dict[str, int]: """{relative path: size} of every file under ``directory`` but the record.""" root = Path(directory) diff --git a/yeto/rl/engine/sampler_checks.py b/yeto/rl/engine/sampler_checks.py new file mode 100644 index 000000000..c29e74db0 --- /dev/null +++ b/yeto/rl/engine/sampler_checks.py @@ -0,0 +1,334 @@ +"""Load a release through the ``SamplerEngine`` port and check what the engine +serves before it may generate (S25 N2.3; rl-hetero-sampler-engines D9; A-INF §3). + +A load that returns without error does not prove that the engine serves the +right weights: the reports list loads that skip expert weights or misread them +silently (A-INF items 1, 2, 3, 5). :class:`PortReleaseLoader` is the +``PolicyFollower`` load hook of an engine sampler island. For each release it: + +1. refuses a release whose scheme is not the one the engine serves + (:class:`~.policy_release.ReleaseNotLoaded`, reason ``scheme_mismatch``; + the engine keeps its version); +2. counts the routed-expert matrices in the release index and compares them + with the expected count (release record, else the check config); +3. loads the release (``engine.load``; a retry after a failed check restarts + the engine with ``engine.start``); +4. runs the three post-load checks of :class:`PostLoadChecks`: + + * ``expert_count``: the engine's loaded expert-matrix count equals the + expected count (GLM5.3-Flash: 36,288); when the engine cannot tell, its + weight memory per card is compared with the expected value; when neither + is known the check fails unless ``allow_unverified_count`` is set (and + the tape says so); + * ``fixed_prompts``: greedy answers to fixed prompts (4 GSM8K questions by + default) are non-empty, contain the expected answer, carry no forbidden + marker and almost no unreadable characters; + * ``logprob_probe``: teacher-forced log-probs of fixed continuations; the + mean absolute change against the previous version (or against a stored + round-0 reference) must not exceed ``max_probe_jump`` (None: recorded only; + the threshold is set after the first run on cards, A-INF §3); + +5. checks that the engine reports the release version (``served_version``). + +Only then is the release served. A failed step writes +``sampler_load_check_failed``, marks the engine not ready (the loop generates +nothing) and raises ``ReleaseNotLoaded`` so the follower tries again; after +``max_consecutive_failures`` failures it raises :class:`SamplerLoadFatal` and +the island stops with the reason. +""" + +from __future__ import annotations + +import json +import math +import re +import time +from collections.abc import Callable, Sequence +from dataclasses import dataclass, field +from pathlib import Path +from typing import Any + +from .policy_release import PolicyRelease, ReleaseNotLoaded +from .sampler_engine import SamplerEngine, SamplerEngineError, SamplingParams + +# Routed-expert matrices of GLM5.3-Flash as the inference report counts them +# (SESSION25/reports/glm-inference-ppu-910b.md, A-INF §3 item 1). +GLM53_FLASH_EXPERT_MATRICES = 36_288 +# Names of routed-expert weight tensors in an HF / compressed-tensors index: +# "...mlp.experts.._proj.weight" or ".weight_packed". +EXPERT_WEIGHT_PATTERN = r"\.experts\.\d+\.\w+_proj\.weight(?:_packed)?$" +INDEX_FILE = "model.safetensors.index.json" +BF16 = "bf16" + +LOAD_CHECK_FAILED_EVENT = "sampler_load_check_failed" +LOAD_CHECKED_EVENT = "rl_sampler_load_checked" + + +class SamplerLoadFatal(RuntimeError): + """Post-load checks failed too often; the island must stop.""" + + +@dataclass(frozen=True) +class FixedPrompt: + prompt: str + answer: str # must appear in the greedy output + + +GSM8K_FIXED_PROMPTS = ( + FixedPrompt("Question: Natalia sold clips to 48 of her friends in April, and then she sold " + "half as many clips in May. How many clips did Natalia sell altogether in April " + "and May?\nAnswer:", "72"), + FixedPrompt("Question: Weng earns $12 an hour for babysitting. Yesterday, she just did 50 " + "minutes of babysitting. How much did she earn?\nAnswer:", "10"), + FixedPrompt("Question: Betty is saving money for a new wallet which costs $100. Betty has only " + "half of the money she needs. Her parents decided to give her $15 for that " + "purpose, and her grandparents twice as much as her parents. How much more money " + "does Betty need to buy the wallet?\nAnswer:", "5"), + FixedPrompt("Question: Julie is reading a 120-page book. Yesterday, she was able to read 12 " + "pages and today, she read twice as many pages as yesterday. If she wants to read " + "half of the remaining pages tomorrow, how many pages should she read?\nAnswer:", + "42"), +) + + +def release_scheme(release: PolicyRelease | None) -> str: + """The scheme of a release: ``derived_from.scheme`` of a derived release, else bf16.""" + derived = None if release is None else release.derived_from + return str((derived or {}).get("scheme") or BF16) + + +def count_release_experts(directory: Path, pattern: str = EXPERT_WEIGHT_PATTERN) -> int | None: + """Routed-expert weight tensors named in the release's safetensors index + (None: the release has no index file).""" + index = Path(directory) / INDEX_FILE + if not index.is_file(): + return None + weight_map = json.loads(index.read_text(encoding="utf-8")).get("weight_map") or {} + matcher = re.compile(pattern) + return sum(1 for name in weight_map if matcher.search(name)) + + +def release_quant_format(directory: Path) -> tuple[str | None, str | None]: + """``(quant_method, format)`` of ``config.json``'s ``quantization_config`` + (``(None, None)``: no quantization config, a BF16 checkpoint).""" + config = Path(directory) / "config.json" + if not config.is_file(): + return None, None + quant = json.loads(config.read_text(encoding="utf-8")).get("quantization_config") or {} + method = quant.get("quant_method") + fmt = quant.get("format") + return (None if method is None else str(method)), (None if fmt is None else str(fmt)) + + +def _unreadable_share(text: str) -> float: + if not text: + return 1.0 + bad = sum(1 for c in text if c == "�" or (ord(c) < 32 and c not in "\n\r\t")) + return bad / len(text) + + +def _contains_answer(text: str, answer: str) -> bool: + """A numeric answer must stand alone ("5" is not found in "15").""" + if answer.isdigit(): + return re.search(rf"(? list[tuple[str, str]]: + if self.probe_pairs is not None: + return list(self.probe_pairs) + return [(p.prompt, f" {p.answer}") for p in self.fixed_prompts] + + # ------------------------------------------------------------------ checks + def expert_count(self, engine: SamplerEngine, expected: int | None) -> dict[str, Any]: + if expected is None: + return {"name": "expert_count", "passed": True, "detail": "no expected count (bf16 engine)"} + loaded = engine.loaded_expert_matrices() + if loaded is not None: + return {"name": "expert_count", "passed": int(loaded) == int(expected), + "loaded": int(loaded), "expected": int(expected)} + gib = engine.weight_gib_per_card() + want = self.expected_weight_gib_per_card + if gib is not None and want is not None: + ok = abs(float(gib) - want) <= self.weight_gib_tolerance * want + return {"name": "expert_count", "passed": ok, "weight_gib_per_card": float(gib), + "expected_weight_gib_per_card": want, "tolerance": self.weight_gib_tolerance} + return {"name": "expert_count", "passed": bool(self.allow_unverified_count), + "unverified": True, "expected": int(expected), + "detail": ("the engine reports neither its expert count nor its weight memory; " + + ("allowed by allow_unverified_count" if self.allow_unverified_count + else "set expected_weight_gib_per_card or allow_unverified_count"))} + + def fixed_prompt_answers(self, engine: SamplerEngine) -> dict[str, Any]: + if not self.fixed_prompts: + return {"name": "fixed_prompts", "passed": True, "detail": "no fixed prompts"} + greedy = SamplingParams(max_tokens=int(self.fixed_max_tokens), temperature=0.0, seed=0).check() + groups = engine.generate([p.prompt for p in self.fixed_prompts], greedy) + answers = [] + for item, group in zip(self.fixed_prompts, groups): + text = group[0].text if group else "" + problems = [] + if not text.strip(): + problems.append("empty") + if not _contains_answer(text, item.answer): + 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]}) + return {"name": "fixed_prompts", "passed": all(a["ok"] for a in answers) + and len(groups) == len(self.fixed_prompts), "answers": answers} + + def logprob_probe(self, engine: SamplerEngine) -> dict[str, Any]: + means = [] + for prompt, continuation in self.pairs(): + logprobs = engine.score(prompt, continuation) + if not logprobs or not all(math.isfinite(x) for x in logprobs): + return {"name": "logprob_probe", "passed": False, + "detail": f"empty or non-finite log-probs for {continuation!r}"} + means.append(sum(logprobs) / len(logprobs)) + against, label = ((self.reference_probe, "reference") if self.reference_probe is not None + else (self._last_probe, "previous")) + out: dict[str, Any] = {"name": "logprob_probe", "means": [round(m, 6) for m in means], + "against": label if against is not None else None} + jump = None + if against is not None: + if len(against) != len(means): + return {**out, "passed": False, "detail": "probe pair count differs from the reference"} + jump = sum(abs(a - b) for a, b in zip(means, against)) / len(means) + out["mean_abs_jump"] = round(jump, 6) + out["max_probe_jump"] = self.max_probe_jump + out["passed"] = jump is None or self.max_probe_jump is None or jump <= self.max_probe_jump + if out["passed"]: + self._last_probe = means + return out + + def run(self, engine: SamplerEngine, *, expected_experts: int | None) -> list[dict[str, Any]]: + results = [self.expert_count(engine, expected_experts)] + results.append(self.fixed_prompt_answers(engine)) + results.append(self.logprob_probe(engine)) + return results + + +class PortReleaseLoader: + """``PolicyFollower`` load hook over a :class:`SamplerEngine` (module doc). + + ``emit(event, **fields)`` writes the tape. ``expected_experts`` for a release + is ``derived_from.expert_matrices`` when the release records it, else the + check config's ``expected_expert_matrices`` (bf16 engines: None).""" + + def __init__(self, engine: SamplerEngine, checks: PostLoadChecks, emit: Callable[..., None], *, + max_consecutive_failures: int = 3, + clock: Callable[[], float] = time.monotonic) -> None: + if int(max_consecutive_failures) < 1: + raise ValueError("max_consecutive_failures must be >= 1") + self.engine = engine + self.checks = checks + self.emit = emit + self.max_consecutive_failures = int(max_consecutive_failures) + self.clock = clock + self.ready = False # True after a load whose checks all passed + self.failures = 0 + self.last_load: dict[str, Any] | None = None # reload_mode, pause_s, checks + + @property + def scheme(self) -> str: + return self.engine.identity.quant_scheme + + def _expected(self, release: PolicyRelease | None) -> int | None: + recorded = None if release is None else (release.derived_from or {}).get("expert_matrices") + if recorded is not None: + return int(recorded) + return None if self.scheme == BF16 else self.checks.expected_expert_matrices + + def start(self, directory: Path | None, version: int, release: PolicyRelease | None = None) -> None: + """Start the engine (base weights when ``directory`` is None) and check it.""" + self._serve(release, directory, int(version), first=True) + + def __call__(self, release: PolicyRelease, directory: Path) -> None: + scheme = release_scheme(release) + if scheme != self.scheme: + raise ReleaseNotLoaded( + f"release r{release.seq}-v{release.version} has scheme {scheme}; this engine " + f"serves {self.scheme}; keeping v{self.engine.served_version()}", + reason="scheme_mismatch") + self._serve(release, Path(directory), int(release.version), first=False) + + def _fail(self, version: int, step: str, detail: Any) -> None: + self.ready = False + self.failures += 1 + self.emit(LOAD_CHECK_FAILED_EVENT, policy_version=version, step=step, detail=detail, + failures=self.failures, max_consecutive_failures=self.max_consecutive_failures) + if self.failures >= self.max_consecutive_failures: + raise SamplerLoadFatal(f"post-load checks failed {self.failures} times in a row " + f"(last: {step} for v{version}): {str(detail)[:300]}") + raise ReleaseNotLoaded(f"v{version}: {step} failed: {str(detail)[:300]}", + reason=f"check_failed:{step}") + + def _serve(self, release: PolicyRelease | None, directory: Path | None, version: int, + *, first: bool) -> None: + expected = self._expected(release) + if directory is not None and expected is not None: + in_release = count_release_experts(directory) + if in_release is not None and in_release != expected: + self._fail(version, "release_expert_count", + {"in_release": in_release, "expected": expected}) + restart = first or self.failures > 0 # after a failed check: restart, never patch in place + started = self.clock() + try: + report = (self.engine.start(directory, version) if restart + else self.engine.load(directory, version)) + except SamplerEngineError as error: + self._fail(version, "load", f"{type(error).__name__}: {error}") + results = self.checks.run(self.engine, expected_experts=expected) + served = self.engine.served_version() + results.append({"name": "served_version", "passed": served == version, + "served": served, "expected": version}) + failed = [r for r in results if not r["passed"]] + if failed: + self._fail(version, failed[0]["name"], failed) + self.ready = True + self.failures = 0 + self.last_load = {"reload_mode": report.reload_mode, "pause_s": round(report.pause_s, 3), + "load_and_check_s": round(self.clock() - started, 3), + "scheme": self.scheme, "checks": results, **dict(report.detail)} + self.emit(LOAD_CHECKED_EVENT, policy_version=version, **self.last_load) + + def loaded_fields(self) -> dict[str, Any]: + """Extra fields of ``rl_sampler_policy_loaded`` (``SamplerLoop.loaded_fields``).""" + if self.last_load is None: + return {} + return {k: self.last_load[k] for k in ("reload_mode", "pause_s", "scheme")} + + +def check_sampling_fits(sampling: SamplingParams, engine: SamplerEngine, *, + max_prompt_tokens: int) -> None: + """Refuse a run whose ``max_tokens`` this engine cannot serve: every source + must truncate at the same length (A-INF 5.1), so an engine with a lower + limit may not silently cut earlier.""" + need = int(sampling.max_tokens) + int(max_prompt_tokens) + if need > int(engine.max_model_len): + raise SamplerEngineError( + f"{engine.identity.engine} serves max_model_len {engine.max_model_len}, but the run " + f"needs max_tokens {sampling.max_tokens} + prompt {max_prompt_tokens} = {need}; raise " + "the engine limit or lower max_tokens for every source of the run") + diff --git a/yeto/rl/engine/sampler_engine.py b/yeto/rl/engine/sampler_engine.py index 431a2fc25..dd663471b 100644 --- a/yeto/rl/engine/sampler_engine.py +++ b/yeto/rl/engine/sampler_engine.py @@ -119,12 +119,15 @@ def replace(self, **changes: Any) -> "SamplingParams": @dataclass(frozen=True) class Completion: """One completion: response tokens, the log-probability of each sampled - token (same length), the response text and why it ended.""" + token (same length), the response text and why it ended. + ``prompt_token_ids``: the prompt as the engine tokenized it (empty when the + engine does not return it; the trainer then tokenizes the prompt text).""" token_ids: tuple[int, ...] logprobs: tuple[float, ...] text: str finish_reason: str + prompt_token_ids: tuple[int, ...] = () def __post_init__(self) -> None: if len(self.token_ids) != len(self.logprobs): diff --git a/yeto/rl/engine/sampler_loop.py b/yeto/rl/engine/sampler_loop.py index 117696570..a4cca960a 100644 --- a/yeto/rl/engine/sampler_loop.py +++ b/yeto/rl/engine/sampler_loop.py @@ -12,6 +12,13 @@ With ``require_release`` (a LoRA sampler: its engines serve no adapter until the first release is loaded) step 2 also waits until a release is loaded. +A load hook may refuse a release without stopping the island +(:class:`~.policy_release.ReleaseNotLoaded`: a scheme the engine does not serve, +a failed post-load check). The loop writes ``rl_sampler_release_not_loaded`` and +goes on; with ``ready`` (the engine sampler: False while the engine holds +weights that failed a check) it waits instead of generating, so no sample of +unchecked weights reaches the buffer. + After the last round it marks the data end, so the trainer stops after reading the rest. The loop is engine-neutral: ``generate(rollout_id)`` runs one round. """ @@ -22,7 +29,7 @@ from collections.abc import Callable from typing import Any -from .policy_release import PolicyFollower +from .policy_release import PolicyFollower, ReleaseNotLoaded DEFAULT_POLL_S = 5.0 @@ -36,7 +43,8 @@ def __init__(self, *, follower: PolicyFollower, buffer: Any, min_round_s: float | None = None, clock: Callable[[], float] = time.monotonic, require_release: bool = False, - loaded_fields: Callable[[], dict[str, Any]] = dict) -> None: + loaded_fields: Callable[[], dict[str, Any]] = dict, + ready: Callable[[], bool] = lambda: True) -> None: if int(max_backlog_groups) < 1: raise ValueError("max_backlog_groups must be >= 1") self.follower = follower @@ -52,6 +60,8 @@ def __init__(self, *, follower: PolicyFollower, buffer: Any, self.clock = clock self.require_release = bool(require_release) self.loaded_fields = loaded_fields # extra fields of rl_sampler_policy_loaded + self.ready = ready # False: the engine may not generate (failed post-load check) + self._not_loaded: tuple[int, str] | None = None # (seq, reason) last reported def backlog(self) -> int: """Groups in the buffer after the trainer's read position of its last release @@ -70,7 +80,16 @@ def backlog(self) -> int: return max(0, len(self.buffer) - read) def _load_newest(self) -> None: - loaded = self.follower.refresh() + try: + loaded = self.follower.refresh() + except ReleaseNotLoaded as error: + seen = self.follower.last_seen + mark = (None if seen is None else seen.seq, error.reason) + if mark != self._not_loaded: # one event per release and cause, not per poll + self._not_loaded = mark + self.emit("rl_sampler_release_not_loaded", seq=mark[0], reason=error.reason, + detail=str(error)[:300], policy_version=self.follower.version) + return if loaded is not None: self.emit("rl_sampler_policy_loaded", seq=loaded.seq, policy_version=loaded.version, policy_hash=loaded.policy_hash, buffer_selected=loaded.buffer_selected, @@ -82,7 +101,7 @@ def _admit(self, rollout_id: int) -> bool: waiting_policy_since = None while True: self._load_newest() - if self.require_release and self.follower.release is None: + if not self.ready() or (self.require_release and self.follower.release is None): if waiting_policy_since is None: waiting_policy_since = time.monotonic() self.emit("rl_sampler_waiting_policy", rollout_id=rollout_id) From dc2aa6148ae52049c1fe247cb264085bc1fdd284 Mon Sep 17 00:00:00 2001 From: michael <78113554+chowhsu@users.noreply.github.com> Date: Mon, 19 Oct 2026 08:21:39 +0000 Subject: [PATCH 09/34] Expect the source field in the unavailable resource sample (S25 N2.2) Co-Authored-By: Claude Opus 5.5 --- tests/test_dashboard_startup_steps.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/tests/test_dashboard_startup_steps.py b/tests/test_dashboard_startup_steps.py index ebe779b69..efc466ce2 100644 --- a/tests/test_dashboard_startup_steps.py +++ b/tests/test_dashboard_startup_steps.py @@ -56,7 +56,8 @@ def no_nvml(): hb = [e for e in events if e["event"] == T.HEARTBEAT_EVENT] assert hb and hb[0]["phase"] == "startup" and hb[0]["startup_step"] == "begin" and hb[0]["rollout_id"] is None res = [e for e in events if e["event"] == T.RESOURCE_EVENT] - assert res == [{"event": T.RESOURCE_EVENT, "available": False, "t": 100.0, "reason": "NoNvml: no gpu here", + assert res == [{"event": T.RESOURCE_EVENT, "available": False, "t": 100.0, "source": "nvml", + "reason": "NoNvml: no gpu here", "phase": "startup", "island_label": "x"}] steps = [e for e in events if e["event"] == T.STARTUP_STEP_EVENT] assert [(e["step"], e["seconds"], e["step_s"]) for e in steps] == [ From bb657d68c305231c83964ea96f1fc9abae2dfa4b Mon Sep 17 00:00:00 2001 From: michael <78113554+chowhsu@users.noreply.github.com> Date: Mon, 19 Oct 2026 09:09:02 +0000 Subject: [PATCH 10/34] Treat engine errors in post-load checks as failed checks (S25 N2.6) A SamplerEngineError raised by a check (fixed prompts, log-prob probe) escaped the loader, so it was not counted and the loader kept ready=True from the previous load while the engine held unchecked weights. The loader now marks the engine not ready before each load and counts such an error as a failed check. The probe base is updated only after every check of the load passed. Co-Authored-By: Claude Opus 5.5 --- tests/test_rl_hetero_sampler_engines.py | 34 +++++++++++++++++++++++++ yeto/rl/engine/loop_guard.py | 2 +- yeto/rl/engine/sampler_checks.py | 17 ++++++++++--- 3 files changed, 49 insertions(+), 4 deletions(-) diff --git a/tests/test_rl_hetero_sampler_engines.py b/tests/test_rl_hetero_sampler_engines.py index 0121c4a80..721a84edc 100644 --- a/tests/test_rl_hetero_sampler_engines.py +++ b/tests/test_rl_hetero_sampler_engines.py @@ -297,6 +297,40 @@ def test_a_failed_check_marks_the_engine_not_ready_and_restarts_next_time(tmp_pa loader(release, releases.resolve(release)) +def test_an_engine_error_inside_a_check_is_a_failed_check(tmp_path): + from yeto.rl.engine.policy_release import ReleaseNotLoaded + from yeto.rl.engine.sampler_engine import SamplerEngineError + + releases = PolicyReleaseStore(CheckpointStore(tmp_path)) + engine = FakeEngine(identity=W4A8, answers=ANSWERS, expert_matrices=4) + events = [] + loader = _loader(engine, events) + loader.start(None, 0) + + def broken_score(prompt, continuation): + raise SamplerEngineError("HTTP 500") + + engine.score = broken_score + release = _w4a8_release(releases, 1) + with pytest.raises(ReleaseNotLoaded, match="checks"): + loader(release, releases.resolve(release)) + assert not loader.ready # the engine now holds v1 weights that were never checked + assert [f["step"] for n, f in events if n == LOAD_CHECK_FAILED_EVENT] == ["checks"] + + +def test_probe_of_a_load_that_failed_another_check_is_not_the_next_base(tmp_path): + releases = PolicyReleaseStore(CheckpointStore(tmp_path)) + engine = FakeEngine(identity=W4A8, answers=ANSWERS, expert_matrices=4) + loader = _loader(engine, [], max_probe_jump=0.5) + loader.start(None, 0) + base = list(loader.checks._last_probe) + engine.answers = {} # fixed prompts fail; the probe itself passes + release = _w4a8_release(releases, 1) + with pytest.raises(Exception, match="fixed_prompts"): + loader(release, releases.resolve(release)) + assert loader.checks._last_probe == base + + def test_engine_that_reports_no_count_needs_memory_or_an_explicit_waiver(tmp_path): engine = FakeEngine(identity=W4A8, answers=ANSWERS, expert_matrices=None) checks = PostLoadChecks(expected_expert_matrices=4) diff --git a/yeto/rl/engine/loop_guard.py b/yeto/rl/engine/loop_guard.py index a10e03b66..3c6228d80 100644 --- a/yeto/rl/engine/loop_guard.py +++ b/yeto/rl/engine/loop_guard.py @@ -48,7 +48,7 @@ def apply(self, completion: Completion) -> Completion: if keep is None or keep >= len(completion.token_ids): return completion # The text of the kept tokens is not known without the tokenizer; the - # sample's text is the engine text marked as cut (training reads tokens). + # sample keeps the full engine text; finish_reason marks the cut (training reads tokens). return Completion(token_ids=completion.token_ids[:keep], logprobs=completion.logprobs[:keep], text=completion.text, finish_reason=FINISH_LOOP_ABORT, prompt_token_ids=completion.prompt_token_ids) diff --git a/yeto/rl/engine/sampler_checks.py b/yeto/rl/engine/sampler_checks.py index c29e74db0..498dc8723 100644 --- a/yeto/rl/engine/sampler_checks.py +++ b/yeto/rl/engine/sampler_checks.py @@ -151,6 +151,7 @@ class PostLoadChecks: max_probe_jump: float | None = None reference_probe: Sequence[float] | None = None # round-0 per-pair mean log-probs _last_probe: list[float] | None = field(default=None, init=False, repr=False) + _pending_probe: list[float] | None = field(default=None, init=False, repr=False) def pairs(self) -> list[tuple[str, str]]: if self.probe_pairs is not None: @@ -218,10 +219,15 @@ def logprob_probe(self, engine: SamplerEngine) -> dict[str, Any]: out["mean_abs_jump"] = round(jump, 6) out["max_probe_jump"] = self.max_probe_jump out["passed"] = jump is None or self.max_probe_jump is None or jump <= self.max_probe_jump - if out["passed"]: - self._last_probe = means + self._pending_probe = means if out["passed"] else None return out + def accept(self) -> None: + """Make the last probe the next comparison base; only after every check of the load passed + (a probe of weights that failed another check must not become the base).""" + if self._pending_probe is not None: + self._last_probe, self._pending_probe = self._pending_probe, None + def run(self, engine: SamplerEngine, *, expected_experts: int | None) -> list[dict[str, Any]]: results = [self.expert_count(engine, expected_experts)] results.append(self.fixed_prompt_answers(engine)) @@ -293,19 +299,24 @@ def _serve(self, release: PolicyRelease | None, directory: Path | None, version: self._fail(version, "release_expert_count", {"in_release": in_release, "expected": expected}) restart = first or self.failures > 0 # after a failed check: restart, never patch in place + self.ready = False # the engine holds unchecked weights from here until the checks pass started = self.clock() try: report = (self.engine.start(directory, version) if restart else self.engine.load(directory, version)) except SamplerEngineError as error: self._fail(version, "load", f"{type(error).__name__}: {error}") - results = self.checks.run(self.engine, expected_experts=expected) + try: + results = self.checks.run(self.engine, expected_experts=expected) + except SamplerEngineError as error: # an engine call of a check failed: a failed check + self._fail(version, "checks", f"{type(error).__name__}: {error}") served = self.engine.served_version() results.append({"name": "served_version", "passed": served == version, "served": served, "expected": version}) failed = [r for r in results if not r["passed"]] if failed: self._fail(version, failed[0]["name"], failed) + self.checks.accept() self.ready = True self.failures = 0 self.last_load = {"reload_mode": report.reload_mode, "pause_s": round(report.pause_s, 3), From 7aabbe4b94978e7432db1cebcb22edc334a824cc Mon Sep 17 00:00:00 2001 From: michael <78113554+chowhsu@users.noreply.github.com> Date: Mon, 19 Oct 2026 08:33:52 +0000 Subject: [PATCH 11/34] Add the records proxy and GLM chat template alignment check (S25 N10.5/N10.6) - SamplerEngine port: optional TokenSamplerEngine.generate_tokens (token-id prompts, stop token ids); SGLang (input_ids), vLLM (id-list prompts) and the fake engine implement it. - RecordingBackend: a gateway SessionBackend that samples the engine directly and keeps one append-only token stream per session, with loss mask, log-probs, policy version and provenance per turn; SGLang-style meta_info in the reply; per-turn render drift check; optional CompactionRL context compaction reusing yeto.rl.compaction. - RecordsProxy: the existing Gateway (Responses/Chat translation, chains, signed sampling) over the records backend, with an aiohttp shell (Responses SSE reuses the Codex bridge encoder). - chat_templates: template registry (glm5, qwen3), output parsers and the token alignment check against apply_chat_template. GLM5.3-Flash tokenizer is not available; glm5 is checked with the GLM-4.7-Flash tokenizer. Co-Authored-By: Claude Opus 5.5 --- tests/test_records_proxy.py | 408 ++++++++++++++++++ yeto/rl/adapters/sampler_engines/fake.py | 33 +- .../adapters/sampler_engines/sglang_server.py | 25 +- .../adapters/sampler_engines/vllm_server.py | 16 +- yeto/rl/engine/sampler_engine.py | 16 + yeto/rl/harness/chat_templates.py | 262 +++++++++++ yeto/rl/harness/gateway/proxy.py | 148 +++++++ yeto/rl/harness/gateway/records.py | 329 ++++++++++++++ 8 files changed, 1233 insertions(+), 4 deletions(-) create mode 100644 tests/test_records_proxy.py create mode 100644 yeto/rl/harness/chat_templates.py create mode 100644 yeto/rl/harness/gateway/proxy.py create mode 100644 yeto/rl/harness/gateway/records.py diff --git a/tests/test_records_proxy.py b/tests/test_records_proxy.py new file mode 100644 index 000000000..8c39833bf --- /dev/null +++ b/tests/test_records_proxy.py @@ -0,0 +1,408 @@ +"""S25 N10.5 records proxy and N10.6 template token alignment (CPU). + +The model is a scripted :class:`FakeEngine`: each reply is the text the chat +template itself renders for a chosen assistant message, so a correct proxy +shows no render drift. Tokenizers come from the local HF cache (skip if absent): +Qwen3-0.6B (synced to the remote test box) and GLM-4.7-Flash (stand-in for the +GLM5.3-Flash tokenizer, which is not available: GLM5.3 itself is not verified). +""" + +from __future__ import annotations + +import asyncio +import copy +import json +from typing import Any + +import pytest + +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.gateway import Gateway, GatewayConfig +from yeto.rl.harness.gateway.proxy import RecordsProxy, chat_tools +from yeto.rl.harness.gateway.records import RecordingBackend, RecordsMismatch + +SIGNED = {"temperature": 1.0, "top_p": 1.0, "max_tokens": 256} +TOOLS = [{"type": "function", "name": "shell", "description": "run a command", + "parameters": {"type": "object", "properties": {"cmd": {"type": "string"}}, "required": ["cmd"]}}] + + +def _template(name: str) -> ChatTemplate: + try: + return ChatTemplate.load(name) + except (OSError, ValueError, ImportError) as exc: # tokenizer not in the local cache + pytest.skip(f"{name} tokenizer not available: {exc}") + + +@pytest.fixture(scope="module") +def qwen() -> ChatTemplate: + return _template("qwen3") + + +@pytest.fixture(scope="module") +def glm() -> ChatTemplate: + return _template("glm5") + + +def call(call_id: str, cmd: str) -> dict[str, Any]: + return {"id": call_id, "type": "function", "function": {"name": "shell", "arguments": json.dumps({"cmd": cmd})}} + + +def assistant(reasoning: str, content: str = "", calls: list[dict[str, Any]] | None = None) -> dict[str, Any]: + message = {"role": "assistant", "content": content, "reasoning_content": reasoning} + if calls: + message["tool_calls"] = calls + return message + + +class ScriptedModel: + """``token_reply`` for FakeEngine: replies with the template's own rendering + of the next scripted assistant message, plus the stop token.""" + + def __init__(self, template: ChatTemplate, replies: list[dict[str, Any]], *, return_stop: bool = True) -> None: + self.template = template + self.replies = list(replies) + self.return_stop = return_stop + + def __call__(self, prompt: tuple[int, ...], index: int) -> list[int]: + message = self.replies.pop(0) + reply = model_text(self.template, message) + ids = self.template.encode(reply) + return ids + ([self.template.encode(turn_end(self.template, message))[0]] if self.return_stop else []) + + +def turn_end(template: ChatTemplate, message: dict[str, Any]) -> str: + """The stop token the model emits after ``message``: GLM opens the next turn + (tool output or user), Qwen3 closes its own turn.""" + if template.spec.name == "glm5": + return "<|observation|>" if message.get("tool_calls") else "<|user|>" + return "<|im_end|>" + + +def model_text(template: ChatTemplate, message: dict[str, Any]) -> str: + """The text of ``message`` as the template renders it after a generation + prompt, without the turn end.""" + rendered = template.render([{"role": "user", "content": "x"}, message], add_generation_prompt=False) + opener = template.render([{"role": "user", "content": "x"}], add_generation_prompt=True) + body = rendered[len(opener):] + for stop in template.spec.stop_tokens: + body = body.split(stop, 1)[0] + return body.rstrip("\n") if template.spec.name == "glm5" else body + + +def engine_for(model: ScriptedModel, **kw: Any) -> FakeEngine: + engine = FakeEngine(token_reply=model, **kw) + engine.start(None, 3) + return engine + + +def responses_body(items: list[dict[str, Any]], **kw: Any) -> dict[str, Any]: + return {"model": "policy", "store": False, "instructions": "You are a coding agent.", "tools": TOOLS, + "input": items, **kw} + + +def user_item(text: str) -> dict[str, Any]: + return {"type": "message", "role": "user", "content": [{"type": "input_text", "text": text}]} + + +def echo(output: list[dict[str, Any]]) -> list[dict[str, Any]]: + """What Codex sends back: the output items as input items.""" + items = [] + for item in output: + item = {k: v for k, v in item.items() if k not in ("id", "status")} + if item["type"] == "message": + item["content"] = [{"type": "output_text", "text": c["text"]} for c in item["content"]] + items.append(item) + return items + + +def tool_output(call_id: str, text: str) -> dict[str, Any]: + return {"type": "function_call_output", "call_id": call_id, "output": text} + + +SCRIPT = [assistant("look first", calls=[call("x", "ls")]), + assistant("now read", calls=[call("x", "cat a.txt")]), + assistant("done", content="The file says hello.")] + + +def run_codex(proxy: Any, trajectory_id: str, *, turns: int = 3) -> list[dict[str, Any]]: + """Drive a Codex-like Responses loop: echo output, answer each tool call.""" + items = [user_item("read a.txt")] + outputs = [] + for turn in range(turns): + response = asyncio.run(proxy.handle("responses", responses_body(copy.deepcopy(items)), trajectory_id=trajectory_id)) + outputs.append(response) + items += echo(response["output"]) + for out in response["output"]: + if out["type"] == "function_call": + items.append(tool_output(out["call_id"], f"result {turn}")) + return outputs + + +# ---------------------------------------------------------------------- N10.5 records +@pytest.mark.parametrize("name", ["qwen3", "glm5"]) +def test_codex_responses_loop_records_one_aligned_stream(name, request): + template = request.getfixturevalue("qwen" if name == "qwen3" else "glm") + 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"] + assert outputs[-1]["output"][-1]["content"][0]["text"] == "The file says hello." + + records = proxy.records("t1") + assert records["gateway"]["chains_total"] == 1 and records["gateway"]["invalidated"] is None + (session,) = records["sessions"] + (sample,) = session["samples"] + ids, mask, logprobs = sample["token_ids"], sample["loss_mask"], sample["logprobs"] + assert len(ids) == len(mask) == len(logprobs) + turns = sample["turns"] + assert [t["kind"] for t in turns] == ["action"] * 3 + # Each completion sits at mask 1 right after its prompt ids, with its log-probs. + pos = 0 + for turn in turns: + pos += len(turn["stop_added"]) + len(turn["prompt_added"]) + n = len(turn["completion"]) + assert ids[pos:pos + n] == list(turn["completion"]) and mask[pos:pos + n] == [1] * n + assert logprobs[pos:pos + n] == list(turn["logprobs"]) + pos += n + assert pos == len(ids) and sum(mask) == sum(len(t["completion"]) for t in turns) + # The prompts the engine got are prefixes of the final stream (append-only). + prompts = [p for p, *_ in engine.token_requests] + assert all(list(p[0]) == ids[:len(p[0])] for p in prompts) + # N10.6: the stream equals a fresh tokenization of the rendering on every turn. + assert records["counters"].get("render_drift_turns", 0) == 0 + assert all(t["render_diff"] is None for t in turns) + # Source tag and policy version on every turn. + assert {t["policy_version"] for t in turns} == {3} + assert turns[0]["provenance"]["engine"] == "fake" and turns[0]["provenance"]["sampler_id"] == "records-proxy" + assert turns[0]["provenance"]["sampling_digest"] == SamplingParams(max_tokens=256).digest() + + +def test_proxy_matches_the_existing_gateway_on_the_same_conversation(qwen): + """Pair test: the gateway with the scripted session backend (TB2 path) and + the records proxy give the same Responses output and the same chains.""" + + class SessionBackend: + def __init__(self, replies): + self.replies = list(replies) + + async def create_session(self, trajectory_id): + return "s0" + + async def chat(self, session_id, messages, sampling): + message = self.replies.pop(0) + return {"choices": [{"index": 0, "message": message, "finish_reason": "stop"}], + "usage": {"prompt_tokens": 0, "completion_tokens": 0, "total_tokens": 0}} + + proxy = RecordsProxy(engine_for(ScriptedModel(qwen, SCRIPT)), qwen, model="policy", signed_sampling=SIGNED) + ours = run_codex(proxy, "t") + # The reference backend replies with the messages the proxy parsed (same call ids). + parsed = [qwen.parse(qwen.decode(qwen.encode(model_text(qwen, m)))) for m in SCRIPT] + reference = Gateway(GatewayConfig(model="policy", signed_sampling=SIGNED, + keeps_history_reasoning=qwen.spec.keeps_history_reasoning), + SessionBackend(parsed)) + theirs = run_codex(reference, "t") + + def strip(response): + return [{k: v for k, v in item.items() if k != "id"} for item in response["output"]] + + assert [strip(r) for r in ours] == [strip(r) for r in theirs] + a, b = proxy.gateway.snapshot("t"), reference.snapshot("t") + for snap in (a, b): + for chain in snap["chains"]: + chain.pop("session_id") + assert a == b + + +def test_meta_info_is_the_sglang_shape_the_miles_session_server_reads(qwen): + engine = engine_for(ScriptedModel(qwen, SCRIPT[:1])) + backend = RecordingBackend(engine, qwen) + + async def go(): + sid = await backend.create_session("t") + return await backend.chat(sid, [{"role": "user", "content": "hi"}], SIGNED) + + completion = asyncio.run(go()) + meta = completion["choices"][0]["meta_info"] + pairs = meta["output_token_logprobs"] + assert len(pairs) == meta["completion_tokens"] == completion["usage"]["completion_tokens"] + assert all(isinstance(lp, float) and isinstance(tid, int) for lp, tid, _ in pairs) + assert pairs[-1][1] == qwen.encode("<|im_end|>")[0] + + +def test_stop_token_not_returned_by_the_engine_is_added_with_mask_zero(glm): + engine = engine_for(ScriptedModel(glm, SCRIPT, return_stop=False)) + proxy = RecordsProxy(engine, glm, model="policy", signed_sampling=SIGNED) + run_codex(proxy, "t") + records = proxy.records("t") + sample = records["sessions"][0]["samples"][0] + added = [t["stop_added"] for t in sample["turns"]] + assert [tuple(a) for a in added] == [(glm.encode("<|observation|>")[0],)] * 2 + [()] + assert records["counters"]["stop_added"] == 2 + assert records["counters"].get("render_drift_turns", 0) == 0 + for turn in sample["turns"]: + assert turn["completion"][-1] not in glm.stop_token_ids + + +def test_length_truncated_turn_and_changed_tools_fail_closed(qwen): + engine = engine_for(ScriptedModel(qwen, SCRIPT * 2)) + proxy = RecordsProxy(engine, qwen, model="policy", signed_sampling={**SIGNED, "max_tokens": 3}) + first = asyncio.run(proxy.handle("responses", responses_body([user_item("go")]), trajectory_id="t")) + assert first["status"] == "completed" + items = [user_item("go")] + echo(first["output"]) + [user_item("more")] + with pytest.raises(RecordsMismatch, match="length-truncated"): + asyncio.run(proxy.handle("responses", responses_body(items), trajectory_id="t")) + with pytest.raises(RecordsMismatch, match="tool schemas changed"): + asyncio.run(proxy.handle("responses", responses_body([user_item("go")], tools=[]), trajectory_id="t")) + + +def test_backend_refuses_an_engine_without_token_input(qwen): + class TextOnly: + pass + + assert isinstance(FakeEngine(), TokenSamplerEngine) + with pytest.raises(TypeError, match="generate_tokens"): + RecordingBackend(TextOnly(), qwen) # type: ignore[arg-type] + + +def _compaction_proxy(template: ChatTemplate) -> tuple[RecordsProxy, Any]: + summary = assistant("check", content="1. listed files") + dry = engine_for(ScriptedModel(template, SCRIPT)) + run_codex(RecordsProxy(dry, template, model="policy", signed_sampling=SIGNED), "dry") + third_prompt = len(dry.token_requests[2][0][0]) + # Budget so the third request (after two tool steps) triggers, not the second. + cfg = CompactionConfig(context_budget=third_prompt, t_comp=1, keep_recent_steps=1, max_compactions=1) + engine = engine_for(ScriptedModel(template, [SCRIPT[0], SCRIPT[1], summary, SCRIPT[2]])) + return RecordsProxy(engine, template, model="policy", signed_sampling=SIGNED, compaction=cfg), engine + + +def test_compaction_summarizes_and_starts_a_new_segment(glm): + """GLM with clear_thinking=False keeps every turn's reasoning, so the + summary request appends to the same stream.""" + proxy, _ = _compaction_proxy(glm) + outputs = run_codex(proxy, "t") + assert outputs[-1]["output"][-1]["content"][0]["text"] == "The file says hello." + records = proxy.records("t") + assert records["counters"]["compactions"] == 1 + first, second = records["sessions"][0]["samples"] + assert [t["kind"] for t in first["turns"]] == ["action", "action", "summary"] + assert first["summary_ok"] is True and second["kept_steps"] in (0, 1) + assert [t["kind"] for t in second["turns"]] == ["action"] + second_prompt = glm.decode(second["token_ids"]) + assert "1. listed files" in second_prompt and "continues from an earlier one" in second_prompt + assert "result 0" not in second_prompt # the first step was summarized away + + +def test_compaction_on_a_template_that_drops_history_reasoning_fails_closed(qwen): + """Qwen3 drops earlier reasoning once a user message follows, so the + summary request cannot extend the recorded stream token-exactly.""" + proxy, _ = _compaction_proxy(qwen) + with pytest.raises(RecordsMismatch, match="rendering of the earlier turns changed"): + run_codex(proxy, "t") + + +# ---------------------------------------------------------------------- N10.6 alignment +CONVERSATION = [ + {"role": "system", "content": "You are a coding agent."}, + {"role": "user", "content": "列出文件 and read them"}, + assistant("plan", calls=[call("c1", "ls -la")]), + {"role": "tool", "tool_call_id": "c1", "content": "a.txt\nb.txt"}, + assistant("next", calls=[call("c2", "cat a.txt"), call("c3", "cat b.txt")]), + {"role": "tool", "tool_call_id": "c2", "content": "hello"}, + {"role": "tool", "tool_call_id": "c3", "content": "{\"k\": 1}"}, + assistant("done", content="Both files read."), + {"role": "user", "content": "thanks"}, +] + + +@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") + tools = chat_tools("responses", {"tools": TOOLS}) + for end in range(2, len(CONVERSATION) + 1): + report = check_alignment(template, CONVERSATION[:end], tools=tools) + assert report.aligned, (end, report) + # A text prompt tokenized with special tokens would not double the template's BOS. + assert report.engine_text_extra == 0 + + +def test_glm_parser_round_trips_its_own_rendering(glm): + for message in (SCRIPT[0], assistant("r", content="text", calls=[call("c", "echo 1")]), + assistant("r", calls=[{"id": "c", "type": "function", + "function": {"name": "f", "arguments": json.dumps({"n": 3, "s": "x"})}}])): + parsed = parse_glm(model_text(glm, message)) + assert parsed["reasoning_content"] == message["reasoning_content"] + assert parsed["content"] == message["content"] + assert [json.loads(c["function"]["arguments"]) for c in parsed.get("tool_calls", [])] == \ + [json.loads(c["function"]["arguments"]) for c in message.get("tool_calls", [])] + + +def test_parsers_fail_closed_on_malformed_tool_calls(): + with pytest.raises(TemplateError): + parse_qwen3("\nnot json\n") + with pytest.raises(TemplateError): + parse_glm("ra1") + + +# ---------------------------------------------------------------------- adapters +def test_sglang_and_vllm_adapters_send_token_id_prompts(tmp_path): + from test_rl_hetero_sampler_engines import _config, _Http, _Process, _vllm + + from yeto.rl.adapters.sampler_engines.sglang_server import SGLangServerEngine + from yeto.rl.adapters.sampler_engines.vllm_server import VllmServerEngine + + def sglang(body): + return [{"text": "x", "meta_info": {"output_token_logprobs": [[-0.1, 7, None], [-0.2, 2, None]], + "finish_reason": {"type": "stop"}}} for _ in body["input_ids"]] + + http = _Http({("GET", "/health"): {}, ("POST", "/generate"): sglang}) + engine = SGLangServerEngine(_config(tmp_path), spawn=lambda a, e: _Process(), http=http, sleep=lambda s: None) + engine.start(None, 0) + (group,) = engine.generate_tokens([[1, 2, 3]], SamplingParams(max_tokens=4, n=2), stop_token_ids=[2]) + sent = [b for m, p, b in http.calls if p == "/generate"][-1] + assert sent["input_ids"] == [[1, 2, 3], [1, 2, 3]] and sent["sampling_params"]["stop_token_ids"] == [2] + assert group[0].token_ids == (7, 2) and group[0].prompt_token_ids == (1, 2, 3) + + def completions(body): + return {"choices": [{"index": i, "text": "t", "finish_reason": "stop", + "logprobs": {"tokens": ["token_id:4", "token_id:9"], "token_logprobs": [-0.3, -0.4]}} + for i in range(len(body["prompt"]) * body["n"])]} + + http = _Http({("GET", "/health"): {}, ("POST", "/v1/completions"): completions}) + engine = VllmServerEngine(_vllm(tmp_path), spawn=lambda a, e: _Process(), http=http, sleep=lambda s: None) + engine.start(None, 0) + (group,) = engine.generate_tokens([[5, 6]], SamplingParams(max_tokens=4), stop_token_ids=[9]) + sent = [b for m, p, b in http.calls if p == "/v1/completions"][-1] + assert sent["prompt"] == [[5, 6]] and sent["stop_token_ids"] == [9] + assert group[0].token_ids == (4, 9) and group[0].prompt_token_ids == (5, 6) + + +# ---------------------------------------------------------------------- HTTP shell +def test_http_shell_serves_codex_sse_and_records(qwen): + from aiohttp.test_utils import TestClient, TestServer + + proxy = RecordsProxy(engine_for(ScriptedModel(qwen, SCRIPT)), qwen, model="policy", signed_sampling=SIGNED) + + async def go(): + async with TestClient(TestServer(proxy.app())) as client: + created = await (await client.post("/trajectories")).json() + tid = created["trajectory_id"] + assert created["base_url"].endswith(f"/t/{tid}/v1") + reply = await client.post(f"/t/{tid}/v1/responses", json=responses_body([user_item("go")], stream=True)) + assert reply.status == 200 and reply.headers["Content-Type"].startswith("text/event-stream") + text = await reply.text() + assert "event: response.completed" in text and text.endswith("data: [DONE]\n\n") + bad = await client.post(f"/t/{tid}/v1/responses", json=responses_body([user_item("go")], store=True)) + assert bad.status == 400 + override = await client.post(f"/t/{tid}/v1/chat/completions", + json={"model": "policy", "temperature": 0.1, + "tools": chat_tools("responses", {"tools": TOOLS}), + "messages": [{"role": "user", "content": "x"}]}) + assert override.status == 400 + records = await (await client.get(f"/t/{tid}/records")).json() + assert records["sessions"][0]["samples"][0]["loss_mask"].count(1) > 0 + assert (await client.post(f"/t/{tid}/close")).status == 200 + + asyncio.run(go()) diff --git a/yeto/rl/adapters/sampler_engines/fake.py b/yeto/rl/adapters/sampler_engines/fake.py index 927d0fb6e..fb406cd90 100644 --- a/yeto/rl/adapters/sampler_engines/fake.py +++ b/yeto/rl/adapters/sampler_engines/fake.py @@ -43,7 +43,8 @@ def __init__(self, *, identity: EngineIdentity = FAKE_IDENTITY, max_model_len: i in_place: bool = True, expert_matrices: int | None = 36_288, weight_gib: float | None = None, answers: Mapping[str, str] | None = None, tokens_for: Callable[[str, int], int] = lambda prompt, i: 4 + i, - token_ids_for: Callable[[str, int, int], list[int]] | None = None) -> None: + token_ids_for: Callable[[str, int, int], list[int]] | None = None, + token_reply: Callable[[tuple[int, ...], int], Sequence[int]] | None = None) -> None: self._identity = identity self._max_model_len = int(max_model_len) self.in_place = bool(in_place) @@ -52,6 +53,8 @@ def __init__(self, *, identity: EngineIdentity = FAKE_IDENTITY, max_model_len: i self.answers = dict(answers or {}) self.tokens_for = tokens_for self.token_ids_for = token_ids_for + self.token_reply = token_reply # generate_tokens: (prompt ids, sample index) -> reply ids + self.token_requests: list[tuple[tuple[tuple[int, ...], ...], SamplingParams, tuple[int, ...]]] = [] self.probe_shift = 0.0 # added to every score log-prob (a "silently wrong" load) self.fail_load: str | None = None # raise this message on the next load self._version: int | None = None @@ -111,6 +114,34 @@ def generate(self, prompts: Sequence[str], sampling: SamplingParams) -> list[lis out.append(group) return out + def generate_tokens(self, prompts: Sequence[Sequence[int]], sampling: SamplingParams, *, + stop_token_ids: Sequence[int] = ()) -> list[list[Completion]]: + """Token-id prompts. The reply is ``token_reply(prompt, i)`` cut at the + first stop token (kept) and at ``max_tokens``; text is left empty, so a + caller must decode the ids itself (as the records proxy does).""" + if self._version is None or self.stopped: + raise SamplerEngineError("fake engine is not serving") + stops = {int(t) for t in stop_token_ids} + keyed = [tuple(int(t) for t in p) for p in prompts] + self.token_requests.append((tuple(keyed), sampling, tuple(sorted(stops)))) + out = [] + for prompt in keyed: + group = [] + for i in range(int(sampling.n)): + reply = (list(self.token_reply(prompt, i)) if self.token_reply + else [int(_unit(prompt, i, k) * 1000) for k in range(int(self.tokens_for(str(prompt), i)))]) + ids: list[int] = [] + for token in reply[: sampling.max_tokens]: + ids.append(int(token)) + if token in stops: + break + finish = FINISH_STOP if (ids and ids[-1] in stops) or len(ids) < sampling.max_tokens else FINISH_LENGTH + logprobs = tuple(-_unit(prompt, i, k, self._version) for k in range(len(ids))) + group.append(Completion(token_ids=tuple(ids), logprobs=logprobs, text="", + finish_reason=finish, prompt_token_ids=prompt)) + out.append(group) + return out + def score(self, prompt: str, continuation: str) -> list[float]: if self._version is None: raise SamplerEngineError("fake engine is not serving") diff --git a/yeto/rl/adapters/sampler_engines/sglang_server.py b/yeto/rl/adapters/sampler_engines/sglang_server.py index ed879083e..1be20274e 100644 --- a/yeto/rl/adapters/sampler_engines/sglang_server.py +++ b/yeto/rl/adapters/sampler_engines/sglang_server.py @@ -162,9 +162,30 @@ def _generate_raw(self, texts: Sequence[str], sampling: SamplingParams) -> list[ def generate(self, prompts: Sequence[str], sampling: SamplingParams) -> list[list[Completion]]: n = int(sampling.n) texts = [p for p in prompts for _ in range(n)] # one request per sample: order is explicit - items = self._generate_raw(texts, sampling.replace(n=1)) + return self._completions(self._generate_raw(texts, sampling.replace(n=1)), len(prompts), n) + + def generate_tokens(self, prompts: Sequence[Sequence[int]], sampling: SamplingParams, *, + stop_token_ids: Sequence[int] = ()) -> list[list[Completion]]: + """Token-id prompts (``input_ids``), for the records proxy (S25 N10.5).""" + n = int(sampling.n) + inputs = [[int(t) for t in p] for p in prompts for _ in range(n)] + params = self._params(sampling.replace(n=1)) + if stop_token_ids: + params["stop_token_ids"] = [int(t) for t in stop_token_ids] + body = self.call("POST", "/generate", {"input_ids": inputs, "sampling_params": params, + "return_logprob": True}) + items = body if isinstance(body, list) else [body] + if len(items) != len(inputs): + raise SamplerEngineError(f"sglang returned {len(items)} outputs for {len(inputs)} inputs") + completions = self._completions(items, len(prompts), n) + return [[c if c.prompt_token_ids else Completion(c.token_ids, c.logprobs, c.text, c.finish_reason, + prompt_token_ids=tuple(inputs[i * n])) + for c in group] for i, group in enumerate(completions)] + + @staticmethod + def _completions(items: Sequence[dict[str, Any]], prompts: int, n: int) -> list[list[Completion]]: out: list[list[Completion]] = [] - for i in range(len(prompts)): + for i in range(prompts): group = [] for item in items[i * n:(i + 1) * n]: meta = item.get("meta_info") or {} diff --git a/yeto/rl/adapters/sampler_engines/vllm_server.py b/yeto/rl/adapters/sampler_engines/vllm_server.py index af5f9067e..f3ebcd339 100644 --- a/yeto/rl/adapters/sampler_engines/vllm_server.py +++ b/yeto/rl/adapters/sampler_engines/vllm_server.py @@ -110,8 +110,22 @@ def _body(self, prompt: Any, sampling: SamplingParams, **extra: Any) -> dict[str return body def generate(self, prompts: Sequence[str], sampling: SamplingParams) -> list[list[Completion]]: + return self._complete(list(prompts), sampling) + + def generate_tokens(self, prompts: Sequence[Sequence[int]], sampling: SamplingParams, *, + stop_token_ids: Sequence[int] = ()) -> list[list[Completion]]: + """Token-id prompts (a list of id lists as ``prompt``), for the records + proxy (S25 N10.5). vLLM does not re-tokenize an id prompt, so the + template's own BOS tokens are not doubled.""" + ids = [[int(t) for t in p] for p in prompts] + extra = {"stop_token_ids": [int(t) for t in stop_token_ids]} if stop_token_ids else {} + out = self._complete(ids, sampling, **extra) + return [[Completion(c.token_ids, c.logprobs, c.text, c.finish_reason, prompt_token_ids=tuple(ids[i])) + for c in group] for i, group in enumerate(out)] + + def _complete(self, prompts: list[Any], sampling: SamplingParams, **extra: Any) -> list[list[Completion]]: n = int(sampling.n) - body = self.call("POST", "/v1/completions", self._body(list(prompts), sampling)) + body = self.call("POST", "/v1/completions", self._body(prompts, sampling, **extra)) choices = sorted(body.get("choices") or [], key=lambda c: int(c["index"])) if len(choices) != len(prompts) * n: raise SamplerEngineError(f"vllm returned {len(choices)} choices for {len(prompts)} x {n}") diff --git a/yeto/rl/engine/sampler_engine.py b/yeto/rl/engine/sampler_engine.py index dd663471b..a476679e0 100644 --- a/yeto/rl/engine/sampler_engine.py +++ b/yeto/rl/engine/sampler_engine.py @@ -215,6 +215,22 @@ def stop(self) -> None: ... +@runtime_checkable +class TokenSamplerEngine(SamplerEngine, Protocol): + """Optional capability: generate from prompt token ids (S25 N10.5). + + The records proxy renders the chat template once and sends the ids, so the + engine never re-tokenizes the prompt and the recorded prompt tokens are the + tokens the engine conditioned on. ``stop_token_ids``: ids that end a turn in + the served chat template; an adapter returns the stop token in + ``token_ids`` when the engine reports it (the proxy handles both cases).""" + + def generate_tokens(self, prompts: Sequence[Sequence[int]], sampling: SamplingParams, *, + stop_token_ids: Sequence[int] = ()) -> list[list[Completion]]: + """``sampling.n`` completions for each token-id prompt, in prompt order.""" + ... + + # ---------------------------------------------------------------------- provenance UNKNOWN = "unknown" diff --git a/yeto/rl/harness/chat_templates.py b/yeto/rl/harness/chat_templates.py new file mode 100644 index 000000000..13ea1883d --- /dev/null +++ b/yeto/rl/harness/chat_templates.py @@ -0,0 +1,262 @@ +"""Chat template registry and the token alignment check (S25 N10.5/N10.6). + +The records proxy (:mod:`yeto.rl.harness.gateway.records`) renders canonical +chat messages (the gateway form, :mod:`.gateway.translate`) with the served +model's own chat template, sends the token ids to the engine and parses the +generated text back into an assistant message. A :class:`TemplateSpec` holds +the model-specific facts for that: tokenizer, template kwargs, the tokens that +end a turn, whether history reasoning is kept, and the output parser. + +Registered templates: + +* ``glm5``: the GLM XML tool-call format (``namek + v``), ``clear_thinking=False`` so every + turn keeps its reasoning (the stream stays append-only). The GLM5.3-Flash + tokenizer is not on this machine: the default tokenizer is the GLM-4.7-Flash + 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. + +:func:`check_alignment` is the N10.6 check: the template rendering, tokenized +the way the proxy sends it, must equal the tokenizer's own +``apply_chat_template`` ids token by token, and it reports how many tokens an +engine would add if it got the text and tokenized it with special tokens. +""" + +from __future__ import annotations + +import hashlib +import json +import re +from collections.abc import Callable, Mapping, Sequence +from dataclasses import dataclass, field +from typing import Any + +Message = dict[str, Any] + + +class TemplateError(ValueError): + """The template cannot render or parse a turn (fail closed).""" + + +# ---------------------------------------------------------------------- parsers +def _call_id(text: str, index: int) -> str: + return f"call_{hashlib.sha256(text.encode('utf-8')).hexdigest()[:12]}_{index}" + + +def _tool_call(call_id: str, name: str, arguments: Mapping[str, Any]) -> Message: + # Key order is kept: the template renders the arguments in this order. + return {"id": call_id, "type": "function", + "function": {"name": name, "arguments": json.dumps(dict(arguments), ensure_ascii=False)}} + + +def _assistant(content: str, reasoning: str | None, calls: list[Message]) -> Message: + message: Message = {"role": "assistant", "content": content} + if reasoning is not None: + message["reasoning_content"] = reasoning + if calls: + message["tool_calls"] = calls + return message + + +def _split_reasoning(text: str, *, opened: bool) -> tuple[str | None, str]: + """``(reasoning, rest)``. ``opened``: the generation prompt already wrote + ````. Without ```` an opened block is all reasoning.""" + if "" in text: + head, rest = text.split("", 1) + return head.split("", 1)[-1].strip("\n"), rest + if opened: + return text.strip("\n"), "" + return None, text + + +_GLM_CALL = re.compile(r"(.*?)", re.S) +_GLM_ARG = re.compile(r"(.*?)\s*(.*?)", re.S) + + +def _glm_value(raw: str) -> Any: + # The template writes strings raw and other values as JSON (``tojson``). + try: + value = json.loads(raw) + except ValueError: + return raw + return raw if isinstance(value, str) else value + + +def parse_glm(text: str) -> Message: + reasoning, rest = _split_reasoning(text, opened=True) + calls = [] + for index, match in enumerate(_GLM_CALL.finditer(rest)): + body = match.group(1) + name = body.split("", 1)[0].strip() + if not name: + raise TemplateError(f"GLM tool call without a name: {body[:80]!r}") + calls.append(_tool_call(_call_id(text, index), name, + {k.strip(): _glm_value(v) for k, v in _GLM_ARG.findall(body)})) + content = rest.split("", 1)[0].strip() + return _assistant(content, reasoning, calls) + + +_QWEN_CALL = re.compile(r"\s*(.*?)\s*", re.S) + + +def parse_qwen3(text: str) -> Message: + reasoning, rest = _split_reasoning(text, opened=False) + calls = [] + for index, match in enumerate(_QWEN_CALL.finditer(rest)): + try: + payload = json.loads(match.group(1)) + name, arguments = payload["name"], payload.get("arguments", {}) + except (ValueError, KeyError, TypeError) as exc: + raise TemplateError(f"Qwen3 tool call is not {{name, arguments}} JSON: {match.group(1)[:80]!r}") from exc + if not isinstance(arguments, Mapping): + raise TemplateError("Qwen3 tool call arguments are not an object") + calls.append(_tool_call(_call_id(text, index), str(name), arguments)) + content = rest.split("", 1)[0].strip("\n") + return _assistant(content, reasoning, calls) + + +# ---------------------------------------------------------------------- registry +@dataclass(frozen=True) +class TemplateSpec: + """Model-specific template facts. ``stop_tokens``: special tokens that end + 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).""" + + name: str + tokenizer: str + parser: Callable[[str], Message] + stop_tokens: tuple[str, ...] + template_kwargs: Mapping[str, Any] = field(default_factory=dict) + keeps_history_reasoning: bool = True + verified: str = "not verified" + + +TEMPLATES: dict[str, TemplateSpec] = {} + + +def register_template(spec: TemplateSpec) -> TemplateSpec: + if spec.name in TEMPLATES and TEMPLATES[spec.name] != spec: + raise TemplateError(f"template {spec.name!r} is already registered") + TEMPLATES[spec.name] = spec + return spec + + +register_template(TemplateSpec( + name="glm5", tokenizer="zai-org/GLM-4.7-Flash", parser=parse_glm, + stop_tokens=("<|endoftext|>", "<|user|>", "<|observation|>"), + template_kwargs={"clear_thinking": False}, + verified="GLM-4.7-Flash tokenizer as a stand-in; GLM5.3-Flash tokenizer not verified")) +register_template(TemplateSpec( + name="qwen3", tokenizer="Qwen/Qwen3-0.6B", parser=parse_qwen3, + stop_tokens=("<|im_end|>", "<|endoftext|>"), + template_kwargs={"enable_thinking": True}, + # 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)")) + + +# ---------------------------------------------------------------------- template +def _render_ready(messages: Sequence[Message]) -> list[Message]: + """Canonical messages -> template input: tool-call arguments as objects.""" + out = [] + for message in messages: + message = dict(message) + if message.get("tool_calls"): + calls = [] + for call in message["tool_calls"]: + function = dict(call["function"]) + if isinstance(function.get("arguments"), str): + try: + function["arguments"] = json.loads(function["arguments"] or "{}") + except ValueError as exc: + raise TemplateError(f"tool call arguments are not JSON: {function['arguments'][:80]!r}") from exc + calls.append({**call, "function": function}) + message["tool_calls"] = calls + out.append(message) + return out + + +class ChatTemplate: + """A :class:`TemplateSpec` bound to a tokenizer (HF ``PreTrainedTokenizer``).""" + + def __init__(self, spec: TemplateSpec, tokenizer: Any) -> None: + self.spec = spec + self.tokenizer = tokenizer + ids = [] + for token in spec.stop_tokens: + token_ids = tokenizer.encode(token, add_special_tokens=False) + if len(token_ids) != 1: + raise TemplateError(f"stop token {token!r} is not one token in this tokenizer") + ids.append(token_ids[0]) + self.stop_token_ids: tuple[int, ...] = tuple(ids) + self._stop_by_id = dict(zip(ids, spec.stop_tokens)) + + @classmethod + def load(cls, name: str, *, tokenizer_path: str | None = None, local_files_only: bool = True) -> "ChatTemplate": + try: + spec = TEMPLATES[name] + except KeyError: + raise TemplateError(f"unknown chat template {name!r} (known: {sorted(TEMPLATES)})") from None + from transformers import AutoTokenizer + + tokenizer = AutoTokenizer.from_pretrained(tokenizer_path or spec.tokenizer, local_files_only=local_files_only) + return cls(spec, tokenizer) + + 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)) + + def encode(self, text: str) -> list[int]: + """Token ids of rendered text; the template writes its own BOS tokens.""" + return list(self.tokenizer.encode(text, add_special_tokens=False)) + + def decode(self, ids: Sequence[int]) -> str: + return self.tokenizer.decode(list(ids), skip_special_tokens=False) + + def stop_text(self, token_id: int) -> str | None: + return self._stop_by_id.get(int(token_id)) + + def parse(self, text: str) -> Message: + return self.spec.parser(text) + + +# ---------------------------------------------------------------------- N10.6 check +@dataclass(frozen=True) +class AlignmentReport: + """``aligned``: the proxy's ids equal ``apply_chat_template`` ids. + ``first_diff``: first differing position (None when aligned). + ``engine_text_extra``: tokens an engine adds when it tokenizes the rendered + text with special tokens (a text prompt would carry them twice).""" + + aligned: bool + first_diff: int | None + tokens: int + engine_text_extra: int + + +def first_difference(a: Sequence[int], b: Sequence[int]) -> int | None: + for i, (x, y) in enumerate(zip(a, b)): + if x != y: + return i + return None if len(a) == len(b) else min(len(a), len(b)) + + +def check_alignment(template: ChatTemplate, messages: Sequence[Message], *, + tools: Sequence[Message] | None = None, add_generation_prompt: bool = True) -> AlignmentReport: + 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, + add_generation_prompt=add_generation_prompt, **dict(template.spec.template_kwargs)) + if isinstance(reference, Mapping): + reference = reference["input_ids"] + engine_text = template.tokenizer.encode(text, add_special_tokens=True) + diff = first_difference(ours, list(reference)) + return AlignmentReport(aligned=diff is None, first_diff=diff, tokens=len(ours), + engine_text_extra=len(engine_text) - len(ours)) diff --git a/yeto/rl/harness/gateway/proxy.py b/yeto/rl/harness/gateway/proxy.py new file mode 100644 index 000000000..2b8c832b2 --- /dev/null +++ b/yeto/rl/harness/gateway/proxy.py @@ -0,0 +1,148 @@ +"""Records proxy: the gateway plus :class:`.records.RecordingBackend`, with an +HTTP shell (S25 N10.5). + +Sits in front of a sampler engine. An agent that only speaks OpenAI (Harbor's +built-in codex agent uses Responses) gets a per-trajectory base URL:: + + POST /trajectories -> {"trajectory_id", "base_url"} + POST /t/{trajectory_id}/v1/responses (Codex; stream:true answers SSE) + POST /t/{trajectory_id}/v1/chat/completions + GET /t/{trajectory_id}/records token ids, loss mask, log-probs, source + POST /t/{trajectory_id}/close + +The sampling fields stay signed by the gateway (a request may not change them). +Errors fail closed: an invalid trajectory answers 409, an unsupported request +400, a draining member 503. +""" + +from __future__ import annotations + +import json +import secrets +from typing import Any + +from yeto.rl.compaction import CompactionConfig +from yeto.rl.engine.sampler_engine import TokenSamplerEngine +from yeto.rl.harness.chat_templates import ChatTemplate, TemplateError + +from .core import AdmissionClosed, Gateway, GatewayConfig, GatewayError, SamplingOverrideError, TrajectoryInvalid +from .records import RecordingBackend +from .translate import 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).""" + tools = body.get("tools") or [] + if entry != "responses": + return list(tools) + out = [] + for tool in tools: + 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} + out.append({"type": "function", "function": function}) + return out + + +class RecordsProxy: + """Gateway + records backend for one engine and one chat template.""" + + def __init__(self, engine: TokenSamplerEngine, template: ChatTemplate, *, model: str, + signed_sampling: dict[str, Any], sampler_id: str = "records-proxy", + compaction: CompactionConfig | None = None, max_chains: int | None = None, + harness_board: Any = None, member: str | None = None) -> None: + self.backend = RecordingBackend(engine, template, sampler_id=sampler_id, compaction=compaction) + self.gateway = Gateway(GatewayConfig( + model=model, signed_sampling=dict(signed_sampling), max_chains=max_chains, + keeps_history_reasoning=template.spec.keeps_history_reasoning, + harness_board=harness_board, member=member), self.backend) + self._responses = 0 + + def open(self) -> str: + return f"traj-{secrets.token_hex(6)}" + + 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) + + def records(self, trajectory_id: str) -> dict[str, Any]: + out = self.backend.records(trajectory_id) + out["gateway"] = self.gateway.snapshot(trajectory_id) + return out + + def close(self, trajectory_id: str) -> None: + self.gateway.close(trajectory_id) + + # ------------------------------------------------------------------ HTTP + def app(self, *, public_url: str = "") -> Any: + from aiohttp import web + + async def create(request: web.Request) -> web.Response: + trajectory_id = self.open() + base = public_url or f"{request.scheme}://{request.host}" + return web.json_response({"trajectory_id": trajectory_id, "base_url": f"{base}/t/{trajectory_id}/v1"}) + + def entry_route(entry: str) -> Any: + async def route(request: web.Request) -> web.StreamResponse: + trajectory_id = request.match_info["tid"] + try: + body = await request.json() + result = await self.handle(entry, body, trajectory_id=trajectory_id) + except (UnsupportedShape, SamplingOverrideError, TemplateError, json.JSONDecodeError) as exc: + return _error(400, exc) + except AdmissionClosed as exc: + return _error(503, exc) + except (TrajectoryInvalid, GatewayError) as exc: + return _error(409, exc) + if not body.get("stream"): + return web.json_response(result) + self._responses += 1 + payload = _responses_sse(body, result, self._responses) if entry == "responses" else _chat_sse(result) + return web.Response(body=payload, headers={"Content-Type": "text/event-stream", "Cache-Control": "no-cache"}) + return route + + async def records(request: web.Request) -> web.Response: + return web.json_response(self.records(request.match_info["tid"])) + + async def close(request: web.Request) -> web.Response: + self.close(request.match_info["tid"]) + return web.json_response({"closed": True}) + + app = web.Application(client_max_size=64 * 1024**2) + app.router.add_post("/trajectories", create) + app.router.add_post("/t/{tid}/v1/responses", entry_route("responses")) + app.router.add_post("/t/{tid}/v1/chat/completions", entry_route("chat")) + app.router.add_get("/t/{tid}/records", records) + app.router.add_post("/t/{tid}/close", close) + return app + + +def _error(status: int, exc: Exception) -> Any: + from aiohttp import web + + return web.json_response({"error": {"type": type(exc).__name__, "message": str(exc)}}, status=status) + + +def _responses_sse(request: dict[str, Any], response: dict[str, Any], number: int) -> bytes: + # Reuse the Codex bridge's event sequence (created, item added/done, completed). + from yeto.rl.harness.codex.codex_harness_agent import _sse_response + + return _sse_response(request={"model": request.get("model"), "tools": request.get("tools") or [], + "prompt_cache_key": request.get("prompt_cache_key")}, + output=response["output"], usage=response["usage"], + response_number=number) + + +def _chat_sse(completion: dict[str, Any]) -> bytes: + """One chunk with the whole message, then ``[DONE]`` (no token streaming).""" + choice = completion["choices"][0] + delta = {k: v for k, v in choice["message"].items() if k in ("role", "content", "reasoning_content", "tool_calls")} + if delta.get("tool_calls"): + delta["tool_calls"] = [{"index": i, **call} for i, call in enumerate(delta["tool_calls"])] + chunk = {"id": completion.get("id", "chatcmpl"), "object": "chat.completion.chunk", + "choices": [{"index": 0, "delta": delta, "finish_reason": choice["finish_reason"]}], + "usage": completion.get("usage")} + return b"data: " + json.dumps(chunk).encode() + b"\n\ndata: [DONE]\n\n" + diff --git a/yeto/rl/harness/gateway/records.py b/yeto/rl/harness/gateway/records.py new file mode 100644 index 000000000..61a610ff2 --- /dev/null +++ b/yeto/rl/harness/gateway/records.py @@ -0,0 +1,329 @@ +"""Records backend: the gateway ``SessionBackend`` that samples a ``SamplerEngine`` +directly and records every token (S25 N10.5). + +The gateway (:class:`.core.Gateway`) keeps doing what it does for the Miles +session server: Responses/Chat/Messages -> canonical chat, chain location, +signed sampling. This backend replaces the session server behind it: + +* one backend session = one append-only token stream (TITO: token in, token + out). The first turn renders the chat template; a later turn appends only the + ids of the new messages after the tokens the engine generated, so the + recorded ids are the ids the engine saw and sampled; +* per token: id, loss mask (1 = sampled by the policy), log-probability, and + per turn the policy version and the source (:class:`Provenance`); +* N10.6 drift check: each turn compares the stream with a fresh tokenization of + the full rendering and records the first differing position. The stream is + what trains; a drift means the template does not reproduce the generated + text (counted, not fatal); +* optional context compaction (CompactionRL, :mod:`yeto.rl.compaction`): when + ``C - |h| < T_comp`` the policy writes a summary in the current segment + (trainable), and the next segment starts from ``system + u_resume(S) + last + k steps``. One training sample per segment, as in the Codex bridge. + +The chat completion carries SGLang's ``meta_info.output_token_logprobs``, so a +Miles session server can also sit behind this backend. +""" + +from __future__ import annotations + +import asyncio +import itertools +from collections import Counter +from collections.abc import Mapping, Sequence +from dataclasses import dataclass, field +from typing import Any + +from yeto.rl.compaction import RESUME_TEMPLATE, SUMMARY_PROMPT, CompactionConfig, extract_summary, should_compact +from yeto.rl.engine.sampler_engine import ( + FINISH_LENGTH, + Completion, + Provenance, + SamplingParams, + TokenSamplerEngine, +) +from yeto.rl.harness.chat_templates import ChatTemplate, first_difference + +from .core import GatewayError + +Message = dict[str, Any] + +KIND_ACTION = "action" +KIND_SUMMARY = "summary" + + +class RecordsMismatch(GatewayError): + """The stream cannot be continued token-exactly (fail closed).""" + + +class ContextFull(GatewayError): + """The prompt leaves no room for a reply within ``max_model_len``.""" + + +def sampling_params(signed: Mapping[str, Any]) -> SamplingParams: + """Gateway signed sampling (OpenAI field names) -> :class:`SamplingParams`.""" + if int(signed.get("n", 1)) != 1: + raise GatewayError("the records proxy samples one completion per request (n=1)") + max_tokens = next((signed[k] for k in ("max_tokens", "max_completion_tokens", "max_output_tokens") if k in signed), None) + if max_tokens is None: + raise GatewayError("signed sampling must set max_tokens") + stop = signed.get("stop") or () + return SamplingParams( + max_tokens=int(max_tokens), temperature=float(signed.get("temperature", 1.0)), + top_p=float(signed.get("top_p", 1.0)), top_k=int(signed.get("top_k", -1)), + seed=None if signed.get("seed") is None else int(signed["seed"]), + stop=(stop,) if isinstance(stop, str) else tuple(stop), + frequency_penalty=float(signed.get("frequency_penalty", 0.0)), + presence_penalty=float(signed.get("presence_penalty", 0.0)), + ).check() + + +@dataclass +class TurnRecord: + """One engine call. ``prompt_added``: mask-0 ids appended before it; + ``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).""" + + turn: int + kind: str + prompt_added: tuple[int, ...] + completion: tuple[int, ...] + logprobs: tuple[float, ...] + finish_reason: str + policy_version: int | None + provenance: dict[str, str] + stop_added: tuple[int, ...] = () + render_diff: int | None = None + + +@dataclass +class Segment: + """One training sample. Context = ``head + history[base:]``; ``n_context``: + context messages through the last generated assistant message.""" + + index: int + head: list[Message] + base: int + token_ids: list[int] = field(default_factory=list) + loss_mask: list[int] = field(default_factory=list) + logprobs: list[float] = field(default_factory=list) + turns: list[TurnRecord] = field(default_factory=list) + n_context: int = 0 + last_stop: str | None = None # stop token text of the last turn (None: not returned) + last_finish: str | None = None + kept_steps: int | None = None + summary_ok: bool | None = None + + def extend(self, ids: Sequence[int], *, mask: int, logprobs: Sequence[float] | None = None) -> None: + self.token_ids.extend(int(t) for t in ids) + self.loss_mask.extend([mask] * len(ids)) + self.logprobs.extend([0.0] * len(ids) if logprobs is None else [float(x) for x in logprobs]) + + def sample(self) -> dict[str, Any]: + return {"segment": self.index, "token_ids": list(self.token_ids), "loss_mask": list(self.loss_mask), + "logprobs": list(self.logprobs), "kept_steps": self.kept_steps, "summary_ok": self.summary_ok, + "turns": [vars(t).copy() for t in self.turns]} + + +@dataclass(frozen=True) +class _Prompt: + """Ids to append before the next sample. ``stop``: a stop token the engine + ended on but did not return (mask 0); ``tail``: the new messages (mask 0); + ``render_diff``: N10.6 drift position (None: aligned or not checked).""" + + stop: tuple[int, ...] + tail: tuple[int, ...] + render_diff: int | None = None + + def __len__(self) -> int: + return len(self.stop) + len(self.tail) + + +@dataclass +class _Session: + trajectory_id: str + segments: list[Segment] + lock: asyncio.Lock = field(default_factory=asyncio.Lock) + turns: int = 0 + + +def _atomic_step_starts(messages: Sequence[Message]) -> list[int]: + """Indexes of the assistant messages (a step = assistant + its tool outputs).""" + return [i for i, m in enumerate(messages) if m.get("role") == "assistant"] + + +class RecordingBackend: + """``SessionBackend`` over a :class:`TokenSamplerEngine` (module doc).""" + + def __init__(self, engine: TokenSamplerEngine, template: ChatTemplate, *, sampler_id: str = "records-proxy", + compaction: CompactionConfig | None = None, check_render: bool = True) -> None: + if not isinstance(engine, TokenSamplerEngine): + raise TypeError(f"{type(engine).__name__} has no generate_tokens (TokenSamplerEngine)") + self.engine = engine + self.template = template + self.sampler_id = sampler_id + self.compaction = compaction + self.check_render = check_render + self.sessions: dict[str, _Session] = {} + self.tools: dict[str, list[Message]] = {} + self.counters: dict[str, Counter[str]] = {} + self._ids = itertools.count() + + # -------------------------------------------------------------- gateway seam + def declare_tools(self, trajectory_id: str, tools: Sequence[Message] | None) -> None: + """Tool schemas of a trajectory (rendered into the prompt); fixed once set.""" + tools = list(tools or []) + known = self.tools.setdefault(trajectory_id, tools) + if known != tools: + raise RecordsMismatch("tool schemas changed within a trajectory; the rendered prompt would change") + + async def create_session(self, trajectory_id: str) -> str: + session_id = f"rec{next(self._ids)}" + self.sessions[session_id] = _Session(trajectory_id, [Segment(0, head=[], base=0)]) + self.counters.setdefault(trajectory_id, Counter()) + return session_id + + async def chat(self, session_id: str, messages: list[Message], sampling: dict[str, Any]) -> dict[str, Any]: + session = self.sessions[session_id] + params = sampling_params(sampling) + async with session.lock: + segment = session.segments[-1] + context = segment.head + list(messages[segment.base:]) + pending = self._prompt(session, segment, context) + if self._should_compact(session, segment, pending): + await self._compact(session, segment, context, params) + segment = session.segments[-1] + context = segment.head + list(messages[segment.base:]) + pending = self._prompt(session, segment, context) + completion = await self._sample(session, segment, pending, params, kind=KIND_ACTION) + segment.n_context = len(context) + 1 + message = self._parse(completion) + return self._chat_response(session_id, message, completion, prompt_tokens=len(segment.token_ids) - len(completion.token_ids)) + + # -------------------------------------------------------------- records + def records(self, trajectory_id: str) -> dict[str, Any]: + sessions = [(sid, s) for sid, s in self.sessions.items() if s.trajectory_id == trajectory_id] + return {"trajectory_id": trajectory_id, "template": self.template.spec.name, + "template_verified": self.template.spec.verified, + "sessions": [{"session_id": sid, "samples": [seg.sample() for seg in s.segments]} for sid, s in sessions], + "counters": dict(self.counters.get(trajectory_id, Counter()))} + + # -------------------------------------------------------------- internals + def _prompt(self, session: _Session, segment: Segment, context: list[Message]) -> _Prompt: + """What to append before the next sample (the whole prompt on a segment's first turn).""" + tools = self.tools.get(session.trajectory_id) + full = self.template.render(context, tools=tools, add_generation_prompt=True) + if not segment.turns: + return _Prompt(stop=(), tail=tuple(self.template.encode(full))) + if segment.last_finish == FINISH_LENGTH: + raise RecordsMismatch("cannot continue a stream after a length-truncated turn") + if len(context) < segment.n_context: + raise RecordsMismatch("history is shorter than the recorded stream") + start = self.template.render(context[: segment.n_context - 1], tools=tools, add_generation_prompt=True) + if not full.startswith(start): + raise RecordsMismatch("the rendering of the earlier turns changed (template rewrites history)") + stop_ids: tuple[int, ...] = () + if segment.last_stop is None: # the engine ended on a stop token but did not return it + found = sorted((full.find(s, len(start)), s) for s in self.template.spec.stop_tokens) + found = [(pos, s) for pos, s in found if pos >= 0] + if not found: + raise RecordsMismatch("no stop token after the last assistant turn in the rendering") + pos, stop = found[0] + stop_ids = tuple(self.template.encode(stop)) + else: + stop = segment.last_stop + pos = full.find(stop, len(start)) + if pos < 0: + raise RecordsMismatch(f"the rendering has no {stop!r} after the last assistant turn") + tail = tuple(self.template.encode(full[pos + len(stop):])) + diff = None + if self.check_render: + diff = first_difference([*segment.token_ids, *stop_ids, *tail], self.template.encode(full)) + if diff is not None: + self.counters[session.trajectory_id]["render_drift_turns"] += 1 + return _Prompt(stop=stop_ids, tail=tail, render_diff=diff) + + def _should_compact(self, session: _Session, segment: Segment, pending: _Prompt) -> bool: + cfg = self.compaction + if cfg is None or len(session.segments) > cfg.max_compactions: + return False + if not any(t.kind == KIND_ACTION for t in segment.turns): + return False # no step to summarize yet + return should_compact(len(segment.token_ids) + len(pending), cfg) + + async def _compact(self, session: _Session, segment: Segment, context: list[Message], + params: SamplingParams) -> None: + """Summary turn in this segment, then open the next one (CompactionRL eq. 8/9). + + ``history`` is the agent's own messages in this segment's context + (``messages[segment.base:]``); the next segment keeps the last ``k`` + whole steps (assistant + its tool outputs), ``k`` reduced while the + rebuilt context would trigger again (as :func:`yeto.rl.compaction.rebuild_context`).""" + cfg = self.compaction + assert cfg is not None + summary_context = context + [{"role": "user", "content": SUMMARY_PROMPT}] + completion = await self._sample(session, segment, self._prompt(session, segment, summary_context), + params, kind=KIND_SUMMARY) + summary, segment.summary_ok = extract_summary(self.template.decode(self._strip_stop(completion.token_ids))) + self.counters[session.trajectory_id]["compactions"] += 1 + + systems = list(itertools.takewhile(lambda m: m.get("role") == "system", context)) + head = systems + [{"role": "user", "content": RESUME_TEMPLATE.format(summary=summary)}] + history = context[len(segment.head):] + steps = _atomic_step_starts(history) + tools = self.tools.get(session.trajectory_id) + kept, first = 0, len(history) + for k in range(min(cfg.keep_recent_steps, len(steps)), 0, -1): + rebuilt = head + history[steps[-k]:] + if not should_compact(len(self.template.encode(self.template.render(rebuilt, tools=tools, add_generation_prompt=True))), cfg): + kept, first = k, steps[-k] + break + session.segments.append(Segment(len(session.segments), head=head, base=segment.base + first, kept_steps=kept)) + + async def _sample(self, session: _Session, segment: Segment, pending: _Prompt, params: SamplingParams, + *, kind: str) -> Completion: + 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}") + params = params.replace(max_tokens=min(params.max_tokens, room)) + (group,) = await asyncio.to_thread(self.engine.generate_tokens, [prompt], params, + stop_token_ids=self.template.stop_token_ids) + completion = group[0] + if completion.prompt_token_ids and tuple(completion.prompt_token_ids) != tuple(prompt): + raise RecordsMismatch("the engine reports other prompt ids than the proxy sent") + if pending.stop: + segment.turns[-1].stop_added = pending.stop + self.counters[session.trajectory_id]["stop_added"] += 1 + segment.extend(pending.stop, mask=0) + segment.extend(pending.tail, mask=0) + segment.extend(completion.token_ids, mask=1, logprobs=completion.logprobs) + last = completion.token_ids[-1] if completion.token_ids else None + segment.last_stop = None if last is None else self.template.stop_text(last) + segment.last_finish = completion.finish_reason + provenance = Provenance.of(self.engine.identity, sampler_id=self.sampler_id, sampling=params) + segment.turns.append(TurnRecord( + turn=session.turns, kind=kind, prompt_added=pending.tail, completion=tuple(completion.token_ids), + logprobs=tuple(completion.logprobs), finish_reason=completion.finish_reason, + policy_version=self.engine.served_version(), provenance=provenance.to_record(), render_diff=pending.render_diff)) + session.turns += 1 + return completion + + def _strip_stop(self, ids: Sequence[int]) -> list[int]: + ids = list(ids) + return ids[:-1] if ids and self.template.stop_text(ids[-1]) is not None else ids + + def _parse(self, completion: Completion) -> Message: + return self.template.parse(self.template.decode(self._strip_stop(completion.token_ids))) + + @staticmethod + def _chat_response(session_id: str, message: Message, completion: Completion, *, prompt_tokens: int) -> dict[str, Any]: + finish = "length" if completion.finish_reason == FINISH_LENGTH else ("tool_calls" if message.get("tool_calls") else "stop") + n = len(completion.token_ids) + return {"id": f"chatcmpl-{session_id}", "object": "chat.completion", + "choices": [{"index": 0, "message": message, "finish_reason": finish, + "meta_info": {"output_token_logprobs": [[lp, tid, None] for lp, tid in + zip(completion.logprobs, completion.token_ids)], + "prompt_tokens": prompt_tokens, "completion_tokens": n}}], + "usage": {"prompt_tokens": prompt_tokens, "completion_tokens": n, "total_tokens": prompt_tokens + n}} From 042a12bf71ba6bb55ddab85f5137e9e456d7a48a Mon Sep 17 00:00:00 2001 From: michael <78113554+chowhsu@users.noreply.github.com> Date: Mon, 19 Oct 2026 09:09:09 +0000 Subject: [PATCH 12/34] Fail closed on a reload during a records-proxy turn (S25 N10.5) The turn's policy version was read after generate_tokens returned, so a reload during the call recorded the new version for tokens the old one may have sampled. The backend now reads the version before and after the call and refuses the turn when it changed. Signed repetition_penalty is passed to SamplingParams so check() refuses it instead of dropping it silently, and prompt_tokens is computed under the session lock. Co-Authored-By: Claude Opus 5.5 --- tests/test_records_proxy.py | 22 ++++++++++++++++++++++ yeto/rl/harness/gateway/records.py | 10 ++++++++-- 2 files changed, 30 insertions(+), 2 deletions(-) diff --git a/tests/test_records_proxy.py b/tests/test_records_proxy.py index 8c39833bf..e9a77ace3 100644 --- a/tests/test_records_proxy.py +++ b/tests/test_records_proxy.py @@ -258,6 +258,28 @@ def test_length_truncated_turn_and_changed_tools_fail_closed(qwen): asyncio.run(proxy.handle("responses", responses_body([user_item("go")], tools=[]), trajectory_id="t")) +def test_signed_repetition_penalty_is_refused_not_dropped(): + from yeto.rl.engine.sampler_engine import SamplingConfigError + from yeto.rl.harness.gateway.records import sampling_params + + with pytest.raises(SamplingConfigError, match="repetition_penalty"): + sampling_params({**SIGNED, "repetition_penalty": 1.1}) + + +def test_a_reload_during_a_turn_fails_closed(qwen): + model = ScriptedModel(qwen, SCRIPT) + engine = engine_for(model) + + def reply_then_reload(prompt, index): + engine.start(None, 4) # the engine serves a new version while the turn runs + return model(prompt, index) + + engine.token_reply = reply_then_reload + proxy = RecordsProxy(engine, qwen, model="policy", signed_sampling=SIGNED) + with pytest.raises(RecordsMismatch, match="served version changed"): + asyncio.run(proxy.handle("responses", responses_body([user_item("go")]), trajectory_id="t")) + + def test_backend_refuses_an_engine_without_token_input(qwen): class TextOnly: pass diff --git a/yeto/rl/harness/gateway/records.py b/yeto/rl/harness/gateway/records.py index 61a610ff2..39ab8312c 100644 --- a/yeto/rl/harness/gateway/records.py +++ b/yeto/rl/harness/gateway/records.py @@ -72,6 +72,7 @@ def sampling_params(signed: Mapping[str, Any]) -> SamplingParams: top_p=float(signed.get("top_p", 1.0)), top_k=int(signed.get("top_k", -1)), seed=None if signed.get("seed") is None else int(signed["seed"]), stop=(stop,) if isinstance(stop, str) else tuple(stop), + repetition_penalty=float(signed.get("repetition_penalty", 1.0)), frequency_penalty=float(signed.get("frequency_penalty", 0.0)), presence_penalty=float(signed.get("presence_penalty", 0.0)), ).check() @@ -198,8 +199,9 @@ async def chat(self, session_id: str, messages: list[Message], sampling: dict[st pending = self._prompt(session, segment, context) 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) - return self._chat_response(session_id, message, completion, prompt_tokens=len(segment.token_ids) - len(completion.token_ids)) + return self._chat_response(session_id, message, completion, prompt_tokens=prompt_tokens) # -------------------------------------------------------------- records def records(self, trajectory_id: str) -> dict[str, Any]: @@ -288,8 +290,12 @@ async def _sample(self, session: _Session, segment: Segment, pending: _Prompt, p if room < 1: raise ContextFull(f"prompt has {len(prompt)} tokens; max_model_len is {self.engine.max_model_len}") 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, stop_token_ids=self.template.stop_token_ids) + if self.engine.served_version() != version: # a reload during the call: version unknown + raise RecordsMismatch(f"served version changed from {version} during the turn; " + "the turn's behavior version is unknown") completion = group[0] if completion.prompt_token_ids and tuple(completion.prompt_token_ids) != tuple(prompt): raise RecordsMismatch("the engine reports other prompt ids than the proxy sent") @@ -306,7 +312,7 @@ async def _sample(self, session: _Session, segment: Segment, pending: _Prompt, p segment.turns.append(TurnRecord( turn=session.turns, kind=kind, prompt_added=pending.tail, completion=tuple(completion.token_ids), logprobs=tuple(completion.logprobs), finish_reason=completion.finish_reason, - policy_version=self.engine.served_version(), provenance=provenance.to_record(), render_diff=pending.render_diff)) + policy_version=version, provenance=provenance.to_record(), render_diff=pending.render_diff)) session.turns += 1 return completion From 56616519ae789c9b8c504fdfcf3e8fb247d7b6fd Mon Sep 17 00:00:00 2001 From: michael <78113554+chowhsu@users.noreply.github.com> Date: Mon, 19 Oct 2026 12:13:59 +0000 Subject: [PATCH 13/34] Pick learner device and collective backend from the device family (S25 N7.7) The Megatron learner, the Miles value island and the FSDP2 plan no longer hard-code torch.cuda and NCCL. yeto.accel maps the inventory device family (cuda, ppu, npu, cpu) to a torch device type and a collective backend; PPU runs as cuda + NCCL over ACCL-P and NPU as torch_npu + HCCL (neither verified on its machine). The SSH island launcher exports YETO_DEVICE_FAMILY from the node inventory so a PPU rank does not rely on auto-detection. Co-Authored-By: Claude Opus 5.5 --- tests/test_accel.py | 65 +++++++++++++++++++ tests/test_megatron_learner_device_family.py | 66 +++++++++++++++++++ tests/test_rl_fullparam_fsdp2.py | 20 +++++- tests/test_rl_ssh_trainer_island.py | 1 + yeto/accel.py | 68 +++++++++++++++++++- yeto/gpu_spec.py | 3 + yeto/launch/ssh_trainer_island.py | 4 +- yeto/megatron/learner.py | 23 +++++-- yeto/megatron/miles_value_island.py | 4 +- yeto/rl/fullparam/fsdp2.py | 27 +++++++- 10 files changed, 267 insertions(+), 14 deletions(-) create mode 100644 tests/test_megatron_learner_device_family.py diff --git a/tests/test_accel.py b/tests/test_accel.py index 190a25c36..eca47d52e 100644 --- a/tests/test_accel.py +++ b/tests/test_accel.py @@ -293,3 +293,68 @@ def test_fork_rng_without_a_device_forks_no_cards_on_cpu( with accel.fork_rng(None): pass assert recorded_fork == [{"devices": []}] + + +# --- device family (S25 N7.7) -------------------------------------------- + + +@pytest.fixture +def fake_cuda(monkeypatch): + stub = FakeAccelerator() + monkeypatch.setattr(torch, "cuda", stub) + return stub + + +@pytest.mark.parametrize(("family", "device_type", "backend"), [ + ("cuda", "cuda", "nccl"), + ("ppu", "cuda", "nccl"), # CUDA device name, NCCL API over ACCL-P (not verified) + ("npu", "npu", "hccl"), + ("cpu", "cpu", "gloo"), +]) +def test_family_maps_to_device_type_and_backend(family, device_type, backend): + assert accel.family_device_type(family) == device_type + assert accel.family_dist_backend(family) == backend + + +def test_unknown_family_is_refused(monkeypatch): + with pytest.raises(ValueError, match="unknown device family"): + accel.family_dist_backend("tpu") + monkeypatch.setenv(accel.FAMILY_ENV, "tpu") + with pytest.raises(ValueError, match="unknown device family"): + accel.device_family() + + +def test_device_family_precedence(no_cuda, npu, monkeypatch): + monkeypatch.delenv(accel.FAMILY_ENV, raising=False) + assert accel.device_family() == "npu" # auto-detect + monkeypatch.setenv(accel.FAMILY_ENV, "ppu") + assert accel.device_family() == "ppu" # the launcher's inventory value wins + assert accel.device_family("cuda") == "cuda" # an explicit flag wins over both + + +def test_ppu_binds_a_cuda_card(fake_cuda, monkeypatch): + monkeypatch.setenv("LOCAL_RANK", "2") + bound = accel.bind_family_device("ppu") + assert bound == torch.device("cuda", 2) + assert fake_cuda.bound == [torch.device("cuda", 2)] + assert accel.current_device("ppu") == torch.device("cuda", 3) + + +def test_npu_binds_through_torch_npu(npu, monkeypatch): + monkeypatch.setattr(torch, "device", device) # "npu" is unregistered without torch_npu + bound = accel.bind_family_device("npu", local_rank=1) + assert (bound.type, bound.index) == ("npu", 1) + assert npu.bound == [bound] + assert (accel.current_device("npu").type, accel.current_device("npu").index) == ("npu", 3) + + +def test_npu_without_torch_npu_is_a_clear_error(monkeypatch): + monkeypatch.delattr(torch, "npu", raising=False) + monkeypatch.setattr(accel, "register_backends", lambda: None) + with pytest.raises(RuntimeError, match="torch_npu"): + accel.bind_family_device("npu", local_rank=0) + + +def test_cpu_family_binds_nothing(fake_cuda): + assert accel.bind_family_device("cpu", local_rank=5) == torch.device("cpu") + assert fake_cuda.bound == [] diff --git a/tests/test_megatron_learner_device_family.py b/tests/test_megatron_learner_device_family.py new file mode 100644 index 000000000..a4a7cc2e8 --- /dev/null +++ b/tests/test_megatron_learner_device_family.py @@ -0,0 +1,66 @@ +"""S25 N7.7: the Megatron learner takes its card and collective from the device family.""" + +from __future__ import annotations + +import sys +from types import ModuleType, SimpleNamespace + +import pytest +import torch +import torch.distributed as dist + +from yeto import accel +from yeto.megatron import learner + + +@pytest.fixture +def megatron_stub(monkeypatch): + """The two Megatron-Core entry points _init_distributed calls, recorded.""" + calls: list[tuple] = [] + parallel_state = SimpleNamespace(initialize_model_parallel=lambda **kw: calls.append(("mp", kw))) + random = SimpleNamespace(model_parallel_cuda_manual_seed=lambda seed: calls.append(("seed", seed))) + modules = { + "megatron": ModuleType("megatron"), + "megatron.core": SimpleNamespace(parallel_state=parallel_state), + "megatron.core.tensor_parallel": ModuleType("megatron.core.tensor_parallel"), + "megatron.core.tensor_parallel.random": random, + } + for name, module in modules.items(): + monkeypatch.setitem(sys.modules, name, module) + monkeypatch.setattr(dist, "init_process_group", lambda **kw: calls.append(("pg", kw))) + monkeypatch.setattr(dist, "get_rank", lambda: 0) + monkeypatch.setattr(dist, "get_world_size", lambda: 1) + return calls + + +def _args(family): + return SimpleNamespace(device_family=family, tensor_parallel=1, pipeline_parallel=1, + expert_parallel=1, seed=7) + + +@pytest.mark.parametrize(("family", "backend"), [("cuda", "nccl"), ("ppu", "nccl"), ("npu", "hccl")]) +def test_init_distributed_follows_the_family(megatron_stub, monkeypatch, family, backend): + bound = [] + monkeypatch.setenv("LOCAL_RANK", "3") + monkeypatch.setattr(accel, "bind_family_device", + lambda fam, rank: bound.append((fam, rank)) or SimpleNamespace(type=fam, index=rank)) + rank, world, device = learner._init_distributed(_args(family)) + assert (rank, world, device.index) == (0, 1, 3) + assert bound == [(family, 3)] + assert megatron_stub[0] == ("pg", {"backend": backend}) + assert megatron_stub[-1] == ("seed", 7) + + +def test_family_comes_from_the_launcher_env(megatron_stub, monkeypatch): + monkeypatch.setenv(accel.FAMILY_ENV, "npu") + monkeypatch.setattr(accel, "bind_family_device", lambda fam, rank: torch.device("cpu")) + learner._init_distributed(_args(None)) + assert megatron_stub[0] == ("pg", {"backend": "hccl"}) + + +def test_cli_accepts_only_known_families(): + base = ["--model", "m", "--data", "d"] + assert learner.parse_args(base).device_family is None + assert learner.parse_args([*base, "--device-family", "ppu"]).device_family == "ppu" + with pytest.raises(SystemExit): + learner.parse_args([*base, "--device-family", "tpu"]) diff --git a/tests/test_rl_fullparam_fsdp2.py b/tests/test_rl_fullparam_fsdp2.py index aacd2d817..9a1ef2518 100644 --- a/tests/test_rl_fullparam_fsdp2.py +++ b/tests/test_rl_fullparam_fsdp2.py @@ -17,6 +17,7 @@ FullParamConfigError, check_ulysses, decoder_layers, + init_distributed, shard_model, train_step, verl_overrides, @@ -69,7 +70,7 @@ def build(config: dict) -> torch.nn.Module: def _worker(rank: int, world: int, init_file: str, config: dict, offload: str, out: str) -> None: torch.set_num_threads(1) - dist.init_process_group("gloo", init_method=f"file://{init_file}", rank=rank, world_size=world) + init_distributed("cpu", init_method=f"file://{init_file}", rank=rank, world_size=world) try: torch.manual_seed(0) model = build(config) @@ -160,3 +161,20 @@ def test_verl_overrides_full_parameter(): assert "actor_rollout_ref.actor.fsdp_config.fsdp_size=8" in ov assert "actor_rollout_ref.actor.fsdp_config.optimizer_offload=false" in ov assert "actor_rollout_ref.actor.fsdp_config.fsdp_size=-1" in verl_overrides(FullParamConfig()) + + +def test_backends_follow_the_device_family(): + from yeto.rl.fullparam.fsdp2 import BACKENDS + + assert BACKENDS == {"cpu": "gloo", "cuda": "nccl", "ppu": "nccl", "npu": "hccl"} + + +def test_init_distributed_passes_the_family_backend_and_card(monkeypatch): + import yeto.rl.fullparam.fsdp2 as fsdp2 + + calls = [] + monkeypatch.setattr(fsdp2.accel, "bind_family_device", lambda family: torch.device("cuda", 1)) + monkeypatch.setattr(dist, "init_process_group", lambda **kw: calls.append(kw)) + device = fsdp2.init_distributed("ppu", rank=0, world_size=1) + assert device == torch.device("cuda", 1) + assert calls == [{"backend": "nccl", "device_id": torch.device("cuda", 1), "rank": 0, "world_size": 1}] diff --git a/tests/test_rl_ssh_trainer_island.py b/tests/test_rl_ssh_trainer_island.py index 74d9f1cd5..1c50402fc 100644 --- a/tests/test_rl_ssh_trainer_island.py +++ b/tests/test_rl_ssh_trainer_island.py @@ -89,6 +89,7 @@ def test_start_script_runs_torchrun_with_rank_rendezvous_and_comm_env(): assert "--nnodes=2" in script and "--master-addr=10.0.0.1" in script and "--nproc-per-node=8" in script assert "export NCCL_SOCKET_IFNAME=eth0" in script and "export NCCL_IB_DISABLE=0" in script assert "export YETO_X=1" in script and "setsid bash -c" in script + assert f"export YETO_DEVICE_FAMILY={s.nodes[0].device_family}" in script assert "--node-rank=0" in script0 and "--node-rank=1" in script1 assert "yeto.rl.quant.release" in script0 and "--loop" in script0 and "--no-awq" in script0 assert "yeto.rl.quant.release" not in script1 diff --git a/yeto/accel.py b/yeto/accel.py index b2b4e1485..4750384ab 100644 --- a/yeto/accel.py +++ b/yeto/accel.py @@ -19,6 +19,8 @@ import torch +from .gpu_spec import DEVICE_FAMILY_ENV + # Visible-card environment variables, in accelerator auto-detection priority. _VISIBLE_DEVICE_ENVS = { "cuda": "CUDA_VISIBLE_DEVICES", @@ -26,9 +28,19 @@ } _ACCELERATORS = tuple(_VISIBLE_DEVICE_ENVS) -# torch.distributed backend per device family. +# torch.distributed backend per torch device type. _DIST_BACKENDS = {"cuda": "nccl", "npu": "hccl", "cpu": "gloo"} +# Inventory device family (gpu_spec.DEVICE_FAMILIES, plus "cpu") -> the torch +# device type its learner runs on. A PPU keeps the CUDA device name and the +# NCCL API (served by ACCL-P), so it shares the "cuda" runtime: an inference +# from the vendor documentation, not verified on a PPU machine. NPU uses +# torch_npu and HCCL. The collective backend follows from the device type. +FAMILY_DEVICE_TYPES = {"cuda": "cuda", "ppu": "cuda", "npu": "npu", "cpu": "cpu"} + +# Set by the island launcher from the node inventory; overrides auto-detection. +FAMILY_ENV = DEVICE_FAMILY_ENV + def register_backends() -> None: """Register out-of-tree device types. Idempotent; safe without the card.""" @@ -92,6 +104,60 @@ def dist_backend(device: torch.device | None = None) -> str: return _DIST_BACKENDS.get(device_type, "gloo") +def device_family(explicit: str | None = None) -> str: + """The device family of this rank: ``explicit``, else ``$YETO_DEVICE_FAMILY``, + else the detected torch device type (a PPU then reads as ``"cuda"``, + which selects the same runtime).""" + family = explicit or os.environ.get(FAMILY_ENV) or available_type() + if family not in FAMILY_DEVICE_TYPES: + raise ValueError(f"unknown device family {family!r} (known: {sorted(FAMILY_DEVICE_TYPES)})") + return family + + +def family_device_type(family: str) -> str: + """torch device type of a device family (``"ppu"`` -> ``"cuda"``).""" + try: + return FAMILY_DEVICE_TYPES[family] + except KeyError: + raise ValueError(f"unknown device family {family!r} (known: {sorted(FAMILY_DEVICE_TYPES)})") from None + + +def family_dist_backend(family: str) -> str: + """torch.distributed backend of a device family (nccl, hccl or gloo).""" + return _DIST_BACKENDS[family_device_type(family)] + + +def bind_family_device(family: str, local_rank: int | None = None) -> torch.device: + """Select this rank's local card of ``family`` and return it. + + The CPU family returns ``cpu`` and binds nothing. ``local_rank`` defaults + to ``$LOCAL_RANK``. + """ + device_type = family_device_type(family) + if device_type == "cpu": + return torch.device("cpu") + register_backends() + module = getattr(torch, device_type, None) + if module is None: + raise RuntimeError(f"device family {family!r} needs torch.{device_type}; " + + ("install the torch_npu build matching this torch and CANN release" + if device_type == "npu" else "this torch build has no CUDA runtime")) + if local_rank is None: + local_rank = int(os.environ.get("LOCAL_RANK", 0)) + device = torch.device(device_type, local_rank) + module.set_device(device) + return device + + +def current_device(family: str) -> torch.device: + """The card this rank already selected in ``family`` (``cpu`` for the CPU).""" + device_type = family_device_type(family) + if device_type == "cpu": + return torch.device("cpu") + register_backends() + return torch.device(device_type, getattr(torch, device_type).current_device()) + + def detect(explicit: str | None) -> torch.device: """Resolve the training device, binding this rank to its local card.""" register_backends() diff --git a/yeto/gpu_spec.py b/yeto/gpu_spec.py index 4a4e78e65..33b5b24bb 100644 --- a/yeto/gpu_spec.py +++ b/yeto/gpu_spec.py @@ -61,6 +61,9 @@ # communication variables. It is not the torch device type: a PPU card has # family ``ppu`` and device type ``cuda``. DEVICE_FAMILIES = ("cuda", "npu", "ppu") +# Environment variable the island launcher sets to the node's device family; +# yeto.accel.device_family reads it (torch cannot tell a PPU from a GPU). +DEVICE_FAMILY_ENV = "YETO_DEVICE_FAMILY" _CARD_DEVICE_FAMILY = {**{card: "npu" for card in _CARD_DEVICE_TYPE}, "810E": "ppu", "M890": "ppu"} diff --git a/yeto/launch/ssh_trainer_island.py b/yeto/launch/ssh_trainer_island.py index 72ebd423d..7dee9fe84 100644 --- a/yeto/launch/ssh_trainer_island.py +++ b/yeto/launch/ssh_trainer_island.py @@ -52,6 +52,7 @@ from yeto.rl.engine.trainer_lifecycle import CutRef from yeto.rl.engine.trainer_remote import island_env +from ..gpu_spec import DEVICE_FAMILY_ENV from .node_inventory import Node from .topology import island_comm_env from .trainer_island import ( @@ -164,7 +165,8 @@ def start_script(spec: SshTrainerSpec, rank: int, env: Mapping[str, str], paths: "set -eu\n" f"mkdir -p \"{paths['dir']}\"\n" f"cd {shlex.quote(spec.workdir)}\n" - + _exports({**comm, **spec.env, **env}) + # The learner reads the family to pick its card type and collective backend (accel). + + _exports({DEVICE_FAMILY_ENV: node.device_family, **comm, **spec.env, **env}) + f"rm -f \"{paths['rc']}\"\n" # Job control is off in a non-interactive shell, so the background setsid # is not a group leader: it does not fork, and $! is the new group's id. diff --git a/yeto/megatron/learner.py b/yeto/megatron/learner.py index 1f2752dcd..cb040dfda 100644 --- a/yeto/megatron/learner.py +++ b/yeto/megatron/learner.py @@ -105,6 +105,9 @@ def parse_args(argv=None): p.add_argument("--output-dir", default="~/yeto-output") # Megatron parallelism. p.add_argument("--island-backend", default="megatron") + p.add_argument("--device-family", choices=("cuda", "ppu", "npu"), default=None, + help="card family (default: $YETO_DEVICE_FAMILY, else auto-detect); " + "selects the torch device type and collective backend") p.add_argument("--expert-parallel", type=int, default=1) p.add_argument("--tensor-parallel", type=int, default=1) p.add_argument("--pipeline-parallel", type=int, default=1) @@ -116,15 +119,22 @@ def parse_args(argv=None): def _init_distributed(args): """torch.distributed first, then Megatron parallel state, then the - model-parallel RNG (the order Megatron-Core asserts).""" - import torch + model-parallel RNG (the order Megatron-Core asserts). + + The device family picks the card type and collective backend (accel). + PPU runs as "cuda" + NCCL over ACCL-P; NPU as torch_npu + HCCL, where + Megatron's CUDA-named RNG helper relies on torch_npu's CUDA-to-NPU patch + (MindSpeed). Neither family has run this learner yet (not verified).""" import torch.distributed as dist from megatron.core import parallel_state from megatron.core.tensor_parallel.random import model_parallel_cuda_manual_seed + from .. import accel + local_rank = int(os.environ.get("LOCAL_RANK", "0")) - torch.cuda.set_device(local_rank) - dist.init_process_group(backend="nccl") + family = accel.device_family(args.device_family) + device = accel.bind_family_device(family, local_rank) + dist.init_process_group(backend=accel.family_dist_backend(family)) parallel_state.initialize_model_parallel( tensor_model_parallel_size=args.tensor_parallel, pipeline_model_parallel_size=args.pipeline_parallel, @@ -139,7 +149,7 @@ def _init_distributed(args): "of linear_in/linear_out and PP>1 a cross-stage gather before the " "DiLoCo push — implement those before enabling." ) - return dist.get_rank(), dist.get_world_size(), local_rank + return dist.get_rank(), dist.get_world_size(), device def _build_model(args, device): @@ -359,7 +369,7 @@ def main(argv=None): args = parse_args(argv) logging.basicConfig(level=logging.INFO, format="%(asctime)s megatron %(levelname)s %(message)s") - rank, world, local_rank = _init_distributed(args) + rank, world, device = _init_distributed(args) from ..provenance import ( pin_distributed_runtime_provenance, verify_distributed_source_tree_sha256, @@ -371,7 +381,6 @@ def main(argv=None): world=world, ) pin_distributed_runtime_provenance(args, rank=rank, world=world) - device = torch.device("cuda", local_rank) model, bridge = _build_model(args, device) diff --git a/yeto/megatron/miles_value_island.py b/yeto/megatron/miles_value_island.py index d43b92981..287a5c08a 100644 --- a/yeto/megatron/miles_value_island.py +++ b/yeto/megatron/miles_value_island.py @@ -23,6 +23,7 @@ import torch import torch.distributed as dist +from .. import accel from ..fragments import ( FRAGMENT_PATTERNS, MERGE_AVG, @@ -652,7 +653,8 @@ def __init__(self, args: Any, model: Any, optimizer: Any): self.parallel = get_parallel_state() self.global_rank = dist.get_rank() self.world_size = dist.get_world_size() - self.device = torch.device("cuda", torch.cuda.current_device()) + # The card Miles already bound for this rank, in its device family. + self.device = accel.current_device(accel.device_family()) self._validate_topology() self.is_leader = self.parallel.cp.rank == 0 and self.parallel.tp.rank == 0 diff --git a/yeto/rl/fullparam/fsdp2.py b/yeto/rl/fullparam/fsdp2.py index 6208d624c..470e8687d 100644 --- a/yeto/rl/fullparam/fsdp2.py +++ b/yeto/rl/fullparam/fsdp2.py @@ -36,11 +36,13 @@ import torch from torch import nn +from ... import accel from .memory import OFFLOAD_MODES -# The collective backend of each device family (D4). PPU keeps the CUDA device -# name and NCCL API (over ACCL-P; not verified on a PPU machine). -BACKENDS = {"cpu": "gloo", "cuda": "nccl", "ppu": "nccl", "npu": "hccl"} +# The collective backend of each device family (D4), from yeto.accel. PPU keeps +# the CUDA device name and NCCL API (over ACCL-P); NPU uses torch_npu + HCCL. +# Neither is verified on its machine. +BACKENDS = {family: accel.family_dist_backend(family) for family in accel.FAMILY_DEVICE_TYPES} class FullParamConfigError(ValueError): @@ -97,6 +99,25 @@ def decoder_layers(model: nn.Module) -> list[nn.Module]: return layers +def init_distributed(family: str | None = None, **kwargs: Any) -> torch.device: + """Bind this rank's card and join the default process group. + + ``family``: ``"cuda"``, ``"ppu"``, ``"npu"`` or ``"cpu"``; default + ``$YETO_DEVICE_FAMILY``, else auto-detected (accel.device_family). The + backend follows the family (:data:`BACKENDS`); an accelerator rank passes + its card as ``device_id``. ``kwargs`` go to ``init_process_group`` + (``init_method``, ``rank``, ``world_size``). Returns the bound device. + """ + import torch.distributed as dist + + family = accel.device_family(family) + device = accel.bind_family_device(family) + if accel.is_accelerator(device): + kwargs.setdefault("device_id", device) + dist.init_process_group(backend=BACKENDS[family], **kwargs) + return device + + def shard_model(model: nn.Module, cfg: FullParamConfig, mesh: Any = None) -> nn.Module: """Apply FSDP2 to every decoder layer, then to the root (in place). From 1ced1cacb98817f34e57b72c92a0ae8090566cf7 Mon Sep 17 00:00:00 2001 From: michael <78113554+chowhsu@users.noreply.github.com> Date: Mon, 19 Oct 2026 12:19:00 +0000 Subject: [PATCH 14/34] Add verl/FSDP2 cut save and resume with the resume assertions (S25 N7.6) verl had no cut path, so its resume assertions were never applied. The verl adapter now writes a neutral CutManifest (same progress/data/ledger/ outer and resume sections as Miles) over one yeto-owned file per FSDP rank (local DTensor shards, optimizer state, LR scheduler, RNG). Restore needs a fresh trainer at the same world size, checks each file checksum and the loaded-state digest, then the resume assertions; the first round must advance every Adam step by the optimizer steps per round and the scheduler by one. verl itself is unchanged. Review fixes on the same assertions: - Miles compared the data watermark with a value read from the same cut (T-REV2). The watermark is now compared with the data cursor that actually landed: the driver after its seek, the trainer rebuild after its live re-read (confirm_resumed_data_cursor). The first step after a restore refuses an unconfirmed cursor. - A DP-change restore (restore_cut_resharded) now checks the run-wide counters: global step, every rank's Adam step, scheduler progress and LRs against the cut, plus the first-step and data-cursor checks. Neutral: restore_problems(check_data_watermark=), data_watermark_problems, run_wide_reference, first_step_problems(optimizer_steps=); a cut may name fewer data cursor fields (cursor_fields), sample_offset stays required. CPU tests: FSDP2 over gloo with 2 ranks, 3 rounds + cut + restore + 2 rounds equals 5 rounds bit for bit; a reset Adam step on one rank is refused on every rank; watermark, first-round and refusal cases. Co-Authored-By: Claude Opus 5.5 --- tests/test_rl_resume_manifest.py | 35 ++ tests/test_rl_trainer_reshard.py | 23 ++ tests/test_rl_verl_cut.py | 273 +++++++++++++ yeto/rl/adapters/miles/cut_plugin.py | 3 + yeto/rl/adapters/miles/trainer.py | 56 ++- yeto/rl/adapters/miles/trainer_rebuild.py | 13 +- yeto/rl/adapters/verl/cut.py | 455 ++++++++++++++++++++++ yeto/rl/adapters/verl/ports_impl.py | 34 ++ yeto/rl/adapters/verl/resume_check.py | 56 ++- yeto/rl/engine/cut.py | 6 +- yeto/rl/engine/driver.py | 3 + yeto/rl/engine/resume_manifest.py | 60 ++- 12 files changed, 991 insertions(+), 26 deletions(-) create mode 100644 tests/test_rl_verl_cut.py create mode 100644 yeto/rl/adapters/verl/cut.py diff --git a/tests/test_rl_resume_manifest.py b/tests/test_rl_resume_manifest.py index 21219e77d..d972c4e33 100644 --- a/tests/test_rl_resume_manifest.py +++ b/tests/test_rl_resume_manifest.py @@ -344,10 +344,44 @@ def test_miles_cut_records_and_checks_the_resume_manifest(tmp_path): state = next(iter(resume.ranks.values())) assert {s.max for s in state.adam_steps.values()} == {2} and state.moments train_step(fresh, torch.randn(3, 6)) + with pytest.raises(ResumeAssertionError, match="never compared"): + trainer._check_first_step_after_resume(1) + trainer._resume_reference = resume # the refused check consumed nothing else + trainer.confirm_resumed_data_cursor({"sample_offset": 4}) # the cursor the driver landed on trainer._check_first_step_after_resume(1) # step 2 -> 3, scheduler 2*GBS -> 3*GBS assert trainer._resume_reference is None +def test_miles_restore_compares_the_landed_cursor_not_the_cut(tmp_path): + """T-REV2: the watermark is compared with the cursor the driver actually restored.""" + from tests.rl_cut_fakes import make_rank + + RankGroup, _context, _expect, _trained_rank, _trainer = _miles() + _trainer(RankGroup([_trained_rank()])).save_cut(epoch=1, context=_context(tmp_path)) + trainer = _trainer(RankGroup([make_rank(9)])) + trainer.restore_cut("cut-a", epoch=1, root=str(tmp_path), expect=_expect()) + with pytest.raises(ResumeAssertionError, match="data watermark 0 != saved 4"): + trainer.confirm_resumed_data_cursor({"sample_offset": 0}) # data source restarted at 0 + trainer.confirm_resumed_data_cursor({"sample_offset": 4}) + + +def test_run_wide_reference_for_a_dp_change(): + from yeto.rl.engine.resume_manifest import GroupSteps, OptimizerResumeState, restore_problems, run_wide_reference + + def state(step, progress=6): + return OptimizerResumeState({"named": GroupSteps(step, step)}, progress, (1e-3,), {}) + + saved = ResumeManifest(3, 12, 2, {"dp0": state(3), "dp1": state(3)}) + ref = run_wide_reference(saved, ["a", "b", "c"]) + assert sorted(ref.ranks) == ["a", "b", "c"] and ref.global_step == 3 + now = ResumeManifest(3, 12, 2, {"a": state(3), "b": state(0), "c": state(3, progress=0)}) + problems = restore_problems(ref, now, check_data_watermark=False) + assert any("b/named: Adam step 0 != expected 3 (Adam step was reset)" in p for p in problems) + assert any("c: scheduler progress 0 != saved 6 (scheduler was reset)" in p for p in problems) + with pytest.raises(ResumeAssertionError, match="disagree"): + run_wide_reference(ResumeManifest(3, 12, 2, {"dp0": state(3), "dp1": state(2)}), ["a"]) + + def test_miles_restore_refuses_a_reset_adam_step(tmp_path): from tests.rl_cut_fakes import make_rank @@ -374,6 +408,7 @@ def test_miles_first_step_without_advance_is_refused(tmp_path): fresh = make_rank(9) trainer = _trainer(RankGroup([fresh])) trainer.restore_cut("cut-a", epoch=1, root=str(tmp_path), expect=_expect()) + trainer.confirm_resumed_data_cursor({"sample_offset": 4}) with pytest.raises(ResumeAssertionError, match="Adam step 2 != expected 3"): trainer._check_first_step_after_resume(1) # no optimizer step was taken diff --git a/tests/test_rl_trainer_reshard.py b/tests/test_rl_trainer_reshard.py index 4fc3e8ca2..d2aaea624 100644 --- a/tests/test_rl_trainer_reshard.py +++ b/tests/test_rl_trainer_reshard.py @@ -354,3 +354,26 @@ def test_dp_shards_must_agree_on_the_miles_counters(tmp_path): _trainer(new).restore_cut_resharded("c", epoch=1, root=str(tmp_path), expect=_expect(2), plan=_plan(2, 1), certified=[SPEC_SHA]) assert new[0].weight_updater.weight_version == 0 # refused before any write + + +@pytest.mark.parametrize("src,dst", [(1, 2), (2, 1)]) +def test_reshard_checks_the_run_wide_resume_counters(tmp_path, src, dst): + """S25 N7.6: a DP change still checks global step, every rank's Adam step and the + scheduler against the cut, and the first step and data cursor after it.""" + from yeto.rl.engine.resume_manifest import ResumeAssertionError + + _trainer(_trained(src)).save_cut(epoch=1, context=_context(tmp_path, "c")) + new = make_world(dst, seed=99) + trainer = _trainer(new) + trainer.restore_cut_resharded("c", epoch=1, root=str(tmp_path), expect=_expect(src), + plan=_plan(src, dst), certified=[SPEC_SHA]) + reference = trainer._resume_reference + assert len(reference.ranks) == dst and reference.global_step == 2 + assert {s.max for st in reference.ranks.values() for s in st.adam_steps.values()} == {2} + with pytest.raises(ResumeAssertionError, match="never compared"): + trainer._check_first_step_after_resume(1) + trainer._resume_reference = reference + trainer.confirm_resumed_data_cursor({"sample_offset": reference.data_watermark}) + train_step(new, _batches()[2]) + trainer._check_first_step_after_resume(1) + diff --git a/tests/test_rl_verl_cut.py b/tests/test_rl_verl_cut.py new file mode 100644 index 000000000..0736c287a --- /dev/null +++ b/tests/test_rl_verl_cut.py @@ -0,0 +1,273 @@ +"""S25 N7.6: verl/FSDP2 cut save + restore with the resume assertions. CPU only. + +A verl round is ``K`` optimizer steps (PPO mini-batches) and one LR-scheduler +step, as in verl's ``update_actor``. Checked: + +* FSDP2 over gloo, 2 ranks: "3 rounds, cut, fresh trainer, restore, 2 rounds" + equals "5 rounds" bit for bit (parameters and Adam state on every rank); +* a restored Adam step reset to 0 on one rank is refused on every rank; +* single process through :class:`VerlCutTrainer`: the data watermark is compared + with the cursor that landed, the first round must confirm it and advance every + Adam step by ``K`` and the scheduler by 1; a cut without a resume section, a + non-fresh trainer and a DP change are refused. +""" + +from __future__ import annotations + +import hashlib +import json +import os +import socket +from types import SimpleNamespace + +import pytest +import torch + +from yeto.rl.adapters.verl import cut as vcut +from yeto.rl.adapters.verl.resume_check import TrainHandles, handles_rank_state +from yeto.rl.engine.cut import ( + MANIFEST_NAME, + AlgorithmIdentity, + CutContext, + CutError, + CutProgress, + RestoreExpectation, + canonical_json, + commit_manifest, + load_manifest, +) +from yeto.rl.engine.resume_manifest import ResumeAssertionError + +TOTAL, CUT, K = 5, 3, 2 # rounds, cut after round CUT, optimizer steps per round +ALGO = AlgorithmIdentity("a" * 64) +FINGERPRINT = "verl@test" +PROMPTS_PER_ROUND = 4 + + +def _model() -> torch.nn.Module: + torch.manual_seed(0) + return torch.nn.Sequential(torch.nn.Linear(8, 16), torch.nn.Tanh(), torch.nn.Linear(16, 1)) + + +def _handles(model) -> TrainHandles: + opt = torch.optim.AdamW(model.parameters(), lr=1e-2, weight_decay=0.01) + return TrainHandles(model, opt, torch.optim.lr_scheduler.LambdaLR(opt, lambda s: 1.0 / (1 + s))) + + +def _round(h: TrainHandles, r: int, rank: int = 0) -> None: + for mini in range(K): + g = torch.Generator().manual_seed(1000 * rank + K * r + mini) + x, y = torch.randn(4, 8, generator=g), torch.randn(4, 1, generator=g) + h.optimizer.zero_grad(set_to_none=True) + torch.nn.functional.mse_loss(h.module(x), y).backward() + h.optimizer.step() + h.scheduler.step() + + +def _context(root, rounds=CUT, cut_id="cut-v") -> CutContext: + return CutContext( + root=str(root), cut_id=cut_id, backend_fingerprint=FINGERPRINT, + progress=CutProgress(local_step=rounds, scheduler_samples=rounds, global_batch_size=1, + next_rollout_id=rounds, policy_version=rounds, policy_hash="h"), + algorithm=ALGO, data={"sample_offset": PROMPTS_PER_ROUND * rounds}, + ledger={"carried_over": 0, "ready_unconsumed": 0}, outer={"settled": True}) + + +def _expect(world=1, rounds=CUT) -> RestoreExpectation: + return RestoreExpectation(algorithm=ALGO, layout=vcut.layout_for(world), backend_fingerprint=FINGERPRINT, + local_step=rounds, policy_version=rounds, epoch=1) + + +def _local(t): + return t.to_local() if hasattr(t, "to_local") else t + + +def _assert_same_train_state(a: TrainHandles, b: TrainHandles) -> None: + for (name, p), (_, q) in zip(a.module.named_parameters(), b.module.named_parameters()): + assert torch.equal(_local(p), _local(q)), name + for p, q in zip(a.module.parameters(), b.module.parameters()): + for key in ("exp_avg", "exp_avg_sq", "step"): + assert torch.equal(_local(a.optimizer.state[p][key]), _local(b.optimizer.state[q][key])), key + assert a.scheduler.last_epoch == b.scheduler.last_epoch + + +def _reset_steps(h: TrainHandles) -> None: + for state in h.optimizer.state.values(): + state["step"] = torch.zeros_like(state["step"]) + + +# ---------------------------------------------------------------- FSDP2, 2 ranks + + +def _free_port() -> int: + with socket.socket() as s: + s.bind(("127.0.0.1", 0)) + return s.getsockname()[1] + + +def _fsdp2_worker(rank: int, world: int, port: int, root: str) -> None: + import torch.distributed as dist + from torch.distributed.fsdp import fully_shard + + os.environ.update(MASTER_ADDR="127.0.0.1", MASTER_PORT=str(port)) + dist.init_process_group("gloo", rank=rank, world_size=world) + try: + def build() -> TrainHandles: + model = _model() + for layer in model: + if isinstance(layer, torch.nn.Linear): + fully_shard(layer) + fully_shard(model) + return _handles(model) + + def gathered(value): + out = [None] * world + dist.all_gather_object(out, value) + return out + + def trainer_over(h: TrainHandles, *, restore=vcut.restore_rank) -> vcut.VerlCutTrainer: + # call_ranks: every rank runs fn on its own handles, all see every result + return vcut.VerlCutTrainer( + lambda fn, *args: gathered(fn(h, *args)), optimizer_steps_per_round=K, + fns={"restore": lambda hh, d, f, c: restore(hh, directory=d, files=f, cut_id=c), + "resume": vcut.resume_rank, "layout": lambda hh: vcut.remote_rank_layout()}) + + continuous = build() + for r in range(TOTAL): + _round(continuous, r, rank) + + first = build() + for r in range(CUT): + _round(first, r, rank) + context = _context(root) + results = gathered(vcut.save_rank(first, directory=os.path.join(root, context.cut_id), + cut_id=context.cut_id)) + if rank == 0: + commit_manifest(root, vcut.build_manifest(context, results, epoch=1)) + dist.barrier() + + resumed = build() + trainer = trainer_over(resumed) + trainer.restore_cut(context.cut_id, epoch=1, root=root, expect=_expect(world)) + trainer.confirm_resumed_data_cursor({"sample_offset": PROMPTS_PER_ROUND * CUT}) + _round(resumed, CUT, rank) + trainer.check_first_step(1) + for r in range(CUT + 1, TOTAL): + _round(resumed, r, rank) + _assert_same_train_state(continuous, resumed) + + def resetting_on_rank0(h, *, directory, files, cut_id): + result = vcut.restore_rank(h, directory=directory, files=files, cut_id=cut_id) + if rank == 0: # the PPU fault, on one rank only + _reset_steps(h) + result["resume"] = handles_rank_state(h).to_dict() + return result + + with pytest.raises(ResumeAssertionError, match=r"rank0/group0: Adam step 0 != expected 6 \(Adam step was reset\)"): + trainer_over(build(), restore=resetting_on_rank0).restore_cut( + context.cut_id, epoch=1, root=root, expect=_expect(world)) + open(os.path.join(root, f"ok{rank}"), "w").close() + finally: + dist.destroy_process_group() + + +def test_fsdp2_gloo_two_ranks_cut_resume_equals_continuous(tmp_path): + import torch.multiprocessing as mp + + pytest.importorskip("torch.distributed.fsdp") + if not torch.distributed.is_gloo_available(): + pytest.skip("gloo backend not built") + world = 2 + mp.spawn(_fsdp2_worker, args=(world, _free_port(), str(tmp_path)), nprocs=world, join=True) + assert sorted(p for p in os.listdir(tmp_path) if p.startswith("ok")) == [f"ok{r}" for r in range(world)] + manifest = load_manifest(tmp_path, "cut-v") + assert [f.path for f in manifest.files] == ["rank00000.pt", "rank00001.pt"] + assert manifest.runtime["layout"] == {"world": 2, "dp": 2} + assert manifest.resume["global_step"] == CUT and manifest.resume["data_watermark"] == PROMPTS_PER_ROUND * CUT + + +# ---------------------------------------------------------------- single process, VerlCutTrainer + + +def _single(tmp_path): + """(cut written after CUT rounds, a fresh trainer over a fresh worker, the fresh handles).""" + trained = _handles(_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)) + fresh = _handles(_model()) + fresh_worker = SimpleNamespace(module=fresh.module, optimizer=fresh.optimizer, lr_scheduler=fresh.scheduler) + return vcut.VerlCutTrainer(lambda fn, *a: [fn(fresh_worker, *a)], optimizer_steps_per_round=K), fresh + + +def test_single_rank_resume_checks_watermark_and_first_round(tmp_path): + trainer, fresh = _single(tmp_path) + trainer.restore_cut("cut-v", epoch=1, root=str(tmp_path), expect=_expect()) + with pytest.raises(ResumeAssertionError, match="data watermark 8 != saved 12"): + trainer.confirm_resumed_data_cursor({"sample_offset": 8}) + _round(fresh, CUT) + with pytest.raises(ResumeAssertionError, match="never compared"): + trainer.check_first_step(1) + trainer.confirm_resumed_data_cursor({"sample_offset": 12}) + trainer.check_first_step(1) # Adam 6 -> 8, scheduler 3 -> 4 + trainer.check_first_step(1) # nothing pending any more + + +def test_single_rank_first_round_must_advance_by_k(tmp_path): + trainer, fresh = _single(tmp_path) + trainer.restore_cut("cut-v", epoch=1, root=str(tmp_path), expect=_expect()) + trainer.confirm_resumed_data_cursor({"sample_offset": 12}) + fresh.scheduler.step() # the scheduler moved, no optimizer step was taken + with pytest.raises(ResumeAssertionError, match="Adam step 6 != expected 8"): + trainer.check_first_step(1) + + +def test_single_rank_refusals(tmp_path): + trainer, fresh = _single(tmp_path) + with pytest.raises(CutError, match="parallel layout"): + trainer.restore_cut("cut-v", epoch=1, root=str(tmp_path), expect=_expect(world=2)) + trainer.restore_cut("cut-v", epoch=1, root=str(tmp_path), expect=_expect()) + with pytest.raises(CutError, match="freshly built"): # the optimizer now has state + trainer.restore_cut("cut-v", epoch=1, root=str(tmp_path), expect=_expect()) + + path = tmp_path / "cut-v" / MANIFEST_NAME + raw = json.loads(path.read_text()) + raw.pop("resume") + body = {k: v for k, v in raw.items() if k != "manifest_sha256"} + raw["manifest_sha256"] = hashlib.sha256(canonical_json(body).encode()).hexdigest() + path.write_text(json.dumps(raw)) + again, _ = _single(tmp_path / "other") + with pytest.raises(CutError, match="no resume manifest"): + again.restore_cut("cut-v", epoch=1, root=str(tmp_path), expect=_expect()) + + +def test_cut_refuses_wrong_progress_units(tmp_path): + h = _handles(_model()) + _round(h, 0) + worker = SimpleNamespace(module=h.module, optimizer=h.optimizer, lr_scheduler=h.scheduler) + context = _context(tmp_path) + bad = CutContext(**{**context.__dict__, "progress": CutProgress( + local_step=2, scheduler_samples=8, global_batch_size=4, next_rollout_id=2, policy_version=2, + policy_hash="h")}) + with pytest.raises(CutError, match="global_batch_size must be 1"): + vcut.VerlCutTrainer(lambda fn, *a: [fn(worker, *a)], optimizer_steps_per_round=K).save_cut( + epoch=1, context=bad) + assert not (tmp_path / "cut-v").exists() + + +def test_optimizer_steps_per_round_from_config(): + config = SimpleNamespace(actor_rollout_ref=SimpleNamespace(actor={"ppo_epochs": 2, "ppo_mini_batch_size": 4}), + data=SimpleNamespace(train_batch_size=16)) + config.actor_rollout_ref.actor = _Cfg(config.actor_rollout_ref.actor) + assert vcut.optimizer_steps_per_round(config) == 8 + config.data.train_batch_size = 10 + with pytest.raises(CutError, match="whole number"): + vcut.optimizer_steps_per_round(config) + + +class _Cfg(dict): + """Attribute + ``get`` access, like an omegaconf node.""" + + __getattr__ = dict.__getitem__ diff --git a/yeto/rl/adapters/miles/cut_plugin.py b/yeto/rl/adapters/miles/cut_plugin.py index 2f17dac02..a1321114c 100644 --- a/yeto/rl/adapters/miles/cut_plugin.py +++ b/yeto/rl/adapters/miles/cut_plugin.py @@ -1035,6 +1035,9 @@ def _restore_resharded(actor, *, directory, files, cut_id, source_dp, rng_policy "rng_digest": state_digest(backend.capture_rng()), "optimizer_names": sorted(export.get("entries", {})), "adapter_names": sorted(names), + # S25 N7.6: run-wide resume assertions (Adam step, scheduler, LRs) at the new DP + "path": shard_name(coord), + "resume": resume_rank_state(actor, export), } diff --git a/yeto/rl/adapters/miles/trainer.py b/yeto/rl/adapters/miles/trainer.py index 3d7e5ebe2..ff548c45a 100644 --- a/yeto/rl/adapters/miles/trainer.py +++ b/yeto/rl/adapters/miles/trainer.py @@ -293,6 +293,9 @@ def __init__( self.optimizer_steps_total = 0 # S25 N7.6: the restored cut's resume manifest until the first step after it is checked self._resume_reference: Any = None + # ... and whether the data cursor the driver restored was compared with it + # (confirm_resumed_data_cursor); the first step refuses an unconfirmed one. + self._resume_cursor_confirmed = True # A LoRA actor installs its step recorders through the state plugins # the publisher calls before the first step; a full-parameter actor # calls none of them, so the trainer installs the recorders itself. @@ -410,9 +413,13 @@ def _rank_resume_states(self, results: list[Mapping[str, Any]], what: str) -> di def _check_first_step_after_resume(self, steps: int) -> None: """The first round after ``restore_cut`` advanced every rank's Adam step and the LR scheduler by exactly ``steps`` optimizer steps (S25 N7.6; A-TRN report #8).""" - from yeto.rl.engine.resume_manifest import ResumeManifest, require_first_step + from yeto.rl.engine.resume_manifest import ResumeAssertionError, ResumeManifest, require_first_step from .cut_plugin import RESUME_STATE + if not self._resume_cursor_confirmed: + raise ResumeAssertionError( + "first step after resume, but the restored data cursor was never compared with the cut's " + "data watermark (confirm_resumed_data_cursor not called); refusing to continue") saved, self._resume_reference = self._resume_reference, None results = [dict(r) for r in self._run(self._actor.run_plugin(RESUME_STATE, {}))] now = ResumeManifest( @@ -423,6 +430,19 @@ def _check_first_step_after_resume(self, steps: int) -> None: ) require_first_step(saved, now, steps=steps) + def confirm_resumed_data_cursor(self, cursor: Mapping[str, int]) -> None: + """Compare the data cursor the run actually resumed from with the restored + cut's data watermark (S25 N7.6). Called by whoever restored the cursor + after ``restore_cut`` / ``restore_cut_resharded``: the driver after its + seek landed (restart), or the trainer rebuild after its live re-read. + A no-op when no restore is pending its first step.""" + from yeto.rl.engine.resume_manifest import require_data_watermark + + if self._resume_reference is None: + return + require_data_watermark(self._resume_reference, int(cursor["sample_offset"])) + self._resume_cursor_confirmed = True + def _install_actor_recorders(self) -> None: """Record every optimizer step's applied LR and grad norm in the actor processes, once per trainer; without them ``_applied_lrs`` finds no @@ -940,12 +960,15 @@ def restore_cut(self, cut_id: str, *, epoch: int, root: str, expect: Any, require(resume_problems, "restored optimizer does not continue the cut; refusing to resume") self._restore_critic_cut(cut_dir(root, cut_id) / "critic", manifest) self._resume_reference = reference + self._resume_cursor_confirmed = False return manifest def _restored_resume(self, manifest: Any, expect: Any, results: list[Mapping[str, Any]]) -> tuple[Any, list[str]]: """The cut's resume manifest and every resume assertion the restored ranks fail - (S25 N7.6; A-TRN report #8: a silently reset Adam step).""" + (S25 N7.6; A-TRN report #8: a silently reset Adam step). The data watermark is + not checked here: the data cursor is restored after this call, by the driver + or the trainer rebuild, which then call :meth:`confirm_resumed_data_cursor`.""" from yeto.rl.engine.cut import CutError from yeto.rl.engine.resume_manifest import ResumeManifest, restore_problems @@ -955,11 +978,30 @@ def _restored_resume(self, manifest: Any, expect: Any, saved = ResumeManifest.from_dict(manifest.resume) now = ResumeManifest( global_step=int(expect.local_step), - data_watermark=int(manifest.data["sample_offset"]), # the cursor the driver restores + data_watermark=saved.data_watermark, # unused: check_data_watermark=False scheduler_increment=saved.scheduler_increment, ranks=self._rank_resume_states(results, "restore_cut"), ) - return saved, restore_problems(saved, now) + return saved, restore_problems(saved, now, check_data_watermark=False) + + def _resharded_resume(self, manifest: Any, expect: Any, results: list[Mapping[str, Any]]) -> Any: + """Resume assertions for a DP change (S25 N7.6). Shard names and optimizer + ranges change, so per-rank moment fingerprints have no counterpart; the + run-wide counters (global step, every rank's Adam step, LR scheduler + progress, applied LRs) must still continue the cut. Moments are covered by + the plugin's slice check; the data watermark by confirm_resumed_data_cursor.""" + from yeto.rl.engine.cut import CutError + from yeto.rl.engine.resume_manifest import ResumeManifest, require, restore_problems, run_wide_reference + + if not manifest.resume: + raise CutError(f"cut {manifest.cut_id!r} has no resume manifest (written before S25 N7.6); " + "refusing to resume without the optimizer-continuity checks") + ranks = self._rank_resume_states(results, "restore_cut_resharded") + reference = run_wide_reference(ResumeManifest.from_dict(manifest.resume), ranks) + now = ResumeManifest(int(expect.local_step), reference.data_watermark, reference.scheduler_increment, ranks) + require(restore_problems(reference, now, check_data_watermark=False), + "resharded optimizer does not continue the cut; refusing to resume") + return reference def restore_cut_resharded(self, cut_id: str, *, epoch: int, root: str, expect: Any, plan: Any, certified: Any = None, shared_filesystem: bool = True) -> dict[str, Any]: @@ -1024,10 +1066,10 @@ def restore_cut_resharded(self, cut_id: str, *, epoch: int, root: str, expect: A c = r["coord"] if expected.get((c["tp"], c["pp"], c["dp"])) != r["rng"]: raise CutError(f"rank {c}: RNG {r['rng']} does not follow the recorded mapping") + reference = self._resharded_resume(manifest, expect, results) self._reshard_plan = plan # every later batch is guarded (batch_guard_problems) - # DP change: shard names and optimizer ranges change, so the per-rank resume - # manifest has no counterpart; the resharded restore keeps its own slice checks. - self._resume_reference = None + self._resume_reference = reference + self._resume_cursor_confirmed = False return {"manifest": manifest, "plan": plan.to_dict(), "rng_mapping": mapping, "full_state_digests": {f"tp{k[0]}_pp{k[1]}": g[0]["full_state_digest"] for k, g in sorted(by_group.items())}, diff --git a/yeto/rl/adapters/miles/trainer_rebuild.py b/yeto/rl/adapters/miles/trainer_rebuild.py index 124141796..8d770c60d 100644 --- a/yeto/rl/adapters/miles/trainer_rebuild.py +++ b/yeto/rl/adapters/miles/trainer_rebuild.py @@ -118,6 +118,14 @@ def live_cursor(rollout: Any, when: str) -> dict[str, int]: return dict(cursor) +def confirm_resumed_cursor(trainer: Any, cursor: Mapping[str, int]) -> None: + """Hand the data cursor the rebuilt trainer continues from to its resume + assertions (S25 N7.6: compared with the cut's data watermark).""" + confirm = getattr(trainer, "confirm_resumed_data_cursor", None) + if callable(confirm): + confirm(cursor) + + def rebuild_preconditions(args: Any) -> list[str]: out = [] if getattr(args, "requested_load", None) is not None: @@ -276,6 +284,7 @@ def rebuild_same_shape( if now != layout: raise RuntimeError(f"layout changed across a same-shape rebuild: {layout} -> {now}") manifest = restore() + confirm_resumed_cursor(trainer, after) except Exception as error: # noqa: BLE001 attempts.append({"attempt": attempt, "stage": "after_rebuild", "error": repr(error)}) raise RecoveryRequired(f"trainer unusable after rebuild: {error}", attempts=attempts) from error @@ -388,7 +397,9 @@ def _attempt(stage: str, args: Any, view: Any, layout: Mapping[str, int], restor now = trainer.actual_layout() if now != dict(layout): raise RuntimeError(f"{stage}: running layout {now} != expected {dict(layout)}") - return generation, restore() + restored = restore() + confirm_resumed_cursor(trainer, after) + return generation, restored fallback_view = old_view try: diff --git a/yeto/rl/adapters/verl/cut.py b/yeto/rl/adapters/verl/cut.py new file mode 100644 index 000000000..785c479e3 --- /dev/null +++ b/yeto/rl/adapters/verl/cut.py @@ -0,0 +1,455 @@ +"""verl / FSDP2 cut: save at a safe point, restore into a fresh trainer, prove the +optimizer continues (S25 N7.6; same cut semantics as the Miles path). + +A verl cut is a neutral :class:`~yeto.rl.engine.cut.CutManifest` (same +``progress`` / ``data`` / ``ledger`` / ``outer`` sections, same ``resume`` +section, same immutability and checksum rules) over one yeto-owned file per +FSDP rank. verl itself is not changed and its own checkpoint manager is not +used: a rank file holds exactly the train state yeto resumes, so the CPU test +(FSDP2 over gloo) runs the same code as a GPU rank. + +Rank file ``rank.pt`` (``torch.save``): + +* ``model`` -- trainable parameters (frozen base weights are skipped: LoRA), as + the rank's local FSDP2 shards (``DTensor.to_local``) on CPU; +* ``optim`` -- the optimizer state dict keyed by parameter name, DTensor + entries stored as local shards; ``dtensor`` lists which entries were DTensors; +* ``scheduler`` -- ``lr_scheduler.state_dict()``; ``rng`` -- CPU (and CUDA) RNG. + +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 +(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 +``update_actor`` call). verl steps the LR scheduler once per round and the +optimizer once per PPO mini-batch, so ``CutProgress`` uses +``global_batch_size=1`` (scheduler progress == rounds) and the first step after +a restore expects every Adam step to advance by ``optimizer_steps_per_round``. + +Driver-side functions are pure (lists of rank results in, manifest/problems +out); :class:`VerlCutTrainer` binds them to a ``call_ranks`` callable (verl: +``VerlIsland.call_workers``). +""" + +from __future__ import annotations + +import hashlib +import json +import os +from collections.abc import Callable, Mapping, Sequence +from pathlib import Path +from typing import Any + +from yeto.rl.engine.cut import ( + CutContext, + CutError, + CutFile, + CutManifest, + RestoreExpectation, + commit_manifest, + context_problems, + cut_dir, + fsync_dir, + fsync_file, + sha256_file, + verify_cut, +) +from yeto.rl.engine.resume_manifest import ( + OptimizerResumeState, + ResumeAssertionError, + ResumeManifest, + require, + require_data_watermark, + require_first_step, + require_saveable, + restore_problems, + scheduler_progress_of, +) + +from .resume_check import TrainHandles, handles_rank_state, rank_label, train_handles + +RANK_FILE_SCHEMA = "yeto.verl_cut_rank/v1" +RNG_POLICY = "restore_per_rank" +DATA_CURSOR_FIELDS = ("sample_offset",) # verl data cursor (data_cursor.VerlDataCursor) + + +def rank_file(rank: int) -> str: + return f"rank{int(rank):05d}.pt" + + +def layout_for(world: int) -> dict[str, int]: + """FSDP2 actor layout: every rank is a data-parallel shard.""" + return {"world": int(world), "dp": int(world)} + + +# -------------------------------------------------------------------------- +# Rank side (inside each FSDP rank) +# -------------------------------------------------------------------------- + + +def _state_options() -> Any: + from torch.distributed.checkpoint.state_dict import StateDictOptions + + return StateDictOptions(full_state_dict=False, ignore_frozen_params=True) + + +def _is_dtensor(value: Any) -> bool: + return hasattr(value, "to_local") and hasattr(value, "placements") + + +def _to_local(value: Any) -> Any: + import torch + + if _is_dtensor(value): + value = value.to_local() + return value.detach().to("cpu").clone() if isinstance(value, torch.Tensor) else value + + +def _local_state(handles: TrainHandles) -> tuple[dict[str, Any], dict[str, Any], list[str]]: + """(model, optim, dtensor keys): the rank's local trainable state, CPU tensors only.""" + from torch.distributed.checkpoint.state_dict import get_state_dict + + model_sd, optim_sd = get_state_dict(handles.module, handles.optimizer, options=_state_options()) + dtensors = sorted(f"model/{k}" for k, v in model_sd.items() if _is_dtensor(v)) + model = {k: _to_local(v) for k, v in model_sd.items()} + state = {} + for name, entry in optim_sd.get("state", {}).items(): + state[name] = {k: _to_local(v) for k, v in entry.items()} + dtensors += sorted(f"optim/{name}/{k}" for k, v in entry.items() if _is_dtensor(v)) + optim = {"state": state, "param_groups": [dict(g) for g in optim_sd.get("param_groups", [])]} + 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).""" + 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()) + return h.hexdigest() + + +def _rng_state() -> dict[str, Any]: + import torch + + rng = {"cpu": torch.get_rng_state()} + if torch.cuda.is_available(): + rng["cuda"] = torch.cuda.get_rng_state() + return rng + + +def _set_rng_state(rng: Mapping[str, Any]) -> None: + import torch + + torch.set_rng_state(rng["cpu"]) + if "cuda" in rng: + if not torch.cuda.is_available(): + raise CutError("cut rank file carries a CUDA RNG state but this rank has no CUDA device") + torch.cuda.set_rng_state(rng["cuda"]) + + +def _rank_and_world() -> tuple[int, int]: + import torch.distributed as dist + + if dist.is_available() and dist.is_initialized(): + return dist.get_rank(), dist.get_world_size() + return 0, 1 + + +def _rank_summary(handles: TrainHandles, *, rank: int, world: int, digest: str) -> dict[str, Any]: + resume = handles_rank_state(handles) + return { + "coord": {"rank": rank, "world": world}, + "scheduler_samples": scheduler_progress_of(handles.scheduler), + "has_optimizer_state": any(s.max > 0 for s in resume.adam_steps.values()) or bool(resume.moments), + "has_rng": True, + "state_digest": digest, + "resume": resume.to_dict(), + } + + +def save_rank(handles: TrainHandles, *, directory: str, cut_id: str) -> dict[str, Any]: + """Write this rank's file (fsynced) and return its summary for the manifest.""" + import torch + + rank, world = _rank_and_world() + model, optim, dtensors = _local_state(handles) + name = rank_file(rank) + path = Path(directory) / name + path.parent.mkdir(parents=True, exist_ok=True) + tmp = path.with_name(name + ".tmp") + torch.save({"schema": RANK_FILE_SCHEMA, "cut_id": cut_id, "rank": rank, "world": world, + "model": model, "optim": optim, "dtensor": dtensors, + "scheduler": handles.scheduler.state_dict(), "rng": _rng_state()}, tmp) + fsync_file(tmp) + os.replace(tmp, path) + fsync_dir(path.parent) + summary = _rank_summary(handles, rank=rank, world=world, digest=_digest(model, optim)) + return {"path": name, "sha256": sha256_file(path), "bytes": path.stat().st_size, **summary} + + +def _rebuild(local: Any, template: Any) -> Any: + """A saved local shard as a DTensor shaped like ``template`` (the fresh module's entry).""" + from torch.distributed.tensor import DTensor + + local = local.to(template.to_local().device) + return DTensor.from_local(local, template.device_mesh, template.placements, run_check=False, + shape=template.shape, stride=template.stride()) + + +def restore_rank(handles: TrainHandles, *, directory: str, files: Sequence[Mapping[str, Any]], + cut_id: str) -> dict[str, Any]: + """Load this rank's file into a freshly built module/optimizer; return the summary + of the loaded state (digest + resume state) for the driver-side comparison.""" + import torch + from torch.distributed.checkpoint.state_dict import get_state_dict, set_state_dict + + rank, world = _rank_and_world() + entry = next((CutFile.from_dict(f) for f in files if f["path"] == rank_file(rank)), None) + if entry is None: + raise CutError(f"cut {cut_id!r} has no file for rank {rank}") + if len(files) != world: + raise CutError(f"cut has {len(files)} rank files, this run has {world} ranks (DP change unsupported)") + path = Path(directory) / entry.path + if sha256_file(path) != entry.sha256: + raise CutError(f"rank {rank}: {entry.path} checksum differs from the manifest") + if any(state for state in handles.optimizer.state.values()): + raise CutError("restore_cut needs a freshly built trainer (the optimizer already has state)") + saved = torch.load(path, map_location="cpu", weights_only=False) + if saved.get("schema") != RANK_FILE_SCHEMA or saved.get("cut_id") != cut_id or saved.get("rank") != rank: + raise CutError(f"rank {rank}: {entry.path} is not this rank's file of cut {cut_id!r}") + + template, _ = get_state_dict(handles.module, handles.optimizer, options=_state_options()) + if set(template) != set(saved["model"]): + raise CutError(f"rank {rank}: trainable parameters differ from the cut: " + f"{sorted(set(template) ^ set(saved['model']))[:4]}") + marked = set(saved["dtensor"]) + + def value(key: str, local: Any, like: str) -> Any: + if key not in marked: + return local + if not _is_dtensor(template[like]): + raise CutError(f"rank {rank}: {key} was sharded at save but {like} is not sharded now") + return _rebuild(local, template[like]) + + model_sd = {k: value(f"model/{k}", v, k) for k, v in saved["model"].items()} + optim_sd = { + "state": {n: {k: value(f"optim/{n}/{k}", v, n) for k, v in e.items()} + for n, e in saved["optim"]["state"].items()}, + "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()) + 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))} + + +def resume_rank(handles: TrainHandles) -> dict[str, Any]: + """This rank's resume state of the live trainer (first-step check).""" + rank, _ = _rank_and_world() + return {"path": rank_file(rank), "resume": handles_rank_state(handles).to_dict()} + + + + +# Ray entry points: ``VerlIsland.call_workers(fn, *args)`` calls ``fn(worker, *args)``. +def remote_save_rank(worker: Any, directory: str, cut_id: str) -> dict[str, Any]: + return save_rank(train_handles(worker), directory=directory, cut_id=cut_id) + + +def remote_restore_rank(worker: Any, directory: str, files: list, cut_id: str) -> dict[str, Any]: + return restore_rank(train_handles(worker), directory=directory, files=files, cut_id=cut_id) + + +def remote_resume_rank(worker: Any) -> dict[str, Any]: + return resume_rank(train_handles(worker)) + + +def remote_rank_layout(worker: Any = None) -> dict[str, int]: # noqa: ARG001 - call_workers shape + _, world = _rank_and_world() + return layout_for(world) + + +def common_layout(layouts: Sequence[Mapping[str, int]]) -> dict[str, int]: + distinct = {json.dumps(dict(x), sort_keys=True) for x in layouts} + if len(distinct) != 1: + raise CutError(f"ranks report different layouts: {sorted(distinct)}") + return json.loads(distinct.pop()) + + +# -------------------------------------------------------------------------- +# Driver side (pure) +# -------------------------------------------------------------------------- + + +def _rank_states(results: Sequence[Mapping[str, Any]]) -> dict[str, OptimizerResumeState]: + """Rank results keyed ``rank`` (from the rank file name every result carries).""" + return {rank_label(int(Path(r["path"]).stem.removeprefix("rank"))): OptimizerResumeState.from_dict(r["resume"]) + for r in results} + + +def optimizer_steps_per_round(config: Any) -> int: + """Optimizer steps one verl ``update_actor`` takes: ``ppo_epochs`` x PPO mini-batches + (``data.train_batch_size / actor.ppo_mini_batch_size``, both in prompts). Read from + the verl config; not yet checked against a live verl run.""" + actor = config.actor_rollout_ref.actor + epochs = int(actor.get("ppo_epochs", 1) or 1) + batch, mini = int(config.data.train_batch_size), int(actor.ppo_mini_batch_size) + if mini <= 0 or batch % mini: + raise CutError(f"train_batch_size {batch} is not a whole number of ppo_mini_batch_size {mini}") + return epochs * (batch // mini) + + +def cut_resume_manifest(context: CutContext, results: Sequence[Mapping[str, Any]]) -> ResumeManifest: + manifest = ResumeManifest(global_step=int(context.progress.local_step), + data_watermark=int(context.data["sample_offset"]), + scheduler_increment=1, ranks=_rank_states(results)) + require_saveable(manifest) + return manifest + + +def build_manifest(context: CutContext, results: Sequence[Mapping[str, Any]], *, epoch: int) -> CutManifest: + """The cut manifest over every rank's summary (refuses an incomplete or inconsistent cut).""" + worlds = {int(r["coord"]["world"]) for r in results} + if len(worlds) != 1 or len(results) != next(iter(worlds)): + raise CutError(f"{len(results)} rank results for world sizes {sorted(worlds)}") + summaries = [{k: v for k, v in r.items() if k not in ("sha256", "bytes")} for r in results] + return CutManifest( + cut_id=context.cut_id, epoch=int(epoch), + runtime={"backend": "verl", "backend_fingerprint": context.backend_fingerprint, + "layout": layout_for(next(iter(worlds))), "rng_policy": RNG_POLICY}, + progress=context.progress, algorithm=context.algorithm, + data={**dict(context.data), "cursor_fields": list(DATA_CURSOR_FIELDS)}, + ledger=context.ledger, outer=context.outer, + in_flight=tuple(context.in_flight or ()), + resume=cut_resume_manifest(context, results).to_dict(), + files=tuple(CutFile(r["path"], r["sha256"], int(r["bytes"]), dict(r["coord"])) + for r in sorted(results, key=lambda r: r["path"])), + rank_summaries=tuple(sorted(summaries, key=lambda r: r["path"])), + ) + + +def restored_problems(manifest: CutManifest, expect: RestoreExpectation, + results: Sequence[Mapping[str, Any]]) -> tuple[ResumeManifest, list[str]]: + """The cut's resume manifest and every problem of the restored ranks: file set, + loaded-state digest, scheduler progress and the resume assertions (data + watermark excluded: compared later with the cursor that actually landed).""" + if not manifest.resume: + raise CutError(f"cut {manifest.cut_id!r} has no resume manifest; refusing to resume") + saved = ResumeManifest.from_dict(manifest.resume) + recorded = {s["path"]: s for s in manifest.rank_summaries} + problems = [] + if sorted(r["path"] for r in results) != sorted(recorded): + problems.append(f"restored ranks {sorted(r['path'] for r in results)} != cut {sorted(recorded)}") + for r in results: + s = recorded.get(r["path"]) + if s is None: + continue + if r["state_digest"] != s["state_digest"]: + problems.append(f"{r['path']}: loaded train state differs from the cut") + if r["scheduler_samples"] != s["scheduler_samples"]: + problems.append(f"{r['path']}: scheduler at {r['scheduler_samples']}, cut {s['scheduler_samples']}") + now = ResumeManifest(int(expect.local_step), saved.data_watermark, saved.scheduler_increment, + _rank_states(results)) + return saved, problems + restore_problems(saved, now, check_data_watermark=False) + + +def save_cut(call_ranks: Callable[..., list], *, epoch: int, context: CutContext, layout: Mapping[str, int], + save_fn: Callable = remote_save_rank) -> CutManifest: + """Check the context (nothing written on a problem), write every rank file, commit the manifest.""" + problems = context_problems( + cut_id=context.cut_id, progress=context.progress, algorithm=context.algorithm, + data={**dict(context.data), "cursor_fields": list(DATA_CURSOR_FIELDS)}, ledger=context.ledger, + outer=context.outer, runtime={"backend_fingerprint": context.backend_fingerprint, + "layout": dict(layout), "rng_policy": RNG_POLICY}, + in_flight=tuple(context.in_flight or ())) + if int(context.progress.global_batch_size) != 1: + problems.append("verl cut progress counts scheduler steps: global_batch_size must be 1") + if problems: + raise CutError("refusing a cut: " + "; ".join(problems)) + directory = cut_dir(context.root, context.cut_id) + results = [dict(r) for r in call_ranks(save_fn, str(directory), context.cut_id)] + manifest = build_manifest(context, results, epoch=epoch) + commit_manifest(context.root, manifest, check_files=context.shared_filesystem) + return manifest + + +class VerlCutTrainer: + """Cut verbs of the verl trainer group (``CuttableTrainer`` save/restore half). + + ``call_ranks(fn, *args)`` runs ``fn(worker, *args)`` on every FSDP rank and + returns their results. After :meth:`restore_cut` the caller that restores + the data cursor calls :meth:`confirm_resumed_data_cursor`; the trainer group + calls :meth:`check_first_step` after the first round.""" + + def __init__(self, call_ranks: Callable[..., list], *, optimizer_steps_per_round: int, + fns: Mapping[str, Callable] | None = None) -> None: + if int(optimizer_steps_per_round) <= 0: + raise ValueError("optimizer_steps_per_round must be positive") + self.call_ranks = call_ranks + self.optimizer_steps_per_round = int(optimizer_steps_per_round) + self.fns = {"save": remote_save_rank, "restore": remote_restore_rank, "resume": remote_resume_rank, + "layout": remote_rank_layout, **dict(fns or {})} + self.resume_reference: ResumeManifest | None = None + self.cursor_confirmed = True + + def actual_layout(self) -> dict[str, int]: + return common_layout(self.call_ranks(self.fns["layout"])) + + def save_cut(self, *, epoch: int, context: CutContext) -> str: + save_cut(self.call_ranks, epoch=epoch, context=context, layout=self.actual_layout(), + save_fn=self.fns["save"]) + return context.cut_id + + def restore_cut(self, cut_id: str, *, epoch: int, root: str, expect: RestoreExpectation, + shared_filesystem: bool = True) -> CutManifest: + """Restore into a freshly built trainer; any exception means RECOVERY_REQUIRED.""" + manifest = verify_cut(root, cut_id, expect, check_files=shared_filesystem) + if manifest.runtime.get("backend") != "verl": + raise CutError(f"cut {cut_id!r} was written by {manifest.runtime.get('backend')!r}, not verl") + if manifest.epoch > epoch: + raise CutError(f"cut epoch {manifest.epoch} is newer than the restoring epoch {epoch}") + files = [f.to_dict() for f in manifest.files] + results = [dict(r) for r in self.call_ranks(self.fns["restore"], str(cut_dir(root, cut_id)), files, cut_id)] + saved, problems = restored_problems(manifest, expect, results) + require(problems, "restored verl trainer does not continue the cut; refusing to resume") + self.resume_reference, self.cursor_confirmed = saved, False + return manifest + + def confirm_resumed_data_cursor(self, cursor: Mapping[str, int]) -> None: + """The data cursor the run resumed from vs the cut's data watermark.""" + if self.resume_reference is None: + return + require_data_watermark(self.resume_reference, int(cursor["sample_offset"])) + self.cursor_confirmed = True + + def check_first_step(self, rounds: int = 1) -> None: + """After the first ``rounds`` rounds on the restored state (no-op without a pending restore).""" + if self.resume_reference is None: + return + if not self.cursor_confirmed: + raise ResumeAssertionError("first round after resume, but the restored data cursor was never " + "compared with the cut's data watermark; refusing to continue") + saved, self.resume_reference = self.resume_reference, None + results = [dict(r) for r in self.call_ranks(self.fns["resume"])] + now = ResumeManifest(saved.global_step + rounds, saved.data_watermark, saved.scheduler_increment, + _rank_states(results)) + require_first_step(saved, now, steps=rounds, optimizer_steps=rounds * self.optimizer_steps_per_round) diff --git a/yeto/rl/adapters/verl/ports_impl.py b/yeto/rl/adapters/verl/ports_impl.py index dad0f2801..0f83b3f75 100644 --- a/yeto/rl/adapters/verl/ports_impl.py +++ b/yeto/rl/adapters/verl/ports_impl.py @@ -269,6 +269,38 @@ class VerlTrainerGroup: def __init__(self, island: VerlIsland): self.island = island self._metrics = None + self._cut_trainer = None + + # ---- cut verbs (S25 N7.6; .cut) ---- + def _cuts(self): + if self._cut_trainer is None: + from .cut import VerlCutTrainer, optimizer_steps_per_round + + self._cut_trainer = VerlCutTrainer( + self.island.call_workers, + optimizer_steps_per_round=optimizer_steps_per_round(self.island.trainer.config)) + return self._cut_trainer + + def layout(self) -> dict[str, int]: + return self.actual_layout() + + def actual_layout(self) -> dict[str, int]: + from .cut import common_layout, remote_rank_layout + + return common_layout(self.island.call_workers(remote_rank_layout)) + + def save_cut(self, *, epoch: int, context) -> str: + return self._cuts().save_cut(epoch=epoch, context=context) + + def restore_cut(self, cut_id: str, *, epoch: int, root: str, expect, shared_filesystem: bool = True): + return self._cuts().restore_cut(cut_id, epoch=epoch, root=root, expect=expect, + shared_filesystem=shared_filesystem) + + def confirm_resumed_data_cursor(self, cursor) -> None: + self._cuts().confirm_resumed_data_cursor(cursor) + + def reshard_problems(self, plan, **_: object) -> list[str]: + return ["verl FSDP2 cuts restore at the same world size only (no DP change)"] def onload(self) -> None: return None @@ -319,6 +351,8 @@ def train_step(self, batch: RolloutBatchHandle) -> LocalStepReceipt: train_step=batch.rollout_id, 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 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, diff --git a/yeto/rl/adapters/verl/resume_check.py b/yeto/rl/adapters/verl/resume_check.py index 8de6cf605..f4f008f05 100644 --- a/yeto/rl/adapters/verl/resume_check.py +++ b/yeto/rl/adapters/verl/resume_check.py @@ -1,10 +1,16 @@ """verl / FSDP2 binding of the resume assertions (S25 N7.6; A-TRN report #8). The checks live in :mod:`yeto.rl.engine.resume_manifest`; this module only maps -a verl FSDP actor worker onto them. verl's ``ActorRolloutRefWorker`` holds the -FSDP2-wrapped actor, its ``torch.optim`` optimizer and LR scheduler as -``actor_module_fsdp``, ``actor_optimizer`` and ``actor_lr_scheduler`` (attribute -names from verl's fsdp worker; not yet checked against a live verl build). +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`` + (``workers/engine/fsdp/transformer_impl.py``); +* the legacy fsdp worker: ``actor_module_fsdp`` / ``actor_optimizer`` / + ``actor_lr_scheduler``. + +Attribute names are read from the verl source; not yet checked on a live verl build. +The cut save/restore path that uses these handles is :mod:`.cut`. Each rank reads its own FSDP2 shard (``DTensor.to_local``), so the per-rank state is keyed by rank. Use: @@ -19,6 +25,7 @@ from __future__ import annotations +from dataclasses import dataclass from typing import Any from yeto.rl.engine.resume_manifest import ( @@ -31,21 +38,50 @@ torch_optimizer_state, ) +ENGINE_ATTRS = ("module", "optimizer", "lr_scheduler") WORKER_ATTRS = ("actor_module_fsdp", "actor_optimizer", "actor_lr_scheduler") +@dataclass(frozen=True) +class TrainHandles: + """The FSDP2 actor module, its ``torch.optim`` optimizer and LR scheduler on one rank.""" + + module: Any + optimizer: Any + scheduler: Any + + +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 + + +def train_handles(worker: Any) -> TrainHandles: + """Resolve the train handles of a verl worker (engine or legacy shape, see module doc).""" + if all(getattr(worker, a, None) is not None for a in WORKER_ATTRS): + return TrainHandles(*(getattr(worker, a) for a in WORKER_ATTRS)) + engine = _engine_of(worker) + missing = [a for a in ENGINE_ATTRS if getattr(engine, a, None) is None] + if missing: + raise ResumeAssertionError(f"verl worker {type(worker).__name__} lacks {missing} (engine) and " + f"{list(WORKER_ATTRS)} (legacy); cannot check the resumed optimizer") + return TrainHandles(*(getattr(engine, a) for a in ENGINE_ATTRS)) + + def rank_label(rank: int) -> str: return f"rank{int(rank)}" +def handles_rank_state(handles: TrainHandles, *, samples: int = DEFAULT_SAMPLES) -> OptimizerResumeState: + return torch_optimizer_state(handles.optimizer, handles.module.named_parameters(), handles.scheduler, + samples=samples) + + def worker_rank_state(worker: Any, *, samples: int = DEFAULT_SAMPLES) -> OptimizerResumeState: """This rank's optimizer/scheduler state of a verl FSDP actor worker.""" - missing = [a for a in WORKER_ATTRS if getattr(worker, a, None) is None] - if missing: - raise ResumeAssertionError(f"verl worker {type(worker).__name__} lacks {missing}; " - "cannot check the resumed optimizer") - return torch_optimizer_state(worker.actor_optimizer, worker.actor_module_fsdp.named_parameters(), - worker.actor_lr_scheduler, samples=samples) + return handles_rank_state(train_handles(worker), samples=samples) def worker_manifest(worker: Any, *, rank: int, global_step: int, data_watermark: int, diff --git a/yeto/rl/engine/cut.py b/yeto/rl/engine/cut.py index 0d6a645b9..cfaeabd35 100644 --- a/yeto/rl/engine/cut.py +++ b/yeto/rl/engine/cut.py @@ -373,7 +373,11 @@ def context_problems( out.append(f"invalid cut_id {cut_id!r}") if not algorithm.algorithm_spec_sha256: out.append("algorithm: algorithm_spec_sha256 missing") - for key in DATA_CURSOR_FIELDS: + # A backend whose data source has fewer position fields (verl: sample_offset + # only) names them in ``cursor_fields``; sample_offset (the data watermark) is + # always required. + fields = tuple(data.get("cursor_fields") or DATA_CURSOR_FIELDS) + for key in dict.fromkeys(("sample_offset", *fields)): if not isinstance(data.get(key), int): out.append(f"data: cursor field {key!r} missing") buffered = data.get("buffer_length") diff --git a/yeto/rl/engine/driver.py b/yeto/rl/engine/driver.py index 7b6c34236..0f8d28384 100644 --- a/yeto/rl/engine/driver.py +++ b/yeto/rl/engine/driver.py @@ -1741,6 +1741,9 @@ def _restore_data_cursor(self, start_rollout_id: int) -> None: if landed != cut_cursor: raise DriverError(f"restart at rollout {start_rollout_id}: data source seek " f"to the cut's cursor {cut_cursor} landed on {landed}") + confirm = getattr(self.trainer, "confirm_resumed_data_cursor", None) + if callable(confirm): # S25 N7.6: the landed cursor vs the cut's resume manifest + confirm(landed) self.emit("rl_data_cursor_restored", rollout_id=start_rollout_id, source="round_cut", data_cursor=dict(landed)) return diff --git a/yeto/rl/engine/resume_manifest.py b/yeto/rl/engine/resume_manifest.py index bd6785139..18680922a 100644 --- a/yeto/rl/engine/resume_manifest.py +++ b/yeto/rl/engine/resume_manifest.py @@ -277,13 +277,17 @@ def save_problems(manifest: ResumeManifest) -> list[str]: return out -def restore_problems(saved: ResumeManifest, now: ResumeManifest) -> list[str]: - """Restored state vs the saved one; every item must match exactly (LR within 1e-12 relative).""" +def restore_problems(saved: ResumeManifest, now: ResumeManifest, *, check_data_watermark: bool = True) -> list[str]: + """Restored state vs the saved one; every item must match exactly (LR within 1e-12 relative). + + ``check_data_watermark=False`` is for a backend whose data cursor is restored + after the optimizer (Miles: the driver seeks the rollout data source later); + it must then call :func:`data_watermark_problems` with the cursor that landed.""" out = [] if now.global_step != saved.global_step: out.append(f"global step {now.global_step} != saved {saved.global_step}") - if now.data_watermark != saved.data_watermark: - out.append(f"data watermark {now.data_watermark} != saved {saved.data_watermark}") + if check_data_watermark: + out += data_watermark_problems(saved, now.data_watermark) if now.scheduler_increment != saved.scheduler_increment: out.append(f"scheduler increment {now.scheduler_increment} != saved {saved.scheduler_increment}") problems, pairs = _rank_pairs(saved, now) @@ -306,16 +310,52 @@ def restore_problems(saved: ResumeManifest, now: ResumeManifest) -> list[str]: return out +def data_watermark_problems(saved: ResumeManifest, restored_offset: int) -> list[str]: + """The data position the run actually resumed from (``sample_offset`` of the + restored data cursor) vs the one recorded at the cut.""" + if int(restored_offset) != saved.data_watermark: + return [f"data watermark {int(restored_offset)} != saved {saved.data_watermark}"] + return [] + + +def run_wide_reference(saved: ResumeManifest, labels: Iterable[str]) -> ResumeManifest: + """Reference for a restore at a different data-parallel size (ranks re-sliced). + + Per-rank moment fingerprints have no counterpart there, but the run-wide + counters do: every saved rank must agree on the Adam step of each group, the + scheduler progress and the applied LRs, and every restored rank (``labels``) + must then show those values. Moments are left out (checked by the backend's + own slice comparison). Raises when the saved ranks disagree.""" + states = list(saved.ranks.values()) + if not states: + raise ResumeAssertionError("resume manifest has no rank state") + first = states[0] + for state in states[1:]: + same = (state.adam_steps == first.adam_steps and state.scheduler_progress == first.scheduler_progress + and len(state.lrs) == len(first.lrs) and all(_lr_close(a, b) for a, b in zip(state.lrs, first.lrs))) + if not same: + raise ResumeAssertionError("saved ranks disagree on Adam step / scheduler progress / LRs; " + "no run-wide reference for a resharded restore") + shared = OptimizerResumeState(first.adam_steps, first.scheduler_progress, first.lrs, {}) + return ResumeManifest(saved.global_step, saved.data_watermark, saved.scheduler_increment, + {str(label): shared for label in labels}, saved.schema) + + def first_step_problems(saved: ResumeManifest, now: ResumeManifest, *, steps: int = 1, + optimizer_steps: int | None = None, expected_lrs: Mapping[str, Sequence[float]] | None = None) -> list[str]: - """After ``steps`` optimizer steps on the restored state: every counter advanced by exactly ``steps``.""" + """After ``steps`` run steps on the restored state: the global step advanced by + ``steps``, the LR scheduler by ``steps * scheduler_increment`` and every Adam step + by ``optimizer_steps`` (default ``steps``; verl takes several optimizer steps, + one per mini-batch, per scheduler step).""" + adam_advance = steps if optimizer_steps is None else int(optimizer_steps) out = [] if now.global_step != saved.global_step + steps: out.append(f"global step {now.global_step} != saved {saved.global_step} + {steps}") problems, pairs = _rank_pairs(saved, now) out += problems for rank, s, n in pairs: - out += _step_problems(rank, s.adam_steps, n.adam_steps, steps) + out += _step_problems(rank, s.adam_steps, n.adam_steps, adam_advance) want = s.scheduler_progress + steps * saved.scheduler_increment if n.scheduler_progress != want: out.append(f"{rank}: scheduler progress {n.scheduler_progress} != expected {want}") @@ -338,9 +378,15 @@ def require_restored(saved: ResumeManifest, now: ResumeManifest) -> None: require(restore_problems(saved, now), "restored optimizer does not continue the cut; refusing to resume") +def require_data_watermark(saved: ResumeManifest, restored_offset: int) -> None: + require(data_watermark_problems(saved, restored_offset), "restored data cursor does not continue the cut") + + def require_first_step(saved: ResumeManifest, now: ResumeManifest, *, steps: int = 1, + optimizer_steps: int | None = None, expected_lrs: Mapping[str, Sequence[float]] | None = None) -> None: - require(first_step_problems(saved, now, steps=steps, expected_lrs=expected_lrs), + require(first_step_problems(saved, now, steps=steps, optimizer_steps=optimizer_steps, + expected_lrs=expected_lrs), "first step after resume did not advance the optimizer by exactly the steps taken") From ce429e9db1c0daa21cf91ea64b12d64e44615311 Mon Sep 17 00:00:00 2001 From: michael <78113554+chowhsu@users.noreply.github.com> Date: Mon, 19 Oct 2026 12:19:14 +0000 Subject: [PATCH 15/34] Quantize per trainer rank from stacked layouts; scheme-named releases (S25 N8.2, N8.3, N8.5b) - quant/stacked.py: trainer-native stacked experts (gate_up_proj [E,2I,H], down_proj [E,H,I]) map to the per-expert source names; the converter reads a stacked source directory directly, slicing each expert on read. - quant/shards.py: each rank quantizes the rows it holds into a part file; ranks merge disjoint output shards; rank 0 writes index and metadata. The result is byte-equal to the single-process converter (W4A8 and MXFP4). - adapters/miles/quant_shards.py (Megatron grouped and sequential expert names, EP offset, ETP=1) and adapters/verl/quant_shards.py (FSDP2 Shard(0) DTensors) produce the neutral source slices. - quant/schemes.py: one place for scheme names and release_dir_for(); MXFP4 releases go to policy-releases-mxfp4-g32-rtn, W4A8 keeps policy-releases-w4a8. Records add awq_source and expert_matrices; the releaser refuses a target directory of another scheme; PolicyFollower (scheme=...) refuses a release of another scheme before reading it. Co-Authored-By: Claude Opus 5.5 --- tests/test_rl_quant_shards.py | 382 +++++++++++++++++++++++++ tests/test_rl_ssh_trainer_island.py | 4 +- yeto/launch/ssh_trainer_island.py | 5 +- yeto/rl/adapters/miles/quant_shards.py | 74 +++++ yeto/rl/adapters/verl/quant_shards.py | 72 +++++ yeto/rl/engine/policy_release.py | 41 ++- yeto/rl/quant/__init__.py | 7 + yeto/rl/quant/release.py | 76 +++-- yeto/rl/quant/requant.py | 84 ++++-- yeto/rl/quant/schemes.py | 42 +++ yeto/rl/quant/shards.py | 203 +++++++++++++ yeto/rl/quant/stacked.py | 111 +++++++ 12 files changed, 1046 insertions(+), 55 deletions(-) create mode 100644 tests/test_rl_quant_shards.py create mode 100644 yeto/rl/adapters/miles/quant_shards.py create mode 100644 yeto/rl/adapters/verl/quant_shards.py create mode 100644 yeto/rl/quant/schemes.py create mode 100644 yeto/rl/quant/shards.py create mode 100644 yeto/rl/quant/stacked.py diff --git a/tests/test_rl_quant_shards.py b/tests/test_rl_quant_shards.py new file mode 100644 index 000000000..0350ab94f --- /dev/null +++ b/tests/test_rl_quant_shards.py @@ -0,0 +1,382 @@ +"""Stacked native source (N8.5b), per-rank re-quantization (N8.2) and the trainer adapters. + +Both schemes use the tiny artefacts of tests/test_rl_quant_requant.py (W4A8, R2 +AWQ) and tests/test_rl_quant_mxfp4.py (MXFP4 RTN). The reference output is the +single-process converter on the per-expert BF16 source; every other path must +write the same bytes. +""" + +from __future__ import annotations + +import json +import os +import re +import shutil +import subprocess +import sys +from pathlib import Path + +import pytest +import torch +import torch.distributed as dist +import torch.multiprocessing as mp + +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.verl.quant_shards import fsdp2_slices, shard0_rows +from yeto.rl.engine.policy_release import ( + RELEASE_DIR, + W4A8_RELEASE_DIR, + PolicyFollower, + PolicyReleaseStore, + SchemeMismatch, + release_scheme, +) +from yeto.rl.engine.resume import CheckpointStore +from yeto.rl.quant.awq import AwqSmoothing +from yeto.rl.quant.release import RequantReleaser +from yeto.rl.quant.requant import ( + INDEX_FILE, + MXFP4_G32_RTN, + W4A8_INT4_CHANNEL, + CheckpointLayout, + LayoutError, + finalize_checkpoint, + requantize_checkpoint, +) +from yeto.rl.quant.schemes import BF16, release_dir_for +from yeto.rl.quant.shards import ( + PART_DIR, + SourceSlice, + merge_parts, + requantize_distributed, + write_part, +) +from yeto.rl.quant.stacked import StackedSourceIndex, unstack_experts + +_PER_EXPERT = re.compile(r"^(?P.+\.mlp\.experts)\.(?P\d+)\.(?Pgate|up|down)_proj\.weight$") + + +def _build(tmp: Path, scheme: str) -> dict: + """Template, smoothing and the per-expert BF16 source of one scheme's tiny artefact.""" + tmp.mkdir(parents=True, exist_ok=True) + original = tmp / "original" + if scheme == "w4a8": + s = (0.8 + 0.4 * torch.rand(w4.H, generator=torch.Generator().manual_seed(7))).double() + scales = tmp / "awq-scales" + scales.mkdir() + torch.save(s, scales / f"{w4.NORM}.pt") + src = w4._source_tensors() + w4._write_ckpt(original, w4._reference_artefact(src, s), w4._shard_of) + (original / "config.json").write_text(json.dumps({"model_type": "glm5_next", "quantization_config": w4.QCONFIG})) + (original / "recipe.yaml").write_text(w4.RECIPE) + smoothing = AwqSmoothing.from_artefact(original / "recipe.yaml", scales) + else: + src = mx._source_tensors() + mx._write_ckpt(original, mx._reference_artefact(src), mx._shard_of) + (original / "config.json").write_text(json.dumps({"model_type": "glm5_next", "quantization_config": mx.QCONFIG})) + smoothing = AwqSmoothing.none() + template = tmp / "template" + CheckpointLayout.from_checkpoint(original).save_template(original, template) + source = tmp / "bf16" + w4._write_ckpt(source, src, lambda n: "model-00001-of-00001.safetensors") + reference = tmp / "reference" + requantize_checkpoint(source, template, reference, smoothing) + return {"src": src, "template": template, "smoothing": smoothing, "reference": reference, "tmp": tmp} + + +@pytest.fixture(params=["w4a8", "mxfp4"]) +def world(request, tmp_path): + return _build(tmp_path / request.param, request.param) + + +def _stack(src: dict[str, torch.Tensor]) -> dict[str, torch.Tensor]: + """The trainer-native layout: experts.gate_up_proj [E, 2I, H] and experts.down_proj [E, H, I].""" + groups: dict[str, dict[int, dict[str, torch.Tensor]]] = {} + out = {} + for name, t in src.items(): + m = _PER_EXPERT.match(name) + if m is None: + out[name] = t + else: + groups.setdefault(m["prefix"], {}).setdefault(int(m["e"]), {})[m["proj"]] = t + for prefix, experts in groups.items(): + order = sorted(experts) + assert order == list(range(len(order))) + out[prefix + ".gate_up_proj"] = torch.stack([torch.cat([experts[e]["gate"], experts[e]["up"]]) for e in order]) + out[prefix + ".down_proj"] = torch.stack([experts[e]["down"] for e in order]) + return out + + +def _assert_same_files(a: Path, b: Path) -> None: + names = sorted(p.name for p in a.iterdir()) + assert names == sorted(p.name for p in b.iterdir()) + for name in names: + assert (a / name).read_bytes() == (b / name).read_bytes(), name + + +# -- N8.5b: stacked native source -------------------------------------------------- + +def test_unstack_round_trips_the_per_expert_tensors(world): + stacked = _stack(world["src"]) + got = {} + for name, t in stacked.items(): + if name.endswith((".gate_up_proj", ".down_proj")): + got.update(unstack_experts(name, t)) + expected = {n: t for n, t in world["src"].items() if _PER_EXPERT.match(n)} + assert got.keys() == expected.keys() and all(torch.equal(got[n], expected[n]) for n in got) + + +def test_stacked_native_source_converts_to_the_same_bytes(world): + native = world["tmp"] / "native" + w4._write_ckpt(native, _stack(world["src"]), lambda n: "model-00001-of-00002.safetensors" + if "gate_up" in n else "model-00002-of-00002.safetensors") + out = world["tmp"] / "from-native" + requantize_checkpoint(native, world["template"], out, world["smoothing"]) + _assert_same_files(out, world["reference"]) + + +def test_source_with_stacked_and_per_expert_experts_is_rejected(): + names = ["a.mlp.experts.down_proj", "a.mlp.experts.0.down_proj.weight"] + with pytest.raises(ValueError, match="both stacked and per-expert"): + StackedSourceIndex.from_names(names, lambda n: (1, 4, 4)) + + +# -- N8.2: parts and merge ----------------------------------------------------------- + +def _row_split(src: dict[str, torch.Tensor], ranks: int) -> list[list[SourceSlice]]: + """Per-expert matrices round robin; other tensors of dim >= 2 split by rows; small ones on rank 0.""" + out: list[list[SourceSlice]] = [[] for _ in range(ranks)] + for i, (name, t) in enumerate(sorted(src.items())): + if _PER_EXPERT.match(name): + out[i % ranks].append(SourceSlice(name, t)) + elif t.dim() >= 2 and t.shape[0] >= ranks: + for r in range(ranks): + a, b = shard0_rows(t.shape[0], r, ranks) + out[r].append(SourceSlice(name, t[a:b], a)) + else: + out[0].append(SourceSlice(name, t)) + return out + + +@pytest.mark.parametrize("ranks", [2, 3]) +def test_parts_merge_to_the_single_process_bytes(world, ranks): + out = world["tmp"] / f"merged-{ranks}" + parts = out / PART_DIR + for r, slices in enumerate(_row_split(world["src"], ranks)): + write_part(parts, r, slices, world["template"], world["smoothing"]) + shards = list(CheckpointLayout.load(world["template"]).shards()) + for r in range(ranks): # each "rank" merges its own subset of shards + merge_parts(parts, world["template"], out, shards[r::ranks]) + finalize_checkpoint(world["template"], out) + shutil.rmtree(parts) + _assert_same_files(out, world["reference"]) + + +def test_merge_rejects_gaps_overlaps_and_missing_tensors(world): + src, tpl, sm = world["src"], world["template"], world["smoothing"] + big = next(n for n, t in sorted(src.items()) if t.dim() == 2 and t.shape[0] >= 4 and not _PER_EXPERT.match(n)) + rest = [SourceSlice(n, t) for n, t in src.items() if n != big] + rows = src[big].shape[0] + cases = { + "gap": [SourceSlice(big, src[big][: rows // 2])], + "overlap": [SourceSlice(big, src[big]), SourceSlice(big, src[big][1:], 1)], + "no rank gave": [], + } + for i, (msg, extra) in enumerate(cases.items()): + parts = world["tmp"] / f"bad-{i}" + write_part(parts, 0, rest, tpl, sm) + if extra: + write_part(parts, 1, extra, tpl, sm) + with pytest.raises(LayoutError, match=msg): + merge_parts(parts, tpl, world["tmp"] / f"bad-out-{i}") + + +def test_a_slice_that_does_not_fit_is_rejected_on_the_rank(world): + name, t = next((n, t) for n, t in world["src"].items() if t.dim() == 2 and not _PER_EXPERT.match(n)) + with pytest.raises(LayoutError, match="does not fit"): + write_part(world["tmp"] / "p", 0, [SourceSlice(name, t, row_start=1)], world["template"], world["smoothing"]) + + +# -- Megatron adapter ---------------------------------------------------------------- + +def _megatron_rank(src: dict[str, torch.Tensor], ep_rank: int, num_local: int, sequential: bool): + """Megatron names for the experts of one EP rank (layers keep their global index; offset 0).""" + params = [] + for name, t in src.items(): + m = _PER_EXPERT.match(name) + if m is None or m["proj"] == "up": + continue + e = int(m["e"]) + if e // num_local != ep_rank: + continue + layer = re.search(r"layers\.(\d+)\.", name).group(1) + local = e % num_local + fc = "1" if m["proj"] == "gate" else "2" + weight = torch.cat([t, src[name.replace("gate_proj", "up_proj")]]) if fc == "1" else t + key = (f"module.decoder.layers.{layer}.mlp.experts.local_experts.{local}.linear_fc{fc}.weight" if sequential + else f"module.decoder.layers.{layer}.mlp.experts.linear_fc{fc}.weight{local}") + params.append((key, weight)) + params.append(("module.decoder.layers.0.self_attention.linear_qkv.weight", torch.zeros(2, 2))) # skipped + return params + + +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)] + 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)) + if ep_rank == 0: + slices += list(hf_slices(non_expert)) + write_part(parts, ep_rank, slices, world["template"], world["smoothing"]) + merge_parts(parts, world["template"], out) + finalize_checkpoint(world["template"], out) + shutil.rmtree(parts) + _assert_same_files(out, world["reference"]) + + +def test_megatron_adapter_refuses_expert_tensor_parallelism_and_bad_indices(): + with pytest.raises(ValueError, match="expert tensor parallel"): + list(megatron_expert_slices([], ep_rank=0, num_local_experts=1, expert_tp_size=2)) + with pytest.raises(ValueError, match="local expert 3"): + list(megatron_expert_slices([("decoder.layers.0.mlp.experts.linear_fc2.weight3", torch.zeros(2, 2))], + ep_rank=0, num_local_experts=2)) + + +def test_megatron_names_map_to_global_layer_and_expert(): + fc1 = torch.arange(8.0).reshape(4, 2) + got = list(megatron_expert_slices([("decoder.layers.1.mlp.experts.linear_fc1.weight1", fc1)], + ep_rank=2, num_local_experts=4, layer_offset=10)) + assert [s.name for s in got] == ["model.language_model.layers.11.mlp.experts.9.gate_proj.weight", + "model.language_model.layers.11.mlp.experts.9.up_proj.weight"] + assert torch.equal(got[0].tensor, fc1[:2]) and torch.equal(got[1].tensor, fc1[2:]) + + +def test_shard0_rows_follows_torch_chunk(): + for rows in (0, 1, 5, 288): + for world_size in (1, 2, 3, 7): + chunks = torch.arange(rows).chunk(world_size) if rows else () + got = [shard0_rows(rows, r, world_size) for r in range(world_size)] + starts = [(int(c[0]), int(c[-1]) + 1) for c in chunks] + assert [g for g in got if g[1] > g[0]] == starts + + +# -- 2 gloo ranks, FSDP2 DTensors ------------------------------------------------------ + +def _gloo_worker(rank: int, world_size: int, init: str, case: str, out: str) -> None: + from torch.distributed.device_mesh import init_device_mesh + from torch.distributed.tensor import DTensor, Shard + + dist.init_process_group("gloo", init_method=init, rank=rank, world_size=world_size) + try: + w = _build(Path(case), Path(case).name) if rank == 0 else None + dist.barrier() + if rank != 0: # rank 0 built the files; the other ranks load the same inputs + w = _load(Path(case)) + mesh = init_device_mesh("cpu", (world_size,)) + params = [] + for name, t in sorted(_stack(w["src"]).items()): + if t.dim() >= 2: + a, b = shard0_rows(t.shape[0], rank, world_size) + local = t[a:b].clone() + params.append(("_fsdp." + name, DTensor.from_local(local, mesh, [Shard(0)], run_check=False, + shape=t.shape, stride=t.stride()))) + else: + params.append(("_fsdp." + name, t)) # replicated: rank 0 gives it + slices = fsdp2_slices(params, rank=rank, rename=lambda n: n.removeprefix("_fsdp.")) + requantize_distributed(slices, w["template"], Path(out), w["smoothing"], rank=rank, + world_size=world_size, barrier=dist.barrier) + finally: + dist.destroy_process_group() + + +def _load(case: Path) -> dict: + from safetensors.torch import load_file + + scheme = case.name + original = case / "original" + smoothing = (AwqSmoothing.from_artefact(original / "recipe.yaml", case / "awq-scales") if scheme == "w4a8" + else AwqSmoothing.none()) + return {"src": load_file(str(case / "bf16" / "model-00001-of-00001.safetensors")), + "template": case / "template", "smoothing": smoothing} + + +@pytest.mark.skipif(not dist.is_gloo_available(), reason="gloo is not available") +@pytest.mark.parametrize("scheme", ["w4a8", "mxfp4"]) +def test_two_gloo_ranks_write_the_single_process_bytes(tmp_path, scheme): + case, out = tmp_path / scheme, tmp_path / "release" + init = f"file://{tmp_path / 'rendezvous'}" + mp.start_processes(_gloo_worker, args=(2, init, str(case), str(out)), nprocs=2, join=True, start_method="spawn") + assert not (out / PART_DIR).exists() + assert json.loads((out / INDEX_FILE).read_text())["weight_map"] + _assert_same_files(out, case / "reference") + + +# -- N8.3: release directories, records and the reader check --------------------------- + +def test_release_directories_are_named_by_scheme_with_the_old_w4a8_name(): + assert release_dir_for(BF16) == RELEASE_DIR == "policy-releases" + assert release_dir_for(W4A8_INT4_CHANNEL.name) == W4A8_RELEASE_DIR == "policy-releases-w4a8" + assert release_dir_for(MXFP4_G32_RTN.name) == "policy-releases-mxfp4-g32-rtn" + with pytest.raises(ValueError, match="unknown weight scheme"): + release_dir_for("fp8") + + +def _bf16_store(tmp: Path, world: dict): + store = CheckpointStore(tmp / "store") + source = world["tmp"] / "bf16" + PolicyReleaseStore(store).release(1, "h1", lambda d: [os.link(p, d / p.name) for p in source.iterdir()]) + return store + + +def test_mxfp4_release_cli_writes_its_own_directory_and_record(tmp_path): + w = _build(tmp_path / "mxfp4", "mxfp4") + store = _bf16_store(tmp_path, w) + done = subprocess.run([sys.executable, "-m", "yeto.rl.quant.release", "--store", str(store.root), + "--template", str(w["template"]), "--no-awq"], capture_output=True, text=True, + env={**os.environ, "PYTHONPATH": str(Path(__file__).parents[1])}) + assert done.returncode == 0, done.stderr + assert not (store.root / "policy-releases-w4a8").exists() + rec = PolicyReleaseStore(store, release_dir="policy-releases-mxfp4-g32-rtn").latest() + assert rec.version == 1 and rec.policy_hash == "h1+mxfp4-g32-rtn" + assert {k: rec.derived_from[k] for k in ("release_dir", "version", "scheme", "awq", "awq_source")} == { + "release_dir": "policy-releases", "version": 1, "scheme": "mxfp4-g32-rtn", "awq": "none", "awq_source": None} + assert rec.derived_from["expert_matrices"] == len(CheckpointLayout.load(w["template"]).quantized_modules) + out = store.root / "policy-releases-mxfp4-g32-rtn" / rec.path + for f in w["reference"].glob("*.safetensors"): + assert (out / f.name).read_bytes() == f.read_bytes() + + +def test_releaser_refuses_a_directory_of_another_scheme_and_bad_awq_records(tmp_path): + store = CheckpointStore(tmp_path / "s") + bf16 = PolicyReleaseStore(store) + w4a8 = PolicyReleaseStore(store, release_dir="policy-releases-w4a8") + noop = lambda src, out: None # noqa: E731 + with pytest.raises(ValueError, match="mxfp4-g32-rtn releases go to"): + RequantReleaser(bf16, w4a8, noop, scheme="mxfp4-g32-rtn") + with pytest.raises(ValueError, match="not a re-quantized scheme"): + RequantReleaser(bf16, w4a8, noop, scheme="bf16") + with pytest.raises(ValueError, match="give a source exactly"): + RequantReleaser(bf16, w4a8, noop, awq="fixed-artefact") + + +def test_follower_refuses_a_release_of_another_scheme_before_reading_it(tmp_path): + store = CheckpointStore(tmp_path / "s") + # A W4A8 release written into the MXFP4 directory (a misconfigured releaser or copy). + wrong = PolicyReleaseStore(store, release_dir="policy-releases-mxfp4-g32-rtn") + rec = wrong.release(3, "h+w4a8-int4-channel", lambda d: (d / "model.safetensors").write_text("x"), + derived_from={"scheme": "w4a8-int4-channel", "version": 3}) + assert release_scheme(rec) == "w4a8-int4-channel" + loaded = [] + follower = PolicyFollower(wrong, lambda r, d: loaded.append(r), scheme="mxfp4-g32-rtn") + with pytest.raises(SchemeMismatch, match="w4a8-int4-channel"): + follower.refresh() + assert loaded == [] and follower.release is None + assert PolicyFollower(wrong, lambda r, d: loaded.append(r), scheme="w4a8-int4-channel").refresh().version == 3 + bf16 = PolicyReleaseStore(store) + 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() diff --git a/tests/test_rl_ssh_trainer_island.py b/tests/test_rl_ssh_trainer_island.py index 74d9f1cd5..4667411fa 100644 --- a/tests/test_rl_ssh_trainer_island.py +++ b/tests/test_rl_ssh_trainer_island.py @@ -257,7 +257,8 @@ def test_each_bf16_release_becomes_one_w4a8_release_the_sampler_follows(tmp_path assert rec.version == 1 and rec.policy_hash == "ha+w4a8-int4-channel" and rec.buffer_selected == 10 assert rec.buffer_shards == {"npu": 4} assert rec.derived_from == {"release_dir": "policy-releases", "seq": 1, "version": 1, "policy_hash": "ha", - "scheme": "w4a8-int4-channel", "awq": "none"} + "scheme": "w4a8-int4-channel", "awq": "none", "awq_source": None, + "expert_matrices": None} assert releaser.sync() is None and len(calls) == 1 # converted once loaded = [] follower = PolicyFollower(PolicyReleaseStore(store, release_dir=W4A8_RELEASE_DIR), @@ -303,6 +304,7 @@ def test_release_cli_one_pass_with_the_real_converter(world): assert done.returncode == 0, done.stderr rec = PolicyReleaseStore(store, release_dir=W4A8_RELEASE_DIR).latest() assert rec.version == 1 and rec.derived_from["awq"] == "fixed-artefact" + assert "recipe.yaml" in rec.derived_from["awq_source"] and rec.derived_from["expert_matrices"] == 6 out = store.root / W4A8_RELEASE_DIR / rec.path shards = sorted(p.name for p in original.glob("*.safetensors")) assert shards and sorted(p.name for p in out.glob("*.safetensors")) == shards diff --git a/yeto/launch/ssh_trainer_island.py b/yeto/launch/ssh_trainer_island.py index 72ebd423d..31f49d74a 100644 --- a/yeto/launch/ssh_trainer_island.py +++ b/yeto/launch/ssh_trainer_island.py @@ -17,9 +17,10 @@ under ``torchrun`` with the incarnation environment (``trainer_remote.island_env``), the collective-communication variables of its family and fabric (``topology.island_comm_env``) and the run's - environment. Node 0 also runs the W4A8 re-quantizer side process + environment. Node 0 also runs the re-quantizer side process (``yeto.rl.quant.release --loop``) when the spec has one: each BF16 release - becomes a W4A8 release in the store, which W4A8 samplers follow; + becomes a release of the template's scheme (W4A8 or MXFP4) in that scheme's + directory of the store, which samplers of that scheme follow; 3. the SSH session streams node 0's log to the head (``IslandEvents`` reads the trainer events from it); the other nodes' logs go to local files; 4. release kills each node's process group (TERM, then KILL after a grace diff --git a/yeto/rl/adapters/miles/quant_shards.py b/yeto/rl/adapters/miles/quant_shards.py new file mode 100644 index 000000000..b9c7c6be7 --- /dev/null +++ b/yeto/rl/adapters/miles/quant_shards.py @@ -0,0 +1,74 @@ +"""Megatron (Miles) expert weights -> re-quantization source slices (S25 N8.2 / N8.5b adapter). + +Megatron holds the routed experts of a rank as one weight per local expert: + +* grouped GEMM (TE ``TEGroupedMLP``): ``decoder.layers..mlp.experts.linear_fc1.weight`` + ``[2I, H]`` and ``...linear_fc2.weight`` ``[H, I]``; +* sequential MLP (the ``--transformer-impl local`` fallback, N7.5): + ``decoder.layers..mlp.experts.local_experts..linear_fc{1,2}.weight``. + +```` is the layer index inside the pipeline stage and ```` the local +expert index; the global expert is ``ep_rank * num_local_experts + e``. +``linear_fc1`` holds the gate rows, then the up rows (SwiGLU ``chunk(2)``). +This adapter renames them to the per-expert HF source names of the +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 +:mod:`yeto.rl.quant.shards`, "Ownership"). + +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. +""" + +from __future__ import annotations + +import re +from collections.abc import Iterable, Iterator + +import torch + +from yeto.rl.quant.shards import SourceSlice +from yeto.rl.quant.stacked import DOWN, expert_weight_name, split_gate_up + +HF_LAYER_PREFIX = "model.language_model.layers." +_EXPERT = re.compile( + r"^(?:module\.)*decoder\.layers\.(?P\d+)\.mlp\.experts\." + r"(?:linear_fc(?P[12])\.weight(?P\d+)|local_experts\.(?P\d+)\.linear_fc(?P[12])\.weight)$" +) + + +def megatron_expert_slices(named_params: Iterable[tuple[str, torch.Tensor]], *, ep_rank: int, + num_local_experts: int, layer_offset: int = 0, expert_tp_size: int = 1, + hf_layer_prefix: str = HF_LAYER_PREFIX) -> Iterator[SourceSlice]: + """Source slices of this rank's routed experts; other parameters are skipped. + + ``layer_offset``: the global index of the stage's first layer (pipeline + parallelism).""" + if int(expert_tp_size) != 1: + raise ValueError(f"expert tensor parallel size {expert_tp_size} is not supported; use 1") + if int(num_local_experts) < 1 or int(ep_rank) < 0: + raise ValueError(f"bad expert parallel position: ep_rank={ep_rank}, num_local_experts={num_local_experts}") + for name, tensor in named_params: + match = _EXPERT.match(name) + if match is None: + continue + local = int(match["e"] if match["e"] is not None else match["e2"]) + if local >= num_local_experts: + 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() + if (match["fc"] or match["fc2"]) == "1": + for hf_name, part in split_gate_up(prefix, expert, weight): + yield SourceSlice(hf_name, part) + else: + yield SourceSlice(expert_weight_name(prefix, expert, DOWN), weight) + + +def hf_slices(named_tensors: Iterable[tuple[str, torch.Tensor]]) -> Iterator[SourceSlice]: + """Whole tensors already in HF names (Miles ``convert_to_hf`` output).""" + for name, tensor in named_tensors: + yield SourceSlice(name, tensor.detach()) diff --git a/yeto/rl/adapters/verl/quant_shards.py b/yeto/rl/adapters/verl/quant_shards.py new file mode 100644 index 000000000..e9b5d3461 --- /dev/null +++ b/yeto/rl/adapters/verl/quant_shards.py @@ -0,0 +1,72 @@ +"""FSDP2 (verl / HF) parameters -> re-quantization source slices (S25 N8.2 / N8.5b adapter). + +FSDP2 shards every parameter on dim 0 over a 1-D mesh (``Shard(0)``): rank +``r`` holds rows ``[r * c, min((r + 1) * c, N))`` with ``c = ceil(N / world)`` +(``torch.chunk``). The HF model keeps the routed experts stacked +(``mlp.experts.gate_up_proj`` ``[E, 2I, H]``, ``mlp.experts.down_proj`` +``[E, H, I]``), so a dim-0 shard of them is a range of whole experts. This +adapter turns the local shards into :class:`~yeto.rl.quant.shards.SourceSlice` +items: stacked experts into per-expert matrices (:mod:`yeto.rl.quant.stacked`), +other parameters into row slices. No parameter is gathered. + +``rename`` maps a module parameter name to the checkpoint name of the layout +(default: identity; e.g. strip a wrapper prefix). A replicated parameter (a +plain tensor or a ``Replicate`` DTensor) is given by mesh coordinate 0 only. + +Limits (explicit errors): DTensors on a mesh of more than one dimension, or +with a placement other than ``Shard(0)`` / ``Replicate``. Status: CPU tests with +real DTensors on 2 gloo ranks; not run inside verl. +""" + +from __future__ import annotations + +from collections.abc import Callable, Iterable, Iterator + +import torch + +from yeto.rl.quant.shards import SourceSlice +from yeto.rl.quant.stacked import STACKED_EXPERTS, unstack_experts + + +def shard0_rows(rows: int, coordinate: int, world: int) -> tuple[int, int]: + """The rows ``[start, stop)`` a ``Shard(0)`` coordinate holds (``torch.chunk`` split).""" + chunk = -(-int(rows) // int(world)) if rows else 0 + start = min(int(coordinate) * chunk, rows) + return start, min(start + chunk, rows) + + +def fsdp2_slices(named_params: Iterable[tuple[str, torch.Tensor]], *, rank: int = 0, + rename: Callable[[str], str] = lambda name: name) -> Iterator[SourceSlice]: + """Source slices of the parameters this rank holds. ``rank`` decides who gives + plain (non-DTensor) tensors: rank 0 only.""" + from torch.distributed.tensor import DTensor, Replicate, Shard + + for module_name, param in named_params: + name = rename(module_name) + if isinstance(param, DTensor): + mesh, placements = param.device_mesh, tuple(param.placements) + if mesh.ndim != 1: + raise ValueError(f"{module_name}: {mesh.ndim}-D device mesh is not supported (FSDP2 uses 1-D)") + coordinate = mesh.get_local_rank() + if placements == (Replicate(),): + if coordinate == 0: + yield from _slices(name, param.to_local().detach(), 0) + continue + if placements != (Shard(0),): + raise ValueError(f"{module_name}: placement {placements} is not supported (Shard(0) or Replicate)") + local = param.to_local().detach() + start, stop = shard0_rows(param.shape[0], coordinate, mesh.size()) + if local.shape[0] != stop - start: + raise ValueError(f"{module_name}: local shard has {local.shape[0]} rows, expected {stop - start}") + if stop > start: + yield from _slices(name, local, start) + elif int(rank) == 0: + yield from _slices(name, param.detach(), 0) + + +def _slices(name: str, tensor: torch.Tensor, row_start: int) -> Iterator[SourceSlice]: + if STACKED_EXPERTS.match(name): + for expert_name, matrix in unstack_experts(name, tensor, first_expert=row_start): + yield SourceSlice(expert_name, matrix) + else: + yield SourceSlice(name, tensor, row_start) diff --git a/yeto/rl/engine/policy_release.py b/yeto/rl/engine/policy_release.py index 4b65583d1..634495f13 100644 --- a/yeto/rl/engine/policy_release.py +++ b/yeto/rl/engine/policy_release.py @@ -62,14 +62,15 @@ from pathlib import Path from typing import Any +from yeto.rl.quant.schemes import BF16, RELEASE_DIR, W4A8_INT4_CHANNEL, release_dir_for + from .object_store import ObjectStore, download_file, object_size from .resume import CheckpointStore, _write_json_atomic SCHEMA = "yeto.rl.policy-release/v1" -RELEASE_DIR = "policy-releases" -# W4A8 releases re-quantized from the BF16 ones (yeto.rl.quant.release); a W4A8 -# sampler follows this directory, a BF16 sampler the default one. -W4A8_RELEASE_DIR = "policy-releases-w4a8" +# Releases re-quantized from the BF16 ones (yeto.rl.quant.release) live in +# release_dir_for(scheme); a sampler follows the directory of the scheme it serves. +W4A8_RELEASE_DIR = release_dir_for(W4A8_INT4_CHANNEL) LATEST = "LATEST.json" RECORD = "RELEASE.json" _RELEASE_NAME = re.compile(r"r(\d+)-v\d+") @@ -109,6 +110,24 @@ def from_record(cls, record: dict[str, Any]) -> "PolicyRelease": return cls(**{k: record[k] for k in record if k in cls.__dataclass_fields__}) +def release_scheme(release: PolicyRelease) -> str: + """The weight scheme of a release: ``derived_from.scheme`` of a derived + release, else :data:`~yeto.rl.quant.schemes.BF16` (trained weights).""" + return str((release.derived_from or {}).get("scheme") or BF16) + + +class SchemeMismatch(PolicyReleaseError): + """A release whose weight scheme is not the one the reader serves.""" + + +def check_release_scheme(release: PolicyRelease, expected: str) -> None: + """Raise :class:`SchemeMismatch` when ``release`` is not of scheme ``expected``.""" + actual = release_scheme(release) + if actual != expected: + raise SchemeMismatch(f"release r{release.seq}-v{release.version} has weight scheme {actual!r}; " + f"this reader serves {expected!r} (it follows the wrong release directory?)") + + def directory_digest(directory: Path) -> str: """sha256 over every file under ``directory`` except the release record: relative path, size and bytes, in sorted path order.""" @@ -169,8 +188,8 @@ def check_release_files(release: PolicyRelease, directory: Path) -> None: class PolicyReleaseStore: """Writes (trainer) and reads (sampler) the releases in one store. ``cache_dir``: where a store with an object API downloads a release to. - ``release_dir``: the release directory in the store (``W4A8_RELEASE_DIR`` - for re-quantized releases).""" + ``release_dir``: the release directory in the store + (:func:`~yeto.rl.quant.schemes.release_dir_for` of the scheme).""" def __init__(self, store: CheckpointStore, *, keep: int = 2, clock: Callable[[], float] = time.time, @@ -300,11 +319,15 @@ class PolicyFollower: on every engine (raises on any failure; the served release is then unchanged). The weight files are checked against the record's ``sha256`` first. ``version`` is the version the engines serve before any release (the base - policy, 0).""" + policy, 0). ``scheme``: the weight scheme the engines serve; a release of + another scheme raises :class:`SchemeMismatch` before any file is read + (None: no check, for callers that check in ``load``).""" def __init__(self, releases: PolicyReleaseStore, - load: Callable[[PolicyRelease, Path], None], *, version: int = 0) -> None: + load: Callable[[PolicyRelease, Path], None], *, version: int = 0, + scheme: str | None = None) -> None: self.releases = releases + self.scheme = scheme self._load = load self.version = int(version) self.release: PolicyRelease | None = None # the release the engines serve @@ -319,6 +342,8 @@ def refresh(self) -> PolicyRelease | None: self.last_seen = latest if latest is None or (self.release is not None and latest.seq == self.release.seq): return None + if self.scheme is not None: + check_release_scheme(latest, self.scheme) try: directory = self.releases.resolve(latest) except ReleaseIncomplete as error: diff --git a/yeto/rl/quant/__init__.py b/yeto/rl/quant/__init__.py index 61357cde1..5addfcd12 100644 --- a/yeto/rl/quant/__init__.py +++ b/yeto/rl/quant/__init__.py @@ -28,4 +28,11 @@ * :mod:`.awq`: the fixed AWQ smoothing (R2) and the no-smoothing fallback (R1). * :mod:`.requant`: the checkpoint layout and its scheme, the per-tensor transform, the shard-by-shard converter, and the tensor-by-tensor comparison. +* :mod:`.schemes`: the scheme names and their release directories (no torch). +* :mod:`.stacked`: the trainer-native stacked expert layout -> per-expert names. +* :mod:`.shards`: per-rank quantization into part files and their merge (R3). +* :mod:`.release`: the side process that turns BF16 releases into quantized ones. + +Trainer adapters: :mod:`yeto.rl.adapters.miles.quant_shards` (Megatron) and +:mod:`yeto.rl.adapters.verl.quant_shards` (FSDP2). """ diff --git a/yeto/rl/quant/release.py b/yeto/rl/quant/release.py index 2bdfd52f8..0eae27629 100644 --- a/yeto/rl/quant/release.py +++ b/yeto/rl/quant/release.py @@ -1,22 +1,30 @@ -"""Turn each BF16 policy release into a W4A8 release for W4A8 samplers (S25 N7.3 + N8). +"""Turn each BF16 policy release into a re-quantized release for quantized samplers (S25 N7.3 + N8). The trainer releases BF16 weights into ``policy-releases/`` (``ReleasingPublisher``). :class:`RequantReleaser` follows that directory and, for each new BF16 release, writes the re-quantized checkpoint (:func:`yeto.rl.quant.requant.requantize_checkpoint`) -as a release in ``policy-releases-w4a8/`` with the same commit order (weights, -``RELEASE.json``, then ``LATEST.json``). A W4A8 sampler follows that directory with -the usual ``PolicyFollower``; it never sees a partial W4A8 checkpoint. - -The W4A8 record keeps the BF16 policy version (the version the sampler stamps on -its samples), the trainer's buffer positions, and ``derived_from`` (the BF16 -release's directory, seq, version and hash, the scheme and the AWQ source). Its -policy hash is the BF16 hash plus the scheme, so a restarted releaser that sees -the same BF16 release writes nothing. +as a release in the directory of its scheme (:func:`~.schemes.release_dir_for`: +``policy-releases-w4a8/`` for W4A8, ``policy-releases-mxfp4-g32-rtn/`` for +MXFP4) with the same commit order (weights, ``RELEASE.json``, then +``LATEST.json``). A quantized sampler follows that directory with the usual +``PolicyFollower`` (``scheme=`` checks each record); it never sees a partial +checkpoint. + +The record keeps the BF16 policy version (the version the sampler stamps on +its samples), the trainer's buffer positions, and ``derived_from``: the BF16 +release's directory, seq, ``version`` (the BF16 version) and hash, ``scheme``, +``awq`` (:data:`~.schemes.AWQ_NONE` or :data:`~.schemes.AWQ_FIXED_ARTEFACT`), +``awq_source`` (the recipe and factor files of R2, else None) and +``expert_matrices`` (quantized expert matrices in the layout, the sampler's +post-load count check; None when unknown). Its policy hash is the BF16 hash +plus the scheme, so a restarted releaser that sees the same BF16 release +writes nothing. Where it runs: as a side process on the first trainer node (``yeto.launch.ssh_trainer_island``), which reads the BF16 release from the -shared store. Per-rank quantization of the trainer's own shards (N8.2) would -replace the whole-checkpoint read; not done. +shared store. Per-rank quantization of the trainer's own shards +(:mod:`.shards`, N8.2) can replace the whole-checkpoint read; it is not wired +into a trainer's release path yet. Status: CPU tests with a fake re-quantizer and with the real converter on a tiny layout (tests/test_rl_quant_release.py); not run on a real release. @@ -39,35 +47,44 @@ from yeto.rl.engine.policy_release import ( FULL, - RELEASE_DIR, - W4A8_RELEASE_DIR, PolicyRelease, PolicyReleaseError, PolicyReleaseStore, ReleaseIncomplete, ) -SCHEME = "w4a8-int4-channel" +from .schemes import AWQ_FIXED_ARTEFACT, AWQ_NONE, BF16, QUANTIZED_SCHEMES, W4A8_INT4_CHANNEL, release_dir_for Requantize = Callable[[Path, Path], Any] # (bf16 directory, output directory) -> stats class RequantReleaser: - """Writes one W4A8 release per BF16 release (module doc). + """Writes one re-quantized release per BF16 release (module doc). ``requantize(source_dir, out_dir)`` writes a complete loadable checkpoint into ``out_dir`` (default: :func:`requantizer` over a template).""" def __init__(self, source: PolicyReleaseStore, target: PolicyReleaseStore, requantize: Requantize, - *, scheme: str = SCHEME, awq: str = "none", + *, scheme: str = W4A8_INT4_CHANNEL, awq: str = AWQ_NONE, awq_source: str | None = None, + expert_matrices: int | None = None, log: Callable[[str], None] = lambda m: print(m, flush=True)) -> None: if source.release_dir == target.release_dir: raise ValueError(f"source and target are the same release directory {source.release_dir!r}") + if scheme not in QUANTIZED_SCHEMES: + raise ValueError(f"scheme {scheme!r} is not a re-quantized scheme {QUANTIZED_SCHEMES}") + if source.release_dir != release_dir_for(BF16): + raise ValueError(f"source {source.release_dir!r} is not the BF16 release directory {release_dir_for(BF16)!r}") + if target.release_dir != release_dir_for(scheme): + raise ValueError(f"{scheme} releases go to {release_dir_for(scheme)!r}, not {target.release_dir!r}") + if awq not in (AWQ_NONE, AWQ_FIXED_ARTEFACT) or (awq == AWQ_NONE) != (awq_source is None): + raise ValueError(f"awq {awq!r} with source {awq_source!r}: give a source exactly for {AWQ_FIXED_ARTEFACT!r}") self.source = source self.target = target self.requantize = requantize self.scheme = scheme self.awq = awq + self.awq_source = awq_source + self.expert_matrices = None if expert_matrices is None else int(expert_matrices) self.log = log self.waiting: str | None = None # why the newest BF16 release is not converted yet @@ -78,7 +95,7 @@ def done_for(self, release: PolicyRelease) -> bool: and latest.version == release.version) def sync(self) -> PolicyRelease | None: - """Convert the newest BF16 release once; returns the W4A8 release written, + """Convert the newest BF16 release once; returns the release written, or None (no BF16 release, already converted, or not complete yet).""" src = self.source.latest() if src is None or self.done_for(src): @@ -93,7 +110,8 @@ def sync(self) -> PolicyRelease | None: return None self.waiting = None derived = {"release_dir": self.source.release_dir, "seq": src.seq, "version": src.version, - "policy_hash": src.policy_hash, "scheme": self.scheme, "awq": self.awq} + "policy_hash": src.policy_hash, "scheme": self.scheme, "awq": self.awq, "awq_source": self.awq_source, + "expert_matrices": self.expert_matrices} stats: dict[str, Any] = {} def write(out: Path) -> None: @@ -135,26 +153,28 @@ def main(argv: list[str] | None = None) -> int: p.add_argument("--recipe") p.add_argument("--awq-scales") p.add_argument("--no-awq", action="store_true") - p.add_argument("--source-dir", default=RELEASE_DIR) - p.add_argument("--target-dir", default=W4A8_RELEASE_DIR) 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) a = p.parse_args(argv) if a.no_awq: - smoothing, awq = AwqSmoothing.none(), "none" + smoothing, awq = AwqSmoothing.none(), AWQ_NONE elif a.recipe and a.awq_scales: - smoothing, awq = AwqSmoothing.from_artefact(Path(a.recipe), Path(a.awq_scales)), "fixed-artefact" + smoothing, awq = AwqSmoothing.from_artefact(Path(a.recipe), Path(a.awq_scales)), AWQ_FIXED_ARTEFACT else: p.error("give --recipe and --awq-scales (option R2) or --no-awq (option R1)") from .requant import CheckpointLayout - # The scheme tag follows the template format (W4A8 int4 or MXFP4), so readers can check it (N8.3). - scheme = CheckpointLayout.load(Path(a.template)).scheme.name + # The scheme follows the template format (W4A8 int4 or MXFP4); it names the + # target directory and is written into each record, so readers can check it (N8.3). + layout = CheckpointLayout.load(Path(a.template)) + scheme = layout.scheme.name store = store_for(a.store) - releaser = RequantReleaser(PolicyReleaseStore(store, keep=a.keep, release_dir=a.source_dir), - PolicyReleaseStore(store, keep=a.keep, release_dir=a.target_dir), - requantizer(Path(a.template), smoothing), scheme=scheme, awq=awq) + 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, + awq_source=smoothing.source if smoothing.enabled else None, + expert_matrices=len(layout.quantized_modules)) if a.loop: releaser.run(threading.Event(), a.poll_seconds) else: diff --git a/yeto/rl/quant/requant.py b/yeto/rl/quant/requant.py index eaba2d783..44fced0e2 100644 --- a/yeto/rl/quant/requant.py +++ b/yeto/rl/quant/requant.py @@ -25,7 +25,9 @@ Work split: :func:`requantize_checkpoint` takes a subset of output shards, so several workers can each write their own shards. :func:`transform_tensor` works on any row slice of a 2-D weight with the same result rows, which is the -basis for quantizing on the trainer ranks (D7, task N8.2; not done here). +basis for quantizing on the trainer ranks (D7, task N8.2, :mod:`.shards`). +The source may be per-expert (HF checkpoint) or the trainer's stacked native +layout (:mod:`.stacked`, task N8.5b). Status: W4A8 transform rules checked bit for bit against the original artefact on layer 30 (router, shared experts, norm, one routed expert); MXFP4 rules @@ -58,6 +60,9 @@ from .awq import AwqSmoothing from .int4 import QMAX, QMIN, quantize_per_channel from .mxfp4 import GROUP_SIZE, quantize_mxfp4, unpack_fp4 +from .schemes import MXFP4_G32_RTN as MXFP4_NAME +from .schemes import W4A8_INT4_CHANNEL as W4A8_NAME +from .stacked import StackedSourceIndex INDEX_FILE = "model.safetensors.index.json" LAYOUT_FILE = "layout.json" @@ -123,7 +128,7 @@ def check_values(self, name: str, data: torch.Tensor) -> None: class _Int4Channel(QuantScheme): - name = "w4a8-int4-channel" + name = W4A8_NAME rtn_only = False format = "int-quantized" data_suffix = WEIGHT_SUFFIX @@ -148,7 +153,7 @@ def check_values(self, name: str, data: torch.Tensor) -> None: class _Mxfp4Group32(QuantScheme): - name = "mxfp4-g32-rtn" + name = MXFP4_NAME rtn_only = True format = "mxfp4-pack-quantized" data_suffix = PACKED_SUFFIX @@ -269,6 +274,11 @@ def source_names(self) -> set[str]: return {self.source_of(n) for n in self.tensors} def shards(self) -> dict[str, list[str]]: + """Output shard -> its tensor names (sorted); do not modify the result.""" + return self._shards + + @cached_property + def _shards(self) -> dict[str, list[str]]: out: dict[str, list[str]] = defaultdict(list) for name, spec in self.tensors.items(): out[spec.shard].append(name) @@ -315,23 +325,53 @@ def transform_tensor( return {module + scheme.data_suffix: data, module + SCALE_SUFFIX: scale} +def check_smoothing(layout: CheckpointLayout, smoothing: AwqSmoothing) -> None: + """Refuse AWQ smoothing for a round-to-nearest-only scheme (MXFP4).""" + if layout.scheme.rtn_only and smoothing.enabled: + raise LayoutError(f"{layout.scheme.name} is round-to-nearest only; use no AWQ smoothing") + + class _SourceReader: - """Lazy reader over a sharded safetensors source checkpoint.""" + """Lazy reader over a sharded safetensors source checkpoint. + + A source in the trainer's native stacked layout (``mlp.experts.gate_up_proj`` + ``[E, 2I, H]`` and ``mlp.experts.down_proj`` ``[E, H, I]``, :mod:`.stacked`) + is read under the per-expert names of the target layout; each per-expert + matrix is sliced from the stacked tensor on read. + """ def __init__(self, source_dir: Path) -> None: self.dir = Path(source_dir) index = self.dir / INDEX_FILE if index.is_file(): - self.weight_map: dict[str, str] = json.loads(index.read_text())["weight_map"] + stored: dict[str, str] = json.loads(index.read_text())["weight_map"] else: - self.weight_map = {n: p.name for p in sorted(self.dir.glob("*.safetensors")) for n in read_header(p)} + stored = {n: p.name for p in sorted(self.dir.glob("*.safetensors")) for n in read_header(p)} + headers: dict[str, dict[str, Any]] = {} + + def shape_of(name: str) -> list[int]: + shard = stored[name] + if shard not in headers: + headers[shard] = read_header(self.dir / shard) + return headers[shard][name]["shape"] + + self.stacked = StackedSourceIndex.from_names(stored, shape_of) + self.stored = stored + self.weight_map: dict[str, str] = {n: s for n, s in stored.items() if n not in self.stacked.stacked_names} + for name, view in self.stacked.views.items(): + self.weight_map[name] = stored[view.stored] self._open: dict[str, Any] = {} - def get(self, name: str) -> torch.Tensor: - shard = self.weight_map[name] + def _file(self, shard: str) -> Any: if shard not in self._open: self._open[shard] = safe_open(str(self.dir / shard), framework="pt", device="cpu") - return self._open[shard].get_tensor(name) + return self._open[shard] + + def get(self, name: str) -> torch.Tensor: + view = self.stacked.views.get(name) + if view is None: + return self._file(self.weight_map[name]).get_tensor(name) + return view.read(self._file(self.stored[view.stored]).get_slice(view.stored)) def close(self) -> None: self._open.clear() @@ -354,8 +394,7 @@ def requantize_checkpoint( The source must provide exactly the layout's source tensors. """ layout = CheckpointLayout.load(template_dir) - if layout.scheme.rtn_only and smoothing.enabled: - raise LayoutError(f"{layout.scheme.name} is round-to-nearest only; use no AWQ smoothing") + check_smoothing(layout, smoothing) source = _SourceReader(source_dir) expected = layout.source_names() missing, extra = expected - set(source.weight_map), set(source.weight_map) - expected @@ -381,11 +420,8 @@ def requantize_checkpoint( if name not in produced: src = layout.source_of(name) produced.update(transform_tensor(src, source.get(src), layout, smoothing)) - tensors = {n: produced[n].contiguous() for n in all_shards[shard]} - _check_against_layout(tensors, layout) - tmp = out_dir / (shard + ".tmp") - save_file(tensors, str(tmp), metadata={"format": "pt"}) - os.replace(tmp, out_dir / shard) + tensors = {n: produced[n] for n in all_shards[shard]} + write_output_shard(out_dir, shard, tensors, layout) stats["shards"] += 1 stats["quantized"] += sum(n in quantized for n in tensors) stats["tensors"] += len(tensors) @@ -396,6 +432,22 @@ def requantize_checkpoint( return stats +def write_output_shard(out_dir: Path, shard: str, tensors: Mapping[str, torch.Tensor], layout: CheckpointLayout) -> None: + """Check ``tensors`` (exactly the layout's tensors of ``shard``) and write the shard file atomically. + + The single-process converter and the merge of per-rank parts + (:mod:`.shards`) both write through here, so equal tensors give equal bytes. + """ + expected = layout.shards().get(shard) + if expected is None or set(tensors) != set(expected): + raise LayoutError(f"{shard}: tensors differ from the layout's tensors of this shard") + tensors = {n: t.contiguous() for n, t in tensors.items()} + _check_against_layout(tensors, layout) + tmp = Path(out_dir) / (shard + ".tmp") + save_file(tensors, str(tmp), metadata={"format": "pt"}) + os.replace(tmp, Path(out_dir) / shard) + + def finalize_checkpoint(template_dir: Path, out_dir: Path) -> None: """Check that every output shard exists, then write the metadata files and, last, the index.""" layout = CheckpointLayout.load(template_dir) diff --git a/yeto/rl/quant/schemes.py b/yeto/rl/quant/schemes.py new file mode 100644 index 000000000..86d85a6e7 --- /dev/null +++ b/yeto/rl/quant/schemes.py @@ -0,0 +1,42 @@ +"""Names of the weight schemes a policy release can carry (S25 N8.3). + +This module is the one place for the scheme tags. The writer side +(:mod:`.requant`, :mod:`.release`) and the reader side (the sampler's +``PolicyFollower``) both import them from here. It does not import torch, so a +sampler process can use it without the quantization code. + +* :data:`BF16`: trained weights as the trainer releases them. +* :data:`W4A8_INT4_CHANNEL`: routed experts int4 per output channel, BF16 + scales, 8-bit dynamic activations (the original W4A8 artefact, N8.1). +* :data:`MXFP4_G32_RTN`: routed experts MXFP4, one E8M0 scale per 32 inputs, + round to nearest without AWQ (N8.5). + +Release directories are named by scheme (:func:`release_dir_for`): +``policy-releases`` for BF16, ``policy-releases-`` for a re-quantized +scheme. The W4A8 scheme keeps its older name ``policy-releases-w4a8``, so +stores written before N8.3 stay readable. +""" + +from __future__ import annotations + +BF16 = "bf16" +W4A8_INT4_CHANNEL = "w4a8-int4-channel" +MXFP4_G32_RTN = "mxfp4-g32-rtn" +QUANTIZED_SCHEMES = (W4A8_INT4_CHANNEL, MXFP4_G32_RTN) +ALL_SCHEMES = (BF16, *QUANTIZED_SCHEMES) + +RELEASE_DIR = "policy-releases" +_LEGACY_DIRS = {W4A8_INT4_CHANNEL: "policy-releases-w4a8"} + +# AWQ source tags of a release record (``derived_from.awq``). +AWQ_NONE = "none" # option R1, and every MXFP4 release +AWQ_FIXED_ARTEFACT = "fixed-artefact" # option R2: fixed factors from the original artefact + + +def release_dir_for(scheme: str) -> str: + """The store directory that holds the releases of ``scheme``.""" + if scheme == BF16: + return RELEASE_DIR + if scheme not in QUANTIZED_SCHEMES: + raise ValueError(f"unknown weight scheme {scheme!r}; known: {', '.join(ALL_SCHEMES)}") + return _LEGACY_DIRS.get(scheme, f"{RELEASE_DIR}-{scheme}") diff --git a/yeto/rl/quant/shards.py b/yeto/rl/quant/shards.py new file mode 100644 index 000000000..145441e07 --- /dev/null +++ b/yeto/rl/quant/shards.py @@ -0,0 +1,203 @@ +"""Distributed re-quantization: each trainer rank quantizes the weights it holds (S25 N8.2, option R3). + +No rank gathers the whole BF16 model. Each rank gives the source tensors it +holds as :class:`SourceSlice` items (per-expert source names of the target +layout, see :mod:`.stacked`; a slice is a range of rows) and writes one part +file with the quantized result of those rows (:func:`write_part`). Then the +parts are merged into the output shards of the layout (:func:`merge_parts`), +each rank merging a disjoint subset of the shards, and the index and metadata +are written last (:func:`~.requant.finalize_checkpoint`). + +Why row slices work: both schemes quantize each output row on its own (W4A8: +one scale per row; MXFP4: one scale per 32 inputs of a row), AWQ balance +factors scale columns, and casts work per element. So the rows of a slice +quantize to the same rows as in the whole tensor, and the merged shard is byte +for byte the file :func:`~.requant.requantize_checkpoint` writes from the whole +source (both write through :func:`~.requant.write_output_shard`). Limits: an AWQ +smooth layer (1-D) must be given whole; 0-D tensors must be given whole. + +Ownership: every row of every source tensor must come from exactly one rank; +the merge fails on a gap or an overlap. Replicated parameters (for example the +non-expert weights on every expert-parallel rank of Megatron) must be given by +one rank only; the trainer adapters decide which. + +Part file: ``/part-.safetensors``; tensor key +``@``. + +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. +""" + +from __future__ import annotations + +import os +import shutil +from collections import defaultdict +from collections.abc import Callable, Iterable +from dataclasses import dataclass +from pathlib import Path +from typing import Any + +import torch +from safetensors import safe_open +from safetensors.torch import save_file + +from .awq import AwqSmoothing +from .requant import ( + CheckpointLayout, + LayoutError, + check_smoothing, + finalize_checkpoint, + read_header, + transform_tensor, + write_output_shard, +) + +PART_DIR = ".requant-parts" +_KEY_SEP = "@" + + +@dataclass(frozen=True) +class SourceSlice: + """Rows ``[row_start, row_start + len(tensor))`` of the source tensor ``name``.""" + + name: str + tensor: torch.Tensor + row_start: int = 0 + + +def part_path(part_dir: Path, rank: int) -> Path: + return Path(part_dir) / f"part-{int(rank):05d}.safetensors" + + +def write_part(part_dir: Path, rank: int, slices: Iterable[SourceSlice], template_dir: Path, + smoothing: AwqSmoothing) -> dict[str, int]: + """Quantize this rank's slices and write its part file; returns counts.""" + layout = CheckpointLayout.load(template_dir) + check_smoothing(layout, smoothing) + out: dict[str, torch.Tensor] = {} + sources = 0 + for item in slices: + start = int(item.row_start) + if start < 0 or (item.tensor.dim() == 0 and start): + raise LayoutError(f"{item.name}: bad first row {start}") + for target, tensor in transform_tensor(item.name, item.tensor, layout, smoothing).items(): + spec = layout.tensors[target] + if tensor.dim() != len(spec.shape) or tuple(tensor.shape[1:]) != spec.shape[1:] \ + or (tensor.dim() and start + tensor.shape[0] > spec.shape[0]): + raise LayoutError(f"{target}: slice {tuple(tensor.shape)} at row {start} " + f"does not fit the layout shape {spec.shape}") + key = f"{target}{_KEY_SEP}{start}" + if key in out: + raise LayoutError(f"{target}: rows from {start} given twice on rank {rank}") + # A cast to the same dtype returns the caller's tensor, possibly a view + # that shares storage with another slice; safetensors refuses shared storage. + out[key] = tensor.clone(memory_format=torch.contiguous_format) + sources += 1 + path = part_path(part_dir, rank) + path.parent.mkdir(parents=True, exist_ok=True) + tmp = path.with_name(path.name + ".tmp") + save_file(out, str(tmp), metadata={"format": "pt"}) + os.replace(tmp, path) + return {"sources": sources, "tensors": len(out)} + + +def _part_index(part_dir: Path) -> dict[str, list[tuple[int, Path, str]]]: + """target -> [(first row, part file, key)] over every part file, sorted by row.""" + index: dict[str, list[tuple[int, Path, str]]] = defaultdict(list) + for path in sorted(Path(part_dir).glob("part-*.safetensors")): + for key in read_header(path): + target, _, start = key.rpartition(_KEY_SEP) + index[target].append((int(start), path, key)) + return {t: sorted(v, key=lambda x: x[0]) for t, v in index.items()} + + +def merge_parts(part_dir: Path, template_dir: Path, out_dir: Path, + shards: Iterable[str] | None = None) -> dict[str, int]: + """Write the given output shards (default: all) from the part files. + + Fails when a tensor of these shards has a gap, an overlap, or a part names a + tensor the layout does not have.""" + layout = CheckpointLayout.load(template_dir) + all_shards = layout.shards() + selected = list(all_shards) if shards is None else list(shards) + if unknown := set(selected) - set(all_shards): + raise LayoutError(f"unknown output shards {sorted(unknown)}") + index = _part_index(part_dir) + if stray := sorted(set(index) - set(layout.tensors)): + raise LayoutError(f"parts hold {len(stray)} tensors the layout does not have, e.g. {stray[:3]}") + out_dir = Path(out_dir) + out_dir.mkdir(parents=True, exist_ok=True) + files: dict[Path, Any] = {} + stats = {"shards": 0, "tensors": 0, "slices": 0} + try: + for shard in selected: + tensors: dict[str, torch.Tensor] = {} + for name in all_shards[shard]: + pieces = index.get(name) + if not pieces: + raise LayoutError(f"{name}: no rank gave this tensor") + read = [] + for start, path, key in pieces: + if path not in files: + files[path] = safe_open(str(path), framework="pt", device="cpu") + read.append((start, files[path].get_tensor(key))) + tensors[name] = _join_rows(name, read, layout.tensors[name].shape) + stats["slices"] += len(read) + write_output_shard(out_dir, shard, tensors, layout) + stats["shards"] += 1 + stats["tensors"] += len(tensors) + finally: + files.clear() + return stats + + +def _join_rows(name: str, pieces: list[tuple[int, torch.Tensor]], shape: tuple[int, ...]) -> torch.Tensor: + if not shape: + if len(pieces) != 1: + raise LayoutError(f"{name}: scalar given {len(pieces)} times") + return pieces[0][1] + row = 0 + for start, piece in pieces: + if start != row: + kind = "overlap" if start < row else "gap" + raise LayoutError(f"{name}: rows {min(start, row)}..{max(start, row)} are an {kind} between ranks") + row += piece.shape[0] + if row != shape[0]: + raise LayoutError(f"{name}: rows {row}..{shape[0]} are a gap (no rank gave them)") + return pieces[0][1] if len(pieces) == 1 else torch.cat([p for _, p in pieces], dim=0) + + +def requantize_distributed(slices: Iterable[SourceSlice], template_dir: Path, out_dir: Path, + smoothing: AwqSmoothing, *, rank: int, world_size: int, + barrier: Callable[[], None]) -> dict[str, Any]: + """Every rank calls this with its own slices; the complete checkpoint is in + ``out_dir`` when it returns on any rank. + + ``out_dir`` must be on storage all ranks share (the release directory). + ``barrier()`` synchronizes the ranks (:func:`torch_barrier` under + ``torch.distributed``). Steps: write the part; barrier; merge shards + ``rank::world_size``; barrier; rank 0 writes the index and metadata and + removes the parts; barrier.""" + rank, world_size = int(rank), int(world_size) + if not 0 <= rank < world_size: + raise ValueError(f"rank {rank} not in a world of {world_size}") + template_dir, out_dir = Path(template_dir), Path(out_dir) + part_dir = out_dir / PART_DIR + stats: dict[str, Any] = {"part": write_part(part_dir, rank, slices, template_dir, smoothing)} + barrier() + shards = list(CheckpointLayout.load(template_dir).shards())[rank::world_size] + stats["merge"] = merge_parts(part_dir, template_dir, out_dir, shards) + barrier() + if rank == 0: + finalize_checkpoint(template_dir, out_dir) + shutil.rmtree(part_dir) + barrier() + return stats + + +def torch_barrier() -> None: + """``torch.distributed.barrier`` of the default group (the trainer's).""" + import torch.distributed as dist + + dist.barrier() diff --git a/yeto/rl/quant/stacked.py b/yeto/rl/quant/stacked.py new file mode 100644 index 000000000..f7a59d676 --- /dev/null +++ b/yeto/rl/quant/stacked.py @@ -0,0 +1,111 @@ +"""Trainer-native stacked expert layout -> per-expert source names (S25 N8.5b). + +The quantization layout (:mod:`.requant`) names one BF16 source matrix per +routed expert, as the HF per-expert checkpoint does:: + + .mlp.experts..gate_proj.weight [I, H] + .mlp.experts..up_proj.weight [I, H] + .mlp.experts..down_proj.weight [H, I] + +Trainers hold the experts stacked. The HF native layout (FSDP2 / verl, and the +copy in Modal ``glm-rl-tglm:/native/apex-flash-1-cecb5eeb``) stores, per layer:: + + .mlp.experts.gate_up_proj [E, 2I, H] rows 0..I-1 gate, I..2I-1 up + .mlp.experts.down_proj [E, H, I] + +Megatron (Miles) stores one ``linear_fc1`` ``[2I, H]`` (gate rows, then up +rows) and one ``linear_fc2`` ``[H, I]`` per local expert; its adapter +(:mod:`yeto.rl.adapters.miles.quant_shards`) maps them with +:func:`split_gate_up`. + +This module only renames and slices: the values are not changed, so a stacked +source quantizes to the same bytes as the per-expert one. +""" + +from __future__ import annotations + +import re +from collections.abc import Callable, Iterator, Mapping, Sequence +from dataclasses import dataclass + +import torch + +STACKED_EXPERTS = re.compile(r"^(?P.+\.mlp\.experts)\.(?Pgate_up_proj|down_proj)$") +GATE, UP, DOWN = "gate_proj", "up_proj", "down_proj" + + +def expert_weight_name(experts_prefix: str, expert: int, proj: str) -> str: + """``...weight`` (``prefix`` ends in ``.mlp.experts``).""" + return f"{experts_prefix}.{int(expert)}.{proj}.weight" + + +def split_gate_up(experts_prefix: str, expert: int, gate_up: torch.Tensor) -> Iterator[tuple[str, torch.Tensor]]: + """The gate and up matrices of one expert's merged ``[2I, H]`` weight (gate rows first).""" + if gate_up.dim() != 2 or gate_up.shape[0] % 2: + raise ValueError(f"{experts_prefix}.{expert}: merged gate/up weight must be [2I, H], got {tuple(gate_up.shape)}") + inner = gate_up.shape[0] // 2 + yield expert_weight_name(experts_prefix, expert, GATE), gate_up[:inner] + yield expert_weight_name(experts_prefix, expert, UP), gate_up[inner:] + + +def unstack_experts(name: str, tensor: torch.Tensor, first_expert: int = 0) -> Iterator[tuple[str, torch.Tensor]]: + """Per-expert matrices of a stacked ``gate_up_proj`` / ``down_proj`` (or of its + dim-0 slice that starts at expert ``first_expert``).""" + match = STACKED_EXPERTS.match(name) + if match is None: + raise ValueError(f"{name}: not a stacked expert tensor") + if tensor.dim() != 3: + raise ValueError(f"{name}: stacked experts must be 3-D, got {tuple(tensor.shape)}") + prefix = match["prefix"] + for j in range(tensor.shape[0]): + if match["kind"] == "down_proj": + yield expert_weight_name(prefix, first_expert + j, DOWN), tensor[j] + else: + yield from split_gate_up(prefix, first_expert + j, tensor[j]) + + +@dataclass(frozen=True) +class ExpertView: + """Rows ``[row_start, row_stop)`` of expert ``expert`` in the stored stacked tensor ``stored``.""" + + stored: str + expert: int + row_start: int + row_stop: int + + def read(self, stacked_slice) -> torch.Tensor: + """Read the view from a safetensors slice of ``stored`` (only these rows are read).""" + return stacked_slice[self.expert, self.row_start:self.row_stop] + + +@dataclass(frozen=True) +class StackedSourceIndex: + """Per-expert names of the stacked tensors in a checkpoint index.""" + + views: Mapping[str, ExpertView] + stacked_names: frozenset[str] + + @classmethod + def from_names(cls, names: Sequence[str] | Mapping[str, object], + shape_of: Callable[[str], Sequence[int]]) -> "StackedSourceIndex": + """``shape_of(name)`` is called for stacked tensors only (a header read).""" + views: dict[str, ExpertView] = {} + stacked = [n for n in names if STACKED_EXPERTS.match(n)] + for name in stacked: + match = STACKED_EXPERTS.match(name) + shape = tuple(shape_of(name)) + if len(shape) != 3: + raise ValueError(f"{name}: stacked experts must be 3-D, got {shape}") + experts, rows = shape[0], shape[1] + if match["kind"] == "down_proj": + parts = ((DOWN, 0, rows),) + else: + if rows % 2: + raise ValueError(f"{name}: merged gate/up has an odd row count {rows}") + parts = ((GATE, 0, rows // 2), (UP, rows // 2, rows)) + for e in range(experts): + for proj, a, b in parts: + views[expert_weight_name(match["prefix"], e, proj)] = ExpertView(name, e, a, b) + if clash := sorted(set(views) & set(names)): + raise ValueError(f"source has both stacked and per-expert experts, e.g. {clash[:3]}") + return cls(views, frozenset(stacked)) From 3e7cb32a5664df638ad1ee7465a3235d8705c4bc Mon Sep 17 00:00:00 2001 From: michael <78113554+chowhsu@users.noreply.github.com> Date: Mon, 19 Oct 2026 12:20:52 +0000 Subject: [PATCH 16/34] Keep serving the checked weights when a release has the wrong expert count (S25 N2) The release expert count is read from the release index before the engine loads anything, so a mismatch leaves the engine on its checked version. The loader used to mark the engine not ready anyway, which stopped generation and forced a restart on the next load. Now the release is refused: the engine stays ready, the failure event carries engine_touched=false and the kept version, and the next good release loads in place. Refusals still count toward max_consecutive_failures, so a run that only publishes bad releases stops. Co-Authored-By: Claude Opus 5.5 --- tests/test_rl_hetero_sampler_engines.py | 33 ++++++++++++++++++-- yeto/rl/engine/sampler_checks.py | 40 ++++++++++++++++++------- 2 files changed, 61 insertions(+), 12 deletions(-) diff --git a/tests/test_rl_hetero_sampler_engines.py b/tests/test_rl_hetero_sampler_engines.py index 721a84edc..e727958c0 100644 --- a/tests/test_rl_hetero_sampler_engines.py +++ b/tests/test_rl_hetero_sampler_engines.py @@ -271,7 +271,6 @@ def test_loader_serves_a_release_only_after_all_checks_pass(tmp_path): ("silent_skip", "expert_count"), # the engine loaded fewer experts (A-INF item 1) ("wrong_answers", "fixed_prompts"), ("probe_jump", "logprob_probe"), - ("release_short", "release_expert_count"), ]) def test_a_failed_check_marks_the_engine_not_ready_and_restarts_next_time(tmp_path, fault, step): from yeto.rl.engine.policy_release import ReleaseNotLoaded @@ -281,7 +280,7 @@ def test_a_failed_check_marks_the_engine_not_ready_and_restarts_next_time(tmp_pa events = [] loader = _loader(engine, events, max_probe_jump=0.5) loader.start(None, 0) - release = _w4a8_release(releases, 1, experts=3 if fault == "release_short" else 4) + release = _w4a8_release(releases, 1) if fault == "silent_skip": engine.expert_matrices = 0 elif fault == "wrong_answers": @@ -297,6 +296,36 @@ def test_a_failed_check_marks_the_engine_not_ready_and_restarts_next_time(tmp_pa loader(release, releases.resolve(release)) +def test_a_release_with_the_wrong_expert_count_keeps_the_served_weights(tmp_path): + """The release is refused before the engine is touched: v0 stays served and ready.""" + from yeto.rl.engine.policy_release import ReleaseNotLoaded + + releases = PolicyReleaseStore(CheckpointStore(tmp_path)) + engine = FakeEngine(identity=W4A8, answers=ANSWERS, expert_matrices=4) + events = [] + loader = _loader(engine, events) + loader.start(None, 0) + bad = _w4a8_release(releases, 1, experts=3) + with pytest.raises(ReleaseNotLoaded, match="release_expert_count") as refused: + loader(bad, releases.resolve(bad)) + assert refused.value.reason == "release_refused:release_expert_count" + assert loader.ready and engine.served_version() == 0 + (failed,) = [f for n, f in events if n == LOAD_CHECK_FAILED_EVENT] + assert failed["engine_touched"] is False and failed["kept_version"] == 0 + assert failed["detail"] == {"in_release": 3, "expected": 4} + + good = _w4a8_release(releases, 2) # the engine was never touched: an in-place load follows + loader(good, releases.resolve(good)) + assert loader.ready and engine.served_version() == 2 and loader.failures == 0 + checked = [f for n, f in events if n == "rl_sampler_load_checked"] + assert checked[-1]["reload_mode"] == "in_place" + + with pytest.raises(ReleaseNotLoaded): + loader(bad, releases.resolve(bad)) + with pytest.raises(SamplerLoadFatal): # still counted: a run that only publishes bad releases stops + loader(bad, releases.resolve(bad)) + + def test_an_engine_error_inside_a_check_is_a_failed_check(tmp_path): from yeto.rl.engine.policy_release import ReleaseNotLoaded from yeto.rl.engine.sampler_engine import SamplerEngineError diff --git a/yeto/rl/engine/sampler_checks.py b/yeto/rl/engine/sampler_checks.py index 498dc8723..36d702748 100644 --- a/yeto/rl/engine/sampler_checks.py +++ b/yeto/rl/engine/sampler_checks.py @@ -31,10 +31,14 @@ 5. checks that the engine reports the release version (``served_version``). Only then is the release served. A failed step writes -``sampler_load_check_failed``, marks the engine not ready (the loop generates -nothing) and raises ``ReleaseNotLoaded`` so the follower tries again; after -``max_consecutive_failures`` failures it raises :class:`SamplerLoadFatal` and -the island stops with the reason. +``sampler_load_check_failed`` and raises ``ReleaseNotLoaded`` so the follower +tries again; after ``max_consecutive_failures`` failures in a row it raises +:class:`SamplerLoadFatal` and the island stops with the reason. A step that +failed after the engine was touched (load, post-load checks) marks the engine +not ready (the loop generates nothing) and the next load restarts it. A +release refused before the engine was touched (scheme, release expert count) +leaves the engine serving its checked version (``ready`` unchanged; the event +carries ``engine_touched: false`` and ``kept_version``). """ from __future__ import annotations @@ -254,6 +258,7 @@ def __init__(self, engine: SamplerEngine, checks: PostLoadChecks, emit: Callable self.clock = clock self.ready = False # True after a load whose checks all passed self.failures = 0 + self._restart = False # a failed load or check left the engine in an unknown state self.last_load: dict[str, Any] | None = None # reload_mode, pause_s, checks @property @@ -280,15 +285,29 @@ def __call__(self, release: PolicyRelease, directory: Path) -> None: self._serve(release, Path(directory), int(release.version), first=False) def _fail(self, version: int, step: str, detail: Any) -> None: + """A load or check failed: the engine holds unchecked weights, so it stops serving.""" self.ready = False + self._restart = True + self._count_failure(version, step, detail, engine_touched=True) + raise ReleaseNotLoaded(f"v{version}: {step} failed: {str(detail)[:300]}", + reason=f"check_failed:{step}") + + def _refuse(self, version: int, step: str, detail: Any) -> None: + """The release failed a check before the engine was touched: the engine + keeps serving its checked version (``ready`` is unchanged).""" + kept = self.engine.served_version() if self.ready else None + self._count_failure(version, step, detail, engine_touched=False, kept_version=kept) + raise ReleaseNotLoaded(f"v{version}: {step} failed: {str(detail)[:300]}; " + f"keeping v{kept}", reason=f"release_refused:{step}") + + def _count_failure(self, version: int, step: str, detail: Any, **fields: Any) -> None: self.failures += 1 self.emit(LOAD_CHECK_FAILED_EVENT, policy_version=version, step=step, detail=detail, - failures=self.failures, max_consecutive_failures=self.max_consecutive_failures) + failures=self.failures, max_consecutive_failures=self.max_consecutive_failures, + **fields) if self.failures >= self.max_consecutive_failures: raise SamplerLoadFatal(f"post-load checks failed {self.failures} times in a row " f"(last: {step} for v{version}): {str(detail)[:300]}") - raise ReleaseNotLoaded(f"v{version}: {step} failed: {str(detail)[:300]}", - reason=f"check_failed:{step}") def _serve(self, release: PolicyRelease | None, directory: Path | None, version: int, *, first: bool) -> None: @@ -296,9 +315,9 @@ def _serve(self, release: PolicyRelease | None, directory: Path | None, version: if directory is not None and expected is not None: in_release = count_release_experts(directory) if in_release is not None and in_release != expected: - self._fail(version, "release_expert_count", - {"in_release": in_release, "expected": expected}) - restart = first or self.failures > 0 # after a failed check: restart, never patch in place + self._refuse(version, "release_expert_count", + {"in_release": in_release, "expected": expected}) + restart = first or self._restart # after a failed load or check: restart, never patch in place self.ready = False # the engine holds unchecked weights from here until the checks pass started = self.clock() try: @@ -319,6 +338,7 @@ def _serve(self, release: PolicyRelease | None, directory: Path | None, version: self.checks.accept() self.ready = True self.failures = 0 + self._restart = False self.last_load = {"reload_mode": report.reload_mode, "pause_s": round(report.pause_s, 3), "load_and_check_s": round(self.clock() - started, 3), "scheme": self.scheme, "checks": results, **dict(report.detail)} From 00df5975a1b1cca08d675d988a5d897ef6cae4bb Mon Sep 17 00:00:00 2001 From: michael <78113554+chowhsu@users.noreply.github.com> Date: Mon, 19 Oct 2026 12:21:50 +0000 Subject: [PATCH 17/34] Keep the transfer error when the boat staging cleanup also fails (T-REV3) download_bytes and download_dir removed the VM staging files in a finally block, so a failing cleanup replaced the transfer's own error. The cleanup now runs after the transfer; when the transfer raised, a cleanup failure is logged and the transfer's error propagates. Also drop a duplicate import and give os.environ.get a string default in the Miles agent function (PLW1508). Co-Authored-By: Claude Opus 5.5 --- tests/test_harbor_boat.py | 48 ++++++++++++++++++++++++++ yeto/rl/adapters/miles/harbor_agent.py | 2 +- yeto/rl/harness/harbor/boat_sandbox.py | 34 ++++++++++++------ 3 files changed, 73 insertions(+), 11 deletions(-) diff --git a/tests/test_harbor_boat.py b/tests/test_harbor_boat.py index 5d0a38a95..c6001a8c2 100644 --- a/tests/test_harbor_boat.py +++ b/tests/test_harbor_boat.py @@ -285,6 +285,54 @@ async def go(): assert any("rm -rf" in c for c in fake.commands) # staging is cleaned up +class _CleanupFails(FakeBoat): + """The staging cleanup command cannot start (boat API error); ``transfer_fails``: get_file too.""" + + def __init__(self, *, transfer_fails: bool) -> None: + super().__init__() + self.transfer_fails = transfer_fails + self.cleanups = 0 + + def start_command(self, sandbox_id: str, command: str) -> int: + if command.startswith("sudo -n rm -"): + self.cleanups += 1 + raise BoatError("command API: 503") + return super().start_command(sandbox_id, command) + + def get_file(self, sandbox_id: str, path: str) -> bytes: + if self.transfer_fails: + raise BoatError("file API: 502") + return super().get_file(sandbox_id, path) + + +def _download(fake: FakeBoat, tmp_path: Path, which: str) -> None: + sandbox = _sandbox(fake) + + async def go(): + await sandbox.start(SPEC) + try: + if which == "bytes": + await sandbox.download_bytes("/logs/verifier/reward.txt") + else: + await sandbox.download_dir("/logs/verifier", tmp_path / "out") + finally: + await sandbox.close() + + asyncio.run(go()) + + +@pytest.mark.parametrize("which", ["bytes", "dir"]) +def test_a_failed_download_keeps_its_error_when_the_cleanup_also_fails(tmp_path: Path, which: str): + fake = _CleanupFails(transfer_fails=True) + with pytest.raises(BoatError, match="file API: 502"): # not the cleanup's 503 + _download(fake, tmp_path, which) + assert fake.cleanups == 1 and fake.deleted # cleanup tried, VM deleted + + fake = _CleanupFails(transfer_fails=False) # after a good transfer the cleanup error is raised + with pytest.raises(BoatError, match="command API: 503"): + _download(fake, tmp_path, which) + + def test_pack_and_extract_round_trip(tmp_path: Path): (tmp_path / "a" / "empty").mkdir(parents=True) (tmp_path / "a" / "f").write_text("x") diff --git a/yeto/rl/adapters/miles/harbor_agent.py b/yeto/rl/adapters/miles/harbor_agent.py index 4ebd2ce11..7506cf571 100644 --- a/yeto/rl/adapters/miles/harbor_agent.py +++ b/yeto/rl/adapters/miles/harbor_agent.py @@ -82,7 +82,7 @@ def task_from_metadata(metadata: dict[str, Any]) -> AgentTask: defaults = { "agent_name": DEFAULT_AGENT, "provider": os.environ.get("HARBOR_ENV_TYPE", "").strip().lower(), - "wall_clock_s": float(os.environ.get("AGENT_TRIAL_TIMEOUT", 7200)), + "wall_clock_s": float(os.environ.get("AGENT_TRIAL_TIMEOUT", "7200")), "model_name": os.environ.get("AGENT_MODEL_NAME", "model"), "agent_timeout_s": os.environ.get("AGENT_TIMEOUT") or None, # as the example reads it } diff --git a/yeto/rl/harness/harbor/boat_sandbox.py b/yeto/rl/harness/harbor/boat_sandbox.py index 4dfc78b7c..990bab720 100644 --- a/yeto/rl/harness/harbor/boat_sandbox.py +++ b/yeto/rl/harness/harbor/boat_sandbox.py @@ -20,19 +20,20 @@ import asyncio import contextlib -import contextlib import io import logging import shlex import tarfile import uuid -from collections.abc import Callable, Mapping +from collections.abc import Awaitable, Callable, Mapping from dataclasses import dataclass from pathlib import Path +from typing import TypeVar from .boat_client import FAILED_STATES, READY_STATES, BoatClient, BoatError, BoatQuota logger = logging.getLogger(__name__) +T = TypeVar("T") CONTAINER = "main" STAGING_DIR = "/tmp/yeto-xfer" # the boat file API only reaches /home/user and /tmp @@ -258,10 +259,8 @@ async def download_bytes(self, source_path: str) -> bytes: s = shlex.quote(staged) await self._checked(f"mkdir -p {STAGING_DIR} && {DOCKER} cp -L {CONTAINER}:{shlex.quote(source_path)} {s} " f"&& sudo -n chmod a+r {s}", f"download {source_path}") - try: - return await asyncio.to_thread(self._client.get_file, sandbox_id, staged) - finally: - await self.vm_exec(f"sudo -n rm -f {s}", timeout_s=60) + return await self._then_clean(asyncio.to_thread(self._client.get_file, sandbox_id, staged), + f"sudo -n rm -f {s}") async def download_dir(self, source_dir: str, target_dir: Path) -> None: sandbox_id = self._require_started() @@ -271,12 +270,27 @@ async def download_dir(self, source_dir: str, target_dir: Path) -> None: f"mkdir -p {s} && {DOCKER} cp {CONTAINER}:{shlex.quote(source_dir.rstrip('/') + '/.')} {s} && " f"sudo -n tar -czf {s}.tar.gz -C {s} . && sudo -n chmod a+r {s}.tar.gz", f"download dir {source_dir}") - try: - archive = await asyncio.to_thread(self._client.get_file, sandbox_id, f"{staged}.tar.gz") - finally: - await self.vm_exec(f"sudo -n rm -rf {s} {s}.tar.gz", timeout_s=60) + archive = await self._then_clean(asyncio.to_thread(self._client.get_file, sandbox_id, f"{staged}.tar.gz"), + f"sudo -n rm -rf {s} {s}.tar.gz") await asyncio.to_thread(extract_dir, archive, target_dir) + async def _then_clean(self, work: Awaitable[T], cleanup: str) -> T: + """Await ``work``, then run the VM ``cleanup`` command. + + When ``work`` raises, a failing cleanup is logged and the original + error propagates (a ``finally`` would replace it with the cleanup's). + """ + try: + result = await work + except BaseException: + try: + await self.vm_exec(cleanup, timeout_s=60) + except Exception: # the work's error is the one to report + logger.warning("boat staging cleanup failed after an error: %s", cleanup, exc_info=True) + raise + await self.vm_exec(cleanup, timeout_s=60) + return result + def _require_started(self) -> str: if self.sandbox_id is None: raise BoatError("boat sandbox is not started (or already deleted)") From 204e7327e99f9755f18bab52e83320b41c783f2a Mon Sep 17 00:00:00 2001 From: michael <78113554+chowhsu@users.noreply.github.com> Date: Mon, 19 Oct 2026 12:22:48 +0000 Subject: [PATCH 18/34] Define PPU and NPU engine images, cpu_arch node selection, Megatron memory estimate S25 N0.2, N1.6, N7.1b. Nothing is built: both images are written only. N0.2: docker/ppu-train (x86_64, PPU trainer + SGLang sampler, on the Alibaba inference-xpu-pytorch-opt v0.5.17 base, plus Megatron-LM, mbridge, Miles glm5_next and verl) and docker/npu-sampler (aarch64, vllm-ascend 0.30.0rc1 on CANN 9.x). Each has a versions.lock.json that carries the versions and the decisions behind them, and docker/image-assert/assert_image_versions.py checks the image against its lock at every container start (and at build time for ppu-train), exiting 1 before the engine starts. The npu1 CANN gap (host 8.0.1 vs 9.x) is handled by the container's own CANN user-space, with a host driver upgrade kept as the fallback; the driver minimum is recorded, not enforced, until the first NPU probe measures it. N1.6: the node inventory requires cpu_arch, accepts a top-level images map (cpu_arch -> image), and refuses a node whose image is a yeto catalog image of another architecture or device family. The SSH trainer island compares uname -m with cpu_arch before any start. N7.1b: the memory estimator gains framework="megatron" (TP/PP/EP/CP, distributed optimizer by group, PP in-flight micro batches, mHC residual streams, the DSA dense fallback) and labels every output as derived, not measured. Co-Authored-By: Claude Opus 5.5 --- docker/README.md | 43 +++ docker/image-assert/assert_image_versions.py | 180 ++++++++++ docker/image-assert/entrypoint.sh | 8 + docker/npu-sampler/Dockerfile | 60 ++++ docker/npu-sampler/versions.lock.json | 47 +++ docker/ppu-train/Dockerfile | 108 ++++++ docker/ppu-train/requirements.txt | 56 ++++ docker/ppu-train/versions.lock.json | 44 +++ tests/test_engine_images.py | 216 ++++++++++++ tests/test_rl_fullparam_memory.py | 75 +++++ tests/test_rl_ssh_node_pool.py | 68 +++- tests/test_rl_ssh_trainer_island.py | 24 +- yeto/launch/images.py | 336 ++++--------------- yeto/launch/node_inventory.py | 90 ++++- yeto/launch/ssh_trainer_island.py | 15 +- yeto/rl/fullparam/memory.py | 148 +++++++- 16 files changed, 1225 insertions(+), 293 deletions(-) create mode 100644 docker/README.md create mode 100755 docker/image-assert/assert_image_versions.py create mode 100755 docker/image-assert/entrypoint.sh create mode 100644 docker/npu-sampler/Dockerfile create mode 100644 docker/npu-sampler/versions.lock.json create mode 100644 docker/ppu-train/Dockerfile create mode 100644 docker/ppu-train/requirements.txt create mode 100644 docker/ppu-train/versions.lock.json create mode 100644 tests/test_engine_images.py diff --git a/docker/README.md b/docker/README.md new file mode 100644 index 000000000..b6f98407c --- /dev/null +++ b/docker/README.md @@ -0,0 +1,43 @@ +# Engine images + +| Directory | Image | CPU arch | Device family | Roles | Status | +|---|---|---|---|---|---| +| `miles-ports/` | `yeto-miles-ports` | x86_64 | cuda | trainer + sampler | built and pushed (see the Dockerfile header) | +| `yeto-release/` | release layer on `yeto-miles-ports` | x86_64 | cuda | — | built once (2026-10-19) | +| `ppu-train/` | `yeto-ppu-train` | x86_64 | ppu | trainer + sampler | written, never built (S25 N0.2) | +| `npu-sampler/` | `yeto-npu-sampler` | aarch64 | npu | sampler | written, never built (S25 N0.2) | + +`yeto/launch/images.py` is the catalog the launcher reads. A node of the SSH +inventory declares `cpu_arch`; a node with no `image` takes the inventory's +`images[cpu_arch]`, and a node whose image is a catalog image of another +architecture or device family is refused before any SSH (N1.6). + +## Version locks + +Each image directory has `versions.lock.json`. It is the single source of the +versions the image asserts, and it holds the decisions behind them (the +`decisions` list). `docker/image-assert/assert_image_versions.py` reads it: + +* the image entrypoint (`docker/image-assert/entrypoint.sh`) runs it at every + container start and exits 1 before the engine starts when a check fails; +* `yeto-ppu-train` also runs it at build time, so an image whose own lock does + not hold is never pushed. `yeto-npu-sampler` cannot: `import torch_npu` + needs the NPU driver, which no build host has. + +Checked: CPU architecture, Python, package versions and imports, build-time +commits (`/opt/yeto/image/commits.json`), the CANN toolkit version, the host +driver version and required files. An entry with `"version": null` is recorded +in the report but not compared; `"optional": true` reports a missing package or +failed import without failing (TransformerEngine and tilelang on PPU, whose +support the first PPU probe decides). + +## Open points (both new images) + +1. Neither image was built: no base digest, no `pip check` output, no real + version report. Every version in the two locks comes from a survey or from + upstream documentation, not from a built image. +2. `yeto-ppu-train`: the base image's registry path and digest are unknown + (N0.1 read the tag `inference-xpu-pytorch-opt:v0.5.17` on ppu1 only), and + `MEGATRON_COMMIT` has no value yet. +3. `yeto-npu-sampler`: the ModelArts fixes come from a build context that does + not exist in this repository yet, and the arm64 release layer is a follow-up. diff --git a/docker/image-assert/assert_image_versions.py b/docker/image-assert/assert_image_versions.py new file mode 100755 index 000000000..2b7c7b71f --- /dev/null +++ b/docker/image-assert/assert_image_versions.py @@ -0,0 +1,180 @@ +#!/usr/bin/env python3 +"""Check an engine image against its version lock at container start (S25 N0.2). + +Usage: ``assert_image_versions.py LOCK.json`` (the image entrypoint runs it +before the command). Exit 0 when every check passes, 1 otherwise; the report +goes to stderr, one line per check, and to ``--report`` as JSON when given. + +The lock (``docker//versions.lock.json``) lists: + +* ``cpu_arch``: must equal ``platform.machine()``; +* ``python``: ``major.minor`` prefix of the interpreter version (null: record); +* ``packages``: ``{dist: {"version", "import", "optional"}}``. ``version`` is + compared with the public part of the installed version (``2.10.0+abc`` -> + ``2.10.0``); null means "record only". ``import`` names a module the image + must import. ``optional`` true turns a missing package or failed import into + a report line instead of a failure (used where the probe on the real machine + decides, e.g. TransformerEngine on PPU); +* ``commits``: ``{name: sha}`` compared with ``/opt/yeto/image/commits.json`` + written at build time (null: record); +* ``cann``: ``{"version", "version_files"}``: the toolkit version read from the + first existing file (``version=`` / ``Version=`` line) must start with + ``version``; +* ``driver``: ``{"min", "enforce", "version_file"}``: the host driver mounted + into the container; ``min`` null means record only; ``enforce`` false + reports a driver below ``min`` without failing (kept while a probe on the + real host decides, see the NPU lock); +* ``required_files``: paths that must exist (patches, stubs). + +Stdlib only: it runs before anything else in the image is trusted. +""" + +from __future__ import annotations + +import argparse +import importlib +import importlib.metadata +import json +import platform +import re +import sys +from pathlib import Path + +COMMITS_FILE = Path("/opt/yeto/image/commits.json") +_VERSION_LINE = re.compile(r"^\s*(?:version|Version)\s*=\s*(\S+)", re.MULTILINE) + + +class Report: + def __init__(self) -> None: + self.rows: list[dict] = [] + + def add(self, check: str, ok: bool, detail: str, *, optional: bool = False) -> None: + status = "ok" if ok else ("not-enforced" if optional else "FAIL") + self.rows.append({"check": check, "status": status, "detail": detail}) + print(f"[image-assert] {status:16} {check}: {detail}", file=sys.stderr) + + @property + def failed(self) -> list[dict]: + return [row for row in self.rows if row["status"] == "FAIL"] + + +def public_version(version: str) -> str: + return version.split("+", 1)[0] + + +def _version_tuple(version: str) -> tuple[int, ...]: + return tuple(int(part) for part in re.findall(r"\d+", version)) + + +def check_packages(packages: dict, report: Report) -> None: + for dist, spec in packages.items(): + optional = bool(spec.get("optional")) + want = spec.get("version") + try: + have = importlib.metadata.version(dist) + except importlib.metadata.PackageNotFoundError: + report.add(f"package {dist}", False, "not installed", optional=optional) + continue + if want is None: + report.add(f"package {dist}", True, f"{have} (recorded, not pinned)") + else: + ok = public_version(have) == want + report.add(f"package {dist}", ok, f"{have} (lock {want})", optional=optional) + if not ok: + continue + module = spec.get("import") + if module: + try: + importlib.import_module(module) + except Exception as exc: # noqa: BLE001 - any import failure is the finding + report.add(f"import {module}", False, f"{type(exc).__name__}: {exc}", optional=optional) + else: + report.add(f"import {module}", True, "imported") + + +def check_commits(commits: dict, report: Report, commits_file: Path) -> None: + if not commits: + return + try: + built = json.loads(commits_file.read_text()) + except (OSError, ValueError) as exc: + report.add("commits", False, f"cannot read {commits_file}: {exc}") + return + for name, want in commits.items(): + have = built.get(name) + if want is None: + report.add(f"commit {name}", have is not None, f"{have} (recorded, not pinned)") + else: + report.add(f"commit {name}", have == want, f"{have} (lock {want})") + + +def _read_version(files: list[str]) -> tuple[str | None, str | None]: + for name in files: + path = Path(name) + if path.is_file(): + match = _VERSION_LINE.search(path.read_text(errors="replace")) + if match: + return match.group(1), name + return None, None + + +def check_cann(cann: dict | None, report: Report) -> None: + if not cann: + return + have, source = _read_version(cann["version_files"]) + if have is None: + report.add("cann", False, f"no version in {cann['version_files']}") + else: + report.add("cann", have.startswith(cann["version"]), f"{have} from {source} (lock {cann['version']})") + + +def check_driver(driver: dict | None, report: Report) -> None: + if not driver: + return + have, source = _read_version([driver["version_file"]]) + want = driver.get("min") + if have is None: + report.add("driver", False, f"no version in {driver['version_file']} (is the host driver mounted?)") + elif want is None: + report.add("driver", True, f"{have} from {source} (recorded; no minimum in the lock)") + else: + report.add("driver", _version_tuple(have) >= _version_tuple(want), f"{have} (minimum {want})", + optional=not driver.get("enforce", True)) + + +def run(lock: dict, *, commits_file: Path = COMMITS_FILE, machine: str | None = None) -> Report: + report = Report() + machine = machine or platform.machine() + report.add("cpu_arch", machine == lock["cpu_arch"], f"{machine} (lock {lock['cpu_arch']})") + want_py = lock.get("python") + have_py = platform.python_version() + report.add("python", want_py is None or have_py.startswith(want_py + ".") or have_py == want_py, + f"{have_py} (lock {want_py})") + check_packages(lock.get("packages", {}), report) + check_commits(lock.get("commits", {}), report, commits_file) + check_cann(lock.get("cann"), report) + check_driver(lock.get("driver"), report) + for name in lock.get("required_files", []): + report.add(f"file {name}", Path(name).exists(), "present" if Path(name).exists() else "absent") + return report + + +def main(argv: list[str] | None = None) -> int: + ap = argparse.ArgumentParser(description=__doc__.splitlines()[0]) + ap.add_argument("lock", type=Path) + ap.add_argument("--report", type=Path, help="also write the rows as JSON here") + ap.add_argument("--commits-file", type=Path, default=COMMITS_FILE) + args = ap.parse_args(argv) + lock = json.loads(args.lock.read_text()) + report = run(lock, commits_file=args.commits_file) + if args.report: + args.report.write_text(json.dumps({"image": lock.get("image"), "rows": report.rows}, indent=2)) + if report.failed: + print(f"[image-assert] {len(report.failed)} check(s) failed; refusing to start " + f"image {lock.get('image')}", file=sys.stderr) + return 1 + return 0 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/docker/image-assert/entrypoint.sh b/docker/image-assert/entrypoint.sh new file mode 100755 index 000000000..cdcd03de6 --- /dev/null +++ b/docker/image-assert/entrypoint.sh @@ -0,0 +1,8 @@ +#!/bin/sh +# Engine image entrypoint (S25 N0.2): check the image against its version lock, +# then run the command. A failed check stops the container (exit 1) before any +# engine starts. The report is kept at /tmp/yeto-image-assert.json. +set -e +python3 /opt/yeto/image/assert_image_versions.py /opt/yeto/image/versions.lock.json \ + --report /tmp/yeto-image-assert.json +exec "$@" diff --git a/docker/npu-sampler/Dockerfile b/docker/npu-sampler/Dockerfile new file mode 100644 index 000000000..6ffbce804 --- /dev/null +++ b/docker/npu-sampler/Dockerfile @@ -0,0 +1,60 @@ +# yeto-npu-sampler: Ascend 910B (A2) vllm-ascend sampler image, linux/arm64 +# (S25 N0.2, stage 2). WRITTEN, NEVER BUILT. +# +# Versions: versions.lock.json in this directory is the single source. +# +# docker build --platform linux/arm64 -f docker/npu-sampler/Dockerfile \ +# --build-arg BASE_IMAGE=quay.io/ascend/vllm-ascend:v0.30.0rc1@sha256: \ +# --build-context fixes= . +# +# Base: the official vllm-ascend v0.30.0rc1 A2 image. It carries the CANN +# user-space (toolkit, nnal), torch / torch_npu, vLLM and vllm-ascend; the host +# provides only the driver and firmware, mounted at run time +# (/usr/local/Ascend/driver). The user ran GLM-5.3-Flash W4A8 on ModelArts with +# vllm-ascend 0.30.0rc1 (as a rootfs, CANN 9.1); the vllm-ascend compatibility +# matrix (mkdocs.yml, commit 0ca6a8b) names CANN 9.2.0 and Ascend HDK 26.0.RC1 +# for that release. Host npu1 has driver 24.1.0.3 and CANN 8.0.1: see +# versions.lock.json "decisions" for how that gap is handled. +# +# yeto code is not copied: the sampler adapter runs from /opt/yeto/src (the +# release layer, as for the other images). The release build script is +# linux/amd64 only today, so an arm64 release layer is a follow-up. +ARG BASE_IMAGE +FROM ${BASE_IMAGE} + +ARG BASE_IMAGE +RUN case "${BASE_IMAGE}" in *@sha256:*) ;; \ + *) echo "BASE_IMAGE must be pinned by digest: '${BASE_IMAGE}'" >&2; exit 1;; esac + +# Freeze the base exactly as docker/ppu-train does: nothing below may replace a +# base package (torch, torch_npu, vllm, vllm_ascend, triton_ascend, ...). +RUN mkdir -p /opt/yeto/image \ + && pip freeze > /opt/yeto/image/base-freeze.txt \ + && python3 -c 'import importlib.metadata as m; print("\n".join(sorted({f"{d.name}=={d.version}" for d in m.distributions()})))' \ + > /opt/yeto/image/base-constraints.txt +ENV PIP_CONSTRAINT=/opt/yeto/image/base-constraints.txt + +# yeto sampler-side runtime (HTTP client, weight follower, post-load checks). +RUN pip install --no-cache-dir cloudpickle "huggingface-hub>=0.34" "safetensors>=0.4" pyyaml "httpx[http2]" + +# The two run-time fixes the user needed on ModelArts (A-INF items 15 and 17): +# the libprofapi symbol stub (CANN 9.1 HCCL hang; whether CANN 9.2 still needs +# it is unknown) and the W4A8 host-memory staging change (the adapter turns it +# on with VLLM_ASCEND_W4A8_CPU_STAGE=1). Their source is the user's ModelArts +# environment, not this repository, so they arrive as a build context: +# apply.sh installs them and writes /opt/yeto/image/fixes.json +# ({"fixes": [{"name", "sha256"}, ...]}). The PT_INTERP change is for the +# no-root rootfs only; a container brings its own loader. The build stops when +# the context is missing or apply.sh writes nothing. +COPY --from=fixes . /tmp/yeto-fixes +RUN sh /tmp/yeto-fixes/apply.sh \ + && test -s /opt/yeto/image/fixes.json \ + && rm -rf /tmp/yeto-fixes + +COPY docker/npu-sampler/versions.lock.json /opt/yeto/image/versions.lock.json +COPY docker/image-assert/assert_image_versions.py docker/image-assert/entrypoint.sh /opt/yeto/image/ +ENV PYTHONPATH=/opt/yeto/src +# No build-time check here: torch_npu imports need the driver, which no build +# host has. The entrypoint runs the whole check on the node instead. +ENTRYPOINT ["/opt/yeto/image/entrypoint.sh"] +CMD ["bash"] diff --git a/docker/npu-sampler/versions.lock.json b/docker/npu-sampler/versions.lock.json new file mode 100644 index 000000000..5431d5db9 --- /dev/null +++ b/docker/npu-sampler/versions.lock.json @@ -0,0 +1,47 @@ +{ + "image": "yeto-npu-sampler", + "cpu_arch": "aarch64", + "device_family": "npu", + "base_image": { + "name": "quay.io/ascend/vllm-ascend:v0.30.0rc1", + "reference": null, + "note": "digest unknown: no pull was made from this host (npu1 reaches hf-mirror, quay was not tested)" + }, + "python": "3.12", + "packages": { + "torch": {"version": "2.10.0", "import": "torch"}, + "torch-npu": {"version": "2.10.0.post4", "import": "torch_npu"}, + "vllm": {"version": "0.30.0", "import": "vllm"}, + "vllm-ascend": {"version": "0.30.0rc1", "import": "vllm_ascend"}, + "triton-ascend": {"version": "3.2.2", "import": null, "optional": true}, + "transformers": {"version": null, "import": "transformers"}, + "compressed-tensors": {"version": null, "import": "compressed_tensors"}, + "safetensors": {"version": null, "import": "safetensors"} + }, + "commits": {}, + "cann": { + "version": "9.", + "version_files": [ + "/usr/local/Ascend/ascend-toolkit/latest/version.cfg", + "/usr/local/Ascend/ascend-toolkit/latest/aarch64-linux/ascend_toolkit_install.info" + ] + }, + "driver": { + "min": "26.0", + "enforce": false, + "version_file": "/usr/local/Ascend/driver/version.info" + }, + "required_files": [ + "/usr/local/Ascend/driver/version.info", + "/opt/yeto/image/fixes.json" + ], + "decisions": [ + "CANN gap, chosen option A: the container carries its own CANN user-space (the official vllm-ascend image) and uses the host driver as it is. Reasons: (1) the user already ran this stack on ModelArts without touching a host driver; (2) npu1 is shared with a colleague and is in use, so a driver upgrade needs the owner's consent and a reboot; (3) a container can be removed, a driver cannot. The host keeps CANN 8.0.1; nothing in the container reads it.", + "CANN gap, option B (not chosen, kept as the fallback): upgrade the npu1 driver and firmware to the Ascend HDK release that matches CANN 9.2 (26.0.RC1 per the vllm-ascend matrix) and then run either the container or a rootfs. Take this only if option A fails on the driver, and only with the machine owner's agreement.", + "driver minimum: recorded, not enforced (enforce false). The container's CANN user-space usually runs on an older driver, but the exact floor for CANN 9.2 on driver 24.1.0.3 is unknown. The first NPU probe reads the driver version and tries `import torch_npu` plus one small device operation; if it fails, option B applies and enforce becomes true with the measured floor.", + "cann version: the check only requires a 9.x toolkit in the container, because the base image may ship 9.1 or 9.2 depending on when it was built.", + "triton-ascend: optional; the GLM5.3 sampler path uses vllm-ascend's own kernels, and this entry only records what the base carries.", + "W4A8 checkpoint: the sampler loads the int-quantized `checkpoint/` directory, not `checkpoint-packed/` (the user's 10-15 result). This is an engine-adapter rule, not an image check.", + "transformers / compressed-tensors: recorded, not pinned, until the base image's versions are read." + ] +} diff --git a/docker/ppu-train/Dockerfile b/docker/ppu-train/Dockerfile new file mode 100644 index 000000000..49755e8ec --- /dev/null +++ b/docker/ppu-train/Dockerfile @@ -0,0 +1,108 @@ +# yeto-ppu-train: PPU trainer (Megatron + Miles glm5_next, or verl FSDP2) and +# SGLang sampler image, linux/amd64 (S25 N0.2). WRITTEN, NEVER BUILT. +# +# Versions: versions.lock.json in this directory is the single source; the +# build args below must equal it (tests/test_engine_images.py checks this). +# +# docker build --platform linux/amd64 -f docker/ppu-train/Dockerfile \ +# --build-arg BASE_IMAGE=/inference-xpu-pytorch-opt@sha256: \ +# --build-arg MEGATRON_COMMIT= . +# +# Build context: the yeto repository root. yeto code is NOT copied here: the +# release layer (scripts/build_yeto_release_image.sh, docker/yeto-release) puts +# one clean commit at /opt/yeto/src on top of this image. +# +# Base: the Alibaba PPU image the user runs as `ppu_sgl` on ppu1-4 +# (inference-xpu-pytorch-opt:v0.5.17: sglang 0.5.17, torch 2.10.0, triton 3.6.0, +# fla 1.0.0, transformers 5.12.1, ray 2.31.0; SESSION25/N0.1-ppu-survey.md). +# Its registry path and digest are not known yet (read them on ppu1 with +# `docker inspect`), so BASE_IMAGE has no default and must be a digest. +ARG BASE_IMAGE +FROM ${BASE_IMAGE} + +ARG BASE_IMAGE +RUN case "${BASE_IMAGE}" in *@sha256:*) ;; \ + *) echo "BASE_IMAGE must be pinned by digest: '${BASE_IMAGE}'" >&2; exit 1;; esac + +# Freeze the base: every later pip install runs with these constraints, so a +# dependency that would replace a base package fails the build. ray is left +# out on purpose (Miles needs >= 2.56; versions.lock.json "decisions"). +# Built from package metadata, not `pip freeze`, so editable and URL installs +# (the base's sglang may be one) are constrained too. +RUN mkdir -p /opt/yeto/image \ + && pip freeze > /opt/yeto/image/base-freeze.txt \ + && python3 -c 'import importlib.metadata as m; print("\n".join(sorted({f"{d.name}=={d.version}" for d in m.distributions() if d.name.lower() != "ray"})))' \ + > /opt/yeto/image/base-constraints.txt + +ENV PIP_CONSTRAINT=/opt/yeto/image/base-constraints.txt + +COPY docker/ppu-train/requirements.txt /opt/yeto/image/requirements.txt +RUN pip install --no-cache-dir -r /opt/yeto/image/requirements.txt + +# Megatron-LM: radixark fork, branch miles-main (the only one with +# megatron/core/transformer/hyper_connection.py, needed by glm5_next mHC). +# No commit is pinned yet (Miles' own image follows the branch head), so the +# build refuses to run without one. +ARG MEGATRON_REPOSITORY=https://github.com/radixark/Megatron-LM +ARG MEGATRON_COMMIT +RUN test -n "${MEGATRON_COMMIT}" || { echo "MEGATRON_COMMIT is required" >&2; exit 1; } \ + && git init -q /opt/Megatron-LM && cd /opt/Megatron-LM \ + && git remote add origin "${MEGATRON_REPOSITORY}" \ + && git fetch -q --depth 1 origin "${MEGATRON_COMMIT}" \ + && git checkout -q --detach FETCH_HEAD \ + && test "$(git rev-parse HEAD)" = "${MEGATRON_COMMIT}" \ + && pip install --no-cache-dir --no-deps -e /opt/Megatron-LM + +# mbridge (HF <-> Megatron weights; Miles glm5_next bridge plugin builds on it), +# Megatron-Bridge and two Megatron data extras, all without deps as upstream. +ARG MBRIDGE_COMMIT=89eb10887887bc74853f89a4de258c0702932a1c +ARG MEGATRON_BRIDGE_COMMIT=2e09c234a3272285140224d0d698593418b55ba5 +RUN pip install --no-cache-dir --no-deps \ + "git+https://github.com/ISEEKYAN/mbridge.git@${MBRIDGE_COMMIT}" \ + "git+https://github.com/radixark/Megatron-Bridge.git@${MEGATRON_BRIDGE_COMMIT}" \ + megatron-energon multi-storage-client + +# tilelang: the DSA SparseMLA kernels of Miles glm5_next need it. Whether it +# runs on PPU is the first PPU probe (N7.5); the lock marks it optional. +ARG TILELANG_VERSION=0.1.14 +RUN pip install --no-cache-dir --no-deps "tilelang==${TILELANG_VERSION}" apache-tvm-ffi==0.1.11 \ + || echo "tilelang ${TILELANG_VERSION} did not install; DSA falls back to torch (N7.5)" >&2 + +# Miles at the GLM 5.3 Flash commit (radixark/miles #2786). T-N77 builds the +# yeto fork branch on this commit; switch MILES_* to it when that is pushed. +ARG MILES_REPOSITORY=https://github.com/radixark/miles +ARG MILES_COMMIT=cc76e23915b2132ecfff65b97331fd83b02637e6 +RUN git init -q /opt/miles && cd /opt/miles \ + && git remote add origin "${MILES_REPOSITORY}" \ + && git fetch -q --depth 1 origin "${MILES_COMMIT}" \ + && git checkout -q --detach FETCH_HEAD \ + && test "$(git rev-parse HEAD)" = "${MILES_COMMIT}" \ + && pip install --no-cache-dir --no-deps -e /opt/miles + +# verl fork at the commit yeto pins (yeto/rl/adapters/verl/pins.py). Its +# runtime deps are in requirements.txt; vllm is not installed (the PPU +# sampler is SGLang). +ARG VERL_REPOSITORY=https://github.com/michaellchung/verl +ARG VERL_COMMIT=acad9875a8bdfc81afbcb0a50d146b2630f44093 +RUN git init -q /opt/verl && cd /opt/verl \ + && git remote add origin "${VERL_REPOSITORY}" \ + && git fetch -q --depth 1 origin "${VERL_COMMIT}" \ + && git checkout -q --detach FETCH_HEAD \ + && test "$(git rev-parse HEAD)" = "${VERL_COMMIT}" \ + && pip install --no-cache-dir --no-deps -e /opt/verl + +# Commits the start check compares with the lock. +RUN printf '{"megatron": "%s", "mbridge": "%s", "megatron_bridge": "%s", "miles": "%s", "verl": "%s"}\n' \ + "${MEGATRON_COMMIT}" "${MBRIDGE_COMMIT}" "${MEGATRON_BRIDGE_COMMIT}" "${MILES_COMMIT}" "${VERL_COMMIT}" \ + > /opt/yeto/image/commits.json \ + && pip check || echo "pip check reported conflicts (kept in the build log)" >&2 + +COPY docker/ppu-train/versions.lock.json /opt/yeto/image/versions.lock.json +COPY docker/image-assert/assert_image_versions.py docker/image-assert/entrypoint.sh /opt/yeto/image/ +ENV PYTHONPATH=/opt/yeto/src +# Fail at build time too: an image whose own lock does not hold is never pushed. +# (Imports that need a PPU device are not in the lock; the start check repeats +# everything on the node.) +RUN python3 /opt/yeto/image/assert_image_versions.py /opt/yeto/image/versions.lock.json +ENTRYPOINT ["/opt/yeto/image/entrypoint.sh"] +CMD ["bash"] diff --git a/docker/ppu-train/requirements.txt b/docker/ppu-train/requirements.txt new file mode 100644 index 000000000..32d08651d --- /dev/null +++ b/docker/ppu-train/requirements.txt @@ -0,0 +1,56 @@ +# Python packages yeto-ppu-train adds on top of the PPU base image (S25 N0.2). +# Installed with `-c /opt/yeto/image/base-constraints.txt`: a package that would +# replace a base package (torch, triton, sglang, fla, transformers, ...) fails +# the build instead of changing the base. ray is the one deliberate exception +# (see versions.lock.json "decisions"). +# +# Miles cc76e2391 requirements.txt, without the CUDA-only entries the PPU base +# cannot use: nvidia-resiliency-ext, torchft-nightly (pins its own torch), +# ring_flash_attn (needs a flash_attn CUDA build). +accelerate +blake3==1.0.9 +blobfile +datasets +httpx[http2] +kubernetes_asyncio +mcp[cli] +memray +omegaconf +onnxscript +openai +pillow +polars==1.42.1 +psycopg[binary] +pybase64 +pylatexenc +pyyaml +qwen_vl_utils +ray[default]==2.56.0 +safetensors>=0.8.0 +sglang-router>=0.2.3 +tensorboard +tqdm +typing_extensions>=4.5 +wandb +xxhash==3.7.1 +zstandard==0.25.0 +# verl acad9875 requirements.txt, runtime part (no dev tools, no TransferQueue). +codetiming +dill +fastapi +hydra-core +latex2sympy2_extended +math_verify +numpy>=2.0.0 +pandas +pyarrow>=19.0.0 +tensordict>=0.8.0,<=0.10.0,!=0.9.0 +torchdata +uvicorn +# yeto runtime (pyproject dependencies without torch / transformers, which the +# base pins; yeto itself runs from source, see the Dockerfile). +cloudpickle +huggingface-hub>=0.34 +peft==0.20.0 +protobuf +sentencepiece diff --git a/docker/ppu-train/versions.lock.json b/docker/ppu-train/versions.lock.json new file mode 100644 index 000000000..011829c22 --- /dev/null +++ b/docker/ppu-train/versions.lock.json @@ -0,0 +1,44 @@ +{ + "image": "yeto-ppu-train", + "cpu_arch": "x86_64", + "device_family": "ppu", + "base_image": { + "name": "inference-xpu-pytorch-opt:v0.5.17", + "reference": null, + "note": "registry path and digest unknown: read on ppu1 with docker inspect (N0.1 saw the tag only)" + }, + "python": null, + "packages": { + "torch": {"version": "2.10.0", "import": "torch"}, + "triton": {"version": "3.6.0", "import": "triton"}, + "flash-linear-attention": {"version": "1.0.0", "import": "fla"}, + "sglang": {"version": "0.5.17", "import": "sglang"}, + "transformers": {"version": "5.12.1", "import": "transformers"}, + "ray": {"version": "2.56.0", "import": "ray"}, + "peft": {"version": "0.20.0", "import": "peft"}, + "compressed-tensors": {"version": null, "import": "compressed_tensors"}, + "megatron-core": {"version": null, "import": "megatron.core.transformer.hyper_connection"}, + "mbridge": {"version": null, "import": "mbridge"}, + "miles": {"version": null, "import": "miles"}, + "verl": {"version": null, "import": "verl"}, + "transformer-engine": {"version": null, "import": "transformer_engine.pytorch", "optional": true}, + "tilelang": {"version": "0.1.14", "import": "tilelang", "optional": true} + }, + "commits": { + "megatron": null, + "mbridge": "89eb10887887bc74853f89a4de258c0702932a1c", + "megatron_bridge": "2e09c234a3272285140224d0d698593418b55ba5", + "miles": "cc76e23915b2132ecfff65b97331fd83b02637e6", + "verl": "acad9875a8bdfc81afbcb0a50d146b2630f44093" + }, + "required_files": [], + "decisions": [ + "ray: the base has 2.31.0, Miles cc76e2391 requires >= 2.56.0. The image installs ray 2.56.0 (Miles floor). Not checked: sglang 0.5.17 with ray 2.56. The NPU sampler talks HTTP, not Ray, so the two images need not share a ray version.", + "fla: the base 1.0.0 is kept; Miles pins 0.5.2 plus a KDA patch. Compatibility of Miles glm5_next kda.py with fla 1.0.0 is the first PPU probe (N7.5).", + "TransformerEngine: recorded, optional. Megatron runs with TE by default; --transformer-impl local is the fallback (main agent 10-19). The PPU probe decides.", + "tilelang: optional; the DSA SparseMLA path needs it, its PPU support is unknown.", + "megatron commit: not pinned (Miles' image follows radixark/Megatron-LM miles-main); the build requires MEGATRON_COMMIT and records it.", + "compressed-tensors: recorded, not pinned, until the base's version is read (A-INF item 6 asks for a pin).", + "yeto: not installed as a package (pyproject pins torch < 2.9 and transformers 5.13); it runs from /opt/yeto/src via the release layer." + ] +} diff --git a/tests/test_engine_images.py b/tests/test_engine_images.py new file mode 100644 index 000000000..52c9d0d79 --- /dev/null +++ b/tests/test_engine_images.py @@ -0,0 +1,216 @@ +"""S25 N0.2: the image catalog, the version locks and the start check.""" + +from __future__ import annotations + +import importlib.util +import json +import re +from pathlib import Path + +import pytest + +from yeto.launch import images +from yeto.rl.adapters.miles import pins as miles_pins +from yeto.rl.adapters.verl import pins as verl_pins + +DOCKER = Path(__file__).resolve().parents[1] / "docker" +ASSERT_SCRIPT = DOCKER / "image-assert" / "assert_image_versions.py" + + +def _load_assert_module(): + spec = importlib.util.spec_from_file_location("assert_image_versions", ASSERT_SCRIPT) + module = importlib.util.module_from_spec(spec) + spec.loader.exec_module(module) + return module + + +checker = _load_assert_module() + + +def _lock(name: str) -> dict: + return json.loads(images.IMAGES[name].lock_path.read_text()) + + +def _dockerfile(name: str) -> str: + return (DOCKER / name / "Dockerfile").read_text() + + +# --- catalog ---------------------------------------------------------------- + + +def test_every_catalog_image_has_a_dockerfile_and_a_lock_that_matches_it(): + for name, spec in images.IMAGES.items(): + assert (DOCKER / name / "Dockerfile").is_file(), name + lock = _lock(name) + assert lock["image"] == spec.repository + assert (lock["cpu_arch"], lock["device_family"]) == (spec.cpu_arch, spec.device_family) + assert spec.cpu_arch in images.CPU_ARCHES + assert lock["decisions"], f"{name}: the lock must record the decisions behind its versions" + + +def test_canonical_cpu_arch_accepts_the_docker_spellings_and_refuses_others(): + assert images.canonical_cpu_arch("amd64") == "x86_64" + assert images.canonical_cpu_arch("ARM64") == "aarch64" + assert images.canonical_cpu_arch(" x86_64\n") == "x86_64" + with pytest.raises(images.UnknownCpuArch, match="cpu_arch 'ppc64le' is unknown"): + images.canonical_cpu_arch("ppc64le") + + +@pytest.mark.parametrize("reference, expected", [ + ("reg.example/yeto-ppu-train:1", "ppu-train"), + ("reg.example/a/b/yeto-npu-sampler:tag@sha256:ab", "npu-sampler"), + ("reg.example/yeto-ppu-train", "ppu-train"), + ("reg.example/site-own:1", None), + ("yeto-ppu-train-extra:1", None), +]) +def test_catalog_image_recognizes_a_reference_by_repository(reference, expected): + spec = images.catalog_image(reference) + assert (spec.name if spec else None) == expected + + +def test_image_mismatch_only_judges_catalog_images(): + assert images.image_mismatch("site/own:1", cpu_arch="aarch64", device_family="npu") is None + assert images.image_mismatch("r/yeto-npu-sampler:1", cpu_arch="aarch64", device_family="npu") is None + assert "the node is x86_64" in images.image_mismatch( + "r/yeto-npu-sampler:1", cpu_arch="x86_64", device_family="npu") + assert "device_family cuda" in images.image_mismatch( + "r/yeto-ppu-train:1", cpu_arch="x86_64", device_family="cuda") + + +# --- the locks against the Dockerfiles and the yeto pins -------------------- + + +def test_ppu_lock_matches_the_dockerfile_build_args_and_the_yeto_pins(): + lock, text = _lock("ppu-train"), _dockerfile("ppu-train") + args = dict(re.findall(r"^ARG (\w+)=(\S+)$", text, re.MULTILINE)) + assert lock["commits"]["miles"] == args["MILES_COMMIT"] + assert lock["commits"]["verl"] == args["VERL_COMMIT"] == verl_pins.VERL_COMMIT + assert lock["commits"]["mbridge"] == args["MBRIDGE_COMMIT"] + assert lock["commits"]["megatron_bridge"] == args["MEGATRON_BRIDGE_COMMIT"] + assert lock["packages"]["tilelang"]["version"] == args["TILELANG_VERSION"] + assert lock["packages"]["peft"]["version"] == miles_pins.MILES_PEFT_VERSION + # No commit for Megatron yet: the build demands one and records it. + assert lock["commits"]["megatron"] is None and "MEGATRON_COMMIT is required" in text + + +def test_ppu_lock_keeps_the_ppu_base_versions_the_survey_read(): + packages = _lock("ppu-train")["packages"] + assert packages["torch"]["version"] == "2.10.0" + assert packages["triton"]["version"] == "3.6.0" + assert packages["flash-linear-attention"]["version"] == "1.0.0" + assert packages["sglang"]["version"] == "0.5.17" + # TransformerEngine and tilelang must not fail the start check before the + # first PPU probe says whether they work (N7.5). + assert packages["transformer-engine"]["optional"] and packages["tilelang"]["optional"] + assert packages["megatron-core"]["import"].endswith("hyper_connection") # mHC + + +def test_npu_lock_matches_the_vllm_ascend_release_line(): + lock = _lock("npu-sampler") + assert lock["packages"]["vllm-ascend"]["version"] == "0.30.0rc1" + assert lock["packages"]["torch-npu"]["version"] == "2.10.0.post4" + assert lock["cann"]["version"] == "9." # 9.x in the container, not the host's 8.0.1 + assert lock["driver"]["enforce"] is False + assert any("option A" in d for d in lock["decisions"]) + assert any("option B" in d for d in lock["decisions"]) + + +def test_requirements_of_the_ppu_image_leave_the_base_packages_alone(): + text = (DOCKER / "ppu-train" / "requirements.txt").read_text() + pinned = {line.split("=")[0].split(">")[0].split("[")[0].strip() + for line in text.splitlines() if line and not line.startswith("#")} + assert not pinned & {"torch", "triton", "sglang", "flash-linear-attention", "transformers"} + assert "ray[default]==2.56.0" in text # the one deliberate base change + + +def test_both_new_images_run_the_start_check_from_their_entrypoint(): + for name in ("ppu-train", "npu-sampler"): + text = _dockerfile(name) + assert 'ENTRYPOINT ["/opt/yeto/image/entrypoint.sh"]' in text, name + assert "docker/image-assert/assert_image_versions.py" in text, name + assert "BASE_IMAGE must be pinned by digest" in text, name + entrypoint = (DOCKER / "image-assert" / "entrypoint.sh").read_text() + assert "set -e" in entrypoint and 'exec "$@"' in entrypoint + + +# --- the start check -------------------------------------------------------- + + +BASE_LOCK = {"image": "test", "cpu_arch": "x86_64", "python": None, "packages": {}} + + +def _rows(report) -> dict[str, str]: + return {row["check"]: row["status"] for row in report.rows} + + +def test_start_check_passes_on_a_matching_image(): + lock = {**BASE_LOCK, "packages": {"pytest": {"version": pytest.__version__, "import": "pytest"}}} + report = checker.run(lock, machine="x86_64") + assert not report.failed and _rows(report)["import pytest"] == "ok" + + +def test_start_check_fails_on_the_wrong_arch_version_or_import(): + report = checker.run({**BASE_LOCK, "cpu_arch": "aarch64"}, machine="x86_64") + assert [row["check"] for row in report.failed] == ["cpu_arch"] + lock = {**BASE_LOCK, "packages": { + "pytest": {"version": "0.0.1"}, + "no-such-dist-yeto": {"version": "1.0"}, + "pytest-asyncio": {"version": None, "import": "yeto_no_such_module"}, + }} + failed = [row["check"] for row in checker.run(lock, machine="x86_64").failed] + assert failed == ["package pytest", "package no-such-dist-yeto", "import yeto_no_such_module"] + + +def test_start_check_records_an_unpinned_version_and_tolerates_an_optional_gap(): + lock = {**BASE_LOCK, "packages": { + "pytest": {"version": None}, + "no-such-dist-yeto": {"version": "1.0", "optional": True}, + "pytest-asyncio": {"version": None, "import": "yeto_no_such_module", "optional": True}, + }} + report = checker.run(lock, machine="x86_64") + assert not report.failed + assert _rows(report)["package no-such-dist-yeto"] == "not-enforced" + + +def test_start_check_compares_the_public_version_only(): + assert checker.public_version("2.10.0+abc") == "2.10.0" + + +def test_start_check_reads_commits_cann_driver_and_required_files(tmp_path): + (tmp_path / "commits.json").write_text(json.dumps({"miles": "abc", "verl": "def", + "megatron": "012"})) + (tmp_path / "version.cfg").write_text("package=Ascend-cann-toolkit\nversion=9.2.0\n") + (tmp_path / "driver.info").write_text("Version=26.0.RC1\n") + lock = {**BASE_LOCK, + "commits": {"miles": "abc", "verl": "xxx", "megatron": None, "absent": None}, + "cann": {"version": "9.", "version_files": [str(tmp_path / "missing"), + str(tmp_path / "version.cfg")]}, + "driver": {"min": "24.1", "version_file": str(tmp_path / "driver.info")}, + "required_files": [str(tmp_path / "commits.json"), str(tmp_path / "absent")]} + report = checker.run(lock, machine="x86_64", commits_file=tmp_path / "commits.json") + rows = _rows(report) + assert rows["commit miles"] == "ok" and rows["commit verl"] == "FAIL" + # Recorded, not pinned: present passes, absent fails (the build must write it). + assert rows["commit megatron"] == "ok" and rows["commit absent"] == "FAIL" + assert rows["cann"] == "ok" and rows["driver"] == "ok" + assert rows[f"file {tmp_path / 'absent'}"] == "FAIL" + + +def test_start_check_fails_a_driver_below_the_minimum_unless_the_lock_says_otherwise(tmp_path): + info = tmp_path / "driver.info" + info.write_text("Version=24.1.0.3\n") + lock = {**BASE_LOCK, "driver": {"min": "26.0", "version_file": str(info)}} + assert [row["check"] for row in checker.run(lock, machine="x86_64").failed] == ["driver"] + lock["driver"]["enforce"] = False + assert not checker.run(lock, machine="x86_64").failed + + +def test_start_check_exit_code_and_report_file(tmp_path): + lock_path = tmp_path / "lock.json" + lock_path.write_text(json.dumps({**BASE_LOCK, "cpu_arch": "aarch64"})) + report = tmp_path / "report.json" + assert checker.main([str(lock_path), "--report", str(report)]) == 1 + assert json.loads(report.read_text())["rows"][0]["check"] == "cpu_arch" + lock_path.write_text(json.dumps({**BASE_LOCK, "cpu_arch": images.canonical_cpu_arch( + __import__("platform").machine())})) + assert checker.main([str(lock_path)]) == 0 diff --git a/tests/test_rl_fullparam_memory.py b/tests/test_rl_fullparam_memory.py index d706d7e34..ff8710060 100644 --- a/tests/test_rl_fullparam_memory.py +++ b/tests/test_rl_fullparam_memory.py @@ -116,3 +116,78 @@ def test_shape_from_hf_config_counts_a_tiny_model(): assert shape.largest_layer_params == layer assert shape.params == 2 * layer + 2 * 64 * 32 + 32 # embed + lm_head + final norm assert shape.head_counts == (2, 4) and shape.dsa_layers == 0 and GIB == 1024**3 + + +# --- N7.1b Megatron variant ------------------------------------------------- + +# Synthetic MoE: 1000 params, 600 of them in 8 routed experts, 4 layers. +MOE = ModelShape(params=1000, largest_layer_params=250, num_layers=4, hidden_size=8, vocab_size=16, + residual_streams=4, head_counts=(4,), expert_params=600, num_experts=8) +MEG = TrainPlan(cards=16, cards_per_node=8, hbm_gib=144, host_memory_gib=1024, seq_len=64, + framework="megatron", tp=2, pp=2, ep=4, chunked_loss=True, activation_margin_gib=0.0) + + +def test_megatron_state_splits_dense_by_tp_pp_and_experts_by_ep_pp(): + est = estimate(MOE, MEG) + dense, experts = 400 / 2 / 2, 600 / 4 / 2 # (P - E) / TP / PP, E / EP / PP + assert est.items["params_local_bf16"].bytes == pytest.approx(2 * (dense + experts)) + assert est.items["main_grads_local"].bytes == pytest.approx(4 * (dense + experts)) + # 4 B master + 8 B Adam; dense DP = 16 / (2*2) = 4, expert DP = 16 / (4*2) = 2 + assert est.items["optimizer_shard"].bytes == pytest.approx(12 * (dense / 4 + experts / 2)) + plain = estimate(MOE, replace(MEG, distributed_optimizer=False)) + assert plain.items["optimizer_local"].bytes == pytest.approx(12 * (dense + experts)) + assert "gathered_layers" not in est.items and est.total("host") == 0 + + +def test_megatron_activations_count_pp_micro_batches_and_mhc_streams(): + est = estimate(MOE, MEG) + # 2 stage layers x 2 in flight x 64/(TP 2 x CP 1) tokens x 4 streams x 8 hidden x 2 B + assert est.items["layer_inputs"].bytes == 2 * 2 * 32 * 4 * 8 * 2 + one_stream = estimate(replace(MOE, residual_streams=1), MEG) + assert est.items["layer_inputs"].bytes == 4 * one_stream.items["layer_inputs"].bytes + + +def test_megatron_uneven_pipeline_uses_the_largest_stage(): + est = estimate(replace(MOE, num_layers=5, largest_layer_params=200), replace(MEG, pp=2)) + assert "stage of 3/5 layers" in est.items["params_local_bf16"].formula + + +def test_megatron_output_is_labelled_derived(): + summary = estimate(MOE, MEG).summary() + assert summary["basis"] == "推算" and summary["framework"] == "megatron" + assert estimate(GLM, PPU).summary()["basis"] == "推算" + + +@pytest.mark.parametrize("change, match", [ + ({"ep": 3}, "ep x pp = 6 does not divide cards 16"), + ({"tp": 3}, "tp x pp x cp = 6 does not divide"), + ({"offload": "full"}, "only offload none"), + ({"ulysses": 2}, "use cp instead of ulysses"), + ({"pp": 8, "ep": 2}, "pp 8 exceeds the 4 layers"), + ({"framework": "deepspeed"}, "framework must be one of"), +]) +def test_malformed_megatron_plans_are_refused(change, match): + with pytest.raises(MemoryPlanError, match=match): + estimate(MOE, replace(MEG, **change)) + + +def test_megatron_refuses_cp_with_dsa_and_fsdp2_refuses_megatron_knobs(): + with pytest.raises(MemoryPlanError, match="DSA layer does not support cp"): + estimate(replace(MOE, dsa_layers=1, index_heads=4), replace(MEG, cp=2, pp=1, ep=8)) + with pytest.raises(MemoryPlanError, match="megatron only"): + estimate(GLM, replace(PPU, tp=2)) + with pytest.raises(MemoryPlanError, match="go together"): + replace(MOE, num_experts=0) + + +def test_megatron_glm_on_32_ppu_cards_needs_bf16_optimizer_state(): + """GLM5.3-Flash on 4 PPU nodes (32 x 144 GiB), TP4 EP32, 32K tokens. The + routed-expert share (95 % of parameters) is an assumption of this test, not + counted: expert DP is 32 / (EP 32 x PP 1) = 1, so FP32 Adam on the experts + is not divided by the distributed optimizer.""" + glm = replace(GLM, expert_params=int(GLM.params * 0.95), num_experts=288) + plan = TrainPlan(cards=32, cards_per_node=8, hbm_gib=144, host_memory_gib=2048, seq_len=32768, + framework="megatron", tp=4, ep=32, dsa_path="sparse_kernel", chunked_loss=True) + fp32 = estimate(glm, plan) + assert not fp32.fits and "optimizer_shard" in fp32.problems()[0] + assert estimate(glm, replace(plan, param_bytes=2, state_bytes=4)).fits diff --git a/tests/test_rl_ssh_node_pool.py b/tests/test_rl_ssh_node_pool.py index 6366cafe7..9130d375d 100644 --- a/tests/test_rl_ssh_node_pool.py +++ b/tests/test_rl_ssh_node_pool.py @@ -73,7 +73,7 @@ def npu_smi_info(count: int, hbm_mb: int = 65536, card: str = "910B4") -> str: def _inventory_node(**overrides) -> dict: node = {"name": "ascend00", "host": "user@172.28.11.10", "device_family": "npu", - "card": "910B4", "cards_per_node": 8, "hbm_gib": 64, + "cpu_arch": "aarch64", "card": "910B4", "cards_per_node": 8, "hbm_gib": 64, "host_memory_gib": 1024, "fabric": "roce", "interface": "eth0", "image": "registry.example/yeto-ascend:1"} node.update(overrides) @@ -88,7 +88,7 @@ def test_inventory_reads_every_declared_field(tmp_path): path.write_text(json.dumps({"nodes": [ _inventory_node(proxy_jump="user@bastion", mounts={"/mnt/shared": "/mnt/shared"}, price_per_card_hour=1.5, roles=["trainer"]), - _inventory_node(name="ppu00", host="user@ppu00", device_family="ppu", card="810E", + _inventory_node(name="ppu00", host="user@ppu00", device_family="ppu", cpu_arch="x86_64", card="810E", cards_per_node=16, hbm_gib=64, host_memory_gib=2048, fabric="roce", interface="hpn0", image="registry.example/yeto-ppu:1"), @@ -135,6 +135,68 @@ def test_inventory_refuses_a_bad_field(overrides, message): parse_inventory({"nodes": [_inventory_node(**overrides)]}) +# --- N1.6 cpu_arch and image selection ------------------------------------- + + +def test_inventory_requires_cpu_arch_and_normalizes_docker_spellings(): + with pytest.raises(InventoryError, match="ascend00: cpu_arch is required"): + parse_inventory({"nodes": [_inventory_node(cpu_arch=None)]}) + with pytest.raises(InventoryError, match="cpu_arch 'riscv64' is unknown"): + parse_inventory({"nodes": [_inventory_node(cpu_arch="riscv64")]}) + node = parse_inventory({"nodes": [_inventory_node(cpu_arch="arm64")]})[0] + assert node.cpu_arch == "aarch64" and node.record()["cpu_arch"] == "aarch64" + + +def test_inventory_selects_the_image_of_the_node_arch(): + images = {"arm64": "reg.example/yeto-npu-sampler:1", "x86_64": "reg.example/yeto-ppu-train:1"} + nodes = parse_inventory({"images": images, "nodes": [ + _inventory_node(image=None), + _inventory_node(name="ppu00", host="user@ppu00", device_family="ppu", cpu_arch="x86_64", + card="810E", image=None), + _inventory_node(name="ascend01", host="user@h1", image="site/own-ascend:2"), + ]}) + assert [n.image for n in nodes] == ["reg.example/yeto-npu-sampler:1", + "reg.example/yeto-ppu-train:1", "site/own-ascend:2"] + + +@pytest.mark.parametrize("overrides, message", [ + ({"image": "reg.example/yeto-ppu-train:1@sha256:ab"}, "x86_64 image ppu-train; the node is aarch64"), + ({"cpu_arch": "x86_64", "image": "reg.example/yeto-npu-sampler:1"}, + "aarch64 image npu-sampler; the node is x86_64"), + ({"image": None}, "set image, or images.aarch64"), +]) +def test_inventory_refuses_an_image_of_another_arch(overrides, message): + document = {"images": {"x86_64": "reg.example/yeto-ppu-train:1"}, + "nodes": [_inventory_node(**overrides)]} + with pytest.raises(InventoryError, match=message): + parse_inventory(document) + + +def test_inventory_refuses_a_catalog_image_of_another_device_family(): + with pytest.raises(InventoryError, match="ppu image ppu-train; the node is device_family cuda"): + parse_inventory({"nodes": [_inventory_node(device_family="cuda", cpu_arch="x86_64", + card="H200", image="r/yeto-ppu-train:1")]}) + + +@pytest.mark.parametrize("document, message", [ + ({"images": {"sparc": "r/x:1"}, "nodes": []}, "images: cpu_arch 'sparc' is unknown"), + ({"images": {"arm64": "r/a:1", "aarch64": "r/b:1"}, "nodes": []}, "aarch64 is listed twice"), + ({"images": ["r/x:1"], "nodes": []}, "images must be a map"), + ({"image": "r/x:1", "nodes": []}, "unknown top-level inventory field"), +]) +def test_inventory_refuses_a_bad_images_map(document, message): + with pytest.raises(InventoryError, match=message): + parse_inventory(document) + + +def test_node_checks_uname_against_cpu_arch(): + node = parse_inventory({"nodes": [_inventory_node()]})[0] + node.check_cpu_arch("aarch64\n") + node.check_cpu_arch("arm64") + with pytest.raises(InventoryError, match="uname -m reports x86_64, the inventory declares cpu_arch aarch64"): + node.check_cpu_arch("x86_64\n") + + def test_inventory_accepts_a_rootfs_node_without_an_image(): node = parse_inventory({"nodes": [_inventory_node(runtime="rootfs", image=None)]})[0] assert (node.runtime, node.image) == ("rootfs", None) @@ -466,7 +528,7 @@ def test_every_ssh_and_copy_to_a_jump_host_node_goes_through_the_bastion(): def test_a_node_without_a_jump_host_uses_plain_ssh(): - node = Node(name="n0", host="user@n0", device_family="cuda", card="H200", + node = Node(name="n0", host="user@n0", device_family="cuda", cpu_arch="x86_64", card="H200", cards_per_node=8, hbm_gib=141, host_memory_gib=2048) assert node.ssh_options() == [] assert node.preflight_commands()[0][0] == "nvidia-smi" diff --git a/tests/test_rl_ssh_trainer_island.py b/tests/test_rl_ssh_trainer_island.py index 74d9f1cd5..87c008b1b 100644 --- a/tests/test_rl_ssh_trainer_island.py +++ b/tests/test_rl_ssh_trainer_island.py @@ -46,7 +46,7 @@ def node(name="ppu0", host="u@10.0.0.1", **kw) -> Node: - base = dict(device_family="ppu", card="PPU", cards_per_node=8, hbm_gib=144.0, host_memory_gib=2048.0, + base = dict(device_family="ppu", cpu_arch="x86_64", card="PPU", cards_per_node=8, hbm_gib=144.0, host_memory_gib=2048.0, runtime="rootfs", fabric="roce", interface="eth0") return Node(name=name, host=host, **(base | kw)) @@ -115,11 +115,13 @@ def terminate(self): class FakeOps(SshOps): - def __init__(self, card_rc=0): - self.runs, self.starts, self.card_rc = [], [], card_rc + def __init__(self, card_rc=0, uname="x86_64"): + self.runs, self.starts, self.card_rc, self.uname = [], [], card_rc, uname def run(self, n, script, *, timeout=300.0): self.runs.append((n.name, script)) + if script == "uname -m": + return subprocess.CompletedProcess([], 0, stdout=self.uname + "\n", stderr="") rc = self.card_rc if "smi" in script else 0 return subprocess.CompletedProcess([], rc, stdout="cards", stderr="") @@ -157,6 +159,22 @@ def test_preflight_fails_when_no_card_tool_answers(tmp_path, checked_cards): islands.launch(1, None) +@pytest.mark.parametrize("uname", ["aarch64", "ppc64le"]) +def test_preflight_refuses_a_node_of_another_cpu_arch(tmp_path, checked_cards, uname): + ops = FakeOps(uname=uname) + islands = SshTrainerIslands(spec(), name="t", jobs_path=tmp_path / "j.json", ops=ops) + with pytest.raises(SshIslandError, match=f"ppu0: uname -m reports '?{uname}'?.*cpu_arch x86_64"): + islands.launch(1, None) + assert not ops.starts and not checked_cards + + +def test_preflight_accepts_the_docker_spelling_of_an_arch(tmp_path, checked_cards): + islands = SshTrainerIslands(spec(nodes=[node(cpu_arch="aarch64")]), name="t", + jobs_path=tmp_path / "j.json", ops=FakeOps(uname="aarch64")) + islands.preflight() + assert islands.checked + + def test_ssh_binding_reads_a_sharded_store_by_key(tmp_path, checked_cards): args = SimpleNamespace(rollout_batch_size=2, rl_checkpoint_store="oss://bkt/runs/a", rl_trainer_command_poll_seconds=None, tuning="full") diff --git a/yeto/launch/images.py b/yeto/launch/images.py index 9d9ebdff8..d98b9340c 100644 --- a/yeto/launch/images.py +++ b/yeto/launch/images.py @@ -1,287 +1,95 @@ -"""Learner image selection and the Miles source setup.""" +"""Engine images by CPU architecture (S25 N0.2, N1.6). -from __future__ import annotations - -import shlex -import sys - -from ..gpu_spec import ClusterSpec -from .rl_backend import _miles_overlay - - -# AWS keeps this SSM parameter pointing at the CURRENT Deep Learning Base -# OSS-NVIDIA-driver AMI per region (open kernel modules, Blackwell-ready), -# so resolving at launch time needs no region x AMI table of our own. -_DLAMI_SSM_PARAM = ( - "/aws/service/deeplearning/ami/x86_64/" - "base-oss-nvidia-driver-gpu-ubuntu-22.04/latest/ami-id" -) - - -def resolve_blackwell_image(region: str) -> str | None: - """Current Blackwell-capable DLAMI for a region, or None (callers then - rely on the setup-time driver remediation instead).""" - try: - import boto3 +yeto defines one image per (device family, CPU architecture) pair. Each image +has a Dockerfile and a version lock under ``docker//``; the lock is the +single source of the versions the image asserts at start +(``docker/image-assert/assert_image_versions.py``). - ssm = boto3.client("ssm", region_name=region) - return ssm.get_parameter(Name=_DLAMI_SSM_PARAM)["Parameter"]["Value"] - except Exception as e: # noqa: BLE001 - degrade to driver remediation - print( - f"[launcher] could not resolve a Blackwell DLAMI for {region} ({e}); " - "relying on setup-time driver install", - file=sys.stderr, - ) - return None +A node of the SSH node pool declares its ``cpu_arch``. The launcher then uses +this catalog in two ways: +* a node that names no image gets the inventory's image for its architecture; +* a node whose image is a catalog image of another architecture or another + device family is refused before any SSH (an aarch64 node cannot run an + x86_64 image, and a PPU image has no Ascend runtime). -# Internal image-override table: (cloud, GPU) pairs whose provider-default -# image is known stale/broken, mapped to a `region -> image id` resolver. -# First entry earned in production: sky's pinned AMI ships driver 535, -# which never binds to SM100 silicon. Second class earned the same way: -# NVSwitch instances (p4d/p4de A100, p5 H100, p5e H200) need -# nvidia-fabricmanager running before CUDA will initialize at all -# (cudaGetDeviceCount -> Error 802), and sky's pinned AMI does not ship it; -# the DL Base GPU AMI does, preinstalled and enabled. Extend here as new -# GPU generations outpace provider image pins; an explicit --learner-image -# always wins, and a resolver returning None degrades to the setup-time -# driver install. -GPU_IMAGE_OVERRIDES: dict[tuple[str, str], object] = { - ("aws", "B200"): resolve_blackwell_image, - # NVSwitch (fabric manager required); same current DL Base GPU AMI. - ("aws", "A100"): resolve_blackwell_image, - ("aws", "A100-80GB"): resolve_blackwell_image, - ("aws", "H100"): resolve_blackwell_image, - ("aws", "H200"): resolve_blackwell_image, -} +An image reference that is not in the catalog (a site image) is accepted as +it is: yeto cannot tell its architecture from the name, and the node +preflight (``uname -m``) still checks the node itself. +""" +from __future__ import annotations -# The Megatron backend runs INSIDE an NGC container that ships the whole -# stack prebuilt (torch + Transformer Engine + megatron-core + megatron-bridge) -# — the pip-on-DLAMI install proved intractable (two shakedowns; see -# docs/MEGATRON.md). sky runs the task in this container via `image_id: -# docker:...`; the host still supplies the GPU kernel driver, so B200 needs -# sky's default AWS docker host AMI to carry driver >=570 open-kernel (the one -# integration unknown that still needs a live B200 check). NeMo images ship -# megatron-bridge; pull requires nvcr.io auth (an NGC key on the host). -MEGATRON_IMAGE = "docker:nvcr.io/nvidia/nemo:25.09" +from dataclasses import dataclass +from pathlib import Path +# The two architectures yeto runs on. Aliases are the names Docker and Go +# use for the same machines. +CPU_ARCHES = ("x86_64", "aarch64") +_ARCH_ALIASES = {"amd64": "x86_64", "x86-64": "x86_64", "arm64": "aarch64"} -def learner_image_for(args, spec: ClusterSpec, learner_id: int | None = None): - """The image for a learner cluster: explicit flag > megatron container > - internal override table > None (provider default + setup-time remediation).""" - explicit = parse_image_spec(getattr(args, "learner_image", None)) - if explicit is not None: - if isinstance(explicit, dict) and learner_id is not None: - for key in (str(learner_id), f"l{learner_id}"): - if key in explicit: - return explicit[key] - return explicit - if getattr(args, "island_backend", "torch") == "megatron": - return MEGATRON_IMAGE - resolver = GPU_IMAGE_OVERRIDES.get((spec.cloud, spec.gpu)) - if resolver is None or not spec.region: - return None - image = resolver(spec.region) - if image: - print(f"[launcher] {spec.gpu} learner: pinning image {image} ({spec.region})") - return image +class UnknownCpuArch(ValueError): + """The CPU architecture is not one yeto runs on.""" -def parse_image_spec(value: str | None): - """--learner-image: a single image id/tag applied everywhere, or - comma-separated region=id pairs -> the region dict sky expects. Numeric - keys are learner ids and are resolved before region keys; this covers - providers where a saved OS volume is not a reusable image.""" - if not value: - return None - if "=" not in value: - return value - images = {} - for pair in value.split(","): - key, _, image = pair.partition("=") - if not key or not image: - raise ValueError( - f"bad --learner-image entry {pair!r}; expected region=image-id " - "or learner-id=image-id" - ) - images[key.strip()] = image.strip() - return images +def canonical_cpu_arch(value: str) -> str: + """``uname -m`` spelling of ``value`` (``amd64`` -> ``x86_64``).""" + arch = str(value).strip().lower() + arch = _ARCH_ALIASES.get(arch, arch) + if arch not in CPU_ARCHES: + raise UnknownCpuArch(f"cpu_arch {value!r} is unknown (known: {list(CPU_ARCHES)})") + return arch -# Exit codes of the image-manifest check (:func:`image_manifest_check`). -MANIFEST_MATCH, MANIFEST_ABSENT, MANIFEST_MISMATCH = 0, 3, 4 -# Stdlib-only: reading the manifest must not import the packages it describes -# (``import peft`` ~27 s, ``import sglang`` ~18 s; coldstart-gpu-review 6.1). -_MANIFEST_CHECK_PY = """\ -import json, os, sys -path, package, key, want, want_path = sys.argv[1:6] -try: - entry = json.load(open(path)).get(package) or {} -except FileNotFoundError: - sys.exit(3) -have = entry.get(key) -if have is None: - sys.exit(3) -if have != want: - print(f"[yeto-setup] image manifest {path}: {package} {key} {have} " - f"!= pinned {want}", file=sys.stderr) - sys.exit(4) -if want_path and os.path.realpath(entry.get("path", "")) != os.path.realpath(want_path): - print(f"[yeto-setup] image manifest {path}: {package} path {entry.get('path')} " - f"!= {want_path}", file=sys.stderr) - sys.exit(4) -""" +@dataclass(frozen=True) +class ImageSpec: + """One yeto engine image. ``repository`` is the last path segment of the + image reference (``REGISTRY/:TAG``).""" -def image_manifest_check(package: str, key: str, want: str, *, fallback: str, - path: str = "", on_pass: str = ":") -> str: - """Shell lines: check ``package`` against the image manifest, without importing it. + name: str + repository: str + cpu_arch: str + device_family: str + roles: tuple[str, ...] - Manifest entry equal to ``want`` (and at ``path`` when given): run - ``on_pass``. Manifest or entry absent: record ``image_manifest_fallback`` - and run ``fallback`` (the old import check). Different value: fail the - setup (before Ray starts), naming the package and both values. - """ + @property + def lock_path(self) -> Path: + """``docker//versions.lock.json`` of a source checkout.""" + return Path(__file__).resolve().parents[2] / "docker" / self.name / "versions.lock.json" - from ..rl import MILES_NEXT_IMAGE_MANIFEST - from ..rl.engine.start_phases import shell_event - argv = " ".join(shlex.quote(a) for a in (MILES_NEXT_IMAGE_MANIFEST, package, key, want, path)) - event = shell_event("image_manifest_fallback", - {"package": package, "manifest": MILES_NEXT_IMAGE_MANIFEST, - "check": "import"}) - return ( - f"rc=0; python3 -c {shlex.quote(_MANIFEST_CHECK_PY)} {argv} || rc=$?\n" - 'case "$rc" in\n' - f"{MANIFEST_MATCH}) echo '[yeto-setup] image manifest: {package} {key} {want}'; {on_pass} ;;\n" - f"{MANIFEST_ABSENT})\n{event}\n{fallback} ;;\n" - f"*) echo \"[yeto-setup] image manifest check of {package} failed (exit $rc); " - "refusing to start\" >&2; exit 1 ;;\n" - "esac" - ) +IMAGES: dict[str, ImageSpec] = {spec.name: spec for spec in ( + # PPU trainer and SGLang sampler (stage 1): Alibaba PPU nodes are x86_64. + ImageSpec("ppu-train", "yeto-ppu-train", "x86_64", "ppu", ("trainer", "sampler")), + # Ascend 910B sampler (stage 2): Kunpeng hosts are aarch64. + ImageSpec("npu-sampler", "yeto-npu-sampler", "aarch64", "npu", ("sampler",)), +)} -def _miles_source_setup(overlay: str | None = None) -> tuple[str, str]: - """Return the (miles_setup, sglang_setup) remote steps of the pinned Miles/SGLang forks. +def _repository(reference: str) -> str: + """``REGISTRY/a/b:tag@sha256:..`` -> ``b``.""" + path = reference.split("@", 1)[0] + last = path.rsplit("/", 1)[-1] + return last.split(":", 1)[0] - ``overlay`` (default None = off): a :mod:`yeto.rl.adapters.miles.overlay` - patch applied to the image's ~/miles after the checkout checks. - """ - from ..rl import ( - MILES_NEXT_COMMIT, - MILES_NEXT_IMAGE_MANIFEST, - MILES_NEXT_IMAGE_SGLANG_ROOT, - MILES_NEXT_REPOSITORY, - MILES_PEFT_VERSION, - SGLANG_NEXT_COMMIT, - SGLANG_NEXT_REPOSITORY, - ) +def catalog_image(reference: str) -> ImageSpec | None: + """The catalog image ``reference`` points at, or None for a site image.""" + repository = _repository(reference) + return next((spec for spec in IMAGES.values() if spec.repository == repository), None) - def checkout( - path: str, repository: str, commit: str, refreshed_flag: str = "" - ) -> str: - repo = shlex.quote(repository) - # ``refreshed_flag``: shell variable set to 1 when the checkout - # had to move to the pin (so the caller knows the image's - # install of it can no longer be trusted as-is). - mark = f"{refreshed_flag}=1\n" if refreshed_flag else "" - return ( - f"if [ ! -d {path}/.git ]; then git clone --no-checkout " - f"{repo} {path}; fi\n" - f"git -C {path} remote set-url origin {repo}\n" - f'test "$(git -C {path} config --get remote.origin.url)" = ' - f"{repo}\n" - # MILES_NEXT_IMAGE already has the pinned fork checked out - # at ~/miles (/root/miles); only fetch when it is not HEAD. - f'if [ "$(git -C {path} rev-parse HEAD 2>/dev/null)" != ' - f"{commit} ]; then\n" - f"git -C {path} fetch --depth 1 origin {commit}\n" - f"git -C {path} checkout --detach {commit}\n" - + mark - + "fi\n" - f'test "$(git -C {path} rev-parse HEAD)" = {commit}\n' - f'test "$(git -C {path} rev-parse --abbrev-ref HEAD)" = HEAD\n' - f'test -z "$(git -C {path} status --porcelain ' - '--untracked-files=all)"\n' - ) - # MILES_NEXT_IMAGE ships the pinned SGLang fork installed editable - # at /sgl-workspace/sglang: ~/sglang (first on the island's - # PYTHONPATH) becomes a link to it. Any other image, or an existing - # ~/sglang directory, gets the clone + install at ~/sglang as before. - image_root = shlex.quote(MILES_NEXT_IMAGE_SGLANG_ROOT) - sglang_checkout_in_image = ( - f'{{ [ ! -e ~/sglang ] || [ "$(readlink ~/sglang)" = {image_root} ]; }} && ' - f'[ "$(git -C {image_root} rev-parse HEAD 2>/dev/null)" = ' - f"{SGLANG_NEXT_COMMIT} ] && " - f'[ "$(git -C {image_root} config --get remote.origin.url)" = ' - f"{shlex.quote(SGLANG_NEXT_REPOSITORY)} ] && " - f'[ -z "$(git -C {image_root} status --porcelain ' - '--untracked-files=all)" ]' - ) - # Fallback for an image without a manifest entry: the old import check. - sglang_import_ok = ( - "python3 -c 'import os, sys, sglang; sys.exit(0 if os.path.realpath(" - f'sglang.__file__).startswith("{MILES_NEXT_IMAGE_SGLANG_ROOT}/python/' - "\") else 1)'" - ) - # MILES_NEXT_IMAGE also ships the pinned Miles fork installed - # editable at /root/miles (= ~/miles) and records it in - # MILES_NEXT_IMAGE_MANIFEST. When ~/miles was already at the pin - # (no fetch above) and the manifest names that commit at that path, - # the image's install is kept: the editable re-install is skipped - # and only a missing/other peft is (re)installed. Any other image, - # a moved checkout, or a manifest mismatch gets the full install. - manifest = shlex.quote(MILES_NEXT_IMAGE_MANIFEST) - miles_in_image = ( - '[ "$MILES_REFRESHED" = 0 ] && ' - f"python3 -c 'import json, os, sys; " - f'm = json.load(open("{MILES_NEXT_IMAGE_MANIFEST}"))["miles"]; ' - f'sys.exit(0 if m["commit"] == "{MILES_NEXT_COMMIT}" and ' - 'os.path.realpath(m["path"]) == ' - "os.path.realpath(os.path.expanduser(\"~/miles\")) else 1)' " - "2>/dev/null" - ) - peft_ok = ( - "python3 -c 'import peft, sys; " - f"sys.exit(0 if peft.__version__ == \"{MILES_PEFT_VERSION}\" else 1)' " - "2>/dev/null" - ) - peft_check = image_manifest_check( - "peft", "version", MILES_PEFT_VERSION, - fallback=f"{peft_ok} || python3 -m pip install -q 'peft=={MILES_PEFT_VERSION}'", - ) - sglang_check = image_manifest_check( - "sglang", "commit", SGLANG_NEXT_COMMIT, path=MILES_NEXT_IMAGE_SGLANG_ROOT, - on_pass="SGLANG_IN_IMAGE=1", - fallback=f"if {sglang_import_ok}; then SGLANG_IN_IMAGE=1; fi", - ) - return ( - "set -e\n" - "MILES_REFRESHED=0\n" - + checkout( - "~/miles", MILES_NEXT_REPOSITORY, MILES_NEXT_COMMIT, "MILES_REFRESHED" - ) - + f"if {miles_in_image}; then\n" - f"echo \"[yeto-setup] image provides miles {MILES_NEXT_COMMIT} " - f'({manifest}); editable install kept"\n' - f"{peft_check}\n" - "else\n" - "python3 -m pip install -q --no-deps -e ~/miles " - f"'peft=={MILES_PEFT_VERSION}'\n" - "fi" - + (f"\n{_miles_overlay().overlay_setup(overlay)}" if overlay else ""), - "SGLANG_IN_IMAGE=0\n" - f"if {sglang_checkout_in_image}; then\n" - f"{sglang_check}\n" - "fi\n" - 'if [ "$SGLANG_IN_IMAGE" = 1 ]; then\n' - f"ln -sfn {image_root} ~/sglang\n" - f"echo '[yeto-setup] image provides sglang {SGLANG_NEXT_COMMIT}'\n" - "else\n" - + checkout("~/sglang", SGLANG_NEXT_REPOSITORY, SGLANG_NEXT_COMMIT) - + "python3 -m pip install -q --no-deps -e ~/sglang/python\n" - "fi", - ) +def image_mismatch(reference: str, *, cpu_arch: str, device_family: str) -> str | None: + """Why ``reference`` cannot run on a node of this kind, or None when it can + (or when the reference is not a catalog image).""" + spec = catalog_image(reference) + if spec is None: + return None + if spec.cpu_arch != cpu_arch: + return (f"image {reference} is the {spec.cpu_arch} image {spec.name}; " + f"the node is {cpu_arch}") + if spec.device_family != device_family: + return (f"image {reference} is the {spec.device_family} image {spec.name}; " + f"the node is device_family {device_family}") + return None diff --git a/yeto/launch/node_inventory.py b/yeto/launch/node_inventory.py index 6a31652ac..7277d930f 100644 --- a/yeto/launch/node_inventory.py +++ b/yeto/launch/node_inventory.py @@ -1,18 +1,28 @@ """Node inventory for machines yeto reaches by SSH (rl-ssh-node-pool D1). One file describes every node the launcher may use, on-premises or rented: -the SSH target, an optional jump host, the device family, the card model and -count, the HBM size per card, the host memory (the trainer offload plan reads +the SSH target, an optional jump host, the device family, the CPU +architecture, the card model and count, the HBM size per card, the host memory (the trainer offload plan reads it), the fabric, shared mounts, the container image and an optional declared price. The core stays neutral: provider differences live in the file. +``cpu_arch`` (``x86_64`` or ``aarch64``) selects the image (S25 N1.6): a node +without ``image`` takes ``images[cpu_arch]`` from the top of the file, and a +node whose image is a yeto catalog image of another architecture or device +family is refused (``yeto.launch.images``). The trainer preflight also +compares ``uname -m`` with ``cpu_arch``. + Example:: + images: + aarch64: REGISTRY/yeto-npu-sampler:TAG + x86_64: REGISTRY/yeto-ppu-train:TAG nodes: - name: ascend00 host: user@172.28.11.10 proxy_jump: user@bastion.example device_family: npu + cpu_arch: aarch64 card: 910B4 cards_per_node: 8 hbm_gib: 64 @@ -20,7 +30,6 @@ fabric: roce interface: eth0 mounts: {/mnt/shared: /mnt/shared} - image: REGISTRY/yeto-ascend:TAG runtime: docker price_per_card_hour: 0.0 """ @@ -33,6 +42,7 @@ from ..gpu_spec import DEVICE_FAMILIES, UnknownCard, canonical_card, device_family_of from ..hw import card_tools +from . import images as image_catalog # Fabrics an inventory may declare. "tcp" is Ethernet without RDMA; "roce" and # "ib" are RDMA over Ethernet and InfiniBand; "hccs" is the Ascend on-node @@ -41,7 +51,7 @@ FABRICS = ("tcp", "roce", "ib", "hccs") RUNTIMES = ("docker", "rootfs") -_REQUIRED = ("name", "host", "device_family", "card", "cards_per_node", "hbm_gib", +_REQUIRED = ("name", "host", "device_family", "cpu_arch", "card", "cards_per_node", "hbm_gib", "host_memory_gib") _KNOWN = _REQUIRED + ("proxy_jump", "fabric", "interface", "mounts", "image", "runtime", "price_per_card_hour", "roles") @@ -58,6 +68,7 @@ class Node: name: str host: str # SSH target, "user@host" or a host alias device_family: str # "cuda", "npu" or "ppu" + cpu_arch: str # "x86_64" or "aarch64" (``uname -m``) card: str # canonical catalog name, e.g. "910B4" cards_per_node: int hbm_gib: float @@ -92,9 +103,21 @@ def check_cards(self, text: str) -> list[card_tools.CardRecord]: count=self.cards_per_node, hbm_gib=self.hbm_gib, node=self.name) + def check_cpu_arch(self, uname_m: str) -> None: + """Compare the node's ``uname -m`` output with ``cpu_arch``.""" + try: + reported = image_catalog.canonical_cpu_arch(uname_m) + except image_catalog.UnknownCpuArch: + raise InventoryError(f"node {self.name}: uname -m reports {uname_m.strip()!r}, " + f"the inventory declares cpu_arch {self.cpu_arch}") from None + if reported != self.cpu_arch: + raise InventoryError(f"node {self.name}: uname -m reports {reported}, the inventory " + f"declares cpu_arch {self.cpu_arch}") + def record(self) -> dict[str, Any]: """Manifest / event-tape row of this node (no SSH user, no mounts).""" - return {"name": self.name, "device_family": self.device_family, "card": self.card, + return {"name": self.name, "device_family": self.device_family, + "cpu_arch": self.cpu_arch, "card": self.card, "cards_per_node": self.cards_per_node, "hbm_gib": self.hbm_gib, "host_memory_gib": self.host_memory_gib, "fabric": self.fabric, "runtime": self.runtime, "behind_jump_host": self.proxy_jump is not None} @@ -112,8 +135,38 @@ def _require_text(value: Any, node: str, key: str) -> str: return value.strip() -def parse_node(raw: Any, *, index: int) -> Node: - """One inventory entry as a :class:`Node`; any gap refuses the launch.""" +def _cpu_arch(value: Any, where: str) -> str: + try: + return image_catalog.canonical_cpu_arch(_require_text(value, where, "cpu_arch")) + except image_catalog.UnknownCpuArch as exc: + raise InventoryError(f"node {where}: {exc}") from None + + +def parse_images(raw: Any) -> dict[str, str]: + """The top-level ``images:`` map (CPU architecture -> image reference).""" + if raw is None: + return {} + if not isinstance(raw, dict): + raise InventoryError(f"images must be a map of cpu_arch to image, got {raw!r}") + out: dict[str, str] = {} + for key, value in raw.items(): + try: + arch = image_catalog.canonical_cpu_arch(key) + except image_catalog.UnknownCpuArch as exc: + raise InventoryError(f"images: {exc}") from None + if arch in out: + raise InventoryError(f"images: cpu_arch {arch} is listed twice") + if not isinstance(value, str) or not value.strip(): + raise InventoryError(f"images: {key} must name an image, got {value!r}") + out[arch] = value.strip() + return out + + +def parse_node(raw: Any, *, index: int, images: Mapping[str, str] | None = None) -> Node: + """One inventory entry as a :class:`Node`; any gap refuses the launch. + + ``images`` is the inventory's ``images:`` map; a node without ``image`` + takes the entry for its ``cpu_arch``.""" if not isinstance(raw, dict): raise InventoryError(f"inventory entry {index} must be a mapping, got {type(raw).__name__}") name = raw.get("name") @@ -130,6 +183,7 @@ def parse_node(raw: Any, *, index: int) -> Node: if family not in DEVICE_FAMILIES: raise InventoryError(f"node {name}: device_family {family!r} is unknown " f"(known: {list(DEVICE_FAMILIES)})") + cpu_arch = _cpu_arch(raw["cpu_arch"], name) try: card = canonical_card(_require_text(raw["card"], name, "card")) except UnknownCard as exc: @@ -149,8 +203,14 @@ def parse_node(raw: Any, *, index: int) -> Node: if runtime not in RUNTIMES: raise InventoryError(f"node {name}: runtime {runtime!r} is unknown " f"(known: {list(RUNTIMES)})") - if runtime == "docker" and raw.get("image") is None: - raise InventoryError(f"node {name}: runtime docker needs an image") + image = (_require_text(raw["image"], name, "image") if raw.get("image") is not None + else (images or {}).get(cpu_arch)) + if runtime == "docker" and image is None: + raise InventoryError(f"node {name}: runtime docker needs an image (set image, or " + f"images.{cpu_arch} at the top of the inventory)") + if image is not None and (why := image_catalog.image_mismatch( + image, cpu_arch=cpu_arch, device_family=family)): + raise InventoryError(f"node {name}: {why}") mounts = raw.get("mounts") or {} if not isinstance(mounts, dict) or not all( isinstance(k, str) and isinstance(v, str) for k, v in mounts.items()): @@ -168,6 +228,7 @@ def parse_node(raw: Any, *, index: int) -> Node: name=name, host=_require_text(raw["host"], name, "host"), device_family=family, + cpu_arch=cpu_arch, card=card, cards_per_node=cards, hbm_gib=_require_positive(raw["hbm_gib"], name, "hbm_gib"), @@ -178,8 +239,7 @@ def parse_node(raw: Any, *, index: int) -> Node: interface=(_require_text(raw["interface"], name, "interface") if raw.get("interface") is not None else None), mounts=dict(mounts), - image=(_require_text(raw["image"], name, "image") - if raw.get("image") is not None else None), + image=image, runtime=runtime, price_per_card_hour=(float(price) if price is not None else None), roles=tuple(roles), @@ -188,16 +248,22 @@ def parse_node(raw: Any, *, index: int) -> Node: def parse_inventory(document: Any) -> list[Node]: """The ``nodes:`` list of a loaded inventory document.""" + images: dict[str, str] = {} if isinstance(document, list): entries = document elif isinstance(document, dict): + unknown = sorted(set(document) - {"nodes", "images"}) + if unknown: + raise InventoryError(f"unknown top-level inventory field(s) {unknown} " + f"(known: ['images', 'nodes'])") entries = document.get("nodes") + images = parse_images(document.get("images")) else: raise InventoryError(f"the inventory must be a mapping with a nodes list, got " f"{type(document).__name__}") if not isinstance(entries, list) or not entries: raise InventoryError("the inventory must contain at least one node under nodes:") - nodes = [parse_node(entry, index=index) for index, entry in enumerate(entries)] + nodes = [parse_node(entry, index=index, images=images) for index, entry in enumerate(entries)] for key, label in (("name", "name"), ("host", "SSH target")): seen: dict[str, str] = {} for node in nodes: diff --git a/yeto/launch/ssh_trainer_island.py b/yeto/launch/ssh_trainer_island.py index 72ebd423d..1fc7c3179 100644 --- a/yeto/launch/ssh_trainer_island.py +++ b/yeto/launch/ssh_trainer_island.py @@ -9,7 +9,8 @@ One incarnation = one ``torchrun`` per node, started over SSH: -1. before the first start, every node runs its device family's card tool and +1. before the first start, every node's ``uname -m`` must equal its + ``cpu_arch`` (``Node.check_cpu_arch``) and its device family's card tool must match the inventory (``Node.check_cards``); the memory estimate of the plan (:mod:`yeto.rl.fullparam.memory`) is checked when the spec is built (:func:`ssh_trainer_spec`), so a plan that does not fit never reaches a node; @@ -52,7 +53,7 @@ from yeto.rl.engine.trainer_lifecycle import CutRef from yeto.rl.engine.trainer_remote import island_env -from .node_inventory import Node +from .node_inventory import InventoryError, Node from .topology import island_comm_env from .trainer_island import ( JOBS_FILE, @@ -230,8 +231,16 @@ def __init__(self, spec: SshTrainerSpec, *, name: str, jobs_path: Path, ops: Ssh # -- preflight ---------------------------------------------------------------- def preflight(self) -> None: - """Every node's cards must match the inventory (raises before any start).""" + """Every node's CPU architecture and cards must match the inventory + (raises before any start).""" for node in self.spec.nodes: + done = self.ops.run(node, "uname -m", timeout=60.0) + if done.returncode != 0: + raise SshIslandError(f"node {node.name}: uname -m failed (rc {done.returncode})") + try: + node.check_cpu_arch(done.stdout) + except InventoryError as exc: + raise SshIslandError(str(exc)) from None for argv in node.preflight_commands(): done = self.ops.run(node, shlex.join(argv), timeout=120.0) if done.returncode == 0: diff --git a/yeto/rl/fullparam/memory.py b/yeto/rl/fullparam/memory.py index 675157f71..c1a65778f 100644 --- a/yeto/rl/fullparam/memory.py +++ b/yeto/rl/fullparam/memory.py @@ -1,4 +1,4 @@ -"""Memory estimate of one full-parameter FSDP2 training step (S25 N7.1, design D2). +"""Memory estimate of one full-parameter training step (S25 N7.1 FSDP2, N7.1b Megatron). Every number here is derived from formulas, not measured. The launcher uses the estimate to refuse a plan before it reaches any node: a plan whose @@ -34,6 +34,30 @@ source read, not run at this size). A sparse DSA kernel removes them; none exists for training on PPU or NPU yet (N7.0). +Megatron (``framework="megatron"``, N7.1b; the PPU stage-1 backend, A-TRN §4). +``W`` cards, tensor parallel ``TP`` (sequence parallel assumed on), pipeline +``PP``, expert parallel ``EP`` (expert tensor parallel 1), context parallel +``CP``. Figures are for the worst pipeline stage (``ceil(layers / PP)`` +layers; every parameter is spread over the layers, embedding and output-layer +skew is not modelled): + +* local parameters: dense ``(P - E) x f / TP`` plus experts ``E x f / EP`` + (``E`` = routed-expert parameters, ``f`` = the stage's layer fraction); +* BF16 parameters (2 B) and the main gradients (``grad_bytes``) on every + local parameter; +* optimizer: ``param_bytes`` 4 adds a 4 B FP32 master copy to the + ``state_bytes`` Adam state; with the distributed optimizer this is divided + by the data-parallel size of its group (dense ``W / (TP x PP)``, CP + included; experts ``W / (EP x PP)``), without it it is per local parameter; +* activations with full recomputation: stage 0 of 1F1B holds ``PP`` micro + batches, so ``ceil(layers / PP) x PP x mbs x T / (TP x CP) x streams x hidden + x 2 B``; mHC keeps ``hc_mult`` (4) residual streams, so this term is 4 times + that of a one-stream model; the last stage adds the logits split by ``TP``; +* Miles' DSA layer does not support CP (``dsa.py`` NotImplementedError), so a + DSA model with ``CP > 1`` is refused. + +Every output carries ``basis: "推算"`` (derived, not measured). + Not modelled: allocator fragmentation, collective buffers, the runtime's own reserve (covered by ``device_usable_fraction``), CPU optimizer step time. @@ -53,18 +77,20 @@ import sys from collections import Counter from collections.abc import Iterable, Mapping, Sequence -from dataclasses import asdict, dataclass, field, replace +from dataclasses import MISSING, asdict, dataclass, field, replace from typing import Any GIB = 1024**3 OFFLOAD_MODES = ("none", "optimizer", "full") DSA_PATHS = ("none", "hf_dense", "sparse_kernel") +FRAMEWORKS = ("fsdp2", "megatron") PLAN_STATUS = { "none": "FSDP2 standard path", "optimizer": "needs a host optimizer step; not implemented in yeto or verl", "full": "FSDP2 CPUOffloadPolicy (optimizer step on the host CPU)", } LABEL = "derived from formulas, not measured" +BASIS = "推算" # the dashboard / report word for LABEL class MemoryPlanError(ValueError): @@ -78,7 +104,8 @@ class ModelShape: ``params``: trained parameters; ``largest_layer_params``: the largest decoder layer (the unit FSDP2 gathers); ``head_counts``: every head count Ulysses must divide; ``dsa_layers`` / ``index_heads`` / ``index_kpool``: - the DSA indexer (0 layers: no DSA).""" + the DSA indexer (0 layers: no DSA); ``expert_params`` / ``num_experts``: + the routed experts that Megatron splits by expert parallel (0: dense).""" params: int largest_layer_params: int @@ -90,6 +117,8 @@ class ModelShape: dsa_layers: int = 0 index_heads: int = 0 index_kpool: int = 1 + expert_params: int = 0 + num_experts: int = 0 source: str = "manual" def __post_init__(self) -> None: @@ -99,12 +128,18 @@ def __post_init__(self) -> None: raise MemoryPlanError(f"model shape: {name} must be positive, got {getattr(self, name)}") if self.largest_layer_params > self.params: raise MemoryPlanError("model shape: largest_layer_params exceeds params") + if not 0 <= self.expert_params <= self.params or self.num_experts < 0: + raise MemoryPlanError("model shape: expert_params must be in [0, params], num_experts >= 0") + if bool(self.expert_params) != bool(self.num_experts): + raise MemoryPlanError("model shape: expert_params and num_experts go together") @dataclass(frozen=True) class TrainPlan: """One trainer island and how it trains. ``fsdp_size`` 0 means all cards in - one shard group; a smaller divisor replicates groups (HSDP).""" + one shard group; a smaller divisor replicates groups (HSDP). ``tp`` / + ``pp`` / ``ep`` / ``cp`` / ``distributed_optimizer`` apply to + ``framework="megatron"`` only.""" cards: int cards_per_node: int @@ -126,14 +161,57 @@ class TrainPlan: device_usable_fraction: float = 0.92 host_usable_fraction: float = 0.85 host_reserve_gib: float = 64.0 # per node: OS, data loader, the requant sidecar + framework: str = "fsdp2" + tp: int = 1 + pp: int = 1 + ep: int = 1 + cp: int = 1 + distributed_optimizer: bool = True @property def shard_group(self) -> int: return self.fsdp_size or self.cards +def _check_megatron_form(plan: TrainPlan, shape: ModelShape) -> None: + for name in ("tp", "pp", "ep", "cp"): + if int(getattr(plan, name)) <= 0: + raise MemoryPlanError(f"{name} must be positive, got {getattr(plan, name)}") + if plan.offload != "none": + raise MemoryPlanError("megatron: only offload none is modelled (Megatron optimizer CPU " + "offload is not)") + if plan.ulysses != 1 or plan.fsdp_size: + raise MemoryPlanError("megatron: use cp instead of ulysses, and no fsdp_size") + W = plan.cards + if W % (plan.tp * plan.pp * plan.cp): + raise MemoryPlanError(f"megatron: tp x pp x cp = {plan.tp * plan.pp * plan.cp} does not divide " + f"cards {W}") + if W % (plan.ep * plan.pp): + raise MemoryPlanError(f"megatron: ep x pp = {plan.ep * plan.pp} does not divide cards {W}") + if plan.pp > shape.num_layers: + raise MemoryPlanError(f"megatron: pp {plan.pp} exceeds the {shape.num_layers} layers") + if plan.ep > 1 and not shape.num_experts: + raise MemoryPlanError("megatron: ep > 1 needs a model with routed experts") + if shape.num_experts and shape.num_experts % plan.ep: + raise MemoryPlanError(f"megatron: ep {plan.ep} does not divide {shape.num_experts} experts") + if plan.seq_len % (plan.tp * plan.cp): + raise MemoryPlanError(f"megatron: tp x cp = {plan.tp * plan.cp} does not divide seq_len " + f"{plan.seq_len} (sequence parallel)") + for heads in shape.head_counts: + if heads % plan.tp: + raise MemoryPlanError(f"megatron: tp {plan.tp} does not divide a head count {heads}") + if shape.dsa_layers and plan.cp > 1: + raise MemoryPlanError("megatron: the Miles DSA layer does not support cp > 1") + + def check_plan_form(plan: TrainPlan, shape: ModelShape) -> None: """Refuse a plan that is not well formed (raises :class:`MemoryPlanError`).""" + if plan.framework not in FRAMEWORKS: + raise MemoryPlanError(f"framework must be one of {FRAMEWORKS}, got {plan.framework!r}") + if plan.framework == "megatron": + _check_megatron_form(plan, shape) + elif (plan.tp, plan.pp, plan.ep, plan.cp) != (1, 1, 1, 1): + raise MemoryPlanError("tp / pp / ep / cp apply to framework megatron only") if plan.offload not in OFFLOAD_MODES: raise MemoryPlanError(f"offload must be one of {OFFLOAD_MODES}, got {plan.offload!r}") if plan.dsa_path not in DSA_PATHS: @@ -204,7 +282,10 @@ def fits(self) -> bool: def summary(self) -> dict[str, Any]: return { "label": self.label, - "plan_status": PLAN_STATUS[self.plan.offload], + "basis": BASIS, + "framework": self.plan.framework, + "plan_status": (PLAN_STATUS[self.plan.offload] if self.plan.framework == "fsdp2" + else "Megatron (no offload)"), "device_gib_per_card": round(self.total("device") / GIB, 1), "device_usable_gib": round(self.device_limit / GIB, 1), "host_gib_per_node": round(self.total("host") / GIB, 1), @@ -222,6 +303,8 @@ def to_dict(self) -> dict[str, Any]: def estimate(shape: ModelShape, plan: TrainPlan) -> Estimate: """The per-card device and per-node host memory of one training step.""" check_plan_form(plan, shape) + if plan.framework == "megatron": + return _estimate_megatron(shape, plan) est = Estimate(shape, plan) P, N, per_node = shape.params, plan.shard_group, plan.cards_per_node sharded = P / N @@ -267,6 +350,50 @@ def add(name: str, value: float, formula: str, where: str = "device") -> None: return est +def _estimate_megatron(shape: ModelShape, plan: TrainPlan) -> Estimate: + """Megatron variant of :func:`estimate` (module doc, N7.1b).""" + est = Estimate(shape, plan) + + def add(name: str, value: float, formula: str) -> None: + est.items[name] = Item(float(value), formula, "device") + + W, tp, pp, ep, cp = plan.cards, plan.tp, plan.pp, plan.ep, plan.cp + stage_layers = -(-shape.num_layers // pp) + f = stage_layers / shape.num_layers + dense = (shape.params - shape.expert_params) * f / tp + experts = shape.expert_params * f / ep + local = dense + experts + where = f"stage of {stage_layers}/{shape.num_layers} layers" + add("params_local_bf16", 2 * local, f"2 B x (dense/TP({tp}) + experts/EP({ep})), {where}") + gb = plan.grad_bytes + add("main_grads_local", gb * local, f"{gb} B x local params") + per_param = (4 if plan.param_bytes == 4 else 0) + plan.state_bytes + if plan.distributed_optimizer: + dp_dense, dp_expert = W // (tp * pp), W // (ep * pp) + add("optimizer_shard", per_param * (dense / dp_dense + experts / dp_expert), + f"{per_param} B x (dense local / DP({dp_dense}) + expert local / expert DP({dp_expert})), " + f"distributed optimizer") + else: + add("optimizer_local", per_param * local, f"{per_param} B x local params (no distributed optimizer)") + tokens = plan.seq_len * plan.micro_batch_seqs / (tp * cp) + add("layer_inputs", stage_layers * pp * tokens * shape.residual_streams * shape.hidden_size * 2, + f"stage layers({stage_layers}) x in-flight micro batches PP({pp}) x T*mbs/(TP*CP)({tokens:.0f}) " + f"x streams({shape.residual_streams}, mHC) x hidden({shape.hidden_size}) x 2 B") + if not plan.chunked_loss: + lb = 4 if plan.logits_fp32 else 2 + logit_tokens = plan.seq_len * plan.micro_batch_seqs / cp + add("logits_and_grad", logit_tokens * shape.vocab_size / tp * lb * 2, + f"T*mbs/CP x vocab({shape.vocab_size})/TP x {lb} B x 2 (last stage)") + add("working_set_margin", plan.activation_margin_gib * GIB, + "configured: recomputed layer working set (KDA, MoE all-to-all, attention); not measured") + if shape.dsa_layers and plan.dsa_path == "hf_dense": + T, b = plan.seq_len, plan.micro_batch_seqs + dense_terms = b * T * T * 5 + b * T * shape.index_heads * -(-T // shape.index_kpool) * 4 + add("dsa_dense_terms", dense_terms, + "mask mbs*T*T*5 B + indexer scores mbs*T*heads*T/kpool*4 B (dense DSA fallback)") + return est + + def check(shape: ModelShape, plan: TrainPlan) -> Estimate: """:func:`estimate`, raising :class:`MemoryPlanError` when the plan does not fit.""" est = estimate(shape, plan) @@ -300,7 +427,7 @@ def shape_from_hf_config(config: Mapping[str, Any], *, include_vision: bool = Fa except ValueError: # a multimodal wrapper (e.g. glm5_next): count its language model model = transformers.AutoModelForImageTextToText.from_config(hf) tc = _text_config(hf) - total, per_layer = 0, Counter() + total, experts, per_layer = 0, 0, Counter() for name, p in model.named_parameters(): parts = name.split(".") is_vision = "language_model" not in parts and "lm_head" not in parts and any( @@ -308,6 +435,8 @@ def shape_from_hf_config(config: Mapping[str, Any], *, include_vision: bool = Fa if is_vision and not include_vision: continue total += p.numel() + if "experts" in parts and "shared_experts" not in parts: + experts += p.numel() if "layers" in parts and not is_vision: per_layer[int(parts[parts.index("layers") + 1])] += p.numel() layer_types = list(getattr(tc, "layer_types", None) or []) @@ -323,6 +452,9 @@ def shape_from_hf_config(config: Mapping[str, Any], *, include_vision: bool = Fa dsa_layers=layer_types.count("deepseek_sparse_attention"), index_heads=int(getattr(tc, "index_n_heads", 0) or 0), index_kpool=int(getattr(tc, "index_kpool", 1) or 1), + expert_params=experts, + num_experts=int(getattr(tc, "n_routed_experts", None) or getattr(tc, "num_experts", None) or 0) + if experts else 0, source=f"{type(model).__name__} on meta (transformers {transformers.__version__})", ) @@ -361,8 +493,8 @@ def main(argv: list[str] | None = None) -> int: flag = "--" + name.replace("_", "-") if f.type == "bool": ap.add_argument(flag, type=lambda s: s.lower() in ("1", "true", "yes"), default=f.default) - elif name in ("cards", "cards_per_node"): - ap.add_argument(flag, type=int, required=True) + elif f.default is MISSING: # cards, cards_per_node, hbm_gib, host_memory_gib + ap.add_argument(flag, type=types[f.type], required=True) elif name == "seq_len": ap.add_argument(flag, type=int, default=None) else: From 9aa6514f909db9b10cf9578ba3a5bae501eba51c Mon Sep 17 00:00:00 2001 From: michael <78113554+chowhsu@users.noreply.github.com> Date: Mon, 19 Oct 2026 12:23:18 +0000 Subject: [PATCH 19/34] Rewrap the node inventory module docstring Co-Authored-By: Claude Opus 5.5 --- yeto/launch/node_inventory.py | 7 ++++--- 1 file changed, 4 insertions(+), 3 deletions(-) diff --git a/yeto/launch/node_inventory.py b/yeto/launch/node_inventory.py index 7277d930f..3c7d8bad8 100644 --- a/yeto/launch/node_inventory.py +++ b/yeto/launch/node_inventory.py @@ -2,9 +2,10 @@ One file describes every node the launcher may use, on-premises or rented: the SSH target, an optional jump host, the device family, the CPU -architecture, the card model and count, the HBM size per card, the host memory (the trainer offload plan reads -it), the fabric, shared mounts, the container image and an optional declared -price. The core stays neutral: provider differences live in the file. +architecture, the card model and count, the HBM size per card, the host +memory (the trainer offload plan reads it), the fabric, shared mounts, the +container image and an optional declared price. The core stays neutral: +provider differences live in the file. ``cpu_arch`` (``x86_64`` or ``aarch64``) selects the image (S25 N1.6): a node without ``image`` takes ``images[cpu_arch]`` from the top of the file, and a From 302ef1e1b18346de0bd6631d2074a9150a06ac56 Mon Sep 17 00:00:00 2001 From: michael <78113554+chowhsu@users.noreply.github.com> Date: Mon, 19 Oct 2026 12:24:42 +0000 Subject: [PATCH 20/34] Move the engine-image catalog to yeto/launch/engine_images.py The previous commit overwrote yeto/launch/images.py (learner image selection and the Miles source setup), which broke every import of yeto.launcher. That file is restored unchanged and the new catalog lives in its own module. Co-Authored-By: Claude Opus 5.5 --- docker/README.md | 4 +- tests/test_engine_images.py | 2 +- yeto/launch/engine_images.py | 95 ++++++++++ yeto/launch/images.py | 336 ++++++++++++++++++++++++++-------- yeto/launch/node_inventory.py | 4 +- 5 files changed, 364 insertions(+), 77 deletions(-) create mode 100644 yeto/launch/engine_images.py diff --git a/docker/README.md b/docker/README.md index b6f98407c..936a24eb0 100644 --- a/docker/README.md +++ b/docker/README.md @@ -7,8 +7,8 @@ | `ppu-train/` | `yeto-ppu-train` | x86_64 | ppu | trainer + sampler | written, never built (S25 N0.2) | | `npu-sampler/` | `yeto-npu-sampler` | aarch64 | npu | sampler | written, never built (S25 N0.2) | -`yeto/launch/images.py` is the catalog the launcher reads. A node of the SSH -inventory declares `cpu_arch`; a node with no `image` takes the inventory's +`yeto/launch/engine_images.py` is the catalog the launcher reads. A node of +the SSH inventory declares `cpu_arch`; a node with no `image` takes the inventory's `images[cpu_arch]`, and a node whose image is a catalog image of another architecture or device family is refused before any SSH (N1.6). diff --git a/tests/test_engine_images.py b/tests/test_engine_images.py index 52c9d0d79..eee334860 100644 --- a/tests/test_engine_images.py +++ b/tests/test_engine_images.py @@ -9,7 +9,7 @@ import pytest -from yeto.launch import images +from yeto.launch import engine_images as images from yeto.rl.adapters.miles import pins as miles_pins from yeto.rl.adapters.verl import pins as verl_pins diff --git a/yeto/launch/engine_images.py b/yeto/launch/engine_images.py new file mode 100644 index 000000000..d98b9340c --- /dev/null +++ b/yeto/launch/engine_images.py @@ -0,0 +1,95 @@ +"""Engine images by CPU architecture (S25 N0.2, N1.6). + +yeto defines one image per (device family, CPU architecture) pair. Each image +has a Dockerfile and a version lock under ``docker//``; the lock is the +single source of the versions the image asserts at start +(``docker/image-assert/assert_image_versions.py``). + +A node of the SSH node pool declares its ``cpu_arch``. The launcher then uses +this catalog in two ways: + +* a node that names no image gets the inventory's image for its architecture; +* a node whose image is a catalog image of another architecture or another + device family is refused before any SSH (an aarch64 node cannot run an + x86_64 image, and a PPU image has no Ascend runtime). + +An image reference that is not in the catalog (a site image) is accepted as +it is: yeto cannot tell its architecture from the name, and the node +preflight (``uname -m``) still checks the node itself. +""" + +from __future__ import annotations + +from dataclasses import dataclass +from pathlib import Path + +# The two architectures yeto runs on. Aliases are the names Docker and Go +# use for the same machines. +CPU_ARCHES = ("x86_64", "aarch64") +_ARCH_ALIASES = {"amd64": "x86_64", "x86-64": "x86_64", "arm64": "aarch64"} + + +class UnknownCpuArch(ValueError): + """The CPU architecture is not one yeto runs on.""" + + +def canonical_cpu_arch(value: str) -> str: + """``uname -m`` spelling of ``value`` (``amd64`` -> ``x86_64``).""" + arch = str(value).strip().lower() + arch = _ARCH_ALIASES.get(arch, arch) + if arch not in CPU_ARCHES: + raise UnknownCpuArch(f"cpu_arch {value!r} is unknown (known: {list(CPU_ARCHES)})") + return arch + + +@dataclass(frozen=True) +class ImageSpec: + """One yeto engine image. ``repository`` is the last path segment of the + image reference (``REGISTRY/:TAG``).""" + + name: str + repository: str + cpu_arch: str + device_family: str + roles: tuple[str, ...] + + @property + def lock_path(self) -> Path: + """``docker//versions.lock.json`` of a source checkout.""" + return Path(__file__).resolve().parents[2] / "docker" / self.name / "versions.lock.json" + + +IMAGES: dict[str, ImageSpec] = {spec.name: spec for spec in ( + # PPU trainer and SGLang sampler (stage 1): Alibaba PPU nodes are x86_64. + ImageSpec("ppu-train", "yeto-ppu-train", "x86_64", "ppu", ("trainer", "sampler")), + # Ascend 910B sampler (stage 2): Kunpeng hosts are aarch64. + ImageSpec("npu-sampler", "yeto-npu-sampler", "aarch64", "npu", ("sampler",)), +)} + + +def _repository(reference: str) -> str: + """``REGISTRY/a/b:tag@sha256:..`` -> ``b``.""" + path = reference.split("@", 1)[0] + last = path.rsplit("/", 1)[-1] + return last.split(":", 1)[0] + + +def catalog_image(reference: str) -> ImageSpec | None: + """The catalog image ``reference`` points at, or None for a site image.""" + repository = _repository(reference) + return next((spec for spec in IMAGES.values() if spec.repository == repository), None) + + +def image_mismatch(reference: str, *, cpu_arch: str, device_family: str) -> str | None: + """Why ``reference`` cannot run on a node of this kind, or None when it can + (or when the reference is not a catalog image).""" + spec = catalog_image(reference) + if spec is None: + return None + if spec.cpu_arch != cpu_arch: + return (f"image {reference} is the {spec.cpu_arch} image {spec.name}; " + f"the node is {cpu_arch}") + if spec.device_family != device_family: + return (f"image {reference} is the {spec.device_family} image {spec.name}; " + f"the node is device_family {device_family}") + return None diff --git a/yeto/launch/images.py b/yeto/launch/images.py index d98b9340c..9d9ebdff8 100644 --- a/yeto/launch/images.py +++ b/yeto/launch/images.py @@ -1,95 +1,287 @@ -"""Engine images by CPU architecture (S25 N0.2, N1.6). +"""Learner image selection and the Miles source setup.""" -yeto defines one image per (device family, CPU architecture) pair. Each image -has a Dockerfile and a version lock under ``docker//``; the lock is the -single source of the versions the image asserts at start -(``docker/image-assert/assert_image_versions.py``). +from __future__ import annotations -A node of the SSH node pool declares its ``cpu_arch``. The launcher then uses -this catalog in two ways: +import shlex +import sys -* a node that names no image gets the inventory's image for its architecture; -* a node whose image is a catalog image of another architecture or another - device family is refused before any SSH (an aarch64 node cannot run an - x86_64 image, and a PPU image has no Ascend runtime). +from ..gpu_spec import ClusterSpec +from .rl_backend import _miles_overlay -An image reference that is not in the catalog (a site image) is accepted as -it is: yeto cannot tell its architecture from the name, and the node -preflight (``uname -m``) still checks the node itself. -""" -from __future__ import annotations +# AWS keeps this SSM parameter pointing at the CURRENT Deep Learning Base +# OSS-NVIDIA-driver AMI per region (open kernel modules, Blackwell-ready), +# so resolving at launch time needs no region x AMI table of our own. +_DLAMI_SSM_PARAM = ( + "/aws/service/deeplearning/ami/x86_64/" + "base-oss-nvidia-driver-gpu-ubuntu-22.04/latest/ami-id" +) + + +def resolve_blackwell_image(region: str) -> str | None: + """Current Blackwell-capable DLAMI for a region, or None (callers then + rely on the setup-time driver remediation instead).""" + try: + import boto3 + + ssm = boto3.client("ssm", region_name=region) + return ssm.get_parameter(Name=_DLAMI_SSM_PARAM)["Parameter"]["Value"] + except Exception as e: # noqa: BLE001 - degrade to driver remediation + print( + f"[launcher] could not resolve a Blackwell DLAMI for {region} ({e}); " + "relying on setup-time driver install", + file=sys.stderr, + ) + return None -from dataclasses import dataclass -from pathlib import Path -# The two architectures yeto runs on. Aliases are the names Docker and Go -# use for the same machines. -CPU_ARCHES = ("x86_64", "aarch64") -_ARCH_ALIASES = {"amd64": "x86_64", "x86-64": "x86_64", "arm64": "aarch64"} +# Internal image-override table: (cloud, GPU) pairs whose provider-default +# image is known stale/broken, mapped to a `region -> image id` resolver. +# First entry earned in production: sky's pinned AMI ships driver 535, +# which never binds to SM100 silicon. Second class earned the same way: +# NVSwitch instances (p4d/p4de A100, p5 H100, p5e H200) need +# nvidia-fabricmanager running before CUDA will initialize at all +# (cudaGetDeviceCount -> Error 802), and sky's pinned AMI does not ship it; +# the DL Base GPU AMI does, preinstalled and enabled. Extend here as new +# GPU generations outpace provider image pins; an explicit --learner-image +# always wins, and a resolver returning None degrades to the setup-time +# driver install. +GPU_IMAGE_OVERRIDES: dict[tuple[str, str], object] = { + ("aws", "B200"): resolve_blackwell_image, + # NVSwitch (fabric manager required); same current DL Base GPU AMI. + ("aws", "A100"): resolve_blackwell_image, + ("aws", "A100-80GB"): resolve_blackwell_image, + ("aws", "H100"): resolve_blackwell_image, + ("aws", "H200"): resolve_blackwell_image, +} -class UnknownCpuArch(ValueError): - """The CPU architecture is not one yeto runs on.""" +# The Megatron backend runs INSIDE an NGC container that ships the whole +# stack prebuilt (torch + Transformer Engine + megatron-core + megatron-bridge) +# — the pip-on-DLAMI install proved intractable (two shakedowns; see +# docs/MEGATRON.md). sky runs the task in this container via `image_id: +# docker:...`; the host still supplies the GPU kernel driver, so B200 needs +# sky's default AWS docker host AMI to carry driver >=570 open-kernel (the one +# integration unknown that still needs a live B200 check). NeMo images ship +# megatron-bridge; pull requires nvcr.io auth (an NGC key on the host). +MEGATRON_IMAGE = "docker:nvcr.io/nvidia/nemo:25.09" -def canonical_cpu_arch(value: str) -> str: - """``uname -m`` spelling of ``value`` (``amd64`` -> ``x86_64``).""" - arch = str(value).strip().lower() - arch = _ARCH_ALIASES.get(arch, arch) - if arch not in CPU_ARCHES: - raise UnknownCpuArch(f"cpu_arch {value!r} is unknown (known: {list(CPU_ARCHES)})") - return arch +def learner_image_for(args, spec: ClusterSpec, learner_id: int | None = None): + """The image for a learner cluster: explicit flag > megatron container > + internal override table > None (provider default + setup-time remediation).""" + explicit = parse_image_spec(getattr(args, "learner_image", None)) + if explicit is not None: + if isinstance(explicit, dict) and learner_id is not None: + for key in (str(learner_id), f"l{learner_id}"): + if key in explicit: + return explicit[key] + return explicit + if getattr(args, "island_backend", "torch") == "megatron": + return MEGATRON_IMAGE + resolver = GPU_IMAGE_OVERRIDES.get((spec.cloud, spec.gpu)) + if resolver is None or not spec.region: + return None + image = resolver(spec.region) + if image: + print(f"[launcher] {spec.gpu} learner: pinning image {image} ({spec.region})") + return image -@dataclass(frozen=True) -class ImageSpec: - """One yeto engine image. ``repository`` is the last path segment of the - image reference (``REGISTRY/:TAG``).""" +def parse_image_spec(value: str | None): + """--learner-image: a single image id/tag applied everywhere, or + comma-separated region=id pairs -> the region dict sky expects. Numeric + keys are learner ids and are resolved before region keys; this covers + providers where a saved OS volume is not a reusable image.""" + if not value: + return None + if "=" not in value: + return value + images = {} + for pair in value.split(","): + key, _, image = pair.partition("=") + if not key or not image: + raise ValueError( + f"bad --learner-image entry {pair!r}; expected region=image-id " + "or learner-id=image-id" + ) + images[key.strip()] = image.strip() + return images - name: str - repository: str - cpu_arch: str - device_family: str - roles: tuple[str, ...] - @property - def lock_path(self) -> Path: - """``docker//versions.lock.json`` of a source checkout.""" - return Path(__file__).resolve().parents[2] / "docker" / self.name / "versions.lock.json" +# Exit codes of the image-manifest check (:func:`image_manifest_check`). +MANIFEST_MATCH, MANIFEST_ABSENT, MANIFEST_MISMATCH = 0, 3, 4 +# Stdlib-only: reading the manifest must not import the packages it describes +# (``import peft`` ~27 s, ``import sglang`` ~18 s; coldstart-gpu-review 6.1). +_MANIFEST_CHECK_PY = """\ +import json, os, sys +path, package, key, want, want_path = sys.argv[1:6] +try: + entry = json.load(open(path)).get(package) or {} +except FileNotFoundError: + sys.exit(3) +have = entry.get(key) +if have is None: + sys.exit(3) +if have != want: + print(f"[yeto-setup] image manifest {path}: {package} {key} {have} " + f"!= pinned {want}", file=sys.stderr) + sys.exit(4) +if want_path and os.path.realpath(entry.get("path", "")) != os.path.realpath(want_path): + print(f"[yeto-setup] image manifest {path}: {package} path {entry.get('path')} " + f"!= {want_path}", file=sys.stderr) + sys.exit(4) +""" -IMAGES: dict[str, ImageSpec] = {spec.name: spec for spec in ( - # PPU trainer and SGLang sampler (stage 1): Alibaba PPU nodes are x86_64. - ImageSpec("ppu-train", "yeto-ppu-train", "x86_64", "ppu", ("trainer", "sampler")), - # Ascend 910B sampler (stage 2): Kunpeng hosts are aarch64. - ImageSpec("npu-sampler", "yeto-npu-sampler", "aarch64", "npu", ("sampler",)), -)} +def image_manifest_check(package: str, key: str, want: str, *, fallback: str, + path: str = "", on_pass: str = ":") -> str: + """Shell lines: check ``package`` against the image manifest, without importing it. + Manifest entry equal to ``want`` (and at ``path`` when given): run + ``on_pass``. Manifest or entry absent: record ``image_manifest_fallback`` + and run ``fallback`` (the old import check). Different value: fail the + setup (before Ray starts), naming the package and both values. + """ -def _repository(reference: str) -> str: - """``REGISTRY/a/b:tag@sha256:..`` -> ``b``.""" - path = reference.split("@", 1)[0] - last = path.rsplit("/", 1)[-1] - return last.split(":", 1)[0] + from ..rl import MILES_NEXT_IMAGE_MANIFEST + from ..rl.engine.start_phases import shell_event + argv = " ".join(shlex.quote(a) for a in (MILES_NEXT_IMAGE_MANIFEST, package, key, want, path)) + event = shell_event("image_manifest_fallback", + {"package": package, "manifest": MILES_NEXT_IMAGE_MANIFEST, + "check": "import"}) + return ( + f"rc=0; python3 -c {shlex.quote(_MANIFEST_CHECK_PY)} {argv} || rc=$?\n" + 'case "$rc" in\n' + f"{MANIFEST_MATCH}) echo '[yeto-setup] image manifest: {package} {key} {want}'; {on_pass} ;;\n" + f"{MANIFEST_ABSENT})\n{event}\n{fallback} ;;\n" + f"*) echo \"[yeto-setup] image manifest check of {package} failed (exit $rc); " + "refusing to start\" >&2; exit 1 ;;\n" + "esac" + ) -def catalog_image(reference: str) -> ImageSpec | None: - """The catalog image ``reference`` points at, or None for a site image.""" - repository = _repository(reference) - return next((spec for spec in IMAGES.values() if spec.repository == repository), None) +def _miles_source_setup(overlay: str | None = None) -> tuple[str, str]: + """Return the (miles_setup, sglang_setup) remote steps of the pinned Miles/SGLang forks. -def image_mismatch(reference: str, *, cpu_arch: str, device_family: str) -> str | None: - """Why ``reference`` cannot run on a node of this kind, or None when it can - (or when the reference is not a catalog image).""" - spec = catalog_image(reference) - if spec is None: - return None - if spec.cpu_arch != cpu_arch: - return (f"image {reference} is the {spec.cpu_arch} image {spec.name}; " - f"the node is {cpu_arch}") - if spec.device_family != device_family: - return (f"image {reference} is the {spec.device_family} image {spec.name}; " - f"the node is device_family {device_family}") - return None + ``overlay`` (default None = off): a :mod:`yeto.rl.adapters.miles.overlay` + patch applied to the image's ~/miles after the checkout checks. + """ + + from ..rl import ( + MILES_NEXT_COMMIT, + MILES_NEXT_IMAGE_MANIFEST, + MILES_NEXT_IMAGE_SGLANG_ROOT, + MILES_NEXT_REPOSITORY, + MILES_PEFT_VERSION, + SGLANG_NEXT_COMMIT, + SGLANG_NEXT_REPOSITORY, + ) + + def checkout( + path: str, repository: str, commit: str, refreshed_flag: str = "" + ) -> str: + repo = shlex.quote(repository) + # ``refreshed_flag``: shell variable set to 1 when the checkout + # had to move to the pin (so the caller knows the image's + # install of it can no longer be trusted as-is). + mark = f"{refreshed_flag}=1\n" if refreshed_flag else "" + return ( + f"if [ ! -d {path}/.git ]; then git clone --no-checkout " + f"{repo} {path}; fi\n" + f"git -C {path} remote set-url origin {repo}\n" + f'test "$(git -C {path} config --get remote.origin.url)" = ' + f"{repo}\n" + # MILES_NEXT_IMAGE already has the pinned fork checked out + # at ~/miles (/root/miles); only fetch when it is not HEAD. + f'if [ "$(git -C {path} rev-parse HEAD 2>/dev/null)" != ' + f"{commit} ]; then\n" + f"git -C {path} fetch --depth 1 origin {commit}\n" + f"git -C {path} checkout --detach {commit}\n" + + mark + + "fi\n" + f'test "$(git -C {path} rev-parse HEAD)" = {commit}\n' + f'test "$(git -C {path} rev-parse --abbrev-ref HEAD)" = HEAD\n' + f'test -z "$(git -C {path} status --porcelain ' + '--untracked-files=all)"\n' + ) + + # MILES_NEXT_IMAGE ships the pinned SGLang fork installed editable + # at /sgl-workspace/sglang: ~/sglang (first on the island's + # PYTHONPATH) becomes a link to it. Any other image, or an existing + # ~/sglang directory, gets the clone + install at ~/sglang as before. + image_root = shlex.quote(MILES_NEXT_IMAGE_SGLANG_ROOT) + sglang_checkout_in_image = ( + f'{{ [ ! -e ~/sglang ] || [ "$(readlink ~/sglang)" = {image_root} ]; }} && ' + f'[ "$(git -C {image_root} rev-parse HEAD 2>/dev/null)" = ' + f"{SGLANG_NEXT_COMMIT} ] && " + f'[ "$(git -C {image_root} config --get remote.origin.url)" = ' + f"{shlex.quote(SGLANG_NEXT_REPOSITORY)} ] && " + f'[ -z "$(git -C {image_root} status --porcelain ' + '--untracked-files=all)" ]' + ) + # Fallback for an image without a manifest entry: the old import check. + sglang_import_ok = ( + "python3 -c 'import os, sys, sglang; sys.exit(0 if os.path.realpath(" + f'sglang.__file__).startswith("{MILES_NEXT_IMAGE_SGLANG_ROOT}/python/' + "\") else 1)'" + ) + # MILES_NEXT_IMAGE also ships the pinned Miles fork installed + # editable at /root/miles (= ~/miles) and records it in + # MILES_NEXT_IMAGE_MANIFEST. When ~/miles was already at the pin + # (no fetch above) and the manifest names that commit at that path, + # the image's install is kept: the editable re-install is skipped + # and only a missing/other peft is (re)installed. Any other image, + # a moved checkout, or a manifest mismatch gets the full install. + manifest = shlex.quote(MILES_NEXT_IMAGE_MANIFEST) + miles_in_image = ( + '[ "$MILES_REFRESHED" = 0 ] && ' + f"python3 -c 'import json, os, sys; " + f'm = json.load(open("{MILES_NEXT_IMAGE_MANIFEST}"))["miles"]; ' + f'sys.exit(0 if m["commit"] == "{MILES_NEXT_COMMIT}" and ' + 'os.path.realpath(m["path"]) == ' + "os.path.realpath(os.path.expanduser(\"~/miles\")) else 1)' " + "2>/dev/null" + ) + peft_ok = ( + "python3 -c 'import peft, sys; " + f"sys.exit(0 if peft.__version__ == \"{MILES_PEFT_VERSION}\" else 1)' " + "2>/dev/null" + ) + peft_check = image_manifest_check( + "peft", "version", MILES_PEFT_VERSION, + fallback=f"{peft_ok} || python3 -m pip install -q 'peft=={MILES_PEFT_VERSION}'", + ) + sglang_check = image_manifest_check( + "sglang", "commit", SGLANG_NEXT_COMMIT, path=MILES_NEXT_IMAGE_SGLANG_ROOT, + on_pass="SGLANG_IN_IMAGE=1", + fallback=f"if {sglang_import_ok}; then SGLANG_IN_IMAGE=1; fi", + ) + return ( + "set -e\n" + "MILES_REFRESHED=0\n" + + checkout( + "~/miles", MILES_NEXT_REPOSITORY, MILES_NEXT_COMMIT, "MILES_REFRESHED" + ) + + f"if {miles_in_image}; then\n" + f"echo \"[yeto-setup] image provides miles {MILES_NEXT_COMMIT} " + f'({manifest}); editable install kept"\n' + f"{peft_check}\n" + "else\n" + "python3 -m pip install -q --no-deps -e ~/miles " + f"'peft=={MILES_PEFT_VERSION}'\n" + "fi" + + (f"\n{_miles_overlay().overlay_setup(overlay)}" if overlay else ""), + "SGLANG_IN_IMAGE=0\n" + f"if {sglang_checkout_in_image}; then\n" + f"{sglang_check}\n" + "fi\n" + 'if [ "$SGLANG_IN_IMAGE" = 1 ]; then\n' + f"ln -sfn {image_root} ~/sglang\n" + f"echo '[yeto-setup] image provides sglang {SGLANG_NEXT_COMMIT}'\n" + "else\n" + + checkout("~/sglang", SGLANG_NEXT_REPOSITORY, SGLANG_NEXT_COMMIT) + + "python3 -m pip install -q --no-deps -e ~/sglang/python\n" + "fi", + ) diff --git a/yeto/launch/node_inventory.py b/yeto/launch/node_inventory.py index 3c7d8bad8..bdb55277b 100644 --- a/yeto/launch/node_inventory.py +++ b/yeto/launch/node_inventory.py @@ -10,7 +10,7 @@ ``cpu_arch`` (``x86_64`` or ``aarch64``) selects the image (S25 N1.6): a node without ``image`` takes ``images[cpu_arch]`` from the top of the file, and a node whose image is a yeto catalog image of another architecture or device -family is refused (``yeto.launch.images``). The trainer preflight also +family is refused (``yeto.launch.engine_images``). The trainer preflight also compares ``uname -m`` with ``cpu_arch``. Example:: @@ -43,7 +43,7 @@ from ..gpu_spec import DEVICE_FAMILIES, UnknownCard, canonical_card, device_family_of from ..hw import card_tools -from . import images as image_catalog +from . import engine_images as image_catalog # Fabrics an inventory may declare. "tcp" is Ethernet without RDMA; "roce" and # "ib" are RDMA over Ethernet and InfiniBand; "hccs" is the Ascend on-node From d857d9eb329dd11dba6cc077aba3ceec876438f6 Mon Sep 17 00:00:00 2001 From: michael <78113554+chowhsu@users.noreply.github.com> Date: Mon, 19 Oct 2026 12:25:04 +0000 Subject: [PATCH 21/34] Add the stage-1 PPU plan: one file expands and checks the GLM5.3-Flash run yeto-stage1 plan expands examples/stage1/glm53-flash-ppu.yaml and its node inventory into the Miles/Megatron layout and arguments, the verl FSDP2 overrides, the SGLang TP8 server command (built by the sglang-server adapter), the MXFP4 re-quantizer argv, the OSS paths and the ordered per-node steps, and checks that card counts, scheme and release-dir names, revisions, max_tokens and penalties agree. It opens no SSH connection. yeto-stage1 run refuses while a check error or an open gap (trainer image, sampler entry, glm5_next recipe, SSH launch path, head access to OSS) remains, and needs --confirm-machines-free. Co-Authored-By: Claude Opus 5.5 --- examples/stage1/glm53-flash-ppu.yaml | 116 +++++ examples/stage1/ppu-inventory.yaml | 58 +++ pyproject.toml | 1 + tests/test_stage1_plan.py | 176 +++++++ yeto/launch/stage1.py | 735 +++++++++++++++++++++++++++ 5 files changed, 1086 insertions(+) create mode 100644 examples/stage1/glm53-flash-ppu.yaml create mode 100644 examples/stage1/ppu-inventory.yaml create mode 100644 tests/test_stage1_plan.py create mode 100644 yeto/launch/stage1.py diff --git a/examples/stage1/glm53-flash-ppu.yaml b/examples/stage1/glm53-flash-ppu.yaml new file mode 100644 index 000000000..20dd25bbd --- /dev/null +++ b/examples/stage1/glm53-flash-ppu.yaml @@ -0,0 +1,116 @@ +# Stage 1 of the S25 goal: GLM5.3-Flash (glm5_next) bf16 full-parameter RL on +# 4 PPU nodes; the trainer and the sampler both run on PPU, the head runs on +# this server, the cross-site store is Alibaba OSS. +# +# yeto-stage1 plan examples/stage1/glm53-flash-ppu.yaml # expand + check, no SSH +# yeto-stage1 run examples/stage1/glm53-flash-ppu.yaml --confirm-machines-free +# +# Every number marked "estimate" is computed, not measured +# (SESSION25/analysis/A-TRN-ppu-training.md, A-INF-sampler-integration.md). +# Nothing in this file was run on a PPU. + +run_id: s25-stage1-glm53-ppu +inventory: ppu-inventory.yaml # relative to this file +rounds: 3 # the stage-1 closed loop: 3 RL rounds +seed: 17 + +# Rendezvous (eth0) address of each node. The order ppu1..ppu4 = 10.0.0.74..77 +# is assumed from the survey listing; runbook step 1 checks it. +rendezvous: + ppu1: 10.0.0.74 + ppu2: 10.0.0.75 + ppu3: 10.0.0.76 + ppu4: 10.0.0.77 + +roles: + trainer: [ppu1, ppu2, ppu3] # 24 cards, one pipeline stage per node + sampler: [ppu4] # 8 cards, one SGLang TP8 engine + +model: + name: GLM5.3-Flash + architecture: glm5_next + num_layers: 45 # 34 KDA + 11 DSA MLA; layer 45 carries MTP + params_b: 321 # 642,652,083,504 B bf16 index / 2 + expert_params_b: 304 # A-QNT estimate + bf16_source: + revision: cecb5eeb9c6b32404a0dd930df81de2c239bdd84 + 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 + sampler_initial: # round 0: the HF MXFP4 export, not our re-quantizer + repo: yetalabs/apex-flash-1-mxfp4-experts + revision: 70e86a0 + staged: oss://oss-pai-ircyeaqtnfjw11wby2-cn-wulanchabu/yeto/models/apex-flash-1-mxfp4-experts-70e86a0/ + local_dir: /data/models/apex-flash-1-mxfp4-experts-70e86a0 + +quant: # rounds >= 1: N8.5 re-quantizes each bf16 release + scheme: mxfp4-g32-rtn + format: mxfp4-pack-quantized + release_dir: policy-releases-mxfp4-g32-rtn + expected_expert_matrices: 36288 # A-INF post-load check 1 + +store: oss://oss-pai-ircyeaqtnfjw11wby2-cn-wulanchabu/yeto/s25-stage1 + +images: + sampler: inference-xpu-pytorch-opt:v0.5.17 + trainer: null # N0.2 (T-N02) builds it; until then `run` refuses + docker_device_args: [] # PPU device flags of `docker run` (N0.2b, not known yet) + +trainer: + backend: miles # first backend; verl below is written, not first + miles: + commit: cc76e2391 + megatron: radixark/Megatron-LM@miles-main + mbridge: 89eb108 + workdir: /data/yeto/miles + tp: 4 + sequence_parallel: true + pp: 3 + pp_layers: [15, 15, 15] # first / middle / last stage + cp: 1 + ep: 8 # inside one node + etp: 1 + optimizer_offload_fraction: 1.0 # Adam state on the host (estimate: 1.2 TiB/node of 2266 GiB) + transformer_impl: transformer_engine # main-agent ruling 10-19 (3): TE stays default; `local` is the N7.5 fallback + recompute: full + max_tokens_per_gpu: 8192 + lr: 1.0e-6 + verl: + commit: null # pinned by N7.2/N10.3 work; not chosen yet + strategy: fsdp2 + fsdp_size: 24 + ulysses: 1 # glm5_next has no Ulysses path (N7.1) + param_offload: false + optimizer_offload: true + lr: 1.0e-6 + +sampler: + engine: sglang + engine_version: 0.5.17 + topology: tp8 + tp: 8 + precision: mxfp4 + max_model_len: 8192 + port: 30000 + moe_runner_backend: triton # report item 4: flashinfer cutlass MoE crashes in graph capture + hierarchical_cache: false # report item 5 + nextn: false # A-INF items 2 and 8 + renames_weight_packed: true # MXFP4 stores experts as *.weight_packed; runbook step S5 checks the count + in_place_reload: false # restart per version until runbook step S6 passes + +sampling: # one value for every source of the run + max_tokens: 4096 + max_prompt_len: 1024 + temperature: 1.0 + top_p: 1.0 + top_k: -1 + repetition_penalty: 1.0 # neutral: penalties change the behaviour distribution + frequency_penalty: 0.0 + presence_penalty: 0.0 + +task: # main-agent ruling: GSM8K first, TB2/Codex with N10.9 + data: zhuzilin/gsm8k + data_revision: 0cbd9f31d91ac21a7613dcbc7fef992adac459ae + reward: yeto.rl.gsm8k_reward:score + rollout_batch_size: 16 + n_samples_per_prompt: 8 diff --git a/examples/stage1/ppu-inventory.yaml b/examples/stage1/ppu-inventory.yaml new file mode 100644 index 000000000..fbc21ea33 --- /dev/null +++ b/examples/stage1/ppu-inventory.yaml @@ -0,0 +1,58 @@ +# Stage-1 node inventory: 4 x PPU nodes, 8 x ZW-M890P 144 GiB each +# (SESSION25/N0.1-ppu-survey.md, read-only survey 2026-10-16). +# host = the SSH alias in ~/.ssh/config of the head (this server). The aliases +# do not resolve on the nodes; the stage-1 plan gives each node's rendezvous +# address (eth0) separately. +nodes: + - name: ppu1 + host: ppu1 + device_family: ppu + card: M890 + cards_per_node: 8 + hbm_gib: 144 + host_memory_gib: 2266 + fabric: roce + interface: eth0 + mounts: {/data: /data} + image: inference-xpu-pytorch-opt:v0.5.17 + runtime: docker + roles: [trainer] + - name: ppu2 + host: ppu2 + device_family: ppu + card: M890 + cards_per_node: 8 + hbm_gib: 144 + host_memory_gib: 2266 + fabric: roce + interface: eth0 + mounts: {/data: /data} + image: inference-xpu-pytorch-opt:v0.5.17 + runtime: docker + roles: [trainer] + - name: ppu3 + host: ppu3 + device_family: ppu + card: M890 + cards_per_node: 8 + hbm_gib: 144 + host_memory_gib: 2266 + fabric: roce + interface: eth0 + mounts: {/data: /data} + image: inference-xpu-pytorch-opt:v0.5.17 + runtime: docker + roles: [trainer] + - name: ppu4 + host: ppu4 + device_family: ppu + card: M890 + cards_per_node: 8 + hbm_gib: 144 + host_memory_gib: 2266 + fabric: roce + interface: eth0 + mounts: {/data: /data} + image: inference-xpu-pytorch-opt:v0.5.17 + runtime: docker + roles: [sampler] diff --git a/pyproject.toml b/pyproject.toml index 1ca14343d..c556b1277 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -62,6 +62,7 @@ yeto-diffusion-export = "yeto.diffusion.export:main" yeto-diffusion-sample = "yeto.diffusion.sample:main" yeto-rl-export = "yeto.rl.export:main" yeto-rl-ssh = "yeto.rl.ssh_harness:main" +yeto-stage1 = "yeto.launch.stage1:main" yeto-protenix-export-batch = "yeto.diffusion.protenix_export:main" yeto-hunyuan3d-export-batch = "yeto.diffusion.hunyuan3d_export:main" diff --git a/tests/test_stage1_plan.py b/tests/test_stage1_plan.py new file mode 100644 index 000000000..f30d90b29 --- /dev/null +++ b/tests/test_stage1_plan.py @@ -0,0 +1,176 @@ +"""Stage-1 plan expansion and consistency checks (S25 T-STAGE1). CPU only, no SSH.""" + +from __future__ import annotations + +import copy +import shlex +from pathlib import Path + +import pytest +import yaml + +from yeto.launch import stage1 + +REPO = Path(__file__).resolve().parents[1] +PLAN = REPO / "examples" / "stage1" / "glm53-flash-ppu.yaml" + + +@pytest.fixture(scope="module") +def shipped(): + return stage1.expand(stage1.load_plan(PLAN)) + + +def _variant(tmp_path: Path, edit) -> stage1.Expanded: + """Expand a copy of the shipped plan after ``edit(doc)`` changed it.""" + doc = yaml.safe_load(PLAN.read_text()) + edit(doc) + (tmp_path / "ppu-inventory.yaml").write_text((PLAN.parent / "ppu-inventory.yaml").read_text()) + path = tmp_path / "plan.yaml" + path.write_text(yaml.safe_dump(doc)) + return stage1.expand(stage1.load_plan(path)) + + +def test_shipped_plan_is_consistent(shipped): + assert shipped.errors == [] + assert shipped.trainer["cards"] == 24 + assert shipped.trainer["miles"]["layout"] == {"tp": 4, "pp": 3, "cp": 1, "ep": 8, "etp": 1, "dp": 2, + "expert_dp": 1} + assert [e["node"] for e in shipped.sampler["engines"]] == ["ppu4"] + + +def test_names_agree_across_roles(shipped): + run = shipped.store["run"] + assert run.endswith("/s25-stage1-glm53-ppu") and run.startswith("oss://") + req = shipped.requant["argv"] + assert req[req.index("--store") + 1] == run + assert req[req.index("--target-dir") + 1] == "policy-releases-mxfp4-g32-rtn" + island = shipped.trainer["miles"]["island_argv"] + assert island[island.index("--rl-resume-store") + 1] == run + for step in shipped.steps: + if step.script.startswith("docker run"): + assert f"YETO_RL_STORE={run}" in step.script + assert "yeto-s25-stage1-glm53-ppu-" in step.script + + +def test_sampler_command_comes_from_the_adapter(shipped): + engine = shipped.sampler["engines"][0] + argv = engine["argv"] + assert argv[argv.index("--tp") + 1] == "8" + assert argv[argv.index("--model-path") + 1].endswith("mxfp4-experts-70e86a0") + assert "--speculative-algorithm" not in argv # no NEXTN in stage 1 + assert argv[argv.index("--moe-runner-backend") + 1] == "triton" + assert engine["env"]["MX_PATCH"] == "1" + assert engine["url"] == "http://10.0.0.77:30000" + + +def test_max_tokens_is_one_value_everywhere(shipped): + island = shipped.trainer["miles"]["island_argv"] + margs = shipped.trainer["miles"]["megatron_args"] + assert island[island.index("--rollout-max-response-len") + 1] == "4096" + assert margs[margs.index("--rollout-max-response-len") + 1] == "4096" + assert "data.max_response_length=4096" in shipped.trainer["verl"]["overrides"] + assert shipped.sampler["sampling"]["max_tokens"] == 4096 + + +def test_trainer_is_bf16_with_host_optimizer(shipped): + margs = " ".join(shipped.trainer["miles"]["megatron_args"]) + assert "--bf16" in margs and "fp8" not in margs + assert "--optimizer-cpu-offload" in margs + assert "--decoder-first-pipeline-num-layers 15 --decoder-last-pipeline-num-layers 15" in margs + assert shipped.trainer["miles"]["drop_model_args"] == [] # TE stays the default (ruling 10-19 (3)) + + +def test_island_argv_parses_with_the_island_parser(shipped): + assert stage1._check_island_argv(shipped.trainer["miles"]["island_argv"]) is None + + +def test_steps_order_and_staging(shipped): + names = [s.name for s in shipped.steps] + assert names[:4] == [f"preflight-ppu{i}" for i in range(1, 5)] + assert names.index("stage-sampler_initial-ppu1") < names.index("requant-sidecar") # template on node 0 + assert names.index("sampler-ppu4") < names.index("trainer-island") + assert names[-1] == "requant-sidecar" + for step in shipped.steps: + if step.script.startswith("docker run"): + shlex.split(step.script) # quotable as one command + + +def test_gaps_block_run_until_wired(shipped, capsys): + assert any("N0.2" in g for g in shipped.gaps) + assert any("OSS" in g for g in shipped.gaps) + assert stage1.main(["run", str(PLAN), "--confirm-machines-free"]) == 2 + assert "run refused" in capsys.readouterr().err + + +def test_plan_command_prints_without_ssh(capsys): + assert stage1.main(["plan", str(PLAN), "--json"]) == 0 + assert '"schema": "yeto.stage1/v1"' in capsys.readouterr().out + + +def test_run_needs_confirmation(monkeypatch, shipped, capsys): + clean = copy.copy(shipped) + clean.gaps = [] + monkeypatch.setattr(stage1, "expand", lambda plan, hf_config=None: clean) + calls = [] + monkeypatch.setattr(stage1, "run_steps", lambda out: calls.append(out) or 0) + assert stage1.main(["run", str(PLAN)]) == 2 + assert "--confirm-machines-free" in capsys.readouterr().err + assert stage1.main(["run", str(PLAN), "--confirm-machines-free"]) == 0 + assert calls == [clean] + + +def test_run_steps_stops_at_first_failure(shipped): + seen = [] + + class Done: + def __init__(self, rc): + self.returncode = rc + + def runner(argv, input=None, text=True): + seen.append(argv) + return Done(1 if len(seen) == 2 else 0) + + assert stage1.run_steps(shipped, runner=runner) == 1 + assert len(seen) == 2 and seen[0][-3:] == ["ppu1", "bash", "-s"] + + +@pytest.mark.parametrize("edit, message", [ + (lambda d: d["trainer"]["miles"].update(pp_layers=[15, 15, 14]), "sum to 44"), + (lambda d: d["trainer"]["miles"].update(pp_layers=[23, 22]), "2 stages"), + (lambda d: d["sampling"].update(repetition_penalty=1.1), "repetition_penalty"), + (lambda d: d["quant"].update(release_dir="policy-releases-w4a8"), "release_dir"), + (lambda d: d["quant"].update(scheme="w4a8-int4-channel"), "is scheme mxfp4-g32-rtn"), + (lambda d: d["sampler"].update(renames_weight_packed=False), "weight_packed"), + (lambda d: d["sampler"].update(engine_version="0.5.18"), "engine version"), + (lambda d: d["sampler"].update(max_model_len=4096), "max_model_len"), + (lambda d: d["model"]["sampler_initial"].update(revision="aeb9adf"), "revision aeb9adf"), + (lambda d: d["model"]["bf16_source"].update(local_dir="/data/models/bf16"), "bf16_source revision"), + (lambda d: d["roles"].update(sampler=["ppu3", "ppu4"]), "both trainer and sampler"), + (lambda d: d["rendezvous"].update(ppu2="ppu2"), "not an IP address"), + (lambda d: d["trainer"]["miles"].update(optimizer_offload_fraction=0.0), "activations"), + (lambda d: d.update(store="s3://bucket/x"), "oss://"), +]) +def test_contradictions_are_errors(tmp_path, edit, message): + out = _variant(tmp_path, edit) + assert any(message in e for e in out.errors), out.errors + + +def test_layout_megatron_refuses(tmp_path): + with pytest.raises(stage1.Stage1Error, match="inside one node"): + _variant(tmp_path, lambda d: d["trainer"]["miles"].update(ep=12, pp=1, tp=1, pp_layers=[45])) + with pytest.raises(stage1.Stage1Error, match="multiple"): + _variant(tmp_path, lambda d: d["trainer"]["miles"].update(tp=5)) + + +def test_te_fallback_drops_te_only_model_args(tmp_path): + out = _variant(tmp_path, lambda d: d["trainer"]["miles"].update(transformer_impl="local")) + assert out.errors == [] + assert out.trainer["miles"]["drop_model_args"] == ["--moe-grouped-gemm", "--moe-permute-fusion"] + assert "--no-gradient-accumulation-fusion" in out.trainer["miles"]["megatron_args"] + + +def test_verl_backend_expands(tmp_path): + out = _variant(tmp_path, lambda d: d["trainer"].update(backend="verl")) + assert out.errors == [] + assert "trainer.nnodes=3" in out.trainer["verl"]["overrides"] + assert any("verl" in g for g in out.gaps) diff --git a/yeto/launch/stage1.py b/yeto/launch/stage1.py new file mode 100644 index 000000000..92ba9d962 --- /dev/null +++ b/yeto/launch/stage1.py @@ -0,0 +1,735 @@ +"""Stage 1 of the S25 goal as one plan: GLM5.3-Flash bf16 full-parameter RL on +PPU nodes (trainer + SGLang sampler on PPU, head on this server, OSS store). + +``yeto-stage1 plan PLAN.yaml`` expands the plan file and its node inventory +(``--rl-node-inventory`` format, :mod:`.node_inventory`) into every command, +environment and path the run uses, checks that they agree with each other, and +prints the result. It never opens an SSH connection. + +``yeto-stage1 run PLAN.yaml --confirm-machines-free`` runs the same steps in +order over SSH. It refuses while the plan has a check error or an open gap +(an entry point, image or path that does not exist yet), so it cannot start a +half-wired run on machines other people use. + +What the plan expands (all derived from one place, so the names agree): + +* nodes and roles: trainer nodes, sampler nodes, card counts, rendezvous + addresses; +* trainer, first backend Miles/Megatron: the parallel layout (TP x PP x EP, + the pipeline layer split), the Megatron arguments the glm5_next recipe must + give, the yeto Miles island argv; second backend verl/FSDP2: the trainer + overrides (written, not the first backend); +* sampler: the SGLang TP8 server argv and environment from the + ``sglang-server`` adapter itself (:mod:`yeto.rl.adapters.sampler_engines`); +* re-quantizer: the :mod:`yeto.rl.quant.release` argv (N8.5 MXFP4 RTN), with + the release directory named after the scheme; +* store: the run directory on OSS and the OSS staging paths of the weights; +* steps: the ordered per-node scripts of the run (preflight, staging, sampler, + trainer). + +Check errors are contradictions inside the plan (card counts, scheme names, +revisions, sampling values). Gaps are things the plan needs that the code base +or the machines do not provide yet; they are expected until the stage-1 +wiring is finished and are listed so the runbook can close them one by one. + +Nothing here was run against a PPU node. +""" + +from __future__ import annotations + +import argparse +import importlib +import importlib.util +import ipaddress +import json +import shlex +import subprocess +import sys +from collections.abc import Mapping, Sequence +from dataclasses import dataclass, field +from pathlib import Path +from typing import Any + +from .node_inventory import InventoryError, Node, load_inventory +from .topology import island_comm_env + +PLAN_SCHEMA = "yeto.stage1/v1" +BACKENDS = ("miles", "verl") +NEUTRAL_PENALTIES = {"repetition_penalty": 1.0, "frequency_penalty": 0.0, "presence_penalty": 0.0} +RAY_PORT = 6379 +HOST_USABLE_FRACTION = 0.85 # as yeto.rl.fullparam.memory.TrainPlan.host_usable_fraction +RELEASE_DIR_PREFIX = "policy-releases-" + + +class Stage1Error(ValueError): + """The plan file cannot be read or expanded.""" + + +# --------------------------------------------------------------------- plan file + +def _section(doc: Mapping[str, Any], key: str) -> dict[str, Any]: + value = doc.get(key) + if not isinstance(value, dict): + raise Stage1Error(f"plan: section {key!r} must be a mapping, got {type(value).__name__}") + return value + + +def _get(section: Mapping[str, Any], key: str, where: str) -> Any: + if key not in section: + raise Stage1Error(f"plan: {where}.{key} is required") + return section[key] + + +@dataclass(frozen=True) +class Stage1Plan: + """The plan file, read and with its inventory loaded.""" + + path: Path + doc: Mapping[str, Any] + nodes: tuple[Node, ...] + + @property + def run_id(self) -> str: + return str(_get(self.doc, "run_id", "plan")) + + def section(self, key: str) -> dict[str, Any]: + return _section(self.doc, key) + + def nodes_of(self, role: str) -> tuple[Node, ...]: + by_name = {n.name: n for n in self.nodes} + names = _section(self.doc, "roles").get(role) or [] + missing = [n for n in names if n not in by_name] + if missing: + raise Stage1Error(f"plan: roles.{role} names {missing}, which the inventory does not list") + return tuple(by_name[n] for n in names) + + +def load_plan(path: str | Path) -> Stage1Plan: + import yaml # noqa: PLC0415 - optional at import time, like node_inventory + + file = Path(path) + try: + doc = yaml.safe_load(file.read_text()) + except (OSError, yaml.YAMLError) as exc: + raise Stage1Error(f"cannot read the plan {file}: {exc}") from None + if not isinstance(doc, dict): + raise Stage1Error(f"{file}: the plan must be a mapping") + inventory = file.parent / str(_get(doc, "inventory", "plan")) + try: + nodes = tuple(load_inventory(inventory)) + except InventoryError as exc: + raise Stage1Error(str(exc)) from None + return Stage1Plan(path=file, doc=doc, nodes=nodes) + + +# ----------------------------------------------------------------- expansion + +@dataclass +class Step: + """One step of the run: a script for one node (``node`` None = the head).""" + + name: str + node: str | None + script: str + purpose: str + + def record(self) -> dict[str, Any]: + return {"name": self.name, "node": self.node, "purpose": self.purpose, "script": self.script} + + +@dataclass +class Expanded: + plan: Stage1Plan + trainer: dict[str, Any] = field(default_factory=dict) + sampler: dict[str, Any] = field(default_factory=dict) + requant: dict[str, Any] = field(default_factory=dict) + store: dict[str, Any] = field(default_factory=dict) + steps: list[Step] = field(default_factory=list) + errors: list[str] = field(default_factory=list) + gaps: list[str] = field(default_factory=list) + + def record(self) -> dict[str, Any]: + return { + "schema": PLAN_SCHEMA, + "run_id": self.plan.run_id, + "plan": str(self.plan.path), + "nodes": [n.record() for n in self.plan.nodes], + "roles": {r: [n.name for n in self.plan.nodes_of(r)] for r in ("trainer", "sampler")}, + "store": self.store, + "trainer": self.trainer, + "sampler": self.sampler, + "requant": self.requant, + "steps": [s.record() for s in self.steps], + "errors": list(self.errors), + "gaps": list(self.gaps), + } + + +def run_store(plan: Stage1Plan) -> str: + """The run directory on the store: ``/``.""" + return f"{str(_get(plan.doc, 'store', 'plan')).rstrip('/')}/{plan.run_id}" + + +def megatron_layout(miles: Mapping[str, Any], trainer_cards: int, cards_per_node: int) -> dict[str, int]: + """Data-parallel sizes of a TP x PP x EP layout (Megatron rules); raises on a + layout Megatron would refuse.""" + tp, pp, cp, ep, etp = (int(miles[k]) for k in ("tp", "pp", "cp", "ep", "etp")) + if trainer_cards % (tp * pp * cp): + raise Stage1Error(f"trainer cards {trainer_cards} are not a multiple of TP*PP*CP = {tp * pp * cp}") + dp = trainer_cards // (tp * pp * cp) + if (dp * tp * cp) % (ep * etp): + raise Stage1Error(f"EP*ETP = {ep * etp} must divide DP*TP*CP = {dp * tp * cp}") + if ep * etp > cards_per_node or tp > cards_per_node: + raise Stage1Error(f"EP*ETP = {ep * etp} and TP = {tp} must stay inside one node " + f"({cards_per_node} cards; A-TRN: EP and TP do not cross RoCE)") + return {"tp": tp, "pp": pp, "cp": cp, "ep": ep, "etp": etp, "dp": dp, + "expert_dp": dp * tp * cp // (ep * etp)} + + +def megatron_static_gib(model: Mapping[str, Any], layout: Mapping[str, int]) -> dict[str, float]: + """Per-card static memory of the layout, estimate (A-TRN §2.1 method): + bf16 weights + fp32 main grads = 6 B/param on the card; Adam + fp32 master + = 12 B/param, sharded over the data-parallel group, on the host when + offloaded. Activations are not included.""" + experts = float(model["expert_params_b"]) * 1e9 + dense = float(model["params_b"]) * 1e9 - experts + per_card = experts / (layout["pp"] * layout["ep"] * layout["etp"]) + dense / (layout["pp"] * layout["tp"]) + optimizer_card = (experts / (layout["pp"] * layout["ep"] * layout["etp"] * layout["expert_dp"]) + + dense / (layout["pp"] * layout["tp"] * layout["dp"])) * 12 + gib = 2.0 ** 30 + return {"params_per_card_b": round(per_card / 1e9, 2), "weights_grads_gib": round(per_card * 6 / gib, 1), + "optimizer_gib_per_card": round(optimizer_card / gib, 1)} + + +def megatron_args(model: Mapping[str, Any], miles: Mapping[str, Any], layout: Mapping[str, int], + sampling: Mapping[str, Any]) -> list[str]: + """The Megatron arguments the Miles glm5_next recipe must pass on top of the + model arguments (Miles ``scripts/models/glm5.3-flash.py``).""" + split = [int(x) for x in miles["pp_layers"]] + args = ["--tensor-model-parallel-size", str(layout["tp"]), + "--pipeline-model-parallel-size", str(layout["pp"]), + "--context-parallel-size", str(layout["cp"]), + "--expert-model-parallel-size", str(layout["ep"]), + "--expert-tensor-parallel-size", str(layout["etp"])] + if miles.get("sequence_parallel"): + args.append("--sequence-parallel") + if layout["pp"] > 1: + args += ["--decoder-first-pipeline-num-layers", str(split[0]), + "--decoder-last-pipeline-num-layers", str(split[-1])] + args += ["--recompute-granularity", str(miles["recompute"]), "--recompute-method", "uniform", + "--recompute-num-layers", "1", "--micro-batch-size", "1", + "--max-tokens-per-gpu", str(int(miles["max_tokens_per_gpu"])), + "--optimizer", "adam", "--lr", repr(float(miles["lr"])), "--lr-decay-style", "constant", + "--use-distributed-optimizer", "--bf16", + "--rollout-max-response-len", str(int(sampling["max_tokens"])), + "--transformer-impl", str(miles["transformer_impl"])] + fraction = float(miles.get("optimizer_offload_fraction") or 0.0) + if fraction > 0: + args += ["--optimizer-cpu-offload", "--optimizer-offload-fraction", repr(fraction), + "--overlap-cpu-optimizer-d2h-h2d"] + if miles["transformer_impl"] == "local": # N7.5 fallback: no TE kernels + args.append("--no-gradient-accumulation-fusion") + return args + + +TE_ONLY_MODEL_ARGS = ("--moe-grouped-gemm", "--moe-permute-fusion") # N7.5 E6, E7 + + +def verl_overrides(verl: Mapping[str, Any], trainer_nodes: Sequence[Node], task: Mapping[str, Any], + sampling: Mapping[str, Any], model_dir: str) -> list[str]: + """verl trainer overrides of the FSDP2 backend (rollout runs on the yeto + sampler island, not inside verl).""" + actor = "actor_rollout_ref.actor" + return [f"actor_rollout_ref.model.path={model_dir}", + f"{actor}.strategy={verl['strategy']}", + f"{actor}.fsdp_config.fsdp_size={int(verl['fsdp_size'])}", + f"{actor}.fsdp_config.param_offload={bool(verl['param_offload'])}", + f"{actor}.fsdp_config.optimizer_offload={bool(verl['optimizer_offload'])}", + f"{actor}.ulysses_sequence_parallel_size={int(verl['ulysses'])}", + f"{actor}.optim.lr={float(verl['lr'])}", + f"actor_rollout_ref.rollout.n={int(task['n_samples_per_prompt'])}", + f"data.train_batch_size={int(task['rollout_batch_size'])}", + f"data.max_prompt_length={int(sampling['max_prompt_len'])}", + f"data.max_response_length={int(sampling['max_tokens'])}", + f"trainer.nnodes={len(trainer_nodes)}", + f"trainer.n_gpus_per_node={trainer_nodes[0].cards_per_node}", + "trainer.device=cuda"] # PPU runs the CUDA compatibility layer (rl-ssh-node-pool D2) + + +def _sampler_engine(plan: Stage1Plan, node: Node): + from yeto.rl.adapters.sampler_engines.server import ServerConfig # noqa: PLC0415 + from yeto.rl.adapters.sampler_engines.sglang_server import SGLangOptions, SGLangServerEngine # noqa: PLC0415 + + s, q = plan.section("sampler"), plan.section("quant") + extra = ["--moe-runner-backend", str(s["moe_runner_backend"])] + if s.get("hierarchical_cache"): + extra.append("--enable-hierarchical-cache") + config = ServerConfig( + base_dir=Path(plan.section("model")["sampler_initial"]["local_dir"]), + card_model=node.card, device_family=node.device_family, + engine_version=str(s["engine_version"]), precision=str(s["precision"]), + 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"))) + options = SGLangOptions(tp=int(s["tp"]), nextn=bool(s.get("nextn")), + renames_weight_packed=bool(s.get("renames_weight_packed"))) + return SGLangServerEngine(config, options) + + +def _docker(plan: Stage1Plan, node: Node, role: str, image: str | None, command: Sequence[str], + env: Mapping[str, str]) -> str: + images = plan.section("images") + argv = ["docker", "run", "--rm", "--name", f"yeto-{plan.run_id}-{role}-{node.name}", + "--network", "host", "--ipc", "host", *[str(a) for a in images.get("docker_device_args") or []]] + for host_path, container_path in sorted(node.mounts.items()): + argv += ["-v", f"{host_path}:{container_path}"] + for key, value in sorted(env.items()): + argv += ["-e", f"{key}={value}"] + return shlex.join([*argv, image or "", *command]) + + +def _module_gap(module: str, what: str, attr: str = "main") -> str | None: + try: + spec = importlib.util.find_spec(module) + except (ImportError, ValueError): + spec = None + if spec is None: + return f"{what}: module {module} does not exist in this tree" + try: + loaded = importlib.import_module(module) + except Exception as exc: # noqa: BLE001 - any import failure is a gap, reported as such + return f"{what}: module {module} does not import here ({type(exc).__name__}: {exc})" + if not callable(getattr(loaded, attr, None)): + return f"{what}: module {module} has no {attr}() entry point" + return None + + +def _provenance_args(reward: str) -> list[str]: + """``--reward-sha256`` / ``--source-sha256`` of this tree, as the launcher + computes them (the code pack the nodes run must be this tree).""" + from yeto.provenance import python_spec_sha256, source_tree_sha256 # noqa: PLC0415 + + repo = Path(__file__).resolve().parents[2] + try: + reward_sha = python_spec_sha256(reward, base_dir=repo) + except (OSError, ValueError) as exc: + raise Stage1Error(f"plan: task.reward {reward!r} does not resolve in this tree: {exc}") from None + return ["--reward-sha256", reward_sha, "--source-sha256", source_tree_sha256(repo / "yeto")] + + +class _ParseError(Exception): + pass + + +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 + + parser = build_parser() + + def fail(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 + + +def expand(plan: Stage1Plan, *, hf_config: str | Path | None = None) -> Expanded: + """Every command, environment and path of the plan, plus the check results.""" + out = Expanded(plan=plan) + model, quant, sampling = plan.section("model"), plan.section("quant"), plan.section("sampling") + task, trainer_cfg, sampler_cfg = plan.section("task"), plan.section("trainer"), plan.section("sampler") + images, rendezvous = plan.section("images"), plan.section("rendezvous") + trainers, samplers = plan.nodes_of("trainer"), plan.nodes_of("sampler") + store = run_store(plan) + staging = {k: model[k]["staged"] for k in ("bf16_source", "sampler_initial")} + out.store = {"run": store, "release_dir": f"{store}/{quant['release_dir']}", + "bf16_release_dir": f"{store}/policy-releases", "staging": staging} + + # Trainer ------------------------------------------------------------- + backend = str(trainer_cfg["backend"]) + if backend not in BACKENDS: + raise Stage1Error(f"plan: trainer.backend {backend!r} is not one of {BACKENDS}") + cards = sum(n.cards_per_node for n in trainers) + per_node = trainers[0].cards_per_node if trainers else 0 + miles = _section(trainer_cfg, "miles") + layout = megatron_layout(miles, cards, per_node) + common_env = {"YETO_RL_DEVICE_FAMILY": trainers[0].device_family, "YETO_RUN_ID": plan.run_id, + "YETO_RL_STORE": store, "HF_HUB_OFFLINE": "1"} + comm = island_comm_env(trainers[0].device_family, trainers[0].fabric, trainers[0].interface, + num_nodes=len(trainers)) + trainer_env = {**common_env, **comm, "TORCHINDUCTOR_COMPILE_THREADS": "1"} + model_dir = str(model["bf16_source"]["local_dir"]) + state_dir = f"/data/yeto/{plan.run_id}" + 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"]), + "--reward-function", str(task["reward"]), + "--samples-per-group", str(int(task["n_samples_per_prompt"])), + "--rollout-max-response-len", str(int(sampling["max_tokens"])), + "--seq-len", str(int(sampling["max_prompt_len"]) + int(sampling["max_tokens"])), + "--inner-lr", repr(float(miles["lr"])), "--seed", str(int(plan.doc.get("seed", 0))), + "--parameter-mode", "full", "--rl-sample-role", "trainer", + "--rl-resume-store", store, + "--tensor-parallel", str(layout["tp"]), "--pipeline-parallel", str(layout["pp"]), + "--expert-parallel", str(layout["ep"]), + "--miles-root", str(miles["workdir"]), + "--trust-remote-code", + "--learner-id", "0", "--num-learners", "1", "--rl-single-island-no-sync", + "--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"])), + "--optimizer-steps", "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"])), + "--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"]))] + out.trainer = { + "backend": backend, "nodes": [n.name for n in trainers], "cards": cards, + "image": images.get("trainer"), "env": trainer_env, + "miles": {"commit": miles["commit"], "megatron": miles["megatron"], "mbridge": miles["mbridge"], + "layout": layout, "static_memory_estimate": megatron_static_gib(model, layout), + "megatron_args": megatron_args(model, miles, layout, sampling), + "drop_model_args": list(TE_ONLY_MODEL_ARGS) if miles["transformer_impl"] == "local" else [], + "island_argv": island_argv}, + "verl": {"commit": _section(trainer_cfg, "verl").get("commit"), + "overrides": verl_overrides(_section(trainer_cfg, "verl"), trainers, task, sampling, model_dir)}, + } + if hf_config is not None: + out.trainer["verl"]["memory_estimate"] = _verl_memory(trainers, plan, sampling, hf_config) + + # Sampler ------------------------------------------------------------- + sampler_runs = [] + for node in samplers: + engine = _sampler_engine(plan, node) + sampler_runs.append({"node": node.name, "argv": engine.argv(engine.config.base_dir), + "env": {**common_env, "YETO_RL_DEVICE_FAMILY": node.device_family, **engine.env()}, + "url": f"http://{rendezvous.get(node.name)}:{engine.config.port}"}) + out.sampler = {"engine": sampler_cfg["engine"], "engine_version": sampler_cfg["engine_version"], + "image": images.get("sampler"), "topology": sampler_cfg["topology"], + "precision": sampler_cfg["precision"], "quant_scheme": quant["scheme"], + "initial": {k: model["sampler_initial"][k] for k in ("repo", "revision", "local_dir")}, + "sampling": dict(sampling), "engines": sampler_runs} + + # Re-quantizer (node 0 of the trainer, a side process of the island) --- + from .ssh_trainer_island import requant_argv # noqa: PLC0415 + + argv = [*requant_argv(store, str(model["sampler_initial"]["local_dir"]), python="python3"), + "--target-dir", str(quant["release_dir"])] + out.requant = {"node": trainers[0].name if trainers else None, "scheme": quant["scheme"], + "format": quant["format"], "argv": argv, "loop_argv": [*argv, "--loop"]} + + out.steps = _steps(plan, out, trainers, samplers) + out.errors = check(plan, out) + out.gaps = gaps(plan, out) + return out + + +def _verl_memory(trainers: Sequence[Node], plan: Stage1Plan, sampling: Mapping[str, Any], + hf_config: str | Path) -> dict[str, Any]: + from yeto.rl.fullparam.memory import estimate, plan_for_nodes, shape_from_hf_config # noqa: PLC0415 + + verl = _section(plan.section("trainer"), "verl") + shape = shape_from_hf_config(json.loads(Path(hf_config).read_text())) + # yeto's FSDP2 config has no optimizer-only host offload (N7.1); the optimizer + # offload of the plan is estimated as the full CPU offload policy. + offload = "full" if verl["param_offload"] or verl["optimizer_offload"] else "none" + train_plan = plan_for_nodes(list(trainers), seq_len=int(sampling["max_prompt_len"]) + int(sampling["max_tokens"]), + fsdp_size=int(verl["fsdp_size"]), ulysses=int(verl["ulysses"]), offload=offload) + return estimate(shape, train_plan).summary() + + +def _steps(plan: Stage1Plan, out: Expanded, trainers: Sequence[Node], samplers: Sequence[Node]) -> list[Step]: + from yeto.hw.card_tools import preflight_commands # noqa: PLC0415 + + rendezvous = plan.section("rendezvous") + steps: list[Step] = [] + for node in plan.nodes: + tools = "\n".join(shlex.join(c) for c in preflight_commands(node.device_family)) + steps.append(Step(f"preflight-{node.name}", node.name, + f"set -eu\n{tools}\nip -4 -o addr show {node.interface or 'eth0'}\n" + "df -h /data\nfree -g\n", + "cards match the inventory; rendezvous address and /data present")) + # The re-quantizer on trainer node 0 reads the MXFP4 export as its template. + for key, nodes in (("bf16_source", trainers), ("sampler_initial", (*samplers, *trainers[:1]))): + src = plan.section("model")[key] + for node in nodes: + steps.append(Step(f"stage-{key}-{node.name}", node.name, + f"set -eu\nmkdir -p {shlex.quote(src['local_dir'])}\n" + 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']})")) + for run in out.sampler["engines"]: + node = next(n for n in samplers if n.name == run["node"]) + steps.append(Step(f"sampler-{node.name}", node.name, + _docker(plan, node, "sampler", out.sampler["image"], run["argv"], run["env"]), + "SGLang TP8 engine on the round-0 MXFP4 weights (the yeto engine sampler " + "island drives it)")) + head = trainers[0] if trainers else None + for rank, node in enumerate(trainers): + ip = rendezvous.get(node.name) + ray = (["ray", "start", "--head", f"--port={RAY_PORT}", f"--node-ip-address={ip}", + f"--num-gpus={node.cards_per_node}", "--block"] if rank == 0 else + ["ray", "start", f"--address={rendezvous.get(head.name)}:{RAY_PORT}", + f"--node-ip-address={ip}", f"--num-gpus={node.cards_per_node}", "--block"]) + steps.append(Step(f"trainer-ray-{node.name}", node.name, + _docker(plan, node, "trainer-ray", out.trainer["image"], ray, out.trainer["env"]), + "Ray node of the Miles trainer island")) + if head is not None: + island = out.trainer["miles"]["island_argv"] + steps.append(Step("trainer-island", head.name, + _docker(plan, head, "trainer", out.trainer["image"], island, out.trainer["env"]), + "yeto Miles island entry, sample role trainer; reads samples from the store")) + steps.append(Step("requant-sidecar", head.name, + _docker(plan, head, "requant", out.trainer["image"], out.requant["loop_argv"], + out.trainer["env"]), + "re-quantize every bf16 release to MXFP4 RTN under the scheme's release dir")) + return steps + + +# ------------------------------------------------------------------ checks + +def check(plan: Stage1Plan, out: Expanded) -> list[str]: + """Contradictions inside the expanded plan. Empty list = consistent.""" + errors: list[str] = [] + model, quant, sampling = plan.section("model"), plan.section("quant"), plan.section("sampling") + sampler_cfg, trainer_cfg = plan.section("sampler"), plan.section("trainer") + trainers, samplers = plan.nodes_of("trainer"), plan.nodes_of("sampler") + miles = _section(trainer_cfg, "miles") + + # Nodes and roles. + t_names, s_names = {n.name for n in trainers}, {n.name for n in samplers} + if t_names & s_names: + errors.append(f"roles: nodes {sorted(t_names & s_names)} are both trainer and sampler") + if unused := sorted({n.name for n in plan.nodes} - t_names - s_names): + errors.append(f"roles: inventory nodes {unused} have no role") + for node in (*trainers, *samplers): + role = "trainer" if node.name in t_names else "sampler" + if node.roles and role not in node.roles: + errors.append(f"{node.name}: the plan gives role {role}, the inventory says {list(node.roles)}") + ip = plan.section("rendezvous").get(node.name) + try: + ipaddress.ip_address(str(ip)) + except ValueError: + errors.append(f"{node.name}: rendezvous address {ip!r} is not an IP address") + if len({(n.device_family, n.card, n.cards_per_node) for n in trainers}) > 1: + errors.append("trainer nodes differ in device family, card or card count") + + # Trainer layout. + layout = out.trainer["miles"]["layout"] + split = [int(x) for x in miles["pp_layers"]] + if len(split) != layout["pp"]: + errors.append(f"miles.pp_layers has {len(split)} stages, PP = {layout['pp']}") + if sum(split) != int(model["num_layers"]): + errors.append(f"miles.pp_layers sum to {sum(split)}, the model has {model['num_layers']} layers") + if layout["pp"] > 2 and len(set(split[1:-1])) > 1: + errors.append(f"miles.pp_layers middle stages differ {split[1:-1]}; Megatron splits them evenly") + if layout["tp"] * layout["pp"] * layout["dp"] * layout["cp"] != out.trainer["cards"]: + errors.append("TP*PP*DP*CP does not equal the trainer card count") + if miles["transformer_impl"] not in ("transformer_engine", "local"): + errors.append(f"miles.transformer_impl {miles['transformer_impl']!r} is not transformer_engine or local") + flat = " ".join(out.trainer["miles"]["megatron_args"]) + if "fp8" in flat or "fp4" in flat: + errors.append("trainer arguments enable FP8/FP4; training is bf16 only (user 10-19, report #1)") + verl = _section(trainer_cfg, "verl") + if int(verl["fsdp_size"]) not in (0, out.trainer["cards"]) and out.trainer["cards"] % int(verl["fsdp_size"]): + errors.append(f"verl.fsdp_size {verl['fsdp_size']} does not divide {out.trainer['cards']} cards") + + # Sampler. + for node in samplers: + if int(sampler_cfg["tp"]) != node.cards_per_node: + errors.append(f"{node.name}: sampler TP {sampler_cfg['tp']} != {node.cards_per_node} cards " + "(stage 1: one TP engine per node)") + 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']}"): + 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") + + # Quantization scheme, names and revisions. + errors += _check_scheme(quant, sampler_cfg) + initial = model["sampler_initial"] + if str(initial["revision"]) not in str(initial["local_dir"]) or str(initial["revision"]) not in str(initial["staged"]): + errors.append(f"sampler_initial revision {initial['revision']} is not in its local_dir/staged names") + bf16 = model["bf16_source"] + short = str(bf16["revision"])[:7] + if short not in str(bf16["local_dir"]) or short not in str(bf16["staged"]) or str(bf16["revision"]) not in str(bf16["origin"]): + errors.append(f"bf16_source revision {bf16['revision']} is not in its origin/local_dir/staged names") + requant = out.requant["argv"] + if requant[requant.index("--store") + 1] != out.store["run"]: + errors.append("re-quantizer --store is not the run store") + if requant[requant.index("--target-dir") + 1] != quant["release_dir"]: + errors.append("re-quantizer --target-dir is not quant.release_dir") + + # Store. + store = out.store["run"] + if not store.startswith("oss://"): + errors.append(f"store {store!r} is not an oss:// URI (stage 1: Alibaba OSS)") + 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 + for key, uri in out.store["staging"].items(): + if bucket and not str(uri).startswith(f"oss://{bucket}/"): + errors.append(f"staging path of {key} is not in the run's bucket {bucket}") + + # Sampling: one value for every source, neutral penalties. + for key, neutral in NEUTRAL_PENALTIES.items(): + if float(sampling.get(key, neutral)) != neutral: + errors.append(f"sampling.{key} = {sampling[key]}; rollouts use the neutral value {neutral} (A-INF §5)") + island = out.trainer["miles"]["island_argv"] + if island[island.index("--rollout-max-response-len") + 1] != str(int(sampling["max_tokens"])): + errors.append("trainer --rollout-max-response-len differs from sampling.max_tokens") + margs = out.trainer["miles"]["megatron_args"] + if margs[margs.index("--rollout-max-response-len") + 1] != str(int(sampling["max_tokens"])): + errors.append("Megatron --rollout-max-response-len differs from sampling.max_tokens") + verl_max = [o for o in out.trainer["verl"]["overrides"] if o.startswith("data.max_response_length=")] + if verl_max != [f"data.max_response_length={int(sampling['max_tokens'])}"]: + errors.append("verl data.max_response_length differs from sampling.max_tokens") + + # Static memory estimate must leave room for activations (A-TRN §2.3). + mem = out.trainer["miles"]["static_memory_estimate"] + hbm = trainers[0].hbm_gib if trainers else 0 + on_card = mem["weights_grads_gib"] + (0 if float(miles.get("optimizer_offload_fraction") or 0) >= 1 + else mem["optimizer_gib_per_card"]) + if on_card > 0.75 * hbm: + errors.append(f"static memory estimate {on_card:.1f} GiB/card leaves under 25 % of {hbm:g} GiB " + "for activations") + fraction = float(miles.get("optimizer_offload_fraction") or 0) + if trainers and fraction > 0: + host = mem["optimizer_gib_per_card"] * fraction * trainers[0].cards_per_node + if host > HOST_USABLE_FRACTION * trainers[0].host_memory_gib: + errors.append(f"offloaded optimizer estimate {host:.0f} GiB/node exceeds " + f"{HOST_USABLE_FRACTION:.0%} of {trainers[0].host_memory_gib:g} GiB host memory") + return errors + + +def _check_scheme(quant: Mapping[str, Any], sampler_cfg: Mapping[str, Any]) -> list[str]: + from yeto.rl.quant.requant import SCHEMES # noqa: PLC0415 + + errors = [] + scheme = SCHEMES.get(str(quant["format"])) + if scheme is None: + return [f"quant.format {quant['format']!r} is not a re-quantizer scheme (known: {sorted(SCHEMES)})"] + if scheme.name != quant["scheme"]: + errors.append(f"quant.format {quant['format']} is scheme {scheme.name}, the plan says {quant['scheme']}") + if quant["release_dir"] != RELEASE_DIR_PREFIX + scheme.name: + errors.append(f"quant.release_dir {quant['release_dir']!r} is not {RELEASE_DIR_PREFIX + scheme.name!r}") + if not scheme.name.startswith(str(sampler_cfg["precision"])): + errors.append(f"sampler precision {sampler_cfg['precision']} does not match scheme {scheme.name}") + if scheme.data_suffix.endswith("weight_packed") and not sampler_cfg.get("renames_weight_packed"): + errors.append(f"scheme {scheme.name} stores *.weight_packed; the sglang-server adapter refuses it " + "unless sampler.renames_weight_packed (A-INF item 1)") + return errors + + +def _miles_recipes() -> list[str]: + from yeto.rl.adapters.miles.island_entry import build_parser # noqa: PLC0415 + + for action in build_parser()._actions: # noqa: SLF001 - argparse has no public lookup + if "--rl-model-recipe" in action.option_strings: + return sorted(action.choices or ()) + return [] + + +def gaps(plan: Stage1Plan, out: Expanded) -> list[str]: + """What the run needs that does not exist yet (each one blocks ``run``).""" + found: list[str] = [] + if not plan.section("images").get("trainer"): + found.append("images.trainer is unset: the PPU trainer image is task N0.2") + if not plan.section("images").get("docker_device_args"): + found.append("images.docker_device_args is empty: the PPU device flags of docker run are N0.2b") + for module, what in (("yeto.rl.engine.engine_sampler", "sampler island entry (N2 wiring)"),): + if gap := _module_gap(module, what): + found.append(gap) + if out.trainer["backend"] == "miles": + if gap := _check_island_argv(out.trainer["miles"]["island_argv"]): + found.append(gap) + recipes = _miles_recipes() + if not any("glm5" in r for r in recipes): + found.append(f"Miles island: no glm5_next model recipe passes trainer.miles.megatron_args " + f"(--rl-model-recipe choices: {recipes})") + found.append("head launcher: `yeto launch` does not start an SSH trainer island from the " + "inventory yet (N1.3); the trainer-island step runs the island entry directly") + else: + if not out.trainer["verl"]["commit"]: + found.append("trainer.verl.commit is unset (verl pin)") + found.append("verl backend: no SSH launch path for the verl trainer island and no cut save path (N7.6 verl)") + store = out.store["run"] + if store.startswith("oss://"): + found.append("head on this server cannot reach the intranet-only OSS bucket (PublicEndpointForbidden): " + "needs a CNAME/public endpoint or a relay through ppu1 (runbook N3.5)") + return found + + +# --------------------------------------------------------------------- CLI + +def run_steps(out: Expanded, *, ssh: Sequence[str] = ("ssh", "-o", "BatchMode=yes"), + runner=subprocess.run) -> int: + """Run the steps in order; stop at the first failure.""" + hosts = {n.name: n.host for n in out.plan.nodes} + for step in out.steps: + print(f"[yeto-stage1] step {step.name} on {step.node or 'head'}: {step.purpose}", flush=True) + argv = ["bash", "-c", step.script] if step.node is None else [*ssh, hosts[step.node], "bash", "-s"] + rc = runner(argv, input=None if step.node is None else step.script, text=True).returncode + if rc != 0: + print(f"[yeto-stage1] step {step.name} failed (exit {rc}); later steps not run", file=sys.stderr) + return rc + return 0 + + +def main(argv: list[str] | None = None) -> int: + parser = argparse.ArgumentParser(prog="yeto-stage1", description=__doc__.splitlines()[0]) + sub = parser.add_subparsers(dest="command", required=True) + for name, text in (("plan", "expand and check the plan; no SSH"), + ("run", "run the steps over SSH (refused while errors or gaps remain)")): + p = sub.add_parser(name, help=text) + p.add_argument("plan") + p.add_argument("--hf-config", help="GLM config.json for the verl/FSDP2 memory estimate") + p.add_argument("--backend", choices=BACKENDS, help="override trainer.backend") + sub.choices["plan"].add_argument("--json", action="store_true", help="print the full expansion as JSON") + sub.choices["run"].add_argument("--confirm-machines-free", action="store_true", + help="the machine owner said the nodes are free for this run") + args = parser.parse_args(argv) + try: + plan = load_plan(args.plan) + if args.backend: + doc = dict(plan.doc) + doc["trainer"] = {**plan.section("trainer"), "backend": args.backend} + plan = Stage1Plan(path=plan.path, doc=doc, nodes=plan.nodes) + out = expand(plan, hf_config=args.hf_config) + except Stage1Error as exc: + print(f"[yeto-stage1] plan rejected: {exc}", file=sys.stderr) + return 2 + if args.command == "plan": + if args.json: + print(json.dumps(out.record(), indent=2, sort_keys=False)) + else: + for step in out.steps: + print(f"# --- {step.name} ({step.node or 'head'}): {step.purpose}\n{step.script.rstrip()}\n") + for e in out.errors: + print(f"[yeto-stage1] ERROR {e}", file=sys.stderr) + for g in out.gaps: + print(f"[yeto-stage1] GAP {g}", file=sys.stderr) + return 1 if out.errors else 0 + if out.errors or out.gaps: + print(f"[yeto-stage1] run refused: {len(out.errors)} error(s), {len(out.gaps)} gap(s); " + "see `yeto-stage1 plan`", file=sys.stderr) + return 2 + if not args.confirm_machines_free: + print("[yeto-stage1] run refused: pass --confirm-machines-free after the machine owner " + "said the nodes are free", file=sys.stderr) + return 2 + return run_steps(out) + + +if __name__ == "__main__": + raise SystemExit(main()) From 6ee1e99eedea48cdbb96bcfb42c177c312039ef7 Mon Sep 17 00:00:00 2001 From: michael <78113554+chowhsu@users.noreply.github.com> Date: Mon, 19 Oct 2026 12:28:56 +0000 Subject: [PATCH 22/34] Keep ABORTED placeholders out of verl GRPO group baselines (S25 N10.3b) verl cannot drop one sample, so HarborAgentLoop returns an ABORTED trajectory as a one-token placeholder with mask 0 and reward 0, and verl's GRPO counted it in its prompt group's mean and std. The new yeto_grpo estimator, registered through verl's estimator registry, moves every row without a loss token into a group of its own and then calls verl's GRPO function unchanged: the placeholder's advantage is 0 and its siblings' baseline is theirs alone. verl is not modified. Miles keeps its check_no_aborted filter (drops the group, samples a new one); both backends now hold the same rule that an ABORTED sample never enters a group mean. The test compiles verl's own GRPO function from the pinned fork source. Co-Authored-By: Claude Opus 5.5 --- tests/test_verl_yeto_grpo.py | 95 +++++++++++++++++++++ yeto/rl/adapters/verl/advantage.py | 49 +++++++++++ yeto/rl/adapters/verl/fully_async_runner.py | 1 + yeto/rl/adapters/verl/run_config_rules.py | 5 ++ yeto/rl/adapters/verl/trainer.py | 2 + yeto/rl/agentic/groups.py | 34 ++++++++ 6 files changed, 186 insertions(+) create mode 100644 tests/test_verl_yeto_grpo.py create mode 100644 yeto/rl/adapters/verl/advantage.py create mode 100644 yeto/rl/agentic/groups.py diff --git a/tests/test_verl_yeto_grpo.py b/tests/test_verl_yeto_grpo.py new file mode 100644 index 000000000..12c0ab5c3 --- /dev/null +++ b/tests/test_verl_yeto_grpo.py @@ -0,0 +1,95 @@ +"""S25 N10.3b: the ``yeto_grpo`` estimator keeps ABORTED placeholders out of GRPO groups (CPU). + +verl is not installed on the test machines, so ``verl.trainer.ppo.core_algos`` +is a stub whose ``compute_grpo_outcome_advantage`` is verl's own function, +compiled from the pinned fork's source (``verl-fork-acad9875``; skip if the +checkout is absent). +""" + +from __future__ import annotations + +import ast +import importlib +import sys +import types +from collections import defaultdict +from pathlib import Path + +import pytest + +torch = pytest.importorskip("torch") +np = pytest.importorskip("numpy") + +from yeto.rl.agentic.groups import ISOLATED, isolate_rows # noqa: E402 + +VERL_CORE_ALGOS = Path("/home/michael/work/verl-fork-acad9875/verl/trainer/ppo/core_algos.py") + + +def _verl_grpo(): + if not VERL_CORE_ALGOS.exists(): + pytest.skip(f"verl fork checkout not found: {VERL_CORE_ALGOS}") + tree = ast.parse(VERL_CORE_ALGOS.read_text()) + (fn,) = [n for n in tree.body if isinstance(n, ast.FunctionDef) and n.name == "compute_grpo_outcome_advantage"] + fn.decorator_list = [] + namespace = {"torch": torch, "np": np, "defaultdict": defaultdict, "Optional": __import__("typing").Optional, + "AlgoConfig": object} + exec(compile(ast.Module(body=[fn], type_ignores=[]), str(VERL_CORE_ALGOS), "exec"), namespace) + return namespace["compute_grpo_outcome_advantage"] + + +@pytest.fixture +def advantage(monkeypatch): + registry = {} + core = types.ModuleType("verl.trainer.ppo.core_algos") + core.compute_grpo_outcome_advantage = _verl_grpo() + + def register_adv_est(name): + def deco(fn): + registry[name] = fn + return fn + return deco + + core.register_adv_est, core.registry = register_adv_est, registry + for name in ("verl", "verl.trainer", "verl.trainer.ppo"): + monkeypatch.setitem(sys.modules, name, types.ModuleType(name)) + sys.modules["verl.trainer.ppo"].core_algos = core + monkeypatch.setitem(sys.modules, "verl.trainer.ppo.core_algos", core) + monkeypatch.delitem(sys.modules, "yeto.rl.adapters.verl.advantage", raising=False) + module = importlib.import_module("yeto.rl.adapters.verl.advantage") + yield module, core + sys.modules.pop("yeto.rl.adapters.verl.advantage", None) + + +def test_isolate_rows_moves_only_the_excluded_rows(): + assert isolate_rows(["a", "a", "b"], [False, True, False]) == ["a", (ISOLATED, 1), "b"] + with pytest.raises(ValueError): + isolate_rows(["a"], []) + + +def test_aborted_placeholder_does_not_move_its_siblings_baseline(advantage): + module, core = advantage + assert core.registry["yeto_grpo"] is module.compute_yeto_grpo_advantage + # group "p": rewards 1, 0, 1 plus an ABORTED placeholder (reward 0, mask all 0); group "q": 1, 0 + rewards = torch.tensor([[0.0, 1.0], [0.0, 0.0], [1.0, 0.0], [0.0, 0.0], [1.0, 0.0], [0.0, 0.0]]) + mask = torch.tensor([[1, 1], [1, 1], [1, 0], [0, 0], [1, 1], [1, 1]], dtype=torch.float32) + index = np.array(["p", "p", "p", "p", "q", "q"], dtype=object) + config = {"norm_adv_by_std_in_grpo": True} + + ours, _ = module.compute_yeto_grpo_advantage(rewards, mask, index, config=config) + # reference: verl GRPO on the batch without the placeholder row + keep = [0, 1, 2, 4, 5] + ref, _ = core.compute_grpo_outcome_advantage(rewards[keep], mask[keep], index[keep]) + assert torch.allclose(ours[keep], ref) + assert torch.equal(ours[3], torch.zeros(2)) # the placeholder trains nothing + # plain verl GRPO would count the placeholder: group "p" mean 0.5 instead of 2/3 + plain, _ = core.compute_grpo_outcome_advantage(rewards, mask, index) + assert not torch.allclose(plain[keep[:3]], ref[:3]) + + +def test_std_normalization_follows_the_config(advantage): + module, core = advantage + rewards = torch.tensor([[1.0], [0.0], [0.0]]) + mask = torch.ones(3, 1) + index = np.array(["p", "p", "p"], dtype=object) + ours, _ = module.compute_yeto_grpo_advantage(rewards, mask, index, config={"norm_adv_by_std_in_grpo": False}) + assert torch.allclose(ours[:, 0], torch.tensor([2 / 3, -1 / 3, -1 / 3])) diff --git a/yeto/rl/adapters/verl/advantage.py b/yeto/rl/adapters/verl/advantage.py new file mode 100644 index 000000000..2fac5523a --- /dev/null +++ b/yeto/rl/adapters/verl/advantage.py @@ -0,0 +1,49 @@ +"""The ``yeto_grpo`` advantage estimator: verl GRPO without placeholder rows (S25 N10.3b). + +verl cannot drop one sample of a batch, so ``HarborAgentLoop`` returns an +ABORTED trajectory as a placeholder: one token with loss mask 0 and reward 0. +verl's GRPO would still count it in its prompt group's mean and std. This +estimator moves every row without a loss token to a group of its own +(:func:`yeto.rl.agentic.groups.isolate_rows`), then calls verl's own GRPO +function unchanged. A one-row group has mean 0 and std 1 in verl, so the +placeholder's advantage is ``0 * mask = 0`` and its siblings' baseline is +theirs alone. verl is not modified: the estimator goes through verl's +registry (``core_algos.register_adv_est``) and is selected with +``algorithm.adv_estimator=yeto_grpo``. + +"No loss token" is the test, not the ABORTED flag: verl passes the estimator +only rewards, mask and group ids. A row whose whole response is masked +(an ABORTED placeholder, or a sequence a rollout correction rejected in full) +trains nothing, so leaving its reward out of the baseline is correct for it too. + +verl's multi-trajectory GRPO path (``compute_advantage_for_multi_trajectories``) +only runs for the built-in ``grpo`` name; ``HarborAgentLoop`` returns one +output per sample, so that path is not needed here. +""" + +from __future__ import annotations + +from typing import Any + +import numpy as np +from verl.trainer.ppo import core_algos + +from yeto.rl.agentic.groups import isolate_rows + +from .run_config_rules import AGENTIC_ADV_ESTIMATOR as YETO_GRPO + + +@core_algos.register_adv_est(YETO_GRPO) +def compute_yeto_grpo_advantage(token_level_rewards: Any, response_mask: Any, index: np.ndarray, + epsilon: float = 1e-6, norm_adv_by_std_in_grpo: bool | None = None, + config: Any = None, **_: Any) -> tuple[Any, Any]: + """verl GRPO over ``index`` with every all-masked row isolated (module doc).""" + if norm_adv_by_std_in_grpo is None: # the registry path passes only the config + norm_adv_by_std_in_grpo = True if config is None else bool(config.get("norm_adv_by_std_in_grpo", True)) + empty = (response_mask.sum(dim=-1) == 0).tolist() + groups = np.empty(len(empty), dtype=object) # filled one by one: numpy would unpack the tuples + for row, group in enumerate(isolate_rows(list(index), empty)): + groups[row] = group + return core_algos.compute_grpo_outcome_advantage( + token_level_rewards=token_level_rewards, response_mask=response_mask, index=groups, + epsilon=epsilon, norm_adv_by_std_in_grpo=norm_adv_by_std_in_grpo, config=config) diff --git a/yeto/rl/adapters/verl/fully_async_runner.py b/yeto/rl/adapters/verl/fully_async_runner.py index 2b414f522..e43c15b2d 100644 --- a/yeto/rl/adapters/verl/fully_async_runner.py +++ b/yeto/rl/adapters/verl/fully_async_runner.py @@ -30,6 +30,7 @@ from verl.experimental.fully_async_policy import fully_async_main as fa_main from verl.experimental.fully_async_policy.fully_async_trainer import FullyAsyncTrainer +from . import advantage # noqa: F401 - registers the yeto_grpo estimator in the trainer process from .fully_async_round import RESUME_CALLS_KEY _TrainerBase = getattr(FullyAsyncTrainer, "__ray_actor_class__", FullyAsyncTrainer) diff --git a/yeto/rl/adapters/verl/run_config_rules.py b/yeto/rl/adapters/verl/run_config_rules.py index a785decc9..e402126df 100644 --- a/yeto/rl/adapters/verl/run_config_rules.py +++ b/yeto/rl/adapters/verl/run_config_rules.py @@ -39,6 +39,11 @@ def check_over_sampling(rollout_batch_size: int, over_sampling_batch_size: int) f"{over_sampling_batch_size} 必须等于 --rollout-batch-size {rollout_batch_size}") +# S25 N10.3b: the advantage estimator the verl agentic path needs +# (``algorithm.adv_estimator``; registered by ``.advantage`` in the trainer process). +AGENTIC_ADV_ESTIMATOR = "yeto_grpo" + + # S25 N10.4: the Miles agentic wiring has no verl spelling. The verl agentic # path is the yeto ``HarborAgentLoop`` (agent loop ``yeto_harbor``), selected per # data row, not these flags; passing them used to be ignored without a word. diff --git a/yeto/rl/adapters/verl/trainer.py b/yeto/rl/adapters/verl/trainer.py index f0ddda693..72e7859e0 100644 --- a/yeto/rl/adapters/verl/trainer.py +++ b/yeto/rl/adapters/verl/trainer.py @@ -16,6 +16,8 @@ from verl.trainer.ppo.v1.trainer_base import register_trainer from verl.trainer.ppo.v1.trainer_sync import PPOTrainerSync +from . import advantage # noqa: F401 - registers the yeto_grpo estimator in the trainer process + @register_trainer("yeto_sync") class YetoSyncTrainer(PPOTrainerSync): diff --git a/yeto/rl/agentic/groups.py b/yeto/rl/agentic/groups.py new file mode 100644 index 000000000..d601281f3 --- /dev/null +++ b/yeto/rl/agentic/groups.py @@ -0,0 +1,34 @@ +"""Keep ABORTED samples out of group baselines (S25 N10.3b). + +An ABORTED trajectory (infrastructure failure, misaligned tokens) is never +trained: it carries no loss tokens and reward 0. A group-relative estimator +(GRPO) still averages every row of a prompt group, so a placeholder row +would pull its siblings' baseline toward 0 and inflate their advantages. + +The backends drop it differently, with the same invariant: an ABORTED +sample never enters a group mean. + +* Miles: the dynamic-sampling filter ``check_no_aborted`` drops a group that + holds an ABORTED sample, and Miles samples a new group in its place + (``harbor_agent.train_args``). +* verl: a batch has no per-sample drop, so the placeholder stays in the batch + and :func:`isolate_rows` gives it a group of its own before the advantage + step (``adapters.verl.advantage``). Its siblings keep their group. +""" + +from __future__ import annotations + +from collections.abc import Hashable, Sequence + +ISOLATED = "yeto-isolated" + + +def isolate_rows(group_ids: Sequence[Hashable], excluded: Sequence[bool]) -> list[Hashable]: + """``group_ids`` with each excluded row moved to a one-row group of its own. + + The new id ``(ISOLATED, row)`` cannot equal an original id (verl uses + strings), so no other row joins it. + """ + if len(group_ids) != len(excluded): + raise ValueError(f"{len(group_ids)} group ids but {len(excluded)} exclusion flags") + return [(ISOLATED, row) if out else gid for row, (gid, out) in enumerate(zip(group_ids, excluded))] From 80ecd8d3b1003a0a286c7ab95faeb123d685e8b9 Mon Sep 17 00:00:00 2001 From: michael <78113554+chowhsu@users.noreply.github.com> Date: Mon, 19 Oct 2026 12:28:56 +0000 Subject: [PATCH 23/34] Share one records-proxy recorder per event loop in verl's Harbor loop (S25 N10.5) verl builds a new agent loop for every sample, and HarborAgentLoop called the YETO_AGENTIC_RECORDER factory in each constructor. The loop now keeps the recorder per event loop (LoopLocal) and calls the factory once per loop. ProxyTokenRecorder is the TokenRecorder over the N10.5 records proxy: it serves the proxy over HTTP in the worker, refuses trainer sampling that differs from the proxy's signed sampling, turns the recorded segments into TokenTurns with policy version and source, and forgets the trajectory after close. recorder_from_env builds it from YETO_RECORDS_PROXY. The loop also opens one recorder session per sample (verl repeats a row n times with the same trajectory id), refuses an estimator other than yeto_grpo, puts the sample's sampler source into extra_fields, turns an invalidated recording into ABORTED, and no longer lets a failing recorder close hide the trial's own error. Co-Authored-By: Claude Opus 5.5 --- tests/test_agentic_harbor.py | 6 +- tests/test_verl_records_recorder.py | 199 +++++++++++++++++++++ yeto/rl/adapters/verl/harbor_agent_loop.py | 100 +++++++++-- yeto/rl/agentic/types.py | 9 +- yeto/rl/harness/gateway/core.py | 5 + yeto/rl/harness/gateway/proxy.py | 5 + yeto/rl/harness/gateway/recorder.py | 163 +++++++++++++++++ yeto/rl/harness/gateway/records.py | 7 + 8 files changed, 475 insertions(+), 19 deletions(-) create mode 100644 tests/test_verl_records_recorder.py create mode 100644 yeto/rl/harness/gateway/recorder.py diff --git a/tests/test_agentic_harbor.py b/tests/test_agentic_harbor.py index c3178396b..e3039cf71 100644 --- a/tests/test_agentic_harbor.py +++ b/tests/test_agentic_harbor.py @@ -430,8 +430,10 @@ def test_verl_loop_token_mask_alignment_with_fake_trial(verl_loop): loop = _make_loop(loop_mod, recorder, harness) out = _run(loop.run({"temperature": 1.0, "top_p": 1.0}, **_ROW)) - assert recorder.opened == [("traj-1", {"temperature": 1.0, "top_p": 1.0})] and recorder.closed == ["traj-1"] - assert _FakeTrial.configs == [{"url": "http://recorder/traj-1/v1"}] # agent talks to the recorder + ((session, sampling),) = recorder.opened + assert session.startswith("traj-1-") and sampling == {"temperature": 1.0, "top_p": 1.0} + assert recorder.closed == [session] # one session per sample: verl repeats a row n times + assert _FakeTrial.configs == [{"url": f"http://recorder/{session}/v1"}] # agent talks to the recorder assert out.prompt_ids == [5, 6, 7] assert out.response_ids == [40, 41, 42, 90, 91, 43, 2] assert out.response_mask == [1, 1, 1, 0, 0, 1, 1] diff --git a/tests/test_verl_records_recorder.py b/tests/test_verl_records_recorder.py new file mode 100644 index 000000000..0d45d96ac --- /dev/null +++ b/tests/test_verl_records_recorder.py @@ -0,0 +1,199 @@ +"""S25 N10.5 on verl: HarborAgentLoop with the records proxy as its token recorder (CPU). + +Real pieces: ``HarborAgentLoop.run``, the outer layer, ``ProxyTokenRecorder`` +serving a ``RecordsProxy`` over HTTP on 127.0.0.1, the gateway and the +recording backend. Fakes: verl's agent-loop module (stub, as in +``test_agentic_harbor``), the engine (scripted ``FakeEngine``) and the agent +(a harness that calls the recorder URL over OpenAI Responses, as Harbor's +codex agent would). Tokenizer: Qwen3-0.6B from the local HF cache (skip if absent). +""" + +from __future__ import annotations + +import asyncio +import copy +import sys +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, + echo, + engine_for, + responses_body, + tool_output, + user_item, +) + +from yeto.rl.agentic import OUTCOME_KEY, HarnessConfigError, TrialVerdict +from yeto.rl.harness.gateway.proxy import RecordsProxy +from yeto.rl.harness.gateway.recorder import ProxyTokenRecorder, recorder_from_env, turns_from_records + +pytest.importorskip("aiohttp") + +verl_loop = agentic.verl_loop # the verl agent-loop stub fixture + +ROW = {"extra_info": {"agent_task": {"task_id": "fix-git", "agent": "codex", "provider": "boat", + "trajectory_id": "traj-1"}}} +VERL_SAMPLING = {"temperature": 1.0, "top_p": 1.0, "top_k": -1, "repetition_penalty": 1.0, "logprobs": True} + + +@pytest.fixture(scope="module") +def qwen(): + return _template("qwen3") + + +class CodexLikeAgent: + """Harness stand-in: a Codex-like Responses loop against the recorder URL (as Harbor's codex agent).""" + + def __init__(self, turns: int = 3) -> None: + self.turns = turns + self.endpoints: list[str] = [] + + async def run(self, task: Any, endpoint: str) -> TrialVerdict: + import aiohttp + + self.endpoints.append(endpoint) + items = [user_item("read a.txt")] + async with aiohttp.ClientSession() as http: + for turn in range(self.turns): + async with http.post(f"{endpoint}/responses", json=responses_body(copy.deepcopy(items))) as reply: + assert reply.status == 200, await reply.text() + 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}) + + +def _proxy(template, script=SCRIPT) -> RecordsProxy: + return RecordsProxy(engine_for(ScriptedModel(template, copy.deepcopy(script))), template, + model="policy", signed_sampling=SIGNED, max_chains=1) + + +def test_recorder_feeds_the_verl_loop_the_tokens_the_engine_sampled(verl_loop, qwen): + loop_mod, _ = verl_loop + proxy = _proxy(qwen) + recorder = ProxyTokenRecorder(proxy) + agent = CodexLikeAgent() + loop = _make_loop(loop_mod, recorder, agent, prompt_length=4096, response_length=4096) + + async def go(): + try: + out = await loop.run(VERL_SAMPLING, **ROW) + finally: + await recorder.aclose() + return out + + out = 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 endpoint.startswith("http://127.0.0.1:") and "/t/traj-1-" in endpoint + # prompt + response is the engine's last prompt + its completion: trained == sampled + engine = proxy.backend.engine + last_prompt = list(engine.token_requests[-1][0][0]) + assert out.prompt_ids + out.response_ids[:len(last_prompt) - len(out.prompt_ids)] == last_prompt + assert len(out.response_ids) == len(out.response_mask) == len(out.response_logprobs) + 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 proxy.backend.sessions == {} and proxy.gateway.registries == {} # forgotten after close + + +def test_one_recorder_per_event_loop_from_the_env_factory(verl_loop, qwen, monkeypatch): + """verl builds one agent loop per sample; the factory runs once per event loop.""" + loop_mod, _ = verl_loop + built: list[ProxyTokenRecorder] = [] + module = type(sys)("t_n2w_recorder_factory") + + def make() -> ProxyTokenRecorder: + built.append(recorder_from_env()) + return built[-1] + + module.make, module.proxy = make, lambda: _proxy(qwen, SCRIPT * 2) + monkeypatch.setitem(sys.modules, module.__name__, module) + monkeypatch.setenv(loop_mod.RECORDER_ENV, f"{module.__name__}:make") + monkeypatch.setenv("YETO_RECORDS_PROXY", f"{module.__name__}:proxy") + loop_mod._recorders.clear() + agent = CodexLikeAgent() + cls = type("Loop", (loop_mod.HarborAgentLoop,), {"harness_factory": staticmethod(lambda: agent)}) + rollout = agentic.SimpleNamespace(prompt_length=4096, response_length=4096) + cfg = agentic.SimpleNamespace(actor_rollout_ref=agentic.SimpleNamespace(rollout=rollout)) + tok = agentic.SimpleNamespace(pad_token_id=0, eos_token_id=2) + + async def two_samples(): + loops = [cls(cfg, None, tok, None, None, None) for _ in range(2)] # one per sample, as verl does + try: + return await asyncio.gather(*(loop.run(VERL_SAMPLING, **ROW) for loop in loops)) + finally: + for recorder in built: + await recorder.aclose() + + outs = asyncio.run(two_samples()) + assert [o.extra_fields[OUTCOME_KEY] for o in outs] == ["ok", "ok"] + assert len(built) == 1 # shared: one proxy, one HTTP server for the loop + assert len(set(agent.endpoints)) == 2 # same task row, two recorder sessions + loop_mod._recorders.clear() + + +def test_recorder_refuses_other_sampling_and_unknown_sessions(qwen): + recorder = ProxyTokenRecorder(_proxy(qwen)) + + async def go(): + with pytest.raises(HarnessConfigError, match="temperature"): + await recorder.open("s1", {**VERL_SAMPLING, "temperature": 0.7}) + with pytest.raises(HarnessConfigError, match="not open"): + await recorder.close("s2") + await recorder.open("s3", VERL_SAMPLING) + with pytest.raises(HarnessConfigError, match="already open"): + await recorder.open("s3", VERL_SAMPLING) + assert await recorder.close("s3") == [] # the agent never called: no turns + await recorder.aclose() + + asyncio.run(go()) + + +def test_estimator_other_than_yeto_grpo_is_refused(verl_loop): + loop_mod, _ = verl_loop + for estimator, ok in (("grpo", False), ("yeto_grpo", True)): + cfg = agentic.SimpleNamespace(algorithm=agentic.SimpleNamespace(adv_estimator=estimator)) + if ok: + loop_mod.check_estimator(cfg) + else: + with pytest.raises(HarnessConfigError, match="yeto_grpo"): + loop_mod.check_estimator(cfg) + loop_mod.check_estimator(agentic.SimpleNamespace()) # stub config: not checked + + +def test_turns_from_records_refuses_a_stream_no_turn_accounts_for(): + sample = {"segment": 0, "token_ids": [1, 2, 3, 4], "loss_mask": [0, 1, 0, 0], + "turns": [{"turn": 0, "prompt_added": [1], "completion": [2], "logprobs": [-0.1], + "stop_added": []}]} + from yeto.rl.agentic import TokenAlignmentError + + with pytest.raises(TokenAlignmentError, match="no turn accounts for"): + turns_from_records({"sessions": [{"samples": [sample]}]}) + sample = {**sample, "token_ids": [1, 2], "loss_mask": [0, 1]} + (turn,) = turns_from_records({"sessions": [{"samples": [sample]}]}) + assert (turn.prompt_ids, turn.completion_ids, turn.completion_logprobs) == ((1,), (2,), (-0.1,)) + + +def test_a_trial_error_is_not_hidden_by_a_failing_recorder_close(verl_loop): + loop_mod, _ = verl_loop + + class BrokenClose(agentic._Recorder): + async def close(self, trajectory_id): + raise RuntimeError("recorder gone") + + class Crash: + async def run(self, task, endpoint): + raise KeyboardInterrupt # not an Exception: the outer layer lets it through + + loop = _make_loop(loop_mod, BrokenClose([]), Crash()) + with pytest.raises(KeyboardInterrupt): + asyncio.run(loop.run({}, **ROW)) diff --git a/yeto/rl/adapters/verl/harbor_agent_loop.py b/yeto/rl/adapters/verl/harbor_agent_loop.py index ec110a462..9093839da 100644 --- a/yeto/rl/adapters/verl/harbor_agent_loop.py +++ b/yeto/rl/adapters/verl/harbor_agent_loop.py @@ -10,7 +10,8 @@ 1. reads the :class:`~yeto.rl.agentic.AgentTask` from the row (``extra_info.agent_task``); 2. opens a token-recorder session (N10.5) with verl's sampling parameters and - hands its OpenAI URL to the agent; + hands its OpenAI URL to the agent (the session id is the task's + ``trajectory_id`` plus a per-sample suffix: verl repeats a row ``n`` times); 3. runs one Harbor ``Trial`` through ``run_agent_task``; 4. joins the recorded calls into prompt / response / mask / logprobs (``join_turns``: a call that does not extend the stream is ABORTED, so the @@ -18,14 +19,17 @@ ABORTED on verl: the batch has no per-sample drop, so an ABORTED sample is a one-token placeholder with mask 0 and reward 0, flagged in ``extra_fields``. -It adds no policy-gradient tokens, but it still counts in its GRPO group's -reward mean; a group-level drop on verl is not implemented (not verified). +It adds no policy-gradient tokens. The ``yeto_grpo`` estimator +(:mod:`.advantage`) keeps it out of its GRPO group's mean and std (N10.3b); +the loop refuses to start under another estimator. Configuration (worker environment; pass it to the Ray workers, e.g. through ``ray_kwargs.ray_init.runtime_env.env_vars``): - YETO_AGENTIC_RECORDER ``module:callable`` returning a TokenRecorder (required). - verl builds the agent loop once per sample, so the - callable is called per sample; return a shared object. + YETO_AGENTIC_RECORDER ``module:callable`` returning a TokenRecorder (required), + e.g. ``yeto.rl.harness.gateway.recorder:recorder_from_env``. + verl builds the agent loop once per sample; the loop + calls the callable once per event loop and shares + the recorder (``LoopLocal``). HARBOR_TASKS_DIR Harbor task directories (required) HARBOR_TRIALS_DIR trial output (default /tmp/harbor_trials) HARBOR_ENV_TYPE sandbox provider of rows without ``agent_task.provider`` @@ -43,7 +47,7 @@ ``YETO_BOAT_MAX_ACTIVE``; keeping ``num_workers * YETO_BOAT_MAX_ACTIVE`` within the plan (box_20: 100 active, 12 starts per minute) is the launcher's job. -The admission gate and the boat quota hold asyncio primitives and are kept +The admission gate, the recorder and the boat quota hold asyncio state and are kept per event loop (:class:`~yeto.rl.agentic.loop_local.LoopLocal`). verl runs ``AgentLoopWorker`` as a Ray async actor (one loop per process), so in a worker they are per process. @@ -52,8 +56,10 @@ from __future__ import annotations import importlib +import logging import os import uuid +from collections.abc import Mapping, Sequence from functools import partial from typing import Any @@ -74,9 +80,14 @@ run_agent_task, ) from yeto.rl.agentic.loop_local import LoopLocal +from yeto.rl.agentic.types import TokenTurn from yeto.rl.harness.harbor import HarborHarness, HarborSettings, build_trial_config from yeto.rl.harness.harbor.boat_client import check_in_flight_cap +from .run_config_rules import AGENTIC_ADV_ESTIMATOR as YETO_GRPO + +logger = logging.getLogger(__name__) + AGENT_LOOP_NAME = "yeto_harbor" RECORDER_ENV = "YETO_AGENTIC_RECORDER" PROVIDER_ENV = "HARBOR_ENV_TYPE" @@ -104,6 +115,33 @@ def _recorder_from_env() -> TokenRecorder: return getattr(importlib.import_module(module), attr)() +_recorders: LoopLocal[TokenRecorder] = LoopLocal(_recorder_from_env) + + +def _shared_recorder() -> TokenRecorder: + """The recorder of this worker's event loop, built once from ``YETO_AGENTIC_RECORDER``.""" + return _recorders.get() + + +def recorder_session_id(task: AgentTask) -> str: + """One recorder session per sample: verl repeats a row ``n`` times with the same task.""" + suffix = uuid.uuid4().hex[:12] + return f"{task.trajectory_id}-{suffix}" if task.trajectory_id else uuid.uuid4().hex + + +def check_estimator(config: Any) -> None: + """Refuse an advantage estimator that would count ABORTED placeholders (N10.3b). + + A config without ``algorithm.adv_estimator`` (a test stub) is not checked. + """ + algorithm = getattr(config, "algorithm", None) + estimator = getattr(algorithm, "adv_estimator", None) + if estimator is not None and str(getattr(estimator, "value", estimator)) != YETO_GRPO: + raise HarnessConfigError( + f"HarborAgentLoop needs algorithm.adv_estimator={YETO_GRPO} (got {estimator}): another " + "estimator counts ABORTED placeholder samples in their group's baseline") + + def _harness_from_env() -> AgentHarness: return HarborHarness(partial(build_trial_config, settings=HarborSettings.from_env())) @@ -139,24 +177,32 @@ class HarborAgentLoop(AgentLoopBase): """One Harbor trial per sample; tokens from the recorder, verdict from Harbor.""" # Factories are class attributes so a test (or another wiring) can replace them. - recorder_factory = staticmethod(_recorder_from_env) + recorder_factory = staticmethod(_shared_recorder) harness_factory = staticmethod(_harness_from_env) def __init__(self, *args: Any, **kwargs: Any) -> None: super().__init__(*args, **kwargs) + check_estimator(self.config) self.prompt_length = int(self.rollout_config.prompt_length) self.response_length = int(self.rollout_config.response_length) - self._recorder = self.recorder_factory() self._harness = self.harness_factory() async def run(self, sampling_params: dict[str, Any], **kwargs: Any) -> AgentLoopOutput: task = task_from_row(kwargs) - trajectory_id = task.trajectory_id or uuid.uuid4().hex - endpoint = await self._recorder.open(trajectory_id, sampling_params) + recorder = self.recorder_factory() + session = recorder_session_id(task) + endpoint = await recorder.open(session, sampling_params) try: trajectory = await run_agent_task(task, self._harness, endpoint, _gate()) - finally: - turns = await self._recorder.close(trajectory_id) + except BaseException: + await _close_after_error(recorder, session) + raise + turns: Sequence[TokenTurn] = () + try: + turns = await recorder.close(session) + except TokenAlignmentError as exc: + if trajectory.outcome is not Outcome.ABORTED: + trajectory = _abort(trajectory, f"token_alignment: {exc}") segments: TokenSegments | None = None if trajectory.outcome is not Outcome.ABORTED: @@ -166,9 +212,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, n_calls=len(turns)) + return self._output(trajectory, segments, turns) - def _output(self, trajectory: AgentTrajectory, segments: TokenSegments | None, *, n_calls: int) -> AgentLoopOutput: + def _output(self, trajectory: AgentTrajectory, segments: TokenSegments | None, + turns: Sequence[TokenTurn]) -> AgentLoopOutput: if segments is None: segments = self._placeholder() metadata = trajectory.to_metadata() @@ -178,7 +225,7 @@ def _output(self, trajectory: AgentTrajectory, segments: TokenSegments | None, * response_mask=list(segments.response_mask), response_logprobs=list(segments.response_logprobs), reward_score=trajectory.reward, - num_turns=2 * n_calls, + num_turns=2 * len(turns), metrics=AgentLoopMetrics(), extra_fields={ OUTCOME_KEY: metadata[OUTCOME_KEY], @@ -186,6 +233,7 @@ def _output(self, trajectory: AgentTrajectory, segments: TokenSegments | None, * "policy_version": trajectory.policy_version, "truncated": segments.truncated, "agent_metrics": metadata["agent_metrics"], + "provenance": turn_provenance(turns), }, ) @@ -199,3 +247,23 @@ 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, verdict=trajectory.verdict, policy_version=trajectory.policy_version) + + +def turn_provenance(turns: Sequence[TokenTurn]) -> Mapping[str, str] | None: + """The sampler source of the sample: the one every call shares, else None. + + Per-source metrics (``offpolicy_source``) group by it; a sample whose calls + came from different sources is counted under ``unknown``. + """ + records = {tuple(sorted(t.provenance.items())) for t in turns if t.provenance} + if len(records) != 1 or any(not t.provenance for t in turns): + return None + return dict(next(iter(records))) + + +async def _close_after_error(recorder: TokenRecorder, session: str) -> None: + """Close the recorder session after the trial raised; the trial's error is the one to report.""" + try: + await recorder.close(session) + except Exception: # the trial's error propagates + logger.warning("recorder close failed for %s after a trial error", session, exc_info=True) diff --git a/yeto/rl/agentic/types.py b/yeto/rl/agentic/types.py index 3290af13e..2bd677890 100644 --- a/yeto/rl/agentic/types.py +++ b/yeto/rl/agentic/types.py @@ -118,11 +118,18 @@ class TrialVerdict: @dataclass(frozen=True) class TokenTurn: - """One model call as the token recorder saw it (prompt rendered by the sampler).""" + """One model call as the token recorder saw it (prompt rendered by the sampler). + + ``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. + """ prompt_ids: tuple[int, ...] completion_ids: tuple[int, ...] completion_logprobs: tuple[float, ...] + policy_version: int | None = None + provenance: Mapping[str, str] | None = None def __post_init__(self) -> None: if len(self.completion_ids) != len(self.completion_logprobs): diff --git a/yeto/rl/harness/gateway/core.py b/yeto/rl/harness/gateway/core.py index 9167f0a06..045fb6ca9 100644 --- a/yeto/rl/harness/gateway/core.py +++ b/yeto/rl/harness/gateway/core.py @@ -167,6 +167,11 @@ def close(self, trajectory_id: str) -> None: for chain in registry.chains: self._board("exit_session", chain.session_id) + def forget(self, trajectory_id: str) -> None: + """Drop the trajectory's chain state (after :meth:`close`; a long-lived gateway would grow).""" + self.registries.pop(trajectory_id, None) + self.invalidated.pop(trajectory_id, None) + def snapshot(self, trajectory_id: str) -> dict[str, Any]: snap = self.registry(trajectory_id).snapshot() snap["invalidated"] = self.invalidated.get(trajectory_id) diff --git a/yeto/rl/harness/gateway/proxy.py b/yeto/rl/harness/gateway/proxy.py index 2b8c832b2..f414fa3ae 100644 --- a/yeto/rl/harness/gateway/proxy.py +++ b/yeto/rl/harness/gateway/proxy.py @@ -75,6 +75,11 @@ def records(self, trajectory_id: str) -> dict[str, Any]: def close(self, trajectory_id: str) -> None: self.gateway.close(trajectory_id) + def forget(self, trajectory_id: str) -> None: + """Drop the trajectory's state in the gateway and the backend (after close and records).""" + self.gateway.forget(trajectory_id) + self.backend.forget(trajectory_id) + # ------------------------------------------------------------------ HTTP def app(self, *, public_url: str = "") -> Any: from aiohttp import web diff --git a/yeto/rl/harness/gateway/recorder.py b/yeto/rl/harness/gateway/recorder.py new file mode 100644 index 000000000..2aadeabcf --- /dev/null +++ b/yeto/rl/harness/gateway/recorder.py @@ -0,0 +1,163 @@ +"""The records proxy as a :class:`~yeto.rl.agentic.TokenRecorder` (S25 N10.5, N10.3). + +verl's ``HarborAgentLoop`` asks a token recorder for one OpenAI base URL per +sample and, after the trial, for the sample's model calls. +:class:`ProxyTokenRecorder` is that recorder in front of a sampler engine: it +serves a :class:`~.proxy.RecordsProxy` over HTTP in the worker process (the +sandboxed agent calls ``/t//v1``) and turns the proxy's records +into :class:`~yeto.rl.agentic.TokenTurn` values. + +One recorder per process, not per sample: verl builds a new agent loop for +every sample, so ``HarborAgentLoop`` keeps the recorder per event loop +(:class:`~yeto.rl.agentic.loop_local.LoopLocal`) and calls the +``YETO_AGENTIC_RECORDER`` callable once per loop. :func:`recorder_from_env` +is that callable for this recorder; it builds the proxy with the deployment's +own ``module:callable`` (the engine and the chat template are deployment +facts, not code here). + +Sampling: the proxy signs its sampling settings once (gateway config). A +session whose trainer sampling parameters differ from them is refused +(:class:`~yeto.rl.agentic.HarnessConfigError`) instead of sampling with other +settings than the trainer assumes. + +Miles does not use this recorder: its Harbor path records through Miles' +own session server (GOAL §9). + +Environment of :func:`recorder_from_env`: + YETO_RECORDS_PROXY ``module:callable`` returning a RecordsProxy (required) + 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) +""" + +from __future__ import annotations + +import asyncio +import importlib +import os +from collections.abc import Mapping, Sequence +from typing import Any + +from yeto.rl.agentic import HarnessConfigError, TokenAlignmentError, TokenTurn + +from .proxy import RecordsProxy +from .records import sampling_params + +PROXY_ENV = "YETO_RECORDS_PROXY" +HOST_ENV = "YETO_RECORDS_PROXY_HOST" +PORT_ENV = "YETO_RECORDS_PROXY_PORT" +PUBLIC_HOST_ENV = "YETO_RECORDS_PROXY_PUBLIC_HOST" + +# verl sampling keys -> SamplingParams fields; other keys (``logprobs``) do not change sampling. +CHECKED_SAMPLING = ("temperature", "top_p", "top_k", "repetition_penalty") + + +def turns_from_records(records: Mapping[str, Any]) -> list[TokenTurn]: + """The proxy records of one trajectory as model calls in call order. + + Each recorded turn's prompt is the segment stream before its completion, + so :func:`~yeto.rl.agentic.join_turns` sees what the engine saw. Calls of + later sessions or compaction segments follow the first ones; they do not + extend its stream, so the join refuses them (one verl sample is one stream). + """ + turns: list[TokenTurn] = [] + for session in records.get("sessions") or (): + for sample in session["samples"]: + ids, mask = list(sample["token_ids"]), list(sample["loss_mask"]) + cursor, pending_stop = 0, 0 + for turn in sample["turns"]: + cursor += pending_stop + len(turn["prompt_added"]) + completion = tuple(int(t) for t in turn["completion"]) + end = cursor + len(completion) + if tuple(ids[cursor:end]) != completion or any(m != 1 for m in mask[cursor:end]): + raise TokenAlignmentError(f"records turn {turn['turn']} does not sit in its segment stream") + 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"))) + 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} " + "tokens no turn accounts for") + return turns + + +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: + self.proxy = proxy + self.host = host + self.port = int(port) + self.public_host = public_host or host + self._signed = sampling_params(proxy.gateway.config.signed_sampling) + self._runner: Any = None + self._base_url: str | None = None + self._start_lock = asyncio.Lock() + self._open: set[str] = set() + + async def _base(self) -> str: + async with self._start_lock: + if self._base_url is None: + from aiohttp import web + + runner = web.AppRunner(self.proxy.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}" + return self._base_url + + def check_sampling(self, sampling: Mapping[str, Any]) -> None: + """Refuse trainer sampling settings the proxy does not sample with.""" + differ = {key: (sampling[key], getattr(self._signed, key)) for key in CHECKED_SAMPLING + if key in sampling and sampling[key] is not None + and float(sampling[key]) != float(getattr(self._signed, key))} + if differ: + raise HarnessConfigError(f"trainer sampling differs from the records proxy's signed sampling " + f"(key: (trainer, proxy)): {differ}") + + async def open(self, trajectory_id: str, sampling_params: Mapping[str, Any]) -> str: + self.check_sampling(sampling_params) + if trajectory_id in self._open: + raise HarnessConfigError(f"recorder session {trajectory_id!r} is already open") + base = await self._base() + self._open.add(trajectory_id) + return f"{base}/t/{trajectory_id}/v1" + + async def close(self, trajectory_id: str) -> Sequence[TokenTurn]: + if trajectory_id not in self._open: + raise HarnessConfigError(f"recorder session {trajectory_id!r} is not open") + self._open.discard(trajectory_id) + try: + records = self.proxy.records(trajectory_id) + self.proxy.close(trajectory_id) + finally: + self.proxy.forget(trajectory_id) + invalid = (records.get("gateway") or {}).get("invalidated") + if invalid: + raise TokenAlignmentError(f"records proxy invalidated the trajectory: {invalid}") + return turns_from_records(records) + + async def aclose(self) -> None: + """Stop the HTTP server (tests; a worker process keeps it for its life).""" + if self._runner is not None: + await self._runner.cleanup() + self._runner, self._base_url = None, None + + +def recorder_from_env() -> ProxyTokenRecorder: + """``YETO_AGENTIC_RECORDER=yeto.rl.harness.gateway.recorder:recorder_from_env`` (module doc).""" + spec = os.environ.get(PROXY_ENV, "") + module, _, attr = spec.partition(":") + if not module or not attr: + raise HarnessConfigError(f"set {PROXY_ENV}=module:callable returning a RecordsProxy") + proxy = getattr(importlib.import_module(module), attr)() + 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") + return ProxyTokenRecorder(proxy, host=host, port=int(os.environ.get(PORT_ENV, "0")), + public_host=os.environ.get(PUBLIC_HOST_ENV) or host) diff --git a/yeto/rl/harness/gateway/records.py b/yeto/rl/harness/gateway/records.py index 39ab8312c..c9b9f95da 100644 --- a/yeto/rl/harness/gateway/records.py +++ b/yeto/rl/harness/gateway/records.py @@ -211,6 +211,13 @@ def records(self, trajectory_id: str) -> dict[str, Any]: "sessions": [{"session_id": sid, "samples": [seg.sample() for seg in s.segments]} for sid, s in sessions], "counters": dict(self.counters.get(trajectory_id, Counter()))} + def forget(self, trajectory_id: str) -> None: + """Drop everything recorded for the trajectory (after its records were read).""" + for session_id in [sid for sid, s in self.sessions.items() if s.trajectory_id == trajectory_id]: + del self.sessions[session_id] + self.tools.pop(trajectory_id, None) + self.counters.pop(trajectory_id, None) + # -------------------------------------------------------------- internals def _prompt(self, session: _Session, segment: Segment, context: list[Message]) -> _Prompt: """What to append before the next sample (the whole prompt on a segment's first turn).""" From 77d0c2bfaef0f04ee36313b9ee3326049b19e067 Mon Sep 17 00:00:00 2001 From: michael <78113554+chowhsu@users.noreply.github.com> Date: Mon, 19 Oct 2026 12:28:59 +0000 Subject: [PATCH 24/34] Let the E2 harness name the replayed cursor; load verl rank files weights-only The E2 harness replays the frozen batch drawn at the cut after arm A moved the live cursor, so rebuild_same_shape takes the resumed cursor from its caller when given (default: the live cursor). verl rank files hold only tensors and plain containers, so torch.load uses weights_only=True. The driver reads its trainer defensively for test doubles. Co-Authored-By: Claude Opus 5.5 --- yeto/rl/adapters/miles/e2_harness.py | 2 ++ yeto/rl/adapters/miles/trainer_rebuild.py | 8 +++++++- yeto/rl/adapters/verl/cut.py | 2 +- yeto/rl/engine/driver.py | 2 +- 4 files changed, 11 insertions(+), 3 deletions(-) diff --git a/yeto/rl/adapters/miles/e2_harness.py b/yeto/rl/adapters/miles/e2_harness.py index 78cda53bf..68b0a626f 100644 --- a/yeto/rl/adapters/miles/e2_harness.py +++ b/yeto/rl/adapters/miles/e2_harness.py @@ -288,6 +288,8 @@ def rebuild() -> Any: trainer, args=ctx.miles_args, rollout_executor=ctx.rollout_executor, actor=ctx.actor, run=ctx.runner.run, restore=restore_with_rejections, rollout=driver.rollout, worker_manager=ctx.worker_manager, rebuild=ctx.rebuild, + # arm B replays the frozen batch B3, drawn at the cut's data position + resumed_cursor=dict(manifest.data), ) if colocated: diff --git a/yeto/rl/adapters/miles/trainer_rebuild.py b/yeto/rl/adapters/miles/trainer_rebuild.py index 8d770c60d..5a101e812 100644 --- a/yeto/rl/adapters/miles/trainer_rebuild.py +++ b/yeto/rl/adapters/miles/trainer_rebuild.py @@ -230,9 +230,15 @@ def rebuild_same_shape( max_attempts: int = 2, critic: SwappableActor | None = None, critic_trainer_id: str = "critic", + resumed_cursor: Mapping[str, int] | None = None, ) -> RebuildResult: """Dispose + rebuild the trainer (same shape) and restore the cut via ``restore()``. + ``resumed_cursor`` is the data position the next batch starts from, compared + with the cut's data watermark after the restore (S25 N7.6); default: the + live data cursor read after the rebuild. Only a caller that replays a batch + drawn at the cut (the E2 harness) passes the cut's cursor. + ``trainer`` is the :class:`~.trainer.MilesTrainerGroup`; its layout is read back from the running ranks before and after. ``restore`` calls ``trainer.restore_cut(...)`` for the cut taken before this call. @@ -284,7 +290,7 @@ def rebuild_same_shape( if now != layout: raise RuntimeError(f"layout changed across a same-shape rebuild: {layout} -> {now}") manifest = restore() - confirm_resumed_cursor(trainer, after) + confirm_resumed_cursor(trainer, after if resumed_cursor is None else resumed_cursor) except Exception as error: # noqa: BLE001 attempts.append({"attempt": attempt, "stage": "after_rebuild", "error": repr(error)}) raise RecoveryRequired(f"trainer unusable after rebuild: {error}", attempts=attempts) from error diff --git a/yeto/rl/adapters/verl/cut.py b/yeto/rl/adapters/verl/cut.py index 785c479e3..5a3e2d2de 100644 --- a/yeto/rl/adapters/verl/cut.py +++ b/yeto/rl/adapters/verl/cut.py @@ -231,7 +231,7 @@ def restore_rank(handles: TrainHandles, *, directory: str, files: Sequence[Mappi raise CutError(f"rank {rank}: {entry.path} checksum differs from the manifest") if any(state for state in handles.optimizer.state.values()): raise CutError("restore_cut needs a freshly built trainer (the optimizer already has state)") - saved = torch.load(path, map_location="cpu", weights_only=False) + saved = torch.load(path, map_location="cpu", weights_only=True) if saved.get("schema") != RANK_FILE_SCHEMA or saved.get("cut_id") != cut_id or saved.get("rank") != rank: raise CutError(f"rank {rank}: {entry.path} is not this rank's file of cut {cut_id!r}") diff --git a/yeto/rl/engine/driver.py b/yeto/rl/engine/driver.py index 0f8d28384..c4ecee2c4 100644 --- a/yeto/rl/engine/driver.py +++ b/yeto/rl/engine/driver.py @@ -1741,7 +1741,7 @@ def _restore_data_cursor(self, start_rollout_id: int) -> None: if landed != cut_cursor: raise DriverError(f"restart at rollout {start_rollout_id}: data source seek " f"to the cut's cursor {cut_cursor} landed on {landed}") - confirm = getattr(self.trainer, "confirm_resumed_data_cursor", None) + confirm = getattr(getattr(self, "trainer", None), "confirm_resumed_data_cursor", None) if callable(confirm): # S25 N7.6: the landed cursor vs the cut's resume manifest confirm(landed) self.emit("rl_data_cursor_restored", rollout_id=start_rollout_id, source="round_cut", From 722d8098ed9a6818363a57a73aacf5880116db5f Mon Sep 17 00:00:00 2001 From: michael <78113554+chowhsu@users.noreply.github.com> Date: Mon, 19 Oct 2026 12:33:51 +0000 Subject: [PATCH 25/34] Move the stage-1 planner out of the launch layer yeto/launch must not import the adapters (import boundaries); the planner builds the SGLang argv with the sglang-server adapter and parses the island argv with the Miles island parser, so it lives in yeto/rl. Co-Authored-By: Claude Opus 5.5 --- pyproject.toml | 2 +- tests/test_stage1_plan.py | 2 +- yeto/{launch => rl}/stage1.py | 6 +++--- 3 files changed, 5 insertions(+), 5 deletions(-) rename yeto/{launch => rl}/stage1.py (99%) diff --git a/pyproject.toml b/pyproject.toml index c556b1277..0f46c88cb 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -62,7 +62,7 @@ yeto-diffusion-export = "yeto.diffusion.export:main" yeto-diffusion-sample = "yeto.diffusion.sample:main" yeto-rl-export = "yeto.rl.export:main" yeto-rl-ssh = "yeto.rl.ssh_harness:main" -yeto-stage1 = "yeto.launch.stage1:main" +yeto-stage1 = "yeto.rl.stage1:main" yeto-protenix-export-batch = "yeto.diffusion.protenix_export:main" yeto-hunyuan3d-export-batch = "yeto.diffusion.hunyuan3d_export:main" diff --git a/tests/test_stage1_plan.py b/tests/test_stage1_plan.py index f30d90b29..c9f29ecc6 100644 --- a/tests/test_stage1_plan.py +++ b/tests/test_stage1_plan.py @@ -9,7 +9,7 @@ import pytest import yaml -from yeto.launch import stage1 +from yeto.rl import stage1 REPO = Path(__file__).resolve().parents[1] PLAN = REPO / "examples" / "stage1" / "glm53-flash-ppu.yaml" diff --git a/yeto/launch/stage1.py b/yeto/rl/stage1.py similarity index 99% rename from yeto/launch/stage1.py rename to yeto/rl/stage1.py index 92ba9d962..5902cb3db 100644 --- a/yeto/launch/stage1.py +++ b/yeto/rl/stage1.py @@ -50,8 +50,8 @@ from pathlib import Path from typing import Any -from .node_inventory import InventoryError, Node, load_inventory -from .topology import island_comm_env +from yeto.launch.node_inventory import InventoryError, Node, load_inventory +from yeto.launch.topology import island_comm_env PLAN_SCHEMA = "yeto.stage1/v1" BACKENDS = ("miles", "verl") @@ -419,7 +419,7 @@ def expand(plan: Stage1Plan, *, hf_config: str | Path | None = None) -> Expanded "sampling": dict(sampling), "engines": sampler_runs} # Re-quantizer (node 0 of the trainer, a side process of the island) --- - from .ssh_trainer_island import requant_argv # noqa: PLC0415 + from yeto.launch.ssh_trainer_island import requant_argv # noqa: PLC0415 argv = [*requant_argv(store, str(model["sampler_initial"]["local_dir"]), python="python3"), "--target-dir", str(quant["release_dir"])] From 3be3981d1b36de3fe1ee97be744e2586285b853b Mon Sep 17 00:00:00 2001 From: michael <78113554+chowhsu@users.noreply.github.com> Date: Mon, 19 Oct 2026 12:34:02 +0000 Subject: [PATCH 26/34] Set the island device family from the node inventory (S25 N2.2) The SSH trainer island now exports YETO_DEVICE_FAMILY from its inventory nodes, so the resource sampler picks the card collector of the node's family (npu-smi, ppu-smi, NVML) instead of defaulting to NVML. The variable is the one T-N77 introduced for yeto.accel (defined in gpu_spec); the telemetry reads it under that name and the old YETO_RL_DEVICE_FAMILY name is gone. A run env that sets another family than the inventory is refused. Co-Authored-By: Claude Opus 5.5 --- tests/test_rl_ssh_trainer_island.py | 9 +++++++++ yeto/gpu_spec.py | 3 +++ yeto/launch/ssh_trainer_island.py | 12 +++++++++--- yeto/rl/engine/telemetry.py | 4 ++-- 4 files changed, 23 insertions(+), 5 deletions(-) diff --git a/tests/test_rl_ssh_trainer_island.py b/tests/test_rl_ssh_trainer_island.py index 74d9f1cd5..34db2f85e 100644 --- a/tests/test_rl_ssh_trainer_island.py +++ b/tests/test_rl_ssh_trainer_island.py @@ -89,12 +89,21 @@ def test_start_script_runs_torchrun_with_rank_rendezvous_and_comm_env(): assert "--nnodes=2" in script and "--master-addr=10.0.0.1" in script and "--nproc-per-node=8" in script assert "export NCCL_SOCKET_IFNAME=eth0" in script and "export NCCL_IB_DISABLE=0" in script assert "export YETO_X=1" in script and "setsid bash -c" in script + assert f"export YETO_DEVICE_FAMILY={s.nodes[0].device_family}" in script assert "--node-rank=0" in script0 and "--node-rank=1" in script1 assert "yeto.rl.quant.release" in script0 and "--loop" in script0 and "--no-awq" in script0 assert "yeto.rl.quant.release" not in script1 assert "'a b'" in script0 # arguments are quoted +def test_env_cannot_override_the_inventory_device_family(): + family = spec().device_family + other = "npu" if family != "npu" else "ppu" + with pytest.raises(SshIslandError, match="inventory decides"): + spec(env={"YETO_DEVICE_FAMILY": other}) + assert spec(env={"YETO_DEVICE_FAMILY": family}).device_family == family + + def test_requant_argv_needs_both_awq_inputs(): assert "--recipe" in requant_argv("s", "t", recipe="r.yaml", awq_scales="d") with pytest.raises(SshIslandError, match="both"): diff --git a/yeto/gpu_spec.py b/yeto/gpu_spec.py index 4a4e78e65..33b5b24bb 100644 --- a/yeto/gpu_spec.py +++ b/yeto/gpu_spec.py @@ -61,6 +61,9 @@ # communication variables. It is not the torch device type: a PPU card has # family ``ppu`` and device type ``cuda``. DEVICE_FAMILIES = ("cuda", "npu", "ppu") +# Environment variable the island launcher sets to the node's device family; +# yeto.accel.device_family reads it (torch cannot tell a PPU from a GPU). +DEVICE_FAMILY_ENV = "YETO_DEVICE_FAMILY" _CARD_DEVICE_FAMILY = {**{card: "npu" for card in _CARD_DEVICE_TYPE}, "810E": "ppu", "M890": "ppu"} diff --git a/yeto/launch/ssh_trainer_island.py b/yeto/launch/ssh_trainer_island.py index 72ebd423d..bfefe7cf8 100644 --- a/yeto/launch/ssh_trainer_island.py +++ b/yeto/launch/ssh_trainer_island.py @@ -16,8 +16,9 @@ 2. each node runs, in a new session (one process group), the trainer command under ``torchrun`` with the incarnation environment (``trainer_remote.island_env``), the collective-communication variables of - its family and fabric (``topology.island_comm_env``) and the run's - environment. Node 0 also runs the W4A8 re-quantizer side process + its family and fabric (``topology.island_comm_env``), the run's + environment and the inventory's device family (``YETO_DEVICE_FAMILY``: + the telemetry's card collector and the learner's device follow it). Node 0 also runs the W4A8 re-quantizer side process (``yeto.rl.quant.release --loop``) when the spec has one: each BF16 release becomes a W4A8 release in the store, which W4A8 samplers follow; 3. the SSH session streams node 0's log to the head (``IslandEvents`` reads the @@ -52,6 +53,7 @@ from yeto.rl.engine.trainer_lifecycle import CutRef from yeto.rl.engine.trainer_remote import island_env +from ..gpu_spec import DEVICE_FAMILY_ENV from .node_inventory import Node from .topology import island_comm_env from .trainer_island import ( @@ -108,6 +110,9 @@ def __post_init__(self) -> None: "(the trainer image's device flags are task N0.2b); use runtime: rootfs") if len({n.cards_per_node for n in self.nodes}) != 1: raise SshIslandError("trainer nodes must have the same number of cards") + if self.env.get(DEVICE_FAMILY_ENV, self.device_family) != self.device_family: + raise SshIslandError(f"env sets {DEVICE_FAMILY_ENV}={self.env[DEVICE_FAMILY_ENV]}, but the " + f"inventory nodes are {self.device_family}; the inventory decides") @property def device_family(self) -> str: @@ -164,7 +169,8 @@ def start_script(spec: SshTrainerSpec, rank: int, env: Mapping[str, str], paths: "set -eu\n" f"mkdir -p \"{paths['dir']}\"\n" f"cd {shlex.quote(spec.workdir)}\n" - + _exports({**comm, **spec.env, **env}) + # The learner reads the family to pick its card type and collective backend (accel). + + _exports({DEVICE_FAMILY_ENV: node.device_family, **comm, **spec.env, **env}) + f"rm -f \"{paths['rc']}\"\n" # Job control is off in a non-interactive shell, so the background setsid # is not a group leader: it does not fork, and $! is the new group's id. diff --git a/yeto/rl/engine/telemetry.py b/yeto/rl/engine/telemetry.py index db6ba30f5..743cfaa70 100644 --- a/yeto/rl/engine/telemetry.py +++ b/yeto/rl/engine/telemetry.py @@ -22,6 +22,7 @@ import time from typing import Any, Callable, Mapping +from yeto.gpu_spec import DEVICE_FAMILY_ENV from yeto.hw.card_telemetry import CardCollector, collector_for HEARTBEAT_EVENT = "rl_heartbeat" @@ -31,7 +32,6 @@ # Per-GPU keys of an available ``rl_resource_sample`` (None = not readable). GPU_SAMPLE_KEYS = ("index", "uuid", "util_pct", "mem_used_mb", "mem_total_mb", "power_w") # The island's launcher device family (cuda / npu / ppu): selects the card collector. -DEVICE_FAMILY_ENV = "YETO_RL_DEVICE_FAMILY" Emit = Callable[..., None] @@ -130,7 +130,7 @@ class ResourceSampler(_Periodic): The card source is a :class:`yeto.hw.card_telemetry.CardCollector` of the island's device family (S25 N2.2): ``collector`` when given, else the one of - ``$YETO_RL_DEVICE_FAMILY`` (``cuda`` when unset: NVML through + ``$YETO_DEVICE_FAMILY`` (``cuda`` when unset: NVML through ``nvml_loader``, as before). A collector that cannot open writes one ``available: false`` record that names the reason, and no thread starts.""" From f87a2cad35c93a05de9692dfb7f689b246419bae Mon Sep 17 00:00:00 2001 From: michael <78113554+chowhsu@users.noreply.github.com> Date: Mon, 19 Oct 2026 12:34:02 +0000 Subject: [PATCH 27/34] Wire engine-sampler payloads and per-source metrics into the trainers (S25 N2.3, N2.5) Readers: an engine that does not return prompt ids (SGLang text input) writes response-only tokens with prompt_tokens_known=False. The Miles buffer reader now rebuilds the prompt ids with the trainer's tokenizer (complete_prompt_tokens, loaded on first need) and records that the trainer tokenized them; without a tokenizer such a sample is refused instead of training the response as the prompt. The source record stays in the sample metadata. offpolicy_source: verl emits the event per step from the trainer batch (old and rollout log-probs, response mask, per-row provenance from the HarborAgentLoop extra field; rows without it group under unknown), in both the sync and the fully async trainer. Miles' training process does not hand its log-probs back, so the Miles reader reports the per-source sample, truncation and loop-abort shares as rollout metrics, without ratio fields. Co-Authored-By: Claude Opus 5.5 --- tests/test_rl_n2_trainer_wiring.py | 118 ++++++++++++++++++ .../rl/adapters/miles/buffer_rollout_miles.py | 21 +++- yeto/rl/adapters/verl/fully_async_ports.py | 5 + yeto/rl/adapters/verl/fully_async_runner.py | 22 +++- yeto/rl/adapters/verl/ports_impl.py | 44 ++++++- yeto/rl/engine/buffer_payload.py | 37 +++++- yeto/rl/engine/offpolicy_source.py | 56 +++++++++ 7 files changed, 299 insertions(+), 4 deletions(-) create mode 100644 tests/test_rl_n2_trainer_wiring.py diff --git a/tests/test_rl_n2_trainer_wiring.py b/tests/test_rl_n2_trainer_wiring.py new file mode 100644 index 000000000..be3b8a907 --- /dev/null +++ b/tests/test_rl_n2_trainer_wiring.py @@ -0,0 +1,118 @@ +"""S25 N2 wiring on the trainer side (CPU): engine-sampler payloads in the +trainer readers, and the per-source ``offpolicy_source`` event in Miles and verl. + +* ``complete_prompt_tokens``: a payload whose engine gave no prompt ids + (``prompt_tokens_known`` False) gets its prompt ids from the trainer's + tokenizer, or is refused; +* Miles reader (``buffer_rollout_miles._BufferRound``, stub Miles modules): + rebuilt tokens, provenance kept in the metadata, per-source metrics; +* verl (``ports_impl.source_event``): padded trainer tensors + per-row + provenance -> one entry per source, masked tokens only. +""" + +from __future__ import annotations + +import math +import types + +import pytest +from test_rl_buffer_rollout import _import_plugin_with_stub_miles + +from yeto.rl.adapters.sampler_engines.fake import FAKE_IDENTITY +from yeto.rl.engine.buffer_payload import PROMPT_TOKENIZED_BY, SampleBufferRolloutError, complete_prompt_tokens +from yeto.rl.engine.engine_sampler import PromptItem, sample_record +from yeto.rl.engine.offpolicy_source import ( + flat_metrics, + masked_source_samples, + offpolicy_source_event, + payload_source_sample, +) +from yeto.rl.engine.sampler_engine import FINISH_LENGTH, FINISH_STOP, Completion, Provenance, SamplingParams + +SOURCE = Provenance.of(FAKE_IDENTITY, sampler_id="s1", sampling=SamplingParams(max_tokens=8)).to_record() +KEY = Provenance.from_record(SOURCE).key() + + +def _record(prompt_ids=(), finish=FINISH_STOP, index=0): + completion = Completion(token_ids=(40, 41), logprobs=(-0.1, -0.2), text="ok", finish_reason=finish, + prompt_token_ids=tuple(prompt_ids)) + return sample_record(PromptItem("p0", "2+2="), completion, reward=1.0, version=3, provenance=SOURCE, + group_index=0, index=index) + + +def test_prompt_ids_come_from_the_trainer_tokenizer_when_the_engine_gave_none(): + sample = _record() + assert sample["metadata"]["prompt_tokens_known"] is False and sample["tokens"] == [40, 41] + done = complete_prompt_tokens(sample, lambda text: [7, 8, 9] if text == "2+2=" else []) + assert done["tokens"] == [7, 8, 9, 40, 41] and done["response_length"] == 2 + assert done["metadata"][PROMPT_TOKENIZED_BY] == "trainer" + assert done["metadata"]["provenance"] == SOURCE # the source record travels on + assert sample["tokens"] == [40, 41] # the payload sample itself is not changed + + with pytest.raises(SampleBufferRolloutError, match="no tokenizer"): + complete_prompt_tokens(sample, None) + with pytest.raises(SampleBufferRolloutError, match="no ids"): + complete_prompt_tokens(sample, lambda text: []) + with pytest.raises(SampleBufferRolloutError, match="response_length"): + complete_prompt_tokens({**sample, "tokens": [1, 40, 41]}, lambda text: [7]) + + known = _record(prompt_ids=(7, 8)) + assert complete_prompt_tokens(known, None)["tokens"] == [7, 8, 40, 41] # engine ids are kept + + +def test_miles_reader_rebuilds_prompts_and_reports_per_source_metrics(monkeypatch): + plugin = _import_plugin_with_stub_miles(monkeypatch, {}) + monkeypatch.setattr(plugin, "Sample", types.SimpleNamespace(from_dict=dict)) + monkeypatch.setattr(plugin, "RolloutFnTrainOutput", lambda samples, metrics: (samples, metrics)) + reader = plugin._BufferRound.__new__(plugin._BufferRound) + 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}) + 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 + + samples, metrics = reader(types.SimpleNamespace(evaluation=False, rollout_id=2, weight_version=3)) + ((first, second),) = samples + assert first["tokens"] == [7, 8, 40, 41] and first["group_index"] == 5 and second["index"] == 11 + assert first["metadata"]["provenance"] == SOURCE + assert metrics["sample_buffer/groups"] == 1.0 + assert metrics[f"offpolicy_source/{KEY}/samples"] == 2.0 + assert metrics[f"offpolicy_source/{KEY}/truncated_share"] == 0.5 + # Miles' trainer does not hand its log-probs back: no ratio metrics on this side + assert not any(name.endswith("/ratio_p50") for name in metrics) + + +def test_payload_source_sample_counts_without_train_logprobs(): + event = offpolicy_source_event([payload_source_sample(_record())], rollout_id=0, trainer_version=0, + backend="miles") + (entry,) = event["sources"] + assert entry["key"] == KEY and entry["samples"] == 1 and entry["tokens"] == 0 and entry["ratio_p50"] is None + sample = payload_source_sample(_record(), train_logprobs=[-0.1, -0.3]) + assert list(sample.rollout_logprobs) == [-0.1, -0.2] + + +def test_verl_source_event_groups_padded_rows_by_their_provenance(): + torch = pytest.importorskip("torch") + from yeto.rl.adapters.verl.ports_impl import source_event + + train = torch.tensor([[-0.1, -0.2, 0.0], [-1.0, 0.0, 0.0], [-0.5, -0.5, -0.5]]) + rollout = torch.tensor([[-0.1, -0.2, 9.0], [-2.0, 0.0, 0.0], [-0.5, -0.5, -0.5]]) + mask = torch.tensor([[1, 1, 0], [1, 0, 0], [0, 0, 0]]) # row 2: an ABORTED placeholder + event = source_event(train, rollout, mask, [SOURCE, SOURCE, None], "ok", + rollout_id=4, trainer_version=4, tis_upper=2.0) + sources = {s["key"]: s for s in event["sources"]} + assert event["consumed"] == 3 and event["provenance_read"] == "ok" and event["backend"] == "verl" + ours = sources[KEY] + assert ours["samples"] == 2 and ours["tokens"] == 3 # the padded token (rollout 9.0) is not read + assert math.isclose(ours["ratio_max"], math.e) # row 1: exp(-1.0 - -2.0) + assert sources["unknown/unknown/unknown@unknown/unknown/unknown"]["tokens"] == 0 + assert flat_metrics(event)[f"offpolicy_source/{KEY}/samples"] == 2.0 + + +def test_masked_source_samples_checks_row_counts(): + with pytest.raises(ValueError, match="rows"): + 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]) diff --git a/yeto/rl/adapters/miles/buffer_rollout_miles.py b/yeto/rl/adapters/miles/buffer_rollout_miles.py index a2ffb173f..44959814b 100644 --- a/yeto/rl/adapters/miles/buffer_rollout_miles.py +++ b/yeto/rl/adapters/miles/buffer_rollout_miles.py @@ -22,6 +22,9 @@ from miles.utils.function_registry import load_function from miles.utils.types import Sample +from yeto.rl.engine.buffer_payload import complete_prompt_tokens +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, @@ -97,6 +100,7 @@ def __init__(self, args: Any, *, create_index: bool = True, merge_shards: bool = wait_s=float(os.environ.get(WAIT_ENV, DEFAULT_WAIT_S)), poll_s=float(os.environ.get(POLL_ENV, DEFAULT_POLL_S))) register_reader(self.reader) + self._tokenizer: Any = None def __call__(self, input: RolloutFnInput) -> RolloutFnTrainOutput: if input.evaluation: @@ -104,11 +108,26 @@ def __call__(self, input: RolloutFnInput) -> RolloutFnTrainOutput: "set --eval-function-path or disable eval") groups, report = self.reader.take(input.rollout_id, self.groups, current_version=input.weight_version) + groups = [(position, [complete_prompt_tokens(d, self._tokenize) for d in group]) + for position, group in groups] + source = offpolicy_source_event( + (payload_source_sample(d) for _, group in groups for d in group), + rollout_id=input.rollout_id, backend="miles", + trainer_version=-1 if input.weight_version is None else int(input.weight_version)) 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] self._run_round_hooks(samples) - return RolloutFnTrainOutput(samples=samples, metrics=report.metrics()) + return RolloutFnTrainOutput(samples=samples, metrics={**report.metrics(), **flat_metrics(source)}) + + def _tokenize(self, text: str) -> list[int]: + """Prompt ids for a sample whose engine gave none (loads the tokenizer on first use).""" + if self._tokenizer is None: + from miles.utils.processing_utils import load_tokenizer + + self._tokenizer = load_tokenizer(self.args.hf_checkpoint, chat_template_path=self.args.chat_template_path, + trust_remote_code=True) + return list(self._tokenizer(text, add_special_tokens=True)["input_ids"]) def _run_round_hooks(self, samples: list[list[Any]]) -> None: """Call the end-of-round hooks Miles' generating rollout calls, so the yeto diff --git a/yeto/rl/adapters/verl/fully_async_ports.py b/yeto/rl/adapters/verl/fully_async_ports.py index 0f59ff8b2..061fb920c 100644 --- a/yeto/rl/adapters/verl/fully_async_ports.py +++ b/yeto/rl/adapters/verl/fully_async_ports.py @@ -34,6 +34,7 @@ from typing import Any from yeto.rl.contracts import InferencePublicationManifest, LocalStepReceipt +from yeto.rl.engine.offpolicy_source import OFFPOLICY_SOURCE_EVENT from yeto.rl.engine.ports import PlacementDescription, PublicationResult, RolloutBatchHandle from . import publish as pub @@ -180,6 +181,10 @@ def train_step(self, batch: RolloutBatchHandle) -> LocalStepReceipt: if isl.last_criteria: isl.append_jsonl(f"verl-criteria-{isl.learner_id}.jsonl", isl.last_criteria) isl.emit("rl_mismatch_criteria", **isl.last_criteria) + if source := out.get("offpolicy_source"): + source = {**source, "rollout_id": int(batch.rollout_id)} + isl.append_jsonl(f"verl-offpolicy-source-{isl.learner_id}.jsonl", source) + isl.emit(OFFPOLICY_SOURCE_EVENT, **source) def metric(*suffixes): for key, value in metrics.items(): diff --git a/yeto/rl/adapters/verl/fully_async_runner.py b/yeto/rl/adapters/verl/fully_async_runner.py index e43c15b2d..0d07f881e 100644 --- a/yeto/rl/adapters/verl/fully_async_runner.py +++ b/yeto/rl/adapters/verl/fully_async_runner.py @@ -120,6 +120,24 @@ def _yeto_criteria(self, batch) -> dict: crit["param_version"] = int(self.current_param_version) return crit + def _yeto_offpolicy_source(self, batch) -> dict: + """Per-source log-prob mismatch of the step (S25 N2.5); ``rollout_id`` is set by the driver.""" + from .ports_impl import source_error, source_event + + version = int(self.current_param_version) + try: + train_lp, infer_lp = batch.batch["old_log_probs"], batch.batch["rollout_log_probs"] + mask = batch.batch["response_mask"] + width = min(train_lp.shape[-1], infer_lp.shape[-1], mask.shape[-1]) + provenance = batch.non_tensor_batch.get("provenance") + return source_event(train_lp[..., :width], infer_lp[..., :width], mask[..., :width], + None if provenance is None else list(provenance), + "ok" if provenance is not None else "missing: no provenance field", + rollout_id=-1, trainer_version=version, + tis_upper=float(getattr(self, "yeto_tis_upper", 2.0))) + except Exception as exc: # noqa: BLE001 - archival, training continues + return source_error(exc, rollout_id=-1, trainer_version=version) + def yeto_configure(self, tis_upper: float, thresholds_key) -> None: self.yeto_tis_upper = float(tis_upper) self.yeto_thresholds_key = tuple(thresholds_key) @@ -132,6 +150,7 @@ def yeto_train_step(self): batch = self._fit_compute_reward(batch) batch = self._fit_compute_log_prob(batch) criteria = self._yeto_criteria(batch) # 6.4b item 6: re-hooked after old log-probs + source = self._yeto_offpolicy_source(batch) batch = self._fit_compute_ref_log_prob(batch) batch = self._fit_compute_critic(batch) batch = self._fit_compute_advantage(batch) @@ -146,7 +165,8 @@ def yeto_train_step(self): self.global_steps += 1 numeric = {k: float(v) for k, v in self.metrics.items() if isinstance(v, (int, float)) and not isinstance(v, bool)} - return {"metrics": numeric, "criteria": criteria, "param_version": int(self.current_param_version)} + return {"metrics": numeric, "criteria": criteria, "offpolicy_source": source, + "param_version": int(self.current_param_version)} async def yeto_push_weights(self): t0 = time.time() diff --git a/yeto/rl/adapters/verl/ports_impl.py b/yeto/rl/adapters/verl/ports_impl.py index dad0f2801..4392c447e 100644 --- a/yeto/rl/adapters/verl/ports_impl.py +++ b/yeto/rl/adapters/verl/ports_impl.py @@ -34,7 +34,7 @@ from typing import Any from yeto.rl.contracts import InferencePublicationManifest, LocalStepReceipt -from yeto.rl.engine import mismatch_criteria +from yeto.rl.engine import mismatch_criteria, offpolicy_source from yeto.rl.engine.ports import GroupMetadata, PlacementDescription, PublicationResult, RolloutBatchHandle from yeto.rl.engine.trainable_state import TrainableState @@ -477,3 +477,45 @@ def criteria_from_trainer_batch(isl: VerlIsland, kv_batch) -> None: isl.last_criteria = crit isl.append_jsonl(f"verl-criteria-{isl.learner_id}.jsonl", crit) isl.emit("rl_mismatch_criteria", **{k: v for k, v in crit.items()}) + if "error" not in crit: + try: + provenances, read = _row_field(kv_batch, "provenance") + event = source_event(train_lp, infer_lp, mask, provenances, read, + rollout_id=step, trainer_version=step, tis_upper=isl.tis_upper) + except Exception as exc: # noqa: BLE001 - archival, training continues + event = source_error(exc, rollout_id=step, trainer_version=step) + _emit_offpolicy_source(isl, event) + + +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.""" + 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" + except Exception as exc: # noqa: BLE001 - the rows then group under ``unknown`` + return None, f"missing: {type(exc).__name__}: {str(exc)[:200]}" + + +def source_event(train_lp, infer_lp, mask, provenances, provenance_read: str, *, rollout_id: int, + trainer_version: int, tis_upper: float) -> dict: + """The ``offpolicy_source`` event of one verl step (padded tensors, same width).""" + event = offpolicy_source.offpolicy_source_event( + offpolicy_source.masked_source_samples(train_lp.float().tolist(), infer_lp.float().tolist(), + mask.tolist(), provenances=provenances), + rollout_id=rollout_id, trainer_version=trainer_version, backend="verl", tis_upper=tis_upper) + event["provenance_read"] = provenance_read + return event + + +def source_error(exc: Exception, *, rollout_id: int, trainer_version: int) -> dict: + return {"schema": offpolicy_source.SCHEMA, "rollout_id": int(rollout_id), "trainer_version": int(trainer_version), + "backend": "verl", "error": f"{type(exc).__name__}: {str(exc)[:300]}", "sources": []} + + +def _emit_offpolicy_source(isl: VerlIsland, event: dict) -> None: + isl.append_jsonl(f"verl-offpolicy-source-{isl.learner_id}.jsonl", event) + isl.emit(offpolicy_source.OFFPOLICY_SOURCE_EVENT, **event) diff --git a/yeto/rl/engine/buffer_payload.py b/yeto/rl/engine/buffer_payload.py index dd6cb3292..23bb1f8c8 100644 --- a/yeto/rl/engine/buffer_payload.py +++ b/yeto/rl/engine/buffer_payload.py @@ -13,7 +13,7 @@ import base64 import json -from collections.abc import Mapping, Sequence +from collections.abc import Callable, Mapping, Sequence from typing import Any from .policy_versions import engine_version_number @@ -97,6 +97,41 @@ def group_entry_fields(samples: Sequence[Mapping[str, Any]]) -> tuple[int, int]: return int(oldest), tokens +PROMPT_TOKENIZED_BY = "prompt_tokenized_by" # metadata: who produced the prompt ids + + +def complete_prompt_tokens(sample: Mapping[str, Any], + tokenize: Callable[[str], Sequence[int]] | None) -> dict[str, Any]: + """Trainer side of ``metadata.prompt_tokens_known`` (S25 N2.3 payload). + + An engine that does not return the prompt ids (SGLang text input) writes + ``tokens`` = response only and ``prompt_tokens_known`` False. The trainer + then tokenizes the prompt text itself: ``tokens`` becomes + ``tokenize(prompt) + response`` and ``metadata.prompt_tokenized_by`` is + ``trainer``. ``tokenize`` must match the engine's own text tokenization + (SGLang and vLLM default to the HF tokenizer with special tokens); that + match is not verified on a real engine. Without ``tokenize`` such a + sample is refused: training on response-only ``tokens`` would read the + response as the prompt. A sample with known prompt ids is returned as is. + """ + metadata = sample.get("metadata") or {} + if metadata.get("prompt_tokens_known", True): + return dict(sample) + where = f"sample {sample.get('index')}" + if tokenize is None: + raise SampleBufferRolloutError(f"{where}: the engine gave no prompt ids and this reader has no " + "tokenizer to rebuild them") + response = list(sample.get("tokens") or ()) + if len(response) != int(sample.get("response_length") or 0): + raise SampleBufferRolloutError(f"{where}: prompt_tokens_known is False but tokens has " + f"{len(response)} ids for response_length {sample.get('response_length')}") + prompt = [int(t) for t in tokenize(str(sample.get("prompt") or ""))] + if not prompt: + raise SampleBufferRolloutError(f"{where}: the prompt text tokenizes to no ids") + return {**sample, "tokens": prompt + response, + "metadata": {**metadata, PROMPT_TOKENIZED_BY: "trainer"}} + + class GroupWriter: """Writes finished groups to the buffer: payload first, then the index entry.""" diff --git a/yeto/rl/engine/offpolicy_source.py b/yeto/rl/engine/offpolicy_source.py index dc8b7fb98..d0974e4b4 100644 --- a/yeto/rl/engine/offpolicy_source.py +++ b/yeto/rl/engine/offpolicy_source.py @@ -160,3 +160,59 @@ def offpolicy_source_event(samples: Iterable[SourceSample], *, rollout_id: int, return {"schema": SCHEMA, "rollout_id": int(rollout_id), "trainer_version": int(trainer_version), "backend": str(backend), "tis_upper": float(tis_upper), "consumed": consumed, "sources": sources} + + +def payload_source_sample(sample: Mapping[str, Any], *, train_logprobs: Sequence[float] = (), + rollout_logprobs: Sequence[float] | None = None) -> SourceSample: + """A :class:`SourceSample` from a buffer payload sample (``metadata.provenance`` + and ``metadata.finish_reason``, as the engine sampler island writes them). + + Without ``train_logprobs`` the sample counts in ``samples`` and ``finish`` + only (its ratio fields stay empty): the reader side of a backend whose + training process does not hand its log-probs back (Miles). + """ + metadata = sample.get("metadata") or {} + rollout: list[float] = [] + if train_logprobs: + rollout = list(sample.get("rollout_log_probs") or () if rollout_logprobs is None else rollout_logprobs) + return SourceSample(provenance=metadata.get("provenance"), train_logprobs=list(train_logprobs), + rollout_logprobs=rollout, finish_reason=metadata.get("finish_reason")) + + +def flat_metrics(event: Mapping[str, Any], prefix: str = OFFPOLICY_SOURCE_EVENT) -> dict[str, float]: + """The numeric per-source fields of an event as ``//`` metrics.""" + out: dict[str, float] = {} + for entry in event["sources"]: + for name in ("samples", "tokens", "ratio_p50", "ratio_p99", "tis_clipfrac", "abs_diff_mean", + "k3", "truncated_share", "loop_abort_share"): + value = entry.get(name) + if isinstance(value, (int, float)): + out[f"{prefix}/{entry['key']}/{name}"] = float(value) + return out + + +def masked_source_samples(train_logprobs: Sequence[Sequence[float]], rollout_logprobs: Sequence[Sequence[float]], + masks: Sequence[Sequence[float]], *, + provenances: Sequence[Mapping[str, str] | None] | None = None, + finish_reasons: Sequence[str | None] | None = None) -> list[SourceSample]: + """One :class:`SourceSample` per row of a padded trainer batch (nested lists, + e.g. ``tensor.tolist()``), keeping only the tokens whose mask is set. + + ``provenances`` / ``finish_reasons``: per row, or None when the batch does + not carry them (the rows then group under ``unknown``). + """ + rows = len(masks) + if len(train_logprobs) != rows or len(rollout_logprobs) != rows: + raise ValueError(f"{len(train_logprobs)} train rows, {len(rollout_logprobs)} rollout rows, {rows} mask rows") + for name, values in (("provenances", provenances), ("finish_reasons", finish_reasons)): + if values is not None and len(values) != rows: + raise ValueError(f"{len(values)} {name} for {rows} rows") + out = [] + for row in range(rows): + keep = [i for i, m in enumerate(masks[row]) if m] + out.append(SourceSample( + provenance=None if provenances is None else provenances[row], + train_logprobs=[float(train_logprobs[row][i]) for i in keep], + rollout_logprobs=[float(rollout_logprobs[row][i]) for i in keep], + finish_reason=None if finish_reasons is None else finish_reasons[row])) + return out From 526cf078159506419578bd77411dd825ca4f015b Mon Sep 17 00:00:00 2001 From: michael <78113554+chowhsu@users.noreply.github.com> Date: Mon, 19 Oct 2026 12:42:33 +0000 Subject: [PATCH 28/34] Load the Miles reader's tokenizer through the one allowed Miles load The prompt-id rebuild loaded the tokenizer with its own trust_remote_code=True call, which the production-tree check refuses. It now uses rollout_meta_hook's Miles tokenizer load, the single allowed one. Co-Authored-By: Claude Opus 5.5 --- yeto/rl/adapters/miles/buffer_rollout_miles.py | 7 +++---- 1 file changed, 3 insertions(+), 4 deletions(-) diff --git a/yeto/rl/adapters/miles/buffer_rollout_miles.py b/yeto/rl/adapters/miles/buffer_rollout_miles.py index 44959814b..e9f6bad87 100644 --- a/yeto/rl/adapters/miles/buffer_rollout_miles.py +++ b/yeto/rl/adapters/miles/buffer_rollout_miles.py @@ -122,11 +122,10 @@ def __call__(self, input: RolloutFnInput) -> RolloutFnTrainOutput: def _tokenize(self, text: str) -> list[int]: """Prompt ids for a sample whose engine gave none (loads the tokenizer on first use).""" - if self._tokenizer is None: - from miles.utils.processing_utils import load_tokenizer + if self._tokenizer is None: # the one Miles tokenizer load yeto allows (test_provenance) + from .rollout_meta_hook import _load_miles_tokenizer - self._tokenizer = load_tokenizer(self.args.hf_checkpoint, chat_template_path=self.args.chat_template_path, - trust_remote_code=True) + self._tokenizer = _load_miles_tokenizer(self.args) return list(self._tokenizer(text, add_special_tokens=True)["input_ids"]) def _run_round_hooks(self, samples: list[list[Any]]) -> None: From 05f91cbe7df14f1c45cd281c2043687700b47f95 Mon Sep 17 00:00:00 2001 From: michael <78113554+chowhsu@users.noreply.github.com> Date: Mon, 19 Oct 2026 13:23:40 +0000 Subject: [PATCH 29/34] Close stage-1 code gaps: sampler entry, glm5_next recipe, SSH docker islands, OSS relay - engine_sampler gains main(): builds the engine sampler island from a JSON spec (engine factory, prompts JSONL, reward, store, post-load checks). Adapters are built by yeto.rl.adapters.sampler_engines.from_spec, so the neutral core names no adapter. - Miles adapter: --rl-model-recipe glm5_next with --rl-transformer-impl (transformer_engine default, local = N7.5 TE-free fallback); the recipe's model arguments replace the Bridge geometry. - SshTrainerSpec: run_script mode with the sky task variables, runtime docker (image, mounts, device args; refused without device args), rendezvous from the inventory address instead of the SSH alias. The launcher builds the held-back ssh: trainer island from --rl-node-inventory. - Head OSS access through a relay node (ossutil over SSH), switched by YETO_RL_STORE_OSS_RELAY; unset keeps the direct path. - yeto-stage1: sampler spec step, trainer steps from the same SshTrainerSpec, store_access relay check, service steps, and a new gap for the Miles full-parameter one-node DP=1 rule (shared run_config function). Co-Authored-By: Claude Opus 5.5 --- examples/stage1/glm53-flash-ppu.yaml | 14 + examples/stage1/ppu-inventory.yaml | 8 +- tests/test_rl_ssh_trainer_island.py | 6 +- tests/test_run_flags.py | 2 + tests/test_stage1_gaps.py | 276 +++++++++++++++++ tests/test_stage1_plan.py | 48 ++- yeto/cloud/oss_relay.py | 158 ++++++++++ yeto/launch/node_inventory.py | 27 +- yeto/launch/rl_island_task.py | 5 + yeto/launch/ssh_trainer_island.py | 256 ++++++++++++---- yeto/launch/trainer_island.py | 12 +- yeto/launch/trainer_loop.py | 13 + yeto/launcher.py | 12 +- yeto/rl/adapters/miles/config.py | 15 +- yeto/rl/adapters/miles/models/glm5_next.py | 129 ++++++++ yeto/rl/adapters/miles/run_flags.py | 9 +- yeto/rl/adapters/sampler_engines/from_spec.py | 54 ++++ yeto/rl/engine/engine_sampler.py | 15 +- yeto/rl/engine/engine_sampler_entry.py | 274 +++++++++++++++++ yeto/rl/engine/run_config.py | 59 +++- yeto/rl/stage1.py | 277 ++++++++++++++---- 21 files changed, 1528 insertions(+), 141 deletions(-) create mode 100644 tests/test_stage1_gaps.py create mode 100644 yeto/cloud/oss_relay.py create mode 100644 yeto/rl/adapters/miles/models/glm5_next.py create mode 100644 yeto/rl/adapters/sampler_engines/from_spec.py create mode 100644 yeto/rl/engine/engine_sampler_entry.py diff --git a/examples/stage1/glm53-flash-ppu.yaml b/examples/stage1/glm53-flash-ppu.yaml index 20dd25bbd..74b3e14cd 100644 --- a/examples/stage1/glm53-flash-ppu.yaml +++ b/examples/stage1/glm53-flash-ppu.yaml @@ -50,6 +50,12 @@ quant: # rounds >= 1: N8.5 re-quantizes each bf expected_expert_matrices: 36288 # A-INF post-load check 1 store: oss://oss-pai-ircyeaqtnfjw11wby2-cn-wulanchabu/yeto/s25-stage1 +# The bucket accepts intranet endpoints only. The head on this server reaches +# it through ppu1: ossutil on ppu1 over SSH (yeto.cloud.oss_relay; main-agent +# ruling 10-19). With a public endpoint or CNAME, set head: direct. +store_access: + head: relay + relay: ppu1 images: sampler: inference-xpu-pytorch-opt:v0.5.17 @@ -114,3 +120,11 @@ task: # main-agent ruling: GSM8K first, TB2/Co reward: yeto.rl.gsm8k_reward:score rollout_batch_size: 16 n_samples_per_prompt: 8 + # The engine sampler island reads prompts from a local JSONL on the sampler + # node (staged from the dataset revision above; runbook S5). Column names + # follow the Miles GSM8K examples (--input-key messages --label-key label), + # not checked against this revision. + prompts_jsonl: /data/datasets/gsm8k-0cbd9f3/train.jsonl + input_key: messages + label_key: label + prompt_format: chat diff --git a/examples/stage1/ppu-inventory.yaml b/examples/stage1/ppu-inventory.yaml index fbc21ea33..4bde13cdb 100644 --- a/examples/stage1/ppu-inventory.yaml +++ b/examples/stage1/ppu-inventory.yaml @@ -1,11 +1,12 @@ # Stage-1 node inventory: 4 x PPU nodes, 8 x ZW-M890P 144 GiB each # (SESSION25/N0.1-ppu-survey.md, read-only survey 2026-10-16). # host = the SSH alias in ~/.ssh/config of the head (this server). The aliases -# do not resolve on the nodes; the stage-1 plan gives each node's rendezvous -# address (eth0) separately. +# do not resolve on the nodes; address = the node's eth0 IP, which the other +# nodes dial (the plan's rendezvous section must agree). nodes: - name: ppu1 host: ppu1 + address: 10.0.0.74 # eth0; order assumed, runbook S1 checks it device_family: ppu card: M890 cards_per_node: 8 @@ -19,6 +20,7 @@ nodes: roles: [trainer] - name: ppu2 host: ppu2 + address: 10.0.0.75 # eth0; order assumed, runbook S1 checks it device_family: ppu card: M890 cards_per_node: 8 @@ -32,6 +34,7 @@ nodes: roles: [trainer] - name: ppu3 host: ppu3 + address: 10.0.0.76 # eth0; order assumed, runbook S1 checks it device_family: ppu card: M890 cards_per_node: 8 @@ -45,6 +48,7 @@ nodes: roles: [trainer] - name: ppu4 host: ppu4 + address: 10.0.0.77 # eth0; order assumed, runbook S1 checks it device_family: ppu card: M890 cards_per_node: 8 diff --git a/tests/test_rl_ssh_trainer_island.py b/tests/test_rl_ssh_trainer_island.py index 74d9f1cd5..9bfedfeaf 100644 --- a/tests/test_rl_ssh_trainer_island.py +++ b/tests/test_rl_ssh_trainer_island.py @@ -60,8 +60,8 @@ def spec(nodes=None, **kw) -> SshTrainerSpec: def test_spec_refuses_mixed_families_docker_and_empty(): with pytest.raises(SshIslandError, match="device family"): spec([node(), node("n1", "h1", device_family="npu")]) - with pytest.raises(SshIslandError, match="docker"): - spec([node(runtime="docker")]) + with pytest.raises(SshIslandError, match="docker_args"): + spec([node(runtime="docker", image="img")]) with pytest.raises(SshIslandError, match="at least one node"): SshTrainerSpec(nodes=(), run_id="r", command=("x",), workdir="/") @@ -73,7 +73,7 @@ def test_spec_is_refused_when_the_memory_plan_does_not_fit(): with pytest.raises(MemoryPlanError, match="does not fit"): ssh_trainer_spec(small, shape=shape, plan_overrides={"seq_len": 8192}, run_id="r", command=("x",), workdir="/") - ok = ssh_trainer_spec([node(name=f"n{i}", host=f"h{i}") for i in range(8)], shape=shape, + ok = ssh_trainer_spec([node(name=f"n{i}", host=f"h{i}", address=f"10.0.0.{i + 1}") for i in range(8)], shape=shape, plan_overrides={"seq_len": 8192, "offload": "full"}, run_id="r", command=("x",), workdir="/") assert ok.memory["fits"] and ok.memory["label"].startswith("derived") diff --git a/tests/test_run_flags.py b/tests/test_run_flags.py index 3eef1ea57..afc9b9ea6 100644 --- a/tests/test_run_flags.py +++ b/tests/test_run_flags.py @@ -270,6 +270,8 @@ def test_forward_tables_name_flags_of_both_layers(): # 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", + # docker run of a runtime: docker SSH trainer node (ssh_trainer_island.py, node_inventory.py example) + "--device", "--network", }) diff --git a/tests/test_stage1_gaps.py b/tests/test_stage1_gaps.py new file mode 100644 index 000000000..3969df259 --- /dev/null +++ b/tests/test_stage1_gaps.py @@ -0,0 +1,276 @@ +"""S25 stage 1, gaps 3-6 (T-GAPS): the engine sampler process entry, the +glm5_next Miles recipe, SSH trainer islands on docker nodes from the inventory, +and the head's OSS access through a relay node. CPU only, no SSH.""" + +from __future__ import annotations + +import json +import subprocess +from pathlib import Path +from types import SimpleNamespace + +import pytest + +from yeto.cloud.oss_relay import RELAY_ENV, Relay, RelayOssObjectStore, relay_from_env +from yeto.gpu_spec import ClusterSpec +from yeto.launch.node_inventory import InventoryError, Node, parse_inventory +from yeto.launch.ssh_trainer_island import ( + SshIslandError, + SshTrainerSpec, + container_name, + inventory_trainer_nodes, + island_spec_from_inventory, + kill_script, + remote_paths, + start_script, +) +from yeto.launch.trainer_island import head_store_objects +from yeto.rl.adapters.miles.models import glm5_next +from yeto.rl.adapters.sampler_engines.sglang_server import SGLangServerEngine +from yeto.rl.engine import engine_sampler, engine_sampler_entry as entry +from yeto.rl.engine.object_store import ObjectStoreError +from yeto.rl.engine.resume import CheckpointStore +from yeto.rl.engine.run_config import full_parameter_layout_problem +from yeto.rl.engine.sample_buffer import FileSampleBuffer +from yeto.rl.engine.sampler_checks import GSM8K_FIXED_PROMPTS +from yeto.rl import stage1 + +PLAN = Path(__file__).resolve().parents[1] / "examples" / "stage1" / "glm53-flash-ppu.yaml" + + +# ------------------------------------------------------------ gap 3: sampler entry +def _sampler_spec(tmp_path: Path, **over) -> dict: + prompts = tmp_path / "prompts.jsonl" + prompts.write_text("".join(json.dumps({"q": f"question {i}", "a": str(i)}) + "\n" for i in range(3))) + spec = { + "sampler_id": "ppu4", + "engine": {"factory": "yeto.rl.adapters.sampler_engines.from_spec:build", "adapter": "fake", "kwargs": { + "answers": {p.prompt: f"The answer is {p.answer}." for p in GSM8K_FIXED_PROMPTS}, + "max_model_len": 64}}, + "sampling": {"max_tokens": 6, "n": 2}, + "num_rollouts": 2, "max_backlog_groups": 100, "max_prompt_tokens": 32, + "store": str(tmp_path / "store"), + "prompts": {"path": str(prompts), "input_key": "q", "label_key": "a", "groups_per_round": 2}, + "reward": "tests.test_stage1_gaps:half_reward", + "event_tape": str(tmp_path / "tape" / "events.jsonl"), + } + return {**spec, **over} + + +def half_reward(item, completion) -> float: + return 0.5 + + +def test_engine_sampler_entry_runs_the_island_from_a_spec(tmp_path): + assert entry.run(_sampler_spec(tmp_path), resource_sampler=lambda emit: None) == 2 + 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()] + names = [e["event"] for e in events] + assert names[0] == "rl_sampler_start" and names.count("rl_sampler_engine_round") == 2 + assert all(e["sampler_id"] == "ppu4" for e in events) + + +def test_prompts_wrap_and_carry_labels(tmp_path): + prompts = entry.build_prompts(_sampler_spec(tmp_path)) + assert [p.prompt for p in prompts(1)] == ["question 2", "question 0"] + assert prompts(0)[1].label == "1" + + +def test_main_rejects_a_bad_spec(tmp_path, capsys): + spec = tmp_path / "spec.json" + spec.write_text(json.dumps(_sampler_spec( + tmp_path, engine={"factory": "yeto.rl.adapters.sampler_engines.from_spec:build", "adapter": "nope"}))) + assert engine_sampler.main(["--spec", str(spec)]) == 2 + assert "unknown sampler engine" in capsys.readouterr().err + + +def test_miles_reward_wraps_an_async_score(): + async def score(args, sample): + return 1.0 if sample.response.endswith(sample.label) else 0.0 + + reward = entry.miles_reward(score) + item = engine_sampler.PromptItem("p", "q", label="7") + assert reward(item, SimpleNamespace(text="it is 7")) == 1.0 + + +def test_stage1_sampler_spec_builds_the_same_sglang_command(): + out = stage1.expand(stage1.load_plan(PLAN)) + run = out.sampler["engines"][0] + engine = entry.build_engine(run["spec"]) + assert isinstance(engine, SGLangServerEngine) + assert engine.argv(engine.config.base_dir) == run["engine_argv"] + assert entry.island_config(run["spec"]).sampling.max_tokens == 4096 + assert run["spec"]["release_dir"] == "policy-releases-mxfp4-g32-rtn" + + +# ------------------------------------------------------------ gap 4: glm5_next recipe +def test_glm5_next_model_args_follow_transformer_impl(): + te = glm5_next.model_argv() + local = glm5_next.model_argv(transformer_impl="local") + assert "--moe-grouped-gemm" in te and "--transformer-impl" not in te + assert "--moe-grouped-gemm" not in local and "--no-gradient-accumulation-fusion" in local + assert te[te.index("--moe-layer-freq") + 1].count("1") == 42 + with pytest.raises(ValueError, match="transformer_impl"): + glm5_next.model_argv(transformer_impl="te") + + +def test_glm5_next_recipe_replaces_the_bridge_geometry(): + generic = ["train.py", "--num-layers", "4", "--num-query-groups", "8", "--hidden-size", "64", + "--no-gradient-accumulation-fusion", "--bf16"] + te = glm5_next.apply_recipe(generic, glm5_next.model_argv()) + assert te[0] == "train.py" and "--num-query-groups" not in te + assert te.count("--num-layers") == 1 and te[te.index("--num-layers") + 1] == "45" + assert "--no-gradient-accumulation-fusion" in te and "--bf16" in te + local = glm5_next.apply_recipe(generic, glm5_next.model_argv(transformer_impl="local")) + assert local.count("--no-gradient-accumulation-fusion") == 1 + + +def test_island_parser_takes_the_recipe_and_transformer_impl(): + from yeto.rl.adapters.miles.island_entry import build_parser + + argv = list(stage1.expand(stage1.load_plan(PLAN)).trainer["miles"]["island_argv"][3:]) + argv[argv.index("--rl-transformer-impl") + 1] = "local" + args, unknown = build_parser().parse_known_args(argv) + assert unknown == [] + assert (args.rl_model_recipe, args.rl_transformer_impl) == ("glm5_next", "local") + + +def test_full_parameter_layout_rule_is_shared(): + assert full_parameter_layout_problem(data_parallel=2, actor_num_nodes=3, rollout_gpus=0, + rollout_gpus_per_engine=1) == "dense full-parameter GRPO requires DP=1" + assert full_parameter_layout_problem(data_parallel=1, actor_num_nodes=1, rollout_gpus=8, + rollout_gpus_per_engine=8) is None + + +# ------------------------------------------------------------ gap 5: SSH trainer island +def _node(name="ppu1", **kw) -> Node: + base = dict(host=name, device_family="ppu", card="M890", cards_per_node=8, hbm_gib=144.0, + host_memory_gib=2266.0, fabric="roce", interface="eth0", runtime="docker", + image="img:1", mounts={"/data": "/data"}, docker_args=("--device", "/dev/x")) + return Node(name=name, **(base | kw)) + + +def test_inventory_reads_address_and_docker_args(): + raw = {"name": "n", "host": "ppu1", "device_family": "ppu", "card": "M890", "cards_per_node": 8, + "hbm_gib": 144, "host_memory_gib": 2266, "image": "i", "address": "10.0.0.74", + "docker_args": ["--device", "/dev/a"]} + node = parse_inventory({"nodes": [raw]})[0] + assert (node.address, node.docker_args) == ("10.0.0.74", ("--device", "/dev/a")) + with pytest.raises(InventoryError, match="IP address"): + parse_inventory({"nodes": [{**raw, "address": "ppu1"}]}) + + +def test_docker_island_needs_device_args_and_alias_nodes_need_an_address(): + with pytest.raises(SshIslandError, match="docker_args"): + SshTrainerSpec(nodes=(_node(docker_args=()),), run_id="r", workdir="/w", command=("x",)) + with pytest.raises(SshIslandError, match="rendezvous address"): + SshTrainerSpec(nodes=(_node(), _node("ppu2")), run_id="r", workdir="/w", command=("x",)) + with pytest.raises(SshIslandError, match="one trainer command or one run script"): + SshTrainerSpec(nodes=(_node(),), run_id="r", workdir="/w") + + +def test_docker_run_script_uses_inventory_addresses_and_the_container_is_removed(): + nodes = (_node(address="10.0.0.74"), _node("ppu2", address="10.0.0.75")) + spec = SshTrainerSpec(nodes=nodes, run_id="r", workdir="/w", run_script='echo "$SKYPILOT_NODE_RANK"', + env={"SECRET": "s3cr3t"}, image="trainer:2") + paths = remote_paths("r", "trainer", 1) + name = container_name("r", "trainer", 1, 1) + script = start_script(spec, 1, {}, paths, container=name) + subprocess.run(["bash", "-n"], input=script, text=True, check=True) + assert "export SKYPILOT_NODE_IPS='10.0.0.74\n10.0.0.75'" in script and "SKYPILOT_NODE_RANK=1" in script + assert f"--name {name}" in script and "--device /dev/x" in script and "trainer:2" in script + assert "-e SECRET" in script and "SECRET=s3cr3t -" not in script.split("setsid", 1)[1] + assert "ppu1" not in script.split("setsid", 1)[1] # no SSH alias as rendezvous + assert f"docker rm -f {name}" in kill_script(paths, container=name) + with pytest.raises(SshIslandError, match="container name"): + start_script(spec, 0, {}, paths) + + +def test_run_script_mode_runs_locally_with_the_sky_variables(tmp_path): + spec = SshTrainerSpec(nodes=(_node(runtime="rootfs", docker_args=(), host="127.0.0.1"),), run_id="r", + workdir=str(tmp_path), run_script=f'echo "$SKYPILOT_NUM_NODES $SKYPILOT_NODE_IPS" > {tmp_path}/o\nexit 4') + paths = {k: v.replace("$HOME", str(tmp_path)) for k, v in remote_paths("r", "t", 0).items()} + script = start_script(spec, 0, {}, paths) + subprocess.run(["bash", "-c", script], check=False, timeout=30) + assert (tmp_path / "o").read_text().strip() == "1 127.0.0.1" + assert Path(paths["rc"]).read_text().strip() == "4" + + +def test_launcher_builds_the_ssh_island_from_the_inventory(): + nodes = [_node("ppu1", address="10.0.0.74", roles=("trainer",)), _node("ppu2", address="10.0.0.75", + roles=("trainer",)), _node("ppu4", roles=("sampler",))] + args = SimpleNamespace(rl_node_inventory_nodes=nodes, rl_ssh_trainer_workdir="/data/yeto/code", + cluster_prefix="run7") + island = ClusterSpec(cloud="ssh", region=None, num_nodes=2, gpus_per_node=8, gpu="M890") + task = SimpleNamespace(setup="pip check", run="echo run", envs={"A": "1"}, secrets={"B": "2"}) + spec = island_spec_from_inventory(args, island, task) + assert [n.name for n in spec.nodes] == ["ppu1", "ppu2"] and spec.addresses() == ["10.0.0.74", "10.0.0.75"] + assert spec.env == {"A": "1", "B": "2"} and spec.run_script.endswith("echo run") + with pytest.raises(SshIslandError, match="3 node"): + inventory_trainer_nodes(args, ClusterSpec(cloud="ssh", region=None, num_nodes=3, gpus_per_node=8, gpu="M890")) + with pytest.raises(SshIslandError, match="workdir"): + island_spec_from_inventory(SimpleNamespace(**{**vars(args), "rl_ssh_trainer_workdir": None}), island, task) + + +# ------------------------------------------------------------ gap 6: OSS through a relay +class FakeRelay: + """``ssh RELAY bash -c SCRIPT`` against a dict, by the ossutil verb.""" + + def __init__(self): + self.objects: dict[str, bytes] = {} + self.calls: list[str] = [] + + def __call__(self, argv, data): + script = argv[-1] + self.calls.append(script) + url = next(t for t in script.replace("'", " ").split() if t.startswith("oss://")) + if " cat " in script: + return (0, self.objects[url], b"") if url in self.objects else (1, b"", b"Error: NoSuchKey") + if " cp -f " in script: + self.objects[url] = data + return 0, b"", b"" + if " stat " in script: + return (0, b"", b"") if url in self.objects else (1, b"", b"StatusCode=404") + if " ls -s " in script: + listing = "".join(f"{k}\n" for k in sorted(self.objects) if k.startswith(url)) + return 0, (listing + "\nObject Number is: 2\n").encode(), b"" + if " rm " in script: + self.objects.pop(url, None) + return 0, b"", b"" + return 1, b"", b"permission denied" + + +def test_relay_store_reads_and_writes_through_ossutil(): + fake = FakeRelay() + store = RelayOssObjectStore("oss://bkt/yeto/run", Relay("ppu1"), runner=fake) + assert store.read("cuts/a.json") is None and not store.exists("cuts/a.json") + store.write("cuts/a.json", b"{}") + assert store.create("cuts/b.json", b"1") and not store.create("cuts/b.json", b"2") + assert store.read("cuts/a.json") == b"{}" and store.read("cuts/b.json") == b"1" + assert store.list("cuts/") == ["cuts/a.json", "cuts/b.json"] + assert store.list("cuts/", start_after="cuts/a.json") == ["cuts/b.json"] + store.delete("cuts/a.json") + assert store.list("cuts/") == ["cuts/b.json"] + assert all("ossutil" in c for c in fake.calls) + + +def test_relay_store_reports_other_failures(): + store = RelayOssObjectStore("oss://bkt", Relay("ppu1"), runner=lambda argv, data: (1, b"", b"AccessDenied")) + with pytest.raises(ObjectStoreError, match="AccessDenied"): + store.read("x") + + +def test_head_store_switches_between_relay_and_direct(): + assert relay_from_env({}) is None + store = head_store_objects("oss://bkt/p", environ={RELAY_ENV: "ppu1"}) + assert isinstance(store, RelayOssObjectStore) and store.relay.host == "ppu1" + + +def test_stage1_plan_has_no_code_gaps_left_and_runs_the_relay_check(): + out = stage1.expand(stage1.load_plan(PLAN)) + assert out.store["head_access"]["env"] == {RELAY_ENV: "ppu1"} + assert not any("OSS" in g or "engine_sampler" in g or "recipe" in g or "yeto launch" in g for g in out.gaps) + relay = next(s for s in out.steps if s.name == "head-store-relay") + assert relay.node is None and "ossutil ls -s" in relay.script diff --git a/tests/test_stage1_plan.py b/tests/test_stage1_plan.py index c9f29ecc6..6bfd9d7fe 100644 --- a/tests/test_stage1_plan.py +++ b/tests/test_stage1_plan.py @@ -5,6 +5,7 @@ import copy import shlex from pathlib import Path +from types import SimpleNamespace import pytest import yaml @@ -54,13 +55,14 @@ def test_names_agree_across_roles(shipped): def test_sampler_command_comes_from_the_adapter(shipped): engine = shipped.sampler["engines"][0] - argv = engine["argv"] + argv = engine["engine_argv"] # the island entry starts this server assert argv[argv.index("--tp") + 1] == "8" assert argv[argv.index("--model-path") + 1].endswith("mxfp4-experts-70e86a0") assert "--speculative-algorithm" not in argv # no NEXTN in stage 1 assert argv[argv.index("--moe-runner-backend") + 1] == "triton" assert engine["env"]["MX_PATCH"] == "1" assert engine["url"] == "http://10.0.0.77:30000" + assert engine["argv"][-2:] == ["--spec", engine["spec_path"]] def test_max_tokens_is_one_value_everywhere(shipped): @@ -87,9 +89,11 @@ def test_island_argv_parses_with_the_island_parser(shipped): def test_steps_order_and_staging(shipped): names = [s.name for s in shipped.steps] assert names[:4] == [f"preflight-ppu{i}" for i in range(1, 5)] - assert names.index("stage-sampler_initial-ppu1") < names.index("requant-sidecar") # template on node 0 - assert names.index("sampler-ppu4") < names.index("trainer-island") - assert names[-1] == "requant-sidecar" + assert names.index("stage-sampler_initial-ppu1") < names.index("trainer-ppu1") # requant template on node 0 + assert names.index("sampler-spec-ppu4") < names.index("sampler-ppu4") < names.index("trainer-ppu1") + assert names[-3:] == ["trainer-ppu1", "trainer-ppu2", "trainer-ppu3"] + assert [s.name for s in shipped.steps if s.service] == ["sampler-ppu4", *names[-3:]] + assert "yeto.rl.quant.release" in shipped.steps[-3].script # node 0 side process for step in shipped.steps: if step.script.startswith("docker run"): shlex.split(step.script) # quotable as one command @@ -97,7 +101,7 @@ def test_steps_order_and_staging(shipped): def test_gaps_block_run_until_wired(shipped, capsys): assert any("N0.2" in g for g in shipped.gaps) - assert any("OSS" in g for g in shipped.gaps) + assert any("DP=1" in g for g in shipped.gaps) # Miles full-parameter rule (T-GAPS finding) assert stage1.main(["run", str(PLAN), "--confirm-machines-free"]) == 2 assert "run refused" in capsys.readouterr().err @@ -134,6 +138,26 @@ def runner(argv, input=None, text=True): assert len(seen) == 2 and seen[0][-3:] == ["ppu1", "bash", "-s"] +def test_run_steps_starts_services_and_waits_for_the_trainer(shipped): + started = [] + + class Proc: + stdin = None + + def __init__(self, argv): + self.argv = argv + + def wait(self): + return 0 if "ppu1" in self.argv else 9 + + def runner(argv, input=None, text=True): + return SimpleNamespace(returncode=0) + + assert stage1.run_steps(shipped, runner=runner, + starter=lambda argv, script: started.append(argv) or Proc(argv)) == 0 + assert len(started) == 4 # the sampler and the three trainer nodes + + @pytest.mark.parametrize("edit, message", [ (lambda d: d["trainer"]["miles"].update(pp_layers=[15, 15, 14]), "sum to 44"), (lambda d: d["trainer"]["miles"].update(pp_layers=[23, 22]), "2 stages"), @@ -147,6 +171,7 @@ def runner(argv, input=None, text=True): (lambda d: d["model"]["bf16_source"].update(local_dir="/data/models/bf16"), "bf16_source revision"), (lambda d: d["roles"].update(sampler=["ppu3", "ppu4"]), "both trainer and sampler"), (lambda d: d["rendezvous"].update(ppu2="ppu2"), "not an IP address"), + (lambda d: d["rendezvous"].update(ppu2="10.0.0.99"), "differs from the inventory address"), (lambda d: d["trainer"]["miles"].update(optimizer_offload_fraction=0.0), "activations"), (lambda d: d.update(store="s3://bucket/x"), "oss://"), ]) @@ -166,7 +191,18 @@ def test_te_fallback_drops_te_only_model_args(tmp_path): out = _variant(tmp_path, lambda d: d["trainer"]["miles"].update(transformer_impl="local")) assert out.errors == [] assert out.trainer["miles"]["drop_model_args"] == ["--moe-grouped-gemm", "--moe-permute-fusion"] - assert "--no-gradient-accumulation-fusion" in out.trainer["miles"]["megatron_args"] + model_args = out.trainer["miles"]["model_args"] + assert "--no-gradient-accumulation-fusion" in model_args and "--moe-grouped-gemm" not in model_args + island = out.trainer["miles"]["island_argv"] + assert island[island.index("--rl-transformer-impl") + 1] == "local" + + +def test_direct_head_access_keeps_the_oss_gap(tmp_path): + out = _variant(tmp_path, lambda d: d.update(store_access={"head": "direct"})) + assert any("OSS" in g for g in out.gaps) + assert not any(s.name == "head-store-relay" for s in out.steps) + with pytest.raises(stage1.Stage1Error, match="not an inventory node"): + _variant(tmp_path, lambda d: d.update(store_access={"head": "relay", "relay": "boat"})) def test_verl_backend_expands(tmp_path): diff --git a/yeto/cloud/oss_relay.py b/yeto/cloud/oss_relay.py new file mode 100644 index 000000000..433699382 --- /dev/null +++ b/yeto/cloud/oss_relay.py @@ -0,0 +1,158 @@ +"""Cloud layer: the head's object API of an ``oss://`` store through a relay node +(S25 stage 1, gap 6; main-agent ruling 10-19: the head on this server reads and +writes OSS through ppu1). + +The stage-1 bucket accepts only Alibaba intranet endpoints +(``PublicEndpointForbidden`` from this server). The islands on the PPU nodes +reach it directly (:mod:`yeto.cloud.oss_store`). The head instead runs +``ossutil`` on a relay node over SSH: every :class:`RelayOssObjectStore` call is +one ``ssh RELAY bash -c SCRIPT`` with one ``ossutil`` command; object bodies +travel on the SSH stdin and stdout. ``ossutil`` on the relay uses the relay's own +configuration (``~/.ossutilconfig`` with the intranet endpoint, runbook S4); no +credential passes through the head. + +Switch: :data:`RELAY_ENV` names the relay's SSH target. Unset, the head talks +to OSS directly (``oss_store.head_objects``), so a public endpoint or CNAME, +once it exists, needs no code change. + +Limits (first implementation): + +* ``create`` is a stat followed by a copy, not an atomic conditional write; + the head is the only writer of the keys it creates (trainer commands), and + the sample buffer keeps one shard per sampler (``buffer_shards``), so no + record depends on it; +* one SSH round trip per call (the head reads small objects: cuts, commands, + index records); bulk weights never go through the head; +* the ``ossutil`` output formats parsed here (``ls -s`` lines, ``NoSuchKey`` / + ``404`` on a missing key) follow ossutil 1.7 documentation only; not run + against OSS. +""" + +from __future__ import annotations + +import shlex +import subprocess +from collections.abc import Callable, Mapping, Sequence +from dataclasses import dataclass + +from yeto.cloud.oss_store import SCHEME, parse_oss_uri +from yeto.rl.engine.object_store import ObjectStoreError, check_key + +RELAY_ENV = "YETO_RL_STORE_OSS_RELAY" # SSH target of the relay node; unset = direct +OSSUTIL_ENV = "YETO_RL_STORE_OSS_RELAY_OSSUTIL" # ossutil command on the relay (default "ossutil") +SSH_OPTIONS = ("-o", "BatchMode=yes", "-o", "ConnectTimeout=20") +_MISSING_MARKERS = ("NoSuchKey", "StatusCode=404", "404 Not Found") + +# (argv, stdin bytes) -> (exit code, stdout bytes, stderr bytes) +Runner = Callable[[Sequence[str], bytes], tuple[int, bytes, bytes]] + + +def subprocess_runner(timeout_s: float = 600.0) -> Runner: + def run(argv: Sequence[str], data: bytes) -> tuple[int, bytes, bytes]: + done = subprocess.run(list(argv), input=data, capture_output=True, timeout=timeout_s) # noqa: S603 + return done.returncode, done.stdout, done.stderr + + return run + + +@dataclass(frozen=True) +class Relay: + """The relay node: its SSH target and the ``ossutil`` command there.""" + + host: str + ossutil: str = "ossutil" + + def argv(self, script: str) -> list[str]: + """``ssh`` that runs ``script`` on the relay (ssh joins its arguments into + one remote shell line, hence the quoting).""" + return ["ssh", *SSH_OPTIONS, self.host, f"bash -c {shlex.quote(script)}"] + + +class RelayOssObjectStore: + """:class:`~yeto.rl.engine.object_store.ObjectStore` over ``oss://BUCKET[/PREFIX]`` + through ``ossutil`` on a relay node (module doc).""" + + scheme = SCHEME + buffer_shards = True # as OssObjectStore: no record depends on create + + def __init__(self, uri: str, relay: Relay, *, runner: Runner | None = None) -> None: + self.bucket, self.prefix = parse_oss_uri(uri) + self.uri = f"{SCHEME}://{self.bucket}" + (f"/{self.prefix}" if self.prefix else "") + self.relay = relay + self.runner = runner or subprocess_runner() + + # ---------------------------------------------------------------- helpers + def _url(self, key: str) -> str: + key = check_key(key) + return f"{SCHEME}://{self.bucket}/" + (f"{self.prefix}/{key}" if self.prefix else key) + + def _run(self, script: str, data: bytes = b"") -> tuple[int, bytes, str]: + """(exit code, stdout, error text); ``data`` is the script's stdin.""" + rc, out, err = self.runner(self.relay.argv(script), data) + detail = "" if rc == 0 else err.decode(errors="replace") + out.decode(errors="replace")[-2000:] + return rc, out, detail + + def _fail(self, what: str, key: str, detail: str) -> ObjectStoreError: + return ObjectStoreError(f"oss relay {self.relay.host}: {what} {self._url(key)} failed: {detail.strip()[-500:]}") + + @staticmethod + def _missing(detail: str) -> bool: + return any(marker in detail for marker in _MISSING_MARKERS) + + # ------------------------------------------------------------ ObjectStore + def read(self, key: str) -> bytes | None: + rc, out, detail = self._run(f"exec {self.relay.ossutil} cat {shlex.quote(self._url(key))}") + if rc == 0: + return out + if self._missing(detail): + return None + raise self._fail("read", key, detail) + + def write(self, key: str, data: bytes) -> None: + script = ('t=$(mktemp) && trap \'rm -f "$t"\' EXIT\n' + 'cat > "$t"\n' + f'{self.relay.ossutil} cp -f "$t" {shlex.quote(self._url(key))} >/dev/null') + rc, _, detail = self._run(script, data) + if rc: + raise self._fail("write", key, detail) + + def exists(self, key: str) -> bool: + rc, _, detail = self._run(f"exec {self.relay.ossutil} stat {shlex.quote(self._url(key))} >/dev/null") + if rc == 0: + return True + if self._missing(detail): + return False + raise self._fail("stat", key, detail) + + def create(self, key: str, data: bytes) -> bool: + """Write ``key`` unless it exists; not atomic (module doc).""" + if self.exists(key): + return False + self.write(key, data) + return True + + def list(self, prefix: str, *, start_after: str = "") -> list[str]: + full = f"{SCHEME}://{self.bucket}/" + (f"{self.prefix}/{prefix}" if self.prefix else prefix) + rc, out, detail = self._run(f"exec {self.relay.ossutil} ls -s {shlex.quote(full)}") + if rc: + raise ObjectStoreError(f"oss relay {self.relay.host}: list {full} failed: {detail.strip()[-500:]}") + root = f"{SCHEME}://{self.bucket}/" + (f"{self.prefix}/" if self.prefix else "") + keys = sorted(line[len(root):] for line in out.decode().splitlines() + if line.startswith(root) and not line.endswith("/")) + return [k for k in keys if k > start_after] if start_after else keys + + def delete(self, key: str) -> None: + rc, _, detail = self._run(f"exec {self.relay.ossutil} rm {shlex.quote(self._url(key))} >/dev/null") + if rc and not self._missing(detail): + raise self._fail("delete", key, detail) + + +def relay_from_env(environ: Mapping[str, str]) -> Relay | None: + host = environ.get(RELAY_ENV) + return Relay(host, environ.get(OSSUTIL_ENV) or "ossutil") if host else None + + +def relay_shell(relay: Relay, command: str) -> str: + """``command`` as one shell line that runs on the relay (for plan steps the + head runs, e.g. a listing of the run directory).""" + return shlex.join(["ssh", *SSH_OPTIONS, relay.host, command]) diff --git a/yeto/launch/node_inventory.py b/yeto/launch/node_inventory.py index 104c7d9e5..1a3c67980 100644 --- a/yeto/launch/node_inventory.py +++ b/yeto/launch/node_inventory.py @@ -19,14 +19,23 @@ host_memory_gib: 1024 fabric: roce interface: eth0 + address: 172.28.11.10 mounts: {/mnt/shared: /mnt/shared} image: REGISTRY/yeto-ascend:TAG runtime: docker + docker_args: [--device, /dev/davinci0, ...] price_per_card_hour: 0.0 + +``address`` is the node's IP on the cluster network (the ``interface``), which +the other nodes dial for rendezvous; ``host`` may be an SSH alias that only the +head resolves. ``docker_args`` are the device flags of ``docker run`` for the +node's device family (a ``runtime: docker`` trainer island refuses to start +without them). """ from __future__ import annotations +import ipaddress from dataclasses import dataclass, field from pathlib import Path from typing import Any, Mapping @@ -44,8 +53,8 @@ _REQUIRED = ("name", "host", "device_family", "card", "cards_per_node", "hbm_gib", "host_memory_gib") -_KNOWN = _REQUIRED + ("proxy_jump", "fabric", "interface", "mounts", "image", "runtime", - "price_per_card_hour", "roles") +_KNOWN = _REQUIRED + ("proxy_jump", "fabric", "interface", "address", "mounts", "image", "runtime", + "docker_args", "price_per_card_hour", "roles") class InventoryError(ValueError): @@ -66,9 +75,11 @@ class Node: proxy_jump: str | None = None fabric: str = "tcp" interface: str | None = None + address: str | None = None # IP on the cluster network (rendezvous); None: not declared mounts: Mapping[str, str] = field(default_factory=dict) image: str | None = None runtime: str = "docker" + docker_args: tuple[str, ...] = () # device flags of `docker run` (runtime docker) price_per_card_hour: float | None = None roles: tuple[str, ...] = () @@ -162,6 +173,16 @@ def parse_node(raw: Any, *, index: int) -> Node: isinstance(k, str) and isinstance(v, str) for k, v in mounts.items()): raise InventoryError(f"node {name}: mounts must be a map of host path to " f"container path, got {mounts!r}") + address = raw.get("address") + if address is not None: + try: + address = str(ipaddress.ip_address(_require_text(address, name, "address"))) + except ValueError: + raise InventoryError(f"node {name}: address must be an IP address (the rendezvous " + f"address on the cluster network), got {address!r}") from None + docker_args = raw.get("docker_args") or () + if not isinstance(docker_args, (list, tuple)) or not all(isinstance(a, str) for a in docker_args): + raise InventoryError(f"node {name}: docker_args must be a list of strings, got {docker_args!r}") price = raw.get("price_per_card_hour") if price is not None and (isinstance(price, bool) or not isinstance(price, (int, float)) or price < 0): @@ -183,10 +204,12 @@ def parse_node(raw: Any, *, index: int) -> Node: fabric=fabric, interface=(_require_text(raw["interface"], name, "interface") if raw.get("interface") is not None else None), + address=address, mounts=dict(mounts), image=(_require_text(raw["image"], name, "image") if raw.get("image") is not None else None), runtime=runtime, + docker_args=tuple(docker_args), price_per_card_hour=(float(price) if price is not None else None), roles=tuple(roles), ) diff --git a/yeto/launch/rl_island_task.py b/yeto/launch/rl_island_task.py index a4a2681ec..a5b443aa5 100644 --- a/yeto/launch/rl_island_task.py +++ b/yeto/launch/rl_island_task.py @@ -103,6 +103,11 @@ def make_miles_island_task( flags = ( f" --model {shlex.quote(args.model)}" f" --rl-model-recipe {shlex.quote(args.rl_model_recipe)}" + + ( + f" --rl-transformer-impl {args.rl_transformer_impl}" + if getattr(args, "rl_transformer_impl", "transformer_engine") != "transformer_engine" + else "" + ) + ( f" --rl-training-attention-backend {args.rl_training_attention_backend}" if getattr(args, "rl_training_attention_backend", "auto") not in (None, "auto") diff --git a/yeto/launch/ssh_trainer_island.py b/yeto/launch/ssh_trainer_island.py index 72ebd423d..cca124765 100644 --- a/yeto/launch/ssh_trainer_island.py +++ b/yeto/launch/ssh_trainer_island.py @@ -7,19 +7,27 @@ it over nodes of the SSH node inventory (:mod:`.node_inventory`), e.g. PPU or Ascend machines that yeto reaches by SSH and does not provision. -One incarnation = one ``torchrun`` per node, started over SSH: +One incarnation = one trainer process per node, started over SSH: 1. before the first start, every node runs its device family's card tool and must match the inventory (``Node.check_cards``); the memory estimate of the plan (:mod:`yeto.rl.fullparam.memory`) is checked when the spec is built (:func:`ssh_trainer_spec`), so a plan that does not fit never reaches a node; -2. each node runs, in a new session (one process group), the trainer command - under ``torchrun`` with the incarnation environment - (``trainer_remote.island_env``), the collective-communication variables of - its family and fabric (``topology.island_comm_env``) and the run's - environment. Node 0 also runs the W4A8 re-quantizer side process +2. each node runs, in a new session (one process group), the trainer with the + incarnation environment (``trainer_remote.island_env``), the + collective-communication variables of its family and fabric + (``topology.island_comm_env``) and the run's environment. The trainer is + either ``command`` under ``torchrun`` or ``run_script``, a bash script that + reads the sky task variables (``SKYPILOT_NODE_RANK``, ``SKYPILOT_NODE_IPS``, + ``SKYPILOT_NUM_NODES``, ``SKYPILOT_NUM_GPUS_PER_NODE``), so the run script + of a sky island (e.g. Ray head on rank 0, Ray workers on the others) runs + unchanged. Node 0 also runs the re-quantizer side process (``yeto.rl.quant.release --loop``) when the spec has one: each BF16 release - becomes a W4A8 release in the store, which W4A8 samplers follow; + becomes a quantized release in the store, which quantized samplers follow; + on ``runtime: docker`` nodes the trainer runs in a container + (``docker run --network host``) of the spec's image with the node's mounts + and the device flags of its family (``docker_args``; a docker island + without them is refused: the flags are not known for every family yet); 3. the SSH session streams node 0's log to the head (``IslandEvents`` reads the trainer events from it); the other nodes' logs go to local files; 4. release kills each node's process group (TERM, then KILL after a grace @@ -27,18 +35,23 @@ an orphan incarnation (the PID file path is derived from the run id, island name and incarnation, recorded in the journal before any start). +Rendezvous: the other nodes dial node 0 at its ``address`` from the inventory +(the IP on the cluster network); an SSH alias resolves on the head only. A +node without an ``address`` is accepted only when its SSH target is an IP. + The head reads the cuts, the commands and the sample buffer through the store's object API (``s3://`` or ``oss://``); with a sharded buffer it also counts the records the trainer has not merged yet (``UnmergedShards``, N3.6). -Not done here: the launcher does not yet hold back an SSH island from ``--gpu`` -(the SSH node pool launch path, N1.3); ``runtime: docker`` nodes are refused -(the per-family device flags of the trainer image are N0.2b). Nothing here was -run against a real node. +The launcher holds back an ``ssh:`` island of ``--gpu`` as the trainer island +and builds its spec from ``--rl-node-inventory`` (:func:`island_spec_from_inventory`). +Nothing here was run against a real node. """ from __future__ import annotations +import ipaddress +import re import shlex import subprocess import threading @@ -69,6 +82,7 @@ KILL_GRACE_S = 30 SSH_OPTIONS = ("-o", "BatchMode=yes", "-o", "ConnectTimeout=20", "-o", "ServerAliveInterval=30") STATE_DIR = ".yeto/trainer-islands" # under the remote $HOME +CONTAINER_STATE_DIR = "/yeto-island-state" # the incarnation's state dir inside a docker trainer class SshIslandError(RuntimeError): @@ -80,39 +94,69 @@ class SshTrainerSpec: """What every incarnation runs. ``command``: the trainer entry, as ``torchrun`` arguments after ``--`` - (e.g. ``["-m", "glm_rl.verl_ext.train", ...]``); ``workdir``: the remote - directory it runs in; ``requant``: argv of the W4A8 re-quantizer side - process on node 0 (None: BF16 samplers only).""" + (e.g. ``["-m", "glm_rl.verl_ext.train", ...]``); or ``run_script``: a bash + script run on every node with the sky task variables (module doc); exactly + one of the two. ``workdir``: the directory it runs in (inside the container + on docker nodes); ``requant``: argv of the re-quantizer side process on + node 0 (None: BF16 samplers only). ``image`` / ``docker_args`` replace the + nodes' inventory values on docker nodes (e.g. a trainer image).""" nodes: tuple[Node, ...] run_id: str - command: tuple[str, ...] workdir: str + command: tuple[str, ...] = () + run_script: str | None = None env: Mapping[str, str] = field(default_factory=dict) requant: tuple[str, ...] | None = None requant_poll_s: float = 10.0 master_port: int = 29500 torchrun: str = "torchrun" + image: str | None = None + docker_args: tuple[str, ...] | None = None memory: Mapping[str, Any] | None = None # the estimate summary the spec was checked with def __post_init__(self) -> None: if not self.nodes: raise SshIslandError("an SSH trainer island needs at least one node") - if not self.command: - raise SshIslandError("an SSH trainer island needs a trainer command") + if bool(self.command) == bool(self.run_script): + raise SshIslandError("an SSH trainer island needs one trainer command or one run script") families = {n.device_family for n in self.nodes} if len(families) != 1: raise SshIslandError(f"one trainer island spans one device family, got {sorted(families)}") - if docker := [n.name for n in self.nodes if n.runtime != "rootfs"]: - raise SshIslandError(f"nodes {docker}: runtime docker is not wired for trainer islands yet " - "(the trainer image's device flags are task N0.2b); use runtime: rootfs") if len({n.cards_per_node for n in self.nodes}) != 1: raise SshIslandError("trainer nodes must have the same number of cards") + if len({n.runtime for n in self.nodes}) != 1: + raise SshIslandError("trainer nodes must have the same runtime (docker or rootfs)") + for node in self.nodes: + if node.runtime != "docker": + continue + if not (self.image or node.image): + raise SshIslandError(f"node {node.name}: runtime docker needs an image") + if not self.docker_args_of(node): + raise SshIslandError( + f"node {node.name}: runtime docker needs docker_args, the device flags of `docker run` " + f"for device family {node.device_family} (inventory docker_args; PPU: task N0.2b)") + if len(self.nodes) > 1: + for node in self.nodes: + rendezvous_address(node) # raises on a node the others cannot dial @property def device_family(self) -> str: return self.nodes[0].device_family + @property + def docker(self) -> bool: + return self.nodes[0].runtime == "docker" + + def docker_args_of(self, node: Node) -> tuple[str, ...]: + return tuple(self.docker_args if self.docker_args is not None else node.docker_args) + + def addresses(self) -> list[str]: + """Rendezvous address of every node (one node: loopback when it has none).""" + if len(self.nodes) == 1 and self.nodes[0].address is None and not _ip_host(self.nodes[0]): + return ["127.0.0.1"] + return [rendezvous_address(n) for n in self.nodes] + def ssh_trainer_spec(nodes: Sequence[Node], *, shape: Any, plan_overrides: Mapping[str, Any], **spec: Any) -> SshTrainerSpec: @@ -124,10 +168,22 @@ def ssh_trainer_spec(nodes: Sequence[Node], *, shape: Any, plan_overrides: Mappi return SshTrainerSpec(nodes=tuple(nodes), memory=estimate.summary(), **spec) -def master_address(node: Node) -> str: - """The rendezvous address the other nodes dial: the SSH target without the - user (an SSH alias must also resolve on the nodes).""" - return node.host.rpartition("@")[2] +def _ip_host(node: Node) -> str | None: + host = node.host.rpartition("@")[2] + try: + return str(ipaddress.ip_address(host)) + except ValueError: + return None + + +def rendezvous_address(node: Node) -> str: + """The address the other nodes dial: the inventory ``address``, else the + SSH target when it is an IP. An SSH alias resolves on the head only.""" + address = node.address or _ip_host(node) + if address is None: + raise SshIslandError(f"node {node.name}: no rendezvous address; set `address` (its IP on the " + f"cluster network) in the inventory (SSH target {node.host!r} is not an IP)") + return address def remote_paths(run_id: str, island: str, incarnation: int) -> dict[str, str]: @@ -136,35 +192,86 @@ def remote_paths(run_id: str, island: str, incarnation: int) -> dict[str, str]: "requant_log": f"{base}/requant.log"} +def container_name(run_id: str, island: str, incarnation: int, rank: int) -> str: + """The docker container of one node of one incarnation (the kill script removes it).""" + raw = f"yeto-{run_id}-{island}-i{int(incarnation)}-n{int(rank)}" + return re.sub(r"[^A-Za-z0-9_.-]", "-", raw) + + def _exports(env: Mapping[str, str]) -> str: return "".join(f"export {k}={shlex.quote(str(v))}\n" for k, v in sorted(env.items())) -def start_script(spec: SshTrainerSpec, rank: int, env: Mapping[str, str], paths: Mapping[str, str]) -> str: - """The bash script one node runs: start the island in its own session, then - stream its log until it exits (node 0's stream carries the trainer events).""" +def skypilot_env(spec: SshTrainerSpec, rank: int) -> dict[str, str]: + """The sky task variables a run script reads (as ``modal_runner.skypilot_env``).""" + addresses = spec.addresses() + return {"SKYPILOT_NODE_IPS": "\n".join(addresses), "SKYPILOT_NUM_NODES": str(len(addresses)), + "SKYPILOT_NODE_RANK": str(int(rank)), + "SKYPILOT_NUM_GPUS_PER_NODE": str(spec.nodes[rank].cards_per_node)} + + +def _trainer_body(spec: SshTrainerSpec, rank: int, requant_log: str) -> str: + """The node's job as bash; it exits with the trainer's exit code.""" node = spec.nodes[rank] - torchrun = [spec.torchrun, f"--nnodes={len(spec.nodes)}", f"--node-rank={rank}", - f"--nproc-per-node={node.cards_per_node}", f"--master-addr={master_address(spec.nodes[0])}", - f"--master-port={spec.master_port}", *spec.command] + if spec.run_script: + trainer = f"(\n{spec.run_script.rstrip()}\n)" # a subshell: the script's exit and exec stay inside + else: + trainer = shlex.join([spec.torchrun, f"--nnodes={len(spec.nodes)}", f"--node-rank={rank}", + f"--nproc-per-node={node.cards_per_node}", + f"--master-addr={spec.addresses()[0]}", + f"--master-port={spec.master_port}", *spec.command]) if rank == 0 and spec.requant: # The loop converts releases while the trainer runs; after the trainer # ends, one more pass converts the last BF16 release. - log = f"\"{paths['requant_log']}\"" - body = (f"{shlex.join(spec.requant)} --loop --poll-seconds {spec.requant_poll_s:g} > {log} 2>&1 &\n" + return (f"{shlex.join(spec.requant)} --loop --poll-seconds {spec.requant_poll_s:g} > {requant_log} 2>&1 &\n" "req=$!\n" - f"{shlex.join(torchrun)} && rc=0 || rc=$?\n" + f"{trainer} && rc=0 || rc=$?\n" "kill $req 2>/dev/null; wait $req 2>/dev/null\n" - f"{shlex.join(spec.requant)} >> {log} 2>&1 || {{ [ $rc -ne 0 ] || rc=1; }}\n" - f"echo $rc > \"{paths['rc']}\"\n") - else: - body = f"{shlex.join(torchrun)} && rc=0 || rc=$?\necho $rc > \"{paths['rc']}\"\n" + f"{shlex.join(spec.requant)} >> {requant_log} 2>&1 || {{ [ $rc -ne 0 ] || rc=1; }}\n" + "exit $rc\n") + return f"{trainer}\n" + + +def docker_argv(spec: SshTrainerSpec, rank: int, name: str, env_names: Sequence[str]) -> list[str]: + """``docker run`` of one node, the image last; the environment goes by name + (``-e KEY``), so no value appears in the process list.""" + node = spec.nodes[rank] + argv = ["docker", "run", "--rm", "--name", name, "--network", "host", "--ipc", "host", + *spec.docker_args_of(node)] + for host_path, container_path in sorted(node.mounts.items()): + argv += ["-v", f"{host_path}:{container_path}"] + for key in sorted(env_names): + argv += ["-e", key] + return [*argv, "-w", spec.workdir, str(spec.image or node.image)] + + +def start_script(spec: SshTrainerSpec, rank: int, env: Mapping[str, str], paths: Mapping[str, str], + *, container: str | None = None) -> str: + """The bash script one node runs: start the island in its own session, then + stream its log until it exits (node 0's stream carries the trainer events). + ``container``: the docker container name (docker nodes).""" + node = spec.nodes[rank] comm = island_comm_env(node.device_family, node.fabric, node.interface, num_nodes=len(spec.nodes)) + full_env = {**comm, **spec.env, **env, **(skypilot_env(spec, rank) if spec.run_script else {})} + if spec.docker: + if container is None: + raise SshIslandError("a docker node needs its container name") + # The state dir is mounted so the re-quantizer log lands next to the trainer log. + inner = _trainer_body(spec, rank, f"{CONTAINER_STATE_DIR}/requant.log") + *options, image = docker_argv(spec, rank, container, list(full_env)) + mount = f"-v \"{paths['dir']}:{CONTAINER_STATE_DIR}\"" # $HOME expands on the node + job = f"{shlex.join(options)} {mount} {shlex.quote(image)} bash -c {shlex.quote(inner)}" + enter = "" + else: + requant_log = f"\"{paths['requant_log']}\"" + job = f"bash -c {shlex.quote(_trainer_body(spec, rank, requant_log))}" + enter = f"cd {shlex.quote(spec.workdir)}\n" + body = f"{job} && rc=0 || rc=$?\necho $rc > \"{paths['rc']}\"\n" return ( "set -eu\n" f"mkdir -p \"{paths['dir']}\"\n" - f"cd {shlex.quote(spec.workdir)}\n" - + _exports({**comm, **spec.env, **env}) + + enter + + _exports(full_env) + f"rm -f \"{paths['rc']}\"\n" # Job control is off in a non-interactive shell, so the background setsid # is not a group leader: it does not fork, and $! is the new group's id. @@ -176,14 +283,18 @@ def start_script(spec: SshTrainerSpec, rank: int, env: Mapping[str, str], paths: ) -def kill_script(paths: Mapping[str, str], grace_s: int = KILL_GRACE_S) -> str: - """TERM the island's process group, KILL it after ``grace_s``; exit 0 when it is gone.""" +def kill_script(paths: Mapping[str, str], grace_s: int = KILL_GRACE_S, *, container: str | None = None) -> str: + """TERM the island's process group, KILL it after ``grace_s``; exit 0 when it + is gone. ``container``: remove that docker container first (a killed + ``docker run`` client does not stop its container).""" + remove = f"docker rm -f {shlex.quote(container)} >/dev/null 2>&1 || true\n" if container else "" return ( - f"pid=$(cat \"{paths['pid']}\" 2>/dev/null) || exit 0\n" - "kill -TERM -- -$pid 2>/dev/null || exit 0\n" - f"for _ in $(seq {int(grace_s)}); do kill -0 -- -$pid 2>/dev/null || exit 0; sleep 1; done\n" - "kill -KILL -- -$pid 2>/dev/null || true\n" - "sleep 1; ! kill -0 -- -$pid 2>/dev/null\n" + remove + + f"pid=$(cat \"{paths['pid']}\" 2>/dev/null) || exit 0\n" + + "kill -TERM -- -$pid 2>/dev/null || exit 0\n" + + f"for _ in $(seq {int(grace_s)}); do kill -0 -- -$pid 2>/dev/null || exit 0; sleep 1; done\n" + + "kill -KILL -- -$pid 2>/dev/null || true\n" + + "sleep 1; ! kill -0 -- -$pid 2>/dev/null\n" ) @@ -253,7 +364,8 @@ def launch(self, incarnation: int, resume_from: CutRef | None) -> _SshJob: IslandEvents(incarnation_tape(self.events_dir, self.name, incarnation))) try: for rank, node in enumerate(self.spec.nodes): - job.procs.append(self.ops.start(node, start_script(self.spec, rank, env, paths))) + job.procs.append(self.ops.start(node, start_script( + self.spec, rank, env, paths, container=self._container(incarnation, rank)))) except Exception: self.release(job) raise @@ -293,10 +405,14 @@ def start_events(self) -> list[Mapping[str, Any]]: return [e for job in self.jobs for e in job.events.all] # -- internals ------------------------------------------------------------------ + def _container(self, incarnation: int, rank: int) -> str | None: + return container_name(self.spec.run_id, self.name, incarnation, rank) if self.spec.docker else None + def _kill(self, incarnation: int) -> bool: - script = kill_script(remote_paths(self.spec.run_id, self.name, incarnation), self.kill_grace_s) + paths = remote_paths(self.spec.run_id, self.name, incarnation) ok = True - for node in self.spec.nodes: + for rank, node in enumerate(self.spec.nodes): + script = kill_script(paths, self.kill_grace_s, container=self._container(incarnation, rank)) try: done = self.ops.run(node, script, timeout=self.kill_grace_s + 60.0) ok &= done.returncode == 0 @@ -357,3 +473,45 @@ def requant_argv(store: str, template: str, *, recipe: str | None = None, awq_sc else: argv.append("--no-awq") return tuple(argv) + + +def inventory_trainer_nodes(args: Any, island: Any) -> tuple[Node, ...]: + """The inventory nodes of an ``ssh:`` trainer island (``island``: its + ``gpu_spec.ClusterSpec``): the nodes with role ``trainer`` (or no roles), + as many as the island has, with its card and card count. Raises before any SSH.""" + nodes = getattr(args, "rl_node_inventory_nodes", None) + if not nodes: + raise SshIslandError("an ssh: trainer island needs --rl-node-inventory (the nodes to start it on)") + from .node_inventory import island_nodes + + trainers = island_nodes(list(nodes), "trainer") + if len(trainers) != island.num_nodes: + raise SshIslandError(f"the ssh: trainer island has {island.num_nodes} node(s); the inventory lists " + f"{len(trainers)} trainer node(s) {[n.name for n in trainers]}") + wrong = [n.name for n in trainers + if n.cards_per_node != island.gpus_per_node or n.card.lower() != str(island.gpu).lower()] + if wrong: + raise SshIslandError(f"inventory trainer nodes {wrong} do not have {island.gpus_per_node} x {island.gpu} " + "as the ssh: island says") + return tuple(trainers) + + +def island_spec_from_inventory(args: Any, island: Any, task: Any) -> SshTrainerSpec: + """The :class:`SshTrainerSpec` of the held-back ``ssh:`` trainer island: the + inventory's trainer nodes run the island's sky task (setup, then run) as + ``run_script`` with the task's environment, in ``--rl-ssh-trainer-workdir`` + (the code pack must be there already: SSH nodes get no sky workdir sync).""" + workdir = getattr(args, "rl_ssh_trainer_workdir", None) + if not workdir: + raise SshIslandError("an ssh: trainer island needs --rl-ssh-trainer-workdir (the yeto checkout on " + "the nodes; SSH nodes get no sky workdir sync)") + run = getattr(task, "run", None) + if not isinstance(run, str) or not run.strip(): + raise SshIslandError("the ssh: trainer island's task has no run script") + setup = getattr(task, "setup", None) or "" + env = {**(getattr(task, "envs", None) or {}), **(getattr(task, "secrets", None) or {})} + return SshTrainerSpec(nodes=inventory_trainer_nodes(args, island), run_id=str(args.cluster_prefix), + workdir=str(workdir), + # sky runs setup and run as two scripts; a failed setup ends the job + run_script=f"bash -c {shlex.quote(setup)} || exit $?\n{run}" if setup.strip() else run, + env={k: str(v) for k, v in env.items() if v is not None}) diff --git a/yeto/launch/trainer_island.py b/yeto/launch/trainer_island.py index b98ea3408..26680f8dc 100644 --- a/yeto/launch/trainer_island.py +++ b/yeto/launch/trainer_island.py @@ -29,6 +29,7 @@ import io import json +import os import threading import time from collections.abc import Callable, Iterable, Mapping @@ -555,11 +556,16 @@ def object_store_binding(args: Any, files: ObjectStore, islands: Any, *, return _binding(args, files, islands, journal_dir=journal_dir, buffer=buffer, unmerged=unmerged) -def head_store_objects(uri: str) -> ObjectStore: - """The head's object API of an ``s3://`` or ``oss://`` checkpoint store.""" - from yeto.cloud import oss_store +def head_store_objects(uri: str, environ: Mapping[str, str] | None = None) -> ObjectStore: + """The head's object API of an ``s3://`` or ``oss://`` checkpoint store. An + ``oss://`` store goes through the relay node of ``oss_relay.RELAY_ENV`` when + it is set (a bucket the head cannot reach directly), else directly.""" + from yeto.cloud import oss_relay, oss_store if oss_store.is_oss_uri(uri): + relay = oss_relay.relay_from_env(os.environ if environ is None else environ) + if relay is not None: + return oss_relay.RelayOssObjectStore(uri, relay) return oss_store.head_objects(uri) from yeto.cloud.s3_store import head_objects diff --git a/yeto/launch/trainer_loop.py b/yeto/launch/trainer_loop.py index 33ddbf63b..68396c6c3 100644 --- a/yeto/launch/trainer_loop.py +++ b/yeto/launch/trainer_loop.py @@ -97,6 +97,10 @@ def add_arguments(group) -> None: "and stops; the fleet does not start it. Every island must reach the " "--rl-checkpoint-store: modal-volume:// on Modal islands, s3:// or oss:// on sky " "islands (not Verda)") + group.add_argument( + "--rl-ssh-trainer-workdir", default=None, metavar="PATH", + help="with an ssh: trainer island: the yeto checkout on the inventory's trainer nodes " + "(inside the container on runtime: docker nodes); the island's run script runs there") group.add_argument( "--rl-trainer-command-poll-seconds", type=float, default=None, help="seconds between two reads of the command file on the trainer island and of " @@ -234,6 +238,15 @@ def _check_trainer_island(args) -> None: if not 0 <= m < len(specs): raise SystemExit(f"--rl-trainer-island {m}: --gpu has {len(specs)} island(s)") cloud = specs[m].cloud + if cloud == "ssh": + from .ssh_trainer_island import SshIslandError, inventory_trainer_nodes + + try: + inventory_trainer_nodes(args, specs[m]) + except SshIslandError as error: + raise SystemExit(f"--rl-trainer-island {m}: {error}") from None + if not getattr(args, "rl_ssh_trainer_workdir", None): + raise SystemExit(f"--rl-trainer-island {m} is an ssh: island: it needs --rl-ssh-trainer-workdir") if cloud == "verda": raise SystemExit(f"--rl-trainer-island {m} is a Verda island; the scheduler launches " "a cluster per trainer start, which the Verda launch path does not do") diff --git a/yeto/launcher.py b/yeto/launcher.py index dca725899..7955bea99 100644 --- a/yeto/launcher.py +++ b/yeto/launcher.py @@ -191,6 +191,15 @@ def _island_task(task_factory, args, spec, m, num_learners, syncer_addr, overrid return task +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": + return None + from .launch.ssh_trainer_island import island_spec_from_inventory + + return island_spec_from_inventory(args, spec, task) + + def warn_if_model_wont_fit(args, specs: list[ClusterSpec]) -> None: weight_gb = MODEL_WEIGHT_GB.get(args.model) if weight_gb is None: @@ -659,7 +668,8 @@ def drain_tapes(limit: float = NO_SYNC_EVENT_DRAIN_S) -> None: args, controller.run, island=None if held_name is None else trainer_loop.TrainerIsland( held_name, specs[held].cloud, held_cfg, modal_ops, sky_task=held_task, - spec=specs[held], events_dir=events_dir), + spec=specs[held], events_dir=events_dir, + ssh=_ssh_trainer_spec(args, specs[held], held_task)), abort=controller.abort) except trainer_loop.TrainerLoopGaveUp as error: # the trainer scheduler released the trainer and stopped; the samplers diff --git a/yeto/rl/adapters/miles/config.py b/yeto/rl/adapters/miles/config.py index 409c802c0..2202d1bc2 100644 --- a/yeto/rl/adapters/miles/config.py +++ b/yeto/rl/adapters/miles/config.py @@ -327,6 +327,7 @@ def _check_sglang_tp(value, config): "model_recipe.training_attention_backend": _ok, "model_recipe.gdn.gated_delta_net": _ok, "model_recipe.gdn.qkv_format": _ok, + "model_recipe.transformer_impl": _ok, **{f"geometry.{name}": _ok for name in _GEOMETRY_FIELDS}, **{ f"geometry.moe.{name}": _ok @@ -771,8 +772,14 @@ def translate_run_config( ) dynamic_filter = _algorithm_filter(config, algorithm) - from yeto.rl.engine.run_config import RECIPE_DEEPSEEK_V4_FLASH, RECIPE_QWEN3_5, RECIPE_QWEN3_8_NEXT + from yeto.rl.engine.run_config import ( + RECIPE_DEEPSEEK_V4_FLASH, + RECIPE_GLM5_NEXT, + RECIPE_QWEN3_5, + RECIPE_QWEN3_8_NEXT, + ) from .lr_schedule import lr_schedule_argv + from .models import glm5_next from .models.deepseek_v4 import ports_recipe as deepseek_v4 geometry = config.geometry @@ -840,6 +847,8 @@ def translate_run_config( "--hidden-dropout", "0.0", ] + elif recipe.name == RECIPE_GLM5_NEXT: + model_name = glm5_next.MODEL_NAME # the model flags replace the geometry below elif recipe.name == RECIPE_DEEPSEEK_V4_FLASH: model_name = deepseek_v4.MODEL_NAME model_recipe_values = deepseek_v4.recipe_argv(trainable) @@ -1130,6 +1139,10 @@ def translate_run_config( for flag in ("--sglang-tp-size", "--sglang-ep-size"): if flag not in values: values.extend((flag, engine)) + if recipe.name == RECIPE_GLM5_NEXT: + # the plugin spec replaces the Bridge-geometry model flags + values = glm5_next.apply_recipe(values, glm5_next.model_argv( + num_layers=geometry.num_layers, 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/models/glm5_next.py b/yeto/rl/adapters/miles/models/glm5_next.py new file mode 100644 index 000000000..cade9faac --- /dev/null +++ b/yeto/rl/adapters/miles/models/glm5_next.py @@ -0,0 +1,129 @@ +"""Miles argv of the GLM5.3-Flash (``glm5_next``) recipe (S25 stage 1, gap 4). + +Model arguments: Miles ``scripts/models/glm5.3-flash.py`` ``model_args()`` at +the pinned Miles commit (cc76e2391, with the s25-no-te fallback a38decedf), +copied as one tuple so the launch layer does not import a Miles script. The +model trains through the Miles plugin spec +``miles_plugins.models.glm5_next.glm5_next:get_glm5_next_spec`` (KDA + DSA MLA +layers, 288 routed experts), not through a Bridge-derived geometry: the +recipe replaces every geometry flag the generic translation derives. + +``transformer_impl``: + +* ``transformer_engine`` (default, main-agent ruling 10-19 (3)): TE layers, + grouped GEMM experts and fused permutation; +* ``local``: the TE-free fallback of N7.5 (Miles ``glm5_next/local_layers.py``); + the TE-only flags are dropped (Megatron raises on ``--moe-permute-fusion`` + without TE) and gradient-accumulation fusion is off. + +Not run on a PPU; the argv was compared with the Miles script by reading only. +""" + +from __future__ import annotations + +from collections.abc import Iterator, Sequence + +MODEL_NAME = "glm5_next" +TRANSFORMER_ENGINE = "transformer_engine" +LOCAL = "local" +TRANSFORMER_IMPLS = (TRANSFORMER_ENGINE, LOCAL) +SPEC = ("miles_plugins.models.glm5_next.glm5_next", "get_glm5_next_spec") +NUM_LAYERS = 45 +FIRST_K_DENSE_REPLACE = 3 + +# N7.5 E6/E7: flags that exist only with TransformerEngine. +TE_ONLY_ARGV: tuple[str, ...] = ("--moe-grouped-gemm", "--moe-permute-fusion") +LOCAL_ARGV: tuple[str, ...] = ("--transformer-impl", LOCAL, "--no-gradient-accumulation-fusion") + +_MODEL_ARGV: tuple[str, ...] = ( + "--hidden-size", "4096", + "--ffn-hidden-size", "12288", + "--num-attention-heads", "64", + "--multi-latent-attention", + "--q-lora-rank", "1536", + "--kv-lora-rank", "512", + "--qk-head-dim", "256", + "--qk-pos-emb-head-dim", "0", + "--v-head-dim", "256", + "--kv-channels", "256", + "--qk-layernorm", + "--num-experts", "288", + "--moe-router-topk", "8", + "--moe-router-score-function", "sigmoid", + "--moe-router-pre-softmax", + "--moe-router-enable-expert-bias", + "--moe-router-bias-update-rate", "0", + "--moe-router-topk-scaling-factor", "2.5", + "--moe-ffn-hidden-size", "2048", + "--moe-shared-expert-intermediate-size", "2048", + "--moe-router-dtype", "fp32", + "--vocab-size", "154880", + "--rotary-base", "10000", + "--rope-type", "rope", + "--no-rope-fusion", + "--position-embedding-type", "rope", + "--normalization", "RMSNorm", + "--norm-epsilon", "1e-5", + "--swiglu", + # the checkpoint sets swiglu_limit=10, which SGLang applies to dense, shared and routed MLPs + "--activation-func-clamp-value", "10", + "--disable-bias-linear", + "--untie-embeddings-and-output-weights", + "--make-vocab-size-divisible-by", "16", + "--enable-experimental", +) + +# scripts/run_glm5_3_flash.py engine arguments the model needs (packed sequences). +RUN_ARGV: tuple[str, ...] = ("--qkv-format", "thd") + +# Flags of the generic translation that the plugin spec does not take (the +# geometry the Bridge derives for other models). +DROPPED_FLAGS = frozenset({"--num-query-groups", "--max-position-embeddings", "--rotary-percent", + "--add-qkv-bias"}) + + +def moe_layer_freq(num_layers: int = NUM_LAYERS, first_k_dense_replace: int = FIRST_K_DENSE_REPLACE) -> str: + """Megatron's ``--moe-layer-freq`` list: dense for the first layers, MoE after + (Miles ``model_args_utils.moe_layer_freq``).""" + if not 0 <= first_k_dense_replace <= num_layers: + raise ValueError(f"first_k_dense_replace {first_k_dense_replace} is outside 0..{num_layers}") + return "[" + ",".join(["0"] * first_k_dense_replace + ["1"] * (num_layers - first_k_dense_replace)) + "]" + + +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}") + return transformer_impl + + +def model_argv(*, num_layers: int = NUM_LAYERS, first_k_dense_replace: int = FIRST_K_DENSE_REPLACE, + 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`).""" + impl = check_transformer_impl(transformer_impl) + return ["--spec", *SPEC, "--num-layers", str(int(num_layers)), *_MODEL_ARGV, + "--moe-layer-freq", moe_layer_freq(num_layers, first_k_dense_replace), + *(TE_ONLY_ARGV if impl == TRANSFORMER_ENGINE else LOCAL_ARGV), *RUN_ARGV] + + +def _flag_groups(argv: Sequence[str]) -> Iterator[list[str]]: + group: list[str] = [] + for token in argv: + if token.startswith("--") and group: + yield group + group = [] + group.append(token) + if group: + yield group + + +def apply_recipe(argv: Sequence[str], recipe: Sequence[str]) -> list[str]: + """``argv`` (``argv[0]`` the script) without every flag the recipe sets or + drops, then the recipe. ``--transformer-impl transformer_engine`` stays + implicit (Megatron's default), so a generic argv that also leaves it out + is not changed by the default recipe.""" + override = {group[0] for group in _flag_groups(recipe)} | DROPPED_FLAGS + if "--transformer-impl" in recipe: # local: gradient-accumulation fusion goes off as well + 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] diff --git a/yeto/rl/adapters/miles/run_flags.py b/yeto/rl/adapters/miles/run_flags.py index 9ec90e7f2..c44008683 100644 --- a/yeto/rl/adapters/miles/run_flags.py +++ b/yeto/rl/adapters/miles/run_flags.py @@ -39,10 +39,17 @@ ), shared_flag( "--rl-model-recipe", - choices=["generic", "deepseek-v4-flash"], + choices=["generic", "deepseek-v4-flash", "glm5_next"], default="generic", help="model-specific Miles memory/kernel contract", ), + shared_flag( + "--rl-transformer-impl", + choices=["transformer_engine", "local"], + default="transformer_engine", + help="Megatron layer implementation of the glm5_next recipe: transformer_engine " + "(default) or local, the TE-free fallback (no grouped GEMM, no permute fusion)", + ), shared_flag( "--rl-training-attention-backend", choices=["auto", "flash", "fused", "unfused"], diff --git a/yeto/rl/adapters/sampler_engines/from_spec.py b/yeto/rl/adapters/sampler_engines/from_spec.py new file mode 100644 index 000000000..1b4486a89 --- /dev/null +++ b/yeto/rl/adapters/sampler_engines/from_spec.py @@ -0,0 +1,54 @@ +"""Engine adapters from a sampler spec's ``engine`` section (the engine sampler +process entry, :mod:`yeto.rl.engine.engine_sampler_entry`, calls :func:`build` +through ``engine.factory``). + +Section forms: + +* ``{"adapter": NAME, "server": {...}, "options": {...}}``: an HTTP-server + adapter; ``server`` = :class:`~.server.ServerConfig` fields, ``options`` = the + adapter's options dataclass (``SGLangOptions``, ``VllmOptions``); +* ``{"adapter": NAME, "kwargs": {...}}``: an adapter built from keywords (the + CPU ``fake`` engine). + +``NAME`` is a key of :data:`~yeto.rl.adapters.sampler_engines.ENGINE_ADAPTERS`. +""" + +from __future__ import annotations + +import dataclasses +import typing +from collections.abc import Mapping +from pathlib import Path +from typing import Any + +from . import engine_class +from .server import ServerConfig + +FACTORY = "yeto.rl.adapters.sampler_engines.from_spec:build" # value of the spec's engine.factory + + +def options_type(cls: type) -> type | None: + """The options dataclass an adapter's ``__init__(config, options)`` takes.""" + hint = typing.get_type_hints(cls.__init__).get("options") + found = [m for m in (typing.get_args(hint) or (hint,)) if dataclasses.is_dataclass(m)] + return found[0] if found else None + + +def build(section: Mapping[str, Any]) -> Any: + """The adapter of one ``engine`` section (module doc). Raises ValueError / + KeyError / TypeError on a section the adapter does not accept.""" + cls = engine_class(str(section["adapter"])) + if "server" not in section: + return cls(**dict(section.get("kwargs") or {})) + server = dict(section["server"]) + server["base_dir"] = Path(server["base_dir"]) + server["extra_args"] = tuple(server.get("extra_args") or ()) + options = dict(section.get("options") or {}) + options_cls = options_type(cls) + if options_cls is None: + if options: + raise ValueError(f"adapter {section['adapter']!r} takes no options, the section gives {sorted(options)}") + return cls(ServerConfig(**server)) + # JSON has no tuples; the options dataclasses hold tuples + options = {k: tuple(v) if isinstance(v, list) else v for k, v in options.items()} + return cls(ServerConfig(**server), options_cls(**options)) diff --git a/yeto/rl/engine/engine_sampler.py b/yeto/rl/engine/engine_sampler.py index 1ec63a1bf..7ef136bea 100644 --- a/yeto/rl/engine/engine_sampler.py +++ b/yeto/rl/engine/engine_sampler.py @@ -21,7 +21,8 @@ a failure there stops the island), the follower and the loop. The caller starts the resource sampler of its device family (``telemetry.ResourceSampler`` with ``card_telemetry.collector_for``) and -calls ``island.run()``. +calls ``island.run()``. The process entry (``python -m yeto.rl.engine.engine_sampler +--spec SPEC.json``, :func:`main`) does both from a spec file (:mod:`.engine_sampler_entry`). """ from __future__ import annotations @@ -196,3 +197,15 @@ def build_engine_sampler(config: EngineSamplerConfig, *, engine: SamplerEngine, 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) + + +def main(argv: Sequence[str] | None = None) -> int: + """Process entry: build the island from a spec file and run it + (:mod:`.engine_sampler_entry`).""" + from .engine_sampler_entry import main as entry_main # noqa: PLC0415 - keeps this module's imports lean + + return entry_main(argv) + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/yeto/rl/engine/engine_sampler_entry.py b/yeto/rl/engine/engine_sampler_entry.py new file mode 100644 index 000000000..1eda00f9b --- /dev/null +++ b/yeto/rl/engine/engine_sampler_entry.py @@ -0,0 +1,274 @@ +"""Process entry of the engine sampler island (``python -m +yeto.rl.engine.engine_sampler --spec SPEC.json``; S25 stage 1, gap 3). + +:func:`~.engine_sampler.build_engine_sampler` wires the island from objects; +this module builds those objects from one JSON spec file that the launch +layer writes (the stage-1 plan writes it with ``yeto-stage1 plan``), then runs +the island to its data end. It adds no sampling logic of its own. + +The spec (:class:`SamplerSpec`, every key below is required unless marked): + +* ``sampler_id``: the sampler's buffer shard and sample id prefix; +* ``engine``: ``{"factory": "module:callable", ...}``; the factory (adapter + layer, e.g. ``yeto.rl.adapters.sampler_engines.from_spec:build``) builds the + engine from this section; +* ``sampling``: :class:`~.sampler_engine.SamplingParams` fields; +* ``num_rollouts``, ``max_backlog_groups``, ``max_prompt_tokens``; + optional ``base_version`` (0), ``base_dir`` (the engine's base weights), + ``min_round_s``, ``poll_s``, ``max_consecutive_load_failures``; +* ``store``: the run's checkpoint store (``store_for``; an ``oss://`` store + needs the store environment of :mod:`yeto.cloud.oss_store`); + ``release_dir``: the release directory the sampler follows (optional, + default ``policy-releases``); ``buffer_root``: the store holding + ``sample-buffer/`` (optional, default ``store``); +* ``prompts``: ``{"path": JSONL, "input_key", "label_key", "groups_per_round", + "format": "text" | "chat"}`` -- round ``r`` takes the + ``groups_per_round`` rows after row ``r * groups_per_round`` (wrapping); + ``chat`` applies the tokenizer's chat template of ``tokenizer`` (default + ``base_dir``) to a message list (``trust_remote_code``: optional, false); +* ``reward``: ``"module:function"`` called as ``function(item, completion)`` + -> float, or a Miles reward (``async score(args, sample)``) when + ``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. + +Nothing here was run on a PPU or NPU node. +""" + +from __future__ import annotations + +import argparse +import asyncio +import importlib +import inspect +import json +import sys +import time +import types +from collections.abc import Callable, Mapping, Sequence +from pathlib import Path +from typing import Any + +from .engine_sampler import EngineSamplerConfig, PromptItem, Reward, build_engine_sampler +from .sampler_checks import PostLoadChecks +from .sampler_engine import Completion, SamplerEngine, SamplingParams + +PROMPT_FORMATS = ("text", "chat") +REWARD_KINDS = ("neutral", "miles") + + +class SamplerSpecError(ValueError): + """The sampler spec file is incomplete or contradicts itself.""" + + +def _require(doc: Mapping[str, Any], key: str, where: str = "spec") -> Any: + if key not in doc: + raise SamplerSpecError(f"{where}.{key} is required") + return doc[key] + + +def _import(target: str) -> Any: + module, _, attr = str(target).partition(":") + if not module or not attr: + raise SamplerSpecError(f"{target!r} is not 'module:attribute'") + return getattr(importlib.import_module(module), attr) + + +# ------------------------------------------------------------------- engine + +def build_engine(spec: Mapping[str, Any]) -> SamplerEngine: + """The engine adapter of ``spec["engine"]``: ``factory`` (``module:callable``, + an adapter-layer function) builds it from the engine section. The neutral + core names no adapter (import boundaries, ``yeto/rl/engine/README.md``).""" + engine = _require(spec, "engine") + factory = _import(_require(engine, "factory", "engine")) + try: + return factory(engine) + except (KeyError, TypeError, ValueError) as error: + raise SamplerSpecError(f"engine: {error}") from None + + +# ------------------------------------------------------------------ prompts + +class JsonlPrompts: + """``prompts(rollout_id)`` over a local JSONL file (module doc).""" + + def __init__(self, path: str | Path, *, input_key: str, label_key: str | None, groups_per_round: int, + render: Callable[[Any], str] = str) -> None: + self.path = Path(path) + self.rows = [json.loads(line) for line in self.path.read_text(encoding="utf-8").splitlines() + if line.strip()] + if not self.rows: + raise SamplerSpecError(f"prompts file {self.path} has no rows") + if int(groups_per_round) < 1: + raise SamplerSpecError("prompts.groups_per_round must be >= 1") + self.input_key, self.label_key = input_key, label_key + self.groups_per_round = int(groups_per_round) + self.render = render + + def __call__(self, rollout_id: int) -> list[PromptItem]: + start = int(rollout_id) * self.groups_per_round + items = [] + for offset in range(self.groups_per_round): + index = (start + offset) % len(self.rows) + row = self.rows[index] + if self.input_key not in row: + raise SamplerSpecError(f"{self.path} row {index} has no {self.input_key!r}") + items.append(PromptItem(prompt_id=f"row{index}", prompt=self.render(row[self.input_key]), + label=row.get(self.label_key) if self.label_key else None, + metadata={"row": index})) + return items + + +def chat_renderer(tokenizer_dir: str | Path, *, remote_code: bool = False) -> Callable[[Any], str]: + from transformers import AutoTokenizer # noqa: PLC0415 - only the chat format needs it + + tokenizer = AutoTokenizer.from_pretrained(str(tokenizer_dir), trust_remote_code=remote_code) + + def render(messages: Any) -> str: + if isinstance(messages, str): + messages = [{"role": "user", "content": messages}] + return tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True) + + return render + + +def build_prompts(spec: Mapping[str, Any]) -> JsonlPrompts: + prompts = _require(spec, "prompts") + fmt = str(prompts.get("format", "text")) + if fmt not in PROMPT_FORMATS: + raise SamplerSpecError(f"prompts.format {fmt!r} is not one of {PROMPT_FORMATS}") + render: Callable[[Any], str] = str + if fmt == "chat": + tokenizer = prompts.get("tokenizer") or spec.get("base_dir") \ + or (spec.get("engine") or {}).get("server", {}).get("base_dir") + if not tokenizer: + raise SamplerSpecError("prompts.format chat needs prompts.tokenizer or a base_dir") + render = chat_renderer(tokenizer, remote_code=bool(prompts.get("trust_remote_code", False))) + return JsonlPrompts(_require(prompts, "path", "prompts"), input_key=str(_require(prompts, "input_key", "prompts")), + label_key=prompts.get("label_key"), + groups_per_round=int(_require(prompts, "groups_per_round", "prompts")), render=render) + + +# ------------------------------------------------------------------- reward + +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.""" + + def reward(item: PromptItem, completion: Completion) -> float: + sample = types.SimpleNamespace(prompt=item.prompt, response=completion.text, label=item.label, + metadata=dict(item.metadata)) + value = score(None, sample) + return float(asyncio.run(value) if inspect.isawaitable(value) else value) + + return reward + + +def build_reward(spec: Mapping[str, Any]) -> Reward: + kind = str(spec.get("reward_kind", "neutral")) + if kind not in REWARD_KINDS: + raise SamplerSpecError(f"reward_kind {kind!r} is not one of {REWARD_KINDS}") + function = _import(_require(spec, "reward")) + return miles_reward(function) if kind == "miles" else function + + +# -------------------------------------------------------------------- store + +def build_stores(spec: Mapping[str, Any], environ: Mapping[str, str] | None = None) -> tuple[Any, Any]: + """(release store, sample buffer) of the spec.""" + from .policy_release import RELEASE_DIR, PolicyReleaseStore # noqa: PLC0415 + from .resume import store_for # noqa: PLC0415 + from .sample_buffer import INDEX, FileSampleBuffer, needs_shards, shard_index # noqa: PLC0415 + + store = str(_require(spec, "store")) + releases = PolicyReleaseStore(store_for(store, environ=environ), + release_dir=str(spec.get("release_dir") or RELEASE_DIR)) + buffer_store = store_for(str(spec.get("buffer_root") or store), environ=environ) + sampler_id = str(_require(spec, "sampler_id")) + index = shard_index(sampler_id) if needs_shards(buffer_store) else INDEX + return releases, FileSampleBuffer(buffer_store, batch_size=1, index=index) # as open_buffer(sampler=True) + + +def tape_writer(path: str | Path, sampler_id: str) -> Callable[..., None]: + from yeto.rl.event_echo import append_record # noqa: PLC0415 + + tape = Path(path).expanduser() + tape.parent.mkdir(parents=True, exist_ok=True) + + def emit(event: str, **fields: Any) -> None: + append_record(tape, {"sampler_id": sampler_id, "time_unix": time.time(), "event": event, **fields}) + + return emit + + +# ---------------------------------------------------------------------- run + +def island_config(spec: Mapping[str, Any]) -> EngineSamplerConfig: + sampling = dict(_require(spec, "sampling")) + sampling["stop"] = tuple(sampling.get("stop") or ()) + base_dir = spec.get("base_dir") + return EngineSamplerConfig( + sampler_id=str(_require(spec, "sampler_id")), num_rollouts=int(_require(spec, "num_rollouts")), + max_backlog_groups=int(_require(spec, "max_backlog_groups")), + sampling=SamplingParams(**sampling), max_prompt_tokens=int(_require(spec, "max_prompt_tokens")), + base_version=int(spec.get("base_version", 0)), base_dir=None if base_dir is None else Path(base_dir), + min_round_s=spec.get("min_round_s"), poll_s=float(spec.get("poll_s", 5.0)), + max_consecutive_load_failures=int(spec.get("max_consecutive_load_failures", 3))) + + +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).""" + config = island_config(spec) + emit = tape_writer(_require(spec, "event_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 {})) + engine = build_engine(spec) + if resource_sampler is None: + from .telemetry import ResourceSampler # noqa: PLC0415 + + def resource_sampler(sink: Callable[..., None]) -> Any: + return ResourceSampler(lambda event, **fields: sink(event, **fields)).start() + + telemetry = resource_sampler(emit) + try: + island = build_engine_sampler(config, engine=engine, releases=releases, buffer=buffer, emit=emit, + prompts=prompts, reward=reward, checks=checks) + return island.run() + finally: + try: + engine.stop() + finally: + if telemetry is not None: + telemetry.stop() + + +def load_spec(path: str | Path) -> dict[str, Any]: + try: + doc = json.loads(Path(path).read_text(encoding="utf-8")) + except (OSError, ValueError) as error: + raise SamplerSpecError(f"cannot read the sampler spec {path}: {error}") from None + if not isinstance(doc, dict): + raise SamplerSpecError(f"{path}: the sampler spec must be a JSON object") + return doc + + +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") + parser.add_argument("--spec", required=True, help="the sampler spec JSON (module doc)") + args = parser.parse_args(argv) + try: + spec = load_spec(args.spec) + rounds = run(spec) + except SamplerSpecError as error: + print(f"[engine-sampler] spec rejected: {error}", file=sys.stderr) + return 2 + print(f"[engine-sampler] done: {rounds} round(s)", flush=True) + return 0 diff --git a/yeto/rl/engine/run_config.py b/yeto/rl/engine/run_config.py index 7e105a183..7b596bbda 100644 --- a/yeto/rl/engine/run_config.py +++ b/yeto/rl/engine/run_config.py @@ -37,6 +37,10 @@ # Qwen3.8-Flash-Next (qwen4_exp): GDN + QSA + 512-expert MoE through the native # Miles plugin (miles_plugins/models/qwen3_8_next), never the Qwen3.5 spec. RECIPE_QWEN3_8_NEXT = "qwen3_8_next" +# GLM5.3-Flash (glm5_next): KDA + DSA MLA + 288-expert MoE through the Miles +# plugin spec (adapters/miles/models/glm5_next.py), full-parameter. +RECIPE_GLM5_NEXT = "glm5_next" +TRANSFORMER_IMPLS = ("transformer_engine", "local") GATED_DELTA_NET = "gated_delta_net" @@ -132,10 +136,13 @@ class GdnRecipe: @dataclass(frozen=True) class ModelRecipe: - name: str # RECIPE_GENERIC | RECIPE_QWEN3_5 | RECIPE_DEEPSEEK_V4_FLASH | RECIPE_QWEN3_8_NEXT + name: str # RECIPE_GENERIC | RECIPE_QWEN3_5 | RECIPE_DEEPSEEK_V4_FLASH | RECIPE_QWEN3_8_NEXT | RECIPE_GLM5_NEXT provider_class: str training_attention_backend: str gdn: GdnRecipe + # Megatron --transformer-impl of a recipe that chooses it (glm5_next: the + # N7.5 TE-free fallback is "local"); the other recipes keep their own. + transformer_impl: str = "transformer_engine" @dataclass(frozen=True) @@ -533,6 +540,25 @@ def _rules(): return backends.module("run_config_rules", backends.DEFAULT_BACKEND) +def full_parameter_layout_problem(*, data_parallel: int, actor_num_nodes: int, rollout_gpus: Any, + rollout_gpus_per_engine: Any) -> str | None: + """Why the Miles full-parameter mode refuses this layout (Milestone 1: one + node, DP=1, dedicated evenly partitioned engines); None when it accepts it.""" + if data_parallel != 1: + return "dense full-parameter GRPO requires DP=1" + if ( + actor_num_nodes != 1 + or type(rollout_gpus) is not int + or rollout_gpus < 1 + or type(rollout_gpus_per_engine) is not int + or rollout_gpus_per_engine < 1 + or rollout_gpus % rollout_gpus_per_engine + ): + return ("Milestone-1 dense full-parameter mode requires one node and " + "dedicated, evenly partitioned inference engines") + return None + + def _resolve_ref_load(args, model_path) -> str: """Reference-model path; the checkpoint-format rules belong to the backend (4.2).""" return _rules().resolve_ref_load(args, model_path) @@ -649,24 +675,14 @@ def resolve_rl_run_config( "expert-full tuning requires the DeepSeek V4 recipe and attention LoRA" ) data_parallel = actor_gpus // model_parallel - if parameter_mode == "full" and data_parallel != 1: - raise ValueError("dense full-parameter GRPO requires DP=1") dedicated_rollout_gpus: int | None = None if parameter_mode == "full": rollout_gpus = getattr(args, "rollout_num_gpus", None) rollout_gpus_per_engine = getattr(args, "rollout_num_gpus_per_engine", None) - if ( - args.actor_num_nodes != 1 - or type(rollout_gpus) is not int - or rollout_gpus < 1 - or type(rollout_gpus_per_engine) is not int - or rollout_gpus_per_engine < 1 - or rollout_gpus % rollout_gpus_per_engine - ): - raise ValueError( - "Milestone-1 dense full-parameter mode requires one node and " - "dedicated, evenly partitioned inference engines" - ) + if problem := full_parameter_layout_problem( + data_parallel=data_parallel, actor_num_nodes=args.actor_num_nodes, + rollout_gpus=rollout_gpus, rollout_gpus_per_engine=rollout_gpus_per_engine): + raise ValueError(problem) if qwen35_recipe and ( rollout_gpus_per_engine != 1 or getattr(args, "sglang_tp_size", None) not in {None, 1} @@ -736,6 +752,11 @@ def resolve_rl_run_config( "DeepSeek V4 Flash recipe requires the 43-layer MoE/MLA provider" ) recipe_name, attention_backend = RECIPE_DEEPSEEK_V4_FLASH, "flash" + elif recipe == RECIPE_GLM5_NEXT: + if parameter_mode != "full": + raise ValueError("the glm5_next recipe is full-parameter only") + # 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: if parameter_mode != "lora": raise ValueError("the Qwen3.8-Flash-Next recipe is LoRA-only") @@ -774,6 +795,13 @@ def resolve_rl_run_config( ) attention_backend = backend_override + transformer_impl = str(getattr(args, "rl_transformer_impl", None) or TRANSFORMER_IMPLS[0]) + if transformer_impl not in TRANSFORMER_IMPLS: + raise ValueError(f"--rl-transformer-impl {transformer_impl!r} is not one of {TRANSFORMER_IMPLS}") + if transformer_impl != TRANSFORMER_IMPLS[0] and recipe_name != RECIPE_GLM5_NEXT: + raise ValueError(f"--rl-transformer-impl {transformer_impl} is supported by the " + f"{RECIPE_GLM5_NEXT} recipe only") + _validate_callable_spec(args.reward_function) eval_config = _resolve_eval( @@ -885,6 +913,7 @@ def resolve_rl_run_config( provider_class=type(provider).__name__, training_attention_backend=attention_backend, gdn=gdn, + transformer_impl=transformer_impl, ), geometry=geometry, parallel=ParallelLayout( diff --git a/yeto/rl/stage1.py b/yeto/rl/stage1.py index 5902cb3db..05d7a6a30 100644 --- a/yeto/rl/stage1.py +++ b/yeto/rl/stage1.py @@ -16,14 +16,20 @@ * nodes and roles: trainer nodes, sampler nodes, card counts, rendezvous addresses; * trainer, first backend Miles/Megatron: the parallel layout (TP x PP x EP, - the pipeline layer split), the Megatron arguments the glm5_next recipe must - give, the yeto Miles island argv; second backend verl/FSDP2: the trainer - overrides (written, not the first backend); -* sampler: the SGLang TP8 server argv and environment from the - ``sglang-server`` adapter itself (:mod:`yeto.rl.adapters.sampler_engines`); + the pipeline layer split), the model arguments of the yeto glm5_next recipe + (:mod:`yeto.rl.adapters.miles.models.glm5_next`), the parallel and optimizer + arguments, the yeto Miles island argv, and the island as an SSH trainer + island (:class:`~yeto.launch.ssh_trainer_island.SshTrainerSpec`: Ray on + every node in a docker container, the island entry on node 0); second + backend verl/FSDP2: the trainer overrides (written, not the first backend); +* sampler: the engine sampler island's spec file and command + (``python -m yeto.rl.engine.engine_sampler``), which starts the SGLang TP8 + server of the ``sglang-server`` adapter (:mod:`yeto.rl.adapters.sampler_engines`); * re-quantizer: the :mod:`yeto.rl.quant.release` argv (N8.5 MXFP4 RTN), with the release directory named after the scheme; -* store: the run directory on OSS and the OSS staging paths of the weights; +* store: the run directory on OSS, the OSS staging paths of the weights, and + how the head reaches the store (``store_access``: through a relay node, + :mod:`yeto.cloud.oss_relay`, or directly); * steps: the ordered per-node scripts of the run (preflight, staging, sampler, trainer). @@ -38,6 +44,7 @@ from __future__ import annotations import argparse +import dataclasses import importlib import importlib.util import ipaddress @@ -51,7 +58,9 @@ from typing import Any from yeto.launch.node_inventory import InventoryError, Node, load_inventory -from yeto.launch.topology import island_comm_env +from yeto.rl.adapters.miles.models import glm5_next +from yeto.rl.adapters.sampler_engines.from_spec import FACTORY as ENGINE_FACTORY +from yeto.rl.engine.telemetry import DEVICE_FAMILY_ENV PLAN_SCHEMA = "yeto.stage1/v1" BACKENDS = ("miles", "verl") @@ -59,6 +68,8 @@ RAY_PORT = 6379 HOST_USABLE_FRACTION = 0.85 # as yeto.rl.fullparam.memory.TrainPlan.host_usable_fraction RELEASE_DIR_PREFIX = "policy-releases-" +HEAD_ACCESS = ("relay", "direct") # how the head reaches the store (store_access.head) +RAY_JOIN_TIMEOUT_S = 600 class Stage1Error(ValueError): @@ -132,9 +143,11 @@ class Step: node: str | None script: str purpose: str + service: bool = False # runs for the whole run (started, not waited for) def record(self) -> dict[str, Any]: - return {"name": self.name, "node": self.node, "purpose": self.purpose, "script": self.script} + return {"name": self.name, "node": self.node, "purpose": self.purpose, "script": self.script, + "service": self.service} @dataclass @@ -144,6 +157,7 @@ class Expanded: sampler: dict[str, Any] = field(default_factory=dict) requant: dict[str, Any] = field(default_factory=dict) store: dict[str, Any] = field(default_factory=dict) + trainer_spec: Any = None # the SshTrainerSpec the trainer steps come from steps: list[Step] = field(default_factory=list) errors: list[str] = field(default_factory=list) gaps: list[str] = field(default_factory=list) @@ -203,8 +217,8 @@ def megatron_static_gib(model: Mapping[str, Any], layout: Mapping[str, int]) -> def megatron_args(model: Mapping[str, Any], miles: Mapping[str, Any], layout: Mapping[str, int], sampling: Mapping[str, Any]) -> list[str]: - """The Megatron arguments the Miles glm5_next recipe must pass on top of the - model arguments (Miles ``scripts/models/glm5.3-flash.py``).""" + """The parallel, recompute and optimizer arguments of the run, on top of the + model arguments of the glm5_next recipe (``glm5_next.model_argv``).""" split = [int(x) for x in miles["pp_layers"]] args = ["--tensor-model-parallel-size", str(layout["tp"]), "--pipeline-model-parallel-size", str(layout["pp"]), @@ -221,20 +235,14 @@ def megatron_args(model: Mapping[str, Any], miles: Mapping[str, Any], layout: Ma "--max-tokens-per-gpu", str(int(miles["max_tokens_per_gpu"])), "--optimizer", "adam", "--lr", repr(float(miles["lr"])), "--lr-decay-style", "constant", "--use-distributed-optimizer", "--bf16", - "--rollout-max-response-len", str(int(sampling["max_tokens"])), - "--transformer-impl", str(miles["transformer_impl"])] + "--rollout-max-response-len", str(int(sampling["max_tokens"]))] fraction = float(miles.get("optimizer_offload_fraction") or 0.0) if fraction > 0: args += ["--optimizer-cpu-offload", "--optimizer-offload-fraction", repr(fraction), "--overlap-cpu-optimizer-d2h-h2d"] - if miles["transformer_impl"] == "local": # N7.5 fallback: no TE kernels - args.append("--no-gradient-accumulation-fusion") return args -TE_ONLY_MODEL_ARGS = ("--moe-grouped-gemm", "--moe-permute-fusion") # N7.5 E6, E7 - - def verl_overrides(verl: Mapping[str, Any], trainer_nodes: Sequence[Node], task: Mapping[str, Any], sampling: Mapping[str, Any], model_dir: str) -> list[str]: """verl trainer overrides of the FSDP2 backend (rollout runs on the yeto @@ -359,13 +367,12 @@ def expand(plan: Stage1Plan, *, hf_config: str | Path | None = None) -> Expanded per_node = trainers[0].cards_per_node if trainers else 0 miles = _section(trainer_cfg, "miles") layout = megatron_layout(miles, cards, per_node) - common_env = {"YETO_RL_DEVICE_FAMILY": trainers[0].device_family, "YETO_RUN_ID": plan.run_id, + common_env = {DEVICE_FAMILY_ENV: trainers[0].device_family, "YETO_RUN_ID": plan.run_id, "YETO_RL_STORE": store, "HF_HUB_OFFLINE": "1"} - comm = island_comm_env(trainers[0].device_family, trainers[0].fabric, trainers[0].interface, - num_nodes=len(trainers)) - trainer_env = {**common_env, **comm, "TORCHINDUCTOR_COMPILE_THREADS": "1"} + trainer_env = {**common_env, "TORCHINDUCTOR_COMPILE_THREADS": "1"} model_dir = str(model["bf16_source"]["local_dir"]) - state_dir = f"/data/yeto/{plan.run_id}" + state_dir = state_dir_of(plan) + impl = str(miles["transformer_impl"]) 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"]), @@ -376,6 +383,7 @@ def expand(plan: Stage1Plan, *, hf_config: str | Path | None = None) -> Expanded "--inner-lr", repr(float(miles["lr"])), "--seed", str(int(plan.doc.get("seed", 0))), "--parameter-mode", "full", "--rl-sample-role", "trainer", "--rl-resume-store", store, + "--rl-model-recipe", glm5_next.MODEL_NAME, "--rl-transformer-impl", impl, "--tensor-parallel", str(layout["tp"]), "--pipeline-parallel", str(layout["pp"]), "--expert-parallel", str(layout["ep"]), "--miles-root", str(miles["workdir"]), @@ -391,13 +399,18 @@ 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"]))] + try: + model_args = glm5_next.model_argv(num_layers=int(model["num_layers"]), transformer_impl=impl) + except ValueError as error: + raise Stage1Error(f"plan: trainer.miles.transformer_impl: {error}") from None out.trainer = { "backend": backend, "nodes": [n.name for n in trainers], "cards": cards, "image": images.get("trainer"), "env": trainer_env, "miles": {"commit": miles["commit"], "megatron": miles["megatron"], "mbridge": miles["mbridge"], "layout": layout, "static_memory_estimate": megatron_static_gib(model, layout), + "model_args": model_args, "megatron_args": megatron_args(model, miles, layout, sampling), - "drop_model_args": list(TE_ONLY_MODEL_ARGS) if miles["transformer_impl"] == "local" else [], + "drop_model_args": list(glm5_next.TE_ONLY_ARGV) if impl == glm5_next.LOCAL else [], "island_argv": island_argv}, "verl": {"commit": _section(trainer_cfg, "verl").get("commit"), "overrides": verl_overrides(_section(trainer_cfg, "verl"), trainers, task, sampling, model_dir)}, @@ -409,8 +422,12 @@ def expand(plan: Stage1Plan, *, hf_config: str | Path | None = None) -> Expanded sampler_runs = [] for node in samplers: engine = _sampler_engine(plan, node) - sampler_runs.append({"node": node.name, "argv": engine.argv(engine.config.base_dir), - "env": {**common_env, "YETO_RL_DEVICE_FAMILY": node.device_family, **engine.env()}, + env = {**common_env, DEVICE_FAMILY_ENV: node.device_family, **engine.env()} + sampler_runs.append({"node": node.name, "engine_argv": engine.argv(engine.config.base_dir), "env": env, + "spec_path": f"{state_dir}/sampler-{node.name}.json", + "spec": sampler_spec(plan, node, engine), + "argv": ["python3", "-m", "yeto.rl.engine.engine_sampler", + "--spec", f"{state_dir}/sampler-{node.name}.json"], "url": f"http://{rendezvous.get(node.name)}:{engine.config.port}"}) out.sampler = {"engine": sampler_cfg["engine"], "engine_version": sampler_cfg["engine_version"], "image": images.get("sampler"), "topology": sampler_cfg["topology"], @@ -426,12 +443,119 @@ def expand(plan: Stage1Plan, *, hf_config: str | Path | None = None) -> Expanded out.requant = {"node": trainers[0].name if trainers else None, "scheme": quant["scheme"], "format": quant["format"], "argv": argv, "loop_argv": [*argv, "--loop"]} + out.store["head_access"] = head_store_access(plan) + out.trainer_spec = trainer_island_spec(plan, out, trainers) if backend == "miles" and trainers else None out.steps = _steps(plan, out, trainers, samplers) out.errors = check(plan, out) out.gaps = gaps(plan, out) return out +def state_dir_of(plan: Stage1Plan) -> str: + """The run's state directory on every node (``/data`` is on every node).""" + return f"/data/yeto/{plan.run_id}" + + +def sampler_spec(plan: Stage1Plan, node: Node, engine: Any) -> dict[str, Any]: + """The engine sampler island's spec file (``engine_sampler_entry``).""" + sampling, task, quant = plan.section("sampling"), plan.section("task"), plan.section("quant") + config, options = engine.config, engine.options + server = {f.name: getattr(config, f.name) for f in dataclasses.fields(config)} + server["base_dir"] = str(config.base_dir) + server["extra_args"] = list(config.extra_args) + server["extra_env"] = dict(config.extra_env) + state_dir = state_dir_of(plan) + return { + "sampler_id": f"{plan.run_id}-{node.name}", + "engine": {"factory": ENGINE_FACTORY, "adapter": f"{plan.section('sampler')['engine']}-server", + "server": server, + "options": {k: list(v) if isinstance(v, tuple) else v + for k, v in dataclasses.asdict(options).items()}}, + "sampling": {"max_tokens": int(sampling["max_tokens"]), "temperature": float(sampling["temperature"]), + "top_p": float(sampling["top_p"]), "top_k": int(sampling["top_k"]), + "n": int(task["n_samples_per_prompt"]), "seed": int(plan.doc.get("seed", 0)), + **{k: float(sampling.get(k, v)) for k, v in NEUTRAL_PENALTIES.items()}}, + "num_rollouts": int(plan.doc["rounds"]), + # decision T-GAPS 10-19: one round ahead of the trainer at most (stage 1 runs on-policy) + "max_backlog_groups": int(task["rollout_batch_size"]), + "max_prompt_tokens": int(sampling["max_prompt_len"]), + "base_version": 0, "base_dir": str(config.base_dir), + "store": run_store(plan), "release_dir": str(quant["release_dir"]), + "prompts": {"path": str(_get(task, "prompts_jsonl", "task")), "input_key": str(_get(task, "input_key", "task")), + "label_key": str(_get(task, "label_key", "task")), + "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"])}, + "event_tape": f"{state_dir}/sampler-{node.name}/events.jsonl", + } + + +def head_store_access(plan: Stage1Plan) -> dict[str, Any]: + """How the head on this server reaches the run store (``store_access``): + ``relay`` = ``ossutil`` on a relay node over SSH (:mod:`yeto.cloud.oss_relay`; + main-agent ruling 10-19), ``direct`` = the OSS endpoint from this server.""" + from yeto.cloud.oss_relay import RELAY_ENV, Relay, relay_shell # noqa: PLC0415 + + access = plan.doc.get("store_access") or {"head": "direct"} + mode = str(access.get("head", "direct")) + if mode not in HEAD_ACCESS: + raise Stage1Error(f"plan: store_access.head {mode!r} is not one of {HEAD_ACCESS}") + if mode == "direct": + return {"mode": mode, "env": {}} + name = str(access.get("relay") or "") + node = next((n for n in plan.nodes if n.name == name), None) + if node is None: + raise Stage1Error(f"plan: store_access.relay {name!r} is not an inventory node") + relay = Relay(node.host, str(access.get("ossutil") or "ossutil")) + return {"mode": mode, "relay": node.name, "relay_host": node.host, "env": {RELAY_ENV: node.host}, + "check": relay_shell(relay, f"{relay.ossutil} ls -s {shlex.quote(run_store(plan) + '/')} | tail -n 3")} + + +def trainer_island_spec(plan: Stage1Plan, out: Expanded, trainers: Sequence[Node]) -> Any: + """The Miles trainer island as an SSH trainer island (``SshTrainerSpec``): + Ray on every trainer node, the island entry on node 0 after every node + joined, the re-quantizer as node 0's side process; one docker container per + node. Unset image or device flags show as placeholders (gaps N0.2/N0.2b).""" + from yeto.launch.ssh_trainer_island import SshTrainerSpec # noqa: PLC0415 + + images, rendezvous = plan.section("images"), plan.section("rendezvous") + nodes = tuple(n if n.address else dataclasses.replace(n, address=str(rendezvous.get(n.name))) + for n in trainers) + island = shlex.join(out.trainer["miles"]["island_argv"]) + join = (f"python3 -c {shlex.quote(_RAY_JOIN)} \"$SKYPILOT_NUM_NODES\" {RAY_JOIN_TIMEOUT_S}") + script = ( + "set -eu\n" + "head=$(echo \"$SKYPILOT_NODE_IPS\" | sed -n 1p)\n" + "me=$(echo \"$SKYPILOT_NODE_IPS\" | sed -n \"$((SKYPILOT_NODE_RANK + 1))p\")\n" + "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" + f" exec {island}\n" + "fi\n" + f"exec ray start --address=\"$head:{RAY_PORT}\" --node-ip-address=\"$me\" " + "--num-gpus=\"$SKYPILOT_NUM_GPUS_PER_NODE\" --block\n") + try: + return SshTrainerSpec( + nodes=nodes, run_id=plan.run_id, workdir=str(_section(plan.section("trainer"), "miles")["workdir"]), + run_script=script, env=out.trainer["env"], requant=tuple(out.requant["argv"]), + image=images.get("trainer") or "", + docker_args=tuple(images.get("docker_device_args") or ("",))) + except Exception as error: # SshIslandError: reported as a check error, not a crash + raise Stage1Error(f"trainer island: {error}") from None + + +# Rank 0 waits until every trainer node joined Ray (argv: node count, timeout seconds). +_RAY_JOIN = ("import sys, time, ray\n" + "ray.init(address=\"auto\")\n" + "want, deadline = int(sys.argv[1]), time.time() + float(sys.argv[2])\n" + "alive = lambda: sum(n[\"Alive\"] for n in ray.nodes())\n" + "while alive() < want:\n" + " if time.time() > deadline:\n" + " sys.exit(f\"only {alive()} of {want} nodes joined Ray\")\n" + " time.sleep(5)\n") + + def _verl_memory(trainers: Sequence[Node], plan: Stage1Plan, sampling: Mapping[str, Any], hf_config: str | Path) -> dict[str, Any]: from yeto.rl.fullparam.memory import estimate, plan_for_nodes, shape_from_hf_config # noqa: PLC0415 @@ -448,8 +572,8 @@ def _verl_memory(trainers: Sequence[Node], plan: Stage1Plan, sampling: Mapping[s def _steps(plan: Stage1Plan, out: Expanded, trainers: Sequence[Node], samplers: Sequence[Node]) -> list[Step]: from yeto.hw.card_tools import preflight_commands # noqa: PLC0415 + from yeto.launch.ssh_trainer_island import container_name, remote_paths, start_script # noqa: PLC0415 - rendezvous = plan.section("rendezvous") steps: list[Step] = [] for node in plan.nodes: tools = "\n".join(shlex.join(c) for c in preflight_commands(node.device_family)) @@ -457,6 +581,10 @@ def _steps(plan: Stage1Plan, out: Expanded, trainers: Sequence[Node], samplers: f"set -eu\n{tools}\nip -4 -o addr show {node.interface or 'eth0'}\n" "df -h /data\nfree -g\n", "cards match the inventory; rendezvous address and /data present")) + access = out.store["head_access"] + if access["mode"] == "relay": + steps.append(Step("head-store-relay", None, f"set -eu\n{access['check']}\n", + f"the head reads the run store through {access['relay']} (ossutil over SSH)")) # The re-quantizer on trainer node 0 reads the MXFP4 export as its template. for key, nodes in (("bf16_source", trainers), ("sampler_initial", (*samplers, *trainers[:1]))): src = plan.section("model")[key] @@ -468,29 +596,27 @@ def _steps(plan: Stage1Plan, out: Expanded, trainers: Sequence[Node], samplers: f"{key} weights on the node's /data (revision {src['revision']})")) 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) + steps.append(Step(f"sampler-spec-{node.name}", node.name, + f"set -eu\nmkdir -p {shlex.quote(str(Path(run['spec_path']).parent))}\n" + f"cat > {shlex.quote(run['spec_path'])} <<'YETO_SAMPLER_SPEC'\n{spec_json}\n" + "YETO_SAMPLER_SPEC\n", + "the engine sampler island's spec file")) steps.append(Step(f"sampler-{node.name}", node.name, _docker(plan, node, "sampler", out.sampler["image"], run["argv"], run["env"]), - "SGLang TP8 engine on the round-0 MXFP4 weights (the yeto engine sampler " - "island drives it)")) - head = trainers[0] if trainers else None - for rank, node in enumerate(trainers): - ip = rendezvous.get(node.name) - ray = (["ray", "start", "--head", f"--port={RAY_PORT}", f"--node-ip-address={ip}", - f"--num-gpus={node.cards_per_node}", "--block"] if rank == 0 else - ["ray", "start", f"--address={rendezvous.get(head.name)}:{RAY_PORT}", - f"--node-ip-address={ip}", f"--num-gpus={node.cards_per_node}", "--block"]) - steps.append(Step(f"trainer-ray-{node.name}", node.name, - _docker(plan, node, "trainer-ray", out.trainer["image"], ray, out.trainer["env"]), - "Ray node of the Miles trainer island")) - if head is not None: - island = out.trainer["miles"]["island_argv"] - steps.append(Step("trainer-island", head.name, - _docker(plan, head, "trainer", out.trainer["image"], island, out.trainer["env"]), - "yeto Miles island entry, sample role trainer; reads samples from the store")) - steps.append(Step("requant-sidecar", head.name, - _docker(plan, head, "requant", out.trainer["image"], out.requant["loop_argv"], - out.trainer["env"]), - "re-quantize every bf16 release to MXFP4 RTN under the scheme's release dir")) + "engine sampler island: starts SGLang TP8 on the round-0 MXFP4 weights, runs the " + "post-load checks, writes groups to the store, follows the MXFP4 releases", + service=True)) + spec = out.trainer_spec + if spec is not None: + paths = remote_paths(plan.run_id, "trainer", 0) + for rank, node in enumerate(spec.nodes): + steps.append(Step(f"trainer-{node.name}", node.name, + start_script(spec, rank, {}, paths, + container=container_name(plan.run_id, "trainer", 0, rank)), + ("Ray head, then the yeto Miles island entry (sample role trainer) and the " + "MXFP4 re-quantizer side process" if rank == 0 else "Ray worker of the trainer island"), + service=True)) return steps @@ -519,6 +645,8 @@ def check(plan: Stage1Plan, out: Expanded) -> list[str]: ipaddress.ip_address(str(ip)) except ValueError: errors.append(f"{node.name}: rendezvous address {ip!r} is not an IP address") + if node.address is not None and node.address != str(ip): + errors.append(f"{node.name}: rendezvous address {ip} differs from the inventory address {node.address}") if len({(n.device_family, n.card, n.cards_per_node) for n in trainers}) > 1: errors.append("trainer nodes differ in device family, card or card count") @@ -654,36 +782,71 @@ def gaps(plan: Stage1Plan, out: Expanded) -> list[str]: if gap := _check_island_argv(out.trainer["miles"]["island_argv"]): found.append(gap) recipes = _miles_recipes() - if not any("glm5" in r for r in recipes): - found.append(f"Miles island: no glm5_next model recipe passes trainer.miles.megatron_args " + if glm5_next.MODEL_NAME not in recipes: + found.append(f"Miles island: no {glm5_next.MODEL_NAME} model recipe " f"(--rl-model-recipe choices: {recipes})") - found.append("head launcher: `yeto launch` does not start an SSH trainer island from the " - "inventory yet (N1.3); the trainer-island step runs the island entry directly") + if gap := _full_parameter_layout_gap(plan, out): + found.append(gap) else: if not out.trainer["verl"]["commit"]: found.append("trainer.verl.commit is unset (verl pin)") found.append("verl backend: no SSH launch path for the verl trainer island and no cut save path (N7.6 verl)") - store = out.store["run"] - if store.startswith("oss://"): + if out.store["run"].startswith("oss://") and out.store["head_access"]["mode"] == "direct": found.append("head on this server cannot reach the intranet-only OSS bucket (PublicEndpointForbidden): " - "needs a CNAME/public endpoint or a relay through ppu1 (runbook N3.5)") + "set store_access.head: relay (through ppu1) or give a public endpoint/CNAME (runbook N3.5)") return found +def _full_parameter_layout_gap(plan: Stage1Plan, out: Expanded) -> str | None: + """The Miles island's own full-parameter layout rule (``run_config``).""" + from yeto.rl.engine.run_config import full_parameter_layout_problem # noqa: PLC0415 + + argv = out.trainer["miles"]["island_argv"] + 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_per_engine=1) # the island's default; it gives no engine + if problem is None: + return None + return (f"Miles island: the full-parameter mode refuses this layout ({problem}; run_config Milestone 1): " + f"{len(plan.nodes_of('trainer'))} trainer nodes, DP {out.trainer['miles']['layout']['dp']}, " + "samples from the store instead of colocated engines") + + # --------------------------------------------------------------------- CLI def run_steps(out: Expanded, *, ssh: Sequence[str] = ("ssh", "-o", "BatchMode=yes"), - runner=subprocess.run) -> int: - """Run the steps in order; stop at the first failure.""" + runner=subprocess.run, starter=None) -> int: + """Run the steps in order; stop at the first failure. A service step is + started (``starter``, default ``subprocess.Popen``) and not waited for; after + the last step the run waits for the first trainer node's service, whose + exit is the run's exit (the other services are then left to the operator: + runbook S11 stops them).""" + starter = starter or (lambda argv, script: subprocess.Popen( # noqa: S603 + argv, stdin=subprocess.PIPE, text=True)) hosts = {n.name: n.host for n in out.plan.nodes} + services: list[tuple[Step, Any]] = [] for step in out.steps: print(f"[yeto-stage1] step {step.name} on {step.node or 'head'}: {step.purpose}", flush=True) argv = ["bash", "-c", step.script] if step.node is None else [*ssh, hosts[step.node], "bash", "-s"] + if step.service: + proc = starter(argv, step.script) + if step.node is not None and getattr(proc, "stdin", None) is not None: + proc.stdin.write(step.script) + proc.stdin.close() + services.append((step, proc)) + continue rc = runner(argv, input=None if step.node is None else step.script, text=True).returncode if rc != 0: print(f"[yeto-stage1] step {step.name} failed (exit {rc}); later steps not run", file=sys.stderr) return rc - return 0 + trainer = next((p for step, p in services if step.name.startswith("trainer-")), None) + if trainer is None: + return 0 + rc = trainer.wait() + print(f"[yeto-stage1] trainer island ended (exit {rc})", file=sys.stderr if rc else sys.stdout, flush=True) + return rc def main(argv: list[str] | None = None) -> int: From c013dcd6543b1e4eabacb97fb67eaeb162cd0afc Mon Sep 17 00:00:00 2001 From: michael <78113554+chowhsu@users.noreply.github.com> Date: Mon, 19 Oct 2026 13:26:24 +0000 Subject: [PATCH 30/34] Use one device-family constant and derive the stage-1 release directory from the scheme gpu_spec.DEVICE_FAMILY_ENV is the only name of YETO_DEVICE_FAMILY: drop the accel alias and import it directly in stage1 and the tests. Co-Authored-By: Claude Opus 5.5 --- tests/test_accel.py | 7 ++++--- tests/test_megatron_learner_device_family.py | 3 ++- yeto/accel.py | 4 +--- yeto/rl/stage1.py | 2 +- 4 files changed, 8 insertions(+), 8 deletions(-) diff --git a/tests/test_accel.py b/tests/test_accel.py index eca47d52e..037be1371 100644 --- a/tests/test_accel.py +++ b/tests/test_accel.py @@ -16,6 +16,7 @@ import torch from yeto import accel +from yeto.gpu_spec import DEVICE_FAMILY_ENV class FakeOutOfMemory(RuntimeError): @@ -319,15 +320,15 @@ def test_family_maps_to_device_type_and_backend(family, device_type, backend): def test_unknown_family_is_refused(monkeypatch): with pytest.raises(ValueError, match="unknown device family"): accel.family_dist_backend("tpu") - monkeypatch.setenv(accel.FAMILY_ENV, "tpu") + monkeypatch.setenv(DEVICE_FAMILY_ENV, "tpu") with pytest.raises(ValueError, match="unknown device family"): accel.device_family() def test_device_family_precedence(no_cuda, npu, monkeypatch): - monkeypatch.delenv(accel.FAMILY_ENV, raising=False) + monkeypatch.delenv(DEVICE_FAMILY_ENV, raising=False) assert accel.device_family() == "npu" # auto-detect - monkeypatch.setenv(accel.FAMILY_ENV, "ppu") + monkeypatch.setenv(DEVICE_FAMILY_ENV, "ppu") assert accel.device_family() == "ppu" # the launcher's inventory value wins assert accel.device_family("cuda") == "cuda" # an explicit flag wins over both diff --git a/tests/test_megatron_learner_device_family.py b/tests/test_megatron_learner_device_family.py index a4a7cc2e8..8846476cf 100644 --- a/tests/test_megatron_learner_device_family.py +++ b/tests/test_megatron_learner_device_family.py @@ -10,6 +10,7 @@ import torch.distributed as dist from yeto import accel +from yeto.gpu_spec import DEVICE_FAMILY_ENV from yeto.megatron import learner @@ -52,7 +53,7 @@ def test_init_distributed_follows_the_family(megatron_stub, monkeypatch, family, def test_family_comes_from_the_launcher_env(megatron_stub, monkeypatch): - monkeypatch.setenv(accel.FAMILY_ENV, "npu") + monkeypatch.setenv(DEVICE_FAMILY_ENV, "npu") monkeypatch.setattr(accel, "bind_family_device", lambda fam, rank: torch.device("cpu")) learner._init_distributed(_args(None)) assert megatron_stub[0] == ("pg", {"backend": "hccl"}) diff --git a/yeto/accel.py b/yeto/accel.py index 4750384ab..e7a504ac3 100644 --- a/yeto/accel.py +++ b/yeto/accel.py @@ -38,8 +38,6 @@ # torch_npu and HCCL. The collective backend follows from the device type. FAMILY_DEVICE_TYPES = {"cuda": "cuda", "ppu": "cuda", "npu": "npu", "cpu": "cpu"} -# Set by the island launcher from the node inventory; overrides auto-detection. -FAMILY_ENV = DEVICE_FAMILY_ENV def register_backends() -> None: @@ -108,7 +106,7 @@ def device_family(explicit: str | None = None) -> str: """The device family of this rank: ``explicit``, else ``$YETO_DEVICE_FAMILY``, else the detected torch device type (a PPU then reads as ``"cuda"``, which selects the same runtime).""" - family = explicit or os.environ.get(FAMILY_ENV) or available_type() + family = explicit or os.environ.get(DEVICE_FAMILY_ENV) # set by the island launcher or available_type() if family not in FAMILY_DEVICE_TYPES: raise ValueError(f"unknown device family {family!r} (known: {sorted(FAMILY_DEVICE_TYPES)})") return family diff --git a/yeto/rl/stage1.py b/yeto/rl/stage1.py index 1411e3164..67da5d6f7 100644 --- a/yeto/rl/stage1.py +++ b/yeto/rl/stage1.py @@ -57,10 +57,10 @@ from pathlib import Path from typing import Any +from yeto.gpu_spec import DEVICE_FAMILY_ENV from yeto.launch.node_inventory import InventoryError, Node, load_inventory from yeto.rl.adapters.miles.models import glm5_next from yeto.rl.adapters.sampler_engines.from_spec import FACTORY as ENGINE_FACTORY -from yeto.rl.engine.telemetry import DEVICE_FAMILY_ENV from yeto.rl.quant.schemes import release_dir_for PLAN_SCHEMA = "yeto.stage1/v1" From 2f318587f5530d239bb5d39c2a5219b3d493b237 Mon Sep 17 00:00:00 2001 From: michael <78113554+chowhsu@users.noreply.github.com> Date: Mon, 19 Oct 2026 13:30:28 +0000 Subject: [PATCH 31/34] Accept multi-node, DP>1 full-parameter Miles islands that read samples from the store A full-parameter island with --rl-sample-role trainer is trainer-only: a fixed partition with no rollout GPU, which Miles runs as --debug-train-only (no engine, no weight sync to engines). Any node count and DP are accepted. Islands with engines keep the Milestone-1 rule (one node, DP=1, dedicated engines); a full-parameter trainer never shares cards with engines. Co-Authored-By: Claude Opus 5.5 --- examples/stage1/ppu-inventory.yaml | 4 +++ tests/test_rl_hetero_sampler_engines.py | 2 +- tests/test_rl_miles_adapter_config.py | 12 ++++++++ tests/test_rl_miles_adapter_placement.py | 35 ++++++++++++++++++++++++ tests/test_stage1_gaps.py | 15 +++++++--- tests/test_stage1_plan.py | 3 +- yeto/rl/adapters/miles/config.py | 3 ++ yeto/rl/adapters/miles/placement.py | 19 ++++++++++++- yeto/rl/engine/run_config.py | 35 +++++++++++++++++++----- yeto/rl/stage1.py | 8 +++--- 10 files changed, 118 insertions(+), 18 deletions(-) diff --git a/examples/stage1/ppu-inventory.yaml b/examples/stage1/ppu-inventory.yaml index 4bde13cdb..f05a67757 100644 --- a/examples/stage1/ppu-inventory.yaml +++ b/examples/stage1/ppu-inventory.yaml @@ -8,6 +8,7 @@ nodes: host: ppu1 address: 10.0.0.74 # eth0; order assumed, runbook S1 checks it device_family: ppu + cpu_arch: x86_64 # assumed (PPU base image is x86_64); trainer preflight checks uname -m card: M890 cards_per_node: 8 hbm_gib: 144 @@ -22,6 +23,7 @@ nodes: host: ppu2 address: 10.0.0.75 # eth0; order assumed, runbook S1 checks it device_family: ppu + cpu_arch: x86_64 # assumed (PPU base image is x86_64); trainer preflight checks uname -m card: M890 cards_per_node: 8 hbm_gib: 144 @@ -36,6 +38,7 @@ nodes: host: ppu3 address: 10.0.0.76 # eth0; order assumed, runbook S1 checks it device_family: ppu + cpu_arch: x86_64 # assumed (PPU base image is x86_64); trainer preflight checks uname -m card: M890 cards_per_node: 8 hbm_gib: 144 @@ -50,6 +53,7 @@ nodes: host: ppu4 address: 10.0.0.77 # eth0; order assumed, runbook S1 checks it device_family: ppu + cpu_arch: x86_64 # assumed (PPU base image is x86_64); trainer preflight checks uname -m card: M890 cards_per_node: 8 hbm_gib: 144 diff --git a/tests/test_rl_hetero_sampler_engines.py b/tests/test_rl_hetero_sampler_engines.py index e727958c0..7f2a5a442 100644 --- a/tests/test_rl_hetero_sampler_engines.py +++ b/tests/test_rl_hetero_sampler_engines.py @@ -73,7 +73,7 @@ def test_compat_group_carries_family_and_hbm(): def _inventory(**node): - base = {"name": "npu1", "host": "u@h", "device_family": "npu", "card": "910B4", + base = {"name": "npu1", "host": "u@h", "device_family": "npu", "cpu_arch": "aarch64", "card": "910B4", "cards_per_node": 8, "hbm_gib": 64, "host_memory_gib": 1024, "image": "img"} return parse_inventory({"nodes": [{**base, **node}]}) diff --git a/tests/test_rl_miles_adapter_config.py b/tests/test_rl_miles_adapter_config.py index fdfa4a0e3..9e4d92a96 100644 --- a/tests/test_rl_miles_adapter_config.py +++ b/tests/test_rl_miles_adapter_config.py @@ -237,6 +237,18 @@ def test_full_parameter_mode_emits_no_lora_flags(): +def test_trainer_only_full_parameter_island_starts_no_engine(): + # S25 T-INTEG2: a full-parameter trainer that reads the store (DP>1, no engine). + cfg = sub(make_config(), "trainable", parameter_mode="full", target_modules=()) + cfg = sub(cfg, "parallel", dedicated_rollout_gpus=0, visible_gpus_per_node=2) + assert not cfg.parallel.colocated + request = mc.placement_request(cfg) + assert request.trainer_only + argv = mc.translate_run_config(cfg, AlgorithmSpec()).argv + assert "--debug-train-only" in argv + assert "--colocate" not in argv and "--rollout-num-gpus" not in argv + + def test_full_parameter_mode_turns_the_distributed_optimizer_off(): # Miles' set_default_megatron_args forces the distributed optimizer on for # Adam, which leaves each FP32 master a flat slice; the full-parameter diff --git a/tests/test_rl_miles_adapter_placement.py b/tests/test_rl_miles_adapter_placement.py index 9d157d43a..db5f5f7e6 100644 --- a/tests/test_rl_miles_adapter_placement.py +++ b/tests/test_rl_miles_adapter_placement.py @@ -157,3 +157,38 @@ def test_standby_placement_describes_m1_roles_and_detects_rewrites(): groups["standby"] = (None, [0, 3], [0, 3]) with pytest.raises(PlacementRewriteError, match="overlapping"): MilesPlacement(req, groups) + + +# -- trainer-only fixed partition (S25 T-INTEG2: full-parameter trainer reading the store) -- +TRAINER_ONLY = PlacementRequest("fixed-partition", trainer_gpus=24, rollout_gpus=0, gpus_per_engine=1) + + +def test_trainer_only_request_has_no_rollout_bundles(): + assert TRAINER_ONLY.trainer_only and not PARTITION.trainer_only + assert TRAINER_ONLY.expected_layout() == (24, 24) + assert TRAINER_ONLY.role_bundles()["rollout"] == () + with pytest.raises(ValueError, match="positive int"): + PlacementRequest("colocated", trainer_gpus=4, rollout_gpus=0, gpus_per_engine=1) + with pytest.raises(ValueError, match="trainer-only"): + PlacementRequest("fixed-partition", trainer_gpus=4, rollout_gpus=0, gpus_per_engine=1, + standby_gpus=4) + + +def test_trainer_only_needs_debug_train_only(): + ok = args(colocate=False, actor_num_nodes=3, actor_num_gpus_per_node=8, rollout_num_gpus=0, + rollout_num_gpus_per_engine=1, debug_train_only=True) + check_placement_not_rewritten(TRAINER_ONLY, ok) + with pytest.raises(PlacementRewriteError, match="debug_train_only=True"): + check_placement_not_rewritten(TRAINER_ONLY, args(**{**vars(ok), "debug_train_only": False})) + # A partition with engines still refuses debug_train_only. + with pytest.raises(PlacementRewriteError, match="debug_train_only=True"): + check_placement_not_rewritten(PARTITION, args(colocate=False, actor_num_gpus_per_node=2, + rollout_num_gpus=2, rollout_num_gpus_per_engine=1, + debug_train_only=True)) + + +def test_trainer_only_description(): + placement = MilesPlacement(TRAINER_ONLY, {"actor": pg(range(24), range(24)), "rollout": pg([], [])}, + logical=True) + d = placement.describe() + assert (d.kind, len(d.trainer_gpus), d.rollout_gpus) == ("fixed-partition", 24, ()) diff --git a/tests/test_stage1_gaps.py b/tests/test_stage1_gaps.py index 3969df259..9ff34c178 100644 --- a/tests/test_stage1_gaps.py +++ b/tests/test_stage1_gaps.py @@ -138,22 +138,29 @@ def test_island_parser_takes_the_recipe_and_transformer_impl(): def test_full_parameter_layout_rule_is_shared(): - assert full_parameter_layout_problem(data_parallel=2, actor_num_nodes=3, rollout_gpus=0, - rollout_gpus_per_engine=1) == "dense full-parameter GRPO requires DP=1" + # Engines on the island (Milestone 1): one node, DP 1, dedicated engines. + assert "requires DP=1" in full_parameter_layout_problem( + data_parallel=2, actor_num_nodes=3, rollout_gpus=0, rollout_gpus_per_engine=1) assert full_parameter_layout_problem(data_parallel=1, actor_num_nodes=1, rollout_gpus=8, rollout_gpus_per_engine=8) is None + # Samples from the store: any node count and DP, but no engine on the island. + assert full_parameter_layout_problem(data_parallel=2, actor_num_nodes=3, rollout_gpus=0, + rollout_gpus_per_engine=1, sample_role="trainer") is None + assert "must be 0" in full_parameter_layout_problem( + data_parallel=2, actor_num_nodes=3, rollout_gpus=8, rollout_gpus_per_engine=8, + sample_role="trainer") # ------------------------------------------------------------ gap 5: SSH trainer island def _node(name="ppu1", **kw) -> Node: - base = dict(host=name, device_family="ppu", card="M890", cards_per_node=8, hbm_gib=144.0, + base = dict(host=name, device_family="ppu", cpu_arch="x86_64", card="M890", cards_per_node=8, hbm_gib=144.0, host_memory_gib=2266.0, fabric="roce", interface="eth0", runtime="docker", image="img:1", mounts={"/data": "/data"}, docker_args=("--device", "/dev/x")) return Node(name=name, **(base | kw)) def test_inventory_reads_address_and_docker_args(): - raw = {"name": "n", "host": "ppu1", "device_family": "ppu", "card": "M890", "cards_per_node": 8, + raw = {"name": "n", "host": "ppu1", "device_family": "ppu", "cpu_arch": "x86_64", "card": "M890", "cards_per_node": 8, "hbm_gib": 144, "host_memory_gib": 2266, "image": "i", "address": "10.0.0.74", "docker_args": ["--device", "/dev/a"]} node = parse_inventory({"nodes": [raw]})[0] diff --git a/tests/test_stage1_plan.py b/tests/test_stage1_plan.py index 1cb225da6..5eb478cc2 100644 --- a/tests/test_stage1_plan.py +++ b/tests/test_stage1_plan.py @@ -101,7 +101,8 @@ def test_steps_order_and_staging(shipped): def test_gaps_block_run_until_wired(shipped, capsys): assert any("N0.2" in g for g in shipped.gaps) - assert any("DP=1" in g for g in shipped.gaps) # Miles full-parameter rule (T-GAPS finding) + assert not any("full-parameter" in g for g in shipped.gaps) # trainer-only island accepted (T-INTEG2) + assert len(shipped.gaps) == 2 # only N0.2 (image) and N0.2b (docker device flags) need the machines assert stage1.main(["run", str(PLAN), "--confirm-machines-free"]) == 2 assert "run refused" in capsys.readouterr().err diff --git a/yeto/rl/adapters/miles/config.py b/yeto/rl/adapters/miles/config.py index 2202d1bc2..eb414e458 100644 --- a/yeto/rl/adapters/miles/config.py +++ b/yeto/rl/adapters/miles/config.py @@ -812,6 +812,9 @@ def translate_run_config( if request.kind == "colocated": placement_values = ["--colocate"] + elif request.trainer_only: + # Samples come from the store: Miles starts no engine and syncs no weights. + placement_values = ["--debug-train-only"] else: placement_values = ["--rollout-num-gpus", str(request.rollout_gpus)] if trainable.parameter_mode == "lora": diff --git a/yeto/rl/adapters/miles/placement.py b/yeto/rl/adapters/miles/placement.py index dcd398fa9..6e15fb6f2 100644 --- a/yeto/rl/adapters/miles/placement.py +++ b/yeto/rl/adapters/miles/placement.py @@ -2,7 +2,9 @@ R0 supports two placements fixed at startup: ``colocated`` (trainer and rollout share the same GPUs) and ``fixed-partition`` (disjoint trainer and -rollout GPU sets). Upstream Miles normalizes its arguments (``--colocate`` +rollout GPU sets). A fixed partition with no rollout GPU is *trainer-only* +(:attr:`PlacementRequest.trainer_only`): the island reads its samples from the +store and starts no engine (Miles ``--debug-train-only``). Upstream Miles normalizes its arguments (``--colocate`` forces ``rollout_num_gpus``; ``--debug-rollout-only`` turns colocate off; external engines drop rollout GPUs; ...). The spec forbids running with a silently rewritten placement, so :func:`check_placement_not_rewritten` @@ -61,7 +63,11 @@ def __post_init__(self) -> None: for name in ("trainer_gpus", "rollout_gpus", "gpus_per_engine"): value = getattr(self, name) if isinstance(value, bool) or not isinstance(value, int) or value < 1: + if name == "rollout_gpus" and value == 0 and self.kind == "fixed-partition": + continue # trainer-only raise ValueError(f"placement {name} must be a positive int") + if self.trainer_only and (self.standby_gpus or self.rollout_cell_names or self.bundle_map is not None): + raise ValueError("a trainer-only placement has no standby GPUs, rollout cells or bundle map") if self.kind == "colocated" and self.rollout_gpus != self.trainer_gpus: raise ValueError("colocated placement shares one GPU set") if self.rollout_gpus % self.gpus_per_engine: @@ -93,6 +99,11 @@ def __post_init__(self) -> None: if self.gpus_per_node is not None: self._check_nodes() + @property + def trainer_only(self) -> bool: + """A fixed partition without rollout GPUs: the island starts no engine.""" + return self.kind == "fixed-partition" and self.rollout_gpus == 0 + @property def in_node_parallel(self) -> int: """Trainer group that must stay inside one node: ``node_parallel`` (tp*cp) @@ -264,6 +275,8 @@ def _parsed_layout(args: Any) -> tuple[int, int] | None: trainer = int(args.actor_num_nodes) * int(args.actor_num_gpus_per_node) rollout = int(getattr(args, "rollout_num_gpus", 0) or 0) + if getattr(args, "debug_train_only", False): + return trainer, trainer # upstream: no rollout bundles (eval_num_gpus is refused) if getattr(args, "colocate", False): return max(trainer, rollout), 0 return trainer + rollout, trainer @@ -276,8 +289,12 @@ def check_placement_not_rewritten(request: PlacementRequest, args: Any) -> None: return # test-only rollout-dump replay (config._apply_test_debug_rollout): no engines by design conflicts: list[str] = [] for name in ("debug_rollout_only", "debug_train_only", "rollout_external"): + if name == "debug_train_only" and request.trainer_only: + continue if getattr(args, name, False): conflicts.append(f"{name}=True") + if request.trainer_only and not getattr(args, "debug_train_only", False): + conflicts.append("a trainer-only placement needs debug_train_only=True (no engine)") if int(getattr(args, "eval_num_gpus", 0) or 0): conflicts.append(f"eval_num_gpus={args.eval_num_gpus} adds GPUs to the rollout partition") colocate = bool(getattr(args, "colocate", False)) diff --git a/yeto/rl/engine/run_config.py b/yeto/rl/engine/run_config.py index 7b596bbda..548a902da 100644 --- a/yeto/rl/engine/run_config.py +++ b/yeto/rl/engine/run_config.py @@ -541,11 +541,27 @@ def _rules(): def full_parameter_layout_problem(*, data_parallel: int, actor_num_nodes: int, rollout_gpus: Any, - rollout_gpus_per_engine: Any) -> str | None: - """Why the Miles full-parameter mode refuses this layout (Milestone 1: one - node, DP=1, dedicated evenly partitioned engines); None when it accepts it.""" + rollout_gpus_per_engine: Any, sample_role: str = sample_roles.BOTH + ) -> str | None: + """Why the Miles full-parameter mode refuses this layout; None when it accepts it. + + Two layouts are accepted: + + * samples from the store (``--rl-sample-role trainer``): a trainer-only + island. Any node count and any DP; the island starts no inference engine + (``--rollout-num-gpus`` absent or 0), so no card is shared with one; + * engines on the island (Milestone 1): one node, DP=1 and dedicated, evenly + partitioned engines. A full-parameter trainer never shares its cards with + engines (no colocated placement). + """ + if sample_role == sample_roles.TRAINER: + if rollout_gpus not in (None, 0): + return ("a full-parameter trainer that reads samples from the store starts no " + "inference engine: --rollout-num-gpus must be 0 or absent") + return None if data_parallel != 1: - return "dense full-parameter GRPO requires DP=1" + return ("dense full-parameter GRPO with engines on the island requires DP=1 " + "(read samples from the store, --rl-sample-role trainer, for DP>1)") if ( actor_num_nodes != 1 or type(rollout_gpus) is not int @@ -554,8 +570,9 @@ def full_parameter_layout_problem(*, data_parallel: int, actor_num_nodes: int, r or rollout_gpus_per_engine < 1 or rollout_gpus % rollout_gpus_per_engine ): - return ("Milestone-1 dense full-parameter mode requires one node and " - "dedicated, evenly partitioned inference engines") + return ("full-parameter mode with engines on the island requires one node and " + "dedicated, evenly partitioned inference engines (never colocated); " + "read samples from the store (--rl-sample-role trainer) for more nodes") return None @@ -679,10 +696,14 @@ def resolve_rl_run_config( if parameter_mode == "full": rollout_gpus = getattr(args, "rollout_num_gpus", None) rollout_gpus_per_engine = getattr(args, "rollout_num_gpus_per_engine", None) + role = sample_roles.sample_role(args) if problem := full_parameter_layout_problem( data_parallel=data_parallel, actor_num_nodes=args.actor_num_nodes, - rollout_gpus=rollout_gpus, rollout_gpus_per_engine=rollout_gpus_per_engine): + rollout_gpus=rollout_gpus, rollout_gpus_per_engine=rollout_gpus_per_engine, + sample_role=role): raise ValueError(problem) + if role == sample_roles.TRAINER: + rollout_gpus = 0 # trainer-only island (Miles --debug-train-only): no engine if qwen35_recipe and ( rollout_gpus_per_engine != 1 or getattr(args, "sglang_tp_size", None) not in {None, 1} diff --git a/yeto/rl/stage1.py b/yeto/rl/stage1.py index 67da5d6f7..7a4c8f09e 100644 --- a/yeto/rl/stage1.py +++ b/yeto/rl/stage1.py @@ -806,12 +806,12 @@ def _full_parameter_layout_gap(plan: Stage1Plan, out: Expanded) -> str | None: 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_per_engine=1) # the island's default; it gives no engine + 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: return None - return (f"Miles island: the full-parameter mode refuses this layout ({problem}; run_config Milestone 1): " - f"{len(plan.nodes_of('trainer'))} trainer nodes, DP {out.trainer['miles']['layout']['dp']}, " - "samples from the store instead of colocated engines") + return (f"Miles island: the full-parameter mode refuses this layout ({problem}): " + f"{len(plan.nodes_of('trainer'))} trainer nodes, DP {out.trainer['miles']['layout']['dp']}") # --------------------------------------------------------------------- CLI From 47ef79930c110ca06a30d609fec02aefad5f1005 Mon Sep 17 00:00:00 2001 From: michael <78113554+chowhsu@users.noreply.github.com> Date: Mon, 19 Oct 2026 13:34:42 +0000 Subject: [PATCH 32/34] Model Megatron optimizer host offload and check the stage-1 memory figures with it (N7.1c) The Megatron estimate takes optimizer_offload_fraction (Adam state, and the FP32 master unless offload_master is off) and reports the offloaded bytes as host memory per node. yeto-stage1 plan now uses the estimator instead of its own static formula, checks the plan's 78.7 GiB/card and 1196 GiB/node figures (both derived) against it, and reports a gap: the Miles island runs full-parameter mode without the distributed optimizer and without optimizer offload, which the estimator puts at 243 GiB/card. Co-Authored-By: Claude Opus 5.5 --- examples/stage1/glm53-flash-ppu.yaml | 5 +- tests/test_rl_fullparam_memory.py | 20 ++++- tests/test_stage1_plan.py | 20 ++++- yeto/rl/adapters/miles/config.py | 6 +- yeto/rl/adapters/miles/models/glm5_next.py | 9 +- yeto/rl/fullparam/memory.py | 54 +++++++++--- yeto/rl/stage1.py | 99 +++++++++++++++------- 7 files changed, 165 insertions(+), 48 deletions(-) diff --git a/examples/stage1/glm53-flash-ppu.yaml b/examples/stage1/glm53-flash-ppu.yaml index 74b3e14cd..5f23b3a1d 100644 --- a/examples/stage1/glm53-flash-ppu.yaml +++ b/examples/stage1/glm53-flash-ppu.yaml @@ -76,7 +76,10 @@ trainer: cp: 1 ep: 8 # inside one node etp: 1 - optimizer_offload_fraction: 1.0 # Adam state on the host (estimate: 1.2 TiB/node of 2266 GiB) + optimizer_offload_fraction: 1.0 # Adam state + FP32 master on the host + memory_estimate: # 推算 (T-STAGE1); `plan` checks them against the estimator (N7.1c) + weights_grads_gib_per_card: 78.7 # bf16 weights + fp32 main grads, worst stage + host_optimizer_gib_per_node: 1196 # 12 B/param x owned share x 8 cards transformer_impl: transformer_engine # main-agent ruling 10-19 (3): TE stays default; `local` is the N7.5 fallback recompute: full max_tokens_per_gpu: 8192 diff --git a/tests/test_rl_fullparam_memory.py b/tests/test_rl_fullparam_memory.py index ff8710060..ef664cf52 100644 --- a/tests/test_rl_fullparam_memory.py +++ b/tests/test_rl_fullparam_memory.py @@ -161,7 +161,8 @@ def test_megatron_output_is_labelled_derived(): @pytest.mark.parametrize("change, match", [ ({"ep": 3}, "ep x pp = 6 does not divide cards 16"), ({"tp": 3}, "tp x pp x cp = 6 does not divide"), - ({"offload": "full"}, "only offload none"), + ({"offload": "full"}, "offload is an FSDP2 mode"), + ({"optimizer_offload_fraction": 1.5}, r"must be in \[0, 1\]"), ({"ulysses": 2}, "use cp instead of ulysses"), ({"pp": 8, "ep": 2}, "pp 8 exceeds the 4 layers"), ({"framework": "deepspeed"}, "framework must be one of"), @@ -171,11 +172,28 @@ def test_malformed_megatron_plans_are_refused(change, match): estimate(MOE, replace(MEG, **change)) +def test_megatron_optimizer_offload_moves_state_and_master_to_the_host(): + # N7.1c: fraction F of (8 B Adam + 4 B master) leaves the card; host = per card x cards_per_node. + dense, experts = 400 / 2 / 2, 600 / 4 / 2 + owned = dense / 4 + experts / 2 + full = estimate(MOE, replace(MEG, optimizer_offload_fraction=1.0)) + assert full.items["optimizer_shard"].bytes == pytest.approx(0) + assert full.items["host_optimizer_offload"].bytes == pytest.approx(12 * owned * 8) + assert full.items["host_optimizer_offload"].where == "host" + half = estimate(MOE, replace(MEG, optimizer_offload_fraction=0.5, offload_master=False)) + assert half.items["optimizer_shard"].bytes == pytest.approx((12 - 4) * owned) + assert half.total("host") == pytest.approx(4 * owned * 8) + assert "optimizer offload 0.5" in half.summary()["plan_status"] + assert "incl. FP32 master" in full.summary()["plan_status"] + + def test_megatron_refuses_cp_with_dsa_and_fsdp2_refuses_megatron_knobs(): with pytest.raises(MemoryPlanError, match="DSA layer does not support cp"): estimate(replace(MOE, dsa_layers=1, index_heads=4), replace(MEG, cp=2, pp=1, ep=8)) with pytest.raises(MemoryPlanError, match="megatron only"): estimate(GLM, replace(PPU, tp=2)) + with pytest.raises(MemoryPlanError, match="optimizer_offload_fraction applies to framework megatron"): + estimate(GLM, replace(PPU, optimizer_offload_fraction=0.5)) with pytest.raises(MemoryPlanError, match="go together"): replace(MOE, num_experts=0) diff --git a/tests/test_stage1_plan.py b/tests/test_stage1_plan.py index 5eb478cc2..9965e963f 100644 --- a/tests/test_stage1_plan.py +++ b/tests/test_stage1_plan.py @@ -102,7 +102,9 @@ def test_steps_order_and_staging(shipped): def test_gaps_block_run_until_wired(shipped, capsys): assert any("N0.2" in g for g in shipped.gaps) assert not any("full-parameter" in g for g in shipped.gaps) # trainer-only island accepted (T-INTEG2) - assert len(shipped.gaps) == 2 # only N0.2 (image) and N0.2b (docker device flags) need the machines + # Left: N0.2 (image), N0.2b (docker device flags), and N7.1d: the island runs + # full-parameter Megatron without the distributed optimizer or optimizer offload. + assert len(shipped.gaps) == 3 and "N7.1d" in shipped.gaps[2] assert stage1.main(["run", str(PLAN), "--confirm-machines-free"]) == 2 assert "run refused" in capsys.readouterr().err @@ -211,3 +213,19 @@ def test_verl_backend_expands(tmp_path): assert out.errors == [] assert "trainer.nnodes=3" in out.trainer["verl"]["overrides"] assert any("verl" in g for g in out.gaps) + + +def test_memory_figures_of_the_plan_match_the_estimator(shipped): + # N7.1c: the plan's 78.7 GiB/card and 1196 GiB/node (推算) agree with the estimator. + mem = shipped.trainer["miles"]["memory_estimate"] + assert mem["plan"]["basis"] == "推算" and mem["plan"]["fits"] + assert mem["plan"]["weights_grads_gib_per_card"] == pytest.approx(78.7, abs=0.1) + assert mem["plan"]["host_optimizer_gib_per_node"] == pytest.approx(1196, rel=0.02) + assert not mem["island"]["fits"] # no distributed optimizer, no offload (gap N7.1d) + + +def test_memory_figure_that_disagrees_is_an_error(tmp_path): + def edit(doc): + doc["trainer"]["miles"]["memory_estimate"]["host_optimizer_gib_per_node"] = 600 + out = _variant(tmp_path, edit) + assert any("host_optimizer_gib_per_node = 600" in e for e in out.errors) diff --git a/yeto/rl/adapters/miles/config.py b/yeto/rl/adapters/miles/config.py index eb414e458..1562325fd 100644 --- a/yeto/rl/adapters/miles/config.py +++ b/yeto/rl/adapters/miles/config.py @@ -1217,8 +1217,10 @@ def _full_parameter_optimizer_attrs(trainable: Any) -> dict[str, Any]: flat slice of the owned range even at DP=1, and the full-parameter shard manifest (Miles ``full_parameter_state._master_parameter``) refuses it as an incomplete master. The attribute is applied in ``preprocess_args``, - after Miles' defaults and before Megatron's ``validate_args``. Full mode is - DP=1 only (``resolve_rl_run_config``), so no master is DP-sharded either way. + after Miles' defaults and before Megatron's ``validate_args``. With engines + on the island full mode is DP=1; a trainer-only island (samples from the + store) may run DP>1, and then every DP replica keeps the whole optimizer + state of its local parameters (more device memory; stage-1 gap N7.1d). """ if trainable.parameter_mode != "full": diff --git a/yeto/rl/adapters/miles/models/glm5_next.py b/yeto/rl/adapters/miles/models/glm5_next.py index cade9faac..5ab9b8c5e 100644 --- a/yeto/rl/adapters/miles/models/glm5_next.py +++ b/yeto/rl/adapters/miles/models/glm5_next.py @@ -30,13 +30,16 @@ SPEC = ("miles_plugins.models.glm5_next.glm5_next", "get_glm5_next_spec") NUM_LAYERS = 45 FIRST_K_DENSE_REPLACE = 3 +HIDDEN_SIZE = 4096 +VOCAB_SIZE = 154880 +NUM_EXPERTS = 288 # N7.5 E6/E7: flags that exist only with TransformerEngine. TE_ONLY_ARGV: tuple[str, ...] = ("--moe-grouped-gemm", "--moe-permute-fusion") LOCAL_ARGV: tuple[str, ...] = ("--transformer-impl", LOCAL, "--no-gradient-accumulation-fusion") _MODEL_ARGV: tuple[str, ...] = ( - "--hidden-size", "4096", + "--hidden-size", str(HIDDEN_SIZE), "--ffn-hidden-size", "12288", "--num-attention-heads", "64", "--multi-latent-attention", @@ -47,7 +50,7 @@ "--v-head-dim", "256", "--kv-channels", "256", "--qk-layernorm", - "--num-experts", "288", + "--num-experts", str(NUM_EXPERTS), "--moe-router-topk", "8", "--moe-router-score-function", "sigmoid", "--moe-router-pre-softmax", @@ -57,7 +60,7 @@ "--moe-ffn-hidden-size", "2048", "--moe-shared-expert-intermediate-size", "2048", "--moe-router-dtype", "fp32", - "--vocab-size", "154880", + "--vocab-size", str(VOCAB_SIZE), "--rotary-base", "10000", "--rope-type", "rope", "--no-rope-fusion", diff --git a/yeto/rl/fullparam/memory.py b/yeto/rl/fullparam/memory.py index c1a65778f..120ddfb8a 100644 --- a/yeto/rl/fullparam/memory.py +++ b/yeto/rl/fullparam/memory.py @@ -49,6 +49,13 @@ ``state_bytes`` Adam state; with the distributed optimizer this is divided by the data-parallel size of its group (dense ``W / (TP x PP)``, CP included; experts ``W / (EP x PP)``), without it it is per local parameter; +* optimizer host offload (N7.1c; Megatron ``--optimizer-cpu-offload + --optimizer-offload-fraction F``): the fraction ``F`` of the Adam state + moves to the host, and with ``offload_master`` also that fraction of the + FP32 master. Host memory per node is the offloaded bytes of one card times + ``cards_per_node`` (every card of a node offloads the same share). Not + modelled: FP32 gradient copies or pinned staging buffers that the CPU + optimizer may keep on the host (Megatron source not checked for this); * activations with full recomputation: stage 0 of 1F1B holds ``PP`` micro batches, so ``ceil(layers / PP) x PP x mbs x T / (TP x CP) x streams x hidden x 2 B``; mHC keeps ``hc_mult`` (4) residual streams, so this term is 4 times @@ -138,7 +145,8 @@ def __post_init__(self) -> None: class TrainPlan: """One trainer island and how it trains. ``fsdp_size`` 0 means all cards in one shard group; a smaller divisor replicates groups (HSDP). ``tp`` / - ``pp`` / ``ep`` / ``cp`` / ``distributed_optimizer`` apply to + ``pp`` / ``ep`` / ``cp`` / ``distributed_optimizer`` / + ``optimizer_offload_fraction`` / ``offload_master`` apply to ``framework="megatron"`` only.""" cards: int @@ -167,6 +175,8 @@ class TrainPlan: ep: int = 1 cp: int = 1 distributed_optimizer: bool = True + optimizer_offload_fraction: float = 0.0 # share of the optimizer on the host (Megatron) + offload_master: bool = True # the offloaded share includes the FP32 master copy @property def shard_group(self) -> int: @@ -178,8 +188,11 @@ def _check_megatron_form(plan: TrainPlan, shape: ModelShape) -> None: if int(getattr(plan, name)) <= 0: raise MemoryPlanError(f"{name} must be positive, got {getattr(plan, name)}") if plan.offload != "none": - raise MemoryPlanError("megatron: only offload none is modelled (Megatron optimizer CPU " - "offload is not)") + raise MemoryPlanError("megatron: offload is an FSDP2 mode; use optimizer_offload_fraction " + "(Megatron --optimizer-offload-fraction)") + if not 0.0 <= plan.optimizer_offload_fraction <= 1.0: + raise MemoryPlanError(f"megatron: optimizer_offload_fraction must be in [0, 1], " + f"got {plan.optimizer_offload_fraction}") if plan.ulysses != 1 or plan.fsdp_size: raise MemoryPlanError("megatron: use cp instead of ulysses, and no fsdp_size") W = plan.cards @@ -212,6 +225,9 @@ def check_plan_form(plan: TrainPlan, shape: ModelShape) -> None: _check_megatron_form(plan, shape) elif (plan.tp, plan.pp, plan.ep, plan.cp) != (1, 1, 1, 1): raise MemoryPlanError("tp / pp / ep / cp apply to framework megatron only") + if plan.framework != "megatron" and plan.optimizer_offload_fraction: + raise MemoryPlanError("optimizer_offload_fraction applies to framework megatron only " + "(FSDP2: offload)") if plan.offload not in OFFLOAD_MODES: raise MemoryPlanError(f"offload must be one of {OFFLOAD_MODES}, got {plan.offload!r}") if plan.dsa_path not in DSA_PATHS: @@ -285,7 +301,7 @@ def summary(self) -> dict[str, Any]: "basis": BASIS, "framework": self.plan.framework, "plan_status": (PLAN_STATUS[self.plan.offload] if self.plan.framework == "fsdp2" - else "Megatron (no offload)"), + else _megatron_status(self.plan)), "device_gib_per_card": round(self.total("device") / GIB, 1), "device_usable_gib": round(self.device_limit / GIB, 1), "host_gib_per_node": round(self.total("host") / GIB, 1), @@ -354,8 +370,8 @@ def _estimate_megatron(shape: ModelShape, plan: TrainPlan) -> Estimate: """Megatron variant of :func:`estimate` (module doc, N7.1b).""" est = Estimate(shape, plan) - def add(name: str, value: float, formula: str) -> None: - est.items[name] = Item(float(value), formula, "device") + def add(name: str, value: float, formula: str, where: str = "device") -> None: + est.items[name] = Item(float(value), formula, where) W, tp, pp, ep, cp = plan.cards, plan.tp, plan.pp, plan.ep, plan.cp stage_layers = -(-shape.num_layers // pp) @@ -367,14 +383,23 @@ def add(name: str, value: float, formula: str) -> None: add("params_local_bf16", 2 * local, f"2 B x (dense/TP({tp}) + experts/EP({ep})), {where}") gb = plan.grad_bytes add("main_grads_local", gb * local, f"{gb} B x local params") - per_param = (4 if plan.param_bytes == 4 else 0) + plan.state_bytes if plan.distributed_optimizer: dp_dense, dp_expert = W // (tp * pp), W // (ep * pp) - add("optimizer_shard", per_param * (dense / dp_dense + experts / dp_expert), - f"{per_param} B x (dense local / DP({dp_dense}) + expert local / expert DP({dp_expert})), " - f"distributed optimizer") + owned = dense / dp_dense + experts / dp_expert + owned_text = (f"(dense local / DP({dp_dense}) + expert local / expert DP({dp_expert})), " + "distributed optimizer") else: - add("optimizer_local", per_param * local, f"{per_param} B x local params (no distributed optimizer)") + owned, owned_text = local, "local params (no distributed optimizer)" + master = 4 if plan.param_bytes == 4 else 0 + F = plan.optimizer_offload_fraction + host_per_param = F * (plan.state_bytes + (master if plan.offload_master else 0)) + card_per_param = plan.state_bytes + master - host_per_param + name = "optimizer_shard" if plan.distributed_optimizer else "optimizer_local" + add(name, card_per_param * owned, f"{card_per_param:g} B x {owned_text}, on the card") + if F: + what = "Adam state + FP32 master" if plan.offload_master and master else "Adam state" + add("host_optimizer_offload", host_per_param * owned * plan.cards_per_node, + f"{F:g} x {what} ({host_per_param:g} B) x {owned_text} x cards_per_node", "host") tokens = plan.seq_len * plan.micro_batch_seqs / (tp * cp) add("layer_inputs", stage_layers * pp * tokens * shape.residual_streams * shape.hidden_size * 2, f"stage layers({stage_layers}) x in-flight micro batches PP({pp}) x T*mbs/(TP*CP)({tokens:.0f}) " @@ -394,6 +419,13 @@ def add(name: str, value: float, formula: str) -> None: return est +def _megatron_status(plan: TrainPlan) -> str: + if not plan.optimizer_offload_fraction: + return "Megatron (no optimizer offload)" + master = " incl. FP32 master" if plan.offload_master and plan.param_bytes == 4 else "" + return f"Megatron (optimizer offload {plan.optimizer_offload_fraction:g}{master} on the host)" + + def check(shape: ModelShape, plan: TrainPlan) -> Estimate: """:func:`estimate`, raising :class:`MemoryPlanError` when the plan does not fit.""" est = estimate(shape, plan) diff --git a/yeto/rl/stage1.py b/yeto/rl/stage1.py index 7a4c8f09e..192e0e7f7 100644 --- a/yeto/rl/stage1.py +++ b/yeto/rl/stage1.py @@ -67,7 +67,7 @@ BACKENDS = ("miles", "verl") NEUTRAL_PENALTIES = {"repetition_penalty": 1.0, "frequency_penalty": 0.0, "presence_penalty": 0.0} RAY_PORT = 6379 -HOST_USABLE_FRACTION = 0.85 # as yeto.rl.fullparam.memory.TrainPlan.host_usable_fraction +MEMORY_TOLERANCE = 0.02 # plan memory figures vs the estimator (both 推算) HEAD_ACCESS = ("relay", "direct") # how the head reaches the store (store_access.head) RAY_JOIN_TIMEOUT_S = 600 @@ -200,19 +200,49 @@ def megatron_layout(miles: Mapping[str, Any], trainer_cards: int, cards_per_node "expert_dp": dp * tp * cp // (ep * etp)} -def megatron_static_gib(model: Mapping[str, Any], layout: Mapping[str, int]) -> dict[str, float]: - """Per-card static memory of the layout, estimate (A-TRN §2.1 method): - bf16 weights + fp32 main grads = 6 B/param on the card; Adam + fp32 master - = 12 B/param, sharded over the data-parallel group, on the host when - offloaded. Activations are not included.""" - experts = float(model["expert_params_b"]) * 1e9 - dense = float(model["params_b"]) * 1e9 - experts - per_card = experts / (layout["pp"] * layout["ep"] * layout["etp"]) + dense / (layout["pp"] * layout["tp"]) - optimizer_card = (experts / (layout["pp"] * layout["ep"] * layout["etp"] * layout["expert_dp"]) - + dense / (layout["pp"] * layout["tp"] * layout["dp"])) * 12 - gib = 2.0 ** 30 - return {"params_per_card_b": round(per_card / 1e9, 2), "weights_grads_gib": round(per_card * 6 / gib, 1), - "optimizer_gib_per_card": round(optimizer_card / gib, 1)} +def megatron_memory(model: Mapping[str, Any], miles: Mapping[str, Any], layout: Mapping[str, int], + trainers: Sequence[Node], sampling: Mapping[str, Any], + hf_config: str | Path | None = None) -> dict[str, Any]: + """Memory of the Miles layout from the estimator (``yeto.rl.fullparam.memory``, + framework megatron, N7.1c), all figures "推算". + + ``plan``: the layout as the plan writes it (distributed optimizer, the + plan's optimizer offload fraction, FP32 master offloaded with the Adam + state). ``island``: what the Miles island runs today in full-parameter mode + (distributed optimizer off, no optimizer offload; ``island_optimizer_gap``). + Without ``hf_config`` the shape takes hidden size and vocabulary from the + glm5_next recipe and one residual stream; the static items (weights, grads, + optimizer) do not depend on them, the activation items do.""" + from yeto.rl.fullparam.memory import ModelShape, TrainPlan, estimate, shape_from_hf_config # noqa: PLC0415 + + params = int(float(model["params_b"]) * 1e9) + experts = int(float(model["expert_params_b"]) * 1e9) + 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) + 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") + 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, + host_memory_gib=node.host_memory_gib, + seq_len=int(miles["max_tokens_per_gpu"]), tp=layout["tp"], pp=layout["pp"], + ep=layout["ep"], cp=layout["cp"], dsa_path="sparse_kernel", # Miles DSA layer, no dense mask + optimizer_offload_fraction=float(miles.get("optimizer_offload_fraction") or 0.0)) + out = {} + for key, plan in (("plan", base), + ("island", dataclasses.replace(base, distributed_optimizer=False, + optimizer_offload_fraction=0.0))): + est = estimate(shape, plan) + static = sum(est.items[k].bytes for k in ("params_local_bf16", "main_grads_local")) / 2**30 + out[key] = {**est.summary(), "weights_grads_gib_per_card": round(static, 1), + "host_optimizer_gib_per_node": round(est.total("host") / 2**30, 1)} + out["shape_source"] = shape.source + return out def megatron_args(model: Mapping[str, Any], miles: Mapping[str, Any], layout: Mapping[str, int], @@ -407,7 +437,7 @@ def expand(plan: Stage1Plan, *, hf_config: str | Path | None = None) -> Expanded "backend": backend, "nodes": [n.name for n in trainers], "cards": cards, "image": images.get("trainer"), "env": trainer_env, "miles": {"commit": miles["commit"], "megatron": miles["megatron"], "mbridge": miles["mbridge"], - "layout": layout, "static_memory_estimate": megatron_static_gib(model, layout), + "layout": layout, "memory_estimate": megatron_memory(model, miles, layout, trainers, sampling, hf_config), "model_args": model_args, "megatron_args": megatron_args(model, miles, layout, sampling), "drop_model_args": list(glm5_next.TE_ONLY_ARGV) if impl == glm5_next.LOCAL else [], @@ -722,20 +752,15 @@ def check(plan: Stage1Plan, out: Expanded) -> list[str]: if verl_max != [f"data.max_response_length={int(sampling['max_tokens'])}"]: errors.append("verl data.max_response_length differs from sampling.max_tokens") - # Static memory estimate must leave room for activations (A-TRN §2.3). - mem = out.trainer["miles"]["static_memory_estimate"] - hbm = trainers[0].hbm_gib if trainers else 0 - on_card = mem["weights_grads_gib"] + (0 if float(miles.get("optimizer_offload_fraction") or 0) >= 1 - else mem["optimizer_gib_per_card"]) - if on_card > 0.75 * hbm: - errors.append(f"static memory estimate {on_card:.1f} GiB/card leaves under 25 % of {hbm:g} GiB " - "for activations") - fraction = float(miles.get("optimizer_offload_fraction") or 0) - if trainers and fraction > 0: - host = mem["optimizer_gib_per_card"] * fraction * trainers[0].cards_per_node - if host > HOST_USABLE_FRACTION * trainers[0].host_memory_gib: - errors.append(f"offloaded optimizer estimate {host:.0f} GiB/node exceeds " - f"{HOST_USABLE_FRACTION:.0%} of {trainers[0].host_memory_gib:g} GiB host memory") + # Memory (N7.1c): the estimator decides; the plan's own figures must agree with it. + mem = out.trainer["miles"]["memory_estimate"]["plan"] + errors += [f"Miles memory estimate (推算): {p}" for p in mem["problems"]] + stated = miles.get("memory_estimate") or {} + for key, ours in (("weights_grads_gib_per_card", mem["weights_grads_gib_per_card"]), + ("host_optimizer_gib_per_node", mem["host_optimizer_gib_per_node"])): + if key in stated and abs(float(stated[key]) - ours) > MEMORY_TOLERANCE * max(ours, 1.0): + errors.append(f"trainer.miles.memory_estimate.{key} = {stated[key]} GiB, the estimator gives " + f"{ours} GiB (推算; tolerance {MEMORY_TOLERANCE:.0%})") return errors @@ -787,6 +812,8 @@ def gaps(plan: Stage1Plan, out: Expanded) -> list[str]: f"(--rl-model-recipe choices: {recipes})") if gap := _full_parameter_layout_gap(plan, out): found.append(gap) + if gap := _island_optimizer_gap(out): + found.append(gap) else: if not out.trainer["verl"]["commit"]: found.append("trainer.verl.commit is unset (verl pin)") @@ -814,6 +841,20 @@ def _full_parameter_layout_gap(plan: Stage1Plan, out: Expanded) -> str | None: f"{len(plan.nodes_of('trainer'))} trainer nodes, DP {out.trainer['miles']['layout']['dp']}") +def _island_optimizer_gap(out: Expanded) -> str | None: + """The Miles island in full-parameter mode turns the distributed optimizer off + (``adapters.miles.config._full_parameter_optimizer_attrs``: whole FP32 + masters for the shard manifest) and has no optimizer-offload flag, so the + plan's ``megatron_args`` offload is not what the island runs.""" + island = out.trainer["miles"]["memory_estimate"]["island"] + if island["fits"]: + return None + return ("Miles island runs full-parameter Megatron without the distributed optimizer and without " + "optimizer CPU offload (no island flag; the plan's megatron_args are not passed): estimator " + f"(推算) {island['device_gib_per_card']} GiB/card > {island['device_usable_gib']} GiB usable " + "(N7.1d)") + + # --------------------------------------------------------------------- CLI def run_steps(out: Expanded, *, ssh: Sequence[str] = ("ssh", "-o", "BatchMode=yes"), From 1d8d98c03ac1b68835c1d556799b2761c9afbd43 Mon Sep 17 00:00:00 2001 From: michael <78113554+chowhsu@users.noreply.github.com> Date: Mon, 19 Oct 2026 13:43:09 +0000 Subject: [PATCH 33/34] Fix the device-family fallback and the stage-1 test expectations Co-Authored-By: Claude Opus 5.5 --- tests/test_stage1_plan.py | 4 ++-- yeto/accel.py | 3 ++- 2 files changed, 4 insertions(+), 3 deletions(-) diff --git a/tests/test_stage1_plan.py b/tests/test_stage1_plan.py index 9965e963f..a69123d04 100644 --- a/tests/test_stage1_plan.py +++ b/tests/test_stage1_plan.py @@ -101,7 +101,7 @@ def test_steps_order_and_staging(shipped): def test_gaps_block_run_until_wired(shipped, capsys): assert any("N0.2" in g for g in shipped.gaps) - assert not any("full-parameter" in g for g in shipped.gaps) # trainer-only island accepted (T-INTEG2) + assert not any("refuses this layout" in g for g in shipped.gaps) # trainer-only island accepted (T-INTEG2) # Left: N0.2 (image), N0.2b (docker device flags), and N7.1d: the island runs # full-parameter Megatron without the distributed optimizer or optimizer offload. assert len(shipped.gaps) == 3 and "N7.1d" in shipped.gaps[2] @@ -175,7 +175,7 @@ def runner(argv, input=None, text=True): (lambda d: d["roles"].update(sampler=["ppu3", "ppu4"]), "both trainer and sampler"), (lambda d: d["rendezvous"].update(ppu2="ppu2"), "not an IP address"), (lambda d: d["rendezvous"].update(ppu2="10.0.0.99"), "differs from the inventory address"), - (lambda d: d["trainer"]["miles"].update(optimizer_offload_fraction=0.0), "activations"), + (lambda d: d["trainer"]["miles"].update(optimizer_offload_fraction=0.0), "exceeds the usable"), (lambda d: d.update(store="s3://bucket/x"), "oss://"), ]) def test_contradictions_are_errors(tmp_path, edit, message): diff --git a/yeto/accel.py b/yeto/accel.py index e7a504ac3..2e3461f0f 100644 --- a/yeto/accel.py +++ b/yeto/accel.py @@ -106,7 +106,8 @@ def device_family(explicit: str | None = None) -> str: """The device family of this rank: ``explicit``, else ``$YETO_DEVICE_FAMILY``, else the detected torch device type (a PPU then reads as ``"cuda"``, which selects the same runtime).""" - family = explicit or os.environ.get(DEVICE_FAMILY_ENV) # set by the island launcher or available_type() + # The island launcher sets the variable from the node inventory. + family = explicit or os.environ.get(DEVICE_FAMILY_ENV) or available_type() if family not in FAMILY_DEVICE_TYPES: raise ValueError(f"unknown device family {family!r} (known: {sorted(FAMILY_DEVICE_TYPES)})") return family From a8864cd56b7fb8e7237735d366b5b6e2dafe8406 Mon Sep 17 00:00:00 2001 From: michael <78113554+chowhsu@users.noreply.github.com> Date: Mon, 19 Oct 2026 14:26:39 +0000 Subject: [PATCH 34/34] Run the stage-1 Miles island with the distributed optimizer and CPU offload Close stage-1 gap N7.1d. The full-parameter Miles island turned the distributed optimizer off and had no optimizer-offload flag, so the plan's TP4 x PP3 x EP8 x DP2 layout ran at about 243 GiB/card (estimate) against 132.5 GiB usable. - New island/launcher flag --rl-optimizer-offload-fraction (TrainMemoryConfig.optimizer_offload_fraction). It translates to --use-distributed-optimizer --use-precision-aware-optimizer --optimizer-cpu-offload --overlap-cpu-optimizer-d2h-h2d --optimizer-offload-fraction F; Megatron's validate_args requires the precision-aware optimizer with CPU offload. - Whole FP32 masters are forced only for the full-parameter sync, which reads them; a no-sync trainer island keeps Miles' distributed optimizer. Offload together with the full-parameter sync is refused. - Cut plugin: accept the precision-aware optimizer when it is the FP32 CPU-offload path (needs Miles s25-distopt for the HybridDeviceOptimizer state); the GPU path stays refused. - yeto-stage1: the island argv carries the plan's offload fraction; the island memory estimate is derived from the island argv; plan errors when the island optimizer differs from megatron_args. plan now reports only N0.2 and N0.2b. - Tests: translation, flag forwarding, stage-1 gaps, and the DistOpt PP2/EP2 cut tests run against both the image pin and s25-distopt, plus per-rank resume-manifest sampling under the sharded optimizer. Co-Authored-By: Claude Opus 5.5 --- tests/test_rl_cut_distopt_pp_ep.py | 66 ++++++++++++++++++++++---- tests/test_rl_miles_adapter_config.py | 42 ++++++++++++++++ tests/test_rl_miles_cut_plugin.py | 9 ++++ tests/test_stage1_plan.py | 29 +++++++++-- yeto/launch/rl_island_task.py | 2 + yeto/rl/adapters/miles/config.py | 59 ++++++++++++++++------- yeto/rl/adapters/miles/cut_plugin.py | 24 +++++++++- yeto/rl/adapters/miles/island_entry.py | 3 +- yeto/rl/adapters/miles/run_flags.py | 8 ++++ yeto/rl/engine/run_config.py | 12 +++-- yeto/rl/stage1.py | 57 +++++++++++++++------- 11 files changed, 260 insertions(+), 51 deletions(-) diff --git a/tests/test_rl_cut_distopt_pp_ep.py b/tests/test_rl_cut_distopt_pp_ep.py index 3de8fabda..d687e8d38 100644 --- a/tests/test_rl_cut_distopt_pp_ep.py +++ b/tests/test_rl_cut_distopt_pp_ep.py @@ -3,7 +3,8 @@ CPU protocol test, NOT GPU evidence. Every rank is simulated in one process with duck-typed Megatron DistOpt leaves (``gbuf_ranges``, main-param shards, Adam moments) and the REAL fork-M5 helpers (``dp_invariant_state`` at the -Miles image pin, read from a local Miles checkout; skipped when absent). +Miles image pin and at the stage-1 branch s25-distopt, read from a local Miles +checkout; skipped when absent). save -> fresh rebuild -> restore -> re-export must be element-wise equal, and the full FP32 masters gathered for the republish (state_plugin.full_masters) must equal the pre-cut ones. A different DP/EP layout is refused. @@ -11,6 +12,7 @@ from __future__ import annotations +import dataclasses import importlib.util import os import subprocess @@ -23,26 +25,42 @@ from yeto.rl.adapters.miles import cut_plugin as cp from yeto.rl.adapters.miles import state_plugin as sp -MILES_CHECKOUT = os.environ.get("YETO_MILES_CHECKOUT", "/home/michael/work/miles-m3") -MILES_PIN = "c35702e" DPS_PATH = "miles/backends/megatron_utils/lora/dp_invariant_state.py" +# The fork-M5 helpers at the Miles image pin, and at the stage-1 fork branch that +# carries the CPU-offload (HybridDeviceOptimizer) state (S25 N7.1d, not pushed). +DPS_SOURCES = { + "m5-c35702e": (os.environ.get("YETO_MILES_CHECKOUT", "/home/michael/work/miles-m3"), "c35702e"), + "s25-distopt": (os.environ.get("YETO_MILES_DISTOPT_CHECKOUT", "/home/michael/work/miles-next"), "s25-distopt"), +} -def _load_dps(): +def _load_dps(checkout, ref): try: - src = subprocess.run(["git", "-C", MILES_CHECKOUT, "show", f"{MILES_PIN}:{DPS_PATH}"], + src = subprocess.run(["git", "-C", checkout, "show", f"{ref}:{DPS_PATH}"], capture_output=True, check=True, text=True).stdout except (OSError, subprocess.CalledProcessError): return None - spec = importlib.util.spec_from_loader("_yeto_test_dps", loader=None) + spec = importlib.util.spec_from_loader(f"_yeto_test_dps_{ref.replace('-', '_')}", loader=None) module = importlib.util.module_from_spec(spec) sys.modules[spec.name] = module exec(compile(src, DPS_PATH, "exec"), module.__dict__) # noqa: S102 - pinned local source return module -DPS = _load_dps() -needs_dps = pytest.mark.skipif(DPS is None, reason="Miles checkout with the fork-M5 helpers not available") +_LOADED = {label: _load_dps(*source) for label, source in DPS_SOURCES.items()} +DPS = None # set per test by the ``dps`` fixture + + +@pytest.fixture(params=sorted(DPS_SOURCES)) +def dps(request, monkeypatch): + module = _LOADED[request.param] + if module is None: + pytest.skip(f"Miles checkout with the fork-M5 helpers at {DPS_SOURCES[request.param]} not available") + monkeypatch.setattr(sys.modules[__name__], "DPS", module) + return module + + +needs_dps = pytest.mark.usefixtures("dps") class _Range: @@ -255,6 +273,38 @@ def test_save_restore_reexport_is_elementwise_equal(tmp_path, layout): assert torch.equal(a, b) +@needs_dps +@pytest.mark.parametrize("layout", ["pp2dp2", "ep2"]) +def test_resume_manifest_samples_each_rank_own_ranges_and_survives_the_cut(tmp_path, layout): + # N7.6 under the sharded optimizer (N7.1d): every rank records the moments of the + # parameters it owns a range of, and the restored rank reproduces them exactly. + from yeto.rl.engine.resume_manifest import OptimizerResumeState, restore_problems, ResumeManifest + + ranks = _world(layout) + _train(ranks) + files = _save(ranks, tmp_path) + saved = {f["path"]: OptimizerResumeState.from_dict(f["resume"]) for f in files} + for rank, f in zip(ranks, files, strict=True): + owned = set(_export(rank)["entries"]) + assert saved[f["path"]].moments and set(saved[f["path"]].moments) <= owned + fresh = _world(layout) + restored = {} + for rank in fresh: + out = cp.restore_cut_shard(rank, directory=str(tmp_path), files=files, cut_id="c1") + assert "refused" not in out, out + restored[out["path"]] = OptimizerResumeState.from_dict(out["resume"]) + + def manifest(states): + return ResumeManifest(global_step=1, data_watermark=8, scheduler_increment=8, ranks=states) + + assert restore_problems(manifest(saved), manifest(restored)) == [] + # a rank that records another rank's moments (here: none of its own) is caught + first, last = sorted(saved)[0], sorted(saved)[-1] + swapped = {**restored, first: dataclasses.replace(restored[first], moments=restored[last].moments)} + assert saved[first].moments != saved[last].moments + assert any(first in p for p in restore_problems(manifest(saved), manifest(swapped))) + + @needs_dps def test_ep_layout_change_is_refused(tmp_path): ranks = _world("ep2") diff --git a/tests/test_rl_miles_adapter_config.py b/tests/test_rl_miles_adapter_config.py index 9e4d92a96..503fe2e28 100644 --- a/tests/test_rl_miles_adapter_config.py +++ b/tests/test_rl_miles_adapter_config.py @@ -262,6 +262,44 @@ def test_full_parameter_mode_turns_the_distributed_optimizer_off(): assert namespace.use_distributed_optimizer is False +def test_no_sync_full_parameter_trainer_keeps_the_distributed_optimizer(): + # N7.1d: only the full-parameter sync reads whole masters; a store-fed trainer + # island shards the optimizer state over DP and may offload it to the host. + cfg = sub(make_config(), "trainable", parameter_mode="full", target_modules=()) + cfg = dataclasses.replace(cfg, train_memory=rc.TrainMemoryConfig(optimizer_offload_fraction=1.0)) + launch = mc.translate_run_config(cfg, AlgorithmSpec(), full_parameter_sync=False) + assert "use_distributed_optimizer" not in launch.runtime_attrs + joined = " ".join(launch.argv) + assert ("--use-distributed-optimizer --use-precision-aware-optimizer --optimizer-cpu-offload " + "--overlap-cpu-optimizer-d2h-h2d --optimizer-offload-fraction 1.0") in joined + mc.validate_parsed_args(_parsed(use_distributed_optimizer=True), launch, num_cells=lambda a: 1) + + +def test_optimizer_offload_with_the_full_parameter_sync_is_refused(): + cfg = sub(make_config(), "trainable", parameter_mode="full", target_modules=()) + cfg = dataclasses.replace(cfg, train_memory=rc.TrainMemoryConfig(optimizer_offload_fraction=0.5)) + with pytest.raises(mc.MilesConfigError, match="needs the distributed optimizer"): + mc.translate_run_config(cfg, AlgorithmSpec()) + + +def test_build_ports_launch_marks_a_no_sync_trainer_island(monkeypatch): + from yeto.rl.adapters.miles import island_entry + + seen = {} + + def translate(config, algorithm, **kw): + seen.update(kw) + return "launch" + + monkeypatch.setattr(mc, "translate_run_config", translate) + args = SimpleNamespace(rl_single_island_no_sync=True, rl_sample_role="trainer", rl_max_policy_age=0) + assert island_entry.build_ports_launch(args, make_config()) == "launch" + assert seen["full_parameter_sync"] is False + args.rl_single_island_no_sync = False + island_entry.build_ports_launch(args, make_config()) + assert seen["full_parameter_sync"] is True + + def test_lora_mode_leaves_the_distributed_optimizer_to_miles(): launch = mc.translate_run_config(make_config(), AlgorithmSpec()) assert "use_distributed_optimizer" not in launch.runtime_attrs @@ -433,6 +471,10 @@ def test_train_memory_from_args_refuses_a_budget_below_seq_len(): args.rl_max_tokens_per_gpu = 32768 assert rc._train_memory(args) == rc.TrainMemoryConfig(True, 32768) assert rc._train_memory(SimpleNamespace(seq_len=8)) == rc.TrainMemoryConfig() + assert rc._train_memory(SimpleNamespace(seq_len=8, rl_optimizer_offload_fraction=1.0)) == \ + rc.TrainMemoryConfig(optimizer_offload_fraction=1.0) + with pytest.raises(ValueError, match="offload-fraction"): + rc._train_memory(SimpleNamespace(seq_len=8, rl_optimizer_offload_fraction=1.5)) @pytest.mark.ray_local # upstream parse_args can start Ray on this machine diff --git a/tests/test_rl_miles_cut_plugin.py b/tests/test_rl_miles_cut_plugin.py index 5fc33a259..17cd4ee06 100644 --- a/tests/test_rl_miles_cut_plugin.py +++ b/tests/test_rl_miles_cut_plugin.py @@ -147,6 +147,15 @@ def test_unsupported_configurations(args, match): assert any(match in p for p in config_problems(args)) +def test_the_fp32_cpu_offload_path_of_the_precision_aware_optimizer_is_carried(): + # N7.1d: Megatron's --optimizer-cpu-offload requires --use-precision-aware-optimizer. + offload = dict(use_precision_aware_optimizer=True, optimizer_cpu_offload=True, use_distributed_optimizer=True) + assert not config_problems(SimpleNamespace(**offload)) + assert not config_problems(SimpleNamespace(**offload, main_params_dtype="torch.float32")) + low = config_problems(SimpleNamespace(**offload, exp_avg_dtype="torch.bfloat16")) + assert any("exp_avg_dtype" in p for p in low) + + def test_ep_is_carried_by_same_shape_cuts_but_not_by_resharding(): args = SimpleNamespace(expert_model_parallel_size=2, use_distributed_optimizer=True) assert not config_problems(args) diff --git a/tests/test_stage1_plan.py b/tests/test_stage1_plan.py index a69123d04..c2639852e 100644 --- a/tests/test_stage1_plan.py +++ b/tests/test_stage1_plan.py @@ -78,6 +78,10 @@ def test_trainer_is_bf16_with_host_optimizer(shipped): margs = " ".join(shipped.trainer["miles"]["megatron_args"]) assert "--bf16" in margs and "fp8" not in margs assert "--optimizer-cpu-offload" in margs + assert "--use-precision-aware-optimizer" in margs # Megatron asserts it with CPU offload + island = shipped.trainer["miles"]["island_argv"] + assert island[island.index("--rl-optimizer-offload-fraction") + 1] == "1.0" + assert stage1.island_optimizer(island) == {"distributed_optimizer": True, "optimizer_offload_fraction": 1.0} assert "--decoder-first-pipeline-num-layers 15 --decoder-last-pipeline-num-layers 15" in margs assert shipped.trainer["miles"]["drop_model_args"] == [] # TE stays the default (ruling 10-19 (3)) @@ -102,9 +106,9 @@ def test_steps_order_and_staging(shipped): def test_gaps_block_run_until_wired(shipped, capsys): assert any("N0.2" in g for g in shipped.gaps) assert not any("refuses this layout" in g for g in shipped.gaps) # trainer-only island accepted (T-INTEG2) - # Left: N0.2 (image), N0.2b (docker device flags), and N7.1d: the island runs - # full-parameter Megatron without the distributed optimizer or optimizer offload. - assert len(shipped.gaps) == 3 and "N7.1d" in shipped.gaps[2] + # Left: N0.2 (image) and N0.2b (docker device flags); N7.1d closed (T-N71D: + # the island keeps the distributed optimizer and offloads Adam to the host). + assert len(shipped.gaps) == 2 and not any("N7.1d" in g for g in shipped.gaps) assert stage1.main(["run", str(PLAN), "--confirm-machines-free"]) == 2 assert "run refused" in capsys.readouterr().err @@ -221,7 +225,24 @@ def test_memory_figures_of_the_plan_match_the_estimator(shipped): assert mem["plan"]["basis"] == "推算" and mem["plan"]["fits"] assert mem["plan"]["weights_grads_gib_per_card"] == pytest.approx(78.7, abs=0.1) assert mem["plan"]["host_optimizer_gib_per_node"] == pytest.approx(1196, rel=0.02) - assert not mem["island"]["fits"] # no distributed optimizer, no offload (gap N7.1d) + assert mem["island"] == mem["plan"] # N7.1d: the island runs the plan's optimizer + + +def test_an_island_without_optimizer_offload_is_a_gap(tmp_path): + def edit(doc): + doc["trainer"]["miles"]["optimizer_offload_fraction"] = 0.0 + doc["trainer"]["miles"].pop("memory_estimate", None) + out = _variant(tmp_path, edit) + assert "--rl-optimizer-offload-fraction" not in out.trainer["miles"]["island_argv"] + assert any("N7.1d" in g and "does not fit" in g for g in out.gaps) + + +def test_island_optimizer_follows_the_island_translation(): + full = ["--parameter-mode", "full"] + assert stage1.island_optimizer(full) == {"distributed_optimizer": False, "optimizer_offload_fraction": 0.0} + no_sync = [*full, "--rl-single-island-no-sync", "--rl-sample-role", "trainer"] + assert stage1.island_optimizer(no_sync)["distributed_optimizer"] + assert stage1.island_optimizer(["--rl-single-island-no-sync"])["distributed_optimizer"] # LoRA def test_memory_figure_that_disagrees_is_an_error(tmp_path): diff --git a/yeto/launch/rl_island_task.py b/yeto/launch/rl_island_task.py index a5b443aa5..1185f4c9a 100644 --- a/yeto/launch/rl_island_task.py +++ b/yeto/launch/rl_island_task.py @@ -155,6 +155,8 @@ def make_miles_island_task( + (" --rl-activation-recompute" if getattr(args, "rl_activation_recompute", False) else "") + (f" --rl-max-tokens-per-gpu {int(args.rl_max_tokens_per_gpu)}" if getattr(args, "rl_max_tokens_per_gpu", None) is not None else "") + + (f" --rl-optimizer-offload-fraction {float(args.rl_optimizer_offload_fraction)!r}" + if float(getattr(args, "rl_optimizer_offload_fraction", 0.0) or 0.0) > 0 else "") + f" --seed {args.seed}" f" --wan-streams {args.wan_streams}" " --miles-root ~/miles" diff --git a/yeto/rl/adapters/miles/config.py b/yeto/rl/adapters/miles/config.py index 1562325fd..b6ead7169 100644 --- a/yeto/rl/adapters/miles/config.py +++ b/yeto/rl/adapters/miles/config.py @@ -471,6 +471,7 @@ def _check_sglang_tp(value, config): # translated below (Megatron full recompute / Miles dynamic batching) "train_memory.activation_recompute": _ok, "train_memory.max_tokens_per_gpu": _ok, + "train_memory.optimizer_offload_fraction": _ok, **{f"sample_buffer.{name}": _ok for name in ("role", "store", "max_backlog_groups", "wait_s", "min_round_s")}, } @@ -715,6 +716,7 @@ def translate_run_config( extra_argv: Sequence[str] = (), tito_parser_resolver: Callable[[str], tuple[str | None, str | None]] | None = None, max_policy_age: int = 0, + full_parameter_sync: bool = True, ) -> MilesLaunchArgs: """Pure translation of an ``RLRunConfig`` into upstream Miles argv. @@ -724,7 +726,11 @@ def translate_run_config( partial rollout (carry-over at the cut-off); a partitioned island overlaps rollout with training through Miles' fully_async (P0-S), so the all-samples metadata hook and the carry-over buffer filter are not set: - the trained-groups hook reports the round instead.""" + the trained-groups hook reports the round instead. + + ``full_parameter_sync`` (default True): a full-parameter trainer feeds the + ports full-parameter sync, which needs whole FP32 masters; False for a + no-sync trainer island (``_full_parameter_optimizer_attrs``).""" if not isinstance(algorithm, AlgorithmSpec): raise TypeError("algorithm must be an AlgorithmSpec") @@ -1160,7 +1166,8 @@ def translate_run_config( rollout_overlap=rollout_overlap, runtime_attrs={ **algorithm.to_legacy_runtime_attrs(), - **_full_parameter_optimizer_attrs(trainable), + **_full_parameter_optimizer_attrs(trainable, config.train_memory, + full_parameter_sync=full_parameter_sync), # rl-algo-critic-family 4.1: receipt init source (warm-up product) **({"yeto_rl_critic_init_sha256": config.algorithm.critic.init_sha256} if getattr(config.algorithm, "critic", None) is not None @@ -1180,9 +1187,18 @@ def translate_run_config( SAMPLER_ONLY_FLAG = "--debug-rollout-only" +# Megatron's --optimizer-cpu-offload runs a HybridDeviceOptimizer under the +# distributed optimizer and requires the precision-aware optimizer path (its +# validate_args asserts it); the FP32 masters and moments stay FP32 (Megatron's +# default dtypes), which the cut's name-keyed optimizer export carries. +OPTIMIZER_OFFLOAD_ARGV = ("--use-distributed-optimizer", "--use-precision-aware-optimizer", + "--optimizer-cpu-offload", "--overlap-cpu-optimizer-d2h-h2d") + + def train_memory_argv(memory: Any) -> tuple[str, ...]: - """Megatron full uniform recompute (one layer per chunk) and Miles dynamic - batching; the same values the validated Flash-Next recipe uses.""" + """Megatron full uniform recompute (one layer per chunk), Miles dynamic + batching (the same values the validated Flash-Next recipe uses) and the + optimizer host offload.""" values: list[str] = [] if memory.activation_recompute: @@ -1191,6 +1207,9 @@ def train_memory_argv(memory: Any) -> tuple[str, ...]: if memory.max_tokens_per_gpu is not None: values.extend(("--use-dynamic-batch-size", "--max-tokens-per-gpu", str(memory.max_tokens_per_gpu))) + fraction = float(getattr(memory, "optimizer_offload_fraction", 0.0) or 0.0) + if fraction > 0: + values.extend((*OPTIMIZER_OFFLOAD_ARGV, "--optimizer-offload-fraction", repr(fraction))) return tuple(values) @@ -1208,23 +1227,29 @@ def sample_buffer_argv(role: str, *, max_policy_age: int) -> tuple[str, ...]: return (*argv, "--save-hf", RELEASE_EXPORTER_SAVE_HF) -def _full_parameter_optimizer_attrs(trainable: Any) -> dict[str, Any]: - """Keep every FP32 master whole in full-parameter mode. +def _full_parameter_optimizer_attrs(trainable: Any, memory: Any, *, full_parameter_sync: bool) -> dict[str, Any]: + """Keep every FP32 master whole when the full-parameter sync reads them. - Miles' ``set_default_megatron_args`` turns the distributed optimizer on for - Adam regardless of argv, so ``--no-use-distributed-optimizer`` would be - overwritten. Under the distributed optimizer each ``param.main_param`` is a - flat slice of the owned range even at DP=1, and the full-parameter shard - manifest (Miles ``full_parameter_state._master_parameter``) refuses it as - an incomplete master. The attribute is applied in ``preprocess_args``, - after Miles' defaults and before Megatron's ``validate_args``. With engines - on the island full mode is DP=1; a trainer-only island (samples from the - store) may run DP>1, and then every DP replica keeps the whole optimizer - state of its local parameters (more device memory; stage-1 gap N7.1d). + The ports full-parameter sync (presets ``full-parameter`` / + ``sao-streaming-full``) builds its shard manifest from whole masters (Miles + ``full_parameter_state._master_parameter`` refuses the flat owned-range + slice that ``param.main_param`` is under the distributed optimizer, even at + DP=1). Miles' ``set_default_megatron_args`` turns the distributed optimizer + on for Adam regardless of argv, so the attribute is applied in + ``preprocess_args``, after Miles' defaults and before Megatron's + ``validate_args``. + + A no-sync trainer island (samples from the store, L4.6) has no such reader: + it keeps Miles' distributed optimizer, so DP replicas shard the optimizer + state, and its cuts carry each rank's owned ranges (``cut_plugin``, N7.1d). """ - if trainable.parameter_mode != "full": + if trainable.parameter_mode != "full" or not full_parameter_sync: return {} + if float(getattr(memory, "optimizer_offload_fraction", 0.0) or 0.0) > 0: + raise MilesConfigError( + "--rl-optimizer-offload-fraction needs the distributed optimizer; the full-parameter sync " + "keeps whole FP32 masters (use a no-sync trainer island)") return {"use_distributed_optimizer": False} diff --git a/yeto/rl/adapters/miles/cut_plugin.py b/yeto/rl/adapters/miles/cut_plugin.py index a1321114c..bb22edb07 100644 --- a/yeto/rl/adapters/miles/cut_plugin.py +++ b/yeto/rl/adapters/miles/cut_plugin.py @@ -26,7 +26,7 @@ (DistributedOptimizer) shard is written back only onto identical per-parameter ranges of the rebuilt optimizer -- no peer shard is read (TP/PP>1 and EP>1 included). Refused like fork-M5: fp16 (loss-scaler state not saved), -precision-aware optimizer, partial DistOpt instances, CP>1, and a trainer +precision-aware optimizer (except as the FP32 CPU-offload path), partial DistOpt instances, CP>1, and a trainer whose trainable parameters mix LoRA adapters with others; the DP-resharding restore also refuses EP>1 and TP/PP>1 with DistOpt. A full-parameter shard holds the whole model copy on every DP rank plus this rank's optimizer range @@ -232,7 +232,7 @@ def config_problems(args: Any, coord: Mapping[str, int] | None = None, *, reshar if getattr(args, "fp16", False): out.append("--fp16 keeps a dynamic loss-scaler state that the cut does not carry") if getattr(args, "use_precision_aware_optimizer", False): - out.append("--use-precision-aware-optimizer stores scaled/low-precision state (not supported)") + out += _precision_aware_problems(args) if int(getattr(args, "num_distributed_optimizer_instances", 1) or 1) > 1: out.append("--num-distributed-optimizer-instances > 1 (partial DP optimizer groups)") cp = int((coord or {}).get("cp_size") or getattr(args, "context_parallel_size", 1) or 1) @@ -249,6 +249,26 @@ def config_problems(args: Any, coord: Mapping[str, int] | None = None, *, reshar return out +_FP32_NAMES = ("fp32", "float32", "torch.float32") + + +def _precision_aware_problems(args: Any) -> list[str]: + """The precision-aware optimizer is carried only as Megatron's CPU-offload path with FP32 state. + + ``--optimizer-cpu-offload`` requires ``--use-precision-aware-optimizer``; + its HybridDeviceOptimizer keeps standalone FP32 masters and moments, which + the fork's name-keyed export reads unscaled (Miles ``dp_invariant_state``, + s25-distopt). On the device the precision-aware path keeps scaled or + low-precision state, which a cut does not represent. + """ + if not getattr(args, "optimizer_cpu_offload", False): + return ["--use-precision-aware-optimizer without --optimizer-cpu-offload stores scaled/low-precision " + "state (not supported)"] + low = [name for name in ("main_params_dtype", "exp_avg_dtype", "exp_avg_sq_dtype") + if str(getattr(args, name, "fp32")) not in _FP32_NAMES] + return [f"--optimizer-cpu-offload with non-FP32 {low} (not supported)"] if low else [] + + def _require(actor: Any, backend: Any, *, reshard: bool = False) -> dict[str, int]: coord = backend.coord() problems = config_problems(actor.args, coord, reshard=reshard) diff --git a/yeto/rl/adapters/miles/island_entry.py b/yeto/rl/adapters/miles/island_entry.py index 27150383d..9a031e335 100644 --- a/yeto/rl/adapters/miles/island_entry.py +++ b/yeto/rl/adapters/miles/island_entry.py @@ -150,7 +150,8 @@ def build_ports_launch(args, run_config, extra_argv=()): ), ) return translate_run_config(run_config, base_algorithm, extra_argv=tuple(extra_argv), - max_policy_age=int(getattr(args, "rl_max_policy_age", 0) or 0)) + max_policy_age=int(getattr(args, "rl_max_policy_age", 0) or 0), + full_parameter_sync=not _buffer_island(args)) _ELASTIC_COMPANIONS = rl_option_pairs( diff --git a/yeto/rl/adapters/miles/run_flags.py b/yeto/rl/adapters/miles/run_flags.py index c44008683..a6a9c6dad 100644 --- a/yeto/rl/adapters/miles/run_flags.py +++ b/yeto/rl/adapters/miles/run_flags.py @@ -175,6 +175,14 @@ "at most this many tokens per GPU (must be >= --seq-len). Default: fixed " "micro batches", ), + shared_flag( + "--rl-optimizer-offload-fraction", + type=float, + default=0.0, + help="ports: share (0, 1] of the Adam state (FP32 masters and moments) " + "that Megatron keeps and updates on the host (--optimizer-cpu-offload, " + "with the distributed optimizer). Default 0: all on the device", + ), shared_flag( "--rl-distributed-timeout-minutes", type=int, diff --git a/yeto/rl/engine/run_config.py b/yeto/rl/engine/run_config.py index 548a902da..b2c82c4ea 100644 --- a/yeto/rl/engine/run_config.py +++ b/yeto/rl/engine/run_config.py @@ -244,15 +244,17 @@ class BatchConfig: @dataclass(frozen=True) class TrainMemoryConfig: - """Trainer activation-memory knobs (--rl-activation-recompute, - --rl-max-tokens-per-gpu). Defaults keep the historical argv: no - recompute, fixed micro batches.""" + """Trainer memory knobs (--rl-activation-recompute, --rl-max-tokens-per-gpu, + --rl-optimizer-offload-fraction). Defaults keep the historical argv: no + recompute, fixed micro batches, optimizer state on the device.""" # Recompute every layer's activations in backward (keep only layer inputs). activation_recompute: bool = False # Token budget of one packed micro batch per GPU (dynamic batching); # None = fixed micro batches of --micro-batch-size samples. max_tokens_per_gpu: int | None = None + # Share of the optimizer state kept and updated on the host (0 = none). + optimizer_offload_fraction: float = 0.0 @dataclass(frozen=True) @@ -1078,9 +1080,13 @@ def _train_memory(args) -> TrainMemoryConfig: raise ValueError( f"--rl-max-tokens-per-gpu {max_tokens} is below --seq-len {args.seq_len}" ) + offload = float(getattr(args, "rl_optimizer_offload_fraction", 0.0) or 0.0) + if not 0.0 <= offload <= 1.0: + raise ValueError(f"--rl-optimizer-offload-fraction {offload} is not in [0, 1]") return TrainMemoryConfig( activation_recompute=bool(getattr(args, "rl_activation_recompute", False)), max_tokens_per_gpu=int(max_tokens) if max_tokens is not None else None, + optimizer_offload_fraction=offload, ) diff --git a/yeto/rl/stage1.py b/yeto/rl/stage1.py index 192e0e7f7..4825e9caf 100644 --- a/yeto/rl/stage1.py +++ b/yeto/rl/stage1.py @@ -200,16 +200,32 @@ def megatron_layout(miles: Mapping[str, Any], trainer_cards: int, cards_per_node "expert_dp": dp * tp * cp // (ep * etp)} +def island_optimizer(argv: Sequence[str]) -> dict[str, Any]: + """The optimizer the Miles island runs for this argv (``adapters.miles.config``): + a no-sync trainer island in full-parameter mode keeps Miles' distributed + optimizer (only the full-parameter sync needs whole FP32 masters, N7.1d); + ``--rl-optimizer-offload-fraction`` sets the host share of the Adam state.""" + def value(flag: str, default: str) -> str: + return argv[argv.index(flag) + 1] if flag in argv else default + + no_sync_island = "--rl-single-island-no-sync" in argv and value("--rl-sample-role", "both") != "both" + whole_masters = value("--parameter-mode", "lora") == "full" and not no_sync_island + return {"distributed_optimizer": not whole_masters, + "optimizer_offload_fraction": float(value("--rl-optimizer-offload-fraction", "0"))} + + def megatron_memory(model: Mapping[str, Any], miles: Mapping[str, Any], layout: Mapping[str, int], trainers: Sequence[Node], sampling: Mapping[str, Any], - hf_config: str | Path | None = None) -> dict[str, Any]: + hf_config: str | Path | None = None, *, + island: Mapping[str, Any] | None = None) -> dict[str, Any]: """Memory of the Miles layout from the estimator (``yeto.rl.fullparam.memory``, framework megatron, N7.1c), all figures "推算". ``plan``: the layout as the plan writes it (distributed optimizer, the plan's optimizer offload fraction, FP32 master offloaded with the Adam - state). ``island``: what the Miles island runs today in full-parameter mode - (distributed optimizer off, no optimizer offload; ``island_optimizer_gap``). + state). ``island``: the optimizer the Miles island argv runs + (:func:`island_optimizer`; default: the plan's); ``island_optimizer_gap`` + blocks the run when it does not fit. Without ``hf_config`` the shape takes hidden size and vocabulary from the glm5_next recipe and one residual stream; the static items (weights, grads, optimizer) do not depend on them, the activation items do.""" @@ -235,8 +251,7 @@ def megatron_memory(model: Mapping[str, Any], miles: Mapping[str, Any], layout: optimizer_offload_fraction=float(miles.get("optimizer_offload_fraction") or 0.0)) out = {} for key, plan in (("plan", base), - ("island", dataclasses.replace(base, distributed_optimizer=False, - optimizer_offload_fraction=0.0))): + ("island", dataclasses.replace(base, **(island or {})))): est = estimate(shape, plan) static = sum(est.items[k].bytes for k in ("params_local_bf16", "main_grads_local")) / 2**30 out[key] = {**est.summary(), "weights_grads_gib_per_card": round(static, 1), @@ -268,8 +283,10 @@ def megatron_args(model: Mapping[str, Any], miles: Mapping[str, Any], layout: Ma "--rollout-max-response-len", str(int(sampling["max_tokens"]))] fraction = float(miles.get("optimizer_offload_fraction") or 0.0) if fraction > 0: - args += ["--optimizer-cpu-offload", "--optimizer-offload-fraction", repr(fraction), - "--overlap-cpu-optimizer-d2h-h2d"] + from yeto.rl.adapters.miles.config import OPTIMIZER_OFFLOAD_ARGV # noqa: PLC0415 + + args += [a for a in OPTIMIZER_OFFLOAD_ARGV if a not in args] + args += ["--optimizer-offload-fraction", repr(fraction)] return args @@ -403,6 +420,7 @@ def expand(plan: Stage1Plan, *, hf_config: str | Path | None = None) -> Expanded model_dir = str(model["bf16_source"]["local_dir"]) state_dir = state_dir_of(plan) impl = str(miles["transformer_impl"]) + 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"]), @@ -426,6 +444,7 @@ def expand(plan: Stage1Plan, *, hf_config: str | Path | None = None) -> Expanded "--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", "--event-tape", f"{state_dir}/events.jsonl", "--audit-dir", f"{state_dir}/audit", *_provenance_args(str(task["reward"]))] @@ -437,7 +456,8 @@ def expand(plan: Stage1Plan, *, hf_config: str | Path | None = None) -> Expanded "backend": backend, "nodes": [n.name for n in trainers], "cards": cards, "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), + "layout": layout, "memory_estimate": megatron_memory(model, miles, layout, trainers, sampling, hf_config, + island=island_optimizer(island_argv)), "model_args": model_args, "megatron_args": megatron_args(model, miles, layout, sampling), "drop_model_args": list(glm5_next.TE_ONLY_ARGV) if impl == glm5_next.LOCAL else [], @@ -752,6 +772,12 @@ def check(plan: Stage1Plan, out: Expanded) -> list[str]: if verl_max != [f"data.max_response_length={int(sampling['max_tokens'])}"]: errors.append("verl data.max_response_length differs from sampling.max_tokens") + # N7.1d: the island runs the plan's optimizer (megatron_args are the reference). + runs = island_optimizer(island) + planned = {"distributed_optimizer": "--use-distributed-optimizer" in margs, + "optimizer_offload_fraction": float(miles.get("optimizer_offload_fraction") or 0.0)} + if runs != planned: + errors.append(f"Miles island optimizer {runs} differs from the plan's megatron_args {planned}") # Memory (N7.1c): the estimator decides; the plan's own figures must agree with it. mem = out.trainer["miles"]["memory_estimate"]["plan"] errors += [f"Miles memory estimate (推算): {p}" for p in mem["problems"]] @@ -842,17 +868,16 @@ def _full_parameter_layout_gap(plan: Stage1Plan, out: Expanded) -> str | None: def _island_optimizer_gap(out: Expanded) -> str | None: - """The Miles island in full-parameter mode turns the distributed optimizer off - (``adapters.miles.config._full_parameter_optimizer_attrs``: whole FP32 - masters for the shard manifest) and has no optimizer-offload flag, so the - plan's ``megatron_args`` offload is not what the island runs.""" + """The optimizer the Miles island argv runs (:func:`island_optimizer`) must fit: + the plan's ``megatron_args`` are the reference, the island translates its own + flags (``--rl-optimizer-offload-fraction``; the distributed optimizer stays + on for a no-sync trainer island).""" island = out.trainer["miles"]["memory_estimate"]["island"] if island["fits"]: return None - return ("Miles island runs full-parameter Megatron without the distributed optimizer and without " - "optimizer CPU offload (no island flag; the plan's megatron_args are not passed): estimator " - f"(推算) {island['device_gib_per_card']} GiB/card > {island['device_usable_gib']} GiB usable " - "(N7.1d)") + runs = island_optimizer(out.trainer["miles"]["island_argv"]) + return (f"Miles island optimizer {runs} does not fit: estimator (推算) {island['device_gib_per_card']} " + f"GiB/card > {island['device_usable_gib']} GiB usable (N7.1d)") # --------------------------------------------------------------------- CLI