diff --git a/docker/miles-ports/Dockerfile b/docker/miles-ports/Dockerfile index 940711fd7..28ac792b8 100644 --- a/docker/miles-ports/Dockerfile +++ b/docker/miles-ports/Dockerfile @@ -19,12 +19,12 @@ FROM docker.io/radixark/miles@sha256:4d69750720eb1a4b99976a3fed45ed018e1fc4fff3d240667fa17d258dfc6e0c ARG MILES_REPOSITORY=https://github.com/michaellchung/miles -ARG MILES_COMMIT=44d41efd0a0428cb12b4ee34422de6f5e1ccb9d9 +ARG MILES_COMMIT=9462de532f747cf688e090bbaf6b24f80e378fe8 ARG SGLANG_REPOSITORY=https://github.com/michaellchung/sglang -ARG SGLANG_COMMIT=2fa880182eefdb31e64d8ea70317be27f9f32121 -# setuptools-scm version of SGLANG_COMMIT (describe v0.5.20-69-g4e4148f1b); +ARG SGLANG_COMMIT=754605964776322226a073763a6f636aa8233932 +# setuptools-scm version of SGLANG_COMMIT (describe v0.5.20-72-g754605964); # pinned because the one-commit clone has no tags to describe. -ARG SGLANG_VERSION=0.5.21.dev70+g2fa8801 +ARG SGLANG_VERSION=0.5.21.dev72+g7546059 # Miles fork -> /root/miles (same path the base installed editable). RUN rm -rf /root/miles \ diff --git a/docker/miles-ports/make_layer.py b/docker/miles-ports/make_layer.py index 31cdb726c..3801deeb6 100755 --- a/docker/miles-ports/make_layer.py +++ b/docker/miles-ports/make_layer.py @@ -173,7 +173,8 @@ def main() -> None: p.add_argument("--out", type=Path, required=True) a = p.parse_args() - check_pure_python(a.miles_src, a.miles_base_commit, a.miles_commit, ("miles/", "miles_plugins/")) + # tools/: stand-alone scripts run from the checkout (e.g. convert_hf_to_fp8.py), nothing built + check_pure_python(a.miles_src, a.miles_base_commit, a.miles_commit, ("miles/", "miles_plugins/", "tools/")) check_pure_python(a.sglang_src, a.sglang_base_commit, a.sglang_commit, ("python/sglang/",)) entries: dict[str, tuple[str, bytes | None, int, str | None]] = {} diff --git a/openspec/changes/rl-infra-spec/evidence/ports-image/2026-10-21-9462de5-10afe11/build-record.json b/openspec/changes/rl-infra-spec/evidence/ports-image/2026-10-21-9462de5-10afe11/build-record.json new file mode 100644 index 000000000..b85f0b0d3 --- /dev/null +++ b/openspec/changes/rl-infra-spec/evidence/ports-image/2026-10-21-9462de5-10afe11/build-record.json @@ -0,0 +1,12 @@ +{ + "tag": "ghcr.io/michaellchung/yeto-miles-ports:9462de5-10afe11", + "digest": "sha256:6e0a97d203cdd062c9f2fa35b6b834575cf05ee1f71fb45e921d8803b620f072", + "layer_tar_sha256": "08874741ace9733a576a2fa2b87736a6ebc46c8896806a3e025768b84e27e4ed", + "base": "docker.io/radixark/miles@sha256:4d69750720eb1a4b99976a3fed45ed018e1fc4fff3d240667fa17d258dfc6e0c", + "base_index": "sha256:90940828dcd4d54fd907ff668b43537cbd94778047580e4160d6560af548b74d", + "miles_commit": "9462de532f747cf688e090bbaf6b24f80e378fe8", + "sglang_commit": "10afe11a7ef2b53cedd9895cc4450474d85108b9", + "sglang_version": "0.5.21.dev71+g10afe11", + "peft_version": "0.20.0", + "peft_wheel_sha256": "0fbba16ffebfad3de96e06f2da6860fd860292324b85b6141909fa1e26ea9233" +} \ No newline at end of file diff --git a/openspec/changes/rl-infra-spec/evidence/ports-image/2026-10-21-9462de5-10afe11/image-manifest.json b/openspec/changes/rl-infra-spec/evidence/ports-image/2026-10-21-9462de5-10afe11/image-manifest.json new file mode 100644 index 000000000..9689afaf5 --- /dev/null +++ b/openspec/changes/rl-infra-spec/evidence/ports-image/2026-10-21-9462de5-10afe11/image-manifest.json @@ -0,0 +1,50 @@ +{ + "base": { + "index_digest": "sha256:90940828dcd4d54fd907ff668b43537cbd94778047580e4160d6560af548b74d", + "manifest_digest": "sha256:4d69750720eb1a4b99976a3fed45ed018e1fc4fff3d240667fa17d258dfc6e0c", + "miles_commit": "9e4260de047a704208535c0e90c531929879ab40", + "platform": "linux/amd64", + "repository": "docker.io/radixark/miles", + "sglang_commit": "e5bba1f9c433da5352e521a8bb826bab960e7836" + }, + "builder": "scripts/build_miles_ports_image.sh (crane append); docker/miles-ports/Dockerfile equivalent", + "image": "ghcr.io/michaellchung/yeto-miles-ports", + "megatron_bridge": { + "commit": "8cd3466d14d2337c8492827b3712482c2b3e4866", + "patches": [ + "lazy-modelopt (docker/miles-ports/megatron_bridge_init.py)" + ] + }, + "megatron_lm": { + "commit": "a84b105473eca9c9c6e49dbf1c8aa460c5cf796e", + "patches": [ + "lazy-modelopt (docker/miles-ports/megatron_lm_lazy_modelopt.patch)" + ], + "path": "/root/Megatron-LM", + "repository": "https://github.com/radixark/Megatron-LM" + }, + "miles": { + "commit": "9462de532f747cf688e090bbaf6b24f80e378fe8", + "install": "base's editable install (pip install -e /root/miles --no-deps) kept (same path); /root is on the base's sys.path, so without PYTHONPATH import miles is a namespace package over the /root/miles checkout -- yeto runs with PYTHONPATH=$HOME/miles (regular package)", + "path": "/root/miles", + "repository": "https://github.com/michaellchung/miles", + "upstream_commit": "9e4260de047a704208535c0e90c531929879ab40" + }, + "peft": { + "install": "wheel unpacked into the base's site-packages (pip install --no-deps equivalent)", + "version": "0.20.0", + "wheel": "peft-0.20.0-py3-none-any.whl", + "wheel_sha256": "0fbba16ffebfad3de96e06f2da6860fd860292324b85b6141909fa1e26ea9233" + }, + "pure_python_overlay": true, + "schema": 1, + "sglang": { + "commit": "10afe11a7ef2b53cedd9895cc4450474d85108b9", + "install": "pip install -e /sgl-workspace/sglang/python[all] --no-deps (base's editable install, same path)", + "path": "/sgl-workspace/sglang", + "repository": "https://github.com/michaellchung/sglang", + "upstream_commit": "571212b636baca45e10fa3b4da11a289123f3235", + "version": "0.5.21.dev71+g10afe11" + }, + "tag": "9462de5-10afe11" +} \ No newline at end of file diff --git a/openspec/changes/rl-infra-spec/evidence/ports-image/2026-10-21-9462de5-7546059/build-record.json b/openspec/changes/rl-infra-spec/evidence/ports-image/2026-10-21-9462de5-7546059/build-record.json new file mode 100644 index 000000000..2c4499f94 --- /dev/null +++ b/openspec/changes/rl-infra-spec/evidence/ports-image/2026-10-21-9462de5-7546059/build-record.json @@ -0,0 +1,12 @@ +{ + "tag": "ghcr.io/michaellchung/yeto-miles-ports:9462de5-7546059", + "digest": "sha256:37acaca3e1eee8adc61bf3edbbd706ece36e40e9b7fd5463e54f56e3f4cc9b2b", + "layer_tar_sha256": "4eee8887e4d85fe79fad9da706bbcef5b734d0deb74206427651a9fd8cf270b4", + "base": "docker.io/radixark/miles@sha256:4d69750720eb1a4b99976a3fed45ed018e1fc4fff3d240667fa17d258dfc6e0c", + "base_index": "sha256:90940828dcd4d54fd907ff668b43537cbd94778047580e4160d6560af548b74d", + "miles_commit": "9462de532f747cf688e090bbaf6b24f80e378fe8", + "sglang_commit": "754605964776322226a073763a6f636aa8233932", + "sglang_version": "0.5.21.dev72+g7546059", + "peft_version": "0.20.0", + "peft_wheel_sha256": "0fbba16ffebfad3de96e06f2da6860fd860292324b85b6141909fa1e26ea9233" +} \ No newline at end of file diff --git a/openspec/changes/rl-infra-spec/evidence/ports-image/2026-10-21-9462de5-7546059/image-manifest.json b/openspec/changes/rl-infra-spec/evidence/ports-image/2026-10-21-9462de5-7546059/image-manifest.json new file mode 100644 index 000000000..7a0edf600 --- /dev/null +++ b/openspec/changes/rl-infra-spec/evidence/ports-image/2026-10-21-9462de5-7546059/image-manifest.json @@ -0,0 +1,50 @@ +{ + "base": { + "index_digest": "sha256:90940828dcd4d54fd907ff668b43537cbd94778047580e4160d6560af548b74d", + "manifest_digest": "sha256:4d69750720eb1a4b99976a3fed45ed018e1fc4fff3d240667fa17d258dfc6e0c", + "miles_commit": "9e4260de047a704208535c0e90c531929879ab40", + "platform": "linux/amd64", + "repository": "docker.io/radixark/miles", + "sglang_commit": "e5bba1f9c433da5352e521a8bb826bab960e7836" + }, + "builder": "scripts/build_miles_ports_image.sh (crane append); docker/miles-ports/Dockerfile equivalent", + "image": "ghcr.io/michaellchung/yeto-miles-ports", + "megatron_bridge": { + "commit": "8cd3466d14d2337c8492827b3712482c2b3e4866", + "patches": [ + "lazy-modelopt (docker/miles-ports/megatron_bridge_init.py)" + ] + }, + "megatron_lm": { + "commit": "a84b105473eca9c9c6e49dbf1c8aa460c5cf796e", + "patches": [ + "lazy-modelopt (docker/miles-ports/megatron_lm_lazy_modelopt.patch)" + ], + "path": "/root/Megatron-LM", + "repository": "https://github.com/radixark/Megatron-LM" + }, + "miles": { + "commit": "9462de532f747cf688e090bbaf6b24f80e378fe8", + "install": "base's editable install (pip install -e /root/miles --no-deps) kept (same path); /root is on the base's sys.path, so without PYTHONPATH import miles is a namespace package over the /root/miles checkout -- yeto runs with PYTHONPATH=$HOME/miles (regular package)", + "path": "/root/miles", + "repository": "https://github.com/michaellchung/miles", + "upstream_commit": "9e4260de047a704208535c0e90c531929879ab40" + }, + "peft": { + "install": "wheel unpacked into the base's site-packages (pip install --no-deps equivalent)", + "version": "0.20.0", + "wheel": "peft-0.20.0-py3-none-any.whl", + "wheel_sha256": "0fbba16ffebfad3de96e06f2da6860fd860292324b85b6141909fa1e26ea9233" + }, + "pure_python_overlay": true, + "schema": 1, + "sglang": { + "commit": "754605964776322226a073763a6f636aa8233932", + "install": "pip install -e /sgl-workspace/sglang/python[all] --no-deps (base's editable install, same path)", + "path": "/sgl-workspace/sglang", + "repository": "https://github.com/michaellchung/sglang", + "upstream_commit": "571212b636baca45e10fa3b4da11a289123f3235", + "version": "0.5.21.dev72+g7546059" + }, + "tag": "9462de5-7546059" +} \ No newline at end of file diff --git a/tests/data/qwen3_4b_instruct_2507_weights.json b/tests/data/qwen3_4b_instruct_2507_weights.json new file mode 100644 index 000000000..1b3794d97 --- /dev/null +++ b/tests/data/qwen3_4b_instruct_2507_weights.json @@ -0,0 +1,410 @@ +{ +"source": "Qwen/Qwen3-4B-Instruct-2507@cdbee75f17c01a7cc42f958dc650907174af0554 model.safetensors.index.json", +"tie_word_embeddings": true, +"num_hidden_layers": 36, +"hidden_size": 2560, +"num_attention_heads": 32, +"num_key_value_heads": 8, +"head_dim": 128, +"intermediate_size": 9728, +"weight_names": [ +"model.embed_tokens.weight", +"model.layers.0.input_layernorm.weight", +"model.layers.0.mlp.down_proj.weight", +"model.layers.0.mlp.gate_proj.weight", +"model.layers.0.mlp.up_proj.weight", +"model.layers.0.post_attention_layernorm.weight", +"model.layers.0.self_attn.k_norm.weight", +"model.layers.0.self_attn.k_proj.weight", +"model.layers.0.self_attn.o_proj.weight", +"model.layers.0.self_attn.q_norm.weight", +"model.layers.0.self_attn.q_proj.weight", +"model.layers.0.self_attn.v_proj.weight", +"model.layers.1.input_layernorm.weight", +"model.layers.1.mlp.down_proj.weight", +"model.layers.1.mlp.gate_proj.weight", +"model.layers.1.mlp.up_proj.weight", +"model.layers.1.post_attention_layernorm.weight", +"model.layers.1.self_attn.k_norm.weight", +"model.layers.1.self_attn.k_proj.weight", +"model.layers.1.self_attn.o_proj.weight", +"model.layers.1.self_attn.q_norm.weight", +"model.layers.1.self_attn.q_proj.weight", +"model.layers.1.self_attn.v_proj.weight", +"model.layers.10.input_layernorm.weight", +"model.layers.10.mlp.down_proj.weight", +"model.layers.10.mlp.gate_proj.weight", +"model.layers.10.mlp.up_proj.weight", +"model.layers.10.post_attention_layernorm.weight", +"model.layers.10.self_attn.k_norm.weight", +"model.layers.10.self_attn.k_proj.weight", +"model.layers.10.self_attn.o_proj.weight", +"model.layers.10.self_attn.q_norm.weight", +"model.layers.10.self_attn.q_proj.weight", +"model.layers.10.self_attn.v_proj.weight", +"model.layers.11.input_layernorm.weight", +"model.layers.11.mlp.down_proj.weight", +"model.layers.11.mlp.gate_proj.weight", +"model.layers.11.mlp.up_proj.weight", +"model.layers.11.post_attention_layernorm.weight", +"model.layers.11.self_attn.k_norm.weight", +"model.layers.11.self_attn.k_proj.weight", +"model.layers.11.self_attn.o_proj.weight", +"model.layers.11.self_attn.q_norm.weight", +"model.layers.11.self_attn.q_proj.weight", +"model.layers.11.self_attn.v_proj.weight", +"model.layers.12.input_layernorm.weight", +"model.layers.12.mlp.down_proj.weight", +"model.layers.12.mlp.gate_proj.weight", +"model.layers.12.mlp.up_proj.weight", +"model.layers.12.post_attention_layernorm.weight", +"model.layers.12.self_attn.k_norm.weight", +"model.layers.12.self_attn.k_proj.weight", +"model.layers.12.self_attn.o_proj.weight", +"model.layers.12.self_attn.q_norm.weight", +"model.layers.12.self_attn.q_proj.weight", +"model.layers.12.self_attn.v_proj.weight", +"model.layers.13.input_layernorm.weight", +"model.layers.13.mlp.down_proj.weight", +"model.layers.13.mlp.gate_proj.weight", +"model.layers.13.mlp.up_proj.weight", +"model.layers.13.post_attention_layernorm.weight", +"model.layers.13.self_attn.k_norm.weight", +"model.layers.13.self_attn.k_proj.weight", +"model.layers.13.self_attn.o_proj.weight", +"model.layers.13.self_attn.q_norm.weight", +"model.layers.13.self_attn.q_proj.weight", +"model.layers.13.self_attn.v_proj.weight", +"model.layers.14.input_layernorm.weight", +"model.layers.14.mlp.down_proj.weight", +"model.layers.14.mlp.gate_proj.weight", +"model.layers.14.mlp.up_proj.weight", +"model.layers.14.post_attention_layernorm.weight", +"model.layers.14.self_attn.k_norm.weight", +"model.layers.14.self_attn.k_proj.weight", +"model.layers.14.self_attn.o_proj.weight", +"model.layers.14.self_attn.q_norm.weight", +"model.layers.14.self_attn.q_proj.weight", +"model.layers.14.self_attn.v_proj.weight", +"model.layers.15.input_layernorm.weight", +"model.layers.15.mlp.down_proj.weight", +"model.layers.15.mlp.gate_proj.weight", +"model.layers.15.mlp.up_proj.weight", +"model.layers.15.post_attention_layernorm.weight", +"model.layers.15.self_attn.k_norm.weight", +"model.layers.15.self_attn.k_proj.weight", +"model.layers.15.self_attn.o_proj.weight", +"model.layers.15.self_attn.q_norm.weight", +"model.layers.15.self_attn.q_proj.weight", +"model.layers.15.self_attn.v_proj.weight", +"model.layers.16.input_layernorm.weight", +"model.layers.16.mlp.down_proj.weight", +"model.layers.16.mlp.gate_proj.weight", +"model.layers.16.mlp.up_proj.weight", +"model.layers.16.post_attention_layernorm.weight", +"model.layers.16.self_attn.k_norm.weight", +"model.layers.16.self_attn.k_proj.weight", +"model.layers.16.self_attn.o_proj.weight", +"model.layers.16.self_attn.q_norm.weight", +"model.layers.16.self_attn.q_proj.weight", +"model.layers.16.self_attn.v_proj.weight", +"model.layers.17.input_layernorm.weight", +"model.layers.17.mlp.down_proj.weight", +"model.layers.17.mlp.gate_proj.weight", +"model.layers.17.mlp.up_proj.weight", +"model.layers.17.post_attention_layernorm.weight", +"model.layers.17.self_attn.k_norm.weight", +"model.layers.17.self_attn.k_proj.weight", +"model.layers.17.self_attn.o_proj.weight", +"model.layers.17.self_attn.q_norm.weight", +"model.layers.17.self_attn.q_proj.weight", +"model.layers.17.self_attn.v_proj.weight", +"model.layers.18.input_layernorm.weight", +"model.layers.18.mlp.down_proj.weight", +"model.layers.18.mlp.gate_proj.weight", +"model.layers.18.mlp.up_proj.weight", +"model.layers.18.post_attention_layernorm.weight", +"model.layers.18.self_attn.k_norm.weight", +"model.layers.18.self_attn.k_proj.weight", +"model.layers.18.self_attn.o_proj.weight", +"model.layers.18.self_attn.q_norm.weight", +"model.layers.18.self_attn.q_proj.weight", +"model.layers.18.self_attn.v_proj.weight", +"model.layers.19.input_layernorm.weight", +"model.layers.19.mlp.down_proj.weight", +"model.layers.19.mlp.gate_proj.weight", +"model.layers.19.mlp.up_proj.weight", +"model.layers.19.post_attention_layernorm.weight", +"model.layers.19.self_attn.k_norm.weight", +"model.layers.19.self_attn.k_proj.weight", +"model.layers.19.self_attn.o_proj.weight", +"model.layers.19.self_attn.q_norm.weight", +"model.layers.19.self_attn.q_proj.weight", +"model.layers.19.self_attn.v_proj.weight", +"model.layers.2.input_layernorm.weight", +"model.layers.2.mlp.down_proj.weight", +"model.layers.2.mlp.gate_proj.weight", +"model.layers.2.mlp.up_proj.weight", +"model.layers.2.post_attention_layernorm.weight", +"model.layers.2.self_attn.k_norm.weight", +"model.layers.2.self_attn.k_proj.weight", +"model.layers.2.self_attn.o_proj.weight", +"model.layers.2.self_attn.q_norm.weight", +"model.layers.2.self_attn.q_proj.weight", +"model.layers.2.self_attn.v_proj.weight", +"model.layers.20.input_layernorm.weight", +"model.layers.20.mlp.down_proj.weight", +"model.layers.20.mlp.gate_proj.weight", +"model.layers.20.mlp.up_proj.weight", +"model.layers.20.post_attention_layernorm.weight", +"model.layers.20.self_attn.k_norm.weight", +"model.layers.20.self_attn.k_proj.weight", +"model.layers.20.self_attn.o_proj.weight", +"model.layers.20.self_attn.q_norm.weight", +"model.layers.20.self_attn.q_proj.weight", +"model.layers.20.self_attn.v_proj.weight", +"model.layers.21.input_layernorm.weight", +"model.layers.21.mlp.down_proj.weight", +"model.layers.21.mlp.gate_proj.weight", +"model.layers.21.mlp.up_proj.weight", +"model.layers.21.post_attention_layernorm.weight", +"model.layers.21.self_attn.k_norm.weight", +"model.layers.21.self_attn.k_proj.weight", +"model.layers.21.self_attn.o_proj.weight", +"model.layers.21.self_attn.q_norm.weight", +"model.layers.21.self_attn.q_proj.weight", +"model.layers.21.self_attn.v_proj.weight", +"model.layers.22.input_layernorm.weight", +"model.layers.22.mlp.down_proj.weight", +"model.layers.22.mlp.gate_proj.weight", +"model.layers.22.mlp.up_proj.weight", +"model.layers.22.post_attention_layernorm.weight", +"model.layers.22.self_attn.k_norm.weight", +"model.layers.22.self_attn.k_proj.weight", +"model.layers.22.self_attn.o_proj.weight", +"model.layers.22.self_attn.q_norm.weight", +"model.layers.22.self_attn.q_proj.weight", +"model.layers.22.self_attn.v_proj.weight", +"model.layers.23.input_layernorm.weight", +"model.layers.23.mlp.down_proj.weight", +"model.layers.23.mlp.gate_proj.weight", +"model.layers.23.mlp.up_proj.weight", +"model.layers.23.post_attention_layernorm.weight", +"model.layers.23.self_attn.k_norm.weight", +"model.layers.23.self_attn.k_proj.weight", +"model.layers.23.self_attn.o_proj.weight", +"model.layers.23.self_attn.q_norm.weight", +"model.layers.23.self_attn.q_proj.weight", +"model.layers.23.self_attn.v_proj.weight", +"model.layers.24.input_layernorm.weight", +"model.layers.24.mlp.down_proj.weight", +"model.layers.24.mlp.gate_proj.weight", +"model.layers.24.mlp.up_proj.weight", +"model.layers.24.post_attention_layernorm.weight", +"model.layers.24.self_attn.k_norm.weight", +"model.layers.24.self_attn.k_proj.weight", +"model.layers.24.self_attn.o_proj.weight", +"model.layers.24.self_attn.q_norm.weight", +"model.layers.24.self_attn.q_proj.weight", +"model.layers.24.self_attn.v_proj.weight", +"model.layers.25.input_layernorm.weight", +"model.layers.25.mlp.down_proj.weight", +"model.layers.25.mlp.gate_proj.weight", +"model.layers.25.mlp.up_proj.weight", +"model.layers.25.post_attention_layernorm.weight", +"model.layers.25.self_attn.k_norm.weight", +"model.layers.25.self_attn.k_proj.weight", +"model.layers.25.self_attn.o_proj.weight", +"model.layers.25.self_attn.q_norm.weight", +"model.layers.25.self_attn.q_proj.weight", +"model.layers.25.self_attn.v_proj.weight", +"model.layers.26.input_layernorm.weight", +"model.layers.26.mlp.down_proj.weight", +"model.layers.26.mlp.gate_proj.weight", +"model.layers.26.mlp.up_proj.weight", +"model.layers.26.post_attention_layernorm.weight", +"model.layers.26.self_attn.k_norm.weight", +"model.layers.26.self_attn.k_proj.weight", +"model.layers.26.self_attn.o_proj.weight", +"model.layers.26.self_attn.q_norm.weight", +"model.layers.26.self_attn.q_proj.weight", +"model.layers.26.self_attn.v_proj.weight", +"model.layers.27.input_layernorm.weight", +"model.layers.27.mlp.down_proj.weight", +"model.layers.27.mlp.gate_proj.weight", +"model.layers.27.mlp.up_proj.weight", +"model.layers.27.post_attention_layernorm.weight", +"model.layers.27.self_attn.k_norm.weight", +"model.layers.27.self_attn.k_proj.weight", +"model.layers.27.self_attn.o_proj.weight", +"model.layers.27.self_attn.q_norm.weight", +"model.layers.27.self_attn.q_proj.weight", +"model.layers.27.self_attn.v_proj.weight", +"model.layers.28.input_layernorm.weight", +"model.layers.28.mlp.down_proj.weight", +"model.layers.28.mlp.gate_proj.weight", +"model.layers.28.mlp.up_proj.weight", +"model.layers.28.post_attention_layernorm.weight", +"model.layers.28.self_attn.k_norm.weight", +"model.layers.28.self_attn.k_proj.weight", +"model.layers.28.self_attn.o_proj.weight", +"model.layers.28.self_attn.q_norm.weight", +"model.layers.28.self_attn.q_proj.weight", +"model.layers.28.self_attn.v_proj.weight", +"model.layers.29.input_layernorm.weight", +"model.layers.29.mlp.down_proj.weight", +"model.layers.29.mlp.gate_proj.weight", +"model.layers.29.mlp.up_proj.weight", +"model.layers.29.post_attention_layernorm.weight", +"model.layers.29.self_attn.k_norm.weight", +"model.layers.29.self_attn.k_proj.weight", +"model.layers.29.self_attn.o_proj.weight", +"model.layers.29.self_attn.q_norm.weight", +"model.layers.29.self_attn.q_proj.weight", +"model.layers.29.self_attn.v_proj.weight", +"model.layers.3.input_layernorm.weight", +"model.layers.3.mlp.down_proj.weight", +"model.layers.3.mlp.gate_proj.weight", +"model.layers.3.mlp.up_proj.weight", +"model.layers.3.post_attention_layernorm.weight", +"model.layers.3.self_attn.k_norm.weight", +"model.layers.3.self_attn.k_proj.weight", +"model.layers.3.self_attn.o_proj.weight", +"model.layers.3.self_attn.q_norm.weight", +"model.layers.3.self_attn.q_proj.weight", +"model.layers.3.self_attn.v_proj.weight", +"model.layers.30.input_layernorm.weight", +"model.layers.30.mlp.down_proj.weight", +"model.layers.30.mlp.gate_proj.weight", +"model.layers.30.mlp.up_proj.weight", +"model.layers.30.post_attention_layernorm.weight", +"model.layers.30.self_attn.k_norm.weight", +"model.layers.30.self_attn.k_proj.weight", +"model.layers.30.self_attn.o_proj.weight", +"model.layers.30.self_attn.q_norm.weight", +"model.layers.30.self_attn.q_proj.weight", +"model.layers.30.self_attn.v_proj.weight", +"model.layers.31.input_layernorm.weight", +"model.layers.31.mlp.down_proj.weight", +"model.layers.31.mlp.gate_proj.weight", +"model.layers.31.mlp.up_proj.weight", +"model.layers.31.post_attention_layernorm.weight", +"model.layers.31.self_attn.k_norm.weight", +"model.layers.31.self_attn.k_proj.weight", +"model.layers.31.self_attn.o_proj.weight", +"model.layers.31.self_attn.q_norm.weight", +"model.layers.31.self_attn.q_proj.weight", +"model.layers.31.self_attn.v_proj.weight", +"model.layers.32.input_layernorm.weight", +"model.layers.32.mlp.down_proj.weight", +"model.layers.32.mlp.gate_proj.weight", +"model.layers.32.mlp.up_proj.weight", +"model.layers.32.post_attention_layernorm.weight", +"model.layers.32.self_attn.k_norm.weight", +"model.layers.32.self_attn.k_proj.weight", +"model.layers.32.self_attn.o_proj.weight", +"model.layers.32.self_attn.q_norm.weight", +"model.layers.32.self_attn.q_proj.weight", +"model.layers.32.self_attn.v_proj.weight", +"model.layers.33.input_layernorm.weight", +"model.layers.33.mlp.down_proj.weight", +"model.layers.33.mlp.gate_proj.weight", +"model.layers.33.mlp.up_proj.weight", +"model.layers.33.post_attention_layernorm.weight", +"model.layers.33.self_attn.k_norm.weight", +"model.layers.33.self_attn.k_proj.weight", +"model.layers.33.self_attn.o_proj.weight", +"model.layers.33.self_attn.q_norm.weight", +"model.layers.33.self_attn.q_proj.weight", +"model.layers.33.self_attn.v_proj.weight", +"model.layers.34.input_layernorm.weight", +"model.layers.34.mlp.down_proj.weight", +"model.layers.34.mlp.gate_proj.weight", +"model.layers.34.mlp.up_proj.weight", +"model.layers.34.post_attention_layernorm.weight", +"model.layers.34.self_attn.k_norm.weight", +"model.layers.34.self_attn.k_proj.weight", +"model.layers.34.self_attn.o_proj.weight", +"model.layers.34.self_attn.q_norm.weight", +"model.layers.34.self_attn.q_proj.weight", +"model.layers.34.self_attn.v_proj.weight", +"model.layers.35.input_layernorm.weight", +"model.layers.35.mlp.down_proj.weight", +"model.layers.35.mlp.gate_proj.weight", +"model.layers.35.mlp.up_proj.weight", +"model.layers.35.post_attention_layernorm.weight", +"model.layers.35.self_attn.k_norm.weight", +"model.layers.35.self_attn.k_proj.weight", +"model.layers.35.self_attn.o_proj.weight", +"model.layers.35.self_attn.q_norm.weight", +"model.layers.35.self_attn.q_proj.weight", +"model.layers.35.self_attn.v_proj.weight", +"model.layers.4.input_layernorm.weight", +"model.layers.4.mlp.down_proj.weight", +"model.layers.4.mlp.gate_proj.weight", +"model.layers.4.mlp.up_proj.weight", +"model.layers.4.post_attention_layernorm.weight", +"model.layers.4.self_attn.k_norm.weight", +"model.layers.4.self_attn.k_proj.weight", +"model.layers.4.self_attn.o_proj.weight", +"model.layers.4.self_attn.q_norm.weight", +"model.layers.4.self_attn.q_proj.weight", +"model.layers.4.self_attn.v_proj.weight", +"model.layers.5.input_layernorm.weight", +"model.layers.5.mlp.down_proj.weight", +"model.layers.5.mlp.gate_proj.weight", +"model.layers.5.mlp.up_proj.weight", +"model.layers.5.post_attention_layernorm.weight", +"model.layers.5.self_attn.k_norm.weight", +"model.layers.5.self_attn.k_proj.weight", +"model.layers.5.self_attn.o_proj.weight", +"model.layers.5.self_attn.q_norm.weight", +"model.layers.5.self_attn.q_proj.weight", +"model.layers.5.self_attn.v_proj.weight", +"model.layers.6.input_layernorm.weight", +"model.layers.6.mlp.down_proj.weight", +"model.layers.6.mlp.gate_proj.weight", +"model.layers.6.mlp.up_proj.weight", +"model.layers.6.post_attention_layernorm.weight", +"model.layers.6.self_attn.k_norm.weight", +"model.layers.6.self_attn.k_proj.weight", +"model.layers.6.self_attn.o_proj.weight", +"model.layers.6.self_attn.q_norm.weight", +"model.layers.6.self_attn.q_proj.weight", +"model.layers.6.self_attn.v_proj.weight", +"model.layers.7.input_layernorm.weight", +"model.layers.7.mlp.down_proj.weight", +"model.layers.7.mlp.gate_proj.weight", +"model.layers.7.mlp.up_proj.weight", +"model.layers.7.post_attention_layernorm.weight", +"model.layers.7.self_attn.k_norm.weight", +"model.layers.7.self_attn.k_proj.weight", +"model.layers.7.self_attn.o_proj.weight", +"model.layers.7.self_attn.q_norm.weight", +"model.layers.7.self_attn.q_proj.weight", +"model.layers.7.self_attn.v_proj.weight", +"model.layers.8.input_layernorm.weight", +"model.layers.8.mlp.down_proj.weight", +"model.layers.8.mlp.gate_proj.weight", +"model.layers.8.mlp.up_proj.weight", +"model.layers.8.post_attention_layernorm.weight", +"model.layers.8.self_attn.k_norm.weight", +"model.layers.8.self_attn.k_proj.weight", +"model.layers.8.self_attn.o_proj.weight", +"model.layers.8.self_attn.q_norm.weight", +"model.layers.8.self_attn.q_proj.weight", +"model.layers.8.self_attn.v_proj.weight", +"model.layers.9.input_layernorm.weight", +"model.layers.9.mlp.down_proj.weight", +"model.layers.9.mlp.gate_proj.weight", +"model.layers.9.mlp.up_proj.weight", +"model.layers.9.post_attention_layernorm.weight", +"model.layers.9.self_attn.k_norm.weight", +"model.layers.9.self_attn.k_proj.weight", +"model.layers.9.self_attn.o_proj.weight", +"model.layers.9.self_attn.q_norm.weight", +"model.layers.9.self_attn.q_proj.weight", +"model.layers.9.self_attn.v_proj.weight", +"model.norm.weight" +] +} \ No newline at end of file diff --git a/tests/test_modal_store_write_guard.py b/tests/test_modal_store_write_guard.py new file mode 100644 index 000000000..c79051562 --- /dev/null +++ b/tests/test_modal_store_write_guard.py @@ -0,0 +1,113 @@ +"""ModalVolumeStore: a reload never runs while a writer of the same Volume holds +``writing()``, in this process or another one of the container (S26 E6: rollout +actors' reloads dropped a half-written tier 2 release).""" + +from __future__ import annotations + +import subprocess +import sys +import threading + +import pytest + +from yeto.cloud.modal_ckpt_store import ModalVolumeStore +from yeto.rl.engine.policy_release import PolicyReleaseStore + + +class _Volume: + def __init__(self) -> None: + self.reloads = 0 + self.commits = 0 + + def reload(self) -> None: + self.reloads += 1 + + def commit(self) -> None: + self.commits += 1 + + +@pytest.fixture(autouse=True) +def _guard_dir(tmp_path, monkeypatch): + monkeypatch.setenv("YETO_MODAL_VOLUME_GUARD_DIR", str(tmp_path)) + + +def _store(tmp_path, volume: _Volume, name: str = "v") -> ModalVolumeStore: + return ModalVolumeStore(tmp_path, f"{name}-{id(volume)}", volume_factory=lambda _n: volume) + + +def test_reload_is_skipped_while_a_writer_of_the_same_volume_holds_the_guard(tmp_path): + volume = _Volume() + writer, reader = _store(tmp_path, volume), _store(tmp_path, volume) # two store objects, one Volume + with writer.writing(): + reader.reload() + assert volume.reloads == 0 + reader.reload() + assert volume.reloads == 1 + + +def test_a_release_written_while_another_thread_reloads_is_never_reloaded_over(tmp_path): + volume = _Volume() + store = _store(tmp_path, volume) + releases = PolicyReleaseStore(store) + inside = threading.Event() + seen: list[int] = [] + + def write(directory): + (directory / "model.safetensors").write_bytes(b"x") + inside.set() + before = volume.reloads + reloader = threading.Thread(target=_store(tmp_path, volume).reload) + reloader.start() + reloader.join(5) + seen.append(volume.reloads - before) + + record = releases.release(1, "h1", write) + assert inside.is_set() and seen == [0] + assert record is not None and (tmp_path / releases.release_dir / record.path / "model.safetensors").is_file() + + +def test_a_writer_waits_for_a_reload_in_flight(tmp_path): + started, release = threading.Event(), threading.Event() + + class SlowVolume(_Volume): + def reload(self) -> None: + started.set() + release.wait(5) + super().reload() + + volume = SlowVolume() + store = _store(tmp_path, volume) + reloader = threading.Thread(target=store.reload) + reloader.start() + started.wait(5) + entered = threading.Event() + + def writer(): + with store.writing(): + entered.set() + + t = threading.Thread(target=writer) + t.start() + assert not entered.wait(0.2) # blocked behind the reload + release.set() + reloader.join(5) + assert entered.wait(5) + t.join(5) + + +def test_a_writer_in_another_process_blocks_the_reload(tmp_path): + volume = _Volume() + store = _store(tmp_path, volume) + code = ("import os, sys, time; from yeto.cloud.modal_ckpt_store import ModalVolumeStore\n" + "s = ModalVolumeStore(sys.argv[1], sys.argv[2], volume_factory=lambda n: None)\n" + "with s.writing():\n print('held', flush=True); sys.stdin.readline()\n") + proc = subprocess.Popen([sys.executable, "-c", code, str(tmp_path), store.volume_name], + stdin=subprocess.PIPE, stdout=subprocess.PIPE, text=True) + try: + assert proc.stdout.readline().strip() == "held" + store.reload() + assert volume.reloads == 0 + finally: + proc.communicate("\n", timeout=30) + store.reload() + assert volume.reloads == 1 diff --git a/tests/test_rl_cut_distopt_pp_ep.py b/tests/test_rl_cut_distopt_pp_ep.py index d687e8d38..5a42c4ef8 100644 --- a/tests/test_rl_cut_distopt_pp_ep.py +++ b/tests/test_rl_cut_distopt_pp_ep.py @@ -31,6 +31,8 @@ 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"), + # S26 E6 fork branch: side-effect-free reads plus the ported CPU-offload state + "s26-e6-fp8": (os.environ.get("YETO_MILES_E6_CHECKOUT", "/home/michael/work/miles-s26-e6-fp8"), "ed5f64fd7"), } diff --git a/tests/test_rl_ports_image.py b/tests/test_rl_ports_image.py index 15585c2c3..38671095a 100644 --- a/tests/test_rl_ports_image.py +++ b/tests/test_rl_ports_image.py @@ -70,7 +70,7 @@ def test_build_inputs_match_the_pins(): def test_commits_and_digest_are_pinned_together_with_the_build_record(): """Commit pins, tag, Dockerfile ARGs and the image digest move together: the latest build record must name exactly the pinned commits and digest.""" - record_dir = REPO / "openspec/changes/rl-infra-spec/evidence/ports-image/2026-10-14-44d41ef-2fa8801" + record_dir = REPO / "openspec/changes/rl-infra-spec/evidence/ports-image/2026-10-21-9462de5-7546059" record = json.loads((record_dir / "build-record.json").read_text()) assert record["miles_commit"] == rl.MILES_NEXT_COMMIT assert record["sglang_commit"] == rl.SGLANG_NEXT_COMMIT diff --git a/tests/test_rl_quant_engine_sampler.py b/tests/test_rl_quant_engine_sampler.py new file mode 100644 index 000000000..f1e60c43c --- /dev/null +++ b/tests/test_rl_quant_engine_sampler.py @@ -0,0 +1,98 @@ +"""Rollout quant tiers on the engine sampler island (S26 E6): the spec written by the +launch layer, the settings filled on the node, the dense expert count.""" + +from __future__ import annotations + +from pathlib import Path + +import pytest + +from tests.test_rl_quant_tiers import write_checkpoint +from tests.test_rl_quant_scope import dense_shapes, moe_shapes +from yeto.rl.adapters.sampler_engines import quant_tier +from yeto.rl.adapters.sampler_engines.fake import FakeEngine +from yeto.rl.engine.engine_sampler_entry import SamplerSpecError, apply_rollout_quant +from yeto.rl.engine.sampler_checks import PostLoadChecks +from yeto.rl.quant import tiers +from yeto.rl.quant.schemes import FP8_FFN_CHANNEL, release_dir_for + + +def _spec(base: Path, tier: int, adapter: str = "sglang-server", **server) -> dict: + return {"engine": {"factory": "x:y", "adapter": adapter, + "server": {"base_dir": str(base), "quant_scheme": "bf16", "precision": "bf16", + "extra_args": [], "extra_env": {}, **server}}, + "base_dir": str(base), "release_dir": "policy-releases", "checks": {}, + "event_tape": str(base.parent / "tape" / "events.jsonl"), + **({"rollout_quant": s} if (s := quant_tier.spec_section(tier)) else {})} + + +@pytest.fixture +def dense(tmp_path): + return write_checkpoint(tmp_path / "dense", dense_shapes(), {"num_hidden_layers": 2}) + + +def test_spec_section_and_tier_0_is_unchanged(dense): + assert quant_tier.spec_section(0) is None + assert quant_tier.spec_section(1) == {"tier": 1, "apply": quant_tier.APPLY} + with pytest.raises(quant_tier.QuantTierSpecError): + quant_tier.spec_section(3) + spec = _spec(dense, 0) + assert apply_rollout_quant(spec) is spec + + +def test_tier_1_fills_sglang_flags_and_ignore_list(dense): + out = apply_rollout_quant(_spec(dense, 1)) + server = out["engine"]["server"] + assert server["extra_args"] == ["--quantization", "fp8"] + assert "model.layers.0.self_attn.q_proj" in server["extra_env"][tiers.SGLANG_IGNORE_ENV] + assert server["quant_scheme"] == "bf16" and server["precision"] == "fp8" + assert out["release_dir"] == "policy-releases" and out["rollout_quant"]["plan"]["quantized_modules"] == 6 + + +def test_tier_1_on_vllm_and_double_quantization_are_refused(dense): + with pytest.raises(SamplerSpecError, match="no ignore list"): + apply_rollout_quant(_spec(dense, 1, adapter="vllm-server")) + with pytest.raises(SamplerSpecError, match="already sets"): + apply_rollout_quant(_spec(dense, 1, extra_args=["--quantization", "w8a8_fp8"])) + with pytest.raises(SamplerSpecError, match="already sets"): + apply_rollout_quant(_spec(dense, 2, quant_scheme="mxfp4-g32-rtn")) + + +@pytest.mark.parametrize("kind", ["dense", "moe"]) +def test_tier_2_converts_the_base_and_sets_scheme_format_and_expert_count(tmp_path, kind): + shapes, config = (dense_shapes(), {"num_hidden_layers": 2}) if kind == "dense" else \ + (moe_shapes(), {"num_hidden_layers": 2, "num_experts": 2}) + base = write_checkpoint(tmp_path / kind, shapes, config) + out = apply_rollout_quant(_spec(base, 2, adapter="vllm-server")) + server = out["engine"]["server"] + converted = Path(server["base_dir"]) + assert converted == tmp_path / "tape" / "fp8-base" and out["base_dir"] == str(converted) + assert tiers.check_release_scope(converted) == 6 + assert (server["quant_scheme"], server["expected_format"]) == (FP8_FFN_CHANNEL, "float-quantized") + assert out["release_dir"] == release_dir_for(FP8_FFN_CHANNEL) + assert out["checks"]["expected_expert_matrices"] == (0 if kind == "dense" else 6) + assert out["rollout_quant"]["bf16_base"] == str(base) + + +def test_tier_2_refuses_a_contradicting_expert_count(dense): + spec = _spec(dense, 2) + spec["checks"] = {"expected_expert_matrices": 36288} + with pytest.raises(SamplerSpecError, match="36288"): + apply_rollout_quant(spec) + + +def test_dense_expert_count_zero_passes_without_engine_report(): + engine = FakeEngine(expert_matrices=None) + result = PostLoadChecks().expert_count(engine, 0) + assert result["passed"] and "dense" in result["detail"] + assert not PostLoadChecks().expert_count(engine, 6)["passed"] # MoE still needs a report + + +def test_stage1_writes_the_tier_into_the_sampler_spec(): + from yeto.rl import stage1 + + doc = {"rollout_quant_tier": 1} + plan = type("P", (), {"doc": doc})() + assert stage1.quant_spec_section(plan) == {"tier": 1, "apply": quant_tier.APPLY} + plan.doc = {} + assert stage1.quant_spec_section(plan) is None diff --git a/tests/test_rl_quant_fp8.py b/tests/test_rl_quant_fp8.py new file mode 100644 index 000000000..11c685420 --- /dev/null +++ b/tests/test_rl_quant_fp8.py @@ -0,0 +1,79 @@ +"""FP8 E4M3 per-channel quantization (S26 E6 tier 2): rule, error bound, reference.""" + +from __future__ import annotations + +import pytest +import torch + +from yeto.rl.quant.fp8 import FP8_DTYPE, FP8_MAX, dequantize_fp8_per_channel, quantize_fp8_per_channel + + +def _weight(rows: int = 16, cols: int = 64, seed: int = 0) -> torch.Tensor: + g = torch.Generator().manual_seed(seed) + w = torch.randn(rows, cols, generator=g) * torch.logspace(-3, 1, rows).unsqueeze(1) + return w.to(torch.bfloat16) + + +def test_scale_is_row_amax_over_448_and_values_follow_the_cast(): + w = _weight() + q, scale = quantize_fp8_per_channel(w) + assert q.dtype == FP8_DTYPE and scale.dtype == torch.float32 and scale.shape == (16, 1) + amax = w.float().abs().amax(dim=1, keepdim=True) + assert torch.equal(scale, amax / FP8_MAX) + expected = torch.clamp(w.float() / scale, -FP8_MAX, FP8_MAX).to(FP8_DTYPE) + assert torch.equal(q.view(torch.uint8), expected.view(torch.uint8)) + # every row reaches the FP8 maximum exactly once its amax is scaled + assert torch.equal(q.float().abs().amax(dim=1), torch.full((16,), FP8_MAX)) + + +def test_round_trip_error_is_within_half_an_e4m3_step(): + w = _weight(seed=3) + q, scale = quantize_fp8_per_channel(w) + back = dequantize_fp8_per_channel(q, scale) + x = w.float() / scale + # E4M3 has 3 mantissa bits: a normal value rounds within 2**-4 relative; below the + # smallest normal (2**-6) the step is 2**-9 absolute. + bound = torch.maximum(x.abs() * 2.0**-4, torch.full_like(x, 2.0**-10)) * scale + assert torch.all((back - w.float()).abs() <= bound + 1e-12) + + +def test_rows_are_independent_so_row_slices_match(): + w = _weight(seed=5) + q, scale = quantize_fp8_per_channel(w) + q2, scale2 = quantize_fp8_per_channel(w[4:9]) + assert torch.equal(q[4:9].view(torch.uint8), q2.view(torch.uint8)) and torch.equal(scale[4:9], scale2) + + +def test_zero_row_gets_epsilon_scale_and_zero_codes(): + w = _weight() + w[2] = 0 + q, scale = quantize_fp8_per_channel(w) + assert scale[2, 0].item() == torch.finfo(torch.float32).eps + assert torch.all(q[2].float() == 0) + + +def test_bad_input_is_refused(): + with pytest.raises(ValueError, match="2-D"): + quantize_fp8_per_channel(torch.zeros(4)) + with pytest.raises(TypeError, match="floating"): + quantize_fp8_per_channel(torch.zeros(2, 2, dtype=torch.int8)) + w = _weight() + w[0, 0] = float("nan") + with pytest.raises(ValueError, match="NaN"): + quantize_fp8_per_channel(w) + + +def test_equals_compressed_tensors_where_installed(): + ct = pytest.importorskip("compressed_tensors") + from compressed_tensors.quantization import QuantizationArgs + from compressed_tensors.quantization.lifecycle.forward import quantize + from compressed_tensors.quantization.utils import calculate_qparams + + args = QuantizationArgs(num_bits=8, type="float", symmetric=True, strategy="channel") + w = _weight(rows=32, cols=128, seed=7) + w[1] = 0 + ref_scale, zp = calculate_qparams(w.float().amin(dim=1, keepdim=True), w.float().amax(dim=1, keepdim=True), args) + ref_q = quantize(w.float(), ref_scale, zp, args, dtype=FP8_DTYPE) + q, scale = quantize_fp8_per_channel(w) + assert torch.equal(scale, ref_scale), ct.__version__ + assert torch.equal(q.view(torch.uint8), ref_q.view(torch.uint8)) diff --git a/tests/test_rl_quant_gate.py b/tests/test_rl_quant_gate.py new file mode 100644 index 000000000..48bf3cec8 --- /dev/null +++ b/tests/test_rl_quant_gate.py @@ -0,0 +1,142 @@ +"""Gate between rollout quantization tiers (S26 E6): statistics, verdict, collection, CLI.""" + +from __future__ import annotations + +import json +import math +from pathlib import Path + +import pytest + +from yeto.rl.engine.sampler_engine import FINISH_STOP, Completion, SamplingParams +from yeto.rl.quant import gate +from yeto.rl.quant.gate import GateCriteria, Pair + + +def test_statistics_on_hand_made_pairs(): + pairs = [Pair((-1.0, -2.0), (-1.0, -2.5)), Pair((-0.1,), (-0.1 + math.log(3.0),)), Pair((-1.0,), (-1.0, -1.0))] + stats = gate.compare(pairs, GateCriteria(tis_upper=2.0)) + assert stats["completions"] == 3 and stats["unaligned_completions"] == 1 and stats["tokens"] == 3 + diffs = [0.0, 0.5, -math.log(3.0)] + assert stats["mean_diff"] == pytest.approx(sum(diffs) / 3) + assert stats["mean_abs_diff"] == pytest.approx(sum(map(abs, diffs)) / 3) + assert stats["max_abs_diff"] == pytest.approx(math.log(3.0)) + assert stats["p50_abs_diff"] == pytest.approx(0.5) + # TIS weight exp(bf16 - quant): 1, exp(-0.5), 3 -> one of three above 2 + assert stats["tis_clip_high_share"] == pytest.approx(1 / 3) and stats["tis_clip_low_share"] == 0 + low = gate.compare(pairs, GateCriteria(tis_upper=2.0, tis_lower=0.7)) + assert low["tis_clip_low_share"] == pytest.approx(1 / 3) # exp(-0.5) = 0.61 < 0.7 + + +def test_quantile_interpolates(): + assert gate.quantile([3.0, 1.0, 2.0, 4.0], 0.5) == 2.5 + assert gate.quantile([1.0], 0.99) == 1.0 + assert math.isnan(gate.quantile([], 0.5)) + + +def test_verdict_applies_only_given_criteria_and_fails_on_missing_values(): + stats = {"tokens": 10, "mean_abs_diff": 0.02, "p99_abs_diff": 0.3, "tis_clip_share": 0.001, "unaligned_share": 0.0} + ok = gate.verdict(stats, None, GateCriteria(max_mean_abs_diff=0.05, max_tis_clip_share=0.01)) + assert ok["passed"] and set(ok["checks"]) == {"max_mean_abs_diff", "max_tis_clip_share", "max_unaligned_share"} + bad = gate.verdict(stats, None, GateCriteria(max_p99_abs_diff=0.1)) + assert not bad["passed"] and not bad["checks"]["max_p99_abs_diff"]["passed"] + no_speed = gate.verdict(stats, None, GateCriteria(min_speedup=1.1)) + assert not no_speed["passed"] + fast = gate.verdict(stats, {"speedup": 1.3}, GateCriteria(min_speedup=1.1)) + assert fast["passed"] + empty = gate.verdict({"tokens": 0, "mean_abs_diff": math.nan, "unaligned_share": math.nan}, None, + GateCriteria(max_mean_abs_diff=1.0)) + assert not empty["passed"] + + +def test_criteria_file_rejects_unknown_keys_and_bad_bounds(): + with pytest.raises(ValueError, match="unknown gate criteria"): + GateCriteria.from_mapping({"max_kl": 1}) + with pytest.raises(ValueError, match="TIS bounds"): + GateCriteria.from_mapping({"tis_upper": 0.9}) + + +class StubEngine: + """Generates fixed completions; scores them with a fixed shift (the BF16 engine).""" + + def __init__(self, shift: float, tokens: int = 4, tick: float = 1.0): + self.shift, self.tokens, self.tick, self.calls = shift, tokens, tick, [] + + def start(self, release_dir, version): + self.calls.append("start") + + def stop(self): + self.calls.append("stop") + + def generate(self, prompts, sampling): + self.calls.append("generate") + return [[Completion(tuple(range(self.tokens)), tuple(-0.5 for _ in range(self.tokens)), f"t{p}", + FINISH_STOP) for _ in range(sampling.n)] for p in prompts] + + def score(self, prompt, continuation): + return [-0.5 + self.shift] * self.tokens + + +def test_collect_runs_engines_one_after_the_other(): + now = iter([0.0, 2.0, 10.0, 14.0]) + quant, bf16 = StubEngine(0.0), StubEngine(-0.1) + pairs, speed, raw = gate.collect(quant, bf16, ["a", "b"], SamplingParams(max_tokens=8, n=2), clock=lambda: next(now)) + assert quant.calls == ["start", "generate", "stop"] and bf16.calls == ["start", "generate", "stop"] + assert len(pairs) == 4 and len(raw) == 4 and raw[0]["bf16_logprobs"] == [-0.6] * 4 + assert speed["quant"]["tokens_per_s"] == 16 / 2 and speed["bf16"]["tokens_per_s"] == 16 / 4 + assert speed["speedup"] == 2.0 + stats = gate.compare(pairs, GateCriteria()) + assert stats["mean_diff"] == pytest.approx(0.1) + + +def test_records_cli_writes_the_report_and_sets_the_exit_code(tmp_path: Path): + records = tmp_path / "r.jsonl" + records.write_text("\n".join(json.dumps({"rollout_log_probs": [-1.0, -2.0], "log_probs": [-1.01, -2.02]}) + for _ in range(3))) + criteria = tmp_path / "c.json" + criteria.write_text(json.dumps({"max_mean_abs_diff": 0.05, "max_tis_clip_share": 0.0})) + out = tmp_path / "gate" / "report.json" + assert gate.main(["records", "--records", str(records), "--criteria", str(criteria), "--out", str(out), + "--label", "tier0->tier1"]) == 0 + saved = json.loads(out.read_text()) + assert saved["schema"] == gate.SCHEMA and saved["verdict"]["passed"] and saved["meta"]["label"] == "tier0->tier1" + assert saved["stats"]["tokens"] == 6 + criteria.write_text(json.dumps({"max_mean_abs_diff": 0.001})) + assert gate.main(["records", "--records", str(records), "--criteria", str(criteria), "--out", str(out)]) == 1 + assert not json.loads(out.read_text())["verdict"]["passed"] + + +class TokenScoringStub(StubEngine): + """Text scoring drops the trailing EOS (one token short); token scoring does not.""" + + def score(self, prompt, continuation): + return [-0.5] * (self.tokens - 1) + + def score_tokens(self, prompt, token_ids): + self.calls.append(("score_tokens", tuple(token_ids))) + return [-0.5] * len(token_ids) + + +def test_collect_scores_by_token_ids_when_the_engine_can(): + quant, bf16 = StubEngine(0.0), TokenScoringStub(0.0) + pairs, _speed, _raw = gate.collect(quant, bf16, ["a"], SamplingParams(max_tokens=8, n=1), clock=lambda: 1.0) + assert ("score_tokens", (0, 1, 2, 3)) in bf16.calls + assert gate.compare(pairs, GateCriteria())["unaligned_completions"] == 0 + + +def test_sglang_score_tokens_sends_the_exact_ids(monkeypatch): + from yeto.rl.adapters.sampler_engines.sglang_server import SGLangServerEngine + + sent = [] + + def call(self, method, path, body): + sent.append(body) + if "text" in body: # the prompt: two tokens + return {"meta_info": {"input_token_logprobs": [[None, 7, None], [-0.1, 8, None]]}} + return {"meta_info": {"input_token_logprobs": [[None, 7, None], [-0.1, 8, None], + [-0.2, 5, None], [-0.3, 151645, None]]}} + + monkeypatch.setattr(SGLangServerEngine, "call", call) + engine = SGLangServerEngine.__new__(SGLangServerEngine) + assert engine.score_tokens("p", [5, 151645]) == [-0.2, -0.3] + assert sent[1]["input_ids"] == [7, 8, 5, 151645] diff --git a/tests/test_rl_quant_scope.py b/tests/test_rl_quant_scope.py new file mode 100644 index 000000000..0612bd89e --- /dev/null +++ b/tests/test_rl_quant_scope.py @@ -0,0 +1,168 @@ +"""Scope rule of quantized rollout (S26 E6): which linears may be quantized, and the +tier 1 engine ignore lists that follow from it.""" + +from __future__ import annotations + +import pytest + +from yeto.rl.quant import scope +from yeto.rl.quant.scope import DENSE, MOE, QuantScopeError + +H, I = 8, 16 + + +def dense_shapes(layers: int = 2) -> dict[str, list[int]]: + """Qwen3 dense names (Qwen3-4B layout, tied lm_head omitted like the HF checkpoint).""" + out = {"model.embed_tokens.weight": [32, H], "model.norm.weight": [H]} + for i in range(layers): + p = f"model.layers.{i}" + out.update({f"{p}.input_layernorm.weight": [H], f"{p}.post_attention_layernorm.weight": [H], + f"{p}.self_attn.q_proj.weight": [H, H], f"{p}.self_attn.k_proj.weight": [H, H], + f"{p}.self_attn.v_proj.weight": [H, H], f"{p}.self_attn.o_proj.weight": [H, H], + f"{p}.self_attn.q_norm.weight": [H], f"{p}.self_attn.k_norm.weight": [H], + f"{p}.mlp.gate_proj.weight": [I, H], f"{p}.mlp.up_proj.weight": [I, H], + f"{p}.mlp.down_proj.weight": [H, I]}) + return out + + +def moe_shapes(experts: int = 2) -> dict[str, list[int]]: + """Qwen3-MoE / Qwen2-MoE names: router, routed experts, shared expert + its gate, + one dense FFN layer (layer 0, as DeepSeek's first dense layers) and lm_head.""" + out = {"model.embed_tokens.weight": [32, H], "model.norm.weight": [H], "lm_head.weight": [32, H]} + out.update({"model.layers.0.self_attn.q_proj.weight": [H, H], "model.layers.0.self_attn.o_proj.weight": [H, H], + "model.layers.0.mlp.gate_proj.weight": [I, H], "model.layers.0.mlp.up_proj.weight": [I, H], + "model.layers.0.mlp.down_proj.weight": [H, I]}) + p = "model.layers.1" + out.update({f"{p}.self_attn.q_proj.weight": [H, H], f"{p}.self_attn.o_proj.weight": [H, H], + f"{p}.mlp.gate.weight": [experts, H], f"{p}.mlp.shared_expert_gate.weight": [1, H], + f"{p}.mlp.shared_expert.gate_proj.weight": [I, H], f"{p}.mlp.shared_expert.up_proj.weight": [I, H], + f"{p}.mlp.shared_expert.down_proj.weight": [H, I]}) + for e in range(experts): + out.update({f"{p}.mlp.experts.{e}.gate_proj.weight": [I, H], f"{p}.mlp.experts.{e}.up_proj.weight": [I, H], + f"{p}.mlp.experts.{e}.down_proj.weight": [H, I]}) + return out + + +def test_model_kind_reads_the_expert_count(): + assert scope.model_kind({"num_hidden_layers": 36}) == DENSE + assert scope.model_kind({"num_experts": 128}) == MOE + assert scope.model_kind({"n_routed_experts": 256}) == MOE + assert scope.model_kind({"text_config": {"num_local_experts": 8}}) == MOE + assert scope.model_kind({"num_experts": 0}) == DENSE + + +def test_dense_scope_is_the_ffn_only(): + q = [m for m in scope.linear_modules(dense_shapes()) if scope.in_scope(m, DENSE)] + assert q == sorted(f"model.layers.{i}.mlp.{p}" for i in range(2) for p in ("down_proj", "gate_proj", "up_proj")) + assert scope.in_scope("model.layers.0.mlp.gate_up_proj", DENSE) # fused engine name + for name in ("model.layers.0.self_attn.q_proj", "lm_head", "model.embed_tokens", "model.norm", + "model.layers.0.mlp.shared_expert.up_proj", "visual.blocks.0.mlp.fc1"): + assert not scope.in_scope(name, DENSE), name + + +def test_moe_scope_is_the_routed_experts_only(): + q = [m for m in scope.linear_modules(moe_shapes()) if scope.in_scope(m, MOE)] + assert q == sorted(f"model.layers.1.mlp.experts.{e}.{p}" for e in range(2) + for p in ("down_proj", "gate_proj", "up_proj")) + assert scope.in_scope("model.layers.3.mlp.experts.gate_up_proj", MOE) # stacked trainer layout + assert scope.in_scope("model.language_model.layers.3.mlp.experts.7.w2", MOE) + reasons = scope.violations(scope.linear_modules(moe_shapes()), MOE) + assert reasons["model.layers.1.mlp.gate"] == "router" + assert reasons["model.layers.1.mlp.shared_expert_gate"] == "shared_expert" + assert reasons["model.layers.1.mlp.shared_expert.up_proj"] == "shared_expert" + assert reasons["model.layers.0.mlp.up_proj"] == "not_routed_expert" # dense layer of a MoE model + assert reasons["model.layers.0.self_attn.q_proj"] == "attention" + assert reasons["lm_head"] == "lm_head" + + +def test_check_quantized_refuses_out_of_scope_modules(): + scope.check_quantized(["model.layers.0.mlp.up_proj"], DENSE, what="t") + with pytest.raises(QuantScopeError, match=r"1 module\(s\) outside the dense scope.*self_attn.o_proj \(attention\)"): + scope.check_quantized(["model.layers.0.mlp.up_proj", "model.layers.0.self_attn.o_proj"], DENSE, what="t") + with pytest.raises(QuantScopeError, match="router"): + scope.check_quantized(["model.layers.1.mlp.gate"], MOE, what="t") + with pytest.raises(ValueError, match="unknown model kind"): + scope.in_scope("x", "hybrid") + + +@pytest.mark.parametrize("kind,shapes", [(DENSE, dense_shapes()), (MOE, moe_shapes())]) +def test_engine_ignore_list_quantizes_exactly_the_scope(kind, shapes): + ignore = scope.engine_ignore_list(shapes, kind) + quantized = scope.check_online_plan(shapes, kind, ignore) + assert quantized == [m for m in scope.linear_modules(shapes) if scope.in_scope(m, kind)] + # the path segments alone must never cover an in-scope module (no overmatch) + segments = sorted(set(ignore) - set(scope.ignored_modules(shapes, kind))) + assert all(not scope.path_match(s, m) for s in segments for m in quantized) + + +def test_segments_also_cover_renamed_prefixes(): + """A module whose engine prefix differs from the checkpoint name (composite models) + is still kept BF16 by the segment entries.""" + ignore = scope.engine_ignore_list(dense_shapes(), DENSE) + renamed = ["language_model.model.layers.0.self_attn.qkv_proj", "language_model.model.layers.0.mlp.gate_up_proj"] + assert scope.online_quantized(renamed, ignore) == ["language_model.model.layers.0.mlp.gate_up_proj"] + + +def test_online_plan_refuses_an_ignore_list_that_misses_a_module(): + shapes = dense_shapes() + with pytest.raises(QuantScopeError, match="attention"): + scope.check_online_plan(shapes, DENSE, ["lm_head"]) + with pytest.raises(QuantScopeError, match="quantizes nothing"): + scope.check_online_plan(shapes, DENSE, scope.linear_modules(shapes)) + + +def test_path_match_follows_sglang_segments(): + assert scope.path_match("self_attn", "model.layers.0.self_attn.qkv_proj") + assert scope.path_match("mlp.gate", "model.layers.3.mlp.gate") + assert not scope.path_match("mlp.gate", "model.layers.3.mlp.gate_proj") + assert not scope.path_match("mlp.gate", "model.layers.3.mlp.gate_up_proj") + assert scope.path_match("model.layers.0", "model.layers.0.mlp.up_proj") + assert not scope.path_match("model.layers.1", "model.layers.10.mlp.up_proj") + + +def test_path_match_equals_sglang_where_installed(): + utils = pytest.importorskip("sglang.srt.layers.quantization.utils") + cases = [("self_attn", "model.layers.0.self_attn.qkv_proj"), ("mlp.gate", "model.layers.3.mlp.gate_proj"), + ("mlp.gate", "model.layers.3.mlp.gate"), ("model.layers.1", "model.layers.10.mlp.up_proj"), + ("mlp.shared_expert", "model.layers.1.mlp.shared_expert_gate")] + for ignored, prefix in cases: + assert scope.path_match(ignored, prefix) == utils._module_path_match(ignored, prefix), (ignored, prefix) + + +@pytest.mark.parametrize("kind,shapes,engine_modules", [ + (DENSE, dense_shapes(), {"model.layers.0.self_attn.qkv_proj": True, "model.layers.0.self_attn.o_proj": True, + "model.layers.0.mlp.gate_up_proj": False, "model.layers.0.mlp.down_proj": False}), + (MOE, moe_shapes(), {"model.layers.1.self_attn.qkv_proj": True, "model.layers.1.mlp.gate": True, + "model.layers.1.mlp.shared_expert.gate_up_proj": True, + "model.layers.1.mlp.shared_expert_gate": True, "model.layers.0.mlp.gate_up_proj": True, + "model.layers.0.mlp.down_proj": True, "model.layers.1.mlp.experts": False}), +]) +def test_sglang_is_layer_skipped_keeps_out_of_scope_engine_modules(kind, shapes, engine_modules): + """SGLang's own skip rule (fused qkv_proj / gate_up_proj shards, FusedMoE prefix) on the + tier 1 ignore list: True = stays BF16.""" + utils = pytest.importorskip("sglang.srt.layers.quantization.utils") + ignore = scope.engine_ignore_list(shapes, kind) + got = {prefix: bool(utils.is_layer_skipped(prefix, ignore)) for prefix in engine_modules} + assert got == engine_modules + + +def test_with_fused_aliases_adds_only_complete_groups(): + got = scope.with_fused_aliases(["a.self_attn.q_proj", "a.self_attn.k_proj", "a.self_attn.v_proj", + "b.mlp.gate_proj", "lm_head"]) + assert "a.self_attn.qkv_proj" in got and "b.mlp.gate_up_proj" not in got and "lm_head" in got + + +@pytest.mark.parametrize("kind,shapes,engine_modules", [ + (DENSE, dense_shapes(), {"model.layers.0.self_attn.qkv_proj": True, "model.layers.0.self_attn.o_proj": True, + "model.layers.0.mlp.gate_up_proj": False, "model.layers.0.mlp.down_proj": False}), + (MOE, moe_shapes(), {"model.layers.1.self_attn.o_proj": True, "model.layers.1.mlp.gate": True, + "model.layers.1.mlp.shared_expert.gate_up_proj": True, + "model.layers.0.mlp.gate_up_proj": True}), +]) +def test_sglang_compressed_tensors_ignore_without_packed_mapping(kind, shapes, engine_modules): + """Tier 2 checkpoint ignore list under SGLang's compressed-tensors rule with an empty + ``packed_modules_mapping`` (Qwen3ForCausalLM sets none; e6-gate-t2 on the card).""" + ct_utils = pytest.importorskip("sglang.srt.layers.quantization.compressed_tensors.utils") + ignore = scope.with_fused_aliases([*scope.ignored_modules(shapes, kind), "lm_head"]) + got = {p: bool(ct_utils.should_ignore_layer(p, ignore=ignore, fused_mapping={})) for p in engine_modules} + assert got == engine_modules diff --git a/tests/test_rl_quant_tiers.py b/tests/test_rl_quant_tiers.py new file mode 100644 index 000000000..e53f64f47 --- /dev/null +++ b/tests/test_rl_quant_tiers.py @@ -0,0 +1,561 @@ +"""Rollout quantization tiers (S26 E6): launch refusals, tier 1 engine plans, tier 2 +scoped FP8 layout, conversion, release scope check and the releaser.""" + +from __future__ import annotations + +import json +import os +from pathlib import Path +from types import SimpleNamespace + +import pytest +import torch +from safetensors import safe_open +from safetensors.torch import save_file + +from tests.test_rl_quant_scope import H, dense_shapes, moe_shapes +from yeto.rl.engine.policy_release import PolicyReleaseStore +from yeto.rl.engine.resume import CheckpointStore +from yeto.rl.quant import tiers +from yeto.rl.quant.awq import AwqSmoothing +from yeto.rl.quant.fp8 import quantize_fp8_per_channel +from yeto.rl.quant.release import scoped_releaser +from yeto.rl.quant.requant import ( + CheckpointLayout, + LayoutError, + requantize_checkpoint, + save_scoped_template, + scoped_layout, +) +from yeto.rl.quant.schemes import FP8_FFN_CHANNEL, release_dir_for +from yeto.rl.quant.scope import QuantScopeError +from yeto.rl.quant.tiers import QuantTierError + + +def write_checkpoint(directory: Path, shapes: dict[str, list[int]], config: dict, seed: int = 0) -> Path: + """Two shards (layer 0 + embeddings, the rest), an index and ``config.json``.""" + g = torch.Generator().manual_seed(seed) + directory.mkdir(parents=True, exist_ok=True) + shards: dict[str, dict[str, torch.Tensor]] = {"model-00001-of-00002.safetensors": {}, + "model-00002-of-00002.safetensors": {}} + for name, shape in shapes.items(): + shard = "model-00001-of-00002.safetensors" if ".layers.1." not in name else "model-00002-of-00002.safetensors" + shards[shard][name] = torch.randn(*shape, generator=g).to(torch.bfloat16) + weight_map = {} + for shard, tensors in shards.items(): + save_file(tensors, str(directory / shard), metadata={"format": "pt"}) + weight_map.update({n: shard for n in tensors}) + (directory / "model.safetensors.index.json").write_text(json.dumps({"metadata": {}, "weight_map": weight_map})) + (directory / "config.json").write_text(json.dumps(config)) + (directory / "tokenizer.json").write_text("{}") + return directory + + +@pytest.fixture +def dense(tmp_path: Path) -> Path: + return write_checkpoint(tmp_path / "dense", dense_shapes(), {"num_hidden_layers": 2, "tie_word_embeddings": True}) + + +@pytest.fixture +def moe(tmp_path: Path) -> Path: + return write_checkpoint(tmp_path / "moe", moe_shapes(), {"num_hidden_layers": 2, "num_experts": 2}) + + +def _read(directory: Path) -> dict[str, torch.Tensor]: + out = {} + for shard in sorted(directory.glob("*.safetensors")): + with safe_open(str(shard), framework="pt") as f: + out.update({k: f.get_tensor(k) for k in f.keys()}) + return out + + +# ---------------------------------------------------------------- launch refusals +def _args(**kw): + base = {"rl_rollout_quant_tier": 1, "gpu": "modal:1xh100", "tuning": "full", "rl_island_sample_role": None} + return SimpleNamespace(**{**base, **kw}) + + +def test_tier_0_and_1_pass_and_tier_3_is_refused(): + assert tiers.validate_launch(_args(rl_rollout_quant_tier=0, gpu="aws:1xa100")) == 0 + assert tiers.validate_launch(_args()) == 1 + assert tiers.validate_launch(_args(gpu="aws:1xl4@us-east-1,nebius:2x8xh200", tuning="lora")) == 1 + assert tiers.validate_launch(_args(gpu=None)) == 1 # auto fleet: the planner picks the card + with pytest.raises(QuantTierError, match="interface only"): + tiers.validate_launch(_args(rl_rollout_quant_tier=3)) + with pytest.raises(QuantTierError, match="expected one of"): + tiers.validate_launch(_args(rl_rollout_quant_tier=4)) + + +@pytest.mark.parametrize("gpu", ["aws:1xa100", "modal:1xa10g", "modal:1xh100,aws:2xa100-80gb", "ssh:1x8x910b4"]) +def test_cards_without_fp8_are_refused(gpu): + with pytest.raises(QuantTierError, match="FP8 tensor cores"): + tiers.validate_launch(_args(gpu=gpu)) + + +def test_tier_2_needs_sampler_islands_and_full_parameters(): + with pytest.raises(QuantTierError, match="--rl-island-sample-role"): + tiers.validate_launch(_args(rl_rollout_quant_tier=2)) + with pytest.raises(QuantTierError, match="--tuning full"): + tiers.validate_launch(_args(rl_rollout_quant_tier=2, rl_island_sample_role=["1:sampler"], tuning="lora")) + assert tiers.validate_launch(_args(rl_rollout_quant_tier=2, rl_island_sample_role=["1:sampler"])) == 2 + for role in ("sampler", "trainer"): # one launch per role on a shared store (S25 C3 shape) + assert tiers.validate_launch(_args(rl_rollout_quant_tier=2, rl_sample_role=role)) == 2 + with pytest.raises(QuantTierError, match="store-coupled"): + tiers.validate_launch(_args(rl_rollout_quant_tier=2, rl_sample_role="both")) + assert tiers.release_scheme(2) == FP8_FFN_CHANNEL and tiers.release_scheme(1) == "bf16" + + +def test_launcher_runs_the_tier_check_before_anything_else(): + from yeto.launch.rl_args import _prepare_rl_args + + args = SimpleNamespace(training_mode="rl", rl_rollout_quant_tier=3, rl_image=None, gpu="modal:1xh100") + with pytest.raises(QuantTierError, match="interface only"): + _prepare_rl_args(args) + + +def test_flag_is_one_shared_definition_for_cli_and_island(): + from yeto.rl.engine.run_flags import CLI, ISLAND, rl_flags + + (flag,) = [f for f in rl_flags() if f.options == (tiers.FLAG,)] + assert flag.layers == frozenset({CLI, ISLAND}) and not flag.island + assert flag.kwargs["choices"] == tiers.TIERS and flag.kwargs["default"] == 0 + + +# ---------------------------------------------------------------- tier 1 plans +def test_online_plan_on_a_dense_checkpoint(dense): + plan = tiers.online_plan(dense) + assert plan.kind == "dense" + assert plan.quantized == tuple(f"model.layers.{i}.mlp.{p}" for i in range(2) + for p in ("down_proj", "gate_proj", "up_proj")) + assert "model.layers.0.self_attn.q_proj" in plan.explicit and "self_attn" in plan.segments + env = tiers.sglang_online_env(plan) + assert set(env) == {"SGLANG_FP8_IGNORED_LAYERS"} + entries = env["SGLANG_FP8_IGNORED_LAYERS"].split(",") + assert "model.layers.1.self_attn.o_proj" in entries and not any("mlp.up_proj" in e for e in entries) + assert tiers.sglang_online_argv() == ["--sglang-quantization", "fp8"] + + +def test_vllm_overrides_carry_verl_fp8_config_with_our_ignore_list(dense): + plan = tiers.online_plan(dense) + overrides = tiers.vllm_online_overrides(plan, tiers.num_hidden_layers(dense)) + assert overrides[0] == "actor_rollout_ref.rollout.quantization=fp8" + assert not any("kv_cache_dtype" in o for o in overrides) # KV cache stays BF16 (gate, 10-21) + key, value = overrides[1].split("=", 1) + assert key == "++actor_rollout_ref.rollout.engine_kwargs.vllm.hf_overrides" + assert value.startswith("'") and value.endswith("'") + config = json.loads(value[1:-1])["quantization_config"] + assert {k: config[k] for k in tiers.VLLM_FP8_CONFIG} == tiers.VLLM_FP8_CONFIG + ignored = set(config["ignored_layers"]) + assert {"model.layers.0.self_attn.q_proj", "model.layers.1.self_attn.v_proj", "model.layers.0.mlp.gate"} <= ignored + assert not {m for m in ignored if m.endswith(("gate_proj", "up_proj", "down_proj"))} + + +def test_online_plan_refuses_a_model_with_an_unknown_linear(tmp_path): + shapes = {**dense_shapes(), "visual.blocks.0.mlp.fc1.weight": [16, 8]} + ckpt = write_checkpoint(tmp_path / "vl", shapes, {"num_hidden_layers": 2}) + plan = tiers.online_plan(ckpt) # the explicit list keeps the vision linear BF16 + assert "visual.blocks.0.mlp.fc1" in plan.explicit and "visual.blocks.0.mlp.fc1" not in plan.quantized + + +def test_miles_island_env_and_argv_for_tier_1(dense): + from yeto.rl.adapters.miles.entry import rollout_quant_env + + assert rollout_quant_env(None) == {} + assert rollout_quant_env(SimpleNamespace(yeto_rollout_quant_tier=0, hf_checkpoint=str(dense))) == {} + env = rollout_quant_env(SimpleNamespace(yeto_rollout_quant_tier=1, hf_checkpoint=str(dense))) + assert "model.layers.0.self_attn.k_proj" in env["SGLANG_FP8_IGNORED_LAYERS"] + + +def test_miles_connect_island_ray_forwards_the_ignore_list(dense): + from yeto.rl.adapters.miles.entry import connect_island_ray + + seen = {} + + class FakeRay: + @staticmethod + def is_initialized(): + return False + + @staticmethod + def init(address, runtime_env): + seen.update(address=address, env=runtime_env["env_vars"]) + + args = SimpleNamespace(yeto_rollout_quant_tier=1, hf_checkpoint=str(dense)) + connect_island_ray(environ={"RAY_ADDRESS": "10.0.0.1:6379"}, ray_module=FakeRay, miles_args=args) + assert "model.layers.1.self_attn.q_proj" in seen["env"]["SGLANG_FP8_IGNORED_LAYERS"] + environ: dict[str, str] = {} + assert connect_island_ray(environ=environ, ray_module=FakeRay, miles_args=args) is None + assert "SGLANG_FP8_IGNORED_LAYERS" in environ + + +def test_verl_tier_1_is_refused_on_a_colocated_island_and_at_launch(dense): + from yeto.rl.adapters.verl.island_entry import rollout_quant_overrides + + assert rollout_quant_overrides(0, dense, colocated=True) == ([], {"tier": 0, "name": "bf16"}) + assert rollout_quant_overrides(1, dense, colocated=False)[0] == [] + with pytest.raises(SystemExit, match="not available on verl"): + rollout_quant_overrides(1, dense, colocated=True) + with pytest.raises(tiers.QuantTierError, match="not available on verl"): + tiers.validate_launch(_args(rl_backend="verl")) + assert tiers.validate_launch(_args(rl_backend="miles")) == 1 + + +# ---------------------------------------------------------------- tier 2 layout and conversion +def test_scoped_layout_quantizes_only_the_dense_ffn(dense, tmp_path): + layout = save_scoped_template(dense, tmp_path / "t", FP8_FFN_CHANNEL) + assert layout.scheme.name == FP8_FFN_CHANNEL + assert layout.quantized_modules == frozenset(f"model.layers.{i}.mlp.{p}" for i in range(2) + for p in ("gate_proj", "up_proj", "down_proj")) + assert layout.tensors["model.layers.0.mlp.up_proj.weight"].dtype == "F8_E4M3" + assert layout.tensors["model.layers.0.mlp.up_proj.weight_scale"].shape == (16, 1) + assert layout.tensors["model.layers.1.mlp.up_proj.weight_scale"].shard == "model-00002-of-00002.safetensors" + assert layout.tensors["model.layers.0.self_attn.q_proj.weight"].dtype == "BF16" + config = json.loads((tmp_path / "t" / "config.json").read_text()) + q = config["quantization_config"] + assert (q["quant_method"], q["format"]) == ("compressed-tensors", "float-quantized") + assert {"model.layers.0.self_attn.o_proj", "model.embed_tokens", "lm_head"} <= set(q["ignore"]) # lm_head tied, absent + # fused alias for engines without a packed_modules_mapping (SGLang Qwen3); FFN stays quantized + assert "model.layers.0.self_attn.qkv_proj" in q["ignore"] + assert not any(name.endswith("gate_up_proj") for name in q["ignore"]) + assert CheckpointLayout.load(tmp_path / "t").quantized_modules == layout.quantized_modules + assert (tmp_path / "t" / "tokenizer.json").is_file() + + +def test_scoped_layout_quantizes_only_moe_routed_experts(moe): + layout, kind = scoped_layout(moe, FP8_FFN_CHANNEL) + assert kind == "moe" + assert layout.quantized_modules == frozenset(f"model.layers.1.mlp.experts.{e}.{p}" for e in range(2) + for p in ("gate_proj", "up_proj", "down_proj")) + ignore = set(layout.quantization_config["ignore"]) + assert {"model.layers.1.mlp.gate", "model.layers.1.mlp.shared_expert.up_proj", "lm_head", + "model.layers.0.mlp.up_proj", "model.layers.0.mlp.gate_up_proj", + "model.layers.1.mlp.shared_expert.gate_up_proj"} <= ignore + assert not any(".experts." in name for name in ignore) + + +def test_scoped_layout_refuses_unknown_scheme_and_empty_scope(tmp_path): + ckpt = write_checkpoint(tmp_path / "attn", {k: v for k, v in dense_shapes().items() if ".mlp." not in k}, + {"num_hidden_layers": 2}) + with pytest.raises(LayoutError, match="not a scoped scheme"): + scoped_layout(ckpt, "w4a8-int4-channel") + with pytest.raises(LayoutError, match="no in-scope linear"): + scoped_layout(ckpt, FP8_FFN_CHANNEL) + + +def test_conversion_matches_the_reference_and_keeps_the_rest_bit_equal(dense, tmp_path): + save_scoped_template(dense, tmp_path / "t", FP8_FFN_CHANNEL) + stats = requantize_checkpoint(dense, tmp_path / "t", tmp_path / "out", AwqSmoothing.none()) + assert stats["quantized"] == 6 # FP8 weight tensors + src, out = _read(dense), _read(tmp_path / "out") + for name, tensor in src.items(): + module = name[: -len(".weight")] + if ".mlp." in name: + q, scale = quantize_fp8_per_channel(tensor) + assert torch.equal(out[name].view(torch.uint8), q.view(torch.uint8)), name + assert torch.equal(out[module + ".weight_scale"], scale), name + else: + assert out[name].dtype == torch.bfloat16 and torch.equal(out[name], tensor), name + assert tiers.check_release_scope(tmp_path / "out") == 6 + assert json.loads((tmp_path / "out" / "config.json").read_text())["quantization_config"]["format"] == "float-quantized" + + +def test_release_scope_check_refuses_fp8_attention_and_unlisted_bf16(dense, tmp_path): + save_scoped_template(dense, tmp_path / "t", FP8_FFN_CHANNEL) + requantize_checkpoint(dense, tmp_path / "t", tmp_path / "out", AwqSmoothing.none()) + bad = tmp_path / "out" / "model-00001-of-00002.safetensors" + tensors = _read(tmp_path / "out") + shard0 = {k: v for k, v in tensors.items() if ".layers.1." not in k} + shard0["model.layers.0.self_attn.q_proj.weight"] = quantize_fp8_per_channel(shard0["model.layers.0.self_attn.q_proj.weight"])[0] + save_file(shard0, str(bad), metadata={"format": "pt"}) + with pytest.raises(QuantScopeError, match=r"self_attn.q_proj \(attention\)"): + tiers.check_release_scope(tmp_path / "out") + + config_path = tmp_path / "out" / "config.json" + requantize_checkpoint(dense, tmp_path / "t", tmp_path / "out", AwqSmoothing.none()) + config = json.loads(config_path.read_text()) + config["quantization_config"]["ignore"].remove("model.layers.0.self_attn.q_proj") + config_path.write_text(json.dumps(config)) + with pytest.raises(QuantScopeError, match="not in the config ignore list"): + tiers.check_release_scope(tmp_path / "out") + config["quantization_config"]["format"] = "int-quantized" + config_path.write_text(json.dumps(config)) + with pytest.raises(QuantScopeError, match="expected compressed-tensors/float-quantized"): + tiers.check_release_scope(tmp_path / "out") + + +def test_scoped_releaser_converts_each_bf16_release(dense, tmp_path): + store = CheckpointStore(tmp_path / "store") + trained = write_checkpoint(tmp_path / "trained", dense_shapes(), {"num_hidden_layers": 2}, seed=9) + PolicyReleaseStore(store).release(3, "h3", lambda d: [os.link(p, d / p.name) for p in trained.iterdir()]) + releaser = scoped_releaser(store, dense, tmp_path / "template", FP8_FFN_CHANNEL, log=lambda m: None) + record = releaser.sync() + assert record is not None and record.version == 3 and record.policy_hash == f"h3+{FP8_FFN_CHANNEL}" + assert record.derived_from["scheme"] == FP8_FFN_CHANNEL and record.derived_from["expert_matrices"] == 0 # dense: no routed expert + target = PolicyReleaseStore(store, release_dir=release_dir_for(FP8_FFN_CHANNEL)) + out = target.resolve(target.latest()) + assert tiers.check_release_scope(out) == 6 + q, _ = quantize_fp8_per_channel(_read(trained)["model.layers.1.mlp.down_proj.weight"]) + assert torch.equal(_read(out)["model.layers.1.mlp.down_proj.weight"].view(torch.uint8), q.view(torch.uint8)) + assert releaser.sync() is None # converted once + assert release_dir_for(FP8_FFN_CHANNEL) == "policy-releases-fp8-w8a8-channel" + + +def test_port_loader_refuses_an_out_of_scope_fp8_release(dense, tmp_path): + from yeto.rl.adapters.sampler_engines.fake import FAKE_IDENTITY, FakeEngine + from yeto.rl.engine.policy_release import ReleaseNotLoaded + from yeto.rl.engine.sampler_checks import PortReleaseLoader, PostLoadChecks + + save_scoped_template(dense, tmp_path / "t", FP8_FFN_CHANNEL) + requantize_checkpoint(dense, tmp_path / "t", tmp_path / "out", AwqSmoothing.none()) + config_path = tmp_path / "out" / "config.json" + config = json.loads(config_path.read_text()) + config["quantization_config"]["ignore"] = [] + config_path.write_text(json.dumps(config)) + from dataclasses import replace + + engine = FakeEngine(identity=replace(FAKE_IDENTITY, quant_scheme=FP8_FFN_CHANNEL), expert_matrices=None) + events = [] + loader = PortReleaseLoader(engine, PostLoadChecks(), lambda e, **f: events.append((e, f))) + release = SimpleNamespace(seq=1, version=1, derived_from={"scheme": FP8_FFN_CHANNEL, "expert_matrices": 0}) + with pytest.raises(ReleaseNotLoaded, match="quant_scope"): + loader._serve(release, tmp_path / "out", 1, first=False) + assert engine.loads == [] and events[-1][1]["step"] == "quant_scope" + + +def test_moe_release_expert_count_matches_the_sampler_index_check(moe, tmp_path): + from yeto.rl.engine.sampler_checks import count_release_experts + from yeto.rl.quant.release import routed_expert_matrices + + layout = save_scoped_template(moe, tmp_path / "t", FP8_FFN_CHANNEL) + requantize_checkpoint(moe, tmp_path / "t", tmp_path / "out", AwqSmoothing.none()) + assert routed_expert_matrices(layout.quantized_modules) == count_release_experts(tmp_path / "out") == 6 + assert tiers.check_release_scope(tmp_path / "out") == 6 + out = _read(tmp_path / "out") + assert out["model.layers.1.mlp.gate.weight"].dtype == torch.bfloat16 + assert out["model.layers.1.mlp.shared_expert.up_proj.weight"].dtype == torch.bfloat16 + assert out["model.layers.1.mlp.experts.1.up_proj.weight"].dtype == torch.float8_e4m3fn + + +def test_miles_translation_adds_sglang_flags_only_for_tier_1(): + from tests.test_rl_miles_adapter_config import make_config, sub + from yeto.rl.adapters.miles import config as mc + from yeto.rl.engine.algorithm import AlgorithmSpec + + base = mc.translate_run_config(make_config(), AlgorithmSpec()) + assert "--sglang-quantization" not in base.argv and "yeto_rollout_quant_tier" not in base.runtime_attrs + tier1 = mc.translate_run_config(sub(make_config(), "serving", rollout_quant_tier=1), AlgorithmSpec()) + argv = list(tier1.argv) + assert argv[argv.index("--sglang-quantization") + 1] == "fp8" + assert "--sglang-kv-cache-dtype" not in argv + assert tier1.runtime_attrs["yeto_rollout_quant_tier"] == 1 + tier2 = mc.translate_run_config(sub(make_config(), "serving", rollout_quant_tier=2), AlgorithmSpec()) + assert "--sglang-quantization" not in tier2.argv and tier2.runtime_attrs["yeto_rollout_quant_tier"] == 2 + + +@pytest.mark.parametrize("tier", [0, 1]) +def test_island_task_forwards_the_tier_only_when_set(monkeypatch, tier): + import sys + import types + + from tests import test_rl_launcher as tl + from yeto.gpu_spec import parse_gpu_spec + from yeto.launch.rl_args import _prepare_rl_args + from yeto.launcher import make_miles_island_task + + monkeypatch.setitem(sys.modules, "sky", types.SimpleNamespace( + Task=tl._Task, Resources=tl._Resources, Storage=tl._Storage, StorageMode=tl._StorageMode)) + args = tl._args(("--gpu", "aws:1xh100@us-east-1", "--rl-rollout-quant-tier", str(tier))) + args.model_revision, args.data_revision = "a" * 40, "b" * 40 + args.source_sha256, args.reward_sha256 = "c" * 64, "d" * 64 + _prepare_rl_args(args) + task = make_miles_island_task(args, parse_gpu_spec(args.gpu)[0], 0, 1, "127.0.0.1:29400") + assert ("--rl-rollout-quant-tier 1" in task.run) is (tier == 1) + # Miles tier 1 defaults to the trainer source; tier 0 always uses the engine. + assert ("--rl-rollout-quant-fp8-source trainer" in task.run) is (tier == 1) + if tier == 1: + args.rl_rollout_quant_fp8_source = "engine" + task = make_miles_island_task(args, parse_gpu_spec(args.gpu)[0], 0, 1, "127.0.0.1:29400") + assert "--rl-rollout-quant-fp8-source" not in task.run + + +MILES_NEXT = Path("/home/michael/work/miles-next") + + +@pytest.mark.skipif(not (MILES_NEXT / "miles/utils/workers/process_utils.py").is_file(), reason="no miles-next checkout") +def test_ignore_list_reaches_the_sglang_server_process(dense, tmp_path, monkeypatch): + """Tier 1, item 4 of the GPU review: the last hop. Miles' CommandActor starts + ``sglang.launch_server`` with ``process_utils.launch_bound_subprocess(["/bin/sh", "-c", + "exec "], envs={})`` (ray_worker_manager.py:754, command_actor.py:22), whose child env + is ``{**os.environ, **envs}``. The actor's os.environ holds the job runtime_env entries + (connect_island_ray, tested above). This runs Miles' own function and reads the + variable from the started process.""" + import importlib.util + import sys + + from yeto.rl.adapters.miles.entry import rollout_quant_env + + monkeypatch.setenv("PYTHONPATH", f"{MILES_NEXT}:{os.environ.get('PYTHONPATH', '')}") + for key, value in rollout_quant_env(SimpleNamespace(yeto_rollout_quant_tier=1, hf_checkpoint=str(dense))).items(): + monkeypatch.setenv(key, value) # what the Ray worker's environment holds + # loaded from its file, so no real ``miles`` package lands in sys.modules (other tests fake it) + spec = importlib.util.spec_from_file_location("miles_process_utils", + MILES_NEXT / "miles/utils/workers/process_utils.py") + process_utils = importlib.util.module_from_spec(spec) + spec.loader.exec_module(process_utils) + out = tmp_path / "seen.txt" + cmd = f"exec {sys.executable} -c \"import os; open('{out}', 'w').write(os.environ['SGLANG_FP8_IGNORED_LAYERS'])\"" + process = process_utils.launch_bound_subprocess(["/bin/sh", "-c", cmd], envs={}) + assert process.wait(timeout=60) == 0 + assert out.read_text() == os.environ["SGLANG_FP8_IGNORED_LAYERS"] + assert "model.layers.1.self_attn.v_proj" in out.read_text().split(",") + + +QWEN3_4B = json.loads((Path(__file__).parent / "data" / "qwen3_4b_instruct_2507_weights.json").read_text()) + + +@pytest.mark.skipif(not (MILES_NEXT / "miles/backends/megatron_utils/megatron_to_hf/qwen2.py").is_file(), + reason="no miles-next checkout") +def test_miles_raw_export_names_equal_the_hf_base_names(): + """Item 5 of the GPU review (raw path): every Megatron parameter of Qwen3-4B (TE layout, + tied embeddings: no output_layer) through Miles' own ``convert_qwen2_to_hf`` gives exactly + the tensor names of the HF base index. The Qwen3-4B launch uses Miles' bridge export + (Megatron-Bridge, not in this checkout); that path is checked on the first release.""" + import importlib.util + + spec = importlib.util.spec_from_file_location( + "miles_qwen2_to_hf", MILES_NEXT / "miles/backends/megatron_utils/megatron_to_hf/qwen2.py") + module = importlib.util.module_from_spec(spec) + spec.loader.exec_module(module) + c = QWEN3_4B + head_dim, groups, heads = 4, 2, 4 + args = SimpleNamespace(kv_channels=head_dim, hidden_size=8, num_attention_heads=heads, num_query_groups=groups) + params = {"module.module.embedding.word_embeddings.weight": torch.zeros(4, 8), + "module.module.decoder.final_layernorm.weight": torch.zeros(8)} + for i in range(c["num_hidden_layers"]): + p = f"module.module.decoder.layers.{i}." + params.update({p + "self_attention.linear_qkv.weight": torch.zeros(groups * (heads // groups + 2) * head_dim, 8), + p + "self_attention.linear_qkv.layer_norm_weight": torch.zeros(8), + p + "self_attention.linear_proj.weight": torch.zeros(8, 8), + p + "self_attention.q_layernorm.weight": torch.zeros(head_dim), + p + "self_attention.k_layernorm.weight": torch.zeros(head_dim), + p + "mlp.linear_fc1.weight": torch.zeros(16, 8), + p + "mlp.linear_fc1.layer_norm_weight": torch.zeros(8), + p + "mlp.linear_fc2.weight": torch.zeros(8, 8)}) + names = {hf for name, param in params.items() for hf, _ in module.convert_qwen2_to_hf(args, name, param)} + assert c["tie_word_embeddings"] and "lm_head.weight" not in c["weight_names"] + assert names == set(c["weight_names"]) + + +def test_tied_lm_head_in_an_export_is_dropped(tmp_path): + """Item 5 fallback: an exporter that writes the tied head anyway still converts.""" + shapes = dense_shapes() + base = write_checkpoint(tmp_path / "base", shapes, {"num_hidden_layers": 2, "tie_word_embeddings": True}) + export = write_checkpoint(tmp_path / "export", {**shapes, "lm_head.weight": [32, H]}, + {"num_hidden_layers": 2, "tie_word_embeddings": True}) + save_scoped_template(base, tmp_path / "t", FP8_FFN_CHANNEL) + stats = requantize_checkpoint(export, tmp_path / "t", tmp_path / "out", AwqSmoothing.none()) + assert stats["dropped_tied_lm_head"] == 1 and "lm_head.weight" not in _read(tmp_path / "out") + untied = write_checkpoint(tmp_path / "base2", shapes, {"num_hidden_layers": 2}) + save_scoped_template(untied, tmp_path / "t2", FP8_FFN_CHANNEL) + with pytest.raises(LayoutError, match="1 extra"): + requantize_checkpoint(export, tmp_path / "t2", tmp_path / "out2", AwqSmoothing.none()) + + +# ---------------------------------------------------------------- tier 1, trainer FP8 source (Miles native) +def test_tier_1_defaults_to_the_trainer_source_on_miles_only(): + miles = _args(rl_rollout_quant_fp8_source=None, rl_backend="miles") + assert tiers.validate_launch(miles) == 1 + assert miles.rl_rollout_quant_fp8_source == "trainer" + explicit = _args(rl_rollout_quant_fp8_source="engine", rl_backend="miles") + tiers.validate_launch(explicit) + assert explicit.rl_rollout_quant_fp8_source == "engine" + tier0 = _args(rl_rollout_quant_tier=0, rl_rollout_quant_fp8_source=None, rl_backend="miles") + tiers.validate_launch(tier0) + assert tier0.rl_rollout_quant_fp8_source == "engine" + assert tiers.default_fp8_source(_args(rl_backend="verl")) == "engine" + + +def test_trainer_source_is_tier_1_only(): + assert tiers.validate_launch(_args(rl_rollout_quant_fp8_source="trainer")) == 1 + with pytest.raises(QuantTierError, match="needs --rl-rollout-quant-tier 1"): + tiers.validate_launch(_args(rl_rollout_quant_tier=0, rl_rollout_quant_fp8_source="trainer")) + with pytest.raises(QuantTierError, match="tier 1 only"): + tiers.validate_launch(_args(rl_rollout_quant_tier=2, rl_island_sample_role=["0:sampler"], + rl_rollout_quant_fp8_source="trainer")) + + +def test_miles_translation_of_the_trainer_source_has_no_sglang_flag(): + from tests.test_rl_miles_adapter_config import make_config, sub + from yeto.rl.adapters.miles import config as mc + from yeto.rl.engine.algorithm import AlgorithmSpec + + native = mc.translate_run_config( + sub(make_config(), "serving", rollout_quant_tier=1, rollout_quant_fp8_source="trainer"), AlgorithmSpec()) + assert "--sglang-quantization" not in native.argv + assert native.runtime_attrs["yeto_rollout_quant_fp8_source"] == "trainer" + online = mc.translate_run_config(sub(make_config(), "serving", rollout_quant_tier=1), AlgorithmSpec()) + assert "yeto_rollout_quant_fp8_source" not in online.runtime_attrs + + +def _fake_converter(fp8_modules: list[str]): + """Stands in for Miles' convert_hf_to_fp8.py (needs CUDA): FP8 for ``fp8_modules``.""" + calls = [] + + def run(argv, **_kw): + calls.append(argv) + src, out = Path(argv[argv.index("--model-dir") + 1]), Path(argv[argv.index("--save-dir") + 1]) + tensors = _read(src) + for module in fp8_modules: + tensors[f"{module}.weight"] = tensors[f"{module}.weight"].to(torch.float8_e4m3fn) + out.mkdir(parents=True, exist_ok=True) + save_file(tensors, str(out / "model.safetensors"), metadata={"format": "pt"}) + config = json.loads((src / "config.json").read_text()) + ignored = json.loads(Path(argv[argv.index("--ignored-layers-file") + 1]).read_text()) + config["quantization_config"] = {"quant_method": "fp8", "fmt": "e4m3", "activation_scheme": "dynamic", + "weight_block_size": [128, 128], "modules_to_not_convert": ignored} + (out / "config.json").write_text(json.dumps(config)) + return run, calls + + +def test_native_base_converts_with_our_ignore_list_and_points_the_engines_at_it(dense, tmp_path): + from yeto.rl.adapters.miles.entry import prepare_native_fp8_base + + ffn = [f"model.layers.{i}.mlp.{p}" for i in range(2) for p in ("gate_proj", "up_proj", "down_proj")] + run, calls = _fake_converter(ffn) + args = SimpleNamespace(yeto_rollout_quant_tier=1, hf_checkpoint=str(dense), rollout_hf_checkpoint=None) + out = prepare_native_fp8_base(args, tiers.online_plan(dense), tmp_path / "work", run=run) + assert args.rollout_hf_checkpoint == str(out) and args.hf_checkpoint == str(dense) + argv = calls[0] + assert argv[argv.index("--strategy") + 1] == "block" and argv[argv.index("--block-size") + 1:][:2] == ["128", "128"] + ignored = json.loads((tmp_path / "work" / "fp8-native-ignore.json").read_text()) + assert "model.layers.0.self_attn.q_proj" in ignored and "lm_head" in ignored + assert not any("mlp.gate_proj" in name for name in ignored) + + +def test_native_base_with_fp8_attention_is_refused(dense, tmp_path): + from yeto.rl.adapters.miles.entry import prepare_native_fp8_base + + run, _ = _fake_converter(["model.layers.0.mlp.up_proj", "model.layers.0.self_attn.q_proj"]) + args = SimpleNamespace(yeto_rollout_quant_tier=1, hf_checkpoint=str(dense), rollout_hf_checkpoint=None) + with pytest.raises(QuantScopeError, match="attention"): + prepare_native_fp8_base(args, tiers.online_plan(dense), tmp_path / "work", run=run) + assert args.rollout_hf_checkpoint is None + + +def test_native_base_needs_the_miles_rollout_checkpoint_flag(dense, tmp_path): + from yeto.rl.adapters.miles.entry import prepare_native_fp8_base + + with pytest.raises(RuntimeError, match="--rollout-hf-checkpoint"): + prepare_native_fp8_base(SimpleNamespace(hf_checkpoint=str(dense)), tiers.online_plan(dense), tmp_path, + run=_fake_converter([])[0]) + + +def test_miles_island_env_of_the_trainer_source_is_empty_after_the_conversion(dense, tmp_path, monkeypatch): + from yeto.rl.adapters.miles import entry + + seen = [] + monkeypatch.setattr(entry, "prepare_native_fp8_base", lambda args, plan, work: seen.append(work)) + args = SimpleNamespace(yeto_rollout_quant_tier=1, yeto_rollout_quant_fp8_source="trainer", + hf_checkpoint=str(dense)) + assert entry.rollout_quant_env(args, {"YETO_RL_WORK_DIR": str(tmp_path)}) == {} + assert seen == [tmp_path] diff --git a/tests/test_rl_resume_from_checkpoint.py b/tests/test_rl_resume_from_checkpoint.py index cf6b44a44..927cb310e 100644 --- a/tests/test_rl_resume_from_checkpoint.py +++ b/tests/test_rl_resume_from_checkpoint.py @@ -310,6 +310,33 @@ def test_colocated_trainer_is_woken_for_the_cut(tmp_path): assert calls == ["on", "off"] and info["trainer_onloaded_for_cut"] is True +def test_a_volume_reload_never_runs_inside_a_round_cut(tmp_path, monkeypatch): + """S26 E6 (e6-t2-pair-e): another process of the container reloaded the Modal + Volume while the trainer ranks wrote a cut; the remount made the shard's + makedirs fail (EPERM) and the island exited 134. The cut holds the store's + writer guard from the first shard to the commit, so the reload is skipped.""" + monkeypatch.setenv(mcs.GUARD_DIR_ENV, str(tmp_path / "guard")) + (tmp_path / "guard").mkdir() + vol = _Vol() + store = _store(tmp_path / "store", vol) + other = _store(tmp_path / "store", vol) # e.g. the sample buffer's store object + _r, _c, driver, cp = _trained_island(tmp_path, store, 2) + save_cut, seen = cp.trainer.save_cut, [] + + def save_with_a_concurrent_reload(**kw): + before = vol.reloads + other.reload() # what a reload loop in another thread/process does mid-cut + seen.append(vol.reloads - before) + return save_cut(**kw) + + cp.trainer.save_cut = save_with_a_concurrent_reload + info = cp.save(driver, rollout_id=2) + assert seen == [0] and info["store_synced"] + before = vol.reloads + other.reload() # after the commit the reload runs again + assert vol.reloads == before + 1 + + def test_local_only_sync_stop_after_and_final_cut(): sync = LocalOnlySync(6) sync.stop_after = 3 diff --git a/tests/test_rl_sampler_snapshot.py b/tests/test_rl_sampler_snapshot.py index 4ab5bff0e..de3199281 100644 --- a/tests/test_rl_sampler_snapshot.py +++ b/tests/test_rl_sampler_snapshot.py @@ -2,6 +2,7 @@ from __future__ import annotations +import dataclasses from types import SimpleNamespace import pytest @@ -24,12 +25,29 @@ def snapshot_config(*, colocated: bool = False): return sub(cfg, "serving", offload_train=False, sampler_snapshot=True) +# The G-L8 passes were recorded on SGLang 2fa880182 (S24). The pin moved to 10afe11a7 +# and 754605964 (S26 E6: the FP8 activation-cast fix in Qwen2/Qwen3, the compressed-tensors +# FP8 in-place reload); until G-L8 is re-run on it +# the snapshot mode is refused, as for any commit without a record. +VERIFIED_SGLANG_COMMIT = "2fa880182eefdb31e64d8ea70317be27f9f32121" + + +def _records_at_the_pin(): + return tuple(dataclasses.replace(r, sglang_commit=SGLANG_NEXT_COMMIT) for r in ss.load_compat_records()) + + def test_shipped_records_hold_the_l8_1_tensor_pass(): - record = ss.require_recorded_pass(ss.UPDATE_PATH_TENSOR) - assert record.sglang_commit == SGLANG_NEXT_COMMIT + record = ss.require_recorded_pass(ss.UPDATE_PATH_TENSOR, sglang_commit=VERIFIED_SGLANG_COMMIT) assert record.check == "G-L8.1" +def test_the_pinned_commit_waits_for_its_own_record(): + if SGLANG_NEXT_COMMIT == VERIFIED_SGLANG_COMMIT: + pytest.skip("the pin is the verified commit") + with pytest.raises(ss.SamplerSnapshotRefused, match="no recorded pass"): + ss.require_recorded_pass(ss.UPDATE_PATH_TENSOR) + + def test_refused_without_recorded_pass_names_the_missing_check(): with pytest.raises(ss.SamplerSnapshotRefused) as exc: ss.require_recorded_pass(ss.UPDATE_PATH_DISTRIBUTED, records=()) @@ -99,8 +117,10 @@ def test_cli_flag_is_off_by_default_and_parses(): assert cli.build_parser().parse_args([*base, "--rl-sampler-snapshot"]).rl_sampler_snapshot is True -def test_shipped_records_hold_the_l8_2_distributed_pass(): - record = ss.require_recorded_pass(ss.UPDATE_PATH_DISTRIBUTED) +def test_shipped_records_hold_the_l8_2_distributed_pass(monkeypatch): + record = ss.require_recorded_pass(ss.UPDATE_PATH_DISTRIBUTED, sglang_commit=VERIFIED_SGLANG_COMMIT) assert record.check == "G-L8.2" + at_the_pin = _records_at_the_pin() + monkeypatch.setattr(ss, "load_compat_records", lambda text=None: at_the_pin) argv = mc.translate_run_config(snapshot_config(), AlgorithmSpec()).argv assert "--sglang-config" in argv diff --git a/tests/test_rl_sampler_trainer_pair_launch.py b/tests/test_rl_sampler_trainer_pair_launch.py index 0cbd7e69c..6f986c1eb 100644 --- a/tests/test_rl_sampler_trainer_pair_launch.py +++ b/tests/test_rl_sampler_trainer_pair_launch.py @@ -113,7 +113,7 @@ def test_full_parameter_without_a_buffer_role_is_still_refused_on_the_island(): from yeto.rl.adapters.miles import island_entry args = types.SimpleNamespace(parameter_mode="full", rl_single_island_no_sync=True, - rl_sample_role="both") + rl_sample_role="both", rl_placement="colocated") with pytest.raises(ValueError, match="full-parameter sync preset"): island_entry._full_parameter_mode(args, "strict-avg") @@ -153,6 +153,43 @@ def test_tuning_full_outside_buffer_islands_is_refused(): _args(("--tuning", "full"), gpu="modal:1xh100", base=()) +DEDICATED_ENGINE = ("--tuning", "full", "--rl-single-island-no-sync", + "--rl-placement", "fixed-partition", "--rl-rollout-gpus", "1") + + +def test_tuning_full_on_one_island_with_a_dedicated_engine_card_is_accepted(): + """S26 E6: one no-sync island, 1 training card + 1 engine card; the island takes + full parameter mode and the placement (Miles' own full-parameter update).""" + from yeto.launcher import dry_run_plan + from yeto.rl.engine import sample_roles + + args = _args(DEDICATED_ENGINE, gpu="modal:2xh100", base=()) + assert sample_roles.dedicated_engine_full_parameter(args) + tokens = _island_tokens(dry_run_plan(args)["island_requests"][0]["learner_command"]) + assert _flag(tokens, "--parameter-mode") == "full" + assert _flag(tokens, "--rl-placement") == "fixed-partition" + assert "--rl-single-island-no-sync" in tokens + + +def test_dedicated_engine_island_command_passes_the_island_checks(): + from yeto.launcher import dry_run_plan + from yeto.rl.adapters.miles import island_entry + + args = _args(DEDICATED_ENGINE, gpu="modal:2xh100", base=()) + tokens = _island_tokens(dry_run_plan(args)["island_requests"][0]["learner_command"]) + parsed = island_entry.parse_args(tokens) + island_entry._check_single_island_no_sync(parsed) + assert island_entry._full_parameter_mode(parsed, parsed.sync_preset) is True + assert (parsed.rl_placement, parsed.rollout_num_gpus) == ("fixed-partition", 1) + + +@pytest.mark.parametrize("drop", [("--rl-single-island-no-sync",), + ("--rl-placement", "fixed-partition", "--rl-rollout-gpus", "1")]) +def test_tuning_full_on_one_island_needs_no_sync_and_dedicated_engine_cards(drop): + with pytest.raises(ValueError, match="requires --tuning lora"): + _args(_replace(DEDICATED_ENGINE, drop, ()), gpu="modal:2xh100", base=()) + + def test_held_trainer_drains_the_buffer_after_the_fleet_ended(tmp_path, monkeypatch): """When run() returns (the sampler ended), the loop sees data_ended and keeps ticking until it finished; then run_with_trainer_loop returns.""" diff --git a/tools/probes/s26_e6/algo-tis2.json b/tools/probes/s26_e6/algo-tis2.json new file mode 100644 index 000000000..f45d2e970 --- /dev/null +++ b/tools/probes/s26_e6/algo-tis2.json @@ -0,0 +1,53 @@ +{ + "advantage": { + "estimator": "grpo", + "reward_binary": false, + "reward_postprocess": null, + "rewards_normalization": true, + "std_normalization": true, + "whiten": false + }, + "correction": { + "batch_normalize": false, + "function": null, + "level": "token", + "lower": 0.0, + "method": "tis", + "mismatch_metrics": false, + "opsm_delta": null, + "ratio_mask": "none", + "upper": 2.0, + "use_rollout_logprobs": false, + "veto": null + }, + "entropy_coef": 0.0, + "execution": { + "group_required": true, + "max_policy_staleness": 2, + "needs_critic": false, + "needs_rollout_logprobs": false + }, + "kl": { + "coef": null, + "estimator": null, + "placement": "none", + "unbiased": false + }, + "loss": { + "aggregation": "default", + "custom_loss": null, + "eps_clip": 0.2, + "eps_clip_c": null, + "eps_clip_high": 0.28, + "reducer": null, + "variant": "policy_loss" + }, + "plugins": [], + "sampling": { + "filter": null, + "max_replacements": null, + "over_sampling_batch_size": null, + "overlong_filter": false + }, + "schema": "yeto-rl-algorithm-spec-v2" +} \ No newline at end of file diff --git a/tools/probes/s26_e6/gate_modal.py b/tools/probes/s26_e6/gate_modal.py new file mode 100644 index 000000000..86a729d3b --- /dev/null +++ b/tools/probes/s26_e6/gate_modal.py @@ -0,0 +1,217 @@ +"""S26 E6 step 3a: the rollout quant gate on one Modal ``H100!`` (Qwen3-4B-Instruct-2507). + +One container in the Miles ports image (SGLang 2fa88018) runs +``python -m yeto.rl.quant.gate collect`` for one tier: + +* the quantized engine section is the production one: a BF16 sampler spec goes + through ``sampler_engines.quant_tier.apply`` (tier 1: SGLang online FP8 flags, + FP8 KV cache and the ignore list; tier 2: the base converted to the FP8 W8A8 + channel checkpoint in the container); +* the BF16 engine section is the same spec without the tier. + +Prompts: the first ``prompts`` GSM8K train questions, chat template with +thinking off. Raw per-token data (``report-raw.json``), the report, the engine +logs and NVML go to the results Volume ``/`` and are pulled to +``--out-dir``. The container asserts every GPU is an H100 before anything else. + +Run (client venv with modal; the launcher stays on this machine):: + + python tools/probes/s26_e6/gate_modal.py --run e6-gate-t1 --tier 1 --out-dir DIR +""" + +from __future__ import annotations + +import argparse +import json +import os +import subprocess +import sys +import time +from pathlib import Path +from typing import Any + +REPO = Path(__file__).resolve().parents[3] +WORKDIR = "/root/sky_workdir" +RESULTS = "/mnt/results" +RESULTS_VOLUME = "yeto-eval-results" +MODEL = "Qwen/Qwen3-4B-Instruct-2507" +MODEL_REV = "cdbee75f17c01a7cc42f958dc650907174af0554" +DATA = "zhuzilin/gsm8k" +DATA_REV = "0cbd9f31d91ac21a7613dcbc7fef992adac459ae" +GPU = "H100!" +GPU_NAME = "H100" +# GOAL / gpu-review E6 ยง3a suggested criteria (not loosened); speed is recorded only +CRITERIA = {"tis_upper": 2.0, "max_mean_abs_diff": 0.05, "max_p99_abs_diff": 0.5, + "max_tis_clip_share": 0.01, "max_unaligned_share": 0.05} +IGNORE = [".git", ".venv", ".claude", "**/__pycache__", "**/*.pyc", "s1-runs", "syncer/target", + "evidence*", "openspec", "docs", "tutorial", "data"] + + +def gate(run: str, cfg: dict) -> dict: + """In the container: GPU check, prompts, engine sections, ``gate collect``.""" + import modal + + sys.path.insert(0, WORKDIR) + out = Path(RESULTS) / run + out.mkdir(parents=True, exist_ok=True) + meta: dict[str, Any] = {"run": run, "cfg": cfg, "container": os.environ.get("MODAL_TASK_ID"), + "start_unix": time.time()} + smi = subprocess.run(["nvidia-smi", "--query-gpu=name", "--format=csv,noheader"], + capture_output=True, text=True, timeout=60).stdout.split("\n") + names = [n.strip() for n in smi if n.strip()] + meta["gpu"] = names + (out / "nvidia-smi.txt").write_text(subprocess.run(["nvidia-smi"], capture_output=True, text=True).stdout) + nvml = subprocess.Popen(["bash", "-c", "nvidia-smi --query-gpu=timestamp,name,utilization.gpu,memory.used," + f"power.draw --format=csv -l 2 >> {out / 'nvml.csv'}"]) + try: + if not names or not all(GPU_NAME in n for n in names): + raise RuntimeError(f"asked for {GPU_NAME}, got {names}") + import pandas as pd + import sglang + from huggingface_hub import hf_hub_download, snapshot_download + from transformers import AutoTokenizer + + from yeto.cloud.modal_eval_island import _image_paths + from yeto.rl.adapters.sampler_engines import quant_tier + + _image_paths() + base = snapshot_download(MODEL, revision=MODEL_REV, local_dir="/tmp/qwen3-4b") + tok = AutoTokenizer.from_pretrained(base) + frame = pd.read_parquet(hf_hub_download(DATA, "train.parquet", repo_type="dataset", revision=DATA_REV)) + prompts = out / "prompts.jsonl" + with prompts.open("w") as handle: + for _, row in frame.head(cfg["prompts"]).iterrows(): + messages = [{"role": m["role"], "content": m["content"]} for m in row["messages"]] + handle.write(json.dumps({"prompt": tok.apply_chat_template( + messages, tokenize=False, add_generation_prompt=True)}) + "\n") + server = {"base_dir": base, "card_model": names[0], "device_family": "cuda", + "engine_version": sglang.__version__, "precision": "bf16", "quant_scheme": "bf16", + "topology": "tp1", "max_model_len": 2048, "port": 30000, "python": sys.executable, + "start_timeout_s": 1200.0, "in_place_reload": True, + "extra_args": ["--mem-fraction-static", "0.8"]} + bf16 = {"factory": "yeto.rl.adapters.sampler_engines.from_spec:build", "adapter": "sglang-server", + "server": server, "options": {"tp": 1, "served_model_name": "qwen3-4b"}} + spec = {"engine": bf16, "base_dir": base, "checks": {}, + "rollout_quant": quant_tier.spec_section(cfg["tier"])} + began = time.time() + quant = quant_tier.apply(spec, work_dir=Path("/tmp/e6-quant")) + meta["prepare_s"] = round(time.time() - began, 1) + if cfg.get("variant") in ("miles-native-block", "miles-native-block-ignore"): + # Miles' own FP8 path: decoder linears block-quantized 128x128 by Miles' converter + # (the format quantize_params_fp8 sends). "-ignore": the production tier 1 trainer + # source (entry.prepare_native_fp8_base): our ignore list stays BF16 (attention, ...) + native = Path("/tmp/e6-quant/miles-fp8-block") + convert = [sys.executable, "/root/miles/tools/convert_hf_to_fp8.py", "--model-dir", base, + "--save-dir", str(native), "--strategy", "block", "--block-size", "128", "128"] + if cfg["variant"] == "miles-native-block-ignore": + from yeto.rl.quant import tiers # noqa: PLC0415 + + ignore_file = Path("/tmp/e6-quant/native-ignore.json") + ignore_file.parent.mkdir(parents=True, exist_ok=True) + ignore_file.write_text(json.dumps(tiers.native_ignore_list(tiers.online_plan(Path(base))))) + convert += ["--ignored-layers-file", str(ignore_file)] + subprocess.run(convert, check=True, cwd="/root/miles", stdout=subprocess.DEVNULL) + if cfg["variant"] == "miles-native-block-ignore": + meta["native_fp8_modules"] = tiers.check_native_base(native) + from yeto.rl.engine.sampler_checks import release_quant_format # noqa: PLC0415 + + _method, fmt = release_quant_format(native) + (out / "quant-config.json").write_text((native / "config.json").read_text()) + quant = {"engine": {**bf16, "server": {**server, "base_dir": str(native), "precision": "fp8", + "quant_scheme": "fp8-block-native", "expected_format": fmt}}, + "rollout_quant": {"variant": cfg["variant"]}} + meta["rollout_quant"] = quant["rollout_quant"] + elif cfg.get("variant") == "sglang-online-all": # SGLang online FP8 without an ignore list + quant = {"engine": {**bf16, "server": {**server, "precision": "fp8", + "extra_args": [*server["extra_args"], "--quantization", "fp8"]}}, + "rollout_quant": {"variant": "sglang-online-all"}} + meta["rollout_quant"] = quant["rollout_quant"] + if cfg.get("bf16_kv"): # diagnosis only: tier 1 without its FP8 KV cache + args = list(quant["engine"]["server"]["extra_args"]) + i = args.index("--kv-cache-dtype") + quant["engine"]["server"]["extra_args"] = args[:i] + args[i + 2:] + meta["rollout_quant"] = quant["rollout_quant"] + qconfig = Path(quant["engine"]["server"]["base_dir"]) / "config.json" + (out / "quant-config.json").write_text(qconfig.read_text()) + engines = {"quant": {**quant["engine"], "server": {**quant["engine"]["server"], "port": 30001}}, + "bf16": bf16} + (out / "engines.json").write_text(json.dumps(engines, indent=1)) + (out / "criteria.json").write_text(json.dumps(CRITERIA, indent=1)) + argv = [sys.executable, "-m", "yeto.rl.quant.gate", "collect", "--engines", str(out / "engines.json"), + "--prompts", str(prompts), "--max-tokens", str(cfg["max_tokens"]), "--temperature", "1.0", + "--n", str(cfg["n"]), "--seed", "1", "--criteria", str(out / "criteria.json"), + "--out", str(out / "report.json"), "--label", f"tier0->tier{cfg['tier']} {MODEL}@{MODEL_REV[:8]}"] + meta["argv"] = argv + with (out / "gate.log").open("w") as log: + meta["rc"] = subprocess.run(argv, stdout=log, stderr=subprocess.STDOUT, cwd=WORKDIR, + env={**os.environ, "PYTHONPATH": WORKDIR}, + timeout=cfg["timeout_s"] - 300).returncode + except BaseException as exc: # noqa: BLE001 - recorded, then re-raised + meta["error"] = f"{type(exc).__name__}: {exc}" + raise + finally: + nvml.terminate() + meta["end_unix"] = time.time() + (out / "meta.json").write_text(json.dumps(meta, indent=1, default=str)) + modal.Volume.from_name(RESULTS_VOLUME).commit() + return meta + + +def main() -> int: + p = argparse.ArgumentParser(description=__doc__.splitlines()[0]) + p.add_argument("--run", required=True) + p.add_argument("--tier", type=int, choices=(0, 1, 2), required=True, + help="0 = control: BF16 generation scored by BF16 (the floor of the measurement)") + p.add_argument("--out-dir", required=True, type=Path) + p.add_argument("--prompts", type=int, default=64) + p.add_argument("--n", type=int, default=4) + p.add_argument("--max-tokens", type=int, default=512) + p.add_argument("--timeout-s", type=int, default=3000) + p.add_argument("--variant", choices=("miles-native-block", "miles-native-block-ignore", "sglang-online-all"), + default=None, + help="framework-native FP8 instead of the tier (with --tier 0)") + p.add_argument("--bf16-kv", action="store_true", help="diagnosis: drop the tier 1 FP8 KV cache flag") + p.add_argument("--sglang-cast-overlay", action="store_true", + help="patch the SGLang fork's FP8 activation cast in the image (sglang_fp8_cast_overlay.py)") + a = p.parse_args() + import modal + + sys.path.insert(0, str(REPO)) + from modal._vendor import cloudpickle + + from yeto.rl.adapters.miles.pins import MILES_NEXT_IMAGE + + cloudpickle.register_pickle_by_value(sys.modules[__name__]) + cfg = {"tier": a.tier, "prompts": a.prompts, "n": a.n, "max_tokens": a.max_tokens, "timeout_s": a.timeout_s, + "sglang_cast_overlay": a.sglang_cast_overlay, "bf16_kv": a.bf16_kv, "variant": a.variant} + image = (modal.Image.from_registry(MILES_NEXT_IMAGE.removeprefix("docker:")) + .env({"HOME": "/root", "PYTHONUNBUFFERED": "1"})) + if a.sglang_cast_overlay: # sglang_fp8_cast_overlay.py (the fork casts BF16 activations to FP8) + overlay = Path(__file__).with_name("sglang_fp8_cast_overlay.py") + image = (image.add_local_file(str(overlay), "/opt/e6/overlay.py", copy=True) + .run_commands("/opt/sglang/bin/python /opt/e6/overlay.py /sgl-workspace/sglang/python")) + image = image.add_local_dir(str(REPO), WORKDIR, copy=False, ignore=IGNORE) + app = modal.App(f"yeto-{a.run}") + results = modal.Volume.from_name(RESULTS_VOLUME, create_if_missing=True, version=2) + fn = app.function(image=image, gpu=GPU, cpu=8.0, memory=65536, timeout=a.timeout_s, retries=0, + volumes={RESULTS: results}, name="gate", serialized=True)(gate) + a.out_dir.mkdir(parents=True, exist_ok=True) + began = time.time() + with modal.enable_output(), app.run(detach=True): # a killed client never cancels a spawned call + (a.out_dir / "app_id.txt").write_text(f"{app.app_id}\n") + call = fn.spawn(a.run, cfg) + try: + result = {"ok": call.get(timeout=a.timeout_s + 600)} + except Exception as exc: # noqa: BLE001 - recorded; the raw data is pulled anyway + result = {"error": f"{type(exc).__name__}: {exc}"} + result["wall_s"] = round(time.time() - began, 1) + pulled = subprocess.run(["modal", "volume", "get", "--force", RESULTS_VOLUME, a.run, str(a.out_dir)], + capture_output=True, text=True, timeout=1800) + result["pull_rc"] = pulled.returncode + (a.out_dir / "driver.json").write_text(json.dumps(result, indent=1, default=str)) + print(json.dumps(result, default=str)[:2000]) + return 0 if "ok" in result else 1 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/tools/probes/s26_e6/image_check_modal.py b/tools/probes/s26_e6/image_check_modal.py new file mode 100644 index 000000000..a1c1ecb16 --- /dev/null +++ b/tools/probes/s26_e6/image_check_modal.py @@ -0,0 +1,81 @@ +"""S26 E6: check the pinned Miles ports image on a Modal CPU container (no GPU). + +Asserts in the image: the SGLang and Miles checkouts are at the pins, the +image manifest names the same commits, the SGLang FP8 activation-cast fix is +in Qwen2 MLP / Qwen3 attention, the compressed-tensors FP8 in-place reload +is there and its unit test passes, and Miles has the FP8 ignore list. + + python tools/probes/s26_e6/image_check_modal.py +""" + +from __future__ import annotations + +import json +import sys +from pathlib import Path + +REPO = Path(__file__).resolve().parents[3] + + +def check(pins: dict) -> dict: + import subprocess + + def head(path: str) -> str: + return subprocess.run(["git", "-C", path, "rev-parse", "HEAD"], capture_output=True, text=True, + check=True).stdout.strip() + + manifest = json.load(open("/opt/yeto/image-manifest.json")) + models = Path("/sgl-workspace/sglang/python/sglang/srt/models") + result = { + "sglang_head": head("/sgl-workspace/sglang"), + "miles_head": head("/root/miles"), + "manifest_sglang": manifest["sglang"]["commit"], + "manifest_miles": manifest["miles"]["commit"], + "cast_fix": {name: "dense_activation_cast_dtype(" in (models / name).read_text() + for name in ("qwen2.py", "qwen3.py")}, + "miles_fp8_ignore": Path("/root/miles/miles/utils/fp8_ignore.py").is_file() + and "is_ignored(key, ignored)" in Path("/root/miles/tools/convert_hf_to_fp8.py").read_text(), + } + scheme = Path("/sgl-workspace/sglang/python/sglang/srt/layers/quantization/compressed_tensors/" + "schemes/compressed_tensors_w8a8_fp8.py") + result["fp8_reload_fix"] = "def restore_weights_before_loading" in scheme.read_text() + reload_test = subprocess.run( + ["/opt/sglang/bin/python", "-m", "pytest", "-q", "-p", "no:cacheprovider", + "test/registered/unit/layers/quantization/test_compressed_tensors_fp8_reload.py"], + cwd="/sgl-workspace/sglang", capture_output=True, text=True) + result["fp8_reload_test_rc"] = reload_test.returncode + result["fp8_reload_test_tail"] = reload_test.stdout.strip().splitlines()[-1:] if reload_test.stdout else [] + helper = subprocess.run( + ["/opt/sglang/bin/python", "-c", + "import torch; from sglang.srt.true_on_policy.config import dense_activation_cast_dtype as f; " + "print(f.__module__)"], capture_output=True, text=True) + result["helper_import_rc"] = helper.returncode + result["ok"] = (result["sglang_head"] == result["manifest_sglang"] == pins["sglang"] + and result["miles_head"] == result["manifest_miles"] == pins["miles"] + and all(result["cast_fix"].values()) and result["miles_fp8_ignore"] + and result["fp8_reload_fix"] and reload_test.returncode == 0 + and helper.returncode == 0) + return result + + +def main() -> int: + import modal + + sys.path.insert(0, str(REPO)) + from modal._vendor import cloudpickle + + from yeto.rl.adapters.miles.pins import MILES_NEXT_COMMIT, MILES_NEXT_IMAGE, SGLANG_NEXT_COMMIT + + cloudpickle.register_pickle_by_value(sys.modules[__name__]) + app = modal.App("yeto-e6-image-check") + image = modal.Image.from_registry(MILES_NEXT_IMAGE.removeprefix("docker:")) + fn = app.function(image=image, cpu=2.0, memory=8192, timeout=900, serialized=True, name="check")(check) + with app.run(): + result = fn.remote({"sglang": SGLANG_NEXT_COMMIT, "miles": MILES_NEXT_COMMIT}) + result["image"] = MILES_NEXT_IMAGE + print(json.dumps(result, indent=1)) + return 0 if result["ok"] else 1 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/tools/probes/s26_e6/launch_locked.sh b/tools/probes/s26_e6/launch_locked.sh new file mode 100755 index 000000000..800246de0 --- /dev/null +++ b/tools/probes/s26_e6/launch_locked.sh @@ -0,0 +1,18 @@ +#!/bin/bash +# Start one E6 Modal driver under GPU-LAUNCH.lock and release the lock once the +# detached app exists (app_id.txt) plus 60 s; the driver keeps running unlocked. +# usage: launch_locked.sh OUT_DIR LOG -- driver argv... +set -u +OUT=$1; LOG=$2; shift 3 +LOCK=/home/michael/work/infra-drafts/GPU-LAUNCH.lock +exec 9>"$LOCK" +flock 9 +T=$(ps -u michael -L --no-headers | wc -l) +while [ "$T" -ge 2950 ]; do echo "waiting: $T threads" >> "$LOG"; sleep 30; T=$(ps -u michael -L --no-headers | wc -l); done +mkdir -p "$OUT" +"$@" 9>&- >> "$LOG" 2>&1 & +PID=$! +for _ in $(seq 1 120); do [ -s "$OUT/app_id.txt" ] && break; kill -0 $PID 2>/dev/null || break; sleep 5; done +sleep 60 +flock -u 9; exec 9>&- +wait $PID diff --git a/tools/probes/s26_e6/miles_local.sh b/tools/probes/s26_e6/miles_local.sh new file mode 100644 index 000000000..923eb73c0 --- /dev/null +++ b/tools/probes/s26_e6/miles_local.sh @@ -0,0 +1,109 @@ +#!/bin/bash +# S26 E6 copy of tools/eval/miles_local.sh (S26 E0.1, s26-e01-effeval 5a01732c): one yeto +# launcher run on Modal (`ARGS`), optionally a second sampler-only launch (`SARGS`). +# Derived from s1-runs/s25-c1/run.sh (S25 C1) and s1-runs/s24-l5.4/run.sh (S24 L5.4). +# +# env: P (run id; app yeto-$P, store prefix $P), R (raw dir), REPO (yeto checkout), +# ARGS (launcher args; the store is modal-volume://$STOREVOL/$P), SARGS (optional), +# ISLANDS (containers before the GPU lock is released), HARD (s), STOREVOL, TAPEVOL, +# RECLAIM_LABEL (e.g. "l0") + RECLAIM_ROUND (rollout id written before the stop), +# PY (python with modal + the launcher deps), DRY (non-empty: --dry-run only) +set -u +: "${P:?}" "${R:?}" "${REPO:?}" "${ARGS:?}" "${STOREVOL:?}" "${PY:?}" +HARD=${HARD:-3000}; ISLANDS=${ISLANDS:-1}; TAPEVOL=${TAPEVOL:-yeto-event-tapes} +MODAL="$PY -m modal"; APP=yeto-$P; SAPP=yeto-$P-s; LOCK=/home/michael/work/infra-drafts/GPU-LAUNCH.lock +CLONE=$R/yeto +T=$(ps -u michael -L --no-headers | wc -l); [ -n "${DRY:-}" ] || [ "$T" -lt 3000 ] || { echo "abort: $T threads"; exit 3; } +pgrep -x raylet >/dev/null && { echo "abort: local Ray"; exit 3; } +[ -e $R/start_utc.txt ] && { echo "abort: $R used"; exit 67; } +mkdir -p $R/tapes $R/store +# a real clone: the code pack needs a git checkout +git clone -q --no-local $REPO $CLONE && git -C $REPO rev-parse HEAD > $R/yeto_sha.txt +cp $REPO/yeto/rl/gsm8k_reward.py $CLONE/gsm8k_reward.py +(cd $CLONE && git add -f gsm8k_reward.py && git -c user.email=eval@yeto -c user.name=t-eval commit -qm 'run inputs') +eval "$(/usr/bin/python3 - <<'PYC' +import shlex, tomllib +t = tomllib.load(open("/home/michael/.modal.toml", "rb")) +prof = next((v for v in t.values() if isinstance(v, dict) and v.get("active")), None) or next(v for v in t.values() if isinstance(v, dict)) +print(f"export MODAL_TOKEN_ID={shlex.quote(prof['token_id'])} MODAL_TOKEN_SECRET={shlex.quote(prof['token_secret'])}") +PYC +)" +echo "$ARGS" > $R/args.txt; echo "${SARGS:-}" > $R/sargs.txt +if [ -n "${DRY:-}" ]; then + (cd $CLONE && eval "PYTHONPATH=$CLONE YETO_RUNS_DIR=$R/runs $PY -m yeto.cli $ARGS --dry-run") > $R/plan.json 2> $R/plan.err + echo "dry-run rc=$?"; tail -3 $R/plan.err + if [ -n "${SARGS:-}" ]; then + (cd $CLONE && eval "PYTHONPATH=$CLONE YETO_RUNS_DIR=$R/sruns $PY -m yeto.cli $SARGS --dry-run") > $R/splan.json 2> $R/splan.err + echo "sampler dry-run rc=$?"; tail -3 $R/splan.err + fi + exit 0 +fi +stop_apps(){ for a in $APP $SAPP; do case $a in yeto-$P|yeto-$P-s) (cd /tmp && timeout 120 $MODAL app stop -y $a);; esac; done > $R/final-stop.out 2>&1; } +trap stop_apps EXIT +setsid nohup bash -c "sleep $(( HARD + 300 )); cd /tmp; $MODAL app stop -y $APP; $MODAL app stop -y $SAPP; touch $R/WATCHDOG_FIRED" >/dev/null 2>&1 < /dev/null & +echo $! > $R/watchdog.pid +# container tape + NVML puller (both apps) +( while sleep 20; do c=$(cd /tmp && timeout 60 $MODAL container list --json 2>/dev/null | python3 -c "import json,sys; print(' '.join(x['container_id'] for x in json.load(sys.stdin) if x.get('app_name') in ('$APP','$SAPP')))" 2>/dev/null) + echo "$(date -u +%FT%TZ) $c" >> $R/containers.log + for x in $c; do + (cd /tmp && timeout 60 $MODAL container exec $x -- sh -c "cat /root/yeto-output/rl-island-*.jsonl 2>/dev/null") > $R/tapes/.$x 2>/dev/null && [ -s $R/tapes/.$x ] && grep -q '"event"' $R/tapes/.$x && mv $R/tapes/.$x $R/tapes/$x.jsonl + (cd /tmp && timeout 30 $MODAL container exec $x -- sh -c "nvidia-smi --query-gpu=timestamp,name,utilization.gpu,memory.used,power.draw --format=csv,noheader") 2>/dev/null | sed "s/^/$x,/" >> $R/nvml.csv; done + done ) 9>&- & +PULLER=$! +exec 9>$LOCK; flock 9 +date -u +%FT%TZ > $R/start_utc.txt +( cd $CLONE && eval "PYTHONPATH=$CLONE YETO_RUNS_DIR=$R/runs timeout $HARD $PY -m yeto.cli $ARGS" 2>&1 | awk '{ print strftime("%FT%TZ", systime(), 1) " " $0; fflush() }' > $R/launch.ts.log; echo ${PIPESTATUS[0]} > $R/rc ) 9>&- & +SUB=$! +SSUB="" +if [ -n "${SARGS:-}" ]; then + ( cd $CLONE && eval "PYTHONPATH=$CLONE YETO_RUNS_DIR=$R/sruns timeout $HARD $PY -m yeto.cli $SARGS" 2>&1 | awk '{ print strftime("%FT%TZ", systime(), 1) " " $0; fflush() }' > $R/slaunch.ts.log; echo ${PIPESTATUS[0]} > $R/src ) 9>&- & + SSUB=$! +fi +# stall stop: no new trained round (trainer) or sampler round for STALL_S s -> stop both apps +( last=$(date +%s); prev=-1 + while kill -0 $SUB 2>/dev/null; do sleep 30 + n=$(cat $R/launch.ts.log $R/slaunch.ts.log 2>/dev/null | grep -aEc '"event":"(rl_round_trained|rl_sampler_round)"') + [ "$n" != "$prev" ] && { prev=$n; last=$(date +%s); } + if [ $(( $(date +%s) - last )) -ge ${STALL_S:-900} ]; then + echo "$(date -u +%FT%TZ) no new step for ${STALL_S:-900} s (events=$n): stopping" >> $R/STALL_STOP + (cd /tmp && timeout 120 $MODAL app stop -y $APP; timeout 120 $MODAL app stop -y $SAPP) >> $R/STALL_STOP 2>&1; exit 0 + fi + done ) 9>&- & +# release the launch lock once every island has a container (or the launch ended) +for i in $(seq 1 120); do + n=$(cat $R/launch.ts.log $R/slaunch.ts.log 2>/dev/null | grep -ac 'rank 0 container') + [ "$n" -ge "$ISLANDS" ] && break + kill -0 $SUB 2>/dev/null || break; sleep 10 +done +exec 9>&- +# one reclaim: island $RECLAIM_LABEL wrote rollout $RECLAIM_ROUND -> stop its container +if [ -n "${RECLAIM_LABEL:-}" ]; then +( cd /tmp + for i in $(seq 1 600); do + kill -0 $SUB 2>/dev/null || { echo "$(date -u +%FT%TZ) launch ended before the reclaim" >> $R/reclaim.log; exit 0; } + if grep -a -- "-$RECLAIM_LABEL-modal" $R/launch.ts.log | grep -a '"event":"rl_sampler_round"' | grep -aq "\"rollout_id\":$RECLAIM_ROUND[,}]"; then + CID=$(grep -a -oP "\[$P-$RECLAIM_LABEL-modal\] \[modal-island \d+\] rank 0 container \K\S+" $R/launch.ts.log | tail -1) + echo "$(date -u +%FT%TZ) trigger container=$CID" >> $R/reclaim.log + [ -n "$CID" ] && timeout 120 $MODAL container stop --yes $CID >> $R/reclaim.log 2>&1; echo "stop rc=$? $(date -u +%FT%TZ) $(date +%s.%N)" >> $R/reclaim.log + exit 0 + fi + sleep 5 + done ) 9>&- & +fi +wait $SUB +# a sampler-only launch ends at its backlog cap once the trainer is gone: stop it +if [ -n "$SSUB" ]; then sleep 30; (cd /tmp && timeout 120 $MODAL app stop -y $SAPP) >> $R/final-stop.out 2>&1; wait $SSUB; fi +date -u +%FT%TZ > $R/end_utc.txt +sed 's/^[^ ]* //' $R/launch.ts.log > $R/launch.log +kill $PULLER 2>/dev/null; kill $(cat $R/watchdog.pid) 2>/dev/null +stop_apps; trap - EXIT; sleep 10 +(cd /tmp && timeout 120 $MODAL app list --json) > $R/final-applist.json 2>&1 +mkdir -p $R/tape-direct +for a in $APP $SAPP; do + (cd /tmp && timeout 600 $MODAL volume get --force $TAPEVOL $a $R/tape-direct) >> $R/tape-direct.out 2>&1 + AID=$(/usr/bin/python3 -c "import json,sys;print(next((a['app_id'] for a in json.load(open(sys.argv[1])) if a.get('description')==sys.argv[2]),''))" $R/final-applist.json $a 2>/dev/null) + [ -n "$AID" ] && (cd /tmp && timeout 180 $MODAL app logs --timestamps $AID > $R/modal-app-logs-$a.txt 2>&1) +done +(cd /tmp && for d in sample-buffer policy-releases round-cuts trainer-control; do timeout 900 $MODAL volume get --force $STOREVOL $P/$d $R/store; done + timeout 300 $MODAL volume ls --json $STOREVOL $P) > $R/store.out 2>&1 +echo "$P rc=$(cat $R/rc) src=$(cat $R/src 2>/dev/null)" diff --git a/tools/probes/s26_e6/parse_miles.py b/tools/probes/s26_e6/parse_miles.py new file mode 100644 index 000000000..3f2b86a39 --- /dev/null +++ b/tools/probes/s26_e6/parse_miles.py @@ -0,0 +1,94 @@ +"""S26 E6: per-round metrics of Miles RL runs (rl_miles.sh) from the launcher logs. + +For each run directory: reward mean, rollout/trainer mismatch (Miles ``mismatch``: +train_rollout_kl, tis_abs = mean |ratio - 1|, tis_clipfrac), rollout seconds and +generated tokens (colocated runs), round wall seconds. Then per group the +[min, max] over its runs of the last-3-round reward mean (the range criterion, +gpu-review/E6.md ยง4). + + python tools/probes/s26_e6/parse_miles.py /home/michael/work/s1-runs/s26-e6 > summary.json +""" + +from __future__ import annotations + +import json +import re +import sys +from pathlib import Path +from statistics import mean + +EVENT = re.compile(r"YETO_RL_EVENT (\{.*\})\s*$") + + +def events(log: Path, name: str) -> list[dict]: + out = [] + if not log.is_file(): + return out + for line in log.read_text(errors="replace").splitlines(): + m = EVENT.search(line) + if m and f'"event":"{name}"' in m.group(1): + try: + out.append(json.loads(m.group(1))) + except json.JSONDecodeError: + pass + return out + + +def run_rounds(run: Path) -> list[dict]: + log = run / "launch.ts.log" + trained = {e.get("rollout_id", i): e for i, e in enumerate(events(log, "rl_round_trained"))} + local = {e["local_round_id"] - 1: e for e in events(log, "rl_local_round") if "local_round_id" in e} + rows = [] + for rid in sorted(trained): + t, lr = trained[rid], local.get(rid, {}) + mm = t.get("mismatch") or {} + rows.append({ + "round": rid, + "reward_mean": lr.get("reward_mean", t.get("reward_mean")), + "train_rollout_kl": mm.get("train_rollout_kl"), + "tis_abs": mm.get("tis_abs"), + "tis_clipfrac": mm.get("tis_clipfrac"), + "rollout_seconds": lr.get("rollout_seconds"), + "action_tokens": lr.get("action_tokens"), + "train_seconds": lr.get("train_seconds"), + "resp_len_mean": t.get("resp_len_mean"), + }) + return rows + + +def summary(rows: list[dict]) -> dict: + def col(key): + return [r[key] for r in rows if r.get(key) is not None] + + rewards = col("reward_mean") + tok, sec = col("action_tokens"), col("rollout_seconds") + return { + "rounds": len(rows), + "reward_last3_mean": mean(rewards[-3:]) if len(rewards) >= 3 else None, + "reward_all": [round(x, 4) for x in rewards], + "kl_mean": mean(col("train_rollout_kl")) if col("train_rollout_kl") else None, + "tis_abs_mean": mean(col("tis_abs")) if col("tis_abs") else None, + "tis_clipfrac_mean": mean(col("tis_clipfrac")) if col("tis_clipfrac") else None, + "tis_clipfrac_max": max(col("tis_clipfrac")) if col("tis_clipfrac") else None, + "rollout_tokens_per_s": (sum(tok) / sum(sec)) if tok and sec and len(tok) == len(sec) else None, + } + + +def main(base: str) -> None: + runs = {} + for run in sorted(Path(base).glob("e6-*")): + if run.is_dir() and "fail" not in run.name and "dry" not in run.name: + rows = run_rounds(run) + runs[run.name] = {"rc": (run / "rc").read_text().strip() if (run / "rc").is_file() else None, + "summary": summary(rows), "rounds": rows} + groups: dict[str, list[float]] = {} + for name, doc in runs.items(): + value = doc["summary"]["reward_last3_mean"] + if value is not None: + groups.setdefault(name.rsplit("-", 1)[0], []).append(value) + print(json.dumps({"groups_last3_reward_range": {g: [min(v), max(v)] for g, v in groups.items()}, + "runs": runs}, indent=1, default=str)) + + +if __name__ == "__main__": + main(sys.argv[1]) diff --git a/tools/probes/s26_e6/parse_verl.py b/tools/probes/s26_e6/parse_verl.py new file mode 100644 index 000000000..99104bd1e --- /dev/null +++ b/tools/probes/s26_e6/parse_verl.py @@ -0,0 +1,51 @@ +"""Per-step metrics of an E6 verl run from its island log (verl console metrics). + +Per step: reward mean (``critic/score/mean``), rollout-vs-trainer log-prob |diff| +(``yeto_mismatch/abs_diff``), TIS clip share (``yeto_mismatch/tis_clipfrac``), +response tokens (``response_length/mean`` x samples), the step's wall time +(between consecutive step lines; includes generation and training) and its +output tokens per second. Usage: ``python parse_verl.py RUN_DIR... > out.json``. +""" + +from __future__ import annotations + +import json +import re +import sys +from pathlib import Path + +KEYS = ("critic/score/mean", "yeto_mismatch/abs_diff", "yeto_mismatch/tis_clipfrac", "yeto_mismatch/signed_mean", + "response_length/mean", "training/rollout_probs_diff_mean", "timing_s/gen", "timing_s/step") +NUM = r"(?:np\.float64\()?(-?[0-9.]+(?:e-?[0-9]+)?)" + + +def parse(log: Path, samples: int) -> list[dict]: + steps, last_t = [], None + for line in log.read_text(errors="replace").splitlines(): + m = re.search(r"\bstep:(\d+) ", line) + if not m or "critic/score/mean" not in line: + continue + row: dict = {"step": int(m.group(1))} + for key in KEYS: + hit = re.search(re.escape(key) + ":" + NUM, line) + if hit: + row[key] = float(hit.group(1)) + steps.append(row) + return steps + + +def main() -> None: + out = {} + for run_dir in sys.argv[1:]: + run_dir = Path(run_dir) + log = next(run_dir.rglob("island.log")) + meta = json.loads(next(run_dir.rglob("meta.json")).read_text()) + cfg = meta["cfg"] + rows = parse(log, cfg["samples"]) + out[run_dir.name] = {"rc": meta.get("rc"), "seconds": meta.get("seconds"), "gpu": meta.get("gpu"), + "groups": cfg["groups"], "samples": cfg["samples"], "steps": rows} + json.dump(out, sys.stdout, indent=1) + + +if __name__ == "__main__": + main() diff --git a/tools/probes/s26_e6/rl_miles.sh b/tools/probes/s26_e6/rl_miles.sh new file mode 100755 index 000000000..e1662d4bd --- /dev/null +++ b/tools/probes/s26_e6/rl_miles.sh @@ -0,0 +1,59 @@ +#!/bin/bash +# S26 E6 step 3 (Miles): one RL run of a group on Modal H100!, Qwen3-4B-Instruct-2507 on GSM8K, +# 10 rounds x 32 prompts x 4 samples, TIS token-level upper 2.0 (algo-tis2.json, as E0.1). +# Review: infra-drafts/SESSION26/gpu-review/E6.md. usage: rl_miles.sh GROUP RUN_ID +# t0-lora / t1-lora 1 colocated island, LoRA r=32, tier 0 / tier 1 (SGLang online FP8 + ignore list) +# t1n-lora 1 colocated island, LoRA r=32, tier 1 trainer source (engines on Miles' +# block-FP8 base + ignore list) +# t0-full / t1n-full 1 island, full parameters, 1 training card + 1 dedicated engine card +# (engines on the island need DP=1; Adam state on the host, 200 GiB) +# (fixed partition): Miles' own full-parameter update every round; tier 0 / +# tier 1 trainer source (Miles block-quantizes each update, ignore list) +# t0-pair / t2-pair trainer island (2 H100) + sampler island (1 actor + 1 engine H100), full +# parameters, tier 0 / tier 2 +# env: DRY=1 (launcher --dry-run only), SEED (default 17; the repeat of a group uses another +# seed, since Miles is nearly deterministic here and two equal-seed runs give one point) +set -u +G=${1:?group}; P=${2:?run id}; case $P in e6-*) ;; *) echo "abort: run id must start with e6-"; exit 65;; esac +HERE=$(cd "$(dirname "$0")" && pwd); REPO=$(cd "$HERE/../../.." && pwd) +R=/home/michael/work/s1-runs/s26-e6/$P +STOREVOL=yeto-eval-store +COMMON="launch --training-mode rl --rl-engine ports --preflight-threads off --modal-gpu-exact --modal-retries 0 \ +--modal-cpu 8 --modal-tape-volume yeto-event-tapes --controller local --rl-observe-timeline \ +--model Qwen/Qwen3-4B-Instruct-2507 --model-revision cdbee75f17c01a7cc42f958dc650907174af0554 \ +--data zhuzilin/gsm8k --data-revision 0cbd9f31d91ac21a7613dcbc7fef992adac459ae \ +--rl-checkpoint-store modal-volume://$STOREVOL/$P --reward-function gsm8k_reward:score --fragments 1 --pipeline 1 \ +--rollout-batch-size 32 --n-samples-per-prompt 4 --rollout-max-response-len 512 --seq-len 1024 --inner-lr 1e-5 \ +--seed ${SEED:-17} --total-steps 10 --trust-remote-code --rl-algorithm-spec tools/probes/s26_e6/algo-tis2.json" +LORA="--tuning lora --lora-r 32 --modal-memory 64" +ISLANDS=1 +case $G in + t0-lora) ARGS="$COMMON --cluster-prefix $P --gpu modal:1xh100 --rl-single-island-no-sync $LORA" ;; + t1-lora) ARGS="$COMMON --cluster-prefix $P --gpu modal:1xh100 --rl-single-island-no-sync $LORA --rl-rollout-quant-tier 1" ;; + t1n-lora) ARGS="$COMMON --cluster-prefix $P --gpu modal:1xh100 --rl-single-island-no-sync $LORA --rl-rollout-quant-tier 1 --rl-rollout-quant-fp8-source trainer" ;; + t0-full|t1n-full) + Q=""; [ $G = t1n-full ] && Q="--rl-rollout-quant-tier 1 --rl-rollout-quant-fp8-source trainer" + ARGS="$COMMON --cluster-prefix $P --gpu modal:2xh100 --rl-single-island-no-sync --tuning full \ +--modal-memory 200 --rl-optimizer-offload-fraction 1.0 --rl-placement fixed-partition --rl-rollout-gpus 1 $Q" ;; + t0-pair|t2-pair) + # Two launches on one store (as S25 C3): the trainer island reads samples and starts no + # engine (1 card); a full-parameter sampler island needs dedicated engines (fixed + # partition, 1 actor + 1 engine card). One launch with --rl-island-sample-role cannot + # give the two islands different placements for full parameters. + ROLE="$COMMON --rl-single-island-no-sync --tuning full --modal-memory 128 --rl-max-policy-age 1" + Q=""; [ $G = t2-pair ] && Q="--rl-rollout-quant-tier 2" + # full-parameter Adam state of 4B does not fit next to activations on one H100 + # (e6-t0-pair-a: CUDA OOM in round 2), and with it on the host the worker was killed + # (e6-t0-pair-c): 2 training cards shard it (distributed optimizer) + # 3c rerun (E6-5): a round cut is ~64 GB and ~4.7 min, so cut every 5 rounds; both tiers + # use the same cadence (a cut only reads the state, the training numbers do not change) + ARGS="$ROLE --cluster-prefix $P --gpu modal:2xh100 --rl-sample-role trainer --rl-sample-wait-s 900 \ +--rl-cut-every ${CUT_EVERY:-5} $Q" + SARGS="$ROLE --cluster-prefix $P-s --gpu modal:2xh100 --rl-placement fixed-partition --rl-rollout-gpus 1 \ +--rl-sample-role sampler --rl-sampler-max-backlog-rounds 2 --total-steps 12 $Q" + ISLANDS=2 ;; + *) echo "abort: group $G"; exit 65 ;; +esac +mkdir -p "$(dirname $R)" +export P R REPO ARGS ${SARGS:+SARGS} STOREVOL ISLANDS HARD=${HARD:-3600} PY=/home/michael/work/gpu-head/venv/bin/python +exec bash "$HERE/miles_local.sh" diff --git a/tools/probes/s26_e6/rl_verl_modal.py b/tools/probes/s26_e6/rl_verl_modal.py new file mode 100644 index 000000000..30977e91b --- /dev/null +++ b/tools/probes/s26_e6/rl_verl_modal.py @@ -0,0 +1,133 @@ +"""S26 E6 step 3b (verl): tier 0 against tier 1 RL on one colocated verl island, Modal ``H100!``. + +One container in the verl image (fork acad9875, vLLM 0.29) per run, as the +T-EVAL / E0.1 yeto+verl arm: the GPU type is asserted first (``modal_verl._gpu``), +NVML is sampled every 5 s, a local Ray head starts, and +``yeto.rl.adapters.verl.island_entry --rl-single-island-no-sync`` runs Qwen3-4B-Instruct-2507 +with LoRA r=32 on GSM8K (``rounds`` x ``groups`` x ``samples``; TIS token-level, upper 2.0). +Tier 1 adds ``--rl-rollout-quant-tier 1`` (vLLM online FP8 with yeto's ignore list +and FP8 KV cache, ``yeto.rl.quant.tiers.vllm_online_overrides``); nothing else changes. + +The island log (verl metrics per step), the event tape and NVML go to the results +Volume ``/island/`` and are pulled to ``--out-dir``. + +Run (client venv with modal; the launcher stays on this machine):: + + python tools/probes/s26_e6/rl_verl_modal.py --run e6-t1-verl-a --tier 1 --out-dir DIR +""" + +from __future__ import annotations + +import argparse +import json +import os +import subprocess +import sys +import time +from pathlib import Path +from typing import Any + +REPO = Path(__file__).resolve().parents[3] +sys.path.insert(0, str(REPO)) + +from tools.eval import modal_verl as mv # noqa: E402 + +MODEL = "Qwen/Qwen3-4B-Instruct-2507" +MODEL_REV = "cdbee75f17c01a7cc42f958dc650907174af0554" +GPU = "H100!" + + +def island(run: str, cfg: dict) -> dict: + """In the container: one verl island run (module doc).""" + sys.path.insert(0, mv.WORKDIR) + out = Path(mv.RESULTS) / run / "island" + out.mkdir(parents=True, exist_ok=True) + meta: dict[str, Any] = {"run": run, "cfg": cfg, "container": os.environ.get("MODAL_TASK_ID")} + meta["gpu"] = mv._gpu(out, {"gpu": GPU}) + nvml = mv._nvml(out / "nvml.csv") + began = time.time() + meta["start_unix"] = began + tape = out / "island-tape.jsonl" + try: + from yeto.rl.engine.algorithm import AlgorithmSpec, CorrectionSpec + + spec = AlgorithmSpec(correction=CorrectionSpec(method="tis", tis_clip=2.0, tis_clip_low=0.0)) + Path("/tmp/algo.json").write_text(json.dumps(spec.to_dict())) + env = mv._island_env({"gpu": GPU}, 0) + argv = [mv.VERL_PY, "-m", "yeto.rl.adapters.verl.island_entry", "--model", MODEL, + "--model-revision", MODEL_REV, "--data", mv.DATA, "--data-revision", mv.DATA_REV, + "--learner-id", "0", "--reward-function", "gsm8k_reward:score", + "--global-rounds", str(cfg["rounds"]), "--groups-per-round", str(cfg["groups"]), + "--samples-per-group", str(cfg["samples"]), "--rollout-max-response-len", str(cfg["max_tokens"]), + "--event-tape", str(tape), "--inner-lr", str(cfg["lr"]), "--lora-r", "32", + "--lora-targets", "all-linear", "--seq-len", str(cfg["seq_len"]), "--seed", str(cfg["seed"]), + "--rl-algorithm-spec", "/tmp/algo.json", "--rl-single-island-no-sync", + "--apply-chat-template-kwargs", '{"enable_thinking": false}'] + if cfg["tier"]: + argv += ["--rl-rollout-quant-tier", str(cfg["tier"])] + meta["argv"] = argv + with (out / "island.log").open("w") as log: + mv._ray_start(env, log) + meta["rc"] = subprocess.run(argv, env=env, stdout=log, stderr=subprocess.STDOUT, cwd="/workspace/verl", + timeout=cfg["timeout_s"] - 300).returncode + except BaseException as exc: # noqa: BLE001 - recorded, then re-raised + meta["error"] = f"{type(exc).__name__}: {exc}" + raise + finally: + nvml.terminate() + mv._finish(out, meta, began, {"yeto-output": Path("~/yeto-output").expanduser()}) + return meta + + +def main() -> int: + p = argparse.ArgumentParser(description=__doc__.splitlines()[0]) + p.add_argument("--run", required=True) + p.add_argument("--tier", type=int, choices=(0, 1), required=True) + p.add_argument("--out-dir", required=True, type=Path) + p.add_argument("--rounds", type=int, default=10) + p.add_argument("--groups", type=int, default=32) + p.add_argument("--samples", type=int, default=4) + p.add_argument("--max-tokens", type=int, default=512) + p.add_argument("--seq-len", type=int, default=1024) + p.add_argument("--lr", type=float, default=1e-5) + p.add_argument("--seed", type=int, default=17) + p.add_argument("--timeout-s", type=int, default=3600) + a = p.parse_args() + import modal + from modal._vendor import cloudpickle + + from yeto.rl.adapters.verl.image import modal_image + + cloudpickle.register_pickle_by_value(sys.modules[__name__]) + cloudpickle.register_pickle_by_value(mv) + cfg = {"tier": a.tier, "rounds": a.rounds, "groups": a.groups, "samples": a.samples, + "max_tokens": a.max_tokens, "seq_len": a.seq_len, "lr": a.lr, "seed": a.seed, "timeout_s": a.timeout_s} + ignore = [".git", ".venv", ".claude", "**/__pycache__", "**/*.pyc", "s1-runs", "syncer/target", + "evidence*", "openspec", "docs", "tutorial", "data"] + image = (modal_image(modal).env({"PYTHONUNBUFFERED": "1"}) + .add_local_dir(str(REPO), mv.WORKDIR, copy=False, ignore=ignore)) + app = modal.App(f"yeto-{a.run}") + results = modal.Volume.from_name(mv.RESULTS_VOLUME, create_if_missing=True, version=2) + fn = app.function(image=image, gpu=GPU, cpu=8.0, memory=65536, timeout=a.timeout_s, retries=0, + volumes={mv.RESULTS: results}, single_use_containers=True, name="island", + serialized=True)(island) + a.out_dir.mkdir(parents=True, exist_ok=True) + began = time.time() + with modal.enable_output(), app.run(detach=True): # a killed client never cancels a spawned call + (a.out_dir / "app_id.txt").write_text(f"{app.app_id}\n") + call = fn.spawn(a.run, cfg) + try: + result = {"ok": call.get(timeout=a.timeout_s + 600)} + except Exception as exc: # noqa: BLE001 - recorded; the raw data is pulled anyway + result = {"error": f"{type(exc).__name__}: {exc}"} + result["wall_s"] = round(time.time() - began, 1) + pulled = subprocess.run(["modal", "volume", "get", "--force", mv.RESULTS_VOLUME, a.run, str(a.out_dir)], + capture_output=True, text=True, timeout=1800) + result["pull_rc"] = pulled.returncode + (a.out_dir / "driver.json").write_text(json.dumps(result, indent=1, default=str)) + print(json.dumps(result, default=str)[:2000]) + return 0 if "ok" in result else 1 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/tools/probes/s26_e6/sglang_fp8_cast_overlay.py b/tools/probes/s26_e6/sglang_fp8_cast_overlay.py new file mode 100644 index 000000000..d82f71f2d --- /dev/null +++ b/tools/probes/s26_e6/sglang_fp8_cast_overlay.py @@ -0,0 +1,42 @@ +"""Overlay for the SGLang fork 2fa88018 in the Miles ports image (S26 E6 probe only). + +Qwen2 MLP and Qwen3 attention cast the activation to the projection's weight dtype +when they differ (true-on-policy: fp32 RMSNorm output -> bf16 weights). With an FP8 +projection (online FP8 or a compressed-tensors FP8 checkpoint) the weight dtype is +float8_e4m3fn, so the BF16 activation is cast to FP8 and ``apply_fp8_linear`` then +asserts on a "pre-quantized" input without a scale (e6-gate-t1 / e6-gate-t2, +2026-10-21). The overlay limits the cast to weights of two or more bytes. + +Usage (in the image): ``python sglang_fp8_cast_overlay.py /sgl-workspace/sglang/python``. +Refuses to run when a pattern is not found exactly once (another SGLang commit). +""" + +import sys +from pathlib import Path + +EDITS = { + "sglang/srt/models/qwen2.py": [ + ("or x.dtype != self.gate_up_proj.weight.dtype\n", + "or (x.dtype != self.gate_up_proj.weight.dtype and self.gate_up_proj.weight.dtype.itemsize > 1)\n"), + ], + "sglang/srt/models/qwen3.py": [ + ("or hidden_states.dtype != self.qkv_proj.weight.dtype\n", + "or (hidden_states.dtype != self.qkv_proj.weight.dtype and self.qkv_proj.weight.dtype.itemsize > 1)\n"), + ], +} + + +def main(root: str) -> None: + for rel, edits in EDITS.items(): + path = Path(root) / rel + text = path.read_text() + for old, new in edits: + if text.count(old) != 1: + raise SystemExit(f"{path}: expected one {old!r}, found {text.count(old)}") + text = text.replace(old, new) + path.write_text(text) + print(f"patched {path}") + + +if __name__ == "__main__": + main(sys.argv[1]) diff --git a/tools/probes/s26_e6/verl_hf_overrides_probe.py b/tools/probes/s26_e6/verl_hf_overrides_probe.py new file mode 100644 index 000000000..59e5336cb --- /dev/null +++ b/tools/probes/s26_e6/verl_hf_overrides_probe.py @@ -0,0 +1,114 @@ +"""S26 E6 probe (Modal CPU, verl image): what vLLM receives from the tier 1 verl overrides. + +The overrides come from ``yeto.rl.quant.tiers.vllm_online_overrides`` for the real +Qwen3-4B-Instruct-2507 checkpoint (computed on the client, passed as JSON). In the +pinned verl image (fork acad9875, vLLM 0.29) the container: + +1. composes verl's ``ppo_trainer`` config with Hydra and these overrides; +2. turns the rollout section into verl's ``RolloutConfig`` (``omega_conf_to_dataclass``) + and builds the vLLM argument dict as ``vLLMHttpServer`` does (verl's own FP8 + ``hf_overrides`` first, then ``**engine_kwargs``), then the CLI list with verl's + ``build_cli_args_from_config``; +3. parses that list with vLLM's ``AsyncEngineArgs`` parser and reads + ``hf_overrides`` / ``kv_cache_dtype``; +4. builds vLLM's ``Fp8Config`` from the received ``quantization_config`` and asks + vLLM's ``is_layer_skipped`` for the fused Qwen3 modules. + +No GPU, no model weights. Run: ``modal run tools/probes/s26_e6/verl_hf_overrides_probe.py +--overrides-json FILE --out FILE`` (client venv with modal). +""" + +from __future__ import annotations + +import json +import sys +from pathlib import Path + +import modal + +app = modal.App("yeto-s26-e6-verl-hf-overrides-probe") +if modal.is_local(): # the image is defined on the client; the container only runs probe() + sys.path.insert(0, str(Path(__file__).resolve().parents[3])) + from yeto.rl.adapters.verl.image import modal_image + from yeto.rl.adapters.verl.pins import VERL_ROOT + + image = modal_image(modal) +else: + VERL_ROOT = "/workspace/verl" # yeto.rl.adapters.verl.pins.VERL_ROOT + image = modal.Image.debian_slim() + +CONTAINER = r''' +import json, sys +from hydra import compose, initialize_config_dir +from omegaconf import OmegaConf +overrides = ["actor_rollout_ref.rollout.name=vllm", *json.loads(sys.argv[1])]; layers = int(sys.argv[2]) +with initialize_config_dir(config_dir=VERL_ROOT + "/verl/trainer/config", version_base=None): + cfg = compose(config_name="ppo_trainer", overrides=overrides) +from verl.utils.config import omega_conf_to_dataclass +from verl.workers.config.rollout import RolloutConfig +OmegaConf.resolve(cfg) +section = cfg.actor_rollout_ref.rollout +try: + rollout = omega_conf_to_dataclass(section, dataclass_type=RolloutConfig) + source = "RolloutConfig" +except Exception as error: # the composed-but-unwired config may miss fields the trainer fills + print("RolloutConfig failed: " + repr(error)[:300], file=sys.stderr) + rollout, source = section, "DictConfig" +engine_kwargs = rollout.get("engine_kwargs", {}).get("vllm", {}) or {} +engine_kwargs = {k: v for k, v in engine_kwargs.items() if v is not None} +verl_fp8 = {"activation_scheme": "dynamic", "fmt": "e4m3", "quant_method": "fp8", "weight_block_size": [128, 128], + "ignored_layers": [f"model.layers.{i}.mlp.gate" for i in range(layers)]} +args = {"quantization": rollout.quantization, "hf_overrides": {"quantization_config": verl_fp8}, **engine_kwargs} +from verl.workers.rollout.vllm_rollout.utils import build_cli_args_from_config +cli = build_cli_args_from_config(args) +# CPU container: vLLM cannot infer a device type; DeviceConfig's check only matters for +# an engine, not for parsing arguments, so it is skipped here (parsing is what is probed). +import vllm.config.device as _device +_device.DeviceConfig.__post_init__ = lambda self: None +from vllm.engine.arg_utils import AsyncEngineArgs +from vllm.utils.argparse_utils import FlexibleArgumentParser +parser = AsyncEngineArgs.add_cli_args(FlexibleArgumentParser()) +ns = parser.parse_args(["--model", "/nonexistent", *cli]) +ea = AsyncEngineArgs.from_cli_args(ns) +qc = ea.hf_overrides["quantization_config"] +from vllm.model_executor.layers.quantization.fp8 import Fp8Config +from vllm.model_executor.layers.quantization.utils.quant_utils import is_layer_skipped +fp8 = Fp8Config.from_config(qc) +fused = {"qkv_proj": ["q_proj", "k_proj", "v_proj"], "gate_up_proj": ["gate_proj", "up_proj"]} +probe = {p: bool(is_layer_skipped(p, fp8.ignored_layers, fused)) for p in + ["model.layers.0.self_attn.qkv_proj", "model.layers.0.self_attn.o_proj", + "model.layers.35.self_attn.qkv_proj", "model.layers.0.mlp.gate_up_proj", "model.layers.35.mlp.down_proj"]} +import vllm +print("RESULT " + json.dumps({ + "vllm": vllm.__version__, "rollout_config_source": source, "rollout_quantization": rollout.quantization, + "engine_kwargs_types": {k: type(v).__name__ for k, v in engine_kwargs.items()}, + "cli_flags": [c for c in cli if c.startswith("--")], + "kv_cache_dtype": ea.kv_cache_dtype, "quantization": ea.quantization, + "received_ignored_layers": sorted(fp8.ignored_layers), "weight_block_size": fp8.weight_block_size, + "is_layer_skipped": probe})) +'''.replace("VERL_ROOT", repr(VERL_ROOT)) + + +@app.function(image=image, cpu=2, memory=8192, timeout=900) +def probe(overrides: list[str], layers: int) -> str: + import subprocess + + done = subprocess.run([f"{VERL_ROOT}/.venv/bin/python", "-c", CONTAINER, json.dumps(overrides), str(layers)], + capture_output=True, text=True, cwd=VERL_ROOT) + return json.dumps({"rc": done.returncode, "stdout": done.stdout[-20000:], "stderr": done.stderr[-8000:]}) + + +@app.local_entrypoint() +def main(overrides_json: str, out: str) -> None: + spec = json.loads(Path(overrides_json).read_text()) + raw = json.loads(probe.remote(spec["overrides"], int(spec["num_layers"]))) + Path(out).write_text(json.dumps(raw, indent=1)) # raw first + line = next((l for l in raw["stdout"].splitlines() if l.startswith("RESULT ")), None) + if raw["rc"] != 0 or line is None: + print(raw["stderr"][-3000:]) + raise SystemExit(f"probe failed rc={raw['rc']}") + result = json.loads(line[len("RESULT "):]) + expected = sorted(json.loads(spec["overrides"][2].split("=", 1)[1][1:-1])["quantization_config"]["ignored_layers"]) + result["ignored_layers_equal_sent"] = result["received_ignored_layers"] == expected + Path(out).with_suffix(".result.json").write_text(json.dumps(result, indent=1)) + print(json.dumps({k: v for k, v in result.items() if k != "received_ignored_layers"}, indent=1)) diff --git a/tools/probes/s26_e6/vllm_fp8_probe.py b/tools/probes/s26_e6/vllm_fp8_probe.py new file mode 100644 index 000000000..8cafd6d44 --- /dev/null +++ b/tools/probes/s26_e6/vllm_fp8_probe.py @@ -0,0 +1,113 @@ +"""S26 E6 probe (Modal H100!, verl image, vLLM 0.29): does vLLM generate sensibly with +the tier 1 verl settings, without verl's weight sync? + +e6-t1-verl-a-r2 generated 512-token garbage with NaN log-probs (score 0) on step 1. +This probe starts vLLM ``LLM`` on Qwen3-4B-Instruct-2507 with exactly the +``quantization`` and ``hf_overrides`` yeto passes for tier 1 +(``tiers.vllm_online_overrides``), with and without ``enable_lora``, generates 16 +GSM8K answers at temperature 0 and records lengths, a text sample and whether +any log-prob is NaN. A BF16 engine is the control. + +Run: ``python tools/probes/s26_e6/vllm_fp8_probe.py --out DIR`` (client venv with modal). +""" + +from __future__ import annotations + +import argparse +import json +import subprocess +import sys +from pathlib import Path + +REPO = Path(__file__).resolve().parents[3] +sys.path.insert(0, str(REPO)) + +from tools.eval import modal_verl as mv # noqa: E402 + +MODEL = "Qwen/Qwen3-4B-Instruct-2507" +MODEL_REV = "cdbee75f17c01a7cc42f958dc650907174af0554" + +CONTAINER = r''' +import json, math, sys +from huggingface_hub import snapshot_download +from vllm import LLM, SamplingParams +case = json.loads(sys.argv[1]) +base = snapshot_download(case["model"], revision=case["rev"]) +kw = dict(model=base, max_model_len=1024, gpu_memory_utilization=0.6, enforce_eager=False) +if case["fp8"]: + kw.update(quantization="fp8") + if case.get("hf_overrides"): + kw.update(hf_overrides=case["hf_overrides"]) +if case["lora"]: + kw.update(enable_lora=True, max_loras=1, max_lora_rank=32) +llm = LLM(**kw) +tok = llm.get_tokenizer() +prompts = [tok.apply_chat_template([{"role": "user", "content": q}], tokenize=False, add_generation_prompt=True) + for q in case["questions"]] +outs = llm.generate(prompts, SamplingParams(temperature=0.0, max_tokens=512, logprobs=0)) +rows = [] +for o in outs: + c = o.outputs[0] + lps = [list(d.values())[0].logprob for d in (c.logprobs or [])] + rows.append({"tokens": len(c.token_ids), "nan": sum(1 for x in lps if not math.isfinite(x)), "text": c.text[:200]}) +print("RESULT " + json.dumps(rows)) +''' + + +def probe(cases: list[dict]) -> list[dict]: + out = [] + mv._gpu(Path("/tmp"), {"gpu": "H100!"}) + for case in cases: + proc = subprocess.run([mv.VERL_PY, "-c", CONTAINER, json.dumps(case)], capture_output=True, text=True, + timeout=1500) + line = next((ln for ln in proc.stdout.splitlines() if ln.startswith("RESULT ")), None) + out.append({"name": case["name"], "rc": proc.returncode, + "rows": json.loads(line[7:]) if line else None, "stderr_tail": proc.stderr[-3000:]}) + return out + + +def main() -> int: + p = argparse.ArgumentParser() + p.add_argument("--out", required=True, type=Path) + p.add_argument("--round", type=int, default=1) + a = p.parse_args() + import modal + from modal._vendor import cloudpickle + + from yeto.rl.adapters.verl.image import modal_image + + cloudpickle.register_pickle_by_value(sys.modules[__name__]) + cloudpickle.register_pickle_by_value(mv) + # the exact value of e6-t1-verl-a-r2 (``tiers.vllm_online_overrides``), saved from its argv + hf = json.loads(Path(a.out / "hf_overrides.json").read_text()) + questions = ["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?", + "Weng earns $12 an hour for babysitting. Yesterday, she just did 50 minutes of babysitting. " + "How much did she earn?"] * 8 + common = {"model": MODEL, "rev": MODEL_REV, "questions": questions, "hf_overrides": hf} + verl_native = {"quantization_config": {"activation_scheme": "dynamic", "fmt": "e4m3", "quant_method": "fp8", + "weight_block_size": [128, 128], + "ignored_layers": [f"model.layers.{i}.mlp.gate" for i in range(36)]}} + if a.round == 1: + cases = [{**common, "name": "bf16", "fp8": False, "lora": False}, + {**common, "name": "fp8", "fp8": True, "lora": False}, + {**common, "name": "fp8+lora", "fp8": True, "lora": True}] + else: # round 2: verl's own FP8 config (no yeto ignore list) and vLLM online FP8 (no hf_overrides) + cases = [{**common, "name": "verl-native-config", "fp8": True, "lora": False, "hf_overrides": verl_native}, + {**common, "name": "vllm-online", "fp8": True, "lora": False, "hf_overrides": None}, + {**common, "name": "vllm-online+lora", "fp8": True, "lora": True, "hf_overrides": None}] + image = modal_image(modal).env({"PYTHONUNBUFFERED": "1"}) + app = modal.App("yeto-e6-vllm-fp8-probe") + fn = app.function(image=image, gpu="H100!", cpu=8.0, memory=65536, timeout=3000, retries=0, + name="probe", serialized=True)(probe) + with modal.enable_output(), app.run(detach=True): + (a.out / "app_id.txt").write_text(f"{app.app_id}\n") + result = fn.spawn(cases).get(timeout=3600) + (a.out / f"result-{a.round}.json").write_text(json.dumps(result, indent=1)) + for r in result: + print(r["name"], r["rc"], [(x["tokens"], x["nan"]) for x in (r["rows"] or [])][:4]) + return 0 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/yeto/cloud/modal_ckpt_store.py b/yeto/cloud/modal_ckpt_store.py index f7a4bf3fa..820e9c2b4 100644 --- a/yeto/cloud/modal_ckpt_store.py +++ b/yeto/cloud/modal_ckpt_store.py @@ -11,10 +11,11 @@ from __future__ import annotations +import contextlib import logging import os import time -from collections.abc import Callable, Mapping +from collections.abc import Callable, Iterator, Mapping from typing import Any from yeto.rl.engine.resume import CheckpointStore @@ -24,6 +25,40 @@ STORE_IMPL = "yeto.cloud.modal_ckpt_store:store_from_env" # value of resume.STORE_IMPL_ENV +GUARD_DIR_ENV = "YETO_MODAL_VOLUME_GUARD_DIR" # tests; default /tmp (container-local, not the Volume) + + +@contextlib.contextmanager +def _volume_lock(volume: str, mode: int) -> Iterator[bool]: + """``flock`` on a container-local file per Volume, shared by every process of the + container (the trainer driver, its Ray actors and threads each open their own + descriptor). Yields whether the lock was taken (False only for a non-blocking try). + + A reload remounts the Volume for the whole container and drops local files not + committed yet, so writers hold it shared and a reload tries it exclusive and is + skipped while any writer holds it (reload is best effort). Seen on Modal + (e6-t2-pair-a/-c, 2026-10-21): rollout actors' reloads removed the tier 2 releaser + thread's half-written FP8 release (FileNotFoundError on RELEASE.json.tmp / LATEST.json.tmp).""" + import fcntl + import re + + root = os.environ.get(GUARD_DIR_ENV) or "/tmp" + path = os.path.join(root, "yeto-modal-volume-" + re.sub(r"[^A-Za-z0-9_.-]", "_", volume) + ".lock") + fd = os.open(path, os.O_RDWR | os.O_CREAT, 0o666) + try: + try: + fcntl.flock(fd, mode) + except BlockingIOError: + yield False + return + try: + yield True + finally: + fcntl.flock(fd, fcntl.LOCK_UN) + finally: + os.close(fd) + + class ModalVolumeStore(CheckpointStore): """A Modal Volume (v1) mounted into the island container. Writes become visible to another container only after ``commit()`` (Modal docs: background commits every @@ -72,14 +107,30 @@ def _runner_call(self, verb: str) -> None: commits = 0 + @contextlib.contextmanager + def writing(self) -> Iterator[None]: + import fcntl + + with _volume_lock(self.volume_name, fcntl.LOCK_SH): + yield + def reload(self) -> None: - try: - if self._factory is None and os.environ.get("YETO_MODAL_PYTHON"): - self._runner_call("reload") - else: - self._volume().reload() - except Exception as exc: # noqa: BLE001 - reload is best effort (open files block it) - logging.getLogger(__name__).warning("modal volume reload failed: %r", exc) + import fcntl + + with _volume_lock(self.volume_name, fcntl.LOCK_EX | fcntl.LOCK_NB) as taken: + if not taken: # a reload now would drop a writer's uncommitted files + logging.getLogger(__name__).info("modal volume reload skipped: a writer holds %s", + self.volume_name) + return + try: + if self._factory is None and os.environ.get("YETO_MODAL_PYTHON"): + self._runner_call("reload") + else: + self._volume().reload() + except Exception as exc: # noqa: BLE001 - reload is best effort (open files block it) + logging.getLogger(__name__).warning("modal volume reload failed: %r", exc) + else: # logged so a remount can be lined up with a writer's failure (e6-t2-pair-e) + logging.getLogger(__name__).info("modal volume reloaded: %s", self.volume_name) def describe(self) -> dict[str, Any]: return {**super().describe(), "volume": self.volume_name} diff --git a/yeto/launch/rl_args.py b/yeto/launch/rl_args.py index 0c21669b9..b29fa247b 100644 --- a/yeto/launch/rl_args.py +++ b/yeto/launch/rl_args.py @@ -212,6 +212,9 @@ def _prepare_rl_args( raise ValueError("--codex-reasoning-effort requires --training-mode rl") return + from yeto.rl.quant.tiers import validate_launch as _validate_quant_tier + + _validate_quant_tier(args) # S26 E6: refused before any cloud spend _prepare_node_inventory(args) from ..adapter_lifecycle import directory_sha256, selected_parent @@ -266,10 +269,11 @@ def _prepare_rl_args( raise ValueError("RL v0 supports only causal language models") from ..rl.engine import sample_roles - if args.tuning != "lora" and not sample_roles.buffer_full_parameter(args): + if args.tuning != "lora" and not sample_roles.no_sync_full_parameter(args): raise ValueError("RL v0 requires --tuning lora (--tuning full: only sampler / trainer " - "islands, --rl-sample-role or --rl-island-sample-role, with " - "--rl-single-island-no-sync)") + "islands, --rl-sample-role or --rl-island-sample-role, or one island " + "with dedicated engine cards, --rl-placement fixed-partition; both " + "with --rl-single-island-no-sync)") if args.lora_r <= 0: raise ValueError("RL v0 requires a positive LoRA rank") if args.total_steps <= 0: diff --git a/yeto/launch/rl_island_task.py b/yeto/launch/rl_island_task.py index 1185f4c9a..e5af8234a 100644 --- a/yeto/launch/rl_island_task.py +++ b/yeto/launch/rl_island_task.py @@ -220,6 +220,13 @@ def make_miles_island_task( value = getattr(args, name, None) if value is not None: flags += f" {flag} {shlex.quote(str(value))}" + if getattr(args, "rl_rollout_quant_tier", 0): # S26 E6 (validated by the launcher) + flags += f" --rl-rollout-quant-tier {int(args.rl_rollout_quant_tier)}" + from yeto.rl.quant.tiers import FP8_SOURCE_ENGINE, fp8_source + + source = fp8_source(args) + if source != FP8_SOURCE_ENGINE: + flags += f" --rl-rollout-quant-fp8-source {shlex.quote(source)}" if args.use_rollout_routing_replay: flags += " --use-rollout-routing-replay" if getattr(args, "apply_chat_template_kwargs", None): diff --git a/yeto/rl/adapters/miles/config.py b/yeto/rl/adapters/miles/config.py index dadbb341d..6ae0f8e9a 100644 --- a/yeto/rl/adapters/miles/config.py +++ b/yeto/rl/adapters/miles/config.py @@ -37,6 +37,7 @@ normalize_filter_name, AlgorithmSpec, ) +from yeto.rl.quant.tiers import FP8_SOURCE_ENGINE, TIER_ONLINE_FP8, sglang_online_argv from .lr_schedule import LR_SCHEDULE_FLAGS from .algorithm_flags import ( OBJECTIVE_FLAGS, @@ -427,6 +428,8 @@ def _check_sglang_tp(value, config): "serving.max_running_requests": _ok, "serving.chunked_prefill_size": _ok, "serving.sampler_snapshot": _ok, # gated in translate_run_config (sampler_snapshot) + "serving.rollout_quant_tier": _ok, # S26 E6: SGLang flags in translate_run_config + "serving.rollout_quant_fp8_source": _ok, # S26 E6: tier 1 engine vs trainer (Miles block FP8) "ports.rollout_engine_base_port": _must_be_none( "upstream Miles has no --rollout-engine-base-port (fork-only port isolation, task 1.4b)" ), @@ -1004,6 +1007,12 @@ def translate_run_config( values.extend((flag, str(value))) if serving.deterministic_inference: values.append("--sglang-enable-deterministic-inference") + if getattr(serving, "rollout_quant_tier", 0) == TIER_ONLINE_FP8 and \ + getattr(serving, "rollout_quant_fp8_source", FP8_SOURCE_ENGINE) == FP8_SOURCE_ENGINE: + # S26 E6 tier 1: SGLang quantizes the in-scope linears itself; the ignore list + # travels in the workers' environment (entry.connect_island_ray). The trainer + # source needs no SGLang flag: the FP8 checkpoint's config carries the method. + values.extend(sglang_online_argv()) if getattr(serving, "sampler_snapshot", False): try: values.extend(sampler_snapshot.snapshot_argv(request, extra_argv=extra_argv)) @@ -1174,6 +1183,11 @@ def translate_run_config( rollout_overlap=rollout_overlap, runtime_attrs={ **algorithm.to_legacy_runtime_attrs(), + # S26 E6: only a quantized tier adds the attribute (tier 0 keeps the argv golden) + **({"yeto_rollout_quant_tier": int(serving.rollout_quant_tier)} + if getattr(serving, "rollout_quant_tier", 0) else {}), + **({"yeto_rollout_quant_fp8_source": serving.rollout_quant_fp8_source} + if getattr(serving, "rollout_quant_fp8_source", FP8_SOURCE_ENGINE) != FP8_SOURCE_ENGINE else {}), **_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) diff --git a/yeto/rl/adapters/miles/entry.py b/yeto/rl/adapters/miles/entry.py index f05f8ac4e..18c121841 100644 --- a/yeto/rl/adapters/miles/entry.py +++ b/yeto/rl/adapters/miles/entry.py @@ -22,6 +22,8 @@ import logging from pathlib import Path import os +import subprocess +import sys import time from collections.abc import Awaitable, Callable, Mapping, Sequence from typing import Any @@ -1300,8 +1302,10 @@ def connect_island_ray(*, environ=None, ray_module=None, miles_args=None) -> str """ environ = os.environ if environ is None else environ + quant_env = rollout_quant_env(miles_args, environ) # S26 E6 tier 1: checked before any engine starts address = environ.get("RAY_ADDRESS") if not address: + environ.update(quant_env) # a Ray the driver starts inherits the driver's environment return None if ray_module is None: import ray as ray_module @@ -1361,10 +1365,57 @@ def connect_island_ray(*, environ=None, ray_module=None, miles_args=None) -> str from yeto.rl.harness.codex.preflight import worker_runtime_env env_vars.update(worker_runtime_env(miles_args, environ)) + env_vars.update(quant_env) ray_module.init(address=address, runtime_env={"env_vars": env_vars}) return address +def rollout_quant_env(miles_args: Any, environ: Mapping[str, str] | None = None) -> dict[str, str]: + """S26 E6 tier 1: the SGLang ignore list for ``miles_args.hf_checkpoint``, after the + scope check (:func:`yeto.rl.quant.tiers.online_plan` raises before any engine starts). + Empty for every other tier and for a trainer-only island (no rollout engine). + + With the trainer FP8 source (``yeto_rollout_quant_fp8_source``) the engines start on + Miles' block-FP8 conversion instead (:func:`prepare_native_fp8_base`, which sets + ``miles_args.rollout_hf_checkpoint``); its config carries the ignore list, so no + environment entry is needed.""" + from yeto.rl.quant import tiers + + if miles_args is None or getattr(miles_args, "yeto_rollout_quant_tier", 0) != tiers.TIER_ONLINE_FP8: + return {} + plan = tiers.online_plan(Path(miles_args.hf_checkpoint)) + print(f"[rl] {tiers.describe(tiers.TIER_ONLINE_FP8)}: {json.dumps(plan.summary())}", flush=True) + if getattr(miles_args, "yeto_rollout_quant_fp8_source", tiers.FP8_SOURCE_ENGINE) == tiers.FP8_SOURCE_TRAINER: + environ = os.environ if environ is None else environ + prepare_native_fp8_base(miles_args, plan, Path(environ.get("YETO_RL_WORK_DIR") or Path.home() / "yeto-rl")) + return {} + return tiers.sglang_online_env(plan) + + +def prepare_native_fp8_base(miles_args: Any, plan: Any, work_dir: Path, *, run: Any = subprocess.run) -> Path: + """S26 E6 tier 1, trainer source: convert the BF16 base with Miles' own converter + (``tools/convert_hf_to_fp8.py``, 128x128 blocks) keeping the plan's ignore list in BF16, + check the result against the scope rule, and point the engines at it + (``--rollout-hf-checkpoint``). The trainer keeps ``hf_checkpoint`` (BF16) and Miles + block-quantizes each weight update by the converted config (``quantize_params_fp8``).""" + from yeto.rl.quant import tiers + + if not hasattr(miles_args, "rollout_hf_checkpoint"): + raise RuntimeError("this Miles has no --rollout-hf-checkpoint (needs the S26 E6 fork commit)") + work_dir.mkdir(parents=True, exist_ok=True) + ignore_file, out = work_dir / "fp8-native-ignore.json", work_dir / "fp8-native-base" + ignore_file.write_text(json.dumps(tiers.native_ignore_list(plan))) + converter = Path.home() / "miles" / "tools" / "convert_hf_to_fp8.py" + run([sys.executable, str(converter), "--model-dir", str(miles_args.hf_checkpoint), "--save-dir", str(out), + "--strategy", "block", "--block-size", *map(str, tiers.NATIVE_BLOCK), + "--ignored-layers-file", str(ignore_file)], check=True, stdout=subprocess.DEVNULL) + count = tiers.check_native_base(out) + print(f"[rl] {tiers.describe(tiers.TIER_ONLINE_FP8)} trainer source: {miles_args.hf_checkpoint} -> {out} " + f"({count} FP8 modules)", flush=True) + miles_args.rollout_hf_checkpoint = str(out) + return out + + # E2 plan-v2 ยง0 determinism environment (with Megatron --deterministic-mode); # the table row lives in the core (decoupling 4.9, yeto.rl.engine.determinism). # It reaches every rank through connect_island_ray (A2 follow-up, L-2.3). diff --git a/yeto/rl/adapters/miles/island_entry.py b/yeto/rl/adapters/miles/island_entry.py index 48585224f..d80f2dac9 100644 --- a/yeto/rl/adapters/miles/island_entry.py +++ b/yeto/rl/adapters/miles/island_entry.py @@ -697,11 +697,17 @@ def _full_parameter_identity(model_config_sha256: str) -> tuple[str, str]: def _buffer_island(args) -> bool: - """A no-sync sampler / trainer island (L4.6): it has no outer sync, so full - parameter mode runs without the full-parameter sync preset.""" + """A no-sync island with no outer sync, so full parameter mode runs without the + full-parameter sync preset: a sampler / trainer island (L4.6), or one island + whose dedicated engine cards take Miles' own full-parameter weight update + (S26 E6; colocated is refused).""" from yeto.rl.engine.sample_roles import BOTH, sample_role - return bool(getattr(args, "rl_single_island_no_sync", False)) and sample_role(args) != BOTH + if not getattr(args, "rl_single_island_no_sync", False): + return False + if sample_role(args) != BOTH: + return True + return getattr(args, "rl_placement", "colocated") == "fixed-partition" def _full_parameter_mode(args, sync_preset: str) -> bool: diff --git a/yeto/rl/adapters/miles/pins.py b/yeto/rl/adapters/miles/pins.py index 2bc5533b6..b22cba820 100644 --- a/yeto/rl/adapters/miles/pins.py +++ b/yeto/rl/adapters/miles/pins.py @@ -53,7 +53,9 @@ class MilesRevisionPins(NamedTuple): # import stays light, so RayWorkerManager no longer loads megatron.bridge # (rl-trainer-cold-start-cost L2.2)). 17aeb5cb5 = 386ae3b07 + light Ray actors # stay off megatron_utils (L2.1) + the master port is held until use (S23 defect 9). -MILES_NEXT_COMMIT = "44d41efd0a0428cb12b4ee34422de6f5e1ccb9d9" +# yeto/s26-e6-fp8-ignore (9462de532 = 44d41efd0 + quantize_params_fp8 and +# convert_hf_to_fp8 honor the FP8 ignore list; --rollout-hf-checkpoint; S26 E6). +MILES_NEXT_COMMIT = "9462de532f747cf688e090bbaf6b24f80e378fe8" # sgl-project/sglang ``sglang-miles`` head when radixark/miles@9e4260d was # committed (upstream's Dockerfile follows that branch unpinned). SGLANG_NEXT_REPOSITORY = "https://github.com/michaellchung/sglang" @@ -63,7 +65,11 @@ class MilesRevisionPins(NamedTuple): # m3-qwen4exp-lora (4e4148f1b = a1240c530 + Qwen4ExpForConditionalGeneration LoRA hooks). # n17-qwen3coder-single-call (2fa880182 = 4e4148f1b + qwen3_coder honors # parallel_tool_calls=false under tool_choice=auto: at most one tool call; S17 N17). -SGLANG_NEXT_COMMIT = "2fa880182eefdb31e64d8ea70317be27f9f32121" +# yeto/fp8-act-cast (10afe11a7 = 2fa880182 + Qwen2 MLP / Qwen3 attention leave the +# activation alone in front of an FP8 projection; S26 E6). +# yeto/fp8-act-cast (754605964 = 10afe11a7 + compressed-tensors TENSOR/CHANNEL FP8 linears +# reload in place through update_weights_from_disk; S26 E6 tier 2 sampler). +SGLANG_NEXT_COMMIT = "754605964776322226a073763a6f636aa8233932" # Fork lineage on michaellchung/miles yeto/ports, oldest first; the single # source for every "this fork commit carries X" set (rl-code-quality-cleanup # 0.7). Each entry is the previous one plus the change named here; the newest @@ -108,6 +114,8 @@ class MilesRevisionPins(NamedTuple): "17aeb5cb51635b9ba6ef4ca10db95762695c7e03", # 44d41efd0 = + megatron_utils runtime patches moved out of the package init (S24 L2.2); no loss path. "44d41efd0a0428cb12b4ee34422de6f5e1ccb9d9", + # 9462de532 = + FP8 ignore list in quantize_params_fp8, --rollout-hf-checkpoint (S26 E6); no loss path. + "9462de532f747cf688e090bbaf6b24f80e378fe8", ) @@ -146,7 +154,13 @@ def miles_fork_lineage_from(commit: str) -> frozenset[str]: # every SHA. PUBLIC on ghcr.io/michaellchung since 2026-10-07 (anonymous # pull); a private image needs SKYPILOT_DOCKER_{USERNAME,PASSWORD,SERVER} or # --rl-image-private (yeto.launcher.registry_login_for; read:packages token). -# Tag 44d41ef-2fa8801; (S24 L2.2: Miles s24-miles-import2 and the lazy-modelopt +# Tag 9462de5-7546059; (S26 E6: SGLang compressed-tensors FP8 in-place reload; +# previous 9462de5-10afe11 @sha256:6e0a97d2...; build record +# openspec/changes/rl-infra-spec/evidence/ports-image/2026-10-21-9462de5-7546059/). +# Earlier: tag 9462de5-10afe11; (S26 E6: Miles FP8 ignore list + SGLang FP8 activation +# cast fix; previous 44d41ef-2fa8801 @sha256:5b9dc2a9...; build record +# openspec/changes/rl-infra-spec/evidence/ports-image/2026-10-21-9462de5-10afe11/). +# Earlier: tag 44d41ef-2fa8801; (S24 L2.2: Miles s24-miles-import2 and the lazy-modelopt # Megatron-LM patch in the layer; previous 17aeb5c-2fa8801 @sha256:b304f857...; # build record openspec/changes/rl-infra-spec/evidence/ports-image/ # 2026-10-14-44d41ef-2fa8801/). Earlier: tag 17aeb5c-2fa8801; (S24 L2.1-L2.4: Miles s24-miles-import, peft 0.20.0 and the @@ -155,7 +169,7 @@ def miles_fork_lineage_from(commit: str) -> frozenset[str]: # ports-image/2026-10-14-17aeb5c-2fa8801/). Earlier: tag 386ae3b-2fa8801; (S23 P0-1.11 GPU run: Miles s23-sao-streaming-roles; previous 64b591a-2fa8801 @sha256:fa2413be...; build record openspec/changes/rl-infra-spec/evidence/ports-image/2026-10-14-386ae3b-2fa8801/). Earlier: (S19 #1: critic family + algo-supplement metrics baked in; previous ddce209-2fa8801 @sha256:9c252c38...) (S18 ARU-3: Miles s18-agentic-suspend = 2f7871fb2 + --agentic-suspend-between-turns; previous 2f7871f-2fa8801 @sha256:62b4f164..., efbbc63-2fa8801 @sha256:a7990076..., 8bc5223-4e4148f @sha256:4aeafd77...) ghcr.io/michaellchung (public package); linux/amd64 only. MILES_NEXT_IMAGE = ( "docker:ghcr.io/michaellchung/yeto-miles-ports@sha256:" - "5b9dc2a9e6e315d73dabfbc6faae428b59cc01aabec6189645a5c8c80d6e3721" + "37acaca3e1eee8adc61bf3edbbd706ece36e40e9b7fd5463e54f56e3f4cc9b2b" ) MILES_NEXT_IMAGE_MANIFEST = "/opt/yeto/image-manifest.json" # Nebius VM images whose /var/lib/docker already holds a docker image's diff --git a/yeto/rl/adapters/miles/sampler_island.py b/yeto/rl/adapters/miles/sampler_island.py index 22da4c7d9..4ad097849 100644 --- a/yeto/rl/adapters/miles/sampler_island.py +++ b/yeto/rl/adapters/miles/sampler_island.py @@ -48,10 +48,14 @@ PolicyFollower, PolicyRelease, PolicyReleaseStore, + ReleaseNotLoaded, ) from yeto.rl.engine.merged_release import MergedReleaseGuard, is_publisher, merged_version from yeto.rl.engine.policy_versions import engine_weight_version from yeto.rl.engine.sample_roles import CONVERTIBLE, SAMPLER, TRAINER +from yeto.rl.quant.schemes import BF16, FP8_FFN_CHANNEL, release_dir_for +from yeto.rl.quant.scope import QuantScopeError +from yeto.rl.quant.tiers import TIER_TRAINER_FP8, check_release_scope, describe, release_scheme from .buffer_rollout import ( SampleBufferRolloutError, @@ -86,10 +90,48 @@ def _buffer_config(miles_args: Any, role: str) -> Any: return config if config is not None and config.role == role else None -def release_store(config: Any, environ: Any = None) -> PolicyReleaseStore: +def release_store(config: Any, environ: Any = None, *, scheme: str = BF16) -> PolicyReleaseStore: from yeto.rl.engine.resume import store_for - return PolicyReleaseStore(store_for(config.store, environ=environ)) + if scheme == BF16: + return PolicyReleaseStore(store_for(config.store, environ=environ)) + return PolicyReleaseStore(store_for(config.store, environ=environ), release_dir=release_dir_for(scheme)) + + +def _quant_tier(miles_args: Any) -> int: + return int(getattr(miles_args, "yeto_rollout_quant_tier", 0) or 0) + + +def prepare_quantized_base(miles_args: Any, work_dir: Path) -> Path: + """S26 E6 tier 2, sampler side: the engines start on the FP8 conversion of the BF16 + base (``miles_args.hf_checkpoint``), so later FP8 releases load into FP8 parameters + of the same layout. The conversion is checked against the scope rule. Returns the + converted directory, which replaces ``miles_args.hf_checkpoint``.""" + from yeto.rl.quant.awq import AwqSmoothing + from yeto.rl.quant.requant import requantize_checkpoint, save_scoped_template + + base = Path(miles_args.hf_checkpoint) + template, out = work_dir / "fp8-template", work_dir / "fp8-base" + save_scoped_template(base, template, FP8_FFN_CHANNEL) + stats = requantize_checkpoint(base, template, out, AwqSmoothing.none()) + check_release_scope(out) + print(f"[rl] {describe(TIER_TRAINER_FP8)}: base {base} -> {out} {stats}", flush=True) + miles_args.hf_checkpoint = str(out) + return out + + +def scope_checked(load: Callable[[PolicyRelease, Path], None]) -> Callable[[PolicyRelease, Path], None]: + """A ``PolicyFollower`` load hook that refuses an FP8 release outside the scope rule + before the engines are touched (S26 E6 tier 2 post-load check, step one).""" + def hook(release: PolicyRelease, directory: Path) -> None: + try: + check_release_scope(Path(directory)) + except QuantScopeError as error: + raise ReleaseNotLoaded(f"r{release.seq}-v{release.version}: {error}", + reason="release_refused:quant_scope") from error + load(release, directory) + + return hook # ---------------------------------------------------------------------- sampler side @@ -310,6 +352,10 @@ def emit(event: str, **fields: Any) -> None: incarnation, resume, emit) progress = {"rollout_id": None, "rounds": 0} # the round in flight, rounds done here + if _quant_tier(miles_args) == TIER_TRAINER_FP8: + if lora_mode(miles_args): + raise SampleBufferRolloutError("rollout quant tier 2 needs full-parameter releases") + prepare_quantized_base(miles_args, Path(env.get("YETO_RL_WORK_DIR") or Path.home() / "yeto-rl")) connect_island_ray(miles_args=miles_args) # The rollout process's metadata hook pushes each round's group records to this # named actor (on a trainer island the driver creates it and consumes them). @@ -355,8 +401,10 @@ def leave_reclaimed() -> None: loader = EngineReleaseLoader(controller, executor, runner.run, lora=lora) if not lora: # LoRA: no adapter is served before the first release (require_release) loader.stamp(BASE_POLICY_VERSION) # until the first release is loaded - follower = PolicyFollower(release_store(config, environ), loader, - version=BASE_POLICY_VERSION) + scheme = release_scheme(_quant_tier(miles_args)) + follower = PolicyFollower(release_store(config, environ, scheme=scheme), + loader if scheme == BF16 else scope_checked(loader), + version=BASE_POLICY_VERSION, scheme=scheme) loop = SamplerLoop( follower=follower, buffer=buffer, emit=emit, generate=lambda rollout_id: runner.run(generate(rollout_id)), @@ -554,4 +602,23 @@ def export(version: int, directory: Path) -> None: read_shards=lambda: in_executor(executor_read_shards), emit=driver.emit, weights=LORA if lora else FULL, **synced) driver.publisher = publisher + if _quant_tier(miles_args) == TIER_TRAINER_FP8: + start_trainer_releaser(miles_args, config, environ) return publisher + + +def start_trainer_releaser(miles_args: Any, config: Any, environ: Any = None) -> Any: + """S26 E6 tier 2, trainer side: a thread of this driver converts every BF16 release + into an FP8 release (``yeto.rl.quant.release.scoped_releaser``) for the sampler islands.""" + import threading + + from yeto.rl.engine.resume import store_for + from yeto.rl.quant.release import scoped_releaser, start_scoped_releaser + + if lora_mode(miles_args): + raise SampleBufferRolloutError("rollout quant tier 2 needs full-parameter releases") + env = os.environ if environ is None else environ + template = Path(env.get("YETO_RL_WORK_DIR") or Path.home() / "yeto-rl") / "fp8-template" + releaser = scoped_releaser(store_for(config.store, environ=environ), Path(miles_args.hf_checkpoint), + template, FP8_FFN_CHANNEL) + return start_scoped_releaser(releaser, threading.Event()) diff --git a/yeto/rl/adapters/sampler_engines/quant_tier.py b/yeto/rl/adapters/sampler_engines/quant_tier.py new file mode 100644 index 000000000..5355fdac7 --- /dev/null +++ b/yeto/rl/adapters/sampler_engines/quant_tier.py @@ -0,0 +1,94 @@ +"""Rollout quantization tier of an engine sampler island (S26 E6; adapter layer). + +The launch layer writes ``"rollout_quant": {"tier": N, "apply": APPLY}`` into the +sampler spec (:func:`spec_section`); the engine sampler entry calls ``APPLY`` on +the sampler node before it builds anything (the checkpoint is only there). The +function returns a new spec with the engine settings of the tier filled in: + +* tier 1, ``sglang-server`` only: ``--quantization fp8`` (KV cache BF16) and ``SGLANG_FP8_IGNORED_LAYERS`` from the checked online plan of + ``base_dir`` (:func:`yeto.rl.quant.tiers.online_plan`). ``vllm-server`` is + refused: stand-alone vLLM online FP8 has no ignore list, so it would quantize + attention (verl passes one through its own FP8 path instead); +* tier 2, both adapters: ``base_dir`` is converted to the FP8 scheme in + ``work_dir`` (the engine must start on the layout later releases load into); + ``quant_scheme`` / ``expected_format`` / ``release_dir`` follow the FP8 + scheme; ``checks.expected_expert_matrices`` is the routed-expert count of the + layout (0 for a dense model, which the post-load check accepts as "no routed + experts" once the release index also counts 0). + +A spec that already sets a quantized scheme, an engine ``--quantization`` flag +or another expected expert count is refused (one source of truth per setting). +""" + +from __future__ import annotations + +import copy +from collections.abc import Mapping +from pathlib import Path +from typing import Any + +from yeto.rl.quant import tiers +from yeto.rl.quant.schemes import BF16, FP8_FFN_CHANNEL, release_dir_for + +APPLY = "yeto.rl.adapters.sampler_engines.quant_tier:apply" +FP8_FORMAT = "float-quantized" +SGLANG = "sglang-server" +VLLM = "vllm-server" + + +class QuantTierSpecError(ValueError): + """The sampler spec cannot carry the requested tier.""" + + +def spec_section(tier: int) -> dict[str, Any] | None: + """What the launch layer writes into a sampler spec (None for tier 0).""" + if tier == tiers.TIER_BF16: + return None + if tier not in (tiers.TIER_ONLINE_FP8, tiers.TIER_TRAINER_FP8): + raise QuantTierSpecError(f"{tiers.FLAG} {tier}: engine sampler islands run tier 0, 1 or 2") + return {"tier": int(tier), "apply": APPLY} + + +def apply(spec: Mapping[str, Any], *, work_dir: Path | None = None) -> dict[str, Any]: + """The spec with the tier's engine settings (module doc).""" + tier = int((spec.get("rollout_quant") or {}).get("tier", tiers.TIER_BF16)) + out = copy.deepcopy(dict(spec)) + if tier == tiers.TIER_BF16: + return out + engine = out["engine"] + server = engine["server"] + adapter = str(engine.get("adapter")) + if server.get("quant_scheme", BF16) != BF16 or "--quantization" in (server.get("extra_args") or ()): + raise QuantTierSpecError(f"{tiers.FLAG} {tier}: the spec already sets quant_scheme " + f"{server.get('quant_scheme')!r} / extra_args {server.get('extra_args')}") + base = Path(server["base_dir"]) + if tier == tiers.TIER_ONLINE_FP8: + if adapter != SGLANG: + raise QuantTierSpecError(f"{tiers.FLAG} 1 on {adapter}: stand-alone vLLM online FP8 has no ignore " + "list and would quantize attention; use sglang-server or tier 2") + plan = tiers.online_plan(base) + server["extra_args"] = [*(server.get("extra_args") or ()), "--quantization", "fp8"] + server["extra_env"] = {**(server.get("extra_env") or {}), **tiers.sglang_online_env(plan)} + server["precision"] = "fp8" # provenance of the samples; the release scheme stays bf16 + out["rollout_quant"] = {**out["rollout_quant"], "plan": plan.summary()} + return out + from yeto.rl.quant.awq import AwqSmoothing # noqa: PLC0415 - torch only for tier 2 + from yeto.rl.quant.release import routed_expert_matrices # noqa: PLC0415 + from yeto.rl.quant.requant import requantize_checkpoint, save_scoped_template # noqa: PLC0415 + + work = Path(work_dir) if work_dir is not None else Path(str(out.get("event_tape", "."))).expanduser().parent + template, converted = work / "fp8-template", work / "fp8-base" + layout = save_scoped_template(base, template, FP8_FFN_CHANNEL) + requantize_checkpoint(base, template, converted, AwqSmoothing.none()) + tiers.check_release_scope(converted) + experts = routed_expert_matrices(layout.quantized_modules) + checks = dict(out.get("checks") or {}) + if "expected_expert_matrices" in checks and checks["expected_expert_matrices"] != experts: + raise QuantTierSpecError(f"checks.expected_expert_matrices {checks['expected_expert_matrices']} but the " + f"FP8 layout has {experts} routed-expert matrices") + checks["expected_expert_matrices"] = experts + server.update(base_dir=str(converted), quant_scheme=FP8_FFN_CHANNEL, expected_format=FP8_FORMAT, + precision="fp8") + out.update(base_dir=str(converted), release_dir=release_dir_for(FP8_FFN_CHANNEL), checks=checks) + out["rollout_quant"] = {**out["rollout_quant"], "bf16_base": str(base), "routed_expert_matrices": experts} + return out diff --git a/yeto/rl/adapters/sampler_engines/sglang_server.py b/yeto/rl/adapters/sampler_engines/sglang_server.py index cdf57ddb1..dfb4f8aae 100644 --- a/yeto/rl/adapters/sampler_engines/sglang_server.py +++ b/yeto/rl/adapters/sampler_engines/sglang_server.py @@ -237,14 +237,19 @@ def _completions(items: Sequence[dict[str, Any]], prompts: int, n: int) -> list[ out.append(group) return out - def _input_logprobs(self, text: str) -> list[float | None]: + def _input_token_logprobs(self, prompt: dict[str, Any]) -> list[tuple[float | None, int]]: + """``(logprob, token id)`` of every input token; ``prompt`` is ``{"text": ...}`` or + ``{"input_ids": [...]}``.""" body = self.call("POST", "/generate", { - "text": text, "sampling_params": {"max_new_tokens": 0, "temperature": 0.0}, + **prompt, "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]) + return [(None if t[0] is None else float(t[0]), int(t[1])) for t in (body.get("meta_info") or {}).get("input_token_logprobs") or []] + def _input_logprobs(self, text: str) -> list[float | None]: + return [lp for lp, _ in self._input_token_logprobs({"text": text})] + def score(self, prompt: str, continuation: str) -> list[float]: prompt_len = len(self._input_logprobs(prompt)) full = self._input_logprobs(prompt + continuation) @@ -252,3 +257,13 @@ def score(self, prompt: str, continuation: str) -> list[float]: 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] + + def score_tokens(self, prompt: str, token_ids: Sequence[int]) -> list[float]: + """Teacher-forced log-probs of exactly ``token_ids`` after ``prompt`` (no re-tokenizing + of the continuation, so a trailing EOS that has no text is scored too).""" + prompt_ids = [tid for _, tid in self._input_token_logprobs({"text": prompt})] + full = self._input_token_logprobs({"input_ids": [*prompt_ids, *map(int, token_ids)]}) + tail = [lp for lp, _ in full[len(prompt_ids):]] + if len(tail) != len(token_ids) or any(x is None for x in tail): + raise SamplerEngineError(f"sglang scored {len(tail)} of {len(token_ids)} continuation tokens") + return [float(x) for x in tail] diff --git a/yeto/rl/adapters/verl/config.py b/yeto/rl/adapters/verl/config.py index 8c0b6a815..8505c2396 100644 --- a/yeto/rl/adapters/verl/config.py +++ b/yeto/rl/adapters/verl/config.py @@ -64,6 +64,9 @@ class VerlRunConfig: lora_merge: bool = False engine_commit: str = VERL_COMMIT gpu_memory_utilization: float = 0.5 + # vLLM context length; None = the model's max_position_embeddings (262144 for + # Qwen3-4B-Instruct-2507, whose KV cache then does not fit next to the actor) + max_model_len: int | None = None n_gpus: int = 1 model_dtype: str = "fp32" reward_path: str = "" @@ -161,6 +164,7 @@ def build_overrides(cfg: VerlRunConfig) -> list[str]: "actor_rollout_ref.rollout.name=vllm", "actor_rollout_ref.rollout.tensor_model_parallel_size=1", f"actor_rollout_ref.rollout.gpu_memory_utilization={cfg.gpu_memory_utilization}", + *([f"actor_rollout_ref.rollout.max_model_len={cfg.max_model_len}"] if cfg.max_model_len else []), f"actor_rollout_ref.rollout.n={cfg.samples_per_group}", f"actor_rollout_ref.rollout.temperature={cfg.temperature}", f"actor_rollout_ref.rollout.top_p={cfg.top_p}", diff --git a/yeto/rl/adapters/verl/island_entry.py b/yeto/rl/adapters/verl/island_entry.py index ab6f640f8..5b1125f72 100644 --- a/yeto/rl/adapters/verl/island_entry.py +++ b/yeto/rl/adapters/verl/island_entry.py @@ -23,10 +23,13 @@ import shlex import subprocess import sys +import threading import time from pathlib import Path from types import SimpleNamespace +from yeto.rl.quant.tiers import TIER_TRAINER_FP8 + from . import config as vconf from . import run_config_rules from .pins import VERL_COMMIT, VERL_ROOT, expected_versions @@ -83,6 +86,8 @@ def parse_args(argv): p.add_argument("--rl-cut-keep", type=int, default=None) p.add_argument("--rl-stop-after-rounds", type=int, default=None) p.add_argument("--rl-resume-allow-config-change", action="store_true") + # S26 E6: rollout quantization tier (yeto.rl.quant.tiers; the launcher validated it) + p.add_argument("--rl-rollout-quant-tier", type=int, choices=(0, 1, 2, 3), default=0) args, ignored = p.parse_known_args(argv) return args, ignored @@ -318,6 +323,7 @@ def main(argv=None) -> int: # fully_async_plan (lower > 0 warns there, D3); the sync path keeps its check. tis_lower=0.0 if limit else algo.get("tis_lower", 0.0), tis_upper=algo["tis_upper"], gpu_memory_utilization=float(os.environ.get("YETO_VERL_GPU_MEM", "0.5")), + max_model_len=max_prompt + args.rollout_max_response_len, reward_path=str(Path(__file__).with_name("reward_fn.py")), reward_name=function_for(args.reward_function), chat_template_kwargs=chat, experiment_name=f"yeto-verl-{lid}") @@ -340,6 +346,9 @@ def main(argv=None) -> int: except ValueError as exc: raise SystemExit(str(exc)) from exc asserted_keys = vconf.ASSERTED_KEYS + quant_overrides, quant_summary = rollout_quant_overrides( + args.rl_rollout_quant_tier, Path(model_path), colocated=buffer_section is None) + overrides += quant_overrides if limit: # agentic-rollout-utilization 6.4b: fully_async path overrides, asserted_keys = fully_async_plan(limit, overrides, spec_dict, args, sync) append_correction_record(args, spec_dict) @@ -367,6 +376,7 @@ def main(argv=None) -> int: "resume": resume, "versions": manifest["versions"], **({"sample_buffer": buffer_section} if buffer_section else {}), + "rollout_quant": quant_summary, } plan_path = work / f"plan-{lid}.json" plan_path.write_text(json.dumps(plan, indent=1, default=str)) @@ -386,17 +396,53 @@ def main(argv=None) -> int: cmd = [sys.executable, "-c", f"from yeto.rl.adapters.verl.verl_main import {entry}; {entry}()", *overrides] print("[yeto-verl] " + " ".join(shlex.quote(c) for c in cmd), flush=True) log = OUTPUT / f"verl-train-{lid}.log" - with log.open("a") as handle: - proc = subprocess.Popen(cmd, env=env, cwd=VERL_ROOT if Path(VERL_ROOT).is_dir() else None, - stdout=subprocess.PIPE, stderr=subprocess.STDOUT, text=True, bufsize=1) - for line in proc.stdout: - sys.stdout.write(line) - handle.write(line) - code = proc.wait() + stop_releaser = threading.Event() + if args.rl_rollout_quant_tier == TIER_TRAINER_FP8: # S26 E6 tier 2: BF16 releases -> FP8 releases + start_trainer_releaser(buffer_section, Path(model_path), work / "fp8-template", stop_releaser) + try: + with log.open("a") as handle: + proc = subprocess.Popen(cmd, env=env, cwd=VERL_ROOT if Path(VERL_ROOT).is_dir() else None, + stdout=subprocess.PIPE, stderr=subprocess.STDOUT, text=True, bufsize=1) + for line in proc.stdout: + sys.stdout.write(line) + handle.write(line) + code = proc.wait() + finally: + stop_releaser.set() print(f"[yeto-verl] island {lid} exit {code} after {time.time() - t0:.0f}s", flush=True) return code +def rollout_quant_overrides(tier: int, model_path: Path, *, colocated: bool) -> tuple[list[str], dict]: + """S26 E6: rollout quantization on a verl island; returns the vLLM overrides (none) and + the plan record. Tier 1 on a colocated island is refused before verl starts + (``tiers.VERL_TIER1_REFUSAL``: LoRA syncs never re-quantize the base). On a buffer + trainer the rollout runs on the sampler island, so nothing is added; tier 2 starts + the releaser instead. ``model_path`` is kept for the tier 1 plan once a full-parameter + verl path exists (``tiers.vllm_online_overrides``).""" + from yeto.rl.quant import tiers + + summary: dict = {"tier": tier, "name": tiers.TIER_NAMES[tier]} + if tier == tiers.TIER_ONLINE_FP8 and colocated: + raise SystemExit(f"[yeto-verl] {tiers.VERL_TIER1_REFUSAL}") + return [], summary + + +def start_trainer_releaser(buffer_section: dict | None, base_dir: Path, template_dir: Path, + stop: threading.Event) -> threading.Thread: + """S26 E6 tier 2 on a buffer trainer: convert every BF16 release of the store into an + FP8 release (``yeto.rl.quant.release.scoped_releaser``) in a thread of this process.""" + from yeto.rl.engine.resume import store_for + from yeto.rl.quant.release import scoped_releaser, start_scoped_releaser + from yeto.rl.quant.schemes import FP8_FFN_CHANNEL + + if not buffer_section: + raise SystemExit("--rl-rollout-quant-tier 2 needs --rl-sample-role trainer on a verl island") + store = store_for(buffer_section.get("release_store") or buffer_section["store"]) + releaser = scoped_releaser(store, base_dir, template_dir, FP8_FFN_CHANNEL) + return start_scoped_releaser(releaser, stop) + + def sample_buffer_section(args) -> dict | None: """The plan's ``sample_buffer`` section of a buffer trainer (S25 C2); None for role ``both``. A verl island is never a sampler: a sampler island runs the diff --git a/yeto/rl/engine/engine_sampler_entry.py b/yeto/rl/engine/engine_sampler_entry.py index f11c5307c..892992cec 100644 --- a/yeto/rl/engine/engine_sampler_entry.py +++ b/yeto/rl/engine/engine_sampler_entry.py @@ -31,6 +31,8 @@ ``reward_kind`` is ``"miles"``; * ``checks``: :class:`~.sampler_checks.PostLoadChecks` keyword overrides (optional; e.g. ``expected_expert_matrices``); +* ``rollout_quant``: optional ``{"tier", "apply"}`` (S26 E6), applied first + (:func:`apply_rollout_quant`); * ``event_tape``: the JSONL tape file of the island's events; each incarnation writes its own file next to it (``events..jsonl``), so a replacement container never overwrites an earlier one's tape. @@ -99,6 +101,19 @@ def build_engine(spec: Mapping[str, Any]) -> SamplerEngine: raise SamplerSpecError(f"engine: {error}") from None +def apply_rollout_quant(spec: Mapping[str, Any]) -> Mapping[str, Any]: + """S26 E6: ``spec["rollout_quant"] = {"tier", "apply": "module:function"}`` (written by + the launch layer) fills the tier's engine settings on this node; the adapter-layer + function returns the new spec. No section: the spec is used as written.""" + section = spec.get("rollout_quant") + if not section: + return spec + try: + return _import(_require(section, "apply", "rollout_quant"))(spec) + except ValueError as error: + raise SamplerSpecError(f"rollout_quant: {error}") from None + + # ------------------------------------------------------------------ prompts class JsonlPrompts: @@ -271,6 +286,7 @@ def run(spec: Mapping[str, Any], *, environ: Mapping[str, str] | None = None, ) env = os.environ if environ is None else environ + spec = apply_rollout_quant(spec) config = island_config(spec) incarnation = incarnation or new_incarnation() tape = incarnation_tape(_require(spec, "event_tape"), incarnation) diff --git a/yeto/rl/engine/policy_release.py b/yeto/rl/engine/policy_release.py index 45aeef17d..16355686f 100644 --- a/yeto/rl/engine/policy_release.py +++ b/yeto/rl/engine/policy_release.py @@ -297,6 +297,15 @@ def release(self, version: int, policy_hash: str, write: Callable[[Path], None], if latest is not None and (latest.version, latest.policy_hash) == (version, policy_hash): return None seq = 1 if latest is None else latest.seq + 1 + with self.store.writing(): # no reload drops the files before they are committed + return self._write_release(seq, version, policy_hash, write, buffer_selected=buffer_selected, + weights=weights, engine_checksum=engine_checksum, + buffer_shards=buffer_shards, derived_from=derived_from, origin=origin) + + def _write_release(self, seq: int, version: int, policy_hash: str, write: Callable[[Path], None], *, + buffer_selected: int | None, weights: str, engine_checksum: str | None, + buffer_shards: dict[str, int] | None, derived_from: dict[str, Any] | None, + origin: dict[str, Any] | None) -> PolicyRelease: name = f"r{seq}-v{version}" directory = self.root / name if directory.exists(): # a partial release of an earlier attempt (no LATEST names it) diff --git a/yeto/rl/engine/resume.py b/yeto/rl/engine/resume.py index bfa3894e9..0f06893fb 100644 --- a/yeto/rl/engine/resume.py +++ b/yeto/rl/engine/resume.py @@ -30,6 +30,7 @@ from __future__ import annotations +import contextlib import hashlib import json import logging @@ -145,6 +146,12 @@ def commit(self) -> float: def reload(self) -> None: return None + def writing(self) -> contextlib.AbstractContextManager[None]: + """Held while a writer builds files it has not committed yet. A store whose + reload can drop uncommitted local writes (a Modal Volume) does not reload + while any writer of the same volume holds it.""" + return contextlib.nullcontext() + def describe(self) -> dict[str, Any]: return {"kind": self.kind, "root": str(self.root)} diff --git a/yeto/rl/engine/round_cut.py b/yeto/rl/engine/round_cut.py index cadc5047e..ded1b97a7 100644 --- a/yeto/rl/engine/round_cut.py +++ b/yeto/rl/engine/round_cut.py @@ -239,8 +239,6 @@ def save(self, driver: Any, *, rollout_id: int, final: bool = False, a stop the trainer scheduler ordered); a periodic cut takes none.""" import time - from yeto.rl.engine.resume import directory_bytes - if trigger is not None and (not final or trigger not in FINAL_TRIGGERS): raise RoundCutError(f"cut trigger {trigger!r} needs final=True and one of {FINAL_TRIGGERS}") ctl = self.controller @@ -259,6 +257,23 @@ def save(self, driver: Any, *, rollout_id: int, final: bool = False, # colocated + offloaded trainer (1 GPU): its weights/optimizer are not on the # device between publish and the next train step; wake it for the cut only. offloaded = bool(getattr(driver, "_trainer_offloaded", False)) + # Hold the store's writer guard from the first shard to the commit: a Modal + # Volume reload in between remounts the Volume and drops uncommitted files. + # Seen on Modal (e6-t2-pair-e, 2026-10-21): a reload during the shard + # snapshot made the rank's makedirs of the store dir fail with EPERM; the + # cell died and the island exited 134. + with ctl.store.writing(): + return self._save_guarded(driver, cut_id=cut_id, context=context, rollout_id=rollout_id, + final=final, trigger=trigger, started=started, offloaded=offloaded) + + def _save_guarded(self, driver: Any, *, cut_id: str, context: Any, rollout_id: int, final: bool, + trigger: str | None, started: float, offloaded: bool) -> dict[str, Any]: + """:meth:`save` under the store's writer guard: the trainer cut, the pointer, the store sync.""" + import time + + from yeto.rl.engine.resume import directory_bytes + + ctl = self.controller if offloaded: self.trainer.onload() try: diff --git a/yeto/rl/engine/run_config.py b/yeto/rl/engine/run_config.py index 616e588a5..d979b4658 100644 --- a/yeto/rl/engine/run_config.py +++ b/yeto/rl/engine/run_config.py @@ -420,6 +420,10 @@ class ServingConfig: # rl-trainer-cold-start-cost L8.2: opt-in SGLang snapshot flags on the # sampler (rollout) engines; refused without a recorded compatibility pass. sampler_snapshot: bool = False + # S26 E6: rollout quantization tier (yeto.rl.quant.tiers; 0 = BF16) + rollout_quant_tier: int = 0 + # tier 1: who quantizes ("engine" online FP8, or "trainer": Miles block FP8 updates) + rollout_quant_fp8_source: str = "engine" @dataclass(frozen=True) @@ -1033,6 +1037,8 @@ def resolve_rl_run_config( max_running_requests=getattr(args, "sglang_max_running_requests", None), chunked_prefill_size=getattr(args, "sglang_chunked_prefill_size", None), sampler_snapshot=bool(getattr(args, "rl_sampler_snapshot", False)), + rollout_quant_tier=int(getattr(args, "rl_rollout_quant_tier", 0) or 0), + rollout_quant_fp8_source=_fp8_source(args), ), ports=NetworkPorts( rollout_engine_base_port=getattr(args, "rollout_engine_base_port", None), @@ -1080,6 +1086,12 @@ def resolve_rl_run_config( ) +def _fp8_source(args) -> str: + from ..quant.tiers import fp8_source + + return fp8_source(args) + + def _train_memory(args) -> TrainMemoryConfig: max_tokens = getattr(args, "rl_max_tokens_per_gpu", None) if max_tokens is not None and int(max_tokens) < int(args.seq_len): diff --git a/yeto/rl/engine/run_flags.py b/yeto/rl/engine/run_flags.py index b4de91487..de0537c78 100644 --- a/yeto/rl/engine/run_flags.py +++ b/yeto/rl/engine/run_flags.py @@ -793,6 +793,19 @@ def rl_forward_pairs(*dests: str | tuple[str, str]) -> tuple[tuple[str, str], .. shared_flag("--use-rollout-routing-replay", action="store_true"), shared_flag("--rollout-max-response-len", type=int, default=32768, island={"default": None, "required": True}), + # S26 E6: quantized rollout tiers (yeto.rl.quant.tiers) + shared_flag("--rl-rollout-quant-tier", type=int, choices=(0, 1, 2, 3), default=0, + help="sampler precision: 0 BF16 (default); 1 the engine quantizes FFN linears " + "(dense) or routed experts (MoE) to FP8 itself, FP8 KV cache; 2 the trainer side " + "converts each release to FP8 W8A8 on the same linears and the sampler island " + "loads it (needs --rl-island-sample-role and --tuning full); 3 4-bit, interface " + "only (refused). Router, lm_head, embeddings, norms and attention stay BF16; the " + "KV cache stays BF16. Needs FP8 cards (H100, H200, L4)"), + shared_flag("--rl-rollout-quant-fp8-source", choices=("engine", "trainer"), default=None, + help="tier 1 only: who quantizes (default: trainer where the backend supports it, else engine). " + "trainer: the engines start on a 128x128 block-FP8 conversion of the base " + "and the trainer block-quantizes every weight update. engine: the rollout engine " + "quantizes at load (online FP8). Same ignore list"), ) _AGENT = ( diff --git a/yeto/rl/engine/sample_roles.py b/yeto/rl/engine/sample_roles.py index de65b83f3..a86ef2fb9 100644 --- a/yeto/rl/engine/sample_roles.py +++ b/yeto/rl/engine/sample_roles.py @@ -248,6 +248,24 @@ def buffer_full_parameter(args: Any) -> bool: and (is_pair(args) or sample_role(args) != BOTH)) +def dedicated_engine_full_parameter(args: Any) -> bool: + """``--tuning full`` on one no-sync island with its own engine cards + (``--rl-placement fixed-partition``, role both): the backend's in-island + full-parameter weight update feeds the engines every round (S26 E6: Miles + block-quantizes each update for the tier-1 trainer FP8 source). A colocated + full-parameter island is not supported.""" + return (getattr(args, "tuning", "lora") == "full" + and bool(getattr(args, "rl_single_island_no_sync", False)) + and not is_pair(args) and sample_role(args) == BOTH + and getattr(args, "rl_placement", "colocated") == "fixed-partition") + + +def no_sync_full_parameter(args: Any) -> bool: + """Full-parameter training without the full-parameter sync: a buffer island or + one island with dedicated engine cards.""" + return buffer_full_parameter(args) or dedicated_engine_full_parameter(args) + + def check_pair(args: Any, num_islands: int) -> None: """Launch checks of a store-coupled pair (nothing outside a pair): every island has a role, exactly one trainer and at least one sampler, diff --git a/yeto/rl/engine/sampler_checks.py b/yeto/rl/engine/sampler_checks.py index 4e0f26d7c..08ac22600 100644 --- a/yeto/rl/engine/sampler_checks.py +++ b/yeto/rl/engine/sampler_checks.py @@ -11,6 +11,8 @@ 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); + an FP8 release (``fp8-w8a8-channel``, S26 E6 tier 2) must quantize only the + in-scope linears (:func:`yeto.rl.quant.tiers.check_release_scope`); 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`: @@ -52,6 +54,8 @@ from pathlib import Path from typing import Any +from yeto.rl.quant.schemes import FP8_FFN_CHANNEL + from .policy_release import PolicyRelease, ReleaseNotLoaded from .sampler_engine import SamplerEngine, SamplerEngineError, SamplingParams @@ -173,6 +177,11 @@ def expert_count(self, engine: SamplerEngine, expected: int | None) -> dict[str, if loaded is not None: return {"name": "expert_count", "passed": int(loaded) == int(expected), "loaded": int(loaded), "expected": int(expected)} + if int(expected) == 0: + # S26 E6: a dense model has no routed experts; the release index count (0) was + # compared before the load, so there is nothing an engine report could add. + return {"name": "expert_count", "passed": True, "expected": 0, + "detail": "no routed experts (dense model); release index count checked"} gib = engine.weight_gib_per_card() want = self.expected_weight_gib_per_card if gib is not None and want is not None: @@ -324,6 +333,15 @@ def _serve(self, release: PolicyRelease | None, directory: Path | None, version: if in_release is not None and in_release != expected: self._refuse(version, "release_expert_count", {"in_release": in_release, "expected": expected}) + if directory is not None and self.scheme == FP8_FFN_CHANNEL: + # S26 E6 tier 2: only in-scope linears may be FP8 (yeto.rl.quant.scope) + from yeto.rl.quant.scope import QuantScopeError + from yeto.rl.quant.tiers import check_release_scope + + try: + check_release_scope(directory) + except QuantScopeError as error: + self._refuse(version, "quant_scope", str(error)) 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() diff --git a/yeto/rl/quant/__init__.py b/yeto/rl/quant/__init__.py index 5addfcd12..c50139f5f 100644 --- a/yeto/rl/quant/__init__.py +++ b/yeto/rl/quant/__init__.py @@ -33,6 +33,18 @@ * :mod:`.shards`: per-rank quantization into part files and their merge (R3). * :mod:`.release`: the side process that turns BF16 releases into quantized ones. +S26 E6 (quantized rollout tiers) adds: + +* :mod:`.scope`: which linears a quantized rollout may quantize (dense FFN, MoE + routed experts; router, lm_head, embeddings, norms, attention stay BF16); +* :mod:`.fp8`: per-channel FP8 E4M3 quantize / dequantize (W8A8, tier 2); +* :mod:`.tiers`: the ``--rl-rollout-quant-tier`` setting, its launch checks and + the SGLang / vLLM online FP8 plans (tier 1); +* :mod:`.gate`: the logprob / TIS / speed gate of a tier against BF16. + +The FP8 scheme derives its target layout from the BF16 checkpoint and the +scope rule (:func:`.requant.scoped_layout`), not from a reference artefact. + 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/fp8.py b/yeto/rl/quant/fp8.py new file mode 100644 index 000000000..42e8098fa --- /dev/null +++ b/yeto/rl/quant/fp8.py @@ -0,0 +1,56 @@ +"""Symmetric per-output-channel FP8 E4M3 weight quantization (S26 E6, tier 2). + +Target format: compressed-tensors ``float-quantized`` with weights +``num_bits=8, type=float, symmetric=True, strategy=channel`` and dynamic +per-token FP8 activations (llm-compressor's ``FP8_DYNAMIC`` preset, W8A8). +vLLM and SGLang both serve it on Hopper (``CompressedTensorsW8A8Fp8``). + +The rule reproduces compressed-tensors 0.18 (``calculate_qparams`` and +``_quantize``) for an FP32 scale: + +* ``scale = amax(|w|, row) / 448`` in FP32 (448 is the E4M3 maximum; the + symmetric range is ``(448 - (-448)) / 2``); a zero row gets the FP32 + machine epsilon, so dequantization stays finite; +* ``q = clamp(w / scale, -448, 448)`` cast to ``float8_e4m3fn`` (round to + nearest even, done by the cast). + +``tests/test_rl_quant_fp8.py`` compares this module with compressed-tensors +itself where it is installed (the ``sglang-next-venv`` has 0.18.0). + +Each output row is quantized on its own, so any row slice of a weight gives +the same rows as the whole weight (per-rank quantization, :mod:`.shards`). +""" + +from __future__ import annotations + +import torch + +FP8_DTYPE = torch.float8_e4m3fn +FP8_MAX = 448.0 # torch.finfo(torch.float8_e4m3fn).max + + +def quantize_fp8_per_channel( + weight: torch.Tensor, scale_dtype: torch.dtype = torch.float32 +) -> tuple[torch.Tensor, torch.Tensor]: + """Quantize a 2-D weight ``[out, in]``; returns ``(q float8_e4m3fn [out, in], scale [out, 1])``.""" + if weight.dim() != 2: + raise ValueError(f"per-channel FP8 needs a 2-D weight, got shape {tuple(weight.shape)}") + if not weight.is_floating_point(): + raise TypeError(f"per-channel FP8 needs a floating-point weight, got {weight.dtype}") + w = weight.to(torch.float32) + if not torch.isfinite(w).all(): + raise ValueError("weight holds NaN or infinity; refusing to quantize it") + scale = w.abs().amax(dim=1, keepdim=True) / FP8_MAX + scale = torch.where(scale == 0, torch.tensor(torch.finfo(torch.float32).eps), scale) + scale = scale.to(scale_dtype) + q = torch.clamp(w / scale.to(torch.float32), -FP8_MAX, FP8_MAX).to(FP8_DTYPE) + return q, scale + + +def dequantize_fp8_per_channel( + q: torch.Tensor, scale: torch.Tensor, dtype: torch.dtype = torch.float32 +) -> torch.Tensor: + """``q * scale`` in ``dtype`` (the reference of what an engine computes with).""" + if q.dtype != FP8_DTYPE: + raise TypeError(f"expected {FP8_DTYPE}, got {q.dtype}") + return (q.to(torch.float32) * scale.to(torch.float32)).to(dtype) diff --git a/yeto/rl/quant/gate.py b/yeto/rl/quant/gate.py new file mode 100644 index 000000000..d27a786f1 --- /dev/null +++ b/yeto/rl/quant/gate.py @@ -0,0 +1,270 @@ +"""Gate between rollout quantization tiers (S26 E6): a quantized sampler against BF16. + +A tier may replace the previous one only when, on one fixed prompt set, its +samples stay close to what the BF16 policy would assign them. For each +completion the quantized engine generates (the behaviour policy of the rollout), +the BF16 engine scores the same text teacher-forced (the policy the trainer +holds). Per token: + +* ``diff = logprob_quant - logprob_bf16`` (signed mean, and the mean and + percentiles p50 / p90 / p99 / max of ``|diff|``); +* the TIS weight ``exp(logprob_bf16 - logprob_quant)`` and the share of tokens + it truncates (above ``tis_upper``, or below ``tis_lower`` when that is > 0), + the same bounds as the run's TIS; +* generation speed: output tokens per second of each engine on the same prompts + (wall clock of one ``generate`` call over the whole set) and their ratio. + +The criteria (:class:`GateCriteria`) are read from a JSON file, so each step +(tier 0 -> 1, 0 -> 2, ...) states its own; a criterion left ``None`` is +reported but not applied. The report is written to disk before the verdict is +printed; the exit code is 0 only when every applied criterion passes. + +Two ways in: + +* ``collect``: two engine sections (``yeto.rl.adapters.sampler_engines.from_spec``) + and a prompt file; the engines run one after the other (one GPU is enough); +* ``records``: per-completion records that already hold both log-prob lists + (e.g. ``rollout_log_probs`` and the trainer's ``log_probs`` of an RL round). + +A completion whose BF16 score has a different token count than the generated +completion (the text tokenized differently) is not compared; the share of such +completions is a criterion of its own (``max_unaligned_share``). + +Status: CPU tests with the fake engine (tests/test_rl_quant_gate.py); not run +on a GPU. +""" + +from __future__ import annotations + +import argparse +import json +import math +import os +import sys +import time +from collections.abc import Iterable, Mapping, Sequence +from dataclasses import asdict, dataclass, fields +from pathlib import Path +from typing import Any + +SCHEMA = "yeto.rl.quant.gate/v1" + + +@dataclass(frozen=True) +class GateCriteria: + """Pass rules of one tier step; ``None`` = reported only.""" + + tis_upper: float = 2.0 + tis_lower: float = 0.0 + max_mean_abs_diff: float | None = None + max_p99_abs_diff: float | None = None + max_tis_clip_share: float | None = None + min_speedup: float | None = None # quant tokens/s / bf16 tokens/s + max_unaligned_share: float | None = 0.05 + + @classmethod + def from_mapping(cls, raw: Mapping[str, Any]) -> "GateCriteria": + known = {f.name for f in fields(cls)} + if unknown := set(raw) - known: + raise ValueError(f"unknown gate criteria {sorted(unknown)}; known: {sorted(known)}") + criteria = cls(**raw) + if not criteria.tis_upper > 1.0 or not 0.0 <= criteria.tis_lower < 1.0: + raise ValueError(f"TIS bounds need 0 <= lower < 1 < upper, got {criteria.tis_lower}, {criteria.tis_upper}") + return criteria + + +@dataclass(frozen=True) +class Pair: + """One completion: the quantized engine's log-probs and the BF16 scores of the same tokens.""" + + quant: Sequence[float] + bf16: Sequence[float] + + +def quantile(values: Sequence[float], q: float) -> float: + """Linear-interpolated quantile of ``values`` (sorted here), ``0 <= q <= 1``.""" + if not values: + return math.nan + ordered = sorted(values) + pos = q * (len(ordered) - 1) + low = math.floor(pos) + high = min(low + 1, len(ordered) - 1) + return ordered[low] + (ordered[high] - ordered[low]) * (pos - low) + + +def compare(pairs: Iterable[Pair], criteria: GateCriteria) -> dict[str, Any]: + """Token statistics over the aligned pairs (module doc).""" + diffs: list[float] = [] + clipped_high = clipped_low = completions = unaligned = 0 + for pair in pairs: + completions += 1 + if len(pair.quant) != len(pair.bf16): + unaligned += 1 + continue + for lq, lb in zip(pair.quant, pair.bf16): + lq, lb = float(lq), float(lb) + if not (math.isfinite(lq) and math.isfinite(lb)): + raise ValueError(f"non-finite log-prob (quant {lq}, bf16 {lb})") + diffs.append(lq - lb) + log_ratio = lb - lq + if log_ratio > math.log(criteria.tis_upper): + clipped_high += 1 + elif criteria.tis_lower > 0 and log_ratio < math.log(criteria.tis_lower): + clipped_low += 1 + absolute = [abs(d) for d in diffs] + tokens = len(diffs) + return { + "completions": completions, + "unaligned_completions": unaligned, + "unaligned_share": unaligned / completions if completions else math.nan, + "tokens": tokens, + "mean_diff": sum(diffs) / tokens if tokens else math.nan, + "mean_abs_diff": sum(absolute) / tokens if tokens else math.nan, + "p50_abs_diff": quantile(absolute, 0.50), + "p90_abs_diff": quantile(absolute, 0.90), + "p99_abs_diff": quantile(absolute, 0.99), + "max_abs_diff": max(absolute) if absolute else math.nan, + "tis_clip_high_share": clipped_high / tokens if tokens else math.nan, + "tis_clip_low_share": clipped_low / tokens if tokens else math.nan, + "tis_clip_share": (clipped_high + clipped_low) / tokens if tokens else math.nan, + } + + +def verdict(stats: Mapping[str, Any], speed: Mapping[str, Any] | None, criteria: GateCriteria) -> dict[str, Any]: + """Each applied criterion with its value and result; ``passed`` is their conjunction. + A criterion whose value is missing or NaN fails (never a silent pass).""" + checks: dict[str, dict[str, Any]] = {} + + def at_most(name: str, value: Any, limit: float | None) -> None: + if limit is not None: + ok = isinstance(value, (int, float)) and math.isfinite(value) and value <= limit + checks[name] = {"value": value, "limit": limit, "passed": bool(ok)} + + at_most("max_mean_abs_diff", stats.get("mean_abs_diff"), criteria.max_mean_abs_diff) + at_most("max_p99_abs_diff", stats.get("p99_abs_diff"), criteria.max_p99_abs_diff) + at_most("max_tis_clip_share", stats.get("tis_clip_share"), criteria.max_tis_clip_share) + at_most("max_unaligned_share", stats.get("unaligned_share"), criteria.max_unaligned_share) + if criteria.min_speedup is not None: + value = None if speed is None else speed.get("speedup") + ok = isinstance(value, (int, float)) and math.isfinite(value) and value >= criteria.min_speedup + checks["min_speedup"] = {"value": value, "limit": criteria.min_speedup, "passed": bool(ok)} + if not stats.get("tokens"): + checks["tokens"] = {"value": stats.get("tokens"), "limit": "> 0", "passed": False} + return {"passed": all(c["passed"] for c in checks.values()), "checks": checks} + + +def report(pairs: Sequence[Pair], criteria: GateCriteria, *, speed: Mapping[str, Any] | None = None, + meta: Mapping[str, Any] | None = None) -> dict[str, Any]: + stats = compare(pairs, criteria) + return {"schema": SCHEMA, "criteria": asdict(criteria), "stats": stats, "speed": dict(speed or {}) or None, + "verdict": verdict(stats, speed, criteria), "meta": dict(meta or {})} + + +# ---------------------------------------------------------------- collection over engines +def _tokens_per_s(completions: Sequence[Sequence[Any]], seconds: float) -> dict[str, Any]: + tokens = sum(len(c.token_ids) for group in completions for c in group) + return {"output_tokens": tokens, "seconds": seconds, "tokens_per_s": tokens / seconds if seconds > 0 else math.nan} + + +def collect(quant: Any, bf16: Any, prompts: Sequence[str], sampling: Any, + clock: Any = time.monotonic) -> tuple[list[Pair], dict[str, Any], list[dict[str, Any]]]: + """Run the quantized engine, then the BF16 engine, on ``prompts`` (module doc). + + Returns the pairs, the speed record and one raw record per completion (to save + before anything is computed from them). Each engine is started on its base weights + (``start(None, 0)``) and stopped afterwards, so one GPU holds one engine at a time. + """ + quant.start(None, 0) + try: + t0 = clock() + generated = quant.generate(list(prompts), sampling) + quant_speed = _tokens_per_s(generated, clock() - t0) + finally: + quant.stop() + bf16.start(None, 0) + try: + t0 = clock() + bf16_generated = bf16.generate(list(prompts), sampling) + bf16_speed = _tokens_per_s(bf16_generated, clock() - t0) + pairs: list[Pair] = [] + raw: list[dict[str, Any]] = [] + for prompt_index, (prompt, group) in enumerate(zip(prompts, generated)): + for sample_index, completion in enumerate(group): + # by token ids when the engine can (a trailing EOS has no text to re-tokenize) + scores = (list(bf16.score_tokens(prompt, completion.token_ids)) if hasattr(bf16, "score_tokens") + else list(bf16.score(prompt, completion.text))) + pairs.append(Pair(tuple(completion.logprobs), tuple(scores))) + raw.append({"prompt_index": prompt_index, "sample": sample_index, + "token_ids": list(completion.token_ids), "quant_logprobs": list(completion.logprobs), + "bf16_logprobs": scores, "finish_reason": completion.finish_reason}) + finally: + bf16.stop() + speedup = (quant_speed["tokens_per_s"] / bf16_speed["tokens_per_s"] + if bf16_speed["tokens_per_s"] and math.isfinite(bf16_speed["tokens_per_s"]) else math.nan) + speed = {"quant": quant_speed, "bf16": bf16_speed, "speedup": speedup} + return pairs, speed, raw + + +def _write_json(path: Path, payload: Any) -> None: + path.parent.mkdir(parents=True, exist_ok=True) + tmp = path.with_name(path.name + ".tmp") + tmp.write_text(json.dumps(payload, indent=1, sort_keys=True, default=str), encoding="utf-8") + os.replace(tmp, path) + + +def _read_records(path: Path, quant_key: str, bf16_key: str) -> list[Pair]: + pairs = [] + for line in Path(path).read_text(encoding="utf-8").splitlines(): + if line.strip(): + row = json.loads(line) + pairs.append(Pair(tuple(row[quant_key]), tuple(row[bf16_key]))) + return pairs + + +def _load_criteria(path: str | None) -> GateCriteria: + return GateCriteria.from_mapping(json.loads(Path(path).read_text())) if path else GateCriteria() + + +def main(argv: list[str] | None = None) -> int: + p = argparse.ArgumentParser(prog="python -m yeto.rl.quant.gate", description=__doc__.splitlines()[0]) + sub = p.add_subparsers(dest="cmd", required=True) + c = sub.add_parser("collect", help="run both engines on a prompt file and gate") + c.add_argument("--engines", required=True, help='JSON file {"quant": , "bf16": }') + c.add_argument("--prompts", required=True, help="JSONL with a 'prompt' field per line") + c.add_argument("--max-tokens", type=int, required=True) + c.add_argument("--temperature", type=float, default=1.0) + c.add_argument("--n", type=int, default=1) + c.add_argument("--seed", type=int, default=1) + r = sub.add_parser("records", help="gate per-completion records that hold both log-prob lists") + r.add_argument("--records", required=True) + r.add_argument("--quant-key", default="rollout_log_probs") + r.add_argument("--bf16-key", default="log_probs") + for q in (c, r): + q.add_argument("--criteria", help="JSON file of GateCriteria fields") + q.add_argument("--out", required=True, help="report JSON (raw records next to it for collect)") + q.add_argument("--label", default="", help="free text stored in the report (e.g. 'tier0->tier1 Qwen3-4B')") + a = p.parse_args(argv) + criteria = _load_criteria(a.criteria) + out = Path(a.out) + meta: dict[str, Any] = {"label": a.label, "cmd": a.cmd, "created_unix": time.time()} + if a.cmd == "collect": + from yeto.rl.adapters.sampler_engines.from_spec import build + from yeto.rl.engine.sampler_engine import SamplingParams + + engines = json.loads(Path(a.engines).read_text()) + prompts = [json.loads(line)["prompt"] for line in Path(a.prompts).read_text().splitlines() if line.strip()] + sampling = SamplingParams(max_tokens=a.max_tokens, temperature=a.temperature, n=a.n, seed=a.seed).check() + pairs, speed, raw = collect(build(engines["quant"]), build(engines["bf16"]), prompts, sampling) + _write_json(out.with_name(out.stem + "-raw.json"), raw) # raw data first + meta.update(prompts=len(prompts), sampling=sampling.to_dict(), engines=engines) + result = report(pairs, criteria, speed=speed, meta=meta) + else: + meta["records"] = a.records + result = report(_read_records(Path(a.records), a.quant_key, a.bf16_key), criteria, meta=meta) + _write_json(out, result) + print(json.dumps({"out": str(out), "verdict": result["verdict"], "stats": result["stats"]}, indent=1, default=str)) + return 0 if result["verdict"]["passed"] else 1 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/yeto/rl/quant/release.py b/yeto/rl/quant/release.py index 82f686c55..541d64c9e 100644 --- a/yeto/rl/quant/release.py +++ b/yeto/rl/quant/release.py @@ -143,6 +143,57 @@ def run(source: Path, out: Path) -> Any: return run +def scoped_releaser(store: Any, base_dir: Path, template_dir: Path, scheme: str, *, keep: int = 2, + log: Callable[[str], None] = lambda m: print(m, flush=True)) -> RequantReleaser: + """A releaser for a scoped scheme (S26 E6 tier 2: FP8 on the in-scope linears only). + + The template is derived from the BF16 base checkpoint ``base_dir`` + (:func:`~.requant.save_scoped_template`) into ``template_dir``, checked against + the scope rule, and every BF16 release of ``store`` is converted with it. Both + trainer backends start it beside their BF16 releases (:func:`start_scoped_releaser`). + """ + from .requant import check_layout_scope, save_scoped_template + from .scope import model_kind_of + + layout = save_scoped_template(Path(base_dir), Path(template_dir), scheme) + check_layout_scope(layout, model_kind_of(Path(base_dir))) + return RequantReleaser(PolicyReleaseStore(store, keep=keep, release_dir=release_dir_for(BF16)), + PolicyReleaseStore(store, keep=keep, release_dir=release_dir_for(scheme)), + requantizer(Path(template_dir), _no_smoothing()), scheme=scheme, + expert_matrices=routed_expert_matrices(layout.quantized_modules), log=log) + + +def routed_expert_matrices(modules: Any) -> int: + """Per-expert routed matrices among ``modules``: the count the sampler's release + check compares with the index (``sampler_checks.EXPERT_WEIGHT_PATTERN``); 0 for a dense model.""" + import re + + return sum(1 for m in modules if re.search(r"\.experts\.\d+\.\w+_proj$", m)) + + +def start_scoped_releaser(releaser: RequantReleaser, stop: threading.Event, + poll_s: float = 10.0) -> threading.Thread: + """Run ``releaser`` in a daemon thread until ``stop`` is set. A failed conversion + ends the thread with the error printed; the sampler then serves no newer version + and its policy age guard stops the run (never a silent BF16 fallback).""" + def body() -> None: + try: + releaser.run(stop, poll_s) + except BaseException as error: # noqa: BLE001 - reported, the thread ends + releaser.log(f"[requant] scoped releaser stopped: {type(error).__name__}: {error}") + raise + + thread = threading.Thread(target=body, name=f"requant-{releaser.scheme}", daemon=True) + thread.start() + return thread + + +def _no_smoothing() -> Any: + from .awq import AwqSmoothing + + return AwqSmoothing.none() + + def main(argv: list[str] | None = None) -> int: from yeto.rl.engine.resume import store_for diff --git a/yeto/rl/quant/requant.py b/yeto/rl/quant/requant.py index 31b61582d..dc1195d9d 100644 --- a/yeto/rl/quant/requant.py +++ b/yeto/rl/quant/requant.py @@ -57,9 +57,12 @@ from safetensors import safe_open from safetensors.torch import save_file +from . import scope as quant_scope from .awq import AwqSmoothing +from .fp8 import FP8_DTYPE, quantize_fp8_per_channel from .int4 import QMAX, QMIN, quantize_per_channel from .mxfp4 import GROUP_SIZE, quantize_mxfp4, unpack_fp4 +from .schemes import FP8_FFN_CHANNEL as FP8_NAME from .schemes import MXFP4_G32_RTN as MXFP4_NAME from .schemes import W4A8_INT4_CHANNEL as W4A8_NAME from .stacked import StackedSourceIndex @@ -80,6 +83,7 @@ "I32": torch.int32, "I64": torch.int64, "U8": torch.uint8, + "F8_E4M3": FP8_DTYPE, "BOOL": torch.bool, } _DTYPE_NAMES = {v: k for k, v in _DTYPES.items()} @@ -170,9 +174,36 @@ def quantize(self, weight: torch.Tensor, scale: TensorSpec) -> tuple[torch.Tenso return quantize_mxfp4(weight) +class _Fp8Channel(QuantScheme): + """FP8 E4M3 weights per output channel, dynamic per-token FP8 activations (W8A8, S26 E6 tier 2).""" + + name = FP8_NAME + rtn_only = True + format = "float-quantized" + data_suffix = WEIGHT_SUFFIX + data_dtype = "F8_E4M3" + weight_scheme = {"num_bits": 8, "type": "float", "symmetric": True, "strategy": "channel", + "group_size": None, "dynamic": False} + + def check_pair(self, module: str, data: TensorSpec, scale: TensorSpec) -> None: + if scale.shape != (data.shape[0], 1) or scale.dtype not in ("F32", "BF16"): + raise LayoutError(f"{module}: scale {scale.dtype}{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_fp8_per_channel(weight, scale_dtype=_DTYPES[scale.dtype]) + + def check_values(self, name: str, data: torch.Tensor) -> None: + if torch.isnan(data.to(torch.float32)).any(): + raise LayoutError(f"{name}: FP8 weight holds NaN") + + W4A8_INT4_CHANNEL = _Int4Channel() MXFP4_G32_RTN = _Mxfp4Group32() -SCHEMES = {s.format: s for s in (W4A8_INT4_CHANNEL, MXFP4_G32_RTN)} +FP8_FFN_CHANNEL = _Fp8Channel() +SCHEMES = {s.format: s for s in (W4A8_INT4_CHANNEL, MXFP4_G32_RTN, FP8_FFN_CHANNEL)} +# Schemes whose target layout is derived from the BF16 checkpoint and the scope +# rule (:func:`scoped_layout`) instead of a reference artefact. +SCOPED_SCHEMES = {FP8_FFN_CHANNEL.name: FP8_FFN_CHANNEL} @dataclass @@ -299,6 +330,86 @@ def read_header(path: Path) -> dict[str, dict[str, Any]]: return header +def _fp8_quantization_config(ignore: list[str]) -> dict[str, Any]: + """The compressed-tensors config of :data:`FP8_FFN_CHANNEL` (llm-compressor ``FP8_DYNAMIC``).""" + scheme = FP8_FFN_CHANNEL + return { + "quant_method": "compressed-tensors", + "format": scheme.format, + "quantization_status": "compressed", + "config_groups": {"group_0": { + "targets": ["Linear"], + "format": scheme.format, + "weights": {**scheme.weight_scheme, "observer": "minmax"}, + "input_activations": {"num_bits": 8, "type": "float", "symmetric": True, "strategy": "token", + "group_size": None, "dynamic": True, "observer": None}, + "output_activations": None, + }}, + "ignore": ignore, + "kv_cache_scheme": None, + } + + +def scoped_layout(source_dir: Path, scheme_name: str) -> tuple[CheckpointLayout, str]: + """The target layout of a scoped scheme, derived from the BF16 checkpoint ``source_dir``. + + Every in-scope linear (:mod:`.scope`: FFN of a dense model, routed experts of a MoE + model) becomes ``.weight`` in the scheme's data dtype plus an FP32 + ``.weight_scale [out, 1]`` in the same shard; every other tensor keeps its + name, dtype, shape and shard. The config ``ignore`` lists every out-of-scope linear by + name, and ``lm_head``. Returns the layout and the model kind. + """ + scheme = SCOPED_SCHEMES.get(scheme_name) + if scheme is None: + raise LayoutError(f"{scheme_name!r} is not a scoped scheme; known: {', '.join(SCOPED_SCHEMES)}") + source = _SourceReader(Path(source_dir)) + if source.stacked.stacked_names: + raise LayoutError(f"{source_dir}: a scoped layout is derived from a per-expert HF checkpoint, " + "not from the stacked trainer layout") + headers: dict[str, dict[str, Any]] = {} + for shard in sorted(set(source.weight_map.values())): + headers.update(read_header(Path(source_dir) / shard)) + kind = quant_scope.model_kind_of(Path(source_dir)) + shapes = {n: headers[n]["shape"] for n in source.weight_map} + quantized = [m for m in quant_scope.linear_modules(shapes) if quant_scope.in_scope(m, kind)] + if not quantized: + raise LayoutError(f"{source_dir}: no in-scope linear module for a {kind} model") + quant_scope.check_quantized(quantized, kind, what=f"the {scheme.name} layout") + tensors: dict[str, TensorSpec] = {} + for name, shard in source.weight_map.items(): + tensors[name] = TensorSpec(headers[name]["dtype"], tuple(headers[name]["shape"]), shard) + for module in quantized: + data = tensors[module + WEIGHT_SUFFIX] + tensors[module + WEIGHT_SUFFIX] = TensorSpec(scheme.data_dtype, data.shape, data.shard) + if module + SCALE_SUFFIX in tensors: + raise LayoutError(f"{module}: the BF16 source already has a weight_scale") + tensors[module + SCALE_SUFFIX] = TensorSpec("F32", (data.shape[0], 1), data.shard) + # ``lm_head`` is listed even when tied to the embeddings and absent from the + # checkpoint (llm-compressor does the same): no loader may treat it as quantized. + config = _fp8_quantization_config( + quant_scope.with_fused_aliases([*quant_scope.ignored_modules(shapes, kind), "lm_head"])) + layout = CheckpointLayout(tensors, config) + layout.validate() + return layout, kind + + +def save_scoped_template(source_dir: Path, template_dir: Path, scheme_name: str) -> CheckpointLayout: + """Write the template of a scoped scheme: ``layout.json`` plus the source's metadata files, + with ``config.json`` carrying the scheme's ``quantization_config``.""" + layout, _kind = scoped_layout(source_dir, scheme_name) + layout.save_template(Path(source_dir), Path(template_dir)) + config_path = Path(template_dir) / "config.json" + config = json.loads(config_path.read_text()) + config["quantization_config"] = layout.quantization_config + _write_text_atomic(config_path, json.dumps(config, indent=2, sort_keys=True)) + return layout + + +def check_layout_scope(layout: CheckpointLayout, kind: str) -> None: + """Refuse a layout that quantizes a module outside the scope rule (:mod:`.scope`).""" + quant_scope.check_quantized(layout.quantized_modules, kind, what=f"the {layout.scheme.name} layout") + + def metadata_files(ckpt_dir: Path) -> list[Path]: """Non-weight files of a checkpoint directory (config, recipe, tokenizer, ...); the index is rebuilt.""" return sorted( @@ -419,6 +530,10 @@ def requantize_checkpoint( alias = {n: n.replace(_COMPOSITE_PREFIX, _TEXT_PREFIX, 1) for n in missing if n.startswith(_COMPOSITE_PREFIX) and n.replace(_COMPOSITE_PREFIX, _TEXT_PREFIX, 1) in extra} missing, extra = missing - set(alias), extra - set(alias.values()) + # A trainer may export the tied output head that the HF base omits (tie_word_embeddings): + # the layout has no lm_head, the engine ties it to the embeddings, so it is dropped. + tied = {n for n in extra if n.endswith("lm_head.weight")} if _ties_embeddings(template_dir) else set() + extra -= tied base, filled = None, set() if frozen is not None and missing: matcher = re.compile(frozen[1]) @@ -441,7 +556,7 @@ def requantize_checkpoint( out_dir = Path(out_dir) out_dir.mkdir(parents=True, exist_ok=True) stats = {"shards": 0, "tensors": 0, "quantized": 0, "frozen_from_base": len(filled), - "renamed_from_text_prefix": len(alias)} + "renamed_from_text_prefix": len(alias), "dropped_tied_lm_head": len(tied)} def read(name: str) -> torch.Tensor: return base.get(name) if name in filled else source.get(alias.get(name, name)) @@ -469,6 +584,14 @@ def read(name: str) -> torch.Tensor: return stats +def _ties_embeddings(template_dir: Path) -> bool: + config = Path(template_dir) / "config.json" + if not config.is_file(): + return False + doc = json.loads(config.read_text()) + return bool(doc.get("tie_word_embeddings") or (doc.get("text_config") or {}).get("tie_word_embeddings")) + + 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. @@ -595,6 +718,10 @@ def main(argv: list[str] | None = None) -> int: 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") + st = sub.add_parser("scoped-template", help="derive the template of a scoped scheme (FP8 FFN) from a BF16 checkpoint") + st.add_argument("--source", required=True) + st.add_argument("--out", required=True) + st.add_argument("--scheme", default=FP8_NAME, choices=sorted(SCOPED_SCHEMES)) 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) @@ -615,6 +742,10 @@ def main(argv: list[str] | None = None) -> int: if args.cmd == "template": CheckpointLayout.from_checkpoint(Path(args.checkpoint)).save_template(Path(args.checkpoint), Path(args.template_dir)) result: Any = {"template": args.template_dir} + elif args.cmd == "scoped-template": + layout = save_scoped_template(Path(args.source), Path(args.out), args.scheme) + result = {"template": args.out, "scheme": args.scheme, "quantized_modules": len(layout.quantized_modules), + "ignored_linears": len(layout.quantization_config["ignore"])} elif args.cmd == "convert": shards = args.shards.split(",") if args.shards else None result = requantize_checkpoint(Path(args.source), Path(args.template), Path(args.out), _smoothing_from_args(args), shards) diff --git a/yeto/rl/quant/schemes.py b/yeto/rl/quant/schemes.py index 86d85a6e7..2c5bb4cb7 100644 --- a/yeto/rl/quant/schemes.py +++ b/yeto/rl/quant/schemes.py @@ -10,6 +10,9 @@ 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). +* :data:`FP8_FFN_CHANNEL`: FP8 E4M3 per output channel with dynamic per-token + FP8 activations (W8A8) on the in-scope linears only (FFN of a dense model, + routed experts of a MoE model; S26 E6 tier 2, :mod:`.scope`). Release directories are named by scheme (:func:`release_dir_for`): ``policy-releases`` for BF16, ``policy-releases-`` for a re-quantized @@ -22,14 +25,15 @@ BF16 = "bf16" W4A8_INT4_CHANNEL = "w4a8-int4-channel" MXFP4_G32_RTN = "mxfp4-g32-rtn" -QUANTIZED_SCHEMES = (W4A8_INT4_CHANNEL, MXFP4_G32_RTN) +FP8_FFN_CHANNEL = "fp8-w8a8-channel" +QUANTIZED_SCHEMES = (W4A8_INT4_CHANNEL, MXFP4_G32_RTN, FP8_FFN_CHANNEL) 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_NONE = "none" # option R1, and every MXFP4 and FP8 release AWQ_FIXED_ARTEFACT = "fixed-artefact" # option R2: fixed factors from the original artefact diff --git a/yeto/rl/quant/scope.py b/yeto/rl/quant/scope.py new file mode 100644 index 000000000..7e561cd7b --- /dev/null +++ b/yeto/rl/quant/scope.py @@ -0,0 +1,211 @@ +"""Which linear modules a quantized rollout may quantize (S26 E6 scope rule). + +The rule is fixed (GOAL ยง5 E6): + +* a dense model quantizes only its FFN (MLP) linears: + ``mlp.gate_proj``, ``mlp.up_proj``, ``mlp.down_proj`` (and the fused + ``mlp.gate_up_proj``); +* a MoE model quantizes only its routed experts: + ``mlp.experts..{gate_proj,up_proj,down_proj,w1,w2,w3}`` (and the stacked + ``mlp.experts.gate_up_proj`` / ``mlp.experts.down_proj``); its router, its + shared experts and any dense FFN layer stay BF16; +* the router, ``lm_head``, the embeddings, every norm and every attention + module always stay BF16. + +Names are module names as an HF checkpoint stores them (``.weight``). +Anything the rule does not name as in scope is out of scope, so a new module +kind (a vision tower, an MTP head) stays BF16 until the rule names it. + +This module does not import torch: the launcher, the island entries and the +sampler's post-load check use it. + +* :func:`model_kind`: ``dense`` or ``moe`` from an HF ``config.json``; +* :func:`in_scope` / :func:`check_quantized`: the rule and its refusal + (:class:`QuantScopeError`, raised before any engine starts); +* :func:`linear_modules` / :func:`ignored_modules`: the out-of-scope linear + modules of a checkpoint (the engines' ignore lists); +* :func:`engine_ignore_list` / :func:`online_quantized`: the SGLang online FP8 + ignore list and what SGLang would quantize with it (simulated, see + :func:`path_match`). +""" + +from __future__ import annotations + +import json +import re +import struct +from collections.abc import Iterable, Mapping, Sequence +from pathlib import Path +from typing import Any + +DENSE = "dense" +MOE = "moe" +KINDS = (DENSE, MOE) + +_FFN_PROJ = r"(?:gate_proj|up_proj|down_proj|gate_up_proj)" +_DENSE_FFN = re.compile(rf"(?:^|\.)mlp\.{_FFN_PROJ}$") +_ROUTED_EXPERT = re.compile( + r"(?:^|\.)mlp\.experts\.(?:\d+\.(?:gate_proj|up_proj|down_proj|w1|w2|w3)|gate_up_proj|down_proj)$") + +# Out-of-scope reasons, checked in this order (the first match names the reason). +_KEEP_RULES: tuple[tuple[str, re.Pattern[str]], ...] = ( + ("lm_head", re.compile(r"(?:^|\.)lm_head$")), + ("embedding", re.compile(r"(?:^|\.)(?:embed_tokens|embeddings?|wte|word_embeddings)(?:$|\.)")), + ("norm", re.compile(r"(?:norm|layernorm|ln_\w*)$", re.IGNORECASE)), + ("attention", re.compile(r"(?:^|\.)(?:self_attn|linear_attn|attention|attn|self_attention)(?:$|\.)")), + ("shared_expert", re.compile(r"(?:^|\.)mlp\.shared_experts?(?:_gate)?(?:$|\.)")), + ("router", re.compile(r"(?:^|\.)mlp\.(?:gate|router)(?:$|\.)")), +) + +# Config keys that give the routed-expert count of a MoE model. +_EXPERT_KEYS = ("num_experts", "n_routed_experts", "num_local_experts", "moe_num_experts") + +# SGLang online FP8 (``--quantization fp8``) reads its ignore list from +# SGLANG_FP8_IGNORED_LAYERS (sglang 2fa88018 ``layers/quantization/fp8.py`` +# Fp8Config.__init__). These path segments are added to the explicit names, so +# a module whose SGLang prefix differs from its checkpoint name still matches. +_ENGINE_SEGMENTS = { + DENSE: ("self_attn", "linear_attn", "lm_head", "embed_tokens", "mlp.shared_expert", + "mlp.shared_experts", "mlp.shared_expert_gate"), + MOE: ("self_attn", "linear_attn", "lm_head", "embed_tokens", "mlp.shared_expert", + "mlp.shared_experts", "mlp.shared_expert_gate", "mlp.gate", "mlp.gate_proj", + "mlp.up_proj", "mlp.down_proj"), +} + + +class QuantScopeError(ValueError): + """A quantization plan or checkpoint quantizes a module outside the scope rule.""" + + +def model_kind(config: Mapping[str, Any]) -> str: + """``moe`` when the HF config declares routed experts (also in ``text_config``), else ``dense``.""" + for section in (config, config.get("text_config") or {}): + for key in _EXPERT_KEYS: + value = section.get(key) + if isinstance(value, int) and not isinstance(value, bool) and value > 0: + return MOE + return DENSE + + +def model_kind_of(checkpoint: Path) -> str: + return model_kind(json.loads((Path(checkpoint) / "config.json").read_text())) + + +def keep_reason(module: str) -> str | None: + """Why ``module`` must stay BF16 by name (``None``: the name gives no reason).""" + for reason, pattern in _KEEP_RULES: + if pattern.search(module): + return reason + return None + + +def in_scope(module: str, kind: str) -> bool: + """True when the scope rule lets ``kind`` models quantize ``module``.""" + if kind not in KINDS: + raise ValueError(f"unknown model kind {kind!r}; known: {', '.join(KINDS)}") + if kind == MOE: + return bool(_ROUTED_EXPERT.search(module)) + return bool(_DENSE_FFN.search(module)) and keep_reason(module) is None + + +def violations(modules: Iterable[str], kind: str) -> dict[str, str]: + """Out-of-scope modules of ``modules`` and why (``not_ffn`` / ``not_routed_expert`` when no name rule fits).""" + other = "not_routed_expert" if kind == MOE else "not_ffn" + return {m: keep_reason(m) or other for m in sorted(set(modules)) if not in_scope(m, kind)} + + +def check_quantized(modules: Iterable[str], kind: str, *, what: str) -> None: + """Raise :class:`QuantScopeError` when ``modules`` holds an out-of-scope module.""" + bad = violations(modules, kind) + if bad: + sample = ", ".join(f"{m} ({r})" for m, r in list(bad.items())[:5]) + raise QuantScopeError( + f"{what} quantizes {len(bad)} module(s) outside the {kind} scope (only " + f"{'routed experts' if kind == MOE else 'FFN linears'} may be quantized; router, " + f"lm_head, embeddings, norms and attention stay BF16), e.g. {sample}") + + +def linear_modules(tensors: Mapping[str, Sequence[int]]) -> list[str]: + """Modules with a 2-D ``.weight`` (``tensors``: name -> shape).""" + return sorted(n[: -len(".weight")] for n, shape in tensors.items() + if n.endswith(".weight") and len(shape) == 2) + + +def ignored_modules(tensors: Mapping[str, Sequence[int]], kind: str) -> list[str]: + """The out-of-scope modules of ``linear_modules(tensors)``.""" + return [m for m in linear_modules(tensors) if not in_scope(m, kind)] + + +# Fused engine modules and their checkpoint shards (SGLang / vLLM ``packed_modules_mapping``). +FUSED_SHARDS: dict[str, tuple[str, ...]] = { + "qkv_proj": ("q_proj", "k_proj", "v_proj"), + "gate_up_proj": ("gate_proj", "up_proj"), +} + + +def with_fused_aliases(modules: Iterable[str]) -> list[str]: + """``modules`` plus the fused engine name of every group whose shards are all listed + (``...self_attn.q/k/v_proj`` -> ``...self_attn.qkv_proj``). + + SGLang 2fa88018 compressed-tensors matches its ``ignore`` list against the fused + module name and falls back to the shards only through the model class's + ``packed_modules_mapping``, which Qwen3ForCausalLM does not set: without the + alias ``qkv_proj`` was given the FP8 scheme on the card (e6-gate-t2, 10-21).""" + listed = set(modules) + aliases = set() + for module in listed: + parent, _, leaf = module.rpartition(".") + for fused, shards in FUSED_SHARDS.items(): + if leaf in shards and all(f"{parent}.{shard}" in listed for shard in shards): + aliases.add(f"{parent}.{fused}" if parent else fused) + return sorted(listed | aliases) + + +def checkpoint_shapes(checkpoint: Path) -> dict[str, list[int]]: + """Tensor name -> shape of a safetensors checkpoint, from the shard headers only.""" + checkpoint = Path(checkpoint) + index = checkpoint / "model.safetensors.index.json" + shards = (sorted(set(json.loads(index.read_text())["weight_map"].values())) if index.is_file() + else sorted(p.name for p in checkpoint.glob("*.safetensors"))) + if not shards: + raise FileNotFoundError(f"{checkpoint}: no safetensors shards") + shapes: dict[str, list[int]] = {} + for shard in shards: + with open(checkpoint / shard, "rb") as f: + (n,) = struct.unpack(" list[str]: + """SGLang online FP8 ignore list: the explicit out-of-scope names plus the path segments.""" + return sorted(set(ignored_modules(tensors, kind)) | set(_ENGINE_SEGMENTS[kind])) + + +def path_match(ignored: str, prefix: str) -> bool: + """SGLang's ``_module_path_match`` (sglang 2fa88018 ``layers/quantization/utils.py``): + equal, a dotted prefix, or a dotted segment run of ``prefix``.""" + ignored, prefix = ignored.rstrip("."), prefix.rstrip(".") + return ignored == prefix or prefix.startswith(ignored + ".") or f".{ignored}." in f".{prefix}." + + +def online_quantized(modules: Iterable[str], ignore: Sequence[str]) -> list[str]: + """The modules an online FP8 engine quantizes with ``ignore`` (every linear not matched). + Embeddings and ``lm_head`` are not linear layers in SGLang or vLLM, so they are never + quantized there; they are still listed in the ignore lists.""" + return sorted(m for m in modules + if keep_reason(m) not in ("embedding", "lm_head") + and not any(path_match(i, m) for i in ignore)) + + +def check_online_plan(tensors: Mapping[str, Sequence[int]], kind: str, ignore: Sequence[str]) -> list[str]: + """Refuse an ignore list that lets the engine quantize an out-of-scope module, or that + keeps every in-scope module BF16; returns the modules the engine would quantize.""" + quantized = online_quantized(linear_modules(tensors), ignore) + check_quantized(quantized, kind, what="the online FP8 plan") + if not quantized: + raise QuantScopeError(f"the online FP8 plan quantizes nothing ({kind} model; no in-scope linear)") + return quantized diff --git a/yeto/rl/quant/tiers.py b/yeto/rl/quant/tiers.py new file mode 100644 index 000000000..d84e42085 --- /dev/null +++ b/yeto/rl/quant/tiers.py @@ -0,0 +1,317 @@ +"""Rollout quantization tiers: one backend-neutral setting, two engine plans (S26 E6). + +``--rl-rollout-quant-tier N`` picks the precision the sampler generates with. +The trainer always trains in its own precision (BF16 weights); the off-policy +gap the tier adds is corrected by the run's TIS (``--rl-algorithm-spec``). + +====== =============== ===================================================== +tier name what changes +====== =============== ===================================================== +0 ``bf16`` nothing (default) +1 ``online-fp8`` the sampler engine quantizes the in-scope linears to + FP8 itself at load (SGLang ``--quantization fp8``; + vLLM ``quantization=fp8``); KV cache stays BF16: + an FP8 KV cache without calibrated scales failed + the gate on Qwen3-4B (p99 |diff| 0.82 against 0.46 + with a BF16 KV cache, limit 0.5; e6-gate-t1-r3 / + e6-gate-t1-bf16kv, 2026-10-21) +2 ``trainer-fp8`` the trainer side converts each BF16 release to FP8 + W8A8 per channel on the in-scope linears + (:data:`~.schemes.FP8_FFN_CHANNEL`) and the sampler + island loads that checkpoint; KV cache stays BF16 +3 ``4bit`` interface only (large models); refused at launch +====== =============== ===================================================== + +Scope (every tier, :mod:`.scope`): a dense model quantizes only its FFN +linears, a MoE model only its routed experts; router, ``lm_head``, +embeddings, norms and attention stay BF16. An engine plan that would quantize +anything else is refused before the engine starts. + +Where each tier runs: + +* tier 1, colocated or sampler islands, both backends. Miles/SGLang: the + ignore list reaches SGLang through ``SGLANG_FP8_IGNORED_LAYERS`` (sglang + 2fa88018 reads it in ``Fp8Config``); verl/vLLM: the FP8 config of verl's own + FP8 rollout (fork acad9875 ``vllm_async_server._apply_quantization``) with + the out-of-scope linears added to ``ignored_layers``. +* tier 2, store-coupled sampler/trainer islands only + (``--rl-island-sample-role``, full-parameter): a colocated island has no + checkpoint transfer to make smaller, and a LoRA release carries no base + weights to convert. + +FP8 needs compute capability 8.9 or newer (Ada, Hopper, Blackwell); the launcher +accepts only the cards of :data:`FP8_CARDS` for tiers 1 and 2. + +Status: CPU tests only (tests/test_rl_quant_tiers.py). NOT VERIFIED on a GPU: +SGLang online FP8 with an in-place BF16 weight update (Miles colocated), the +vLLM FP8 config through ``engine_kwargs.vllm.hf_overrides`` (verl), the FP8 +compressed-tensors checkpoint through ``update_weights_from_disk``. +""" + +from __future__ import annotations + +import json +import struct +from dataclasses import dataclass +from pathlib import Path +from typing import Any + +from . import scope +from .schemes import BF16, FP8_FFN_CHANNEL + +TIER_BF16 = 0 +TIER_ONLINE_FP8 = 1 +TIER_TRAINER_FP8 = 2 +TIER_4BIT = 3 +TIER_NAMES = {TIER_BF16: "bf16", TIER_ONLINE_FP8: "online-fp8", TIER_TRAINER_FP8: "trainer-fp8", + TIER_4BIT: "4bit"} +TIERS = tuple(TIER_NAMES) + +FLAG = "--rl-rollout-quant-tier" +SOURCE_FLAG = "--rl-rollout-quant-fp8-source" +# Tier 1, who quantizes: the engine at load (online FP8), or the trainer on every weight +# update (Miles quantize_params_fp8, 128x128 blocks) with the engines started on Miles' +# block-FP8 conversion of the base. Both keep the scope rule's modules in BF16. +FP8_SOURCE_ENGINE = "engine" +FP8_SOURCE_TRAINER = "trainer" +NATIVE_BLOCK = [128, 128] +SGLANG_IGNORE_ENV = "SGLANG_FP8_IGNORED_LAYERS" +# verl's own FP8 rollout config (fork acad9875 vllm_async_server.py:1373); the +# ignore list is replaced by ours, which keeps verl's ``mlp.gate`` entries. +VLLM_FP8_CONFIG = {"activation_scheme": "dynamic", "fmt": "e4m3", "quant_method": "fp8", + "weight_block_size": [128, 128]} + +# Cards with FP8 tensor cores (compute capability >= 8.9), by canonical name +# (``yeto.gpu_spec``). Any other card is refused, also NPUs and PPUs (not checked). +FP8_CARDS = frozenset({"L4", "L40S", "RTX-6000-Ada", "RTX-PRO-6000", "H100", "H200", "B200"}) + + +class QuantTierError(ValueError): + """The tier cannot run with this launch (refused before any cloud spend).""" + + +def tier_of(args: Any) -> int: + tier = getattr(args, "rl_rollout_quant_tier", TIER_BF16) + tier = TIER_BF16 if tier is None else int(tier) + if tier not in TIERS: + raise QuantTierError(f"{FLAG} {tier}: expected one of {TIERS}") + return tier + + +def release_scheme(tier: int) -> str: + """The release scheme the sampler of ``tier`` serves (tier 2: the FP8 releases).""" + return FP8_FFN_CHANNEL if tier == TIER_TRAINER_FP8 else BF16 + + +VERL_BACKEND = "verl" +# H100 runs 2026-10-21 (SESSION26/logs/E6): vLLM 0.29 given verl's FP8 config +# (hf_overrides, quant_method fp8) on a BF16 checkpoint treats the weights as FP8 +# serialized and generates garbage (512-token answers, NaN log-probs) until verl's +# weight sync re-quantizes the base; yeto verl is LoRA-only and a LoRA sync never sends +# the base (e6-t1-verl-a-r2, e6-vllm-fp8-probe). Plain vLLM online FP8 works with LoRA +# but has no ignore list: it quantizes attention, which failed the gate as SGLang +# online FP8 of every linear (p99 |diff| 0.60, limit 0.5). +VERL_TIER1_REFUSAL = (f"{FLAG} 1 is not available on verl: verl's FP8 rollout needs a base-weight " + "sync to quantize the base, and yeto verl islands are LoRA-only (the engine " + "would generate garbage); use tier 0 on verl") + + +# Backends whose tier 1 defaults to the trainer FP8 source (framework-native block FP8 +# with the ignore list): its H100 gate margin is larger than the engine's online FP8 +# (p99 |diff| 0.43 vs 0.46-0.49, limit 0.5; SESSION26 E6 round 2). +TRAINER_SOURCE_BACKENDS = frozenset({"miles"}) + + +def default_fp8_source(args: Any) -> str: + """The FP8 source when ``--rl-rollout-quant-fp8-source`` is not given.""" + backend = getattr(args, "rl_backend", None) or "miles" + if tier_of(args) == TIER_ONLINE_FP8 and backend in TRAINER_SOURCE_BACKENDS: + return FP8_SOURCE_TRAINER + return FP8_SOURCE_ENGINE + + +def fp8_source(args: Any) -> str: + return getattr(args, "rl_rollout_quant_fp8_source", None) or default_fp8_source(args) + + +def validate_launch(args: Any) -> int: + """Refuse a tier the launch cannot run (launcher side, no checkpoint needed); returns the tier.""" + tier = tier_of(args) + if getattr(args, "rl_rollout_quant_fp8_source", None) is None and hasattr(args, "__dict__"): + # Resolved once here so the island flags and the run config carry it explicitly. + args.rl_rollout_quant_fp8_source = default_fp8_source(args) + if tier == TIER_BF16: + if fp8_source(args) != FP8_SOURCE_ENGINE: + raise QuantTierError(f"{SOURCE_FLAG} {fp8_source(args)} needs {FLAG} 1") + return tier + if tier == TIER_4BIT: + raise QuantTierError(f"{FLAG} 3 (4-bit) is an interface only in this version (large models; " + "H100 has no FP4 units); use tier 0, 1 or 2") + cards = _cards(getattr(args, "gpu", None)) + if lacking := sorted(cards - FP8_CARDS): + raise QuantTierError(f"{FLAG} {tier} needs FP8 tensor cores (compute capability >= 8.9: " + f"{', '.join(sorted(FP8_CARDS))}); --gpu names {', '.join(lacking)}") + if tier == TIER_ONLINE_FP8 and getattr(args, "rl_backend", None) == VERL_BACKEND: + raise QuantTierError(VERL_TIER1_REFUSAL) + if fp8_source(args) == FP8_SOURCE_TRAINER and tier != TIER_ONLINE_FP8: + raise QuantTierError(f"{SOURCE_FLAG} trainer applies to tier 1 only (got {FLAG} {tier})") + if tier == TIER_TRAINER_FP8: + role = getattr(args, "rl_sample_role", None) or "both" + if not getattr(args, "rl_island_sample_role", None) and role not in ("sampler", "trainer"): + raise QuantTierError( + f"{FLAG} 2 converts each trainer release and ships it to a sampler island: it needs " + "store-coupled islands (--rl-island-sample-role, or one launch per role with " + "--rl-sample-role). Colocated islands use tier 1") + if getattr(args, "tuning", "lora") != "full": + raise QuantTierError(f"{FLAG} 2 needs --tuning full: a LoRA release carries no base weights " + "to convert (tier 1 quantizes the base once in the engine)") + return tier + + +def _cards(gpu: Any) -> set[str]: + """Canonical card names of a ``--gpu`` spec; empty when there is none (auto fleet) + or it does not parse (the launcher reports that itself).""" + if not gpu: + return set() + from yeto.gpu_spec import parse_gpu_spec + + try: + return {cluster.gpu for cluster in parse_gpu_spec(str(gpu))} + except ValueError: + return set() + + +# ---------------------------------------------------------------- tier 1 engine plans +@dataclass(frozen=True) +class OnlinePlan: + """What an online FP8 engine gets (tier 1), checked against the checkpoint. + + ``explicit``: every out-of-scope linear of the checkpoint by name; + ``segments``: path segments SGLang also matches (:func:`~.scope.engine_ignore_list`); + ``quantized``: the modules the engine will quantize (simulated, in scope by check).""" + + kind: str + explicit: tuple[str, ...] + segments: tuple[str, ...] + quantized: tuple[str, ...] + + def summary(self) -> dict[str, Any]: + return {"tier": TIER_ONLINE_FP8, "kind": self.kind, "quantized_modules": len(self.quantized), + "ignored_modules": len(self.explicit), "example_quantized": list(self.quantized[:3])} + + +def online_plan(checkpoint: Path) -> OnlinePlan: + """The tier 1 plan for the BF16 checkpoint ``checkpoint``; raises + :class:`~.scope.QuantScopeError` when the engine would quantize an out-of-scope module + (with either the SGLang list or the explicit vLLM list).""" + kind = scope.model_kind_of(checkpoint) + shapes = scope.checkpoint_shapes(checkpoint) + explicit = scope.ignored_modules(shapes, kind) + segments = sorted(set(scope.engine_ignore_list(shapes, kind)) - set(explicit)) + quantized = scope.check_online_plan(shapes, kind, [*explicit, *segments]) + if scope.check_online_plan(shapes, kind, explicit) != quantized: + raise scope.QuantScopeError("the SGLang and vLLM ignore lists quantize different modules") + return OnlinePlan(kind, tuple(explicit), tuple(segments), tuple(quantized)) + + +def sglang_online_argv() -> list[str]: + """Miles spellings of the SGLang tier 1 server flags (Miles prefixes SGLang's own).""" + return ["--sglang-quantization", "fp8"] + + +def sglang_online_env(plan: OnlinePlan) -> dict[str, str]: + """The SGLang worker environment of a tier 1 plan.""" + entries = [*plan.explicit, *plan.segments] + if any("," in entry for entry in entries): + raise scope.QuantScopeError("an ignore entry holds a comma; SGLANG_FP8_IGNORED_LAYERS splits on commas") + return {SGLANG_IGNORE_ENV: ",".join(entries)} + + +def native_ignore_list(plan: OnlinePlan) -> list[str]: + """Tier 1 trainer source: the modules Miles' converter and ``quantize_params_fp8`` + keep in BF16 (the converter records them in ``modules_to_not_convert``, which + SGLang and Miles both read; SGLang path match, as the online ignore list).""" + return [*plan.explicit, *plan.segments] + + +def check_native_base(directory: Path) -> int: + """Tier 1 trainer source: Miles' block-FP8 conversion must be ``quant_method fp8`` + with 128x128 blocks and hold FP8 weights only in scope. Returns the FP8 module count.""" + directory = Path(directory) + config = json.loads((directory / "config.json").read_text()) + q = config.get("quantization_config") or {} + if q.get("quant_method") != "fp8" or list(q.get("weight_block_size") or ()) != NATIVE_BLOCK: + raise scope.QuantScopeError(f"{directory}: quantization_config is {q.get('quant_method')!r} " + f"blocks {q.get('weight_block_size')!r}, expected fp8 {NATIVE_BLOCK}") + dtypes = _checkpoint_dtypes(directory) + quantized = [n[: -len(".weight")] for n, d in dtypes.items() if d == "F8_E4M3" and n.endswith(".weight")] + if not quantized: + raise scope.QuantScopeError(f"{directory}: an FP8 base without FP8 weights") + scope.check_quantized(quantized, scope.model_kind(config), what=f"FP8 base {directory.name}") + return len(quantized) + + +def vllm_online_overrides(plan: OnlinePlan, num_layers: int) -> list[str]: + """verl Hydra overrides of a tier 1 plan. + + ``rollout.quantization=fp8`` keeps verl's FP8 load path (weights quantized on each + sync, only for modules vLLM built as FP8: ``vllm_quant_utils.is_fp8_weight``). + ``engine_kwargs.vllm.hf_overrides`` replaces verl's FP8 config (``**engine_kwargs`` + comes after it in the vLLM arguments) with the same config whose ``ignored_layers`` + are the plan's explicit names plus verl's ``model.layers..mlp.gate`` entries. + vLLM matches ``ignored_layers`` by exact name. The value is a JSON string, which + verl passes unchanged as ``--hf-overrides``. + """ + gates = [f"model.layers.{i}.mlp.gate" for i in range(int(num_layers))] + config = {**VLLM_FP8_CONFIG, "ignored_layers": sorted(set(plan.explicit) | set(gates))} + hf_overrides = json.dumps({"quantization_config": config}, separators=(",", ":")) + if "'" in hf_overrides: + raise scope.QuantScopeError("a module name holds a single quote; it cannot be a Hydra quoted value") + return ["actor_rollout_ref.rollout.quantization=fp8", + f"++actor_rollout_ref.rollout.engine_kwargs.vllm.hf_overrides='{hf_overrides}'"] + + +def num_hidden_layers(checkpoint: Path) -> int: + config = json.loads((Path(checkpoint) / "config.json").read_text()) + section = config if "num_hidden_layers" in config else (config.get("text_config") or {}) + return int(section["num_hidden_layers"]) + + +# ---------------------------------------------------------------- tier 2 release side +def check_release_scope(directory: Path) -> int: + """Post-load check (before the engine is touched) of an FP8 release: the release + must declare the FP8 format and every FP8 weight must be in scope. Returns the + quantized module count. Reads ``config.json`` and the shard headers only.""" + directory = Path(directory) + config = json.loads((directory / "config.json").read_text()) + q = config.get("quantization_config") or {} + if q.get("quant_method") != "compressed-tensors" or q.get("format") != "float-quantized": + raise scope.QuantScopeError(f"{directory}: quantization_config is {q.get('quant_method')!r}/" + f"{q.get('format')!r}, expected compressed-tensors/float-quantized") + kind = scope.model_kind(config) + dtypes = _checkpoint_dtypes(directory) + quantized = [n[: -len(".weight")] for n, d in dtypes.items() if d == "F8_E4M3" and n.endswith(".weight")] + scope.check_quantized(quantized, kind, what=f"release {directory.name}") + if not quantized: + raise scope.QuantScopeError(f"{directory}: an FP8 release without FP8 weights") + unlisted = [m for m in scope.ignored_modules(scope.checkpoint_shapes(directory), kind) + if m not in set(q.get("ignore") or ())] + if unlisted: + raise scope.QuantScopeError(f"{directory}: {len(unlisted)} BF16 linears are not in the config " + f"ignore list (the engine would expect FP8 weights), e.g. {unlisted[:3]}") + return len(quantized) + + +def _checkpoint_dtypes(directory: Path) -> dict[str, str]: + out: dict[str, str] = {} + for shard in sorted(Path(directory).glob("*.safetensors")): + with open(shard, "rb") as f: + (n,) = struct.unpack(" str: + return f"rollout quant tier {tier} ({TIER_NAMES[tier]})" diff --git a/yeto/rl/stage1.py b/yeto/rl/stage1.py index 67ed28a5e..ffdaa16cd 100644 --- a/yeto/rl/stage1.py +++ b/yeto/rl/stage1.py @@ -629,9 +629,18 @@ def sampler_spec(plan: Stage1Plan, node: Node, engine: Any) -> dict[str, Any]: "checks": {**dict(plan.section("sampler").get("checks") or {}), "expected_expert_matrices": int(quant["expected_expert_matrices"])}, "event_tape": f"{state_dir}/sampler-{node.name}/events.jsonl", + # S26 E6: the rollout quant tier, applied on the node (sampler_engines.quant_tier) + **({"rollout_quant": section} if (section := quant_spec_section(plan)) else {}), } +def quant_spec_section(plan: Stage1Plan) -> dict[str, Any] | None: + """``rollout_quant_tier`` of the plan as the sampler spec section (None for tier 0).""" + from yeto.rl.adapters.sampler_engines.quant_tier import spec_section # noqa: PLC0415 + + return spec_section(int(plan.doc.get("rollout_quant_tier", 0) or 0)) + + 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`; @@ -778,6 +787,13 @@ 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") + tier = int(plan.doc.get("rollout_quant_tier", 0) or 0) + if tier and quant.get("scheme") != "bf16": + # S26 E6: stage 1 serves the quant section's scheme; a tier on top would quantize twice + errors.append(f"rollout_quant_tier {tier} needs quant.scheme bf16 (got {quant.get('scheme')!r})") + if tier == 2: + errors.append("rollout_quant_tier 2 is not wired into the stage-1 re-quantizer " + "(its requant side process is template-based); use quant.scheme instead") 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")