diff --git a/.gitignore b/.gitignore index 2d80f62e..ece43ff8 100644 --- a/.gitignore +++ b/.gitignore @@ -95,3 +95,4 @@ openspec/changes/*/evidence/**/*.pyc openspec/**/*.jsonl openspec/**/*.out openspec/**/*.err +.g1-build/ diff --git a/examples/stage1/glm53-flash-ppu.yaml b/examples/stage1/glm53-flash-ppu.yaml index 5f23b3a1..c212ec8d 100644 --- a/examples/stage1/glm53-flash-ppu.yaml +++ b/examples/stage1/glm53-flash-ppu.yaml @@ -37,6 +37,7 @@ model: origin: modal-volume://apex-w4a8-full/source/cecb5eeb9c6b32404a0dd930df81de2c239bdd84 staged: oss://oss-pai-ircyeaqtnfjw11wby2-cn-wulanchabu/yeto/models/apex-flash-1-bf16-cecb5ee/ local_dir: /data/models/apex-flash-1-bf16-cecb5ee + torch_dist: /data/models/apex-flash-1-bf16-cecb5ee_torch_dist # Miles converter output (convert step) sampler_initial: # round 0: the HF MXFP4 export, not our re-quantizer repo: yetalabs/apex-flash-1-mxfp4-experts revision: 70e86a0 diff --git a/examples/stage1/glm5next-mini-modal-dense.yaml b/examples/stage1/glm5next-mini-modal-dense.yaml new file mode 100644 index 00000000..28230adf --- /dev/null +++ b/examples/stage1/glm5next-mini-modal-dense.yaml @@ -0,0 +1,148 @@ +# S25 T-G1B: glm5next-mini-modal.yaml with two changes, so the weights move and the +# chain "trained weights -> re-quantized MXFP4 experts -> sampler" is checked by value: +# task.reward yeto.rl.word_diversity_reward:score (GSM8K is constant 0 on this model, +# so GRPO advantages and grad_norm were 0 in T-G1 run-17) +# trainer.miles.lr 3.0e-4 (1e-6 Adam steps are below one bf16 ulp of most weights, +# so a release hash could stay the same even with a non-zero gradient) +# Everything else (nodes, store, images, sampler, sampling) is the T-G1 plan. +# +# Stage-1 mini rehearsal (S25 T-G1) on Modal H100: the same plan schema and +# the same generated commands as glm53-flash-ppu.yaml; only the nodes +# (Modal containers) and the store (a Modal Volume through the yeto store +# hook, not OSS) differ. Model: Miles' 4-layer GLM5.3-Flash CI checkpoint cut +# to 32 routed experts per MoE layer (tools/probes/s25_g1/make_mini.py). +# +# yeto-stage1 plan examples/stage1/glm5next-mini-modal.yaml +# modal run tools/probes/s25_g1/modal_app.py::... (the driver reads this plan) + +run_id: s25-g1-mini +inventory: modal-mini-inventory.yaml +rounds: 3 +seed: 17 + +rendezvous: + modal-trainer: 127.0.0.1 + modal-sampler: 127.0.0.1 + +roles: + trainer: [modal-trainer] + sampler: [modal-sampler] + +model: + name: GLM5.3-Flash-4layer-e32 + architecture: glm5_next + num_layers: 4 # KDA, DSA, KDA, DSA + first_k_dense_replace: 1 + num_experts: 32 # cut from 288 (make_mini.py) + params_b: 5.3 # 推算: 10.5 GB bf16 / 2 + expert_params_b: 2.4 # 3 MoE layers x 32 x 3 x 4096 x 2048 + bf16_source: + revision: 7fb7883207fe73ea04430ac392f6ffe2b9f936c4 # CharyZeng/GLM-5.3-Flash-4layer, then cut + origin: hf://CharyZeng/GLM-5.3-Flash-4layer@7fb7883207fe73ea04430ac392f6ffe2b9f936c4 + staged: modal-volume://yeto-s25-g1-models/glm5next-mini-bf16-7fb7883 + local_dir: /data/models/glm5next-mini-bf16-7fb7883 + torch_dist: /data/models/glm5next-mini-bf16-7fb7883_torch_dist # Miles converter output (convert step) + sampler_initial: # round 0: our N8.5 converter on the bf16 mini model + repo: local/glm5next-mini-mxfp4-r0-7fb7883 + revision: 7fb7883 + staged: modal-volume://yeto-s25-g1-models/glm5next-mini-mxfp4-r0-7fb7883 + local_dir: /data/models/glm5next-mini-mxfp4-r0-7fb7883 + +quant: + scheme: mxfp4-g32-rtn + format: mxfp4-pack-quantized + release_dir: policy-releases-mxfp4-g32-rtn + expected_expert_matrices: 288 # 3 MoE layers x 32 experts x (gate, up, down) + frozen_pattern: '^model\.visual\.' # Miles trains the text model only; the vision tower stays the base + +store: /data/store/s25-g1 +store_backend: modal-volume # the only difference from OSS: the store hook and its Volume +store_volume: yeto-s25-g1-store +store_access: + head: direct + +images: + sampler: ghcr.io/michaellchung/yeto-miles-ports@sha256:5b9dc2a9e6e315d73dabfbc6faae428b59cc01aabec6189645a5c8c80d6e3721 + trainer: ghcr.io/michaellchung/yeto-miles-ports@sha256:5b9dc2a9e6e315d73dabfbc6faae428b59cc01aabec6189645a5c8c80d6e3721 + docker_device_args: [--gpus, all] # unused on Modal; the containers get the GPUs + +trainer: + backend: miles + miles: + commit: 1618b5c38db80e17ec928cfa47fe832cb18565fc # s25-distopt, checked out in the image + megatron: radixark/Megatron-LM@a84b105 + mbridge: image + workdir: /root/miles + extra_env: + # Megatron's get_model imports megatron.post_training (modelopt check), which the + # image's editable megatron_core install does not expose (Modal run-10) + PYTHONPATH: /root/Megatron-LM:/root/sky_workdir + tp: 1 + sequence_parallel: false + pp: 1 + pp_layers: [4] + cp: 1 + ep: 2 + etp: 1 + optimizer_offload_fraction: 1.0 + memory_estimate: # 推算; `plan` checks them + weights_grads_gib_per_card: 22.9 + host_optimizer_gib_per_node: 59.2 + transformer_impl: transformer_engine + recompute: full + max_tokens_per_gpu: 8192 + lr: 3.0e-4 # T-G1B: see the header + verl: + commit: null + strategy: fsdp2 + fsdp_size: 2 + ulysses: 1 + param_offload: false + optimizer_offload: true + lr: 1.0e-6 + +sampler: + engine: sglang + engine_version: 0.5.21.dev70+g2fa8801 + python: python3 + topology: tp1 + tp: 1 + precision: mxfp4 + max_model_len: 4096 + port: 30000 + moe_runner_backend: triton + hierarchical_cache: false + nextn: false + renames_weight_packed: true # the yeto hook in the image renames (sglang_load_report) + load_report: true # the hook reports loaded expert matrices (post-load check 1) + in_place_reload: false + # SGLang fuses the BF16 shared expert into the routed FusedMoE as expert 32; under the + # MXFP4 MoE method that slot is packed U8, so the BF16 shared weights do not fit (smoke-2) + extra_args: [--disable-shared-experts-fusion, --cuda-graph-max-bs-decode, "16"] # 16 = 4 prompts x 4 samples; capture took 118 s at the default sizes + extra_env: + PYTHONPATH: /root/sky_workdir + checks: + require_fixed_answers: false # a cut 4-layer model cannot answer GSM8K + fixed_max_tokens: 64 + +sampling: + max_tokens: 512 + max_prompt_len: 512 + temperature: 1.0 + top_p: 1.0 + top_k: -1 + repetition_penalty: 1.0 + frequency_penalty: 0.0 + presence_penalty: 0.0 + +task: + data: zhuzilin/gsm8k + data_revision: 0cbd9f31d91ac21a7613dcbc7fef992adac459ae + reward: yeto.rl.word_diversity_reward:score # T-G1B: see the header + rollout_batch_size: 4 + n_samples_per_prompt: 4 + prompts_jsonl: /data/models/datasets/gsm8k-0cbd9f3/train.jsonl + data_local: /data/models/datasets/gsm8k-0cbd9f3/train.jsonl # the same rows for the trainer (offline nodes) + input_key: messages + label_key: label + prompt_format: chat diff --git a/examples/stage1/glm5next-mini-modal-g1c.yaml b/examples/stage1/glm5next-mini-modal-g1c.yaml new file mode 100644 index 00000000..ad1d7533 --- /dev/null +++ b/examples/stage1/glm5next-mini-modal-g1c.yaml @@ -0,0 +1,154 @@ +# S25 T-G1C: glm5next-mini-modal-dense.yaml (T-G1B) with the reward and lr changed again: +# task.reward yeto.rl.token_diversity_reward:score (the word reward was gamed in g1b-1: +# one 512-token word scored 1.0 for every sample, round 2 had grad_norm 0) +# trainer.miles.lr 1.0e-4 (3e-4 drove the collapse in two steps; 1e-4 Adam steps are +# still about one bf16 ulp of a 0.02 weight, so releases still change) +# +# T-G1B: glm5next-mini-modal.yaml with two changes, so the weights move and the +# chain "trained weights -> re-quantized MXFP4 experts -> sampler" is checked by value: +# task.reward yeto.rl.word_diversity_reward:score (GSM8K is constant 0 on this model, +# so GRPO advantages and grad_norm were 0 in T-G1 run-17) +# trainer.miles.lr 3.0e-4 (1e-6 Adam steps are below one bf16 ulp of most weights, +# so a release hash could stay the same even with a non-zero gradient) +# Everything else (nodes, store, images, sampler, sampling) is the T-G1 plan. +# +# Stage-1 mini rehearsal (S25 T-G1) on Modal H100: the same plan schema and +# the same generated commands as glm53-flash-ppu.yaml; only the nodes +# (Modal containers) and the store (a Modal Volume through the yeto store +# hook, not OSS) differ. Model: Miles' 4-layer GLM5.3-Flash CI checkpoint cut +# to 32 routed experts per MoE layer (tools/probes/s25_g1/make_mini.py). +# +# yeto-stage1 plan examples/stage1/glm5next-mini-modal.yaml +# modal run tools/probes/s25_g1/modal_app.py::... (the driver reads this plan) + +run_id: s25-g1-mini +inventory: modal-mini-inventory.yaml +rounds: 3 +seed: 17 + +rendezvous: + modal-trainer: 127.0.0.1 + modal-sampler: 127.0.0.1 + +roles: + trainer: [modal-trainer] + sampler: [modal-sampler] + +model: + name: GLM5.3-Flash-4layer-e32 + architecture: glm5_next + num_layers: 4 # KDA, DSA, KDA, DSA + first_k_dense_replace: 1 + num_experts: 32 # cut from 288 (make_mini.py) + params_b: 5.3 # 推算: 10.5 GB bf16 / 2 + expert_params_b: 2.4 # 3 MoE layers x 32 x 3 x 4096 x 2048 + bf16_source: + revision: 7fb7883207fe73ea04430ac392f6ffe2b9f936c4 # CharyZeng/GLM-5.3-Flash-4layer, then cut + origin: hf://CharyZeng/GLM-5.3-Flash-4layer@7fb7883207fe73ea04430ac392f6ffe2b9f936c4 + staged: modal-volume://yeto-s25-g1-models/glm5next-mini-bf16-7fb7883 + local_dir: /data/models/glm5next-mini-bf16-7fb7883 + torch_dist: /data/models/glm5next-mini-bf16-7fb7883_torch_dist # Miles converter output (convert step) + sampler_initial: # round 0: our N8.5 converter on the bf16 mini model + repo: local/glm5next-mini-mxfp4-r0-7fb7883 + revision: 7fb7883 + staged: modal-volume://yeto-s25-g1-models/glm5next-mini-mxfp4-r0-7fb7883 + local_dir: /data/models/glm5next-mini-mxfp4-r0-7fb7883 + +quant: + scheme: mxfp4-g32-rtn + format: mxfp4-pack-quantized + release_dir: policy-releases-mxfp4-g32-rtn + expected_expert_matrices: 288 # 3 MoE layers x 32 experts x (gate, up, down) + frozen_pattern: '^model\.visual\.' # Miles trains the text model only; the vision tower stays the base + +store: /data/store/s25-g1 +store_backend: modal-volume # the only difference from OSS: the store hook and its Volume +store_volume: yeto-s25-g1-store +store_access: + head: direct + +images: + sampler: ghcr.io/michaellchung/yeto-miles-ports@sha256:5b9dc2a9e6e315d73dabfbc6faae428b59cc01aabec6189645a5c8c80d6e3721 + trainer: ghcr.io/michaellchung/yeto-miles-ports@sha256:5b9dc2a9e6e315d73dabfbc6faae428b59cc01aabec6189645a5c8c80d6e3721 + docker_device_args: [--gpus, all] # unused on Modal; the containers get the GPUs + +trainer: + backend: miles + miles: + commit: 1618b5c38db80e17ec928cfa47fe832cb18565fc # s25-distopt, checked out in the image + megatron: radixark/Megatron-LM@a84b105 + mbridge: image + workdir: /root/miles + extra_env: + # Megatron's get_model imports megatron.post_training (modelopt check), which the + # image's editable megatron_core install does not expose (Modal run-10) + PYTHONPATH: /root/Megatron-LM:/root/sky_workdir + tp: 1 + sequence_parallel: false + pp: 1 + pp_layers: [4] + cp: 1 + ep: 2 + etp: 1 + optimizer_offload_fraction: 1.0 + memory_estimate: # 推算; `plan` checks them + weights_grads_gib_per_card: 22.9 + host_optimizer_gib_per_node: 59.2 + transformer_impl: transformer_engine + recompute: full + max_tokens_per_gpu: 8192 + lr: 1.0e-4 # T-G1C: see the header + verl: + commit: null + strategy: fsdp2 + fsdp_size: 2 + ulysses: 1 + param_offload: false + optimizer_offload: true + lr: 1.0e-6 + +sampler: + engine: sglang + engine_version: 0.5.21.dev70+g2fa8801 + python: python3 + topology: tp1 + tp: 1 + precision: mxfp4 + max_model_len: 4096 + port: 30000 + moe_runner_backend: triton + hierarchical_cache: false + nextn: false + renames_weight_packed: true # the yeto hook in the image renames (sglang_load_report) + load_report: true # the hook reports loaded expert matrices (post-load check 1) + in_place_reload: false + # SGLang fuses the BF16 shared expert into the routed FusedMoE as expert 32; under the + # MXFP4 MoE method that slot is packed U8, so the BF16 shared weights do not fit (smoke-2) + extra_args: [--disable-shared-experts-fusion, --cuda-graph-max-bs-decode, "16"] # 16 = 4 prompts x 4 samples; capture took 118 s at the default sizes + extra_env: + PYTHONPATH: /root/sky_workdir + checks: + require_fixed_answers: false # a cut 4-layer model cannot answer GSM8K + fixed_max_tokens: 64 + +sampling: + max_tokens: 512 + max_prompt_len: 512 + temperature: 1.0 + top_p: 1.0 + top_k: -1 + repetition_penalty: 1.0 + frequency_penalty: 0.0 + presence_penalty: 0.0 + +task: + data: zhuzilin/gsm8k + data_revision: 0cbd9f31d91ac21a7613dcbc7fef992adac459ae + reward: yeto.rl.token_diversity_reward:score # T-G1C: see the header + rollout_batch_size: 4 + n_samples_per_prompt: 4 + prompts_jsonl: /data/models/datasets/gsm8k-0cbd9f3/train.jsonl + data_local: /data/models/datasets/gsm8k-0cbd9f3/train.jsonl # the same rows for the trainer (offline nodes) + input_key: messages + label_key: label + prompt_format: chat diff --git a/examples/stage1/glm5next-mini-modal.yaml b/examples/stage1/glm5next-mini-modal.yaml new file mode 100644 index 00000000..71e3b4d2 --- /dev/null +++ b/examples/stage1/glm5next-mini-modal.yaml @@ -0,0 +1,140 @@ +# Stage-1 mini rehearsal (S25 T-G1) on Modal H100: the same plan schema and +# the same generated commands as glm53-flash-ppu.yaml; only the nodes +# (Modal containers) and the store (a Modal Volume through the yeto store +# hook, not OSS) differ. Model: Miles' 4-layer GLM5.3-Flash CI checkpoint cut +# to 32 routed experts per MoE layer (tools/probes/s25_g1/make_mini.py). +# +# yeto-stage1 plan examples/stage1/glm5next-mini-modal.yaml +# modal run tools/probes/s25_g1/modal_app.py::... (the driver reads this plan) + +run_id: s25-g1-mini +inventory: modal-mini-inventory.yaml +rounds: 3 +seed: 17 + +rendezvous: + modal-trainer: 127.0.0.1 + modal-sampler: 127.0.0.1 + +roles: + trainer: [modal-trainer] + sampler: [modal-sampler] + +model: + name: GLM5.3-Flash-4layer-e32 + architecture: glm5_next + num_layers: 4 # KDA, DSA, KDA, DSA + first_k_dense_replace: 1 + num_experts: 32 # cut from 288 (make_mini.py) + params_b: 5.3 # 推算: 10.5 GB bf16 / 2 + expert_params_b: 2.4 # 3 MoE layers x 32 x 3 x 4096 x 2048 + bf16_source: + revision: 7fb7883207fe73ea04430ac392f6ffe2b9f936c4 # CharyZeng/GLM-5.3-Flash-4layer, then cut + origin: hf://CharyZeng/GLM-5.3-Flash-4layer@7fb7883207fe73ea04430ac392f6ffe2b9f936c4 + staged: modal-volume://yeto-s25-g1-models/glm5next-mini-bf16-7fb7883 + local_dir: /data/models/glm5next-mini-bf16-7fb7883 + torch_dist: /data/models/glm5next-mini-bf16-7fb7883_torch_dist # Miles converter output (convert step) + sampler_initial: # round 0: our N8.5 converter on the bf16 mini model + repo: local/glm5next-mini-mxfp4-r0-7fb7883 + revision: 7fb7883 + staged: modal-volume://yeto-s25-g1-models/glm5next-mini-mxfp4-r0-7fb7883 + local_dir: /data/models/glm5next-mini-mxfp4-r0-7fb7883 + +quant: + scheme: mxfp4-g32-rtn + format: mxfp4-pack-quantized + release_dir: policy-releases-mxfp4-g32-rtn + expected_expert_matrices: 288 # 3 MoE layers x 32 experts x (gate, up, down) + frozen_pattern: '^model\.visual\.' # Miles trains the text model only; the vision tower stays the base + +store: /data/store/s25-g1 +store_backend: modal-volume # the only difference from OSS: the store hook and its Volume +store_volume: yeto-s25-g1-store +store_access: + head: direct + +images: + sampler: ghcr.io/michaellchung/yeto-miles-ports@sha256:5b9dc2a9e6e315d73dabfbc6faae428b59cc01aabec6189645a5c8c80d6e3721 + trainer: ghcr.io/michaellchung/yeto-miles-ports@sha256:5b9dc2a9e6e315d73dabfbc6faae428b59cc01aabec6189645a5c8c80d6e3721 + docker_device_args: [--gpus, all] # unused on Modal; the containers get the GPUs + +trainer: + backend: miles + miles: + commit: 1618b5c38db80e17ec928cfa47fe832cb18565fc # s25-distopt, checked out in the image + megatron: radixark/Megatron-LM@a84b105 + mbridge: image + workdir: /root/miles + extra_env: + # Megatron's get_model imports megatron.post_training (modelopt check), which the + # image's editable megatron_core install does not expose (Modal run-10) + PYTHONPATH: /root/Megatron-LM:/root/sky_workdir + tp: 1 + sequence_parallel: false + pp: 1 + pp_layers: [4] + cp: 1 + ep: 2 + etp: 1 + optimizer_offload_fraction: 1.0 + memory_estimate: # 推算; `plan` checks them + weights_grads_gib_per_card: 22.9 + host_optimizer_gib_per_node: 59.2 + transformer_impl: transformer_engine + recompute: full + max_tokens_per_gpu: 8192 + lr: 1.0e-6 + verl: + commit: null + strategy: fsdp2 + fsdp_size: 2 + ulysses: 1 + param_offload: false + optimizer_offload: true + lr: 1.0e-6 + +sampler: + engine: sglang + engine_version: 0.5.21.dev70+g2fa8801 + python: python3 + topology: tp1 + tp: 1 + precision: mxfp4 + max_model_len: 4096 + port: 30000 + moe_runner_backend: triton + hierarchical_cache: false + nextn: false + renames_weight_packed: true # the yeto hook in the image renames (sglang_load_report) + load_report: true # the hook reports loaded expert matrices (post-load check 1) + in_place_reload: false + # SGLang fuses the BF16 shared expert into the routed FusedMoE as expert 32; under the + # MXFP4 MoE method that slot is packed U8, so the BF16 shared weights do not fit (smoke-2) + extra_args: [--disable-shared-experts-fusion, --cuda-graph-max-bs-decode, "16"] # 16 = 4 prompts x 4 samples; capture took 118 s at the default sizes + extra_env: + PYTHONPATH: /root/sky_workdir + checks: + require_fixed_answers: false # a cut 4-layer model cannot answer GSM8K + fixed_max_tokens: 64 + +sampling: + max_tokens: 512 + max_prompt_len: 512 + temperature: 1.0 + top_p: 1.0 + top_k: -1 + repetition_penalty: 1.0 + frequency_penalty: 0.0 + presence_penalty: 0.0 + +task: + data: zhuzilin/gsm8k + data_revision: 0cbd9f31d91ac21a7613dcbc7fef992adac459ae + reward: yeto.rl.gsm8k_reward:score + rollout_batch_size: 4 + n_samples_per_prompt: 4 + prompts_jsonl: /data/models/datasets/gsm8k-0cbd9f3/train.jsonl + data_local: /data/models/datasets/gsm8k-0cbd9f3/train.jsonl # the same rows for the trainer (offline nodes) + input_key: messages + label_key: label + prompt_format: chat diff --git a/examples/stage1/modal-mini-inventory.yaml b/examples/stage1/modal-mini-inventory.yaml new file mode 100644 index 00000000..fd434c18 --- /dev/null +++ b/examples/stage1/modal-mini-inventory.yaml @@ -0,0 +1,31 @@ +# Stage-1 mini rehearsal on Modal (S25 T-G1): one trainer container with +# 2 x H100 and one sampler container with 1 x H100. "host" is a label only: +# the Modal driver (tools/probes/s25_g1/modal_app.py) starts the containers, +# not SSH. /data/models and /data/store are Modal Volumes in each container. +nodes: + - name: modal-trainer + host: modal-trainer + address: 127.0.0.1 + device_family: cuda + cpu_arch: x86_64 + card: H100 + cards_per_node: 2 + hbm_gib: 80 + host_memory_gib: 400 + mounts: {/data: /data} + image: ghcr.io/michaellchung/yeto-miles-ports@sha256:5b9dc2a9e6e315d73dabfbc6faae428b59cc01aabec6189645a5c8c80d6e3721 + runtime: docker + roles: [trainer] + - name: modal-sampler + host: modal-sampler + address: 127.0.0.1 + device_family: cuda + cpu_arch: x86_64 + card: H100 + cards_per_node: 1 + hbm_gib: 80 + host_memory_gib: 100 + mounts: {/data: /data} + image: ghcr.io/michaellchung/yeto-miles-ports@sha256:5b9dc2a9e6e315d73dabfbc6faae428b59cc01aabec6189645a5c8c80d6e3721 + runtime: docker + roles: [sampler] diff --git a/scripts/modal_records_endpoint.py b/scripts/modal_records_endpoint.py new file mode 100644 index 00000000..7749929f --- /dev/null +++ b/scripts/modal_records_endpoint.py @@ -0,0 +1,265 @@ +#!/usr/bin/env python +"""A records-proxy model endpoint on Modal for Harbor agent trials (S25 N10.7 smoke). + +One GPU container runs the SGLang server adapter (``SGLangServerEngine``, TP1) +on a checkpoint from a Modal Volume and the N10.5 records proxy in front of it +(``RecordsProxy``: Responses -> Chat, token / loss mask / log-prob records). +The proxy is reachable from outside through a Modal tunnel, under a secret path +prefix only (``/k//...``; the proxy itself has no authentication). The +tunnel URL is put into a Modal Dict; the container serves until the Dict's +``stop`` key is set or ``max_serve_s`` passes, then writes every recorded +trajectory to the results Volume. + + python scripts/modal_records_endpoint.py launch --run-name g2-a --secret-file /tmp/secret + python scripts/modal_records_endpoint.py stop --run-name g2-a + +``launch`` spawns the call and returns (killing a client that blocks on +``.remote()`` cancels the input). The GPU type is pinned (``H100!``) and +asserted inside the container. +""" + +from __future__ import annotations + +import argparse +import json +import os +import subprocess +import sys +import threading +import time +from pathlib import Path +from typing import Any + +APP_NAME = "yeto-records-endpoint" +MODELS_VOLUME = "yeto-models-qwen38" +MODELS_MOUNT = "/mnt/models" +RESULTS_VOLUME = "yeto-s25-records-endpoint" +RESULTS_MOUNT = "/mnt/results" +CONTROL_DICT = "yeto-records-endpoint-control" +CONTAINER_WORKDIR = "/root/sky_workdir" # yeto source in the container (modal_eval_island.WORKDIR_MOUNT) +PROXY_PORT = 8000 +SGLANG_PORT = 30000 + +DEFAULTS: dict[str, Any] = { + "model_subdir": "hf/Qwen3.8-27B/1d4bf0f2", + "served_model_name": "Qwen3.8-27B", + "template": "qwen38", + "template_kwargs": {"reasoning_effort": "medium"}, + "max_model_len": 32768, + "mem_fraction_static": 0.88, + "signed_sampling": {"temperature": 1.0, "top_p": 0.95, "top_k": 20, "max_tokens": 8192}, + "assert_gpu_name": "H100", + "max_serve_s": 3600, +} + +# Model presets: the Modal Volume holding the checkpoint and the DEFAULTS keys that differ. +MODELS: dict[str, dict[str, Any]] = { + "qwen38-27b": {"volume": MODELS_VOLUME, "config": {}}, + # S25 N10.9: full-parameter agentic RL model. Sampling is the trainer's + # (temperature 1, no top-p / top-k truncation), so the pass rate measured + # here is the one the RL rollouts see. + "qwen3-4b-2507": {"volume": "yeto-fn-models", "config": { + "model_subdir": "hf/Qwen3-4B-Instruct-2507/cdbee75f", + "served_model_name": "Qwen3-4B-Instruct-2507", + "template": "qwen3_2507", + "template_kwargs": {}, + "signed_sampling": {"temperature": 1.0, "top_p": 1.0, "top_k": -1, "max_tokens": 4096}, + }}, +} + + +def _gpu_names() -> list[str]: + out = subprocess.run(["nvidia-smi", "--query-gpu=name", "--format=csv,noheader"], + capture_output=True, text=True, timeout=30).stdout + return [n.strip() for n in out.splitlines() if n.strip()] + + +def serve(config: dict[str, Any]) -> dict[str, Any]: + """Inside the container (module doc).""" + import asyncio + import dataclasses + + import modal + + sys.path.insert(0, CONTAINER_WORKDIR) + from yeto.cloud.modal_eval_island import _image_paths + + _image_paths() + from aiohttp import web + + import sglang + + from yeto.rl.adapters.sampler_engines.server import ServerConfig + from yeto.rl.adapters.sampler_engines.sglang_server import SGLangOptions, SGLangServerEngine + from yeto.rl.harness.chat_templates import TEMPLATES, ChatTemplate + from yeto.rl.harness.gateway.proxy import RecordsProxy + + names = _gpu_names() + if not names or not all(config["assert_gpu_name"] in n for n in names): + raise RuntimeError(f"asked for {config['assert_gpu_name']}, got {names}") + run_dir = Path(RESULTS_MOUNT) / config["run_name"] + run_dir.mkdir(parents=True, exist_ok=True) + results = modal.Volume.from_name(RESULTS_VOLUME) + control = modal.Dict.from_name(CONTROL_DICT, create_if_missing=True) + began = time.time() + meta: dict[str, Any] = {"gpu_names": names, "sglang": sglang.__version__, "config": { + k: v for k, v in config.items() if k != "secret"}, "container": os.environ.get("MODAL_TASK_ID")} + + model_dir = Path(MODELS_MOUNT) / config["model_subdir"] + engine = SGLangServerEngine( + ServerConfig(base_dir=model_dir, card_model=names[0], device_family="cuda", + engine_version=sglang.__version__, precision="bf16", quant_scheme="bf16", + topology="tp8", max_model_len=int(config["max_model_len"]), port=SGLANG_PORT, + python=sys.executable, start_timeout_s=1800.0, + extra_args=("--mem-fraction-static", str(config["mem_fraction_static"]))), + SGLangOptions(tp=1, served_model_name=config["served_model_name"])) + report = engine.start(None, 0) + meta["engine_start_s"] = round(report.pause_s, 1) + spec = TEMPLATES[config["template"]] + spec = dataclasses.replace(spec, template_kwargs={**spec.template_kwargs, **config["template_kwargs"]}) + template = ChatTemplate(spec, ChatTemplate.load(spec.name, tokenizer_path=str(model_dir)).tokenizer) + proxy = RecordsProxy(engine, template, model=config["served_model_name"], + signed_sampling=config["signed_sampling"]) + trajectories: set[str] = set() + + @web.middleware + async def note(request: web.Request, handler: Any) -> web.StreamResponse: + tid = request.match_info.get("tid") + if tid: + trajectories.add(tid) + return await handler(request) + + inner = proxy.app() + inner.middlewares.append(note) + outer = web.Application() + outer.add_subapp(f"/k/{config['secret']}/", inner) + + loop = asyncio.new_event_loop() + + async def start_server() -> web.AppRunner: + runner = web.AppRunner(outer) + await runner.setup() + await web.TCPSite(runner, "0.0.0.0", PROXY_PORT).start() + return runner + + async def snapshot() -> str: # runs on the server loop: no concurrent mutation + records = {} + for tid in sorted(trajectories): + try: + records[tid] = proxy.records(tid) + except Exception as exc: # noqa: BLE001 - keep the others + records[tid] = {"error": repr(exc)} + return json.dumps(records) + + def dump() -> None: + text = asyncio.run_coroutine_threadsafe(snapshot(), loop).result(timeout=120) + (run_dir / "records.json").write_text(text) + (run_dir / "meta.json").write_text(json.dumps({**meta, "served_s": round(time.time() - began, 1)})) + results.commit() + + threading.Thread(target=loop.run_forever, daemon=True).start() + runner = asyncio.run_coroutine_threadsafe(start_server(), loop).result(timeout=60) + url_key, stop_key = f"{config['run_name']}/url", f"{config['run_name']}/stop" + try: + with modal.forward(PROXY_PORT) as tunnel: + control[url_key] = tunnel.url + print(f"[records-endpoint] serving, engine start {meta['engine_start_s']} s", flush=True) + deadline = began + float(config["max_serve_s"]) + while time.time() < deadline and not control.get(stop_key, False): + time.sleep(15) + dump() + finally: + control.pop(url_key, None) + dump() + asyncio.run_coroutine_threadsafe(runner.cleanup(), loop).result(timeout=60) + engine.stop() + return {"run_name": config["run_name"], "trajectories": len(trajectories)} + + +def check(config: dict[str, Any]) -> dict[str, Any]: + """CPU container: image imports, checkpoint files and the template, no engine.""" + sys.path.insert(0, CONTAINER_WORKDIR) + from yeto.cloud.modal_eval_island import _image_paths + + _image_paths() + import aiohttp + import sglang + + from yeto.rl.adapters.sampler_engines.sglang_server import SGLangServerEngine # noqa: F401 + from yeto.rl.harness.chat_templates import ChatTemplate, check_alignment + from yeto.rl.harness.gateway.proxy import RecordsProxy # noqa: F401 + + model_dir = Path(MODELS_MOUNT) / config["model_subdir"] + template = ChatTemplate.load(config["template"], tokenizer_path=str(model_dir)) + messages = [{"role": "system", "content": "a"}, {"role": "system", "content": "b"}, + {"role": "user", "content": "hi"}] + return {"sglang": sglang.__version__, "aiohttp": aiohttp.__version__, + "shards": len(list(model_dir.glob("*.safetensors"))), + "stop_token_ids": template.stop_token_ids, + "aligned": check_alignment(template, messages).aligned} + + +def build_app(workdir: str, gpu: str, timeout_s: int, models_volume: str = MODELS_VOLUME) -> tuple[Any, Any, Any]: + import modal + + sys.path.insert(0, workdir) + from yeto.cloud.modal_eval_island import WORKDIR_MOUNT + from yeto.rl.adapters.miles.pins import MILES_NEXT_IMAGE + + assert WORKDIR_MOUNT == CONTAINER_WORKDIR + + image = (modal.Image.from_registry(MILES_NEXT_IMAGE.removeprefix("docker:")) + .env({"HOME": "/root", "PYTHONUNBUFFERED": "1"}) + .add_local_dir(workdir, WORKDIR_MOUNT, copy=False, + ignore=[".git", ".venv", ".claude", "**/__pycache__", "**/*.pyc", "s1-runs", + "syncer/target"])) + app = modal.App(APP_NAME) + models = modal.Volume.from_name(models_volume) + results = modal.Volume.from_name(RESULTS_VOLUME, create_if_missing=True) + fn = app.function(image=image, gpu=gpu, cpu=8.0, memory=65536, timeout=timeout_s, retries=0, + volumes={MODELS_MOUNT: models.read_only(), RESULTS_MOUNT: results}, + name="serve", serialized=True)(serve) + checker = app.function(image=image, cpu=2.0, memory=8192, timeout=600, retries=0, + volumes={MODELS_MOUNT: models.read_only()}, name="check")(check) + return app, fn, checker + + +def main() -> None: + parser = argparse.ArgumentParser(description=__doc__.splitlines()[0]) + parser.add_argument("action", choices=("check", "launch", "stop", "url")) + parser.add_argument("--run-name", required=True) + parser.add_argument("--secret-file", type=Path, help="launch: file holding the path secret") + parser.add_argument("--gpu", default="H100!") + parser.add_argument("--model", choices=sorted(MODELS), default="qwen38-27b") + parser.add_argument("--max-serve-s", type=int, default=DEFAULTS["max_serve_s"]) + parser.add_argument("--workdir", default=str(Path(__file__).resolve().parents[1])) + args = parser.parse_args() + import modal + + control = modal.Dict.from_name(CONTROL_DICT, create_if_missing=True) + if args.action == "stop": + control[f"{args.run_name}/stop"] = True + return + if args.action == "url": + print(control.get(f"{args.run_name}/url", "")) + return + preset = MODELS[args.model] + defaults = {**DEFAULTS, **preset["config"]} + if args.action == "check": + app, _, checker = build_app(args.workdir, args.gpu, timeout_s=600, models_volume=preset["volume"]) + with modal.enable_output(), app.run(): + print(json.dumps(checker.remote(defaults))) + return + secret = args.secret_file.read_text().strip() + if len(secret) < 24: + raise SystemExit("the path secret must have at least 24 characters") + config = {**defaults, "run_name": args.run_name, "secret": secret, "max_serve_s": args.max_serve_s} + app, fn, _ = build_app(args.workdir, args.gpu, timeout_s=args.max_serve_s + 1800, + models_volume=preset["volume"]) + with modal.enable_output(), app.run(detach=True): + call = fn.spawn(config) + print(json.dumps({"function_call_id": call.object_id, "run_name": args.run_name})) + + +if __name__ == "__main__": + main() diff --git a/tests/test_agentic_harbor.py b/tests/test_agentic_harbor.py index e3039cf7..e53420c3 100644 --- a/tests/test_agentic_harbor.py +++ b/tests/test_agentic_harbor.py @@ -10,6 +10,7 @@ import importlib import sys import types +from dataclasses import replace from datetime import datetime, timedelta from pathlib import Path from types import SimpleNamespace @@ -286,6 +287,7 @@ class Trial(_FakeTrial): @classmethod async def create(cls, config): seen["version"] = config.agent.kwargs["version"] + seen["summary"] = config.agent.kwargs["reasoning_summary"] return cls(config) Trial.result = _result({"reward": 1.0}) @@ -295,6 +297,7 @@ async def create(cls, config): out = _run(mod.run("http://ss:30000/sessions/s1", metadata={"instance_id": "fix-git"})) assert example.calls == ["key", ("codex", "http://ss:30000/sessions/s1/v1")] assert seen["version"] == "0.145.0" # Harbor's default would be @latest + assert seen["summary"] == "none" # the gateway refuses reasoning summaries assert (out[OUTCOME_KEY], out["reward"]) == ("ok", 1.0) Trial.result = _result(None, "EnvironmentStartTimeoutError") @@ -422,7 +425,7 @@ def test_verl_loop_token_mask_alignment_with_fake_trial(verl_loop): loop_mod, _ = verl_loop from yeto.rl.harness.harbor import HarborHarness - turns = [_turn([5, 6, 7], [40, 41, 42], [-0.1, -0.2, -0.3]), + turns = [replace(_turn([5, 6, 7], [40, 41, 42], [-0.1, -0.2, -0.3]), finish_reason="length"), _turn([5, 6, 7, 40, 41, 42, 90, 91], [43, 2], [-0.4, -0.5])] recorder = _Recorder(turns) _FakeTrial.result, _FakeTrial.configs = _result({"reward": 1.0}), [] @@ -443,6 +446,7 @@ def test_verl_loop_token_mask_alignment_with_fake_trial(verl_loop): assert out.prompt_ids + out.response_ids == list(turns[-1].prompt_ids + turns[-1].completion_ids) assert out.reward_score == 1.0 and out.num_turns == 4 assert out.extra_fields[OUTCOME_KEY] == "ok" and out.extra_fields["truncated"] is False + assert out.extra_fields["length_capped_turns"] == 1 # the first call hit its token cap def test_verl_loop_aborts_on_misaligned_tokens_and_infra_errors(verl_loop): diff --git a/tests/test_codepack.py b/tests/test_codepack.py index 3d09fb3f..351cbfea 100644 --- a/tests/test_codepack.py +++ b/tests/test_codepack.py @@ -77,6 +77,63 @@ def test_verify_tree_rejects_changed_tree(repo: Path) -> None: codepack.verify_tree(tree, pack.digest) +def test_unpack_keeps_a_tree_another_launch_placed_first(repo: Path, monkeypatch) -> None: + # S25 T-EVAL: two launches from one checkout unpack the same digest at once; the + # later rename finds the earlier tree in place (ENOTEMPTY) and must reuse it + pack, data = codepack.build_pack(repo) + cache = repo.parent / "cache" + first = codepack.unpack(data, pack.digest, cache) + real_is_dir = Path.is_dir + calls = {"n": 0} + + def racing_is_dir(self: Path) -> bool: # the first existence check misses the tree + if self == first and calls["n"] == 0: + calls["n"] += 1 + return False + return real_is_dir(self) + + monkeypatch.setattr(Path, "is_dir", racing_is_dir) + assert codepack.unpack(data, pack.digest, cache) == first + assert codepack.verify_tree(first, pack.digest) == pack.digest + assert [p.name for p in cache.iterdir()] == [first.name] # no staging dir left behind + + +def test_an_unpacked_tree_can_be_zipped_for_a_remote_api_server(repo: Path) -> None: + # T-C1b: the pack's entries carry mtime 0; sky zips the workdir it uploads to a + # remote API server and ZIP refuses times before 1980 + import os + import zipfile + + pack, data = codepack.build_pack(repo) + cache = repo.parent / "cache" + tree = codepack.unpack(data, pack.digest, cache) + files = [p for p in tree.rglob("*") if p.is_file()] + assert files and all(p.stat().st_mtime == codepack.UNPACKED_MTIME for p in files) + with zipfile.ZipFile(repo.parent / "upload.zip", "w") as zf: + for p in files: + zf.write(p, p.relative_to(tree)) + os.utime(files[0], (0, 0)) # a tree unpacked before the fix is stamped on reuse + assert codepack.unpack(data, pack.digest, cache) == tree + assert files[0].stat().st_mtime == codepack.UNPACKED_MTIME + + +def test_a_tree_that_lost_its_executable_bits_still_verifies(repo: Path) -> None: + # T-C1b: sky's upload to a remote API server unpacks the workdir zip without file + # modes; the island restores the bits the pack lists before it checks the digest + import os + + script = repo / "run.sh" + script.write_text("#!/bin/sh\n") + script.chmod(0o755) + subprocess.run(["git", "-C", str(repo), "add", "run.sh"], check=True) + pack, data = codepack.build_pack(repo) + tree = codepack.unpack(data, pack.digest, repo.parent / "cache") + assert "run.sh" in codepack.read_pack_info(tree)["executables"] + os.chmod(tree / "run.sh", 0o644) + assert codepack.verify_tree(tree, pack.digest) == pack.digest + assert os.stat(tree / "run.sh").st_mode & 0o111 + + def test_unpack_rejects_wrong_bytes(repo: Path) -> None: pack, data = codepack.build_pack(repo) with pytest.raises(CodePackError, match="does not have digest"): diff --git a/tests/test_harness_gateway.py b/tests/test_harness_gateway.py index 5d90913d..7b72cb4a 100644 --- a/tests/test_harness_gateway.py +++ b/tests/test_harness_gateway.py @@ -115,7 +115,7 @@ def test_signed_sampling_cannot_be_overridden_and_unsupported_shapes_fail_closed run(gw.handle("chat", {**base, "model": "other"}, trajectory_id="x3")) resp = {"store": False, "input": [{"type": "message", "role": "user", "content": [{"type": "input_text", "text": "u"}]}]} for bad in ({**resp, "store": True}, {**resp, "parallel_tool_calls": True}, {**resp, "reasoning": {"summary": "auto"}}, - {**resp, "tools": [{"type": "web_search"}]}, {**resp, "input": resp["input"] + [{"type": "compaction"}]}): + {**resp, "tools": [{"type": "local_shell"}]}, {**resp, "input": resp["input"] + [{"type": "compaction"}]}): with pytest.raises(UnsupportedShape): run(gw.handle("responses", bad, trajectory_id="x4")) with pytest.raises(UnsupportedShape): diff --git a/tests/test_miles_records_generate.py b/tests/test_miles_records_generate.py new file mode 100644 index 00000000..22a1a86e --- /dev/null +++ b/tests/test_miles_records_generate.py @@ -0,0 +1,102 @@ +"""S25 N10.9: the Miles Harbor path with the records proxy as recorder (CPU). + +Real pieces: ``records_generate.run_sample`` (outer layer, outcome rules, +``join_turns``, sample filling), ``ProxyTokenRecorder`` over HTTP, the gateway +and recording backend, ``SGLangRouterEngine`` (its HTTP call replaced by the +scripted model). Fakes: the agent (Codex-like Responses loop), Miles' Sample. +""" + +from __future__ import annotations + +import asyncio +import copy +from types import SimpleNamespace + +import pytest +from test_records_proxy import SCRIPT, SIGNED, ScriptedModel, _template +from test_verl_rollout_engine import IDENTITY, ExitOnErrorAgent, MalformedSecondReply +from test_verl_records_recorder import CodexLikeAgent + +from yeto.rl.adapters.miles import records_generate as rg +from yeto.rl.agentic import OUTCOME_KEY +from yeto.rl.harness.gateway.proxy import RecordsProxy +from yeto.rl.harness.gateway.recorder import ProxyTokenRecorder + +pytest.importorskip("aiohttp") + +MILES_SAMPLING = {"temperature": 1.0, "top_p": 1.0, "top_k": -1, "max_new_tokens": 256} +ARGS = SimpleNamespace(rollout_max_response_len=4096, n_samples_per_prompt=4, grpo_std_normalization=True) + + +@pytest.fixture(scope="module") +def qwen(): + return _template("qwen3") + + +def _router_http(model: ScriptedModel, calls: list): + def http(url: str, body: dict) -> dict: + calls.append(body) + ids = model(tuple(body["input_ids"]), 0) + return {"text": "", "meta_info": {"output_token_logprobs": [[-0.5, t, None] for t in ids], + "finish_reason": {"type": "stop"}}} + return http + + +def _run(qwen, agent, model, monkeypatch): + monkeypatch.setenv("HARBOR_ENV_TYPE", "boat") + calls: list = [] + engine = rg.SGLangRouterEngine("http://router", max_model_len=8192, identity=IDENTITY, + http=_router_http(model, calls)) + recorder = ProxyTokenRecorder(RecordsProxy(engine, qwen, model="policy", signed_sampling=SIGNED, max_chains=1)) + sample = SimpleNamespace(metadata={"instance_id": "fix-git", "agent_name": "codex"}, tokens=[], + response_length=0, loss_mask=None, rollout_log_probs=None, status=None, + reward=None, remove_sample=False) + + async def go(): + try: + return await rg.run_sample(sample, MILES_SAMPLING, ARGS, recorder=recorder, harness=agent) + finally: + await recorder.aclose() + + return asyncio.run(go()), calls + + +def test_sample_carries_the_tokens_the_router_sampled(qwen, monkeypatch): + out, calls = _run(qwen, CodexLikeAgent(), ScriptedModel(qwen, copy.deepcopy(SCRIPT)), monkeypatch) + assert out.metadata[OUTCOME_KEY] == "ok" and out.reward == 1.0 and out.status.name == "COMPLETED" + assert len(calls) == 3 and calls[0]["return_logprob"] is True + assert calls[0]["sampling_params"]["stop_token_ids"] == list(qwen.stop_token_ids) + assert len(out.loss_mask) == len(out.rollout_log_probs) == out.response_length + assert out.tokens[:len(calls[-1]["input_ids"])] == calls[-1]["input_ids"] # trained == sampled + assert all(lp == -0.5 for lp, m in zip(out.rollout_log_probs, out.loss_mask) if m) + assert out.metadata["model_calls"] == 3 and out.metadata["truncated"] is False + + +def test_malformed_reply_is_a_trained_zero(qwen, monkeypatch): + out, _ = _run(qwen, ExitOnErrorAgent(), MalformedSecondReply(qwen, copy.deepcopy(SCRIPT)), monkeypatch) + assert out.metadata[OUTCOME_KEY] == "policy_limit" and out.reward == 0.0 + assert sum(out.loss_mask) > 0 and not out.remove_sample + + +def test_infrastructure_exit_is_aborted_and_never_trained(qwen, monkeypatch): + out, _ = _run(qwen, ExitOnErrorAgent("EnvironmentStartTimeoutError"), + MalformedSecondReply(qwen, copy.deepcopy(SCRIPT)), monkeypatch) + assert out.metadata[OUTCOME_KEY] == "aborted" and out.status.name == "ABORTED" + assert out.remove_sample and out.loss_mask == [0] and out.reward == 0.0 + + +def _s(reward, outcome="ok"): + return SimpleNamespace(reward=reward, metadata={OUTCOME_KEY: outcome}, remove_sample=False) + + +def test_reward_post_process_leaves_aborted_out_of_the_group_baseline(): + samples = [_s(1.0), _s(0.0), _s(0.0, "aborted"), _s(0.0, "policy_limit"), # group 1 + _s(0.0), _s(0.0), _s(0.0), _s(0.0)] # group 2: no variance + raw, adv = rg.reward_post_process(ARGS, samples) + assert raw == [1.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0] + mean, std = 1 / 3, (((2 / 3) ** 2 + 2 * (1 / 3) ** 2) / 2) ** 0.5 + assert adv[0] == pytest.approx((1 - mean) / (std + 1e-6)) and adv[1] == pytest.approx(-mean / (std + 1e-6)) + assert adv[2] == 0.0 and samples[2].remove_sample and not samples[0].remove_sample + assert adv[4:] == [0.0] * 4 # all-zero group: no gradient (expected for most TB2.1 groups) + with pytest.raises(ValueError, match="whole groups"): + rg.reward_post_process(ARGS, samples[:3]) diff --git a/tests/test_nebius_reclaim.py b/tests/test_nebius_reclaim.py index 43f1a358..9cf6fb40 100644 --- a/tests/test_nebius_reclaim.py +++ b/tests/test_nebius_reclaim.py @@ -8,6 +8,8 @@ import threading from types import SimpleNamespace +import pytest + from yeto import launcher from yeto.cloud import droppable as dr from yeto.cloud import nebius_reclaim as nr @@ -156,3 +158,28 @@ def test_launcher_installs_hook_only_for_nebius_droppable(monkeypatch): assert launcher.maybe_install_nebius_reclaim_hook(SimpleNamespace(envs=envs), (7, "H"), "c") is None assert launcher.maybe_install_nebius_reclaim_hook(neb, (7, None), "c2") is False assert len(seen) == 1 + + +def test_host_runner_takes_credentials_from_a_remote_api_server(monkeypatch, tmp_path): + # T-C1b: with the sky API server on another machine the cluster yaml is not here + pytest.importorskip("sky") + from sky.utils import command_runner + + from yeto.cloud import nebius_reclaim as nr + + asked = [] + key = tmp_path / "c.key" + key.write_text("key\n") + + def server_credentials(cluster): + asked.append(cluster) + return {"ssh_user": "ubuntu", "ssh_private_key": str(key), + "ssh_control_name": "c", "ssh_proxy_command": None, "docker_user": "root"} + + monkeypatch.setattr(nr, "_server_credentials", server_credentials) + handle = SimpleNamespace(cluster_yaml=str(tmp_path / "missing.yml"), cluster_name="c-1", + ssh_user="ubuntu", head_ip="10.0.0.5") + runner = nr._host_runner(handle) + assert asked == ["c-1"] and isinstance(runner, command_runner.SSHCommandRunner) + assert runner.ip == "10.0.0.5" and runner.ssh_user == "ubuntu" + assert getattr(runner, "docker_user", None) is None # the HOST, not the container diff --git a/tests/test_records_proxy.py b/tests/test_records_proxy.py index e9a77ace..918793a7 100644 --- a/tests/test_records_proxy.py +++ b/tests/test_records_proxy.py @@ -12,6 +12,7 @@ import asyncio import copy import json +import os from typing import Any import pytest @@ -19,10 +20,11 @@ from yeto.rl.adapters.sampler_engines.fake import FakeEngine from yeto.rl.compaction import CompactionConfig from yeto.rl.engine.sampler_engine import SamplingParams, TokenSamplerEngine -from yeto.rl.harness.chat_templates import ChatTemplate, TemplateError, check_alignment, parse_glm, parse_qwen3 +from yeto.rl.harness.chat_templates import ChatTemplate, TemplateError, check_alignment, parse_glm, parse_qwen3, parse_qwen38 from yeto.rl.harness.gateway import Gateway, GatewayConfig from yeto.rl.harness.gateway.proxy import RecordsProxy, chat_tools from yeto.rl.harness.gateway.records import RecordingBackend, RecordsMismatch +from yeto.rl.harness.gateway.translate import UnsupportedShape, responses_to_messages SIGNED = {"temperature": 1.0, "top_p": 1.0, "max_tokens": 256} TOOLS = [{"type": "function", "name": "shell", "description": "run a command", @@ -36,6 +38,16 @@ def _template(name: str) -> ChatTemplate: pytest.skip(f"{name} tokenizer not available: {exc}") +# The Qwen3.8-27B tokenizer files (not on the HF hub cache by default). +QWEN38_TOKENIZER = os.environ.get("YETO_QWEN38_TOKENIZER", "") + + +def _template_at(name: str, path: str) -> ChatTemplate: + if not path: + pytest.skip(f"set YETO_QWEN38_TOKENIZER to run the {name} template tests") + return ChatTemplate.load(name, tokenizer_path=path) + + @pytest.fixture(scope="module") def qwen() -> ChatTemplate: return _template("qwen3") @@ -46,6 +58,14 @@ def glm() -> ChatTemplate: return _template("glm5") +@pytest.fixture(scope="module") +def qwen3_2507() -> ChatTemplate: + return _template("qwen3_2507") + + +FIXTURES = {"qwen3": "qwen", "glm5": "glm", "qwen3_2507": "qwen3_2507"} + + def call(call_id: str, cmd: str) -> dict[str, Any]: return {"id": call_id, "type": "function", "function": {"name": "shell", "arguments": json.dumps({"cmd": cmd})}} @@ -142,13 +162,15 @@ def run_codex(proxy: Any, trajectory_id: str, *, turns: int = 3) -> list[dict[st # ---------------------------------------------------------------------- N10.5 records -@pytest.mark.parametrize("name", ["qwen3", "glm5"]) +@pytest.mark.parametrize("name", ["qwen3", "glm5", "qwen3_2507"]) def test_codex_responses_loop_records_one_aligned_stream(name, request): - template = request.getfixturevalue("qwen" if name == "qwen3" else "glm") + template = request.getfixturevalue(FIXTURES[name]) engine = engine_for(ScriptedModel(template, SCRIPT)) proxy = RecordsProxy(engine, template, model="policy", signed_sampling=SIGNED, max_chains=1) outputs = run_codex(proxy, "t1") - assert [o["type"] for o in outputs[0]["output"]] == ["reasoning", "function_call"] + # Qwen3-4B-Instruct-2507 has no reasoning block: its template renders no reasoning. + first = ["function_call"] if name == "qwen3_2507" else ["reasoning", "function_call"] + assert [o["type"] for o in outputs[0]["output"]] == first assert outputs[-1]["output"][-1]["content"][0]["text"] == "The file says hello." records = proxy.records("t1") @@ -339,9 +361,21 @@ def test_compaction_on_a_template_that_drops_history_reasoning_fails_closed(qwen ] -@pytest.mark.parametrize("name", ["qwen3", "glm5"]) +def test_codex_namespace_and_hosted_tools_are_left_out_other_types_refused(): + # Codex 0.145.0 sends its sub-agent tools as a namespace and a hosted web_search tool. + namespace = {"type": "namespace", "name": "multi_agent_v1", "tools": [TOOLS[0]]} + tools = chat_tools("responses", {"tools": [*TOOLS, namespace, {"type": "web_search"}]}) + assert [t["function"]["name"] for t in tools] == ["shell"] + with pytest.raises(UnsupportedShape): + chat_tools("responses", {"tools": [{"type": "custom", "name": "apply_patch"}]}) + body = {"store": False, "parallel_tool_calls": False, "tools": [*TOOLS, namespace, {"type": "web_search"}], + "input": [{"type": "message", "role": "user", "content": [{"type": "input_text", "text": "hi"}]}]} + assert responses_to_messages(body) == [{"role": "user", "content": "hi"}] + + +@pytest.mark.parametrize("name", ["qwen3", "glm5", "qwen3_2507"]) def test_template_rendering_matches_tokenizer_ids_token_by_token(name, request): - template = request.getfixturevalue("qwen" if name == "qwen3" else "glm") + template = request.getfixturevalue(FIXTURES[name]) tools = chat_tools("responses", {"tools": TOOLS}) for end in range(2, len(CONVERSATION) + 1): report = check_alignment(template, CONVERSATION[:end], tools=tools) @@ -366,6 +400,33 @@ def test_parsers_fail_closed_on_malformed_tool_calls(): parse_qwen3("\nnot json\n") with pytest.raises(TemplateError): parse_glm("ra1") + with pytest.raises(TemplateError): + parse_qwen38('r\n\n\n\n{"name": "f", "arguments": {}}\n') + + +def test_qwen38_merges_codex_system_messages_and_template_refusals_fail_closed(): + template = _template_at("qwen38", QWEN38_TOKENIZER) + tools = chat_tools("responses", {"tools": TOOLS}) + # Codex: instructions + a developer message, both "system" in the gateway form. + messages = [{"role": "system", "content": "a"}, {"role": "system", "content": "b"}, {"role": "user", "content": "u"}] + text = template.render(messages, tools=tools, add_generation_prompt=True) + assert "a\n\nb" in text and text.count("<|im_start|>system") == 1 + assert check_alignment(template, messages, tools=tools).aligned + with pytest.raises(TemplateError, match="refused"): + template.render([*messages, {"role": "system", "content": "late"}], add_generation_prompt=True) + + +def test_qwen38_parser_reads_the_xml_tool_call_format(): + # The Qwen3.8 template's own rendering of one assistant turn (after "\n"). + text = ("look first\n\n\nChecking.\n\n\n\n" + "\nls -la\nwc -l x\n\n\n3\n\n" + "\n\n\n\n\n") + parsed = parse_qwen38(text) + assert parsed["reasoning_content"] == "look first" + assert parsed["content"] == "Checking." + assert [c["function"]["name"] for c in parsed["tool_calls"]] == ["exec_command", "noop"] + assert json.loads(parsed["tool_calls"][0]["function"]["arguments"]) == {"cmd": "ls -la\nwc -l x", "n": 3} + assert parse_qwen38("still thinking") == {"role": "assistant", "content": "", "reasoning_content": "still thinking"} # ---------------------------------------------------------------------- adapters diff --git a/tests/test_rl_buffer_rollout.py b/tests/test_rl_buffer_rollout.py index 1079914c..d8febda6 100644 --- a/tests/test_rl_buffer_rollout.py +++ b/tests/test_rl_buffer_rollout.py @@ -370,3 +370,29 @@ def test_full_parameter_trainer_stops_at_a_scheduled_cut_and_continues_without_r assert trained == list(range(2 * rounds)) # every group once, in order: none trained twice for name, value in params(reference).items(): assert torch.equal(value, params(rank_b)[name]), name + + +def test_reader_round_reports_its_stale_groups_as_round_metadata(monkeypatch): + """T-EVAL E4: the groups the reader passed over this round reach the driver as + per-round metadata (driver event rl_buffer_stale_skipped), before the + end-of-round hooks put the round's main record.""" + from types import SimpleNamespace + + calls = [] + registry = {"all": lambda args, samples, source: calls.append(("all",))} + plugin = _import_plugin_with_stub_miles(monkeypatch, registry) + monkeypatch.setattr(plugin, "record_round_metadata", + lambda args, rid, **counts: calls.append(("meta", rid, counts))) + monkeypatch.setattr(plugin, "Sample", SimpleNamespace(from_dict=lambda d: d)) + monkeypatch.setattr(plugin, "RolloutFnTrainOutput", lambda **kw: kw) + monkeypatch.setattr(plugin, "offpolicy_source_event", lambda *a, **kw: {}) + monkeypatch.setattr(plugin, "flat_metrics", lambda source: {}) + group = [sample(0, 0, 4), sample(1, 0, 4)] + report = br.TakeReport(rollout_id=2, groups=1, stale_groups=3, start=0, end=4) + reader = plugin._BufferRound.__new__(plugin._BufferRound) + reader.args = SimpleNamespace(rollout_sample_filter_path=None, rollout_all_samples_process_path="all") + reader.groups, reader.samples_per_group = 1, 2 + reader.reader = SimpleNamespace(take=lambda rid, n, current_version: ([(3, group)], report)) + out = reader(SimpleNamespace(evaluation=False, rollout_id=2, weight_version=4)) + assert calls == [("meta", 2, {"buffer_stale_groups": 3}), ("all",)] + assert out["metrics"]["sample_buffer/stale_groups"] == 3.0 diff --git a/tests/test_rl_hetero_sampler_engines.py b/tests/test_rl_hetero_sampler_engines.py index 7f2a5a44..e6c41ded 100644 --- a/tests/test_rl_hetero_sampler_engines.py +++ b/tests/test_rl_hetero_sampler_engines.py @@ -624,3 +624,43 @@ def test_stop_sends_sigterm(tmp_path): engine.stop() assert process.signals == [signal.SIGTERM] and engine.last_exit["normal"] engine.stop() # idempotent + + +def test_sglang_topology_names_the_real_layout(tmp_path): + """S25 N2-fix: the source label is tp/dp with N = --tp, not tp8 for every layout.""" + spawned = [] + http = _Http({("GET", "/health"): {}}) + one = SGLangServerEngine(_config(tmp_path, topology="tp1"), SGLangOptions(tp=1), + spawn=lambda a, e: spawned.append(a) or _Process(), http=http, sleep=lambda s: None) + assert one.identity.topology == "tp1" + one.start(None, 0) + argv = spawned[0] + assert argv[argv.index("--tp") + 1] == "1" and "--dp" not in argv + dp = SGLangServerEngine(_config(tmp_path, topology="dp4"), SGLangOptions(tp=4), + spawn=lambda a, e: spawned.append(a) or _Process(), http=http, sleep=lambda s: None) + dp.start(None, 0) + assert spawned[1][spawned[1].index("--dp") + 1] == "4" + for topology, tp in (("tp8", 1), ("tp1", 8), ("pd8", 8), ("tp", 8)): + with pytest.raises(SamplerEngineError, match="does not match TP"): + SGLangServerEngine(_config(tmp_path, topology=topology), SGLangOptions(tp=tp)) + + +def test_probe_without_pairs_passes_every_load(): + """S25 N2-fix: no fixed prompts and no probe pairs: the second load divided by zero.""" + checks = PostLoadChecks(expected_expert_matrices=None, fixed_prompts=()) + for _ in range(2): + result = checks.logprob_probe(engine=None) + checks.accept() + assert result == {"name": "logprob_probe", "passed": True, "detail": "no probe pairs"} + + +def test_server_process_ignores_the_reclaim_sigint(): + """T-E2V: a cloud reclaim sends SIGINT to every process of the container; the + engine server must survive it so the island can finish its round and leave.""" + import sys + + from yeto.rl.adapters.sampler_engines.server import popen_spawn + + probe = popen_spawn([sys.executable, "-c", "import signal, sys; " + "sys.exit(7 if signal.getsignal(signal.SIGINT) is signal.SIG_IGN else 3)"], {}) + assert probe.wait(timeout=60) == 7 diff --git a/tests/test_rl_merged_release.py b/tests/test_rl_merged_release.py new file mode 100644 index 00000000..c0b10287 --- /dev/null +++ b/tests/test_rl_merged_release.py @@ -0,0 +1,351 @@ +"""S25 C3: synced trainer islands feed sampler islands (yeto.rl.engine.merged_release). + +Unit level: buffer lanes, the lane cursor log, the release guard, the publisher +wrappers, the role checks and the launch/island flags. The end-to-end case with +the real Rust syncer is in test_rl_synced_trainers_sampler.py.""" + +from __future__ import annotations + +import json +import shlex +from types import SimpleNamespace + +import pytest +from test_rl_algorithm_provenance import _fake_sky, _no_modal_listing # noqa: F401 + +from yeto.rl.adapters.miles import buffer_rollout as br +from yeto.rl.adapters.miles.sampler_island import GateVersionPublisher, ReleasingPublisher +from yeto.rl.engine import sample_roles +from yeto.rl.engine.merged_release import ( + BufferLane, + LaneCursorLog, + MergedReleaseGuard, + ObjectLaneCursorLog, + gate_version, + is_publisher, + lane_cursor_log, + merged_version, + merges_shards, +) +from yeto.rl.engine.policy_release import PolicyReleaseStore +from yeto.rl.engine.resume import CheckpointStore +from yeto.rl.engine.sample_buffer import FileSampleBuffer, SampleBufferError, open_trainer_buffer +from yeto.rl.engine.trainer_lifecycle import SampleIndexEntry + + +def _fill(buffer: FileSampleBuffer, n: int, version: int = 0) -> None: + for i in range(len(buffer), len(buffer) + n): + ref = buffer.write_payload(f"s{i}", br.encode_group([{"tokens": [i]}])) + buffer.append(SampleIndexEntry(f"s{i}", version, ref)) + + +# ---------------------------------------------------------------- lanes +def test_lane_parse_and_ownership(): + assert BufferLane.parse(None) is None and BufferLane.parse("0/1") is None + lane = BufferLane.parse("1/3") + assert lane == BufferLane(1, 3) and lane.env() == "1/3" + assert [p for p in range(9) if lane.owns(p)] == [1, 4, 7] + with pytest.raises(ValueError): + BufferLane(2, 2) + with pytest.raises(ValueError): + BufferLane.parse("x/2") + + +def test_two_lane_readers_split_one_buffer_without_overlap(tmp_path): + writer = FileSampleBuffer(CheckpointStore(tmp_path), batch_size=1) + _fill(writer, 7) + readers = [FileSampleBuffer(CheckpointStore(tmp_path), batch_size=1, lane=BufferLane(m, 2)) + for m in range(2)] + taken = [[s.position for s in r.select("trainer", 10)] for r in readers] + assert taken == [[0, 2, 4, 6], [1, 3, 5]] + assert readers[0].read_cursor == readers[1].read_cursor == 7 # foreign positions are passed + # a limited select stops at the limit and resumes after the last lane position + again = FileSampleBuffer(CheckpointStore(tmp_path), batch_size=1, lane=BufferLane(0, 2)) + assert [s.position for s in again.select("trainer", 2)] == [0, 2] + assert [s.position for s in again.select("trainer", 2)] == [4, 6] + + +def test_only_the_publishers_lane_merges_shards(tmp_path): + assert merges_shards(None) and merges_shards(BufferLane(0, 2)) and not merges_shards(BufferLane(1, 2)) + with pytest.raises(SampleBufferError, match="only the publisher's lane merges"): + FileSampleBuffer(CheckpointStore(tmp_path), batch_size=1, lane=BufferLane(1, 2), + merge_shards=True) + + +# ---------------------------------------------------------------- object stores (S3, OSS) +def _s3_store(tmp_path, s3): + from yeto.cloud import s3_store + from yeto.rl.engine.resume import STORE_IMPL_ENV, store_for + + env = {STORE_IMPL_ENV: s3_store.STORE_IMPL, s3_store.URI_ENV: "s3://bkt/runs/a"} + return store_for(tmp_path / "mount", environ=env, client=s3) + + +def test_lane_readers_and_cursors_on_an_s3_store_use_the_object_api(tmp_path): + from test_rl_object_store import FakeS3 + + s3 = FakeS3() + _fill(FileSampleBuffer(_s3_store(tmp_path, s3), batch_size=1), 6) + + def reader(m): + buffer = open_trainer_buffer(_s3_store(tmp_path, s3), batch_size=1, lane=BufferLane(m, 2)) + return br.BatchReader(buffer, max_policy_age=2, cursor_log=lane_cursor_log(buffer)) + + first = [reader(m) for m in range(2)] + assert all(isinstance(r.cursor_log, ObjectLaneCursorLog) for r in first) + taken = [[p for p, _ in r.take(0, 2, current_version=0)[0]] for r in first] + assert taken == [[0, 2], [1, 3]] + keys = sorted(k for _, k in s3.objects if "/lanes/" in k) + assert keys == ["runs/a/sample-buffer/lanes/0-of-2/0000000000.json", + "runs/a/sample-buffer/lanes/1-of-2/0000000000.json"] + assert not (tmp_path / "mount").exists() # nothing through the mount + # a restarted lane reader continues after its last round (at most once) + again = reader(1) + assert again.buffer.read_cursor == 4 + assert [p for p, _ in again.take(1, 1, current_version=0)[0]] == [5] + assert again.cursor_log.last_cursor() == 6 + assert len([k for _, k in s3.objects if "/lanes/1-of-2/" in k]) == 2 + + +def test_on_an_oss_store_the_publisher_merges_and_the_other_lane_reads_the_merged_index(): + from test_rl_oss_store import _entry, _objects + + from yeto.rl.engine.sample_buffer import shard_index + + objects = _objects() + sampler = FileSampleBuffer(objects, batch_size=1, index=shard_index("s0")) + for i in range(4): + ref = sampler.write_payload(f"s{i}", br.encode_group([{"tokens": [i]}])) + sampler.append(SampleIndexEntry(f"s{i}", 0, ref, appended_at=float(i))) + other = open_trainer_buffer(objects, batch_size=1, lane=BufferLane(1, 2)) + assert other.select("trainer", 10) == [] # nothing merged yet: lane 1 does not merge + publisher = open_trainer_buffer(objects, batch_size=1, lane=BufferLane(0, 2)) + assert [s.position for s in publisher.select("trainer", 10)] == [0, 2] + assert [s.entry.sample_id for s in other.select("trainer", 10)] == ["s1", "s3"] + log = lane_cursor_log(other) + log.append(0, [1, 3], other.read_cursor) + assert lane_cursor_log(open_trainer_buffer(objects, batch_size=1, lane=BufferLane(1, 2))).last_cursor() == 4 + + +def test_reader_continues_after_its_lane_cursor_and_never_retrains(tmp_path): + store = CheckpointStore(tmp_path) + writer = FileSampleBuffer(store, batch_size=1) + _fill(writer, 8) + lane = BufferLane(1, 2) + + def reader(): + buffer = FileSampleBuffer(CheckpointStore(tmp_path), batch_size=1, lane=lane) + return br.BatchReader(buffer, max_policy_age=2, cursor_log=br.open_lane_log(buffer)) + + first = reader() + groups, report = first.take(0, 2, current_version=0) + assert [p for p, _ in groups] == [1, 3] and report.end == 4 + # restart (a re-joined island): the next round continues after position 3 + second = reader() + assert second.buffer.read_cursor == 4 + groups, _ = second.take(1, 2, current_version=1) + assert [p for p, _ in groups] == [5, 7] + lines = [json.loads(line) for line in + (tmp_path / "sample-buffer" / "lanes" / "1-of-2.jsonl").read_text().splitlines()] + assert [(r["rollout_id"], r["positions"], r["cursor"]) for r in lines] == \ + [(0, [1, 3], 4), (1, [5, 7], 8)] + + +def test_lane_cursor_log_skips_a_torn_last_line(tmp_path): + log = LaneCursorLog(tmp_path, BufferLane(0, 2)) + log.append(0, [0, 2], 3) + with log.path.open("a") as fh: + fh.write('{"cursor": 9') + assert log.last_cursor() == 3 + + +def test_gate_uses_the_merged_version_once_set(tmp_path): + writer = FileSampleBuffer(CheckpointStore(tmp_path), batch_size=1) + _fill(writer, 2, version=1) + _fill(writer, 2, version=4) + reader = br.BatchReader(FileSampleBuffer(CheckpointStore(tmp_path), batch_size=1), + max_policy_age=1, wait_s=0.0) + reader.gate_version = 5 # merged version 5: version 1 is 4 behind -> passed over + groups, report = reader.take(0, 2, current_version=1) + assert [p for p, _ in groups] == [2, 3] and report.stale_groups == 2 + assert gate_version(3, None) == 3 and gate_version(3, 7) == 7 + + +def test_executor_gate_version_setter_reaches_the_active_reader(tmp_path, monkeypatch): + monkeypatch.setattr(br, "_ACTIVE", {}) + assert br.executor_set_gate_version(None, 3) is None # no reader yet + reader = br.BatchReader(FileSampleBuffer(CheckpointStore(tmp_path), batch_size=1), + max_policy_age=1) + br.register_reader(reader) + assert br.executor_set_gate_version(None, 3) is None + assert br.executor_set_gate_version(None, 4) == 3 and reader.gate_version == 4 + + +def test_island_env_carries_the_lane_only_for_a_trainer(): + env = br.island_env("trainer", store_root="/s", max_policy_age=2, lane=BufferLane(1, 2)) + assert env[br.LANE_ENV] == "1/2" and br.LANE_ENV in br.ISLAND_ENVS + with pytest.raises(ValueError): + br.island_env("sampler", store_root="/s", sampler_id="a", lane=BufferLane(1, 2)) + + +# ---------------------------------------------------------------- versions and the publisher +def test_publisher_and_merged_version_rules(): + assert is_publisher(0, 2) and not is_publisher(1, 2) + with pytest.raises(ValueError): + is_publisher(2, 2) + state = SimpleNamespace(policy_version=4) + assert merged_version(SimpleNamespace(outer_version=None), state) == 4 + assert merged_version(SimpleNamespace(outer_version=6), state) == 6 + assert merged_version(SimpleNamespace(), state) == 4 + + +def test_guard_releases_increasing_versions_and_reports_gaps(): + gaps = [] + guard = MergedReleaseGuard(lambda a, b: gaps.append((a, b))) + assert guard.check(0, None) == "release" + assert guard.check(1, 0) == "release" + assert guard.check(1, 1) == "stale" and guard.check(0, 1) == "stale" + assert guard.check(4, 1) == "release" and gaps == [(1, 4)] + + +class _State: + def __init__(self, version: int, digest: str) -> None: + self.policy_version = version + self._digest = digest + + def policy_tensor_hash(self) -> str: + return self._digest + + +class _Inner: + def __init__(self) -> None: + self.published: list[int] = [] + + def publish(self, state): + self.published.append(state.policy_version) + return "ok" + + +def test_releasing_publisher_releases_merged_versions_in_order(tmp_path): + events: list[tuple[str, dict]] = [] + releases = PolicyReleaseStore(CheckpointStore(tmp_path)) + driver = SimpleNamespace(outer_version=None) + gates: list[int] = [] + inner = GateVersionPublisher(_Inner(), gates.append, lambda s: merged_version(driver, s)) + publisher = ReleasingPublisher( + inner, releases, lambda v, d: (d / "w.bin").write_bytes(b"%d" % v), lambda: 3, + lambda event, **f: events.append((event, f)), + version_of=lambda s: merged_version(driver, s), + guard=MergedReleaseGuard(lambda a, b: events.append(("gap", {"a": a, "b": b}))), + origin=lambda s: {"kind": "merged", "learner_id": 0, "num_learners": 2, + "local_version": s.policy_version}) + for local, outer, digest in ((0, 0, "h0"), (1, 1, "h1"), (2, 1, "h1"), (3, 3, "h3")): + driver.outer_version = outer + assert publisher.publish(_State(local, digest)) == "ok" + assert gates == [0, 1, 1, 3] + assert [e for e, _ in events] == ["rl_policy_released", "rl_policy_released", + "rl_policy_release_skipped", "gap", "rl_policy_released"] + latest = releases.latest() + assert (latest.seq, latest.version, latest.policy_hash) == (3, 3, "h3") + assert latest.origin == {"kind": "merged", "learner_id": 0, "num_learners": 2, + "local_version": 3} + assert latest.buffer_selected == 3 + + +# ---------------------------------------------------------------- role checks +def _role_args(**kw): + base = dict(rl_sample_role="trainer", rl_single_island_no_sync=False, + rl_checkpoint_store="modal-volume://v/p", rl_max_policy_age=2, tuning="lora") + return SimpleNamespace(**{**base, **kw}) + + +def test_a_trainer_may_sync_a_sampler_may_not(): + sample_roles.check_launch(_role_args()) + sample_roles.check_synced_trainers(_role_args(), 2) + with pytest.raises(ValueError, match="only trainers may sync"): + sample_roles.check_launch(_role_args(rl_sample_role="sampler")) + with pytest.raises(ValueError, match="placement reconfiguration"): + sample_roles.check_launch(_role_args(rl_elastic=True)) + + +@pytest.mark.parametrize("kw,islands,message", [ + ({}, 1, "at least two islands"), + ({"tuning": "full"}, 2, "--tuning lora"), +]) +def test_synced_trainer_refusals(kw, islands, message): + with pytest.raises(ValueError, match=message): + sample_roles.check_synced_trainers(_role_args(**kw), islands) + + +@pytest.mark.parametrize("store", ["s3://b/p", "oss://b/p", "modal-volume://v/p"]) +def test_synced_trainers_take_an_object_store(store): + sample_roles.check_synced_trainers(_role_args(rl_checkpoint_store=store), 2) + + +def test_buffer_lane_of_island_args(): + args = _role_args(learner_id=1, rl_sample_lanes=2) + assert sample_roles.buffer_lane(args) == BufferLane(1, 2) + assert sample_roles.buffer_lane(_role_args(learner_id=0, rl_single_island_no_sync=True, + rl_sample_lanes=2)) is None + resolved = sample_roles.resolve(_role_args(rl_sample_store="/root/store"), groups_per_round=4) + assert resolved.store == "/root/store" and resolved.role == "trainer" + + +# ---------------------------------------------------------------- launch and island flags +SYNCED = ("--rl-sample-role", "trainer", "--rl-max-policy-age", "2", + "--rl-checkpoint-store", "modal-volume://yeto-c3/run-a", "--rl-sample-wait-s", "900", + "--rl-island-scheduling", "elastic", "--cluster-prefix", "c3-synced", + "--rl-algorithm-spec", "examples/rl_algorithms/decoupled-ppo-c2.json") + + +@pytest.fixture(autouse=True) +def _island_hmac_key(monkeypatch): + monkeypatch.setenv("YETO_ISLAND_HMAC_KEY", "c3-test-key") + + +def _synced_plan(extra=()): + from test_rl_sampler_trainer_pair_launch import _args + + from yeto.launcher import dry_run_plan + + args = _args(tuple(extra), base=SYNCED) + return args, dry_run_plan(args) + + +def _island_tokens(command: str, learner: str) -> list[str]: + tail = command.split("-m yeto.rl.adapters.miles.island_entry", 1)[1] + return [{"$LEARNER_ID": learner}.get(t, t) for t in shlex.split(tail)] + + +def test_synced_trainer_plan_forwards_the_sample_store_and_lanes(): + from yeto.rl.adapters.miles import island_entry + + _, plan = _synced_plan() + assert plan["syncer"] is not None + islands = plan["island_requests"] + assert len(islands) == 2 + for m, island in enumerate(islands): + tokens = _island_tokens(island["learner_command"], str(m)) + assert tokens[tokens.index("--rl-sample-role") + 1] == "trainer" + assert tokens[tokens.index("--rl-sample-lanes") + 1] == "2" + assert "--rl-sample-store" in tokens and "--rl-resume-store" not in tokens + parsed = island_entry.parse_args(tokens) + assert sample_roles.buffer_lane(parsed) == BufferLane(m, 2) + + +def test_synced_trainer_launch_refuses_round_cuts(): + with pytest.raises(ValueError, match="no round cuts"): + _synced_plan(("--rl-cut-every", "1")) + + +def test_island_refuses_a_sample_store_outside_a_synced_trainer(): + from yeto.rl.adapters.miles import island_entry + + with pytest.raises(ValueError, match="synced trainer island needs"): + island_entry._check_sample_store(SimpleNamespace( + rl_sample_role="trainer", rl_single_island_no_sync=False, rl_sample_store=None, + rl_sample_lanes=None, rl_resume_store=None)) + with pytest.raises(ValueError, match="need --rl-sample-role trainer"): + island_entry._check_sample_store(SimpleNamespace( + rl_sample_role="sampler", rl_single_island_no_sync=True, rl_sample_store="/s", + rl_sample_lanes=2, rl_resume_store=None)) diff --git a/tests/test_rl_miles_adapter_rollout.py b/tests/test_rl_miles_adapter_rollout.py index 0ecdf8b9..36613882 100644 --- a/tests/test_rl_miles_adapter_rollout.py +++ b/tests/test_rl_miles_adapter_rollout.py @@ -262,7 +262,7 @@ async def get(self, rollout_id): result = await super().get(rollout_id) nonzero = counts[rollout_id] # computed during reward post-processing hook.record_round_metadata(self.args, self.kept, sink=self.sink, - nonzero_advantages=nonzero) + nonzero_advantages=nonzero, buffer_stale_groups=nonzero % 5) return result kept = [group(0, [1.0, 0.0]), group(1, [0.0, 1.0])] @@ -278,7 +278,9 @@ def expected(): seen = {} for rid in (3, 4, 5): current["v"] = next(versions) - seen[rid] = p.generate(rid).nonzero_advantages + batch = p.generate(rid) + seen[rid] = batch.nonzero_advantages + assert batch.buffer_stale_groups == counts[rid] % 5 assert seen == counts diff --git a/tests/test_rl_miles_adapter_trainer_publish.py b/tests/test_rl_miles_adapter_trainer_publish.py index f60893dd..6e4ccbe3 100644 --- a/tests/test_rl_miles_adapter_trainer_publish.py +++ b/tests/test_rl_miles_adapter_trainer_publish.py @@ -574,3 +574,21 @@ def test_gspo_and_gmpo_clip_fraction_unchanged_by_unconditional_drain(variant, m t.train_step(handle()) assert t.step_metrics().clip_fraction == 0.25 assert state_plugin._STEP_LOSSES == [] + + +def test_trainer_only_island_publishes_without_engines(monkeypatch): + """--debug-train-only: no engine update, no stamp; the release is the publication.""" + from yeto.rl.adapters.miles import publish as publish_module + + calls = [] + + async def miles_update_weights(args, actor, executor, ctl, **kw): + calls.append("update_weights") + + monkeypatch.setattr(publish_module, "_default_update_weights", lambda: miles_update_weights) + pub = MilesPublisher( + args=SimpleNamespace(debug_train_only=True), actor_model=PluginActor(calls), rollout_executor=object(), + inference_controller=Controller([]), flatten_checksums=flatten, + ) + result = pub.publish(state()) + assert calls == [] and result.members == frozenset() and pub.last_engine_checksums is None diff --git a/tests/test_rl_n2_trainer_wiring.py b/tests/test_rl_n2_trainer_wiring.py index be3b8a90..4e23caef 100644 --- a/tests/test_rl_n2_trainer_wiring.py +++ b/tests/test_rl_n2_trainer_wiring.py @@ -68,7 +68,8 @@ def test_miles_reader_rebuilds_prompts_and_reports_per_source_metrics(monkeypatc reader.args = types.SimpleNamespace(rollout_sample_filter_path=None, rollout_all_samples_process_path=None) reader.samples_per_group = 2 reader._tokenizer = lambda text, add_special_tokens: {"input_ids": [7, 8]} - report = types.SimpleNamespace(metrics=lambda: {"sample_buffer/groups": 1.0}) + report = types.SimpleNamespace(metrics=lambda: {"sample_buffer/groups": 1.0}, stale_groups=0) + monkeypatch.setattr(plugin, "record_round_metadata", lambda args, rid, **counts: None) group = [_record(index=0), _record(finish=FINISH_LENGTH, index=1)] reader.reader = types.SimpleNamespace(take=lambda rollout_id, groups, current_version: ([(5, group)], report)) reader.groups = 1 @@ -116,3 +117,32 @@ def test_masked_source_samples_checks_row_counts(): masked_source_samples([[0.0]], [[0.0]], [[1], [1]]) with pytest.raises(ValueError, match="provenances"): masked_source_samples([[0.0]], [[0.0]], [[1]], provenances=[None, None]) + + +def test_verl_tq_provenance_is_read_from_the_extra_fields_record(monkeypatch): + """verl's TQ agent loop stores ``AgentLoopOutput.as_dict()``: provenance sits in + ``extra_fields``, not under its own name (S25 T-G4).""" + import sys + import types + + from yeto.rl.adapters.verl import ports_impl + + class Wrapped: # a TransferQueue read can wrap each record (NonTensorData) + def __init__(self, data): + self.data = data + + seen = {} + + def kv_batch_get(*, keys, partition_id, select_fields): + seen["fields"] = select_fields + return {"extra_fields": [Wrapped({"provenance": SOURCE}), {"other": 1}]} + + monkeypatch.setitem(sys.modules, "transfer_queue", types.SimpleNamespace(kv_batch_get=kv_batch_get)) + batch = types.SimpleNamespace(keys=["a", "b"], partition_id="train") + assert ports_impl._row_field(batch, "provenance") == ([SOURCE, None], "ok") + assert seen["fields"] == ["extra_fields"] + + monkeypatch.setitem(sys.modules, "transfer_queue", types.SimpleNamespace( + kv_batch_get=lambda **_: {"extra_fields": [{}, {}]})) + rows, read = ports_impl._row_field(batch, "provenance") + assert rows is None and read.startswith("missing") diff --git a/tests/test_rl_oss_store.py b/tests/test_rl_oss_store.py index 7aab70d3..7f457ad6 100644 --- a/tests/test_rl_oss_store.py +++ b/tests/test_rl_oss_store.py @@ -156,7 +156,7 @@ def test_each_sampler_writes_its_own_shard_and_the_trainer_merges_them(): assert len(again) == 4 # a reader that does not merge (the head) sees the merged index only assert len(FileSampleBuffer(objects, batch_size=1)) == 4 - with pytest.raises(SampleBufferError, match="only the index of an object store merges"): + with pytest.raises(SampleBufferError, match="a shard writer does not merge"): FileSampleBuffer(objects, batch_size=1, index=shard_index("x"), merge_shards=True) diff --git a/tests/test_rl_policy_age.py b/tests/test_rl_policy_age.py index 3f82926f..bc9b5b1e 100644 --- a/tests/test_rl_policy_age.py +++ b/tests/test_rl_policy_age.py @@ -227,3 +227,27 @@ def test_miles_switches_are_derived_from_the_limit(): assert row.translate(AlgorithmSpec(execution=ExecutionSpec(max_policy_staleness=1))) == [] with pytest.raises((AlgorithmSpecError, ValueError), match="rl-max-policy-age"): af.absorb_extra_argv(AlgorithmSpec(), ["--partial-rollout"]) + + +def test_driver_tapes_the_buffer_readers_stale_groups_per_round(tmp_path): + """T-EVAL E4: groups a sample-buffer reader passed over as too old reach the + tape as rl_buffer_stale_skipped (one per round, 0 included); an engine that + does not report them tapes nothing.""" + import dataclasses + + stale = {0: 0, 1: 3, 2: 1} + driver = _driver(tmp_path, 1) + original = driver.rollout.generate + + def gen(rid, **kw): + return dataclasses.replace(original(rid, **kw), buffer_stale_groups=stale[rid]) + + driver.rollout.generate = gen + driver.run() + events = [json.loads(l) for l in (tmp_path / "e.jsonl").read_text().splitlines()] + skipped = [e for e in events if e["event"] == "rl_buffer_stale_skipped"] + assert [(e["rollout_id"], e["stale_groups"], e["max_policy_age"]) for e in skipped] == [ + (0, 0, 1), (1, 3, 1), (2, 1, 1)] + (tmp_path / "x").mkdir() + _driver(tmp_path / "x", 1).run() + assert '"rl_buffer_stale_skipped"' not in (tmp_path / "x" / "e.jsonl").read_text() diff --git a/tests/test_rl_quant_mxfp4.py b/tests/test_rl_quant_mxfp4.py index 34310e09..d7b32150 100644 --- a/tests/test_rl_quant_mxfp4.py +++ b/tests/test_rl_quant_mxfp4.py @@ -337,3 +337,38 @@ def test_cli_template_convert_compare(world, capsys): assert main(["compare", str(out), str(world["original"]), "--match", r"experts\.\d+\."]) == 0 rep = json.loads(capsys.readouterr().out) assert rep["fp4_match_share"] == 1.0 and rep["e8m0_match_share"] == 1.0 and rep["tensors"] == 3 * 3 * E # packed + scale, plus the BF16 MTP experts + + +def test_reference_checkpoint_serves_as_template(world): + # stage 1 passes the round-0 MXFP4 export (no layout.json) as --template + out = world["tmp"] / "out-ckpt-template" + requantize_checkpoint(world["source"], world["original"], out, AwqSmoothing.none()) + rep = compare_checkpoints(out, world["original"]) + assert rep["unequal_tensors"] == 0 and rep["fp4_match_share"] == 1.0 + + +def test_frozen_tensors_come_from_the_base_only_when_named(world): + partial = world["tmp"] / "release-without-lm-head" + _write_ckpt(partial, {n: t for n, t in world["src"].items() if n != "lm_head.weight"}, + lambda n: "model-00001-of-00001.safetensors") + with pytest.raises(LayoutError, match="1 missing"): + requantize_checkpoint(partial, world["template"], world["tmp"] / "a", AwqSmoothing.none()) + with pytest.raises(LayoutError, match="1 missing"): + requantize_checkpoint(partial, world["template"], world["tmp"] / "b", AwqSmoothing.none(), + frozen=(world["source"], r"^model\.visual\.")) + out = world["tmp"] / "c" + stats = requantize_checkpoint(partial, world["template"], out, AwqSmoothing.none(), + frozen=(world["source"], r"^lm_head\.")) + assert stats["frozen_from_base"] == 1 + assert compare_checkpoints(out, world["original"])["unequal_tensors"] == 0 + + +def test_text_model_prefix_of_a_miles_raw_export_is_accepted(world): + # Miles' raw glm5_next export writes "model.layers..."; the layout has "model.language_model.layers..." + raw = world["tmp"] / "raw-export" + _write_ckpt(raw, {n.replace("model.language_model.", "model.", 1): t for n, t in world["src"].items()}, + lambda n: "model-00001-of-00001.safetensors") + out = world["tmp"] / "out-raw" + stats = requantize_checkpoint(raw, world["template"], out, AwqSmoothing.none()) + assert stats["renamed_from_text_prefix"] == sum(n.startswith("model.language_model.") for n in world["src"]) + assert compare_checkpoints(out, world["original"])["unequal_tensors"] == 0 diff --git a/tests/test_rl_quant_shards.py b/tests/test_rl_quant_shards.py index 0350ab94..8868595d 100644 --- a/tests/test_rl_quant_shards.py +++ b/tests/test_rl_quant_shards.py @@ -23,7 +23,7 @@ from tests import test_rl_quant_mxfp4 as mx from tests import test_rl_quant_requant as w4 -from yeto.rl.adapters.miles.quant_shards import hf_slices, megatron_expert_slices +from yeto.rl.adapters.miles.quant_shards import checkpoint_name, hf_slices, megatron_expert_slices from yeto.rl.adapters.verl.quant_shards import fsdp2_slices, shard0_rows from yeto.rl.engine.policy_release import ( RELEASE_DIR, @@ -225,7 +225,9 @@ def _megatron_rank(src: dict[str, torch.Tensor], ep_rank: int, num_local: int, s def test_megatron_expert_parallel_ranks_merge_to_the_same_bytes(world): src, out = world["src"], world["tmp"] / "megatron" parts = out / PART_DIR - non_expert = [(n, t) for n, t in src.items() if not _PER_EXPERT.match(n)] + # non-expert weights as Miles convert_to_hf emits them (SGLang names, no model.language_model.) + non_expert = [(n.replace("model.language_model.", "model.", 1), t) for n, t in src.items() if not _PER_EXPERT.match(n)] + assert any(n.startswith("model.layers.") for n, _ in non_expert) for ep_rank in (0, 1): # E = 2: one local expert per rank; rank 1 uses the sequential-MLP names slices = list(megatron_expert_slices(_megatron_rank(src, ep_rank, 1, sequential=ep_rank == 1), ep_rank=ep_rank, num_local_experts=1)) @@ -380,3 +382,11 @@ def test_follower_refuses_a_release_of_another_scheme_before_reading_it(tmp_path bf16.release(4, "h4", lambda d: (d / "model.safetensors").write_text("y")) with pytest.raises(SchemeMismatch, match="'bf16'"): PolicyFollower(bf16, lambda r, d: None, scheme="w4a8-int4-channel").refresh() + + +def test_miles_convert_to_hf_names_map_to_checkpoint_names(): + assert checkpoint_name("model.layers.3.mlp.gate.weight") == "model.language_model.layers.3.mlp.gate.weight" + assert checkpoint_name("model.embed_tokens.weight") == "model.language_model.embed_tokens.weight" + assert checkpoint_name("model.norm.weight") == "model.language_model.norm.weight" + for unchanged in ("lm_head.weight", "model.language_model.norm.weight", "model.visual.blocks.0.attn.qkv.weight"): + assert checkpoint_name(unchanged) == unchanged diff --git a/tests/test_rl_sampler_trainer_pair_launch.py b/tests/test_rl_sampler_trainer_pair_launch.py index 4debf1e7..0cbd7e69 100644 --- a/tests/test_rl_sampler_trainer_pair_launch.py +++ b/tests/test_rl_sampler_trainer_pair_launch.py @@ -134,7 +134,7 @@ def _replace(base, old: tuple, new: tuple) -> list[str]: (("--rl-single-island-no-sync",), (), "no-sync|single-island"), (("--tuning", "full"), ("--tuning", "full", "--rl-sample-role", "trainer"), "replaces --rl-sample-role"), - (("1:trainer",), ("1:sampler",), "one sampler and one trainer"), + (("1:trainer",), ("1:sampler",), "exactly one trainer"), (("--rl-max-policy-age", "2"), ("--rl-max-policy-age", "0"), "max-policy-age"), (("--rl-checkpoint-store", STORE, "--rl-cut-every", "1"), (), "checkpoint-store"), ]) diff --git a/tests/test_rl_spot_sampler_replace.py b/tests/test_rl_spot_sampler_replace.py new file mode 100644 index 00000000..cbe2e7fa --- /dev/null +++ b/tests/test_rl_spot_sampler_replace.py @@ -0,0 +1,515 @@ +"""S25 C1: a sampler island on spot capacity is reclaimed and replaced. + +CPU only (no Miles, no Ray): the resume point read from the buffer, the +resumable round writer (no round written twice, no prompt drawn twice), several +sampler islands writing their own shards of one shared-directory buffer that +the trainer merges, the launch composition of one trainer and two samplers, and +the Modal runner's reclaim marker for sampler islands. +""" + +from __future__ import annotations + +import shlex +import signal as _signal +from types import SimpleNamespace + +import pytest + +from test_rl_algorithm_provenance import _fake_sky, _launcher_args, _no_modal_listing # noqa: F401 +from tests.buffer_rollout_procs import sample +from yeto.rl.adapters.miles import buffer_rollout as br +from yeto.rl.engine.policy_release import PolicyFollower, PolicyReleaseStore +from yeto.rl.engine.resume import CheckpointStore +from yeto.rl.engine.sample_buffer import FileSampleBuffer, SampleBufferError +from yeto.rl.engine.sampler_loop import SamplerLoop +from yeto.rl.engine.sampler_resume import ( + RECLAIM_MARKER_ENV, + Member, + SamplerResumeError, + SamplerRound, + reclaim_requested, + resume_point, + write_member, +) + + +def _group(g: int, version: int = 1) -> list[dict]: + return [sample(2 * g, g, version), sample(2 * g + 1, g, version)] + + +def _cursor(offset: int) -> dict[str, int]: + return {"sample_offset": offset, "epoch_id": 0, "sample_group_index": offset, "sample_index": 2 * offset} + + +def _round(rid: int, groups: int) -> list[tuple[str, list[dict]]]: + return [(f"r{rid}-g{g}", _group(g)) for g in range(groups)] + + +def _shard_env(store, sampler=None, shards=("island0", "island1")): + env = {br.BUFFER_ROOT_ENV: str(store), br.SHARDS_ENV: ",".join(shards)} + if sampler is not None: + env[br.SAMPLER_ID_ENV] = sampler + return env + + +# ---------------------------------------------------------------- resume point +def test_resume_point_of_an_empty_buffer_starts_at_round_zero(): + resume = resume_point([], "island0") + assert (resume.next_rollout_id, resume.data_cursor, resume.resumed) == (0, None, False) + + +def test_resume_point_continues_after_the_newest_written_round(tmp_path): + buffer = FileSampleBuffer(CheckpointStore(tmp_path), batch_size=1) + writer = br.SamplerRoundWriter(buffer, "island0", incarnation="a") + assert writer.write_round(0, _round(0, 2), _cursor(2)) == 2 + assert writer.write_round(1, _round(1, 2), _cursor(4)) == 2 + other = br.SamplerRoundWriter(buffer, "island9", incarnation="z") + other.write_round(5, _round(5, 1), _cursor(9)) # another sampler: not ours + resume = resume_point(buffer.entries(), "island0") + assert (resume.next_rollout_id, resume.data_cursor) == (2, _cursor(4)) + assert (resume.written_rounds, resume.partial_rounds, resume.groups) == ((0, 1), (), 4) + assert resume.incarnations == ("a",) + + +def test_a_round_cut_short_by_the_reclaim_keeps_its_groups_and_is_not_written_again(tmp_path): + """The island leaves while it writes round 1: the written group stays, the + rest are dropped; the replacement continues at round 2 from round 1's cursor + (round 1's prompts are not drawn again).""" + store = CheckpointStore(tmp_path) + leaving = {"now": False} + old = br.SamplerRoundWriter(FileSampleBuffer(store, batch_size=1), "island0", incarnation="a", + reclaimed=lambda: leaving["now"]) + old.write_round(0, _round(0, 3), _cursor(3)) + + class Leaving(list): + def __iter__(self): # the reclaim lands after the first group of round 1 + for i, item in enumerate(super().__iter__()): + leaving["now"] = i >= 1 + yield item + + assert old.write_round(1, Leaving(_round(1, 3)), _cursor(6)) == 1 + assert old.dropped_groups == 2 + new = br.SamplerRoundWriter(FileSampleBuffer(store, batch_size=1), "island0", incarnation="b") + assert (new.resume.next_rollout_id, new.resume.data_cursor) == (2, _cursor(6)) + assert new.resume.partial_rounds == (1,) + for rid in (0, 1): + with pytest.raises(SamplerResumeError, match="already in the sample buffer"): + new.write_round(rid, _round(rid, 1), _cursor(0)) + new.write_round(2, _round(2, 3), _cursor(9)) + ids = [e.sample_id for e in FileSampleBuffer(store, batch_size=1).entries()] + assert len(ids) == len(set(ids)) == 3 + 1 + 3 + assert [i.split("-")[1] for i in ids] == ["a"] * 4 + ["b"] * 3 + + +def test_a_round_still_generating_at_the_reclaim_is_generated_again(tmp_path): + store = CheckpointStore(tmp_path) + old = br.SamplerRoundWriter(FileSampleBuffer(store, batch_size=1), "island0", incarnation="a", + reclaimed=lambda: True) + assert old.write_round(0, _round(0, 2), _cursor(2)) == 0 # left before the write + new = br.SamplerRoundWriter(FileSampleBuffer(store, batch_size=1), "island0", incarnation="b") + assert (new.resume.next_rollout_id, new.resume.data_cursor) == (0, None) + assert new.write_round(0, _round(0, 2), _cursor(2)) == 2 + + +def test_resume_point_refuses_a_buffer_with_two_cursors_for_one_round(): + def entry(cursor, last): + return SimpleNamespace(sampler_round=SamplerRound("island0", "a", 0, cursor, last).record()) + + with pytest.raises(SamplerResumeError, match="two data cursors"): + resume_point([entry(_cursor(2), False), entry(_cursor(3), True)], "island0") + with pytest.raises(SamplerResumeError, match="two last groups"): + resume_point([entry(_cursor(2), True), entry(_cursor(2), True)], "island0") + + +def test_seek_puts_the_data_source_where_its_load_would_and_only_once(tmp_path): + class Dataset: + def __init__(self): + self.shuffled = [] + + def shuffle(self, epoch): + self.shuffled.append(epoch) + + source = SimpleNamespace(sample_offset=0, epoch_id=0, sample_group_index=0, sample_index=0, + dataset=Dataset()) + store = CheckpointStore(tmp_path) + br.SamplerRoundWriter(FileSampleBuffer(store, batch_size=1), "island0", incarnation="a" + ).write_round(0, _round(0, 1), {**_cursor(5), "epoch_id": 1}) + writer = br.SamplerRoundWriter(FileSampleBuffer(store, batch_size=1), "island0", incarnation="b") + assert writer.seek(SimpleNamespace(__self__=source), shuffle=True)["sample_offset"] == 5 + assert (source.sample_offset, source.epoch_id, source.sample_index) == (5, 1, 10) + assert source.dataset.shuffled == [1] + assert writer.seek(source, shuffle=True) is None # the second round: no seek + with pytest.raises(br.SampleBufferRolloutError, match="cannot seek"): + br.seek_data_source(SimpleNamespace(sample_offset=0), _cursor(1), shuffle=False) + + +def test_a_writer_without_resume_writes_as_before(tmp_path): + """A convertible node keeps its own rollout numbering: no round check.""" + store = CheckpointStore(tmp_path) + writer = br.SamplerRoundWriter(FileSampleBuffer(store, batch_size=1), "n0", incarnation="a", + resume=False) + writer.write_round(0, _round(0, 1), None) + writer.write_round(0, [("r0-again", _round(0, 1)[0][1])], None) # payload names stay unique + assert writer.seek(object(), shuffle=True) is None + assert len(FileSampleBuffer(store, batch_size=1)) == 2 + + +# ---------------------------------------------------------------- two samplers, one trainer +def test_two_samplers_write_shards_the_trainer_merges_and_a_replacement_continues(tmp_path): + a = br.SamplerRoundWriter(br.open_buffer(_shard_env(tmp_path, "island0"), sampler=True), + "island0", incarnation="a1") + b = br.SamplerRoundWriter(br.open_buffer(_shard_env(tmp_path, "island1"), sampler=True), + "island1", incarnation="b1") + assert a.buffer.shard == "island0" and not (tmp_path / "sample-buffer/index.jsonl").exists() + trainer = br.open_buffer(_shard_env(tmp_path), batch_size=2, reader=True) + reader = br.BatchReader(trainer, max_policy_age=2, wait_s=0.0, poll_s=0.0) + a.write_round(0, _round(0, 2), _cursor(2)) + b.write_round(0, _round(0, 2), _cursor(2)) + taken, _ = reader.take(0, 4, current_version=1) + assert len(taken) == 4 + # island0 is reclaimed after one more round; its replacement continues at round 2 + a.write_round(1, _round(1, 2), _cursor(4)) + a2 = br.SamplerRoundWriter(br.open_buffer(_shard_env(tmp_path, "island0"), sampler=True), + "island0", incarnation="a2") + assert (a2.resume.next_rollout_id, a2.resume.incarnations) == (2, ("a1",)) + a2.write_round(2, _round(2, 2), _cursor(6)) + b.write_round(1, _round(1, 2), _cursor(4)) + taken2, _ = reader.take(1, 6, current_version=2) + ids = [e.sample_id for e in trainer.entries()] + assert len(taken2) == 6 and len(ids) == len(set(ids)) == 10 # none repeated, none lost + assert trainer.shard_positions(trainer.read_cursor) == {"island0": 6, "island1": 4} + # the data ends only when every named shard has ended + a2.buffer.mark_data_end() + trainer.refresh() + assert not trainer.data_ended + b.buffer.mark_data_end() + trainer.refresh() + assert trainer.data_ended + + +def test_the_data_does_not_end_before_a_named_shard_exists(tmp_path): + a = br.open_buffer(_shard_env(tmp_path, "island0"), sampler=True) + a.mark_data_end() + trainer = br.open_buffer(_shard_env(tmp_path), reader=True) + assert not trainer.data_ended # island1 has not written anything yet + + +def test_a_restarted_trainer_merges_no_record_twice(tmp_path): + a = br.SamplerRoundWriter(br.open_buffer(_shard_env(tmp_path, "island0"), sampler=True), + "island0", incarnation="a") + a.write_round(0, _round(0, 3), _cursor(3)) + assert len(br.open_buffer(_shard_env(tmp_path), reader=True)) == 3 + a.write_round(1, _round(1, 1), _cursor(4)) + again = br.open_buffer(_shard_env(tmp_path), reader=True) + assert len(again) == 4 + assert [r.get("shard_pos") for r in map(_json, _lines(tmp_path))] == [0, 1, 2, 3] + + +def _lines(root): + return (root / "sample-buffer/index.jsonl").read_text().splitlines() + + +def _json(line): + import json + + return json.loads(line) + + +def test_the_trainer_merges_one_refresh_with_one_store_commit(tmp_path): + class Counting(CheckpointStore): + commits = 0 + + def commit(self): + Counting.commits += 1 + + writer = br.SamplerRoundWriter(FileSampleBuffer(CheckpointStore(tmp_path), batch_size=1, + index="shards/island0.jsonl"), + "island0", incarnation="a") + writer.write_round(0, _round(0, 5), _cursor(5)) + Counting.commits = 0 + trainer = FileSampleBuffer(Counting(tmp_path), batch_size=1, merge_shards=True, + expected_shards=["island0"]) + assert len(trainer) == 5 and Counting.commits == 1 + + +def test_expected_shards_need_a_merging_buffer(tmp_path): + with pytest.raises(SampleBufferError, match="expected_shards needs merge_shards"): + FileSampleBuffer(CheckpointStore(tmp_path), batch_size=1, expected_shards=["a"]) + + +def test_island_env_names_the_shards_and_refuses_a_sampler_outside_them(): + env = br.island_env("sampler", store_root="/s", sampler_id="island1", shards=["island0", "island1"]) + assert env[br.SHARDS_ENV] == "island0,island1" and env[br.SAMPLER_ID_ENV] == "island1" + with pytest.raises(ValueError, match="not one of the buffer shards"): + br.island_env("sampler", store_root="/s", sampler_id="island2", shards=["island0"]) + assert br.SHARDS_ENV in br.ISLAND_ENVS and br.INCARNATION_ENV in br.ISLAND_ENVS + + +# ---------------------------------------------------------------- the sampler loop leaves +def test_the_sampler_loop_stops_on_the_reclaim_and_marks_no_data_end(tmp_path): + store = CheckpointStore(tmp_path) + marker = tmp_path / "reclaim" + env = {RECLAIM_MARKER_ENV: str(marker)} + buffer = FileSampleBuffer(store, batch_size=1) + writer = br.SamplerRoundWriter(buffer, "island0", incarnation="a", + reclaimed=lambda: reclaim_requested(env)) + events = [] + + def generate(rid): + writer.write_round(rid, _round(rid, 2), _cursor(2 * (rid + 1))) + if rid == 1: + marker.write_text("{}") # the container gets SIGINT during round 1 + + loop = SamplerLoop(follower=PolicyFollower(PolicyReleaseStore(store), lambda r, d: None), + buffer=buffer, generate=generate, num_rollouts=5, max_backlog_groups=100, + emit=lambda e, **f: events.append(e), poll_s=0.0, + should_stop=lambda: reclaim_requested(env)) + assert loop.run() == 2 + assert "rl_sampler_stopped" in events and not buffer.data_ended + resume = resume_point(FileSampleBuffer(store, batch_size=1).entries(), "island0") + assert resume.next_rollout_id == 2 + loop.should_stop = lambda: False + marker.unlink() + assert loop.run(start=resume.next_rollout_id) == 3 and buffer.data_ended + + +def test_member_records_are_one_file_per_incarnation(tmp_path): + store = CheckpointStore(tmp_path) + first = write_member(store, Member("island0", "a", 1.0, {"next_rollout_id": 0})) + member = Member("island0", "b", 2.0, {"next_rollout_id": 3}, leave_reason="PREEMPTED") + second = write_member(store, member) + assert first != second and first.exists() and "PREEMPTED" in second.read_text() + + +# ---------------------------------------------------------------- launch composition +GPU3 = "modal:1xh100,modal:1xh100,modal:1xh100" +STORE = "modal-volume://yeto-c1/run-a" +GROUP = ("--rl-single-island-no-sync", "--tuning", "lora", "--lora-r", "16", + "--rl-island-sample-role", "0:sampler", "--rl-island-sample-role", "1:sampler", + "--rl-island-sample-role", "2:trainer", "--rl-max-policy-age", "2", + "--rl-algorithm-spec", "examples/rl_algorithms/decoupled-ppo-c2.json", + "--rl-checkpoint-store", STORE, "--cluster-prefix", "c1-spot", "--total-steps", "4") + + +def _args(extra=()): + from yeto.launcher import _prepare_rl_args + + args = _launcher_args("ports", GROUP + tuple(extra), gpu=GPU3) + args.controller = "local" + _prepare_rl_args(args) + return args + + +def _tokens(command): + return shlex.split(command.split("-m yeto.rl.adapters.miles.island_entry", 1)[1]) + + +def _flag(tokens, flag): + return tokens[tokens.index(flag) + 1] if flag in tokens else None + + +def test_one_trainer_and_two_samplers_get_their_node_and_the_shards(): + from yeto.launcher import dry_run_plan + + plan = dry_run_plan(_args()) + roles = [i["sample_role"] for i in plan["island_requests"]] + assert roles == ["sampler", "sampler", "trainer"] and plan["syncer"] is None + tokens = [_tokens(i["learner_command"]) for i in plan["island_requests"]] + assert [_flag(t, "--rl-sampler-node") for t in tokens] == ["island0", "island1", None] + assert {_flag(t, "--rl-buffer-shards") for t in tokens} == {"island0,island1"} + # the trainer trains --total-steps rounds; each of the two samplers generates half + assert [_flag(t, "--global-rounds") for t in tokens] == ["2", "2", "4"] + assert [_flag(t, "--total-fragment-steps") for t in tokens] == ["2", "2", "4"] + + +def test_the_samplers_split_the_rounds_evenly(): + with pytest.raises(ValueError, match="multiple of the sampler count"): + _args(("--total-steps", "3")) + + +def test_the_island_resolves_its_sampler_id_and_the_shards(): + from yeto.launcher import dry_run_plan + from yeto.rl.adapters.miles import island_entry + + plan = dry_run_plan(_args()) + sampler = island_entry.parse_args([{"$LEARNER_ID": "0"}.get(t, t) for t in + _tokens(plan["island_requests"][1]["learner_command"])]) + env = {} + island_entry.apply_sample_role(sampler, SimpleNamespace(), env) + assert (env[br.SAMPLER_ID_ENV], env[br.SHARDS_ENV]) == ("island1", "island0,island1") + + +def test_only_sampler_islands_get_the_reclaim_marker(): + from yeto.gpu_spec import parse_gpu_spec + from yeto.launch.modal_islands import build_modal_island_config + from yeto.launcher import _island_task, make_miles_island_task + + args = _args() + for m, spec in enumerate(parse_gpu_spec(args.gpu)): + task = _island_task(make_miles_island_task, args, spec, m, 3, "none", {}) + cfg = build_modal_island_config(args, spec, m, task, "none") + assert (RECLAIM_MARKER_ENV in cfg.envs) == (m < 2) + + +def test_two_samplers_refuse_the_trainer_scheduler(): + with pytest.raises(ValueError, match="several sampler islands"): + _args(("--rl-trainer-scheduler", "s1", "--rl-trainer-island", "2")) + + +def test_every_island_needs_a_role_and_one_trainer(): + from yeto.launcher import _prepare_rl_args + + two_trainers = tuple("1:trainer" if t == "1:sampler" else t for t in GROUP) + args = _launcher_args("ports", two_trainers, gpu=GPU3) + args.controller = "local" + with pytest.raises(ValueError, match="exactly one trainer"): + _prepare_rl_args(args) + + +def test_the_modal_runner_turns_the_signal_into_the_marker_on_a_sampler(monkeypatch): + from yeto import modal_runner as mr + + seen = [] + monkeypatch.setattr(mr, "_run_forwarding_to_learner", lambda cmd, env, **kw: seen.append(env) or 0) + + class Sig: + SIGINT, SIGTERM = _signal.SIGINT, _signal.SIGTERM + + @staticmethod + def signal(num, fn): + raise AssertionError("not expected") + + env = {RECLAIM_MARKER_ENV: "/tmp/yeto-reclaim-requested"} + assert mr.run_island_script(["x"], env, popen=lambda c, env, **k: None, signal_mod=Sig) == 0 + assert seen == [env] + + +def test_the_reclaim_watcher_leaves_once_while_the_loop_is_still_in_a_round(tmp_path): + import threading + + from yeto.rl.engine.sampler_resume import ReclaimWatcher, SamplerMembership + + store = CheckpointStore(tmp_path) + events = [] + membership = SamplerMembership(store, "island0", "b", resume_point([], "island0"), + lambda e, **f: events.append((e, f)), clock=lambda: 7.0) + membership.join(latest_release={"seq": 3, "policy_version": 3}) + flag = threading.Event() + watcher = ReclaimWatcher(flag.is_set, lambda: membership.leave("PREEMPTED", rounds=2, + in_flight_rollout_id=4), + poll_s=0.01).start() + flag.set() + assert watcher.fired.wait(2.0) + watcher.stop() + assert membership.leave("FINISHED") is False # the exit path after the watcher: no second LEAVE + names = [e for e, _ in events] + assert names == ["rl_sampler_join", "rl_sampler_leave"] + leave = events[1][1] + assert (leave["reason"], leave["rounds"], leave["in_flight_rollout_id"]) == ("PREEMPTED", 2, 4) + record = (tmp_path / "sample-buffer/members/island0/b.json").read_text() + assert '"leave_reason": "PREEMPTED"' in record and '"joined_at": 7.0' in record + + +def test_the_reclaim_signal_writes_the_marker_instead_of_raising(tmp_path): + from yeto.rl.engine.sampler_resume import install_marker_signals + + handlers = {} + + class Sig: + SIGINT, SIGTERM = _signal.SIGINT, _signal.SIGTERM + + @staticmethod + def signal(num, fn): + handlers[num] = fn + + env = {RECLAIM_MARKER_ENV: str(tmp_path / "marker")} + assert install_marker_signals(env, signal_mod=Sig) == ["SIGINT", "SIGTERM"] + assert not reclaim_requested(env) + handlers[_signal.SIGINT](_signal.SIGINT, None) # Modal's SIGINT reaches the island process + assert reclaim_requested(env) and '"signal": 2' in (tmp_path / "marker").read_text() + + +def test_sampler_spot_bills_only_the_samplers_and_adds_no_spot_volume(monkeypatch): + from yeto.cloud import s3_store + from yeto.gpu_spec import parse_gpu_spec + from yeto.launch.rl_island_task import make_miles_island_task + from yeto.launcher import _island_task, _prepare_rl_args + + monkeypatch.setenv(s3_store.KEY_ID_ENV, "AKIDSTORE") + monkeypatch.setenv(s3_store.SECRET_ENV, "store-secret") + gpu = "nebius:1xh100@eu-north1,nebius:1xh100@eu-north1" + base = ("--rl-single-island-no-sync", "--tuning", "lora", "--rl-island-sample-role", "0:sampler", + "--rl-island-sample-role", "1:trainer", "--rl-max-policy-age", "2", + "--rl-algorithm-spec", "examples/rl_algorithms/decoupled-ppo-c2.json", + "--rl-checkpoint-store", "s3://yeto-ckpt/run-a", "--on-demand", "--rl-sampler-spot") + args = _launcher_args("ports", base, gpu=gpu) + args.controller = "local" + _prepare_rl_args(args) + from yeto.cloud import nebius_reclaim + + spots, volumes, hooks, markers = [], [], [], [] + for m, spec in enumerate(parse_gpu_spec(args.gpu)): + task = _island_task(make_miles_island_task, args, spec, m, 2, "none", {}) + spots.append(task.resources.use_spot) + volumes.append(sorted(getattr(task, "storage_mounts", None) or {})) + hooks.append(nebius_reclaim.wants_host_hook(task.envs)) + markers.append(task.envs.get(RECLAIM_MARKER_ENV)) + # the sampler gets the store mount only, like the trainer: no legacy --spot volume + assert spots == [True, False] and volumes[0] == volumes[1] and len(volumes[0]) == 1 + # Nebius stops the VM without signalling the container: the host unit writes the + # sampler's marker (S25 C1b); the on-demand trainer gets no unit + assert hooks == [True, False] and markers == [nebius_reclaim.NOTICE_FILE, None] + + +def test_an_on_demand_sampler_on_nebius_gets_no_host_unit(): + from yeto.launcher import _spot_sampler_cloud_env + + sampler = SimpleNamespace(rl_sample_role="sampler", rl_sampler_node="island0", spot=False) + assert _spot_sampler_cloud_env(sampler, "nebius") == {} + sampler.spot = True + assert _spot_sampler_cloud_env(sampler, "modal") == {} # Modal signals the container + assert _spot_sampler_cloud_env(sampler, "nebius") + + +def test_the_reclaim_ack_is_written_only_where_a_host_hook_waits(tmp_path): + from yeto.rl.engine.sampler_resume import RECLAIM_ACK_ENV, acknowledge_reclaim, clear_reclaim_files + + assert acknowledge_reclaim({}) is None + env = {RECLAIM_MARKER_ENV: str(tmp_path / "notice"), RECLAIM_ACK_ENV: str(tmp_path / "ack")} + assert acknowledge_reclaim(env) == tmp_path / "ack" and (tmp_path / "ack").exists() + (tmp_path / "notice").write_text("1\n") + clear_reclaim_files(env) # a restarted VM keeps /tmp: the next incarnation starts clean + assert not (tmp_path / "notice").exists() and not (tmp_path / "ack").exists() + + +def test_member_records_go_by_object_on_an_object_store(tmp_path): + from yeto.rl.engine.object_store import DirectoryObjectStore + + class StagedStore(CheckpointStore): # an s3/OSS store: its stage does not carry sample-buffer/ + def __init__(self, root, objects): + super().__init__(root) + self.objects = objects + + objects = DirectoryObjectStore(tmp_path / "bucket") + store = StagedStore(tmp_path / "stage", objects) + write_member(store, Member("island0", "a", 1.0, {"next_rollout_id": 0}, leave_reason="PREEMPTED")) + assert '"leave_reason": "PREEMPTED"' in objects.read("sample-buffer/members/island0/a.json").decode() + assert not (tmp_path / "stage" / "sample-buffer").exists() + + +def test_sampler_spot_needs_a_pair(): + from yeto.launcher import _prepare_rl_args + + args = _launcher_args("ports", ("--rl-single-island-no-sync", "--rl-sampler-spot"), + gpu="modal:1xh100") + args.controller = "local" + with pytest.raises(ValueError, match="--rl-sampler-spot needs --rl-island-sample-role"): + _prepare_rl_args(args) + + +def test_a_failed_marker_write_in_the_signal_handler_does_not_raise(tmp_path, capsys): + from yeto.rl.engine.sampler_resume import write_reclaim_marker + + env = {RECLAIM_MARKER_ENV: str(tmp_path / "missing-dir" / "marker")} + write_reclaim_marker(env, _signal.SIGINT) # T-C1 run b: the runner raced the same temp name + assert "reclaim marker not written" in capsys.readouterr().out diff --git a/tests/test_rl_synced_trainers_sampler.py b/tests/test_rl_synced_trainers_sampler.py new file mode 100644 index 00000000..61f4d096 --- /dev/null +++ b/tests/test_rl_synced_trainers_sampler.py @@ -0,0 +1,203 @@ +"""S25 C3 end to end on CPU: two synced trainer islands, the real Rust syncer +(elastic OuterStep merge), one sampler that follows the releases. + +The trainers are fake ports islands (``yeto.rl.engine.fake``) wired with the +production publisher wrappers (``sampler_island.wire_policy_releases``) and the +production lane reader (``buffer_rollout.BatchReader`` over a lane of one +``FileSampleBuffer``). The sampler is a thread: ``PolicyFollower`` loads each +release and writes groups stamped with the version it serves. + +Checks (the same as the GPU run's, SESSION25/gpu-review/T-C3.md): +1. both trainers apply the same merged policy hash at every outer version; +2. the sampler loads exactly the merged versions 0..R, in order, each with the + merged hash, all released by learner 0; +3. the two lanes are disjoint and no buffer position is trained twice; +4. the syncer ran R OuterSteps. + +Runs when YETO_TEST_ELASTIC_SYNCER points to a syncer binary (as +test_rl_inter_island_ports.py).""" + +from __future__ import annotations + +import asyncio +import json +import os +import subprocess +import threading +import time +from pathlib import Path +from types import SimpleNamespace + +import pytest +import torch + +from test_rl_engine_driver import NAME, _driver, _engine, _strict_config +from test_rl_inter_island_ports import _free_port +from yeto.rl.adapters.miles import buffer_rollout as br +from yeto.rl.adapters.miles.sampler_island import wire_policy_releases +from yeto.rl.elastic_client import ElasticClientConfig, ElasticIslandClient +from yeto.rl.engine.bridges import ElasticAvgSync +from yeto.rl.engine.merged_release import BufferLane +from yeto.rl.engine.policy_release import PolicyFollower, PolicyReleaseStore +from yeto.rl.engine.resume import CheckpointStore +from yeto.rl.engine.sample_buffer import FileSampleBuffer +from yeto.rl.engine.sample_roles import SampleBufferConfig +from yeto.rl.engine.trainer_lifecycle import SampleIndexEntry + +ROUNDS = 3 +MAX_POLICY_AGE = 2 +_EXECUTOR_LOCK = threading.Lock() + + +class _LaneRollout: + """The fake rollout pool of one trainer: each round first takes its groups + from the island's buffer lane (as ``BufferReaderRolloutFn``), then lets the + fake pool build the batch. ``_in_executor`` runs a function as in the Miles + rollout executor, with this island's reader active.""" + + def __init__(self, pool, reader: br.BatchReader, groups: int) -> None: + self.pool, self.reader, self.groups = pool, reader, groups + self.taken: list[tuple[int, list[int], dict[int, int]]] = [] + + def generate(self, rollout_id, **kwargs): + groups, report = self.reader.take(rollout_id, self.groups, current_version=None) + self.taken.append((rollout_id, [p for p, _ in groups], report.behavior_versions)) + return self.pool.generate(rollout_id, **kwargs) + + def _in_executor(self, fn, *args): + with _EXECUTOR_LOCK: + br._ACTIVE[br.TRAINER] = self.reader + return fn(None, *args) + + def __getattr__(self, name): + return getattr(self.pool, name) + + +class _Actor: + """``actor_model.export_hf``: the full weights the engine serves.""" + + def __init__(self, engine) -> None: + self.engine = engine + + async def export_hf(self, *, rollout_id: int, path: str) -> None: + torch.save(self.engine.published_tensors, Path(path) / "model.pt") + + +def _sampler(store_root: Path, stop: threading.Event, loads: list, groups_ahead: int) -> None: + store = CheckpointStore(store_root) + buffer = FileSampleBuffer(store, batch_size=1) + follower = PolicyFollower( + PolicyReleaseStore(store), + lambda release, directory: loads.append( + (release.version, release.policy_hash, release.origin, + torch.load(directory / "model.pt")[NAME].tolist()))) + n = 0 + while not stop.is_set(): + follower.refresh() + release = follower.release + if release is None or len(buffer) - int(release.buffer_selected or 0) >= groups_ahead: + time.sleep(0.005) + continue + ref = buffer.write_payload(f"g{n}", br.encode_group([{"tokens": [n]}])) + buffer.append(SampleIndexEntry(f"g{n}", follower.version, ref)) + n += 1 + buffer.mark_data_end() + + +def _events(path: Path, kind: str) -> list[dict]: + return [r for r in (json.loads(line) for line in path.read_text().splitlines()) + if r.get("event") == kind] + + +@pytest.mark.skipif(not os.environ.get("YETO_TEST_ELASTIC_SYNCER"), + reason="needs YETO_TEST_ELASTIC_SYNCER (syncer binary with elastic mode)") +def test_two_synced_trainers_feed_one_sampler_through_merged_releases(tmp_path): + port, key = _free_port(), "c3-e2e" + out = Path(os.environ.get("YETO_TEST_ELASTIC_OUT", tmp_path)) + out.mkdir(parents=True, exist_ok=True) + for old in out.glob("*.jsonl"): + old.unlink() + store_root = out / "store" + store_root.mkdir(exist_ok=True) + syncer = subprocess.Popen( + [os.environ["YETO_TEST_ELASTIC_SYNCER"], "--port", str(port), "--learners", "2", + "--total-steps", str(ROUNDS), "--event-tape", str(out / "syncer-tape.jsonl"), + "--island-scheduling-mode", "elastic", "--quorum-theta", "0.75", "--carry-gamma", "0.5", + "--soft-deadline-s", "30", "--q-min", "1", "--max-carry-lag", "2", + "--island-lease-s", "6", "--syncer-epoch", "5"], + env=dict(os.environ, YETO_ISLAND_HMAC_KEY=key), + stdout=open(out / "syncer.log", "w"), stderr=subprocess.STDOUT) + stop, loads = threading.Event(), [] + try: + time.sleep(1.0) + engines = (_engine(torch.tensor([1.0, 3.0])), _engine(torch.tensor([3.0, 5.0]))) + sampler = threading.Thread(target=_sampler, daemon=True, + args=(store_root, stop, loads, 4 * engines[0].groups)) + sampler.start() + drivers, rollouts = [], [] + for i, engine in enumerate(engines): + client = ElasticIslandClient(ElasticClientConfig(("127.0.0.1", port), i, lease_s=6.0, + syncer_epoch=5), key.encode()) + cfg = _strict_config(out, engine, learner_id=i, rounds=ROUNDS, tape=f"island-{i}.jsonl") + driver = _driver(engine, ElasticAvgSync(cfg, client=client, base_wait_s=30.0), out, + name=f"events-{i}.jsonl", learner_id=i) + lane = BufferLane(i, 2) + buffer = FileSampleBuffer(CheckpointStore(store_root), batch_size=engine.groups, lane=lane) + reader = br.BatchReader(buffer, max_policy_age=MAX_POLICY_AGE, wait_s=30.0, poll_s=0.01, + cursor_log=br.open_lane_log(buffer)) + rollout = _LaneRollout(driver.rollout, reader, engine.groups) + driver.rollout = rollout + wire_policy_releases( + driver, actor_model=_Actor(engine), runner=SimpleNamespace(run=asyncio.run), + miles_args=SimpleNamespace( + yeto_rl_sample_buffer=SampleBufferConfig(role="trainer", store=str(store_root)), + yeto_rl_sample_lane=lane, lora_rank=0)) + drivers.append(driver) + rollouts.append(rollout) + results, errors = {}, {} + + def run(i, d): + try: + results[i] = d.run() + except BaseException as error: # noqa: BLE001 + errors[i] = error + + threads = [threading.Thread(target=run, args=(i, d), daemon=True) + for i, d in enumerate(drivers)] + for t in threads: + t.start() + for t in threads: + t.join(timeout=60) + assert not t.is_alive() + assert errors == {} + assert syncer.wait(15) == 0 + finally: + stop.set() + if syncer.poll() is None: + syncer.kill() + + # 1. the same merged policy on both trainers at every outer version + applied = [{e["sync/outer_version"]: e["sync/global_policy_hash"] + for e in _events(out / f"events-{i}.jsonl", "rl_policy_apply")} for i in range(2)] + assert sorted(applied[0]) == list(range(ROUNDS + 1)) and applied[0] == applied[1] + assert torch.equal(results[0].tensors[NAME], results[1].tensors[NAME]) + # 2. the sampler loaded the merged versions in order, from the publisher only + assert [v for v, *_ in loads] == list(range(ROUNDS + 1)) + assert all(h == applied[0][v] for v, h, *_ in loads) + assert {o["learner_id"] for _, _, o, _ in loads} == {0} + assert loads[-1][3] == results[0].tensors[NAME].tolist() + released = _events(out / "events-0.jsonl", "rl_policy_released") + assert [e["policy_version"] for e in released] == list(range(ROUNDS + 1)) + assert _events(out / "events-1.jsonl", "rl_policy_released") == [] + # 3. disjoint lanes, no position trained twice + positions = [[p for _, ps, _ in r.taken for p in ps] for r in rollouts] + assert all(p % 2 == i for i, ps in enumerate(positions) for p in ps) + flat = positions[0] + positions[1] + assert len(flat) == len(set(flat)) == 2 * ROUNDS * engines[0].groups + # every trained group is within max_policy_age of the merged version + for r in rollouts: + for rollout_id, _, versions in r.taken: + assert all(rollout_id - v <= MAX_POLICY_AGE for v in versions) + # 4. R OuterSteps + kinds = [json.loads(line).get("kind") for line in (out / "syncer-tape.jsonl").read_text().splitlines()] + assert kinds.count("outer_step") == ROUNDS diff --git a/tests/test_rl_verl_adapter.py b/tests/test_rl_verl_adapter.py index 44f53542..96c58845 100644 --- a/tests/test_rl_verl_adapter.py +++ b/tests/test_rl_verl_adapter.py @@ -355,6 +355,19 @@ def test_image_recipe_is_pinned(): image.build_commands("c" * 40) +def test_image_adds_the_store_client_at_the_yeto_pin(): + import tomllib + from pathlib import Path + + from yeto.rl.adapters.verl import image + from yeto.rl.adapters.verl.pins import STORE_CLIENT_REQUIREMENTS + + project = tomllib.loads((Path(__file__).parents[1] / "pyproject.toml").read_text())["project"] + pins = set(project["dependencies"]).union(*project["optional-dependencies"].values()) + assert set(STORE_CLIENT_REQUIREMENTS) <= pins + assert image.store_client_commands()[0].endswith(" ".join(STORE_CLIENT_REQUIREMENTS)) + + def test_verl_capabilities_accept_default_and_tis_specs_only(): from yeto.rl.engine.algorithm import AlgorithmSpec from yeto.rl.engine.capabilities import CapabilityMismatch diff --git a/tests/test_rl_verl_buffer_trainer.py b/tests/test_rl_verl_buffer_trainer.py new file mode 100644 index 00000000..58dea824 --- /dev/null +++ b/tests/test_rl_verl_buffer_trainer.py @@ -0,0 +1,471 @@ +"""S25 C2: verl trainer island over the sample buffer (CPU, no verl/Ray). + +The trainer actor is a fake reached through the island's ``call``; the buffer, +the release store, the yeto driver, the fully_async rollout port (version map, +policy-age judgement), the feed, the publisher and the merged checkpoint are +the real ones. The end-to-end test runs the real engine sampler island (fake +engine) in a thread against the same store. +""" + +from __future__ import annotations + +import json +import threading +import time +from pathlib import Path + +import pytest +import torch +from safetensors.torch import load_file, save_file + +from yeto.rl.adapters.sampler_engines.fake import FakeEngine +from yeto.rl.adapters.verl import buffer_feed as bf +from yeto.rl.adapters.verl import fully_async_round as far +from yeto.rl.adapters.verl.buffer_ports import (CONSUMED_LOG, RELEASE_LOG, BufferReleasePublisher, + install_buffer_ports, is_base_adapter) +from yeto.rl.adapters.verl.fully_async_translate import VersionMap +from yeto.rl.adapters.verl.merged_release import MergeError, lora_pairs, write_merged_checkpoint +from yeto.rl.engine.algorithm import AlgorithmSpec, CorrectionSpec, ExecutionSpec +from yeto.rl.engine.buffer_payload import GroupWriter +from yeto.rl.engine.engine_sampler import EngineSamplerConfig, PromptItem, build_engine_sampler +from yeto.rl.engine.policy_release import PolicyReleaseStore +from yeto.rl.engine.resume import CheckpointStore +from yeto.rl.engine.sample_buffer import FileSampleBuffer +from yeto.rl.engine.sampler_checks import PostLoadChecks +from yeto.rl.engine.sampler_engine import SamplingParams + +TOLERANT = AlgorithmSpec(execution=ExecutionSpec(max_policy_staleness=1), + correction=CorrectionSpec(method="tis", tis_clip=2.0, tis_clip_low=0.0)) +MODULE = "model.layers.0.q_proj" +NAMES = (f"base_model.model.{MODULE}.lora_A.weight", f"base_model.model.{MODULE}.lora_B.weight") +SHAPE = bf.RowShape(prompt_length=6, response_length=8, pad_token_id=0) + + +def _sample(version, *, prompt=(11, 12), response=(21, 22, 23), reward=1.0, known=True, index=0, + provenance=None): + tokens = (list(prompt) if known else []) + list(response) + return {"index": index, "prompt": "p", "tokens": tokens, "response_length": len(response), + "rollout_log_probs": [-0.5] * len(response), "reward": reward, + "metadata": {"weight_version": str(version), "prompt_tokens_known": known, + "finish_reason": "stop", "provenance": provenance or {"engine": "fake"}}} + + +# ---------------------------------------------------------------------- rows +def test_group_rows_pad_like_the_agent_loop(): + rows = bf.group_rows([_sample(3, index=0), _sample(3, response=(31,), reward=0.0, index=1)], + uid="buf7", param_versions=[2, 2], shape=SHAPE) + t, nt = rows["tensors"], rows["non_tensor"] + assert t["prompts"][0] == [0, 0, 0, 0, 11, 12] + assert t["responses"][0] == [21, 22, 23, 0, 0, 0, 0, 0] + assert t["response_mask"][1] == [1, 0, 0, 0, 0, 0, 0, 0] + assert t["input_ids"][0] == [0, 0, 0, 0, 11, 12, 21, 22, 23, 0, 0, 0, 0, 0] + assert t["attention_mask"][0] == [0] * 4 + [1] * 5 + [0] * 5 + assert t["position_ids"][0] == [0, 0, 0, 0, 0, 1, 2, 3, 4, 4, 4, 4, 4, 4] + assert t["rm_scores"][0][:4] == [0.0, 0.0, 1.0, 0.0] and sum(t["rm_scores"][1]) == 0.0 + assert t["rollout_log_probs"][0][:4] == [-0.5, -0.5, -0.5, 0.0] + assert nt["uid"] == ["buf7", "buf7"] and nt["min_global_steps"] == [2, 2] == nt["max_global_steps"] + assert nt["provenance"][0] == {"engine": "fake"} and nt["finish_reason"] == ["stop", "stop"] + assert len(rows["meta_info"]["metrics"]) == 2 + + +def test_group_rows_rebuild_unknown_prompt_ids_with_the_trainer_tokenizer(): + rows = bf.group_rows([_sample(0, known=False)], uid="u", param_versions=[0], shape=SHAPE, + tokenize=lambda text: [5, 6, 7]) + assert rows["tensors"]["prompts"][0] == [0, 0, 0, 5, 6, 7] + with pytest.raises(Exception, match="no tokenizer"): + bf.group_rows([_sample(0, known=False)], uid="u", param_versions=[0], shape=SHAPE) + + +def test_group_rows_refuse_samples_wider_than_the_rows(): + with pytest.raises(bf.BufferFeedError, match="exceed"): + bf.group_rows([_sample(0, prompt=tuple(range(1, 8)))], uid="u", param_versions=[0], shape=SHAPE) + with pytest.raises(bf.BufferFeedError, match="param versions"): + bf.group_rows([_sample(0)], uid="u", param_versions=[], shape=SHAPE) + + +# ---------------------------------------------------------------------- feed +def _buffer(tmp_path): + return FileSampleBuffer(CheckpointStore(tmp_path), batch_size=1) + + +def test_feed_maps_versions_judges_and_records_each_position_once(tmp_path): + writer = GroupWriter(_buffer(tmp_path), "s0") + for key, version in (("a", 0), ("b", 1), ("c", 5)): + writer.write([_sample(version, index=0), _sample(version, index=1)], group_key=key) + reader, vmap, added, records = _buffer(tmp_path), VersionMap(), [], [] + vmap.record(0, 0) + vmap.record(1, 1) + feed = bf.BufferFeed(reader, vmap=lambda: vmap, shape=SHAPE, add_group=added.append, + record=records.append, sleep=lambda s: None) + collector = far.RoundCollector(vmap, 1, 1, 3) + for _ in range(3): + meta = feed.pull() + keep = collector.offer(far.SampleMeta(**meta)) + feed("yeto_keep_last" if keep else "yeto_drop_last") + assert [r["position"] for r in records] == [0, 1, 2] == feed.pulled_positions + assert [r["kept"] for r in records] == [True, True, False] # v5 was never published here + assert records[2]["param_versions"] == [None, None] + assert [rows["non_tensor"]["uid"][0] for rows in added] == ["buf0", "buf1"] + assert collector.discarded[0].unknown + assert feed.read_cursor() == 3 + + +def test_feed_returns_none_at_the_data_end_and_times_out_without_data(tmp_path): + buffer = _buffer(tmp_path) + clock = iter(range(0, 100, 10)) + feed = bf.BufferFeed(buffer, vmap=VersionMap, shape=SHAPE, add_group=print, record=print, + wait_s=15, clock=lambda: next(clock), sleep=lambda s: None) + with pytest.raises(bf.BufferFeedError, match="no buffer group within 15 s"): + feed.pull() + buffer.mark_data_end() + assert bf.BufferFeed(_buffer(tmp_path), vmap=VersionMap, shape=SHAPE, add_group=print, + record=print).pull() is None + with pytest.raises(bf.BufferFeedError, match="without a pulled group"): + feed.keep() + + +def test_buffer_call_routes_queue_calls_to_the_feed(): + seen = [] + call = bf.BufferCall(lambda name, *a: seen.append(("feed", name)), lambda name, *a, **k: seen.append(("actor", name))) + call("yeto_pull_sample") + call("yeto_train_step") + assert seen == [("feed", "yeto_pull_sample"), ("actor", "yeto_train_step")] + + +# ---------------------------------------------------------------------- merge +def _base(tmp_path, *, dtype=torch.bfloat16): + base = tmp_path / "base" + base.mkdir() + weights = {f"{MODULE}.weight": torch.randn(4, 4).to(dtype), "lm_head.weight": torch.randn(3, 4).to(dtype)} + save_file(weights, str(base / "model.safetensors")) + (base / "config.json").write_text(json.dumps({"model_type": "qwen3"})) + (base / "tokenizer.json").write_text("{}") + return base, weights + + +def test_merge_adds_b_times_a_and_keeps_the_layout(tmp_path): + base, weights = _base(tmp_path, dtype=torch.float32) + a, b = torch.randn(2, 4), torch.randn(4, 2) + summary = write_merged_checkpoint(base, {NAMES[0]: a, NAMES[1]: b}, tmp_path / "out") + out = load_file(str(tmp_path / "out" / "model.safetensors")) + assert torch.allclose(out[f"{MODULE}.weight"], weights[f"{MODULE}.weight"] + b @ a, atol=1e-6) + assert torch.equal(out["lm_head.weight"], weights["lm_head.weight"]) + assert summary["merged_modules"] == 1 and summary["max_relative_change"] > 0 + assert (tmp_path / "out" / "config.json").read_text() == (base / "config.json").read_text() + assert (tmp_path / "out" / "tokenizer.json").is_file() + + +def test_merge_keeps_the_base_dtype_and_refuses_unknown_or_misshaped_modules(tmp_path): + base, _ = _base(tmp_path) + write_merged_checkpoint(base, {NAMES[0]: torch.ones(2, 4), NAMES[1]: torch.ones(4, 2)}, tmp_path / "o1") + assert load_file(str(tmp_path / "o1" / "model.safetensors"))[f"{MODULE}.weight"].dtype == torch.bfloat16 + other = "base_model.model.model.layers.9.q_proj" + with pytest.raises(MergeError, match="no weight for 1 LoRA modules"): + write_merged_checkpoint(base, {f"{other}.lora_A.weight": torch.ones(2, 4), + f"{other}.lora_B.weight": torch.ones(4, 2)}, tmp_path / "o2") + with pytest.raises(MergeError, match="do not fit"): + write_merged_checkpoint(base, {NAMES[0]: torch.ones(2, 3), NAMES[1]: torch.ones(4, 2)}, tmp_path / "o3") + with pytest.raises(MergeError, match="only side"): + lora_pairs({NAMES[0]: torch.ones(2, 4)}) + + +def test_base_adapter_is_version_zero_with_zero_b(): + zero = {NAMES[0]: torch.ones(2, 4), NAMES[1]: torch.zeros(4, 2)} + assert is_base_adapter(0, zero) + assert not is_base_adapter(1, zero) + assert not is_base_adapter(0, {**zero, NAMES[1]: torch.full((4, 2), 1e-3)}) + + +# ---------------------------------------------------------------------- driver end to end +class FakeBufferActor: + """Stands in for YetoBufferTrainer: kept rows, a LoRA pair, a param version.""" + + def __init__(self): + self.tensors = {NAMES[0]: torch.full((2, 4), 0.1), NAMES[1]: torch.zeros(4, 2)} + self.param_version = 0 + self.kept, self.assembled = [], [] + + def __call__(self, name, *args, timeout=None): + return getattr(self, name)(*args) + + def yeto_configure(self, *_): + return None + + def yeto_add_group(self, rows): + self.kept.append(rows) + return len(self.kept) + + def yeto_assemble(self): + self.assembled.append([r["non_tensor"]["uid"][0] for r in self.kept]) + n, self.kept = len(self.kept), [] + return {"samples": n, "queue_size": None} + + def yeto_train_step(self): + self.tensors[NAMES[1]] += 0.01 + self.param_version += 1 + return {"metrics": {"actor/grad_norm": 0.5, "actor/pg_loss": 0.1, "actor/lr": 1e-5}, + "criteria": {}, "param_version": self.param_version} + + def yeto_param_version(self): + return self.param_version + + def yeto_export_lora(self): + return {k: v.clone() for k, v in self.tensors.items()}, {"count": 2} + + def yeto_call_workers(self, fn, tensors, reset): + for k in self.tensors: + self.tensors[k] = tensors[k].clone().float() + return [{"applied": len(tensors)}] + + +def _trainer(tmp_path, store_dir, actor, *, rounds, required, base): + from yeto.rl.adapters.verl.entry import verl_capabilities + from yeto.rl.adapters.verl.fully_async_ports import FullyAsyncIsland, FullyAsyncTrainerGroup + from yeto.rl.adapters.verl.ports_impl import VerlPolicyState + from yeto.rl.core import CanonicalTensorSpec, canonical_layout_hash + from yeto.rl.engine.bridges import LocalOnlySync + from yeto.rl.engine.driver import EventTape, IslandDriver + + specs = (CanonicalTensorSpec(NAMES[0], (2, 4), "float32", 8), CanonicalTensorSpec(NAMES[1], (4, 2), "float32", 8)) + island = FullyAsyncIsland(actor, learner_id=0, limit=1, required=required, out_dir=str(tmp_path / "out"), + readback_dir=str(tmp_path / "rb"), base_model_revision="a" * 40, + lora_config_hash="b" * 64, layout_hash=canonical_layout_hash(specs), + expected_specs=specs) + tape = EventTape(tmp_path / "tape.jsonl", 0) + island.emit = lambda event, **f: tape.append({"event": event, **f}) + store = CheckpointStore(store_dir) + rollout, publisher, placement = install_buffer_ports( + island, buffer=FileSampleBuffer(store, batch_size=1, create_index=False), + releases=PolicyReleaseStore(store), shape=bf.RowShape(64, 16, 0), base_dir=base, + tokenize=lambda text: [7, 8], poll_s=0.01, wait_s=30) + profile = far.execution_profile_for(TOLERANT, 1, groups_per_round=required, samples_per_group=2, sync="none") + driver = IslandDriver( + learner_id=0, rollout=rollout, trainer=FullyAsyncTrainerGroup(island), policy_state=VerlPolicyState(island), + publisher=publisher, placement=placement, + capabilities=verl_capabilities("sha256:" + "0" * 64, fully_async_limit=1), + algorithm=TOLERANT, sync=LocalOnlySync(rounds), events=tape, profile=profile) + return driver, island, publisher + + +def _lines(path): + return [json.loads(line) for line in Path(path).read_text().splitlines()] + + +def test_sampler_island_and_verl_buffer_trainer_run_three_rounds(tmp_path): + store_dir = tmp_path / "store" + store = CheckpointStore(store_dir) + base, weights = _base(tmp_path, dtype=torch.float32) + sampler_events, done = [], threading.Event() + engine = FakeEngine(max_model_len=64) + config = EngineSamplerConfig(sampler_id="island0", num_rollouts=40, max_backlog_groups=2, + sampling=SamplingParams(max_tokens=6, n=2), max_prompt_tokens=32, poll_s=0.01) + sampler = build_engine_sampler( + config, engine=engine, releases=PolicyReleaseStore(store), buffer=FileSampleBuffer(store, batch_size=1), + emit=lambda n, **f: sampler_events.append((n, f)), + prompts=lambda r: [PromptItem(f"p{r}", f"prompt {r}", label="x")], reward=lambda item, c: float(len(c.token_ids) % 2), + checks=PostLoadChecks(expected_expert_matrices=None, fixed_prompts=(), probe_pairs=(("Q", " 7"),)), + should_stop=done.is_set, + sleep=lambda s: time.sleep(0.005)) + thread = threading.Thread(target=sampler.run, daemon=True) + thread.start() + actor = FakeBufferActor() + driver, island, publisher = _trainer(tmp_path, store_dir, actor, rounds=3, required=2, base=base) + try: + state = driver.run() + finally: + done.set() + thread.join(timeout=10) + assert driver.rounds_completed == 3 and state.policy_version == 3 + assert island.vmap.to_dict() == {"0": 0, "1": 1, "2": 2, "3": 3} + # no group consumed twice; trained groups are those handed to the actor + consumed = _lines(tmp_path / "out" / CONSUMED_LOG.format(learner=0)) + positions = [r["position"] for r in consumed] + assert positions == sorted(set(positions)) == list(range(len(positions))) + kept = [r for r in consumed if r["kept"]] + assert len(kept) == 6 and [len(a) for a in actor.assembled] == [2, 2, 2] + assert sorted(u for a in actor.assembled for u in a) == sorted(f"buf{r['position']}" for r in kept) + # versions: v0 (the base) is not released; v1..v3 are, in order + releases = _lines(tmp_path / "out" / RELEASE_LOG.format(learner=0)) + assert [(r["policy_version"], r["released"]) for r in releases] == [(0, False), (1, True), (2, True), (3, True)] + assert [r["seq"] for r in releases if r["released"]] == [1, 2, 3] + latest = PolicyReleaseStore(store).latest() + assert (latest.version, latest.weights, latest.buffer_selected) == (3, "full", releases[-1]["buffer_selected"]) + merged = load_file(str(store_dir / "policy-releases" / latest.path / "model.safetensors")) + a, b = actor.tensors[NAMES[0]], actor.tensors[NAMES[1]] + assert torch.allclose(merged[f"{MODULE}.weight"], weights[f"{MODULE}.weight"] + b @ a, atol=1e-6) + # the sampler loaded the releases in order, and the trained groups respect the age limit + loaded = [f["policy_version"] for n, f in sampler_events if n == "rl_sampler_policy_loaded"] + assert loaded == sorted(set(loaded)) and set(loaded) <= {1, 2, 3} + events = _lines(tmp_path / "tape.jsonl") + carry = [e for e in events if e["event"] == "rl_rollout_carry_over"] + assert [e["rollout_id"] for e in carry] == [0, 1, 2] + rounds = [e for e in events if e["event"] == "rl_verl_fa_round"] + assert sum(e["kept"] for e in rounds) == 6 + assert sum(e["pulled"] for e in rounds) == len(consumed) + + +def test_publisher_releases_the_base_when_its_adapter_is_not_zero(tmp_path): + from yeto.rl.adapters.verl.fully_async_ports import FullyAsyncIsland + + actor = FakeBufferActor() + actor.tensors[NAMES[1]] += 1.0 + island = FullyAsyncIsland(actor, learner_id=0, limit=1, required=1, out_dir=str(tmp_path / "out"), + readback_dir=str(tmp_path / "rb")) + store = CheckpointStore(tmp_path / "store") + feed = bf.BufferFeed(FileSampleBuffer(store, batch_size=1), vmap=lambda: island.vmap, shape=SHAPE, + add_group=print, record=print) + + class State: + policy_version = 0 + + @staticmethod + def policy_tensor_hash(): + return "0" * 64 + + @staticmethod + def to_lora(): + return type("L", (), {"tensors": actor.tensors}) + + calls = [] + publisher = BufferReleasePublisher(island, PolicyReleaseStore(store), feed, base_dir=tmp_path, + merge=lambda base, tensors, d: calls.append(d) or {"merged_modules": 1}) + result = publisher.publish(State()) + assert len(calls) == 1 and island.published == (0, "0" * 64) + assert result.manifest.target_policy_version == 0 and island.vmap.to_dict() == {"0": 0} + + +# ---------------------------------------------------------------------- island entry +def _entry_args(*extra): + from yeto.rl.adapters.verl.island_entry import parse_args + + return parse_args(["--model", "m", "--model-revision", "r", "--data", "d", "--learner-id", "0", + "--reward-function", "gsm8k_reward:score", "--global-rounds", "3", + "--groups-per-round", "2", "--samples-per-group", "2", + "--rollout-max-response-len", "64", "--event-tape", "t", "--inner-lr", "1e-5", + "--seq-len", "512", "--seed", "1", *extra])[0] + + +def test_island_entry_plans_a_buffer_trainer_and_refuses_other_roles(): + from yeto.rl.adapters.verl.island_entry import sample_buffer_section + + assert sample_buffer_section(_entry_args()) is None + section = sample_buffer_section(_entry_args("--rl-sample-role", "trainer", "--rl-resume-store", "/s", + "--rl-single-island-no-sync", + "--rl-max-policy-age", "1", "--rl-sample-wait-s", "60")) + assert section == {"store": "/s", "wait_s": 60.0, "lane": None} + with pytest.raises(SystemExit, match="engine sampler"): + sample_buffer_section(_entry_args("--rl-sample-role", "sampler", "--rl-resume-store", "/s")) + with pytest.raises(SystemExit, match="max-policy-age"): + sample_buffer_section(_entry_args("--rl-sample-role", "trainer", "--rl-resume-store", "/s")) + + +def test_island_entry_plans_a_synced_trainer_lane_and_refuses_partial_flags(): + from yeto.rl.adapters.verl.island_entry import sample_buffer_section + + synced = ("--rl-sample-role", "trainer", "--syncer", "h:1", "--rl-max-policy-age", "2") + section = sample_buffer_section(_entry_args(*synced, "--rl-sample-store", "/s", "--rl-sample-lanes", "2")) + assert section == {"store": "/s", "wait_s": None, "lane": "0/2"} + for extra, message in ((("--rl-sample-store", "/s"), "needs --rl-sample-store and --rl-sample-lanes"), + (("--rl-sample-store", "/s", "--rl-sample-lanes", "2", "--rl-resume-store", "/s"), + "not --rl-resume-store")): + with pytest.raises(SystemExit, match=message): + sample_buffer_section(_entry_args(*synced, *extra)) + with pytest.raises(SystemExit, match="need --rl-sample-role trainer with a syncer"): + sample_buffer_section(_entry_args("--rl-sample-role", "trainer", "--rl-single-island-no-sync", + "--rl-max-policy-age", "1", "--rl-resume-store", "/s", + "--rl-sample-lanes", "2")) + + +# ---------------------------------------------------------------------- synced trainers (S25 C3) +class _State: + """A published policy state: local version, hash, the LoRA tensors.""" + + def __init__(self, version, tensors, label): + import hashlib + + self.policy_version, self._tensors = version, tensors + self.label, self._digest = label, hashlib.sha256(label.encode()).hexdigest() + + def policy_tensor_hash(self): + return self._digest + + def to_lora(self): + return type("L", (), {"tensors": self._tensors}) + + +def _synced_publisher(tmp_path, store, lane, events): + from yeto.rl.adapters.verl.fully_async_ports import FullyAsyncIsland + from yeto.rl.engine.merged_release import BufferLane, lane_cursor_log + + actor = FakeBufferActor() + island = FullyAsyncIsland(actor, learner_id=lane, limit=2, required=1, + out_dir=str(tmp_path / f"out{lane}"), readback_dir=str(tmp_path / "rb")) + (tmp_path / f"out{lane}").mkdir(exist_ok=True) + island.emit = lambda event, **f: events.append((event, f)) + buffer = FileSampleBuffer(store, batch_size=1, lane=BufferLane(lane, 2)) + _, publisher, _ = install_buffer_ports( + island, buffer=buffer, releases=PolicyReleaseStore(store), shape=SHAPE, base_dir=tmp_path, + cursor_log=lane_cursor_log(buffer)) + def merge(base, tensors, directory): + (directory / "model.safetensors").write_bytes(b"w") + return {} + + publisher.merge = merge + driver = type("D", (), {"outer_version": None})() + publisher.bind(driver) + return actor, island, publisher, driver + + +def test_only_the_publisher_releases_each_new_merged_version_once(tmp_path): + store = CheckpointStore(tmp_path / "store") + nonzero = {NAMES[0]: torch.ones(2, 4), NAMES[1]: torch.ones(4, 2)} + zero = {NAMES[0]: torch.ones(2, 4), NAMES[1]: torch.zeros(4, 2)} + events = {0: [], 1: []} + islands = {m: _synced_publisher(tmp_path, store, m, events[m]) for m in (0, 1)} + # (local version, outer version of the applied base): base, a local round without a + # merge, merges 1 and 2, then merge 4 (merge 3 missed: a gap) + steps = [(0, 0, zero), (1, 0, nonzero), (2, 1, nonzero), (3, 2, nonzero), (4, 4, nonzero)] + for local, outer, tensors in steps: + for m, (actor, island, publisher, driver) in islands.items(): + actor.param_version = local + driver.outer_version = outer + publisher.publish(_State(local, tensors, f"h{outer}-{local}")) + released = [f for e, f in events[0] if e == "rl_policy_released"] + assert [(f["policy_version"], f["local_version"]) for f in released] == [(1, 2), (2, 3), (4, 4)] + assert all(f["origin"] == {"kind": "merged", "learner_id": 0, "num_learners": 2, + "local_version": f["local_version"]} for f in released) + assert [f for e, f in events[0] if e == "rl_policy_release_gap"] == [{"previous_version": 2, "policy_version": 4}] + assert not [e for e, _ in events[1] if e == "rl_policy_released"] + assert ("rl_policy_release_role", {"learner_id": 1, "num_learners": 2, "publisher": False}) in events[1] + latest = PolicyReleaseStore(store).latest() + assert (latest.seq, latest.version, latest.policy_hash) == (3, 4, _State(4, {}, "h4-4").policy_tensor_hash()) + # both islands pair each merged version with the first local version that applied it + for _, island, publisher, _ in islands.values(): + assert publisher.feed.merged.to_dict() == {"0": 0, "1": 2, "2": 3, "4": 4} + + +def test_a_synced_feed_reads_its_lane_translates_merged_versions_and_logs_each_round(tmp_path): + store = CheckpointStore(tmp_path / "store") + writer = GroupWriter(FileSampleBuffer(store, batch_size=1), "s0") + for i, version in enumerate((0, 0, 1, 1, 3, 3)): # merged behavior versions + writer.write([_sample(version, index=0)], group_key=f"g{i}") + events = [] + actor, island, publisher, driver = _synced_publisher(tmp_path, store, 1, events) + for local, outer in ((0, 0), (1, 1), (2, 1)): # merge 1 applied at local 1; local 2 trains on + actor.param_version = local + driver.outer_version = outer + publisher.publish(_State(local, {NAMES[1]: torch.zeros(4, 2)}, f"h{local}")) + feed = publisher.feed + pulled = [] + for _ in range(3): + meta = feed.pull() + pulled.append((meta["uid"], meta["min_steps"])) + feed.keep() if None not in meta["min_steps"] else feed.drop() + # lane 1 of 2: positions 1, 3, 5; merged v0 -> local 0, v1 -> local 1, v3 never applied here + assert pulled == [("buf1", (0,)), ("buf3", (1,)), ("buf5", (None,))] + island.call("yeto_assemble") # the round's collection ended: one lane cursor line + lane_file = tmp_path / "store" / "sample-buffer" / "lanes" / "1-of-2.jsonl" + assert _lines(lane_file) == [{"rollout_id": 2, "positions": [1, 3], "cursor": 6}] + # a restarted island continues after the recorded cursor + again = _synced_publisher(tmp_path, store, 1, [])[2].feed + assert again.buffer.read_cursor == 6 diff --git a/tests/test_rl_verl_cut.py b/tests/test_rl_verl_cut.py index 0736c287..390d4351 100644 --- a/tests/test_rl_verl_cut.py +++ b/tests/test_rl_verl_cut.py @@ -271,3 +271,74 @@ class _Cfg(dict): """Attribute + ``get`` access, like an omegaconf node.""" __getattr__ = dict.__getitem__ + + +# ---------------------------------------------------------------- N7.6c: digest over the stepped state + + +def test_restore_whose_optimizer_steps_a_stale_snapshot_is_refused(tmp_path, monkeypatch): + """Inject the T-G3 bug shape: after the load the module holds the cut's weights, but the + optimizer steps a snapshot taken when the fresh trainer was built. The restore digest + reads the stepped tensors, so the restore is refused (a module-side digest would match).""" + import torch.distributed.checkpoint.state_dict as dcp_sd + + real = dcp_sd.set_state_dict + + def load_then_step_stale(module, optimizer, **kwargs): + stale = {id(p): p.detach().clone() for p in module.parameters()} # build-time values + real(module, optimizer, **kwargs) + for group in optimizer.param_groups: + for i, param in enumerate(group["params"]): + copy = torch.nn.Parameter(stale[id(param)]) + optimizer.state[copy] = optimizer.state.pop(param) + group["params"][i] = copy + + monkeypatch.setattr(dcp_sd, "set_state_dict", load_then_step_stale) + trainer, fresh = _single(tmp_path) + saved = torch.load(tmp_path / "cut-v" / vcut.rank_file(0), weights_only=True)["model"] + with pytest.raises(ResumeAssertionError, match="loaded train state differs from the cut"): + trainer.restore_cut("cut-v", epoch=1, root=str(tmp_path), expect=_expect()) + for name, p in fresh.module.named_parameters(): # the module copy alone looks restored + assert torch.equal(p.detach(), saved[name]), name + + +def test_digest_reads_the_tensors_the_optimizer_steps(): + h = _handles(_model()) + _round(h, 0) + before = vcut._digest(h) + first = h.optimizer.param_groups[0]["params"][0] + with torch.no_grad(): + first.add_(1.0) + assert vcut._digest(h) != before + + +def test_frozen_base_weights_are_not_in_the_cut_and_do_not_block_the_restore(tmp_path): + """LoRA shape: frozen base weights are skipped at save; a strict load would report + them missing (GPU T-C4 a, 2026-10-20).""" + + def frozen_model(): + model = _model() + model[0].weight.requires_grad_(False) + return model + + def handles(model): + opt = torch.optim.AdamW([p for p in model.parameters() if p.requires_grad], lr=1e-2) + return TrainHandles(model, opt, torch.optim.lr_scheduler.LambdaLR(opt, lambda s: 1.0 / (1 + s))) + + trained = handles(frozen_model()) + for r in range(CUT): + _round(trained, r) + worker = SimpleNamespace(module=trained.module, optimizer=trained.optimizer, lr_scheduler=trained.scheduler) + vcut.VerlCutTrainer(lambda fn, *a: [fn(worker, *a)], optimizer_steps_per_round=K).save_cut( + epoch=1, context=_context(tmp_path)) + saved = torch.load(tmp_path / "cut-v" / vcut.rank_file(0), weights_only=True)["model"] + assert "0.weight" not in saved and "0.bias" in saved + fresh = handles(frozen_model()) + fresh_worker = SimpleNamespace(module=fresh.module, optimizer=fresh.optimizer, lr_scheduler=fresh.scheduler) + vcut.VerlCutTrainer(lambda fn, *a: [fn(fresh_worker, *a)], optimizer_steps_per_round=K).restore_cut( + "cut-v", epoch=1, root=str(tmp_path), expect=_expect()) + for (name, p), (_, q) in zip(trained.module.named_parameters(), fresh.module.named_parameters()): + if p.requires_grad: # the frozen base keeps the fresh build's values (both from seed 0 here) + assert torch.equal(p, q), name + for key in ("exp_avg", "exp_avg_sq", "step"): + assert torch.equal(trained.optimizer.state[p][key], fresh.optimizer.state[q][key]), key diff --git a/tests/test_rl_verl_round_cuts.py b/tests/test_rl_verl_round_cuts.py new file mode 100644 index 00000000..6653eaa6 --- /dev/null +++ b/tests/test_rl_verl_round_cuts.py @@ -0,0 +1,248 @@ +"""S25 N7.6b (CPU): the verl island's round-cut controller. + +The real cut controller (ResumeController + RoundCutCheckpoint on LocalOnlySync, +wired by ``adapters.verl.round_cuts``) and the real verl cut trainer +(``VerlCutTrainer``: rank file, digests, resume assertions) over a torch model, +behind the fake driver ports. The policy the driver publishes is read from the +torch model, so the resume's policy-hash check compares real restored weights. +""" + +from __future__ import annotations + +from types import SimpleNamespace + +import pytest +import torch +from test_rl_verl_cut import K, _handles, _model, _round + +from yeto.rl.adapters.verl import cut as vcut +from yeto.rl.adapters.verl import round_cuts +from yeto.rl.elastic_client import LeaveReason +from yeto.rl.engine import recovery, trainer_remote +from yeto.rl.engine.algorithm import AlgorithmSpec +from yeto.rl.engine.bridges import LocalOnlySync +from yeto.rl.engine.driver import EventTape, IslandDriver +from yeto.rl.engine.fake import FakeEngine, FakeRolloutPool, FakeTrainerGroup, fake_capabilities +from yeto.rl.engine.recovery import CUT_KEY +from yeto.rl.engine.resume_manifest import ResumeAssertionError + +NAME = "base_model.model.layer.lora_A.weight" +PROMPTS = 4 # prompts drawn per round +ROUNDS = 5 + + +class _Pool(FakeRolloutPool): + """The fake pool with the verl data cursor shape (``sample_offset`` only).""" + + def __init__(self, engine): + super().__init__(engine) + self.drawn = 0 + + def generate(self, rollout_id, **kwargs): + self.drawn += PROMPTS + return super().generate(rollout_id, **kwargs) + + def data_cursor(self): + return {"sample_offset": self.drawn} + + def seek_data_cursor(self, cursor): + target = int(cursor["sample_offset"]) + if target < self.drawn or (target - self.drawn) % PROMPTS: + raise ValueError(f"cannot seek {self.drawn} -> {target}") + self.drawn = target + return self.data_cursor() + + +class _TorchTrainer(FakeTrainerGroup): + """Fake trainer port whose train state is a torch model behind ``VerlCutTrainer``.""" + + def __init__(self, engine, handles, *, skip_optimizer_round=None): + super().__init__(engine) + self.handles = handles + worker = SimpleNamespace(module=handles.module, optimizer=handles.optimizer, + lr_scheduler=handles.scheduler) + self.cuts = vcut.VerlCutTrainer(lambda fn, *a: [fn(worker, *a)], optimizer_steps_per_round=K) + self.skip_optimizer_round = skip_optimizer_round + self._sync_policy() + + def _sync_policy(self): + self.engine.tensors[NAME] = self.handles.module[0].weight.detach()[:1, :2].clone().float() + + def actual_layout(self): + return self.cuts.actual_layout() + + def layout(self): + return self.actual_layout() + + def save_cut(self, *, epoch, context): + return self.cuts.save_cut(epoch=epoch, context=context) + + def restore_cut(self, cut_id, *, epoch, root, expect, shared_filesystem=True): + manifest = self.cuts.restore_cut(cut_id, epoch=epoch, root=root, expect=expect, + shared_filesystem=shared_filesystem) + self._sync_policy() + return manifest + + def confirm_resumed_data_cursor(self, cursor): + self.cuts.confirm_resumed_data_cursor(cursor) + + def train_step(self, batch): + receipt = super().train_step(batch) + if batch.rollout_id == self.skip_optimizer_round: + self.handles.scheduler.step() # injected: the round took no optimizer step + else: + _round(self.handles, batch.rollout_id) + self._sync_policy() + self.cuts.check_first_step(1) + return receipt + + +def _plan(store=None, state_dir=None, **resume): + config = None + if store is not None: + config = {"store": str(store), "state_dir": str(state_dir), "every": 1, "keep": 2, + "allow_config_change": False, "stop_after": None, **resume} + return {"sync": "none", "resume": config, "model_revision": "rev", "lora_targets": "all-linear", + "run_config": {"lr": 1e-2, "seed": 1, "total_rounds": ROUNDS, "out_dir": "/x"}} + + +def _island(tmp_path, plan, *, name, rounds=ROUNDS, **trainer_kwargs): + engine = FakeEngine(tensors={NAME: torch.zeros(1, 2)}, step_delta=1.0) + engine.rollout = _Pool(engine) + handles = _handles(_model()) + engine.trainer = _TorchTrainer(engine, handles, **trainer_kwargs) + resume = round_cuts.resume_wiring_for(plan, runtime_fingerprint="verl@test") + tape = tmp_path / f"{name}.jsonl" + driver = IslandDriver( + learner_id=0, rollout=engine.rollout, trainer=engine.trainer, policy_state=engine.policy_state, + publisher=engine.publisher, placement=engine.placement, algorithm=AlgorithmSpec(), + capabilities=fake_capabilities("miles-like"), sync=LocalOnlySync(rounds), + events=EventTape(tape, 0), ledger=None if resume is None else resume.ledger) + checkpoint = round_cuts.wire_round_cuts(driver, resume, plan=plan, algorithm=AlgorithmSpec(), + install_signals=False) + return driver, handles, checkpoint, resume, tape + + +def _events(path, kind): + import json + + return [e for e in map(json.loads, path.read_text().splitlines()) if e.get("event") == kind] + + +def _params(handles): + return {n: p.detach().clone() for n, p in handles.module.named_parameters()} + + +def test_stop_after_then_resume_equals_continuous_run(tmp_path): + store = tmp_path / "store" + driver, _, checkpoint, _, tape = _island(tmp_path, _plan(store, tmp_path / "a", stop_after=3), name="a") + round_cuts.run_driver(driver, None) + pointer = checkpoint.pointer() + assert pointer["next_rollout_id"] == 3 and pointer["final"] is True + assert pointer[CUT_KEY]["trigger"] == recovery.MANUAL + assert [e["rollout_id"] for e in _events(tape, "rl_cut_saved")] == [1, 2, 3] + + # a fresh container: new state dir, restored from the store's LATEST + driver, resumed, checkpoint, _, tape = _island(tmp_path, _plan(store, tmp_path / "b"), name="b") + round_cuts.run_driver(driver, None) + [resume] = _events(tape, "rl_resume") + assert resume["rollout_id"] == 3 and resume["checks"]["policy_hash"] == "same" + assert _events(tape, "rl_data_cursor_restored")[0]["data_cursor"] == {"sample_offset": 3 * PROMPTS} + assert checkpoint.pointer()["next_rollout_id"] == ROUNDS + assert driver.rounds_completed == ROUNDS - 3 + + driver, continuous, _, _, _ = _island(tmp_path, _plan(), name="c") + round_cuts.run_driver(driver, None) + a, b = _params(resumed), _params(continuous) + assert all(torch.equal(a[n], b[n]) for n in a) + assert resumed.scheduler.last_epoch == continuous.scheduler.last_epoch == ROUNDS + + +def test_cut_every_and_keep(tmp_path): + store = tmp_path / "store" + driver, _, _, _, tape = _island(tmp_path, _plan(store, tmp_path / "a", every=2, keep=1), name="a") + round_cuts.run_driver(driver, None) + assert [e["rollout_id"] for e in _events(tape, "rl_cut_saved")] == [2, 4, 5] # every 2, and the end + + +def test_first_resumed_round_without_an_optimizer_step_is_refused(tmp_path): + store = tmp_path / "store" + driver, _, _, _, _ = _island(tmp_path, _plan(store, tmp_path / "a", stop_after=2), name="a") + round_cuts.run_driver(driver, None) + driver, _, _, _, _ = _island(tmp_path, _plan(store, tmp_path / "b"), name="b", skip_optimizer_round=2) + with pytest.raises(Exception) as info: + round_cuts.run_driver(driver, None) + chain, error = [], info.value + while error is not None: + chain.append(error) + error = error.__cause__ or error.__context__ + assert any(isinstance(e, ResumeAssertionError) and "Adam step" in str(e) for e in chain) + + +def test_scheduler_stop_cuts_with_scheduled_and_resumes(tmp_path): + """The head's trainer scheduler: run 2 rounds, stop (STOPPED -> SCHEDULED cut); + the next incarnation restores that cut.""" + store = tmp_path / "store" + driver, _, checkpoint, resume, tape = _island(tmp_path, _plan(store, tmp_path / "a"), name="a") + files = trainer_remote.LocalStoreFiles(resume.controller.store) + writer = trainer_remote.CommandWriter(files, 0) + writer.send(trainer_remote.RUN, rounds=2) + writer.send(trainer_remote.STOP, reason=LeaveReason.STOPPED) + env = {trainer_remote.CONTROL_ENV: "1", trainer_remote.INCARNATION_ENV: "0", + trainer_remote.RESUME_ENV: "", trainer_remote.POLL_ENV: "0.01"} + round_cuts.run_driver(driver, resume, environ=env) + pointer = checkpoint.pointer() + assert pointer["next_rollout_id"] == 2 and pointer[CUT_KEY]["trigger"] == recovery.SCHEDULED + phases = [e["phase"] for e in _events(tape, trainer_remote.EVENT)] + assert phases == ["started", "ready", "trained", "stopped"] + + driver, _, _, _, tape = _island(tmp_path, _plan(store, tmp_path / "b"), name="b") + round_cuts.run_driver(driver, None) + assert _events(tape, "rl_resume")[0]["rollout_id"] == 2 + + +def test_commanded_run_needs_round_cuts(tmp_path): + driver, _, _, _, _ = _island(tmp_path, _plan(), name="a") + with pytest.raises(RuntimeError, match="needs round cuts"): + round_cuts.run_driver(driver, None, environ={trainer_remote.CONTROL_ENV: "1"}) + + +def _args(**kw): + base = dict(rl_resume_store=None, rl_resume_state_dir=None, rl_cut_every=None, rl_cut_keep=None, + rl_stop_after_rounds=None, rl_resume_allow_config_change=False, + rl_single_island_no_sync=True, rl_max_policy_age=0) + return SimpleNamespace(**{**base, **kw}) + + +def test_resume_config_from_flags(): + assert round_cuts.resume_config(_args()) is None + config = round_cuts.resume_config(_args(rl_resume_store="/s", rl_cut_every=2, rl_stop_after_rounds=3)) + assert config == {"store": "/s", "state_dir": round_cuts.DEFAULT_STATE_DIR, "every": 2, "keep": 2, + "allow_config_change": False, "stop_after": 3} + with pytest.raises(ValueError, match="need --rl-resume-store"): + round_cuts.resume_config(_args(rl_cut_every=2)) + with pytest.raises(ValueError, match="--rl-single-island-no-sync"): + round_cuts.resume_config(_args(rl_resume_store="/s", rl_single_island_no_sync=False)) + with pytest.raises(ValueError, match="fully_async"): + round_cuts.resume_config(_args(rl_resume_store="/s", rl_max_policy_age=2)) + with pytest.raises(ValueError, match="must be >= 1"): + round_cuts.resume_config(_args(rl_resume_store="/s", rl_cut_keep=0)) + + +def test_buffer_trainer_store_is_not_a_cut_store(): + # S25 C2 x N7.6b: a verl buffer trainer's --rl-resume-store is its sample buffer + trainer = dict(rl_sample_role="trainer", rl_resume_store="/s", rl_max_policy_age=1) + assert round_cuts.resume_config(_args(**trainer)) is None + with pytest.raises(ValueError, match="not wired for a verl buffer trainer"): + round_cuts.resume_config(_args(**trainer, rl_cut_every=2)) + + +def test_changed_run_config_is_refused(tmp_path): + store = tmp_path / "store" + driver, _, _, _, _ = _island(tmp_path, _plan(store, tmp_path / "a", stop_after=2), name="a") + round_cuts.run_driver(driver, None) + plan = _plan(store, tmp_path / "b") + plan["run_config"] = {**plan["run_config"], "lr": 2e-2, "total_rounds": 9} # rounds may change + driver, _, _, _, _ = _island(tmp_path, plan, name="b") + with pytest.raises(Exception, match="run configuration differs"): + round_cuts.run_driver(driver, None) diff --git a/tests/test_rl_verl_synced_trainers.py b/tests/test_rl_verl_synced_trainers.py new file mode 100644 index 00000000..85547c6d --- /dev/null +++ b/tests/test_rl_verl_synced_trainers.py @@ -0,0 +1,186 @@ +"""S25 C3V end to end on CPU: two synced verl buffer trainers, the real Rust +syncer (elastic OuterStep merge), one engine sampler island that follows the +merged releases. + +The trainers are the verl buffer trainer of test_rl_verl_buffer_trainer.py +(fake trainer actor; real driver, fully_async rollout port, feed, publisher, +merged checkpoint) over one lane each of one buffer, with ``ElasticAvgSync``. +The sampler is the real engine sampler island (fake engine) in a thread. + +Checks (the same as the GPU run's, SESSION25/gpu-review/T-C3V.md): +1. both trainers apply the same merged policy hash at every outer version; +2. the sampler loads only versions the publisher released, in order, and each + released hash is the merged hash of that outer version; learner 1 releases + nothing; +3. the two lanes are disjoint and no buffer position is trained twice; +4. the syncer ran R OuterSteps. + +Runs when YETO_TEST_ELASTIC_SYNCER points to a syncer binary (as +test_rl_synced_trainers_sampler.py).""" + +from __future__ import annotations + +import json +import os +import subprocess +import threading +import time +from pathlib import Path + +import pytest +import torch + +from test_rl_inter_island_ports import _free_port +from test_rl_verl_buffer_trainer import NAMES, FakeBufferActor, _base +from yeto.rl.adapters.sampler_engines.fake import FakeEngine +from yeto.rl.adapters.verl import buffer_feed as bf +from yeto.rl.adapters.verl import fully_async_round as far +from yeto.rl.adapters.verl.buffer_ports import RELEASE_LOG, install_buffer_ports +from yeto.rl.engine.algorithm import AlgorithmSpec, CorrectionSpec, ExecutionSpec +from yeto.rl.engine.engine_sampler import EngineSamplerConfig, PromptItem, build_engine_sampler +from yeto.rl.engine.merged_release import BufferLane, lane_cursor_log +from yeto.rl.engine.policy_release import PolicyReleaseStore +from yeto.rl.engine.resume import CheckpointStore +from yeto.rl.engine.sample_buffer import FileSampleBuffer, open_trainer_buffer +from yeto.rl.engine.sampler_checks import PostLoadChecks +from yeto.rl.engine.sampler_engine import SamplingParams + +ROUNDS = 3 +GROUPS = 2 +MAX_POLICY_AGE = 2 +TOLERANT = AlgorithmSpec(execution=ExecutionSpec(max_policy_staleness=MAX_POLICY_AGE), + correction=CorrectionSpec(method="tis", tis_clip=2.0, tis_clip_low=0.0)) + + +def _lines(path: Path) -> list[dict]: + return [json.loads(line) for line in path.read_text().splitlines()] if path.is_file() else [] + + +def _island(out: Path, store_root: Path, base: Path, learner_id: int, port: int, key: str): + from yeto.rl.adapters.verl.entry import verl_capabilities + from yeto.rl.adapters.verl.fully_async_ports import FullyAsyncIsland, FullyAsyncTrainerGroup + from yeto.rl.adapters.verl.ports_impl import VerlPolicyState + from yeto.rl.bridge import BridgeConfig + from yeto.rl.core import CanonicalTensorSpec, canonical_layout_hash + from yeto.rl.elastic_client import ElasticClientConfig, ElasticIslandClient + from yeto.rl.engine.bridges import ElasticAvgSync + from yeto.rl.engine.driver import EventTape, IslandDriver + + specs = (CanonicalTensorSpec(NAMES[0], (2, 4), "float32", 8), + CanonicalTensorSpec(NAMES[1], (4, 2), "float32", 8)) + actor = FakeBufferActor() + island = FullyAsyncIsland(actor, learner_id=learner_id, limit=MAX_POLICY_AGE, required=GROUPS, + out_dir=str(out / f"out{learner_id}"), readback_dir=str(out / f"rb{learner_id}"), + base_model_revision="a" * 40, lora_config_hash="b" * 64, + layout_hash=canonical_layout_hash(specs), expected_specs=specs) + (out / f"out{learner_id}").mkdir(exist_ok=True) + tape = EventTape(out / f"tape-{learner_id}.jsonl", learner_id) + island.emit = lambda event, **f: tape.append({"event": event, **f}) + store = CheckpointStore(store_root) + buffer = open_trainer_buffer(store, batch_size=1, create_index=False, lane=BufferLane(learner_id, 2)) + rollout, publisher, placement = install_buffer_ports( + island, buffer=buffer, releases=PolicyReleaseStore(store), shape=bf.RowShape(64, 16, 0), + base_dir=base, tokenize=lambda text: [7, 8], poll_s=0.01, wait_s=30, + cursor_log=lane_cursor_log(buffer)) + bridge = BridgeConfig(syncer_addr=("127.0.0.1", port), learner_id=learner_id, global_rounds=ROUNDS, + groups_per_round=GROUPS, samples_per_group=2, local_optimizer_steps=1, + expected_specs=specs, base_model_revision="a" * 40, lora_config_hash="b" * 64, + layout_hash=canonical_layout_hash(specs), event_tape=str(out / f"bridge-{learner_id}.jsonl"), + wan_streams=0) + client = ElasticIslandClient(ElasticClientConfig(("127.0.0.1", port), learner_id, lease_s=6.0, + syncer_epoch=5), key.encode()) + profile = far.execution_profile_for(TOLERANT, MAX_POLICY_AGE, groups_per_round=GROUPS, samples_per_group=2, + sync="elastic") + driver = IslandDriver( + learner_id=learner_id, rollout=rollout, trainer=FullyAsyncTrainerGroup(island), + policy_state=VerlPolicyState(island), publisher=publisher, placement=placement, + capabilities=verl_capabilities("sha256:" + "0" * 64, fully_async_limit=MAX_POLICY_AGE), + algorithm=TOLERANT, sync=ElasticAvgSync(bridge, client=client, base_wait_s=30.0, syncer_epoch=5), + events=tape, profile=profile) + publisher.bind(driver) + return driver, island + + +@pytest.mark.skipif(not os.environ.get("YETO_TEST_ELASTIC_SYNCER"), + reason="needs YETO_TEST_ELASTIC_SYNCER (syncer binary with elastic mode)") +def test_two_synced_verl_trainers_feed_one_engine_sampler(tmp_path): + port, key = _free_port(), "c3v-e2e" + out = tmp_path + store_root = out / "store" + store = CheckpointStore(store_root) + base, _ = _base(tmp_path, dtype=torch.float32) + syncer = subprocess.Popen( + [os.environ["YETO_TEST_ELASTIC_SYNCER"], "--port", str(port), "--learners", "2", + "--total-steps", str(ROUNDS), "--event-tape", str(out / "syncer-tape.jsonl"), + "--island-scheduling-mode", "elastic", "--quorum-theta", "0.75", "--carry-gamma", "0.5", + "--soft-deadline-s", "30", "--q-min", "1", "--max-carry-lag", "2", + "--island-lease-s", "6", "--syncer-epoch", "5"], + env=dict(os.environ, YETO_ISLAND_HMAC_KEY=key), + stdout=open(out / "syncer.log", "w"), stderr=subprocess.STDOUT) + done, sampler_events = threading.Event(), [] + try: + time.sleep(1.0) + config = EngineSamplerConfig(sampler_id="s0", num_rollouts=200, max_backlog_groups=4 * GROUPS, + sampling=SamplingParams(max_tokens=6, n=2), max_prompt_tokens=32, + poll_s=0.01) + sampler = build_engine_sampler( + config, engine=FakeEngine(max_model_len=64), releases=PolicyReleaseStore(store), + buffer=FileSampleBuffer(store, batch_size=1), emit=lambda n, **f: sampler_events.append((n, f)), + prompts=lambda r: [PromptItem(f"p{r}", f"prompt {r}", label="x")], + reward=lambda item, c: float(len(c.token_ids) % 2), + checks=PostLoadChecks(expected_expert_matrices=None, fixed_prompts=()), # N2-fix: no probe pairs + should_stop=done.is_set, sleep=lambda s: time.sleep(0.005)) + sampler_thread = threading.Thread(target=sampler.run, daemon=True) + sampler_thread.start() + islands = [_island(out, store_root, base, i, port, key) for i in range(2)] + results, errors = {}, {} + + def run(i, driver): + try: + results[i] = driver.run() + except BaseException as error: # noqa: BLE001 + errors[i] = error + + threads = [threading.Thread(target=run, args=(i, d), daemon=True) for i, (d, _) in enumerate(islands)] + for t in threads: + t.start() + for t in threads: + t.join(timeout=90) + assert not t.is_alive() + assert errors == {} + assert syncer.wait(15) == 0 + finally: + done.set() + if syncer.poll() is None: + syncer.kill() + + tapes = [_lines(out / f"tape-{i}.jsonl") for i in range(2)] + # 1. the same merged policy on both trainers at every outer version + applied = [{e["sync/outer_version"]: e["sync/global_policy_hash"] for e in t if e["event"] == "rl_policy_apply"} + for t in tapes] + assert sorted(applied[0]) == list(range(ROUNDS + 1)) and applied[0] == applied[1] + # 2. releases: the publisher only, merged versions increasing, each with the merged hash + released = [e for e in tapes[0] if e["event"] == "rl_policy_released"] + assert [e["policy_version"] for e in released] == list(range(1, ROUNDS + 1)) # v0 is the base + assert all(e["policy_hash"] == applied[0][e["policy_version"]] for e in released) + assert all(e["origin"]["learner_id"] == 0 for e in released) + assert not [e for e in tapes[1] if e["event"] == "rl_policy_released"] + assert not [e for t in tapes for e in t if e["event"] == "rl_policy_release_gap"] + loaded = [f["policy_version"] for n, f in sampler_events if n == "rl_sampler_policy_loaded"] + assert loaded == sorted(set(loaded)) and set(loaded) <= {e["policy_version"] for e in released} + latest = PolicyReleaseStore(store).latest() + assert latest.version == ROUNDS and latest.origin["learner_id"] == 0 + # 3. disjoint lanes, no position trained twice, lane cursor lines per round + lanes = [_lines(store_root / "sample-buffer" / "lanes" / f"{i}-of-2.jsonl") for i in range(2)] + positions = [[p for line in lane for p in line["positions"]] for lane in lanes] + assert all(p % 2 == i for i, ps in enumerate(positions) for p in ps) + flat = positions[0] + positions[1] + assert len(flat) == len(set(flat)) == 2 * ROUNDS * GROUPS + assert [len(lane) for lane in lanes] == [ROUNDS, ROUNDS] + # the islands paired merged versions with local ones identically (one merge per round) + for i in range(2): + records = _lines(out / f"out{i}" / RELEASE_LOG.format(learner=i)) + assert records[-1]["merged_versions"] == {str(v): v for v in range(ROUNDS + 1)} + # 4. R OuterSteps + kinds = [json.loads(line).get("kind") for line in (out / "syncer-tape.jsonl").read_text().splitlines()] + assert kinds.count("outer_step") == ROUNDS diff --git a/tests/test_run_flags.py b/tests/test_run_flags.py index afc9b9ea..9ef5c8ea 100644 --- a/tests/test_run_flags.py +++ b/tests/test_run_flags.py @@ -269,7 +269,7 @@ def test_forward_tables_name_flags_of_both_layers(): "--temp-dir", "--ulimit", "--untracked-files", "--wandb-group", "--yes", # torchrun, tail and the W4A8 re-quantizer CLI in the SSH trainer island (yeto/launch/ssh_trainer_island.py) "--awq-scales", "--loop", "--master-addr", "--master-port", "--no-awq", "--node-rank", "--nproc-per-node", - "--pid", "--poll-seconds", "--recipe", "--store", "--template", + "--pid", "--poll-seconds", "--recipe", "--store", "--template", "--frozen-base", "--frozen-pattern", # docker run of a runtime: docker SSH trainer node (ssh_trainer_island.py, node_inventory.py example) "--device", "--network", }) diff --git a/tests/test_s25_g1_stage1mini.py b/tests/test_s25_g1_stage1mini.py new file mode 100644 index 00000000..e6d2d311 --- /dev/null +++ b/tests/test_s25_g1_stage1mini.py @@ -0,0 +1,262 @@ +"""S25 T-G1: the pieces the stage-1 mini rehearsal added (CPU).""" +from __future__ import annotations + +import json +from pathlib import Path + +import pytest + +from yeto.rl.adapters.miles.models import glm5_next +from yeto.rl.adapters.sampler_engines import sglang_load_report as slr +from yeto.rl.adapters.sampler_engines.server import ServerConfig +from yeto.rl.adapters.sampler_engines.sglang_server import SGLangOptions, SGLangServerEngine +from yeto.rl.engine.sampler_checks import FixedPrompt, PostLoadChecks +from yeto.rl.engine.sampler_engine import Completion + +REPO = Path(__file__).resolve().parents[1] + + +def test_recipe_takes_expert_and_dense_counts_from_the_checkpoint(): + shape = glm5_next.checkpoint_geometry({"text_config": {"num_hidden_layers": 4, "first_k_dense_replace": 1, + "n_routed_experts": 32}}) + assert shape == {"num_layers": 4, "first_k_dense_replace": 1, "num_experts": 32} + argv = glm5_next.model_argv(**shape) + assert argv[argv.index("--num-experts") + 1] == "32" + assert argv[argv.index("--moe-layer-freq") + 1] == "[0,1,1,1]" + full = glm5_next.model_argv() + assert full[full.index("--num-experts") + 1] == "288" and full.count("--num-experts") == 1 + with pytest.raises(ValueError): + glm5_next.model_argv(num_experts=0) + + +class _Param: + def __init__(self, calls: list[str]) -> None: + self.weight_loader = lambda *a, **k: calls.append(a[2] if len(a) > 2 else "?") + + +class _FakeModel: + """Maps names like SGLang's GLM5.3 loader: an expert data tensor reaches a loader + only when its fused name is a parameter.""" + + def __init__(self) -> None: + self.calls: list[str] = [] + self.params = {"model.layers.1.mlp.experts.w13_weight": _Param(self.calls), + "model.layers.1.mlp.experts.w2_weight": _Param(self.calls), + "model.layers.1.mlp.experts.w13_weight_scale": _Param(self.calls)} + + def named_parameters(self): + return iter(self.params.items()) + + def load_weights(self, weights): + for name, tensor in weights: + fused = (name.replace("experts.0.gate_proj.", "experts.w13_").replace("experts.0.up_proj.", "experts.w13_") + .replace("experts.0.down_proj.", "experts.w2_").replace("language_model.", "")) + param = self.params.get(fused) + if param is not None: + param.weight_loader(param, tensor, fused) + + +def _weights(suffix: str): + base = "model.language_model.layers.1.mlp.experts.0." + return [(f"{base}{m}_proj.{suffix}", 0) for m in ("gate", "up", "down")] + \ + [(f"{base}{m}_proj.weight_scale", 0) for m in ("gate", "up", "down")] + + +def test_load_report_counts_only_loaded_expert_matrices(tmp_path, monkeypatch): + cls = type("M", (_FakeModel,), {}) + slr.install(cls) + slr.install(cls) # idempotent + monkeypatch.setenv(slr.REPORT_DIR_ENV, str(tmp_path)) + monkeypatch.delenv(slr.RENAME_ENV, raising=False) + cls().load_weights(_weights("weight_packed")) # no rename: SGLang skips the packed data silently + assert slr.read_reports(tmp_path)["expert_matrices"] == 0 + slr.clear_reports(tmp_path) + assert slr.read_reports(tmp_path) is None + monkeypatch.setenv(slr.RENAME_ENV, "1") + model = cls() + model.load_weights(_weights("weight_packed")) + report = slr.read_reports(tmp_path) + assert report["expert_matrices"] == 3 and report["ranks"] == 1 + assert sorted(set(model.calls)) == ["model.layers.1.mlp.experts.w13_weight", "model.layers.1.mlp.experts.w13_weight_scale", + "model.layers.1.mlp.experts.w2_weight"] + assert list(slr.rename_packed([("a.mlp.experts.3.up_proj.weight_packed", 1), ("x.weight_packed", 2)])) == \ + [("a.mlp.experts.3.up_proj.weight", 1), ("x.weight_packed", 2)] + + +def test_patch_model_file_is_idempotent(tmp_path): + f = tmp_path / "glm5_next.py" + f.write_text("class A: ...\nEntryClass = [A]\n") + assert slr.patch_model_file(f) is True and slr.patch_model_file(f) is False + assert f.read_text().count(slr.MARKER) == 1 + compile(f.read_text(), str(f), "exec") + (tmp_path / "other.py").write_text("x = 1\n") + with pytest.raises(ValueError): + slr.patch_model_file(tmp_path / "other.py") + + +def _engine(tmp_path: Path, tp: int = 1) -> SGLangServerEngine: + config = ServerConfig(base_dir=tmp_path, card_model="H100", device_family="cuda", engine_version="x", + precision="mxfp4", quant_scheme="mxfp4-g32-rtn", topology=f"tp{tp}", max_model_len=64) + return SGLangServerEngine(config, SGLangOptions(tp=tp, renames_weight_packed=True, + load_report_dir=str(tmp_path / "reports"))) + + +def test_adapter_reads_the_report_and_exports_the_hook_env(tmp_path): + engine = _engine(tmp_path, tp=2) + env = engine.env() + assert env[slr.RENAME_ENV] == "1" and env[slr.REPORT_DIR_ENV] == str(tmp_path / "reports") + assert engine.loaded_expert_matrices() is None + reports = tmp_path / "reports" + reports.mkdir() + (reports / "rank0-1.json").write_text(json.dumps({"expert_matrices": 288})) + assert engine.loaded_expert_matrices() == 0 # rank 1 did not report + (reports / "rank1-2.json").write_text(json.dumps({"expert_matrices": 287})) + assert engine.loaded_expert_matrices() == 287 + with pytest.raises(Exception): + SGLangServerEngine(ServerConfig(base_dir=tmp_path, card_model="H100", device_family="cuda", engine_version="x", + precision="mxfp4", quant_scheme="mxfp4-g32-rtn", topology="tp2", + max_model_len=64), SGLangOptions(tp=1)) + + +class _Greedy: + def generate(self, prompts, sampling): + return [[Completion(text="the the the", token_ids=(1, 2, 3), logprobs=(-1.0,) * 3, finish_reason="length")] + for _ in prompts] + + +def test_fixed_answers_may_be_recorded_without_failing(): + prompts = (FixedPrompt("Q: 2+3?", "5"),) + strict = PostLoadChecks(fixed_prompts=prompts).fixed_prompt_answers(_Greedy()) + assert strict["passed"] is False and strict["answers"][0]["problems"] == ["answer missing"] + loose = PostLoadChecks(fixed_prompts=prompts, require_fixed_answers=False).fixed_prompt_answers(_Greedy()) + assert loose["passed"] is True and loose["answers"][0]["answered"] is False + assert loose["answers_required"] is False + + +def test_mini_modal_plan_has_no_errors_and_uses_the_volume_store(): + from yeto.rl.stage1 import expand, load_plan + + out = expand(load_plan(REPO / "examples/stage1/glm5next-mini-modal.yaml")) + assert out.errors == [] and out.gaps == [] + env = out.trainer["env"] + assert env["YETO_RL_STORE_IMPL"] == "yeto.cloud.modal_ckpt_store:store_from_env" + assert env["YETO_RL_STORE_MODAL_VOLUME"] == "yeto-s25-g1-store" + argv = out.trainer["miles"]["island_argv"] + assert argv[argv.index("--miles-commit") + 1] == "1618b5c38db80e17ec928cfa47fe832cb18565fc" + assert out.trainer["miles"]["model_args"][out.trainer["miles"]["model_args"].index("--num-experts") + 1] == "32" + spec = out.sampler["engines"][0]["spec"] + assert spec["checks"] == {"require_fixed_answers": False, "fixed_max_tokens": 64, + "expected_expert_matrices": 288} + assert spec["engine"]["options"]["load_report_dir"].endswith("/load-report") + assert "--frozen-base" in out.requant["argv"] + assert argv[argv.index("--megatron-ref-load") + 1] == out.trainer["torch_dist"] + assert out.trainer["convert_argv"][out.trainer["convert_argv"].index("--save") + 1] == out.trainer["torch_dist"] + assert "--qkv-format" not in out.trainer["convert_argv"] and "--spec" in out.trainer["convert_argv"] + # the PPU plan keeps OSS and passes no --miles-commit for its short commit + ppu = expand(load_plan(REPO / "examples/stage1/glm53-flash-ppu.yaml")) + assert ppu.errors == [] and "--miles-commit" not in ppu.trainer["miles"]["island_argv"] + assert "--frozen-base" not in ppu.requant["argv"] + + +def test_island_argv_check_runs_the_island_combination_rules(): + from yeto.rl.stage1 import _check_island_argv, expand, load_plan + + argv = expand(load_plan(REPO / "examples/stage1/glm5next-mini-modal.yaml")).trainer["miles"]["island_argv"] + assert "--rollout-num-gpus" not in argv and _check_island_argv(argv) is None + problem = _check_island_argv([*argv, "--rollout-num-gpus", "0"]) + assert problem and "fixed-partition" in problem + + +_MINI_TEXT_CONFIG = { + "model_type": "glm5_next_text", "num_hidden_layers": 4, "first_k_dense_replace": 1, "hidden_size": 4096, + "intermediate_size": 12288, "num_attention_heads": 64, "num_key_value_heads": 64, "n_routed_experts": 32, + "n_shared_experts": 1, "moe_intermediate_size": 2048, "num_experts_per_tok": 8, "vocab_size": 154880, + "max_position_embeddings": 1048576, "rms_norm_eps": 1e-5, "q_lora_rank": 1536, "kv_lora_rank": 512, + "qk_head_dim": 256, "qk_rope_head_dim": 0, "v_head_dim": 256, "hidden_act": "silu", "attention_bias": False, +} + + +def test_mini_island_argv_translates_to_a_miles_argv_on_cpu(tmp_path): + """The island path up to the Miles argv (Modal runs 1-3 each failed one step of it).""" + from yeto.rl.adapters.miles.island_entry import build_ports_launch, parse_args + from yeto.rl.engine.run_config import resolve_rl_run_config + from yeto.rl.stage1 import expand, load_plan + + model = tmp_path / "model" + model.mkdir() + (model / "config.json").write_text(json.dumps({"model_type": "glm5_next", "text_config": _MINI_TEXT_CONFIG})) + prompts = tmp_path / "p.jsonl" + prompts.write_text('{"messages": [{"role": "user", "content": "hi"}], "label": "1"}\n') + argv = expand(load_plan(REPO / "examples/stage1/glm5next-mini-modal.yaml")).trainer["miles"]["island_argv"][3:] + argv = [str(model) if a == "/data/models/glm5next-mini-bf16-7fb7883" else a for a in argv] + argv = [str(REPO / a) if a.startswith("examples/rl_algorithms/") else a for a in argv] + args = parse_args(argv) + provider = glm5_next.glm5_next_provider(model) + config = resolve_rl_run_config(args, model_path=model, prompt_path=prompts, provider=provider, + target_modules=[], yeto_policy_sync=False, verify_ref_load=False) + assert config.geometry.num_layers == 4 and config.geometry.moe.num_experts == 32 + miles = list(build_ports_launch(args, config, []).argv) + for flag in ("--debug-train-only", "--use-distributed-optimizer", "--optimizer-cpu-offload", + "--use-precision-aware-optimizer", "--use-tis"): + assert flag in miles + assert miles[miles.index("--num-experts") + 1] == "32" and miles.count("--num-experts") == 1 + assert miles[miles.index("--moe-layer-freq") + 1] == "[0,1,1,1]" + assert miles[miles.index("--megatron-to-hf-mode") + 1] == "raw" + + +def test_align_weight_version_is_a_no_op_on_a_trainer_only_island(): + from types import SimpleNamespace + + from yeto.rl.adapters.miles.state_plugin import align_weight_version + + align_weight_version(SimpleNamespace(weight_updater=None, args=SimpleNamespace(debug_train_only=True)), + next_version=3) + with pytest.raises(RuntimeError, match="no weight_updater"): + align_weight_version(SimpleNamespace(weight_updater=None, args=SimpleNamespace(debug_train_only=False)), + next_version=3) + updater = SimpleNamespace(weight_version=0) + align_weight_version(SimpleNamespace(weight_updater=updater, args=None), next_version=3) + assert updater.weight_version == 2 + + +def test_cut_state_walk_does_not_touch_the_chained_optimizer_property(): + from types import SimpleNamespace + + from yeto.rl.adapters.miles.cut_plugin import _state_dicts, side_effect_free_state + + class Chained: + def __init__(self, members): + self.chained_optimizers = members + self.state = {} + + @property + def optimizer(self): # Megatron asserts len(chained_optimizers) == 1 + raise AssertionError("ChainedOptimizer has more than one optimizer when accessing self.optimizer") + + dense = SimpleNamespace(optimizer=SimpleNamespace(state={"a": {"step": 1}}), state=None) + expert = SimpleNamespace(optimizer=SimpleNamespace(state={}), state=None) + chained = Chained([dense, expert]) + assert len(_state_dicts(chained)) == 3 + with side_effect_free_state(chained): + expert.optimizer.state["new"] = {} + assert "new" not in expert.optimizer.state and dense.optimizer.state == {"a": {"step": 1}} + + +def test_cut_coverage_groups_dense_tensors_over_the_ep_ranks(): + from yeto.rl.adapters.miles.trainer import optimizer_coverage_problems + + dense, expert = "decoder.layers.0.mlp.linear_fc1.weight", "decoder.layers.1.mlp.experts.linear_fc1.weight0" + shared = "decoder.layers.1.mlp.shared_experts.linear_fc1.weight" + + def rank(ep, optimized): + return {"coord": {"tp": 0, "pp": 0, "ep": ep}, "adapter_names": [dense, shared, expert], + "optimizer_names": optimized} + + # EP2 x DP2: the dense DistOpt shards split over both EP ranks; each rank owns its experts + ok = [rank(0, [dense, expert]), rank(1, [shared, expert])] + assert optimizer_coverage_problems(ok) == [] + bad = [rank(0, [dense, expert]), rank(1, [shared])] + problems = optimizer_coverage_problems(bad) + assert problems == [f"tp0/pp0/ep1: no optimizer state for ['{expert}']"] + assert bad[1]["has_optimizer_state"] is False and "has_optimizer_state" not in bad[0] + assert optimizer_coverage_problems([]) == ["no rank reported any adapter tensor"] diff --git a/tests/test_s25_g1c_reward_limits.py b/tests/test_s25_g1c_reward_limits.py new file mode 100644 index 00000000..ca14a026 --- /dev/null +++ b/tests/test_s25_g1c_reward_limits.py @@ -0,0 +1,67 @@ +"""S25 T-G1C: the token-diversity test reward and the G1 driver's wall-clock limits (CPU).""" + +from __future__ import annotations + +import sys +from pathlib import Path + +import pytest + +from yeto.rl import token_diversity_reward as tdr +from yeto.rl.engine import engine_sampler_entry as entry + +sys.path.insert(0, str(Path(__file__).resolve().parents[1] / "tools" / "probes" / "s25_g1")) +from run_limits import Deadline, wait_until # noqa: E402 + + +def test_token_diversity_reward_scores_token_ids_not_words(): + import asyncio + from types import SimpleNamespace + + one_word = SimpleNamespace(tokens=[1, 2] + [7] * 8, response_length=8) # "PrototypePrototype..." + varied = SimpleNamespace(tokens=[1, 2, 3, 4, 3, 5], response_length=4) + assert asyncio.run(tdr.score(None, one_word)) == pytest.approx(1 / 8) + assert asyncio.run(tdr.score(None, varied)) == pytest.approx(3 / 4) + assert tdr.diversity([]) == 0.0 + with pytest.raises(ValueError): + tdr.response_ids(SimpleNamespace(tokens=[1], response_length=3)) + + +def test_miles_reward_wrapper_passes_token_ids(): + from yeto.rl.engine.engine_sampler import PromptItem + from yeto.rl.engine.sampler_engine import Completion + + reward = entry.miles_reward(tdr.score) + done = Completion(text="x", token_ids=(5, 5, 6, 7), logprobs=(0.0,) * 4, finish_reason="length", + prompt_token_ids=(1, 2)) + assert reward(PromptItem("p", "q"), done) == pytest.approx(3 / 4) + + +class _Clock: + def __init__(self) -> None: + self.now = 0.0 + + def __call__(self) -> float: + return self.now + + def sleep(self, s: float) -> None: + self.now += s + + +def test_sampler_wait_after_the_trainers_has_a_limit(): + clock = _Clock() + assert wait_until(lambda: False, 300, poll_s=15, clock=clock, sleep=clock.sleep) is False + assert 300 <= clock.now < 315 + clock.now = 0.0 + assert wait_until(lambda: clock.now >= 60, 300, poll_s=15, clock=clock, sleep=clock.sleep) is True + + +def test_run_deadline_ends_a_wait_early(): + clock = _Clock() + deadline = Deadline(100, clock=clock) + assert not deadline.passed() + assert wait_until(lambda: False, 1000, poll_s=10, clock=clock, sleep=clock.sleep, + deadline=deadline) is False + assert deadline.passed() and clock.now == 100 + with pytest.raises(ValueError): + Deadline(0) diff --git a/tests/test_s25_g1c_sampler_resume.py b/tests/test_s25_g1c_sampler_resume.py new file mode 100644 index 00000000..88ad0fa0 --- /dev/null +++ b/tests/test_s25_g1c_sampler_resume.py @@ -0,0 +1,178 @@ +"""S25 T-G1C: a replaced engine sampler island continues where the buffer says +and never overwrites a payload or an earlier incarnation's tape (CPU).""" + +from __future__ import annotations + +import json +import os +import signal +from pathlib import Path + +import pytest + +from tests.test_stage1_gaps import _sampler_spec +from yeto.rl.engine import engine_sampler_entry as entry +from yeto.rl.engine.resume import CheckpointStore +from yeto.rl.engine.sample_buffer import FileSampleBuffer, SampleBufferError +from yeto.rl.engine.sampler_resume import MEMBERS_DIR, RECLAIM_MARKER_ENV, SamplerResumeError + +STATE_ENV = "YETO_TEST_G1C_REWARD_STATE" # the reward is imported by name: share state via a file + + +def counting_reward(item, completion) -> float: + """Fails (a lost container) or writes the reclaim marker after the limits in the state file.""" + path = Path(os.environ[STATE_ENV]) + state = json.loads(path.read_text()) + state["n"] += 1 + path.write_text(json.dumps(state)) + if state.get("limit") is not None and state["n"] > state["limit"]: + raise RuntimeError("container lost") + if state.get("mark_at") is not None and state["n"] == state["mark_at"]: + Path(state["marker"]).write_text("{}") + return 0.5 + + +class _Calls: + def __init__(self, path: Path) -> None: + self.path = path + self.update(n=0) + + def update(self, **fields) -> None: + state = json.loads(self.path.read_text()) if self.path.exists() else {} + self.path.write_text(json.dumps({**state, **fields})) + + +@pytest.fixture(autouse=True) +def _reset(monkeypatch, tmp_path): + global CALLS + monkeypatch.setenv(STATE_ENV, str(tmp_path / "reward-state.json")) + CALLS = _Calls(tmp_path / "reward-state.json") + monkeypatch.setenv(RECLAIM_MARKER_ENV, str(tmp_path / "reclaim-marker")) + handlers = {s: signal.getsignal(s) for s in (signal.SIGINT, signal.SIGTERM)} + yield + for sig, handler in handlers.items(): + signal.signal(sig, handler) + + +def _spec(tmp_path: Path, rounds: int = 4) -> dict: + return _sampler_spec(tmp_path, num_rollouts=rounds, reward="tests.test_s25_g1c_sampler_resume:counting_reward") + + +def _buffer(tmp_path: Path) -> FileSampleBuffer: + buffer = FileSampleBuffer(CheckpointStore(tmp_path / "store"), batch_size=1) + buffer.refresh() + return buffer + + +def _payloads(tmp_path: Path) -> dict[str, bytes]: + return {p.name: p.read_bytes() for p in (tmp_path / "store" / "sample-buffer" / "payloads").iterdir()} + + +def _tape(path: Path) -> list[dict]: + return [json.loads(line) for line in path.read_text().splitlines()] + + +def test_replacement_continues_after_the_written_rounds_and_overwrites_nothing(tmp_path): + spec = _spec(tmp_path) + CALLS.update(limit=8) # 2 prompts x n=2 per round: the third round fails before it writes + with pytest.raises(RuntimeError, match="container lost"): + entry.run(spec, resource_sampler=lambda emit: None, incarnation="aaaa0001") + first = _payloads(tmp_path) + assert len(first) == 4 and all(name.startswith("ppu4-aaaa0001-r") for name in first) + + CALLS.update(n=0, limit=None) + assert entry.run(spec, resource_sampler=lambda emit: None, incarnation="bbbb0002") == 2 + after = _payloads(tmp_path) + assert {k: after[k] for k in first} == first # the first incarnation's payloads are untouched + assert len(after) == 8 + buffer = _buffer(tmp_path) + entries = buffer.entries() + assert buffer.data_ended and len({e.sample_id for e in entries}) == len(entries) == 8 + rounds = [(e.sampler_round["incarnation"], e.sampler_round["rollout_id"]) for e in entries] + assert rounds == [("aaaa0001", 0)] * 2 + [("aaaa0001", 1)] * 2 + [("bbbb0002", 2)] * 2 \ + + [("bbbb0002", 3)] * 2 + assert all(buffer.read_payload(e.payload_ref) for e in entries) + + tapes = {p.name: _tape(p) for p in (tmp_path / "tape").glob("events.*.jsonl")} + assert set(tapes) == {"events.aaaa0001.jsonl", "events.bbbb0002.jsonl"} + old = {e["event"]: e for e in tapes["events.aaaa0001.jsonl"]} + assert old["rl_sampler_leave"]["reason"] == "INFRA_FAILED" + new = {e["event"]: e for e in tapes["events.bbbb0002.jsonl"]} + assert new["rl_sampler_join"]["next_rollout_id"] == 2 + assert new["rl_sampler_join"]["incarnations"] == ["aaaa0001"] + assert new["rl_sampler_leave"]["reason"] == "FINISHED" + members = sorted(p.name for p in (tmp_path / "store" / MEMBERS_DIR / "ppu4").iterdir()) + assert members == ["aaaa0001.json", "bbbb0002.json"] + + +def test_a_finished_sampler_restarted_generates_nothing(tmp_path): + spec = _spec(tmp_path, rounds=2) + assert entry.run(spec, resource_sampler=lambda emit: None, incarnation="aaaa0001") == 2 + before = _payloads(tmp_path) + assert entry.run(spec, resource_sampler=lambda emit: None, incarnation="bbbb0002") == 0 + assert _payloads(tmp_path) == before + + +def test_a_written_round_is_refused_to_a_new_incarnation(tmp_path): + spec = _spec(tmp_path) + CALLS.update(limit=4) + with pytest.raises(RuntimeError): + entry.run(spec, resource_sampler=lambda emit: None, incarnation="aaaa0001") + from yeto.rl.engine.engine_sampler import EngineRound + from yeto.rl.engine.sampler_resume import resume_point + + resume = resume_point(_buffer(tmp_path).entries(), "ppu4") + assert resume.next_rollout_id == 1 + round_ = EngineRound.__new__(EngineRound) + from yeto.rl.engine.sampler_resume import RoundGuard + + round_.guard = RoundGuard(resume) + with pytest.raises(SamplerResumeError, match="round 0 is already"): + EngineRound.__call__(round_, 0) + + +def test_a_payload_is_never_overwritten(tmp_path): + buffer = FileSampleBuffer(CheckpointStore(tmp_path / "store"), batch_size=1) + buffer.write_payload("ppu4-a-r0-row0", b"one") + with pytest.raises(SampleBufferError, match="never overwritten"): + buffer.write_payload("ppu4-a-r0-row0", b"two") + assert (tmp_path / "store" / "sample-buffer" / "payloads" / "ppu4-a-r0-row0.bin").read_bytes() == b"one" + + +def test_reclaim_marker_stops_writes_and_leaves_preempted(tmp_path): + spec = _spec(tmp_path) + CALLS.update(marker=str(tmp_path / "reclaim-marker"), mark_at=6) # inside round 1 + # round 1 was generated (2 rounds), but none of its groups is written + assert entry.run(spec, resource_sampler=lambda emit: None, incarnation="aaaa0001") == 2 + buffer = _buffer(tmp_path) + assert not buffer.data_ended + assert [e.sampler_round["rollout_id"] for e in buffer.entries()] == [0, 0] + [tape] = (tmp_path / "tape").glob("events.*.jsonl") + leave = [e for e in _tape(tape) if e["event"] == "rl_sampler_leave"] + assert len(leave) == 1 and leave[0]["reason"] == "PREEMPTED" + + +def test_caller_stop_ends_the_island_finished_and_on_event_sees_the_tape(tmp_path): + """T-E2V: the eval harness's end condition (trainers done) stops the loop + without a reclaim: the round in progress is still written and the island + leaves FINISHED; on_event sees every record written to the tape.""" + spec = _spec(tmp_path) + seen = [] + rounds = entry.run(spec, resource_sampler=lambda emit: None, incarnation="aaaa0001", + stop=lambda: any(e == "rl_sampler_round" for e, _ in seen), + on_event=lambda event, **fields: seen.append((event, fields))) + assert rounds == 1 + assert [e.sampler_round["rollout_id"] for e in _buffer(tmp_path).entries()] == [0, 0] + [tape] = (tmp_path / "tape").glob("events.*.jsonl") + assert [e["event"] for e in _tape(tape)] == [e for e, _ in seen] + assert [f["reason"] for e, f in seen if e == "rl_sampler_leave"] == ["FINISHED"] + + +def test_hold_after_reclaim_only_when_the_marker_is_there(tmp_path): + """A reclaimed engine sampler process does not exit cleanly at once (the cloud + could end its input as finished instead of running it again).""" + env = {RECLAIM_MARKER_ENV: str(tmp_path / "marker")} + slept = [] + assert entry.hold_after_reclaim(env, hold_s=5.0, sleep=slept.append) is False and slept == [] + (tmp_path / "marker").write_text("{}") + assert entry.hold_after_reclaim(env, hold_s=5.0, sleep=slept.append) is True and slept == [5.0] diff --git a/tests/test_stage1_gaps.py b/tests/test_stage1_gaps.py index 9ff34c17..0e4a3389 100644 --- a/tests/test_stage1_gaps.py +++ b/tests/test_stage1_gaps.py @@ -66,9 +66,10 @@ def test_engine_sampler_entry_runs_the_island_from_a_spec(tmp_path): reader = FileSampleBuffer(CheckpointStore(tmp_path / "store"), batch_size=1) reader.refresh() assert len(reader.select("t", 10)) == 4 # 2 rounds x 2 prompts - events = [json.loads(line) for line in (tmp_path / "tape" / "events.jsonl").read_text().splitlines()] + [tape] = (tmp_path / "tape").glob("events.*.jsonl") # one tape per incarnation + events = [json.loads(line) for line in tape.read_text().splitlines()] names = [e["event"] for e in events] - assert names[0] == "rl_sampler_start" and names.count("rl_sampler_engine_round") == 2 + assert names[0] == "rl_sampler_start" and "rl_sampler_join" in names and names.count("rl_sampler_engine_round") == 2 assert all(e["sampler_id"] == "ppu4" for e in events) diff --git a/tests/test_tools_eval.py b/tests/test_tools_eval.py new file mode 100644 index 00000000..a0a56d2a --- /dev/null +++ b/tests/test_tools_eval.py @@ -0,0 +1,135 @@ +"""CPU tests of the capability evaluation suite (tools/eval): matrix, verdict +status, evidence loading, a judge on synthetic data, Miles argument building, +and the report.""" + +from __future__ import annotations + +import json +from pathlib import Path + +import pytest + +from tools.eval import __main__ as cli +from tools.eval import cells, miles, report +from tools.eval.verdict import Verdict + + +def test_matrix_covers_every_capability_and_backend(): + assert {(c.capability, c.backend) for c in cells.CELLS.values()} == { + (cap, b) for cap in cells.CAPABILITIES for b in cells.BACKENDS} + assert [c.key for c in cells.CELLS.values() if not c.supported] == [] + with pytest.raises(SystemExit): + cells.get("E9", "verl") + + +@pytest.mark.parametrize("supported,passed,status", [ + (True, True, "pass"), (True, False, "fail"), (False, False, "expected_fail"), (False, True, "unexpected_pass")]) +def test_verdict_status(supported, passed, status): + v = Verdict("E1", "verl", "r", supported) + v.check("a", "a", True) + v.check("b", "b", passed) + assert v.status == status + assert Verdict("E1", "verl", "r", True).status == "error" # no criteria + + +def _island_run(tmp: Path, rounds: int, loss: float = 0.1) -> Path: + run = tmp / "run" + comp = run / "raw" / "run" / "island-c" + (comp / "yeto-output").mkdir(parents=True) + cell = cells.get("E1", "verl") + (run / "cell.json").write_text(json.dumps({"capability": "E1", "backend": "verl", "run_id": "run", + "config": cell.config, "code_sha": "x"})) + (comp / "meta.json").write_text(json.dumps({"rc": 0, "gpu": ["NVIDIA A10"], "seconds": 100.0})) + rows = [] + for r in range(rounds): + rows.append({"event": "rl_round_trained", "rollout_id": r, "time_unix": 10.0 + r}) + rows.append({"event": "rl_local_round", "loss": loss, "grad_norm": 0.2, "reward_mean": 0.5, + "base_policy_version": r, "time_unix": 10.5 + r}) + rows.append({"event": "rl_verl_island_done", "rounds_completed": rounds, "time_unix": 99.0}) + (comp / "yeto-output" / "rl-island-0.jsonl").write_text("\n".join(json.dumps(x) for x in rows) + "\n") + return run + + +def test_e1_judge_passes_and_reports(tmp_path): + run = _island_run(tmp_path, rounds=4) + verdict = cli.judge(run) + assert verdict.status == "pass", [c for c in verdict.criteria if not c.passed] + assert verdict.run["card_seconds"] == 100.0 + verdict.write(run) + md = report.markdown([run], title="t") + assert "| E1 island rollout + trainer | not run | PASS" in md + assert "" in report.html([run], title="t") + + +def test_e1_judge_fails_on_missing_round_and_nan(tmp_path): + run = _island_run(tmp_path, rounds=3, loss=float("nan")) + verdict = cli.judge(run) + failed = {c.id for c in verdict.criteria if not c.passed} + assert verdict.status == "fail" and {"rounds", "finite"} <= failed + + +def test_miles_launcher_args(): + run = "eval-x" + e2 = miles.launcher_args(cells.get("E2", "miles"), run) + assert "--rl-island-sample-role 0:sampler --rl-island-sample-role 1:sampler" in e2["ARGS"] + assert "--rl-island-sample-role 2:trainer" in e2["ARGS"] and "modal:1xa10g,modal:1xa10g,modal:1xa10g" in e2["ARGS"] + assert f"modal-volume://{miles.STORE_VOLUME}/{run}" in e2["ARGS"] + assert (e2["RECLAIM_LABEL"], e2["RECLAIM_ROUND"]) == ("l0", "1") + e4 = miles.launcher_args(cells.get("E4", "miles"), run) + assert "--rl-sample-role trainer" in e4["ARGS"] and "--rl-sample-role sampler" in e4["SARGS"] + assert f"--cluster-prefix {run}-s" in e4["SARGS"] and f"/{run} " in e4["SARGS"] # same store + assert "--rl-trainer-scheduler s1" in miles.launcher_args(cells.get("E5", "miles"), run)["ARGS"] + with pytest.raises(ValueError): + miles.launcher_args(cells.get("E3", "miles"), run) + + +def test_run_refuses_unregistered_run_id(tmp_path, monkeypatch): + ledger = tmp_path / "gpu-spend.md" + ledger.write_text("| other |\n") + monkeypatch.setattr(cli, "LEDGER", ledger) + with pytest.raises(SystemExit, match="not pre-registered"): + cli._preflight("eval-unregistered") + + +class _E4Evidence: + """Just the event reads judge_e4 makes (one Miles trainer, one sampler).""" + backend = "miles" + + def __init__(self, trainer_events, sampler_events): + self.events = {"trainer": trainer_events, "sampler-0": sampler_events} + + def trainers(self): + return ["trainer"] + + def samplers(self): + return ["sampler-0"] + + def of(self, component, name): + return [e for e in self.events.get(component, []) if e["event"] == name] + + +@pytest.mark.parametrize("stale,consumed,checks", [ + ((0, 16, 0, 0), 80, {"stale_event": True, "stale_cross_check": True, "over_age_drop": True}), + ((0, 15, 0, 0), 80, {"stale_event": True, "stale_cross_check": False, "over_age_drop": True}), + ((), 80, {"stale_event": False, "stale_cross_check": False, "over_age_drop": False}), +]) +def test_e4_miles_reads_the_stale_event_and_cross_checks_consumed_minus_trained(monkeypatch, stale, consumed, + checks): + from tools.eval import judges + + for name in ("_rc_ok", "_gpu"): + monkeypatch.setattr(judges, name, lambda *a, **kw: None) + monkeypatch.setattr(judges, "_round_metrics", lambda ev, t: {}) + rounds = 4 + trainer = [{"event": "rl_round_trained", "rollout_id": r, "mismatch": {"tis": 1.0, "tis_abs": 0.01}} + for r in range(rounds)] + trainer += [{"event": "rl_local_round", "base_policy_version": r} for r in range(rounds)] + trainer += [{"event": "rl_round_labels", "rollout_id": r, "rl/groups": 16} for r in range(rounds)] + trainer += [{"event": "rl_buffer_stale_skipped", "rollout_id": r, "stale_groups": n, "max_policy_age": 1} + for r, n in enumerate(stale)] + trainer += [{"event": "rl_policy_released", "policy_version": rounds, "buffer_selected": consumed}] + monkeypatch.setattr(judges, "_trained_rounds", lambda ev, t: list(range(rounds))) + sampler = [{"event": "rl_sampler_policy_loaded", "policy_version": v} for v in (0, 1)] + v = Verdict("E4", "miles", "r", True) + judges.judge_e4(_E4Evidence(trainer, sampler), v, {"rounds": rounds, "max_policy_age": 1}) + assert {c.id: c.passed for c in v.criteria if c.id in checks} == checks diff --git a/tests/test_verl_agentic_main.py b/tests/test_verl_agentic_main.py new file mode 100644 index 00000000..e988ac17 --- /dev/null +++ b/tests/test_verl_agentic_main.py @@ -0,0 +1,31 @@ +"""S25 N10.9: the verl agentic entry's Ray setup hook imports nothing CUDA-related (CPU).""" + +from __future__ import annotations + +import subprocess +import sys +import textwrap + + +def test_setup_hook_defers_the_estimator_until_core_algos_loads(tmp_path): + # Stand-ins for verl's core_algos and yeto's estimator module, in a fresh interpreter. + (tmp_path / "fakeverl").mkdir() + (tmp_path / "fakeverl" / "__init__.py").write_text("") + (tmp_path / "fakeverl" / "core_algos.py").write_text("REG = []\n") + (tmp_path / "fakeest.py").write_text("import fakeverl.core_algos as c\nc.REG.append('yeto_grpo')\n") + code = textwrap.dedent(""" + import sys + from yeto.rl.adapters.verl import agentic_main as m + m.ESTIMATOR_HOST, m.ESTIMATOR_MODULE = "fakeverl.core_algos", "fakeest" + m.setup_worker() + assert "fakeverl.core_algos" not in sys.modules and "torch" not in sys.modules + assert not any(k == "verl" or k.startswith("verl.") for k in sys.modules) + import fakeverl.core_algos as c + assert c.REG == ["yeto_grpo"], c.REG + m.setup_worker() # idempotent once loaded + assert c.REG == ["yeto_grpo"] + print("ok") + """) + out = subprocess.run([sys.executable, "-c", code], cwd=tmp_path, capture_output=True, text=True, + env={"PYTHONPATH": f"{tmp_path}:{__import__('os').getcwd()}"}) + assert out.returncode == 0 and out.stdout.strip() == "ok", out.stderr diff --git a/tests/test_verl_records_recorder.py b/tests/test_verl_records_recorder.py index 0d45d96a..a8add5c5 100644 --- a/tests/test_verl_records_recorder.py +++ b/tests/test_verl_records_recorder.py @@ -102,6 +102,7 @@ async def go(): assert sum(out.response_mask) > 0 and out.response_mask[-1] == 1 assert all(lp == 0.0 for lp, m in zip(out.response_logprobs, out.response_mask) if m == 0) assert out.extra_fields["provenance"]["engine"] == "fake" + assert out.extra_fields["length_capped_turns"] == 0 assert proxy.backend.sessions == {} and proxy.gateway.registries == {} # forgotten after close diff --git a/tests/test_verl_rollout_engine.py b/tests/test_verl_rollout_engine.py new file mode 100644 index 00000000..45c64911 --- /dev/null +++ b/tests/test_verl_rollout_engine.py @@ -0,0 +1,241 @@ +"""S25 N10.9: the records proxy over verl's own rollout client, behind a path secret (CPU). + +Real pieces: ``HarborAgentLoop.run``, ``ProxyTokenRecorder`` serving a +``RecordsProxy`` over HTTP, the gateway and recording backend, and +``VerlRolloutEngine``. Fakes: verl's agent-loop module (stub), verl's +``LLMServerClient`` (an async ``generate`` that replies with the scripted +model's tokens), the agent (Codex-like Responses loop). Tokenizer: Qwen3-0.6B +from the local HF cache (skip if absent). +""" + +from __future__ import annotations + +import asyncio +import copy +from types import SimpleNamespace +from typing import Any + +import pytest +import test_agentic_harbor as agentic +from test_agentic_harbor import _make_loop +from test_records_proxy import SCRIPT, SIGNED, ScriptedModel, _template +from test_verl_records_recorder import ROW, VERL_SAMPLING, CodexLikeAgent + +from yeto.rl.agentic import OUTCOME_KEY, HarnessConfigError +from yeto.rl.engine.sampler_engine import FINISH_LENGTH, FINISH_STOP, EngineIdentity, SamplerEngineError +from yeto.rl.harness.gateway.proxy import RecordsProxy +from yeto.rl.harness.gateway.recorder import ProxyTokenRecorder + +pytest.importorskip("aiohttp") + +verl_loop = agentic.verl_loop # the verl agent-loop stub fixture +SECRET = "s" * 32 +IDENTITY = EngineIdentity(engine="verl-rollout", engine_version="test", device_family="cuda", + card_model="H100", precision="bf16", quant_scheme="bf16", topology="unknown") + + +@pytest.fixture(scope="module") +def qwen(): + return _template("qwen3") + + +class FakeVerlClient: + """verl ``LLMServerClient`` stand-in: replies with the scripted model's tokens.""" + + def __init__(self, model: ScriptedModel, *, aborted: bool = False) -> None: + self.model = model + self.aborted = aborted + self.calls: list[dict[str, Any]] = [] + self.loops: set[int] = set() + + async def generate(self, request_id: str, *, prompt_ids: list[int], sampling_params: dict[str, Any]) -> Any: + self.loops.add(id(asyncio.get_running_loop())) + self.calls.append({"prompt_ids": list(prompt_ids), "sampling_params": dict(sampling_params)}) + if self.aborted: + return SimpleNamespace(token_ids=[], log_probs=None, stop_reason="aborted") + ids = self.model(tuple(prompt_ids), 0) + return SimpleNamespace(token_ids=ids, log_probs=[-0.25] * len(ids), stop_reason="completed", + extra_fields={"global_steps": 3}) + + +def _verl_proxy(template) -> tuple[RecordsProxy, Any]: + from yeto.rl.adapters.verl.rollout_engine import VerlRolloutEngine + + engine = VerlRolloutEngine(max_model_len=8192, identity=IDENTITY) + return RecordsProxy(engine, template, model="policy", signed_sampling=SIGNED, max_chains=1), engine + + +def test_agent_loop_samples_through_verls_client_behind_the_path_secret(verl_loop, qwen): + loop_mod, _ = verl_loop + proxy, engine = _verl_proxy(qwen) + recorder = ProxyTokenRecorder(proxy, path_secret=SECRET) + agent = CodexLikeAgent() + loop = _make_loop(loop_mod, recorder, agent, prompt_length=4096, response_length=4096) + client = FakeVerlClient(ScriptedModel(qwen, copy.deepcopy(SCRIPT))) + loop.server_manager = client # verl's AgentLoopBase sets it; the stub does not + + async def go(): + try: + out = await loop.run(VERL_SAMPLING, **ROW) + import aiohttp + + base = recorder._base_url.split("/k/")[0] + async with aiohttp.ClientSession() as http: # the root is not served, only the secret prefix + async with http.get(f"{base}/t/traj-1/records") as reply: + root_status = reply.status + finally: + await recorder.aclose() + return out, root_status + + out, root_status = asyncio.run(go()) + assert out.extra_fields[OUTCOME_KEY] == "ok" and out.reward_score == 1.0 and out.num_turns == 6 + (endpoint,) = agent.endpoints + assert f"/k/{SECRET}/t/traj-1-" in endpoint and root_status == 404 + assert engine.bound and len(client.calls) == 3 and len(client.loops) == 1 + params = client.calls[0]["sampling_params"] + assert params["logprobs"] is True and params["temperature"] == 1.0 and params["top_k"] == -1 + assert params["stop_token_ids"] == sorted(qwen.stop_token_ids) and params["max_tokens"] == 256 + # trained == sampled: the last prompt verl got is a prefix of prompt + response + last_prompt = client.calls[-1]["prompt_ids"] + assert out.prompt_ids + out.response_ids[:len(last_prompt) - len(out.prompt_ids)] == last_prompt + trained = [lp for lp, m in zip(out.response_logprobs, out.response_mask) if m == 1] + assert trained and all(lp == -0.25 for lp in trained) + assert out.extra_fields["provenance"]["engine"] == "verl-rollout" + assert out.extra_fields["length_capped_turns"] == 0 + # verl V1 needs the weights' step range per sample + assert out.extra_fields["min_global_steps"] == out.extra_fields["max_global_steps"] == 3 + assert engine.pop_steps(endpoint.split("/t/")[1].split("/")[0]) == [] # taken by the loop + + +def test_an_aborted_verl_request_aborts_the_sample(verl_loop, qwen): + loop_mod, _ = verl_loop + proxy, _ = _verl_proxy(qwen) + recorder = ProxyTokenRecorder(proxy) + agent = CodexLikeAgent() + loop = _make_loop(loop_mod, recorder, agent, prompt_length=4096, response_length=4096) + loop.server_manager = FakeVerlClient(ScriptedModel(qwen, copy.deepcopy(SCRIPT)), aborted=True) + + async def go(): + try: + return await loop.run(VERL_SAMPLING, **ROW) + finally: + await recorder.aclose() + + out = asyncio.run(go()) + assert out.extra_fields[OUTCOME_KEY] == "aborted" and sum(out.response_mask) == 0 + assert out.extra_fields["min_global_steps"] == out.extra_fields["max_global_steps"] == 0 # no reply seen + + +def test_unbound_engine_and_short_secret_fail_closed(qwen): + from yeto.rl.adapters.verl.rollout_engine import VerlRolloutEngine, finish_reason + from yeto.rl.engine.sampler_engine import SamplingParams + + engine = VerlRolloutEngine(max_model_len=64, identity=IDENTITY) + with pytest.raises(SamplerEngineError, match="not bound"): + engine.generate_tokens([[1, 2]], SamplingParams(max_tokens=4)) + with pytest.raises(HarnessConfigError, match="at least"): + ProxyTokenRecorder(_verl_proxy(qwen)[0], path_secret="short") + assert finish_reason((5, 7), {7}, 2) == FINISH_STOP + assert finish_reason((5, 6), {7}, 2) == FINISH_LENGTH + assert finish_reason((5,), {7}, 2) == FINISH_STOP + + +class ExitOnErrorAgent(CodexLikeAgent): + """Codex on a 400: the turn fails and the agent process exits non-zero (Harbor's verdict).""" + + def __init__(self, exception_type: str = "NonZeroAgentExitCodeError") -> None: + super().__init__() + self.exception_type = exception_type + self.statuses: list[int] = [] + + async def run(self, task: Any, endpoint: str) -> Any: + import aiohttp + + from test_records_proxy import echo, responses_body, tool_output, user_item + + from yeto.rl.agentic import TrialVerdict + + self.endpoints.append(endpoint) + items = [user_item("read a.txt")] + async with aiohttp.ClientSession() as http: + for turn in range(3): + async with http.post(f"{endpoint}/responses", json=responses_body(copy.deepcopy(items))) as reply: + self.statuses.append(reply.status) + if reply.status != 200: + return TrialVerdict(rewards=None, exception_type=self.exception_type) + response = await reply.json() + items += echo(response["output"]) + items += [tool_output(o["call_id"], f"result {turn}") + for o in response["output"] if o["type"] == "function_call"] + return TrialVerdict(rewards={"reward": 1.0}) + + +class MalformedSecondReply(ScriptedModel): + """The second reply is a tool call that is not JSON (as Qwen3-4B wrote with heredocs).""" + + def __call__(self, prompt: tuple[int, ...], index: int) -> list[int]: + if len(self.replies) == len(SCRIPT) - 1: + self.replies.pop(0) + return self.template.encode('\n{"name": "shell", "arguments": {"cmd": "cat > a\nb"}\n') + \ + [self.template.encode("<|im_end|>")[0]] + return super().__call__(prompt, index) + + +@pytest.mark.parametrize("exception_type,outcome", [("NonZeroAgentExitCodeError", "policy_limit"), + ("EnvironmentStartTimeoutError", "aborted")]) +def test_malformed_tool_call_is_a_trained_zero_not_an_abort(verl_loop, qwen, exception_type, outcome): + loop_mod, _ = verl_loop + proxy, _ = _verl_proxy(qwen) + recorder = ProxyTokenRecorder(proxy) + agent = ExitOnErrorAgent(exception_type) + loop = _make_loop(loop_mod, recorder, agent, prompt_length=4096, response_length=4096) + loop.server_manager = FakeVerlClient(MalformedSecondReply(qwen, copy.deepcopy(SCRIPT))) + + async def go(): + try: + return await loop.run(VERL_SAMPLING, **ROW) + finally: + await recorder.aclose() + + out = asyncio.run(go()) + assert agent.statuses == [200, 400] + assert out.extra_fields[OUTCOME_KEY] == outcome and out.reward_score == 0.0 + if outcome == "policy_limit": + # the malformed reply was sampled and is trained (mask 1) as part of a 0-reward sample + assert out.extra_fields["outcome_reason"].startswith("policy_output_error") + assert sum(out.response_mask) > 0 + else: + assert sum(out.response_mask) == 0 # another exit cause stays an infrastructure abort + + +def test_context_full_is_a_trained_zero_not_an_abort(verl_loop, qwen): + """The conversation outgrows the context: the proxy answers 409, Codex exits; the policy's budget (N10.9).""" + from yeto.rl.adapters.verl.rollout_engine import VerlRolloutEngine + + loop_mod, _ = verl_loop + + def run(max_model_len: int, client: FakeVerlClient): + engine = VerlRolloutEngine(max_model_len=max_model_len, identity=IDENTITY) + recorder = ProxyTokenRecorder(RecordsProxy(engine, qwen, model="policy", signed_sampling=SIGNED, + max_chains=1)) + agent = ExitOnErrorAgent() + loop = _make_loop(loop_mod, recorder, agent, prompt_length=4096, response_length=4096) + loop.server_manager = client + + async def go(): + try: + return await loop.run(VERL_SAMPLING, **ROW) + finally: + await recorder.aclose() + + return asyncio.run(go()), agent + + probe = FakeVerlClient(ScriptedModel(qwen, copy.deepcopy(SCRIPT))) + run(8192, probe) + second = probe.calls[1]["prompt_ids"] + # room for the second turn, not for the third prompt (which adds the tool output) + limit = len(second) + 40 + out, agent = run(limit, FakeVerlClient(ScriptedModel(qwen, copy.deepcopy(SCRIPT)))) + assert 409 in agent.statuses + assert out.extra_fields[OUTCOME_KEY] == "policy_limit" and out.reward_score == 0.0 + assert "context_full" in out.extra_fields["outcome_reason"] and sum(out.response_mask) > 0 diff --git a/tests/test_word_diversity_reward.py b/tests/test_word_diversity_reward.py new file mode 100644 index 00000000..a0c1de2e --- /dev/null +++ b/tests/test_word_diversity_reward.py @@ -0,0 +1,17 @@ +import asyncio +import types + +from yeto.rl.word_diversity_reward import diversity, score + + +def test_diversity_values(): + assert diversity("") == 0.0 + assert diversity("unload unload unload unload") == 0.25 + assert diversity("a b c d") == 1.0 + + +def test_score_has_variance_on_truncated_degenerate_samples(): + texts = ["Ref unload unload unload", "adam unload unloadimers unloadimers"] + values = [asyncio.run(score(None, types.SimpleNamespace(response=t))) for t in texts] + assert values[0] != values[1] + assert asyncio.run(score(None, types.SimpleNamespace(response=None))) == 0.0 diff --git a/tools/eval/__init__.py b/tools/eval/__init__.py new file mode 100644 index 00000000..abfbd6fe --- /dev/null +++ b/tools/eval/__init__.py @@ -0,0 +1,22 @@ +"""yeto RL capability evaluation suite (S25 T-EVAL). + +One repeatable GPU evaluation per capability and backend, one CLI, one verdict +format and one report: + + python -m tools.eval list + python -m tools.eval run --backend miles|verl --run-id [--dry-run] + python -m tools.eval judge + python -m tools.eval report ... --out + +* :mod:`.cells` -- the capability x backend matrix (what each cell runs, its + configuration, its judge, and whether the backend supports it); +* :mod:`.verdict` -- the verdict written by every judge (``verdict.json``); +* :mod:`.judges` -- the pre-registered criteria of each capability; +* :mod:`.modal_verl` -- verl cells: a direct Modal harness (syncer, engine sampler + island, verl islands), derived from the S25 C2/C3V/C4 harnesses; +* ``miles_launch.sh`` -- Miles cells: the yeto launcher on Modal, derived from + the S25 C1/C3 and S24 L5.4 run scripts; +* :mod:`.report` -- Markdown and HTML report over a set of verdicts. + +Every run writes its raw data under ``//raw/`` before judging. +""" diff --git a/tools/eval/__main__.py b/tools/eval/__main__.py new file mode 100644 index 00000000..1d047988 --- /dev/null +++ b/tools/eval/__main__.py @@ -0,0 +1,151 @@ +"""CLI of the evaluation suite (see the package doc). + + python -m tools.eval list + python -m tools.eval run E3 --backend verl --run-id eval-e3-verl-a [--out DIR] [--dry-run] + python -m tools.eval judge DIR/eval-e3-verl-a + python -m tools.eval report DIR/* --out report.md [--html report.html] + +``run`` refuses a run id that is not pre-registered in the GPU ledger (the id +must appear in ``gpu-spend.md``), checks the local thread count and Ray, writes +``cell.json``, runs the cell's harness (raw data into ``raw/``), then judges. +""" + +from __future__ import annotations + +import argparse +import json +import subprocess +import sys +import time +from pathlib import Path + +from . import cells, judges, report +from .verdict import Verdict + +DEFAULT_OUT = Path("/home/michael/work/infra-drafts/SESSION25/logs/T-EVAL") +LEDGER = Path("/home/michael/work/infra-drafts/gpu-spend.md") +THREAD_LIMIT = 3000 + + +def _preflight(run_id: str) -> None: + if not LEDGER.is_file() or run_id not in LEDGER.read_text(): + raise SystemExit(f"run id {run_id} is not pre-registered in {LEDGER}") + threads = subprocess.run("ps -u michael -L --no-headers | wc -l", shell=True, capture_output=True, + text=True).stdout.strip() + if int(threads or 0) >= THREAD_LIMIT: + raise SystemExit(f"{threads} local threads (>= {THREAD_LIMIT}): do not launch") + if subprocess.run(["pgrep", "-x", "raylet"], capture_output=True).returncode == 0: + raise SystemExit("a local Ray is running: do not launch") + + +def cmd_list(_args) -> int: + for key, cell in sorted(cells.CELLS.items()): + flag = "" if cell.supported else " [unsupported: expected to fail]" + print(f"{key:10s} {cell.harness:12s} {cell.cards} card(s) {cell.config.get('gpu')} " + f"{cells.CAPABILITIES[cell.capability]}{flag}") + return 0 + + +def cmd_run(args) -> int: + cell = cells.get(args.capability, args.backend) + run_dir = Path(args.out) / args.run_id + if not args.dry_run: + if (run_dir / "cell.json").exists(): + raise SystemExit(f"{run_dir} already used") + _preflight(args.run_id) + run_dir.mkdir(parents=True, exist_ok=True) + sha = subprocess.run(["git", "-C", str(Path(__file__).resolve().parents[2]), "rev-parse", "HEAD"], + capture_output=True, text=True).stdout.strip() + (run_dir / "cell.json").write_text(json.dumps({ + "capability": cell.capability, "backend": cell.backend, "run_id": args.run_id, "harness": cell.harness, + "cards": cell.cards, "supported": cell.supported, "config": cell.config, "code_sha": sha}, indent=1)) + began = time.time() + if cell.harness == "modal_verl": + if args.dry_run: + print(json.dumps(cell.config, indent=1)) + return 0 + from . import modal_verl + + summary = modal_verl.run(cell, args.run_id, run_dir) + else: + from . import miles + + summary = miles.run(cell, args.run_id, run_dir, dry_run=args.dry_run) + if args.dry_run: + print((run_dir / "dry-run.log").read_text()) + return summary["rc"] + summary["wall_s"] = round(time.time() - began, 1) + (run_dir / "run.json").write_text(json.dumps(summary, indent=1, default=str)) + return cmd_judge(argparse.Namespace(run_dir=str(run_dir))) + + +def judge(run_dir: Path) -> Verdict: + from . import evidence, miles, modal_verl + + cell_doc = json.loads((run_dir / "cell.json").read_text()) + cell = cells.get(cell_doc["capability"], cell_doc["backend"]) + cfg = cell_doc["config"] + verdict = Verdict(capability=cell.capability, backend=cell.backend, run_id=cell_doc["run_id"], + supported=cell.supported) + run_doc = json.loads((run_dir / "run.json").read_text()) if (run_dir / "run.json").is_file() else {} + raw = run_dir / "raw" + seconds = (modal_verl.card_seconds(raw) if cell.backend == "verl" else miles.card_seconds(raw)) + rate = cells.A10G_USD_PER_HOUR if cfg.get("gpu") == "A10G" else 5.3 + verdict.run.update(code_sha=cell_doc.get("code_sha"), wall_s=run_doc.get("wall_s"), + card_seconds=round(seconds, 1), usd_estimate=round(seconds / 3600 * rate, 3), + usd_rate=f"{rate}/container-hour incl. CPU+memory (list-price estimate, not a bill)") + try: + ev = evidence.load(run_dir) + judges.JUDGES[cell.capability](ev, verdict, cfg) + except Exception as exc: # noqa: BLE001 - a judge crash is an error verdict, with the reason + verdict.error = f"{type(exc).__name__}: {exc}" + return verdict + + +def cmd_judge(args) -> int: + run_dir = Path(args.run_dir) + verdict = judge(run_dir) + verdict.write(run_dir) + doc = verdict.to_dict() + print(json.dumps({"run": doc["run_id"], "cell": f"{doc['capability']}-{doc['backend']}", + "status": doc["status"], "failed": [c["id"] for c in doc["criteria"] if not c["passed"]], + "error": doc["error"]})) + return 0 if doc["status"] in ("pass", "expected_fail") else 1 + + +def cmd_report(args) -> int: + dirs = [Path(d) for d in args.run_dirs if (Path(d) / "verdict.json").is_file()] + md = report.markdown(dirs, title=args.title) + Path(args.out).write_text(md) + if args.html: + Path(args.html).write_text(report.html(dirs, title=args.title)) + print(f"{len(dirs)} verdict(s) -> {args.out}") + return 0 + + +def main(argv: list[str] | None = None) -> int: + p = argparse.ArgumentParser(prog="python -m tools.eval") + sub = p.add_subparsers(dest="cmd", required=True) + sub.add_parser("list").set_defaults(fn=cmd_list) + r = sub.add_parser("run") + r.add_argument("capability", choices=sorted(cells.CAPABILITIES)) + r.add_argument("--backend", choices=cells.BACKENDS, required=True) + r.add_argument("--run-id", required=True) + r.add_argument("--out", default=str(DEFAULT_OUT)) + r.add_argument("--dry-run", action="store_true") + r.set_defaults(fn=cmd_run) + j = sub.add_parser("judge") + j.add_argument("run_dir") + j.set_defaults(fn=cmd_judge) + rep = sub.add_parser("report") + rep.add_argument("run_dirs", nargs="+") + rep.add_argument("--out", required=True) + rep.add_argument("--html", default=None) + rep.add_argument("--title", default="yeto RL capability evaluation") + rep.set_defaults(fn=cmd_report) + args = p.parse_args(argv) + return args.fn(args) + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/tools/eval/cells.py b/tools/eval/cells.py new file mode 100644 index 00000000..ec6cb8fe --- /dev/null +++ b/tools/eval/cells.py @@ -0,0 +1,87 @@ +"""The capability x backend matrix of the evaluation suite. + +Model Qwen3-0.6B (pinned revision), GSM8K, LoRA, Modal A10G, the fewest cards +the capability needs. Each cell names its harness, its configuration (what the +judge reads too) and whether the backend supports the capability: an +unsupported cell still runs, and its failure is the expected baseline. +""" + +from __future__ import annotations + +from dataclasses import dataclass, field +from typing import Any + +CAPABILITIES = { + "E1": "island rollout + trainer", + "E2": "pure sampler island reclaimed and replaced", + "E3": "synced trainer islands -> sampler island", + "E4": "async rollout + off-policy correction + over-age drop", + "E5": "round cut: SCHEDULED stop + resume assertions", +} +BACKENDS = ("miles", "verl") +# all-in estimate per A10G container hour: GPU $1.10 + the CPU (8-16 cores) and memory +# (32-64 GiB) the harness requests, which Modal bills too (list prices; not a bill) +A10G_USD_PER_HOUR = 2.0 + + +@dataclass(frozen=True) +class Cell: + capability: str + backend: str + harness: str # "modal_verl" | "miles_local" | "miles_head" + cards: int + config: dict[str, Any] = field(default_factory=dict) + supported: bool = True + note: str = "" + + @property + def key(self) -> str: + return f"{self.capability}-{self.backend}" + + +# ------------------------------------------------------------------ verl (direct Modal harness) +_VERL = {"gpu": "A10G", "rounds": 4, "groups_per_round": 4, "samples_per_group": 4, "max_tokens": 256, + "lr": 1e-4, "lora_r": 32, "seq_len": 768, "seed": 1, "sglang_mem": 0.6, "max_model_len": 1024, + "timeout_s": 2700, "learners": 1, "max_policy_age": 1, "sampler_rounds": 24, "max_backlog_groups": 8, + "reclaim_after_round": None} +# the plain island cells reuse the S25 C4 configuration (8 groups x 4, lr 1e-5) +_VERL_ISLAND = {**_VERL, "kind": "island", "groups_per_round": 8, "lr": 1e-5, "seq_len": 1024, "seed": 17, + "stop_after_rounds": 2} + +# ------------------------------------------------------------------ Miles (yeto launcher on Modal) +_MILES = {"gpu": "A10G", "hard_s": 3000, "rollout_batch_size": 16, "n_samples": 4, "max_response_len": 512, + "seq_len": 1024, "lora_r": 16, "lr": 1e-5, "seed": 17} + +CELLS: dict[str, Cell] = {c.key: c for c in ( + Cell("E1", "verl", "modal_verl", 1, {**_VERL_ISLAND, "phases": ["c"]}), + Cell("E1", "miles", "miles_local", 1, {**_MILES, "rounds": 4, "islands": 1}), + # E2: the S25 C1 shape (two samplers, one reclaimed after it wrote round 1) on Miles; verl's + # sampler island is the neutral engine sampler, which resumes after a reclaim since S25 G1C + Cell("E2", "verl", "modal_verl", 2, + {**_VERL, "kind": "buffer", "max_policy_age": 2, "max_backlog_groups": 12, "reclaim_after_round": 1}), + Cell("E2", "miles", "miles_local", 3, + {**_MILES, "rounds": 8, "islands": 3, "samplers": 2, "max_policy_age": 4, "reclaim_label": "l0", + "reclaim_round": 1, "reclaim_component": "l0"}), + Cell("E3", "verl", "modal_verl", 3, + {**_VERL, "kind": "buffer", "learners": 2, "max_policy_age": 2, "max_backlog_groups": 16, + "sampler_rounds": 20}), + Cell("E3", "miles", "miles_head", 3, {**_MILES, "rounds": 4, "learners": 2, "hard_s": 2700}), + # E4: a sampler island far ahead of the trainer (backlog 3 rounds, short responses) and + # max_policy_age 1: round 2 meets v0 groups, which must be dropped + Cell("E4", "verl", "modal_verl", 2, + {**_VERL, "kind": "buffer", "max_policy_age": 1, "max_backlog_groups": 12, "max_tokens": 128}), + Cell("E4", "miles", "miles_local", 2, + {**_MILES, "rounds": 4, "islands": 2, "max_policy_age": 1, "sampler_rounds": 10, "backlog_rounds": 3, + "max_response_len": 256}), + Cell("E5", "verl", "modal_verl", 1, + {**_VERL_ISLAND, "phases": ["a", "b", "c"], "compare_prompt_sets": True}), + Cell("E5", "miles", "miles_local", 2, + {**_MILES, "rounds": 6, "islands": 2, "max_policy_age": 2, "min_round_s": 120}), +)} + + +def get(capability: str, backend: str) -> Cell: + key = f"{capability.upper()}-{backend}" + if key not in CELLS: + raise SystemExit(f"no cell {key}; known: {', '.join(sorted(CELLS))}") + return CELLS[key] diff --git a/tools/eval/evidence.py b/tools/eval/evidence.py new file mode 100644 index 00000000..be8fda9d --- /dev/null +++ b/tools/eval/evidence.py @@ -0,0 +1,197 @@ +"""Raw run data, read the same way for both backends. + +A run directory holds ``cell.json`` (the cell and run id), ``driver.json`` and +``raw/`` (everything pulled). :class:`Evidence` turns the raw data into: + +* ``islands``: event lists keyed by component (verl: ``trainer0``, ``sampler-0``, + ``island-a``...; Miles: the launcher's island labels ``l0``, ``l1``, and + ``s:l0`` for a separate sampler launch), each sorted by time and de-duplicated + across sources (launch log echo, per-container tapes, the tape Volume); +* ``index``: the sample buffer index (shared-directory JSONL or object files); +* ``lanes``: per-lane read-position logs of synced trainers; +* return codes, GPU names and card seconds. +""" + +from __future__ import annotations + +import csv +import json +import re +from collections import defaultdict +from collections.abc import Iterable, Mapping +from dataclasses import dataclass, field +from pathlib import Path +from typing import Any + +from .verdict import events, find, find_all, jsonl, objects + +TRAINER_EVENTS = ("rl_round_trained", "rl_verl_island_done") +SAMPLER_EVENTS = ("rl_sampler_round", "rl_sampler_start") + + +def _merge(target: list[dict], rows: Iterable[dict], seen: set) -> None: + for row in rows: + key = (row.get("event"), row.get("time_unix"), row.get("seq"), row.get("rollout_id")) + if row.get("event") and key not in seen: + seen.add(key) + target.append(row) + + +@dataclass +class Evidence: + run_dir: Path + backend: str + islands: dict[str, list[dict[str, Any]]] = field(default_factory=dict) + index: list[dict[str, Any]] = field(default_factory=list) + lanes: dict[int, list[dict[str, Any]]] = field(default_factory=dict) + shards: dict[str, list[dict[str, Any]]] = field(default_factory=dict) # sampler shards of the index + rc: dict[str, Any] = field(default_factory=dict) + gpu_names: dict[str, int] = field(default_factory=dict) + store: Path | None = None + files: dict[str, Path] = field(default_factory=dict) + + # ------------------------------------------------------------- queries + def of(self, component: str, name: str) -> list[dict[str, Any]]: + return list(events(self.islands.get(component, []), name)) + + def all(self, name: str) -> list[dict[str, Any]]: + return [e for tape in self.islands.values() for e in events(tape, name)] + + def trainers(self) -> list[str]: + return sorted(k for k, tape in self.islands.items() if any(e.get("event") in TRAINER_EVENTS for e in tape)) + + def samplers(self) -> list[str]: + return sorted(k for k, tape in self.islands.items() if any(e.get("event") in SAMPLER_EVENTS for e in tape)) + + def entries(self) -> list[dict[str, Any]]: + """Index records that are groups (not end markers).""" + return [e for e in self.index if not e.get("end") and ("behavior_version" in e or "sample_id" in e)] + + +# ------------------------------------------------------------------ loaders +def load(run_dir: Path) -> Evidence: + run_dir = Path(run_dir) + cell = json.loads((run_dir / "cell.json").read_text()) + ev = Evidence(run_dir=run_dir, backend=cell["backend"]) + raw = run_dir / "raw" + (_load_verl if ev.backend == "verl" else _load_miles)(ev, raw, cell["run_id"]) + for tape in ev.islands.values(): + tape.sort(key=lambda e: float(e.get("time_unix") or e.get("t") or 0.0)) + if ev.store is not None: + buffer = ev.store / "sample-buffer" + if not buffer.is_dir(): # miles_head.sh pulls index.jsonl and lanes/ into the store dir itself + buffer = ev.store + ev.index = jsonl(buffer / "index.jsonl") or objects(buffer / "index") + for path in sorted((buffer / "shards").glob("*.jsonl")) if (buffer / "shards").is_dir() else []: + ev.shards[path.stem] = jsonl(path) + lanes = buffer / "lanes" + for path in sorted(lanes.glob("*-of-*")) if lanes.is_dir() else []: + m = re.match(r"(\d+)-of-\d+", path.name) + if m: + ev.lanes[int(m.group(1))] = jsonl(path) if path.is_file() else objects(path) + return ev + + +def _load_verl(ev: Evidence, raw: Path, run_id: str) -> None: + root = raw / run_id if (raw / run_id).is_dir() else raw + for comp in sorted(p for p in root.iterdir() if p.is_dir() and p.name != "store") if root.is_dir() else []: + rows: list[dict] = [] + seen: set = set() + for name in ("trainer-tape.jsonl", "sampler-tape.jsonl", "syncer-tape.jsonl"): + _merge(rows, jsonl(comp / name), seen) + for tape in find_all(comp / "yeto-output", "rl-island-*.jsonl"): + _merge(rows, jsonl(tape), seen) + if comp.name == "syncer": + rows = [dict(r, event=r.get("event") or r.get("kind")) for r in jsonl(comp / "syncer-tape.jsonl")] + ev.islands[comp.name] = rows + meta = comp / "meta.json" + if meta.is_file(): + data = json.loads(meta.read_text()) + # a sampler or a phase returns without "rc": its error (if any) is the outcome + ev.rc[comp.name] = data["rc"] if "rc" in data else (data["error"] if data.get("error") else 0) + for name in data.get("gpu") or []: + ev.gpu_names[name] = ev.gpu_names.get(name, 0) + 1 + ev.files[comp.name] = comp + ev.store = root / "store" + + +_EVENT_LINE = r"^(?:\S+ )?\[{prefix}-l(\d+)-modal\] YETO_RL_EVENT (\{{.*\}})\s*$" + + +def _launch_events(path: Path, prefix: str) -> dict[int, list[dict]]: + out: dict[int, list[dict]] = defaultdict(list) + if not path.is_file(): + return out + pattern = re.compile(_EVENT_LINE.format(prefix=re.escape(prefix))) + for line in path.read_text(errors="replace").splitlines(): + m = pattern.match(line) + if m: + try: + out[int(m.group(1))].append(json.loads(m.group(2))) + except ValueError: + pass + return out + + +def _load_miles(ev: Evidence, raw: Path, run_id: str) -> None: + seen: dict[str, set] = defaultdict(set) + for app, tag, log in ((run_id, "", "launch.ts.log"), (f"{run_id}-s", "s:", "slaunch.ts.log")): + for island, rows in _launch_events(raw / log, app).items(): + _merge(ev.islands.setdefault(f"{tag}l{island}", []), rows, seen[f"{tag}l{island}"]) + # the tape Volume holds every incarnation's tape (rl-island-0.inc.jsonl after a + # replacement container); the launch log only echoes the first container + for tape in find_all(raw / "tape-direct" / f"yeto-{app}", "rl-island-*.jsonl"): + label = next((part for part in tape.parts if re.fullmatch(r"l\d+", part)), None) + if label is None: + continue + key = f"{tag}{label}" + _merge(ev.islands.setdefault(key, []), jsonl(tape), seen[key]) + for name, key in (("rc", "launch"), ("src", "sampler_launch")): + path = raw / name + if path.is_file(): + ev.rc[key] = int(path.read_text().strip() or -1) + head_rc = raw / "rc.txt" + if head_rc.is_file(): + m = re.search(r"rc=(\d+)", head_rc.read_text()) + ev.rc["launch"] = int(m.group(1)) if m else -1 + nvml = raw / "nvml.csv" + if nvml.is_file(): + for row in csv.reader(nvml.read_text(errors="replace").splitlines()): + if len(row) > 2: + name = row[2].strip() + ev.gpu_names[name] = ev.gpu_names.get(name, 0) + 1 + store = raw / "store" + # `modal volume get

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

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

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

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

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

{esc(title)}

{head}{''.join(rows)}
Capability
" + f"{''.join(sections)}\n") diff --git a/tools/eval/verdict.py b/tools/eval/verdict.py new file mode 100644 index 00000000..2af2da03 --- /dev/null +++ b/tools/eval/verdict.py @@ -0,0 +1,109 @@ +"""The verdict every capability judge writes (``/verdict.json``). + +Status values: + +* ``pass`` -- every criterion passed; +* ``fail`` -- a criterion failed on a backend that supports the capability; +* ``expected_fail`` -- a criterion failed and the backend does not support the + capability (:attr:`cells.Cell.supported` is False): the run is kept as the + regression baseline of the gap; +* ``unexpected_pass`` -- an unsupported cell passed (update the matrix); +* ``error`` -- the run produced no judgeable data (harness or launch failure). +""" + +from __future__ import annotations + +import json +from collections.abc import Iterable, Mapping +from dataclasses import asdict, dataclass, field +from pathlib import Path +from typing import Any + +PASS, FAIL, EXPECTED_FAIL, UNEXPECTED_PASS, ERROR = "pass", "fail", "expected_fail", "unexpected_pass", "error" +VERDICT_FILE = "verdict.json" + + +@dataclass +class Criterion: + id: str + text: str + passed: bool + evidence: Any = None + + +@dataclass +class Verdict: + capability: str + backend: str + run_id: str + supported: bool + criteria: list[Criterion] = field(default_factory=list) + metrics: dict[str, Any] = field(default_factory=dict) + run: dict[str, Any] = field(default_factory=dict) # wall seconds, card seconds, cost estimate, code sha + error: str | None = None + + def check(self, cid: str, text: str, passed: bool, evidence: Any = None) -> bool: + self.criteria.append(Criterion(cid, text, bool(passed), evidence)) + return bool(passed) + + @property + def status(self) -> str: + if self.error is not None or not self.criteria: + return ERROR + ok = all(c.passed for c in self.criteria) + if self.supported: + return PASS if ok else FAIL + return UNEXPECTED_PASS if ok else EXPECTED_FAIL + + def to_dict(self) -> dict[str, Any]: + return {**asdict(self), "status": self.status} + + def write(self, run_dir: Path) -> Path: + path = Path(run_dir) / VERDICT_FILE + path.write_text(json.dumps(self.to_dict(), indent=1, default=str) + "\n") + return path + + +def load(run_dir: Path) -> dict[str, Any]: + return json.loads((Path(run_dir) / VERDICT_FILE).read_text()) + + +# ---------------------------------------------------------------- raw data readers +def jsonl(path: Path) -> list[dict[str, Any]]: + """The JSON objects of a JSONL file (missing file: []; bad lines skipped).""" + path = Path(path) + if not path.is_file(): + return [] + out = [] + for line in path.read_text(errors="replace").splitlines(): + line = line.strip() + start = line.find("{") + if start < 0: + continue + try: + value = json.loads(line[start:]) + except ValueError: + continue + if isinstance(value, dict): + out.append(value) + return out + + +def find(root: Path, name: str) -> Path: + """The first file called ``name`` under ``root`` (sorted), else ``root/name``.""" + hits = sorted(Path(root).rglob(name)) if Path(root).is_dir() else [] + return hits[0] if hits else Path(root) / name + + +def find_all(root: Path, pattern: str) -> list[Path]: + return sorted(Path(root).rglob(pattern)) if Path(root).is_dir() else [] + + +def events(tape: Iterable[Mapping[str, Any]], name: str) -> list[Mapping[str, Any]]: + return [e for e in tape if e.get("event") == name] + + +def objects(directory: Path) -> list[dict[str, Any]]: + """One JSON object per ``*.json`` file of a directory (object-store layouts).""" + directory = Path(directory) + return [json.loads(p.read_text()) for p in sorted(directory.glob("*.json"))] if directory.is_dir() else [] diff --git a/tools/probes/c4_verl/modal_c4.py b/tools/probes/c4_verl/modal_c4.py new file mode 100644 index 00000000..5b0bbfd2 --- /dev/null +++ b/tools/probes/c4_verl/modal_c4.py @@ -0,0 +1,146 @@ +"""T-C4 (S25 N7.6b) GPU check: the verl island's round cuts on a real verl island. + +Three island runs of ``yeto.rl.adapters.verl.island_entry`` (Qwen3-0.6B, LoRA, +GSM8K, ``--rl-single-island-no-sync``) in the pinned verl image, one GPU each: + +* ``a`` -- commanded by a pre-written trainer-control file (``run 2``, then + ``stop STOPPED``): trains rounds 0-1, the stop cuts with the SCHEDULED trigger + into the store (Modal Volume), the island exits; +* ``b`` -- a fresh container: restores the store into a new state dir, resumes + at rollout 2 (data cursor seek, digest + Adam/scheduler checks), trains 2-3; +* ``c`` -- the continuous control: rounds 0-3, no store (runs beside ``a``). + +usage (after the gpu-spend.md pre-registration): + C4_RUN= [C4_GPU=A10G] modal run --detach modal_c4.py +The entrypoint only spawns the chain (a killed detached client would cancel a +blocking ``.remote()``). Evidence: Volume ``t-c4-s25`` under ``//``. +""" + +from __future__ import annotations + +import json +import os +import shlex +import sys + +import modal + +RUN = os.environ.get("C4_RUN", "unset") +YETO = os.environ.get("C4_YETO", "/home/michael/work/s25-c4-cut-ctl") +GPU = os.environ.get("C4_GPU", "A10G") +# Modal A10G reports "NVIDIA A10" +EXPECT = {"A10G": ("NVIDIA A10", "nvidia-a10g"), "L4": ("L4", "nvidia-l4"), "H100!": ("H100", "nvidia-h100")} +sys.path.insert(0, YETO) +from yeto.rl.adapters.verl.image import modal_image # noqa: E402 + +vol = modal.Volume.from_name("t-c4-s25", create_if_missing=True) +image = (modal_image(modal) + .env({"PYTHONPATH": "/yeto", "HF_HUB_DISABLE_TELEMETRY": "1", "PYTHONUNBUFFERED": "1"}) + .add_local_dir(f"{YETO}/yeto", "/yeto/yeto", copy=False, ignore=["**/__pycache__"])) +app = modal.App(f"t-c4-{RUN}", image=image) + +ROUNDS = 4 +PY = "/workspace/verl/.venv/bin/python" # the pinned verl venv (image.VENV_BIN) +ISLAND_ARGS = [ + "--model", "Qwen/Qwen3-0.6B", "--model-revision", "c1899de289a04d12100db370d81485cdf75e47ca", + "--data", "zhuzilin/gsm8k", "--data-revision", "0cbd9f31d91ac21a7613dcbc7fef992adac459ae", + "--rl-single-island-no-sync", "--learner-id", "0", "--reward-function", "gsm8k_reward:score", + "--global-rounds", str(ROUNDS), "--groups-per-round", "8", "--samples-per-group", "4", + "--rollout-max-response-len", "256", "--event-tape", "/root/yeto-output/rl-island-0.jsonl", + "--lora-r", "32", "--lora-targets", "all-linear", "--inner-lr", "1e-5", "--seq-len", "1024", + "--seed", "17", "--apply-chat-template-kwargs", '{"enable_thinking": false}', +] + + +def _sh(cmd: str, log) -> int: + import subprocess + + log.write(f"$ {cmd}\n") + log.flush() + return subprocess.run(cmd, shell=True, stdout=log, stderr=subprocess.STDOUT).returncode + + +# max_inputs=1: every phase gets a fresh container (a reused warm container kept the +# previous phase's ~/yeto-output and its stale Volume view: run c4-verl-20261020c b). +@app.function(gpu=GPU, cpu=8.0, memory=49152, timeout=40 * 60, volumes={"/vol": vol}, max_inputs=1) +def island(run: str, phase: str, gpu: str) -> dict: + import shutil + import subprocess + import time + + vol.reload() # see the store the previous phase committed + want_name, group = EXPECT[gpu] + out = f"/vol/{run}/{phase}" + os.makedirs(out, exist_ok=True) + done = f"{out}/result.json" + if os.path.isfile(done): # Modal restarts an input after a preemption: never rerun a finished phase + previous = json.load(open(done)) + if previous.get("rc") == 0: + return {**previous, "reused_after_restart": True} + log = open(f"{out}/container.log", "a") + name = subprocess.run(["nvidia-smi", "--query-gpu=name", "--format=csv,noheader"], + capture_output=True, text=True).stdout.strip() + _sh(f"nvidia-smi > {out}/nvidia-smi.txt", log) + if want_name not in name: + log.write(f"GPU ASSERT FAILED: {name!r} is not {want_name}\n") + log.close() + vol.commit() + return {"phase": phase, "rc": 3, "gpu": name} + env = dict(os.environ, YETO_RL_COMPAT_GROUP=group, RAY_ADDRESS="127.0.0.1:6379") + store = f"/vol/{run}/store" + args = list(ISLAND_ARGS) + if phase in ("a", "b"): + args += ["--rl-resume-store", store, "--rl-resume-state-dir", f"/root/state-{phase}"] + if phase == "a": # the head's commands for incarnation 0, written through yeto's own writer + rc = _sh(f"{PY} -c " + shlex.quote( + "import sys; from yeto.rl.engine import trainer_remote as t; " + "from yeto.rl.engine.resume import store_for; from yeto.rl.elastic_client import LeaveReason; " + "w = t.CommandWriter(t.LocalStoreFiles(store_for(sys.argv[1])), 0); " + "w.send(t.RUN, rounds=2); w.send(t.STOP, reason=LeaveReason.STOPPED); " + "print(t.island_env(0, None, poll_s=2.0))") + f" {store}", log) + if rc != 0: + log.close() + vol.commit() + return {"phase": phase, "rc": 4, "gpu": name, "error": "command file not written"} + env.update({"YETO_TRAINER_CONTROL": "1", "YETO_TRAINER_INCARNATION": "0", + "YETO_TRAINER_RESUME": "", "YETO_TRAINER_COMMAND_POLL_S": "2.0"}) + _sh("ray start --head --node-ip-address=127.0.0.1 --port=6379 --num-cpus=32 --num-gpus=1 " + "--include-dashboard=false --temp-dir=/root/miles-ray", log) + t0 = time.time() + with open(f"{out}/island.log", "w") as ilog: + rc = subprocess.run([PY, "-m", "yeto.rl.adapters.verl.island_entry", *args], env=env, + stdout=ilog, stderr=subprocess.STDOUT, cwd="/root").returncode + seconds = time.time() - t0 + _sh("ray stop --force", log) + if os.path.isdir("/root/yeto-output"): + shutil.copytree("/root/yeto-output", f"{out}/yeto-output", dirs_exist_ok=True) + if os.path.isdir(f"/root/state-{phase}"): + shutil.copytree(f"/root/state-{phase}", f"{out}/state-dir", dirs_exist_ok=True) + result = {"phase": phase, "rc": rc, "gpu": name, "island_seconds": round(seconds, 1), "args": args} + with open(f"{out}/result.json", "w") as fh: + json.dump(result, fh, indent=1) + log.close() + vol.commit() + return result + + +@app.function(timeout=90 * 60, volumes={"/vol": vol}) +def chain(run: str, gpu: str) -> dict: + control = island.spawn(run, "c", gpu) + results = {"a": island.remote(run, "a", gpu)} + vol.reload() + if results["a"]["rc"] == 0: + results["b"] = island.remote(run, "b", gpu) + results["c"] = control.get() + vol.reload() + with open(f"/vol/{run}/chain.json", "w") as fh: + json.dump(results, fh, indent=1) + vol.commit() + return results + + +@app.local_entrypoint() +def main(): + assert RUN != "unset", "set C4_RUN" + call = chain.spawn(RUN, GPU) + print("spawned", RUN, call.object_id, flush=True) diff --git a/tools/probes/g3_miles/compare.py b/tools/probes/g3_miles/compare.py new file mode 100644 index 00000000..96dbf180 --- /dev/null +++ b/tools/probes/g3_miles/compare.py @@ -0,0 +1,148 @@ +"""Judge the T-G3 harness outputs (CPU only). usage: python compare.py + +Writes ``/verdict.json`` and prints it. Thresholds are pre-registered in +SESSION25/gpu-review/T-G3.md; the noise floor (TE vs TE repeat, cont5 vs cont5b) +is reported next to every TE-vs-local and resume-vs-continuous difference. +""" + +from __future__ import annotations + +import json +import math +import sys +from pathlib import Path + +import torch + +sys.path.insert(0, str(Path(__file__).resolve().parent)) +from harness import map_local_to_te # noqa: E402 + +LP_MEAN_MAX, LP_ABS_MAX = 5e-3, 0.1 # label log-prob |delta|, nats +GRAD_REL_MAX, GRAD_GLOBAL_REL_MAX = 0.05, 0.02 # per-tensor / all-tensor relative L2 +LOSS_ABS_MAX = 1e-3 + + +def lp_stats(a: dict, b: dict) -> dict: + d = (a["token_loss"][:-1] - b["token_loss"][:-1]).abs().double() + return {"mean": float(d.mean()), "p99": float(d.quantile(0.99)), "max": float(d.max()), + "loss_a": a["loss"], "loss_b": b["loss"], "loss_abs": abs(a["loss"] - b["loss"]), + "logits_max_abs": float((a["logits"] - b["logits"]).abs().max())} + + +def grad_stats(out: Path, tag_a: str, tag_b: str, local_b: bool) -> dict: + per, worst, num, den, ranks = [], None, 0.0, 0.0, 0 + for r in range(8): + fa, fb = out / f"grads-{tag_a}.rank{r}.pt", out / f"grads-{tag_b}.rank{r}.pt" + if not fa.exists(): + break + ranks += 1 + ga = torch.load(fa, weights_only=True) + gb = torch.load(fb, weights_only=True) + if local_b: + mapping, missing, unused = map_local_to_te(gb, set(ga)) + if missing or unused: + return {"error": f"rank{r}: names do not map: missing={missing[:5]} unused={unused[:5]}"} + gb = {mapping[n]: t for n, t in gb.items()} + if set(ga) != set(gb): + return {"error": f"rank{r}: grad name sets differ: {sorted(set(ga) ^ set(gb))[:6]}"} + for n in sorted(ga): + a, b = ga[n].double(), gb[n].double() + na = float(a.norm()) + diff = float((a - b).norm()) + rel = diff / na if na > 0 else (0.0 if diff == 0 else math.inf) + cos = float((a * b).sum() / (a.norm() * b.norm())) if na > 0 and float(b.norm()) > 0 else None + num += diff ** 2 + den += na ** 2 + item = {"rank": r, "name": n, "rel": rel, "cos": cos, "norm": na} + per.append(item) + if worst is None or rel > worst["rel"]: + worst = item + per.sort(key=lambda x: -x["rel"]) + return {"ranks": ranks, "tensors": len(per), "global_rel": math.sqrt(num / den) if den else None, + "worst": per[:8], "over_threshold": [p for p in per if p["rel"] > GRAD_REL_MAX and p["norm"] > 0]} + + +def n75(out: Path) -> dict: + v: dict = {} + load = lambda t: torch.load(out / f"logits-{t}.pt", weights_only=True) # noqa: E731 + have = {t for t in ("te", "te2", "local") if (out / f"logits-{t}.pt").exists()} + if {"te", "te2"} <= have: + v["noise_te_vs_te2"] = {"logprob": lp_stats(load("te"), load("te2")), "grads": grad_stats(out, "te", "te2", False)} + if {"te", "local"} <= have: + lp = lp_stats(load("te"), load("local")) + g = grad_stats(out, "te", "local", True) + ok = (lp["mean"] <= LP_MEAN_MAX and lp["max"] <= LP_ABS_MAX and lp["loss_abs"] <= LOSS_ABS_MAX + and "error" not in g and not g["over_threshold"] and (g["global_rel"] or 0) <= GRAD_GLOBAL_REL_MAX) + v["te_vs_local"] = {"logprob": lp, "grads": g, "pass": ok} + for t in sorted(have): + info = json.loads((out / f"info-{t}.json").read_text()) + v[f"info_{t}"] = [{k: i.get(k) for k in ("impl", "transformer_impl", "dsa_core_attention", "loss", "timing", + "name_map", "num_params_local", "gpu")} for i in info] + v[f"n82_{t}"] = [i.get("n82") for i in info[:1]] + return v + + +def _beyond_noise(c5, c5b, rs) -> list: + """Parameters whose final norm after resume differs from cont5 by more than 5x the cont5/cont5b noise.""" + a, b, c = (x[0]["final_params"]["norms"] for x in (c5, c5b, rs)) + out = [] + for n in a: + noise = abs(b[n] - a[n]) / max(a[n], 1e-30) + res = abs(c[n] - a[n]) / max(a[n], 1e-30) + if res > 5 * noise + 1e-6: + out.append({"name": n, "resume_rel": res, "noise_rel": noise, "norm": a[n]}) + return sorted(out, key=lambda x: -x["resume_rel"]) + + +def n71d(out: Path) -> dict: + load = lambda m: json.loads((out / f"n71d-{m}.json").read_text()) if (out / f"n71d-{m}.json").exists() else None # noqa: E731 + c5, c5b, r3, rs = load("cont5"), load("cont5b"), load("run3"), load("resume2") + v: dict = {"present": [m for m, x in (("cont5", c5), ("cont5b", c5b), ("run3", r3), ("resume2", rs)) if x]} + if c5: + v["flags"] = c5[0]["flags"] + v["hdo_cont5_final"] = c5[0]["hdo_final"] + if c5 and c5b: + v["noise_cont5_vs_cont5b"] = { + "params_bitwise_equal": all(a["final_params"]["sha256"] == b["final_params"]["sha256"] for a, b in zip(c5, c5b)), + "loss_abs": [abs(a["loss_mean_dp"] - b["loss_mean_dp"]) for a, b in zip(c5[0]["steps"], c5b[0]["steps"])], + "grad_norm_abs": [abs(a["grad_norm"] - b["grad_norm"]) for a, b in zip(c5[0]["steps"], c5b[0]["steps"])]} + if c5 and r3: + v["run3_vs_cont5_steps1to3_loss_abs"] = [abs(a["loss_mean_dp"] - b["loss_mean_dp"]) + for a, b in zip(r3[0]["steps"], c5[0]["steps"][:3])] + if rs: + r0 = rs + v["restore"] = { + "digest_equal_per_rank": [x["restore_digest_equal"] for x in r0], + "restore_problems": r0[0]["restore_problems"], "first_step_problems": r0[0]["first_step_problems"], + "hdo_steps_after_restore": [[l.get("steps") for l in x["hdo_after_restore"]["leaves"]] for x in r0], + "hdo_steps_after_first_step": [[l.get("steps") for l in x["hdo_after_first_step"]["leaves"]] for x in r0], + "diff": [x["restore_summary"].get("diff") for x in r0]} + if c5: + steps = {s["step"]: s for s in c5[0]["steps"]} + v["resume_vs_cont5"] = { + "params_bitwise_equal": all(a["final_params"]["sha256"] == b["final_params"]["sha256"] + for a, b in zip(rs, c5)), + "loss_abs": {s["step"]: abs(s["loss_mean_dp"] - steps[s["step"]]["loss_mean_dp"]) for s in rs[0]["steps"]}, + "grad_norm_abs": {s["step"]: abs(s["grad_norm"] - steps[s["step"]]["grad_norm"]) for s in rs[0]["steps"]}, + "lr": {s["step"]: (s["lr"], steps[s["step"]]["lr"]) for s in rs[0]["steps"]}, + "beyond_5x_noise": _beyond_noise(c5, c5b, rs) if c5b else None, + "param_norm_max_rel": max( + abs(a["final_params"]["norms"][n] - b["final_params"]["norms"][n]) / max(b["final_params"]["norms"][n], 1e-30) + for a, b in zip(rs, c5) for n in b["final_params"]["norms"])} + return v + + +def main() -> None: + out = Path(sys.argv[1]) + verdict = {} + for name, fn, sub in (("n75", n75, "n75"), ("n71d", n71d, "n71d")): + try: + verdict[name] = fn(out / sub) + except Exception as exc: # noqa: BLE001 + verdict[name] = {"error": f"{type(exc).__name__}: {exc}"} + (out / "verdict.json").write_text(json.dumps(verdict, indent=1, default=str)) + print(json.dumps(verdict, indent=1, default=str)[:20000]) + + +if __name__ == "__main__": + main() diff --git a/tools/probes/g3_miles/glm53flash_config.json b/tools/probes/g3_miles/glm53flash_config.json new file mode 100644 index 00000000..f93128cf --- /dev/null +++ b/tools/probes/g3_miles/glm53flash_config.json @@ -0,0 +1,1810 @@ +{ + "architectures": [ + "Glm5NextForConditionalGeneration" + ], + "text_config": { + "attention_bias": false, + "attention_dropout": 0.0, + "dtype": "bfloat16", + "eos_token_id": [ + 154820, + 154827, + 154829 + ], + "first_k_dense_replace": 3, + "hc_eps": 1e-06, + "hc_mult": 4, + "hc_sinkhorn_iters": 20, + "head_dim": 0, + "hidden_act": "silu", + "hidden_size": 4096, + "index_head_dim": 128, + "index_kpool": 4, + "index_kpool_always_select_tail": true, + "index_kpool_compress": true, + "index_n_heads": 32, + "index_topk": 2048, + "index_share_for_mtp_iteration": true, + "indexer_rope_interleave": true, + "indexer_types": [ + "full", + "full", + "full", + "full", + "full", + "full", + "full", + "full", + "full", + "full", + "full", + "full", + "full", + "full", + "full", + "full", + "full", + "full", + "full", + "full", + "full", + "full", + "full", + "full", + "full", + "full", + "full", + "full", + "full", + "full", + "full", + "full", + "full", + "full", + "full", + "full", + "full", + "full", + "full", + "full", + "full", + "full", + "full", + "full", + "full" + ], + "initializer_range": 0.02, + "intermediate_size": 12288, + "kv_lora_rank": 512, + "layer_types": [ + "linear_attention", + "linear_attention", + "linear_attention", + "deepseek_sparse_attention", + "linear_attention", + "linear_attention", + "linear_attention", + "deepseek_sparse_attention", + "linear_attention", + "linear_attention", + "linear_attention", + "deepseek_sparse_attention", + "linear_attention", + "linear_attention", + "linear_attention", + "deepseek_sparse_attention", + "linear_attention", + "linear_attention", + "linear_attention", + "deepseek_sparse_attention", + "linear_attention", + "linear_attention", + "linear_attention", + "deepseek_sparse_attention", + "linear_attention", + "linear_attention", + "linear_attention", + "deepseek_sparse_attention", + "linear_attention", + "linear_attention", + "linear_attention", + "deepseek_sparse_attention", + "linear_attention", + "linear_attention", + "linear_attention", + "deepseek_sparse_attention", + "linear_attention", + "linear_attention", + "linear_attention", + "deepseek_sparse_attention", + "linear_attention", + "linear_attention", + "linear_attention", + "deepseek_sparse_attention", + "linear_attention" + ], + "linear_attn_config": { + "num_heads": 64, + "gate_lower_bound": -5.0, + "head_dim": 128, + "short_conv_kernel_size": 4, + "kda_layers": [ + 0, + 1, + 2, + 4, + 5, + 6, + 8, + 9, + 10, + 12, + 13, + 14, + 16, + 17, + 18, + 20, + 21, + 22, + 24, + 25, + 26, + 28, + 29, + 30, + 32, + 33, + 34, + 36, + 37, + 38, + 40, + 41, + 42, + 44 + ], + "full_attn_layers": [ + 3, + 7, + 11, + 15, + 19, + 23, + 27, + 31, + 35, + 39, + 43 + ] + }, + "max_position_embeddings": 1048576, + "mhc": true, + "mla_use_nope": true, + "mlp_layer_types": [ + "dense", + "dense", + "dense", + "sparse", + "sparse", + "sparse", + "sparse", + "sparse", + "sparse", + "sparse", + "sparse", + "sparse", + "sparse", + "sparse", + "sparse", + "sparse", + "sparse", + "sparse", + "sparse", + "sparse", + "sparse", + "sparse", + "sparse", + "sparse", + "sparse", + "sparse", + "sparse", + "sparse", + "sparse", + "sparse", + "sparse", + "sparse", + "sparse", + "sparse", + "sparse", + "sparse", + "sparse", + "sparse", + "sparse", + "sparse", + "sparse", + "sparse", + "sparse", + "sparse", + "sparse" + ], + "model_type": "glm5_next_text", + "moe_intermediate_size": 2048, + "moe_router_dtype": "float32", + "n_group": 1, + "n_routed_experts": 288, + "n_shared_experts": 1, + "norm_topk_prob": true, + "num_attention_heads": 64, + "num_experts_per_tok": 8, + "num_hidden_layers": 45, + "num_key_value_heads": 64, + "num_nextn_predict_layers": 1, + "output_router_logits": false, + "pad_token_id": 154820, + "q_lora_rank": 1536, + "qk_head_dim": 256, + "qk_nope_head_dim": 256, + "qk_rope_head_dim": 0, + "rms_norm_eps": 1e-05, + "routed_scaling_factor": 2.5, + "router_aux_loss_coef": 0.001, + "scoring_func": "sigmoid", + "swiglu_limit": 10.0, + "tie_word_embeddings": false, + "topk_group": 1, + "topk_method": "noaux_tc", + "use_cache": true, + "v_head_dim": 256, + "vocab_size": 154880 + }, + "vision_config": { + "attention_bias": true, + "attention_dropout": 0.0, + "depth": 24, + "hidden_act": "silu", + "hidden_size": 1024, + "image_size": 448, + "in_channels": 3, + "initializer_range": 0.02, + "intermediate_size": 4096, + "model_type": "glm5_next_vision", + "num_heads": 16, + "out_hidden_size": 4096, + "patch_size": 14, + "projection_intermediate_size": 10240, + "rms_norm_eps": 1e-05, + "spatial_merge_size": 2, + "swiglu_limit": 10.0, + "temporal_patch_size": 2 + }, + "image_token_id": 154854, + "video_token_id": 154855, + "image_start_token_id": 154830, + "image_end_token_id": 154831, + "video_start_token_id": 154832, + "video_end_token_id": 154833, + "tie_word_embeddings": false, + "model_type": "glm5_next", + "transformers_version": "5.16.0", + "quantization_config": { + "activation_scheme": "dynamic", + "fmt": "e4m3", + "modules_to_not_convert": [ + "attn_mha", + "attn_mqa", + "dt_bias", + "hyper_connection", + "lm_head", + "mapping_proj", + "model.embed_tokens", + "model.layers.0.hc_attn_base", + "model.layers.0.hc_attn_fn", + "model.layers.0.hc_attn_scale", + "model.layers.0.hc_ffn_base", + "model.layers.0.hc_ffn_fn", + "model.layers.0.hc_ffn_scale", + "model.layers.0.input_layernorm", + "model.layers.0.post_attention_layernorm", + "model.layers.0.self_attn.A_log", + "model.layers.0.self_attn.b_proj", + "model.layers.0.self_attn.dt_bias", + "model.layers.0.self_attn.f_a_proj", + "model.layers.0.self_attn.f_b_proj", + "model.layers.0.self_attn.fused_qkvbfg_a_proj", + "model.layers.0.self_attn.g_a_proj", + "model.layers.0.self_attn.g_b_proj", + "model.layers.0.self_attn.k_conv1d", + "model.layers.0.self_attn.k_proj", + "model.layers.0.self_attn.o_norm", + "model.layers.0.self_attn.o_proj", + "model.layers.0.self_attn.q_conv1d", + "model.layers.0.self_attn.q_proj", + "model.layers.0.self_attn.qkv_proj", + "model.layers.0.self_attn.v_conv1d", + "model.layers.0.self_attn.v_proj", + "model.layers.1.hc_attn_base", + "model.layers.1.hc_attn_fn", + "model.layers.1.hc_attn_scale", + "model.layers.1.hc_ffn_base", + "model.layers.1.hc_ffn_fn", + "model.layers.1.hc_ffn_scale", + "model.layers.1.input_layernorm", + "model.layers.1.post_attention_layernorm", + "model.layers.1.self_attn.A_log", + "model.layers.1.self_attn.b_proj", + "model.layers.1.self_attn.dt_bias", + "model.layers.1.self_attn.f_a_proj", + "model.layers.1.self_attn.f_b_proj", + "model.layers.1.self_attn.fused_qkvbfg_a_proj", + "model.layers.1.self_attn.g_a_proj", + "model.layers.1.self_attn.g_b_proj", + "model.layers.1.self_attn.k_conv1d", + "model.layers.1.self_attn.k_proj", + "model.layers.1.self_attn.o_norm", + "model.layers.1.self_attn.o_proj", + "model.layers.1.self_attn.q_conv1d", + "model.layers.1.self_attn.q_proj", + "model.layers.1.self_attn.qkv_proj", + "model.layers.1.self_attn.v_conv1d", + "model.layers.1.self_attn.v_proj", + "model.layers.10.hc_attn_base", + "model.layers.10.hc_attn_fn", + "model.layers.10.hc_attn_scale", + "model.layers.10.hc_ffn_base", + "model.layers.10.hc_ffn_fn", + "model.layers.10.hc_ffn_scale", + "model.layers.10.input_layernorm", + "model.layers.10.mlp.gate", + "model.layers.10.mlp.gate.e_score_correction_bias", + "model.layers.10.post_attention_layernorm", + "model.layers.10.self_attn.A_log", + "model.layers.10.self_attn.b_proj", + "model.layers.10.self_attn.dt_bias", + "model.layers.10.self_attn.f_a_proj", + "model.layers.10.self_attn.f_b_proj", + "model.layers.10.self_attn.fused_qkvbfg_a_proj", + "model.layers.10.self_attn.g_a_proj", + "model.layers.10.self_attn.g_b_proj", + "model.layers.10.self_attn.k_conv1d", + "model.layers.10.self_attn.k_proj", + "model.layers.10.self_attn.o_norm", + "model.layers.10.self_attn.o_proj", + "model.layers.10.self_attn.q_conv1d", + "model.layers.10.self_attn.q_proj", + "model.layers.10.self_attn.qkv_proj", + "model.layers.10.self_attn.v_conv1d", + "model.layers.10.self_attn.v_proj", + "model.layers.11.hc_attn_base", + "model.layers.11.hc_attn_fn", + "model.layers.11.hc_attn_scale", + "model.layers.11.hc_ffn_base", + "model.layers.11.hc_ffn_fn", + "model.layers.11.hc_ffn_scale", + "model.layers.11.input_layernorm", + "model.layers.11.mlp.gate", + "model.layers.11.mlp.gate.e_score_correction_bias", + "model.layers.11.post_attention_layernorm", + "model.layers.11.self_attn.indexer.index_kpool_compress_ape", + "model.layers.11.self_attn.indexer.index_kpool_compress_gate", + "model.layers.11.self_attn.indexer.k_norm", + "model.layers.11.self_attn.indexer.k_norm.bias", + "model.layers.11.self_attn.indexer.weights_proj", + "model.layers.11.self_attn.indexer.wk", + "model.layers.11.self_attn.indexer.wq_b", + "model.layers.11.self_attn.kv_a_layernorm", + "model.layers.11.self_attn.kv_b_proj", + "model.layers.11.self_attn.q_a_layernorm", + "model.layers.12.hc_attn_base", + "model.layers.12.hc_attn_fn", + "model.layers.12.hc_attn_scale", + "model.layers.12.hc_ffn_base", + "model.layers.12.hc_ffn_fn", + "model.layers.12.hc_ffn_scale", + "model.layers.12.input_layernorm", + "model.layers.12.mlp.gate", + "model.layers.12.mlp.gate.e_score_correction_bias", + "model.layers.12.post_attention_layernorm", + "model.layers.12.self_attn.A_log", + "model.layers.12.self_attn.b_proj", + "model.layers.12.self_attn.dt_bias", + "model.layers.12.self_attn.f_a_proj", + "model.layers.12.self_attn.f_b_proj", + "model.layers.12.self_attn.fused_qkvbfg_a_proj", + "model.layers.12.self_attn.g_a_proj", + "model.layers.12.self_attn.g_b_proj", + "model.layers.12.self_attn.k_conv1d", + "model.layers.12.self_attn.k_proj", + "model.layers.12.self_attn.o_norm", + "model.layers.12.self_attn.o_proj", + "model.layers.12.self_attn.q_conv1d", + "model.layers.12.self_attn.q_proj", + "model.layers.12.self_attn.qkv_proj", + "model.layers.12.self_attn.v_conv1d", + "model.layers.12.self_attn.v_proj", + "model.layers.13.hc_attn_base", + "model.layers.13.hc_attn_fn", + "model.layers.13.hc_attn_scale", + "model.layers.13.hc_ffn_base", + "model.layers.13.hc_ffn_fn", + "model.layers.13.hc_ffn_scale", + "model.layers.13.input_layernorm", + "model.layers.13.mlp.gate", + "model.layers.13.mlp.gate.e_score_correction_bias", + "model.layers.13.post_attention_layernorm", + "model.layers.13.self_attn.A_log", + "model.layers.13.self_attn.b_proj", + "model.layers.13.self_attn.dt_bias", + "model.layers.13.self_attn.f_a_proj", + "model.layers.13.self_attn.f_b_proj", + "model.layers.13.self_attn.fused_qkvbfg_a_proj", + "model.layers.13.self_attn.g_a_proj", + "model.layers.13.self_attn.g_b_proj", + "model.layers.13.self_attn.k_conv1d", + "model.layers.13.self_attn.k_proj", + "model.layers.13.self_attn.o_norm", + "model.layers.13.self_attn.o_proj", + "model.layers.13.self_attn.q_conv1d", + "model.layers.13.self_attn.q_proj", + "model.layers.13.self_attn.qkv_proj", + "model.layers.13.self_attn.v_conv1d", + "model.layers.13.self_attn.v_proj", + "model.layers.14.hc_attn_base", + "model.layers.14.hc_attn_fn", + "model.layers.14.hc_attn_scale", + "model.layers.14.hc_ffn_base", + "model.layers.14.hc_ffn_fn", + "model.layers.14.hc_ffn_scale", + "model.layers.14.input_layernorm", + "model.layers.14.mlp.gate", + "model.layers.14.mlp.gate.e_score_correction_bias", + "model.layers.14.post_attention_layernorm", + "model.layers.14.self_attn.A_log", + "model.layers.14.self_attn.b_proj", + "model.layers.14.self_attn.dt_bias", + "model.layers.14.self_attn.f_a_proj", + "model.layers.14.self_attn.f_b_proj", + "model.layers.14.self_attn.fused_qkvbfg_a_proj", + "model.layers.14.self_attn.g_a_proj", + "model.layers.14.self_attn.g_b_proj", + "model.layers.14.self_attn.k_conv1d", + "model.layers.14.self_attn.k_proj", + "model.layers.14.self_attn.o_norm", + "model.layers.14.self_attn.o_proj", + "model.layers.14.self_attn.q_conv1d", + "model.layers.14.self_attn.q_proj", + "model.layers.14.self_attn.qkv_proj", + "model.layers.14.self_attn.v_conv1d", + "model.layers.14.self_attn.v_proj", + "model.layers.15.hc_attn_base", + "model.layers.15.hc_attn_fn", + "model.layers.15.hc_attn_scale", + "model.layers.15.hc_ffn_base", + "model.layers.15.hc_ffn_fn", + "model.layers.15.hc_ffn_scale", + "model.layers.15.input_layernorm", + "model.layers.15.mlp.gate", + "model.layers.15.mlp.gate.e_score_correction_bias", + "model.layers.15.post_attention_layernorm", + "model.layers.15.self_attn.indexer.index_kpool_compress_ape", + "model.layers.15.self_attn.indexer.index_kpool_compress_gate", + "model.layers.15.self_attn.indexer.k_norm", + "model.layers.15.self_attn.indexer.k_norm.bias", + "model.layers.15.self_attn.indexer.weights_proj", + "model.layers.15.self_attn.indexer.wk", + "model.layers.15.self_attn.indexer.wq_b", + "model.layers.15.self_attn.kv_a_layernorm", + "model.layers.15.self_attn.kv_b_proj", + "model.layers.15.self_attn.q_a_layernorm", + "model.layers.16.hc_attn_base", + "model.layers.16.hc_attn_fn", + "model.layers.16.hc_attn_scale", + "model.layers.16.hc_ffn_base", + "model.layers.16.hc_ffn_fn", + "model.layers.16.hc_ffn_scale", + "model.layers.16.input_layernorm", + "model.layers.16.mlp.gate", + "model.layers.16.mlp.gate.e_score_correction_bias", + "model.layers.16.post_attention_layernorm", + "model.layers.16.self_attn.A_log", + "model.layers.16.self_attn.b_proj", + "model.layers.16.self_attn.dt_bias", + "model.layers.16.self_attn.f_a_proj", + "model.layers.16.self_attn.f_b_proj", + "model.layers.16.self_attn.fused_qkvbfg_a_proj", + "model.layers.16.self_attn.g_a_proj", + "model.layers.16.self_attn.g_b_proj", + "model.layers.16.self_attn.k_conv1d", + "model.layers.16.self_attn.k_proj", + "model.layers.16.self_attn.o_norm", + "model.layers.16.self_attn.o_proj", + "model.layers.16.self_attn.q_conv1d", + "model.layers.16.self_attn.q_proj", + "model.layers.16.self_attn.qkv_proj", + "model.layers.16.self_attn.v_conv1d", + "model.layers.16.self_attn.v_proj", + "model.layers.17.hc_attn_base", + "model.layers.17.hc_attn_fn", + "model.layers.17.hc_attn_scale", + "model.layers.17.hc_ffn_base", + "model.layers.17.hc_ffn_fn", + "model.layers.17.hc_ffn_scale", + "model.layers.17.input_layernorm", + "model.layers.17.mlp.gate", + "model.layers.17.mlp.gate.e_score_correction_bias", + "model.layers.17.post_attention_layernorm", + "model.layers.17.self_attn.A_log", + "model.layers.17.self_attn.b_proj", + "model.layers.17.self_attn.dt_bias", + "model.layers.17.self_attn.f_a_proj", + "model.layers.17.self_attn.f_b_proj", + "model.layers.17.self_attn.fused_qkvbfg_a_proj", + "model.layers.17.self_attn.g_a_proj", + "model.layers.17.self_attn.g_b_proj", + "model.layers.17.self_attn.k_conv1d", + "model.layers.17.self_attn.k_proj", + "model.layers.17.self_attn.o_norm", + "model.layers.17.self_attn.o_proj", + "model.layers.17.self_attn.q_conv1d", + "model.layers.17.self_attn.q_proj", + "model.layers.17.self_attn.qkv_proj", + "model.layers.17.self_attn.v_conv1d", + "model.layers.17.self_attn.v_proj", + "model.layers.18.hc_attn_base", + "model.layers.18.hc_attn_fn", + "model.layers.18.hc_attn_scale", + "model.layers.18.hc_ffn_base", + "model.layers.18.hc_ffn_fn", + "model.layers.18.hc_ffn_scale", + "model.layers.18.input_layernorm", + "model.layers.18.mlp.gate", + "model.layers.18.mlp.gate.e_score_correction_bias", + "model.layers.18.post_attention_layernorm", + "model.layers.18.self_attn.A_log", + "model.layers.18.self_attn.b_proj", + "model.layers.18.self_attn.dt_bias", + "model.layers.18.self_attn.f_a_proj", + "model.layers.18.self_attn.f_b_proj", + "model.layers.18.self_attn.fused_qkvbfg_a_proj", + "model.layers.18.self_attn.g_a_proj", + "model.layers.18.self_attn.g_b_proj", + "model.layers.18.self_attn.k_conv1d", + "model.layers.18.self_attn.k_proj", + "model.layers.18.self_attn.o_norm", + "model.layers.18.self_attn.o_proj", + "model.layers.18.self_attn.q_conv1d", + "model.layers.18.self_attn.q_proj", + "model.layers.18.self_attn.qkv_proj", + "model.layers.18.self_attn.v_conv1d", + "model.layers.18.self_attn.v_proj", + "model.layers.19.hc_attn_base", + "model.layers.19.hc_attn_fn", + "model.layers.19.hc_attn_scale", + "model.layers.19.hc_ffn_base", + "model.layers.19.hc_ffn_fn", + "model.layers.19.hc_ffn_scale", + "model.layers.19.input_layernorm", + "model.layers.19.mlp.gate", + "model.layers.19.mlp.gate.e_score_correction_bias", + "model.layers.19.post_attention_layernorm", + "model.layers.19.self_attn.indexer.index_kpool_compress_ape", + "model.layers.19.self_attn.indexer.index_kpool_compress_gate", + "model.layers.19.self_attn.indexer.k_norm", + "model.layers.19.self_attn.indexer.k_norm.bias", + "model.layers.19.self_attn.indexer.weights_proj", + "model.layers.19.self_attn.indexer.wk", + "model.layers.19.self_attn.indexer.wq_b", + "model.layers.19.self_attn.kv_a_layernorm", + "model.layers.19.self_attn.kv_b_proj", + "model.layers.19.self_attn.q_a_layernorm", + "model.layers.2.hc_attn_base", + "model.layers.2.hc_attn_fn", + "model.layers.2.hc_attn_scale", + "model.layers.2.hc_ffn_base", + "model.layers.2.hc_ffn_fn", + "model.layers.2.hc_ffn_scale", + "model.layers.2.input_layernorm", + "model.layers.2.post_attention_layernorm", + "model.layers.2.self_attn.A_log", + "model.layers.2.self_attn.b_proj", + "model.layers.2.self_attn.dt_bias", + "model.layers.2.self_attn.f_a_proj", + "model.layers.2.self_attn.f_b_proj", + "model.layers.2.self_attn.fused_qkvbfg_a_proj", + "model.layers.2.self_attn.g_a_proj", + "model.layers.2.self_attn.g_b_proj", + "model.layers.2.self_attn.k_conv1d", + "model.layers.2.self_attn.k_proj", + "model.layers.2.self_attn.o_norm", + "model.layers.2.self_attn.o_proj", + "model.layers.2.self_attn.q_conv1d", + "model.layers.2.self_attn.q_proj", + "model.layers.2.self_attn.qkv_proj", + "model.layers.2.self_attn.v_conv1d", + "model.layers.2.self_attn.v_proj", + "model.layers.20.hc_attn_base", + "model.layers.20.hc_attn_fn", + "model.layers.20.hc_attn_scale", + "model.layers.20.hc_ffn_base", + "model.layers.20.hc_ffn_fn", + "model.layers.20.hc_ffn_scale", + "model.layers.20.input_layernorm", + "model.layers.20.mlp.gate", + "model.layers.20.mlp.gate.e_score_correction_bias", + "model.layers.20.post_attention_layernorm", + "model.layers.20.self_attn.A_log", + "model.layers.20.self_attn.b_proj", + "model.layers.20.self_attn.dt_bias", + "model.layers.20.self_attn.f_a_proj", + "model.layers.20.self_attn.f_b_proj", + "model.layers.20.self_attn.fused_qkvbfg_a_proj", + "model.layers.20.self_attn.g_a_proj", + "model.layers.20.self_attn.g_b_proj", + "model.layers.20.self_attn.k_conv1d", + "model.layers.20.self_attn.k_proj", + "model.layers.20.self_attn.o_norm", + "model.layers.20.self_attn.o_proj", + "model.layers.20.self_attn.q_conv1d", + "model.layers.20.self_attn.q_proj", + "model.layers.20.self_attn.qkv_proj", + "model.layers.20.self_attn.v_conv1d", + "model.layers.20.self_attn.v_proj", + "model.layers.21.hc_attn_base", + "model.layers.21.hc_attn_fn", + "model.layers.21.hc_attn_scale", + "model.layers.21.hc_ffn_base", + "model.layers.21.hc_ffn_fn", + "model.layers.21.hc_ffn_scale", + "model.layers.21.input_layernorm", + "model.layers.21.mlp.gate", + "model.layers.21.mlp.gate.e_score_correction_bias", + "model.layers.21.post_attention_layernorm", + "model.layers.21.self_attn.A_log", + "model.layers.21.self_attn.b_proj", + "model.layers.21.self_attn.dt_bias", + "model.layers.21.self_attn.f_a_proj", + "model.layers.21.self_attn.f_b_proj", + "model.layers.21.self_attn.fused_qkvbfg_a_proj", + "model.layers.21.self_attn.g_a_proj", + "model.layers.21.self_attn.g_b_proj", + "model.layers.21.self_attn.k_conv1d", + "model.layers.21.self_attn.k_proj", + "model.layers.21.self_attn.o_norm", + "model.layers.21.self_attn.o_proj", + "model.layers.21.self_attn.q_conv1d", + "model.layers.21.self_attn.q_proj", + "model.layers.21.self_attn.qkv_proj", + "model.layers.21.self_attn.v_conv1d", + "model.layers.21.self_attn.v_proj", + "model.layers.22.hc_attn_base", + "model.layers.22.hc_attn_fn", + "model.layers.22.hc_attn_scale", + "model.layers.22.hc_ffn_base", + "model.layers.22.hc_ffn_fn", + "model.layers.22.hc_ffn_scale", + "model.layers.22.input_layernorm", + "model.layers.22.mlp.gate", + "model.layers.22.mlp.gate.e_score_correction_bias", + "model.layers.22.post_attention_layernorm", + "model.layers.22.self_attn.A_log", + "model.layers.22.self_attn.b_proj", + "model.layers.22.self_attn.dt_bias", + "model.layers.22.self_attn.f_a_proj", + "model.layers.22.self_attn.f_b_proj", + "model.layers.22.self_attn.fused_qkvbfg_a_proj", + "model.layers.22.self_attn.g_a_proj", + "model.layers.22.self_attn.g_b_proj", + "model.layers.22.self_attn.k_conv1d", + "model.layers.22.self_attn.k_proj", + "model.layers.22.self_attn.o_norm", + "model.layers.22.self_attn.o_proj", + "model.layers.22.self_attn.q_conv1d", + "model.layers.22.self_attn.q_proj", + "model.layers.22.self_attn.qkv_proj", + "model.layers.22.self_attn.v_conv1d", + "model.layers.22.self_attn.v_proj", + "model.layers.23.hc_attn_base", + "model.layers.23.hc_attn_fn", + "model.layers.23.hc_attn_scale", + "model.layers.23.hc_ffn_base", + "model.layers.23.hc_ffn_fn", + "model.layers.23.hc_ffn_scale", + "model.layers.23.input_layernorm", + "model.layers.23.mlp.gate", + "model.layers.23.mlp.gate.e_score_correction_bias", + "model.layers.23.post_attention_layernorm", + "model.layers.23.self_attn.indexer.index_kpool_compress_ape", + "model.layers.23.self_attn.indexer.index_kpool_compress_gate", + "model.layers.23.self_attn.indexer.k_norm", + "model.layers.23.self_attn.indexer.k_norm.bias", + "model.layers.23.self_attn.indexer.weights_proj", + "model.layers.23.self_attn.indexer.wk", + "model.layers.23.self_attn.indexer.wq_b", + "model.layers.23.self_attn.kv_a_layernorm", + "model.layers.23.self_attn.kv_b_proj", + "model.layers.23.self_attn.q_a_layernorm", + "model.layers.24.hc_attn_base", + "model.layers.24.hc_attn_fn", + "model.layers.24.hc_attn_scale", + "model.layers.24.hc_ffn_base", + "model.layers.24.hc_ffn_fn", + "model.layers.24.hc_ffn_scale", + "model.layers.24.input_layernorm", + "model.layers.24.mlp.gate", + "model.layers.24.mlp.gate.e_score_correction_bias", + "model.layers.24.post_attention_layernorm", + "model.layers.24.self_attn.A_log", + "model.layers.24.self_attn.b_proj", + "model.layers.24.self_attn.dt_bias", + "model.layers.24.self_attn.f_a_proj", + "model.layers.24.self_attn.f_b_proj", + "model.layers.24.self_attn.fused_qkvbfg_a_proj", + "model.layers.24.self_attn.g_a_proj", + "model.layers.24.self_attn.g_b_proj", + "model.layers.24.self_attn.k_conv1d", + "model.layers.24.self_attn.k_proj", + "model.layers.24.self_attn.o_norm", + "model.layers.24.self_attn.o_proj", + "model.layers.24.self_attn.q_conv1d", + "model.layers.24.self_attn.q_proj", + "model.layers.24.self_attn.qkv_proj", + "model.layers.24.self_attn.v_conv1d", + "model.layers.24.self_attn.v_proj", + "model.layers.25.hc_attn_base", + "model.layers.25.hc_attn_fn", + "model.layers.25.hc_attn_scale", + "model.layers.25.hc_ffn_base", + "model.layers.25.hc_ffn_fn", + "model.layers.25.hc_ffn_scale", + "model.layers.25.input_layernorm", + "model.layers.25.mlp.gate", + "model.layers.25.mlp.gate.e_score_correction_bias", + "model.layers.25.post_attention_layernorm", + "model.layers.25.self_attn.A_log", + "model.layers.25.self_attn.b_proj", + "model.layers.25.self_attn.dt_bias", + "model.layers.25.self_attn.f_a_proj", + "model.layers.25.self_attn.f_b_proj", + "model.layers.25.self_attn.fused_qkvbfg_a_proj", + "model.layers.25.self_attn.g_a_proj", + "model.layers.25.self_attn.g_b_proj", + "model.layers.25.self_attn.k_conv1d", + "model.layers.25.self_attn.k_proj", + "model.layers.25.self_attn.o_norm", + "model.layers.25.self_attn.o_proj", + "model.layers.25.self_attn.q_conv1d", + "model.layers.25.self_attn.q_proj", + "model.layers.25.self_attn.qkv_proj", + "model.layers.25.self_attn.v_conv1d", + "model.layers.25.self_attn.v_proj", + "model.layers.26.hc_attn_base", + "model.layers.26.hc_attn_fn", + "model.layers.26.hc_attn_scale", + "model.layers.26.hc_ffn_base", + "model.layers.26.hc_ffn_fn", + "model.layers.26.hc_ffn_scale", + "model.layers.26.input_layernorm", + "model.layers.26.mlp.gate", + "model.layers.26.mlp.gate.e_score_correction_bias", + "model.layers.26.post_attention_layernorm", + "model.layers.26.self_attn.A_log", + "model.layers.26.self_attn.b_proj", + "model.layers.26.self_attn.dt_bias", + "model.layers.26.self_attn.f_a_proj", + "model.layers.26.self_attn.f_b_proj", + "model.layers.26.self_attn.fused_qkvbfg_a_proj", + "model.layers.26.self_attn.g_a_proj", + "model.layers.26.self_attn.g_b_proj", + "model.layers.26.self_attn.k_conv1d", + "model.layers.26.self_attn.k_proj", + "model.layers.26.self_attn.o_norm", + "model.layers.26.self_attn.o_proj", + "model.layers.26.self_attn.q_conv1d", + "model.layers.26.self_attn.q_proj", + "model.layers.26.self_attn.qkv_proj", + "model.layers.26.self_attn.v_conv1d", + "model.layers.26.self_attn.v_proj", + "model.layers.27.hc_attn_base", + "model.layers.27.hc_attn_fn", + "model.layers.27.hc_attn_scale", + "model.layers.27.hc_ffn_base", + "model.layers.27.hc_ffn_fn", + "model.layers.27.hc_ffn_scale", + "model.layers.27.input_layernorm", + "model.layers.27.mlp.gate", + "model.layers.27.mlp.gate.e_score_correction_bias", + "model.layers.27.post_attention_layernorm", + "model.layers.27.self_attn.indexer.index_kpool_compress_ape", + "model.layers.27.self_attn.indexer.index_kpool_compress_gate", + "model.layers.27.self_attn.indexer.k_norm", + "model.layers.27.self_attn.indexer.k_norm.bias", + "model.layers.27.self_attn.indexer.weights_proj", + "model.layers.27.self_attn.indexer.wk", + "model.layers.27.self_attn.indexer.wq_b", + "model.layers.27.self_attn.kv_a_layernorm", + "model.layers.27.self_attn.kv_b_proj", + "model.layers.27.self_attn.q_a_layernorm", + "model.layers.28.hc_attn_base", + "model.layers.28.hc_attn_fn", + "model.layers.28.hc_attn_scale", + "model.layers.28.hc_ffn_base", + "model.layers.28.hc_ffn_fn", + "model.layers.28.hc_ffn_scale", + "model.layers.28.input_layernorm", + "model.layers.28.mlp.gate", + "model.layers.28.mlp.gate.e_score_correction_bias", + "model.layers.28.post_attention_layernorm", + "model.layers.28.self_attn.A_log", + "model.layers.28.self_attn.b_proj", + "model.layers.28.self_attn.dt_bias", + "model.layers.28.self_attn.f_a_proj", + "model.layers.28.self_attn.f_b_proj", + "model.layers.28.self_attn.fused_qkvbfg_a_proj", + "model.layers.28.self_attn.g_a_proj", + "model.layers.28.self_attn.g_b_proj", + "model.layers.28.self_attn.k_conv1d", + "model.layers.28.self_attn.k_proj", + "model.layers.28.self_attn.o_norm", + "model.layers.28.self_attn.o_proj", + "model.layers.28.self_attn.q_conv1d", + "model.layers.28.self_attn.q_proj", + "model.layers.28.self_attn.qkv_proj", + "model.layers.28.self_attn.v_conv1d", + "model.layers.28.self_attn.v_proj", + "model.layers.29.hc_attn_base", + "model.layers.29.hc_attn_fn", + "model.layers.29.hc_attn_scale", + "model.layers.29.hc_ffn_base", + "model.layers.29.hc_ffn_fn", + "model.layers.29.hc_ffn_scale", + "model.layers.29.input_layernorm", + "model.layers.29.mlp.gate", + "model.layers.29.mlp.gate.e_score_correction_bias", + "model.layers.29.post_attention_layernorm", + "model.layers.29.self_attn.A_log", + "model.layers.29.self_attn.b_proj", + "model.layers.29.self_attn.dt_bias", + "model.layers.29.self_attn.f_a_proj", + "model.layers.29.self_attn.f_b_proj", + "model.layers.29.self_attn.fused_qkvbfg_a_proj", + "model.layers.29.self_attn.g_a_proj", + "model.layers.29.self_attn.g_b_proj", + "model.layers.29.self_attn.k_conv1d", + "model.layers.29.self_attn.k_proj", + "model.layers.29.self_attn.o_norm", + "model.layers.29.self_attn.o_proj", + "model.layers.29.self_attn.q_conv1d", + "model.layers.29.self_attn.q_proj", + "model.layers.29.self_attn.qkv_proj", + "model.layers.29.self_attn.v_conv1d", + "model.layers.29.self_attn.v_proj", + "model.layers.3.hc_attn_base", + "model.layers.3.hc_attn_fn", + "model.layers.3.hc_attn_scale", + "model.layers.3.hc_ffn_base", + "model.layers.3.hc_ffn_fn", + "model.layers.3.hc_ffn_scale", + "model.layers.3.input_layernorm", + "model.layers.3.mlp.gate", + "model.layers.3.mlp.gate.e_score_correction_bias", + "model.layers.3.post_attention_layernorm", + "model.layers.3.self_attn.indexer.index_kpool_compress_ape", + "model.layers.3.self_attn.indexer.index_kpool_compress_gate", + "model.layers.3.self_attn.indexer.k_norm", + "model.layers.3.self_attn.indexer.k_norm.bias", + "model.layers.3.self_attn.indexer.weights_proj", + "model.layers.3.self_attn.indexer.wk", + "model.layers.3.self_attn.indexer.wq_b", + "model.layers.3.self_attn.kv_a_layernorm", + "model.layers.3.self_attn.kv_b_proj", + "model.layers.3.self_attn.q_a_layernorm", + "model.layers.30.hc_attn_base", + "model.layers.30.hc_attn_fn", + "model.layers.30.hc_attn_scale", + "model.layers.30.hc_ffn_base", + "model.layers.30.hc_ffn_fn", + "model.layers.30.hc_ffn_scale", + "model.layers.30.input_layernorm", + "model.layers.30.mlp.gate", + "model.layers.30.mlp.gate.e_score_correction_bias", + "model.layers.30.post_attention_layernorm", + "model.layers.30.self_attn.A_log", + "model.layers.30.self_attn.b_proj", + "model.layers.30.self_attn.dt_bias", + "model.layers.30.self_attn.f_a_proj", + "model.layers.30.self_attn.f_b_proj", + "model.layers.30.self_attn.fused_qkvbfg_a_proj", + "model.layers.30.self_attn.g_a_proj", + "model.layers.30.self_attn.g_b_proj", + "model.layers.30.self_attn.k_conv1d", + "model.layers.30.self_attn.k_proj", + "model.layers.30.self_attn.o_norm", + "model.layers.30.self_attn.o_proj", + "model.layers.30.self_attn.q_conv1d", + "model.layers.30.self_attn.q_proj", + "model.layers.30.self_attn.qkv_proj", + "model.layers.30.self_attn.v_conv1d", + "model.layers.30.self_attn.v_proj", + "model.layers.31.hc_attn_base", + "model.layers.31.hc_attn_fn", + "model.layers.31.hc_attn_scale", + "model.layers.31.hc_ffn_base", + "model.layers.31.hc_ffn_fn", + "model.layers.31.hc_ffn_scale", + "model.layers.31.input_layernorm", + "model.layers.31.mlp.gate", + "model.layers.31.mlp.gate.e_score_correction_bias", + "model.layers.31.post_attention_layernorm", + "model.layers.31.self_attn.indexer.index_kpool_compress_ape", + "model.layers.31.self_attn.indexer.index_kpool_compress_gate", + "model.layers.31.self_attn.indexer.k_norm", + "model.layers.31.self_attn.indexer.k_norm.bias", + "model.layers.31.self_attn.indexer.weights_proj", + "model.layers.31.self_attn.indexer.wk", + "model.layers.31.self_attn.indexer.wq_b", + "model.layers.31.self_attn.kv_a_layernorm", + "model.layers.31.self_attn.kv_b_proj", + "model.layers.31.self_attn.q_a_layernorm", + "model.layers.32.hc_attn_base", + "model.layers.32.hc_attn_fn", + "model.layers.32.hc_attn_scale", + "model.layers.32.hc_ffn_base", + "model.layers.32.hc_ffn_fn", + "model.layers.32.hc_ffn_scale", + "model.layers.32.input_layernorm", + "model.layers.32.mlp.gate", + "model.layers.32.mlp.gate.e_score_correction_bias", + "model.layers.32.post_attention_layernorm", + "model.layers.32.self_attn.A_log", + "model.layers.32.self_attn.b_proj", + "model.layers.32.self_attn.dt_bias", + "model.layers.32.self_attn.f_a_proj", + "model.layers.32.self_attn.f_b_proj", + "model.layers.32.self_attn.fused_qkvbfg_a_proj", + "model.layers.32.self_attn.g_a_proj", + "model.layers.32.self_attn.g_b_proj", + "model.layers.32.self_attn.k_conv1d", + "model.layers.32.self_attn.k_proj", + "model.layers.32.self_attn.o_norm", + "model.layers.32.self_attn.o_proj", + "model.layers.32.self_attn.q_conv1d", + "model.layers.32.self_attn.q_proj", + "model.layers.32.self_attn.qkv_proj", + "model.layers.32.self_attn.v_conv1d", + "model.layers.32.self_attn.v_proj", + "model.layers.33.hc_attn_base", + "model.layers.33.hc_attn_fn", + "model.layers.33.hc_attn_scale", + "model.layers.33.hc_ffn_base", + "model.layers.33.hc_ffn_fn", + "model.layers.33.hc_ffn_scale", + "model.layers.33.input_layernorm", + "model.layers.33.mlp.gate", + "model.layers.33.mlp.gate.e_score_correction_bias", + "model.layers.33.post_attention_layernorm", + "model.layers.33.self_attn.A_log", + "model.layers.33.self_attn.b_proj", + "model.layers.33.self_attn.dt_bias", + "model.layers.33.self_attn.f_a_proj", + "model.layers.33.self_attn.f_b_proj", + "model.layers.33.self_attn.fused_qkvbfg_a_proj", + "model.layers.33.self_attn.g_a_proj", + "model.layers.33.self_attn.g_b_proj", + "model.layers.33.self_attn.k_conv1d", + "model.layers.33.self_attn.k_proj", + "model.layers.33.self_attn.o_norm", + "model.layers.33.self_attn.o_proj", + "model.layers.33.self_attn.q_conv1d", + "model.layers.33.self_attn.q_proj", + "model.layers.33.self_attn.qkv_proj", + "model.layers.33.self_attn.v_conv1d", + "model.layers.33.self_attn.v_proj", + "model.layers.34.hc_attn_base", + "model.layers.34.hc_attn_fn", + "model.layers.34.hc_attn_scale", + "model.layers.34.hc_ffn_base", + "model.layers.34.hc_ffn_fn", + "model.layers.34.hc_ffn_scale", + "model.layers.34.input_layernorm", + "model.layers.34.mlp.gate", + "model.layers.34.mlp.gate.e_score_correction_bias", + "model.layers.34.post_attention_layernorm", + "model.layers.34.self_attn.A_log", + "model.layers.34.self_attn.b_proj", + "model.layers.34.self_attn.dt_bias", + "model.layers.34.self_attn.f_a_proj", + "model.layers.34.self_attn.f_b_proj", + "model.layers.34.self_attn.fused_qkvbfg_a_proj", + "model.layers.34.self_attn.g_a_proj", + "model.layers.34.self_attn.g_b_proj", + "model.layers.34.self_attn.k_conv1d", + "model.layers.34.self_attn.k_proj", + "model.layers.34.self_attn.o_norm", + "model.layers.34.self_attn.o_proj", + "model.layers.34.self_attn.q_conv1d", + "model.layers.34.self_attn.q_proj", + "model.layers.34.self_attn.qkv_proj", + "model.layers.34.self_attn.v_conv1d", + "model.layers.34.self_attn.v_proj", + "model.layers.35.hc_attn_base", + "model.layers.35.hc_attn_fn", + "model.layers.35.hc_attn_scale", + "model.layers.35.hc_ffn_base", + "model.layers.35.hc_ffn_fn", + "model.layers.35.hc_ffn_scale", + "model.layers.35.input_layernorm", + "model.layers.35.mlp.gate", + "model.layers.35.mlp.gate.e_score_correction_bias", + "model.layers.35.post_attention_layernorm", + "model.layers.35.self_attn.indexer.index_kpool_compress_ape", + "model.layers.35.self_attn.indexer.index_kpool_compress_gate", + "model.layers.35.self_attn.indexer.k_norm", + "model.layers.35.self_attn.indexer.k_norm.bias", + "model.layers.35.self_attn.indexer.weights_proj", + "model.layers.35.self_attn.indexer.wk", + "model.layers.35.self_attn.indexer.wq_b", + "model.layers.35.self_attn.kv_a_layernorm", + "model.layers.35.self_attn.kv_b_proj", + "model.layers.35.self_attn.q_a_layernorm", + "model.layers.36.hc_attn_base", + "model.layers.36.hc_attn_fn", + "model.layers.36.hc_attn_scale", + "model.layers.36.hc_ffn_base", + "model.layers.36.hc_ffn_fn", + "model.layers.36.hc_ffn_scale", + "model.layers.36.input_layernorm", + "model.layers.36.mlp.gate", + "model.layers.36.mlp.gate.e_score_correction_bias", + "model.layers.36.post_attention_layernorm", + "model.layers.36.self_attn.A_log", + "model.layers.36.self_attn.b_proj", + "model.layers.36.self_attn.dt_bias", + "model.layers.36.self_attn.f_a_proj", + "model.layers.36.self_attn.f_b_proj", + "model.layers.36.self_attn.fused_qkvbfg_a_proj", + "model.layers.36.self_attn.g_a_proj", + "model.layers.36.self_attn.g_b_proj", + "model.layers.36.self_attn.k_conv1d", + "model.layers.36.self_attn.k_proj", + "model.layers.36.self_attn.o_norm", + "model.layers.36.self_attn.o_proj", + "model.layers.36.self_attn.q_conv1d", + "model.layers.36.self_attn.q_proj", + "model.layers.36.self_attn.qkv_proj", + "model.layers.36.self_attn.v_conv1d", + "model.layers.36.self_attn.v_proj", + "model.layers.37.hc_attn_base", + "model.layers.37.hc_attn_fn", + "model.layers.37.hc_attn_scale", + "model.layers.37.hc_ffn_base", + "model.layers.37.hc_ffn_fn", + "model.layers.37.hc_ffn_scale", + "model.layers.37.input_layernorm", + "model.layers.37.mlp.gate", + "model.layers.37.mlp.gate.e_score_correction_bias", + "model.layers.37.post_attention_layernorm", + "model.layers.37.self_attn.A_log", + "model.layers.37.self_attn.b_proj", + "model.layers.37.self_attn.dt_bias", + "model.layers.37.self_attn.f_a_proj", + "model.layers.37.self_attn.f_b_proj", + "model.layers.37.self_attn.fused_qkvbfg_a_proj", + "model.layers.37.self_attn.g_a_proj", + "model.layers.37.self_attn.g_b_proj", + "model.layers.37.self_attn.k_conv1d", + "model.layers.37.self_attn.k_proj", + "model.layers.37.self_attn.o_norm", + "model.layers.37.self_attn.o_proj", + "model.layers.37.self_attn.q_conv1d", + "model.layers.37.self_attn.q_proj", + "model.layers.37.self_attn.qkv_proj", + "model.layers.37.self_attn.v_conv1d", + "model.layers.37.self_attn.v_proj", + "model.layers.38.hc_attn_base", + "model.layers.38.hc_attn_fn", + "model.layers.38.hc_attn_scale", + "model.layers.38.hc_ffn_base", + "model.layers.38.hc_ffn_fn", + "model.layers.38.hc_ffn_scale", + "model.layers.38.input_layernorm", + "model.layers.38.mlp.gate", + "model.layers.38.mlp.gate.e_score_correction_bias", + "model.layers.38.post_attention_layernorm", + "model.layers.38.self_attn.A_log", + "model.layers.38.self_attn.b_proj", + "model.layers.38.self_attn.dt_bias", + "model.layers.38.self_attn.f_a_proj", + "model.layers.38.self_attn.f_b_proj", + "model.layers.38.self_attn.fused_qkvbfg_a_proj", + "model.layers.38.self_attn.g_a_proj", + "model.layers.38.self_attn.g_b_proj", + "model.layers.38.self_attn.k_conv1d", + "model.layers.38.self_attn.k_proj", + "model.layers.38.self_attn.o_norm", + "model.layers.38.self_attn.o_proj", + "model.layers.38.self_attn.q_conv1d", + "model.layers.38.self_attn.q_proj", + "model.layers.38.self_attn.qkv_proj", + "model.layers.38.self_attn.v_conv1d", + "model.layers.38.self_attn.v_proj", + "model.layers.39.hc_attn_base", + "model.layers.39.hc_attn_fn", + "model.layers.39.hc_attn_scale", + "model.layers.39.hc_ffn_base", + "model.layers.39.hc_ffn_fn", + "model.layers.39.hc_ffn_scale", + "model.layers.39.input_layernorm", + "model.layers.39.mlp.gate", + "model.layers.39.mlp.gate.e_score_correction_bias", + "model.layers.39.post_attention_layernorm", + "model.layers.39.self_attn.indexer.index_kpool_compress_ape", + "model.layers.39.self_attn.indexer.index_kpool_compress_gate", + "model.layers.39.self_attn.indexer.k_norm", + "model.layers.39.self_attn.indexer.k_norm.bias", + "model.layers.39.self_attn.indexer.weights_proj", + "model.layers.39.self_attn.indexer.wk", + "model.layers.39.self_attn.indexer.wq_b", + "model.layers.39.self_attn.kv_a_layernorm", + "model.layers.39.self_attn.kv_b_proj", + "model.layers.39.self_attn.q_a_layernorm", + "model.layers.4.hc_attn_base", + "model.layers.4.hc_attn_fn", + "model.layers.4.hc_attn_scale", + "model.layers.4.hc_ffn_base", + "model.layers.4.hc_ffn_fn", + "model.layers.4.hc_ffn_scale", + "model.layers.4.input_layernorm", + "model.layers.4.mlp.gate", + "model.layers.4.mlp.gate.e_score_correction_bias", + "model.layers.4.post_attention_layernorm", + "model.layers.4.self_attn.A_log", + "model.layers.4.self_attn.b_proj", + "model.layers.4.self_attn.dt_bias", + "model.layers.4.self_attn.f_a_proj", + "model.layers.4.self_attn.f_b_proj", + "model.layers.4.self_attn.fused_qkvbfg_a_proj", + "model.layers.4.self_attn.g_a_proj", + "model.layers.4.self_attn.g_b_proj", + "model.layers.4.self_attn.k_conv1d", + "model.layers.4.self_attn.k_proj", + "model.layers.4.self_attn.o_norm", + "model.layers.4.self_attn.o_proj", + "model.layers.4.self_attn.q_conv1d", + "model.layers.4.self_attn.q_proj", + "model.layers.4.self_attn.qkv_proj", + "model.layers.4.self_attn.v_conv1d", + "model.layers.4.self_attn.v_proj", + "model.layers.40.hc_attn_base", + "model.layers.40.hc_attn_fn", + "model.layers.40.hc_attn_scale", + "model.layers.40.hc_ffn_base", + "model.layers.40.hc_ffn_fn", + "model.layers.40.hc_ffn_scale", + "model.layers.40.input_layernorm", + "model.layers.40.mlp.gate", + "model.layers.40.mlp.gate.e_score_correction_bias", + "model.layers.40.post_attention_layernorm", + "model.layers.40.self_attn.A_log", + "model.layers.40.self_attn.b_proj", + "model.layers.40.self_attn.dt_bias", + "model.layers.40.self_attn.f_a_proj", + "model.layers.40.self_attn.f_b_proj", + "model.layers.40.self_attn.fused_qkvbfg_a_proj", + "model.layers.40.self_attn.g_a_proj", + "model.layers.40.self_attn.g_b_proj", + "model.layers.40.self_attn.k_conv1d", + "model.layers.40.self_attn.k_proj", + "model.layers.40.self_attn.o_norm", + "model.layers.40.self_attn.o_proj", + "model.layers.40.self_attn.q_conv1d", + "model.layers.40.self_attn.q_proj", + "model.layers.40.self_attn.qkv_proj", + "model.layers.40.self_attn.v_conv1d", + "model.layers.40.self_attn.v_proj", + "model.layers.41.hc_attn_base", + "model.layers.41.hc_attn_fn", + "model.layers.41.hc_attn_scale", + "model.layers.41.hc_ffn_base", + "model.layers.41.hc_ffn_fn", + "model.layers.41.hc_ffn_scale", + "model.layers.41.input_layernorm", + "model.layers.41.mlp.gate", + "model.layers.41.mlp.gate.e_score_correction_bias", + "model.layers.41.post_attention_layernorm", + "model.layers.41.self_attn.A_log", + "model.layers.41.self_attn.b_proj", + "model.layers.41.self_attn.dt_bias", + "model.layers.41.self_attn.f_a_proj", + "model.layers.41.self_attn.f_b_proj", + "model.layers.41.self_attn.fused_qkvbfg_a_proj", + "model.layers.41.self_attn.g_a_proj", + "model.layers.41.self_attn.g_b_proj", + "model.layers.41.self_attn.k_conv1d", + "model.layers.41.self_attn.k_proj", + "model.layers.41.self_attn.o_norm", + "model.layers.41.self_attn.o_proj", + "model.layers.41.self_attn.q_conv1d", + "model.layers.41.self_attn.q_proj", + "model.layers.41.self_attn.qkv_proj", + "model.layers.41.self_attn.v_conv1d", + "model.layers.41.self_attn.v_proj", + "model.layers.42.hc_attn_base", + "model.layers.42.hc_attn_fn", + "model.layers.42.hc_attn_scale", + "model.layers.42.hc_ffn_base", + "model.layers.42.hc_ffn_fn", + "model.layers.42.hc_ffn_scale", + "model.layers.42.input_layernorm", + "model.layers.42.mlp.gate", + "model.layers.42.mlp.gate.e_score_correction_bias", + "model.layers.42.post_attention_layernorm", + "model.layers.42.self_attn.A_log", + "model.layers.42.self_attn.b_proj", + "model.layers.42.self_attn.dt_bias", + "model.layers.42.self_attn.f_a_proj", + "model.layers.42.self_attn.f_b_proj", + "model.layers.42.self_attn.fused_qkvbfg_a_proj", + "model.layers.42.self_attn.g_a_proj", + "model.layers.42.self_attn.g_b_proj", + "model.layers.42.self_attn.k_conv1d", + "model.layers.42.self_attn.k_proj", + "model.layers.42.self_attn.o_norm", + "model.layers.42.self_attn.o_proj", + "model.layers.42.self_attn.q_conv1d", + "model.layers.42.self_attn.q_proj", + "model.layers.42.self_attn.qkv_proj", + "model.layers.42.self_attn.v_conv1d", + "model.layers.42.self_attn.v_proj", + "model.layers.43.hc_attn_base", + "model.layers.43.hc_attn_fn", + "model.layers.43.hc_attn_scale", + "model.layers.43.hc_ffn_base", + "model.layers.43.hc_ffn_fn", + "model.layers.43.hc_ffn_scale", + "model.layers.43.input_layernorm", + "model.layers.43.mlp.gate", + "model.layers.43.mlp.gate.e_score_correction_bias", + "model.layers.43.post_attention_layernorm", + "model.layers.43.self_attn.indexer.index_kpool_compress_ape", + "model.layers.43.self_attn.indexer.index_kpool_compress_gate", + "model.layers.43.self_attn.indexer.k_norm", + "model.layers.43.self_attn.indexer.k_norm.bias", + "model.layers.43.self_attn.indexer.weights_proj", + "model.layers.43.self_attn.indexer.wk", + "model.layers.43.self_attn.indexer.wq_b", + "model.layers.43.self_attn.kv_a_layernorm", + "model.layers.43.self_attn.kv_b_proj", + "model.layers.43.self_attn.q_a_layernorm", + "model.layers.44.hc_attn_base", + "model.layers.44.hc_attn_fn", + "model.layers.44.hc_attn_scale", + "model.layers.44.hc_ffn_base", + "model.layers.44.hc_ffn_fn", + "model.layers.44.hc_ffn_scale", + "model.layers.44.input_layernorm", + "model.layers.44.mlp.gate", + "model.layers.44.mlp.gate.e_score_correction_bias", + "model.layers.44.post_attention_layernorm", + "model.layers.44.self_attn.A_log", + "model.layers.44.self_attn.b_proj", + "model.layers.44.self_attn.dt_bias", + "model.layers.44.self_attn.f_a_proj", + "model.layers.44.self_attn.f_b_proj", + "model.layers.44.self_attn.fused_qkvbfg_a_proj", + "model.layers.44.self_attn.g_a_proj", + "model.layers.44.self_attn.g_b_proj", + "model.layers.44.self_attn.k_conv1d", + "model.layers.44.self_attn.k_proj", + "model.layers.44.self_attn.o_norm", + "model.layers.44.self_attn.o_proj", + "model.layers.44.self_attn.q_conv1d", + "model.layers.44.self_attn.q_proj", + "model.layers.44.self_attn.qkv_proj", + "model.layers.44.self_attn.v_conv1d", + "model.layers.44.self_attn.v_proj", + "model.layers.45.eh_proj", + "model.layers.45.enorm", + "model.layers.45.hnorm", + "model.layers.45.input_layernorm", + "model.layers.45.mlp.gate", + "model.layers.45.mlp.gate.e_score_correction_bias", + "model.layers.45.post_attention_layernorm", + "model.layers.45.self_attn.indexer.index_kpool_compress_ape", + "model.layers.45.self_attn.indexer.index_kpool_compress_gate", + "model.layers.45.self_attn.indexer.k_norm", + "model.layers.45.self_attn.indexer.k_norm.bias", + "model.layers.45.self_attn.indexer.weights_proj", + "model.layers.45.self_attn.indexer.wk", + "model.layers.45.self_attn.indexer.wq_b", + "model.layers.45.self_attn.kv_a_layernorm", + "model.layers.45.self_attn.kv_b_proj", + "model.layers.45.self_attn.q_a_layernorm", + "model.layers.45.shared_head.norm", + "model.layers.5.hc_attn_base", + "model.layers.5.hc_attn_fn", + "model.layers.5.hc_attn_scale", + "model.layers.5.hc_ffn_base", + "model.layers.5.hc_ffn_fn", + "model.layers.5.hc_ffn_scale", + "model.layers.5.input_layernorm", + "model.layers.5.mlp.gate", + "model.layers.5.mlp.gate.e_score_correction_bias", + "model.layers.5.post_attention_layernorm", + "model.layers.5.self_attn.A_log", + "model.layers.5.self_attn.b_proj", + "model.layers.5.self_attn.dt_bias", + "model.layers.5.self_attn.f_a_proj", + "model.layers.5.self_attn.f_b_proj", + "model.layers.5.self_attn.fused_qkvbfg_a_proj", + "model.layers.5.self_attn.g_a_proj", + "model.layers.5.self_attn.g_b_proj", + "model.layers.5.self_attn.k_conv1d", + "model.layers.5.self_attn.k_proj", + "model.layers.5.self_attn.o_norm", + "model.layers.5.self_attn.o_proj", + "model.layers.5.self_attn.q_conv1d", + "model.layers.5.self_attn.q_proj", + "model.layers.5.self_attn.qkv_proj", + "model.layers.5.self_attn.v_conv1d", + "model.layers.5.self_attn.v_proj", + "model.layers.6.hc_attn_base", + "model.layers.6.hc_attn_fn", + "model.layers.6.hc_attn_scale", + "model.layers.6.hc_ffn_base", + "model.layers.6.hc_ffn_fn", + "model.layers.6.hc_ffn_scale", + "model.layers.6.input_layernorm", + "model.layers.6.mlp.gate", + "model.layers.6.mlp.gate.e_score_correction_bias", + "model.layers.6.post_attention_layernorm", + "model.layers.6.self_attn.A_log", + "model.layers.6.self_attn.b_proj", + "model.layers.6.self_attn.dt_bias", + "model.layers.6.self_attn.f_a_proj", + "model.layers.6.self_attn.f_b_proj", + "model.layers.6.self_attn.fused_qkvbfg_a_proj", + "model.layers.6.self_attn.g_a_proj", + "model.layers.6.self_attn.g_b_proj", + "model.layers.6.self_attn.k_conv1d", + "model.layers.6.self_attn.k_proj", + "model.layers.6.self_attn.o_norm", + "model.layers.6.self_attn.o_proj", + "model.layers.6.self_attn.q_conv1d", + "model.layers.6.self_attn.q_proj", + "model.layers.6.self_attn.qkv_proj", + "model.layers.6.self_attn.v_conv1d", + "model.layers.6.self_attn.v_proj", + "model.layers.7.hc_attn_base", + "model.layers.7.hc_attn_fn", + "model.layers.7.hc_attn_scale", + "model.layers.7.hc_ffn_base", + "model.layers.7.hc_ffn_fn", + "model.layers.7.hc_ffn_scale", + "model.layers.7.input_layernorm", + "model.layers.7.mlp.gate", + "model.layers.7.mlp.gate.e_score_correction_bias", + "model.layers.7.post_attention_layernorm", + "model.layers.7.self_attn.indexer.index_kpool_compress_ape", + "model.layers.7.self_attn.indexer.index_kpool_compress_gate", + "model.layers.7.self_attn.indexer.k_norm", + "model.layers.7.self_attn.indexer.k_norm.bias", + "model.layers.7.self_attn.indexer.weights_proj", + "model.layers.7.self_attn.indexer.wk", + "model.layers.7.self_attn.indexer.wq_b", + "model.layers.7.self_attn.kv_a_layernorm", + "model.layers.7.self_attn.kv_b_proj", + "model.layers.7.self_attn.q_a_layernorm", + "model.layers.8.hc_attn_base", + "model.layers.8.hc_attn_fn", + "model.layers.8.hc_attn_scale", + "model.layers.8.hc_ffn_base", + "model.layers.8.hc_ffn_fn", + "model.layers.8.hc_ffn_scale", + "model.layers.8.input_layernorm", + "model.layers.8.mlp.gate", + "model.layers.8.mlp.gate.e_score_correction_bias", + "model.layers.8.post_attention_layernorm", + "model.layers.8.self_attn.A_log", + "model.layers.8.self_attn.b_proj", + "model.layers.8.self_attn.dt_bias", + "model.layers.8.self_attn.f_a_proj", + "model.layers.8.self_attn.f_b_proj", + "model.layers.8.self_attn.fused_qkvbfg_a_proj", + "model.layers.8.self_attn.g_a_proj", + "model.layers.8.self_attn.g_b_proj", + "model.layers.8.self_attn.k_conv1d", + "model.layers.8.self_attn.k_proj", + "model.layers.8.self_attn.o_norm", + "model.layers.8.self_attn.o_proj", + "model.layers.8.self_attn.q_conv1d", + "model.layers.8.self_attn.q_proj", + "model.layers.8.self_attn.qkv_proj", + "model.layers.8.self_attn.v_conv1d", + "model.layers.8.self_attn.v_proj", + "model.layers.9.hc_attn_base", + "model.layers.9.hc_attn_fn", + "model.layers.9.hc_attn_scale", + "model.layers.9.hc_ffn_base", + "model.layers.9.hc_ffn_fn", + "model.layers.9.hc_ffn_scale", + "model.layers.9.input_layernorm", + "model.layers.9.mlp.gate", + "model.layers.9.mlp.gate.e_score_correction_bias", + "model.layers.9.post_attention_layernorm", + "model.layers.9.self_attn.A_log", + "model.layers.9.self_attn.b_proj", + "model.layers.9.self_attn.dt_bias", + "model.layers.9.self_attn.f_a_proj", + "model.layers.9.self_attn.f_b_proj", + "model.layers.9.self_attn.fused_qkvbfg_a_proj", + "model.layers.9.self_attn.g_a_proj", + "model.layers.9.self_attn.g_b_proj", + "model.layers.9.self_attn.k_conv1d", + "model.layers.9.self_attn.k_proj", + "model.layers.9.self_attn.o_norm", + "model.layers.9.self_attn.o_proj", + "model.layers.9.self_attn.q_conv1d", + "model.layers.9.self_attn.q_proj", + "model.layers.9.self_attn.qkv_proj", + "model.layers.9.self_attn.v_conv1d", + "model.layers.9.self_attn.v_proj", + "model.norm", + "model.visual", + "router", + "visual", + "visual.blocks.0.attn.k_norm", + "visual.blocks.0.attn.proj", + "visual.blocks.0.attn.proj.bias", + "visual.blocks.0.attn.q_norm", + "visual.blocks.0.attn.qkv", + "visual.blocks.0.attn.qkv.bias", + "visual.blocks.0.mlp.down_proj", + "visual.blocks.0.mlp.down_proj.bias", + "visual.blocks.0.mlp.gate_proj", + "visual.blocks.0.mlp.gate_proj.bias", + "visual.blocks.0.mlp.up_proj", + "visual.blocks.0.mlp.up_proj.bias", + "visual.blocks.0.norm1", + "visual.blocks.0.norm2", + "visual.blocks.1.attn.k_norm", + "visual.blocks.1.attn.proj", + "visual.blocks.1.attn.proj.bias", + "visual.blocks.1.attn.q_norm", + "visual.blocks.1.attn.qkv", + "visual.blocks.1.attn.qkv.bias", + "visual.blocks.1.mlp.down_proj", + "visual.blocks.1.mlp.down_proj.bias", + "visual.blocks.1.mlp.gate_proj", + "visual.blocks.1.mlp.gate_proj.bias", + "visual.blocks.1.mlp.up_proj", + "visual.blocks.1.mlp.up_proj.bias", + "visual.blocks.1.norm1", + "visual.blocks.1.norm2", + "visual.blocks.10.attn.k_norm", + "visual.blocks.10.attn.proj", + "visual.blocks.10.attn.proj.bias", + "visual.blocks.10.attn.q_norm", + "visual.blocks.10.attn.qkv", + "visual.blocks.10.attn.qkv.bias", + "visual.blocks.10.mlp.down_proj", + "visual.blocks.10.mlp.down_proj.bias", + "visual.blocks.10.mlp.gate_proj", + "visual.blocks.10.mlp.gate_proj.bias", + "visual.blocks.10.mlp.up_proj", + "visual.blocks.10.mlp.up_proj.bias", + "visual.blocks.10.norm1", + "visual.blocks.10.norm2", + "visual.blocks.11.attn.k_norm", + "visual.blocks.11.attn.proj", + "visual.blocks.11.attn.proj.bias", + "visual.blocks.11.attn.q_norm", + "visual.blocks.11.attn.qkv", + "visual.blocks.11.attn.qkv.bias", + "visual.blocks.11.mlp.down_proj", + "visual.blocks.11.mlp.down_proj.bias", + "visual.blocks.11.mlp.gate_proj", + "visual.blocks.11.mlp.gate_proj.bias", + "visual.blocks.11.mlp.up_proj", + "visual.blocks.11.mlp.up_proj.bias", + "visual.blocks.11.norm1", + "visual.blocks.11.norm2", + "visual.blocks.12.attn.k_norm", + "visual.blocks.12.attn.proj", + "visual.blocks.12.attn.proj.bias", + "visual.blocks.12.attn.q_norm", + "visual.blocks.12.attn.qkv", + "visual.blocks.12.attn.qkv.bias", + "visual.blocks.12.mlp.down_proj", + "visual.blocks.12.mlp.down_proj.bias", + "visual.blocks.12.mlp.gate_proj", + "visual.blocks.12.mlp.gate_proj.bias", + "visual.blocks.12.mlp.up_proj", + "visual.blocks.12.mlp.up_proj.bias", + "visual.blocks.12.norm1", + "visual.blocks.12.norm2", + "visual.blocks.13.attn.k_norm", + "visual.blocks.13.attn.proj", + "visual.blocks.13.attn.proj.bias", + "visual.blocks.13.attn.q_norm", + "visual.blocks.13.attn.qkv", + "visual.blocks.13.attn.qkv.bias", + "visual.blocks.13.mlp.down_proj", + "visual.blocks.13.mlp.down_proj.bias", + "visual.blocks.13.mlp.gate_proj", + "visual.blocks.13.mlp.gate_proj.bias", + "visual.blocks.13.mlp.up_proj", + "visual.blocks.13.mlp.up_proj.bias", + "visual.blocks.13.norm1", + "visual.blocks.13.norm2", + "visual.blocks.14.attn.k_norm", + "visual.blocks.14.attn.proj", + "visual.blocks.14.attn.proj.bias", + "visual.blocks.14.attn.q_norm", + "visual.blocks.14.attn.qkv", + "visual.blocks.14.attn.qkv.bias", + "visual.blocks.14.mlp.down_proj", + "visual.blocks.14.mlp.down_proj.bias", + "visual.blocks.14.mlp.gate_proj", + "visual.blocks.14.mlp.gate_proj.bias", + "visual.blocks.14.mlp.up_proj", + "visual.blocks.14.mlp.up_proj.bias", + "visual.blocks.14.norm1", + "visual.blocks.14.norm2", + "visual.blocks.15.attn.k_norm", + "visual.blocks.15.attn.proj", + "visual.blocks.15.attn.proj.bias", + "visual.blocks.15.attn.q_norm", + "visual.blocks.15.attn.qkv", + "visual.blocks.15.attn.qkv.bias", + "visual.blocks.15.mlp.down_proj", + "visual.blocks.15.mlp.down_proj.bias", + "visual.blocks.15.mlp.gate_proj", + "visual.blocks.15.mlp.gate_proj.bias", + "visual.blocks.15.mlp.up_proj", + "visual.blocks.15.mlp.up_proj.bias", + "visual.blocks.15.norm1", + "visual.blocks.15.norm2", + "visual.blocks.16.attn.k_norm", + "visual.blocks.16.attn.proj", + "visual.blocks.16.attn.proj.bias", + "visual.blocks.16.attn.q_norm", + "visual.blocks.16.attn.qkv", + "visual.blocks.16.attn.qkv.bias", + "visual.blocks.16.mlp.down_proj", + "visual.blocks.16.mlp.down_proj.bias", + "visual.blocks.16.mlp.gate_proj", + "visual.blocks.16.mlp.gate_proj.bias", + "visual.blocks.16.mlp.up_proj", + "visual.blocks.16.mlp.up_proj.bias", + "visual.blocks.16.norm1", + "visual.blocks.16.norm2", + "visual.blocks.17.attn.k_norm", + "visual.blocks.17.attn.proj", + "visual.blocks.17.attn.proj.bias", + "visual.blocks.17.attn.q_norm", + "visual.blocks.17.attn.qkv", + "visual.blocks.17.attn.qkv.bias", + "visual.blocks.17.mlp.down_proj", + "visual.blocks.17.mlp.down_proj.bias", + "visual.blocks.17.mlp.gate_proj", + "visual.blocks.17.mlp.gate_proj.bias", + "visual.blocks.17.mlp.up_proj", + "visual.blocks.17.mlp.up_proj.bias", + "visual.blocks.17.norm1", + "visual.blocks.17.norm2", + "visual.blocks.18.attn.k_norm", + "visual.blocks.18.attn.proj", + "visual.blocks.18.attn.proj.bias", + "visual.blocks.18.attn.q_norm", + "visual.blocks.18.attn.qkv", + "visual.blocks.18.attn.qkv.bias", + "visual.blocks.18.mlp.down_proj", + "visual.blocks.18.mlp.down_proj.bias", + "visual.blocks.18.mlp.gate_proj", + "visual.blocks.18.mlp.gate_proj.bias", + "visual.blocks.18.mlp.up_proj", + "visual.blocks.18.mlp.up_proj.bias", + "visual.blocks.18.norm1", + "visual.blocks.18.norm2", + "visual.blocks.19.attn.k_norm", + "visual.blocks.19.attn.proj", + "visual.blocks.19.attn.proj.bias", + "visual.blocks.19.attn.q_norm", + "visual.blocks.19.attn.qkv", + "visual.blocks.19.attn.qkv.bias", + "visual.blocks.19.mlp.down_proj", + "visual.blocks.19.mlp.down_proj.bias", + "visual.blocks.19.mlp.gate_proj", + "visual.blocks.19.mlp.gate_proj.bias", + "visual.blocks.19.mlp.up_proj", + "visual.blocks.19.mlp.up_proj.bias", + "visual.blocks.19.norm1", + "visual.blocks.19.norm2", + "visual.blocks.2.attn.k_norm", + "visual.blocks.2.attn.proj", + "visual.blocks.2.attn.proj.bias", + "visual.blocks.2.attn.q_norm", + "visual.blocks.2.attn.qkv", + "visual.blocks.2.attn.qkv.bias", + "visual.blocks.2.mlp.down_proj", + "visual.blocks.2.mlp.down_proj.bias", + "visual.blocks.2.mlp.gate_proj", + "visual.blocks.2.mlp.gate_proj.bias", + "visual.blocks.2.mlp.up_proj", + "visual.blocks.2.mlp.up_proj.bias", + "visual.blocks.2.norm1", + "visual.blocks.2.norm2", + "visual.blocks.20.attn.k_norm", + "visual.blocks.20.attn.proj", + "visual.blocks.20.attn.proj.bias", + "visual.blocks.20.attn.q_norm", + "visual.blocks.20.attn.qkv", + "visual.blocks.20.attn.qkv.bias", + "visual.blocks.20.mlp.down_proj", + "visual.blocks.20.mlp.down_proj.bias", + "visual.blocks.20.mlp.gate_proj", + "visual.blocks.20.mlp.gate_proj.bias", + "visual.blocks.20.mlp.up_proj", + "visual.blocks.20.mlp.up_proj.bias", + "visual.blocks.20.norm1", + "visual.blocks.20.norm2", + "visual.blocks.21.attn.k_norm", + "visual.blocks.21.attn.proj", + "visual.blocks.21.attn.proj.bias", + "visual.blocks.21.attn.q_norm", + "visual.blocks.21.attn.qkv", + "visual.blocks.21.attn.qkv.bias", + "visual.blocks.21.mlp.down_proj", + "visual.blocks.21.mlp.down_proj.bias", + "visual.blocks.21.mlp.gate_proj", + "visual.blocks.21.mlp.gate_proj.bias", + "visual.blocks.21.mlp.up_proj", + "visual.blocks.21.mlp.up_proj.bias", + "visual.blocks.21.norm1", + "visual.blocks.21.norm2", + "visual.blocks.22.attn.k_norm", + "visual.blocks.22.attn.proj", + "visual.blocks.22.attn.proj.bias", + "visual.blocks.22.attn.q_norm", + "visual.blocks.22.attn.qkv", + "visual.blocks.22.attn.qkv.bias", + "visual.blocks.22.mlp.down_proj", + "visual.blocks.22.mlp.down_proj.bias", + "visual.blocks.22.mlp.gate_proj", + "visual.blocks.22.mlp.gate_proj.bias", + "visual.blocks.22.mlp.up_proj", + "visual.blocks.22.mlp.up_proj.bias", + "visual.blocks.22.norm1", + "visual.blocks.22.norm2", + "visual.blocks.23.attn.k_norm", + "visual.blocks.23.attn.proj", + "visual.blocks.23.attn.proj.bias", + "visual.blocks.23.attn.q_norm", + "visual.blocks.23.attn.qkv", + "visual.blocks.23.attn.qkv.bias", + "visual.blocks.23.mlp.down_proj", + "visual.blocks.23.mlp.down_proj.bias", + "visual.blocks.23.mlp.gate_proj", + "visual.blocks.23.mlp.gate_proj.bias", + "visual.blocks.23.mlp.up_proj", + "visual.blocks.23.mlp.up_proj.bias", + "visual.blocks.23.norm1", + "visual.blocks.23.norm2", + "visual.blocks.3.attn.k_norm", + "visual.blocks.3.attn.proj", + "visual.blocks.3.attn.proj.bias", + "visual.blocks.3.attn.q_norm", + "visual.blocks.3.attn.qkv", + "visual.blocks.3.attn.qkv.bias", + "visual.blocks.3.mlp.down_proj", + "visual.blocks.3.mlp.down_proj.bias", + "visual.blocks.3.mlp.gate_proj", + "visual.blocks.3.mlp.gate_proj.bias", + "visual.blocks.3.mlp.up_proj", + "visual.blocks.3.mlp.up_proj.bias", + "visual.blocks.3.norm1", + "visual.blocks.3.norm2", + "visual.blocks.4.attn.k_norm", + "visual.blocks.4.attn.proj", + "visual.blocks.4.attn.proj.bias", + "visual.blocks.4.attn.q_norm", + "visual.blocks.4.attn.qkv", + "visual.blocks.4.attn.qkv.bias", + "visual.blocks.4.mlp.down_proj", + "visual.blocks.4.mlp.down_proj.bias", + "visual.blocks.4.mlp.gate_proj", + "visual.blocks.4.mlp.gate_proj.bias", + "visual.blocks.4.mlp.up_proj", + "visual.blocks.4.mlp.up_proj.bias", + "visual.blocks.4.norm1", + "visual.blocks.4.norm2", + "visual.blocks.5.attn.k_norm", + "visual.blocks.5.attn.proj", + "visual.blocks.5.attn.proj.bias", + "visual.blocks.5.attn.q_norm", + "visual.blocks.5.attn.qkv", + "visual.blocks.5.attn.qkv.bias", + "visual.blocks.5.mlp.down_proj", + "visual.blocks.5.mlp.down_proj.bias", + "visual.blocks.5.mlp.gate_proj", + "visual.blocks.5.mlp.gate_proj.bias", + "visual.blocks.5.mlp.up_proj", + "visual.blocks.5.mlp.up_proj.bias", + "visual.blocks.5.norm1", + "visual.blocks.5.norm2", + "visual.blocks.6.attn.k_norm", + "visual.blocks.6.attn.proj", + "visual.blocks.6.attn.proj.bias", + "visual.blocks.6.attn.q_norm", + "visual.blocks.6.attn.qkv", + "visual.blocks.6.attn.qkv.bias", + "visual.blocks.6.mlp.down_proj", + "visual.blocks.6.mlp.down_proj.bias", + "visual.blocks.6.mlp.gate_proj", + "visual.blocks.6.mlp.gate_proj.bias", + "visual.blocks.6.mlp.up_proj", + "visual.blocks.6.mlp.up_proj.bias", + "visual.blocks.6.norm1", + "visual.blocks.6.norm2", + "visual.blocks.7.attn.k_norm", + "visual.blocks.7.attn.proj", + "visual.blocks.7.attn.proj.bias", + "visual.blocks.7.attn.q_norm", + "visual.blocks.7.attn.qkv", + "visual.blocks.7.attn.qkv.bias", + "visual.blocks.7.mlp.down_proj", + "visual.blocks.7.mlp.down_proj.bias", + "visual.blocks.7.mlp.gate_proj", + "visual.blocks.7.mlp.gate_proj.bias", + "visual.blocks.7.mlp.up_proj", + "visual.blocks.7.mlp.up_proj.bias", + "visual.blocks.7.norm1", + "visual.blocks.7.norm2", + "visual.blocks.8.attn.k_norm", + "visual.blocks.8.attn.proj", + "visual.blocks.8.attn.proj.bias", + "visual.blocks.8.attn.q_norm", + "visual.blocks.8.attn.qkv", + "visual.blocks.8.attn.qkv.bias", + "visual.blocks.8.mlp.down_proj", + "visual.blocks.8.mlp.down_proj.bias", + "visual.blocks.8.mlp.gate_proj", + "visual.blocks.8.mlp.gate_proj.bias", + "visual.blocks.8.mlp.up_proj", + "visual.blocks.8.mlp.up_proj.bias", + "visual.blocks.8.norm1", + "visual.blocks.8.norm2", + "visual.blocks.9.attn.k_norm", + "visual.blocks.9.attn.proj", + "visual.blocks.9.attn.proj.bias", + "visual.blocks.9.attn.q_norm", + "visual.blocks.9.attn.qkv", + "visual.blocks.9.attn.qkv.bias", + "visual.blocks.9.mlp.down_proj", + "visual.blocks.9.mlp.down_proj.bias", + "visual.blocks.9.mlp.gate_proj", + "visual.blocks.9.mlp.gate_proj.bias", + "visual.blocks.9.mlp.up_proj", + "visual.blocks.9.mlp.up_proj.bias", + "visual.blocks.9.norm1", + "visual.blocks.9.norm2", + "visual.downsample", + "visual.downsample.bias", + "visual.merger.down_proj", + "visual.merger.gate_proj", + "visual.merger.post_projection_norm", + "visual.merger.post_projection_norm.bias", + "visual.merger.proj", + "visual.merger.up_proj", + "visual.patch_embed.proj", + "visual.patch_embed.proj.bias", + "visual.post_layernorm", + "weights_proj" + ], + "quant_method": "fp8", + "weight_block_size": [ + 128, + 128 + ] + } +} diff --git a/tools/probes/g3_miles/harness.py b/tools/probes/g3_miles/harness.py new file mode 100644 index 00000000..84cf8cbf --- /dev/null +++ b/tools/probes/g3_miles/harness.py @@ -0,0 +1,789 @@ +"""S25 T-G3 GPU harness: real Miles/Megatron glm5_next (4-layer mini) on 2 GPUs, one torchrun per phase. + +Phases (``--phase``): + +* ``n75`` -- N7.5: TP2+SP, EP2. Builds the model with the TE spec (``--impl te``) or the TE-free + fallback (``--impl local``), loads the weights saved by the first TE run (names mapped), + runs one forward/backward on a fixed packed sequence, saves the gathered logits, the + label log-probs and every parameter gradient, and times forward+backward (whole step and + the routed-expert module). Then runs N8.2 on the same weights (``n82`` below). +* ``n71d`` -- N7.1d: TP1, DP2, distributed optimizer + precision-aware CPU offload (real + HybridDeviceOptimizer: CUDA streams, pinned host memory, d2h/h2d overlap). Sub-modes + ``cont5`` (5 steps), ``run3`` (3 steps, then the yeto cut plugin saves a cut) and + ``resume2`` (fresh process: the cut plugin restores, N7.6 assertions, 2 steps). + +N8.2 (inside ``n75``): the reference gathers every weight to HF names through Miles' own weight-update +path (``named_params_and_buffers`` -> ``all_gather_params_async`` -> ``convert_to_hf``) and quantizes it +in one process; the distributed path gives each rank's experts through the yeto Megatron adapter +(``megatron_expert_slices``) and the non-expert weights once (rank 0, ``hf_slices``) to +``requantize_distributed``. Both outputs are compared byte for byte, for W4A8 (R2 AWQ) and MXFP4. + +Writes everything under ``--out``. Not a test: run by ``modal_g3.py`` on Modal ``H100!:2``. +""" + +from __future__ import annotations + +import argparse +import hashlib +import json +import os +import re +import sys +import time +import traceback +from pathlib import Path + +import torch +import torch.distributed as dist + +HERE = Path(__file__).resolve().parent +MILES = Path(os.environ.get("G3_MILES", "/g3/miles")) +sys.path.insert(0, str(MILES / "scripts" / "models")) +sys.path.insert(0, str(MILES / "miles" / "utils" / "external_utils")) # model_args_utils + +SEQ = 1024 # tokens of the packed sequence (one sample) +TIMING_SEQ = 4096 +VOCAB = 16384 +HIDDEN = 1024 +EXPERTS = 8 + + +# ---------------------------------------------------------------------------------------------- +# Mini HF config (the 45-layer fixture shrunk: 4 layers = 3 KDA + 1 DSA, 1 dense + 3 MoE) +# ---------------------------------------------------------------------------------------------- + +def write_hf_config(src: Path, dst: Path) -> None: + cfg = json.loads(src.read_text()) + cfg.pop("quantization_config", None) + t = cfg["text_config"] + layer_types = ["linear_attention", "linear_attention", "linear_attention", "deepseek_sparse_attention"] + t.update( + hidden_size=HIDDEN, intermediate_size=3 * HIDDEN, moe_intermediate_size=512, num_hidden_layers=4, + first_k_dense_replace=1, n_routed_experts=EXPERTS, num_experts_per_tok=2, vocab_size=VOCAB, + index_topk=256, layer_types=layer_types, indexer_types=["full"] * 4, + mlp_layer_types=["dense", "sparse", "sparse", "sparse"], num_nextn_predict_layers=0, pad_token_id=0, + ) + t["linear_attn_config"] = dict(t["linear_attn_config"], num_heads=16, kda_layers=[0, 1, 2], full_attn_layers=[3]) + dst.mkdir(parents=True, exist_ok=True) + (dst / "config.json").write_text(json.dumps(cfg, indent=1)) + + +def megatron_argv(impl: str, *, tp: int, ep: int, hf: Path, offload: bool, lr: float) -> list[str]: + import importlib.util + + spec = importlib.util.spec_from_file_location("glm53_flash_args", MILES / "scripts" / "models" / "glm5.3-flash.py") + glm = importlib.util.module_from_spec(spec) + spec.loader.exec_module(glm) + model = glm.model_args(nlayers=4, first_k_dense_replace=1, + transformer_impl="transformer_engine" if impl == "te" else "local").split() + overrides = [ + "--hidden-size", str(HIDDEN), "--ffn-hidden-size", str(3 * HIDDEN), "--num-experts", str(EXPERTS), + "--moe-router-topk", "2", "--moe-ffn-hidden-size", "512", "--moe-shared-expert-intermediate-size", "512", + "--vocab-size", str(VOCAB), + ] + par = ["--tensor-model-parallel-size", str(tp), "--pipeline-model-parallel-size", "1", + "--context-parallel-size", "1", "--expert-model-parallel-size", str(ep), + "--expert-tensor-parallel-size", "1"] + if tp > 1: + par.append("--sequence-parallel") + misc = [ + "--qkv-format", "thd", "--micro-batch-size", "1", "--seq-length", str(TIMING_SEQ), + "--attention-dropout", "0.0", "--hidden-dropout", "0.0", "--tokenizer-type", "NullTokenizer", + "--hf-checkpoint", str(hf), "--load", "/nonexistent-g3", "--seed", "1234", + "--optimizer", "adam", "--lr", str(lr), "--lr-decay-style", "constant", "--min-lr", "0", + "--weight-decay", "0.1", "--adam-beta1", "0.9", "--adam-beta2", "0.98", "--clip-grad", "1.0", + "--attention-backend", "auto", "--no-check-for-nan-in-loss-and-grad", "--distributed-timeout-minutes", "10", + ] + if offload: + misc += ["--optimizer-cpu-offload", "--use-precision-aware-optimizer", "--overlap-cpu-optimizer-d2h-h2d", + "--optimizer-offload-fraction", "1.0"] + return model + overrides + par + misc + + +def init_miles(argv: list[str], *, global_batch: int, rollouts: int): + from megatron.training.arguments import parse_args, validate_args + from miles.backends.megatron_utils.arguments import set_default_megatron_args + from miles.backends.megatron_utils.initialize import init + from miles.utils.arguments import get_miles_extra_args_provider, hf_validate_args + from miles.utils.hf_utils.config import load_hf_config + + sys.argv = ["g3"] + argv + ["--global-batch-size", str(global_batch)] + args = parse_args(extra_args_provider=get_miles_extra_args_provider()) + hf_config = load_hf_config(args.hf_checkpoint) + hf_validate_args(args, hf_config) + text = getattr(hf_config, "text_config", None) or hf_config + args.indexer_rope_interleave = bool(getattr(text, "indexer_rope_interleave", False)) + args.compress_ratios = None + # Miles' scheduler sizing (get_optimizer_param_scheduler): train_iters = rollouts*rbs*n // gbs + args.num_rollout, args.rollout_batch_size, args.n_samples_per_prompt = rollouts, 1, global_batch + args.__dict__.setdefault("decrease_batch_size_if_needed", False) + set_default_megatron_args(args) + dist.init_process_group(backend="nccl") + torch.cuda.set_device(int(os.environ.get("LOCAL_RANK", 0))) + validate_args(args) + init(args) + return args + + +def build(args): + from miles.backends.megatron_utils.model import setup_model_and_optimizer, setup_train_iteration_config + + model, optimizer, scheduler = setup_model_and_optimizer(args, "actor") + setup_train_iteration_config(args, model, optimizer, False) + return model, optimizer, scheduler + + +# ---------------------------------------------------------------------------------------------- +# Data, forward/backward +# ---------------------------------------------------------------------------------------------- + +def tokens_for(step: int, dp_rank: int, n: int) -> torch.Tensor: + g = torch.Generator().manual_seed(1000 * step + dp_rank) + return torch.randint(0, VOCAB, (n,), generator=g) + + +def packed(tokens: torch.Tensor): + from megatron.core.packed_seq_params import PackedSeqParams + + n = tokens.numel() + cu = torch.tensor([0, n], dtype=torch.int32, device="cuda") + psp = PackedSeqParams(cu_seqlens_q=cu, cu_seqlens_kv=cu, max_seqlen_q=n, max_seqlen_kv=n, qkv_format="thd") + return tokens.view(1, -1).cuda(), psp + + +def run_fwd_bwd(model, tokens: torch.Tensor, *, capture: list | None = None, forward_only=False): + """One micro-batch through Megatron's schedule; loss = mean next-token CE over the sequence.""" + from megatron.core import tensor_parallel + from megatron.core.pipeline_parallel import get_forward_backward_func + + ids, psp = packed(tokens) + labels = torch.roll(ids, -1, dims=1) + mask = torch.ones_like(ids, dtype=torch.float32) + mask[:, -1] = 0 + + def loss_func(output): + if capture is not None: + logits = output.detach().float().clone() # vocab_parallel_cross_entropy edits its input in place + token_loss = tensor_parallel.vocab_parallel_cross_entropy(output.float(), labels) + loss = (token_loss * mask).sum() / mask.sum() + if capture is not None: + capture.append((logits, token_loss.detach())) + return loss, {"loss": loss.detach()} + + def step(data_iter, chunk): + next(data_iter) + out = chunk(input_ids=ids, position_ids=None, attention_mask=None, labels=None, packed_seq_params=psp) + return out, loss_func + + fb = get_forward_backward_func() + res = fb(forward_step_func=step, data_iterator=iter([None]), model=model, num_microbatches=1, + seq_length=ids.shape[1], micro_batch_size=1, forward_only=forward_only) + return float(res[0]["loss"]) if res else float("nan") + + +def zero(model, optimizer): + for chunk in model: + chunk.zero_grad_buffer() + if optimizer is not None: + optimizer.zero_grad() + + +# ---------------------------------------------------------------------------------------------- +# Names: TE spec <-> local spec +# ---------------------------------------------------------------------------------------------- + +def named(model) -> dict[str, torch.Tensor]: + out = {} + for n, p in model[0].named_parameters(): + out[n] = p + for n, b in model[0].named_buffers(): + if "expert_bias" in n: + out[n] = b + return out + + +def local_to_te(name: str) -> str: + from miles.backends.megatron_utils.named_weights import local_expert_name_to_grouped + + name = te_canonical(local_expert_name_to_grouped(name), set()) + # dense MLP: TE fuses the pre-MLP RMSNorm into linear_fc1 (layer_norm_weight) + name = re.sub(r"(decoder\.layers\.\d+)\.pre_mlp_layernorm\.weight$", r"\1.PRE_MLP_NORM", name) + return name + + +def te_canonical(name: str, te_names: set[str]) -> str: + return re.sub(r"(decoder\.layers\.\d+)\.mlp\.linear_fc1\.layer_norm_weight$", r"\1.PRE_MLP_NORM", name) + + +def map_local_to_te(local_names, te_names) -> tuple[dict[str, str], list[str], list[str]]: + canon_te = {te_canonical(n, te_names): n for n in te_names} + mapping, missing = {}, [] + for n in local_names: + c = local_to_te(n) + if c in canon_te: + mapping[n] = canon_te[c] + elif c.endswith("PRE_MLP_NORM") and c.replace("PRE_MLP_NORM", "pre_mlp_layernorm.weight") in canon_te: + mapping[n] = canon_te[c.replace("PRE_MLP_NORM", "pre_mlp_layernorm.weight")] # MoE layers keep the name + else: + missing.append(n) + unused = sorted(set(te_names) - set(mapping.values())) + return mapping, missing, unused + + +# ---------------------------------------------------------------------------------------------- +# N7.5 phase +# ---------------------------------------------------------------------------------------------- + +def rank_file(out: Path, stem: str) -> Path: + return out / f"{stem}.rank{dist.get_rank()}.pt" + + +def log(out: Path, msg: str) -> None: + line = f"[{time.strftime('%H:%M:%S')}] rank{dist.get_rank() if dist.is_initialized() else '-'} {msg}" + print(line, flush=True) + if not dist.is_initialized() or dist.get_rank() == 0: + with open(out / "progress.log", "a") as fh: + fh.write(line + "\n") + + +def phase_n75(ns) -> None: + out = Path(ns.out) + out.mkdir(parents=True, exist_ok=True) + args = init_miles(megatron_argv(ns.impl, tp=2, ep=2, hf=Path(ns.hf), offload=False, lr=1e-5), + global_batch=1, rollouts=1) + log(out, f"n75 impl={ns.impl} tag={ns.tag} transformer_impl={args.transformer_impl}") + model, optimizer, _ = build(args) + params = named(model) + info = {"impl": ns.impl, "tag": ns.tag, "transformer_impl": args.transformer_impl, + "module_classes": sorted({type(m).__name__ for m in model[0].modules()}), + "num_params_local": sum(p.numel() for p in params.values()), + "gpu": torch.cuda.get_device_name(), "names": sorted(params)} + core = [type(m).__name__ for n, m in model[0].named_modules() if n.endswith("self_attention.core_attention")] + info["dsa_core_attention"] = core + weights_file = Path(ns.weights_from) / f"weights.rank{dist.get_rank()}.pt" if ns.weights_from else None + if weights_file is None: + torch.save({n: p.detach().cpu() for n, p in params.items()}, rank_file(out, "weights")) + info["weights"] = "initialized here (seed 1234)" + else: + saved = torch.load(weights_file, map_location="cpu", weights_only=True) + mapping, missing, unused = map_local_to_te(params, set(saved)) + info["name_map"] = {"mapped": len(mapping), "missing": missing, "unused_te": unused, + "renamed": {k: v for k, v in mapping.items() if k != v}} + if missing or unused: + raise RuntimeError(f"weight names do not map: missing={missing[:8]} unused={unused[:8]}") + with torch.no_grad(): + for n, p in params.items(): + src = saved[mapping[n]] + if tuple(src.shape) != tuple(p.shape): + raise RuntimeError(f"{n}: shape {tuple(p.shape)} != saved {tuple(src.shape)}") + p.copy_(src.to(p.device, p.dtype)) + # the DDP/fp32 main copies are rebuilt from the model weights by the optimizer + if optimizer is not None and hasattr(optimizer, "reload_model_params"): + optimizer.reload_model_params() + info["weights"] = f"loaded from {ns.weights_from}" + + # one forward/backward (no optimizer step) + zero(model, optimizer) + cap: list = [] + loss = run_fwd_bwd(model, tokens_for(0, 0, SEQ), capture=cap) + torch.cuda.synchronize() + logits_local, token_loss = cap[0] + from megatron.core import parallel_state as mps + from megatron.core.tensor_parallel.mappings import gather_from_tensor_model_parallel_region + + logits = gather_from_tensor_model_parallel_region(logits_local.float()) # [1, T, V] + grads = {} + for n, p in params.items(): + g = getattr(p, "main_grad", None) + if g is None: + g = p.grad + if g is not None: + grads[n] = g.detach().float().cpu() + torch.save(grads, rank_file(out, f"grads-{ns.tag}")) + if dist.get_rank() == 0: + torch.save({"logits": logits[0].cpu(), "token_loss": token_loss.float().cpu().view(-1), "loss": loss}, + out / f"logits-{ns.tag}.pt") + info["loss"] = loss + info["grad_tensors"] = len(grads) + info["tp_rank"], info["ep_rank"] = mps.get_tensor_model_parallel_rank(), mps.get_expert_model_parallel_rank() + log(out, f"n75 {ns.tag}: loss={loss:.6f} grads={len(grads)}") + + # timing: forward+backward on TIMING_SEQ tokens; routed experts timed with CUDA events + events: list = [] + + def pre(mod, inp): + e = torch.cuda.Event(enable_timing=True) + e.record() + mod._g3_t0 = e + + def post(mod, inp, outp): + e = torch.cuda.Event(enable_timing=True) + e.record() + events.append((mod._g3_t0, e)) + + hooks = [] + for n, m in model[0].named_modules(): + if n.endswith("mlp.experts"): + hooks += [m.register_forward_pre_hook(pre), m.register_forward_hook(post)] + steps = [] + for i in range(7): + zero(model, optimizer) + events.clear() + torch.cuda.synchronize() + t0 = time.perf_counter() + run_fwd_bwd(model, tokens_for(100 + i, 0, TIMING_SEQ)) + torch.cuda.synchronize() + steps.append({"step_s": time.perf_counter() - t0, + "experts_fwd_ms": sum(a.elapsed_time(b) for a, b in events)}) + for h in hooks: + h.remove() + info["timing"] = {"tokens": TIMING_SEQ, "warmup": 2, "steps": steps, + "peak_mem_gib": torch.cuda.max_memory_allocated() / 2**30} + zero(model, optimizer) + + pre = [None] * dist.get_world_size() + dist.all_gather_object(pre, info) + if dist.get_rank() == 0: + (out / f"info-{ns.tag}.json").write_text(json.dumps(pre, indent=1, default=str)) + # N8.2 on the same (unchanged) weights + try: + info["n82"] = n82(args, model, out / f"n82-{ns.tag}") + except Exception as exc: # noqa: BLE001 -- recorded; the phase result says what failed + info["n82"] = {"error": f"{type(exc).__name__}: {exc}", "traceback": traceback.format_exc()[-4000:]} + log(out, f"n82 {ns.tag} FAILED: {exc}") + all_info = [None] * dist.get_world_size() + dist.all_gather_object(all_info, info) + if dist.get_rank() == 0: + (out / f"info-{ns.tag}.json").write_text(json.dumps(all_info, indent=1, default=str)) + dist.barrier() + + +# ---------------------------------------------------------------------------------------------- +# N8.2 +# ---------------------------------------------------------------------------------------------- + +def nest_language_model(name: str) -> str: + """GLM-5.3 checkpoint names keep the text model under ``model.language_model.`` (mbridge glm5_next).""" + if name.startswith("model.") and not name.startswith(("model.language_model.", "model.visual.")): + return "model.language_model." + name[len("model."):] + return name + + +EXPERT_HF = re.compile(r"\.mlp\.experts\.\d+\.(gate|up|down)_proj\.weight$") + + +def miles_hf_weights(args, model) -> tuple[dict[str, torch.Tensor], dict[str, torch.Tensor]]: + """Miles' weight-update path: global names, TP all-gather, convert_to_hf. Returns (non-expert on every + rank, this EP rank's experts), HF names as Miles emits them (no nesting).""" + from miles.backends.megatron_utils.megatron_to_hf import convert_to_hf + from miles.backends.megatron_utils.named_weights import named_params_and_buffers + from miles.backends.megatron_utils.update_weight.hf_weight_iterator_direct import ( + all_gather_params_async, + is_routed_expert_param, + ) + from miles.utils.types import ParamInfo + + dense, experts = {}, {} + for name, param in named_params_and_buffers(args, model): + pinfo = ParamInfo(name=name, dtype=param.dtype, shape=param.shape, attrs={}, size=0, src_rank=dist.get_rank()) + full = all_gather_params_async(args, [(pinfo, param)])[0] + for hf_name, t in convert_to_hf(args, "glm5_next", name, full): + (experts if is_routed_expert_param(name) else dense)[hf_name] = t.detach().cpu().clone() + return dense, experts + + +def fake_original(source: dict[str, torch.Tensor], scheme: str, out: Path, smooth: str | None) -> Path: + """An "original" quantized checkpoint (headers matter, values do not) with the source's names.""" + from safetensors.torch import save_file + + out.mkdir(parents=True, exist_ok=True) + tensors = {} + for n, t in source.items(): + if EXPERT_HF.search(n): + mod = n[: -len(".weight")] + rows, cols = t.shape + if scheme == "w4a8": + tensors[n] = torch.zeros(rows, cols, dtype=torch.int8) + tensors[mod + ".weight_scale"] = torch.ones(rows, 1, dtype=torch.bfloat16) + else: + tensors[mod + ".weight_packed"] = torch.zeros(rows, cols // 2, dtype=torch.uint8) + tensors[mod + ".weight_scale"] = torch.zeros(rows, cols // 32, dtype=torch.uint8) + elif n.endswith("mlp.gate.weight"): + tensors[n] = t.float() # the router is stored in F32 + else: + tensors[n] = t + shard = lambda n: ("model-00002-of-00002.safetensors" if re.search(r"layers\.[23]\.mlp\.experts", n) # noqa: E731 + else "model-00001-of-00002.safetensors") + files: dict[str, dict] = {} + for n, t in tensors.items(): + files.setdefault(shard(n), {})[n] = t.contiguous() + for f, ts in files.items(): + save_file(ts, str(out / f), metadata={"format": "pt"}) + (out / "model.safetensors.index.json").write_text(json.dumps({"metadata": {}, "weight_map": {n: shard(n) for n in tensors}})) + if scheme == "w4a8": + q = {"quant_method": "compressed-tensors", "format": "int-quantized", "config_groups": {"group_0": { + "targets": ["re:.*mlp\\.experts\\..*(gate|up|down)_proj$"], + "weights": {"num_bits": 4, "type": "int", "symmetric": True, "strategy": "channel", "group_size": None, + "dynamic": False, "observer": "memoryless_minmax"}, + "input_activations": {"num_bits": 8, "strategy": "token", "dynamic": True}}}} + else: + q = {"quant_method": "compressed-tensors", "format": "mxfp4-pack-quantized", "config_groups": {"MXFP4A16": { + "format": "mxfp4-pack-quantized", "targets": ["re:.*mlp\\.experts\\.\\d+\\.(gate|up|down)_proj$"], + "weights": {"num_bits": 4, "type": "float", "symmetric": True, "strategy": "group", "group_size": 32, + "dynamic": False, "scale_dtype": "torch.uint8", "observer": None}, + "input_activations": None}}, "ignore": ["lm_head", "re:^(?!.*mlp\\.experts\\.\\d+\\.).*$"]} + (out / "config.json").write_text(json.dumps({"model_type": "glm5_next", "quantization_config": q})) + if scheme == "w4a8" and smooth: + layer = smooth.rsplit(".post_attention_layernorm", 1)[0] + esc = re.escape(layer) + (out / "recipe.yaml").write_text( + "default_stage:\n default_modifiers:\n RecordedAWQ:\n mappings:\n" + f" - smooth_layer: {smooth}\n" + f" balance_layers: [{layer}.mlp.gate, 're:{esc}\\.mlp\\.(experts\\.\\d+|shared_experts)\\.(gate|up)_proj$']\n" + " duo_scaling: true\n QuantizationModifier:\n ignore: [lm_head]\n") + return out + + +def prepare_reference(source: dict[str, torch.Tensor], scheme: str, d: Path, smooth: str | None) -> dict: + """Rank 0: template, BF16 source, AWQ factors and the single-process reference output under ``d``.""" + from safetensors.torch import save_file + + from yeto.rl.quant.awq import AwqSmoothing + from yeto.rl.quant.requant import CheckpointLayout, requantize_checkpoint + + orig = fake_original(source, scheme, d / "original", smooth) + CheckpointLayout.from_checkpoint(orig).save_template(orig, d / "template") + (d / "bf16").mkdir(parents=True, exist_ok=True) + save_file({n: t.contiguous() for n, t in source.items()}, str(d / "bf16" / "model.safetensors"), + metadata={"format": "pt"}) + if scheme == "w4a8" and smooth: + (d / "awq-scales").mkdir(exist_ok=True) + factors = (0.8 + 0.4 * torch.rand(source[smooth + ".weight"].numel(), + generator=torch.Generator().manual_seed(7))).double() + torch.save(factors, d / "awq-scales" / f"{smooth}.pt") + smoothing = (AwqSmoothing.from_artefact(d / "original" / "recipe.yaml", d / "awq-scales") + if scheme == "w4a8" and smooth else AwqSmoothing.none()) + t0 = time.perf_counter() + requantize_checkpoint(d / "bf16", d / "template", d / "reference", smoothing) + return {"single_process_s": time.perf_counter() - t0} + + +def n82(args, model, out: Path) -> dict: + from megatron.core import parallel_state as mps + + from yeto.rl.adapters.miles.quant_shards import hf_slices, megatron_expert_slices + from yeto.rl.quant.awq import AwqSmoothing + from yeto.rl.quant.shards import SourceSlice, requantize_distributed + + rank, world = dist.get_rank(), dist.get_world_size() + ep_rank, ep_size = mps.get_expert_model_parallel_rank(), mps.get_expert_model_parallel_world_size() + if mps.get_data_parallel_world_size() != 1: + raise RuntimeError("n82 harness expects DP=1 (every rank holds a distinct EP slice)") + out.mkdir(parents=True, exist_ok=True) + dense, experts = miles_hf_weights(args, model) + gathered = [None] * world + dist.all_gather_object(gathered, experts) + res: dict = {"miles_hf_names_sample": sorted(dense)[:6] + sorted(experts)[:3], + "dense_tensors": len(dense), "expert_tensors_this_rank": len(experts)} + # reference source: nested names (the checkpoint layout); every rank builds it (same bytes) + source = {nest_language_model(n): t for n, t in dense.items()} + for part in gathered: + source.update({nest_language_model(n): t for n, t in part.items()}) + smooth = next((n[: -len(".weight")] for n in source if n.endswith("layers.1.post_attention_layernorm.weight")), None) + res["awq_smooth_layer"] = smooth + local = {n: p for n, p in model[0].named_parameters()} + num_local = args.num_experts // ep_size + for scheme in ("w4a8", "mxfp4"): + d = out / scheme + r: dict = {} + try: + err = None + if rank == 0: + try: + r.update(prepare_reference(source, scheme, d, smooth)) + except Exception as exc: # noqa: BLE001 -- re-raised below after telling the peer + err = exc + if err is not None: + signal_failure() + raise err + checked_barrier() # raises PeerFailed on rank 1 when rank 0 failed above + smoothing = (AwqSmoothing.from_artefact(d / "original" / "recipe.yaml", d / "awq-scales") + if scheme == "w4a8" and smooth else AwqSmoothing.none()) + r["awq"] = bool(scheme == "w4a8" and smooth) + # distributed: experts from the yeto adapter on Megatron-local names; non-experts once (rank 0) + mine = list(megatron_expert_slices(((n.removeprefix("module.").removeprefix("module."), p) + for n, p in local.items()), + ep_rank=ep_rank, num_local_experts=num_local)) + mine_cpu = [SourceSlice(x.name, x.tensor.cpu(), x.row_start) for x in mine] + r["expert_slices_this_rank"] = len(mine) + r["expert_slice_fields"] = sorted(getattr(mine[0], "__dataclass_fields__", {})) if mine else [] + # raw: convert_to_hf names as they are; yeto: the adapter's hf_slices (renames to checkpoint + # names), experts first as the GPU tensors a trainer holds, then copied to the host + for variant in ("raw_convert_to_hf_names", "yeto_hf_slices_gpu_experts", "yeto_hf_slices_cpu_experts"): + dst = d / f"dist-{variant}" + slices = list(mine if variant == "yeto_hf_slices_gpu_experts" else mine_cpu) + if rank == 0: + slices += (list(hf_slices(dense.items())) if variant.startswith("yeto_hf_slices") + else [SourceSlice(n, t) for n, t in dense.items()]) + try: + t0 = time.perf_counter() + stats = requantize_distributed(slices, d / "template", dst, smoothing, rank=rank, + world_size=world, barrier=checked_barrier) + r[variant] = {"ok": True, "s": time.perf_counter() - t0, "part": stats.get("part")} + except PeerFailed as exc: + r[variant] = {"ok": False, "error": f"PeerFailed: {exc}"} + except Exception as exc: # noqa: BLE001 + r[variant] = {"ok": False, "error": f"{type(exc).__name__}: {str(exc)[:600]}"} + signal_failure() # pairs with the peer's next checked_barrier + dist.barrier() + if rank == 0 and r[variant]["ok"]: + r[variant]["byte_equal"] = compare_dirs(d / "reference", dst) + except Exception as exc: # noqa: BLE001 + r["error"] = f"{type(exc).__name__}: {exc}" + r["traceback"] = traceback.format_exc()[-3000:] + res[scheme] = r + return res + + +class PeerFailed(RuntimeError): + pass + + +def checked_barrier() -> None: + """Barrier that also carries "a rank failed": every rank either reaches it or calls signal_failure().""" + flag = torch.zeros(1, device="cuda") + dist.all_reduce(flag, op=dist.ReduceOp.MAX) + if flag.item() > 0: + raise PeerFailed("another rank failed in requantize_distributed") + + +def signal_failure() -> None: + dist.all_reduce(torch.ones(1, device="cuda"), op=dist.ReduceOp.MAX) + + +def compare_dirs(a: Path, b: Path) -> dict: + names_a = sorted(p.name for p in a.iterdir() if p.is_file()) + names_b = sorted(p.name for p in b.iterdir() if p.is_file()) + diff = [n for n in names_a if n in names_b and (a / n).read_bytes() != (b / n).read_bytes()] + return {"files": names_a, "same_file_set": names_a == names_b, "differing": diff, + "equal": names_a == names_b and not diff, + "sha256": {n: hashlib.sha256((a / n).read_bytes()).hexdigest()[:16] for n in names_a}} + + +# ---------------------------------------------------------------------------------------------- +# N7.1d phase +# ---------------------------------------------------------------------------------------------- + +class Actor: + """The attributes of a Miles MegatronTrainRayActor that the yeto cut plugin reads.""" + + def __init__(self, args, model, optimizer, scheduler): + self.args, self.model, self.optimizer, self.opt_param_scheduler = args, model, optimizer, scheduler + + +def hdo_report(optimizer) -> dict: + try: + return _hdo_report(optimizer) + except Exception as exc: # noqa: BLE001 -- diagnostics only + return {"error": f"{type(exc).__name__}: {exc}"} + + +def _hdo_report(optimizer) -> dict: + from megatron.core.optimizer.cpu_offloading import HybridDeviceOptimizer + + leaves = getattr(optimizer, "chained_optimizers", [optimizer]) + rep = [] + for leaf in leaves: + inner = getattr(leaf, "optimizer", None) + item = {"leaf": type(leaf).__name__, "inner": type(inner).__name__, + "is_hdo": isinstance(inner, HybridDeviceOptimizer)} + if isinstance(inner, HybridDeviceOptimizer): + fp32 = list(inner.param_to_fp32_param.values()) if getattr(inner, "param_to_fp32_param", None) else [] + copies = list(getattr(inner, "gpu_params_map_cpu_copy", {}).values()) + item.update( + host_copies=len(copies), host_copies_pinned=sum(bool(c.is_pinned()) for c in copies), + host_copy_dtypes=sorted({str(c.dtype) for c in copies}), + fp32_params=len(fp32), fp32_devices=sorted({str(p.device) for p in fp32}), + fp32_pinned=sum(bool(p.is_pinned()) for p in fp32 if p.device.type == "cpu"), + fp32_dtypes=sorted({str(p.dtype) for p in fp32}), + cpu_optimizers=[type(o).__name__ for o in inner.cpu_optimizers], + overlap=bool(inner.defaults.get("overlap_cpu_optimizer_d2h_h2d")), + state_dtypes=sorted({str(v.dtype) for o in inner.cpu_optimizers for s in o.state.values() + for v in s.values() if torch.is_tensor(v) and v.numel() > 1}), + steps=sorted({float(s["step"]) for o in inner.cpu_optimizers for s in o.state.values() if "step" in s}), + ) + rep.append(item) + return {"leaves": rep} + + +def stale_report(optimizer) -> dict: + """Per HDO parameter: does the tensor the sub-optimizer steps equal the HDO's "master_param"? + (the old export read the master; N7.6c reads the stepped tensor)""" + from miles.backends.megatron_utils.lora import dp_invariant_state as dps + + stale, same = [], 0 + for leaf in getattr(optimizer, "chained_optimizers", [optimizer]): + hdo = getattr(leaf, "optimizer", None) + if not hasattr(hdo, "sub_optimizers"): + continue + for param, state in hdo.state.items(): + master = state.get("master_param") + if not torch.is_tensor(master): + continue + inner = dps._hdo_inner_param(hdo, param) + if torch.equal(inner.detach().cpu().reshape(-1), master.detach().cpu().reshape(-1)): + same += 1 + else: + stale.append({"shape": list(param.shape), "dtype": str(param.dtype)}) + return {"stepped_equals_master": same, "stepped_differs_from_master": len(stale), "differs": stale[:40]} + + +def param_digest(model) -> dict: + h = hashlib.sha256() + norms = {} + for n, p in sorted(model[0].named_parameters()): + t = p.detach().float().cpu().contiguous() + h.update(n.encode()) + h.update(p.detach().cpu().contiguous().view(torch.uint8).numpy().tobytes()) + norms[n] = float(t.norm()) + return {"sha256": h.hexdigest(), "norms": norms} + + +def train_steps(args, model, optimizer, scheduler, first: int, last: int, out: Path) -> list[dict]: + from megatron.core import parallel_state as mps + + dp = mps.get_data_parallel_rank() + rec = [] + for step in range(first, last + 1): + zero(model, optimizer) + t0 = time.perf_counter() + loss = run_fwd_bwd(model, tokens_for(step, dp, SEQ)) + ok, grad_norm, _ = optimizer.step() + assert ok + scheduler.step(increment=args.global_batch_size) + torch.cuda.synchronize() + lt = torch.tensor([loss], device="cuda") + dist.all_reduce(lt) + rec.append({"step": step, "loss_mean_dp": float(lt) / dist.get_world_size(), "loss_rank": loss, + "grad_norm": float(grad_norm), "lr": [g["lr"] for g in optimizer.param_groups], + "scheduler_num_steps": int(scheduler.num_steps), "s": time.perf_counter() - t0}) + log(out, f"step {step} loss={loss:.6f} grad_norm={float(grad_norm):.6f}") + zero(model, optimizer) + return rec + + +def phase_n71d(ns) -> None: + from yeto.rl.adapters.miles import cut_plugin + from yeto.rl.engine.resume_manifest import ( + OptimizerResumeState, + ResumeManifest, + first_step_problems, + require_saveable, + restore_problems, + ) + + out = Path(ns.out) + out.mkdir(parents=True, exist_ok=True) + args = init_miles(megatron_argv("te", tp=1, ep=1, hf=Path(ns.hf), offload=True, lr=1e-4), + global_batch=2, rollouts=10) + flags = {k: getattr(args, k, None) for k in ( + "use_distributed_optimizer", "optimizer_cpu_offload", "use_precision_aware_optimizer", + "overlap_cpu_optimizer_d2h_h2d", "optimizer_offload_fraction", "main_params_dtype", "main_grads_dtype", + "exp_avg_dtype", "exp_avg_sq_dtype", "data_parallel_size", "bf16", "transformer_impl")} + log(out, f"n71d {ns.mode} flags={flags}") + model, optimizer, scheduler = build(args) + actor = Actor(args, model, optimizer, scheduler) + res: dict = {"mode": ns.mode, "flags": {k: str(v) for k, v in flags.items()}, "gpu": torch.cuda.get_device_name()} + res["hdo_before"] = hdo_report(optimizer) + cut_dir = Path(ns.cut_dir) + + def gather(obj): + lst = [None] * dist.get_world_size() + dist.all_gather_object(lst, obj) + return lst + + def manifest(summaries, global_step): + return ResumeManifest(global_step=global_step, data_watermark=global_step * args.global_batch_size, + scheduler_increment=args.global_batch_size, + ranks={str(s["path"]): OptimizerResumeState.from_dict(s["resume"]) for s in summaries}) + + if ns.mode.startswith("cont5"): + res["steps"] = train_steps(args, model, optimizer, scheduler, 1, 5, out) + elif ns.mode.startswith("run3"): + res["steps"] = train_steps(args, model, optimizer, scheduler, 1, 3, out) + summary = cut_plugin.save_cut_shard(actor, directory=str(cut_dir), cut_id="g3-cut-3") + if cut_plugin.REFUSED in summary: + raise RuntimeError(f"save_cut refused: {summary}") + summaries = gather(summary) + res["cut_summary"] = {k: summary[k] for k in ("path", "bytes", "scheduler_samples", "tuning", + "adapter_tensors", "resume")} + if dist.get_rank() == 0: + m = manifest(summaries, 3) + require_saveable(m) + (cut_dir / "files.json").write_text(json.dumps(summaries, default=str)) + (cut_dir / "resume_manifest.json").write_text(json.dumps(m.to_dict())) + res["hdo_after_save"] = hdo_report(optimizer) + elif ns.mode.startswith("resume2"): + files = json.loads((cut_dir / "files.json").read_text()) + saved = ResumeManifest.from_dict(json.loads((cut_dir / "resume_manifest.json").read_text())) + inject = ns.mode.endswith("inj") + if inject: # T-C4 N7.6c: put the T-G3 bug back (the restore leaves the HDO stepping its stale host copy) + from miles.backends.megatron_utils.lora import dp_invariant_state as dps + + dps._bind_masters_to_stepped_tensors = lambda hdo: None + summary = cut_plugin.restore_cut_shard(actor, directory=str(cut_dir), files=files, cut_id="g3-cut-3") + if cut_plugin.REFUSED in summary: + raise RuntimeError(f"restore_cut refused: {summary}") + if inject: + res["injected"] = "dp_invariant_state._bind_masters_to_stepped_tensors = no-op" + res["stale_vs_master"] = stale_report(optimizer) + res["restore_summary"] = {k: summary.get(k) for k in ("path", "scheduler_samples", "resume", "diff", + "optimizer_cut_vs_reexport")} + res["restore_digest_equal"] = summary.get("state_digest") == next( + f for f in files if f["path"] == summary["path"])["state_digest"] + res["hdo_after_restore"] = hdo_report(optimizer) + now = manifest(gather(summary), 3) + res["restore_problems"] = restore_problems(saved, now, check_data_watermark=True) + res["steps"] = train_steps(args, model, optimizer, scheduler, 4, 4, out) + st = cut_plugin.resume_state(actor) + after1 = manifest(gather(st), 4) + res["first_step_problems"] = first_step_problems(saved, after1, steps=1) + res["hdo_after_first_step"] = hdo_report(optimizer) + res["steps"] += train_steps(args, model, optimizer, scheduler, 5, 5, out) + else: + raise ValueError(ns.mode) + res["final_params"] = param_digest(model) + res["hdo_final"] = hdo_report(optimizer) + res["peak_mem_gib"] = torch.cuda.max_memory_allocated() / 2**30 + allres = gather(res) + if dist.get_rank() == 0: + (out / f"n71d-{ns.mode}.json").write_text(json.dumps(allres, indent=1, default=str)) + dist.barrier() + + +def main() -> None: + p = argparse.ArgumentParser() + p.add_argument("--phase", required=True, choices=["hfconfig", "n75", "n71d"]) + p.add_argument("--out", required=True) + p.add_argument("--hf", default="/tmp/g3-hf") + p.add_argument("--impl", default="te", choices=["te", "local"]) + p.add_argument("--tag", default="te") + p.add_argument("--weights-from", default="") + p.add_argument("--mode", default="cont5") + p.add_argument("--cut-dir", default="") + p.add_argument("--fixture", default=str(HERE / "glm53flash_config.json")) + ns = p.parse_args() + if ns.phase == "hfconfig": + write_hf_config(Path(ns.fixture), Path(ns.hf)) + return + try: + (phase_n75 if ns.phase == "n75" else phase_n71d)(ns) + except Exception: + traceback.print_exc() + sys.stdout.flush() + os._exit(1) + dist.destroy_process_group() + + +if __name__ == "__main__": + main() diff --git a/tools/probes/g3_miles/modal_g3.py b/tools/probes/g3_miles/modal_g3.py new file mode 100644 index 00000000..eccadf66 --- /dev/null +++ b/tools/probes/g3_miles/modal_g3.py @@ -0,0 +1,126 @@ +"""Modal launcher for the T-G3 chain (N7.5 / N8.2 / N7.1d) on ``H100!:2``. + +usage (from this directory, after the pre-registration in gpu-spend.md): + G3_RUN= G3_MILES= G3_YETO= \\ + modal run --detach modal_g3.py [--phases all|n75|n71d] + +The entrypoint only ``spawn()``s the chain and returns (killing a detached client would +cancel a blocking ``.remote()``). Results land in the Modal Volume ``t-g3-s25`` under +``//``; pull them with ``modal volume get t-g3-s25 / ``. +""" + +from __future__ import annotations + +import os + +import modal + +RUN = os.environ.get("G3_RUN", "unset") +MILES = os.environ.get("G3_MILES", "/home/michael/work/miles-s25-g3") +YETO = os.environ.get("G3_YETO", "/home/michael/work/s25-g3-miles") +HERE = os.path.dirname(os.path.abspath(__file__)) +IMAGE_REF = "radixark/miles:v0.1.1" # Miles 2806267d + Megatron-LM f148a32b4 (mHC output contraction), TE 2.17 +GPU = os.environ.get("G3_GPU", "H100!:2") # preflight: G3_GPU=L4:2 G3_EXPECT="NVIDIA L4" G3_HARD=900 +EXPECT_GPU = os.environ.get("G3_EXPECT", "NVIDIA H100 80GB HBM3") +HARD_S = int(os.environ.get("G3_HARD", "3600")) + +image = ( + modal.Image.from_registry(IMAGE_REF).entrypoint([]) + .env({"PYTHONPATH": "/g3/miles:/g3/yeto:/root/Megatron-LM", "CUDA_DEVICE_MAX_CONNECTIONS": "1", "G3_MILES": "/g3/miles", + "PYTHONUNBUFFERED": "1"}) + .add_local_dir(MILES, "/g3/miles", copy=False, ignore=[".git", "**/__pycache__", "tests", "docs", "examples"]) + .add_local_dir(os.path.join(YETO, "yeto"), "/g3/yeto/yeto", copy=False, ignore=["**/__pycache__"]) + .add_local_dir(HERE, "/g3/harness", copy=False, ignore=["**/__pycache__"]) +) +vol = modal.Volume.from_name("t-g3-s25", create_if_missing=True) +app = modal.App(f"t-g3-{RUN}") + +CHAIN = r""" +set -uo pipefail +exec > >(tee -a "$OUT/chain.log") 2>&1 +H=/g3/harness/harness.py +run() { # name timeout cmd... + local name=$1 t=$2; shift 2 + echo "=== PHASE $name start $(date -u +%T)" + timeout "$t" "$@" > "$OUT/logs/$name.log" 2>&1; local rc=$? + echo "=== PHASE $name rc=$rc end $(date -u +%T)"; tail -5 "$OUT/logs/$name.log" + echo "$name $rc" >> "$OUT/phases.txt" + return $rc +} +mkdir -p "$OUT/logs" "$OUT/n75" "$OUT/n71d" +nvidia-smi --query-gpu=index,name,driver_version,memory.total --format=csv | tee "$OUT/nvidia-smi.csv" +n=$(nvidia-smi --query-gpu=name --format=csv,noheader | grep -c "^$EXPECT_GPU\$") +[ "$n" = 2 ] || { echo "GPU ASSERT FAILED: want 2x $EXPECT_GPU"; exit 3; } +pip list 2>/dev/null | grep -iE '^(torch|transformer.engine|megatron|flash.linear|tilelang|mbridge) ' > "$OUT/versions.txt" +git -C /root/Megatron-LM rev-parse HEAD >> "$OUT/versions.txt" +python $H --phase hfconfig --out "$OUT" --hf /tmp/g3-hf && cp /tmp/g3-hf/config.json "$OUT/hf-config.json" +TR="torchrun --nproc-per-node 2 --master-port 29511" +if [[ "$PHASES" == pre ]]; then + run pre-n75-te 600 $TR $H --phase n75 --impl te --tag te --out "$OUT/n75" + run pre-n71d 400 $TR $H --phase n71d --mode cont5 --out "$OUT/n71d" --cut-dir "$OUT/n71d/cut" +fi +if [[ "$PHASES" == n71d-noise ]]; then # more samples of the run-to-run spread (cut in a second directory) + export NCCL_ALGO=Ring CUBLAS_WORKSPACE_CONFIG=:4096:8 NVIDIA_TF32_OVERRIDE=0 + for m in cont5c cont5d; do + run n71d-$m 600 $TR $H --phase n71d --mode $m --out "$OUT/n71d" --cut-dir "$OUT/n71d/cut" + done + run n71d-run3b 600 $TR $H --phase n71d --mode run3b --out "$OUT/n71d" --cut-dir "$OUT/n71d/cutb" + run n71d-resume2b 600 $TR $H --phase n71d --mode resume2b --out "$OUT/n71d" --cut-dir "$OUT/n71d/cutb" +fi +if [[ "$PHASES" == c4 ]]; then # T-C4 N7.6c: cut, clean resume, resume with the T-G3 bug injected + export NCCL_ALGO=Ring CUBLAS_WORKSPACE_CONFIG=:4096:8 NVIDIA_TF32_OVERRIDE=0 + run n71d-run3 900 $TR $H --phase n71d --mode run3 --out "$OUT/n71d" --cut-dir "$OUT/n71d/cut" + run n71d-resume2 900 $TR $H --phase n71d --mode resume2 --out "$OUT/n71d" --cut-dir "$OUT/n71d/cut" + run n71d-resume2inj 900 $TR $H --phase n71d --mode resume2inj --out "$OUT/n71d" --cut-dir "$OUT/n71d/cut" + du -sh "$OUT"; echo CHAIN DONE; exit 0 +fi +if [[ "$PHASES" == all || "$PHASES" == n75 ]]; then + run n75-te 1200 $TR $H --phase n75 --impl te --tag te --out "$OUT/n75" + run n75-te2 900 $TR $H --phase n75 --impl te --tag te2 --weights-from "$OUT/n75" --out "$OUT/n75" + run n75-local 1200 $TR $H --phase n75 --impl local --tag local --weights-from "$OUT/n75" --out "$OUT/n75" +fi +if [[ "$PHASES" == all || "$PHASES" == n71d ]]; then + export NCCL_ALGO=Ring CUBLAS_WORKSPACE_CONFIG=:4096:8 NVIDIA_TF32_OVERRIDE=0 + run n71d-cont5 900 $TR $H --phase n71d --mode cont5 --out "$OUT/n71d" --cut-dir "$OUT/n71d/cut" + run n71d-cont5b 900 $TR $H --phase n71d --mode cont5b --out "$OUT/n71d" --cut-dir "$OUT/n71d/cut" + run n71d-run3 900 $TR $H --phase n71d --mode run3 --out "$OUT/n71d" --cut-dir "$OUT/n71d/cut" + run n71d-resume2 900 $TR $H --phase n71d --mode resume2 --out "$OUT/n71d" --cut-dir "$OUT/n71d/cut" +fi +python /g3/harness/compare.py "$OUT" > "$OUT/logs/compare.log" 2>&1; echo "compare rc=$?" +du -sh "$OUT"; echo CHAIN DONE +""" + + +@app.function(image=image, gpu=GPU, cpu=16.0, memory=65536, timeout=HARD_S, volumes={"/vol": vol}) +def chain(run: str, phases: str, expect_gpu: str) -> int: + import subprocess + import threading + import time + + out = f"/vol/{run}" + os.makedirs(out, exist_ok=True) + stop = threading.Event() + + def committer(): # partial evidence survives a hard timeout + while not stop.wait(60): + try: + vol.commit() + except Exception as exc: # noqa: BLE001 + print("commit failed", exc, flush=True) + + threading.Thread(target=committer, daemon=True).start() + t0 = time.time() + rc = subprocess.call(["bash", "-c", CHAIN], env={**os.environ, "OUT": out, "PHASES": phases, + "EXPECT_GPU": expect_gpu}) + with open(f"{out}/chain.rc", "w") as fh: + fh.write(f"{rc} {time.time() - t0:.0f}s\n") + stop.set() + vol.commit() + return rc + + +@app.local_entrypoint() +def main(phases: str = "all"): + assert RUN != "unset", "set G3_RUN" + call = chain.spawn(RUN, phases, EXPECT_GPU) # module globals are re-read in the container + print("spawned", RUN, call.object_id, flush=True) diff --git a/tools/probes/s25_g1/image_probe.py b/tools/probes/s25_g1/image_probe.py new file mode 100644 index 00000000..24701de9 --- /dev/null +++ b/tools/probes/s25_g1/image_probe.py @@ -0,0 +1,59 @@ +"""CPU-only probe of the Miles ports image for the S25 G1 stage-1 mini rehearsal. + +modal run tools/probes/s25_g1/image_probe.py (no GPU; prints one JSON report) +""" +import modal + +# Same digest as yeto.rl.adapters.miles.pins.MILES_NEXT_IMAGE (the container has no yeto). +MILES_NEXT_IMAGE = ("docker:ghcr.io/michaellchung/yeto-miles-ports@sha256:" + "5b9dc2a9e6e315d73dabfbc6faae428b59cc01aabec6189645a5c8c80d6e3721") + +app = modal.App("yeto-s25-g1-probe") +image = modal.Image.from_registry(MILES_NEXT_IMAGE.removeprefix("docker:")) + +PROBE = r''' +import importlib, json, subprocess, sys +out = {} +def tryit(key, fn): + try: out[key] = fn() + except Exception as e: out[key] = f"ERR {type(e).__name__}: {str(e)[:300]}" +def ver(m): return getattr(importlib.import_module(m), "__version__", "?") +for m in ("torch", "transformer_engine", "megatron.core", "sglang", "fla", "tilelang", "triton", + "triton_kernels", "mbridge", "compressed_tensors", "flash_attn", "ray", "transformers"): + tryit("ver:" + m, lambda m=m: ver(m)) +def tcfg(): + from megatron.core.transformer import TransformerConfig + import dataclasses + names = {f.name for f in dataclasses.fields(TransformerConfig)} + return {k: k in names for k in ("enable_hyper_connections", "num_residual_streams", + "mhc_sinkhorn_iterations", "use_fused_mhc", "experimental_attention_variant")} +tryit("megatron_mhc_fields", tcfg) +tryit("sglang_glm5_next", lambda: bool(importlib.import_module("sglang.srt.models.glm5_next"))) +tryit("sglang_mxfp4", lambda: bool(importlib.import_module("sglang.srt.layers.quantization.mxfp4").Mxfp4MoEMethod)) +tryit("fla_kda", lambda: bool(importlib.import_module("fla.ops.kda"))) +def git(d): + return subprocess.run(["git", "-C", d, "log", "-1", "--format=%H %cd %s"], capture_output=True, text=True).stdout.strip() +for d in ("/root/miles", "/root/Megatron-LM", "/sgl-workspace/sglang"): + tryit("git:" + d, lambda d=d: git(d)) +tryit("manifest", lambda: open("/opt/yeto/image-manifest.json").read()[:1500]) +tryit("python", lambda: sys.version) +def sh(c): return subprocess.run(["bash", "-lc", c], capture_output=True, text=True).stdout[-1500:] +tryit("mcore_mhc_grep", lambda: sh("grep -rln 'enable_hyper_connections' /root/Megatron-LM/megatron | head")) +tryit("mcore_remote", lambda: sh("git -C /root/Megatron-LM remote -v | head -2; git -C /root/Megatron-LM log --oneline -3")) +tryit("te_pip", lambda: sh("pip show transformer_engine transformer-engine-cu12 transformer_engine_torch 2>/dev/null | grep -E 'Name|Version'")) +tryit("triton_kernels", lambda: sh("python3 -c 'import triton_kernels,os;print(os.path.dirname(triton_kernels.__file__))'; pip show triton_kernels 2>/dev/null|head -3")) +tryit("mem_cpu", lambda: sh("nproc; free -g | head -2; df -h /root | tail -1")) +print(json.dumps(out, indent=1)) +''' + + +@app.function(image=image, cpu=2, memory=8192, timeout=900) +def probe(): + import subprocess + r = subprocess.run(["python3", "-c", PROBE], capture_output=True, text=True) + return r.stdout + "\nSTDERR-TAIL:\n" + r.stderr[-3000:] + + +@app.local_entrypoint() +def main(): + print(probe.remote()) diff --git a/tools/probes/s25_g1/make_mini.py b/tools/probes/s25_g1/make_mini.py new file mode 100644 index 00000000..8e49f9d7 --- /dev/null +++ b/tools/probes/s25_g1/make_mini.py @@ -0,0 +1,148 @@ +"""Build the S25 G1 rehearsal checkpoints from Miles' 4-layer GLM5.3-Flash CI model. + +Runs inside the prep container (python3 make_mini.py ROOT). Writes, under ROOT: + +* ``src-4layer/``: CharyZeng/GLM-5.3-Flash-4layer at REVISION (288 experts, bf16); +* ``glm5next-mini-bf16/``: the same model with the first KEEP routed experts per + MoE layer (router rows and expert bias cut to match; n_routed_experts=KEEP); +* ``glm5next-mini-mxfp4-template/``: layout.json + metadata for the MXFP4 layout + of yetalabs/apex-flash-1-mxfp4-experts@70e86a0 (routed experts + ``weight_packed`` U8 + E8M0 ``weight_scale``, everything else as the bf16 source), + with that export's quantization_config; +* ``glm5next-mini-mxfp4-r0/``: round 0 for the sampler, converted from the bf16 + mini model by the N8.5 converter (``requantize_checkpoint``); +* ``datasets/gsm8k-0cbd9f3/train.jsonl`` and the HF dataset cache under ``hf/``. + +Prints one JSON summary line (``MINI-SUMMARY {...}``). +""" +from __future__ import annotations + +import json +import os +import re +import shutil +import sys +from pathlib import Path + +REPO = "CharyZeng/GLM-5.3-Flash-4layer" +REVISION = "7fb7883207fe73ea04430ac392f6ffe2b9f936c4" +MXFP4_REPO = "yetalabs/apex-flash-1-mxfp4-experts" +MXFP4_REVISION = "70e86a0" +DATA_REPO = "zhuzilin/gsm8k" +DATA_REVISION = "0cbd9f31d91ac21a7613dcbc7fef992adac459ae" +KEEP = int(os.environ.get("MINI_EXPERTS", "32")) +EXPERT = re.compile(r"\.mlp\.experts\.(\d+)\.") +ROUTED = re.compile(r"^(.*\.mlp\.experts\.\d+\.(?:gate|up|down)_proj)\.weight$") + + +def truncate(src: Path, dst: Path) -> dict: + import torch + from safetensors import safe_open + from safetensors.torch import save_file + + dst.mkdir(parents=True, exist_ok=True) + index = json.loads((src / "model.safetensors.index.json").read_text())["weight_map"] + weight_map, total, cut = {}, 0, 0 + for shard in sorted(set(index.values())): + out = {} + with safe_open(str(src / shard), framework="pt") as f: + for name in f.keys(): + m = EXPERT.search(name) + if m and int(m.group(1)) >= KEEP: + cut += 1 + continue + t = f.get_tensor(name) + if name.endswith(".mlp.gate.weight") or name.endswith(".mlp.gate.e_score_correction_bias"): + t = t[:KEEP].clone() + out[name] = t.contiguous() + if out: + save_file(out, str(dst / shard), metadata={"format": "pt"}) + for name, t in out.items(): + weight_map[name] = shard + total += t.numel() * t.element_size() + print(f"[mini] {shard}: kept {len(out)}", flush=True) + (dst / "model.safetensors.index.json").write_text( + json.dumps({"metadata": {"total_size": total}, "weight_map": weight_map}, indent=2, sort_keys=True)) + for f in src.iterdir(): + if f.is_file() and f.suffix != ".safetensors" and f.name != "model.safetensors.index.json" \ + and not f.name.startswith("."): + shutil.copy2(f, dst / f.name) + config = json.loads((dst / "config.json").read_text()) + config["text_config"]["n_routed_experts"] = KEEP + (dst / "config.json").write_text(json.dumps(config, indent=2)) + return {"tensors": len(weight_map), "bytes": total, "cut_expert_tensors": cut} + + +def make_template(bf16: Path, template: Path, quant_config: dict) -> dict: + from yeto.rl.quant.requant import (LAYOUT_FILE, LAYOUT_SCHEMA, CheckpointLayout, TensorSpec, + metadata_files, read_header) + from yeto.rl.quant.mxfp4 import GROUP_SIZE + + weight_map = json.loads((bf16 / "model.safetensors.index.json").read_text())["weight_map"] + tensors: dict[str, TensorSpec] = {} + headers = {s: read_header(bf16 / s) for s in sorted(set(weight_map.values()))} + for name, shard in weight_map.items(): + info = headers[shard][name] + m = ROUTED.match(name) + if m: + rows, cols = info["shape"] + tensors[m.group(1) + ".weight_packed"] = TensorSpec("U8", (rows, cols // 2), shard) + tensors[m.group(1) + ".weight_scale"] = TensorSpec("U8", (rows, cols // GROUP_SIZE), shard) + else: + tensors[name] = TensorSpec(info["dtype"], tuple(info["shape"]), shard) + layout = CheckpointLayout(tensors, quant_config) + layout.validate() + template.mkdir(parents=True, exist_ok=True) + for f in metadata_files(bf16): + shutil.copy2(f, template / f.name) + config = json.loads((template / "config.json").read_text()) + config["quantization_config"] = quant_config + (template / "config.json").write_text(json.dumps(config, indent=2)) + payload = {"schema": LAYOUT_SCHEMA, "quantization_config": quant_config, + "tensors": {n: {"dtype": t.dtype, "shape": list(t.shape), "shard": t.shard} + for n, t in sorted(tensors.items())}} + (template / LAYOUT_FILE).write_text(json.dumps(payload, sort_keys=True)) + return {"quantized_modules": len(layout.quantized_modules), "tensors": len(tensors)} + + +def main(root: str) -> None: + from huggingface_hub import hf_hub_download, snapshot_download + + root_dir = Path(root) + os.environ.setdefault("HF_HOME", str(root_dir / "hf")) + summary: dict = {"keep_experts": KEEP} + src = Path(snapshot_download(REPO, revision=REVISION, local_dir=str(root_dir / "src-4layer"))) + bf16 = root_dir / "glm5next-mini-bf16-7fb7883" + if not (bf16 / "model.safetensors.index.json").is_file(): + summary["truncate"] = truncate(src, bf16) + mx_config = json.loads(Path(hf_hub_download(MXFP4_REPO, "config.json", revision=MXFP4_REVISION)).read_text()) + quant_config = mx_config["quantization_config"] + template = root_dir / "glm5next-mini-mxfp4-template" + summary["template"] = make_template(bf16, template, quant_config) + from yeto.rl.quant.awq import AwqSmoothing + from yeto.rl.quant.requant import CheckpointLayout, requantize_checkpoint + + r0 = root_dir / "glm5next-mini-mxfp4-r0-7fb7883" + if r0.exists(): + shutil.rmtree(r0) + summary["r0"] = requantize_checkpoint(bf16, template, r0, AwqSmoothing.none()) + check = CheckpointLayout.from_checkpoint(r0) # the stage-1 template path (no layout.json) + summary["r0_quantized_modules"] = len(check.quantized_modules) + data = Path(snapshot_download(DATA_REPO, repo_type="dataset", revision=DATA_REVISION)) + import pandas as pd + + rows = pd.read_parquet(data / "train.parquet") + out = root_dir / "datasets/gsm8k-0cbd9f3" + out.mkdir(parents=True, exist_ok=True) + with open(out / "train.jsonl", "w", encoding="utf-8") as f: + for row in rows.to_dict(orient="records"): + f.write(json.dumps({k: (v.tolist() if hasattr(v, "tolist") else v) for k, v in row.items()}, + default=str) + "\n") + summary["gsm8k_rows"] = len(rows) + summary["gsm8k_columns"] = list(rows.columns) + summary["gsm8k_row0"] = str(rows.iloc[0].to_dict())[:400] + print("MINI-SUMMARY " + json.dumps(summary), flush=True) + + +if __name__ == "__main__": + main(sys.argv[1]) diff --git a/tools/probes/s25_g1/modal_app.py b/tools/probes/s25_g1/modal_app.py new file mode 100644 index 00000000..73f0cf20 --- /dev/null +++ b/tools/probes/s25_g1/modal_app.py @@ -0,0 +1,539 @@ +"""S25 T-G1: stage-1 mini rehearsal on Modal (examples/stage1/glm5next-mini-modal.yaml). + +The containers run the commands ``yeto-stage1 plan`` generates for that plan +(computed inside each container, so ``--source-sha256`` matches the mounted +tree): the Miles trainer-only island with the MXFP4 re-quantizer side process, +and the engine sampler island. Only the nodes (Modal containers instead of SSH +hosts) and the store (Modal Volume through the yeto store hook instead of OSS) +differ from the PPU plan. + + modal run tools/probes/s25_g1/modal_app.py::prep # CPU: checkpoints + data + modal run --detach tools/probes/s25_g1/modal_app.py::smoke # 1 x H100: SGLang + post-load checks + modal run --detach tools/probes/s25_g1/modal_app.py::launch # 3 rounds, cut after 2, resume + +``smoke`` and ``launch`` spawn and return; the local client may exit +(never kill a ``.remote()`` caller: memory modal-detach-kill-cancels). +Raw logs go to the Volume ``yeto-s25-g1-logs`` under ``/``. +""" +from __future__ import annotations + +import json +import os +import shlex +import shutil +import subprocess +import sys +import time +from pathlib import Path + +import modal + + +HERE = Path(__file__).resolve() +REPO = HERE.parents[3] if len(HERE.parents) > 3 else Path("/") # in the container the module sits at /root +WORKDIR = "/root/sky_workdir" +# run_limits.py sits next to this file locally and under the mounted tree in a container +sys.path.insert(0, str(HERE.parent) if (HERE.parent / "run_limits.py").is_file() + else f"{WORKDIR}/tools/probes/s25_g1") +from run_limits import Deadline, wait_until # noqa: E402 +# T-G1B: YETO_G1_PLAN=examples/stage1/glm5next-mini-modal-dense.yaml (set at deploy time; +# the image env carries it into every container, so all of them plan the same file) +PLAN = os.environ.get("YETO_G1_PLAN", "examples/stage1/glm5next-mini-modal.yaml") +IMAGE = ("ghcr.io/michaellchung/yeto-miles-ports@sha256:" + "5b9dc2a9e6e315d73dabfbc6faae428b59cc01aabec6189645a5c8c80d6e3721") +MILES_COMMIT = "1618b5c38db80e17ec928cfa47fe832cb18565fc" # s25-distopt +MODELS_VOLUME, STORE_VOLUME = "yeto-s25-g1-models", "yeto-s25-g1-store" +GPU_NAME = "H100" + +app = modal.App("yeto-s25-g1-stage1mini") +models = modal.Volume.from_name(MODELS_VOLUME, create_if_missing=True, version=2) +store = modal.Volume.from_name(STORE_VOLUME, create_if_missing=True, version=2) +# Logs go to their own Volume: the islands reload the store Volume, and a reload +# drops this container's uncommitted copies there (run-16 lost its log mirror). +logs = modal.Volume.from_name("yeto-s25-g1-logs", create_if_missing=True, version=2) +VOLUMES = {"/data/models": models, "/data/store": store, "/data/logs": logs} + +_PATCH_SGLANG = ( + "import importlib.util, sys; sys.path.insert(0, '/tmp/yeto-slr');" + "import sglang_load_report as r;" + "root = importlib.util.find_spec('sglang').submodule_search_locations[0];" + "f = root + '/srt/models/glm5_next.py'; print('patched', f, r.patch_model_file(f))" +) +image = ( + modal.Image.from_registry(IMAGE) + .add_local_file(str(REPO / ".g1-build/miles-s25-distopt.bundle"), "/tmp/miles.bundle", copy=True) + .add_local_file(str(REPO / "yeto/rl/adapters/sampler_engines/sglang_load_report.py"), + "/tmp/yeto-slr/sglang_load_report.py", copy=True) + .run_commands( + "git -C /root/miles fetch /tmp/miles.bundle s25-distopt", + f"git -C /root/miles checkout --detach {MILES_COMMIT}", + "git -C /root/miles status --porcelain --untracked-files=all | head -5", + f"python3 -c \"{_PATCH_SGLANG}\"", + ) + # not PYTHONPATH: Modal's own runtime path must stay intact; subprocesses get it (_env) + .env({"HOME": "/root", "PYTHONUNBUFFERED": "1", "HF_HOME": "/data/models/hf", "YETO_G1_PLAN": PLAN}) + .add_local_dir(str(REPO), WORKDIR, copy=False, + ignore=[".git", ".g1-build", ".venv", ".claude", "syncer/target", "**/__pycache__", + "**/*.pyc", "evidence*", "openspec", "data"]) +) + + +def _hf_secret() -> modal.Secret: + token = Path("~/.cache/huggingface/token").expanduser() + return modal.Secret.from_dict({"HF_TOKEN": token.read_text().strip()} if token.is_file() else {}) + + +LOCAL_LOGS = Path("/tmp/g1-logs") # container disk: an open file on the Volume blocks every reload + + +def _log_dir(tag: str) -> Path: + path = LOCAL_LOGS / tag + path.mkdir(parents=True, exist_ok=True) + return path + + +def _copy_files(src: Path, dst: Path) -> None: + """Copy file contents only (no copystat: a Volume refuses some metadata calls).""" + for path in src.rglob("*"): + if path.is_file() and path.suffix not in (".sock", ".pid"): + target = dst / path.relative_to(src) + target.parent.mkdir(parents=True, exist_ok=True) + shutil.copyfile(path, target) + + +def _commit(tag: str | None = None) -> None: + """Copy the local logs of ``tag`` to the store Volume (closed files only), then commit.""" + try: + if tag is not None: + _copy_files(LOCAL_LOGS / tag, Path("/data/logs") / tag) + logs.commit() + except Exception as error: # noqa: BLE001 - evidence copy only + print(f"[g1] log mirror/commit failed: {error}", flush=True) + + +def _mirror_every(tag: str, seconds: float = 180.0, state: list | None = None) -> None: + """Mirror the logs (and, once known, ``state[0]``: the island's local state dir with its + tape) every ``seconds``; a cancelled container keeps what was mirrored last.""" + import threading + + def loop() -> None: + while True: + time.sleep(seconds) + if state: + try: + shutil.copytree(state[0], LOCAL_LOGS / tag / "state", dirs_exist_ok=True, + ignore=shutil.ignore_patterns("*.pt", "*.safetensors")) + except Exception as error: # noqa: BLE001 + print(f"[g1] state copy failed: {error}", flush=True) + _commit(tag) + + threading.Thread(target=loop, daemon=True).start() + + +def _sh(cmd: str, log: Path | None = None, env: dict | None = None, check: bool = True) -> str: + print(f"[g1] $ {cmd}", flush=True) + r = subprocess.run(["bash", "-lc", cmd], capture_output=True, text=True, env=env or _env({}), cwd=WORKDIR) + out = r.stdout + r.stderr + if log is not None: + with open(log, "a", encoding="utf-8") as f: + f.write(f"$ {cmd}\n{out}\n") + if check and r.returncode: + raise RuntimeError(f"{cmd} -> {r.returncode}: {out[-2000:]}") + return out + + +def _gpu_facts(log_dir: Path, want: int) -> dict: + log = log_dir / "gpu.txt" + out = _sh("nvidia-smi --query-gpu=index,name,memory.total,driver_version --format=csv,noheader", log) + names = [line.split(",")[1].strip() for line in out.strip().splitlines()] + if len(names) != want or not all(GPU_NAME in n for n in names): + raise RuntimeError(f"GPU assert failed: want {want} x {GPU_NAME}, got {names}") + _sh("nvidia-smi; free -g; nproc; df -h /data/models /data/store /root", log, check=False) + return {"gpus": names} + + +def _plan(log: Path) -> dict: + out = _sh(f"python3 -m yeto.rl.stage1 plan {PLAN} --json 2>/dev/null; true", None, check=False) + start = out.index("{") + plan = json.loads(out[start: out.rindex("}") + 1]) + (log / "plan.json").write_text(json.dumps(plan, indent=1)) + if plan["errors"] or plan["gaps"]: + raise RuntimeError(f"plan errors {plan['errors']} gaps {plan['gaps']}") + return plan + + +def _nvml_sampler(path: Path) -> subprocess.Popen: + return subprocess.Popen( + ["bash", "-c", "nvidia-smi --query-gpu=timestamp,index,utilization.gpu,memory.used,power.draw " + f"--format=csv,noheader -l 5 >> {shlex.quote(str(path))}"]) + + +def _env(extra: dict) -> dict: + path = ":".join(p for p in (WORKDIR, os.environ.get("PYTHONPATH", "")) if p) + # The island's Python cannot import the Modal client; the Volume store commits and + # reloads through this interpreter (yeto.cloud.modal_ckpt_store, as yeto.modal_runner does). + runner = {"YETO_MODAL_PYTHON": sys.executable, + "YETO_MODAL_SYSPATH": os.pathsep.join(p for p in sys.path if p)} + return {**os.environ, "PYTHONPATH": path, **runner, **{k: str(v) for k, v in extra.items()}} + + +# --------------------------------------------------------------------- prep (CPU) +@app.function(image=image, volumes={"/data/models": models}, cpu=8, memory=65536, timeout=7200, + secrets=[_hf_secret()]) +def prep() -> str: + out = subprocess.run(["python3", f"{WORKDIR}/tools/probes/s25_g1/make_mini.py", "/data/models"], + capture_output=True, text=True, cwd=WORKDIR, env=_env({})) + text = out.stdout[-6000:] + "\nSTDERR:\n" + out.stderr[-4000:] + Path("/data/models/prep.log").write_text(text) + models.commit() + return text + + +@app.local_entrypoint() +def prep_main() -> None: + print(prep.remote()) + + +# --------------------------------------------------------------------- sampler smoke +@app.function(image=image, volumes=VOLUMES, gpu=f"{GPU_NAME}!:1", cpu=8, memory=65536, timeout=2400) +def smoke_run(tag: str) -> None: + """Start SGLang on round 0 through the adapter and run the post-load checks once.""" + log = _log_dir(tag) + facts = _gpu_facts(log, 1) + nvml = _nvml_sampler(log / "nvml.csv") + try: + plan = _plan(log) + run = plan["sampler"]["engines"][0] + os.environ.update({k: str(v) for k, v in run["env"].items()}) + sys.path.insert(0, WORKDIR) + from yeto.rl.adapters.sampler_engines.from_spec import build + from yeto.rl.engine.sampler_checks import PostLoadChecks + + engine = build(run["spec"]["engine"]) + checks = PostLoadChecks(**run["spec"]["checks"]) + began = time.time() + report = engine.start(None, 0) + started = time.time() - began + results = checks.run(engine, expected_experts=int(run["spec"]["checks"]["expected_expert_matrices"])) + from yeto.rl.engine.sampler_engine import SamplingParams + + sample = engine.generate(["Question: What is 2 + 3?\nAnswer:"], + SamplingParams(max_tokens=32, temperature=1.0, n=2, seed=1).check()) + from yeto.rl.adapters.sampler_engines.sglang_load_report import read_reports + + record = {**facts, "start_s": round(started, 1), "load_report": report.detail, "checks": results, + "hook_report": read_reports(run["spec"]["engine"]["options"]["load_report_dir"]), + "sample": [[c.text for c in g] for g in sample]} + engine.stop() + except Exception as error: # noqa: BLE001 - the record is the evidence + record = {**facts, "error": f"{type(error).__name__}: {error}"} + finally: + nvml.terminate() + (log / "smoke.json").write_text(json.dumps(record, indent=1, default=str)) + print(json.dumps(record, default=str)[:4000], flush=True) + _commit(tag) + + +@app.local_entrypoint() +def smoke(tag: str = "smoke-1") -> None: + call = smoke_run.spawn(tag) + print(f"spawned smoke_run {call.object_id}") + + +# --------------------------------------------------------------------- the islands +@app.function(image=image, volumes=VOLUMES, gpu=f"{GPU_NAME}!:1", cpu=8, memory=65536, timeout=7200, + max_inputs=1) # a fresh container per call: trainer B must not inherit A's disk or Ray +def sampler_island(tag: str) -> int: + # Modal may run this input again in a new container (a lost worker): each container + # keeps its own log directory, so a replacement never overwrites the first one's logs + container = os.environ.get("MODAL_TASK_ID") or f"t{int(time.time())}" + tag = f"{tag}/{container}" + log = _log_dir(tag) + _gpu_facts(log, 1) + nvml = _nvml_sampler(log / "nvml.csv") + state_dir: list = [] + _mirror_every(tag, 120.0, state_dir) + plan = _plan(log) + run = plan["sampler"]["engines"][0] + state_dir.append(Path(run["spec"]["event_tape"]).parent) + spec_path = Path(run["spec_path"]) + spec_path.parent.mkdir(parents=True, exist_ok=True) + spec_path.write_text(json.dumps(run["spec"], indent=1)) + began = time.time() + with open(log / "sampler.log", "w", encoding="utf-8") as f: + rc = subprocess.run(run["argv"], env=_env(run["env"]), cwd=WORKDIR, stdout=f, stderr=subprocess.STDOUT).returncode + nvml.terminate() + (log / "rc.json").write_text(json.dumps({"rc": rc, "wall_s": round(time.time() - began, 1)})) + state = Path(run["spec"]["event_tape"]).parent + if state.is_dir(): + shutil.copytree(state, log / "state", dirs_exist_ok=True) + _commit(tag) + return rc + + +@app.function(image=image, volumes=VOLUMES, gpu=f"{GPU_NAME}!:2", cpu=12, memory=163840, timeout=3600, + max_inputs=1) # a fresh container per call: trainer B must not inherit A's disk or Ray +def trainer_island(tag: str, stop_after_rounds: int | None = None, convert_only: bool = False) -> int: + log = _log_dir(tag) + _gpu_facts(log, 2) + nvml = _nvml_sampler(log / "nvml.csv") + state_dir: list = [] + _mirror_every(tag, 120.0, state_dir) + plan = _plan(log) + state_dir.append(Path("/data/yeto", plan["run_id"])) + trainer = plan["trainer"] + env = _env(trainer["env"]) + island = list(trainer["miles"]["island_argv"]) + if stop_after_rounds is not None: + island += ["--rl-stop-after-rounds", str(int(stop_after_rounds))] + (log / "argv.json").write_text(json.dumps({"island": island, "requant": plan["requant"]["loop_argv"], + "env": trainer["env"]}, indent=1)) + Path("/data/yeto", plan["run_id"]).mkdir(parents=True, exist_ok=True) + marker = Path(str(trainer["torch_dist"])) / "latest_checkpointed_iteration.txt" + if trainer.get("convert_argv") and not (marker.is_file() and marker.read_text().strip() == "release"): + began_convert = time.time() + with open(log / "convert.log", "w", encoding="utf-8") as f: + path = f"{trainer['convert_env']['PYTHONPATH']}:{env['PYTHONPATH']}" + crc = subprocess.run(trainer["convert_argv"], env={**env, "PYTHONPATH": path}, + cwd=WORKDIR, stdout=f, stderr=subprocess.STDOUT).returncode + (log / "convert.json").write_text(json.dumps({"rc": crc, "wall_s": round(time.time() - began_convert, 1)})) + models.commit() # trainer B (a new container) loads the same torch_dist + if crc != 0: + _commit(tag) + return 90 + if convert_only: + _sh(f"ls -la {shlex.quote(str(trainer['torch_dist']))}; du -sh {shlex.quote(str(trainer['torch_dist']))}", + log / "convert-out.txt", check=False) + nvml.terminate() + _commit(tag) + return 0 + _sh("ray stop --force || true", log / "ray.log", env=env, check=False) + _sh("ray start --head --port=6379 --node-ip-address=127.0.0.1 --num-gpus=2 --disable-usage-stats", + log / "ray.log", env=env) + began = time.time() + # The re-quantizer runs in its own container (requant_island), not as a side process + # here: a Modal Volume reload in this container (its poll) discards this container's + # uncommitted writes, and it removed a release directory while Miles wrote it (run-12). + with open(log / "trainer.log", "w", encoding="utf-8") as f: + # as the stage-1 node script does: the island forwards its env to the Ray actors + rc = subprocess.run(island, env={**env, "RAY_ADDRESS": "127.0.0.1:6379"}, cwd=WORKDIR, stdout=f, + stderr=subprocess.STDOUT).returncode + nvml.terminate() + (log / "rc.json").write_text(json.dumps({"rc": rc, "wall_s": round(time.time() - began, 1)})) + shutil.copytree(f"/data/yeto/{plan['run_id']}", log / "state", dirs_exist_ok=True, + ignore=shutil.ignore_patterns("*.pt", "*.safetensors")) + _sh("cp -r /tmp/ray/session_latest/logs " + shlex.quote(str(log / "ray-logs")) + " || true", check=False) + _commit(tag) + return rc + + +@app.function(image=image, volumes=VOLUMES, cpu=8, memory=65536, timeout=3 * 3600) +def requant_island(tag: str) -> int: + """The plan's re-quantizer loop (node: its own CPU container; same argv).""" + log = _log_dir(tag) + _mirror_every(tag, 60.0) + plan = _plan(log) + with open(log / "requant.log", "w", encoding="utf-8") as f: + rc = subprocess.run(plan["requant"]["loop_argv"], env=_env(plan["trainer"]["env"]), cwd=WORKDIR, + stdout=f, stderr=subprocess.STDOUT).returncode + _commit(tag) + return rc + + +# --------------------------------------------------------------------- release digests (T-G1B) +def _tensor_digests(directory: Path) -> dict[str, str]: + """sha256 of every tensor's bytes in the safetensors files of ``directory``.""" + import hashlib + + out: dict[str, str] = {} + for file in sorted(directory.glob("*.safetensors")): + with open(file, "rb") as f: + size = int.from_bytes(f.read(8), "little") + header = json.loads(f.read(size)) + base = 8 + size + for name, meta in sorted(header.items(), key=lambda kv: kv[1]["data_offsets"][0] + if kv[0] != "__metadata__" else -1): + if name == "__metadata__": + continue + begin, end = meta["data_offsets"] + f.seek(base + begin) + digest, left = hashlib.sha256(), end - begin + while left: + chunk = f.read(min(left, 1 << 24)) + digest.update(chunk) + left -= len(chunk) + out[name] = digest.hexdigest() + return out + + +def _summarize(digests: dict[str, str]) -> dict: + import hashlib + + def combined(names: list[str]) -> str: + h = hashlib.sha256() + for name in names: + h.update(f"{name}={digests[name]}\n".encode()) + return h.hexdigest() + + experts = sorted(n for n in digests if ".mlp.experts." in n) + others = sorted(n for n in digests if ".mlp.experts." not in n) + return {"tensors": len(digests), "expert_tensors": len(experts), "experts_sha256": combined(experts), + "others_sha256": combined(others)} + + +@app.function(image=image, volumes=VOLUMES, cpu=4, memory=8192, timeout=3 * 3600) +def release_hasher(tag: str, base_dir: str, store_root: str, release_dirs: list[str], poll_s: float = 20.0) -> int: + """Digest each release as soon as it is complete (the store keeps only the newest 2): + per tensor sha256, plus combined digests of the routed-expert tensors and of the rest. + Writes ``/release-digests.jsonl`` and ``/tensors//.json``.""" + log = _log_dir(tag) + out = log / "release-digests.jsonl" + seen: set[str] = set() + + def record(kind: str, name: str, directory: Path, extra: dict) -> None: + began = time.time() + try: + digests = _tensor_digests(directory) + except Exception as error: # noqa: BLE001 - pruned or still open: record, go on + row = {"kind": kind, "release": name, "error": repr(error)} + else: + target = log / "tensors" / kind / f"{name}.json" + target.parent.mkdir(parents=True, exist_ok=True) + target.write_text(json.dumps(digests, indent=0, sort_keys=True)) + row = {"kind": kind, "release": name, **extra, **_summarize(digests), + "hash_s": round(time.time() - began, 1)} + row["time_unix"] = time.time() + with open(out, "a", encoding="utf-8") as f: + f.write(json.dumps(row) + "\n") + print(f"[g1b] {json.dumps(row)}", flush=True) + _commit(tag) + + record("base", Path(base_dir).name, Path(base_dir), {}) + while True: + try: + store.reload() + except Exception as error: # noqa: BLE001 - an open file blocks the reload; retry + print(f"[g1b] reload failed: {error}", flush=True) + for kind in release_dirs: + root = Path(store_root) / kind + # RELEASE.json is written after the weights (weights -> record -> commit -> LATEST) + for marker in sorted(root.glob("*/RELEASE.json")) if root.is_dir() else (): + name = marker.parent.name + if f"{kind}/{name}" in seen: + continue + seen.add(f"{kind}/{name}") + meta = json.loads(marker.read_text()) + record(kind, name, root / name, {"version": meta.get("version"), + "policy_hash": meta.get("policy_hash")}) + time.sleep(poll_s) + + +RUN_LIMIT_S = 100 * 60 # whole run (g1b-1: trainers about 40 min, sampler first start about 6 min) +SAMPLER_DRAIN_S = 5 * 60 # after both trainers returned 0: the sampler only ends its last round + + +@app.function(image=image, volumes=VOLUMES, cpu=1, memory=1024, timeout=3 * 3600) +def orchestrate(tag: str, cut_after: int = 2, run_limit_s: float = RUN_LIMIT_S, + sampler_drain_s: float = SAMPLER_DRAIN_S) -> dict: + """Sampler for the whole run; trainer A stops after ``cut_after`` rounds (cut), + trainer B resumes in a new container. A trainer that fails cancels the sampler; + a sampler that fails cancels the running trainer (no card waits for nothing). + Two wall-clock limits (``run_limits``): the whole run (``run_limit_s``: a lost + container that Modal runs again, or a hang, cancels every call), and the sampler + after both trainers returned 0 (``sampler_drain_s``).""" + deadline = Deadline(run_limit_s) + limits: list[str] = [] + run_store = Path("/data/store/s25-g1") # the plan's store: every run starts from an empty one + if run_store.exists(): + shutil.rmtree(run_store) + store.commit() + sampler = sampler_island.spawn(f"{tag}/sampler") + requant = requant_island.spawn(f"{tag}/requant") + plan = _plan(_log_dir(f"{tag}/orchestrate")) + spec = plan["sampler"]["engines"][0]["spec"] + hasher = release_hasher.spawn(f"{tag}/hashes", str(spec["base_dir"]), str(spec["store"]), + ["policy-releases", str(spec["release_dir"])]) + sampler_rc: object = None + + def sampler_done() -> bool: + nonlocal sampler_rc + if sampler_rc is not None: + return True + try: + sampler_rc = sampler.get(timeout=0) + return True + except TimeoutError: + return False + except Exception as error: # noqa: BLE001 - the sampler call failed + sampler_rc = f"error: {error}" + return True + + def run_trainer(sub: str, stop: int | None) -> object: + call = trainer_island.spawn(f"{tag}/{sub}", stop) + while True: + try: + return call.get(timeout=30) + except TimeoutError: + pass + except Exception as error: # noqa: BLE001 + return f"error: {error}" + if deadline.passed(): + call.cancel() + limits.append(f"{sub}: run limit {run_limit_s:g} s passed") + return "cancelled: run limit" + if sampler_done() and sampler_rc != 0: + call.cancel() + return "cancelled: sampler ended with " + str(sampler_rc) + + first = run_trainer("trainer-a", cut_after) + second = run_trainer("trainer-b", None) if first == 0 else None + if first != 0 or second != 0: + if not sampler_done(): + sampler.cancel() + sampler_rc = "cancelled" + elif not wait_until(sampler_done, sampler_drain_s, deadline=deadline): + sampler.cancel() + limits.append(f"sampler: not done {sampler_drain_s:g} s after both trainers returned 0") + sampler_rc = "cancelled: drain limit" + time.sleep(90) # the re-quantizer's last log mirror + requant.cancel() + hasher.cancel() + result = {"trainer_a": first, "trainer_b": second, "sampler": sampler_rc, "limits": limits, + "wall_s": round(run_limit_s - deadline.left(), 1)} + print(json.dumps(result, default=str), flush=True) + out = Path(f"/data/logs/{tag}") + out.mkdir(parents=True, exist_ok=True) + (out / "orchestrate.json").write_text(json.dumps(result, default=str)) + logs.commit() + return result + + +@app.local_entrypoint() +def requant(tag: str = "requant-1") -> None: + """Start (or restart) only the re-quantizer loop of a running rehearsal.""" + call = requant_island.spawn(tag) + print(f"spawned requant {call.object_id}") + + +@app.local_entrypoint() +def convert(tag: str = "convert-1") -> None: + call = trainer_island.spawn(tag, None, True) + print(f"spawned convert {call.object_id}") + + +@app.local_entrypoint() +def launch(tag: str = "run-1", cut_after: int = 2) -> None: + call = orchestrate.spawn(tag, cut_after) + print(f"spawned orchestrate {call.object_id}") + + +@app.function(image=image, volumes={"/data/models": models}, cpu=2, memory=8192, timeout=900) +def cpu_shell(cmd: str) -> str: + """CPU-only inspection of the image (no GPU).""" + r = subprocess.run(["bash", "-lc", cmd], capture_output=True, text=True, cwd=WORKDIR, env=_env({})) + return r.stdout[-20000:] + "\nSTDERR:\n" + r.stderr[-4000:] + + +@app.local_entrypoint() +def shell(cmd: str) -> None: + print(cpu_shell.remote(cmd)) diff --git a/tools/probes/s25_g1/run_limits.py b/tools/probes/s25_g1/run_limits.py new file mode 100644 index 00000000..6433c868 --- /dev/null +++ b/tools/probes/s25_g1/run_limits.py @@ -0,0 +1,47 @@ +"""Wall-clock limits of the S25 G1 rehearsal driver (``modal_app.orchestrate``). + +No card may wait for nothing (T-G1B g1b-1: the sampler held 1 x H100 for about +30 min after both trainers had finished): + +* :class:`Deadline`: the whole run's wall-clock limit. A lost container that + Modal runs again (or one that hangs) can not keep the cards past it. +* :func:`wait_until`: wait for a call to finish, at most ``limit_s``; the driver + uses it for the sampler after every trainer returned rc 0 (the sampler then + has nothing to wait for but its backlog). + +Pure functions over ``done()``, ``clock()`` and ``sleep()``: CPU-tested in +``tests/test_s25_g1c_sampler_resume.py``. +""" +from __future__ import annotations + +import time +from collections.abc import Callable + + +class Deadline: + def __init__(self, limit_s: float, clock: Callable[[], float] = time.monotonic) -> None: + if float(limit_s) <= 0: + raise ValueError("limit_s must be > 0") + self.clock = clock + self.limit_s = float(limit_s) + self.started = clock() + + def left(self) -> float: + return self.limit_s - (self.clock() - self.started) + + def passed(self) -> bool: + return self.left() <= 0 + + +def wait_until(done: Callable[[], bool], limit_s: float, *, poll_s: float = 15.0, + clock: Callable[[], float] = time.monotonic, + sleep: Callable[[float], None] = time.sleep, + deadline: Deadline | None = None) -> bool: + """True once ``done()``; False when ``limit_s`` (or ``deadline``) passed first.""" + started = clock() + while True: + if done(): + return True + if clock() - started >= float(limit_s) or (deadline is not None and deadline.passed()): + return False + sleep(poll_s) diff --git a/yeto/cloud/droppable.py b/yeto/cloud/droppable.py index abc6cb22..abd1e375 100644 --- a/yeto/cloud/droppable.py +++ b/yeto/cloud/droppable.py @@ -126,8 +126,10 @@ def role_args(args, role: str | None): def island_keeps_spot_volume(args) -> bool: """The legacy ``--spot`` per-island volume (completed groups, re-synced on a - rebuild) is a resume from an archive; a droppable island must not get it.""" - return bool(getattr(args, "spot", False)) and getattr(args, "_island_role", None) != DROPPABLE + rebuild) is a resume from an archive; a droppable island must not get it, nor + a spot sampler island (``spot_without_volume``: its state is the buffer).""" + return (bool(getattr(args, "spot", False)) and getattr(args, "_island_role", None) != DROPPABLE + and not getattr(args, "spot_without_volume", False)) def role_env(role: str | None, cloud: str, region: str | None) -> dict[str, str]: diff --git a/yeto/cloud/nebius_reclaim.py b/yeto/cloud/nebius_reclaim.py index 31d8b6b9..6903dd73 100644 --- a/yeto/cloud/nebius_reclaim.py +++ b/yeto/cloud/nebius_reclaim.py @@ -106,20 +106,55 @@ def install_command(container: str = CONTAINER, notice: str = NOTICE_FILE, ack: ) +def spot_sampler_env() -> dict[str, str]: + """The environment of a spot sampler island on Nebius (S25 C1b): the sampler's + reclaim marker is the notice file the host unit writes, and the sampler + writes the ack file once its LEAVE is stored. The two env names of this + module also tell :func:`wants_host_hook` to install the unit.""" + from yeto.rl.engine.sampler_resume import RECLAIM_ACK_ENV, RECLAIM_MARKER_ENV + + assert RECLAIM_ACK_ENV == ACK_ENV + return {RECLAIM_MARKER_ENV: NOTICE_FILE, NOTICE_ENV: NOTICE_FILE, ACK_ENV: ACK_FILE} + + def wants_host_hook(task_envs) -> bool: - """True for a droppable island on Nebius (from the task's YETO_ISLAND_ROLE).""" + """True for a droppable island on Nebius (from the task's YETO_ISLAND_ROLE) and + for a Nebius spot sampler island (:func:`spot_sampler_env`).""" from yeto.cloud import droppable - doc = droppable.env_role(dict(task_envs or {})) + envs = dict(task_envs or {}) + if envs.get(NOTICE_ENV) == NOTICE_FILE and envs.get(ACK_ENV) == ACK_FILE: + return True + doc = droppable.env_role(envs) return bool(doc) and doc.get("role") == droppable.DROPPABLE and doc.get("cloud") == "nebius" +def _server_credentials(cluster: str) -> dict[str, Any]: + """The cluster's SSH credentials from a remote sky API server (the cluster yaml + stays on the server), as ``sky status`` gets them; the private key is written + to the client's per-cluster key file, as the sky CLI does.""" + from sky.client import sdk + from sky.utils import cluster_utils + + records = sdk.stream_and_get(sdk.status([cluster], _include_credentials=True)) + creds = dict(next((r.get("credentials") for r in records if r.get("name") == cluster), None) or {}) + if not creds: + raise RuntimeError(f"the sky API server returned no SSH credentials for {cluster}") + cluster_utils.SSHConfigHelper.generate_local_key_file(cluster, creds) + return creds + + def _host_runner(handle: Any) -> Any: - """A sky SSH runner on the VM HOST (docker_user=None), not in the container.""" + """A sky SSH runner on the VM HOST (docker_user=None), not in the container. + The credentials come from the local cluster yaml (sky API server on this + machine), else from the remote API server (T-C1b).""" from sky.backends import backend_utils from sky.utils import command_runner - creds = backend_utils.ssh_credential_from_yaml(handle.cluster_yaml, None, handle.ssh_user) + if os.path.exists(os.path.expanduser(handle.cluster_yaml)): + creds = backend_utils.ssh_credential_from_yaml(handle.cluster_yaml, None, handle.ssh_user) + else: + creds = _server_credentials(handle.cluster_name) creds.pop("docker_user", None) ip = handle.head_ip return command_runner.SSHCommandRunner((ip, 22), **creds) diff --git a/yeto/codepack.py b/yeto/codepack.py index 73fbc4ac..23ea2a4c 100644 --- a/yeto/codepack.py +++ b/yeto/codepack.py @@ -190,7 +190,8 @@ def build_pack(repo: str | os.PathLike[str]) -> tuple[CodePack, bytes]: files = pack_file_list(root) patch_sha = hashlib.sha256(patch).hexdigest() info = {"format": PACK_FORMAT, "mode": CODE_PACK_MODE_GIT_PATCH, "commit": commit, - "dirty": bool(patch), "patch_sha": patch_sha, "files": files} + "dirty": bool(patch), "patch_sha": patch_sha, "files": files, + "executables": _executables(root, files)} data = _tar_bytes(root, files, info) pack = CodePack(commit=commit, dirty=bool(patch), patch_sha=patch_sha, digest=hashlib.sha256(data).hexdigest()) @@ -221,9 +222,34 @@ def tree_digest(tree: str | os.PathLike[str]) -> str: return hashlib.sha256(_tar_bytes(root, files, info)).hexdigest() +def _executables(root: Path, files: Iterable[str]) -> list[str]: + """The listed regular files with an executable bit (the pack records them as 0o755).""" + return [n for n in files + if not (root / n).is_symlink() and os.stat(root / n).st_mode & 0o111] + + +def restore_executables(tree: str | os.PathLike[str]) -> list[str]: + """Set the executable bit again on the files the pack info lists as executable + and returns the ones it changed. A transport may drop the bit: sky's workdir + upload to a remote API server goes through a zip that the server unpacks + without file modes (T-C1b). A pack without the list (older packs) is left alone.""" + root = Path(tree) + fixed = [] + for name in read_pack_info(root).get("executables") or []: + path = root / name + if path.is_file() and not path.is_symlink(): + mode = os.stat(path).st_mode + if not mode & 0o111: + os.chmod(path, mode | 0o755) + fixed.append(name) + return fixed + + def verify_tree(tree: str | os.PathLike[str], expected: str) -> str: - """Refuse a tree whose rebuilt tar does not have the digest ``expected``.""" + """Refuse a tree whose rebuilt tar does not have the digest ``expected`` + (after :func:`restore_executables`: modes are part of the digest).""" + restore_executables(tree) actual = tree_digest(tree) if actual != expected.lower(): raise CodePackError(f"code pack digest mismatch: run manifest has {expected.lower()}, " @@ -241,9 +267,23 @@ def _safe_member(member: tarfile.TarInfo) -> None: raise CodePackError(f"code pack symlink {member.name!r} leaves the tree") +# The pack's entries carry mtime 0 (a deterministic tar); an unpacked tree gets this +# time instead. sky zips a workdir it uploads to a remote API server, and ZIP refuses +# timestamps before 1980 ("ZIP does not support timestamps before 1980", T-C1b). +UNPACKED_MTIME = 946684800 # 2000-01-01T00:00:00Z + + +def _stamp_times(tree: Path) -> None: + """Give every entry of ``tree`` :data:`UNPACKED_MTIME` (links themselves, not + their targets). The digest covers names and bytes only, so it is unchanged.""" + for path in [tree, *tree.rglob("*")]: + os.utime(path, (UNPACKED_MTIME, UNPACKED_MTIME), follow_symlinks=False) + + def unpack(data: bytes, digest: str, cache_root: str | os.PathLike[str]) -> Path: """Check ``data`` against ``digest`` and unpack it to ``/`` - (reused when already there and intact). Returns the tree.""" + (reused when already there and intact; a tree unpacked before + :data:`UNPACKED_MTIME` existed gets its times now). Returns the tree.""" if hashlib.sha256(data).hexdigest() != digest.lower(): raise CodePackError(f"downloaded code pack does not have digest {digest}") @@ -251,6 +291,7 @@ def unpack(data: bytes, digest: str, cache_root: str | os.PathLike[str]) -> Path if dest.is_dir(): try: verify_tree(dest, digest) + _stamp_times(dest) return dest except CodePackError: shutil.rmtree(dest) @@ -262,8 +303,17 @@ def unpack(data: bytes, digest: str, cache_root: str | os.PathLike[str]) -> Path for member in members: _safe_member(member) tar.extractall(staging, members=members, filter="tar") + _stamp_times(staging) verify_tree(staging, digest) - staging.rename(dest) + try: + staging.rename(dest) + except OSError: + # another process (a second launch from the same checkout) unpacked the + # same digest first: keep its tree when it is intact + if not dest.is_dir(): + raise + verify_tree(dest, digest) + shutil.rmtree(staging, ignore_errors=True) except BaseException: shutil.rmtree(staging, ignore_errors=True) raise diff --git a/yeto/launch/rl_ports.py b/yeto/launch/rl_ports.py index c1f2b369..63b41f24 100644 --- a/yeto/launch/rl_ports.py +++ b/yeto/launch/rl_ports.py @@ -79,15 +79,16 @@ def _prepare_ports_algorithm(args) -> None: no_sync = bool(getattr(args, "rl_single_island_no_sync", False)) from ..rl.engine import sample_roles - # L5.4: a store-coupled sampler/trainer pair is two no-sync islands; a - # convertible pool (L7.5) is two or more (sample_roles holds their own rules). + # L5.4: a store-coupled pair (one trainer, one or more samplers: S25 C1) and a + # convertible pool (L7.5) are two or more no-sync islands (sample_roles holds + # their own rules). no_sync_islands = (islands if sample_roles.is_convertible_pool(args) - else 2 if sample_roles.is_pair(args) else 1) + or sample_roles.is_pair(args) else 1) if no_sync and (islands != no_sync_islands or getattr(args, "rl_sync_preset", "strict-avg") != "strict-avg" or getattr(args, "rl_initial_adapter", None)): raise AlgorithmSpecError( "--rl-single-island-no-sync needs exactly one island (one --gpu entry, no " - "external learners; two with --rl-island-sample-role), the default sync " + "external learners; one per island with --rl-island-sample-role), the default sync " "preset and no initial adapter" ) args.rl_allow_unverified_mechanism = list( @@ -255,6 +256,20 @@ def _check_resume_launch_flags(args) -> bool: raise ValueError(f"{flag} needs --rl-checkpoint-store") if getattr(args, "rl_resume_allow_config_change", False) and not store: raise ValueError("--rl-resume-allow-config-change needs --rl-checkpoint-store") + from ..rl.engine.sample_roles import is_synced_trainer + + if store and not elastic and is_synced_trainer(args): + # S25 C3: the store holds the sample buffer and the releases of synced trainers + # (forwarded as --rl-sample-store); multi-island round cuts do not exist + for name, flag in (("rl_cut_every", "--rl-cut-every"), ("rl_cut_keep", "--rl-cut-keep"), + ("rl_stop_after_rounds", "--rl-stop-after-rounds")): + if getattr(args, name, None) is not None: + raise ValueError(f"{flag}: synced trainers write no round cuts (multi-island " + "resume: rl-resume-from-checkpoint 3.3)") + if getattr(args, "rl_resume_allow_config_change", False): + raise ValueError("--rl-resume-allow-config-change: synced trainers write no round cuts") + rl_checkpoint_store_plan(args) # validates the form + return False if elastic: for name, flag in (("rl_cut_keep", "--rl-cut-keep"), ("rl_stop_after_rounds", "--rl-stop-after-rounds")): if getattr(args, name, None) is not None: @@ -378,7 +393,11 @@ def _check_ports_infra_switches(args) -> None: elif sample_roles.is_pair(args): # L5.4 --rl-island-sample-role: per-island roles sample_roles.check_pair(args, rl_island_count(args)) else: + sample_roles.check_synced_trainers(args, rl_island_count(args)) # S25 C3 sample_roles.check_launch(args) # L4.6 --rl-sample-role (both checks nothing) + if getattr(args, "rl_sampler_spot", False) and not sample_roles.is_pair(args): + raise ValueError("--rl-sampler-spot needs --rl-island-sample-role (it bills the sampler " + "islands of a pair as spot)") from .storage import store_object_envs for cloud in {s.cloud for s in parse_gpu_spec(args.gpu)} if getattr(args, "gpu", None) else (): @@ -397,7 +416,8 @@ def _check_ports_infra_switches(args) -> None: if getattr(args, "rl_elastic_accept_rebind", False): given.append("--rl-elastic-accept-rebind") resume_store = _check_resume_launch_flags(args) - if getattr(args, "rl_checkpoint_store", None) and not resume_store: + if getattr(args, "rl_checkpoint_store", None) and not resume_store \ + and not sample_roles.is_synced_trainer(args): given.append("--rl-checkpoint-store") if getattr(args, "rl_elastic_declare_cells", False): given.append("--rl-elastic-declare-cells") @@ -531,6 +551,7 @@ def _ports_infra_flags(args) -> tuple[str, str]: if getattr(args, "rl_max_policy_age", 0): # 0 (default) adds nothing flags += f" --rl-max-policy-age {int(args.rl_max_policy_age)}" from ..rl.engine.sample_roles import island_flags as sample_role_flags + from ..rl.engine.sample_roles import is_synced_trainer as sample_role_is_synced_trainer flags += sample_role_flags(args) # L4.6: role both (default) adds nothing if (getattr(args, "rl_lr_schedule", None) or "auto") != "auto": # auto adds nothing @@ -542,6 +563,11 @@ def _ports_infra_flags(args) -> tuple[str, str]: resume_store = (rl_checkpoint_store_plan(args) if getattr(args, "rl_checkpoint_store", None) and not getattr(args, "rl_elastic", False) else None) + if resume_store is not None and sample_role_is_synced_trainer(args): + # S25 C3: synced trainers mount the store for the buffer and the releases only + flags += (f" --rl-sample-store {shlex.quote(resume_store[0])}" + f" --rl-sample-lanes {rl_island_count(args)}") + resume_store = None if resume_store is not None: # rl-resume-from-checkpoint: non-elastic resume store flags += f" --rl-resume-store {shlex.quote(resume_store[0])}" for name, flag in (("rl_cut_keep", "--rl-cut-keep"), ("rl_stop_after_rounds", "--rl-stop-after-rounds")): diff --git a/yeto/launch/ssh_trainer_island.py b/yeto/launch/ssh_trainer_island.py index 56b3de31..f4731810 100644 --- a/yeto/launch/ssh_trainer_island.py +++ b/yeto/launch/ssh_trainer_island.py @@ -480,10 +480,16 @@ def ssh_binding(args: Any, island: TrainerIsland, *, journal_dir: Path, files: A def requant_argv(store: str, template: str, *, recipe: str | None = None, awq_scales: str | None = None, - python: str = "python") -> tuple[str, ...]: + python: str = "python", frozen_base: str | None = None, + frozen_pattern: str | None = None) -> tuple[str, ...]: """The one-pass argv of :mod:`yeto.rl.quant.release` (R2 with ``recipe`` and - ``awq_scales``; R1 without). The start script adds ``--loop`` for the side process.""" + ``awq_scales``; R1 without; ``frozen_base``/``frozen_pattern``: untrained + tensors the release omits). The start script adds ``--loop`` for the side process.""" argv = [python, "-m", "yeto.rl.quant.release", "--store", store, "--template", template] + if bool(frozen_base) != bool(frozen_pattern): + raise SshIslandError("give both frozen_base and frozen_pattern, or neither") + if frozen_base: + argv += ["--frozen-base", frozen_base, "--frozen-pattern", frozen_pattern] if recipe or awq_scales: if not (recipe and awq_scales): raise SshIslandError("give both recipe and awq_scales (R2), or neither (R1)") diff --git a/yeto/launch_preflight.py b/yeto/launch_preflight.py index 419432c2..539b07a0 100644 --- a/yeto/launch_preflight.py +++ b/yeto/launch_preflight.py @@ -363,12 +363,17 @@ def add_cli_args(p) -> None: "base, never from a checkpoint store); repeatable; an island without a role is " "an anchor; replaces --spot") g.add_argument("--rl-island-sample-role", action="append", default=None, metavar="ISLAND:ROLE", - help="rl-trainer-start-stop L5.4: one launch of a sampler island and a trainer " + help="rl-trainer-start-stop L5.4: one launch of sampler islands and a trainer " "island over one --rl-checkpoint-store (ROLE sampler or trainer, one entry per " - "--gpu island; exactly two islands). Needs --rl-single-island-no-sync: no " + "--gpu island; exactly one trainer, one or more samplers; several samplers " + "split --total-steps). Needs --rl-single-island-no-sync: no " "syncer, each island runs as a one-learner no-sync island with that " - "--rl-sample-role. Combine with --rl-trainer-scheduler s1 " + "--rl-sample-role. With one sampler, combine with --rl-trainer-scheduler s1 " "--rl-trainer-island M to start and stop the trainer island") + g.add_argument("--rl-sampler-spot", action="store_true", + help="with --rl-island-sample-role: the sampler islands use spot (preemptible) " + "capacity, the trainer island stays on demand; a reclaimed sampler island is " + "started again and continues after its last round in the buffer") g.add_argument("--rl-negative-test-run", action="store_true", help="mark this run as a negative-test run (required by --rl-island-override); " "a normal run refuses to resume its checkpoint store and `yeto merge` refuses " diff --git a/yeto/launcher.py b/yeto/launcher.py index 7955bea9..0abd2f4b 100644 --- a/yeto/launcher.py +++ b/yeto/launcher.py @@ -187,10 +187,24 @@ def _island_task(task_factory, args, spec, m, num_learners, syncer_addr, overrid learner_id, island_learners = sample_roles.island_learner(args, m, num_learners) task = task_factory(iargs, spec, learner_id, island_learners, syncer_addr) io.apply_to_task(task, {**io.island_env(args, m, overrides), - **(droppable.role_env(role, spec.cloud, spec.region) if role else {})}) + **(droppable.role_env(role, spec.cloud, spec.region) if role else {}), + **sample_roles.island_env(iargs), + **_spot_sampler_cloud_env(iargs, getattr(spec, "cloud", None))}) return task +def _spot_sampler_cloud_env(iargs, cloud: str | None) -> dict[str, str]: + """A spot sampler island on Nebius: the VM stop reaches the sampler through the + host unit of ``yeto.cloud.nebius_reclaim`` (Modal signals the container itself).""" + from .rl.engine import sample_roles + + if cloud != "nebius" or not getattr(iargs, "spot", False) or not sample_roles.island_env(iargs): + return {} + from .cloud import nebius_reclaim + + return nebius_reclaim.spot_sampler_env() + + def _ssh_trainer_spec(args, spec: ClusterSpec, task): """The SSH spec of a held-back ``ssh:`` trainer island (None on other clouds).""" if spec.cloud != "ssh": diff --git a/yeto/modal_runner.py b/yeto/modal_runner.py index e2b99737..b25be3c4 100644 --- a/yeto/modal_runner.py +++ b/yeto/modal_runner.py @@ -925,7 +925,9 @@ def run_island_script(cmd: list[str], env: dict, *, popen=subprocess.Popen, sign sig = signal_mod or _signal if not _is_droppable(env): - if env.get("YETO_SPOT_INFLIGHT_SAVE_DIR"): + # in-flight save (4.1) and sampler islands (S25 C1, the launcher sets the marker + # path) both turn the reclaim signal into the marker file + if env.get("YETO_SPOT_INFLIGHT_SAVE_DIR") or env.get("YETO_SPOT_RECLAIM_MARKER"): return _run_forwarding_to_learner(cmd, env, popen=popen, sig=sig) return subprocess.call(cmd, env=env) if popen is subprocess.Popen else popen(cmd, env=env).wait() killpg = killpg or os.killpg diff --git a/yeto/rl/adapters/miles/buffer_rollout.py b/yeto/rl/adapters/miles/buffer_rollout.py index 2207badd..f0c07119 100644 --- a/yeto/rl/adapters/miles/buffer_rollout.py +++ b/yeto/rl/adapters/miles/buffer_rollout.py @@ -28,7 +28,8 @@ Limits (first implementation): -* one trainer reads a buffer; on a shared directory (``index.jsonl``) one +* one trainer reads a buffer, or synced trainers read one lane each + (:mod:`yeto.rl.engine.merged_release`); on a shared directory (``index.jsonl``) one sampler island writes it (``flock`` holds only on one host); * the buffer lives in the checkpoint store (:func:`~yeto.rl.engine.resume.store_for` at ``BUFFER_ROOT_ENV``): a Modal Volume or a shared directory that both islands @@ -59,6 +60,15 @@ encode_group, group_entry_fields, ) +from yeto.rl.engine.merged_release import ( + LANE_ENV, + BufferLane, + LaneCursorLog, + ObjectLaneCursorLog, + gate_version, + lane_cursor_log, + merges_shards, +) SAMPLER, TRAINER, CONVERTIBLE = "sampler", "trainer", "convertible" ROLES = (SAMPLER, TRAINER, CONVERTIBLE) @@ -69,9 +79,14 @@ POLL_ENV = "YETO_SAMPLE_BUFFER_POLL_S" # reader: seconds between index reads while waiting MAX_POLICY_AGE_ENV = "YETO_SAMPLE_BUFFER_MAX_POLICY_AGE" SHARD_ENV = "YETO_SAMPLE_BUFFER_SHARD" # convertible node: its own shard of the index +# S25 C1: several sampler islands on one shared directory (a Modal Volume): each +# sampler writes its own shard, the trainer merges these shards (comma list) +SHARDS_ENV = "YETO_SAMPLE_BUFFER_SHARDS" +INCARNATION_ENV = "YETO_SAMPLER_INCARNATION" # set by the sampler island before Ray starts # the rollout functions run in Ray workers, which inherit the raylet's environment: # connect_island_ray forwards these to every worker -ISLAND_ENVS = (BUFFER_ROOT_ENV, SAMPLER_ID_ENV, WAIT_ENV, POLL_ENV, MAX_POLICY_AGE_ENV, SHARD_ENV) +ISLAND_ENVS = (BUFFER_ROOT_ENV, SAMPLER_ID_ENV, WAIT_ENV, POLL_ENV, MAX_POLICY_AGE_ENV, SHARD_ENV, + SHARDS_ENV, INCARNATION_ENV, LANE_ENV) DEFAULT_WAIT_S = 3600.0 DEFAULT_POLL_S = 2.0 @@ -111,10 +126,14 @@ def rollout_argv(role: str | None, *, max_policy_age: int, fully_async: bool = F def island_env(role: str, *, store_root: str, sampler_id: str | None = None, max_policy_age: int | None = None, wait_s: float | None = None, - shard: str | None = None) -> dict[str, str]: + shard: str | None = None, shards: Sequence[str] | None = None, + lane: BufferLane | None = None) -> dict[str, str]: """The environment the launch layer exports on the island (read in the rollout process). A convertible node writes like a sampler, into its own ``shard`` of - the index, and reads like a trainer.""" + the index, and reads like a trainer. ``shards``: the sampler shards of a + launch with several sampler islands (a sampler writes the one named by its + sampler id, the trainer merges them all). ``lane``: a synced trainer's + buffer lane.""" if role not in ROLES: raise ValueError(f"sample buffer role {role!r}: expected one of {ROLES}") env = {BUFFER_ROOT_ENV: str(store_root)} @@ -132,9 +151,24 @@ def island_env(role: str, *, store_root: str, sampler_id: str | None = None, env[MAX_POLICY_AGE_ENV] = str(int(max_policy_age)) if wait_s is not None: env[WAIT_ENV] = f"{float(wait_s):g}" + if shards: + if role == CONVERTIBLE: + raise ValueError("a convertible pool names its shards per node, not with shards") + if role == SAMPLER and sampler_id not in shards: + raise ValueError(f"sampler {sampler_id!r} is not one of the buffer shards {list(shards)}") + env[SHARDS_ENV] = ",".join(shards) + if lane is not None: + if role != TRAINER: + raise ValueError("only a trainer reads a buffer lane") + env[LANE_ENV] = lane.env() return env +def buffer_shards(environ: Mapping[str, str]) -> tuple[str, ...]: + """The sampler shards of :data:`SHARDS_ENV` (() when unset).""" + return tuple(s.strip() for s in (environ.get(SHARDS_ENV) or "").split(",") if s.strip()) + + def open_buffer(environ: Mapping[str, str] | None = None, *, batch_size: int = 1, shard: str | None = None, create_index: bool = True, sampler: bool = False, reader: bool = False) -> Any: @@ -142,9 +176,11 @@ def open_buffer(environ: Mapping[str, str] | None = None, *, batch_size: int = 1 with ``shard``, the writer of that shard of the index; ``create_index=False`` for a reader whose index another host writes. - On a store that needs one shard per sampler island (``sample_buffer.needs_shards``): + On a store that needs one shard per sampler island (``sample_buffer.needs_shards``), + or with :data:`SHARDS_ENV` (several sampler islands on one shared directory): ``sampler`` opens the shard named by :data:`SAMPLER_ID_ENV`, and ``reader`` - (the trainer) merges the shards into the index.""" + (the trainer) merges the shards into the index (with :data:`SHARDS_ENV`, the + data ends once every named shard has ended).""" from yeto.rl.engine.resume import store_for from yeto.rl.engine.sample_buffer import INDEX, FileSampleBuffer, needs_shards, shard_index @@ -154,15 +190,25 @@ def open_buffer(environ: Mapping[str, str] | None = None, *, batch_size: int = 1 raise SampleBufferRolloutError(f"{BUFFER_ROOT_ENV} is not set") store = store_for(root, environ=env) merge = False - if needs_shards(store) and shard is None: + named = buffer_shards(env) + if (needs_shards(store) or named) and shard is None: if sampler: shard = env.get(SAMPLER_ID_ENV) if not shard: raise SampleBufferRolloutError(f"{SAMPLER_ID_ENV} is not set (the sampler's buffer shard)") merge = reader + lane = BufferLane.parse(env.get(LANE_ENV)) if reader else None + merge = merge and merges_shards(lane) return FileSampleBuffer(store, batch_size=batch_size, index=INDEX if shard is None else shard_index(shard), - create_index=create_index, merge_shards=merge) + create_index=create_index, merge_shards=merge, + expected_shards=named if merge and named else None, + lane=lane) + + +def open_lane_log(buffer: Any) -> LaneCursorLog | ObjectLaneCursorLog | None: + """The lane cursor log of a lane reader (None for a single reader).""" + return lane_cursor_log(buffer) def renumber_group(samples: Sequence[dict[str, Any]], *, position: int, @@ -191,6 +237,85 @@ def renumber_group(samples: Sequence[dict[str, Any]], *, position: int, return out +# ---------------------------------------------------------------------- sampler side +_DATA_CURSOR_FIELDS = ("sample_offset", "epoch_id", "sample_group_index", "sample_index") + + +def seek_data_source(data_source: Any, cursor: Mapping[str, int], *, shuffle: bool) -> None: + """Put a Miles ``RolloutDataSource`` at ``cursor`` the way its ``load`` does + (the four position fields, then the epoch's shuffle). Raises when the source + does not have the fields (a source this resume does not know).""" + source = getattr(data_source, "__self__", data_source) + missing = [f for f in _DATA_CURSOR_FIELDS if not isinstance(getattr(source, f, None), int)] + if missing or set(cursor) != set(_DATA_CURSOR_FIELDS): + raise SampleBufferRolloutError( + f"cannot seek data source {type(source).__name__} to {dict(cursor)} " + f"(missing fields {missing})") + for name in _DATA_CURSOR_FIELDS: + setattr(source, name, int(cursor[name])) + if shuffle and getattr(source, "dataset", None) is not None: + source.dataset.shuffle(int(cursor["epoch_id"])) + + +class SamplerRoundWriter: + """The sampler's buffer writes, resumable across incarnations (S25 C1). + + Opened in the rollout process. On open it reads where this sampler's earlier + incarnations got to (:func:`~yeto.rl.engine.sampler_resume.resume_point`); + :meth:`seek` puts the data source at the resume cursor before the first + round; :meth:`write_round` writes one round's finished groups, each with its + :class:`~yeto.rl.engine.sampler_resume.SamplerRound` record. After the + reclaim marker appears (the island is leaving) nothing more is written: an + old incarnation must not write while its replacement writes the same shard.""" + + def __init__(self, buffer: Any, sampler_id: str, *, incarnation: str, resume: bool = True, + reclaimed: Callable[[], bool] = lambda: False) -> None: + """``resume=False``: write as before (a convertible node: no resume + point, no round check, no seek).""" + from yeto.rl.engine.sampler_resume import RoundGuard, resume_point + + self.buffer = buffer + self.sampler_id = str(sampler_id) + self.incarnation = str(incarnation) + self.reclaimed = reclaimed + self.resume = resume_point(buffer.entries() if resume else (), self.sampler_id) + self.guard = RoundGuard(self.resume) if resume else None + self._writer = GroupWriter(buffer, f"{self.sampler_id}-{self.incarnation}") + self._sought = False + self.dropped_groups = 0 # groups not written because the island is leaving + + def seek(self, data_source: Any, *, shuffle: bool) -> dict[str, int] | None: + """Before the first round only: the resume cursor applied (None: none).""" + if self._sought or self.guard is None: + return None + self._sought = True + cursor = self.resume.data_cursor + if cursor is not None: + seek_data_source(data_source, cursor, shuffle=shuffle) + return cursor + + def write_round(self, rollout_id: int, groups: Sequence[tuple[str, list[dict[str, Any]]]], + data_cursor: Mapping[str, int] | None) -> int: + """Write ``groups`` (``(group key, samples)``, aborted groups already left + out) of round ``rollout_id``; returns the groups written.""" + from yeto.rl.engine.sampler_resume import SamplerRound + + if self.guard is not None: + self.guard.check(rollout_id) + written = 0 + for i, (key, samples) in enumerate(groups): + if self.reclaimed(): + self.dropped_groups += len(groups) - i + break + record = SamplerRound(self.sampler_id, self.incarnation, int(rollout_id), + data_cursor=data_cursor, last=i == len(groups) - 1) + self._writer.write(samples, group_key=key, sampler_round=record.record()) + written += 1 + if written and self.guard is not None: + self.guard.written(rollout_id) + return written + + # ---------------------------------------------------------------------- trainer side @dataclass class TakeReport: @@ -215,7 +340,10 @@ class BatchReader: def __init__(self, buffer: Any, *, max_policy_age: int, wait_s: float = DEFAULT_WAIT_S, poll_s: float = DEFAULT_POLL_S, consumer: str = TRAINER_CONSUMER, clock: Callable[[], float] = time.monotonic, - sleep: Callable[[float], None] = time.sleep) -> None: + sleep: Callable[[float], None] = time.sleep, + cursor_log: LaneCursorLog | ObjectLaneCursorLog | None = None) -> None: + """``cursor_log``: a synced trainer's lane cursor (:mod:`yeto.rl.engine.merged_release`): + the reader continues after its last line and appends one per round.""" if int(max_policy_age) < 1: raise ValueError("max_policy_age must be >= 1") self.buffer = buffer @@ -226,6 +354,11 @@ def __init__(self, buffer: Any, *, max_policy_age: int, wait_s: float = DEFAULT_ self.clock = clock self.sleep = sleep self._ends: dict[int, int] = {} # rollout id -> buffer position after its round + # a synced trainer gates on the merged version it last published (set_gate_version) + self.gate_version: int | None = None + self.cursor_log = cursor_log + if cursor_log is not None: + self.buffer.reset_to(cursor_log.last_cursor()) def take(self, rollout_id: int, groups: int, *, current_version: int | None ) -> tuple[list[tuple[int, list[dict[str, Any]]]], TakeReport]: @@ -240,6 +373,7 @@ def take(self, rollout_id: int, groups: int, *, current_version: int | None out: list[tuple[int, list[dict[str, Any]]]] = [] versions: dict[int, int] = {} stale = 0 + current_version = gate_version(current_version, self.gate_version) floor = None if current_version is None else int(current_version) - self.max_policy_age while len(out) < groups: selected = self.buffer.select(self.consumer, groups - len(out)) @@ -264,6 +398,8 @@ def take(self, rollout_id: int, groups: int, *, current_version: int | None raise SampleBufferRolloutError(f"rollout {rollout_id}: the sample buffer ended " f"with no trainable group (position {end})") self._ends[rollout_id] = end + if self.cursor_log is not None: + self.cursor_log.append(rollout_id, [position for position, _ in out], end) return out, TakeReport(rollout_id=rollout_id, groups=len(out), stale_groups=stale, start=start, end=end, behavior_versions=dict(sorted(versions.items())), wait_s=round(self.clock() - started, 3)) @@ -340,6 +476,17 @@ def executor_read_shards(_executor: Any) -> dict[str, int] | None: return reader.buffer.shard_positions(reader.buffer.read_cursor) or None +def executor_set_gate_version(_executor: Any, version: int) -> int | None: + """Runs INSIDE the rollout executor actor: a synced trainer's merged version, + the version its ``max_policy_age`` gate compares against. Returns the previous + one (None before the reader exists: nothing to set).""" + reader = _ACTIVE.get(TRAINER) + if reader is None: + return None + previous, reader.gate_version = reader.gate_version, int(version) + return previous + + def executor_reset(_executor: Any, watermark: int) -> int: """Runs INSIDE the rollout executor actor (``__ray_call__``).""" return _active_reader().reset(int(watermark)) diff --git a/yeto/rl/adapters/miles/buffer_rollout_miles.py b/yeto/rl/adapters/miles/buffer_rollout_miles.py index e9f6bad8..2f285e2b 100644 --- a/yeto/rl/adapters/miles/buffer_rollout_miles.py +++ b/yeto/rl/adapters/miles/buffer_rollout_miles.py @@ -9,7 +9,6 @@ import asyncio import os -import uuid from typing import Any from miles.rollout.base_types import ( @@ -23,11 +22,13 @@ from miles.utils.types import Sample from yeto.rl.engine.buffer_payload import complete_prompt_tokens +from yeto.rl.engine.sampler_resume import new_incarnation, reclaim_requested from yeto.rl.engine.offpolicy_source import flat_metrics, offpolicy_source_event, payload_source_sample from .buffer_rollout import ( DEFAULT_POLL_S, DEFAULT_WAIT_S, + INCARNATION_ENV, MAX_POLICY_AGE_ENV, POLL_ENV, SAMPLER_ID_ENV, @@ -35,13 +36,15 @@ TRAINER, WAIT_ENV, BatchReader, - GroupWriter, SampleBufferRolloutError, + SamplerRoundWriter, convertible_mode, open_buffer, + open_lane_log, register_reader, renumber_group, ) +from .rollout_meta_hook import data_cursor, record_round_metadata def _flat(group: Any) -> list[Any]: @@ -52,32 +55,47 @@ def _flat(group: Any) -> list[Any]: class BufferWriterRolloutFn(InferenceRolloutFn): - """Sampler: Miles' rollout; every finished group is also written to the buffer.""" + """Sampler: Miles' rollout; every finished group is also written to the buffer. + + A replacement incarnation of the island continues where the buffer says the + earlier ones stopped (:class:`.buffer_rollout.SamplerRoundWriter`): the data + source is put at the saved cursor before the first round, and a round + already in the buffer is refused.""" + + resumable = True # a convertible node keeps its own rollout numbering (no resume) def __init__(self, input: RolloutFnConstructorInput) -> None: super().__init__(input) sampler_id = os.environ.get(SAMPLER_ID_ENV) if not sampler_id: raise SampleBufferRolloutError(f"{SAMPLER_ID_ENV} is not set") - # a per-process suffix: a restarted sampler repeats Miles group indexes, so the - # sample ids stay unique; the trainer renumbers groups by buffer position - self._writer = GroupWriter(self._open_writer_buffer(), f"{sampler_id}-{uuid.uuid4().hex[:8]}") + # the incarnation suffix keeps the sample ids unique: a restarted sampler repeats + # Miles group indexes; the trainer renumbers groups by buffer position + incarnation = os.environ.get(INCARNATION_ENV) or new_incarnation() + self._rounds = SamplerRoundWriter(self._open_writer_buffer(), sampler_id, + incarnation=incarnation, resume=self.resumable, + reclaimed=lambda: reclaim_requested(os.environ)) + self._shuffle = bool(getattr(input.args, "rollout_shuffle", False)) # the seek re-shuffles def _open_writer_buffer(self) -> Any: return open_buffer(sampler=True) async def _call_train(self, input: Any) -> Any: + self._rounds.seek(self.data_source, shuffle=self._shuffle) output = await super()._call_train(input) - await asyncio.to_thread(self._write, input.rollout_id, output.samples) + cursor, _ = data_cursor(self.data_source) + await asyncio.to_thread(self._write, input.rollout_id, output.samples, cursor) return output - def _write(self, rollout_id: int, groups: list[Any]) -> None: + def _write(self, rollout_id: int, groups: list[Any], cursor: dict[str, int] | None) -> None: + finished = [] for group in groups: samples = _flat(group) if any(s.status == Sample.Status.ABORTED for s in samples): continue # an aborted group is regenerated, never trained - key = f"r{int(rollout_id)}-g{samples[0].group_index}" - self._writer.write([s.to_dict() for s in samples], group_key=key) + finished.append((f"r{int(rollout_id)}-g{samples[0].group_index}", + [s.to_dict() for s in samples])) + self._rounds.write_round(rollout_id, finished, cursor) class _BufferRound: @@ -94,8 +112,9 @@ def __init__(self, args: Any, *, create_index: bool = True, merge_shards: bool = self.args = args self.groups = int(args.rollout_batch_size) self.samples_per_group = int(args.n_samples_per_prompt) + buffer = open_buffer(batch_size=self.groups, create_index=create_index, reader=merge_shards) self.reader = BatchReader( - open_buffer(batch_size=self.groups, create_index=create_index, reader=merge_shards), + buffer, cursor_log=open_lane_log(buffer), max_policy_age=int(age), wait_s=float(os.environ.get(WAIT_ENV, DEFAULT_WAIT_S)), poll_s=float(os.environ.get(POLL_ENV, DEFAULT_POLL_S))) @@ -117,6 +136,8 @@ def __call__(self, input: RolloutFnInput) -> RolloutFnTrainOutput: samples = [[Sample.from_dict(d) for d in renumber_group( group, position=position, samples_per_group=self.samples_per_group)] for position, group in groups] + record_round_metadata(self.args, int(input.rollout_id), + buffer_stale_groups=int(report.stale_groups)) self._run_round_hooks(samples) return RolloutFnTrainOutput(samples=samples, metrics={**report.metrics(), **flat_metrics(source)}) @@ -158,6 +179,8 @@ class ConvertibleRolloutFn(BufferWriterRolloutFn): :class:`BufferReaderRolloutFn`. The node switches the mode in this process (``buffer_rollout.executor_set_mode``) when it converts.""" + resumable = False + def __init__(self, input: RolloutFnConstructorInput) -> None: super().__init__(input) # the head merges the shards into the index: this node never writes it diff --git a/yeto/rl/adapters/miles/config.py b/yeto/rl/adapters/miles/config.py index b6ead716..dadbb341 100644 --- a/yeto/rl/adapters/miles/config.py +++ b/yeto/rl/adapters/miles/config.py @@ -25,6 +25,7 @@ import json import os import sys +from pathlib import Path from collections.abc import Callable, Iterator, Mapping, Sequence from dataclasses import dataclass, field from typing import Any @@ -869,7 +870,9 @@ def translate_run_config( "--train-backend", "megatron", "--hf-checkpoint", config.hf_checkpoint, "--ref-load", config.ref_load, - "--megatron-to-hf-mode", "bridge", + # glm5_next: Megatron-Bridge has no glm5_next bridge; Miles loads the torch_dist + # checkpoint and exports through its own converters (raw) + "--megatron-to-hf-mode", "raw" if recipe.name == RECIPE_GLM5_NEXT else "bridge", "--model-name", model_name, *model_recipe_values, "--num-layers", str(geometry.num_layers), @@ -1150,8 +1153,13 @@ def translate_run_config( values.extend((flag, engine)) if recipe.name == RECIPE_GLM5_NEXT: # the plugin spec replaces the Bridge-geometry model flags + # experts and dense-layer count from the checkpoint (a cut-down model differs from GLM5.3-Flash) + hf_config = Path(config.hf_checkpoint) / "config.json" + shape = (glm5_next.checkpoint_geometry(json.loads(hf_config.read_text())) + if hf_config.is_file() else {}) + shape["num_layers"] = geometry.num_layers values = glm5_next.apply_recipe(values, glm5_next.model_argv( - num_layers=geometry.num_layers, transformer_impl=recipe.transformer_impl)) + **shape, transformer_impl=recipe.transformer_impl)) values.extend(extra_argv) reject_fault_tolerance_flags(values[1:]) # generated argv never carries FT flags _refuse_decoupled_ppo_breakers(values) # also when the spec emitted the flag diff --git a/yeto/rl/adapters/miles/cut_plugin.py b/yeto/rl/adapters/miles/cut_plugin.py index bb22edb0..50fe3afd 100644 --- a/yeto/rl/adapters/miles/cut_plugin.py +++ b/yeto/rl/adapters/miles/cut_plugin.py @@ -165,8 +165,13 @@ def _state_dicts(optimizer: Any) -> list[Any]: if obj is None or id(obj) in seen: continue seen.add(id(obj)) - stack.extend(getattr(obj, "chained_optimizers", None) or ()) - stack.append(getattr(obj, "optimizer", None)) + chained = getattr(obj, "chained_optimizers", None) + if chained: + # Megatron's ChainedOptimizer.optimizer asserts a single member (an MoE model + # chains the dense and the expert optimizer; Modal run-14 failed the cut here) + stack.extend(chained) + else: + stack.append(getattr(obj, "optimizer", None)) state = getattr(obj, "state", None) if isinstance(state, Mapping): out.append(state) diff --git a/yeto/rl/adapters/miles/harbor_agent.py b/yeto/rl/adapters/miles/harbor_agent.py index 7506cf57..52caae11 100644 --- a/yeto/rl/adapters/miles/harbor_agent.py +++ b/yeto/rl/adapters/miles/harbor_agent.py @@ -44,7 +44,7 @@ from yeto.rl.agentic.loop_local import LoopLocal from yeto.rl.agentic.trust import check_agent_env from yeto.rl.harness.harbor import CODEX_VERSION, HarborHarness, HarborSettings, build_trial_config -from yeto.rl.harness.harbor.trial import CUSTOM_PROVIDERS, TrialConfigBuilder +from yeto.rl.harness.harbor.trial import CODEX_KWARGS, CUSTOM_PROVIDERS, TrialConfigBuilder EXAMPLE_MODULE = "harbor_agent_function" AGENT_FUNCTION = "yeto.rl.adapters.miles.harbor_agent.run" @@ -105,6 +105,8 @@ def build(task: AgentTask, session_url: str) -> Any: # The example has no codex binding; its default (OPENAI_BASE_URL) is # what Harbor's codex agent reads. Pin the build: Harbor's default is @latest. config.agent.kwargs.setdefault("version", task.agent_version or CODEX_VERSION) + for key, value in CODEX_KWARGS.items(): + config.agent.kwargs.setdefault(key, value) return config return build diff --git a/yeto/rl/adapters/miles/island_entry.py b/yeto/rl/adapters/miles/island_entry.py index 9a031e33..48585224 100644 --- a/yeto/rl/adapters/miles/island_entry.py +++ b/yeto/rl/adapters/miles/island_entry.py @@ -192,6 +192,7 @@ def _check_ports_infra_switches(args) -> None: if getattr(args, "rl_elastic_declare_cells", False) and not getattr(args, "rl_elastic", False): raise ValueError("--rl-elastic-declare-cells needs --rl-elastic") _check_resume_switches(args) + _check_sample_store(args) if not getattr(args, "rl_elastic", False): if given: raise ValueError(", ".join(given) + " need --rl-elastic") @@ -210,6 +211,13 @@ def _check_ports_infra_switches(args) -> None: raise ValueError(f"--{name.replace('_', '-')} must be positive") +def _check_sample_store(args) -> None: + """S25 C3: ``--rl-sample-store`` / ``--rl-sample-lanes`` (engine-neutral rule).""" + from yeto.rl.engine.sample_roles import check_sample_store + + check_sample_store(args) + + def _check_resume_switches(args) -> None: """rl-resume-from-checkpoint: ``--rl-resume-store`` is single-island/no-sync; the cut cadence flags need a store (elastic or resume).""" @@ -258,10 +266,14 @@ def apply_sample_role(args, miles_args, environ=None) -> None: config = sample_roles.resolve(args, groups_per_round=args.groups_per_round) miles_args.yeto_rl_sample_buffer = config node = getattr(args, "rl_convertible_node", None) # L7.5: a convertible node's own shard + # a pair island runs as learner 0: its sampler id is its node name (S25 C1) + sampler = node or getattr(args, "rl_sampler_node", None) or f"island{int(args.learner_id)}" (os.environ if environ is None else environ).update(island_env( - config.role, store_root=config.store, sampler_id=node or f"island{int(args.learner_id)}", + config.role, store_root=config.store, sampler_id=sampler, max_policy_age=int(getattr(args, "rl_max_policy_age", 0) or 0), wait_s=config.wait_s, - shard=node)) + shard=node, shards=getattr(args, "rl_buffer_shards", None), + lane=sample_roles.buffer_lane(args))) + miles_args.yeto_rl_sample_lane = sample_roles.buffer_lane(args) # S25 C3: None off a synced trainer def apply_ports_infra_switches(args, miles_args, environ=None) -> None: @@ -1092,6 +1104,20 @@ def _preflight_codex_openenv_adapter(args, profile_name: str) -> None: forward_legacy_openenv_preflight(args, profile_name) +def miles_revision_pins(args: Any) -> Any: + """The Miles commit the island checks for ``--miles-commit`` (full SHA); None + without it (``verify_miles_revision`` then checks the pinned MILES_NEXT_PINS).""" + from yeto.provenance import is_immutable_commit + from yeto.rl.adapters.miles.pins import MILES_NEXT_REPOSITORY, MilesRevisionPins + + commit = getattr(args, "miles_commit", None) + if not commit: + return None + if not is_immutable_commit(commit) or len(str(commit)) != 40: + raise ValueError(f"--miles-commit must be a full 40-hex commit, got {commit!r}") + return MilesRevisionPins(MILES_NEXT_REPOSITORY, str(commit).lower()) + + def _preflight_codex_harness(args) -> None: """Fail before model allocation if the signed Codex surface has drifted.""" @@ -1787,6 +1813,18 @@ def run_miles( ) attention_specs = specs = () canonical_targets = flash_next.target_modules() + elif getattr(args, "rl_model_recipe", "generic") == "glm5_next": + # AutoBridge cannot read glm5_next (transformers has no such model type); + # Miles trains it through its plugin spec, so the island reads the geometry + # from config.json (models/glm5_next.provider_view). Full-parameter only. + from yeto.rl.adapters.miles.models import glm5_next as glm5_next_recipe + + if not full_parameter: + raise ValueError("the glm5_next recipe is full-parameter only") + model_bridge = None + provider = glm5_next_recipe.glm5_next_provider(model_path) + attention_specs = specs = () + canonical_targets = [] else: from megatron.bridge import AutoBridge @@ -2340,7 +2378,11 @@ def main(argv=None) -> None: miles_root = str(Path(args.miles_root).expanduser().resolve()) if miles_root not in sys.path: sys.path.insert(0, miles_root) - verify_miles_revision(miles_root) + pins = miles_revision_pins(args) + if pins is None: + verify_miles_revision(miles_root) # the pinned MILES_NEXT_PINS + else: + verify_miles_revision(miles_root, expected=pins) _preflight_codex_harness(args) probe_syncers(args) diff --git a/yeto/rl/adapters/miles/models/glm5_next.py b/yeto/rl/adapters/miles/models/glm5_next.py index 5ab9b8c5..0cb3c4d6 100644 --- a/yeto/rl/adapters/miles/models/glm5_next.py +++ b/yeto/rl/adapters/miles/models/glm5_next.py @@ -21,7 +21,10 @@ from __future__ import annotations +import dataclasses +import json from collections.abc import Iterator, Sequence +from pathlib import Path MODEL_NAME = "glm5_next" TRANSFORMER_ENGINE = "transformer_engine" @@ -50,7 +53,6 @@ "--v-head-dim", "256", "--kv-channels", "256", "--qk-layernorm", - "--num-experts", str(NUM_EXPERTS), "--moe-router-topk", "8", "--moe-router-score-function", "sigmoid", "--moe-router-pre-softmax", @@ -93,6 +95,20 @@ def moe_layer_freq(num_layers: int = NUM_LAYERS, first_k_dense_replace: int = FI return "[" + ",".join(["0"] * first_k_dense_replace + ["1"] * (num_layers - first_k_dense_replace)) + "]" +def checkpoint_geometry(config: dict) -> dict[str, int]: + """``num_layers``, ``first_k_dense_replace`` and ``num_experts`` of an HF + ``config.json`` (the text config of ``Glm5NextForConditionalGeneration``), + as :func:`model_argv` keywords; a key the config lacks keeps the default.""" + text = config.get("text_config") or config + out = {} + for key, hf_key in (("num_layers", "num_hidden_layers"), + ("first_k_dense_replace", "first_k_dense_replace"), + ("num_experts", "n_routed_experts")): + if text.get(hf_key) is not None: + out[key] = int(text[hf_key]) + return out + + def check_transformer_impl(transformer_impl: str) -> str: if transformer_impl not in TRANSFORMER_IMPLS: raise ValueError(f"transformer_impl must be one of {TRANSFORMER_IMPLS}, got {transformer_impl!r}") @@ -100,15 +116,27 @@ def check_transformer_impl(transformer_impl: str) -> str: def model_argv(*, num_layers: int = NUM_LAYERS, first_k_dense_replace: int = FIRST_K_DENSE_REPLACE, - transformer_impl: str = TRANSFORMER_ENGINE) -> list[str]: + num_experts: int = NUM_EXPERTS, transformer_impl: str = TRANSFORMER_ENGINE) -> list[str]: """The model arguments of the recipe (module doc) and the run arguments the - model needs (:data:`RUN_ARGV`).""" + model needs (:data:`RUN_ARGV`). ``num_layers``, ``first_k_dense_replace`` and + ``num_experts`` default to GLM5.3-Flash; a cut-down checkpoint (e.g. Miles' + 4-layer CI model) passes its own (:func:`checkpoint_geometry`).""" impl = check_transformer_impl(transformer_impl) + if int(num_experts) < 1: + raise ValueError(f"num_experts must be >= 1, got {num_experts}") return ["--spec", *SPEC, "--num-layers", str(int(num_layers)), *_MODEL_ARGV, + "--num-experts", str(int(num_experts)), "--moe-layer-freq", moe_layer_freq(num_layers, first_k_dense_replace), *(TE_ONLY_ARGV if impl == TRANSFORMER_ENGINE else LOCAL_ARGV), *RUN_ARGV] +def model_argv_without_run_flags(argv: Sequence[str]) -> list[str]: + """:func:`model_argv` without :data:`RUN_ARGV` (Miles' checkpoint converter + parses the model arguments only).""" + run = {group[0] for group in _flag_groups(RUN_ARGV)} + return [token for group in _flag_groups(argv) if group[0] not in run for token in group] + + def _flag_groups(argv: Sequence[str]) -> Iterator[list[str]]: group: list[str] = [] for token in argv: @@ -130,3 +158,103 @@ def apply_recipe(argv: Sequence[str], recipe: Sequence[str]) -> list[str]: override.add("--no-gradient-accumulation-fusion") kept = [token for group in _flag_groups(argv[1:]) if group[0] not in override for token in group] return [*argv[:1], *kept, *recipe] + + +# ---------------------------------------------------------------- provider view +# Megatron-Bridge's AutoBridge cannot read a glm5_next checkpoint: transformers +# 5.12.1 has no such model type (Modal run-3), and Miles trains it through its +# plugin spec, not through the Bridge. The island only needs the geometry that +# ``run_config.resolve_rl_run_config`` reads, so it takes this read-only view of +# the checkpoint's config.json (the same approach as the Flash-Next view). +MODEL_TYPES = frozenset({"glm5_next", "glm5_next_text"}) + + +@dataclasses.dataclass(frozen=True) +class Glm5NextModelProvider: + """Read-only stand-in for a Megatron-Bridge provider (the fields run_config reads).""" + + num_layers: int + hidden_size: int + num_attention_heads: int + num_query_groups: int + kv_channels: int + ffn_hidden_size: int + num_moe_experts: int + moe_ffn_hidden_size: int + moe_router_topk: int + moe_layer_freq: tuple[int, ...] + moe_shared_expert_intermediate_size: int + vocab_size: int + seq_length: int + max_position_embeddings: int + layernorm_epsilon: float + q_lora_rank: int + kv_lora_rank: int + qk_head_dim: int + qk_pos_emb_head_dim: int + v_head_dim: int + rotary_base: int = 10000 # the recipe's --rotary-base (the checkpoint has no RoPE dims: qk_rope_head_dim 0) + rotary_percent: float = 1.0 + share_embeddings_and_output_weights: bool = False + multi_latent_attention: bool = True + qk_layernorm: bool = True + gated_linear_unit: bool = True + add_bias_linear: bool = False + add_qkv_bias: bool = False + normalization: str = "RMSNorm" + position_embedding_type: str = "rope" + + def finalize(self) -> None: # provider API parity; nothing to derive + return None + + +def is_glm5_next(model_path: str | Path) -> bool: + path = Path(model_path) / "config.json" + if not path.is_file(): + return False + try: + return json.loads(path.read_text(encoding="utf-8")).get("model_type") in MODEL_TYPES + except (OSError, ValueError): + return False + + +def provider_view(config: dict) -> Glm5NextModelProvider: + """The provider view of a glm5_next ``config.json`` (module section above).""" + if config.get("model_type") not in MODEL_TYPES: + raise ValueError(f"not a glm5_next config: model_type={config.get('model_type')!r}") + text = config.get("text_config") or config + layers = int(text["num_hidden_layers"]) + dense = int(text.get("first_k_dense_replace", FIRST_K_DENSE_REPLACE)) + if text.get("hidden_act", "silu") != "silu": + raise ValueError(f"glm5_next expects SwiGLU (silu), got {text.get('hidden_act')!r}") + if bool(text.get("attention_bias", False)): + raise ValueError("glm5_next expects bias-free linears") + max_positions = int(text["max_position_embeddings"]) + return Glm5NextModelProvider( + num_layers=layers, + hidden_size=int(text["hidden_size"]), + num_attention_heads=int(text["num_attention_heads"]), + num_query_groups=int(text.get("num_key_value_heads", text["num_attention_heads"])), + kv_channels=int(text["v_head_dim"]), + ffn_hidden_size=int(text["intermediate_size"]), + num_moe_experts=int(text["n_routed_experts"]), + moe_ffn_hidden_size=int(text["moe_intermediate_size"]), + moe_router_topk=int(text["num_experts_per_tok"]), + moe_layer_freq=tuple([0] * dense + [1] * (layers - dense)), + moe_shared_expert_intermediate_size=int(text["moe_intermediate_size"]) * int(text.get("n_shared_experts", 1)), + vocab_size=int(text["vocab_size"]), + seq_length=max_positions, + max_position_embeddings=max_positions, + layernorm_epsilon=float(text["rms_norm_eps"]), + q_lora_rank=int(text["q_lora_rank"]), + kv_lora_rank=int(text["kv_lora_rank"]), + qk_head_dim=int(text["qk_head_dim"]), + qk_pos_emb_head_dim=int(text.get("qk_rope_head_dim", 0)), + v_head_dim=int(text["v_head_dim"]), + share_embeddings_and_output_weights=bool(config.get("tie_word_embeddings", False) + or text.get("tie_word_embeddings", False)), + ) + + +def glm5_next_provider(model_path: str | Path) -> Glm5NextModelProvider: + return provider_view(json.loads((Path(model_path) / "config.json").read_text(encoding="utf-8"))) diff --git a/yeto/rl/adapters/miles/publish.py b/yeto/rl/adapters/miles/publish.py index 29bece01..990843ef 100644 --- a/yeto/rl/adapters/miles/publish.py +++ b/yeto/rl/adapters/miles/publish.py @@ -533,6 +533,13 @@ async def _stamp_version(self, version: int, label: str) -> tuple[list[Any], lis async def _publish( self, version: int, token: str ) -> tuple[frozenset[str], dict[str, Any] | None]: + if getattr(self._args, "debug_train_only", False): + # Trainer-only island (Miles --debug-train-only, samples from the store): + # there is no engine to update or stamp. The release written by the + # wrapping ReleasingPublisher is the publication (Modal run-11 stopped at + # the engine stamp: "no rollout engines are eligible"). + self._published_once = True + return frozenset(), None update_weights = self._update_weights_to(version) offload_rollout = bool(getattr(self._args, "offload_rollout", False)) try: diff --git a/yeto/rl/adapters/miles/quant_shards.py b/yeto/rl/adapters/miles/quant_shards.py index b9c7c6be..fc7d5f9c 100644 --- a/yeto/rl/adapters/miles/quant_shards.py +++ b/yeto/rl/adapters/miles/quant_shards.py @@ -14,13 +14,21 @@ quantization layout; the neutral code (:mod:`yeto.rl.quant.shards`) does the rest. Non-expert weights are not mapped here: Miles already converts them to HF names -for the sampler weight update (``convert_to_hf``); exactly one rank per -replica group passes those with :func:`hf_slices` (module doc of +for the sampler weight update (``convert_to_hf``, SGLang names ``model.layers.*``); +exactly one rank per replica group passes those with :func:`hf_slices`, which +renames them to the checkpoint names (``model.language_model.layers.*``; module doc of :mod:`yeto.rl.quant.shards`, "Ownership"). +Both functions hand host (CPU) copies to the quantizer: the CPU path is the one +checked byte for byte against the single-process converter (GPU T-G3, 2026-10-20); +a GPU expert slice with W4A8 AWQ failed there (factors on the host), and GPU +quantization bytes were not compared for W4A8. + Limits (explicit errors): expert tensor parallelism must be 1 (a column split of ``linear_fc2`` would cut the W4A8 per-row scale). Status: CPU tests with -synthetic Megatron names; not run inside Miles. +synthetic Megatron names; on 2xH100 inside a real Miles/Megatron glm5_next +(TP2, EP2; TE grouped and local sequential experts) the merged W4A8 (R2 AWQ) +and MXFP4 outputs equal the single-process output byte for byte (T-G3). """ from __future__ import annotations @@ -60,7 +68,7 @@ def megatron_expert_slices(named_params: Iterable[tuple[str, torch.Tensor]], *, raise ValueError(f"{name}: local expert {local} >= num_local_experts {num_local_experts}") expert = int(ep_rank) * int(num_local_experts) + local prefix = f"{hf_layer_prefix}{int(layer_offset) + int(match['layer'])}.mlp.experts" - weight = tensor.detach() + weight = tensor.detach().cpu() if (match["fc"] or match["fc2"]) == "1": for hf_name, part in split_gate_up(prefix, expert, weight): yield SourceSlice(hf_name, part) @@ -68,7 +76,21 @@ def megatron_expert_slices(named_params: Iterable[tuple[str, torch.Tensor]], *, yield SourceSlice(expert_weight_name(prefix, expert, DOWN), weight) +def checkpoint_name(name: str) -> str: + """Miles ``convert_to_hf`` name -> GLM-5.3 checkpoint name. + + ``convert_to_hf`` emits the SGLang weight-update names (``model.layers....``, + ``model.embed_tokens``, ``model.norm``); the GLM-5.3 checkpoint, and so the + quantization layout, keeps the text model under ``model.language_model.`` + (the same rule as Miles ``Glm5NextBridge._nest_language_model``). Names that + already carry it, ``model.visual.*`` and top-level names (``lm_head.weight``) + are unchanged.""" + if name.startswith("model.") and not name.startswith(("model.language_model.", "model.visual.")): + return "model.language_model." + name[len("model."):] + return name + + def hf_slices(named_tensors: Iterable[tuple[str, torch.Tensor]]) -> Iterator[SourceSlice]: - """Whole tensors already in HF names (Miles ``convert_to_hf`` output).""" + """Whole tensors from Miles ``convert_to_hf``, renamed to checkpoint names (:func:`checkpoint_name`).""" for name, tensor in named_tensors: - yield SourceSlice(name, tensor.detach()) + yield SourceSlice(checkpoint_name(name), tensor.detach().cpu()) diff --git a/yeto/rl/adapters/miles/rebuild_wiring.py b/yeto/rl/adapters/miles/rebuild_wiring.py index ca2009f5..0f729754 100644 --- a/yeto/rl/adapters/miles/rebuild_wiring.py +++ b/yeto/rl/adapters/miles/rebuild_wiring.py @@ -29,104 +29,17 @@ from collections.abc import Callable, Mapping from typing import Any +from yeto.rl.engine.round_cut import CutSource as _CoreCutSource -def _algorithm_identity(algorithm: Any, *, ref_model: Mapping[str, Any] | None) -> Any: - from yeto.rl.engine.cut import AlgorithmIdentity - - runtime_attrs = None - to_attrs = getattr(algorithm, "to_legacy_runtime_attrs", None) - if callable(to_attrs): - runtime_attrs = dict(to_attrs()) or None - return AlgorithmIdentity.from_spec(algorithm, runtime_attrs=runtime_attrs, ref_model=ref_model) - - -class CutSource: - """What the driver contributes to a cut (E2 ``CutContext``) and what a - restore must match (``RestoreExpectation``); shared by the 4.4 same-shape - rebuilder and E3's ``MilesTrainerOps`` (4.7).""" - - def __init__(self, *, driver: Any, trainer: Any, rollout: Any, ledger: Any, algorithm: Any, - backend_fingerprint: str, cut_root: str, global_batch_size: int, - ref_model: Mapping[str, Any] | None = None, - shared_filesystem: bool = True) -> None: - self.driver_ref = driver # the driver, or a zero-arg callable returning it - self.trainer = trainer - self.rollout = rollout - self.ledger = ledger - self.identity = _algorithm_identity(algorithm, ref_model=ref_model) - self.backend_fingerprint = backend_fingerprint - self.cut_root = cut_root - self.global_batch_size = int(global_batch_size) - self.shared_filesystem = shared_filesystem - - @property - def driver(self) -> Any: - ref = self.driver_ref - return ref() if callable(ref) and not hasattr(ref, "published_state") else ref - - def context(self, cut_id: str) -> Any: - from yeto.rl.engine.controller import RebuildRefused - from yeto.rl.engine.cut import CutProgress - from .trainer import CutContext - - driver = self.driver - state = driver.published_state - if state is None or driver.published_version is None: - raise RebuildRefused("no published policy to cut") - cursor = (self.rollout.data_cursor() - if callable(getattr(self.rollout, "data_cursor", None)) else None) - if cursor is None: - raise RebuildRefused("rollout data cursor unknown (needs the elastic metadata)") - data = dict(cursor) - summary = dict(self.ledger.cut_summary()) if self.ledger is not None else {} - if summary.get("engine_buffer_length") is not None: - data.setdefault("buffer_length", summary["engine_buffer_length"]) - in_flight: tuple = () - limit = int(getattr(getattr(driver, "profile", None), "max_policy_age", 0) or 0) - if limit > 0: - # agentic-rollout-utilization 3.3: the unfinished work the rollout - # process holds NOW (buffered groups, suspended agentic trajectories). - exporter = getattr(self.rollout, "export_in_flight", None) - if not callable(exporter): - raise RebuildRefused("the rollout pool cannot export its in-flight trajectories") - from .in_flight import check_export - - exported = exporter() - problems = check_export(exported) - if problems: - raise RebuildRefused("; ".join(problems)) - in_flight = tuple(exported["entries"]) - data["buffer_length"] = int(exported["buffer_groups"]) - summary["carried_over"] = len(in_flight) - summary["max_policy_age"] = limit - progress = CutProgress( - local_step=int(driver.local_step), - scheduler_samples=int(driver.local_step) * self.global_batch_size, - global_batch_size=self.global_batch_size, - next_rollout_id=int(driver.published_version), - policy_version=int(state.policy_version), - policy_hash=state.policy_tensor_hash(), - ) - return CutContext( - root=self.cut_root, cut_id=cut_id, backend_fingerprint=self.backend_fingerprint, - progress=progress, algorithm=self.identity, data=data, ledger=summary, - # The safe point is after boundary() returned non-stop and every - # member acknowledged the policy: the outer commit is settled. - outer={"settled": bool(driver.at_safe_point), "policy_version": state.policy_version, - "policy_token": driver.expected_token}, - shared_filesystem=self.shared_filesystem, - in_flight=in_flight, - ) - - def expectation(self, layout: Mapping[str, int], *, epoch: int) -> Any: - from yeto.rl.engine.cut import RestoreExpectation - - driver = self.driver - return RestoreExpectation( - algorithm=self.identity, layout=dict(layout), - backend_fingerprint=self.backend_fingerprint, local_step=int(driver.local_step), - policy_version=int(driver.published_state.policy_version), epoch=epoch, - ) + +class CutSource(_CoreCutSource): + """:class:`yeto.rl.engine.round_cut.CutSource` with the Miles in-flight export check.""" + + @staticmethod + def check_in_flight(exported: Mapping[str, Any]) -> list[str]: + from .in_flight import check_export + + return check_export(exported) def make_trainer_rebuilder( diff --git a/yeto/rl/adapters/miles/records_generate.py b/yeto/rl/adapters/miles/records_generate.py new file mode 100644 index 00000000..354024ca --- /dev/null +++ b/yeto/rl/adapters/miles/records_generate.py @@ -0,0 +1,327 @@ +"""Miles Harbor path with the records proxy as the token recorder (S25 N10.9). + +Miles' own Harbor example records tokens through Miles' session server, which +speaks Chat Completions only; Harbor's codex agent (0.145.0) speaks Responses +only (T-G2). On this path the yeto records proxy (N10.5) is the recorder, as +on verl: + +* ``--custom-generate-function-path yeto.rl.adapters.miles.records_generate.generate``: + one Harbor trial per sample under the yeto outer layer; the agent calls the + records proxy, the proxy sends token ids to Miles' SGLang router + (:class:`SGLangRouterEngine`), and the recorded calls become the Miles + sample (``join_turns``: tokens, loss mask, rollout log-probs). The reward + is the Harbor verifier's and is set here, so Miles' reward hook is not called. +* ``--custom-reward-post-process-path yeto.rl.adapters.miles.records_generate.reward_post_process``: + GRPO normalisation over each prompt group without its ABORTED members; an + ABORTED sample has advantage 0 and ``remove_sample`` (loss mask 0). Groups + are kept (no dynamic filter), so an ABORTED sample never triggers a new + sandbox. + +A malformed model reply that ends the agent is a trained 0 (POLICY_LIMIT), +not ABORTED (:func:`~yeto.rl.agentic.outcome.policy_output_failure`). + +Rollout-process environment (besides the Harbor/boat variables of +``harbor_agent``): + YETO_RECORDS_PROXY ``yeto.rl.adapters.miles.records_generate:records_proxy_from_env`` + YETO_RECORDS_PROXY_HOST/PORT/PUBLIC_BASE/SECRET_FILE (``gateway.recorder``) + YETO_RECORDS_TEMPLATE, YETO_RECORDS_TOKENIZER, YETO_RECORDS_MODEL, + YETO_RECORDS_SAMPLING, YETO_RECORDS_MAX_MODEL_LEN, YETO_RECORDS_CARD + (``gateway.records_env``) + YETO_SGLANG_ROUTER_URL ``http://:`` of Miles' router (default: from + ``--sglang-router-ip`` / ``--sglang-router-port``) + +Policy version: Miles' synchronous loop updates the engines between +rollouts, never during one; the proxy records version None and the trainer's +rollout id is the version record (as on verl). +""" + +from __future__ import annotations + +import contextlib +import json +import os +from collections.abc import Sequence +from functools import partial +from typing import Any + +from yeto.rl.agentic import ( + OUTCOME_KEY, + AgentTrajectory, + HarnessConfigError, + Outcome, + TokenAlignmentError, + join_turns, + run_agent_task, +) +from yeto.rl.agentic.loop_local import LoopLocal +from yeto.rl.agentic.outcome import policy_output_failure +from yeto.rl.engine.sampler_engine import ( + FINISH_LENGTH, + FINISH_STOP, + UNKNOWN, + Completion, + EngineIdentity, + LoadReport, + SamplerEngineError, + SamplingParams, +) +from yeto.rl.harness.harbor import HarborHarness, HarborSettings, build_trial_config +from yeto.rl.harness.harbor.trial import CUSTOM_PROVIDERS + +from .harbor_agent import gate, task_from_metadata + +ROUTER_URL_ENV = "YETO_SGLANG_ROUTER_URL" +RECORDER_ENV = "YETO_AGENTIC_RECORDER" +DEFAULT_RECORDER = "yeto.rl.harness.gateway.recorder:recorder_from_env" +GENERATE_FUNCTION = "yeto.rl.adapters.miles.records_generate.generate" +REWARD_POST_PROCESS = "yeto.rl.adapters.miles.records_generate.reward_post_process" +HTTP_TIMEOUT_S = 1800.0 + + +# ---------------------------------------------------------------------- engine +class SGLangRouterEngine: + """``TokenSamplerEngine`` over an SGLang router (or server) someone else runs (Miles).""" + + def __init__(self, url: str, *, max_model_len: int, identity: EngineIdentity, http: Any = None) -> None: + self.url = url.rstrip("/") + self._max_model_len = int(max_model_len) + self._identity = identity + self._http = http + + @property + def identity(self) -> EngineIdentity: + return self._identity + + @property + def max_model_len(self) -> int: + return self._max_model_len + + def served_version(self) -> int | None: + return None + + def _post(self, body: dict[str, Any]) -> Any: + if self._http is not None: + return self._http(f"{self.url}/generate", body) + import httpx + + reply = httpx.post(f"{self.url}/generate", json=body, timeout=HTTP_TIMEOUT_S) + reply.raise_for_status() + return reply.json() + + def generate_tokens(self, prompts: Sequence[Sequence[int]], sampling: SamplingParams, *, + stop_token_ids: Sequence[int] = ()) -> list[list[Completion]]: + """Called from a worker thread (the records backend uses ``asyncio.to_thread``).""" + params: dict[str, Any] = { + "temperature": float(sampling.temperature), "top_p": float(sampling.top_p), + "top_k": int(sampling.top_k), "repetition_penalty": float(sampling.repetition_penalty), + "max_new_tokens": int(sampling.max_tokens), "skip_special_tokens": False, + "no_stop_trim": True, + } + if stop_token_ids: + params["stop_token_ids"] = [int(t) for t in stop_token_ids] + out: list[list[Completion]] = [] + for prompt in prompts: + group = [] + for _ in range(int(sampling.n)): + body = self._post({"input_ids": [int(t) for t in prompt], "sampling_params": params, + "return_logprob": True}) + body = body[0] if isinstance(body, list) else body + meta = body.get("meta_info") or {} + pairs = meta.get("output_token_logprobs") or [] + tokens = tuple(int(p[1]) for p in pairs) + finish = (meta.get("finish_reason") or {}).get("type") + if finish == "abort": + raise SamplerEngineError(f"sglang aborted the request: {meta.get('finish_reason')}") + group.append(Completion(token_ids=tokens, logprobs=tuple(float(p[0]) for p in pairs), + text=str(body.get("text", "")), + finish_reason=FINISH_LENGTH if finish == "length" else FINISH_STOP, + prompt_token_ids=tuple(int(t) for t in prompt))) + out.append(group) + return out + + # Miles starts, updates and stops its engines: the proxy never calls these. + def start(self, release_dir: Any, version: int) -> LoadReport: + raise SamplerEngineError("Miles owns its rollout engines") + + def load(self, release_dir: Any, version: int) -> LoadReport: + raise SamplerEngineError("Miles owns its rollout engines") + + def generate(self, prompts: Sequence[str], sampling: SamplingParams) -> list[list[Completion]]: + raise SamplerEngineError("text prompts are not used; the records proxy sends token ids") + + def score(self, prompt: str, continuation: str) -> list[float]: + raise SamplerEngineError("scoring is not used on the Miles rollout") + + def loaded_expert_matrices(self) -> int | None: + return None + + def weight_gib_per_card(self) -> float | None: + return None + + def stop(self) -> None: + return None + + +def records_proxy_from_env() -> Any: + """``YETO_RECORDS_PROXY`` callable: a RecordsProxy over Miles' router (module doc).""" + from yeto.rl.harness.gateway.records_env import ( + CARD_ENV, + MAX_MODEL_LEN_ENV, + MODEL_ENV, + SAMPLING_ENV, + TEMPLATE_ENV, + TOKENIZER_ENV, + ) + from yeto.rl.harness.chat_templates import ChatTemplate + from yeto.rl.harness.gateway.proxy import RecordsProxy + + try: + url = os.environ[ROUTER_URL_ENV] + template = ChatTemplate.load(os.environ[TEMPLATE_ENV], tokenizer_path=os.environ[TOKENIZER_ENV]) + signed = json.loads(os.environ[SAMPLING_ENV]) + engine = SGLangRouterEngine(url, max_model_len=int(os.environ[MAX_MODEL_LEN_ENV]), identity=EngineIdentity( + engine="miles-sglang-router", engine_version=UNKNOWN, device_family="cuda", + card_model=os.environ.get(CARD_ENV, UNKNOWN), precision="bf16", quant_scheme="bf16", + topology=UNKNOWN)) + model = os.environ[MODEL_ENV] + except (KeyError, ValueError) as exc: + raise HarnessConfigError(f"records proxy over Miles' router: bad or missing setting {exc}") from exc + return RecordsProxy(engine, template, model=model, signed_sampling=signed, sampler_id="miles-records") + + +# ---------------------------------------------------------------------- recorder +def _recorder_from_env() -> Any: + import importlib + + spec = os.environ.get(RECORDER_ENV, DEFAULT_RECORDER) + module, _, attr = spec.partition(":") + return getattr(importlib.import_module(module), attr)() + + +_recorders: LoopLocal[Any] = LoopLocal(_recorder_from_env) + + +def _harness(task: Any) -> HarborHarness: + if task.provider not in CUSTOM_PROVIDERS: + raise HarnessConfigError(f"records_generate runs yeto's TrialConfig builder; provider {task.provider!r} " + f"is not one of {sorted(CUSTOM_PROVIDERS)} (use Miles' example path)") + return HarborHarness(partial(build_trial_config, settings=HarborSettings.from_env())) + + +# ---------------------------------------------------------------------- generate +async def generate(input: Any) -> Any: + """Miles ``--custom-generate-function-path`` (module doc).""" + from miles.rollout.base_types import GenerateFnOutput + + args = input.args + # Miles' router address is an argument, not an env variable; the proxy is built from env once per loop. + os.environ.setdefault(ROUTER_URL_ENV, f"http://{args.sglang_router_ip}:{args.sglang_router_port}") + sample = await run_sample(input.sample, input.sampling_params, args, recorder=_recorders.get()) + return GenerateFnOutput(samples=sample) + + +async def run_sample(sample: Any, sampling_params: dict[str, Any], args: Any, *, recorder: Any, + harness: Any = None) -> Any: + """One Harbor trial -> the filled Miles sample (separate from ``generate`` for CPU tests).""" + task = task_from_metadata(dict(sample.metadata or {})) + session = f"miles-{task.task_id}-{os.urandom(6).hex()}" + endpoint = await recorder.open(session, _verl_style_sampling(sampling_params)) + try: + trajectory = await run_agent_task(task, harness or _harness(task), endpoint, gate()) + except BaseException: + with contextlib.suppress(Exception): # the trial's error is the one to report + await recorder.close(session) + raise + turns: Sequence[Any] = () + try: + turns = await recorder.close(session) + except TokenAlignmentError as exc: + trajectory = _reclassify(trajectory, Outcome.ABORTED, f"token_alignment: {exc}") + else: + updated = policy_output_failure(trajectory, [t.policy_error for t in turns if t.policy_error]) + if updated.outcome is not trajectory.outcome: + gate().metrics.reclassify(trajectory.outcome, updated.outcome) + trajectory = updated + segments = None + if trajectory.outcome is not Outcome.ABORTED: + try: + segments = join_turns(turns, max_response_len=int(args.rollout_max_response_len)) + except TokenAlignmentError as exc: + trajectory = _reclassify(trajectory, Outcome.ABORTED, f"token_alignment: {exc}") + return fill_sample(sample, trajectory, segments, turns) + + +def _reclassify(trajectory: AgentTrajectory, outcome: Outcome, reason: str) -> AgentTrajectory: + if trajectory.outcome is not outcome: + gate().metrics.reclassify(trajectory.outcome, outcome) + return AgentTrajectory(task=trajectory.task, outcome=outcome, reward=0.0 if outcome is Outcome.ABORTED + else trajectory.reward, reason=reason, verdict=trajectory.verdict, + policy_version=trajectory.policy_version) + + +def _verl_style_sampling(sampling_params: dict[str, Any]) -> dict[str, Any]: + """Miles' SGLang sampling keys -> the recorder's check keys (temperature, top_p, top_k).""" + return {k: sampling_params[k] for k in ("temperature", "top_p", "top_k", "repetition_penalty") + if k in sampling_params} + + +def fill_sample(sample: Any, trajectory: AgentTrajectory, segments: Any, turns: Sequence[Any]) -> Any: + status = _status() + metadata = {**(sample.metadata or {}), **trajectory.to_metadata(), + "length_capped_turns": sum(t.finish_reason == FINISH_LENGTH for t in turns), + "model_calls": len(turns)} + if segments is None: # ABORTED: one masked placeholder token, never trained + sample.tokens, sample.response_length = [0, 0], 1 + sample.loss_mask, sample.rollout_log_probs = [0], [0.0] + sample.remove_sample = True + sample.status = status.ABORTED + sample.reward = 0.0 + else: + sample.tokens = [*segments.prompt_ids, *segments.response_ids] + sample.response_length = len(segments.response_ids) + sample.loss_mask = list(segments.response_mask) + sample.rollout_log_probs = list(segments.response_logprobs) + sample.status = status.TRUNCATED if segments.truncated else status.COMPLETED + sample.reward = float(trajectory.reward) + metadata["truncated"] = bool(segments.truncated) + sample.metadata = metadata + return sample + + +def _status() -> Any: + try: + from miles.utils.types import Sample + + return Sample.Status + except ImportError: # CPU tests without Miles + from enum import Enum + + return Enum("Status", "PENDING COMPLETED TRUNCATED ABORTED FAILED") + + +# ---------------------------------------------------------------------- rewards +def is_aborted(sample: Any) -> bool: + return (sample.metadata or {}).get(OUTCOME_KEY) == Outcome.ABORTED.value + + +def reward_post_process(args: Any, samples: Sequence[Any]) -> tuple[list[float], list[float]]: + """GRPO over each group of ``n_samples_per_prompt`` without its ABORTED members (module doc).""" + n = int(args.n_samples_per_prompt) + if len(samples) % n: + raise ValueError(f"{len(samples)} samples are not whole groups of {n}") + by_std = bool(getattr(args, "grpo_std_normalization", True)) + raw = [0.0 if is_aborted(s) else float(s.reward) for s in samples] + out = [0.0] * len(samples) + for start in range(0, len(samples), n): + members = [i for i in range(start, start + n) if not is_aborted(samples[i])] + if not members: + continue + values = [raw[i] for i in members] + mean = sum(values) / len(values) + std = (sum((v - mean) ** 2 for v in values) / (len(values) - 1)) ** 0.5 if len(values) > 1 else 0.0 + for i in members: + out[i] = (raw[i] - mean) / (std + 1e-6) if by_std else raw[i] - mean + for sample in samples: + if is_aborted(sample): + sample.remove_sample = True + return raw, out diff --git a/yeto/rl/adapters/miles/rollout.py b/yeto/rl/adapters/miles/rollout.py index f5490f47..df2080d5 100644 --- a/yeto/rl/adapters/miles/rollout.py +++ b/yeto/rl/adapters/miles/rollout.py @@ -476,6 +476,9 @@ def handle_from_metadata( nonzero_advantages=( int(payload["nonzero_advantages"]) if "nonzero_advantages" in payload else None ), + buffer_stale_groups=( + int(payload["buffer_stale_groups"]) if "buffer_stale_groups" in payload else None + ), policy_age_violation=( int(payload["policy_age_violation"]) if "policy_age_violation" in payload else None ), @@ -1068,7 +1071,16 @@ def _offload_after_rollout(self) -> None: def abort(self) -> None: self._run(self._controller.abort_all()) + @property + def engineless(self) -> bool: + """True on a trainer-only island (Miles --debug-train-only): samples come + from the store and no engine serves a policy, so a publication has no + members to acknowledge it (the driver accepts an empty member set).""" + return bool(getattr(self._args, "debug_train_only", False)) + def members(self) -> frozenset[str]: + if self.engineless: + return frozenset() return self._run(running_members(self._controller)) # -- E1 membership verbs (3.4) --------------------------------------------- diff --git a/yeto/rl/adapters/miles/rollout_meta_hook.py b/yeto/rl/adapters/miles/rollout_meta_hook.py index bcb3440d..d34e761d 100644 --- a/yeto/rl/adapters/miles/rollout_meta_hook.py +++ b/yeto/rl/adapters/miles/rollout_meta_hook.py @@ -54,7 +54,9 @@ # rollout's metadata by the driver-side reader, so they are never attributed # to the next round. ROUND_META_SCHEMA = "yeto-rl-round-metadata-v1" -ROUND_METADATA_KEYS = frozenset({"nonzero_advantages"}) +# buffer_stale_groups: groups the sample-buffer reader passed over this round +# because they were older than max_policy_age (consumed, not trained) +ROUND_METADATA_KEYS = frozenset({"nonzero_advantages", "buffer_stale_groups"}) def current_round_id(samples: Sequence[Any] = (), sink: str | None = None) -> int | None: diff --git a/yeto/rl/adapters/miles/round_cut.py b/yeto/rl/adapters/miles/round_cut.py index f43932ac..bfef3fbd 100644 --- a/yeto/rl/adapters/miles/round_cut.py +++ b/yeto/rl/adapters/miles/round_cut.py @@ -1,380 +1,14 @@ -"""rl-multinode-island M4: policy-weight continuation across a machine replacement. - -Without an outer syncer (``--rl-single-island-no-sync``, :class:`~..bridges.LocalOnlySync`) -a restarted learner used to start again at rollout 0 from the startup weights: the -checkpoint store (Q4, C5) carried the journal/ledger, but no LoRA/optimizer state. - -With ``--rl-elastic-checkpoint-store`` set this module keeps a 4.2 cut of the trainer -(every rank's LoRA + optimizer + scheduler + RNG shard, manifest committed last) at -every round-boundary safe point and resumes from the newest one at start: - -* the cut root is ``/round-cuts`` -- the store is visible on every node (sky - Storage MOUNT of the bucket), so the shard of a rank on another node (PP across - nodes) lands where the driver can check it (``shared_filesystem=True``); -* after the manifest is committed, ``round-cut.json`` (the pointer: cut id, next - rollout id, local step, policy hash, writing incarnation) is written into the state - dir and the state dir is synced to the store -- pointer, ledger and journal travel - in ONE store snapshot (``STORE-MANIFEST.json`` last), so a restored pointer never - names a cut newer than the restored ledger; -* at start (in-place restart or a fresh machine after the store restore) the cut the - pointer names is verified (``verify_cut``: algorithm, layout, backend, step, - policy version, every shard digest) and loaded on the fresh trainer - (``MilesTrainerGroup.restore_cut``); the driver then rebases the ledger and seeks - the data source to the pointer's rollout id. Any failure is fatal (no silent - restart at 0 once a pointer exists). -""" +"""Miles round cuts: the cut controller moved to :mod:`yeto.rl.engine.round_cut` +(S25 N7.6b, shared with the verl island); this module keeps the Miles import path.""" from __future__ import annotations -import json -import os -from pathlib import Path -from typing import Any - -from yeto.rl.engine import recovery -from yeto.rl.engine.recovery import CUT_KEY, FINAL_TRIGGERS, MANUAL, PERIODIC, ConsistentCut -from yeto.rl.engine.resume import ROUND_CUTS # under the checkpoint store (not under the state dir) - -ISLAND = "local" # round cuts run without an outer syncer: one island, the local one -POINTER = "round-cut.json" # under the state dir (synced with the ledger/journal) - - -class RoundCutError(RuntimeError): - pass - - -class RoundCutCheckpoint: - """``save`` at a safe point, ``resume`` at sync start (see the module doc). - - rl-resume-from-checkpoint (S17 C4): ``every`` (``--rl-cut-every``) cuts only every N - rounds (plus the last round and our own stop, ``final=True``); the pointer carries - ``run_fingerprint`` / ``lr_at_next_round`` / ``consumed_prompt_ids_digest``; a resume - checks the fingerprint (refused on a difference unless ``allow_config_change``), the - consumed-group digest and, on the first resumed round, the learning rate.""" - - def __init__(self, *, controller: Any, source: Any, trainer: Any, every: int = 1, - fingerprint: dict[str, Any] | None = None, allow_config_change: bool = False, - next_lr: Any = None) -> None: - if getattr(controller, "checkpoint_store", None) is None: - raise ValueError("round cuts need the checkpoint store") - if int(every) < 1: - raise ValueError("--rl-cut-every must be >= 1") - self.controller = controller - self.source = source # rebuild_wiring.CutSource (cut_root is set here) - self.trainer = trainer # MilesTrainerGroup: save_cut / restore_cut / actual_layout - self.root = str(Path(controller.checkpoint_store) / ROUND_CUTS) - source.cut_root = self.root - self.resumed: dict[str, Any] | None = None - self.every = int(every) - self.fingerprint = dict(fingerprint) if fingerprint is not None else None - self.allow_config_change = bool(allow_config_change) - self.next_lr = next_lr # callable(driver) -> list[float] | None - self.pending_lr_check: dict[str, Any] | None = None - # set by resume(): where the rollout data source continues (the cut's cursor) - self.resume_data_cursor: dict[str, int] | None = None - # rl-trainer-start-stop D4: the durable sample buffer's consumed watermark, - # written into the cut record (None: the run reads no sample buffer). - # buffer_watermark(next_rollout_id) commits and returns the position after the - # last trained round; buffer_resume(watermark) continues there after a restore - # (buffer_rollout.wire_sample_buffer sets both). - self.buffer_watermark: Any = None # callable(int) -> int | None - self.buffer_resume: Any = None # callable(int) -> int - - # ---------------------------------------------------------------- pointer - @property - def pointer_path(self) -> Path: - return Path(self.controller.state_dir) / POINTER - - def pointer(self) -> dict[str, Any] | None: - path = self.pointer_path - if not path.is_file(): - return None - return dict(json.loads(path.read_text(encoding="utf-8"))) - - def _write_pointer(self, info: dict[str, Any]) -> None: - tmp = self.pointer_path.with_suffix(".json.tmp") - tmp.write_text(json.dumps(info, sort_keys=True), encoding="utf-8") - with open(tmp, "rb") as fh: - os.fsync(fh.fileno()) - os.replace(tmp, self.pointer_path) - - def _epoch(self) -> int: - return int(self.controller.journal.epochs.config_epoch) - - def _consumed(self) -> tuple[int, str | None]: - ledger = getattr(self.source, "ledger", None) - groups = getattr(ledger, "_consumed_groups", None) - if groups is None: - return 0, None - from yeto.rl.engine.resume import consumed_digest - - return len(groups), consumed_digest(groups) - - # ---------------------------------------------------------------- save - def due(self, rollout_id: int, *, final: bool = False) -> bool: - from yeto.rl.engine.resume import should_cut - - return should_cut(int(rollout_id), every=self.every, final=final) - - def save(self, driver: Any, *, rollout_id: int, final: bool = False, - trigger: str | None = None) -> dict[str, Any] | None: - """Cut the trainer at the safe point before rollout ``rollout_id`` (= the published - version), point the state dir at it, sync the store. Returns the pointer, or None - when nothing was cut (not due, no publication yet, already cut at this rollout, - island RECOVERY_REQUIRED). Raises on a refused/failed cut (the caller logs it). - - ``trigger`` names a final cut's trigger (default ``MANUAL``; ``SCHEDULED`` for - a stop the trainer scheduler ordered); a periodic cut takes none.""" - import time - - from yeto.rl.engine.resume import directory_bytes - - if trigger is not None and (not final or trigger not in FINAL_TRIGGERS): - raise RoundCutError(f"cut trigger {trigger!r} needs final=True and one of {FINAL_TRIGGERS}") - ctl = self.controller - if ctl.recovery_required or driver.published_version is None: - return None - if not self.due(rollout_id, final=final): - return None - if int(driver.published_version) != int(rollout_id): - raise RoundCutError(f"safe point {rollout_id} != published version {driver.published_version}") - current = self.pointer() - if current is not None and int(current["next_rollout_id"]) == int(rollout_id): - return None # just resumed from (or already cut) this point - cut_id = f"r{int(rollout_id):06d}-{ctl.incarnation['id']}" - started = time.monotonic() - context = self.source.context(cut_id) - # colocated + offloaded trainer (1 GPU): its weights/optimizer are not on the - # device between publish and the next train step; wake it for the cut only. - offloaded = bool(getattr(driver, "_trainer_offloaded", False)) - if offloaded: - self.trainer.onload() - try: - self.trainer.save_cut(epoch=self._epoch(), context=context) - finally: - if offloaded: - self.trainer.offload() - save_s = time.monotonic() - started - progress = context.progress - consumed_n, consumed_sha = self._consumed() - next_lr = self.next_lr(driver) if callable(self.next_lr) else None - info = { - "cut_id": cut_id, "root": ROUND_CUTS, "epoch": self._epoch(), - "next_rollout_id": int(progress.next_rollout_id), - "local_step": int(progress.local_step), - "policy_version": int(progress.policy_version), - "policy_hash": progress.policy_hash, - "incarnation": ctl.incarnation["id"], - "incarnation_index": ctl.incarnation.get("index"), - "run_fingerprint": self.fingerprint, - "lr_at_next_round": None if next_lr is None else [float(x) for x in next_lr], - "consumed_groups": consumed_n, - "consumed_prompt_ids_digest": consumed_sha, - "outer_version": getattr(getattr(driver, "sync", None), "outer_version", None), - "final": bool(final), - } - cut = ConsistentCut( - checkpoint_id=cut_id, trigger=(trigger or MANUAL) if final else PERIODIC, - outer_step=info["outer_version"], island_steps={ISLAND: info["local_step"]}, - completed_groups=consumed_n, completed_groups_digest=consumed_sha, - buffer_watermark=(self.buffer_watermark(int(rollout_id)) if callable(self.buffer_watermark) - else None)) - saved = recovery.save(cut, store=ctl.store, write=lambda record: self._write_and_sync(info, record)) - info[CUT_KEY] = saved[CUT_KEY] - synced = saved["synced"] - sync = dict(getattr(ctl, "last_store_sync", None) or {}) - cut_bytes = directory_bytes(Path(self.root) / cut_id) - total_s = time.monotonic() - started - info["store_synced"] = bool(synced) - info["cut_bytes"] = cut_bytes - info["save_s"] = round(save_s, 3) - info["store_copy_s"] = sync.get("copy_s") - info["store_commit_s"] = sync.get("commit_s") - info["total_s"] = round(total_s, 3) - info["write_bytes_per_s"] = (cut_bytes / save_s) if save_s > 0 else None - info["pruned"] = sync.get("pruned") - info["latest_seq"] = sync.get("seq") - info["trainer_onloaded_for_cut"] = offloaded - if not synced: - info["store_error"] = sync.get("error") - return info - - def _write_and_sync(self, info: dict[str, Any], cut_record: dict[str, Any]) -> dict[str, Any]: - """The round cut's ``recovery.save`` writer: the pointer (with the cut record) - into the state dir, then the state dir into the store (which commits).""" - self._write_pointer({**info, CUT_KEY: cut_record}) - synced = self.controller.sync_checkpoint_store("round_cut", skip_if_recovery=True) - # the store sync commits itself (snapshot layout: in write_snapshot; flat layout: - # in the controller, outside its lock), also when it failed: recovery.save must not - return {"synced": bool(synced), "committed": True} - - # ---------------------------------------------------------------- resume - def resume(self, driver: Any) -> dict[str, Any] | None: - """Load the pointed cut into the freshly built trainer and set the driver's local - step. Returns the pointer (``next_rollout_id`` = where the run continues), or None - when this island never cut (fresh run). Raises on any inconsistency.""" - import time - - from yeto.rl.engine.cut import RestoreExpectation - from yeto.rl.engine.resume import check_fingerprint - - info = self.pointer() - if info is None: - if callable(self.buffer_resume): - self.buffer_resume(0) - return None - started = time.monotonic() - rid = int(info["next_rollout_id"]) - checks: dict[str, Any] = {} - diff = {} - if self.fingerprint is not None and info.get("run_fingerprint") is None: - checks["run_fingerprint"] = "absent (pointer written before rl-resume-from-checkpoint)" - elif self.fingerprint is not None: - diff = check_fingerprint(info.get("run_fingerprint"), self.fingerprint, - allow_change=self.allow_config_change) - checks["run_fingerprint"] = "same" if not diff else "changed(allowed)" - ledger = getattr(self.source, "ledger", None) - if ledger is not None and rid > 0: - last = ledger.batch(rid - 1) - if last is None or last.get("state") != "outer_recorded": - raise RoundCutError(f"round cut {info['cut_id']} continues at rollout {rid} but the " - f"ledger has no outer-recorded rollout {rid - 1}: {last}") - checks["ledger_last_round"] = "outer_recorded" - if info.get("consumed_prompt_ids_digest") is not None: - n, sha = self._consumed() - if sha != info["consumed_prompt_ids_digest"] or n != int(info.get("consumed_groups", n)): - raise RoundCutError(f"round cut {info['cut_id']}: the restored ledger's consumed groups " - f"({n}, {sha}) differ from the pointer's " - f"({info.get('consumed_groups')}, {info['consumed_prompt_ids_digest']})") - checks["consumed_prompt_ids_digest"] = "same" - backend_fingerprint = self.source.backend_fingerprint - if self.allow_config_change: - # The backend fingerprint hashes the whole Miles argv, so a longer run - # (--num-rollout, --lr-decay-iters) changes it; allowed changes are recorded. - from yeto.rl.engine.cut import load_manifest - - saved = load_manifest(self.root, info["cut_id"]).runtime.get("backend_fingerprint") - if saved != backend_fingerprint: - checks["backend_fingerprint"] = {"result": "changed(allowed)", "cut": saved, - "run": backend_fingerprint} - backend_fingerprint = saved - expect = RestoreExpectation( - algorithm=self.source.identity, layout=dict(self.trainer.actual_layout()), - backend_fingerprint=backend_fingerprint, - local_step=int(info["local_step"]), policy_version=int(info["policy_version"]), - epoch=self._epoch(), - ) - manifest = self.trainer.restore_cut(info["cut_id"], epoch=self._epoch(), root=self.root, - expect=expect, shared_filesystem=True) - checks["cut_shards_sha256"] = "ok" - if manifest.progress.policy_hash != info["policy_hash"] or manifest.progress.next_rollout_id != rid: - raise RoundCutError(f"round cut {info['cut_id']} manifest disagrees with its pointer") - restored = driver.export_local().policy_tensor_hash() - if restored != info["policy_hash"]: - raise RoundCutError(f"restored trainer holds {restored}, round cut holds {info['policy_hash']}") - checks["policy_hash"] = "same" - checks["policy_version"] = int(info["policy_version"]) - in_flight = self._restore_in_flight(driver, manifest, rid) - if in_flight is not None: - checks["in_flight"] = in_flight - self.resume_data_cursor = _cut_data_cursor(manifest) - if callable(self.buffer_resume): - watermark = (info.get(CUT_KEY) or {}).get("buffer_watermark") - if watermark is None: - raise RoundCutError(f"round cut {info['cut_id']} has no sample buffer watermark " - "but this trainer reads the sample buffer") - checks["buffer_watermark"] = int(self.buffer_resume(int(watermark))) - driver.local_step = int(info["local_step"]) - self.resumed = dict(info) - if info.get("lr_at_next_round") is not None: - self.pending_lr_check = {"rollout_id": rid, "expected": list(info["lr_at_next_round"])} - self.controller._record("round_cut", tx_id=None, action="restore", cut_id=info["cut_id"], - next_rollout_id=rid, local_step=int(info["local_step"]), - policy_hash=info["policy_hash"], cut_incarnation=info["incarnation"], - incarnation=self.controller.incarnation["id"]) - driver.emit("rl_round_cut_restored", rollout_id=rid, cut_id=info["cut_id"], - local_step=int(info["local_step"]), policy_hash=info["policy_hash"], - cut_incarnation=info["incarnation"]) - restored_store = getattr(self.controller, "restored", None) or {} - cut_bytes = sum(int(f.bytes) for f in manifest.files) - driver.emit( - "rl_resume", rollout_id=rid, cut_id=info["cut_id"], - incarnation=self.controller.incarnation.get("index"), - incarnation_id=self.controller.incarnation["id"], - cut_incarnation=info["incarnation"], cut_incarnation_index=info.get("incarnation_index"), - policy_version=int(info["policy_version"]), policy_hash=info["policy_hash"], - local_step=int(info["local_step"]), lr_at_next_round=info.get("lr_at_next_round"), - checks=checks, config_diff=diff or None, - discarded_rounds=list(getattr(self, "discarded_rounds", []) or []), - restore_s=round(time.monotonic() - started, 3), cut_bytes=cut_bytes, - state_restore=restored_store or None, - ) - return info - - def _restore_in_flight(self, driver: Any, manifest: Any, rollout_id: int) -> dict[str, Any] | None: - """agentic-rollout-utilization 3.3: the cut's unfinished trajectories, by the - CURRENT run's limit: within it and not already trained -> continued - (buffered groups back into the data buffer); else discarded and reported. - Suspended agentic trajectories never survive a restart (discarded). - None when the cut carries none.""" - entries = tuple(getattr(manifest, "in_flight", ()) or ()) - if not entries: - return None - limit = int(getattr(getattr(driver, "profile", None), "max_policy_age", 0) or 0) - ledger = getattr(self.source, "ledger", None) - completed = sorted(getattr(ledger, "_consumed_groups", None) or ()) - rollout = getattr(self.source, "rollout", None) - importer = getattr(rollout, "import_in_flight", None) - if not callable(importer): - raise RoundCutError("the cut carries in-flight trajectories but the rollout pool " - "cannot import them") - report = importer(entries, max_policy_age=limit, current_version=int(rollout_id), - completed_group_ids=completed) - driver.emit("rl_in_flight_restored", rollout_id=int(rollout_id), max_policy_age=limit, - entries=len(entries), **report) - return report - - def check_first_round(self, driver: Any, *, rollout_id: int, applied_lrs: Any) -> None: - """Design §4.5, on the boundary of the first resumed round (before it is recorded - or published): the trained learning rate continues bitwise. Raises, unless - ``allow_config_change`` (the event then records the difference).""" - from yeto.rl.engine.resume import check_lr_continues - - pending = self.pending_lr_check - if pending is None or int(rollout_id) != int(pending["rollout_id"]): - return - self.pending_lr_check = None - problem = check_lr_continues(pending["expected"], None if applied_lrs is None else list(applied_lrs)) - driver.emit("rl_resume_check", rollout_id=int(rollout_id), check="lr_at_next_round", - expected=pending["expected"], - applied=None if applied_lrs is None else [float(x) for x in applied_lrs], - ok=problem is None, allowed=problem is not None and self.allow_config_change) - if problem is not None and not self.allow_config_change: - raise RoundCutError(problem) - - -def _cut_data_cursor(manifest: Any) -> dict[str, int] | None: - """The data source position recorded at the cut (None when incomplete). - - Under fully_async the background worker keeps drawing prompts after the - round's trained groups were reported, and the cut puts the untrained ones - back into the data buffer. The position to resume from is therefore the - one read at the cut, not the one the batch ledger recorded for the round.""" - from .rollout_meta_hook import _CURSOR_FIELDS - - data = dict(getattr(manifest, "data", None) or {}) - if not all(isinstance(data.get(f), int) for f in _CURSOR_FIELDS): - return None - return {f: int(data[f]) for f in _CURSOR_FIELDS} - - -def wire_round_cuts(driver: Any, *, controller: Any, source: Any, **options: Any) -> RoundCutCheckpoint | None: - """Give a :class:`~..bridges.LocalOnlySync` driver round cuts when a checkpoint store is - configured (other syncs own the policy: their syncer is authoritative).""" - from yeto.rl.engine.bridges import LocalOnlySync - - if getattr(controller, "checkpoint_store", None) is None or not isinstance(driver.sync, LocalOnlySync): - return None - checkpoint = RoundCutCheckpoint(controller=controller, source=source, trainer=driver.trainer, **options) - driver.sync.round_cuts = checkpoint - return checkpoint +from yeto.rl.engine.round_cut import ( # noqa: F401 - re-exported + ISLAND, + POINTER, + ROUND_CUTS, + RoundCutCheckpoint, + RoundCutError, + _cut_data_cursor, + wire_round_cuts, +) diff --git a/yeto/rl/adapters/miles/run_flags.py b/yeto/rl/adapters/miles/run_flags.py index a6a9c6da..8a9f5202 100644 --- a/yeto/rl/adapters/miles/run_flags.py +++ b/yeto/rl/adapters/miles/run_flags.py @@ -372,6 +372,11 @@ # bucket mount or a shared path); round cuts + a hashed state snapshot + LATEST go there. island_flag("--rl-resume-store", default=None, metavar="PATH"), island_flag("--rl-resume-state-dir", default="~/yeto-rl/resume-state", metavar="PATH"), + # S25 C3: a synced trainer island (--rl-sample-role trainer with a syncer): the + # store of the sample buffer and the policy releases (no round cuts), and the + # number of synced trainers (its buffer lane is learner id of that count) + island_flag("--rl-sample-store", default=None, metavar="PATH"), + island_flag("--rl-sample-lanes", type=int, default=None, metavar="N"), island_flag("--lora-r", type=int, default=None), island_flag( "--lora-targets", @@ -388,6 +393,13 @@ island_flag("--seed", type=int, required=True), island_flag("--wan-streams", type=int, default=4), island_flag("--miles-root", required=True), + island_flag( + "--miles-commit", + default=None, + metavar="SHA", + help="Full commit the --miles-root checkout must be at (default: the pinned " + "MILES_NEXT_COMMIT); a run whose image carries another Miles commit names it here", + ), island_flag( "--rl-expected-algorithm-sha256", default=None, diff --git a/yeto/rl/adapters/miles/sampler_island.py b/yeto/rl/adapters/miles/sampler_island.py index 4c49446a..22da4c7d 100644 --- a/yeto/rl/adapters/miles/sampler_island.py +++ b/yeto/rl/adapters/miles/sampler_island.py @@ -9,7 +9,9 @@ while it trains), every successful publish is also released into the store (:class:`ReleasingPublisher`): a full-parameter trainer exports an HF checkpoint (``actor.export_hf``), a LoRA trainer only its - adapter (:mod:`.release_plugin`), into the release directory. + adapter (:mod:`.release_plugin`), into the release directory. Synced trainers + (several islands and a syncer): only learner 0 releases, with the merged + version (:mod:`yeto.rl.engine.merged_release`). Engine weight load (the order of the publisher, ``publish.MilesPublisher``): under the controller's update lock, every engine loads the release, then the @@ -32,8 +34,10 @@ import asyncio import hashlib +import os import shutil import tempfile +import time from collections.abc import Callable from pathlib import Path from typing import Any @@ -45,10 +49,17 @@ PolicyRelease, PolicyReleaseStore, ) +from yeto.rl.engine.merged_release import MergedReleaseGuard, is_publisher, merged_version from yeto.rl.engine.policy_versions import engine_weight_version from yeto.rl.engine.sample_roles import CONVERTIBLE, SAMPLER, TRAINER -from .buffer_rollout import SampleBufferRolloutError, executor_read_cursor, executor_read_shards, open_buffer +from .buffer_rollout import ( + SampleBufferRolloutError, + executor_read_cursor, + executor_read_shards, + executor_set_gate_version, + open_buffer, +) from .publish import ( _engine_id, _maybe_await, @@ -236,16 +247,42 @@ async def _load(self, path: str | None, version: int) -> list[Any]: return engines -def run_sampler_island(miles_args: Any, *, environ: Any = None) -> int: - """Run the sampler island to its data end; returns the rounds generated.""" +LEAVE_HOLD_S = 300.0 # after a reclaim LEAVE: wait this long for the cloud to stop the container + + +def run_sampler_island(miles_args: Any, *, environ: Any = None, + sleep: Callable[[float], None] = time.sleep) -> int: + """Run the sampler island to its data end; returns the rounds generated. + + A replacement incarnation (S25 C1) continues after the newest round its + island has in the buffer (``sampler_resume``), loads the newest policy + release before its first round, and writes ``rl_sampler_join`` / a member + record. On the reclaim marker (Modal SIGINT/SIGTERM, written by the Modal + runner) it starts no new round, writes no more groups, writes + ``rl_sampler_leave`` with reason PREEMPTED and waits for the cloud to stop + the container (at most :data:`LEAVE_HOLD_S`).""" from miles.ray.placement_group import create_rollout_components from miles.utils.async_utils import Disposer from miles.utils.orchestration_utils import init_orchestration_script + from yeto.rl.elastic_client import LeaveReason from yeto.rl.engine.events import append_tape_event + from yeto.rl.engine.resume import store_for from yeto.rl.engine.sampler_loop import SamplerLoop + from yeto.rl.engine.sampler_resume import ( + RECLAIM_MARKER_ENV, + ReclaimWatcher, + SamplerMembership, + acknowledge_reclaim, + clear_reclaim_files, + install_marker_signals, + new_incarnation, + reclaim_requested, + resume_point, + ) from . import LoopRunner + from .buffer_rollout import INCARNATION_ENV, SAMPLER_ID_ENV from .entry import connect_island_ray from .rollout import RayMetadataSink, _release @@ -254,10 +291,25 @@ def run_sampler_island(miles_args: Any, *, environ: Any = None) -> int: raise SampleBufferRolloutError("run_sampler_island needs --rl-sample-role sampler") if not getattr(miles_args, "debug_rollout_only", False): raise SampleBufferRolloutError("a sampler island runs Miles with --debug-rollout-only") + env = os.environ if environ is None else environ def emit(event: str, **fields: Any) -> None: append_tape_event(miles_args, {"event": event, **fields}) + # one incarnation id for this process and its rollout workers (forwarded to Ray) + sampler_id = env.get(SAMPLER_ID_ENV) + if not sampler_id: + raise SampleBufferRolloutError(f"{SAMPLER_ID_ENV} is not set") + incarnation = env[INCARNATION_ENV] = new_incarnation() + clear_reclaim_files(env) # a restarted machine keeps its /tmp + if env.get(RECLAIM_MARKER_ENV): # set by the launcher on a pair's sampler island + install_marker_signals(env) + buffer = open_buffer(env, sampler=True) + resume = resume_point(buffer.entries(), sampler_id) + membership = SamplerMembership(store_for(config.store, environ=environ), sampler_id, + incarnation, resume, emit) + progress = {"rollout_id": None, "rounds": 0} # the round in flight, rounds done here + connect_island_ray(miles_args=miles_args) # The rollout process's metadata hook pushes each round's group records to this # named actor (on a trainer island the driver creates it and consumes them). @@ -276,11 +328,27 @@ async def init() -> tuple[Any, Any]: return controller, executor async def generate(rollout_id: int) -> None: + progress["rollout_id"] = rollout_id await controller.prepare_rollout(rollout_id) _release(await executor.get(rollout_id)) # the rollout function wrote the groups metadata.discard() # the groups are in the buffer; nobody here reads the records + progress["rounds"] += 1 + progress["rollout_id"] = None + + def leave(reason: LeaveReason, **detail: Any) -> None: + membership.leave(reason.name, rounds=progress["rounds"], + policy_version=None if follower is None else follower.version, + in_flight_rollout_id=progress["rollout_id"], **detail) error: BaseException | None = None + follower = None + reclaimed = lambda: reclaim_requested(env) # noqa: E731 + + def leave_reclaimed() -> None: + leave(LeaveReason.PREEMPTED) + acknowledge_reclaim(env) # a host hook (Nebius) lets the VM stop once the LEAVE is stored + + watcher = ReclaimWatcher(reclaimed, leave_reclaimed).start() try: controller, executor = runner.run(init()) lora = lora_mode(miles_args) @@ -290,20 +358,39 @@ async def generate(rollout_id: int) -> None: follower = PolicyFollower(release_store(config, environ), loader, version=BASE_POLICY_VERSION) loop = SamplerLoop( - follower=follower, buffer=open_buffer(environ, sampler=True), emit=emit, + follower=follower, buffer=buffer, emit=emit, generate=lambda rollout_id: runner.run(generate(rollout_id)), num_rollouts=int(miles_args.num_rollout), max_backlog_groups=int(config.max_backlog_groups), - min_round_s=config.min_round_s, require_release=lora, + min_round_s=config.min_round_s, require_release=lora, should_stop=reclaimed, loaded_fields=lambda: {"served_checksum": loader.last_checksum, "probe": loader.last_probe}) + latest = follower.releases.latest() + membership.join(latest_release=None if latest is None else + {"seq": latest.seq, "policy_version": latest.version}) emit("rl_sampler_start", store=config.store, num_rollouts=int(miles_args.num_rollout), - max_backlog_groups=int(config.max_backlog_groups), weights=loader.weights) - return loop.run() + max_backlog_groups=int(config.max_backlog_groups), weights=loader.weights, + start_rollout_id=resume.next_rollout_id) + if buffer.data_ended: # an earlier incarnation finished and marked the data end + rounds = 0 + else: + rounds = loop.run(start=resume.next_rollout_id) + if reclaimed(): + leave_reclaimed() # the LEAVE is a no-op when the watcher already left + _hold_until_stopped(sleep, LEAVE_HOLD_S) + else: + leave(LeaveReason.FINISHED) + return rounds except BaseException as exc: error = exc + try: + leave(LeaveReason.PREEMPTED if reclaimed() else LeaveReason.INFRA_FAILED, + error=f"{type(exc).__name__}: {exc}"[:300]) + except Exception: # noqa: BLE001 - the original error is the one to report + pass raise finally: + watcher.stop() try: runner.run(disposer.__aexit__(None if error is None else type(error), error, None if error is None else error.__traceback__)) @@ -311,6 +398,15 @@ async def generate(rollout_id: int) -> None: runner.close() +def _hold_until_stopped(sleep: Callable[[float], None], hold_s: float) -> None: + """After a reclaim LEAVE the process stays up until the cloud stops the + container: an early clean exit could end the cloud's input as finished + instead of letting the cloud run the island again.""" + print(f"[sampler-island] reclaimed: waiting up to {hold_s:g} s for the container to stop", + flush=True) + sleep(float(hold_s)) + + # ---------------------------------------------------------------------- trainer side class ReleasingPublisher: """A publisher that also releases each published policy for the samplers. @@ -326,10 +422,19 @@ def __init__(self, inner: Any, releases: PolicyReleaseStore, export: Callable[[int, Path], None], read_position: Callable[[], int | None], emit: Callable[..., None], *, weights: str = FULL, - read_shards: Callable[[], dict[str, int] | None] | None = None) -> None: + read_shards: Callable[[], dict[str, int] | None] | None = None, + version_of: Callable[[Any], int] | None = None, + guard: MergedReleaseGuard | None = None, + origin: Callable[[Any], dict[str, Any]] | None = None) -> None: """``read_shards``: the trainer's read position per sampler shard (None without merged shards), released with every policy for the samplers' - backlog caps.""" + backlog caps. Synced trainers (:mod:`yeto.rl.engine.merged_release`): + ``version_of(state)`` is the merged version to release (default: the + policy version), ``guard`` keeps the release versions increasing and + ``origin(state)`` is the release record's ``origin``.""" + self.version_of = version_of or (lambda state: int(state.policy_version)) + self.guard = guard + self.origin = origin self.inner = inner self.releases = releases self.export = export @@ -346,7 +451,15 @@ def _engine_checksum(self) -> str | None: def publish(self, state: Any, **kwargs: Any) -> Any: result = self.inner.publish(state, **kwargs) - version = int(state.policy_version) + version = int(self.version_of(state)) + if self.guard is not None: + latest = self.releases.latest() + if self.guard.check(version, None if latest is None else latest.version) == "stale": + self.emit("rl_policy_release_skipped", policy_version=version, + local_version=int(state.policy_version), + latest_version=None if latest is None else latest.version, + reason="not_newer_than_latest") + return result position = self.read_position() shards = self.read_shards() if self.read_shards is not None and position is not None else None if position is None: # nothing read yet in this process: keep the last known positions @@ -356,9 +469,11 @@ def publish(self, state: Any, **kwargs: Any) -> Any: release = self.releases.release(version, state.policy_tensor_hash(), lambda directory: self.export(version, directory), buffer_selected=position, weights=self.weights, - engine_checksum=self._engine_checksum(), buffer_shards=shards) + engine_checksum=self._engine_checksum(), buffer_shards=shards, + origin=None if self.origin is None else self.origin(state)) if release is not None: self.emit("rl_policy_released", seq=release.seq, policy_version=version, + local_version=int(state.policy_version), origin=release.origin, policy_hash=release.policy_hash, buffer_selected=release.buffer_selected, weights=release.weights, sha256=release.sha256, engine_checksum=release.engine_checksum) @@ -368,10 +483,35 @@ def __getattr__(self, name: str) -> Any: return getattr(self.inner, name) +class GateVersionPublisher: + """A synced trainer's publisher (:mod:`yeto.rl.engine.merged_release`): after + each publish, the island's buffer reader gates ``max_policy_age`` on the merged + version (``set_gate_version(version)``). Every other attribute is the wrapped + publisher's.""" + + def __init__(self, inner: Any, set_gate_version: Callable[[int], Any], + version_of: Callable[[Any], int]) -> None: + self.inner = inner + self.set_gate_version = set_gate_version + self.version_of = version_of + + def publish(self, state: Any, **kwargs: Any) -> Any: + result = self.inner.publish(state, **kwargs) + self.set_gate_version(int(self.version_of(state))) + return result + + def __getattr__(self, name: str) -> Any: + return getattr(self.inner, name) + + def wire_policy_releases(driver: Any, *, miles_args: Any, actor_model: Any, runner: Any, environ: Any = None) -> ReleasingPublisher | None: """A trainer island (``--rl-sample-role trainer``) releases every publish; so does - the trainer of a convertible node (the other nodes follow the releases).""" + the trainer of a convertible node (the other nodes follow the releases). + + Synced trainers (a buffer lane on ``miles_args.yeto_rl_sample_lane``): every + island gates on the merged version; only the publisher (learner id 0) releases, + with the merged version (:mod:`yeto.rl.engine.merged_release`).""" config = _buffer_config(miles_args, TRAINER) or _buffer_config(miles_args, CONVERTIBLE) if config is None: return None @@ -379,6 +519,27 @@ def wire_policy_releases(driver: Any, *, miles_args: Any, actor_model: Any, runn if not callable(in_executor): raise SampleBufferRolloutError(f"rollout pool {type(driver.rollout).__name__} cannot " "run code in the rollout executor") + lane = getattr(miles_args, "yeto_rl_sample_lane", None) + synced: dict[str, Any] = {} + if lane is not None: + def version_of(state: Any) -> int: + return merged_version(driver, state) + + def origin(state: Any) -> dict[str, Any]: + return {"kind": "merged", "learner_id": lane.index, "num_learners": lane.count, + "local_version": int(state.policy_version)} + + def on_gap(previous: int, version: int) -> None: + driver.emit("rl_policy_release_gap", previous_version=previous, policy_version=version) + + driver.publisher = GateVersionPublisher( + driver.publisher, lambda v: in_executor(executor_set_gate_version, int(v)), version_of) + publisher_island = is_publisher(lane.index, lane.count) + driver.emit("rl_policy_release_role", learner_id=lane.index, num_learners=lane.count, + publisher=publisher_island) + if not publisher_island: + return None + synced = {"version_of": version_of, "guard": MergedReleaseGuard(on_gap), "origin": origin} lora = lora_mode(miles_args) @@ -391,6 +552,6 @@ def export(version: int, directory: Path) -> None: publisher = ReleasingPublisher(driver.publisher, release_store(config, environ), export, read_position=lambda: in_executor(executor_read_cursor), read_shards=lambda: in_executor(executor_read_shards), - emit=driver.emit, weights=LORA if lora else FULL) + emit=driver.emit, weights=LORA if lora else FULL, **synced) driver.publisher = publisher return publisher diff --git a/yeto/rl/adapters/miles/state_plugin.py b/yeto/rl/adapters/miles/state_plugin.py index ea72b068..5125287e 100644 --- a/yeto/rl/adapters/miles/state_plugin.py +++ b/yeto/rl/adapters/miles/state_plugin.py @@ -1828,6 +1828,11 @@ def align_weight_version(actor: Any, *, next_version: int) -> None: updater = getattr(actor, "weight_updater", None) if updater is None: + if getattr(getattr(actor, "args", None), "debug_train_only", False): + # trainer-only island (Miles --debug-train-only): no engine and no weight + # updater; update_weights returns early and the policy release carries + # the version (Modal run-11 failed here) + return raise RuntimeError("the trainer has no weight_updater to carry the policy version") updater.weight_version = int(next_version) - 1 diff --git a/yeto/rl/adapters/miles/trainer.py b/yeto/rl/adapters/miles/trainer.py index ff548c45..54c65ed9 100644 --- a/yeto/rl/adapters/miles/trainer.py +++ b/yeto/rl/adapters/miles/trainer.py @@ -44,6 +44,44 @@ class TrainStepError(RuntimeError): pass + +def _is_expert_tensor(name: str) -> bool: + """A routed-expert tensor (Megatron ``mlp.experts.*``; not ``shared_experts``).""" + return ".mlp.experts." in name + +def optimizer_coverage_problems(summaries: list[dict[str, Any]]) -> list[str]: + """Check that every adapter tensor has optimizer state on at least one rank of + the group that shards it (DistOpt ranges); marks ``has_optimizer_state`` False on + the ranks of a group that fails. Expert tensors are sharded over the expert-DP + ranks of one (tp, pp, ep); every other tensor over all DP ranks of its (tp, pp), + which with EP>1 spans the EP ranks too (Modal run-16: the dense tensors were + checked per EP rank and found missing).""" + problems: list[str] = [] + covered: dict[tuple[int, ...], tuple[set[str], set[str]]] = {} + for s in summaries: + tp, pp, ep = s["coord"]["tp"], s["coord"]["pp"], s["coord"].get("ep", 0) + optimized = set(s.get("optimizer_names") or ()) + for name in s.get("adapter_names") or (): + key = (tp, pp, ep) if _is_expert_tensor(name) else (tp, pp) + names, owned = covered.setdefault(key, (set(), set())) + names.add(name) + if name in optimized: + owned.add(name) + if not covered: + problems.append("no rank reported any adapter tensor") + for key, (adapters, optimized) in sorted(covered.items()): + if not adapters <= optimized: + missing = sorted(adapters - optimized)[:4] + for s in summaries: + if (s["coord"]["tp"], s["coord"]["pp"]) == key[:2] and ( + len(key) == 2 or s["coord"].get("ep", 0) == key[2]): + s["has_optimizer_state"] = False + where = f"tp{key[0]}/pp{key[1]}" + (f"/ep{key[2]}" if len(key) == 3 else " (dense)") + problems.append(f"{where}: no optimizer state for {missing}") + return problems + + + def batch_hash(batch: RolloutBatchHandle) -> str: """Content identity of the trained batch, computed from metadata only.""" @@ -801,21 +839,7 @@ def save_cut(self, *, epoch: int, context: "CutContext") -> str: expected = trainer_workers(self._args) if len(summaries) != expected: raise TrainStepError(f"expected {expected} cut shards (one per rank), got {len(summaries)}") - # has_optimizer_state means: every adapter tensor of each (tp, pp) has - # optimizer state on at least one of its DP ranks (DistOpt ranges). - covered: dict[tuple[int, int], tuple[set[str], set[str]]] = {} - for s in summaries: - key = (s["coord"]["tp"], s["coord"]["pp"], s["coord"].get("ep", 0)) - adapters, optimized = covered.setdefault(key, (set(), set())) - adapters.update(s.get("adapter_names") or ()) - optimized.update(s.get("optimizer_names") or ()) - for key, (adapters, optimized) in sorted(covered.items()): - if not adapters or not adapters <= optimized: - missing = sorted(adapters - optimized)[:4] - for s in summaries: - if (s["coord"]["tp"], s["coord"]["pp"], s["coord"].get("ep", 0)) == key: - s["has_optimizer_state"] = False - problems.append(f"tp{key[0]}/pp{key[1]}/ep{key[2]}: no optimizer state for {missing or 'any adapter'}") + problems += optimizer_coverage_problems(summaries) if problems: raise CutError("refusing an incomplete cut: " + "; ".join(problems)) resume = self._resume_manifest(context, summaries) diff --git a/yeto/rl/adapters/sampler_engines/server.py b/yeto/rl/adapters/sampler_engines/server.py index 5d868c46..ad33b0ac 100644 --- a/yeto/rl/adapters/sampler_engines/server.py +++ b/yeto/rl/adapters/sampler_engines/server.py @@ -58,8 +58,18 @@ def wait(self, timeout: float | None = None) -> int: ... Http = Callable[[str, str, Any, float], tuple[int, Any]] +def _ignore_sigint() -> None: + signal.signal(signal.SIGINT, signal.SIG_IGN) + + def popen_spawn(argv: Sequence[str], env: Mapping[str, str]) -> Process: - return subprocess.Popen(list(argv), env={**os.environ, **env}) # noqa: S603 + """Start the server with SIGINT ignored (inherited by its workers). A cloud + reclaim (Modal ``container stop``) sends SIGINT to every process of the + container: the island turns it into the reclaim marker and leaves after + the round in progress, so the server must not die under that round + (T-E2V eval-e2-verl-e2v-a). ``stop()`` still ends it with SIGTERM.""" + return subprocess.Popen(list(argv), env={**os.environ, **env}, # noqa: S603 + preexec_fn=_ignore_sigint) # noqa: PLW1509 - only resets one signal def httpx_call(method: str, url: str, body: Any, timeout_s: float) -> tuple[int, Any]: diff --git a/yeto/rl/adapters/sampler_engines/sglang_load_report.py b/yeto/rl/adapters/sampler_engines/sglang_load_report.py new file mode 100644 index 00000000..b228e163 --- /dev/null +++ b/yeto/rl/adapters/sampler_engines/sglang_load_report.py @@ -0,0 +1,141 @@ +"""Expert load report and ``weight_packed`` rename inside an SGLang server process +(S25 N2.3 post-load check 1; A-INF section 2.3 option A, item 1). + +SGLang's GLM5.3 ``load_weights`` maps ``experts..gate_proj.weight`` to the +fused ``experts.w13_weight`` parameter. A compressed-tensors MXFP4 release +stores that matrix as ``...gate_proj.weight_packed``; the mapped name +``experts.w13_weight_packed`` is not a parameter, and SGLang skips the tensor +without an error. :func:`install` wraps the model class's ``load_weights``: + +* with ``YETO_SGLANG_RENAME_WEIGHT_PACKED=1`` it renames routed-expert + ``*.weight_packed`` tensors to ``*.weight`` before SGLang maps them (the + MXFP4 MoE method stores the packed bytes in ``w13_weight`` / ``w2_weight``); +* with ``YETO_SGLANG_LOAD_REPORT_DIR`` set it counts the calls of the routed + experts' data-weight loaders (one call = one gate, up or down matrix of one + expert on this rank) and writes ``rank-.json`` into that directory + after each load. A skipped tensor never reaches a loader, so the count is + what the engine holds, not what the checkpoint names. + +The image applies :data:`MODEL_FILE_HOOK` to SGLang's model file +(:func:`patch_model_file`); the server process imports this module from the +yeto tree (``PYTHONPATH``). :func:`read_reports` is the adapter's side +(``SGLangServerEngine.loaded_expert_matrices``). + +This module must not import SGLang at import time (CPU tests use a fake model). +""" + +from __future__ import annotations + +import json +import os +import re +import time +from collections.abc import Iterable, Iterator +from pathlib import Path +from typing import Any + +REPORT_DIR_ENV = "YETO_SGLANG_LOAD_REPORT_DIR" +RENAME_ENV = "YETO_SGLANG_RENAME_WEIGHT_PACKED" +ROUTED_PACKED = re.compile(r"(\.mlp\.experts\.\d+\.(?:gate|up|down)_proj)\.weight_packed$") +# FusedMoE data parameters (not the scales); the routed experts' loaders. +EXPERT_DATA_PARAMS = re.compile(r"\.experts\.(?:w13_weight|w2_weight)$") +MARKER = "# yeto: sglang_load_report hook" +MODEL_FILE_HOOK = f""" + +{MARKER} +import os as _yeto_os +if _yeto_os.environ.get("{REPORT_DIR_ENV}") or _yeto_os.environ.get("{RENAME_ENV}") == "1": + from yeto.rl.adapters.sampler_engines.sglang_load_report import install as _yeto_install + for _yeto_cls in EntryClass: + _yeto_install(_yeto_cls) +""" + + +def rename_packed(weights: Iterable[tuple[str, Any]]) -> Iterator[tuple[str, Any]]: + """Routed-expert ``*.weight_packed`` -> ``*.weight``; every other name unchanged.""" + for name, tensor in weights: + yield ROUTED_PACKED.sub(r"\1.weight", name), tensor + + +def _tp_rank() -> int: + try: + from sglang.srt.distributed import get_tensor_model_parallel_rank # noqa: PLC0415 + return int(get_tensor_model_parallel_rank()) + except Exception: # noqa: BLE001 - outside a parallel group (tests) + return 0 + + +def install(model_cls: type) -> type: + """Wrap ``model_cls.load_weights`` once (module doc).""" + original = model_cls.load_weights + if getattr(original, "_yeto_load_report", False): + return model_cls + + def load_weights(self: Any, weights: Iterable[tuple[str, Any]], *args: Any, **kwargs: Any) -> Any: + if os.environ.get(RENAME_ENV) == "1": + weights = rename_packed(weights) + report_dir = os.environ.get(REPORT_DIR_ENV) + if not report_dir: + return original(self, weights, *args, **kwargs) + calls = {"n": 0} + wrapped = [] + for name, param in self.named_parameters(): + loader = getattr(param, "weight_loader", None) + if EXPERT_DATA_PARAMS.search(name) and loader is not None: + def counting(*a: Any, _loader: Any = loader, **kw: Any) -> Any: + calls["n"] += 1 + return _loader(*a, **kw) + param.weight_loader = counting + wrapped.append((param, loader)) + try: + result = original(self, weights, *args, **kwargs) + finally: + for param, loader in wrapped: + param.weight_loader = loader + rank = _tp_rank() + path = Path(report_dir) / f"rank{rank}-{os.getpid()}.json" + path.parent.mkdir(parents=True, exist_ok=True) + tmp = path.with_suffix(".tmp") + tmp.write_text(json.dumps({"tp_rank": rank, "pid": os.getpid(), "expert_matrices": calls["n"], + "expert_params": len(wrapped), "time": time.time()})) + os.replace(tmp, path) + return result + + load_weights._yeto_load_report = True # type: ignore[attr-defined] + model_cls.load_weights = load_weights + return model_cls + + +def patch_model_file(path: str | os.PathLike[str]) -> bool: + """Append :data:`MODEL_FILE_HOOK` to an SGLang model file that defines + ``EntryClass`` (idempotent; False when already patched).""" + file = Path(path) + text = file.read_text(encoding="utf-8") + if MARKER in text: + return False + if "EntryClass" not in text: + raise ValueError(f"{file} defines no EntryClass") + file.write_text(text + MODEL_FILE_HOOK, encoding="utf-8") + return True + + +def clear_reports(report_dir: str | os.PathLike[str]) -> None: + directory = Path(report_dir) + directory.mkdir(parents=True, exist_ok=True) + for f in directory.glob("rank*.json"): + f.unlink() + + +def read_reports(report_dir: str | os.PathLike[str]) -> dict[str, Any] | None: + """The reports of the current server (None: no rank reported). ``expert_matrices`` + is the smallest rank count, so one rank that skipped experts fails the check.""" + reports = [] + for f in sorted(Path(report_dir).glob("rank*.json")): + try: + reports.append(json.loads(f.read_text(encoding="utf-8"))) + except (OSError, ValueError): + continue + if not reports: + return None + counts = [int(r["expert_matrices"]) for r in reports] + return {"expert_matrices": min(counts), "per_rank": counts, "ranks": len(reports)} diff --git a/yeto/rl/adapters/sampler_engines/sglang_server.py b/yeto/rl/adapters/sampler_engines/sglang_server.py index 1be20274..cdf57ddb 100644 --- a/yeto/rl/adapters/sampler_engines/sglang_server.py +++ b/yeto/rl/adapters/sampler_engines/sglang_server.py @@ -6,7 +6,10 @@ * item 1: SGLang's ``load_weights`` does not map ``weight_packed`` to ``weight``; expert weights stored as ``weight_packed`` are skipped without an error. A release whose index names ``*.weight_packed`` is refused unless the - image carries the rename patch (``renames_weight_packed=True``); + image carries the rename patch (``renames_weight_packed=True``). The yeto + patch (:mod:`.sglang_load_report`, applied to the image's model file) renames + when ``YETO_SGLANG_RENAME_WEIGHT_PACKED=1`` and, with ``load_report_dir``, + reports the expert matrices each rank loaded (post-load check 1); * item 5: an MXFP4 export with ``quantization_config.format == "dense"`` (llmcompressor 0.14) is read wrongly; ``dense`` is refused for every quantized scheme; @@ -31,6 +34,7 @@ from __future__ import annotations import json +import re from collections.abc import Sequence from dataclasses import dataclass from pathlib import Path @@ -45,11 +49,22 @@ SamplingParams, ) +from . import sglang_load_report from .server import ServerConfig, ServerEngine TOPOLOGY_TP8 = "tp8" TOPOLOGY_DP8 = "dp8" -SGLANG_TOPOLOGIES = (TOPOLOGY_TP8, TOPOLOGY_DP8) +_TOPOLOGY = re.compile(r"(tp|dp)([1-9][0-9]*)") + + +def sglang_topology(tp: int, *, dp_attention: bool = False) -> str: + """The topology label of one SGLang instance over ``tp`` cards: ``tp``, or + ``dp`` with DP attention (``--dp N --enable-dp-attention``).""" + return f"{'dp' if dp_attention else 'tp'}{int(tp)}" + + +def is_dp_attention(topology: str) -> bool: + return topology.startswith("dp") @dataclass @@ -59,6 +74,8 @@ class SGLangOptions: renames_weight_packed: bool = False # the image patches the weight_packed -> weight map warmup_batch_sizes: tuple[int, ...] = () served_model_name: str = "yeto-policy" + # the yeto load-report patch writes each rank's loaded expert count here (None: no report) + load_report_dir: str | None = None def _finish(meta: dict[str, Any]) -> str: @@ -80,11 +97,16 @@ class SGLangServerEngine(ServerEngine): def __init__(self, config: ServerConfig, options: SGLangOptions | None = None, **kw: Any) -> None: super().__init__(config, **kw) self.options = options or SGLangOptions() - if config.topology not in SGLANG_TOPOLOGIES: - raise SamplerEngineError(f"sglang topology {config.topology!r} is not supported " - f"(known: {SGLANG_TOPOLOGIES}; PD disaggregation is not used)") - if config.topology == TOPOLOGY_DP8 and self.options.nextn: - raise SamplerEngineError("sglang dp8 + DeepEP cannot run NEXTN (A-INF item 8); " + # The topology is the source label of every sample: it must name the real + # layout, tp or dp with N = the instance's card count (options.tp). + match = _TOPOLOGY.fullmatch(config.topology) + if match is None or int(match.group(2)) != int(self.options.tp): + raise SamplerEngineError( + f"sglang topology {config.topology!r} does not match TP {self.options.tp}: expected " + f"{sglang_topology(self.options.tp)!r} or {sglang_topology(self.options.tp, dp_attention=True)!r} " + "(PD disaggregation is not used)") + if is_dp_attention(config.topology) and self.options.nextn: + raise SamplerEngineError("sglang DP attention + DeepEP cannot run NEXTN (A-INF item 8); " "turn one of them off") def argv(self, model_dir: Path) -> list[str]: @@ -93,7 +115,7 @@ def argv(self, model_dir: Path) -> list[str]: "--host", c.host, "--port", str(c.port), "--tp", str(o.tp), "--context-length", str(c.max_model_len), "--served-model-name", o.served_model_name, "--trust-remote-code"] - if c.topology == TOPOLOGY_DP8: + if is_dp_attention(c.topology): argv += ["--dp", str(o.tp), "--enable-dp-attention", "--moe-a2a-backend", "deepep"] if o.nextn: argv += ["--speculative-algorithm", "NEXTN"] @@ -103,8 +125,27 @@ def env(self) -> dict[str, str]: env = super().env() if self.config.precision == "mxfp4": env["MX_PATCH"] = "1" # item 7 + if self.options.renames_weight_packed: + env[sglang_load_report.RENAME_ENV] = "1" + if self.options.load_report_dir: + env[sglang_load_report.REPORT_DIR_ENV] = str(self.options.load_report_dir) return env + def start(self, release_dir: Path | None, version: int): + if self.options.load_report_dir: # reports of an earlier server must not count for this one + sglang_load_report.clear_reports(self.options.load_report_dir) + return super().start(release_dir, version) + + def loaded_expert_matrices(self) -> int | None: + if not self.options.load_report_dir: + return None + report = sglang_load_report.read_reports(self.options.load_report_dir) + if report is None: + return None + if report["ranks"] < int(self.options.tp): + return 0 # a rank that did not report loaded nothing we can vouch for + return int(report["expert_matrices"]) + def preflight(self, model_dir: Path) -> None: super().preflight(model_dir) index = Path(model_dir) / INDEX_FILE diff --git a/yeto/rl/adapters/verl/agentic_main.py b/yeto/rl/adapters/verl/agentic_main.py new file mode 100644 index 00000000..ea7d69e3 --- /dev/null +++ b/yeto/rl/adapters/verl/agentic_main.py @@ -0,0 +1,108 @@ +"""verl's own synchronous trainer on the Harbor agentic path (S25 N10.9; runs inside the verl image). + +``python -m yeto.rl.adapters.verl.agentic_main `` is verl's +``main_ppo.main`` unchanged, with Ray started here so that every Ray worker: + +* registers the ``yeto_grpo`` advantage estimator right after verl's + ``core_algos`` is imported (Ray ``worker_process_setup_hook`` installs an + import hook). The hook itself must not import verl or torch: it runs + before Ray sets the actor's ``CUDA_VISIBLE_DEVICES``, and an early CUDA + init puts every rank on GPU 0 (NCCL "Duplicate GPU detected"; S24 T-GLM, + S25 T-N109 run a). The ``yeto_harbor`` agent loop is imported by verl + through ``agent_loop_config_path``; +* carries the agentic environment: ``YETO_*`` (recorder, records proxy, boat + caps), ``HARBOR_*`` (task directory, provider) and ``BOAT_API_KEY`` (the + sandbox account, read by ``boat_client``), on top of verl's own runtime env. + +The round loop, weight updates and rollout servers are verl's. The records +proxy runs in each agent-loop worker over verl's rollout client +(:mod:`.rollout_engine`); select it with +``YETO_RECORDS_PROXY=yeto.rl.adapters.verl.rollout_engine:records_proxy_from_env``. +""" + +from __future__ import annotations + +import importlib +import importlib.abc +import importlib.util +import os +import sys +from typing import Any + +FORWARD_ENV_PREFIXES = ("YETO_", "HARBOR_", "HF_", "PYTHONPATH") +FORWARD_ENV_NAMES = ("BOAT_API_KEY",) +SETUP_HOOK = "yeto.rl.adapters.verl.agentic_main.setup_worker" + + +ESTIMATOR_HOST = "verl.trainer.ppo.core_algos" +ESTIMATOR_MODULE = "yeto.rl.adapters.verl.advantage" + + +class _AfterImport(importlib.abc.MetaPathFinder): + """Import ``then`` right after ``target`` has been executed (no import of either before that).""" + + def __init__(self, target: str, then: str) -> None: + self.target, self.then = target, then + + def find_spec(self, name: str, path: Any, target: Any = None) -> Any: + if name != self.target: + return None + sys.meta_path.remove(self) # find the real spec without this finder + try: + spec = importlib.util.find_spec(name) + finally: + sys.meta_path.insert(0, self) + if spec is None or spec.loader is None: + return spec + loader, then = spec.loader, self.then + + class _Loader(importlib.abc.Loader): + def create_module(self, spec: Any) -> Any: + return loader.create_module(spec) + + def exec_module(self, module: Any) -> None: + loader.exec_module(module) + importlib.import_module(then) + + spec.loader = _Loader() + return spec + + +def setup_worker() -> None: + """Ray ``worker_process_setup_hook``: register yeto_grpo once verl's core_algos loads (module doc).""" + if ESTIMATOR_HOST in sys.modules: + importlib.import_module(ESTIMATOR_MODULE) + return + if not any(isinstance(f, _AfterImport) for f in sys.meta_path): + sys.meta_path.insert(0, _AfterImport(ESTIMATOR_HOST, ESTIMATOR_MODULE)) + + +def runtime_env(environ: dict[str, str] | None = None) -> dict: + """verl's PPO runtime env plus the forwarded agentic variables and the setup hook.""" + from verl.trainer.constants_ppo import get_ppo_ray_runtime_env + + environ = dict(os.environ if environ is None else environ) + env = get_ppo_ray_runtime_env(None) + env.setdefault("env_vars", {}) + env["env_vars"]["TRANSFER_QUEUE_ENABLE"] = "1" # verl's V1 trainer runs on TransferQueue + for key, value in environ.items(): + if key.startswith(FORWARD_ENV_PREFIXES) or key in FORWARD_ENV_NAMES: + env["env_vars"][key] = value + env["worker_process_setup_hook"] = SETUP_HOOK + return env + + +def main(argv: list[str] | None = None) -> None: + import ray + import verl.trainer.main_ppo as main_ppo + + importlib.import_module(ESTIMATOR_MODULE) # the driver validates algorithm.adv_estimator too + if not ray.is_initialized(): + ray.init(runtime_env=runtime_env(), namespace="yeto-verl-agentic") + if argv is not None: + sys.argv = [sys.argv[0], *argv] + main_ppo.main() + + +if __name__ == "__main__": + main() diff --git a/yeto/rl/adapters/verl/buffer_feed.py b/yeto/rl/adapters/verl/buffer_feed.py new file mode 100644 index 00000000..cac0808a --- /dev/null +++ b/yeto/rl/adapters/verl/buffer_feed.py @@ -0,0 +1,346 @@ +"""verl trainer island over the durable sample buffer: the feed (S25 C2). + +A verl *buffer trainer* trains on groups that a separate sampler island wrote +to the sample buffer (``yeto.rl.engine.sample_buffer``; the engine sampler +island ``yeto.rl.engine.engine_sampler`` or the Miles sampler island). It reuses +the fully_async consumption unchanged: the yeto ``IslandDriver`` and the ports +of :mod:`.fully_async_ports` (``RoundCollector`` with the ``max_policy_age`` +judgement, carry-over and discard reports, the trainer step with the off-policy +correction). Only the source of the queued samples differs: verl's +``MessageQueue`` (filled by its rollouter) is replaced by :class:`BufferFeed`. + +:class:`BufferFeed` answers the three queue calls of the rollout port +(``yeto_pull_sample`` / ``yeto_keep_last`` / ``yeto_drop_last``) in the driver +process, so everything that decides which group is trained runs here, on plain +data, without verl or Ray: + +* ``pull``: the next buffer entry (one GRPO group), waiting up to ``wait_s`` for + the sampler; its :class:`~.fully_async_round.SampleMeta` with every response's + behavior version translated to verl's ``param_version`` through the island's + :class:`~.fully_async_translate.VersionMap` (a version the trainer never + published here maps to None: the collector discards the group as unknown); +* ``keep``: the group's padded rows (:func:`group_rows`) go to the trainer actor + (``add_group``), which wraps them as a verl ``RolloutSample``; +* ``drop``: nothing reaches the actor; the group stays consumed. + +Every pulled group is recorded once (``record``) with its buffer position, so a +run can show that no group was consumed twice. The buffer positions are those of +the trainer's read cursor (``FileSampleBuffer.select``). + +Synced trainers (S25 C3, ``yeto.rl.engine.merged_release``): the buffer is a +lane reader (positions ``p % N == m`` only) and the feed keeps the lane cursor +(``cursor_log``): it continues after the last recorded round, and +:meth:`BufferFeed.close_round` records the round's trained positions when its +collection ends, before the round trains (at most once, as the Miles reader). +The groups carry the *merged* version as behavior version; +:class:`MergedVersions` turns it into the island's local policy version (the +first local version that applied that merge, whose weights are the merged +policy), and the version map turns that into verl's ``param_version``. So the +``max_policy_age`` judgement counts local versions: equal to merged versions +when every round merges, stricter when an island trains several rounds per +merge (elastic sync). + +:func:`group_rows` builds the rows the way verl's agent loop postprocess pads +them (``AgentLoopWorker._postprocess``): prompts left-padded to the prompt +length, responses right-padded to the response length, the reward on the last +response token (``rm_scores``), the behavior log-probs as ``rollout_log_probs`` +and the sample's source record as ``provenance`` (the per-source off-policy +metrics read it). +""" + +from __future__ import annotations + +import time +from collections.abc import Callable, Mapping, Sequence +from dataclasses import dataclass +from typing import Any + +from yeto.rl.engine.buffer_payload import ( + SampleBufferRolloutError, + complete_prompt_tokens, + decode_group, + group_entry_fields, +) + +from .fully_async_translate import VerlTranslationError, VersionMap + +TRAINER_CONSUMER = "trainer" # the buffer's single consumer name (as the Miles reader) +QUEUE_METHODS = ("yeto_pull_sample", "yeto_keep_last", "yeto_drop_last") +UID_PREFIX = "buf" + + +class BufferFeedError(SampleBufferRolloutError): + """A buffer group cannot become verl rows.""" + + +class MergedVersions: + """Synced trainers: merged (syncer outer) version -> the first local policy + version of this island that applied it (module doc). Merged versions only + increase; a local version pairs with one merged version.""" + + def __init__(self) -> None: + self._local: dict[int, int] = {} + + def record(self, merged: int, local: int) -> bool: + """Record one publication; True when ``merged`` is new on this island.""" + merged, local = int(merged), int(local) + if merged in self._local: + return False + if self._local and (merged < max(self._local) or local <= self._local[max(self._local)]): + raise BufferFeedError(f"merged v{merged} at local v{local} after merged v{max(self._local)} " + f"at local v{self._local[max(self._local)]}: versions must increase") + self._local[merged] = local + return True + + def local(self, merged: int) -> int | None: + return self._local.get(int(merged)) + + def to_dict(self) -> dict[str, int]: + return {str(k): v for k, v in sorted(self._local.items())} + + +@dataclass(frozen=True) +class RowShape: + """The padded widths of the trainer's rows (``data.max_prompt_length`` and + ``data.max_response_length`` of the verl config) and the pad token id.""" + + prompt_length: int + response_length: int + pad_token_id: int + + def __post_init__(self) -> None: + if self.prompt_length < 1 or self.response_length < 1: + raise ValueError(f"row widths must be positive, got {self.prompt_length}/{self.response_length}") + + +def sample_version(sample: Mapping[str, Any]) -> int: + """The oldest engine weight version of one sample (its behavior version).""" + return group_entry_fields([sample])[0] + + +def group_rows(samples: Sequence[Mapping[str, Any]], *, uid: str, param_versions: Sequence[int], + shape: RowShape, tokenize: Callable[[str], Sequence[int]] | None = None) -> dict[str, Any]: + """The verl rows of one group as plain lists (module doc). + + ``param_versions``: verl's ``param_version`` of each sample's behavior + version (``min_global_steps`` = ``max_global_steps``: a buffer sample comes + from one engine version). ``tokenize`` rebuilds the prompt ids of a sample + whose engine did not return them (``complete_prompt_tokens``). Returns + ``{"tensors": {name: rows}, "non_tensor": {name: values}, "meta_info": {...}}``. + Raises :class:`BufferFeedError` when a sample does not fit the widths.""" + if len(param_versions) != len(samples): + raise BufferFeedError(f"{len(param_versions)} param versions for {len(samples)} samples") + tensors: dict[str, list[list[Any]]] = {k: [] for k in ( + "prompts", "responses", "response_mask", "input_ids", "attention_mask", "position_ids", + "rollout_log_probs", "rm_scores")} + provenance, finish = [], [] + pad = int(shape.pad_token_id) + for sample in samples: + sample = complete_prompt_tokens(sample, tokenize) + tokens = [int(t) for t in sample.get("tokens") or ()] + length = int(sample.get("response_length") or 0) + logprobs = [float(x) for x in sample.get("rollout_log_probs") or ()] + where = f"group {uid} sample {sample.get('index')}" + if length < 1 or len(logprobs) != length or len(tokens) <= length: + raise BufferFeedError(f"{where}: {len(tokens)} tokens, response_length {length}, " + f"{len(logprobs)} log-probs") + prompt, response = tokens[:-length], tokens[-length:] + if len(prompt) > shape.prompt_length or length > shape.response_length: + raise BufferFeedError(f"{where}: prompt {len(prompt)} / response {length} tokens exceed " + f"the rows ({shape.prompt_length} / {shape.response_length})") + left = shape.prompt_length - len(prompt) + right = shape.response_length - length + attention = [0] * left + [1] * (len(prompt) + length) + [0] * right + tensors["prompts"].append([pad] * left + prompt) + tensors["responses"].append(response + [pad] * right) + tensors["response_mask"].append([1] * length + [0] * right) + tensors["input_ids"].append([pad] * left + prompt + response + [pad] * right) + tensors["attention_mask"].append(attention) + tensors["position_ids"].append(_position_ids(attention)) + tensors["rollout_log_probs"].append(logprobs + [0.0] * right) + tensors["rm_scores"].append([0.0] * (length - 1) + [float(sample.get("reward") or 0.0)] + + [0.0] * right) + metadata = sample.get("metadata") or {} + provenance.append(dict(metadata.get("provenance") or {}) or None) + finish.append(metadata.get("finish_reason")) + n = len(samples) + versions = [int(v) for v in param_versions] + return { + "tensors": tensors, + "non_tensor": {"uid": [uid] * n, "min_global_steps": versions, "max_global_steps": list(versions), + "provenance": provenance, "finish_reason": finish, "__num_turns__": [1] * n}, + # verl's assemble step pops one ``metrics`` record per row (addition_process) + "meta_info": {"metrics": [{"generate_sequences": 0.0, "tool_calls": 0.0}] * n, + "reward_extra_keys": []}, + } + + +def _position_ids(attention: Sequence[int]) -> list[int]: + """verl's ``compute_position_id_with_mask``: cumulative sum minus one, floored at 0.""" + out, seen = [], 0 + for bit in attention: + seen += int(bit) + out.append(max(seen - 1, 0)) + return out + + +@dataclass +class _Pulled: + """The group the collector is judging: its rows' inputs and its record.""" + + position: int + samples: list[dict[str, Any]] + param_versions: tuple[int | None, ...] + record: dict[str, Any] + + +class BufferFeed: + """The rollout port's queue calls over the sample buffer (module doc). + + ``buffer``: the trainer's ``FileSampleBuffer``; ``vmap``: returns the + island's live version map; ``add_group(rows)``: hands kept rows to the + trainer actor; ``record(fields)``: one record per pulled group.""" + + def __init__(self, buffer: Any, *, vmap: Callable[[], VersionMap], shape: RowShape, + add_group: Callable[[dict[str, Any]], Any], record: Callable[[dict[str, Any]], None], + tokenize: Callable[[str], Sequence[int]] | None = None, wait_s: float = 3600.0, + poll_s: float = 2.0, consumer: str = TRAINER_CONSUMER, + clock: Callable[[], float] = time.monotonic, + sleep: Callable[[float], None] = time.sleep, + cursor_log: Any = None, merged: MergedVersions | None = None) -> None: + """``cursor_log``: a synced trainer's lane cursor log + (``merged_release.lane_cursor_log``); ``merged``: a synced trainer's + merged version pairing (module doc).""" + self.buffer = buffer + self.vmap = vmap + self.shape = shape + self.add_group = add_group + self.record = record + self.tokenize = tokenize + self.wait_s = float(wait_s) + self.poll_s = float(poll_s) + self.consumer = consumer + self.clock = clock + self.sleep = sleep + self._last: _Pulled | None = None + self.pulled_positions: list[int] = [] + self.cursor_log = cursor_log + self.merged = merged + self._round_kept: list[int] = [] # trained positions of the round being collected + if cursor_log is not None: + self.buffer.reset_to(cursor_log.last_cursor()) + + def __call__(self, name: str, *args: Any) -> Any: + """``call(name)`` for the three queue methods (:data:`QUEUE_METHODS`).""" + if name == "yeto_pull_sample": + return self.pull() + if name == "yeto_keep_last": + return self.keep() + if name == "yeto_drop_last": + return self.drop() + raise ValueError(f"{name} is not a queue call of the buffer feed") + + def _next(self) -> Any | None: + """The next selected entry; None after the sampler's data end.""" + started = self.clock() + while True: + selected = self.buffer.select(self.consumer, 1) + if selected: + return selected[0] + if self.buffer.data_ended: + return None + if self.clock() - started > self.wait_s: + raise BufferFeedError(f"no buffer group within {self.wait_s:g} s " + f"(read position {self.buffer.read_cursor})") + self.sleep(self.poll_s) + + def _local(self, version: int) -> int | None: + if self.merged is not None: + local = self.merged.local(version) + if local is None: # a merge this island never applied: unknown + return None + version = local + try: + return self.vmap().local(version) + except VerlTranslationError: + return None + + def pull(self) -> dict[str, Any] | None: + if self._last is not None: + raise BufferFeedError("pull before the previous group was kept or dropped") + item = self._next() + if item is None: + return None + samples = decode_group(self.buffer.read_payload(item.entry.payload_ref)) + if not samples: + raise BufferFeedError(f"buffer position {item.position}: empty group") + versions = [sample_version(s) for s in samples] + local = tuple(self._local(v) for v in versions) + lengths = tuple(int(s.get("response_length") or 0) for s in samples) + position = int(item.position) + self._last = _Pulled(position, samples, local, { + "position": position, "sample_id": item.entry.sample_id, + "behavior_versions": versions, "param_versions": list(local), + "source": (item.entry.provenance or {}).get("sampler_id"), + "samples": len(samples), "response_tokens": sum(lengths)}) + self.pulled_positions.append(position) + return {"uid": f"{UID_PREFIX}{int(item.position)}", "min_steps": local, "max_steps": local, + "calls": (None,) * len(samples), "lengths": lengths, + "rewards": tuple(float(s.get("reward") or 0.0) for s in samples)} + + def keep(self) -> None: + last = self._take_last() + if None in last.param_versions: # the collector keeps only groups whose versions it knows + raise BufferFeedError(f"buffer position {last.position}: kept with an unknown version " + f"{last.param_versions}") + self.add_group(group_rows(last.samples, uid=f"{UID_PREFIX}{last.position}", + param_versions=last.param_versions, shape=self.shape, + tokenize=self.tokenize)) + self._round_kept.append(last.position) + self.record({**last.record, "kept": True}) + + def drop(self) -> None: + self.record({**self._take_last().record, "kept": False}) + + def _take_last(self) -> _Pulled: + if self._last is None: + raise BufferFeedError("keep/drop without a pulled group") + last, self._last = self._last, None + return last + + def close_round(self, rollout_id: int) -> None: + """The round's collection ended (before it trains): record its trained + positions and the read position in the lane cursor log (module doc).""" + if self._last is not None: + raise BufferFeedError("round closed with a pulled group not kept or dropped") + kept, self._round_kept = self._round_kept, [] + if self.cursor_log is not None: + self.cursor_log.append(int(rollout_id), kept, int(self.buffer.read_cursor)) + + def read_cursor(self) -> int: + """The buffer position after the last pulled group (released to the + samplers for their backlog cap).""" + return int(self.buffer.read_cursor) + + +ASSEMBLE_METHOD = "yeto_assemble" # the rollout port's call after a round's collection + + +class BufferCall: + """The island's ``call(method, *args)``: the queue calls go to the feed, every + other call to the trainer actor (``actor_call``). ``round_id()``: the round + being collected; its collection ends at the assemble call, where the feed + closes the round (:meth:`BufferFeed.close_round`).""" + + def __init__(self, feed: BufferFeed, actor_call: Callable[..., Any], + round_id: Callable[[], int] | None = None) -> None: + self.feed = feed + self.actor_call = actor_call + self.round_id = round_id + + def __call__(self, name: str, *args: Any, **kwargs: Any) -> Any: + if name in QUEUE_METHODS: + return self.feed(name, *args) + if name == ASSEMBLE_METHOD and self.round_id is not None: + self.feed.close_round(self.round_id()) + return self.actor_call(name, *args, **kwargs) diff --git a/yeto/rl/adapters/verl/buffer_ports.py b/yeto/rl/adapters/verl/buffer_ports.py new file mode 100644 index 00000000..e8a7ba10 --- /dev/null +++ b/yeto/rl/adapters/verl/buffer_ports.py @@ -0,0 +1,215 @@ +"""verl buffer trainer: the ports that differ from the fully_async path (S25 C2). + +:func:`install_buffer_ports` turns a :class:`~.fully_async_ports.FullyAsyncIsland` +into a trainer island over the sample buffer (``--rl-sample-role trainer``): + +* the island's ``call`` becomes a :class:`~.buffer_feed.BufferCall`: the queue + calls go to the :class:`~.buffer_feed.BufferFeed`, the rest to the trainer + actor. The rollout port, the trainer port and the policy state are the + fully_async ones, unchanged; +* :class:`BufferReleasePublisher` replaces the checkpoint-engine push: the + trained adapter is merged into the base (:mod:`.merged_release`) and released + into the store for the samplers (``PolicyReleaseStore``, full weights) with + the trainer's buffer read position. The pairing ``param_version <-> outer + version`` is recorded at every publication, as on the fully_async path. + +The base policy (version 0, LoRA B all zero: the merged weights are the base) +is not released: the samplers start on the base weights as version 0. A +version-0 state with any non-zero B is released like any other version. + +Synced trainers (S25 C3, ``yeto.rl.engine.merged_release``; the same rules as +the Miles trainer island): every island records, at each publication, the +merged version (``merged_version``: the syncer's outer version of the applied +base) it first applied (:class:`~.buffer_feed.MergedVersions`; the feed turns a +group's merged behavior version into a local one with it). Only the publisher +(learner id 0) releases, once per new merged version, with the merged version +as release version, kept increasing by ``MergedReleaseGuard`` (a skipped +merge is reported as ``rl_policy_release_gap``). A publication without a new +merge (elastic: several local rounds per merge) releases nothing: its weights +are not a merged policy. The driver is bound after it is built +(:meth:`BufferReleasePublisher.bind`). + +The publication is complete when the release record and ``LATEST.json`` are in +the store. Whether a sampler has loaded it is the sampler's report +(``rl_sampler_policy_loaded`` on its tape) and the behavior versions of the +groups it writes next; the trainer does not wait for it (the samplers run +elsewhere and may be stopped). +""" + +from __future__ import annotations + +import hashlib +import json +import time +from collections.abc import Callable, Mapping, Sequence +from pathlib import Path +from typing import Any + +from yeto.rl.contracts import InferencePublicationManifest +from yeto.rl.engine.merged_release import BufferLane, MergedReleaseGuard, is_publisher, merged_version +from yeto.rl.engine.policy_release import FULL, PolicyReleaseStore +from yeto.rl.engine.ports import PlacementDescription, PublicationResult + +from .buffer_feed import BufferCall, BufferFeed, MergedVersions, RowShape +from .fully_async_ports import MEMBER, FullyAsyncIsland, FullyAsyncRolloutPool +from .fully_async_round import FULLY_ASYNC_PLACEMENT +from .merged_release import LORA_B, write_merged_checkpoint + +BASE_POLICY_VERSION = 0 +RELEASE_EVENT = "rl_policy_released" # the Miles trainer island's event name +CONSUMED_LOG = "verl-buffer-consumed-{learner}.jsonl" +RELEASE_LOG = "verl-buffer-release-{learner}.jsonl" + + +def _sha(text: str) -> str: + return hashlib.sha256(text.encode()).hexdigest() + + +def is_base_adapter(version: int, tensors: Mapping[str, Any]) -> bool: + """Version 0 with every LoRA B zero: base + adapter is the base itself.""" + return int(version) == BASE_POLICY_VERSION and all( + not bool(t.detach().abs().max() > 0) for name, t in tensors.items() if name.endswith(LORA_B)) + + +class BufferReleasePublisher: + """yeto publication as a store release of the merged policy (module doc). + + ``feed`` gives the buffer read position; ``base_dir`` is the base HF + checkpoint the adapter trains over; ``lane``: a synced trainer's buffer lane + (module doc; None for a single trainer).""" + + def __init__(self, island: FullyAsyncIsland, releases: PolicyReleaseStore, feed: BufferFeed, + *, base_dir: str | Path, clock: Callable[[], float] = time.monotonic, + merge: Callable[..., dict[str, Any]] = write_merged_checkpoint, + lane: BufferLane | None = None) -> None: + self.island = island + self.releases = releases + self.feed = feed + self.base_dir = Path(base_dir) + self.clock = clock + self.merge = merge + self.lane = lane + self.driver: Any = None + self.guard = MergedReleaseGuard(self._on_gap) + self.releasing = lane is None or is_publisher(lane.index, lane.count) + + def bind(self, driver: Any) -> None: + """The island's driver (synced trainers read its outer version).""" + self.driver = driver + if self.lane is not None: + self.island.emit("rl_policy_release_role", learner_id=self.lane.index, + num_learners=self.lane.count, publisher=self.releasing) + + def _on_gap(self, previous: int, version: int) -> None: + self.island.emit("rl_policy_release_gap", previous_version=previous, policy_version=version) + + def _release_version(self, state: Any) -> int | None: + """The version to release, or None (nothing to release). Single trainer: + the policy version. Synced trainer: the merged version, on the publisher, + the first time this island applies it (module doc).""" + local = int(state.policy_version) + if self.lane is None: + return local + if self.driver is None: + raise RuntimeError("synced buffer trainer: publish before bind(driver)") + merged = merged_version(self.driver, state) + new = self.feed.merged.record(merged, local) + if not (new and self.releasing): + return None + latest = self.releases.latest() + if self.guard.check(merged, None if latest is None else latest.version) == "stale": + self.island.emit("rl_policy_release_skipped", policy_version=merged, local_version=local, + latest_version=latest.version, reason="not_newer_than_latest") + return None + return merged + + def _origin(self, state: Any) -> dict[str, Any] | None: + if self.lane is None: + return None + return {"kind": "merged", "learner_id": self.lane.index, "num_learners": self.lane.count, + "local_version": int(state.policy_version)} + + def publish(self, state: Any) -> PublicationResult: + isl = self.island + version = int(state.policy_version) + param_version = int(isl.call("yeto_param_version")) + isl.vmap.record(param_version, version) # raises on contradiction / regression + digest = state.policy_tensor_hash() + tensors = dict(state.to_lora().tensors) + started = self.clock() + merge: dict[str, Any] = {} + release = None + position = self.feed.read_cursor() + released = self._release_version(state) + if released is not None and not is_base_adapter(released, tensors): + def write(directory: Path) -> None: + merge.update(self.merge(self.base_dir, tensors, directory)) + + release = self.releases.release(released, digest, write, buffer_selected=position, + weights=FULL, origin=self._origin(state)) + record = {"policy_version": version, "param_version": param_version, "policy_hash": digest, + "released": release is not None, "buffer_selected": position, + "seconds": round(self.clock() - started, 3), "merge": merge or None, + "version_map": isl.vmap.to_dict()} + if self.lane is not None: + record.update(merged_version=merged_version(self.driver, state), + merged_versions=self.feed.merged.to_dict()) + if release is not None: + record.update(seq=release.seq, path=release.path, release_version=release.version, + payload_bytes=sum(int(v) for v in (release.files or {}).values())) + isl.emit(RELEASE_EVENT, seq=release.seq, policy_version=release.version, + local_version=version, origin=release.origin, policy_hash=digest, + buffer_selected=position, weights=FULL, sha256=release.sha256, + engine_checksum=None, merged_modules=merge.get("merged_modules"), + max_relative_change=merge.get("max_relative_change"), + seconds=record["seconds"]) + isl.append_jsonl(RELEASE_LOG.format(learner=isl.learner_id), record) + isl.emit("rl_verl_version_map", param_version=param_version, outer_version=version) + isl.published = (version, digest) + isl.hashes[version] = digest + manifest = InferencePublicationManifest( + publication_mode="full", base_policy_version=None, target_policy_version=version, + target_policy_hash=digest, target_manifest_hash=_sha(f"{version}:{digest}"), + payload_hash=_sha(json.dumps({"release": record.get("path"), "hash": digest})), + # the release's files, or (nothing released) the adapter in bf16 as on the fully_async path + payload_bytes=int(record.get("payload_bytes") or sum(int(t.numel()) * 2 for t in tensors.values())), + complete=True) + return PublicationResult(manifest, frozenset({MEMBER})) + + +class BufferPlacement: + """The trainer GPUs only: the samplers run on other islands.""" + + def describe(self) -> PlacementDescription: + return PlacementDescription(FULLY_ASYNC_PLACEMENT, ("cuda:trainer-0",), ("store:sample-buffer",), + extra={"engine": "verl-buffer-trainer"}) + + +def install_buffer_ports(island: FullyAsyncIsland, *, buffer: Any, releases: PolicyReleaseStore, + shape: RowShape, base_dir: str | Path, + tokenize: Callable[[str], Sequence[int]] | None = None, + wait_s: float = 3600.0, poll_s: float = 2.0, cursor_log: Any = None, + **feed_kwargs: Any) -> tuple[FullyAsyncRolloutPool, BufferReleasePublisher, + BufferPlacement]: + """Route the island's queue calls to a buffer feed and return the rollout + port, the publisher and the placement of a buffer trainer (module doc); the + feed is ``publisher.feed``. A lane reader ``buffer`` (``buffer.lane``) makes + a synced trainer: ``cursor_log`` is its lane cursor log + (``merged_release.lane_cursor_log``).""" + lane = getattr(buffer, "lane", None) + if (lane is None) != (cursor_log is None): + raise ValueError("a synced buffer trainer needs both a lane reader and its cursor log") + actor_call = island.call + log = CONSUMED_LOG.format(learner=island.learner_id) + + def record(fields: dict[str, Any]) -> None: + island.append_jsonl(log, fields) + + feed = BufferFeed(buffer, vmap=lambda: island.vmap, shape=shape, + add_group=lambda rows: actor_call("yeto_add_group", rows), record=record, + tokenize=tokenize, wait_s=wait_s, poll_s=poll_s, cursor_log=cursor_log, + merged=None if lane is None else MergedVersions(), **feed_kwargs) + island.call = BufferCall(feed, actor_call, round_id=lambda: int(island.published[0])) + island.start_rollouter = None # no rollouter: the samplers run elsewhere + publisher = BufferReleasePublisher(island, releases, feed, base_dir=base_dir, lane=lane) + return FullyAsyncRolloutPool(island), publisher, BufferPlacement() diff --git a/yeto/rl/adapters/verl/buffer_trainer_runner.py b/yeto/rl/adapters/verl/buffer_trainer_runner.py new file mode 100644 index 00000000..5601eee2 --- /dev/null +++ b/yeto/rl/adapters/verl/buffer_trainer_runner.py @@ -0,0 +1,147 @@ +"""verl buffer trainer island, image side (S25 C2; runs inside the verl image). + +The trainer island of a sampler/trainer split (``--rl-sample-role trainer``): +verl's ``FullyAsyncTrainer`` trains, yeto's driver runs the rounds over the +fully_async ports, and the queued samples come from the durable sample buffer +instead of verl's rollouter (:mod:`.buffer_feed`, :mod:`.buffer_ports`). + +* :class:`YetoBufferTaskRunner` builds only the trainer (tokenizer, worker + mapping, the trainer actor and its workers): no rollouter, no MessageQueue, + no checkpoint engine, no initial weight push. Each publication is a store + release instead (``BufferReleasePublisher``). +* :data:`YetoBufferTrainer` is the fully_async trainer actor + (``_YetoTrainerMixin``: assemble, train step with the mismatch criteria and + the per-source off-policy metrics, LoRA export) plus ``yeto_add_group`` (one + kept group's rows as a verl ``RolloutSample``) and ``yeto_param_version``. + +verl itself is not changed. Only exercised on a GPU; the logic it relies on is +in the CPU-tested :mod:`.buffer_feed` and :mod:`.buffer_ports`. +""" + +from __future__ import annotations + +import os + +import numpy as np +import ray + +from .fully_async_runner import YetoFullyAsyncTaskRunner, _TrainerBase, _YetoTrainerMixin + +FLOAT_ROWS = ("rollout_log_probs", "rm_scores") +INT32_FIELDS = ("__num_turns__",) +SAMPLE_BUFFER_PLAN_KEY = "sample_buffer" # plan section of a buffer trainer (island_entry) + + +def _object_array(values) -> np.ndarray: + out = np.empty(len(values), dtype=object) + out[:] = list(values) + return out + + +def rows_to_dataproto(rows: dict): + """The plain rows of :func:`.buffer_feed.group_rows` as a verl ``DataProto``.""" + import torch + from tensordict import TensorDict + from verl import DataProto + + tensors = rows["tensors"] + n = len(tensors["prompts"]) + batch = TensorDict({name: torch.tensor(values, dtype=torch.float32 if name in FLOAT_ROWS else torch.long) + for name, values in tensors.items()}, batch_size=n) + non_tensor = {name: (np.asarray(values, dtype=np.int32) if name in INT32_FIELDS else _object_array(values)) + for name, values in rows["non_tensor"].items()} + meta = {key: (list(value) if isinstance(value, list) else value) for key, value in rows["meta_info"].items()} + return DataProto(batch=batch, non_tensor_batch=non_tensor, meta_info=meta) + + +class _YetoBufferTrainerMixin(_YetoTrainerMixin): + def yeto_add_group(self, rows: dict) -> int: + from verl.experimental.fully_async_policy.detach_utils import RolloutSample + + if getattr(self, "_yeto_kept", None) is None: + self._yeto_kept = [] + proto = rows_to_dataproto(rows) + self._yeto_kept.append(RolloutSample(full_batch=proto, sample_id=str(rows["non_tensor"]["uid"][0]), + epoch=0, rollout_status={})) + return len(self._yeto_kept) + + def yeto_param_version(self) -> int: + return int(self.current_param_version) + + async def _yeto_queue_size(self): + return None # no MessageQueue: the buffer backlog is the samplers' report + + +YetoBufferTrainer = ray.remote(num_cpus=10)( + type("YetoBufferTrainer", (_YetoBufferTrainerMixin, _TrainerBase), {})) + + +class YetoBufferTaskRunner(YetoFullyAsyncTaskRunner): + """The trainer-only set-up, then the yeto driver over the buffer ports.""" + + trainer_actor = YetoBufferTrainer + + def _initialize_trainer(self, config, total_train_steps: int) -> None: + """verl's ``_initialize_components`` (fork acad9875, fully_async_main.py:51) + up to the trainer, without the rollouter and the MessageQueue.""" + from omegaconf import OmegaConf + from verl.experimental.separation.utils import create_role_worker_mapping + from verl.utils import hf_processor, hf_tokenizer + from verl.utils.fs import copy_to_local + + OmegaConf.resolve(config) + local_path = copy_to_local(config.actor_rollout_ref.model.path, + use_shm=config.actor_rollout_ref.model.get("use_shm", False)) + trust_remote_code = config.data.get("trust_remote_code", False) + self.components["tokenizer"] = hf_tokenizer(local_path, trust_remote_code=trust_remote_code) + self.components["processor"] = hf_processor(local_path, trust_remote_code=trust_remote_code, use_fast=True) + self.components["config"] = config + mapping, group_cls = create_role_worker_mapping(config) + self.components["role_worker_mapping"] = mapping + self.components["ray_worker_group_cls"] = group_cls + self._create_trainer(config) + trainer = self.components["trainer"] + ray.get(trainer.set_total_train_steps.remote(int(total_train_steps))) + ray.get(trainer.load_checkpoint.remote()) + + def run(self, config, plan: dict): + from omegaconf import OmegaConf + + from yeto.island_credential_guard import check_island_credentials + from yeto.rl.engine.policy_release import PolicyReleaseStore + from yeto.rl.engine.resume import store_for + from yeto.rl.engine.merged_release import BufferLane, lane_cursor_log + from yeto.rl.engine.sample_buffer import open_trainer_buffer + + from .buffer_feed import RowShape + from .buffer_ports import install_buffer_ports + from .trainer import run_fully_async_island + + check_island_credentials() # secret-handling-hardening D4 (buffer trainer Ray actor) + section = plan[SAMPLE_BUFFER_PLAN_KEY] + self._initialize_trainer(config, int(plan["global_rounds"])) + trainer = self.components["trainer"] + tokenizer = self.components["tokenizer"] + + def call(name, *args, timeout=None): + return ray.get(getattr(trainer, name).remote(*args), timeout=timeout) + + environ = dict(os.environ) + store = store_for(section["store"], environ=environ) + # S25 C3: a synced trainer reads its lane of the buffer (section "lane": "M/N") + buffer = open_trainer_buffer(store, batch_size=1, create_index=False, + lane=BufferLane.parse(section.get("lane"))) + releases = PolicyReleaseStore(store_for(section.get("release_store") or section["store"], + environ=environ)) + shape = RowShape(int(config.data.max_prompt_length), int(config.data.max_response_length), + int(tokenizer.pad_token_id)) + + def ports(island): + return install_buffer_ports(island, buffer=buffer, releases=releases, shape=shape, + base_dir=plan["model_path"], tokenize=tokenizer.encode, + wait_s=float(section.get("wait_s") or 3600.0), + poll_s=float(section.get("poll_s") or 2.0), + cursor_log=lane_cursor_log(buffer)) + + return run_fully_async_island(call, None, plan, OmegaConf.to_container(config, resolve=True), + buffer_ports=ports) diff --git a/yeto/rl/adapters/verl/cut.py b/yeto/rl/adapters/verl/cut.py index 5a3e2d2d..20b6b3ae 100644 --- a/yeto/rl/adapters/verl/cut.py +++ b/yeto/rl/adapters/verl/cut.py @@ -19,7 +19,8 @@ Restore needs the same world size (FSDP2 shards are per rank; a DP change is refused before any rank writes). Each rank checks its file checksum, loads it into the freshly built module/optimizer (DTensors rebuilt with the fresh -module's mesh and placements), and reports a digest of the loaded local state +module's mesh and placements), and reports a digest of the state its optimizer +now steps (N7.6c: read from ``optimizer.param_groups``, not the module copy) (must equal the one recorded at save) plus its resume state. Units (verl differs from Miles here): the run step is one trainer round (one @@ -89,21 +90,27 @@ def layout_for(world: int) -> dict[str, int]: # -------------------------------------------------------------------------- -def _state_options() -> Any: +def _state_options(*, load: bool = False) -> Any: + """Local shards of the trainable state. A load is not strict: the cut holds no + frozen base weights (LoRA), which a strict load reports as missing (GPU T-C4). + ``restore_rank`` checks the trainable names against the cut before loading.""" from torch.distributed.checkpoint.state_dict import StateDictOptions - return StateDictOptions(full_state_dict=False, ignore_frozen_params=True) + return StateDictOptions(full_state_dict=False, ignore_frozen_params=True, strict=not load) def _is_dtensor(value: Any) -> bool: return hasattr(value, "to_local") and hasattr(value, "placements") +def _local_tensor(value: Any) -> Any: + return value.to_local() if _is_dtensor(value) else value + + def _to_local(value: Any) -> Any: import torch - if _is_dtensor(value): - value = value.to_local() + value = _local_tensor(value) return value.detach().to("cpu").clone() if isinstance(value, torch.Tensor) else value @@ -122,26 +129,45 @@ def _local_state(handles: TrainHandles) -> tuple[dict[str, Any], dict[str, Any], return model, optim, dtensors -def _digest(model: Mapping[str, Any], optim: Mapping[str, Any]) -> str: - """sha256 over the local trainable state (names, dtypes, shapes, bytes; group hyper-parameters).""" +def _stepped_state(handles: TrainHandles) -> tuple[dict[str, dict[str, Any]], list[dict[str, Any]]]: + """The tensors the optimizer actually steps on this rank, keyed by parameter name. + + Read from ``optimizer.param_groups`` and ``optimizer.state`` (local shards), not + from the module's state dict: a digest over the module copy cannot see an + optimizer that steps another tensor (S25 N7.6c; on Miles the CPU-offload host + master went stale after a restore while the cut digest matched, GPU T-G3). A + stepped tensor that is no module parameter is named by its group position.""" + names = {id(p): n for n, p in handles.module.named_parameters()} + state: dict[str, dict[str, Any]] = {} + groups = [] + for gi, group in enumerate(handles.optimizer.param_groups): + groups.append({k: v for k, v in group.items() if k != "params"}) + for pi, param in enumerate(group["params"]): + name = names.get(id(param), f"") + entry = {"param": param} + entry.update({k: v for k, v in handles.optimizer.state.get(param, {}).items()}) + state[name] = entry + return state, groups + + +def _digest(handles: TrainHandles) -> str: + """sha256 over the stepped state (names, dtypes, shapes, local bytes) and the group + hyper-parameters; see :func:`_stepped_state`.""" import torch h = hashlib.sha256() - - def tensor(name: str, t: Any) -> None: - if not isinstance(t, torch.Tensor): - h.update(f"{name}={t!r};".encode()) - return - flat = t.detach().to("cpu").contiguous().reshape(-1) - h.update(f"{name}|{flat.dtype}|{tuple(t.shape)}|".encode()) - h.update(flat.view(torch.uint8).numpy().tobytes()) - - for name in sorted(model): - tensor(f"model/{name}", model[name]) - for name in sorted(optim["state"]): - for key in sorted(optim["state"][name]): - tensor(f"optim/{name}/{key}", optim["state"][name][key]) - h.update(json.dumps(optim["param_groups"], sort_keys=True, default=repr).encode()) + state, groups = _stepped_state(handles) + for name in sorted(state): + for key in sorted(state[name]): + value = state[name][key] + label = f"{name}/{key}" + if not isinstance(value, torch.Tensor): + h.update(f"{label}={value!r};".encode()) + continue + flat = _local_tensor(value).detach().to("cpu").contiguous().reshape(-1) + h.update(f"{label}|{flat.dtype}|{tuple(value.shape)}|".encode()) + h.update(flat.view(torch.uint8).numpy().tobytes()) + h.update(json.dumps(groups, sort_keys=True, default=repr).encode()) return h.hexdigest() @@ -200,7 +226,7 @@ def save_rank(handles: TrainHandles, *, directory: str, cut_id: str) -> dict[str fsync_file(tmp) os.replace(tmp, path) fsync_dir(path.parent) - summary = _rank_summary(handles, rank=rank, world=world, digest=_digest(model, optim)) + summary = _rank_summary(handles, rank=rank, world=world, digest=_digest(handles)) return {"path": name, "sha256": sha256_file(path), "bytes": path.stat().st_size, **summary} @@ -255,11 +281,10 @@ def value(key: str, local: Any, like: str) -> Any: "param_groups": saved["optim"]["param_groups"], } set_state_dict(handles.module, handles.optimizer, model_state_dict=model_sd, optim_state_dict=optim_sd, - options=_state_options()) + options=_state_options(load=True)) handles.scheduler.load_state_dict(saved["scheduler"]) _set_rng_state(saved["rng"]) - model, optim, _ = _local_state(handles) - return {"path": entry.path, **_rank_summary(handles, rank=rank, world=world, digest=_digest(model, optim))} + return {"path": entry.path, **_rank_summary(handles, rank=rank, world=world, digest=_digest(handles))} def resume_rank(handles: TrainHandles) -> dict[str, Any]: diff --git a/yeto/rl/adapters/verl/fully_async_runner.py b/yeto/rl/adapters/verl/fully_async_runner.py index 0d07f881..900ccf6d 100644 --- a/yeto/rl/adapters/verl/fully_async_runner.py +++ b/yeto/rl/adapters/verl/fully_async_runner.py @@ -96,7 +96,11 @@ async def yeto_assemble(self): self._collect_metrics_from_samples(batch, self.metrics) batch.meta_info["temperature"] = self.config.actor_rollout_ref.rollout.temperature self._yeto_batch = batch - return {"samples": len(kept), "queue_size": int(await self.message_queue_client.get_queue_size())} + return {"samples": len(kept), "queue_size": await self._yeto_queue_size()} + + async def _yeto_queue_size(self): + """Samples left in verl's MessageQueue (a buffer trainer has none: None).""" + return int(await self.message_queue_client.get_queue_size()) def _yeto_criteria(self, batch) -> dict: from yeto.rl.engine import mismatch_criteria @@ -196,7 +200,11 @@ async def yeto_export_lora(self): class YetoFullyAsyncTaskRunner(_TaskRunnerBase): """verl's component set-up (trainer, rollouter, MessageQueue, initial weight - sync); then the yeto driver instead of ``_run_training_loop``.""" + sync); then the yeto driver instead of ``_run_training_loop``. + + ``trainer_actor``: the Ray actor class :meth:`_create_trainer` builds.""" + + trainer_actor = YetoFullyAsyncTrainer def run(self, config, plan: dict): from yeto.island_credential_guard import check_island_credentials @@ -240,7 +248,7 @@ def call(name, *args, timeout=None): def _create_trainer(self, config) -> None: """verl's ``_create_trainer`` (fork acad9875, fully_async_main.py:138), building - :data:`YetoFullyAsyncTrainer`. Setting ``fa_main.FullyAsyncTrainer`` did not + :attr:`trainer_actor`. Setting ``fa_main.FullyAsyncTrainer`` did not reach verl's method inside the Ray actor (s19-verl64b-async7-20261010a: the trainer was verl's own class, "no attribute yeto_configure"), so the method is overridden here instead of relying on that module global.""" @@ -250,7 +258,7 @@ def _create_trainer(self, config) -> None: print("[ASYNC MAIN] Starting create trainer (yeto)...") mapping = {role: cls for role, cls in self.components["role_worker_mapping"].items() if role != Role.Rollout} - trainer = YetoFullyAsyncTrainer.remote( + trainer = self.trainer_actor.remote( config=config, tokenizer=self.components["tokenizer"], role_worker_mapping=mapping, @@ -261,4 +269,4 @@ def _create_trainer(self, config) -> None: trainer.yeto_configure # noqa: B018 - fail fast (AttributeError) if this is not our class ray.get(trainer.init_workers.remote()) self.components["trainer"] = trainer - print("[ASYNC MAIN] YetoFullyAsyncTrainer created and initialized successfully") + print("[ASYNC MAIN] yeto trainer actor created and initialized successfully") diff --git a/yeto/rl/adapters/verl/harbor_agent_loop.py b/yeto/rl/adapters/verl/harbor_agent_loop.py index 9093839d..07ce7898 100644 --- a/yeto/rl/adapters/verl/harbor_agent_loop.py +++ b/yeto/rl/adapters/verl/harbor_agent_loop.py @@ -80,10 +80,13 @@ run_agent_task, ) from yeto.rl.agentic.loop_local import LoopLocal +from yeto.rl.agentic.outcome import policy_output_failure from yeto.rl.agentic.types import TokenTurn +from yeto.rl.engine.sampler_engine import FINISH_LENGTH from yeto.rl.harness.harbor import HarborHarness, HarborSettings, build_trial_config from yeto.rl.harness.harbor.boat_client import check_in_flight_cap +from .rollout_engine import bind_rollout_client, weight_steps from .run_config_rules import AGENTIC_ADV_ESTIMATOR as YETO_GRPO logger = logging.getLogger(__name__) @@ -190,6 +193,9 @@ def __init__(self, *args: Any, **kwargs: Any) -> None: async def run(self, sampling_params: dict[str, Any], **kwargs: Any) -> AgentLoopOutput: task = task_from_row(kwargs) recorder = self.recorder_factory() + client = getattr(self, "server_manager", None) + if client is not None: # a recorder over verl's own rollout needs verl's client (N10.9) + bind_rollout_client(recorder, client) session = recorder_session_id(task) endpoint = await recorder.open(session, sampling_params) try: @@ -203,6 +209,8 @@ async def run(self, sampling_params: dict[str, Any], **kwargs: Any) -> AgentLoop except TokenAlignmentError as exc: if trajectory.outcome is not Outcome.ABORTED: trajectory = _abort(trajectory, f"token_alignment: {exc}") + else: + trajectory = _policy_output_failure(trajectory, turns) segments: TokenSegments | None = None if trajectory.outcome is not Outcome.ABORTED: @@ -212,10 +220,10 @@ async def run(self, sampling_params: dict[str, Any], **kwargs: Any) -> AgentLoop trajectory = _abort(trajectory, f"token_alignment: {exc}") if segments is not None and len(segments.prompt_ids) > self.prompt_length: trajectory, segments = _abort(trajectory, "prompt_exceeds_prompt_length"), None - return self._output(trajectory, segments, turns) + return self._output(trajectory, segments, turns, steps=weight_steps(recorder, session)) def _output(self, trajectory: AgentTrajectory, segments: TokenSegments | None, - turns: Sequence[TokenTurn]) -> AgentLoopOutput: + turns: Sequence[TokenTurn], *, steps: tuple[int, int] | None = None) -> AgentLoopOutput: if segments is None: segments = self._placeholder() metadata = trajectory.to_metadata() @@ -232,8 +240,12 @@ def _output(self, trajectory: AgentTrajectory, segments: TokenSegments | None, "outcome_reason": trajectory.reason, "policy_version": trajectory.policy_version, "truncated": segments.truncated, + # Model calls that hit their token cap (the agent then saw a cut reply). + "length_capped_turns": sum(t.finish_reason == FINISH_LENGTH for t in turns), "agent_metrics": metadata["agent_metrics"], "provenance": turn_provenance(turns), + # verl V1 reads both for its staleness metrics; recorders without verl's rollout leave them out + **({"min_global_steps": steps[0], "max_global_steps": steps[1]} if steps else {}), }, ) @@ -243,6 +255,14 @@ def _placeholder(self) -> TokenSegments: return TokenSegments(prompt_ids=(pad,), response_ids=(pad,), response_mask=(0,), response_logprobs=(0.0,)) +def _policy_output_failure(trajectory: AgentTrajectory, turns: Sequence[TokenTurn]) -> AgentTrajectory: + """A malformed model reply that ended the agent is a 0 reward, not ABORTED (N10.9).""" + updated = policy_output_failure(trajectory, [t.policy_error for t in turns if t.policy_error]) + if updated.outcome is not trajectory.outcome: + _gate().metrics.reclassify(trajectory.outcome, updated.outcome) + return updated + + def _abort(trajectory: AgentTrajectory, reason: str) -> AgentTrajectory: _gate().metrics.reclassify(trajectory.outcome, Outcome.ABORTED) return AgentTrajectory(task=trajectory.task, outcome=Outcome.ABORTED, reward=0.0, reason=reason, diff --git a/yeto/rl/adapters/verl/image.py b/yeto/rl/adapters/verl/image.py index e7e9ca14..b096d1b1 100644 --- a/yeto/rl/adapters/verl/image.py +++ b/yeto/rl/adapters/verl/image.py @@ -5,7 +5,8 @@ dependencies from the fork's own ``uv.lock`` (``uv sync --frozen --all-packages --extra vllm --extra fsdp``), plus yeto's one-hook local patch (``patch_verl``). ``--rl-image verl-build:`` selects it; Modal caches -the layers, so only the first run builds. yeto's own image default +the layers, so only the first run builds. yeto's object-store client (boto3, +``STORE_CLIENT_REQUIREMENTS``) is added to the same venv. yeto's own image default (ghcr.io/michaellchung/yeto-miles-ports) is unchanged. """ @@ -13,8 +14,8 @@ from pathlib import Path -from .pins import (VERL_BASE_IMAGE, VERL_COMMIT, VERL_PYTHON, VERL_REPOSITORY, VERL_ROOT, - VERL_UV_VERSION) +from .pins import (STORE_CLIENT_REQUIREMENTS, VERL_BASE_IMAGE, VERL_COMMIT, VERL_PYTHON, VERL_REPOSITORY, + VERL_ROOT, VERL_UV_VERSION) VENV_BIN = f"{VERL_ROOT}/.venv/bin" @@ -30,6 +31,12 @@ def build_commands(commit: str = VERL_COMMIT) -> list[str]: ] +def store_client_commands() -> list[str]: + """An island over an s3:// or oss:// store reads it with boto3 (S25 C3V); its own + layer, so the cached verl build above stays valid.""" + return [f"uv pip install --python {VENV_BIN}/python " + " ".join(STORE_CLIENT_REQUIREMENTS)] + + def image_env() -> dict[str, str]: return { "PATH": f"{VENV_BIN}:/usr/local/nvidia/bin:/usr/local/cuda/bin:/usr/local/sbin:/usr/local/bin:" @@ -46,6 +53,7 @@ def modal_image(modal, commit: str = VERL_COMMIT): .apt_install("git", "build-essential", "curl", "libnuma1") .pip_install(f"uv=={VERL_UV_VERSION}") .run_commands(*build_commands(commit)) + .run_commands(*store_client_commands()) .add_local_file(str(patch), "/workspace/patch_verl.py", copy=True) .run_commands(f"{VENV_BIN}/python /workspace/patch_verl.py {VERL_ROOT}") .env(image_env()) diff --git a/yeto/rl/adapters/verl/island_entry.py b/yeto/rl/adapters/verl/island_entry.py index 0508e441..ab6f640f 100644 --- a/yeto/rl/adapters/verl/island_entry.py +++ b/yeto/rl/adapters/verl/island_entry.py @@ -67,6 +67,22 @@ def parse_args(argv): p.add_argument("--rl-island-lease-s", type=float, default=None) p.add_argument("--rl-echo-events", action="store_true") p.add_argument("--rl-max-policy-age", type=int, default=0) + # S25 C2: a trainer island over the sample buffer (sample_roles; the sampler is + # the engine sampler island, yeto.rl.engine.engine_sampler). --rl-resume-store + # names the buffer's store there, and the round-cut store otherwise (N7.6b) + p.add_argument("--rl-sample-role", default=None) + p.add_argument("--rl-resume-store", default=None) + p.add_argument("--rl-sample-wait-s", type=float, default=None) + # S25 C3: a synced trainer island (syncer, learner m of N) reads lane m of the + # buffer in this store; the launcher forwards both (yeto.rl.engine.merged_release) + p.add_argument("--rl-sample-store", default=None) + p.add_argument("--rl-sample-lanes", type=int, default=None) + # N7.6b round cuts (same flags as the Miles island) + p.add_argument("--rl-resume-state-dir", default=None) + p.add_argument("--rl-cut-every", type=int, default=None) + p.add_argument("--rl-cut-keep", type=int, default=None) + p.add_argument("--rl-stop-after-rounds", type=int, default=None) + p.add_argument("--rl-resume-allow-config-change", action="store_true") args, ignored = p.parse_known_args(argv) return args, ignored @@ -252,6 +268,7 @@ def ensure_ray_cpus(min_cpus: int = MIN_RAY_CPUS) -> dict: def main(argv=None) -> int: args, ignored = parse_args(sys.argv[1:] if argv is None else argv) run_config_rules.refuse_agentic_argv(ignored) # S25 N10.4: never ignore agentic flags + buffer_section = sample_buffer_section(args) # S25 C2: refused roles stop before any work OUTPUT.mkdir(parents=True, exist_ok=True) lid = args.learner_id t0 = time.time() @@ -316,6 +333,12 @@ def main(argv=None) -> int: sync = "strict" if sync != "none" and not args.syncer: raise SystemExit("--syncer is required unless --rl-single-island-no-sync") + from .round_cuts import resume_config + + try: + resume = resume_config(args) + except ValueError as exc: + raise SystemExit(str(exc)) from exc asserted_keys = vconf.ASSERTED_KEYS if limit: # agentic-rollout-utilization 6.4b: fully_async path overrides, asserted_keys = fully_async_plan(limit, overrides, spec_dict, args, sync) @@ -341,7 +364,9 @@ def main(argv=None) -> int: "publish_selftest": os.environ.get("YETO_VERL_PUBLISH_SELFTEST") == "1", "run_config": run.to_dict(), "overrides": overrides, "asserted": {k: _override_value(overrides, k) for k in asserted_keys}, "max_policy_age": limit, "fully_async": bool(limit), + "resume": resume, "versions": manifest["versions"], + **({"sample_buffer": buffer_section} if buffer_section else {}), } plan_path = work / f"plan-{lid}.json" plan_path.write_text(json.dumps(plan, indent=1, default=str)) @@ -357,7 +382,7 @@ def main(argv=None) -> int: env["YETO_VERL_READBACK_DIR"] = str(readback) env["VERL_FILE_LOGGER_PATH"] = str(OUTPUT / f"verl-file-logger-{lid}.jsonl") readback.mkdir(parents=True, exist_ok=True) - entry = "main_fully_async" if limit else "main" + entry = "main_buffer_trainer" if buffer_section else "main_fully_async" if limit else "main" cmd = [sys.executable, "-c", f"from yeto.rl.adapters.verl.verl_main import {entry}; {entry}()", *overrides] print("[yeto-verl] " + " ".join(shlex.quote(c) for c in cmd), flush=True) log = OUTPUT / f"verl-train-{lid}.log" @@ -372,6 +397,29 @@ def main(argv=None) -> int: return code +def sample_buffer_section(args) -> dict | None: + """The plan's ``sample_buffer`` section of a buffer trainer (S25 C2); None for + role ``both``. A verl island is never a sampler: a sampler island runs the + engine sampler (``python -m yeto.rl.engine.engine_sampler``) on any engine.""" + from yeto.rl.engine import sample_roles + + role = sample_roles.sample_role(args) + if role == sample_roles.BOTH: + return None + if role != sample_roles.TRAINER: + raise SystemExit(f"--rl-sample-role {role}: a verl island only trains from the buffer; " + "run the sampler as the engine sampler island (yeto.rl.engine.engine_sampler)") + if int(args.rl_max_policy_age or 0) < 1: + raise SystemExit("--rl-sample-role trainer needs --rl-max-policy-age >= 1 (the fully_async path)") + try: + sample_roles.check_sample_store(args) + lane = sample_roles.buffer_lane(args) + except ValueError as error: + raise SystemExit(str(error)) from None + config = sample_roles.resolve(args, groups_per_round=args.groups_per_round) + return {"store": config.store, "wait_s": config.wait_s, "lane": None if lane is None else lane.env()} + + def fully_async_plan(limit: int, overrides: list[str], spec_dict: dict | None, args, sync: str): """Limit > 0: fully_async overrides + asserted keys; the algorithm contract is checked here (bounded staleness needs ``execution.max_policy_staleness >= limit`` diff --git a/yeto/rl/adapters/verl/merged_release.py b/yeto/rl/adapters/verl/merged_release.py new file mode 100644 index 00000000..a4d829f6 --- /dev/null +++ b/yeto/rl/adapters/verl/merged_release.py @@ -0,0 +1,112 @@ +"""A verl LoRA policy as a full HF checkpoint for a sampler island (S25 C2). + +The engine sampler island loads full checkpoints only (``SamplerEngine.load``: +the SGLang and vllm-ascend adapters reload weights from a model directory). A +verl buffer trainer trains a LoRA adapter over a frozen base, so each release +merges the adapter into the base: ``W' = W + B @ A`` for every LoRA module +(yeto canonical LoRA has no scaling: ``lora_alpha == rank``, +``config.validate``), summed in float32 and stored in the base tensor's dtype. + +The output has the base's file layout (the same safetensors file names, the same +index, the same non-weight files: config, tokenizer, generation config), so the +engine reads it exactly as it reads the base. A module the adapter names that +the base lacks, or an A/B pair with mismatched shapes, is an error: a silently +skipped module would serve the base weights for it. + +Rounding: a bf16 base keeps about 8 bits of mantissa, so a LoRA delta far below +``|W| / 256`` does not change the stored weight. The sampler then serves the +merged weights, not the trainer's exact ``W + BA``; the trainer's log-prob +recomputation shows the difference as off-policy mismatch. +""" + +from __future__ import annotations + +import json +import shutil +from collections.abc import Mapping +from pathlib import Path +from typing import Any + +from .param_names import CANONICAL_PREFIX, ParamNameError + +INDEX_FILE = "model.safetensors.index.json" +SINGLE_FILE = "model.safetensors" +LORA_A, LORA_B = ".lora_A.weight", ".lora_B.weight" + + +class MergeError(ValueError): + """The adapter does not fit the base checkpoint.""" + + +def lora_pairs(lora: Mapping[str, Any]) -> dict[str, tuple[Any, Any]]: + """{base weight name: (A, B)} from canonical LoRA tensors + (``base_model.model..lora_{A,B}.weight``).""" + pairs: dict[str, dict[str, Any]] = {} + for name, tensor in lora.items(): + if not name.startswith(CANONICAL_PREFIX) or not name.endswith((LORA_A, LORA_B)): + raise ParamNameError(f"not a canonical LoRA name: {name!r}") + side = "A" if name.endswith(LORA_A) else "B" + module = name[len(CANONICAL_PREFIX):-len(LORA_A)] + pairs.setdefault(f"{module}.weight", {})[side] = tensor + out = {} + for weight, sides in sorted(pairs.items()): + if set(sides) != {"A", "B"}: + raise MergeError(f"{weight}: LoRA has only side {sorted(sides)}") + out[weight] = (sides["A"], sides["B"]) + return out + + +def weight_files(base_dir: Path) -> dict[str, list[str]]: + """{safetensors file: [tensor names]} of a HF checkpoint directory.""" + index = base_dir / INDEX_FILE + if index.is_file(): + files: dict[str, list[str]] = {} + for name, file in json.loads(index.read_text(encoding="utf-8"))["weight_map"].items(): + files.setdefault(file, []).append(name) + return files + if (base_dir / SINGLE_FILE).is_file(): + from safetensors import safe_open + + with safe_open(str(base_dir / SINGLE_FILE), framework="pt") as handle: + return {SINGLE_FILE: list(handle.keys())} + raise MergeError(f"{base_dir}: no {INDEX_FILE} and no {SINGLE_FILE}") + + +def write_merged_checkpoint(base_dir: str | Path, lora: Mapping[str, Any], out_dir: str | Path) -> dict[str, Any]: + """Write base + LoRA into ``out_dir`` (module doc); returns a summary + (merged modules, files, the largest relative change of a merged weight).""" + import torch + from safetensors.torch import load_file, save_file + + base, out = Path(base_dir), Path(out_dir) + out.mkdir(parents=True, exist_ok=True) + pairs = lora_pairs(lora) + files = weight_files(base) + known = {name for names in files.values() for name in names} + missing = sorted(set(pairs) - known) + if missing: + raise MergeError(f"the base has no weight for {len(missing)} LoRA modules, e.g. {missing[0]}") + merged, max_rel = 0, 0.0 + for file in sorted(files): + tensors = load_file(str(base / file)) + for name in files[file]: + if name not in pairs: + continue + a, b = (t.detach().to("cpu", torch.float32) for t in pairs[name]) + weight = tensors[name] + if a.dim() != 2 or b.dim() != 2 or b.shape[1] != a.shape[0] or \ + (b.shape[0], a.shape[1]) != tuple(weight.shape): + raise MergeError(f"{name}: A {tuple(a.shape)}, B {tuple(b.shape)} do not fit " + f"the weight {tuple(weight.shape)}") + delta = b @ a + new = (weight.to(torch.float32) + delta).to(weight.dtype) + norm = float(weight.to(torch.float32).norm()) + if norm > 0: + max_rel = max(max_rel, float((new.to(torch.float32) - weight.to(torch.float32)).norm()) / norm) + tensors[name] = new.contiguous() + merged += 1 + save_file(tensors, str(out / file), metadata={"format": "pt"}) + for item in sorted(base.iterdir()): + if item.is_file() and item.name not in files and not item.name.startswith("."): + shutil.copy2(item, out / item.name) + return {"merged_modules": merged, "files": sorted(files), "max_relative_change": max_rel} diff --git a/yeto/rl/adapters/verl/pins.py b/yeto/rl/adapters/verl/pins.py index a7922c40..9736ae1e 100644 --- a/yeto/rl/adapters/verl/pins.py +++ b/yeto/rl/adapters/verl/pins.py @@ -20,6 +20,10 @@ "peft": "0.19.1", } +# yeto's object-store client (s3:// and oss:// checkpoint stores: yeto.cloud.s3_store), +# not in the fork's uv.lock; the same pin as yeto's pyproject (a test keeps them equal). +STORE_CLIENT_REQUIREMENTS = ("boto3==1.43.104",) + # --- Ascend NPU stack (rl-verl-backend 3.11/3.14) ---------------------------- # # The NPU stack is a second, separate image: vLLM 0.23 + torch 2.10 + CANN diff --git a/yeto/rl/adapters/verl/ports_impl.py b/yeto/rl/adapters/verl/ports_impl.py index 6bfaf3b0..4952a73c 100644 --- a/yeto/rl/adapters/verl/ports_impl.py +++ b/yeto/rl/adapters/verl/ports_impl.py @@ -216,6 +216,10 @@ def generate(self, rollout_id: int, *, expected_policy_version: str | None = Non padded = data.to_padded_tensor() uids = [str(u) for u in padded.pop("uid").tolist()] prompt_sha = hashlib.sha256(padded["prompts"].cpu().numpy().tobytes()).hexdigest() + # the same digest independent of arrival order (rows sorted): compares the prompts two + # runs drew for one round (T-C4: a resumed island vs a continuous one) + prompt_set_sha = hashlib.sha256(b"".join(sorted( + row.tobytes() for row in padded["prompts"].cpu().numpy()))).hexdigest() gen_steps = sorted({int(tag.get("global_steps", -1)) for tag in batch.tags}) rewards = padded["rm_scores"].float().sum(-1).tolist() lengths = padded["response_mask"].float().sum(-1).tolist() @@ -246,7 +250,8 @@ def generate(self, rollout_id: int, *, expected_policy_version: str | None = Non "rollout_id": rollout_id, "policy_token": token, "groups": len(metas), "samples": len(uids), "reward_mean": float(np.mean(rewards)) if rewards else None, "gen_seconds": gen_seconds, "expected_policy_version": expected_policy_version, - "prompts_sha256": prompt_sha, "sample_submit_global_steps": gen_steps, + "prompts_sha256": prompt_sha, "prompt_set_sha256": prompt_set_sha, + "sample_submit_global_steps": gen_steps, "trainer_global_step": t.global_steps, "rewards": rewards, "uids": uids, "data_cursor": self.data_cursor()}) return RolloutBatchHandle( @@ -352,7 +357,11 @@ def train_step(self, batch: RolloutBatchHandle) -> LocalStepReceipt: applied_lrs=None if lr is None else (lr,), ) if grad_norm is not None and self._cut_trainer is not None: - self._cut_trainer.check_first_step(1) # no-op unless a restore is pending + pending = self._cut_trainer.resume_reference is not None + self._cut_trainer.check_first_step(1) # no-op unless a restore is pending; raises on a problem + if pending: + isl.emit("rl_resume_check", rollout_id=batch.rollout_id, check="first_step", ok=True, + optimizer_steps=self._cut_trainer.optimizer_steps_per_round) ids = tuple(s for g in batch.groups for s in g.sample_ids) return LocalStepReceipt( algorithm="grpo", learner_id=isl.learner_id, learner_generation=0, @@ -522,16 +531,35 @@ def criteria_from_trainer_batch(isl: VerlIsland, kv_batch) -> None: def _row_field(kv_batch, name: str) -> tuple[list | None, str]: - """Per-row values of a non-tensor field (``HarborAgentLoop`` puts ``provenance`` - in ``extra_fields``); ``(None, reason)`` when the batch has none. That verl's - TransferQueue keeps ``extra_fields`` under their own names is not verified.""" + """Per-row values of one ``extra_fields`` entry (``HarborAgentLoop`` puts + ``provenance`` there); ``(None, reason)`` when the batch has none. + + verl's TransferQueue agent loop (``agent_loop_tq._agent_loop_postprocess``) + stores ``AgentLoopOutput.as_dict()``, so the entry sits inside the + ``extra_fields`` record of each row, not under its own name (the classic + agent loop flattens it into ``non_tensor_batch``; the fully async runner + reads that). The batch is read by the record name.""" import transfer_queue as tq try: - data = tq.kv_batch_get(keys=kv_batch.keys, partition_id=kv_batch.partition_id, select_fields=[name]) - return list(data[name]), "ok" + data = tq.kv_batch_get(keys=kv_batch.keys, partition_id=kv_batch.partition_id, + select_fields=["extra_fields"]) + rows = extra_field_values(list(data["extra_fields"]), name) except Exception as exc: # noqa: BLE001 - the rows then group under ``unknown`` return None, f"missing: {type(exc).__name__}: {str(exc)[:200]}" + if all(v is None for v in rows): + return None, f"missing: no {name!r} in extra_fields" + return rows, "ok" + + +def extra_field_values(records: list, name: str) -> list: + """``record[name]`` of each row's ``extra_fields`` (``None`` where absent). A + TransferQueue read can return each record wrapped in a ``NonTensorData``.""" + values = [] + for record in records: + record = getattr(record, "data", record) + values.append(record.get(name) if isinstance(record, dict) else None) + return values def source_event(train_lp, infer_lp, mask, provenances, provenance_read: str, *, rollout_id: int, diff --git a/yeto/rl/adapters/verl/resume_check.py b/yeto/rl/adapters/verl/resume_check.py index f4f008f0..69b52170 100644 --- a/yeto/rl/adapters/verl/resume_check.py +++ b/yeto/rl/adapters/verl/resume_check.py @@ -4,7 +4,7 @@ a verl FSDP actor worker onto them (:func:`train_handles`). Two worker shapes: * the engine worker of the pinned verl (``engine_workers.ActorRolloutRefWorker``): - ``worker.actor.engine`` holds ``module`` / ``optimizer`` / ``lr_scheduler`` + ``worker.actor.engine`` (a bare ``TrainingWorker``: ``worker.engine``) holds ``module`` / ``optimizer`` / ``lr_scheduler`` (``workers/engine/fsdp/transformer_impl.py``); * the legacy fsdp worker: ``actor_module_fsdp`` / ``actor_optimizer`` / ``actor_lr_scheduler``. @@ -55,7 +55,8 @@ def _engine_of(worker: Any) -> Any: inner = getattr(worker, "worker_dict", None) # ray fused worker: {role: worker} if inner: worker = next((w for w in inner.values() if hasattr(w, "actor")), worker) - return getattr(getattr(worker, "actor", None), "engine", None) or worker + actor = getattr(worker, "actor", None) or worker # ActorRolloutRefWorker.actor is a TrainingWorker + return getattr(actor, "engine", None) or worker def train_handles(worker: Any) -> TrainHandles: diff --git a/yeto/rl/adapters/verl/rollout_engine.py b/yeto/rl/adapters/verl/rollout_engine.py new file mode 100644 index 00000000..93db3586 --- /dev/null +++ b/yeto/rl/adapters/verl/rollout_engine.py @@ -0,0 +1,246 @@ +"""verl's own rollout servers behind the records proxy (S25 N10.9). + +In a verl training run the policy is served by verl's rollout replicas (vLLM +or SGLang, started and updated by verl). ``HarborAgentLoop`` needs a token +recorder in front of that policy: the sandboxed agent talks OpenAI to the +records proxy, the proxy renders the chat template and asks an engine for +tokens. :class:`VerlRolloutEngine` is that engine. It owns no server: it +forwards token-id prompts to verl's ``LLMServerClient`` (the agent loop's +``server_manager``) on the agent-loop worker's event loop, so verl keeps +routing, sticky sessions and weight updates. + +Wiring (all in the agent-loop worker process): + +* ``YETO_AGENTIC_RECORDER=yeto.rl.harness.gateway.recorder:recorder_from_env``; +* ``YETO_RECORDS_PROXY=yeto.rl.adapters.verl.rollout_engine:records_proxy_from_env`` + (this module) builds the proxy over a :class:`VerlRolloutEngine`; +* ``HarborAgentLoop.run`` binds the engine to its ``server_manager`` and the + running event loop before the first call (:func:`bind_rollout_client`). + +Environment of :func:`records_proxy_from_env`: + YETO_RECORDS_TEMPLATE chat template name (``chat_templates.TEMPLATES``) + YETO_RECORDS_TOKENIZER tokenizer directory (the checkpoint) + YETO_RECORDS_MODEL model name the agent sends (``AgentTask.model_name``) + YETO_RECORDS_SAMPLING JSON signed sampling (temperature, top_p, top_k, max_tokens) + YETO_RECORDS_MAX_MODEL_LEN context length of verl's rollout (``rollout.max_model_len``) + YETO_RECORDS_CARD GPU name (provenance; default ``unknown``) + +Policy version: verl's synchronous trainer updates the rollout weights only +between steps, never while a batch is generated, so a turn cannot straddle a +reload. Each rollout reply carries the weights' ``global_steps``; the engine +keeps them per trajectory (the records proxy's ``CURRENT_TRAJECTORY``) and +``HarborAgentLoop`` reports their min / max as the sample's +``min_global_steps`` / ``max_global_steps`` (verl's V1 trainer requires both). +:meth:`VerlRolloutEngine.served_version` returns None: the version is known +only after a reply, not before it. + +Finish reason: verl reports ``completed`` for both an end token and the token +cap. A reply that ends without a stop token and fills ``max_tokens`` is +reported as ``length``; everything else as ``stop``. +""" + +from __future__ import annotations + +import asyncio +import json +import os +import uuid +from collections.abc import Sequence +from pathlib import Path +from typing import Any + +from yeto.rl.engine.sampler_engine import ( + FINISH_LENGTH, + FINISH_STOP, + UNKNOWN, + Completion, + EngineIdentity, + LoadReport, + SamplerEngineError, + SamplingParams, +) + +from yeto.rl.harness.gateway.records_env import ( + CARD_ENV, + MAX_MODEL_LEN_ENV, + MODEL_ENV, + SAMPLING_ENV, + TEMPLATE_ENV, + TOKENIZER_ENV, +) + +CALL_TIMEOUT_S = 1800.0 + + +class VerlRolloutEngine: + """``TokenSamplerEngine`` over verl's ``LLMServerClient`` (module doc).""" + + def __init__(self, *, max_model_len: int, identity: EngineIdentity) -> None: + self._max_model_len = int(max_model_len) + self._identity = identity + self._client: Any = None + self._loop: asyncio.AbstractEventLoop | None = None + self._steps: dict[str, list[int]] = {} + self.last_global_steps: int | None = None + + # ------------------------------------------------------------ binding + def bind(self, client: Any, loop: asyncio.AbstractEventLoop) -> None: + """Use ``client`` (verl ``server_manager``) on ``loop``; the same pair for the worker's life.""" + if self._client is not None and (self._client is not client or self._loop is not loop): + raise SamplerEngineError("VerlRolloutEngine is already bound to another rollout client or loop") + self._client, self._loop = client, loop + + @property + def bound(self) -> bool: + return self._client is not None + + def pop_steps(self, trajectory: str) -> list[int]: + """The ``global_steps`` of every reply sampled for ``trajectory`` (and forget them).""" + return self._steps.pop(trajectory, []) + + # ------------------------------------------------------------ engine port + @property + def identity(self) -> EngineIdentity: + return self._identity + + @property + def max_model_len(self) -> int: + return self._max_model_len + + def served_version(self) -> int | None: + return None + + def generate_tokens(self, prompts: Sequence[Sequence[int]], sampling: SamplingParams, *, + stop_token_ids: Sequence[int] = ()) -> list[list[Completion]]: + """Called from a worker thread (the records backend uses ``asyncio.to_thread``).""" + if self._client is None or self._loop is None: + raise SamplerEngineError("VerlRolloutEngine is not bound to verl's rollout client") + from yeto.rl.harness.gateway.proxy import CURRENT_TRAJECTORY + + trajectory = CURRENT_TRAJECTORY.get() # copied into this thread by asyncio.to_thread + stops = {int(t) for t in stop_token_ids} + out: list[list[Completion]] = [] + for prompt in prompts: + group = [] + for _ in range(int(sampling.n)): + future = asyncio.run_coroutine_threadsafe( + self._generate(list(prompt), sampling, sorted(stops), trajectory), self._loop) + group.append(future.result(timeout=CALL_TIMEOUT_S)) + out.append(group) + return out + + async def _generate(self, prompt: list[int], sampling: SamplingParams, stops: list[int], + trajectory: str | None = None) -> Completion: + params: dict[str, Any] = { + "temperature": float(sampling.temperature), "top_p": float(sampling.top_p), + "top_k": int(sampling.top_k), "repetition_penalty": float(sampling.repetition_penalty), + "max_tokens": int(sampling.max_tokens), "logprobs": True, + } + if stops: + params["stop_token_ids"] = stops + output = await self._client.generate(request_id=uuid.uuid4().hex, prompt_ids=prompt, + sampling_params=params) + if getattr(output, "stop_reason", None) == "aborted": + raise SamplerEngineError("verl rollout aborted the request") + steps = (getattr(output, "extra_fields", None) or {}).get("global_steps") + if steps is not None: + self.last_global_steps = int(steps) + if trajectory is not None: + self._steps.setdefault(trajectory, []).append(int(steps)) + tokens = tuple(int(t) for t in output.token_ids) + logprobs = output.log_probs + if logprobs is None or len(logprobs) != len(tokens): + raise SamplerEngineError(f"verl rollout returned {0 if logprobs is None else len(logprobs)} " + f"log-probs for {len(tokens)} tokens") + return Completion(token_ids=tokens, logprobs=tuple(float(x) for x in logprobs), text="", + finish_reason=finish_reason(tokens, stops, int(sampling.max_tokens)), + prompt_token_ids=tuple(prompt)) + + # verl starts, updates and stops its own servers: the proxy never calls these. + def start(self, release_dir: Path | None, version: int) -> LoadReport: + raise SamplerEngineError("verl owns its rollout servers") + + def load(self, release_dir: Path, version: int) -> LoadReport: + raise SamplerEngineError("verl owns its rollout servers") + + def generate(self, prompts: Sequence[str], sampling: SamplingParams) -> list[list[Completion]]: + raise SamplerEngineError("text prompts are not used; the records proxy sends token ids") + + def score(self, prompt: str, continuation: str) -> list[float]: + raise SamplerEngineError("scoring is not used on the verl rollout") + + def loaded_expert_matrices(self) -> int | None: + return None + + def weight_gib_per_card(self) -> float | None: + return None + + def stop(self) -> None: + return None + + +def finish_reason(tokens: Sequence[int], stops: set[int] | Sequence[int], max_tokens: int) -> str: + """``length`` when the reply filled ``max_tokens`` without a stop token (module doc).""" + if tokens and int(tokens[-1]) in set(stops): + return FINISH_STOP + return FINISH_LENGTH if len(tokens) >= max_tokens else FINISH_STOP + + +def _engine_of(recorder: Any) -> Any: + engine = getattr(getattr(getattr(recorder, "proxy", None), "backend", None), "engine", None) + return engine if isinstance(engine, VerlRolloutEngine) else None + + +def weight_steps(recorder: Any, trajectory: str) -> tuple[int, int] | None: + """``(min, max)`` verl ``global_steps`` of the trajectory's replies; None without a verl engine. + + A trajectory with no reply (aborted before its first call) gets the + latest step the engine saw: in the synchronous trainer that is the batch's. + Before the engine saw any reply it gets ``(0, 0)``: verl's V1 trainer + needs an integer, and such a row is an ABORTED placeholder (no loss tokens). + """ + engine = _engine_of(recorder) + if engine is None: + return None + steps = engine.pop_steps(trajectory) or ([engine.last_global_steps] if engine.last_global_steps is not None else []) + return (min(steps), max(steps)) if steps else (0, 0) + + +def bind_rollout_client(recorder: Any, client: Any) -> None: + """Bind the recorder's :class:`VerlRolloutEngine` (if it has one) to verl's client on this loop.""" + proxy = getattr(recorder, "proxy", None) + engine = getattr(getattr(proxy, "backend", None), "engine", None) + if isinstance(engine, VerlRolloutEngine): + engine.bind(client, asyncio.get_running_loop()) + + +def records_proxy_from_env() -> Any: + """``YETO_RECORDS_PROXY`` callable: a RecordsProxy over a :class:`VerlRolloutEngine` (module doc).""" + from yeto.rl.agentic import HarnessConfigError + from yeto.rl.harness.chat_templates import ChatTemplate + from yeto.rl.harness.gateway.proxy import RecordsProxy + + try: + template_name = os.environ[TEMPLATE_ENV] + tokenizer = os.environ[TOKENIZER_ENV] + model = os.environ[MODEL_ENV] + signed = json.loads(os.environ[SAMPLING_ENV]) + max_model_len = int(os.environ[MAX_MODEL_LEN_ENV]) + except (KeyError, ValueError) as exc: + raise HarnessConfigError(f"records proxy over verl rollout: bad or missing setting {exc}") from exc + identity = EngineIdentity(engine="verl-rollout", engine_version=_verl_version(), device_family="cuda", + card_model=os.environ.get(CARD_ENV, UNKNOWN), precision="bf16", + quant_scheme="bf16", topology=UNKNOWN) + engine = VerlRolloutEngine(max_model_len=max_model_len, identity=identity) + template = ChatTemplate.load(template_name, tokenizer_path=tokenizer) + return RecordsProxy(engine, template, model=model, signed_sampling=signed, + sampler_id="verl-rollout-records") + + +def _verl_version() -> str: + try: + import verl + + return str(getattr(verl, "__version__", UNKNOWN)) + except ImportError: + return UNKNOWN diff --git a/yeto/rl/adapters/verl/round_cuts.py b/yeto/rl/adapters/verl/round_cuts.py new file mode 100644 index 00000000..735b282e --- /dev/null +++ b/yeto/rl/adapters/verl/round_cuts.py @@ -0,0 +1,168 @@ +"""verl island round cuts: when to cut, the pointer, stop triggers, resume (S25 N7.6b). + +The verl sync island runs the same cut controller as the Miles island +(:mod:`yeto.rl.engine.round_cut`), wired the same way: with ``--rl-resume-store`` +and no outer syncer (``--rl-single-island-no-sync``), + +* a :class:`~yeto.rl.engine.resume.ResumeController` restores the store's + ``LATEST`` snapshot into the state dir at start (pointer ``round-cut.json``, + batch ledger, journal) and syncs the state dir back after every cut; +* a :class:`~yeto.rl.engine.round_cut.RoundCutCheckpoint` on the + :class:`~yeto.rl.engine.bridges.LocalOnlySync` cuts the trainer + (:class:`~.cut.VerlCutTrainer`: one yeto file per FSDP rank) every + ``--rl-cut-every`` rounds, at the last round and at our own stop + (``--rl-stop-after-rounds``), and resumes from the pointed cut at start; +* triggers of the final cut: ``MANUAL`` (last round, our own stop), + ``SCHEDULED`` (the head's trainer scheduler stopped us: + ``trainer_remote`` commands, :func:`run_driver`); a preemption signal only + writes a marker (``install_preempt_handler``): the next launch resumes from + the newest complete cut. + +The resume assertions of :mod:`.cut` run on this path: the restored rank digests +and Adam steps at restore, the landed data cursor against the cut's data +watermark (driver ``_restore_data_cursor`` -> ``confirm_resumed_data_cursor``), +and the Adam/scheduler advance after the first resumed round. +""" + +from __future__ import annotations + +import os +from collections.abc import Mapping +from dataclasses import dataclass +from typing import Any + +RESUME_FLAGS = ("rl_cut_every", "rl_cut_keep", "rl_stop_after_rounds") +DEFAULT_STATE_DIR = "~/yeto-rl/resume-state" +# VerlRunConfig fields that may change between a cut and the launch resuming it: +# paths of this container, run length (extending a run is the normal resume), +# and settings that do not change the trained numbers. +FINGERPRINT_EXCLUDED = ("out_dir", "train_file", "val_file", "total_rounds", "experiment_name", + "gpu_memory_utilization", "reward_path") + + +def _flag(name: str) -> str: + return "--" + name.replace("_", "-") + + +def resume_config(args: Any) -> dict[str, Any] | None: + """``plan["resume"]`` from the island flags (None: no resume store). Raises on + a flag combination the cut controller cannot run.""" + store = getattr(args, "rl_resume_store", None) + given = [_flag(n) for n in RESUME_FLAGS if getattr(args, n, None) is not None] + if getattr(args, "rl_resume_allow_config_change", False): + given.append("--rl-resume-allow-config-change") + for name in RESUME_FLAGS: + value = getattr(args, name, None) + if value is not None and int(value) < 1: + raise ValueError(f"{_flag(name)} must be >= 1") + if not store: + if given: + raise ValueError(", ".join(given) + " need --rl-resume-store") + return None + from yeto.rl.engine import sample_roles + + if sample_roles.sample_role(args) == sample_roles.TRAINER: + # a buffer trainer (S25 C2) names its sample buffer's store with + # --rl-resume-store; round cuts are not wired on that path + if given: + raise ValueError(", ".join(given) + ": round cuts are not wired for a verl buffer " + "trainer (--rl-sample-role trainer)") + return None + if not getattr(args, "rl_single_island_no_sync", False): + raise ValueError("--rl-resume-store needs --rl-single-island-no-sync (a synced island " + "resumes from the syncer)") + if int(getattr(args, "rl_max_policy_age", 0) or 0): + raise ValueError("--rl-resume-store is not wired for the verl fully_async path " + "(--rl-max-policy-age > 0)") + return { + "store": str(store), + "state_dir": str(getattr(args, "rl_resume_state_dir", None) or DEFAULT_STATE_DIR), + "every": int(getattr(args, "rl_cut_every", None) or 1), + "keep": int(getattr(args, "rl_cut_keep", None) or 2), + "allow_config_change": bool(getattr(args, "rl_resume_allow_config_change", False)), + "stop_after": (None if getattr(args, "rl_stop_after_rounds", None) is None + else int(args.rl_stop_after_rounds)), + } + + +def run_fingerprint_of(plan: Mapping[str, Any], algorithm: Any) -> dict[str, Any]: + """What must be identical between a cut and the launch resuming it.""" + run = {k: v for k, v in dict(plan.get("run_config") or {}).items() if k not in FINGERPRINT_EXCLUDED} + sha = getattr(algorithm, "sha256", None) + return {**run, "model_revision": plan.get("model_revision"), "lora_targets": plan.get("lora_targets"), + "algorithm_spec_sha256": sha() if callable(sha) else None} + + +@dataclass +class ResumeWiring: + """The round-cut controller and batch ledger of one verl island run.""" + + controller: Any # ResumeController + ledger: Any # BatchLedger (in the state dir: restored and synced with the pointer) + config: dict[str, Any] + + +def resume_wiring_for(plan: Mapping[str, Any], *, runtime_fingerprint: str, + environ: Mapping[str, str] | None = None) -> ResumeWiring | None: + """Restore the store into the state dir (before the driver is built: the + ledger the driver gets is the restored one). None without ``plan["resume"]``.""" + config = plan.get("resume") + if not config: + return None + if plan.get("sync") != "none": + raise ValueError("verl round cuts need sync 'none' (--rl-single-island-no-sync)") + if not runtime_fingerprint: + raise ValueError("verl round cuts need a runtime fingerprint (a cut without it is refused)") + from yeto.rl.engine.ledger import BatchLedger + from yeto.rl.engine.resume import ResumeController, store_for + + store = store_for(config["store"], environ=environ) + controller = ResumeController(state_dir=os.path.expanduser(config["state_dir"]), store=store, + runtime_fingerprint=runtime_fingerprint, keep=int(config.get("keep", 2))) + return ResumeWiring(controller=controller, ledger=BatchLedger(controller.state_dir), config=dict(config)) + + +def wire_round_cuts(driver: Any, resume: ResumeWiring | None, *, plan: Mapping[str, Any], + algorithm: Any, install_signals: bool = True) -> Any: + """Give the driver's LocalOnlySync the cut controller (the Miles wiring, verl + units: one scheduler step per round, so the cut's ``global_batch_size`` is 1). + Returns the :class:`~yeto.rl.engine.round_cut.RoundCutCheckpoint` or None.""" + if resume is None: + return None + from yeto.rl.engine import round_cut + from yeto.rl.engine.resume import install_preempt_handler + + controller = resume.controller + source = round_cut.CutSource( + driver=lambda: driver, trainer=driver.trainer, rollout=driver.rollout, ledger=resume.ledger, + algorithm=algorithm, backend_fingerprint=controller.runtime_fingerprint, cut_root="", + global_batch_size=1) + checkpoint = round_cut.wire_round_cuts( + driver, controller=controller, source=source, every=int(resume.config.get("every", 1)), + fingerprint=run_fingerprint_of(plan, algorithm), + allow_config_change=bool(resume.config.get("allow_config_change", False))) + if checkpoint is None: + raise ValueError("verl round cuts need the driver's sync to be LocalOnlySync") + driver.sync.stop_after = resume.config.get("stop_after") + # the handler needs the main thread (verl: the task runner actor's); [] when it is not + signals = install_preempt_handler(controller, emit=driver.emit) if install_signals else [] + driver.emit("rl_round_cuts_wired", backend="verl", every=checkpoint.every, + keep=int(controller.keep), stop_after=driver.sync.stop_after, + store=str(controller.checkpoint_store), preempt_signals=list(signals), + restored_store=controller.restored is not None) + return checkpoint + + +def run_driver(driver: Any, resume: ResumeWiring | None, *, environ: Mapping[str, str] | None = None) -> Any: + """Run the island: as one commanded incarnation when the head's trainer + scheduler starts and stops it (``trainer_remote``: its stop cuts with the + SCHEDULED trigger), else to the end of the run. Returns the final policy state.""" + from yeto.rl.engine import trainer_remote + + if not trainer_remote.commanded(environ): + return driver.run() + if resume is None: + raise RuntimeError("a trainer started by the head's scheduler needs round cuts " + "(--rl-resume-store with --rl-single-island-no-sync)") + trainer_remote.run_commanded_island(driver, store=resume.controller.store, environ=environ) + return driver.published_state diff --git a/yeto/rl/adapters/verl/trainer.py b/yeto/rl/adapters/verl/trainer.py index 72e7859e..6be18c9e 100644 --- a/yeto/rl/adapters/verl/trainer.py +++ b/yeto/rl/adapters/verl/trainer.py @@ -206,6 +206,7 @@ def run_island(trainer, agent_loop_manager, plan: dict, cfg: dict) -> dict: from yeto.rl.engine.driver import EventTape, IslandDriver from yeto.rl.export import adapter_targets, derive_peft_lora_specs + from . import round_cuts from .identity import backend_identity from .ports_impl import (VerlIsland, VerlPlacement, VerlPolicyState, VerlPublisher, VerlRolloutPool, VerlTrainerGroup) @@ -256,16 +257,20 @@ def run_island(trainer, agent_loop_manager, plan: dict, cfg: dict) -> dict: spec = (AlgorithmSpec.from_dict(plan["algorithm_spec"]) if plan.get("algorithm_spec") else AlgorithmSpec()) publisher = VerlPublisher(island, strict=bool(plan.get("publish_strict", True))) + # N7.6b: round cuts (--rl-resume-store): the store is restored before the driver + # is built, so the driver's ledger is the restored one + resume = round_cuts.resume_wiring_for(plan, runtime_fingerprint=_fingerprint(cfg)) driver = IslandDriver( learner_id=learner_id, rollout=VerlRolloutPool(island), trainer=VerlTrainerGroup(island), policy_state=VerlPolicyState(island), publisher=_TestExitPublisher(publisher, plan.get("test_exit_after_version"), tape), placement=VerlPlacement(), capabilities=_capabilities(_fingerprint(cfg)), algorithm=spec, - sync=sync, events=tape) + sync=sync, events=tape, ledger=None if resume is None else resume.ledger) + round_cuts.wire_round_cuts(driver, resume, plan=plan, algorithm=spec) if plan.get("publish_selftest"): tape.append({"event": "rl_verl_publish_selftest", **publish_selftest(island, publisher, VerlPolicyState(island), specs)}) - state = driver.run() + state = round_cuts.run_driver(driver, resume) result = {"learner_id": learner_id, "rounds_completed": driver.rounds_completed, "final_policy_version": state.policy_version, "final_policy_tensor_hash": state.policy_tensor_hash(), @@ -275,11 +280,15 @@ def run_island(trainer, agent_loop_manager, plan: dict, cfg: dict) -> dict: return result -def run_fully_async_island(call, start_rollouter, plan: dict, cfg: dict) -> dict: +def run_fully_async_island(call, start_rollouter, plan: dict, cfg: dict, *, buffer_ports=None) -> dict: """agentic-rollout-utilization 6.4b: the yeto driver over the fully_async ports. ``call(method, *args)`` reaches the trainer actor; ``start_rollouter`` starts the - rollouter's generation loop (called once, after the first verified publication).""" + rollouter's generation loop (called once, after the first verified publication). + + ``buffer_ports(island)``: a buffer trainer (S25 C2, ``buffer_ports.install_buffer_ports`` + bound to the run's buffer and release store) returns ``(rollout, publisher, + placement)``; the trainer port and the policy state stay the fully_async ones.""" from yeto.rl.core import canonical_layout_hash, canonical_lora_config_hash from yeto.rl.engine.algorithm import AlgorithmSpec from yeto.rl.engine.driver import EventTape, IslandDriver @@ -317,18 +326,26 @@ def run_fully_async_island(call, start_rollouter, plan: dict, cfg: dict) -> dict tape.append({"event": "rl_verl_island_start", "plan": plan, "backend_identity": identity.to_dict(), "backend_identity_sha256": identity_sha, "layout_hash": layout_hash, "lora_config_hash": lora_hash, "n_specs": len(specs), "asserted": asserted, - "fully_async": True, "max_policy_age": limit, "init_seconds": time.time() - t0}) + "fully_async": True, "max_policy_age": limit, "init_seconds": time.time() - t0, + **({"sample_role": "trainer"} if buffer_ports is not None else {})}) sync = build_sync(plan, specs, lora_hash=lora_hash, layout_hash=layout_hash, identity_sha256=identity_sha, compat_group=identity.compat_group) spec = AlgorithmSpec.from_dict(plan["algorithm_spec"]) if plan.get("algorithm_spec") else AlgorithmSpec() profile = execution_profile_for(spec, limit, groups_per_round=int(plan["groups_per_round"]), samples_per_group=int(plan["samples_per_group"]), sync=plan["sync"]) - publisher = FullyAsyncPublisher(island, strict=bool(plan.get("publish_strict", True))) + if buffer_ports is None: + rollout, publisher, placement = (FullyAsyncRolloutPool(island), + FullyAsyncPublisher(island, strict=bool(plan.get("publish_strict", True))), + FullyAsyncPlacement()) + else: + rollout, publisher, placement = buffer_ports(island) driver = IslandDriver( - learner_id=learner_id, rollout=FullyAsyncRolloutPool(island), trainer=FullyAsyncTrainerGroup(island), - policy_state=VerlPolicyState(island), publisher=publisher, placement=FullyAsyncPlacement(), + learner_id=learner_id, rollout=rollout, trainer=FullyAsyncTrainerGroup(island), + policy_state=VerlPolicyState(island), publisher=publisher, placement=placement, capabilities=verl_capabilities(_fingerprint(cfg), fully_async_limit=limit), algorithm=spec, sync=sync, events=tape, profile=profile) + if buffer_ports is not None: + publisher.bind(driver) # synced buffer trainers release the driver's merged version state = driver.run() result = {"learner_id": learner_id, "rounds_completed": driver.rounds_completed, "final_policy_version": state.policy_version, diff --git a/yeto/rl/adapters/verl/verl_main.py b/yeto/rl/adapters/verl/verl_main.py index ff0989c1..106a0093 100644 --- a/yeto/rl/adapters/verl/verl_main.py +++ b/yeto/rl/adapters/verl/verl_main.py @@ -140,6 +140,28 @@ def run_ppo(config, task_runner_class=None): fully_async_hydra_entry(fa_main)() +def main_buffer_trainer(argv: list[str] | None = None) -> None: + """S25 C2: the fully_async hydra config and fix-ups as :func:`main_fully_async`, + with :class:`~.buffer_trainer_runner.YetoBufferTaskRunner`: a trainer island + over the sample buffer (the plan's ``sample_buffer`` section), no rollouter.""" + import verl.trainer.main_ppo as main_ppo + from verl.experimental.fully_async_policy import fully_async_main as fa_main + + plan = _plan_and_ray() + + def run_ppo(config, task_runner_class=None): + from yeto.rl.adapters.verl.buffer_trainer_runner import YetoBufferTaskRunner + + runner = ray.remote(num_cpus=1)(YetoBufferTaskRunner) + result = ray.get(runner.remote().run.remote(config, plan)) + print("[yeto-verl] island result " + json.dumps(result, sort_keys=True, default=str), flush=True) + + main_ppo.run_ppo = run_ppo + if argv is not None: + sys.argv = [sys.argv[0], *argv] + fully_async_hydra_entry(fa_main)() + + def fully_async_hydra_entry(fa_main): """``fa_main.main`` re-decorated with an absolute config directory. diff --git a/yeto/rl/agentic/outcome.py b/yeto/rl/agentic/outcome.py index 2b3c34ac..0f9d87ca 100644 --- a/yeto/rl/agentic/outcome.py +++ b/yeto/rl/agentic/outcome.py @@ -13,8 +13,11 @@ from __future__ import annotations +from collections.abc import Sequence +from dataclasses import replace + from .trust import verified_reward -from .types import Outcome, TrialVerdict +from .types import AgentTrajectory, Outcome, TrialVerdict # The policy ran out of its own budget: time, or sequence/context length. POLICY_LIMIT_EXCEPTIONS = frozenset({ @@ -30,6 +33,28 @@ # Reason given when the outer wall-clock cap fires (the harness never returned). WALL_CLOCK_REASON = "wall_clock_cap" +# The agent process exited non-zero. On its own this is ABORTED (could be the +# sandbox); with a recorded malformed model reply it is the policy's failure. +AGENT_EXIT_EXCEPTIONS = frozenset({"NonZeroAgentExitCodeError"}) +POLICY_OUTPUT_ERROR_REASON = "policy_output_error" + + +def policy_output_failure(trajectory: AgentTrajectory, policy_errors: Sequence[str]) -> AgentTrajectory: + """ABORTED by the agent's exit after the recorder refused a malformed model reply -> POLICY_LIMIT, reward 0. + + S25 N10.9 (main agent decisions): a tool call that is not valid JSON, and + a conversation that outgrew the context (the policy's own length budget), + are the policy's failures, trained as a 0 reward, not infrastructure aborts. + Fail closed: only the agent-exit exception, and only with a recorded + refusal, is re-classified; every other ABORTED stays ABORTED. + """ + verdict = trajectory.verdict + if (trajectory.outcome is not Outcome.ABORTED or not policy_errors or verdict is None + or verdict.exception_type not in AGENT_EXIT_EXCEPTIONS): + return trajectory + return replace(trajectory, outcome=Outcome.POLICY_LIMIT, reward=0.0, + reason=f"{POLICY_OUTPUT_ERROR_REASON}: {policy_errors[0][:200]}") + def classify(verdict: TrialVerdict, *, reward_range: tuple[float, float] = (0.0, 1.0)) -> tuple[Outcome, float, str]: """Return ``(outcome, reward, reason)`` for one verdict. diff --git a/yeto/rl/agentic/types.py b/yeto/rl/agentic/types.py index 2bd67789..e6d425f4 100644 --- a/yeto/rl/agentic/types.py +++ b/yeto/rl/agentic/types.py @@ -122,7 +122,12 @@ class TokenTurn: ``policy_version`` and ``provenance`` (the sampler source record, ``Provenance.to_record()``) are what the recorder knows about the call; - None when it does not record them. + None when it does not record them. ``finish_reason``: the engine's reason + (``stop`` / ``length``; ``length`` = the call hit its token cap). + ``policy_error``: the policy's own failure the recorder saw after this + call: a malformed reply (e.g. a tool call that is not valid JSON) or a + conversation that no longer fits the context (``context_full: ...``). + The tokens were sampled and are kept; the agent got an error. """ prompt_ids: tuple[int, ...] @@ -130,6 +135,8 @@ class TokenTurn: completion_logprobs: tuple[float, ...] policy_version: int | None = None provenance: Mapping[str, str] | None = None + finish_reason: str | None = None + policy_error: str | None = None def __post_init__(self) -> None: if len(self.completion_ids) != len(self.completion_logprobs): diff --git a/yeto/rl/engine/buffer_payload.py b/yeto/rl/engine/buffer_payload.py index 23bb1f8c..460831f3 100644 --- a/yeto/rl/engine/buffer_payload.py +++ b/yeto/rl/engine/buffer_payload.py @@ -142,14 +142,18 @@ def __init__(self, buffer: Any, sampler_id: str) -> None: self.sampler_id = str(sampler_id) def write(self, samples: Sequence[Mapping[str, Any]], *, group_key: str, - provenance: Mapping[str, str] | None = None) -> SampleIndexEntry: - """``provenance``: the sampler's source record for the index entry (S25 N2.4).""" + provenance: Mapping[str, str] | None = None, + sampler_round: Mapping[str, Any] | None = None) -> SampleIndexEntry: + """``provenance``: the sampler's source record for the index entry (S25 N2.4); + ``sampler_round``: the round record a replacement sampler resumes from + (``sampler_resume.SamplerRound.record()``, S25 C1).""" behavior_version, tokens = group_entry_fields(samples) sample_id = f"{self.sampler_id}-{group_key}" ref = self.buffer.write_payload(sample_id, encode_group(samples)) entry = SampleIndexEntry(sample_id=sample_id, behavior_version=behavior_version, payload_ref=ref, tokens=tokens, - provenance=None if provenance is None else dict(provenance)) + provenance=None if provenance is None else dict(provenance), + sampler_round=None if sampler_round is None else dict(sampler_round)) self.buffer.append(entry) return entry diff --git a/yeto/rl/engine/driver.py b/yeto/rl/engine/driver.py index c4ecee2c..d815a0a7 100644 --- a/yeto/rl/engine/driver.py +++ b/yeto/rl/engine/driver.py @@ -718,7 +718,10 @@ def publish(self, state: TrainableState, *, rollout_id: int) -> None: f"v{manifest.target_policy_version} {manifest.target_policy_hash}" ) members = frozenset(self.rollout.members()) - if not members or result.members != members: + # an engineless pool (trainer-only island) has no member to acknowledge; + # every other pool must have at least one and all must acknowledge + engineless = bool(getattr(self.rollout, "engineless", False)) + if (not members and not engineless) or result.members != members: missing = sorted(members - result.members) raise PublicationError( f"partial publication of policy v{rollout_id}; unacknowledged " @@ -897,6 +900,12 @@ def _generate(self, rollout_id: int) -> RolloutBatchHandle: f"(policy v{batch.policy_version}) for rollout {rollout_id}" ) batch = self._with_policy_labels(batch) + stale = getattr(batch, "buffer_stale_groups", None) + if stale is not None: + # The sample-buffer reader consumed these groups without training + # them (older than max_policy_age); one event per round, 0 included. + self.emit("rl_buffer_stale_skipped", rollout_id=rollout_id, + stale_groups=int(stale), max_policy_age=self._max_policy_age()) violations = getattr(batch, "policy_age_violation", None) if violations: # IR-3 (age 0): a sample generated by a policy other than the @@ -1694,7 +1703,8 @@ def rebuild_trainer(self, rebuild: Callable[[], Any], *, cut_policy_hash: str) - or manifest.target_policy_hash != cut_policy_hash): raise PublicationError("re-publication after rebuild acknowledged another policy") members = frozenset(self.rollout.members()) - if not members or result_pub.members != members: + engineless = bool(getattr(self.rollout, "engineless", False)) + if (not members and not engineless) or result_pub.members != members: raise PublicationError( f"partial re-publication after rebuild; missing {sorted(members - result_pub.members)}" ) diff --git a/yeto/rl/engine/engine_sampler.py b/yeto/rl/engine/engine_sampler.py index 7ef136be..1d803e2f 100644 --- a/yeto/rl/engine/engine_sampler.py +++ b/yeto/rl/engine/engine_sampler.py @@ -16,6 +16,13 @@ (:class:`~.sampler_engine.Provenance`) in the index entry; 6. write ``rl_sampler_engine_round`` (finish reasons, tokens, provenance key). +Replacement (S25 G1C, the engine-sampler side of C1's :mod:`.sampler_resume`): +each group's sample id carries the incarnation, the round and the prompt row, +so a new incarnation never writes an earlier one's payload name; each index +entry carries a :class:`~.sampler_resume.SamplerRound` record; a round already +in the buffer is never generated or written again (:class:`~.sampler_resume.RoundGuard`); +after the reclaim marker no group is written. + :func:`build_engine_sampler` wires the island: the start checks, the engine start on the base weights with the post-load checks (:mod:`.sampler_checks`; a failure there stops the island), the follower and the loop. The caller @@ -47,6 +54,7 @@ SamplingParams, ) from .sampler_loop import SamplerLoop +from .sampler_resume import RoundGuard, SamplerResume, SamplerRound ROUND_EVENT = "rl_sampler_engine_round" START_EVENT = "rl_sampler_start" @@ -91,7 +99,15 @@ class EngineRound: def __init__(self, engine: SamplerEngine, writer: GroupWriter, *, prompts: Callable[[int], Sequence[PromptItem]], reward: Reward, sampling: SamplingParams, provenance: Provenance, emit: Callable[..., None], - loop_guard: LoopGuard | None = None) -> None: + loop_guard: LoopGuard | None = None, incarnation: str | None = None, + resume: SamplerResume | None = None, + data_cursor: Callable[[int], Mapping[str, int]] | None = None, + reclaimed: Callable[[], bool] = lambda: False) -> None: + """``incarnation``: this process's id (``sampler_resume.new_incarnation``); + with it, sample ids are ``--r-`` and + every entry carries its round record. ``resume``: where the earlier + incarnations got to (its rounds are refused). ``data_cursor(rollout_id)``: + the prompt source's position after the round drew its prompts.""" self.engine = engine self.writer = writer self.prompts = prompts @@ -100,9 +116,24 @@ def __init__(self, engine: SamplerEngine, writer: GroupWriter, *, self.provenance = provenance self.emit = emit self.loop_guard = loop_guard + self.incarnation = None if incarnation is None else str(incarnation) + self.guard = None if resume is None else RoundGuard(resume) + self.data_cursor = data_cursor + self.reclaimed = reclaimed + self.dropped_groups = 0 # groups not written because the island is leaving + + def _round_record(self, rollout_id: int, last: bool) -> dict[str, Any] | None: + if self.incarnation is None: + return None + cursor = None if self.data_cursor is None else self.data_cursor(rollout_id) + return SamplerRound(self.writer.sampler_id, self.incarnation, rollout_id, + data_cursor=cursor, last=last).record() def __call__(self, rollout_id: int) -> None: - items = list(self.prompts(int(rollout_id))) + rollout_id = int(rollout_id) + if self.guard is not None: + self.guard.check(rollout_id) + items = list(self.prompts(rollout_id)) version = self.engine.served_version() if version is None: raise SamplerEngineError("the engine serves no version; it cannot generate") @@ -114,6 +145,7 @@ def __call__(self, rollout_id: int) -> None: record = self.provenance.to_record() finish: Counter[str] = Counter() tokens = 0 + scored = [] for position, (item, group) in enumerate(zip(items, groups)): if len(group) != self.sampling.n: raise SamplerEngineError(f"prompt {item.prompt_id}: {len(group)} completions, " @@ -128,11 +160,23 @@ def __call__(self, rollout_id: int) -> None: item, completion, reward=self.reward(item, completion), version=version, provenance=record, group_index=position, index=position * self.sampling.n + j)) - self.writer.write(samples, group_key=f"r{int(rollout_id)}-{item.prompt_id}", - provenance=record) - self.emit(ROUND_EVENT, rollout_id=int(rollout_id), policy_version=version, - groups=len(items), samples=sum(finish.values()), response_tokens=tokens, - finish=dict(sorted(finish.items())), source=self.provenance.key()) + scored.append((item, samples)) + prefix = "" if self.incarnation is None else f"{self.incarnation}-" + written = 0 + for position, (item, samples) in enumerate(scored): + if self.reclaimed(): # leaving: the replacement owns the shard from now on + self.dropped_groups += len(scored) - position + break + self.writer.write(samples, group_key=f"{prefix}r{rollout_id}-{item.prompt_id}", + provenance=record, + sampler_round=self._round_record(rollout_id, position == len(scored) - 1)) + written += 1 + if written and self.guard is not None: + self.guard.written(rollout_id) + self.emit(ROUND_EVENT, rollout_id=rollout_id, policy_version=version, + groups=written, samples=sum(finish.values()), response_tokens=tokens, + finish=dict(sorted(finish.items())), source=self.provenance.key(), + dropped_groups=len(scored) - written) @dataclass @@ -159,7 +203,12 @@ class EngineSamplerIsland: follower: PolicyFollower round: EngineRound - def run(self, start: int = 0) -> int: + resume: SamplerResume | None = None + + def run(self, start: int | None = None) -> int: + """Rounds from ``start`` (default: the resume point, else 0).""" + if start is None: + start = 0 if self.resume is None else self.resume.next_rollout_id return self.loop.run(start) @@ -168,9 +217,15 @@ def build_engine_sampler(config: EngineSamplerConfig, *, engine: SamplerEngine, prompts: Callable[[int], Sequence[PromptItem]], reward: Reward, checks: PostLoadChecks | None = None, should_stop: Callable[[], bool] = lambda: False, - sleep: Callable[[float], None] | None = None) -> EngineSamplerIsland: + sleep: Callable[[float], None] | None = None, + incarnation: str | None = None, resume: SamplerResume | None = None, + data_cursor: Callable[[int], Mapping[str, int]] | None = None, + reclaimed: Callable[[], bool] | None = None, + ) -> EngineSamplerIsland: """Check the sampling against the engine, start it on the base weights, - run the post-load checks, and return the wired island (module doc).""" + run the post-load checks, and return the wired island (module doc). + ``should_stop``: the loop starts no new round. ``reclaimed``: also stop + writing the current round's groups (default: ``should_stop``).""" sampling = config.sampling.check() check_sampling_fits(sampling, engine, max_prompt_tokens=config.max_prompt_tokens) provenance = Provenance.of(engine.identity, sampler_id=config.sampler_id, sampling=sampling) @@ -188,7 +243,9 @@ def build_engine_sampler(config: EngineSamplerConfig, *, engine: SamplerEngine, follower = PolicyFollower(releases, loader, version=int(config.base_version)) engine_round = EngineRound(engine, GroupWriter(buffer, config.sampler_id), prompts=prompts, reward=reward, sampling=sampling, provenance=provenance, emit=emit, - loop_guard=config.loop_guard) + loop_guard=config.loop_guard, incarnation=incarnation, + resume=resume, data_cursor=data_cursor, + reclaimed=should_stop if reclaimed is None else reclaimed) extra = {} if sleep is None else {"sleep": sleep} loop = SamplerLoop(follower=follower, buffer=buffer, generate=engine_round, num_rollouts=int(config.num_rollouts), @@ -196,7 +253,8 @@ def build_engine_sampler(config: EngineSamplerConfig, *, engine: SamplerEngine, poll_s=config.poll_s, should_stop=should_stop, min_round_s=config.min_round_s, loaded_fields=loader.loaded_fields, ready=lambda: loader.ready, **extra) - return EngineSamplerIsland(loop=loop, loader=loader, follower=follower, round=engine_round) + return EngineSamplerIsland(loop=loop, loader=loader, follower=follower, round=engine_round, + resume=resume) def main(argv: Sequence[str] | None = None) -> int: diff --git a/yeto/rl/engine/engine_sampler_entry.py b/yeto/rl/engine/engine_sampler_entry.py index 1eda00f9..f11c5307 100644 --- a/yeto/rl/engine/engine_sampler_entry.py +++ b/yeto/rl/engine/engine_sampler_entry.py @@ -31,7 +31,17 @@ ``reward_kind`` is ``"miles"``; * ``checks``: :class:`~.sampler_checks.PostLoadChecks` keyword overrides (optional; e.g. ``expected_expert_matrices``); -* ``event_tape``: the JSONL tape file of the island's events. +* ``event_tape``: the JSONL tape file of the island's events; each + incarnation writes its own file next to it (``events..jsonl``), + so a replacement container never overwrites an earlier one's tape. + +Replacement (S25 G1C): the process is one incarnation of the sampler +(``sampler_resume.new_incarnation``). It reads where the earlier incarnations +got to from the buffer (``sampler_resume.resume_point``), starts at the next +round, and writes ``rl_sampler_join`` / ``rl_sampler_leave`` plus one member +record per incarnation in the buffer store. With ``YETO_SPOT_RECLAIM_MARKER`` +set, SIGINT/SIGTERM write the reclaim marker; the island then starts no new +round, writes no more groups and leaves as PREEMPTED. Nothing here was run on a PPU or NPU node. """ @@ -107,6 +117,11 @@ def __init__(self, path: str | Path, *, input_key: str, label_key: str | None, g self.groups_per_round = int(groups_per_round) self.render = render + def cursor(self, rollout_id: int) -> dict[str, int]: + """The data position after round ``rollout_id`` drew its prompts (the + ``sampler_round`` record; rounds draw by rollout id, so it only documents).""" + return {"row_offset": (int(rollout_id) + 1) * self.groups_per_round} + def __call__(self, rollout_id: int) -> list[PromptItem]: start = int(rollout_id) * self.groups_per_round items = [] @@ -155,11 +170,16 @@ def build_prompts(spec: Mapping[str, Any]) -> JsonlPrompts: def miles_reward(score: Callable[..., Any]) -> Reward: """A Miles reward (``async score(args, sample)``) as an engine-sampler reward: - the sample carries ``response``, ``label`` and a ``metadata`` dict.""" + the sample carries ``response``, ``label``, a ``metadata`` dict, ``tokens`` + (prompt ids when the engine gave them, then the response ids) and + ``response_length``, as a Miles ``Sample`` does.""" def reward(item: PromptItem, completion: Completion) -> float: + response = list(getattr(completion, "token_ids", ())) + prompt_ids = list(getattr(completion, "prompt_token_ids", ())) sample = types.SimpleNamespace(prompt=item.prompt, response=completion.text, label=item.label, - metadata=dict(item.metadata)) + metadata=dict(item.metadata), tokens=prompt_ids + response, + response_length=len(response)) value = score(None, sample) return float(asyncio.run(value) if inspect.isawaitable(value) else value) @@ -191,6 +211,12 @@ def build_stores(spec: Mapping[str, Any], environ: Mapping[str, str] | None = No return releases, FileSampleBuffer(buffer_store, batch_size=1, index=index) # as open_buffer(sampler=True) +def incarnation_tape(path: str | Path, incarnation: str) -> Path: + """``events.jsonl`` -> ``events..jsonl`` (one tape per incarnation).""" + tape = Path(path).expanduser() + return tape.with_name(f"{tape.stem}.{incarnation}{tape.suffix or '.jsonl'}") + + def tape_writer(path: str | Path, sampler_id: str) -> Callable[..., None]: from yeto.rl.event_echo import append_record # noqa: PLC0415 @@ -219,16 +245,61 @@ def island_config(spec: Mapping[str, Any]) -> EngineSamplerConfig: def run(spec: Mapping[str, Any], *, environ: Mapping[str, str] | None = None, - resource_sampler: Callable[[Callable[..., None]], Any] | None = None) -> int: - """Build every part of the island from ``spec``, run it, stop the engine; - returns the rounds the loop ran. ``resource_sampler(emit)`` returns a - started telemetry thread (default: ``telemetry.ResourceSampler`` of the - island's device family; None from the factory = no telemetry).""" + resource_sampler: Callable[[Callable[..., None]], Any] | None = None, + incarnation: str | None = None, stop: Callable[[], bool] | None = None, + on_event: Callable[..., None] | None = None) -> int: + """Build every part of the island from ``spec``, run it from the resume + point, stop the engine; returns the rounds this incarnation ran. + ``resource_sampler(emit)`` returns a started telemetry thread (default: + ``telemetry.ResourceSampler`` of the island's device family; None from the + factory = no telemetry). ``stop()``: the caller's own end condition (no new + round starts; the island leaves FINISHED, not PREEMPTED). ``on_event(event, + **fields)`` sees every tape record after it is written.""" + import os # noqa: PLC0415 + + from yeto.rl.elastic_client import LeaveReason # noqa: PLC0415 + + from .sampler_resume import ( # noqa: PLC0415 + RECLAIM_MARKER_ENV, + ReclaimWatcher, + SamplerMembership, + install_marker_signals, + new_incarnation, + reclaim_marker, + reclaim_requested, + resume_point, + ) + + env = os.environ if environ is None else environ config = island_config(spec) - emit = tape_writer(_require(spec, "event_tape"), config.sampler_id) + incarnation = incarnation or new_incarnation() + tape = incarnation_tape(_require(spec, "event_tape"), incarnation) + write = tape_writer(tape, config.sampler_id) + + def emit(event: str, **fields: Any) -> None: + write(event, **fields) + if on_event is not None: + on_event(event, **fields) + releases, buffer = build_stores(spec, environ) prompts, reward = build_prompts(spec), build_reward(spec) checks = PostLoadChecks(**dict(spec.get("checks") or {})) + reclaim_marker(env).unlink(missing_ok=True) # a restarted machine keeps its /tmp + if env.get(RECLAIM_MARKER_ENV): + install_marker_signals(env) + buffer.refresh() + resume = resume_point(buffer.entries(), config.sampler_id) + from .resume import store_for # noqa: PLC0415 + + member_store = store_for(str(spec.get("buffer_root") or spec["store"]), environ=environ) + membership = SamplerMembership(member_store, config.sampler_id, incarnation, resume, emit) + reclaimed = lambda: reclaim_requested(env) # noqa: E731 + island = None + + def leave(reason: LeaveReason, **detail: Any) -> None: + version = None if island is None else island.follower.version + membership.leave(reason.name, policy_version=version, buffer_end=len(buffer), **detail) + engine = build_engine(spec) if resource_sampler is None: from .telemetry import ResourceSampler # noqa: PLC0415 @@ -237,11 +308,29 @@ def resource_sampler(sink: Callable[..., None]) -> Any: return ResourceSampler(lambda event, **fields: sink(event, **fields)).start() telemetry = resource_sampler(emit) + watcher = ReclaimWatcher(reclaimed, lambda: leave(LeaveReason.PREEMPTED)).start() try: island = build_engine_sampler(config, engine=engine, releases=releases, buffer=buffer, emit=emit, - prompts=prompts, reward=reward, checks=checks) - return island.run() + prompts=prompts, reward=reward, checks=checks, + should_stop=lambda: reclaimed() or (stop is not None and stop()), + reclaimed=reclaimed, incarnation=incarnation, resume=resume, + data_cursor=getattr(prompts, "cursor", None)) + latest = releases.latest() + membership.join(tape=str(tape), latest_release=None if latest is None else + {"seq": latest.seq, "policy_version": latest.version}) + rounds = 0 if buffer.data_ended else island.run() + leave(LeaveReason.PREEMPTED if reclaimed() else LeaveReason.FINISHED, rounds=rounds, + dropped_groups=island.round.dropped_groups) + return rounds + except BaseException as error: + try: + leave(LeaveReason.PREEMPTED if reclaimed() else LeaveReason.INFRA_FAILED, + error=f"{type(error).__name__}: {error}"[:300]) + except Exception: # noqa: BLE001 - the original error is the one to report + pass + raise finally: + watcher.stop() try: engine.stop() finally: @@ -259,6 +348,26 @@ def load_spec(path: str | Path) -> dict[str, Any]: return doc +RECLAIM_HOLD_S = 300.0 # after a reclaim LEAVE: wait this long for the cloud to stop the container + + +def hold_after_reclaim(environ: Mapping[str, str] | None = None, *, hold_s: float = RECLAIM_HOLD_S, + sleep: Callable[[float], None] = time.sleep) -> bool: + """After a reclaim the process stays up until the cloud stops the container: + a clean early exit could end the cloud's input as finished instead of letting + the cloud run the island again (as the Miles sampler island does). Returns + whether a reclaim was requested (False: returns at once).""" + import os # noqa: PLC0415 + + from .sampler_resume import reclaim_requested # noqa: PLC0415 + + if not reclaim_requested(os.environ if environ is None else environ): + return False + print(f"[engine-sampler] reclaimed: waiting up to {hold_s:g} s for the container to stop", flush=True) + sleep(hold_s) + return True + + def main(argv: Sequence[str] | None = None) -> int: parser = argparse.ArgumentParser(prog="python -m yeto.rl.engine.engine_sampler", description="run an engine sampler island from a spec file") @@ -271,4 +380,5 @@ def main(argv: Sequence[str] | None = None) -> int: print(f"[engine-sampler] spec rejected: {error}", file=sys.stderr) return 2 print(f"[engine-sampler] done: {rounds} round(s)", flush=True) + hold_after_reclaim() return 0 diff --git a/yeto/rl/engine/merged_release.py b/yeto/rl/engine/merged_release.py new file mode 100644 index 00000000..a6376b75 --- /dev/null +++ b/yeto/rl/engine/merged_release.py @@ -0,0 +1,256 @@ +"""Synced trainer islands feed sampler islands (S25 C3). + +N trainer islands (``--rl-sample-role trainer`` without +``--rl-single-island-no-sync``) train against one syncer (strict or elastic +OuterStep merge) and read their rounds from one durable sample buffer. One or +more sampler islands (separate no-sync launches over the same store) write the +buffer and follow the policy releases. This module holds the engine-neutral +rules; the Miles wiring is in ``adapters.miles.sampler_island`` and +``adapters.miles.buffer_rollout``. + +Who releases: exactly one island, the *publisher*, the trainer with learner id +:data:`PUBLISHER_LEARNER_ID` (0). Every trainer applies the same merged policy +at a merge boundary (the syncer sends one base to all), so one writer is enough, +and one writer keeps ``policy-releases/`` free of write races between hosts (a +Modal Volume or a bucket has no cross-host lock). The other trainers release +nothing. + +When: after every successful publish of the publisher, the same point a no-sync +trainer releases (``ReleasingPublisher``): the merged policy is applied, +verified by hash and loaded on the trainer's own engines before the sampler sees +it. + +Version rule: a release carries the *merged* version, the syncer's outer +version of the applied base (``IslandDriver.outer_version``; the strict syncer +has none, its committed version is the policy version). The publisher releases +only versions above the newest release (:class:`MergedReleaseGuard`), so the +release versions are strictly increasing. They are consecutive (v, v+1, ...) +when the publisher applies every merge; a publisher that misses merges (it was +late or re-joined) releases the next merged version it applies, and the gap is +reported (``rl_policy_release_gap``). The sampler stamps the release version on +its engines, so every buffered group carries the merged version as its behavior +version, and each trainer gates with ``max_policy_age`` against its own merged +version (:func:`gate_version`). + +Buffer: the trainers read disjoint *lanes* of the one buffer: trainer m of N +trains the positions ``p`` with ``p % N == m`` (:class:`BufferLane`), so no +group is trained by two islands. A trainer keeps its lane cursor, one record per +round, written when the round is selected (:func:`lane_cursor_log`), and a +restarted or re-joined trainer continues after it: a group selected by a round +that did not finish is lost, never trained twice (at most once). + +Stores: a shared directory (local disk, Modal Volume) or an object store read +through its object API (``CheckpointStore.objects``: S3, OSS). Each lane has one +writer, so every layout below has one writer per file or key: + +* lane cursor: ``sample-buffer/lanes/-of-.jsonl`` in a directory + (:class:`LaneCursorLog`, appended); ``sample-buffer/lanes/-of-/.json`` + in an object store (:class:`ObjectLaneCursorLog`, one new object per round: + an object store cannot append); +* buffer index: on a store whose samplers write their own shards + (``objects.buffer_shards``, OSS), the index has one writer, the merger + (``sample_buffer`` module doc). With synced trainers the publisher merges + (:func:`merges_shards`) and the other trainers only read the merged index; + so a lane whose publisher is stopped waits for the publisher's return; +* releases: written by the publisher only, through the store mount, in the + order of ``PolicyReleaseStore.release`` (weights, record, commit, then + ``LATEST.json``), so a sampler that reads ``LATEST.json`` by key sees a + complete release. Nothing here needs a cross-host lock or create-if-absent. + +Failures: + +* the publisher island stops: no new release appears. The samplers keep the + last release, up to their backlog cap; the other trainers keep training on + groups within ``max_policy_age`` and pass over older ones (consumed, not + trained). When the publisher re-joins (a replacement island with learner id + 0), it applies the syncer's newest base and releases that version: the gap + is reported, the order stays increasing; +* another trainer stops: its lane is not read until it re-joins; the samplers' + backlog cap follows the publisher's read position, so the buffer grows by at + most one lane share of that cap; +* a sampler stops: the trainers wait up to ``--rl-sample-wait-s`` for groups. +""" + +from __future__ import annotations + +import json +from collections.abc import Callable +from dataclasses import dataclass +from pathlib import Path +from typing import Any + +PUBLISHER_LEARNER_ID = 0 +LANES_DIR = "lanes" +LANE_ENV = "YETO_RL_SAMPLE_BUFFER_LANE" # "M/N" on a synced trainer island + + +@dataclass(frozen=True) +class BufferLane: + """Lane ``index`` of ``count``: the buffer positions ``p % count == index``.""" + + index: int + count: int + + def __post_init__(self) -> None: + if not 0 <= int(self.index) < int(self.count): + raise ValueError(f"buffer lane {self.index}/{self.count}: need 0 <= M < N") + + def owns(self, position: int) -> bool: + return int(position) % self.count == self.index + + def env(self) -> str: + return f"{self.index}/{self.count}" + + @classmethod + def parse(cls, text: str | None) -> "BufferLane | None": + if not text: + return None + index, sep, count = str(text).partition("/") + if not sep or not index.strip().isdigit() or not count.strip().isdigit(): + raise ValueError(f"{LANE_ENV}={text!r}: expected M/N") + lane = cls(int(index), int(count)) + return None if lane.count == 1 else lane + + def name(self) -> str: + return f"{self.index}-of-{self.count}" + + def file_name(self) -> str: + return f"{self.name()}.jsonl" + + +def merges_shards(lane: BufferLane | None) -> bool: + """Whether this trainer merges the sampler shards into the buffer index (on a + store that shards, ``sample_buffer.needs_shards``): a single trainer does, and + of synced trainers only the publisher (module doc: one index writer).""" + return lane is None or lane.index == PUBLISHER_LEARNER_ID + + +def is_publisher(learner_id: int, num_learners: int) -> bool: + """Whether learner ``learner_id`` of ``num_learners`` synced trainers releases.""" + if not 0 <= int(learner_id) < int(num_learners): + raise ValueError(f"learner id {learner_id} outside 0..{int(num_learners) - 1}") + return int(learner_id) == PUBLISHER_LEARNER_ID + + +def merged_version(driver: Any, state: Any) -> int: + """The version a synced trainer releases and gates with: the syncer's outer + version of the applied base when the sync session sets one (elastic), else + the policy version (strict and no-sync: the committed version).""" + outer = getattr(driver, "outer_version", None) + return int(state.policy_version) if outer is None else int(outer) + + +def gate_version(weight_version: int | None, merged: int | None) -> int | None: + """The trainer version the ``max_policy_age`` gate compares against: the + merged version once the island published one, else the engines' version.""" + return merged if merged is not None else weight_version + + +class MergedReleaseGuard: + """Keeps the publisher's release versions strictly increasing. + + ``check(version, latest)`` returns ``"release"``, or ``"stale"`` for a + version at or below the newest release (nothing is released: the samplers + already serve that version or a newer one). A version more than one above + the newest release is released and reported by ``on_gap(previous, version)``.""" + + def __init__(self, on_gap: Callable[[int, int], None]) -> None: + self.on_gap = on_gap + + def check(self, version: int, latest_version: int | None) -> str: + if latest_version is None: + return "release" + if int(version) <= int(latest_version): + return "stale" + if int(version) > int(latest_version) + 1: + self.on_gap(int(latest_version), int(version)) + return "release" + + +class LaneCursorLog: + """The lane cursor of one synced trainer, durable in a shared-directory store. + + One JSON line per selected round: ``{"rollout_id", "positions", "cursor"}`` + (``cursor``: the buffer read position after the round). :meth:`last_cursor` + is where a restarted trainer continues. ``commit`` makes the line durable + (the store's ``commit``; a Modal Volume commit).""" + + def __init__(self, root: str | Path, lane: BufferLane, + commit: Callable[[], None] = lambda: None) -> None: + self.path = Path(root) / LANES_DIR / lane.file_name() + self.commit = commit + + def last_cursor(self) -> int: + if not self.path.is_file(): + return 0 + cursor = 0 + for line in self.path.read_text(encoding="utf-8").splitlines(): + try: + cursor = int(json.loads(line)["cursor"]) + except (ValueError, KeyError, TypeError): + continue # a torn last line (crash during the write) + return cursor + + def append(self, rollout_id: int, positions: list[int], cursor: int) -> None: + self.path.parent.mkdir(parents=True, exist_ok=True) + line = _record(rollout_id, positions, cursor) + with self.path.open("a", encoding="utf-8") as fh: + fh.write(line + "\n") + self.commit() + + +def _record(rollout_id: int, positions: list[int], cursor: int) -> str: + return json.dumps({"rollout_id": int(rollout_id), "positions": [int(p) for p in positions], + "cursor": int(cursor)}, sort_keys=True) + + +class ObjectLaneCursorLog: + """The lane cursor of one synced trainer in an object store (module doc). + + Round ``k`` of this lane (counted over every start of the lane) is the object + ``.json`` with the record of :class:`LaneCursorLog`. One writer per + lane, so the next sequence number is the count of the listed objects; the + write is a plain put (no create-if-absent between hosts needed).""" + + def __init__(self, objects: Any, prefix: str, lane: BufferLane) -> None: + self.objects = objects + self.prefix = f"{prefix.rstrip('/')}/{LANES_DIR}/{lane.name()}/" + self._keys: list[str] | None = None + + def _listed(self) -> list[str]: + if self._keys is None: + self._keys = sorted(self.objects.list(self.prefix)) + return self._keys + + def last_cursor(self) -> int: + for key in reversed(self._listed()): + raw = self.objects.read(key) + try: + return int(json.loads(raw)["cursor"]) + except (TypeError, ValueError, KeyError): + continue # an unreadable object (should not happen: a put is atomic) + return 0 + + def append(self, rollout_id: int, positions: list[int], cursor: int) -> None: + keys = self._listed() + key = f"{self.prefix}{len(keys):010d}.json" + self.objects.write(key, _record(rollout_id, positions, cursor).encode()) + keys.append(key) + + +def lane_cursor_log(buffer: Any) -> LaneCursorLog | ObjectLaneCursorLog | None: + """The lane cursor log of a lane reader (``FileSampleBuffer`` with a lane; None + for a single reader): in the store's object API when it has one, else in the + store directory (committed after each line).""" + lane = getattr(buffer, "lane", None) + if lane is None: + return None + from .resume import CheckpointStore # noqa: PLC0415 + from .sample_buffer import BUFFER_DIR # noqa: PLC0415 - sample_buffer imports this module + + store = buffer.store # a CheckpointStore, or an ObjectStore at the store root + objects = store.objects if isinstance(store, CheckpointStore) else store + if objects is not None: + return ObjectLaneCursorLog(objects, BUFFER_DIR, lane) + return LaneCursorLog(f"{store.root}/{BUFFER_DIR}", lane, commit=store.commit) + diff --git a/yeto/rl/engine/policy_release.py b/yeto/rl/engine/policy_release.py index fbc081d9..45aeef17 100644 --- a/yeto/rl/engine/policy_release.py +++ b/yeto/rl/engine/policy_release.py @@ -40,6 +40,9 @@ GB). The trainer still writes through the mount (each file is new and written once). +Synced trainers (several islands merged by a syncer) release through one +publisher, with the merged version: :mod:`.merged_release`. + A trainer that publishes the same version and hash again (a restart) releases nothing. A trainer restored from an older cut publishes an older version with other weights (a new lineage): that is a new release with a higher ``seq``. @@ -101,6 +104,9 @@ class PolicyRelease: # A release made from another one (a W4A8 re-quantization of a BF16 release): # {"release_dir", "seq", "version", "policy_hash", "scheme", ...}; None: trained weights. derived_from: dict[str, Any] | None = None + # Who released it when the trainers sync (S25 C3, yeto.rl.engine.merged_release): + # {"kind": "merged", "learner_id", "num_learners", "local_version"}; None: one trainer. + origin: dict[str, Any] | None = None schema: str = SCHEMA @classmethod @@ -276,7 +282,8 @@ def release(self, version: int, policy_hash: str, write: Callable[[Path], None], buffer_selected: int | None = None, weights: str = FULL, engine_checksum: str | None = None, buffer_shards: dict[str, int] | None = None, - derived_from: dict[str, Any] | None = None) -> PolicyRelease | None: + derived_from: dict[str, Any] | None = None, + origin: dict[str, Any] | None = None) -> PolicyRelease | None: """Write the weights with ``write(directory)``, then the record (it lists the weight files and their sizes), then LATEST. @@ -303,7 +310,8 @@ def release(self, version: int, policy_hash: str, write: Callable[[Path], None], engine_checksum=engine_checksum or None, files=file_sizes(directory), buffer_shards=None if buffer_shards is None else {str(k): int(v) for k, v in sorted(buffer_shards.items())}, - derived_from=None if derived_from is None else dict(derived_from)) + derived_from=None if derived_from is None else dict(derived_from), + origin=None if origin is None else dict(origin)) _write_json_atomic(directory / RECORD, asdict(record)) self.store.commit() # the weights and the record are durable before LATEST names them _write_json_atomic(self.root / LATEST, asdict(record)) diff --git a/yeto/rl/engine/ports.py b/yeto/rl/engine/ports.py index 9399d2bc..c5a51460 100644 --- a/yeto/rl/engine/ports.py +++ b/yeto/rl/engine/ports.py @@ -103,6 +103,9 @@ class RolloutBatchHandle: # (rl-algo-seq-and-adv R2); None = not reported. Read by per-algorithm # gradient expectations (``gradient_expectation(batch, metrics)``). nonzero_advantages: int | None = None + # Groups a sample-buffer reader passed over this round as older than + # max_policy_age (consumed, never trained); None = not a buffer reader. + buffer_stale_groups: int | None = None # rl-infra-spec 1.7: summed non-generation (tool) time of the rollout's # samples; None = not reported. tool_wait_seconds: float | None = None diff --git a/yeto/rl/engine/round_cut.py b/yeto/rl/engine/round_cut.py new file mode 100644 index 00000000..cadc5047 --- /dev/null +++ b/yeto/rl/engine/round_cut.py @@ -0,0 +1,494 @@ +"""Round cuts of a run without an outer syncer: save at safe points, resume at start. + +Neutral core (moved from the Miles adapter for the verl island, S25 N7.6b): the +backend gives a ``CuttableTrainer`` (``save_cut`` / ``restore_cut`` / +``actual_layout``) and a rollout pool with ``data_cursor`` / ``seek_data_cursor``; +:class:`CutSource` builds the cut context from the driver, and +:class:`RoundCutCheckpoint` is the cut controller wired into +:class:`~.bridges.LocalOnlySync` (:func:`wire_round_cuts`). + +rl-multinode-island M4: policy-weight continuation across a machine replacement. + + +Without an outer syncer (``--rl-single-island-no-sync``, :class:`~..bridges.LocalOnlySync`) +a restarted learner used to start again at rollout 0 from the startup weights: the +checkpoint store (Q4, C5) carried the journal/ledger, but no LoRA/optimizer state. + +With ``--rl-elastic-checkpoint-store`` set this module keeps a 4.2 cut of the trainer +(every rank's LoRA + optimizer + scheduler + RNG shard, manifest committed last) at +every round-boundary safe point and resumes from the newest one at start: + +* the cut root is ``/round-cuts`` -- the store is visible on every node (sky + Storage MOUNT of the bucket), so the shard of a rank on another node (PP across + nodes) lands where the driver can check it (``shared_filesystem=True``); +* after the manifest is committed, ``round-cut.json`` (the pointer: cut id, next + rollout id, local step, policy hash, writing incarnation) is written into the state + dir and the state dir is synced to the store -- pointer, ledger and journal travel + in ONE store snapshot (``STORE-MANIFEST.json`` last), so a restored pointer never + names a cut newer than the restored ledger; +* at start (in-place restart or a fresh machine after the store restore) the cut the + pointer names is verified (``verify_cut``: algorithm, layout, backend, step, + policy version, every shard digest) and loaded on the fresh trainer + (``MilesTrainerGroup.restore_cut``); the driver then rebases the ledger and seeks + the data source to the pointer's rollout id. Any failure is fatal (no silent + restart at 0 once a pointer exists). +""" + +from __future__ import annotations + +import json +import os +from collections.abc import Callable, Mapping +from pathlib import Path +from typing import Any + +from . import recovery +from .recovery import CUT_KEY, FINAL_TRIGGERS, MANUAL, PERIODIC, ConsistentCut +from .resume import ROUND_CUTS # under the checkpoint store (not under the state dir) + +ISLAND = "local" # round cuts run without an outer syncer: one island, the local one +POINTER = "round-cut.json" # under the state dir (synced with the ledger/journal) + + +class RoundCutError(RuntimeError): + pass + + +def _algorithm_identity(algorithm: Any, *, ref_model: Mapping[str, Any] | None) -> Any: + from yeto.rl.engine.cut import AlgorithmIdentity + + runtime_attrs = None + to_attrs = getattr(algorithm, "to_legacy_runtime_attrs", None) + if callable(to_attrs): + runtime_attrs = dict(to_attrs()) or None + return AlgorithmIdentity.from_spec(algorithm, runtime_attrs=runtime_attrs, ref_model=ref_model) + + +class CutSource: + """What the driver contributes to a cut (E2 ``CutContext``) and what a + restore must match (``RestoreExpectation``); shared by the 4.4 same-shape + rebuilder and E3's ``MilesTrainerOps`` (4.7).""" + + # backend check of an in-flight export: problems refuse the cut (None: no check; + # the Miles adapter overrides it with its export check) + check_in_flight: Callable[[Mapping[str, Any]], list[str]] | None = None + + def __init__(self, *, driver: Any, trainer: Any, rollout: Any, ledger: Any, algorithm: Any, + backend_fingerprint: str, cut_root: str, global_batch_size: int, + ref_model: Mapping[str, Any] | None = None, + shared_filesystem: bool = True) -> None: + self.driver_ref = driver # the driver, or a zero-arg callable returning it + self.trainer = trainer + self.rollout = rollout + self.ledger = ledger + self.identity = _algorithm_identity(algorithm, ref_model=ref_model) + self.backend_fingerprint = backend_fingerprint + self.cut_root = cut_root + self.global_batch_size = int(global_batch_size) + self.shared_filesystem = shared_filesystem + + @property + def driver(self) -> Any: + ref = self.driver_ref + return ref() if callable(ref) and not hasattr(ref, "published_state") else ref + + def context(self, cut_id: str) -> Any: + from yeto.rl.engine.controller import RebuildRefused + from yeto.rl.engine.cut import CutProgress + from yeto.rl.engine.cut import CutContext + + driver = self.driver + state = driver.published_state + if state is None or driver.published_version is None: + raise RebuildRefused("no published policy to cut") + cursor = (self.rollout.data_cursor() + if callable(getattr(self.rollout, "data_cursor", None)) else None) + if cursor is None: + raise RebuildRefused("rollout data cursor unknown (needs the elastic metadata)") + data = dict(cursor) + summary = dict(self.ledger.cut_summary()) if self.ledger is not None else {} + if summary.get("engine_buffer_length") is not None: + data.setdefault("buffer_length", summary["engine_buffer_length"]) + in_flight: tuple = () + limit = int(getattr(getattr(driver, "profile", None), "max_policy_age", 0) or 0) + if limit > 0: + # agentic-rollout-utilization 3.3: the unfinished work the rollout + # process holds NOW (buffered groups, suspended agentic trajectories). + exporter = getattr(self.rollout, "export_in_flight", None) + if not callable(exporter): + raise RebuildRefused("the rollout pool cannot export its in-flight trajectories") + exported = exporter() + problems = list(self.check_in_flight(exported)) if self.check_in_flight else [] + if problems: + raise RebuildRefused("; ".join(problems)) + in_flight = tuple(exported["entries"]) + data["buffer_length"] = int(exported["buffer_groups"]) + summary["carried_over"] = len(in_flight) + summary["max_policy_age"] = limit + progress = CutProgress( + local_step=int(driver.local_step), + scheduler_samples=int(driver.local_step) * self.global_batch_size, + global_batch_size=self.global_batch_size, + next_rollout_id=int(driver.published_version), + policy_version=int(state.policy_version), + policy_hash=state.policy_tensor_hash(), + ) + return CutContext( + root=self.cut_root, cut_id=cut_id, backend_fingerprint=self.backend_fingerprint, + progress=progress, algorithm=self.identity, data=data, ledger=summary, + # The safe point is after boundary() returned non-stop and every + # member acknowledged the policy: the outer commit is settled. + outer={"settled": bool(driver.at_safe_point), "policy_version": state.policy_version, + "policy_token": driver.expected_token}, + shared_filesystem=self.shared_filesystem, + in_flight=in_flight, + ) + + def expectation(self, layout: Mapping[str, int], *, epoch: int) -> Any: + from yeto.rl.engine.cut import RestoreExpectation + + driver = self.driver + return RestoreExpectation( + algorithm=self.identity, layout=dict(layout), + backend_fingerprint=self.backend_fingerprint, local_step=int(driver.local_step), + policy_version=int(driver.published_state.policy_version), epoch=epoch, + ) + + +class RoundCutCheckpoint: + """``save`` at a safe point, ``resume`` at sync start (see the module doc). + + rl-resume-from-checkpoint (S17 C4): ``every`` (``--rl-cut-every``) cuts only every N + rounds (plus the last round and our own stop, ``final=True``); the pointer carries + ``run_fingerprint`` / ``lr_at_next_round`` / ``consumed_prompt_ids_digest``; a resume + checks the fingerprint (refused on a difference unless ``allow_config_change``), the + consumed-group digest and, on the first resumed round, the learning rate.""" + + def __init__(self, *, controller: Any, source: Any, trainer: Any, every: int = 1, + fingerprint: dict[str, Any] | None = None, allow_config_change: bool = False, + next_lr: Any = None) -> None: + if getattr(controller, "checkpoint_store", None) is None: + raise ValueError("round cuts need the checkpoint store") + if int(every) < 1: + raise ValueError("--rl-cut-every must be >= 1") + self.controller = controller + self.source = source # rebuild_wiring.CutSource (cut_root is set here) + self.trainer = trainer # MilesTrainerGroup: save_cut / restore_cut / actual_layout + self.root = str(Path(controller.checkpoint_store) / ROUND_CUTS) + source.cut_root = self.root + self.resumed: dict[str, Any] | None = None + self.every = int(every) + self.fingerprint = dict(fingerprint) if fingerprint is not None else None + self.allow_config_change = bool(allow_config_change) + self.next_lr = next_lr # callable(driver) -> list[float] | None + self.pending_lr_check: dict[str, Any] | None = None + # set by resume(): where the rollout data source continues (the cut's cursor) + self.resume_data_cursor: dict[str, int] | None = None + # rl-trainer-start-stop D4: the durable sample buffer's consumed watermark, + # written into the cut record (None: the run reads no sample buffer). + # buffer_watermark(next_rollout_id) commits and returns the position after the + # last trained round; buffer_resume(watermark) continues there after a restore + # (buffer_rollout.wire_sample_buffer sets both). + self.buffer_watermark: Any = None # callable(int) -> int | None + self.buffer_resume: Any = None # callable(int) -> int + + # ---------------------------------------------------------------- pointer + @property + def pointer_path(self) -> Path: + return Path(self.controller.state_dir) / POINTER + + def pointer(self) -> dict[str, Any] | None: + path = self.pointer_path + if not path.is_file(): + return None + return dict(json.loads(path.read_text(encoding="utf-8"))) + + def _write_pointer(self, info: dict[str, Any]) -> None: + tmp = self.pointer_path.with_suffix(".json.tmp") + tmp.write_text(json.dumps(info, sort_keys=True), encoding="utf-8") + with open(tmp, "rb") as fh: + os.fsync(fh.fileno()) + os.replace(tmp, self.pointer_path) + + def _epoch(self) -> int: + return int(self.controller.journal.epochs.config_epoch) + + def _consumed(self) -> tuple[int, str | None]: + ledger = getattr(self.source, "ledger", None) + groups = getattr(ledger, "_consumed_groups", None) + if groups is None: + return 0, None + from yeto.rl.engine.resume import consumed_digest + + return len(groups), consumed_digest(groups) + + # ---------------------------------------------------------------- save + def due(self, rollout_id: int, *, final: bool = False) -> bool: + from yeto.rl.engine.resume import should_cut + + return should_cut(int(rollout_id), every=self.every, final=final) + + def save(self, driver: Any, *, rollout_id: int, final: bool = False, + trigger: str | None = None) -> dict[str, Any] | None: + """Cut the trainer at the safe point before rollout ``rollout_id`` (= the published + version), point the state dir at it, sync the store. Returns the pointer, or None + when nothing was cut (not due, no publication yet, already cut at this rollout, + island RECOVERY_REQUIRED). Raises on a refused/failed cut (the caller logs it). + + ``trigger`` names a final cut's trigger (default ``MANUAL``; ``SCHEDULED`` for + a stop the trainer scheduler ordered); a periodic cut takes none.""" + import time + + from yeto.rl.engine.resume import directory_bytes + + if trigger is not None and (not final or trigger not in FINAL_TRIGGERS): + raise RoundCutError(f"cut trigger {trigger!r} needs final=True and one of {FINAL_TRIGGERS}") + ctl = self.controller + if ctl.recovery_required or driver.published_version is None: + return None + if not self.due(rollout_id, final=final): + return None + if int(driver.published_version) != int(rollout_id): + raise RoundCutError(f"safe point {rollout_id} != published version {driver.published_version}") + current = self.pointer() + if current is not None and int(current["next_rollout_id"]) == int(rollout_id): + return None # just resumed from (or already cut) this point + cut_id = f"r{int(rollout_id):06d}-{ctl.incarnation['id']}" + started = time.monotonic() + context = self.source.context(cut_id) + # colocated + offloaded trainer (1 GPU): its weights/optimizer are not on the + # device between publish and the next train step; wake it for the cut only. + offloaded = bool(getattr(driver, "_trainer_offloaded", False)) + if offloaded: + self.trainer.onload() + try: + self.trainer.save_cut(epoch=self._epoch(), context=context) + finally: + if offloaded: + self.trainer.offload() + save_s = time.monotonic() - started + progress = context.progress + consumed_n, consumed_sha = self._consumed() + next_lr = self.next_lr(driver) if callable(self.next_lr) else None + info = { + "cut_id": cut_id, "root": ROUND_CUTS, "epoch": self._epoch(), + "next_rollout_id": int(progress.next_rollout_id), + "local_step": int(progress.local_step), + "policy_version": int(progress.policy_version), + "policy_hash": progress.policy_hash, + "incarnation": ctl.incarnation["id"], + "incarnation_index": ctl.incarnation.get("index"), + "run_fingerprint": self.fingerprint, + "lr_at_next_round": None if next_lr is None else [float(x) for x in next_lr], + "consumed_groups": consumed_n, + "consumed_prompt_ids_digest": consumed_sha, + "outer_version": getattr(getattr(driver, "sync", None), "outer_version", None), + "final": bool(final), + } + cut = ConsistentCut( + checkpoint_id=cut_id, trigger=(trigger or MANUAL) if final else PERIODIC, + outer_step=info["outer_version"], island_steps={ISLAND: info["local_step"]}, + completed_groups=consumed_n, completed_groups_digest=consumed_sha, + buffer_watermark=(self.buffer_watermark(int(rollout_id)) if callable(self.buffer_watermark) + else None)) + saved = recovery.save(cut, store=ctl.store, write=lambda record: self._write_and_sync(info, record)) + info[CUT_KEY] = saved[CUT_KEY] + synced = saved["synced"] + sync = dict(getattr(ctl, "last_store_sync", None) or {}) + cut_bytes = directory_bytes(Path(self.root) / cut_id) + total_s = time.monotonic() - started + info["store_synced"] = bool(synced) + info["cut_bytes"] = cut_bytes + info["save_s"] = round(save_s, 3) + info["store_copy_s"] = sync.get("copy_s") + info["store_commit_s"] = sync.get("commit_s") + info["total_s"] = round(total_s, 3) + info["write_bytes_per_s"] = (cut_bytes / save_s) if save_s > 0 else None + info["pruned"] = sync.get("pruned") + info["latest_seq"] = sync.get("seq") + info["trainer_onloaded_for_cut"] = offloaded + if not synced: + info["store_error"] = sync.get("error") + return info + + def _write_and_sync(self, info: dict[str, Any], cut_record: dict[str, Any]) -> dict[str, Any]: + """The round cut's ``recovery.save`` writer: the pointer (with the cut record) + into the state dir, then the state dir into the store (which commits).""" + self._write_pointer({**info, CUT_KEY: cut_record}) + synced = self.controller.sync_checkpoint_store("round_cut", skip_if_recovery=True) + # the store sync commits itself (snapshot layout: in write_snapshot; flat layout: + # in the controller, outside its lock), also when it failed: recovery.save must not + return {"synced": bool(synced), "committed": True} + + # ---------------------------------------------------------------- resume + def resume(self, driver: Any) -> dict[str, Any] | None: + """Load the pointed cut into the freshly built trainer and set the driver's local + step. Returns the pointer (``next_rollout_id`` = where the run continues), or None + when this island never cut (fresh run). Raises on any inconsistency.""" + import time + + from yeto.rl.engine.cut import RestoreExpectation + from yeto.rl.engine.resume import check_fingerprint + + info = self.pointer() + if info is None: + if callable(self.buffer_resume): + self.buffer_resume(0) + return None + started = time.monotonic() + rid = int(info["next_rollout_id"]) + checks: dict[str, Any] = {} + diff = {} + if self.fingerprint is not None and info.get("run_fingerprint") is None: + checks["run_fingerprint"] = "absent (pointer written before rl-resume-from-checkpoint)" + elif self.fingerprint is not None: + diff = check_fingerprint(info.get("run_fingerprint"), self.fingerprint, + allow_change=self.allow_config_change) + checks["run_fingerprint"] = "same" if not diff else "changed(allowed)" + ledger = getattr(self.source, "ledger", None) + if ledger is not None and rid > 0: + last = ledger.batch(rid - 1) + if last is None or last.get("state") != "outer_recorded": + raise RoundCutError(f"round cut {info['cut_id']} continues at rollout {rid} but the " + f"ledger has no outer-recorded rollout {rid - 1}: {last}") + checks["ledger_last_round"] = "outer_recorded" + if info.get("consumed_prompt_ids_digest") is not None: + n, sha = self._consumed() + if sha != info["consumed_prompt_ids_digest"] or n != int(info.get("consumed_groups", n)): + raise RoundCutError(f"round cut {info['cut_id']}: the restored ledger's consumed groups " + f"({n}, {sha}) differ from the pointer's " + f"({info.get('consumed_groups')}, {info['consumed_prompt_ids_digest']})") + checks["consumed_prompt_ids_digest"] = "same" + backend_fingerprint = self.source.backend_fingerprint + if self.allow_config_change: + # The backend fingerprint hashes the whole Miles argv, so a longer run + # (--num-rollout, --lr-decay-iters) changes it; allowed changes are recorded. + from yeto.rl.engine.cut import load_manifest + + saved = load_manifest(self.root, info["cut_id"]).runtime.get("backend_fingerprint") + if saved != backend_fingerprint: + checks["backend_fingerprint"] = {"result": "changed(allowed)", "cut": saved, + "run": backend_fingerprint} + backend_fingerprint = saved + expect = RestoreExpectation( + algorithm=self.source.identity, layout=dict(self.trainer.actual_layout()), + backend_fingerprint=backend_fingerprint, + local_step=int(info["local_step"]), policy_version=int(info["policy_version"]), + epoch=self._epoch(), + ) + manifest = self.trainer.restore_cut(info["cut_id"], epoch=self._epoch(), root=self.root, + expect=expect, shared_filesystem=True) + checks["cut_shards_sha256"] = "ok" + if manifest.progress.policy_hash != info["policy_hash"] or manifest.progress.next_rollout_id != rid: + raise RoundCutError(f"round cut {info['cut_id']} manifest disagrees with its pointer") + restored = driver.export_local().policy_tensor_hash() + if restored != info["policy_hash"]: + raise RoundCutError(f"restored trainer holds {restored}, round cut holds {info['policy_hash']}") + checks["policy_hash"] = "same" + checks["policy_version"] = int(info["policy_version"]) + in_flight = self._restore_in_flight(driver, manifest, rid) + if in_flight is not None: + checks["in_flight"] = in_flight + self.resume_data_cursor = _cut_data_cursor(manifest) + if callable(self.buffer_resume): + watermark = (info.get(CUT_KEY) or {}).get("buffer_watermark") + if watermark is None: + raise RoundCutError(f"round cut {info['cut_id']} has no sample buffer watermark " + "but this trainer reads the sample buffer") + checks["buffer_watermark"] = int(self.buffer_resume(int(watermark))) + driver.local_step = int(info["local_step"]) + self.resumed = dict(info) + if info.get("lr_at_next_round") is not None: + self.pending_lr_check = {"rollout_id": rid, "expected": list(info["lr_at_next_round"])} + self.controller._record("round_cut", tx_id=None, action="restore", cut_id=info["cut_id"], + next_rollout_id=rid, local_step=int(info["local_step"]), + policy_hash=info["policy_hash"], cut_incarnation=info["incarnation"], + incarnation=self.controller.incarnation["id"]) + driver.emit("rl_round_cut_restored", rollout_id=rid, cut_id=info["cut_id"], + local_step=int(info["local_step"]), policy_hash=info["policy_hash"], + cut_incarnation=info["incarnation"]) + restored_store = getattr(self.controller, "restored", None) or {} + cut_bytes = sum(int(f.bytes) for f in manifest.files) + driver.emit( + "rl_resume", rollout_id=rid, cut_id=info["cut_id"], + incarnation=self.controller.incarnation.get("index"), + incarnation_id=self.controller.incarnation["id"], + cut_incarnation=info["incarnation"], cut_incarnation_index=info.get("incarnation_index"), + policy_version=int(info["policy_version"]), policy_hash=info["policy_hash"], + local_step=int(info["local_step"]), lr_at_next_round=info.get("lr_at_next_round"), + checks=checks, config_diff=diff or None, + discarded_rounds=list(getattr(self, "discarded_rounds", []) or []), + restore_s=round(time.monotonic() - started, 3), cut_bytes=cut_bytes, + state_restore=restored_store or None, + ) + return info + + def _restore_in_flight(self, driver: Any, manifest: Any, rollout_id: int) -> dict[str, Any] | None: + """agentic-rollout-utilization 3.3: the cut's unfinished trajectories, by the + CURRENT run's limit: within it and not already trained -> continued + (buffered groups back into the data buffer); else discarded and reported. + Suspended agentic trajectories never survive a restart (discarded). + None when the cut carries none.""" + entries = tuple(getattr(manifest, "in_flight", ()) or ()) + if not entries: + return None + limit = int(getattr(getattr(driver, "profile", None), "max_policy_age", 0) or 0) + ledger = getattr(self.source, "ledger", None) + completed = sorted(getattr(ledger, "_consumed_groups", None) or ()) + rollout = getattr(self.source, "rollout", None) + importer = getattr(rollout, "import_in_flight", None) + if not callable(importer): + raise RoundCutError("the cut carries in-flight trajectories but the rollout pool " + "cannot import them") + report = importer(entries, max_policy_age=limit, current_version=int(rollout_id), + completed_group_ids=completed) + driver.emit("rl_in_flight_restored", rollout_id=int(rollout_id), max_policy_age=limit, + entries=len(entries), **report) + return report + + def check_first_round(self, driver: Any, *, rollout_id: int, applied_lrs: Any) -> None: + """Design §4.5, on the boundary of the first resumed round (before it is recorded + or published): the trained learning rate continues bitwise. Raises, unless + ``allow_config_change`` (the event then records the difference).""" + from yeto.rl.engine.resume import check_lr_continues + + pending = self.pending_lr_check + if pending is None or int(rollout_id) != int(pending["rollout_id"]): + return + self.pending_lr_check = None + problem = check_lr_continues(pending["expected"], None if applied_lrs is None else list(applied_lrs)) + driver.emit("rl_resume_check", rollout_id=int(rollout_id), check="lr_at_next_round", + expected=pending["expected"], + applied=None if applied_lrs is None else [float(x) for x in applied_lrs], + ok=problem is None, allowed=problem is not None and self.allow_config_change) + if problem is not None and not self.allow_config_change: + raise RoundCutError(problem) + + +def _cut_data_cursor(manifest: Any) -> dict[str, int] | None: + """The data source position recorded at the cut (None when incomplete). + + Under fully_async the background worker keeps drawing prompts after the + round's trained groups were reported, and the cut puts the untrained ones + back into the data buffer. The position to resume from is therefore the + one read at the cut, not the one the batch ledger recorded for the round.""" + from .cut import DATA_CURSOR_FIELDS + + data = dict(getattr(manifest, "data", None) or {}) + # a backend with fewer position fields (verl: sample_offset) names them in the cut + fields = tuple(data.get("cursor_fields") or DATA_CURSOR_FIELDS) + if not all(isinstance(data.get(f), int) for f in fields): + return None + return {f: int(data[f]) for f in fields} + + +def wire_round_cuts(driver: Any, *, controller: Any, source: Any, **options: Any) -> RoundCutCheckpoint | None: + """Give a :class:`~..bridges.LocalOnlySync` driver round cuts when a checkpoint store is + configured (other syncs own the policy: their syncer is authoritative).""" + from .bridges import LocalOnlySync + + if getattr(controller, "checkpoint_store", None) is None or not isinstance(driver.sync, LocalOnlySync): + return None + checkpoint = RoundCutCheckpoint(controller=controller, source=source, trainer=driver.trainer, **options) + driver.sync.round_cuts = checkpoint + return checkpoint diff --git a/yeto/rl/engine/run_config.py b/yeto/rl/engine/run_config.py index b2c82c4e..616e588a 100644 --- a/yeto/rl/engine/run_config.py +++ b/yeto/rl/engine/run_config.py @@ -778,6 +778,13 @@ def resolve_rl_run_config( elif recipe == RECIPE_GLM5_NEXT: if parameter_mode != "full": raise ValueError("the glm5_next recipe is full-parameter only") + if getattr(args, "megatron_ref_load", None) is None: + # Miles loads an HF checkpoint only through Megatron-Bridge, which has no + # glm5_next bridge (Modal run-7): train from the converted torch_dist checkpoint + raise ValueError( + "the glm5_next recipe needs --megatron-ref-load (Miles " + "tools/convert_hf_to_torch_dist.py output; Megatron-Bridge cannot load glm5_next)" + ) # Megatron picks the backend (Miles run_glm5_3_flash.py passes none) recipe_name, attention_backend = RECIPE_GLM5_NEXT, "auto" elif fn_variant is not None: diff --git a/yeto/rl/engine/run_flags.py b/yeto/rl/engine/run_flags.py index 9b1314c3..b4de9148 100644 --- a/yeto/rl/engine/run_flags.py +++ b/yeto/rl/engine/run_flags.py @@ -588,6 +588,14 @@ def rl_forward_pairs(*dests: str | tuple[str, str]) -> tuple[tuple[str, str], .. shared_flag("--rl-sample-wait-s", type=float, default=None, metavar="S", help="--rl-sample-role trainer: max seconds to wait for one full round from " "the buffer (default 3600)"), + shared_flag("--rl-sampler-node", default=None, metavar="NAME", + help="--rl-sample-role sampler: this island's node name (set by the launcher " + "per island of --rl-island-sample-role; the sampler id that a replacement " + "incarnation of the island resumes under)"), + shared_flag("--rl-buffer-shards", default=None, metavar="A,B", + type=lambda text: [s for s in text.split(",") if s], + help="--rl-island-sample-role with several samplers: the sampler shards of the " + "buffer (set by the launcher); each sampler writes its own, the trainer merges them"), ) _SCHEDULING = ( diff --git a/yeto/rl/engine/sample_buffer.py b/yeto/rl/engine/sample_buffer.py index e4feb402..aa271588 100644 --- a/yeto/rl/engine/sample_buffer.py +++ b/yeto/rl/engine/sample_buffer.py @@ -27,7 +27,8 @@ Arrival notification is by polling the index (``poll`` / ``start_polling``). -Limits of this first implementation: one consumer; with ``index.jsonl``, +Limits of this first implementation: one consumer (or one per lane: synced +trainer islands read disjoint lanes, :mod:`.merged_release`); with ``index.jsonl``, concurrent appenders must share one host (``flock``); the version gate on selection is the consumer's job (P0-S ``max_policy_age``). @@ -67,6 +68,7 @@ from pathlib import Path from typing import Any +from .merged_release import BufferLane, merges_shards from .object_store import ObjectStore from .resume import CheckpointStore from .trainer_lifecycle import BacklogStats, SampleIndexEntry, Selected @@ -96,12 +98,15 @@ def payload_name(sample_id: str) -> str: class JsonlIndex: """The index as ``index.jsonl`` in a shared directory (see the module doc).""" - def __init__(self, store: CheckpointStore, *, index: str = INDEX, create_index: bool = True) -> None: + def __init__(self, store: CheckpointStore, *, index: str = INDEX, create_index: bool = True, + reload: bool = True) -> None: """``index``: the index file under ``sample-buffer/`` (a shard writer passes :func:`shard_index`). ``create_index=False``: a reader that must not write the index, not even an empty one (on a store mount whose commit would - replace the merger's index).""" + replace the merger's index). ``reload=False``: :meth:`read_new` does not + reload the store (a merger's shard reader: the merger reloaded it once).""" self.store = store + self._reload = bool(reload) self.root = Path(store.root) / BUFFER_DIR (self.root / PAYLOADS).mkdir(parents=True, exist_ok=True) self.path = self.root / index @@ -128,22 +133,30 @@ def has_payload(self, ref: str) -> bool: return (self.root / ref).is_file() def read_new(self) -> list[dict[str, Any]]: - self.store.reload() - if not self.path.exists(): # a reader before the first merge + if self._reload: + self.store.reload() + if not self.path.exists(): # a reader before the first merge (or a shard not created yet) return [] with open(self.path, "rb") as fh: return list(self._lines(fh)) def append(self, absorb: Callable[[list[dict[str, Any]]], None], make: Callable[[], dict[str, Any]]) -> None: + self.append_many(absorb, lambda: [make()]) + + def append_many(self, absorb: Callable[[list[dict[str, Any]]], None], + make: Callable[[], list[dict[str, Any]]]) -> None: + """Append ``make()``'s records with one write, one fsync and one store + commit (a merger appends a whole refresh at once: one Modal Volume commit + per record would cost seconds each).""" with open(self.path, "r+b") as fh: fcntl.flock(fh.fileno(), fcntl.LOCK_EX) try: absorb(list(self._lines(fh))) - record = make() + records = make() fh.truncate(self._offset) # drop a torn last line fh.seek(self._offset) - fh.write((json.dumps(record, sort_keys=True) + "\n").encode()) + fh.write(b"".join((json.dumps(r, sort_keys=True) + "\n").encode() for r in records)) fh.flush() os.fsync(fh.fileno()) absorb(list(self._lines(fh))) @@ -151,6 +164,27 @@ def append(self, absorb: Callable[[list[dict[str, Any]]], None], fcntl.flock(fh.fileno(), fcntl.LOCK_UN) self.store.commit() + # ------------------------------------------------------------ shards (merger side) + def shard_names(self) -> list[str]: + """The shard index files under ``sample-buffer/shards/``, sorted.""" + return sorted(p.name[:-len(".jsonl")] for p in (self.root / SHARDS).glob("*.jsonl")) + + def shard_reader(self, shard: str, start: int) -> "JsonlIndex": + """A reader of one shard's index file that skips its first ``start`` + records (merged before).""" + reader = JsonlIndex(self.store, index=shard_index(shard), create_index=False, reload=False) + if start: + skipped = 0 + if reader.path.exists(): + with open(reader.path, "rb") as fh: + for _ in reader._lines(fh): + skipped += 1 + if skipped == start: + break + if skipped < start: + raise SampleBufferError(f"shard {shard}: {skipped} records, but {start} are merged") + return reader + def _lines(self, fh) -> Iterator[dict[str, Any]]: fh.seek(self._offset) data = fh.read() @@ -223,6 +257,13 @@ def append(self, absorb: Callable[[list[dict[str, Any]]], None], return raise SampleBufferError(f"{self.prefix}: lost {_MAX_CREATE_RACES} appends in a row to other writers") + # ------------------------------------------------------------ shards (merger side) + def shard_names(self) -> list[str]: + return shard_names(self.objects) + + def shard_reader(self, shard: str, start: int) -> "ObjectIndex": + return ObjectIndex(self.objects, prefix=shard_prefix(shard), start=start) + def _objects_of(store: CheckpointStore | ObjectStore) -> ObjectStore | None: return store.objects if isinstance(store, CheckpointStore) else store @@ -249,7 +290,8 @@ def _entry(record: dict[str, Any]) -> SampleIndexEntry: return SampleIndexEntry( sample_id=record["sample_id"], behavior_version=int(record["behavior_version"]), payload_ref=record["payload_ref"], tokens=int(record["tokens"]), - appended_at=float(record["appended_at"]), provenance=record.get("provenance")) + appended_at=float(record["appended_at"]), provenance=record.get("provenance"), + sampler_round=record.get("sampler_round")) class UnmergedShards: @@ -312,6 +354,16 @@ def index_for(store: CheckpointStore | ObjectStore, *, index: str = INDEX, return ObjectIndex(objects, prefix=shard_prefix(shard)) +def open_trainer_buffer(store: CheckpointStore | ObjectStore, *, batch_size: int, + lane: BufferLane | None = None, create_index: bool = True, + **kwargs: Any) -> "FileSampleBuffer": + """The buffer of a trainer (the reader): it merges the sampler shards of a + store that shards (:func:`needs_shards`) unless it is a synced trainer other + than the publisher (``merged_release.merges_shards``); ``lane``: its lane.""" + return FileSampleBuffer(store, batch_size=batch_size, create_index=create_index, lane=lane, + merge_shards=needs_shards(store) and merges_shards(lane), **kwargs) + + class FileSampleBuffer: """:class:`~yeto.rl.engine.trainer_lifecycle.DurableSampleBuffer` over a store (a :class:`CheckpointStore`, or an :class:`ObjectStore` at the store root). @@ -321,21 +373,40 @@ class FileSampleBuffer: def __init__(self, store: CheckpointStore | ObjectStore, *, batch_size: int, arrival_window_s: float = 60.0, clock: Callable[[], float] = time.time, - index: str = INDEX, create_index: bool = True, merge_shards: bool = False) -> None: + index: str = INDEX, create_index: bool = True, merge_shards: bool = False, + expected_shards: Iterable[str] | None = None, + lane: BufferLane | None = None) -> None: """``index``/``create_index``: see :class:`JsonlIndex`. ``merge_shards``: - this buffer (the consumer) merges the sampler shards of an object store - into the index on every refresh (see the module doc); it must be the - only one that does.""" + this buffer (the consumer) merges the sampler shards into the index on + every refresh (see the module doc); it must be the only one that does. + On a shared directory the shards are ``shards/.jsonl`` files (one + sampler island writes each; S25 C1: several sampler islands on one + Modal Volume, where ``flock`` does not hold across containers). + ``expected_shards``: the shards that must all end before the merged index + ends (None: the shards seen so far), so an early-ending sampler cannot + end the data before a slower one has written its first record. + ``lane``: this consumer is one of several synced trainers and selects + only the positions of its lane (:mod:`.merged_release`); the read cursor + and the watermark stay positions of the whole index. Of the lane readers + only the publisher's lane merges shards (``merged_release.merges_shards``).""" + if lane is not None and merge_shards and not merges_shards(lane): + raise SampleBufferError(f"buffer lane {lane.env()}: only the publisher's lane merges the " + "sampler shards (yeto.rl.engine.merged_release)") + self.lane = lane if batch_size < 1: raise ValueError("batch_size must be positive") self.store = store self.index = index_for(store, index=index, create_index=create_index) self.shard = _shard_of_index(index) if index != INDEX else None # the shard this buffer writes - if merge_shards and (self.shard is not None or not isinstance(self.index, ObjectIndex)): - raise SampleBufferError("only the index of an object store merges shards") + if merge_shards and self.shard is not None: + raise SampleBufferError("a shard writer does not merge shards (only the consumer's index does)") + if expected_shards is not None and not merge_shards: + raise SampleBufferError("expected_shards needs merge_shards") self._merge = merge_shards + self._expected_shards = (None if expected_shards is None + else frozenset(_shard_name(s) for s in expected_shards)) self._shard_of: list[str | None] = [] # the shard of each position (merged records) - self._shard_readers: dict[str, ObjectIndex] = {} + self._shard_readers: dict[str, JsonlIndex | ObjectIndex] = {} self._shards_ended: set[str] = set() self.batch_size = int(batch_size) self.arrival_window_s = float(arrival_window_s) @@ -355,8 +426,14 @@ def __init__(self, store: CheckpointStore | ObjectStore, *, batch_size: int, # ---------------------------------------------------------------- writing def write_payload(self, sample_id: str, data: bytes) -> str: - """Store one payload; returns the ``payload_ref`` for its index entry.""" - return self.index.write_payload(payload_name(sample_id), data) + """Store one payload; returns the ``payload_ref`` for its index entry. + A payload is written once: a name already in the store is refused (a + replacement sampler must never overwrite a group an index entry names).""" + name = payload_name(sample_id) + if self.index.has_payload(f"{PAYLOADS}/{name}"): + raise SampleBufferError(f"payload {name} is already in the sample buffer; " + "a payload is never overwritten") + return self.index.write_payload(name, data) def read_payload(self, payload_ref: str) -> bytes: return self.index.read_payload(payload_ref) @@ -368,6 +445,8 @@ def append(self, entry: SampleIndexEntry) -> None: entry = dataclasses.replace(entry, appended_at=self.clock()) # ``provenance`` is written only when set, so a line without it reads as before. source = {} if entry.provenance is None else {"provenance": dict(entry.provenance)} + if entry.sampler_round is not None: + source["sampler_round"] = dict(entry.sampler_round) self._write_record(lambda pos: {"schema": SCHEMA, "pos": pos, "sample_id": entry.sample_id, "behavior_version": int(entry.behavior_version), "payload_ref": entry.payload_ref, "tokens": int(entry.tokens), @@ -386,6 +465,22 @@ def record() -> dict: with self._lock: self.index.append(self._absorb, record) + def _write_records(self, records: list[dict[str, Any]]) -> None: + """Append ``records`` (without ``pos``) at the next positions: one write on + a shared directory, one object per record on an object store.""" + def make() -> list[dict]: + if self._ended: + raise SampleBufferError("append after the data end") + start = len(self._entries) + return [{**record, "pos": start + i} for i, record in enumerate(records)] + + with self._lock: + if isinstance(self.index, JsonlIndex): + self.index.append_many(self._absorb, make) + return + for record in records: # an object store: one create-if-absent per record + self._write_record(lambda pos, record=record: {**record, "pos": pos}) + def _absorb(self, records: list[dict[str, Any]]) -> None: for record in records: if record.get("schema") != SCHEMA or int(record["pos"]) != len(self._entries): @@ -407,14 +502,12 @@ def refresh(self) -> int: return len(self._entries) def _merge_shards(self) -> None: - objects = self.index.objects new: list[list[dict[str, Any]]] = [] # per shard, in shard order - for shard in shard_names(objects): + for shard in self.index.shard_names(): reader = self._shard_readers.get(shard) if reader is None: # first seen here: skip what an earlier merger merged done = sum(1 for s in self._shard_of if s == shard) - reader = self._shard_readers[shard] = ObjectIndex(objects, prefix=shard_prefix(shard), - start=done) + reader = self._shard_readers[shard] = self.index.shard_reader(shard, done) records = [] for record in reader.read_new(): if record.get("end"): @@ -430,11 +523,12 @@ def _merge_shards(self) -> None: # records of a shard must stay a prefix of it, because a restarted merger # resumes each shard after its merged count (a clock step back would # otherwise merge a later record first, then skip the earlier one). - merged = heapq.merge(*new, key=lambda r: (float(r.get("appended_at", 0.0)), r["shard"])) - for record in merged: - self._write_record(lambda pos, record=record: {**record, "pos": pos}) - if self._shard_readers and self._shards_ended >= set(self._shard_readers): - self._write_record(lambda pos: {"schema": SCHEMA, "pos": pos, "end": True}) + merged = list(heapq.merge(*new, key=lambda r: (float(r.get("appended_at", 0.0)), r["shard"]))) + expected = set(self._shard_readers) if self._expected_shards is None else self._expected_shards + if expected and self._shards_ended >= expected: + merged.append({"schema": SCHEMA, "end": True}) + if merged: + self._write_records(merged) def shard_positions(self, cursor: int) -> dict[str, int]: """{shard: records of that shard below merged position ``cursor``}: the @@ -449,6 +543,11 @@ def __len__(self) -> int: with self._lock: return len(self._entries) + def entries(self) -> list[SampleIndexEntry]: + """The entries read so far, in buffer order (a copy).""" + with self._lock: + return list(self._entries) + @property def data_ended(self) -> bool: with self._lock: @@ -481,9 +580,13 @@ def select(self, consumer: str, limit: int) -> list[Selected]: elif consumer != self._consumer: raise SampleBufferError(f"buffer is read by {self._consumer!r}, not {consumer!r}") self.refresh() - stop = min(len(self._entries), self._read + max(0, int(limit))) - out = [Selected(pos, self._entries[pos]) for pos in range(self._read, stop)] - self._read = stop + out: list[Selected] = [] + pos = self._read + while len(out) < max(0, int(limit)) and pos < len(self._entries): + if self.lane is None or self.lane.owns(pos): + out.append(Selected(pos, self._entries[pos])) + pos += 1 + self._read = pos return out def commit_consumed(self, cursor: int) -> None: diff --git a/yeto/rl/engine/sample_roles.py b/yeto/rl/engine/sample_roles.py index 10e68c35..de65b83f 100644 --- a/yeto/rl/engine/sample_roles.py +++ b/yeto/rl/engine/sample_roles.py @@ -18,9 +18,19 @@ The sampler and the trainer mount the same checkpoint store (the buffer and the releases live there). They are either separate single-island launches (``--rl-sample-role``), or one launch of two islands, a *store-coupled pair* -(``--rl-island-sample-role 0:sampler --rl-island-sample-role 1:trainer``, L5.4): -there is no syncer, and each island runs as its own one-learner, no-sync island +(``--rl-island-sample-role 0:sampler --rl-island-sample-role 1:trainer``, L5.4), +or one launch of one trainer and several sampler islands (S25 C1: each sampler +writes its own shard of the buffer, :func:`buffer_shards`): there is no syncer, and each island runs as its own one-learner, no-sync island (learner id 0 of 1) whose role comes from its island args (:func:`island_args`). +Synced trainers (S25 C3, :mod:`.merged_release`): a launch of N >= 2 islands +with ``--rl-sample-role trainer`` and a syncer (strict or elastic sync, no +``--rl-single-island-no-sync``): every island is a trainer, learner m of N, and +reads lane m of the buffer; the island with learner id 0 releases the merged +policy. The store reaches the island as ``--rl-sample-store`` (buffer and +releases only: multi-island round cuts do not exist), the trainer count as +``--rl-sample-lanes``. The samplers are separate +no-sync launches over the same store. + This module is engine-neutral; the Miles mapping is in ``adapters.miles.buffer_rollout``. """ @@ -64,6 +74,55 @@ def sample_role(args: Any) -> str: return role +def is_synced_trainer(args: Any) -> bool: + """A trainer island of a synced multi-island launch (module doc).""" + return sample_role(args) == TRAINER and not getattr(args, "rl_single_island_no_sync", False) + + +def buffer_lane(args: Any) -> Any: + """The buffer lane of a synced trainer island (island args: ``learner_id`` and + ``--rl-sample-lanes``, the trainer count the launcher forwards); None for every + other island.""" + from .merged_release import BufferLane + + if not is_synced_trainer(args): + return None + count = int(getattr(args, "rl_sample_lanes", None) or 1) + return None if count == 1 else BufferLane(int(getattr(args, "learner_id", 0) or 0), count) + + +def check_synced_trainers(args: Any, num_islands: int) -> None: + """Launch checks of synced trainers (module doc), before the per-role checks.""" + if not is_synced_trainer(args): + return + if int(num_islands) < 2: + raise ValueError("--rl-sample-role trainer without --rl-single-island-no-sync is a synced " + f"multi-island launch: it needs at least two islands (got {num_islands})") + if is_pair(args) or is_convertible_pool(args): + raise ValueError("synced trainers do not combine with --rl-island-sample-role or " + "convertible nodes") + if getattr(args, "tuning", "lora") != "lora": + raise ValueError("synced trainers merge LoRA deltas: they need --tuning lora") + + +def check_sample_store(args: Any) -> None: + """Island-side check: ``--rl-sample-store`` / ``--rl-sample-lanes`` belong to a + synced trainer island, which needs both (the launcher forwards them) and no + ``--rl-resume-store``; its learner id must name one of the lanes.""" + store, lanes = getattr(args, "rl_sample_store", None), getattr(args, "rl_sample_lanes", None) + if not is_synced_trainer(args): + if store is not None or lanes is not None: + raise ValueError("--rl-sample-store / --rl-sample-lanes need --rl-sample-role trainer " + "with a syncer (no --rl-single-island-no-sync)") + return + if not store or lanes is None or int(lanes) < 2: + raise ValueError("a synced trainer island needs --rl-sample-store and --rl-sample-lanes " + ">= 2 (the launcher forwards both)") + if getattr(args, "rl_resume_store", None): + raise ValueError("a synced trainer island takes --rl-sample-store, not --rl-resume-store") + buffer_lane(args) # learner id inside 0..lanes-1 + + def node_name(island: int) -> str: """The node name of island ``island`` (``--rl-node-class NAME=...``, buffer shard).""" return f"island{int(island)}" @@ -113,10 +172,24 @@ def island_roles(args: Any, num_islands: int) -> dict[int, str]: def is_pair(args: Any) -> bool: - """Whether the launch is a store-coupled sampler/trainer pair.""" + """Whether the launch is a store-coupled group: one trainer island and one or + more sampler islands (``--rl-island-sample-role``; "pair" for short).""" return bool(getattr(args, "rl_island_sample_role", None)) +def sampler_islands(args: Any, num_islands: int) -> list[int]: + """The sampler islands of a pair, sorted.""" + return sorted(m for m, role in island_roles(args, num_islands).items() if role == SAMPLER) + + +def buffer_shards(args: Any, num_islands: int) -> list[str] | None: + """The buffer shards of a pair with several sampler islands (S25 C1): one per + sampler island, named by its node name; None with one sampler (it writes the + index itself, as before).""" + samplers = sampler_islands(args, num_islands) + return [node_name(m) for m in samplers] if len(samplers) > 1 else None + + def is_store_coupled(args: Any) -> bool: """A pair or a convertible pool: no syncer, each island is its own one-learner no-sync island over the shared checkpoint store.""" @@ -133,6 +206,21 @@ def island_args(args: Any, island: int, num_islands: int) -> Any: return args out = copy.copy(args) out.rl_sample_role = roles[island] + if out.rl_sample_role in (SAMPLER, TRAINER): + out.rl_sampler_node = node_name(island) if out.rl_sample_role == SAMPLER else None + out.rl_buffer_shards = buffer_shards(args, num_islands) + samplers = len(sampler_islands(args, num_islands)) + if out.rl_sample_role == SAMPLER and getattr(args, "rl_sampler_spot", False): + # spot billing without the legacy --spot volume: the buffer is the sampler's + # durable state (S25 C1) + out.spot = True + out.spot_without_volume = True + if out.rl_sample_role == SAMPLER and samplers > 1: + # the trainer trains --total-steps rounds; the samplers share them, so each + # one reaches its data end (check_pair needs an even split) + out.total_steps = int(args.total_steps) // samplers + if getattr(args, "rl_total_fragment_steps", None) is not None: + out.rl_total_fragment_steps = out.total_steps * int(getattr(args, "fragments", 1) or 1) if out.rl_sample_role == CONVERTIBLE: out.rl_convertible_node = node_name(island) out.rl_sampler_min_round_seconds = None # a node keeps its backlog cap @@ -161,17 +249,30 @@ def buffer_full_parameter(args: Any) -> bool: def check_pair(args: Any, num_islands: int) -> None: - """Launch checks of a store-coupled pair (nothing outside a pair): exactly two - islands, one sampler and one trainer, ``--rl-single-island-no-sync`` (no - syncer), no ``--rl-sample-role`` and no external learner; then every island's - own role checks on its island args.""" + """Launch checks of a store-coupled pair (nothing outside a pair): every + island has a role, exactly one trainer and at least one sampler, + ``--rl-single-island-no-sync`` (no syncer), no ``--rl-sample-role`` and no + external learner; several samplers split ``--total-steps`` evenly (each + generates its share of the trainer's rounds) and take no trainer scheduler + (the head reads no sampler shard of a shared directory); then every island's own role checks + on its island args.""" roles = island_roles(args, num_islands) if not roles: return - if num_islands != 2 or sorted(roles.values()) != [SAMPLER, TRAINER] or len(roles) != 2: - raise ValueError("--rl-island-sample-role: a pair is two --gpu islands, one " - f"{SAMPLER} and one {TRAINER} (got {num_islands} island(s), roles " - f"{dict(sorted(roles.items()))})") + values = list(roles.values()) + if len(roles) != num_islands or values.count(TRAINER) != 1 or values.count(SAMPLER) < 1: + raise ValueError("--rl-island-sample-role: a pair gives every --gpu island a role, " + f"at least one {SAMPLER} and exactly one {TRAINER} (got {num_islands} " + f"island(s), roles {dict(sorted(roles.items()))})") + samplers = values.count(SAMPLER) + if samplers > 1 and int(getattr(args, "total_steps", 0) or 0) % samplers: + raise ValueError(f"--total-steps {args.total_steps}: {samplers} sampler islands share the " + "trainer's rounds; it must be a multiple of the sampler count") + if getattr(args, "rl_sampler_spot", False) and getattr(args, "spot", False): + raise ValueError("--rl-sampler-spot bills only the samplers as spot; it does not take --spot") + if samplers > 1 and getattr(args, "rl_trainer_scheduler", None) not in (None, "off"): + raise ValueError("--rl-trainer-scheduler with several sampler islands is not supported " + "(the head does not read the sampler shards); drop the scheduler") if sample_role(args) != BOTH: raise ValueError("--rl-island-sample-role replaces --rl-sample-role; give only one") if int(getattr(args, "external_learners", 0) or 0): @@ -221,14 +322,17 @@ def check_launch(args: Any) -> None: if getattr(args, dest, None) is not None: raise ValueError(f"{flag} needs --rl-sample-role sampler or trainer") return - if not getattr(args, "rl_single_island_no_sync", False): + if not getattr(args, "rl_single_island_no_sync", False) and role != TRAINER: raise ValueError(f"--rl-sample-role {role} needs --rl-single-island-no-sync: a sampler " - "and a trainer are separate no-sync islands over one store") + "is a no-sync island over the store (only trainers may sync)") if not getattr(args, "rl_checkpoint_store", None): raise ValueError(f"--rl-sample-role {role} needs --rl-checkpoint-store (the sample " "buffer and the policy releases live in it)") if getattr(args, "rl_elastic", False): - raise ValueError(f"--rl-sample-role {role} does not support --rl-elastic") + # --rl-elastic is placement reconfiguration (T/R cells), not the island sync: + # synced trainers use --rl-island-scheduling elastic (or strict sync) + raise ValueError(f"--rl-sample-role {role} does not support --rl-elastic " + "(placement reconfiguration)") if role in (TRAINER, CONVERTIBLE) and not int(getattr(args, "rl_max_policy_age", 0) or 0) >= 1: raise ValueError(f"--rl-sample-role {role} trains on samples of earlier versions: " "it needs --rl-max-policy-age >= 1") @@ -269,14 +373,15 @@ def _max_backlog_rounds(args: Any, role: str) -> int | None: def resolve(args: Any, *, groups_per_round: int) -> SampleBufferConfig | None: - """The island's buffer config (island side: the store is ``--rl-resume-store``).""" + """The island's buffer config (island side: the store is ``--rl-resume-store``, + or ``--rl-sample-store`` on a synced trainer).""" role = sample_role(args) if role == BOTH: return None - store = getattr(args, "rl_resume_store", None) + store = getattr(args, "rl_sample_store", None) or getattr(args, "rl_resume_store", None) if not store: raise ValueError(f"--rl-sample-role {role} needs the checkpoint store " - "(--rl-resume-store on the island)") + "(--rl-resume-store, or --rl-sample-store on a synced trainer)") rounds = _max_backlog_rounds(args, role) wait_s = getattr(args, "rl_sample_wait_s", None) min_round_s = getattr(args, "rl_sampler_min_round_seconds", None) @@ -288,6 +393,18 @@ def resolve(args: Any, *, groups_per_round: int) -> SampleBufferConfig | None: else None)) +def island_env(args: Any) -> dict[str, str]: + """The environment a pair's sampler island gets from the launcher ({} for + every other island; ``args`` are the island args): the reclaim marker path, + so the Modal runner turns the reclaim signal into that file and the sampler + leaves without being signalled (S25 C1).""" + if sample_role(args) != SAMPLER or not getattr(args, "rl_sampler_node", None): + return {} + from .sampler_resume import DEFAULT_RECLAIM_MARKER, RECLAIM_MARKER_ENV + + return {RECLAIM_MARKER_ENV: DEFAULT_RECLAIM_MARKER} + + def island_flags(args: Any) -> str: """The learner flags the launcher forwards (empty for ``both``); ``args`` are the island args (:func:`island_args`).""" @@ -303,4 +420,8 @@ def island_flags(args: Any) -> str: flags += f" --rl-sampler-min-round-seconds {float(args.rl_sampler_min_round_seconds)!r}" if getattr(args, "rl_sample_wait_s", None) is not None: flags += f" --rl-sample-wait-s {float(args.rl_sample_wait_s)!r}" + if getattr(args, "rl_sampler_node", None): + flags += f" --rl-sampler-node {args.rl_sampler_node}" + if getattr(args, "rl_buffer_shards", None): + flags += f" --rl-buffer-shards {','.join(args.rl_buffer_shards)}" return flags diff --git a/yeto/rl/engine/sampler_checks.py b/yeto/rl/engine/sampler_checks.py index 36d70274..4e0f26d7 100644 --- a/yeto/rl/engine/sampler_checks.py +++ b/yeto/rl/engine/sampler_checks.py @@ -149,6 +149,9 @@ class PostLoadChecks: allow_unverified_count: bool = False fixed_prompts: Sequence[FixedPrompt] = GSM8K_FIXED_PROMPTS fixed_max_tokens: int = 2048 + # False only for a model that cannot answer (a cut-down rehearsal model): a missing + # answer is recorded, the other fixed-prompt problems still fail the check + require_fixed_answers: bool = True forbidden_markers: Sequence[str] = () max_unreadable_share: float = 0.05 probe_pairs: Sequence[tuple[str, str]] | None = None # None: fixed prompts + their answers @@ -193,17 +196,21 @@ def fixed_prompt_answers(self, engine: SamplerEngine) -> dict[str, Any]: problems = [] if not text.strip(): problems.append("empty") - if not _contains_answer(text, item.answer): + answered = _contains_answer(text, item.answer) + if not answered and self.require_fixed_answers: problems.append("answer missing") if _unreadable_share(text) > self.max_unreadable_share: problems.append("unreadable characters") problems += [f"marker {m!r}" for m in self.forbidden_markers if m in text] - answers.append({"answer": item.answer, "ok": not problems, "problems": problems, - "text": text[:160]}) + answers.append({"answer": item.answer, "answered": answered, "ok": not problems, + "problems": problems, "text": text[:160]}) return {"name": "fixed_prompts", "passed": all(a["ok"] for a in answers) - and len(groups) == len(self.fixed_prompts), "answers": answers} + and len(groups) == len(self.fixed_prompts), "answers": answers, + "answers_required": bool(self.require_fixed_answers)} def logprob_probe(self, engine: SamplerEngine) -> dict[str, Any]: + if not self.pairs(): # no fixed prompts and no probe pairs: nothing to score or compare + return {"name": "logprob_probe", "passed": True, "detail": "no probe pairs"} means = [] for prompt, continuation in self.pairs(): logprobs = engine.score(prompt, continuation) diff --git a/yeto/rl/engine/sampler_resume.py b/yeto/rl/engine/sampler_resume.py new file mode 100644 index 00000000..234b06fc --- /dev/null +++ b/yeto/rl/engine/sampler_resume.py @@ -0,0 +1,329 @@ +"""Sampler island replacement: where a new incarnation continues (S25 C1). + +A sampler island on spot capacity can be reclaimed at any time; the cloud (or +the launcher) then starts a new incarnation of the same island. The island's +sample-buffer writes must neither repeat nor lose groups, and the trainer keeps +training on what the buffer holds. + +Each index entry a sampler writes carries a :class:`SamplerRound` record +(``SampleIndexEntry.sampler_round``): the sampler id (the island, stable across +incarnations), the incarnation, the rollout id, the data cursor after the round +drew its prompts, and ``last`` on the round's last group. A Miles round writes +its groups after the whole round is generated, so the buffer itself says how +far the island got: + +* a round with any entry in the buffer is never written again + (:class:`RoundGuard`); the new incarnation continues at the next rollout id + from that round's data cursor, so no prompt is drawn twice; +* a round without its ``last`` entry (reclaimed while it wrote) keeps the + groups it wrote (each whole and written once); its unwritten groups are lost + (``partial_rounds`` counts such rounds); +* a round still generating when the island was reclaimed wrote nothing: the new + incarnation generates that rollout id again from the cursor before it. + +:func:`resume_point` reads the entries; the island loop starts at +``next_rollout_id`` and the rollout function seeks the data source to +``data_cursor``. The incarnation's arrival and departure are written as +``rl_sampler_join`` / ``rl_sampler_leave`` (reason: a ``LeaveReason`` name) and as +one member record per incarnation in the store (:func:`write_member`), which +outlives the island's own tape. Engine-neutral. +""" + +from __future__ import annotations + +import json +import os +import threading +import time +import uuid +from collections.abc import Callable, Iterable, Mapping +from dataclasses import asdict, dataclass, field +from pathlib import Path +from typing import Any + +JOIN_EVENT = "rl_sampler_join" +LEAVE_EVENT = "rl_sampler_leave" +DROPPED_EVENT = "rl_sampler_groups_dropped" # groups not written after the leave +MEMBERS_DIR = "sample-buffer/members" # one record per sampler incarnation, under the store root +RECLAIM_MARKER_ENV = "YETO_SPOT_RECLAIM_MARKER" # written by the Modal runner on the reclaim signal +DEFAULT_RECLAIM_MARKER = "/tmp/yeto-reclaim-requested" +# Set where a host hook waits for the sampler's answer before the VM stops +# (Nebius: yeto.cloud.nebius_reclaim writes the marker and waits for this file). +RECLAIM_ACK_ENV = "YETO_RECLAIM_ACK_FILE" + + +class SamplerResumeError(RuntimeError): + pass + + +def new_incarnation() -> str: + """A fresh incarnation id (unique per island process).""" + return uuid.uuid4().hex[:8] + + +@dataclass(frozen=True) +class SamplerRound: + """The ``sampler_round`` record of one index entry.""" + + sampler: str + incarnation: str + rollout_id: int + data_cursor: Mapping[str, int] | None = None # the data source after the round drew its prompts + last: bool = False # the round's last group + + def record(self) -> dict[str, Any]: + out = {"sampler": self.sampler, "incarnation": self.incarnation, + "rollout_id": int(self.rollout_id), "last": bool(self.last)} + if self.data_cursor is not None: + out["data_cursor"] = {k: int(v) for k, v in self.data_cursor.items()} + return out + + +@dataclass(frozen=True) +class SamplerResume: + """Where the next incarnation of a sampler continues.""" + + next_rollout_id: int # after the newest round with an entry in the buffer + data_cursor: dict[str, int] | None # that round's data cursor (None: start of the data) + written_rounds: tuple[int, ...] # rounds with at least one entry: never written again + partial_rounds: tuple[int, ...] # written rounds without their last group + incarnations: tuple[str, ...] # earlier incarnations, in the order they first wrote + groups: int # every group this sampler has in the buffer + + @property + def resumed(self) -> bool: + return bool(self.incarnations) + + def fields(self) -> dict[str, Any]: + out = asdict(self) + out["written_rounds"] = len(self.written_rounds) + out["partial_rounds"] = list(self.partial_rounds) + out["incarnations"] = list(self.incarnations) + return out + + +def resume_point(entries: Iterable[Any], sampler: str) -> SamplerResume: + """Read the ``sampler_round`` records of ``sampler``'s entries (other samplers' + entries and entries without the record are skipped).""" + cursors: dict[int, dict[str, int] | None] = {} + ended: set[int] = set() + incarnations: list[str] = [] + groups = 0 + for entry in entries: + record = getattr(entry, "sampler_round", None) + if not record or record.get("sampler") != sampler: + continue + groups += 1 + rid = int(record["rollout_id"]) + raw = record.get("data_cursor") + cursor = None if raw is None else {k: int(v) for k, v in raw.items()} + if rid in cursors and cursors[rid] != cursor: + raise SamplerResumeError(f"sampler {sampler}: round {rid} has two data cursors " + f"({cursors[rid]} and {cursor})") + cursors[rid] = cursor + if record.get("incarnation") not in incarnations: + incarnations.append(str(record.get("incarnation"))) + if record.get("last"): + if rid in ended: + raise SamplerResumeError(f"sampler {sampler}: round {rid} has two last groups") + ended.add(rid) + newest = max(cursors) if cursors else None + return SamplerResume(next_rollout_id=0 if newest is None else newest + 1, + data_cursor=None if newest is None else cursors[newest], + written_rounds=tuple(sorted(cursors)), + partial_rounds=tuple(sorted(set(cursors) - ended)), + incarnations=tuple(incarnations), groups=groups) + + +class RoundGuard: + """Write-side check: a round with entries in the buffer is never written by + another incarnation, and an incarnation never writes before its first round.""" + + def __init__(self, resume: SamplerResume) -> None: + self._written = set(resume.written_rounds) + self._first = int(resume.next_rollout_id) + + def check(self, rollout_id: int) -> None: + rid = int(rollout_id) + if rid in self._written or rid < self._first: + raise SamplerResumeError( + f"round {rid} is already in the sample buffer (this incarnation starts at " + f"round {self._first}); a replacement sampler must not write it again") + + def written(self, rollout_id: int) -> None: + self._written.add(int(rollout_id)) + + +def reclaim_marker(environ: Mapping[str, str]) -> Path: + return Path(environ.get(RECLAIM_MARKER_ENV) or DEFAULT_RECLAIM_MARKER) + + +def write_reclaim_marker(environ: Mapping[str, str], signum: int) -> Path: + """Write the reclaim marker (what the Modal runner writes on its signal). The + runner may write it at the same moment: the temp file is per process, and a + failed write is reported, never raised (it runs in a signal handler; the + runner's marker serves as well).""" + path = reclaim_marker(environ) + tmp = path.with_name(f"{path.name}.{os.getpid()}.tmp") + try: + tmp.write_text(json.dumps({"signal": int(signum), "time": time.time()}) + "\n") + tmp.replace(path) + except OSError as error: + print(f"[sampler-island] reclaim marker not written ({error})", flush=True) + return path + + +def install_marker_signals(environ: Mapping[str, str], signal_mod: Any = None) -> list[str]: + """On the island's main thread: SIGINT and SIGTERM write the reclaim marker + instead of raising ``KeyboardInterrupt``. Modal signals the whole container, + so the island process gets the reclaim signal next to the runner; without + this it dies inside its round and leaves as INFRA_FAILED (T-C1 run a). + Returns the signals handled ([] off the main thread).""" + import signal as _signal + + sig = signal_mod or _signal + if threading.current_thread() is not threading.main_thread(): + return [] + done = [] + for name in ("SIGINT", "SIGTERM"): + num = getattr(sig, name, None) + if num is not None: + sig.signal(num, lambda signum, _frame: write_reclaim_marker(environ, signum)) + done.append(name) + return done + + +def clear_reclaim_files(environ: Mapping[str, str]) -> None: + """Remove a reclaim marker and ack left from an earlier boot (a restarted VM keeps /tmp).""" + reclaim_marker(environ).unlink(missing_ok=True) + ack = environ.get(RECLAIM_ACK_ENV) + if ack: + Path(ack).unlink(missing_ok=True) + + +def acknowledge_reclaim(environ: Mapping[str, str]) -> Path | None: + """Write the ack file a host hook waits for (``RECLAIM_ACK_ENV``) once the + LEAVE is durable, so the host lets the VM stop; None when no hook waits.""" + ack = environ.get(RECLAIM_ACK_ENV) + if not ack: + return None + path = Path(ack) + try: + path.write_text(f"{time.time():.3f}\n") + except OSError as error: + print(f"[sampler-island] reclaim ack not written ({error})", flush=True) + return None + return path + + +def reclaim_requested(environ: Mapping[str, str]) -> bool: + """The Modal runner wrote the reclaim marker (the container got SIGINT/SIGTERM).""" + return reclaim_marker(environ).exists() + + +@dataclass +class Member: + """One sampler incarnation's record in the store (``members//.json``).""" + + sampler: str + incarnation: str + joined_at: float + resume: dict[str, Any] + policy_version: int | None = None + left_at: float | None = None + leave_reason: str | None = None + rounds: int | None = None # rounds this incarnation completed + detail: dict[str, Any] = field(default_factory=dict) + + +def member_key(member: Member) -> str: + return f"{MEMBERS_DIR}/{member.sampler}/{member.incarnation}.json" + + +def write_member(store: Any, member: Member) -> Path: + """Write (or rewrite) the incarnation's member record and commit the store. + One file per incarnation: an old and a new incarnation never write the same + file (a Modal Volume keeps the last commit of a file). A store with an object + API (s3, OSS) gets the record by object: its staged mirror does not carry + ``sample-buffer/``, so a file would stay on the island's disk.""" + data = (json.dumps(asdict(member), sort_keys=True) + "\n").encode() + objects = getattr(store, "objects", None) + if objects is not None: + objects.write(member_key(member), data) + return Path(member_key(member)) + path = Path(store.root) / member_key(member) + path.parent.mkdir(parents=True, exist_ok=True) + tmp = path.with_suffix(".json.tmp") + tmp.write_bytes(data) + tmp.replace(path) + store.commit() + return path + + +class SamplerMembership: + """One incarnation's JOIN and LEAVE: the tape events and the member record. + :meth:`leave` acts once (the reclaim watcher and the island's exit path may + both call it); later calls return False.""" + + def __init__(self, store: Any, sampler: str, incarnation: str, resume: SamplerResume, + emit: Callable[..., None], clock: Callable[[], float] = time.time) -> None: + self.store = store + self.emit = emit + self.clock = clock + self.member = Member(sampler=sampler, incarnation=incarnation, joined_at=clock(), + resume=resume.fields()) + self._lock = threading.Lock() + self.left = False + + def join(self, **fields: Any) -> None: + m = self.member + self.emit(JOIN_EVENT, sampler=m.sampler, incarnation=m.incarnation, **m.resume, **fields) + write_member(self.store, m) + + def leave(self, reason: str, *, rounds: int | None = None, policy_version: int | None = None, + **detail: Any) -> bool: + with self._lock: + if self.left: + return False + self.left = True + m = self.member + m.left_at, m.leave_reason, m.rounds = self.clock(), str(reason), rounds + m.policy_version, m.detail = policy_version, dict(detail) + self.emit(LEAVE_EVENT, sampler=m.sampler, incarnation=m.incarnation, reason=str(reason), + rounds=rounds, policy_version=policy_version, **detail) + write_member(self.store, m) + return True + + +class ReclaimWatcher: + """Polls ``reclaimed()`` on a daemon thread and calls ``on_reclaim()`` once, at + once: the cloud stops the container a few seconds after its signal (Modal: + 30 s), while the island's loop may still be inside a round.""" + + def __init__(self, reclaimed: Callable[[], bool], on_reclaim: Callable[[], Any], *, + poll_s: float = 0.5) -> None: + self.reclaimed = reclaimed + self.on_reclaim = on_reclaim + self.poll_s = float(poll_s) + self._stop = threading.Event() + self.fired = threading.Event() + self._thread = threading.Thread(target=self._run, name="sampler-reclaim-watcher", daemon=True) + + def start(self) -> "ReclaimWatcher": + self._thread.start() + return self + + def stop(self) -> None: + self._stop.set() + self._thread.join(timeout=5.0) + + def _run(self) -> None: + while not self._stop.is_set(): + if self.reclaimed(): + self.fired.set() + try: + self.on_reclaim() + except Exception as error: # noqa: BLE001 - the island loop still stops on the marker + print(f"[sampler-island] reclaim handler failed: {error!r}", flush=True) + return + self._stop.wait(self.poll_s) diff --git a/yeto/rl/engine/trainer_lifecycle.py b/yeto/rl/engine/trainer_lifecycle.py index 731ecf41..ec3416ae 100644 --- a/yeto/rl/engine/trainer_lifecycle.py +++ b/yeto/rl/engine/trainer_lifecycle.py @@ -139,7 +139,13 @@ class SampleIndexEntry: record: device family, card, engine and version, precision, quantization scheme, topology); None for entries written before it existed or by a sampler that does not set it (read as ``unknown``). Data only: it never - enters a hash or the version gate (rl-hetero-sampler-engines D3).""" + enters a hash or the version gate (rl-hetero-sampler-engines D3). + + ``sampler_round``: which sampler round wrote the entry + (``sampler_resume.SamplerRound`` record: sampler, incarnation, rollout id, + last group of the round, data cursor after the round); None for entries + written before it existed. A replacement sampler island reads it to resume + after the last complete round (S25 C1).""" sample_id: str behavior_version: int @@ -147,6 +153,7 @@ class SampleIndexEntry: tokens: int = 0 appended_at: float = 0.0 provenance: Mapping[str, str] | None = field(default=None, compare=False, hash=False) + sampler_round: Mapping[str, Any] | None = field(default=None, compare=False, hash=False) @dataclass(frozen=True) diff --git a/yeto/rl/harness/chat_templates.py b/yeto/rl/harness/chat_templates.py index 13ea1883..9662b10c 100644 --- a/yeto/rl/harness/chat_templates.py +++ b/yeto/rl/harness/chat_templates.py @@ -16,7 +16,15 @@ one (same template family). Load the real checkpoint with ``ChatTemplate.load("glm5", tokenizer_path=...)``; whether GLM5.3 uses the same template is not verified; -* ``qwen3``: the Qwen3 hermes tool-call format; the small CPU stand-in. +* ``qwen3``: the Qwen3 hermes tool-call format; the small CPU stand-in; +* ``qwen3_2507``: Qwen3-4B-Instruct-2507 (non-thinking Qwen3, hermes JSON tool + calls, no reasoning block); the Codex system messages are merged into one; +* ``qwen38``: Qwen3.8 (Qwen3.5 family) XML tool calls (`` + v``), thinking kept in every + turn (``preserve_thinking``). Load with the checkpoint's tokenizer + (``tokenizer_path``); a deployment may set ``reasoning_effort`` by + replacing ``template_kwargs`` (xhigh / medium / low; the template's default + is xhigh). :func:`check_alignment` is the N10.6 check: the template rendering, tokenized the way the proxy sends it, must equal the tokenizer's own @@ -117,6 +125,23 @@ def parse_qwen3(text: str) -> Message: return _assistant(content, reasoning, calls) +_QWEN38_CALL = re.compile(r"\s*\n]+)>(.*?)\s*", re.S) +_QWEN38_PARAM = re.compile(r"\n]+)>\n?(.*?)\n?", re.S) + + +def parse_qwen38(text: str) -> Message: + """Qwen3.8 output after the ``\\n`` generation prompt.""" + reasoning, rest = _split_reasoning(text, opened=True) + calls = [] + for index, match in enumerate(_QWEN38_CALL.finditer(rest)): + arguments = {k.strip(): _glm_value(v) for k, v in _QWEN38_PARAM.findall(match.group(2))} + calls.append(_tool_call(_call_id(text, index), match.group(1).strip(), arguments)) + if rest.count("") != len(calls): + raise TemplateError(f"Qwen3.8 tool call is not XML: {rest[rest.find(''):][:80]!r}") + content = rest.split("", 1)[0].strip() + return _assistant(content, reasoning, calls) + + # ---------------------------------------------------------------------- registry @dataclass(frozen=True) class TemplateSpec: @@ -124,7 +149,10 @@ class TemplateSpec: an assistant turn (sent to the engine as stop token ids). ``keeps_history_reasoning``: with ``template_kwargs`` the template renders the reasoning of every earlier turn (gateway ``ChainRegistry`` input). - ``verified``: what the token alignment check ran on (free text).""" + ``verified``: what the token alignment check ran on (free text). + ``merge_leading_system``: join the leading system messages into one (the + template accepts a single system message; Codex sends its instructions + and a developer message, both system in the gateway form).""" name: str tokenizer: str @@ -133,6 +161,7 @@ class TemplateSpec: template_kwargs: Mapping[str, Any] = field(default_factory=dict) keeps_history_reasoning: bool = True verified: str = "not verified" + merge_leading_system: bool = False TEMPLATES: dict[str, TemplateSpec] = {} @@ -157,9 +186,30 @@ def register_template(spec: TemplateSpec) -> TemplateSpec: # The Qwen3 template drops the reasoning of turns before the last user turn. keeps_history_reasoning=False, verified="Qwen3-0.6B tokenizer (CPU stand-in)")) +register_template(TemplateSpec( + name="qwen3_2507", tokenizer="Qwen/Qwen3-4B-Instruct-2507", parser=parse_qwen3, + stop_tokens=("<|im_end|>", "<|endoftext|>"), + verified="Qwen3-4B-Instruct-2507@cdbee75f tokenizer (S25 T-N109)", merge_leading_system=True)) +register_template(TemplateSpec( + name="qwen38", tokenizer="Qwen/Qwen3.8-27B", parser=parse_qwen38, + stop_tokens=("<|im_end|>", "<|endoftext|>"), + template_kwargs={"enable_thinking": True, "preserve_thinking": True}, + verified="Qwen3.8-27B@1d4bf0f2 tokenizer (S25 T-G2)", merge_leading_system=True)) # ---------------------------------------------------------------------- template +def _merge_leading_system(messages: Sequence[Message]) -> list[Message]: + head = 0 + while head < len(messages) and messages[head].get("role") == "system": + head += 1 + if head < 2: + return list(messages) + if any(not isinstance(m.get("content"), str) for m in messages[:head]): + raise TemplateError("cannot merge system messages with non-text content") + merged = {"role": "system", "content": "\n\n".join(m["content"] for m in messages[:head])} + return [merged, *messages[head:]] + + def _render_ready(messages: Sequence[Message]) -> list[Message]: """Canonical messages -> template input: tool-call arguments as objects.""" out = [] @@ -208,9 +258,18 @@ def load(cls, name: str, *, tokenizer_path: str | None = None, local_files_only: def render(self, messages: Sequence[Message], *, tools: Sequence[Message] | None = None, add_generation_prompt: bool) -> str: - return self.tokenizer.apply_chat_template( - _render_ready(messages), tools=list(tools) if tools else None, tokenize=False, - add_generation_prompt=add_generation_prompt, **dict(self.spec.template_kwargs)) + import jinja2 # comes with transformers + + try: + return self.tokenizer.apply_chat_template( + self._ready(messages), tools=list(tools) if tools else None, tokenize=False, + add_generation_prompt=add_generation_prompt, **dict(self.spec.template_kwargs)) + except jinja2.TemplateError as exc: # raise_exception(...): the template refuses the conversation + raise TemplateError(f"{self.spec.name} template refused the conversation: {exc}") from exc + + def _ready(self, messages: Sequence[Message]) -> list[Message]: + ready = _render_ready(messages) + return _merge_leading_system(ready) if self.spec.merge_leading_system else ready def encode(self, text: str) -> list[int]: """Token ids of rendered text; the template writes its own BOS tokens.""" @@ -252,7 +311,7 @@ def check_alignment(template: ChatTemplate, messages: Sequence[Message], *, text = template.render(messages, tools=tools, add_generation_prompt=add_generation_prompt) ours = template.encode(text) reference = template.tokenizer.apply_chat_template( - _render_ready(messages), tools=list(tools) if tools else None, tokenize=True, + template._ready(messages), tools=list(tools) if tools else None, tokenize=True, add_generation_prompt=add_generation_prompt, **dict(template.spec.template_kwargs)) if isinstance(reference, Mapping): reference = reference["input_ids"] diff --git a/yeto/rl/harness/gateway/proxy.py b/yeto/rl/harness/gateway/proxy.py index f414fa3a..ac6e7a34 100644 --- a/yeto/rl/harness/gateway/proxy.py +++ b/yeto/rl/harness/gateway/proxy.py @@ -19,6 +19,7 @@ import json import secrets +from contextvars import ContextVar from typing import Any from yeto.rl.compaction import CompactionConfig @@ -27,18 +28,26 @@ from .core import AdmissionClosed, Gateway, GatewayConfig, GatewayError, SamplingOverrideError, TrajectoryInvalid from .records import RecordingBackend -from .translate import UnsupportedShape +from .translate import OMITTED_TOOL_TYPES, UnsupportedShape + +# The trajectory whose request is being served. Engine calls run in a thread +# started by ``asyncio.to_thread``, which copies this context: an engine that +# owns no version of its own (verl's rollout) can tag what it saw per trajectory. +CURRENT_TRAJECTORY: ContextVar[str | None] = ContextVar("records_proxy_trajectory", default=None) Message = dict[str, Any] def chat_tools(entry: str, body: dict[str, Any]) -> list[Message]: - """The request's tool schemas in Chat form (Responses tools are flat).""" + """The request's tool schemas in Chat form (Responses tools are flat). + Types in :data:`OMITTED_TOOL_TYPES` are left out of the prompt.""" tools = body.get("tools") or [] if entry != "responses": return list(tools) out = [] for tool in tools: + if isinstance(tool, dict) and tool.get("type") in OMITTED_TOOL_TYPES: + continue if not isinstance(tool, dict) or tool.get("type") != "function": raise UnsupportedShape(f"unsupported tool {tool!r}") function = {k: tool[k] for k in ("name", "description", "parameters") if k in tool} @@ -65,7 +74,11 @@ def open(self) -> str: async def handle(self, entry: str, body: dict[str, Any], *, trajectory_id: str) -> dict[str, Any]: self.backend.declare_tools(trajectory_id, chat_tools(entry, body)) - return await self.gateway.handle(entry, body, trajectory_id=trajectory_id) + token = CURRENT_TRAJECTORY.set(trajectory_id) + try: + return await self.gateway.handle(entry, body, trajectory_id=trajectory_id) + finally: + CURRENT_TRAJECTORY.reset(token) def records(self, trajectory_id: str) -> dict[str, Any]: out = self.backend.records(trajectory_id) diff --git a/yeto/rl/harness/gateway/recorder.py b/yeto/rl/harness/gateway/recorder.py index 2aadeabc..4e6b85ef 100644 --- a/yeto/rl/harness/gateway/recorder.py +++ b/yeto/rl/harness/gateway/recorder.py @@ -28,6 +28,11 @@ YETO_RECORDS_PROXY_HOST bind address (default 127.0.0.1) YETO_RECORDS_PROXY_PORT bind port (default 0: a free port) YETO_RECORDS_PROXY_PUBLIC_HOST host the sandbox dials (default: the bind address) + YETO_RECORDS_PROXY_PUBLIC_BASE full URL the sandbox dials instead (a tunnel, e.g. + ``https://.modal.host``); overrides PUBLIC_HOST and the port + YETO_RECORDS_PROXY_SECRET_FILE file holding a path secret: the proxy is served only under + ``/k//`` (a public tunnel never exposes the + unauthenticated root); the secret is never logged """ from __future__ import annotations @@ -35,6 +40,7 @@ import asyncio import importlib import os +from pathlib import Path from collections.abc import Mapping, Sequence from typing import Any @@ -47,6 +53,9 @@ HOST_ENV = "YETO_RECORDS_PROXY_HOST" PORT_ENV = "YETO_RECORDS_PROXY_PORT" PUBLIC_HOST_ENV = "YETO_RECORDS_PROXY_PUBLIC_HOST" +PUBLIC_BASE_ENV = "YETO_RECORDS_PROXY_PUBLIC_BASE" +SECRET_FILE_ENV = "YETO_RECORDS_PROXY_SECRET_FILE" +MIN_SECRET_CHARS = 24 # verl sampling keys -> SamplingParams fields; other keys (``logprobs``) do not change sampling. CHECKED_SAMPLING = ("temperature", "top_p", "top_k", "repetition_penalty") @@ -74,7 +83,9 @@ def turns_from_records(records: Mapping[str, Any]) -> list[TokenTurn]: turns.append(TokenTurn(prompt_ids=tuple(ids[:cursor]), completion_ids=completion, completion_logprobs=tuple(float(x) for x in turn["logprobs"]), policy_version=turn.get("policy_version"), - provenance=turn.get("provenance"))) + provenance=turn.get("provenance"), + finish_reason=turn.get("finish_reason"), + policy_error=turn.get("parse_error"))) cursor, pending_stop = end, len(turn.get("stop_added") or ()) if cursor + pending_stop != len(ids): raise TokenAlignmentError(f"segment {sample['segment']} has {len(ids) - cursor - pending_stop} " @@ -86,11 +97,16 @@ class ProxyTokenRecorder: """:class:`~yeto.rl.agentic.TokenRecorder` over one in-process :class:`RecordsProxy` (module doc).""" def __init__(self, proxy: RecordsProxy, *, host: str = "127.0.0.1", port: int = 0, - public_host: str | None = None) -> None: + public_host: str | None = None, public_base: str | None = None, + path_secret: str | None = None) -> None: + if path_secret is not None and len(path_secret) < MIN_SECRET_CHARS: + raise HarnessConfigError(f"the path secret must have at least {MIN_SECRET_CHARS} characters") self.proxy = proxy self.host = host self.port = int(port) self.public_host = public_host or host + self.public_base = public_base.rstrip("/") if public_base else None + self._prefix = f"/k/{path_secret}" if path_secret else "" self._signed = sampling_params(proxy.gateway.config.signed_sampling) self._runner: Any = None self._base_url: str | None = None @@ -102,13 +118,19 @@ async def _base(self) -> str: if self._base_url is None: from aiohttp import web - runner = web.AppRunner(self.proxy.app()) + app = self.proxy.app() + if self._prefix: + outer = web.Application() + outer.add_subapp(f"{self._prefix}/", app) + app = outer + runner = web.AppRunner(app) await runner.setup() site = web.TCPSite(runner, self.host, self.port) await site.start() port = runner.addresses[0][1] self._runner = runner - self._base_url = f"http://{self.public_host}:{port}" + base = self.public_base or f"http://{self.public_host}:{port}" + self._base_url = f"{base}{self._prefix}" return self._base_url def check_sampling(self, sampling: Mapping[str, Any]) -> None: @@ -159,5 +181,8 @@ def recorder_from_env() -> ProxyTokenRecorder: if not isinstance(proxy, RecordsProxy): raise HarnessConfigError(f"{spec} returned {type(proxy).__name__}, not a RecordsProxy") host = os.environ.get(HOST_ENV, "127.0.0.1") + secret_file = os.environ.get(SECRET_FILE_ENV, "").strip() + secret = Path(secret_file).read_text().strip() if secret_file else None return ProxyTokenRecorder(proxy, host=host, port=int(os.environ.get(PORT_ENV, "0")), - public_host=os.environ.get(PUBLIC_HOST_ENV) or host) + public_host=os.environ.get(PUBLIC_HOST_ENV) or host, + public_base=os.environ.get(PUBLIC_BASE_ENV) or None, path_secret=secret) diff --git a/yeto/rl/harness/gateway/records.py b/yeto/rl/harness/gateway/records.py index c9b9f95d..0a42f789 100644 --- a/yeto/rl/harness/gateway/records.py +++ b/yeto/rl/harness/gateway/records.py @@ -41,7 +41,7 @@ SamplingParams, TokenSamplerEngine, ) -from yeto.rl.harness.chat_templates import ChatTemplate, first_difference +from yeto.rl.harness.chat_templates import ChatTemplate, TemplateError, first_difference from .core import GatewayError @@ -84,7 +84,11 @@ class TurnRecord: ``completion``: the sampled ids (mask 1, stop token included when the engine returned it). ``stop_added``: a stop token the engine did not return, added with mask 0. ``render_diff``: first position where the - stream differs from a fresh tokenization (None: aligned).""" + stream differs from a fresh tokenization (None: aligned). ``parse_error``: + the policy failed after this completion: the completion is not valid + model output for the template (HTTP 400), or the next prompt no longer + fits the context (``context_full: ...``, HTTP 409). The sampled ids stay + in the stream.""" turn: int kind: str @@ -96,6 +100,7 @@ class TurnRecord: provenance: dict[str, str] stop_added: tuple[int, ...] = () render_diff: int | None = None + parse_error: str | None = None @dataclass @@ -200,7 +205,12 @@ async def chat(self, session_id: str, messages: list[Message], sampling: dict[st completion = await self._sample(session, segment, pending, params, kind=KIND_ACTION) segment.n_context = len(context) + 1 prompt_tokens = len(segment.token_ids) - len(completion.token_ids) - message = self._parse(completion) + try: + message = self._parse(completion) + except TemplateError as exc: # malformed model output: the policy's failure (N10.9) + segment.turns[-1].parse_error = str(exc)[:500] + self.counters[session.trajectory_id]["policy_output_errors"] += 1 + raise return self._chat_response(session_id, message, completion, prompt_tokens=prompt_tokens) # -------------------------------------------------------------- records @@ -295,7 +305,13 @@ async def _sample(self, session: _Session, segment: Segment, pending: _Prompt, p prompt = [*segment.token_ids, *pending.stop, *pending.tail] room = int(self.engine.max_model_len) - len(prompt) if room < 1: - raise ContextFull(f"prompt has {len(prompt)} tokens; max_model_len is {self.engine.max_model_len}") + message = f"prompt has {len(prompt)} tokens; max_model_len is {self.engine.max_model_len}" + # The policy filled its own context (a budget, like a token cap): mark it on the last + # sampled turn so the outer layer trains it as a 0 instead of an infrastructure abort. + if segment.turns: + segment.turns[-1].parse_error = f"context_full: {message}" + self.counters[session.trajectory_id]["context_full"] += 1 + raise ContextFull(message) params = params.replace(max_tokens=min(params.max_tokens, room)) version = self.engine.served_version() (group,) = await asyncio.to_thread(self.engine.generate_tokens, [prompt], params, diff --git a/yeto/rl/harness/gateway/records_env.py b/yeto/rl/harness/gateway/records_env.py new file mode 100644 index 00000000..6f434e54 --- /dev/null +++ b/yeto/rl/harness/gateway/records_env.py @@ -0,0 +1,12 @@ +"""Environment names of a records proxy built inside a trainer's rollout process (S25 N10.9). + +Shared by the backend adapters (verl ``rollout_engine``, Miles +``records_generate``); kept here so adapters do not import each other. +""" + +TEMPLATE_ENV = "YETO_RECORDS_TEMPLATE" # chat template name (``chat_templates.TEMPLATES``) +TOKENIZER_ENV = "YETO_RECORDS_TOKENIZER" # tokenizer directory (the checkpoint) +MODEL_ENV = "YETO_RECORDS_MODEL" # model name the agent sends +SAMPLING_ENV = "YETO_RECORDS_SAMPLING" # JSON signed sampling +MAX_MODEL_LEN_ENV = "YETO_RECORDS_MAX_MODEL_LEN" # context length of the rollout engine +CARD_ENV = "YETO_RECORDS_CARD" # GPU name (provenance) diff --git a/yeto/rl/harness/gateway/translate.py b/yeto/rl/harness/gateway/translate.py index a533c4d2..0da3e5e9 100644 --- a/yeto/rl/harness/gateway/translate.py +++ b/yeto/rl/harness/gateway/translate.py @@ -73,6 +73,14 @@ def completion_to_chat(completion: dict[str, Any]) -> dict[str, Any]: # ------------------------------------------------------------------ /v1/responses (Codex) +# Responses tool types the policy is not offered: ``namespace`` (Codex 0.145 +# sends its sub-agent tools so; a sub-agent would open a second token stream) +# and OpenAI hosted tools (the served model cannot call them). Other +# non-function tool types are refused. +OMITTED_TOOL_TYPES = frozenset({"namespace", "web_search", "web_search_preview", "image_generation", + "file_search", "code_interpreter"}) + + def responses_to_messages(body: dict[str, Any]) -> list[dict[str, Any]]: if body.get("store") is not False: raise UnsupportedShape("Responses requests must use store:false (full replay)") @@ -82,6 +90,8 @@ def responses_to_messages(body: dict[str, Any]) -> list[dict[str, Any]]: if summary not in ("none", None): raise UnsupportedShape("reasoning.summary must be none") for tool in body.get("tools") or []: + if isinstance(tool, dict) and tool.get("type") in OMITTED_TOOL_TYPES: + continue if not isinstance(tool, dict) or tool.get("type") != "function": raise UnsupportedShape(f"unsupported tool type {tool.get('type') if isinstance(tool, dict) else tool!r}") messages: list[dict[str, Any]] = [] diff --git a/yeto/rl/harness/harbor/README.md b/yeto/rl/harness/harbor/README.md index 8295646a..8c7195e9 100644 --- a/yeto/rl/harness/harbor/README.md +++ b/yeto/rl/harness/harbor/README.md @@ -1,8 +1,9 @@ # Harbor harness for agentic RL (S25 N10.1-N10.4) Status: Harbor trials with the oracle/nop agents ran on boat sandboxes (see -"boat sandbox provider"). **No model-driven (codex) trajectory has run yet -(not verified).** +"boat sandbox provider"). Codex trajectories with a real model ran on boat +on 2026-10-20 (see "Codex + records proxy"). A training run on either +backend: not verified. ## Who does what @@ -168,3 +169,27 @@ Miles path, 2026-10-19: the Miles agent function rollout worker or session server) plus `reward_func`, `HARBOR_ENV_TYPE=boat`: fix-git and regex-log, oracle reward 1 on both, nop reward 0 on both (outcome ok); 30 billable s. A full Miles training run on boat: not verified. + +## Codex + records proxy (N10.7 smoke, 2026-10-20) + +Codex 0.145.0 refuses `wire_api = "chat"` ("no longer supported"): it only +speaks the Responses API, so the records proxy (Responses -> Chat, N10.5) is +its path to a chat-template engine. Codex first tries a WebSocket on +`/v1/responses` (5 attempts, the proxy answers 405), then uses HTTP POST. + +What the yeto side sets for Codex: `reasoning_summary=none` and +`web_search=disabled` (`trial.CODEX_KWARGS`); the proxy leaves Codex's +`namespace` (sub-agent) tools out of the prompt; the `qwen38` template merges +Codex's two leading system messages (instructions + developer message). + +Run: `scripts/modal_records_endpoint.py` (Qwen3.8-27B bf16, SGLang TP1 on one +H100, records proxy behind a Modal tunnel and a secret path prefix); +Harbor `codex` on boat small VMs; driver and raw data in +`infra-drafts/SESSION25/logs/T-G2/`. Miles path (agent function + reward +hook, proxy in place of the session server): fix-git and +log-summary-date-ranges, reward 1 on both. verl path (`HarborAgentLoop.run`, +verl module stubbed, HTTP recorder over the proxy): openssl-selfsigned-cert 1, +fix-git 1, regex-log 0 (twice), sqlite-db-truncate 0. All 8 outcome ok; every +recorded turn re-tokenizes to the same ids as the stream (render drift 0); +`boat list` empty afterwards. The regex-log runs ended on one reply that hit +the 8192-token cap (`length_capped_turns`). diff --git a/yeto/rl/harness/harbor/trial.py b/yeto/rl/harness/harbor/trial.py index a213c1d3..f506d57e 100644 --- a/yeto/rl/harness/harbor/trial.py +++ b/yeto/rl/harness/harbor/trial.py @@ -42,6 +42,10 @@ # server, which does not authenticate the agent. AGENT_API_KEY = "dummy" +# Harbor codex CLI flags: the gateway serves no reasoning summaries (it refuses +# ``reasoning.summary`` other than none) and no hosted web search. +CODEX_KWARGS: Mapping[str, str] = {"reasoning_summary": "none", "web_search": "disabled"} + @dataclass(frozen=True) class HarborSettings: @@ -112,6 +116,8 @@ def build_trial_config(task: AgentTask, endpoint: str, settings: HarborSettings) version = task.agent_version or (CODEX_VERSION if task.agent == "codex" else None) if version is not None: kwargs["version"] = version + if task.agent == "codex": + kwargs.update(CODEX_KWARGS) if task.budget.max_seq_len is not None and task.agent.startswith("terminus"): kwargs["max_seq_len"] = task.budget.max_seq_len return TrialConfig( diff --git a/yeto/rl/quant/release.py b/yeto/rl/quant/release.py index 0eae2762..82f686c5 100644 --- a/yeto/rl/quant/release.py +++ b/yeto/rl/quant/release.py @@ -132,12 +132,13 @@ def run(self, stop: threading.Event, poll_s: float = 10.0) -> None: stop.wait(poll_s) -def requantizer(template_dir: Path, smoothing: Any) -> Requantize: - """The converter of :mod:`.requant` bound to one template and AWQ smoothing.""" +def requantizer(template_dir: Path, smoothing: Any, frozen: tuple[Path, str] | None = None) -> Requantize: + """The converter of :mod:`.requant` bound to one template and AWQ smoothing + (``frozen``: base checkpoint and name pattern of untrained tensors the release omits).""" from .requant import requantize_checkpoint def run(source: Path, out: Path) -> Any: - return requantize_checkpoint(Path(source), Path(template_dir), Path(out), smoothing) + return requantize_checkpoint(Path(source), Path(template_dir), Path(out), smoothing, frozen=frozen) return run @@ -156,7 +157,12 @@ def main(argv: list[str] | None = None) -> int: p.add_argument("--keep", type=int, default=2) p.add_argument("--loop", action="store_true") p.add_argument("--poll-seconds", type=float, default=10.0) + p.add_argument("--frozen-base", help="BF16 checkpoint that supplies untrained tensors the release omits") + p.add_argument("--frozen-pattern", help="regular expression of those tensor names (with --frozen-base)") a = p.parse_args(argv) + if bool(a.frozen_base) != bool(a.frozen_pattern): + p.error("give --frozen-base and --frozen-pattern together") + frozen = (Path(a.frozen_base), a.frozen_pattern) if a.frozen_base else None if a.no_awq: smoothing, awq = AwqSmoothing.none(), AWQ_NONE elif a.recipe and a.awq_scales: @@ -172,7 +178,7 @@ def main(argv: list[str] | None = None) -> int: store = store_for(a.store) releaser = RequantReleaser(PolicyReleaseStore(store, keep=a.keep, release_dir=release_dir_for(BF16)), PolicyReleaseStore(store, keep=a.keep, release_dir=release_dir_for(scheme)), - requantizer(Path(a.template), smoothing), scheme=scheme, awq=awq, + requantizer(Path(a.template), smoothing, frozen), scheme=scheme, awq=awq, awq_source=smoothing.source if smoothing.enabled else None, expert_matrices=len(layout.quantized_modules)) if a.loop: diff --git a/yeto/rl/quant/requant.py b/yeto/rl/quant/requant.py index 44fced0e..31b61582 100644 --- a/yeto/rl/quant/requant.py +++ b/yeto/rl/quant/requant.py @@ -203,6 +203,11 @@ def from_checkpoint(cls, ckpt_dir: Path) -> "CheckpointLayout": @classmethod def load(cls, template_dir: Path) -> "CheckpointLayout": + """The layout of a template directory (``layout.json``) or, when it has + none, of a reference checkpoint itself (index + shard headers; stage 1 + passes the round-0 MXFP4 export as the template).""" + if not (Path(template_dir) / LAYOUT_FILE).is_file() and (Path(template_dir) / INDEX_FILE).is_file(): + return cls.from_checkpoint(Path(template_dir)) raw = json.loads((Path(template_dir) / LAYOUT_FILE).read_text()) if raw.get("schema") != LAYOUT_SCHEMA: raise LayoutError(f"{template_dir}: unknown layout schema {raw.get('schema')!r}") @@ -377,15 +382,25 @@ def close(self) -> None: self._open.clear() +_COMPOSITE_PREFIX, _TEXT_PREFIX = "model.language_model.", "model." + + def requantize_checkpoint( source_dir: Path, template_dir: Path, out_dir: Path, smoothing: AwqSmoothing, shards: Iterable[str] | None = None, + frozen: tuple[Path, str] | None = None, ) -> dict[str, Any]: """Write the target checkpoint (or the given output shards of it) into ``out_dir``. + ``frozen``: ``(base_dir, pattern)``. A layout tensor that the source lacks and + whose source name matches the regular expression ``pattern`` is read from the + BF16 checkpoint ``base_dir`` (weights the trainer does not train and does not + export, e.g. the vision tower of a text-only Miles run). Every other missing + tensor still fails; ``stats["frozen_from_base"]`` counts the filled ones. + ``template_dir`` holds ``layout.json`` and the metadata files (:meth:`CheckpointLayout.save_template`). Without ``shards`` the whole checkpoint is written, then the index and the metadata files, so a @@ -398,6 +413,21 @@ def requantize_checkpoint( source = _SourceReader(source_dir) expected = layout.source_names() missing, extra = expected - set(source.weight_map), set(source.weight_map) - expected + # Miles' raw glm5_next export names the text model "model.layers..."; the HF + # composite checkpoint (and so the layout) has "model.language_model.layers..." + # (Modal run-14). Only an exact one-to-one prefix match is taken as the same tensor. + alias = {n: n.replace(_COMPOSITE_PREFIX, _TEXT_PREFIX, 1) for n in missing + if n.startswith(_COMPOSITE_PREFIX) and n.replace(_COMPOSITE_PREFIX, _TEXT_PREFIX, 1) in extra} + missing, extra = missing - set(alias), extra - set(alias.values()) + base, filled = None, set() + if frozen is not None and missing: + matcher = re.compile(frozen[1]) + filled = {n for n in missing if matcher.search(n)} + if filled: + base = _SourceReader(Path(frozen[0])) + if absent := filled - set(base.weight_map): + raise LayoutError(f"frozen base {frozen[0]} lacks {len(absent)} tensors, e.g. {sorted(absent)[:3]}") + missing = missing - filled if missing or extra: raise LayoutError( f"source tensors differ from the target layout: {len(missing)} missing (e.g. {sorted(missing)[:3]}), " @@ -410,7 +440,12 @@ def requantize_checkpoint( quantized = set(layout.quantized_weights()) out_dir = Path(out_dir) out_dir.mkdir(parents=True, exist_ok=True) - stats = {"shards": 0, "tensors": 0, "quantized": 0} + stats = {"shards": 0, "tensors": 0, "quantized": 0, "frozen_from_base": len(filled), + "renamed_from_text_prefix": len(alias)} + + def read(name: str) -> torch.Tensor: + return base.get(name) if name in filled else source.get(alias.get(name, name)) + try: for shard in selected: produced: dict[str, torch.Tensor] = {} @@ -419,7 +454,7 @@ def requantize_checkpoint( for name in all_shards[shard]: if name not in produced: src = layout.source_of(name) - produced.update(transform_tensor(src, source.get(src), layout, smoothing)) + produced.update(transform_tensor(src, read(src), layout, smoothing)) tensors = {n: produced[n] for n in all_shards[shard]} write_output_shard(out_dir, shard, tensors, layout) stats["shards"] += 1 @@ -427,6 +462,8 @@ def requantize_checkpoint( stats["tensors"] += len(tensors) finally: source.close() + if base is not None: + base.close() if shards is None: finalize_checkpoint(template_dir, out_dir) return stats diff --git a/yeto/rl/quant/shards.py b/yeto/rl/quant/shards.py index 145441e0..10c6eb28 100644 --- a/yeto/rl/quant/shards.py +++ b/yeto/rl/quant/shards.py @@ -25,7 +25,9 @@ ``@``. Status: CPU tests (2 gloo ranks, both schemes, byte-equal to the single-process -output; tests/test_rl_quant_shards.py). Not run inside a real trainer. +output; tests/test_rl_quant_shards.py). Inside a real Miles trainer (2xH100, +glm5_next 4-layer mini, TP2/EP2, through the Miles adapter): byte-equal for both +schemes (T-G3, 2026-10-20). Not run inside verl. """ from __future__ import annotations diff --git a/yeto/rl/stage1.py b/yeto/rl/stage1.py index 4825e9ca..67ed28a5 100644 --- a/yeto/rl/stage1.py +++ b/yeto/rl/stage1.py @@ -179,6 +179,46 @@ def record(self) -> dict[str, Any]: } +# Where the run store lives: Alibaba OSS (stage 1 on PPU), or a Modal Volume +# mounted at the store path in every container (the Modal rehearsal, S25 T-G1), +# reached through the yeto store hook (resume.STORE_IMPL_ENV). +STORE_BACKENDS = {"oss": None, "modal-volume": "yeto.cloud.modal_ckpt_store:store_from_env"} + + +# GRPO + token-level TIS, execution.max_policy_staleness 2 (relative to the yeto workdir). +DEFAULT_ALGORITHM_SPEC = "examples/rl_algorithms/decoupled-ppo-c2.json" + + +def expected_algorithm_sha256(island_argv: Sequence[str]) -> str: + """The AlgorithmSpec hash the island computes from this argv (the launcher's + ``resolve_ports_algorithm(args).sha256()``). The spec path is resolved against + the yeto tree, where the island runs.""" + from yeto.launch import REPO_ROOT # noqa: PLC0415 + from yeto.rl.adapters.miles.island_entry import parse_args # noqa: PLC0415 + from yeto.rl.engine.algorithm import resolve_ports_algorithm # noqa: PLC0415 + + argv = list(island_argv[3:]) + i = argv.index("--rl-algorithm-spec") + 1 + if not Path(argv[i]).is_absolute(): + argv[i] = str(Path(REPO_ROOT) / argv[i]) + return resolve_ports_algorithm(parse_args(argv)).sha256() + + +def store_env(plan: Stage1Plan) -> dict[str, str]: + """Environment every node needs for the store backend (none for OSS here: the + nodes' ossutil/OSS configuration is set up on the node, runbook S4).""" + backend = str(plan.doc.get("store_backend") or "oss") + if STORE_BACKENDS.get(backend) is None: + return {} + from yeto.cloud.modal_ckpt_store import MODAL_VOLUME_ENV # noqa: PLC0415 + from yeto.rl.engine.resume import STORE_IMPL_ENV # noqa: PLC0415 + + volume = str(plan.doc.get("store_volume") or "") + if not volume: + raise Stage1Error("plan: store_backend modal-volume needs store_volume (the Volume name)") + return {STORE_IMPL_ENV: STORE_BACKENDS[backend], MODAL_VOLUME_ENV: volume} + + def run_store(plan: Stage1Plan) -> str: """The run directory on the store: ``/``.""" return f"{str(_get(plan.doc, 'store', 'plan')).rstrip('/')}/{plan.run_id}" @@ -236,12 +276,13 @@ def megatron_memory(model: Mapping[str, Any], miles: Mapping[str, Any], layout: if hf_config is not None: hf = shape_from_hf_config(json.loads(Path(hf_config).read_text())) shape = dataclasses.replace(hf, params=params, expert_params=experts, - num_experts=hf.num_experts or glm5_next.NUM_EXPERTS) + num_experts=hf.num_experts or int(model.get("num_experts", glm5_next.NUM_EXPERTS))) else: shape = ModelShape(params=params, largest_layer_params=params // int(model["num_layers"]), num_layers=int(model["num_layers"]), hidden_size=glm5_next.HIDDEN_SIZE, vocab_size=glm5_next.VOCAB_SIZE, expert_params=experts, - num_experts=glm5_next.NUM_EXPERTS, source="stage-1 plan + glm5_next recipe") + num_experts=int(model.get("num_experts", glm5_next.NUM_EXPERTS)), + source="stage-1 plan + glm5_next recipe") node = trainers[0] base = TrainPlan(framework="megatron", cards=sum(n.cards_per_node for n in trainers), cards_per_node=node.cards_per_node, hbm_gib=node.hbm_gib, @@ -319,6 +360,7 @@ def _sampler_engine(plan: Stage1Plan, node: Node): extra = ["--moe-runner-backend", str(s["moe_runner_backend"])] if s.get("hierarchical_cache"): extra.append("--enable-hierarchical-cache") + extra += [str(a) for a in s.get("extra_args") or ()] # further engine flags of the plan config = ServerConfig( base_dir=Path(plan.section("model")["sampler_initial"]["local_dir"]), card_model=node.card, device_family=node.device_family, @@ -326,9 +368,12 @@ def _sampler_engine(plan: Stage1Plan, node: Node): quant_scheme=str(q["scheme"]), topology=str(s["topology"]), max_model_len=int(s["max_model_len"]), host="0.0.0.0", port=int(s["port"]), extra_args=tuple(extra), expected_format=str(q["format"]), - in_place_reload=bool(s.get("in_place_reload"))) + in_place_reload=bool(s.get("in_place_reload")), python=str(s.get("python") or "python"), + extra_env={str(k): str(v) for k, v in (s.get("extra_env") or {}).items()}) options = SGLangOptions(tp=int(s["tp"]), nextn=bool(s.get("nextn")), - renames_weight_packed=bool(s.get("renames_weight_packed"))) + renames_weight_packed=bool(s.get("renames_weight_packed")), + load_report_dir=(f"{state_dir_of(plan)}/sampler-{node.name}/load-report" + if s.get("load_report") else None)) return SGLangServerEngine(config, options) @@ -378,20 +423,26 @@ class _ParseError(Exception): def _check_island_argv(argv: Sequence[str]) -> str | None: - """Parse the Miles island argv with the island's own parser.""" - from yeto.rl.adapters.miles.island_entry import build_parser # noqa: PLC0415 + """Parse the Miles island argv with the island's own ``parse_args``: the parser + plus its combination checks (e.g. the ports placement rules; Modal run-1 failed + on one that the bare parser accepted).""" + import argparse # noqa: PLC0415 + from unittest import mock # noqa: PLC0415 - parser = build_parser() + from yeto.rl.adapters.miles.island_entry import build_parser, parse_args # noqa: PLC0415 - def fail(message: str) -> None: + def fail(_parser: Any, message: str) -> None: raise _ParseError(message) - parser.error = fail # type: ignore[method-assign] - no usage dump, no exit - try: - _, unknown = parser.parse_known_args(list(argv[3:])) - except _ParseError as exc: - return f"Miles island argv does not parse: {exc}" - return f"Miles island argv: flags unknown to island_entry: {unknown}" if unknown else None + with mock.patch.object(argparse.ArgumentParser, "error", fail): # no usage dump, no exit + try: + _, unknown = build_parser().parse_known_args(list(argv[3:])) + if unknown: + return f"Miles island argv: flags unknown to island_entry: {unknown}" + parse_args(list(argv[3:])) + except _ParseError as exc: + return f"Miles island argv does not parse: {exc}" + return None def expand(plan: Stage1Plan, *, hf_config: str | Path | None = None) -> Expanded: @@ -415,15 +466,20 @@ def expand(plan: Stage1Plan, *, hf_config: str | Path | None = None) -> Expanded miles = _section(trainer_cfg, "miles") layout = megatron_layout(miles, cards, per_node) common_env = {DEVICE_FAMILY_ENV: trainers[0].device_family, "YETO_RUN_ID": plan.run_id, - "YETO_RL_STORE": store, "HF_HUB_OFFLINE": "1"} - trainer_env = {**common_env, "TORCHINDUCTOR_COMPILE_THREADS": "1"} + "YETO_RL_STORE": store, "HF_HUB_OFFLINE": "1", **store_env(plan)} + trainer_env = {**common_env, "TORCHINDUCTOR_COMPILE_THREADS": "1", + # trainer.miles.extra_env: image-specific settings (e.g. PYTHONPATH) + **{str(k): str(v) for k, v in (_section(trainer_cfg, "miles").get("extra_env") or {}).items()}} model_dir = str(model["bf16_source"]["local_dir"]) state_dir = state_dir_of(plan) impl = str(miles["transformer_impl"]) offload = float(miles.get("optimizer_offload_fraction") or 0.0) island_argv = ["python3", "-m", "yeto.rl.adapters.miles.island_entry", "--model", model_dir, "--model-revision", str(model["bf16_source"]["revision"]), - "--data", str(task["data"]), "--data-revision", str(task["data_revision"]), + # task.data_local: the dataset staged on the node (the nodes run with HF_HUB_OFFLINE=1, + # where a Hub id does not load; Modal run-2); a local path takes no revision + *(("--data", str(task["data_local"])) if task.get("data_local") else + ("--data", str(task["data"]), "--data-revision", str(task["data_revision"]))), "--reward-function", str(task["reward"]), "--samples-per-group", str(int(task["n_samples_per_prompt"])), "--rollout-max-response-len", str(int(sampling["max_tokens"])), @@ -435,25 +491,52 @@ def expand(plan: Stage1Plan, *, hf_config: str | Path | None = None) -> Expanded "--tensor-parallel", str(layout["tp"]), "--pipeline-parallel", str(layout["pp"]), "--expert-parallel", str(layout["ep"]), "--miles-root", str(miles["workdir"]), + *(("--megatron-ref-load", str(model["bf16_source"]["torch_dist"])) + if model["bf16_source"].get("torch_dist") else ()), + *(("--miles-commit", str(miles["commit"])) if len(str(miles["commit"])) == 40 else ()), "--trust-remote-code", "--learner-id", "0", "--num-learners", "1", "--rl-single-island-no-sync", + # samples come from the buffer, one round ahead at most (sampler max_backlog_groups); + # a policy age >= 1 needs an algorithm spec that tolerates it (TIS correction) + "--rl-max-policy-age", "1", + "--rl-algorithm-spec", str(miles.get("algorithm_spec") or DEFAULT_ALGORITHM_SPEC), "--global-rounds", str(int(plan.doc["rounds"])), "--groups-per-round", str(int(task["rollout_batch_size"])), "--over-sampling-batch-size", str(int(task["rollout_batch_size"])), "--optimizer-steps", "1", + # trainer-only island: no --rollout-num-gpus (the ports check refuses any value, + # 0 included, without --rl-placement fixed-partition; Modal run-1) "--actor-num-nodes", str(len(trainers)), "--actor-num-gpus-per-node", str(per_node), - "--rollout-num-gpus", "0", "--rl-activation-recompute", "--rl-max-tokens-per-gpu", str(int(miles["max_tokens_per_gpu"])), *(("--rl-optimizer-offload-fraction", repr(offload)) if offload > 0 else ()), "--completed-groups-path", f"{state_dir}/island-checkpoint.pt", "--event-tape", f"{state_dir}/events.jsonl", "--audit-dir", f"{state_dir}/audit", *_provenance_args(str(task["reward"]))] + # The launcher pins the island's AlgorithmSpec hash (--rl-expected-algorithm-sha256); + # the trainer-only island refuses to start without it (Modal run-5). + island_argv += ["--rl-expected-algorithm-sha256", expected_algorithm_sha256(island_argv)] try: - model_args = glm5_next.model_argv(num_layers=int(model["num_layers"]), transformer_impl=impl) + model_args = glm5_next.model_argv( + num_layers=int(model["num_layers"]), + first_k_dense_replace=int(model.get("first_k_dense_replace", glm5_next.FIRST_K_DENSE_REPLACE)), + num_experts=int(model.get("num_experts", glm5_next.NUM_EXPERTS)), transformer_impl=impl) except ValueError as error: raise Stage1Error(f"plan: trainer.miles.transformer_impl: {error}") from None + torch_dist = model["bf16_source"].get("torch_dist") + convert_argv = None if not torch_dist else [ + "torchrun", "--nproc-per-node", str(per_node), f"{miles['workdir']}/tools/convert_hf_to_torch_dist.py", + # model arguments only: the converter's parser has no run flags (--qkv-format; Modal run-8) + *glm5_next.model_argv_without_run_flags(model_args), "--hf-checkpoint", model_dir, "--save", str(torch_dist), + "--tensor-model-parallel-size", str(layout["tp"]), "--pipeline-model-parallel-size", str(layout["pp"]), + "--expert-model-parallel-size", str(min(layout["ep"], per_node))] out.trainer = { "backend": backend, "nodes": [n.name for n in trainers], "cards": cards, + # Miles trains glm5_next from a Megatron torch_dist checkpoint (raw mode); converted once + # on trainer node 0 (one node's cards; not run for the 45-layer model) + "torch_dist": torch_dist, "convert_argv": convert_argv, + # the converter imports megatron.post_training, which the editable megatron_core + # install does not expose; Miles' own converter call puts the checkout first (run-9) + "convert_env": {"PYTHONPATH": f"{miles.get('megatron_path', '/root/Megatron-LM')}:{miles['workdir']}"}, "image": images.get("trainer"), "env": trainer_env, "miles": {"commit": miles["commit"], "megatron": miles["megatron"], "mbridge": miles["mbridge"], "layout": layout, "memory_estimate": megatron_memory(model, miles, layout, trainers, sampling, hf_config, @@ -490,7 +573,11 @@ def expand(plan: Stage1Plan, *, hf_config: str | Path | None = None) -> Expanded # The release directory follows the template's scheme (quant.schemes); the # plan's release_dir is checked against it, not passed to the re-quantizer. - argv = requant_argv(store, str(model["sampler_initial"]["local_dir"]), python="python3") + # quant.frozen_pattern: tensors the trainer does not export (read from the bf16 source) + frozen = quant.get("frozen_pattern") + argv = requant_argv(store, str(model["sampler_initial"]["local_dir"]), python="python3", + frozen_base=str(model["bf16_source"]["local_dir"]) if frozen else None, + frozen_pattern=str(frozen) if frozen else None) out.requant = {"node": trainers[0].name if trainers else None, "scheme": quant["scheme"], "format": quant["format"], "argv": argv, "loop_argv": [*argv, "--loop"]} @@ -537,7 +624,10 @@ def sampler_spec(plan: Stage1Plan, node: Node, engine: Any) -> dict[str, Any]: "groups_per_round": int(task["rollout_batch_size"]), "format": str(task.get("prompt_format", "chat")), "trust_remote_code": True}, "reward": str(task["reward"]), "reward_kind": "miles", - "checks": {"expected_expert_matrices": int(quant["expected_expert_matrices"])}, + # sampler.checks: further PostLoadChecks overrides (e.g. require_fixed_answers for a + # cut-down rehearsal model); the expected expert count always comes from quant + "checks": {**dict(plan.section("sampler").get("checks") or {}), + "expected_expert_matrices": int(quant["expected_expert_matrices"])}, "event_tape": f"{state_dir}/sampler-{node.name}/events.jsonl", } @@ -582,6 +672,9 @@ def trainer_island_spec(plan: Stage1Plan, out: Expanded, trainers: Sequence[Node "if [ \"$SKYPILOT_NODE_RANK\" = 0 ]; then\n" f" ray start --head --port={RAY_PORT} --node-ip-address=\"$me\" --num-gpus=\"$SKYPILOT_NUM_GPUS_PER_NODE\"\n" f" {join}\n" + # the island forwards its own environment (sample buffer, store) to the Ray + # actors only when RAY_ADDRESS names the cluster (entry.connect_island_ray) + f" export RAY_ADDRESS=\"$me:{RAY_PORT}\"\n" f" exec {island}\n" "fi\n" f"exec ray start --address=\"$head:{RAY_PORT}\" --node-ip-address=\"$me\" " @@ -645,6 +738,14 @@ def _steps(plan: Stage1Plan, out: Expanded, trainers: Sequence[Node], samplers: f"ossutil cp -r -u {shlex.quote(src['staged'])} {shlex.quote(src['local_dir'])}/\n" f"test -f {shlex.quote(src['local_dir'])}/model.safetensors.index.json\n", f"{key} weights on the node's /data (revision {src['revision']})")) + if out.trainer.get("convert_argv") and trainers: + node, dst = trainers[0], str(out.trainer["torch_dist"]) + convert = _docker(plan, node, "convert", out.trainer["image"], out.trainer["convert_argv"], + {**out.trainer["env"], **out.trainer["convert_env"]}) + steps.append(Step(f"convert-torch-dist-{node.name}", node.name, + f"set -eu\nif [ \"$(cat {shlex.quote(dst)}/latest_checkpointed_iteration.txt 2>/dev/null)\" " + f"!= release ]; then\n{convert}\nfi\n", + "bf16 source -> Megatron torch_dist (Miles converter; the island's --megatron-ref-load)")) for run in out.sampler["engines"]: node = next(n for n in samplers if n.name == run["node"]) spec_json = json.dumps(run["spec"], indent=1, sort_keys=True) @@ -729,7 +830,8 @@ def check(plan: Stage1Plan, out: Expanded) -> list[str]: if f"tp{int(sampler_cfg['tp'])}" != sampler_cfg["topology"]: errors.append(f"sampler topology {sampler_cfg['topology']} does not match TP {sampler_cfg['tp']}") image = str(plan.section("images").get("sampler") or "") - if not image.endswith(f":v{sampler_cfg['engine_version']}"): + # a tag must name the engine version; a digest pins the image itself (the version is recorded in the plan) + if not image.endswith(f":v{sampler_cfg['engine_version']}") and "@sha256:" not in image: errors.append(f"sampler image {image!r} does not carry engine version v{sampler_cfg['engine_version']}") if int(sampler_cfg["max_model_len"]) < int(sampling["max_prompt_len"]) + int(sampling["max_tokens"]): errors.append("sampler max_model_len is shorter than max_prompt_len + max_tokens") @@ -749,8 +851,13 @@ def check(plan: Stage1Plan, out: Expanded) -> list[str]: # Store. store = out.store["run"] - if not store.startswith("oss://"): + backend = str(plan.doc.get("store_backend") or "oss") + if backend not in STORE_BACKENDS: + errors.append(f"store_backend {backend!r} is not one of {sorted(STORE_BACKENDS)}") + elif backend == "oss" and not store.startswith("oss://"): errors.append(f"store {store!r} is not an oss:// URI (stage 1: Alibaba OSS)") + elif backend == "modal-volume" and not store.startswith("/"): + errors.append(f"store {store!r}: a modal-volume store is the absolute mount path in the containers") if not store.endswith("/" + plan.run_id): errors.append("store run directory does not end with the run id") bucket = store.split("/")[2] if store.startswith("oss://") else None @@ -830,6 +937,9 @@ def gaps(plan: Stage1Plan, out: Expanded) -> list[str]: if gap := _module_gap(module, what): found.append(gap) if out.trainer["backend"] == "miles": + if not out.trainer.get("torch_dist"): + found.append("model.bf16_source.torch_dist is unset: Miles trains glm5_next only from a " + "Megatron torch_dist checkpoint (Megatron-Bridge cannot load glm5_next)") if gap := _check_island_argv(out.trainer["miles"]["island_argv"]): found.append(gap) recipes = _miles_recipes() @@ -858,7 +968,7 @@ def _full_parameter_layout_gap(plan: Stage1Plan, out: Expanded) -> str | None: problem = full_parameter_layout_problem( data_parallel=out.trainer["miles"]["layout"]["dp"], actor_num_nodes=int(argv[argv.index("--actor-num-nodes") + 1]), - rollout_gpus=int(argv[argv.index("--rollout-num-gpus") + 1]), + rollout_gpus=(int(argv[argv.index("--rollout-num-gpus") + 1]) if "--rollout-num-gpus" in argv else None), rollout_gpus_per_engine=1, # the island's default; it gives no engine sample_role=argv[argv.index("--rl-sample-role") + 1]) if problem is None: diff --git a/yeto/rl/token_diversity_reward.py b/yeto/rl/token_diversity_reward.py new file mode 100644 index 00000000..8547b677 --- /dev/null +++ b/yeto/rl/token_diversity_reward.py @@ -0,0 +1,41 @@ +"""Dense token-diversity reward for rehearsals on models that cannot solve the task. + +A test reward, not a product reward (S25 T-G1C). The word-diversity reward +(:mod:`yeto.rl.word_diversity_reward`) was gamed in T-G1B: after two steps the +cut model wrote one 512-token "word" without spaces, every sample scored 1.0, +the group variance was 0 and round 2 had no gradient. This reward counts token +ids, so one repeated token scores low however it is spelled: + + score = distinct response token ids / response tokens (0.0 when empty) + +The response ids are ``sample.tokens[-response_length:]`` (Miles ``Sample`` +and the engine sampler's reward sample both carry ``tokens`` and +``response_length``). Range (0, 1]; the maximum needs every response token to +differ, which a sampled 512-token response at temperature 1 does not reach. +Launch with + + --reward-function yeto.rl.token_diversity_reward:score +""" + +from __future__ import annotations + +from collections.abc import Sequence +from typing import Any + +__all__ = ["diversity", "response_ids", "score"] + + +def diversity(ids: Sequence[int]) -> float: + return len(set(ids)) / len(ids) if ids else 0.0 + + +def response_ids(sample: Any) -> list[int]: + tokens = list(getattr(sample, "tokens", None) or ()) + length = int(getattr(sample, "response_length", 0) or 0) + if length > len(tokens): + raise ValueError(f"response_length {length} > {len(tokens)} tokens") + return tokens[len(tokens) - length:] if length else [] + + +async def score(args, sample, **kwargs) -> float: + return diversity(response_ids(sample)) diff --git a/yeto/rl/word_diversity_reward.py b/yeto/rl/word_diversity_reward.py new file mode 100644 index 00000000..9eac4b51 --- /dev/null +++ b/yeto/rl/word_diversity_reward.py @@ -0,0 +1,28 @@ +"""Dense word-diversity reward for rehearsals on models that cannot solve the task. + +A test reward, not a product reward. The cut 4-layer GLM5.3-Flash rehearsal model +(S25 T-G1) never answers GSM8K and every response runs to ``max_tokens``: the +GSM8K reward and the length reward are constant within a group, so GRPO +advantages, the loss gradient and grad_norm are all 0 and the weights never move. +This reward depends on the text itself, so samples of one prompt differ even when +all of them are truncated: + + score = distinct whitespace words / words (0.0 for an empty response) + +Range [0, 1]; a response that repeats one word scores low. Launch with + + --reward-function yeto.rl.word_diversity_reward:score +""" + +from __future__ import annotations + +__all__ = ["diversity", "score"] + + +def diversity(text: str) -> float: + words = text.split() + return len(set(words)) / len(words) if words else 0.0 + + +async def score(args, sample, **kwargs) -> float: + return diversity(sample.response or "")