diff --git a/docker/README.md b/docker/README.md new file mode 100644 index 000000000..936a24eb0 --- /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/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). + +## 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/examples/stage1/glm53-flash-ppu.yaml b/examples/stage1/glm53-flash-ppu.yaml new file mode 100644 index 000000000..5f23b3a1d --- /dev/null +++ b/examples/stage1/glm53-flash-ppu.yaml @@ -0,0 +1,133 @@ +# 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 +# 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 + 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 + 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 + 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 + # 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 new file mode 100644 index 000000000..f05a67757 --- /dev/null +++ b/examples/stage1/ppu-inventory.yaml @@ -0,0 +1,66 @@ +# 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; 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 + 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 + 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 + 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 + 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 + 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 + 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 + 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 + 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..0f46c88cb 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.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_accel.py b/tests/test_accel.py index 190a25c36..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): @@ -293,3 +294,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(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(DEVICE_FAMILY_ENV, raising=False) + assert accel.device_family() == "npu" # auto-detect + 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 + + +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_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_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] == [ diff --git a/tests/test_engine_images.py b/tests/test_engine_images.py new file mode 100644 index 000000000..eee334860 --- /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 engine_images as 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_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/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/tests/test_megatron_learner_device_family.py b/tests/test_megatron_learner_device_family.py new file mode 100644 index 000000000..8846476cf --- /dev/null +++ b/tests/test_megatron_learner_device_family.py @@ -0,0 +1,67 @@ +"""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.gpu_spec import DEVICE_FAMILY_ENV +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(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"}) + + +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_records_proxy.py b/tests/test_records_proxy.py new file mode 100644 index 000000000..e9a77ace3 --- /dev/null +++ b/tests/test_records_proxy.py @@ -0,0 +1,430 @@ +"""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_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 + + 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/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_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_fullparam_memory.py b/tests/test_rl_fullparam_memory.py index d706d7e34..ef664cf52 100644 --- a/tests/test_rl_fullparam_memory.py +++ b/tests/test_rl_fullparam_memory.py @@ -116,3 +116,96 @@ 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"}, "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"), +]) +def test_malformed_megatron_plans_are_refused(change, match): + with pytest.raises(MemoryPlanError, match=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) + + +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_hetero_sampler_engines.py b/tests/test_rl_hetero_sampler_engines.py new file mode 100644 index 000000000..7f2a5a442 --- /dev/null +++ b/tests/test_rl_hetero_sampler_engines.py @@ -0,0 +1,626 @@ +"""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", "cpu_arch": "aarch64", "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"), +]) +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) + 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_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 + + 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) + 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/tests/test_rl_miles_adapter_config.py b/tests/test_rl_miles_adapter_config.py index fdfa4a0e3..503fe2e28 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 @@ -250,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 @@ -421,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_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_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_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/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/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_resume_manifest.py b/tests/test_rl_resume_manifest.py new file mode 100644 index 000000000..d972c4e33 --- /dev/null +++ b/tests/test_rl_resume_manifest.py @@ -0,0 +1,433 @@ +"""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)) + 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 + + 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()) + 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 + + +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/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..3251f053e 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)) @@ -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") @@ -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"): @@ -115,11 +124,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 +168,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") @@ -257,7 +284,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 +331,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/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/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..9ff34c178 --- /dev/null +++ b/tests/test_stage1_gaps.py @@ -0,0 +1,283 @@ +"""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(): + # 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", 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", "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] + 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 new file mode 100644 index 000000000..c2639852e --- /dev/null +++ b/tests/test_stage1_plan.py @@ -0,0 +1,252 @@ +"""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 +from types import SimpleNamespace + +import pytest +import yaml + +from yeto.rl 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 "--target-dir" not in req # the directory follows the template's scheme + 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["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): + 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 "--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)) + + +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("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 + + +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) 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 + + +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"] + + +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"), + (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["rendezvous"].update(ppu2="10.0.0.99"), "differs from the inventory address"), + (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): + 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"] + 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): + 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) + + +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 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): + 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/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/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/accel.py b/yeto/accel.py index b2b4e1485..2e3461f0f 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,17 @@ } _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"} + + def register_backends() -> None: """Register out-of-tree device types. Idempotent; safe without the card.""" @@ -92,6 +102,61 @@ 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).""" + # 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 + + +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/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/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/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/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/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/node_inventory.py b/yeto/launch/node_inventory.py index 6a31652ac..7c1d875a0 100644 --- a/yeto/launch/node_inventory.py +++ b/yeto/launch/node_inventory.py @@ -1,38 +1,59 @@ """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 -it), the fabric, shared mounts, the container image and an optional declared -price. The core stays neutral: provider differences live in the file. +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.engine_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 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 from ..gpu_spec import DEVICE_FAMILIES, UnknownCard, canonical_card, device_family_of from ..hw import card_tools +from ..hw.catalog import UnknownCardType, compat_group +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 @@ -41,10 +62,10 @@ 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") +_KNOWN = _REQUIRED + ("proxy_jump", "fabric", "interface", "address", "mounts", "image", "runtime", + "docker_args", "price_per_card_hour", "roles") class InventoryError(ValueError): @@ -58,6 +79,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 @@ -65,9 +87,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, ...] = () @@ -75,6 +99,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. @@ -92,11 +121,23 @@ 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, + "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} @@ -112,8 +153,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 +201,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,13 +221,29 @@ 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()): 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): @@ -164,10 +252,11 @@ 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, + cpu_arch=cpu_arch, card=card, cards_per_node=cards, hbm_gib=_require_positive(raw["hbm_gib"], name, "hbm_gib"), @@ -177,27 +266,39 @@ 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), + image=image, runtime=runtime, + docker_args=tuple(docker_args), 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]: """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/rl_island_task.py b/yeto/launch/rl_island_task.py index a4a2681ec..1185f4c9a 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") @@ -150,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/launch/ssh_trainer_island.py b/yeto/launch/ssh_trainer_island.py index 72ebd423d..56b3de31c 100644 --- a/yeto/launch/ssh_trainer_island.py +++ b/yeto/launch/ssh_trainer_island.py @@ -7,19 +7,31 @@ 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 +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; -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``), 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). 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 release of the template's scheme (W4A8 or MXFP4) in that scheme's + directory of the store, which samplers of that scheme 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 +39,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 @@ -52,7 +69,8 @@ from yeto.rl.engine.trainer_lifecycle import CutRef from yeto.rl.engine.trainer_remote import island_env -from .node_inventory import Node +from ..gpu_spec import DEVICE_FAMILY_ENV +from .node_inventory import InventoryError, Node from .topology import island_comm_env from .trainer_island import ( JOBS_FILE, @@ -69,6 +87,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 +99,72 @@ 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 + 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: 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 +176,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 +200,88 @@ 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)) + # The learner reads the family to pick its card type and collective backend (accel). + full_env = {DEVICE_FAMILY_ENV: node.device_family, **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 +293,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" ) @@ -230,8 +351,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: @@ -253,7 +382,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 +423,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 +491,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/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/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/adapters/miles/buffer_rollout_miles.py b/yeto/rl/adapters/miles/buffer_rollout_miles.py index a2ffb173f..e9f6bad87 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,25 @@ 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: # the one Miles tokenizer load yeto allows (test_provenance) + from .rollout_meta_hook import _load_miles_tokenizer + + 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: """Call the end-of-round hooks Miles' generating rollout calls, so the yeto diff --git a/yeto/rl/adapters/miles/config.py b/yeto/rl/adapters/miles/config.py index 409c802c0..b6ead7169 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 @@ -470,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")}, } @@ -714,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. @@ -723,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") @@ -771,8 +778,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 @@ -805,6 +818,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": @@ -840,6 +856,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 +1148,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 @@ -1144,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 @@ -1164,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: @@ -1175,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) @@ -1192,21 +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. + + 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``. - 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``. Full mode is - DP=1 only (``resolve_rl_run_config``), so no master is DP-sharded either way. + 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 266bc2ead..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) @@ -499,6 +519,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 +659,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 +899,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"]: @@ -984,6 +1055,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/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/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/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/models/glm5_next.py b/yeto/rl/adapters/miles/models/glm5_next.py new file mode 100644 index 000000000..5ab9b8c5e --- /dev/null +++ b/yeto/rl/adapters/miles/models/glm5_next.py @@ -0,0 +1,132 @@ +"""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 +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", str(HIDDEN_SIZE), + "--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", str(NUM_EXPERTS), + "--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", str(VOCAB_SIZE), + "--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/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/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/miles/run_flags.py b/yeto/rl/adapters/miles/run_flags.py index 9ec90e7f2..a6a9c6dad 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"], @@ -168,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/adapters/miles/trainer.py b/yeto/rl/adapters/miles/trainer.py index f8a62240c..ff548c45a 100644 --- a/yeto/rl/adapters/miles/trainer.py +++ b/yeto/rl/adapters/miles/trainer.py @@ -291,6 +291,11 @@ 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 + # ... 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. @@ -380,6 +385,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 +401,48 @@ 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 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( + 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 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 @@ -769,6 +818,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 +832,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 +848,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 +910,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 +941,68 @@ 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 + 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). 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 + + 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=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, 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]: """Load a cut written at ``plan.source`` DP into this freshly built trainer at ``plan.target`` DP (4.6). @@ -954,7 +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) + 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..5a101e812 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: @@ -222,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. @@ -276,6 +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 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 @@ -388,7 +403,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/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..fb406cd90 --- /dev/null +++ b/yeto/rl/adapters/sampler_engines/fake.py @@ -0,0 +1,157 @@ +"""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, + 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) + 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.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 + 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 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") + 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/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/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..1be20274e --- /dev/null +++ b/yeto/rl/adapters/sampler_engines/sglang_server.py @@ -0,0 +1,213 @@ +"""``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 + 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(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..f3ebcd339 --- /dev/null +++ b/yeto/rl/adapters/sampler_engines/vllm_server.py @@ -0,0 +1,160 @@ +"""``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]]: + 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(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}") + 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/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/cut.py b/yeto/rl/adapters/verl/cut.py new file mode 100644 index 000000000..5a3e2d2de --- /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=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}") + + 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/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 2b414f522..0d07f881e 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) @@ -119,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) @@ -131,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) @@ -145,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/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/adapters/verl/ports_impl.py b/yeto/rl/adapters/verl/ports_impl.py index dad0f2801..6bfaf3b0d 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 @@ -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, @@ -477,3 +511,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/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/adapters/verl/resume_check.py b/yeto/rl/adapters/verl/resume_check.py new file mode 100644 index 000000000..f4f008f05 --- /dev/null +++ b/yeto/rl/adapters/verl/resume_check.py @@ -0,0 +1,111 @@ +"""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 (: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: + +* 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 dataclasses import dataclass +from typing import Any + +from yeto.rl.engine.resume_manifest import ( + DEFAULT_SAMPLES, + OptimizerResumeState, + ResumeAssertionError, + ResumeManifest, + require_first_step, + require_restored, + 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.""" + return handles_rank_state(train_handles(worker), 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/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))] 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/engine/buffer_payload.py b/yeto/rl/engine/buffer_payload.py new file mode 100644 index 000000000..23bb1f8c8 --- /dev/null +++ b/yeto/rl/engine/buffer_payload.py @@ -0,0 +1,156 @@ +"""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 Callable, 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 + + +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.""" + + 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/cut.py b/yeto/rl/engine/cut.py index 983da061e..cfaeabd35 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 @@ -367,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..c4ecee2c4 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(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", data_cursor=dict(landed)) return diff --git a/yeto/rl/engine/engine_sampler.py b/yeto/rl/engine/engine_sampler.py new file mode 100644 index 000000000..7ef136bea --- /dev/null +++ b/yeto/rl/engine/engine_sampler.py @@ -0,0 +1,211 @@ +"""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()``. 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 + +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) + + +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/loop_guard.py b/yeto/rl/engine/loop_guard.py new file mode 100644 index 000000000..3c6228d80 --- /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 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/offpolicy_source.py b/yeto/rl/engine/offpolicy_source.py new file mode 100644 index 000000000..d0974e4b4 --- /dev/null +++ b/yeto/rl/engine/offpolicy_source.py @@ -0,0 +1,218 @@ +"""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} + + +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 diff --git a/yeto/rl/engine/policy_release.py b/yeto/rl/engine/policy_release.py index 4b65583d1..fbc081d9e 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.""" @@ -129,6 +148,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) @@ -169,8 +200,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 +331,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 +354,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/engine/resume_manifest.py b/yeto/rl/engine/resume_manifest.py new file mode 100644 index 000000000..18680922a --- /dev/null +++ b/yeto/rl/engine/resume_manifest.py @@ -0,0 +1,420 @@ +"""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, *, 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 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) + 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 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`` 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, 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}") + 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_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, optimizer_steps=optimizer_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) diff --git a/yeto/rl/engine/run_config.py b/yeto/rl/engine/run_config.py index 7e105a183..b2c82c4ea 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) @@ -237,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) @@ -533,6 +542,42 @@ 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, 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 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 + 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 ("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 + + 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 +694,18 @@ 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" - ) + 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, + 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} @@ -736,6 +775,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 +818,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 +936,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( @@ -1028,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/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/sampler_checks.py b/yeto/rl/engine/sampler_checks.py new file mode 100644 index 000000000..36d702748 --- /dev/null +++ b/yeto/rl/engine/sampler_checks.py @@ -0,0 +1,365 @@ +"""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`` 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 + +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 + 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)) + 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._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 + 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: + """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, + **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]}") + + 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._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: + 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}") + 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._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)} + 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 new file mode 100644 index 000000000..a476679e0 --- /dev/null +++ b/yeto/rl/engine/sampler_engine.py @@ -0,0 +1,286 @@ +"""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. + ``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): + 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).""" + ... + + +@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" + + +@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}") 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) diff --git a/yeto/rl/engine/telemetry.py b/yeto/rl/engine/telemetry.py index 8eedbb85d..743cfaa70 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.gpu_spec import DEVICE_FAMILY_ENV +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. 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_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]: 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) 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). diff --git a/yeto/rl/fullparam/memory.py b/yeto/rl/fullparam/memory.py index 675157f71..120ddfb8a 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,37 @@ 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; +* 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 + 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 +84,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 +111,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 +124,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 +135,19 @@ 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`` / + ``optimizer_offload_fraction`` / ``offload_master`` apply to + ``framework="megatron"`` only.""" cards: int cards_per_node: int @@ -126,14 +169,65 @@ 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 + 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: 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: 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 + 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.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: @@ -204,7 +298,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_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), @@ -222,6 +319,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 +366,66 @@ 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, 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) + 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") + if plan.distributed_optimizer: + dp_dense, dp_expert = W // (tp * pp), W // (ep * pp) + owned = dense / dp_dense + experts / dp_expert + owned_text = (f"(dense local / DP({dp_dense}) + expert local / expert DP({dp_expert})), " + "distributed optimizer") + else: + 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}) " + 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 _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) @@ -300,7 +459,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 +467,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 +484,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 +525,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: 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/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 new file mode 100644 index 000000000..f414fa3ae --- /dev/null +++ b/yeto/rl/harness/gateway/proxy.py @@ -0,0 +1,153 @@ +"""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) + + 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 + + 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/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 new file mode 100644 index 000000000..c9b9f95da --- /dev/null +++ b/yeto/rl/harness/gateway/records.py @@ -0,0 +1,342 @@ +"""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), + 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() + + +@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 + prompt_tokens = len(segment.token_ids) - len(completion.token_ids) + message = self._parse(completion) + return self._chat_response(session_id, message, completion, prompt_tokens=prompt_tokens) + + # -------------------------------------------------------------- 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()))} + + 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).""" + 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)) + 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") + 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=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}} 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)") diff --git a/yeto/rl/quant/__init__.py b/yeto/rl/quant/__init__.py index 147a247ec..5addfcd12 100644 --- a/yeto/rl/quant/__init__.py +++ b/yeto/rl/quant/__init__.py @@ -14,10 +14,25 @@ (``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. +* :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/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..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,22 +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 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), 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 0bfb02184..44fced0e2 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 @@ -16,11 +25,16 @@ 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). - -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. +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 +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 +49,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 +59,17 @@ 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" LAYOUT_SCHEMA = "yeto.rl.quant.layout/v1" SCALE_SUFFIX = ".weight_scale" WEIGHT_SUFFIX = ".weight" +PACKED_SUFFIX = ".weight_packed" _DTYPES = { "BF16": torch.bfloat16, @@ -76,6 +96,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_NAME + 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_NAME + 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,38 +225,60 @@ 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]]: + """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) @@ -186,41 +307,71 @@ 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} + + +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() @@ -243,6 +394,7 @@ def requantize_checkpoint( The source must provide exactly the layout's source tensors. """ layout = CheckpointLayout.load(template_dir) + check_smoothing(layout, smoothing) source = _SourceReader(source_dir) expected = layout.source_names() missing, extra = expected - set(source.weight_map), set(source.weight_map) - expected @@ -255,28 +407,23 @@ 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] - 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) + src = layout.source_of(name) + produced.update(transform_tensor(src, source.get(src), layout, smoothing)) + 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) finally: source.close() @@ -285,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) @@ -300,17 +463,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 +485,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 +504,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 +555,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) 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)) diff --git a/yeto/rl/stage1.py b/yeto/rl/stage1.py new file mode 100644 index 000000000..4825e9caf --- /dev/null +++ b/yeto/rl/stage1.py @@ -0,0 +1,964 @@ +"""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 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, 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). + +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 dataclasses +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 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.quant.schemes import release_dir_for + +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 +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 + + +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 + 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, + "service": self.service} + + +@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) + 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) + + 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 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, *, + 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``: 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.""" + 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, **(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), + "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], + sampling: Mapping[str, Any]) -> list[str]: + """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"]), + "--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"]))] + fraction = float(miles.get("optimizer_offload_fraction") or 0.0) + if fraction > 0: + 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 + + +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 = {DEVICE_FAMILY_ENV: trainers[0].device_family, "YETO_RUN_ID": plan.run_id, + "YETO_RL_STORE": store, "HF_HUB_OFFLINE": "1"} + trainer_env = {**common_env, "TORCHINDUCTOR_COMPILE_THREADS": "1"} + 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"]), + "--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, + "--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"]), + "--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"])), + *(("--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"]))] + 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, "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 [], + "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) + 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"], + "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 yeto.launch.ssh_trainer_island import requant_argv # noqa: PLC0415 + + # The release directory follows the template's scheme (quant.schemes); the + # plan's release_dir is checked against it, not passed to the re-quantizer. + argv = requant_argv(store, str(model["sampler_initial"]["local_dir"]), python="python3") + 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 + + 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 + from yeto.launch.ssh_trainer_island import container_name, remote_paths, start_script # noqa: PLC0415 + + 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")) + 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] + 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"]) + 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"]), + "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 + + +# ------------------------------------------------------------------ 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 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") + + # 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") + + # 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") + + # 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"]] + 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 + + +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_for(scheme.name): + errors.append(f"quant.release_dir {quant['release_dir']!r} is not {release_dir_for(scheme.name)!r} " + "(the re-quantizer writes to its scheme's directory)") + 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 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})") + 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)") + found.append("verl backend: no SSH launch path for the verl trainer island and no cut save path (N7.6 verl)") + 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): " + "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 + 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}): " + f"{len(plan.nodes_of('trainer'))} trainer nodes, DP {out.trainer['miles']['layout']['dp']}") + + +def _island_optimizer_gap(out: Expanded) -> str | None: + """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 + 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 + +def run_steps(out: Expanded, *, ssh: Sequence[str] = ("ssh", "-o", "BatchMode=yes"), + 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 + 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: + 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())