From 0f0f2c70e42e73e4244da25a6e7e2b28a47e19cd Mon Sep 17 00:00:00 2001 From: AlexEisie <1987460907@qq.com> Date: Thu, 30 Jul 2026 09:45:33 +0800 Subject: [PATCH 1/6] feat(rl): add pinned Miles LoRA federation Implement the Miles RL v0 path around radixark/miles at dfc66ff38752bfa2c5d325e0037ebc4b537c06de: provision multi-GPU islands, run Miles rollout and GRPO, export complete canonical fp32 LoRA updates, and commit exact-base equal-weight fixed-roster averages through the existing syncer protocol. Add optimizer-state reset with scheduler progress preservation, completed-group recovery and oversampling, custom generation/session-server/TITO forwarding, strict failure handling, standard PEFT export, metrics, documentation, and focused launcher/runtime/Rust coverage. The implementation was validated with dense DP=8, two DP=4 islands, MoE EP=8, direct Miles parity, tool-use sessions, fault injection, syncer restart, and 20 consecutive real-model merges on eight A100 GPUs. Differences from INIT: 1. Use process-local actor/provider adapters and the pinned non-FT private _broadcast path against a clean upstream Miles checkout, instead of maintaining the proposed Miles thin branch and train-loop hook. 2. Keep the authoritative syncer checkpoint on retained local disk; the INIT cross-VM/disk durable checkpoint mount is not implemented. 3. Emit the planned RL and sync metrics to JSONL, but do not enable the INIT dashboard integration. --- README.md | 2 + docs/MILES_RL.md | 558 +++++++++++++++++ pyproject.toml | 1 + syncer/src/main.rs | 16 + syncer/src/server.rs | 344 ++++++++-- syncer/src/state.rs | 73 +++ tests/test_rl_core.py | 1074 ++++++++++++++++++++++++++++++++ tests/test_rl_export.py | 171 +++++ tests/test_rl_integration.py | 576 +++++++++++++++++ tests/test_rl_launcher.py | 719 +++++++++++++++++++++ yeto/cli.py | 39 ++ yeto/export.py | 22 +- yeto/launcher.py | 478 +++++++++++++- yeto/rl/__init__.py | 9 + yeto/rl/bridge.py | 327 ++++++++++ yeto/rl/core.py | 296 +++++++++ yeto/rl/export.py | 223 +++++++ yeto/rl/learner.py | 466 ++++++++++++++ yeto/rl/miles.py | 1136 ++++++++++++++++++++++++++++++++++ 19 files changed, 6476 insertions(+), 54 deletions(-) create mode 100644 docs/MILES_RL.md create mode 100644 tests/test_rl_core.py create mode 100644 tests/test_rl_export.py create mode 100644 tests/test_rl_integration.py create mode 100644 tests/test_rl_launcher.py create mode 100644 yeto/rl/__init__.py create mode 100644 yeto/rl/bridge.py create mode 100644 yeto/rl/core.py create mode 100644 yeto/rl/export.py create mode 100644 yeto/rl/learner.py create mode 100644 yeto/rl/miles.py diff --git a/README.md b/README.md index 558efd7..46e9a6e 100644 --- a/README.md +++ b/README.md @@ -196,6 +196,8 @@ per-seed results for the completed Qwen3.6, LTX-Video, and Wan2.2 benchmarks. [docs/PROTOCOL.md](docs/PROTOCOL.md) — the learner↔syncer wire protocol. [docs/MEGATRON.md](docs/MEGATRON.md) — the Megatron-Core island backend (EP for 1T-class MoE). +[docs/MILES_RL.md](docs/MILES_RL.md) — fixed-roster Miles RL across causal-LM +LoRA islands: rollout/training boundaries, recovery, export, and limitations. [docs/MLX.md](docs/MLX.md) — the Apple-silicon island backend: Macs as learner islands (`yeto launch --external-learners`, cross Mac↔NVIDIA runs). diff --git a/docs/MILES_RL.md b/docs/MILES_RL.md new file mode 100644 index 0000000..febc978 --- /dev/null +++ b/docs/MILES_RL.md @@ -0,0 +1,558 @@ +# Miles RL v0 + +Miles RL runs reinforcement learning in several independent +[Miles](https://github.com/radixark/miles) islands and uses Yeto to turn their +complete local LoRA updates into one committed global policy. Miles owns the +rollout and GRPO loop inside each island; Yeto owns fleet lifecycle, the +cross-island synchronization boundary, recovery, and the final adapter. + +This mode is not Decoupled DiLoCo. It is synchronous, fixed-roster LoRA +FedAvg: every global round waits for one exact-base result from every logical +island and averages them equally. + +> **Status:** the core path is implemented and has passed real multi-GPU dense, +> MoE EP, two-island averaging, recovery, long-task, and 20-merge validation on +> eight A100 GPUs. The 24-hour soak and the operational gaps listed in +> [Validation status](#validation-status) remain outside that evidence. + +## Why Miles is the RL runtime + +An RL learner needs much more than a different loss function. It must keep a +rollout server and trainer coherent while it generates grouped trajectories, +scores responses, computes advantages, and updates the policy. Miles owns +those contracts around colocated SGLang and Megatron-Core. Yeto forwards +Miles' custom-generate and session-server entry points rather than implementing +a second agent or environment runtime. + +Yeto therefore does not reproduce rollout, reward, or GRPO machinery in its +ordinary causal-LM learner. It adds one boundary around a pinned Miles local +round: + +```text + committed global LoRA + │ + ┌─────────────┴─────────────┐ + ▼ ▼ + Miles island 0 Miles island 1 … M islands + SGLang rollout SGLang rollout + reward + GRPO reward + GRPO + Megatron step Megatron step + │ complete local LoRA │ + └─────────────┬───────────────┘ + ▼ + Yeto Rust syncer + fixed roster · f32 mean + checkpoint before broadcast +``` + +This division keeps Miles responsible for policy optimization and Yeto +responsible for distributed policy agreement. + +## Supported contract + +v0 deliberately supports one narrow model-parallel contract while allowing a +Miles island to use its full multi-node GPU allocation: + +| dimension | supported value | +| --- | --- | +| model | causal language model | +| tuning | LoRA | +| local trainer | pinned Miles with Megatron-Core | +| rollout | colocated SGLang | +| island size | one or more nodes, one or more GPUs per node | +| parallelism | TP = PP = CP = 1; dense models use EP = 1; MoE EP is 1 or a divisor of the island world size; Miles supplies the DP layout | +| global state | one complete LoRA fragment | +| wire and merge math | f32, equal-weight AVG | +| membership | fixed logical island IDs `0..M-1`, `M >= 1` | + +Model support is determined by the runtime tensor contract, not a Yeto model +family list. A model is usable when the pinned Miles/Megatron-Bridge stack can +load it, create LoRA, convert every trainable adapter tensor to standard PEFT +names and shapes, and apply the converted values without changing them. +Incompatibility fails during initialization instead of selecting a +model-specific fallback. + +Every canonical policy carries the immutable base-model revision, a hash of +the effective LoRA configuration, the semantic tensor-layout hash, and its +policy version. Tensor specs include name, shape, fp32 dtype, and numel. The +runtime checks the complete identity between base, local, applied, and +broadcast states. The syncer checkpoint persists the layout hash, but not the +base revision or LoRA-config hash; export reconstructs those values from its +required explicit arguments and rejects a layout mismatch. +With EP>1, v0 therefore permits only replicated adapter targets (the normal +MoE `auto`/`attention` path), not LoRA on sharded expert weights. +The replicated path disables Megatron's distributed optimizer so every +trainer rank retains the complete fp32 LoRA masters and optimizer state needed +by the export/apply contract. + +The RL path consumes `--lora-r` and `--lora-targets`. Its effective PEFT +configuration fixes alpha to the rank, dropout to zero, and bias to `none`; +the public RL launcher does not expose a separate alpha setting. + +Not supported in v0: + +- diffusion RL or full-parameter RL; +- stale, asynchronous, or partial-quorum global updates; +- multi-fragment synchronization, server momentum, RDA, HeLoCo, or blending; +- TP>1, PP>1, CP>1, or expert-sharded LoRA adapters; +- cross-island optimizer-state merging; +- Miles' experimental fault-tolerant trainer path; +- recovery of an unfinished trajectory or an unmerged local LoRA; +- RL-specific manifests, attestations, terminal markers, or wire messages. + +The ordinary Yeto SFT and diffusion modes remain separate. Passing +`--training-mode rl` with a diffusion model is rejected. + +## Current integration seams + +The Miles checkout remains a clean detached checkout. The current adapter +installs Yeto export, apply, and optimizer-step methods on the pinned +`MegatronTrainRayActor` at process startup and invokes them through that +commit's non-FT `RayTrainGroup._broadcast` path. Yeto drives Miles rollout and +training primitives directly so the island-local result is not published to +SGLang before the global merge. There is currently no maintained Miles patch +or upstream train-loop synchronization hook. This private compatibility seam +is supported only for the pinned commit and must be revalidated if Miles is +changed. + +The launcher selects strict RL behavior through the existing syncer's general +controls: `--max-base-lag 0`, `--learner-weight equal`, `--quorum M`, and +`--grace-ms 0`. It also supplies the one-fragment, unthrottled AVG settings +listed above and uses the existing `--checkpoint-every 1` control for a +durable cut before each broadcast. There is no RL-specific syncer mode or +wire message. Learners +send `c_steps=1` and `c_tokens=1`; `--learner-weight equal` independently +defines their global contribution. Protocol v4 is unchanged, and omitting the +two exact-base/equal-weight controls preserves SFT behavior. + +## Global-round semantics + +The RL flags describe the work at the Miles boundary: + +- `M`: number of `--gpu` entries and therefore logical islands; +- `N`: Yeto's existing `--total-steps`; +- `G`: `--rollout-batch-size`, the complete GRPO groups per island round; +- `K`: `--n-samples-per-prompt`, the trajectories per group; +- local work: `--local-rl-rounds-per-sync 1` in v0. + +`N`, `G`, and `K` must be positive. `G × K` must be divisible by every +island's Miles data-parallel size. `--over-sampling-batch-size` defaults to +`G` and may be raised to generate extra complete groups; it cannot be smaller +than `G`. Complete extras remain in the same-version queue after the selected +`G` groups are consumed. + +The Miles argument mapping enables `--balance-data`. Multi-rank islands use +Miles' sequence-length partitioner to keep the total token count similar +across DP ranks while preserving equal sample counts. Yeto requires the +sample count to divide evenly across those ranks. + +At committed version `v`, each island follows the same sequence: + +1. Apply the complete global LoRA `theta_v` to the Megatron trainer and + SGLang, then mark rollout policy version `v` active. +2. Generate exactly `G` groups of `K` terminal trajectories (completed or + truncated at the configured limit). Every recorded rollout weight version + must be `v`. +3. Run the one configured Miles GRPO training cycle without publishing the local + result to SGLang. +4. Export the complete local LoRA `theta_i_v` and send + `theta_i_v - theta_v` with base version `v`. +5. Wait for committed version `v + 1` before starting another local round. + +The syncer accepts at most one result per logical island and waits for all +`M` results: + +```text +theta_(v+1) = theta_v + mean(theta_i_v - theta_v) + = mean(theta_i_v), i = 0 .. M-1 +``` + +All merge inputs have unit weight. Prompt counts, response lengths, and token +counts do not change an island's global weight. + +Applying a new global policy removes the existing optimizer state for LoRA +parameters, including their moments and parameter step, then replaces its fp32 +master parameters. It does not rebuild the optimizer or LR scheduler. Applying +committed policy `v` aligns the existing scheduler to +`v * num_steps_per_rollout * global_batch_size`: a replacement advances its +fresh scheduler to the committed progress, an in-process scheduler is already +there, and a scheduler ahead of the committed policy is rejected. This avoids +another warmup while preventing local Adam history from leaking across an +averaging boundary. + +There is no fleet-wide "every island has applied" barrier. Each island has an +exact-base gate of its own, while the next merge still waits for the entire +fixed roster. A faster island may begin first, but it cannot commit without all +other islands at that same base version. + +## Running a fleet + +The launcher creates one Miles island for every `--gpu` entry. An entry may +describe multiple nodes and multiple GPUs per node. The launcher starts one +Ray head and joins the remaining island nodes as workers. A single entry is a +supported parity path; multiple entries enable fixed-roster averaging. +External learner slots are not supported. + +Use a Miles-compatible learner image. The image used for v0 validation is: + +```text +docker:radixark/miles@sha256:95b3afa9ee4313f5633e6ed3779c8276353cc8e24a2462e4f54ec0d5978fbae7 +``` + +The Miles source itself is independently pinned to: + +```text +https://github.com/radixark/miles +dfc66ff38752bfa2c5d325e0037ebc4b537c06de +``` + +The launcher checks out that commit as a detached HEAD and installs the +project's pinned PEFT version. At learner startup Yeto verifies the repository +origin, commit, clean worktree, and imported package path. Runtime +adaptation uses the compatibility seam described above without modifying the +checkout. + +An illustrative two-island run is: + +```bash +yeto launch \ + --training-mode rl \ + --gpu aws:8xa100@us-east-1,aws:8xa100@us-west-2 \ + --rl-runtime miles \ + --rl-image docker:radixark/miles@sha256:95b3afa9ee4313f5633e6ed3779c8276353cc8e24a2462e4f54ec0d5978fbae7 \ + --model HuggingFaceTB/SmolLM2-135M-Instruct \ + --model-revision 12fd25f77366fa6b3b4b768ec3050bf629380bac \ + --data org/long-task-prompts \ + --data-revision 0123456789abcdef0123456789abcdef01234567 \ + --tuning lora \ + --lora-r 4 \ + --lora-targets attention \ + --total-steps 2 \ + --advantage-estimator grpo \ + --rollout-batch-size 16 \ + --over-sampling-batch-size 24 \ + --n-samples-per-prompt 4 \ + --rollout-max-response-len 512 \ + --local-rl-rounds-per-sync 1 \ + --rl-sync-preset strict-avg \ + --rl-policy-version strict \ + --rl-completed-groups-path ~/yeto-rl/island-checkpoint.pt \ + --reward-function project.rewards:score \ + --seq-len 2048 \ + --inner-lr 1e-4 \ + --seed 321 \ + --trust-remote-code \ + --controller local \ + --output ./rl-output +``` + +`--trust-remote-code` is required because the pinned Miles stack enables +remote-code loading in its internal model paths. Continue to pin model and +dataset revisions and enable it only for repositories you trust. + +The example uses a local controller so `./rl-output` is fetched to the +submitting machine, which must remain available for the run. With the default +head controller, use a remote `--output` URI for automatic delivery or +retrieve the retained checkpoint from the head. + +RL uses the public CLI above. `--total-steps` selects the global round count. +Generic DiLoCo controls such as `--fragments`, `--quorum`, `--pipeline`, +`--sync-interval-steps`, `--outer-lr`, `--outer-momentum`, and +`--wire-dtype` are not RL algorithm knobs; the launcher fixes their internal +values to the contract above. + +`--rl-runtime`, `--advantage-estimator`, `--rl-sync-preset`, and +`--rl-policy-version` currently accept only `miles`, `grpo`, `strict-avg`, and +`strict`, respectively. Without `--experimental-rl-sync`, the strict preset +normalizes all generic sync values above. With the flag, the launcher preserves +and forwards working syncer controls such as quorum, grace, pacing, correction, +and outer LR/momentum. The bridge still requires `--fragments 1`, +`--pipeline 1`, `--merge-alpha 0`, and `--wire-dtype f32`; unsupported values +are rejected rather than accepted as ineffective overrides. Exact-base and +equal learner weighting remain explicit. + +For RL, the launcher raises the effective `--seq-len` to at least +`--rollout-max-response-len`; this keeps the RL response default from being +silently clamped by the smaller generic SFT sequence default. Miles deducts the +tokenized prompt from that total rollout/trainer context, so the response value +is still a cap rather than a guaranteed length. At learner startup the actual +Megatron-Bridge provider must advertise a model context limit at least as large +as the effective sequence length. + +### Prompt data + +Each row must provide `messages`, or a string `prompt`/`input` that Yeto can +turn into a user message. `label`, `metadata`, and `tools` are preserved for +Miles and the reward implementation. + +```json +{"messages":[{"role":"user","content":"Give a short proof."}],"label":"proof"} +``` + +RL v0 accepts revision-pinned Hugging Face dataset references. The prepared +prompt file is private to each island. + +Without extra flags the learner uses Miles' default SGLang generation path, in +which one completion produces one trajectory. For environment-driven, +multi-turn or tool-use trajectories, pass an importable Miles generate callable +through `--custom-generate-function-path package.module.function`. Yeto also +forwards `--use-session-server`, optional `--session-server-ip`, and one port or +port range through `--session-server-port`. For a model that needs one of +Miles' model-specific incremental tokenizers, `--tito-model` is forwarded +unchanged and requires the session server. Yeto deliberately does not infer +this value from a model name. Miles continues to own session assembly, +tool/environment calls, terminal status, and weight-version records; Yeto does +not define a second trajectory format. Dataset `tools` and metadata are +preserved for that callable. + +### Reward callable + +`--reward-function` uses `package.module:function` syntax. The module must be +importable in the learner workdir. Its callable follows the pinned Miles +custom-reward API; a minimal implementation is: + +```python +async def score(args, sample, **kwargs) -> float: + del args, kwargs + return 1.0 if "expected phrase" in sample.response else 0.0 +``` + +GRPO needs reward variation within a group to produce a useful advantage. +A run can complete optimizer steps with constant rewards while learning +nothing, so inspect raw reward and advantage metrics during a shakedown. + +Before provisioning, the launcher hashes the selected reward module source. +The learner verifies that digest before importing it. The source must be +inside the Yeto workdir synchronized by SkyPilot. + +## Checkpoints, completion, and export + +Version 0 and every successful global merge are atomically written to the +syncer's checkpoint before the new LoRA is broadcast. That file is the only +authoritative global RL state and includes the canonical layout hash. + +Each Miles island also atomically stores its compatibility configuration, +local round, current policy version, rollout/reward statistics, and complete +unused group queue at `--rl-completed-groups-path`. The compatibility fields +cover model and dataset identifiers/revisions, topology, LoRA identity, +learning rate, sequence length, seed, group and oversampling sizes, optimizer +steps, reward digest, response limit, custom-generate/session mode, and the +selected TITO model. A process restart restores only complete groups whose +policy version and all persisted compatibility fields still match. For Spot RL tasks, +the launcher mounts that file's parent directory from a SkyPilot storage named +for `cluster_prefix + logical island ID`, with reconstruction sync enabled. +The stable per-island name survives a replacement VM without sharing queues +between islands; the non-persistent storage is removed when the task is finally +torn down. Spot paths must name a file inside an absolute or `~/` subdirectory. +On-demand tasks keep the ordinary local path. Groups selected for the current +training batch are removed from the queue; complete oversampling groups that +were not selected remain reusable only while that same global policy is +current. Unfinished groups, cross-version groups, and unmerged local LoRA +values are discarded. + +With the local-controller example above, the launcher fetches the checkpoint +as: + +```text +./rl-output/yeto-state.ckpt +``` + +The default head controller first retains the same file at +`~/yeto-output/yeto-state.ckpt` on the head, then delivers it when `--output` +is a remote URI. + +When global version `N` is committed, the syncer sends its existing final +fragment protocol to every logical learner. Each learner applies the final +LoRA and acknowledges the cut before exiting normally. + +Export the checkpoint with the same base model revision, rank, and target +selection used for training: + +```bash +yeto-rl-export \ + --checkpoint ./rl-output/yeto-state.ckpt \ + --model HuggingFaceTB/SmolLM2-135M-Instruct \ + --model-revision 12fd25f77366fa6b3b4b768ec3050bf629380bac \ + --lora-r 4 \ + --lora-targets attention \ + --output-dir ./adapter +``` + +The result contains standard `adapter_model.safetensors` and +`adapter_config.json` files and can be loaded with ordinary PEFT tooling. The +exporter reconstructs the expected tensor contract from the explicit model +arguments and requires its layout hash to match the checkpoint. It does not +infer configuration or roster membership from the syncer ledger. + +## Failure and recovery + +Recovery always starts from the most recent committed global checkpoint: + +| failure | behavior | +| --- | --- | +| learner exits before its push is accepted | launcher restarts the same logical ID; it reapplies the committed policy and recomputes the round | +| learner exits after its push is accepted | syncer retains that logical learner's result and waits for the missing IDs | +| learner cannot recover | the run fails; the roster is never reduced | +| syncer process exits before checkpoint commit, with its disk retained | partial results are discarded; restart resumes the previous version and learners redo the round | +| syncer process exits after checkpoint commit but before broadcast, with its disk retained | restart loads and broadcasts the newly committed version | +| syncer VM or disk is lost | the current launcher has no durable mount for the global checkpoint, so automatic recovery is not available | +| learner exits while applying a global policy | replacement reapplies the complete committed policy | + +A dead syncer connection makes an island process exit rather than silently +continue local work. The launcher can restart the syncer through its existing +recovery path and restart affected logical learners individually; it does not +need to restart every healthy roster member. Duplicate computation is allowed, +but duplicate merge is not. + +Each island writes policy-apply, optimizer-reset, and `LocalRoundStats` JSONL; +the syncer writes roster, base-version, layout, merge, and responder metrics to +`~/yeto-output/yeto-tape.jsonl`. Miles group-task completion supplies peak +active groups, cancellations, and duration percentiles; +`Sample.non_generation_time` supplies tool wait; grouped raw rewards supply the +zero-variance ratio; and Miles' rank-zero train log supplies KL, ESS, and clip +fraction. The bridge records the actual protocol payload bytes, while each +canonical global apply records its policy hash. +The launcher does not currently enable a Miles or Yeto dashboard for these +records. + +Mixed-version groups, rejected stale updates, layout mismatch, non-finite +deltas, or a policy hash mismatch after apply are deterministic strict +failures. The affected learner or syncer exits nonzero, and the fleet +controller terminates the whole run instead of relaunching the same +deterministic violation. Ordinary process, network, or spot failures remain +recoverable through the committed-checkpoint paths above. + +## Runtime compatibility and diagnostics + +The validated Miles/Transformer Engine stack cannot reliably use its +FlashAttention CUTE GQA kernel on A100. RL islands therefore select unfused +attention before Transformer Engine imports and propagate that choice into +the actual Megatron provider. This is a runtime-wide choice, not a model +family workaround. + +Common startup and progress failures have distinct meanings: + +- **PEFT import failure:** the learner did not run the current Miles setup, + which installs the pinned `peft==0.20.0`. +- **`Operation creation failed` under `flash_attn/cute/pack_gqa.py`:** the + process did not receive the RL launch environment or provider backend. +- **Miles revision/origin/dirty-tree error:** the runtime is not using the + supported checkout; do not bypass this check or patch that tree in place. +- **PEFT/Megatron mapping mismatch:** the model is outside the currently + supported tensor contract. Extend the generic upstream conversion path + rather than adding a model-name branch in Yeto. +- **syncer waits with an incomplete roster:** a logical island is absent or + still recovering. Fixed-roster mode never lowers quorum to make progress. +- **optimizer steps but negligible LoRA change:** first inspect reward + variance, advantages, and the configured LR schedule. + +## Validation status + +The current source has automated coverage for multi-node task construction, +multi-rank actor results, DP rollout-shard collection, EP validation, +single-island and multi-island sync, canonical identity, completed-group +recovery, strict failures, provenance, checkpoint export, and the unchanged +SFT/diffusion defaults. The 2026-07-30 regression passed all 73 focused RL +tests, the full Python suite (`766 passed, 4 skipped`), `cargo fmt --check`, and +all 58 Rust tests. + +Real validation ran on one GCP Spot VM with eight NVIDIA A100-SXM4-40GB GPUs. +It used the pinned Miles commit, PEFT 0.20.0, immutable model and dataset +revisions, real model generations, real rewards, and production Yeto learner +and Rust syncer paths. Observation hooks only captured tensors, tokens, +metrics, and fault windows. + +- A one-island Qwen3-4B DP=8 run completed two global rounds. Each round + trained on 16 trajectories and 8192 action tokens with nonconstant rewards. + Local delta norms were `0.22248` and `0.10623`; both M=1 merges matched the + saved local policy within `3.64e-12`. +- Two concurrent Qwen3-4B DP=4 islands used different prompt-token hashes. + Their local delta norms were `0.223218` and `0.222651`; the committed f32 + policy exactly matched the offline mean, and both islands applied identical + initial and final policies. +- `allenai/OLMoE-1B-7B-0125-Instruct` ran one EP=8 island with replicated + attention LoRA, 16 trajectories, 4096 action tokens, nonconstant rewards, + and a `0.10110` local delta. The merge error was zero. Standard PEFT loaded + the exported adapter, produced finite logits with a nonzero adapter effect, + and completed real generation. +- A direct pinned-Miles round and the production Yeto+Miles M=1 path produced + identical sampled tokens and rewards; their LoRA tensors had max error + `0.0`. The exact trainer-to-SGLang LoRA checksum checker passed throughout. + Trainer-versus-rollout KL stayed around `6e-4` to `1.3e-3` before and after + global applies rather than showing a material token-path mismatch. +- A 20-round Qwen run committed versions `1..20` exactly once, completing 160 + trajectories and 20,480 action tokens. Every round had a nonzero finite + local delta, no mixed-version group was observed, and current-versus-rollout + KL remained between `0.000329` and `0.001067`. +- Learners were killed during rollout, after local train, before push, after + push, before broadcast, and after global apply. Every replacement completed, + and each case produced one committed step. A separate syncer restart resumed + the committed version and completed the next round. +- Completed-group recovery retained one real four-trajectory oversampling + group across learner replacement and selected that group after restart. A + separately replayed real trained delta was accepted once, while a stale + exact-base update caused the strict connection to close. +- A Qwen3-4B session-server run completed two rounds and 16 real environment + tasks. Ten trajectories made actual calculator calls and consumed their + returned values; the second round had nonconstant rewards and a `0.10059` + local update. Fourteen session traces had exact TITO reconstruction. The two + remaining traces reached the 512-token response cap before a terminal token, + were correctly marked truncated, and accounted for the reported 25% TITO + structural mismatch in that round. Trainer-versus-rollout absolute logprob + error remained below `0.01`. + +The following boundaries remain unvalidated or intentionally excluded: + +- the requested 24-hour soak was not run; the 20 consecutive merges are the + bounded-duration stability evidence; +- no physical multi-node island or end-to-end SkyPilot provisioning and Spot + VM replacement was exercised, although multi-node task construction is + covered automatically; +- the syncer checkpoint still has no durable mount for syncer VM or disk loss; +- metrics remain JSONL-only and the launcher enables no dashboard. + +The runtime-injection Miles seam described above remains an implementation +constraint, not a stable upstream interface. + +## Extending v0 safely + +The useful extension boundary is the observable contract, not a model-name +matrix. Preserve these invariants when changing the implementation: + +1. The synchronized state is every trainable LoRA tensor, in standard PEFT + naming and deterministic order, represented as contiguous CPU f32. +2. Every trajectory group records one actual global rollout version. +3. A local result is accepted only for its exact committed base and at most + once per logical island. +4. A merge waits for all fixed logical IDs and uses equal weights. +5. The checkpoint is durable before any corresponding broadcast. +6. Applying global LoRA resets local optimizer history without rebuilding the + optimizer or resetting scheduler progress. +7. Local trainer weights are not published to SGLang between global commits. + +Adding a model should normally require no Yeto change: improve or select the +appropriate Miles/Megatron-Bridge/PEFT conversion and let the existing tensor +checks decide compatibility. Adding TP/PP, expert-sharded LoRA, multiple fragments, a +different aggregation algorithm, Miles FT actors, or Diffusion RL changes the +contract and requires a separate design rather than a compatibility branch. + +For implementation navigation: + +| area | responsibility | +| --- | --- | +| `yeto/rl/core.py` | canonical LoRA state and the single AVG layout | +| `yeto/rl/miles.py` | pinned Miles runtime, trainer/SGLang policy boundary | +| `yeto/rl/bridge.py` | exact-base island loop and protocol interaction | +| `yeto/rl/export.py` | committed checkpoint to PEFT adapter | +| `yeto/rl/learner.py` | island entry point and Miles argument mapping | +| `syncer/src/server.rs` | fixed roster and checkpoint-before-broadcast commit | + +Run the focused and full regressions before a GPU shakedown: + +```bash +python -m pytest -q \ + tests/test_rl_core.py tests/test_rl_export.py \ + tests/test_rl_integration.py tests/test_rl_launcher.py +python -m pytest -q +(cd syncer && cargo fmt --check && cargo test) +``` diff --git a/pyproject.toml b/pyproject.toml index ea036b9..797eced 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -40,6 +40,7 @@ yeto-export = "yeto.export:main" yeto-codex-traces = "yeto.codex_traces:main" yeto-diffusion-export = "yeto.diffusion.export:main" yeto-diffusion-sample = "yeto.diffusion.sample:main" +yeto-rl-export = "yeto.rl.export:main" [tool.setuptools.packages.find] include = ["yeto*"] diff --git a/syncer/src/main.rs b/syncer/src/main.rs index 44d5ff4..ca970ab 100644 --- a/syncer/src/main.rs +++ b/syncer/src/main.rs @@ -81,6 +81,13 @@ struct Args { /// JSONL event tape (one record per merge). #[arg(long)] event_tape: Option, + /// Maximum admitted lag between a round and a learner's base version. + /// Omitted means unbounded (the existing SFT behavior). + #[arg(long)] + max_base_lag: Option, + /// Learner contribution weighting used by AVG/RDA merges. + #[arg(long, default_value = "tokens2-over-steps")] + learner_weight: String, } fn main() -> anyhow::Result<()> { @@ -99,6 +106,13 @@ fn main() -> anyhow::Result<()> { "none" => false, other => anyhow::bail!("--delta-correction must be 'heloco' or 'none', got {other:?}"), }; + let learner_weight = match args.learner_weight.as_str() { + "tokens2-over-steps" => server::LearnerWeight::Tokens2OverSteps, + "equal" => server::LearnerWeight::Equal, + other => { + anyhow::bail!("--learner-weight must be 'tokens2-over-steps' or 'equal', got {other:?}") + } + }; let cfg = server::Config { port: args.port, learners: args.learners, @@ -121,6 +135,8 @@ fn main() -> anyhow::Result<()> { mark_final_checkpoint: args.mark_final_checkpoint, learner_budget_steps: args.learner_budget_steps, event_tape: args.event_tape, + max_base_lag: args.max_base_lag, + learner_weight, }; tokio::runtime::Builder::new_multi_thread() .enable_all() diff --git a/syncer/src/server.rs b/syncer/src/server.rs index 1e25309..6746470 100644 --- a/syncer/src/server.rs +++ b/syncer/src/server.rs @@ -27,6 +27,12 @@ const MAX_PARTIAL_MESSAGES: usize = 64; const WRITE_TIMEOUT: Duration = Duration::from_secs(180); const WRITER_QUEUE: usize = 128; +#[derive(Clone, Copy)] +pub enum LearnerWeight { + Tokens2OverSteps, + Equal, +} + #[derive(Clone)] pub struct Config { pub port: u16, @@ -80,6 +86,10 @@ pub struct Config { pub learner_budget_steps: Option, /// JSONL event tape: one record per merge. pub event_tape: Option, + /// Maximum admitted learner base-version lag; None preserves the + /// existing unbounded behavior. + pub max_base_lag: Option, + pub learner_weight: LearnerWeight, } struct OutFrame { @@ -103,6 +113,7 @@ struct Group { member: Member, dtype: u8, layout: Layout, + layout_fingerprint: [u8; 32], num_streams: u16, max_init_payload: u64, max_push_payload: u64, @@ -185,6 +196,10 @@ impl Group { } enum Event { + Fatal { + metric: &'static str, + message: String, + }, Hello { group: Arc, }, @@ -377,6 +392,7 @@ pub async fn run(cfg: Config) -> Result<()> { let accept_registry = registry.clone(); let accept_session = session.clone(); let expected_learners = cfg.learners; + let strict_layout = cfg.max_base_lag == Some(0); tokio::spawn(async move { loop { match listener.accept().await { @@ -385,8 +401,15 @@ pub async fn run(cfg: Config) -> Result<()> { let session = accept_session.clone(); let tx = event_tx.clone(); tokio::spawn(async move { - if let Err(e) = - handle_connection(stream, reg, session, expected_learners, tx).await + if let Err(e) = handle_connection( + stream, + reg, + session, + expected_learners, + strict_layout, + tx, + ) + .await { warn!(%peer, "connection ended: {e:#}"); } @@ -493,6 +516,7 @@ async fn handle_connection( registry: Registry, session: Session, expected_learners: u32, + strict_layout: bool, event_tx: mpsc::Sender, ) -> Result<()> { stream.set_nodelay(true)?; @@ -560,6 +584,15 @@ async fn handle_connection( }; if let Some(message) = mismatch { send_direct(&mut wr, MSG_ERROR, message.as_bytes()).await?; + if strict_layout { + event_tx + .send(Event::Fatal { + metric: "layout_hash_mismatch", + message: message.clone(), + }) + .await + .ok(); + } bail!(message); } let member = Member { @@ -574,6 +607,7 @@ async fn handle_connection( member, dtype, layout, + layout_fingerprint, num_streams, max_init_payload, max_push_payload, @@ -1025,6 +1059,9 @@ async fn scheduler( ) -> Result<()> { let mut state: Option = None; let mut budget_reports: HashSet = HashSet::new(); + let strict_exact = cfg.max_base_lag == Some(0); + let fixed_roster = strict_exact && cfg.quorum == cfg.learners && cfg.grace_ms == 0; + let checkpoint_each_round = strict_exact && cfg.checkpoint_every == 1; // Phase 1: wait until every fragment is initialized (via INIT_PARAMS or // a resumed checkpoint) and all expected learners have connected (late @@ -1043,6 +1080,10 @@ async fn scheduler( } } match events.recv().await.context("event channel closed")? { + Event::Fatal { metric, message } => { + append_strict_failure(cfg.event_tape.as_deref(), metric, &message); + bail!("RL strict failure {metric}: {message}"); + } Event::Hello { group } => { if state.is_none() { // Layout comes from the HELLO of the first learner. @@ -1050,7 +1091,27 @@ async fn scheduler( let mut st = new_state_for(&group, &cfg)?; if cfg.resume { if let Some(path) = cfg.checkpoint_path.as_ref().filter(|p| p.exists()) { - st.load_checkpoint(path)?; + if let Err(error) = st.load_checkpoint(path) { + if strict_exact { + let message = format!("cannot resume RL checkpoint: {error:#}"); + append_strict_failure( + cfg.event_tape.as_deref(), + "layout_hash_mismatch", + &message, + ); + bail!("RL strict failure layout_hash_mismatch: {message}"); + } + return Err(error); + } + if strict_exact && !st.checkpoint_layout_verified { + let message = "RL checkpoint is missing its canonical layout hash"; + append_strict_failure( + cfg.event_tape.as_deref(), + "layout_hash_mismatch", + message, + ); + bail!("RL strict failure layout_hash_mismatch: {message}"); + } info!(step = st.global_step, "resumed from checkpoint"); } } @@ -1116,6 +1177,16 @@ async fn scheduler( remove_final_marker(cfg.checkpoint_path.as_ref().unwrap())?; } + // Persist a fresh version zero before its first BCAST. A resumed cut is + // already committed and only needs to be rebroadcast. + if checkpoint_each_round { + if let Some(path) = cfg.checkpoint_path.as_ref() { + if !path.exists() { + st.save_checkpoint(path)?; + info!(step = st.global_step, path = %path.display(), "checkpoint committed"); + } + } + } // Send everyone the initial (or resumed) global parameters so all // learners start bit-identical (also serves recovery for late joiners). broadcast_all_fragments(&st, ®istry).await; @@ -1145,7 +1216,7 @@ async fn scheduler( break; } let groups = current_groups(®istry); - if groups.is_empty() { + if groups.is_empty() || (fixed_roster && groups.len() < cfg.quorum as usize) { next_launch = Instant::now() + Duration::from_millis(100); break; } @@ -1217,6 +1288,33 @@ async fn scheduler( RoundAction::Restart => { let r = &mut inflight[i]; let groups = current_groups(®istry); + if fixed_roster { + if groups.len() < cfg.quorum as usize { + r.quorum_deadline = + Instant::now() + Duration::from_secs(cfg.quorum_timeout_s); + i += 1; + continue; + } + warn!( + step = r.t, + responses = r.pushes.len(), + roster = cfg.learners, + "fixed roster incomplete; waiting for missing logical learners" + ); + for g in groups { + if !r + .pushes + .keys() + .any(|member| member.learner_id == g.member.learner_id) + { + let _ = g.send_small(MSG_PULL_REQ, r.pull.clone()).await; + } + } + r.quorum_deadline = + Instant::now() + Duration::from_secs(cfg.quorum_timeout_s); + i += 1; + continue; + } warn!( step = r.t, attempt = r.attempt, @@ -1271,16 +1369,37 @@ async fn scheduler( Err(_) => continue, // deadline hit; loop re-evaluates Ok(None) => bail!("event channel closed"), Ok(Some(ev)) => match ev { + Event::Fatal { metric, message } => { + append_strict_failure(cfg.event_tape.as_deref(), metric, &message); + bail!("RL strict failure {metric}: {message}"); + } Event::Push { member, push } => { let learner_id = member.learner_id; let generation = member.generation; let local_step = push.local_step; let global_step = push.global_step; let fragment_id = push.fragment_id; - match route_push(&mut inflight, member, push) { + let disposition = + route_push(&mut inflight, member, push, cfg.max_base_lag, fixed_roster); + match disposition { PushDisposition::Accepted => { step_rates.note(member, local_step, Instant::now()); } + PushDisposition::Duplicate => warn!( + learner_id, + generation, + step = global_step, + fragment = fragment_id, + "duplicate push rejected" + ), + PushDisposition::StaleBase if strict_exact => { + let metric = "rejected_stale_updates"; + let message = format!( + "learner {learner_id} generation {generation} step {global_step} fragment {fragment_id}: {disposition:?}" + ); + append_strict_failure(cfg.event_tape.as_deref(), metric, &message); + bail!("RL strict failure {metric}: {message}"); + } disposition => warn!( learner_id, generation, @@ -1294,6 +1413,21 @@ async fn scheduler( Event::Hello { group } => { // Rejoining learner: catch it up to the current state. send_all_fragments(&st, &group).await; + if fixed_roster && is_current_member(®istry, group.member) { + for round in &inflight { + let belongs = round + .expected_members + .iter() + .any(|member| member.learner_id == group.member.learner_id); + let answered = round + .pushes + .keys() + .any(|member| member.learner_id == group.member.learner_id); + if belongs && !answered { + let _ = group.send_small(MSG_PULL_REQ, round.pull.clone()).await; + } + } + } } Event::Init { .. } => {} // already initialized; ignore Event::FinalAck { member, .. } => { @@ -1345,20 +1479,21 @@ async fn scheduler( return Ok(()); } - // The outer loop is now quiescent: every launched round has completed. - // Persist this authoritative cut regardless of the periodic checkpoint - // interval so a non-divisible total_steps can never leave a stale final - // checkpoint behind. + // Exact-base runs persist every committed cut before BCAST. Other modes + // still need a final save when total_steps is not divisible by the + // periodic interval. if let Some(path) = &cfg.checkpoint_path { - if cfg.mark_final_checkpoint { - remove_final_marker(path)?; + if !checkpoint_each_round { + if cfg.mark_final_checkpoint { + remove_final_marker(path)?; + } + st.save_checkpoint(path)?; + info!( + step = st.global_step, + path = %path.display(), + "final checkpoint written" + ); } - st.save_checkpoint(path)?; - info!( - step = st.global_step, - path = %path.display(), - "final checkpoint written" - ); if cfg.mark_final_checkpoint { write_final_marker(path, st.global_step)?; } @@ -1370,10 +1505,14 @@ async fn scheduler( // Freeze terminal membership to the live groups at the final cut. // Learners already abandoned by fleet recovery are not valid artifact // producers and must not prevent surviving learners from finalizing. - let final_members: HashSet = current_groups(®istry) - .into_iter() - .map(|group| group.member.learner_id) - .collect(); + let final_members: HashSet = if fixed_roster { + (0..cfg.learners).collect() + } else { + current_groups(®istry) + .into_iter() + .map(|group| group.member.learner_id) + .collect() + }; finalize_learners(&cfg, &st, &mut events, ®istry, &final_members).await?; info!("training complete after {} outer steps", cfg.total_steps); // Give writer tasks a moment to flush the final control frames. @@ -1421,6 +1560,10 @@ async fn collect_budget_reports( .context("learner budget reports require --learner-budget-steps")?; while reports.len() < cfg.learners as usize { match events.recv().await.context("event channel closed")? { + Event::Fatal { metric, message } => { + append_strict_failure(cfg.event_tape.as_deref(), metric, &message); + bail!("RL strict failure {metric}: {message}"); + } Event::BudgetDone { member, local_steps, @@ -1467,6 +1610,7 @@ enum PushDisposition { Accepted, Duplicate, UnexpectedMember, + StaleBase, FutureBase, OutOfRound, } @@ -1499,7 +1643,13 @@ fn fragment_available(rounds: &[Round], fragment_id: usize) -> bool { !rounds.iter().any(|round| round.p == fragment_id) } -fn route_push(rounds: &mut [Round], member: Member, push: Push) -> PushDisposition { +fn route_push( + rounds: &mut [Round], + member: Member, + push: Push, + max_base_lag: Option, + fixed_roster: bool, +) -> PushDisposition { let Some(round) = rounds.iter_mut().find(|round| { round.t == push.global_step && round.p == push.fragment_id as usize @@ -1507,13 +1657,32 @@ fn route_push(rounds: &mut [Round], member: Member, push: Push) -> PushDispositi }) else { return PushDisposition::OutOfRound; }; - if !round.expected_members.contains(&member) { + let expected = if fixed_roster { + round + .expected_members + .iter() + .any(|expected| expected.learner_id == member.learner_id) + } else { + round.expected_members.contains(&member) + }; + if !expected { return PushDisposition::UnexpectedMember; } if push.base_version > round.base_version { return PushDisposition::FutureBase; } - if round.pushes.contains_key(&member) { + if max_base_lag.is_some_and(|limit| round.base_version - push.base_version > limit) { + return PushDisposition::StaleBase; + } + let duplicate = if fixed_roster { + round + .pushes + .keys() + .any(|accepted| accepted.learner_id == member.learner_id) + } else { + round.pushes.contains_key(&member) + }; + if duplicate { return PushDisposition::Duplicate; } round.pushes.insert(member, push); @@ -1561,7 +1730,12 @@ async fn complete_round( ); } outer_gradients.push(push.outer_gradient.as_slice()); - weights.push(crate::merge::learner_weight(push.c_tokens, push.c_steps)); + weights.push(match cfg.learner_weight { + LearnerWeight::Tokens2OverSteps => { + crate::merge::learner_weight(push.c_tokens, push.c_steps) + } + LearnerWeight::Equal => 1.0, + }); responders.push(*member); } let sync_start = Instant::now(); @@ -1574,6 +1748,14 @@ async fn complete_round( st.record_merge(push.learner_id, push.c_steps, push.c_tokens); } + let checkpoint_each_round = cfg.max_base_lag == Some(0) && cfg.checkpoint_every == 1; + if checkpoint_each_round { + if let Some(path) = cfg.checkpoint_path.as_ref() { + st.save_checkpoint(path)?; + info!(step = t, path = %path.display(), "checkpoint committed"); + } + } + // Broadcast the updated fragment. let payload = encode_bcast(st, p)?; for g in current_groups(registry) { @@ -1607,27 +1789,31 @@ async fn complete_round( &pushes, gnorm, ms, + &st.layout_fingerprint, ); } // Consistent cut: this round is fully applied and broadcast, and every // other in-flight round is still gathering (it has not touched state). // A crash-resume loses those gathers; their fragments simply merge on // a later cycle, which the quorum design already tolerates. - if let Some(path) = &cfg.checkpoint_path { - if cfg.checkpoint_every > 0 && t % cfg.checkpoint_every == 0 { - st.save_checkpoint(path)?; - info!(step = t, path = %path.display(), "checkpoint written"); + if !checkpoint_each_round { + if let Some(path) = &cfg.checkpoint_path { + if cfg.checkpoint_every > 0 && t % cfg.checkpoint_every == 0 { + st.save_checkpoint(path)?; + info!(step = t, path = %path.display(), "checkpoint written"); + } } } Ok(()) } fn new_state_for(group: &Arc, cfg: &Config) -> Result { - let mut st = GlobalState::new( + let mut st = GlobalState::new_with_layout_fingerprint( group.layout.clone(), cfg.outer_lr, cfg.outer_momentum, group.dtype, + group.layout_fingerprint, )?; if cfg.delta_correction { st.delta_correction = Some(crate::merge::Heloco::default()); @@ -1768,6 +1954,10 @@ async fn finalize_learners( })? .context("event channel closed during finalization")?; match event { + Event::Fatal { metric, message } => { + append_strict_failure(cfg.event_tape.as_deref(), metric, &message); + bail!("RL strict failure {metric}: {message}"); + } Event::FinalAck { member, global_step, @@ -1906,6 +2096,26 @@ fn encode_final_fragment(st: &GlobalState, p: usize) -> Result { } /// One JSONL record per merge: the event tape. +fn append_strict_failure(path: Option<&std::path::Path>, metric: &str, message: &str) { + use std::io::Write; + + let Some(path) = path else { + return; + }; + let escaped = message.replace('\\', "\\\\").replace('"', "\\\""); + let line = format!( + "{{\"event\":\"rl_strict_failure\",\"metric\":\"{metric}\",\"value\":1,\"error\":\"{escaped}\"}}\n" + ); + let result = std::fs::OpenOptions::new() + .create(true) + .append(true) + .open(path) + .and_then(|mut file| file.write_all(line.as_bytes())); + if let Err(error) = result { + warn!("event tape write failed: {error}"); + } +} + fn append_tape( path: &std::path::Path, step: u64, @@ -1920,6 +2130,7 @@ fn append_tape( pushes: &HashMap, gnorm: f64, ms: u64, + layout_fingerprint: &[u8; 32], ) { use std::io::Write; let mut responded_members: Vec = pushes.keys().copied().collect(); @@ -1967,8 +2178,14 @@ fn append_tape( responders.sort(); let quorum_ms = json_opt_u64(quorum_ms); let grace_ms = json_opt_u64(grace_ms); + let layout_hash: String = layout_fingerprint + .iter() + .map(|byte| format!("{byte:02x}")) + .collect(); + let merge_seconds = sync_ms as f64 / 1000.0; let line = format!( - "{{\"protocol_version\":{PROTOCOL_VERSION},\"delta_semantics\":\"local_minus_raw_anchor\",\"step\":{step},\"fragment\":{fragment},\"launch_base_version\":{launch_base_version},\"attempt\":{attempt},\"gnorm\":{gnorm},\"ms\":{ms},\"quorum\":{quorum},\"expected\":{},\"expected_members\":{},\"responded\":{},\"responded_members\":{},\"missed_grace\":{},\"missed_members\":{},\"quorum_ms\":{quorum_ms},\"grace_ms\":{grace_ms},\"sync_ms\":{sync_ms},\"responders\":[{}]}}\n", + "{{\"protocol_version\":{PROTOCOL_VERSION},\"delta_semantics\":\"local_minus_raw_anchor\",\"sync/layout_hash\":\"{layout_hash}\",\"sync/base_version\":{launch_base_version},\"sync/responders\":{},\"sync/quorum\":{quorum},\"sync/rejected_stale_updates\":0,\"sync/merge_seconds\":{merge_seconds},\"sync/global_delta_norm\":{gnorm},\"step\":{step},\"fragment\":{fragment},\"launch_base_version\":{launch_base_version},\"attempt\":{attempt},\"gnorm\":{gnorm},\"ms\":{ms},\"quorum\":{quorum},\"expected\":{},\"expected_members\":{},\"responded\":{},\"responded_members\":{},\"missed_grace\":{},\"missed_members\":{},\"quorum_ms\":{quorum_ms},\"grace_ms\":{grace_ms},\"sync_ms\":{sync_ms},\"responders\":[{}]}}\n", + responded.len(), json_ids(&expected_learners), json_members(expected_members), json_ids(&responded), @@ -2048,6 +2265,7 @@ mod tests { layout: Layout { fragments: Vec::new(), }, + layout_fingerprint: [0; 32], num_streams: 0, max_init_payload: 0, max_push_payload: 0, @@ -2250,16 +2468,16 @@ mod tests { let captured = member(0, 10); let mut rounds = vec![test_round(vec![captured])]; assert_eq!( - route_push(&mut rounds, member(1, 20), test_push(5)), + route_push(&mut rounds, member(1, 20), test_push(5), None, false), PushDisposition::UnexpectedMember ); assert_eq!( - route_push(&mut rounds, member(0, 11), test_push(5)), + route_push(&mut rounds, member(0, 11), test_push(5), None, false), PushDisposition::UnexpectedMember ); assert!(rounds[0].pushes.is_empty()); assert_eq!( - route_push(&mut rounds, captured, test_push(5)), + route_push(&mut rounds, captured, test_push(5), None, false), PushDisposition::Accepted ); } @@ -2285,26 +2503,71 @@ mod tests { let captured = member(0, 10); let mut rounds = vec![test_round(vec![captured])]; assert_eq!( - route_push(&mut rounds, captured, test_push(6)), + route_push(&mut rounds, captured, test_push(6), None, false), PushDisposition::FutureBase ); assert_eq!( - route_push(&mut rounds, captured, test_push(4)), + route_push(&mut rounds, captured, test_push(4), None, false), PushDisposition::Accepted ); assert_eq!( - route_push(&mut rounds, captured, test_push(4)), + route_push(&mut rounds, captured, test_push(4), None, false), PushDisposition::Duplicate ); let mut wrong_round = test_push(4); wrong_round.global_step = 99; assert_eq!( - route_push(&mut rounds, captured, wrong_round), + route_push(&mut rounds, captured, wrong_round, None, false), PushDisposition::OutOfRound ); assert_eq!(rounds[0].pushes.len(), 1); } + fn test_exact_push(learner_id: u32, base_version: u64) -> Push { + Push { + learner_id, + fragment_id: 1, + global_step: 7, + round_attempt: 1, + base_version, + local_step: 7, + c_steps: 1, + c_tokens: 1, + outer_gradient: vec![1.0], + } + } + + #[test] + fn max_base_lag_zero_is_exact_and_unique_per_logical_learner() { + let original = member(0, 10); + let replacement = member(0, 11); + let second = member(1, 20); + let mut rounds = vec![test_round(vec![original, second])]; + assert_eq!( + route_push( + &mut rounds, + replacement, + test_exact_push(0, 5), + Some(0), + true, + ), + PushDisposition::Accepted + ); + assert_eq!( + route_push(&mut rounds, original, test_exact_push(0, 5), Some(0), true,), + PushDisposition::Duplicate + ); + assert_eq!( + route_push(&mut rounds, second, test_exact_push(1, 4), Some(0), true,), + PushDisposition::StaleBase + ); + assert_eq!( + route_push(&mut rounds, second, test_exact_push(1, 6), Some(0), true,), + PushDisposition::FutureBase + ); + assert_eq!(rounds[0].pushes.len(), 1); + } + #[test] fn quorum_timeout_with_one_of_two_responses_restarts_without_merging() { let first = member(0, 10); @@ -2321,7 +2584,7 @@ mod tests { round.pushes.clear(); let mut rounds = vec![round]; assert_eq!( - route_push(&mut rounds, second, test_push(5)), + route_push(&mut rounds, second, test_push(5), None, false), PushDisposition::OutOfRound ); assert!(rounds[0].pushes.is_empty()); @@ -2366,6 +2629,7 @@ mod tests { &pushes, 0.5, 44, + &[7; 32], ); let text = std::fs::read_to_string(&path).unwrap(); std::fs::remove_file(&path).ok(); @@ -2379,5 +2643,7 @@ mod tests { assert!(text.contains("\"contribution\":1")); assert!(text.contains("\"protocol_version\":4")); assert!(text.contains("\"delta_semantics\":\"local_minus_raw_anchor\"")); + assert!(text.contains(&format!("\"sync/layout_hash\":\"{}\"", "07".repeat(32)))); + assert!(!text.contains("\"layout_hash\":")); } } diff --git a/syncer/src/state.rs b/syncer/src/state.rs index f18382a..4a1e3fd 100644 --- a/syncer/src/state.rs +++ b/syncer/src/state.rs @@ -114,6 +114,10 @@ pub struct LearnerLedger { pub struct GlobalState { pub layout: Layout, + /// Semantic tensor identity supplied in HELLO and persisted in checkpoints. + pub layout_fingerprint: [u8; 32], + /// Whether a loaded checkpoint carried and matched the fingerprint. + pub checkpoint_layout_verified: bool, /// Θ_p, flat f32 per fragment (concatenated tensors in layout order). pub params: Vec>, /// Nesterov momentum buffers, same shape as params. @@ -134,7 +138,18 @@ pub struct GlobalState { } impl GlobalState { + #[cfg(test)] pub fn new(layout: Layout, outer_lr: f32, outer_momentum: f32, wire_dtype: u8) -> Result { + Self::new_with_layout_fingerprint(layout, outer_lr, outer_momentum, wire_dtype, [0; 32]) + } + + pub fn new_with_layout_fingerprint( + layout: Layout, + outer_lr: f32, + outer_momentum: f32, + wire_dtype: u8, + layout_fingerprint: [u8; 32], + ) -> Result { // Validate declared sizes now, but do not allocate model-sized state // from an unauthenticated HELLO. Params arrive in INIT_PARAMS; the // matching momentum is allocated only after that exact-sized payload @@ -165,6 +180,8 @@ impl GlobalState { versions.resize(fragment_count, 0); Ok(Self { layout, + layout_fingerprint, + checkpoint_layout_verified: false, params, momentum, initialized, @@ -386,6 +403,7 @@ impl GlobalState { f.write_all(&l.steps.to_le_bytes())?; f.write_all(&l.tokens.to_le_bytes())?; } + f.write_all(&self.layout_fingerprint)?; f.flush()?; f.get_ref().sync_all()?; } @@ -442,6 +460,17 @@ impl GlobalState { }; self.ledger.insert(id, l); } + self.checkpoint_layout_verified = match r.remaining() { + 0 => false, + 32 => { + let checkpoint_fingerprint: [u8; 32] = r.take(32)?.try_into()?; + if checkpoint_fingerprint != self.layout_fingerprint { + bail!("checkpoint layout fingerprint does not match HELLO"); + } + true + } + remaining => bail!("checkpoint has {remaining} trailing bytes"), + }; Ok(()) } } @@ -585,10 +614,54 @@ mod tests { assert_eq!(st2.versions, vec![7, 0]); assert_eq!(st2.params, st.params); assert!(st2.all_initialized()); + assert!(st2.checkpoint_layout_verified); assert_eq!(st2.ledger.get(&3).unwrap().tokens, 4096); std::fs::remove_file(&path).ok(); } + #[test] + fn checkpoint_layout_fingerprint_is_verified_and_legacy_is_readable() { + let dir = std::env::temp_dir().join(format!("yeto-layout-ckpt-{}", std::process::id())); + std::fs::create_dir_all(&dir).unwrap(); + let path = dir.join("state.ckpt"); + let mut state = GlobalState::new_with_layout_fingerprint( + layout2(), + 0.7, + 0.9, + crate::protocol::DTYPE_F32, + [1; 32], + ) + .unwrap(); + state.init_fragment(0, vec![1.0; 4]).unwrap(); + state.init_fragment(1, vec![2.0; 4]).unwrap(); + state.save_checkpoint(&path).unwrap(); + + let mut mismatched = GlobalState::new_with_layout_fingerprint( + layout2(), + 0.7, + 0.9, + crate::protocol::DTYPE_F32, + [2; 32], + ) + .unwrap(); + assert!(mismatched.load_checkpoint(&path).is_err()); + + let mut bytes = std::fs::read(&path).unwrap(); + bytes.truncate(bytes.len() - 32); + std::fs::write(&path, bytes).unwrap(); + let mut legacy = GlobalState::new_with_layout_fingerprint( + layout2(), + 0.7, + 0.9, + crate::protocol::DTYPE_F32, + [2; 32], + ) + .unwrap(); + legacy.load_checkpoint(&path).unwrap(); + assert!(!legacy.checkpoint_layout_verified); + std::fs::remove_dir_all(&dir).ok(); + } + #[test] fn checkpoint_atomically_replaces_previous_file() { let dir = std::env::temp_dir().join(format!("yeto-atomic-ckpt-{}", std::process::id())); diff --git a/tests/test_rl_core.py b/tests/test_rl_core.py new file mode 100644 index 0000000..49b2fed --- /dev/null +++ b/tests/test_rl_core.py @@ -0,0 +1,1074 @@ +from __future__ import annotations + +import asyncio +import json +import sys +import types +from types import SimpleNamespace + +import pytest +import torch + +import yeto.rl.miles as miles +from yeto.fragments import MERGE_AVG +from yeto.rl.core import ( + CanonicalTensorSpec, + LocalRoundStats, + StrictRlInvariantError, + build_avg_layout, + canonical_layout_hash, + canonical_state, + flat_tensor, + policy_delta, + tensors_from_flat, +) +from yeto.rl.bridge import BridgeConfig, StrictRlBridge +from yeto.protocol import PullRequest +from yeto.rl.export import adapter_targets, derive_peft_lora_specs +from yeto.rl.miles import MilesIslandRuntime + + +def tensors(): + return { + "base_model.model.z.lora_B.weight": torch.tensor([[3.0], [4.0]]), + "base_model.model.a.lora_A.weight": torch.tensor([[1.0, 2.0]]), + } + + +MODEL_REVISION = "a" * 40 +LORA_CONFIG_HASH = "b" * 64 + + +def state(version, values, **kwargs): + return canonical_state( + version, + values, + base_model_revision=kwargs.pop("base_model_revision", MODEL_REVISION), + lora_config_hash=kwargs.pop("lora_config_hash", LORA_CONFIG_HASH), + **kwargs, + ) + + +def test_canonical_lora_is_sorted_f32_cpu_and_one_avg_fragment(): + canonical = state(7, tensors()) + assert [spec.name for spec in canonical.specs] == sorted(tensors()) + assert all(value.dtype == torch.float32 for value in canonical.tensors.values()) + assert all(spec.dtype == "float32" for spec in canonical.specs) + assert canonical.layout_hash == canonical_layout_hash(canonical.specs) + layout = build_avg_layout(canonical.specs) + assert len(layout.fragments) == 1 + assert layout.fragments[0].merge_mode == MERGE_AVG + + +def test_flat_round_trip_and_delta_use_the_exact_tensor_contract(): + base = state(2, tensors()) + flat = flat_tensor(base.tensors, base.specs) + rebuilt = tensors_from_flat(flat, base.specs) + assert all(torch.equal(base.tensors[name], rebuilt[name]) for name in rebuilt) + local = state( + 2, + {name: value + 0.5 for name, value in base.tensors.items()}, + expected_specs=base.specs, + ) + assert torch.equal(policy_delta(local, base), torch.full_like(flat, 0.5)) + + +def test_policy_delta_rejects_canonical_identity_mismatch(): + base = state(2, tensors()) + local = state(2, tensors(), base_model_revision="c" * 40) + with pytest.raises(ValueError, match="identities differ"): + policy_delta(local, base) + with pytest.raises(ValueError, match="layout hash changed"): + state(2, tensors(), layout_hash="d" * 64) + + +@pytest.mark.parametrize("bad", [float("nan"), float("inf"), -float("inf")]) +def test_canonical_lora_rejects_non_finite_values(bad): + values = tensors() + values["base_model.model.a.lora_A.weight"][0, 0] = bad + with pytest.raises(ValueError, match="NaN or Inf"): + state(0, values) + + +def test_canonical_lora_rejects_non_peft_names_and_shape_drift(): + with pytest.raises(ValueError, match="canonical PEFT"): + state(0, {"model.weight": torch.ones(2)}) + canonical = state(0, tensors()) + changed = tensors() + changed["base_model.model.a.lora_A.weight"] = torch.ones(2, 1) + with pytest.raises(ValueError, match="names, shapes, or dtypes"): + state(0, changed, expected_specs=canonical.specs) + + +def test_avg_layout_rejects_duplicate_names(): + spec = CanonicalTensorSpec( + "base_model.model.x.lora_A.weight", (1, 2), "float32", 2 + ) + with pytest.raises(ValueError, match="unique"): + build_avg_layout((spec, spec)) + + +def test_two_model_configs_use_the_same_generic_peft_path(tmp_path): + transformers = pytest.importorskip("transformers") + pytest.importorskip("peft") + configs = ( + transformers.LlamaConfig( + vocab_size=32, + hidden_size=8, + intermediate_size=16, + num_hidden_layers=1, + num_attention_heads=2, + num_key_value_heads=2, + ), + transformers.OPTConfig( + vocab_size=32, + hidden_size=8, + ffn_dim=16, + num_hidden_layers=1, + num_attention_heads=2, + ), + ) + layouts = [] + for index, config in enumerate(configs): + model_dir = tmp_path / str(index) + config.save_pretrained(model_dir) + specs = derive_peft_lora_specs( + str(model_dir), + None, + rank=2, + targets="all-linear", + ) + layouts.append(build_avg_layout(specs)) + assert specs + assert all(layout.fragments[0].merge_mode == MERGE_AVG for layout in layouts) + + +def test_auto_targets_follow_existing_moe_semantics_from_actual_config(tmp_path): + transformers = pytest.importorskip("transformers") + pytest.importorskip("peft") + config = transformers.LlamaConfig( + vocab_size=32, + hidden_size=8, + intermediate_size=16, + num_hidden_layers=1, + num_attention_heads=2, + num_key_value_heads=2, + num_local_experts=2, + ) + model_dir = tmp_path / "moe-marked" + config.save_pretrained(model_dir) + + specs = derive_peft_lora_specs( + str(model_dir), + None, + rank=2, + targets="auto", + ) + + assert set(adapter_targets(specs)) == {"q_proj", "k_proj", "v_proj", "o_proj"} + + +class _VersionedRuntime: + def __init__(self): + self.state = state(0, tensors()) + + def initialize(self): + return self.state + + def apply_global_policy(self, state): + self.state = state + + def run_local_round(self, **kwargs): + return LocalRoundStats( + island_id=0, + local_round_id=1, + base_policy_version=0, + active_groups=1, + completed_groups=1, + cancelled_groups=0, + completed_trajectories=1, + action_tokens=1, + tool_wait_seconds=0.0, + group_p50_seconds=0.1, + group_p95_seconds=0.1, + group_p99_seconds=0.1, + reward_mean=1.0, + reward_std=0.0, + zero_variance_group_ratio=1.0, + mean_kl=None, + ess_ratio=None, + clip_fraction=0.0, + delta_l2_norm=0.0, + rollout_seconds=0.1, + train_seconds=0.1, + ) + + def export_local_policy(self): + return state(self.state.policy_version + 1, self.state.tensors) + + def record_local_round(self, stats): + pass + + def shutdown(self): + pass + + +def _bridge_config(tmp_path="/tmp/yeto-rl-test-events.jsonl"): + canonical = state(0, tensors()) + return BridgeConfig( + syncer_addr=("127.0.0.1", 1), + learner_id=0, + global_rounds=1, + groups_per_round=1, + samples_per_group=1, + local_optimizer_steps=1, + expected_specs=canonical.specs, + base_model_revision=canonical.base_model_revision, + lora_config_hash=canonical.lora_config_hash, + layout_hash=canonical.layout_hash, + event_tape=str(tmp_path), + wan_streams=0, + ) + + +def test_policy_delta_rejects_exported_policy_version_drift(): + runtime = _VersionedRuntime() + bridge = StrictRlBridge(runtime, _bridge_config()) + bridge.current = bridge.initial + with pytest.raises(ValueError, match="versions differ"): + bridge._run_round(PullRequest(0, 1, 1)) + + +def test_local_round_event_contains_every_init_metric(tmp_path): + runtime = _VersionedRuntime() + runtime.export_local_policy = lambda: state( + 0, + {name: value + 1 for name, value in runtime.state.tensors.items()}, + ) + event_tape = tmp_path / "events.jsonl" + bridge = StrictRlBridge(runtime, _bridge_config(event_tape)) + bridge.current = bridge.initial + pushed = [] + bridge.client = SimpleNamespace(push_fragment=lambda *values: pushed.append(values)) + + bridge._run_round(PullRequest(0, 1, 1)) + + event = json.loads(event_tape.read_text()) + expected = { + "rl/active_groups", + "rl/completed_groups", + "rl/cancelled_groups", + "rl/completed_trajectories", + "rl/action_tokens", + "rl/tool_wait_seconds", + "rl/rollout_seconds", + "rl/group_p50_seconds", + "rl/group_p95_seconds", + "rl/group_p99_seconds", + "rl/reward_mean", + "rl/reward_std", + "rl/zero_variance_group_ratio", + "rl/global_policy_version", + "rl/rollout_policy_version", + "rl/mixed_version_group_count", + "rl/local_delta_norm", + "rl/current_vs_rollout_kl", + "rl/ess_ratio", + "rl/clip_fraction", + "sync/bytes_sent", + } + assert expected <= event.keys() + assert event["sync/bytes_sent"] == 48 + len(pushed[0][-1]) + + +def test_miles_admission_pause_aborts_inflight_rollouts(): + modes = [] + + class PauseGeneration: + async def remote(self, mode): + modes.append(mode) + + class Engine: + pause_generation = PauseGeneration() + + runtime = object.__new__(MilesIslandRuntime) + + async def engines(): + return [Engine()] + + runtime._engines = engines + asyncio.run(runtime._pause_rollout()) + assert modes == ["abort"] + + +def test_miles_bridge_propagates_attention_backend(monkeypatch): + class AutoBridge: + def to_megatron_provider(self): + return SimpleNamespace(attention_backend="auto") + + megatron = types.ModuleType("megatron") + bridge = types.ModuleType("megatron.bridge") + training = types.ModuleType("megatron.bridge.training") + config = types.ModuleType("megatron.bridge.training.config") + + class DistributedDataParallelConfig: + def __init__(self, **kwargs): + self.use_distributed_optimizer = kwargs["use_distributed_optimizer"] + + bridge.AutoBridge = AutoBridge + config.DistributedDataParallelConfig = DistributedDataParallelConfig + training.config = config + monkeypatch.setitem(sys.modules, "megatron", megatron) + monkeypatch.setitem(sys.modules, "megatron.bridge", bridge) + monkeypatch.setitem(sys.modules, "megatron.bridge.training", training) + monkeypatch.setitem(sys.modules, "megatron.bridge.training.config", config) + installed = [] + monkeypatch.setattr(miles, "_install_train_metric_capture", lambda: None) + monkeypatch.setattr(miles, "_install_colocated_lora_ipc_sync", lambda: None) + monkeypatch.setattr(miles, "install_miles_actor_adapter", lambda: installed.append(True)) + + args = SimpleNamespace( + attention_backend="unfused", use_distributed_optimizer=True + ) + miles.configure_miles_bridge(args) + + assert AutoBridge().to_megatron_provider().attention_backend == "unfused" + assert not args.use_distributed_optimizer + assert not config.DistributedDataParallelConfig( + use_distributed_optimizer=True + ).use_distributed_optimizer + assert installed == [True] + + +def test_miles_actor_calls_handle_rank0_export_and_all_rank_results(): + runtime = object.__new__(MilesIslandRuntime) + runtime.args = SimpleNamespace(actor_num_nodes=1, actor_num_gpus_per_node=2) + + class Actors: + def __init__(self, results): + self.results = results + + async def _broadcast(self, method, *args): + return self.results + + runtime.actor_model = Actors([{"policy": 1}, None]) + assert asyncio.run(runtime._actor_call("export", rank0=True)) == {"policy": 1} + runtime.actor_model = Actors([(4, "hash"), (4, "hash")]) + assert asyncio.run(runtime._actor_call("apply")) == (4, "hash") + runtime.actor_model = Actors([3, 4]) + with pytest.raises(RuntimeError, match="ranks disagree"): + asyncio.run(runtime._actor_call("steps")) + + +def test_miles_waits_for_the_complete_multinode_ray_island(monkeypatch): + canonical = state(0, tensors()) + args = SimpleNamespace( + actor_num_nodes=2, + actor_num_gpus_per_node=4, + offload_train=True, + yeto_rl_base_model_revision=canonical.base_model_revision, + yeto_rl_lora_config_hash=canonical.lora_config_hash, + yeto_rl_layout_hash=canonical.layout_hash, + ) + resource_counts = iter((4, 8)) + ray = types.ModuleType("ray") + ray.is_initialized = lambda: True + ray.cluster_resources = lambda: {"GPU": next(resource_counts)} + actor = SimpleNamespace() + + async def broadcast(method, *values): + assert method == "yeto_rl_export_policy" + return [canonical.tensors] + [None] * 7 + + async def onload(): + pass + + async def offload(): + pass + + actor._broadcast = broadcast + actor.onload = onload + actor.offload = offload + + async def create_training_models(*values): + return actor, None + + placement = types.ModuleType("miles.ray.placement_group") + placement.create_placement_groups = lambda _args: {"rollout": object()} + placement.create_rollout_manager = lambda *values: (object(), None) + placement.create_training_models = create_training_models + external_miles = types.ModuleType("miles") + ray_package = types.ModuleType("miles.ray") + ray_package.placement_group = placement + external_miles.ray = ray_package + monkeypatch.setitem(sys.modules, "ray", ray) + monkeypatch.setitem(sys.modules, "miles", external_miles) + monkeypatch.setitem(sys.modules, "miles.ray", ray_package) + monkeypatch.setitem(sys.modules, "miles.ray.placement_group", placement) + monkeypatch.setattr(miles, "install_miles_actor_adapter", lambda: None) + + async def no_wait(_seconds): + pass + + monkeypatch.setattr(miles.asyncio, "sleep", no_wait) + runtime = MilesIslandRuntime(args) + try: + initialized = asyncio.run(runtime._initialize()) + finally: + runtime.loop.close() + + assert initialized.layout_hash == canonical.layout_hash + + +def test_miles_keeps_non_offloaded_trainer_resident(): + calls = [] + + async def onload(): + calls.append("onload") + + async def offload(): + calls.append("offload") + + runtime = object.__new__(MilesIslandRuntime) + runtime.args = SimpleNamespace(offload_train=False) + runtime.actor_model = SimpleNamespace(onload=onload, offload=offload) + runtime._trainer_awake = True + + asyncio.run(runtime._onload_trainer()) + asyncio.run(runtime._offload_trainer()) + assert runtime._trainer_awake + assert calls == [] + + +def test_miles_collects_every_data_parallel_rollout_shard(monkeypatch): + monkeypatch.setitem(sys.modules, "ray", SimpleNamespace(get=lambda value: value)) + runtime = object.__new__(MilesIslandRuntime) + runtime.args = SimpleNamespace( + actor_num_nodes=1, + actor_num_gpus_per_node=4, + expert_model_parallel_size=2, + ) + shards = [ + SimpleNamespace(inner={"sample_indices": [0]}), + SimpleNamespace(inner={"sample_indices": [1]}), + SimpleNamespace(inner={"sample_indices": [2]}), + SimpleNamespace(inner={"sample_indices": [3]}), + ] + assert runtime._rollout_batches({"data_ref": shards}) == [ + {"sample_indices": [0]}, + {"sample_indices": [1]}, + {"sample_indices": [2]}, + {"sample_indices": [3]}, + ] + + +def test_miles_rollout_lifecycle_metrics_use_real_task_completion(monkeypatch): + class Sample: + def __init__(self, status): + self.status = SimpleNamespace(value=status) + + class GenerateState: + def __init__(self): + self.pendings = set() + + def submit_generate_tasks(self, groups): + async def finish(group): + await asyncio.sleep(0) + return group + + self.pendings.update(asyncio.create_task(finish(group)) for group in groups) + + upstream = types.ModuleType("miles.rollout.sglang_rollout") + upstream.GenerateState = GenerateState + rollout = types.ModuleType("miles.rollout") + rollout.sglang_rollout = upstream + package = types.ModuleType("miles") + package.rollout = rollout + monkeypatch.setitem(sys.modules, "miles", package) + monkeypatch.setitem(sys.modules, "miles.rollout", rollout) + monkeypatch.setitem(sys.modules, "miles.rollout.sglang_rollout", upstream) + + def generate(*_args, **_kwargs): + async def run(): + state = GenerateState() + state.submit_generate_tasks( + [[Sample("completed")], [Sample("aborted")]] + ) + await asyncio.gather(*state.pendings) + + asyncio.run(run()) + return object() + + _, lifecycle = miles._run_rollout_with_metrics( + generate, SimpleNamespace(), 0, object(), False + ) + assert lifecycle["active"] == 0 + assert lifecycle["peak_active"] == 2 + assert lifecycle["cancelled"] == 1 + assert len(lifecycle["durations"]) == 2 + + +def test_miles_train_metric_capture_returns_rank_zero_values(monkeypatch): + args = SimpleNamespace() + + def log_train_step(*_values, **_kwargs): + return { + "train/train_rollout_kl": 0.1, + "train/ess_ratio": 0.9, + "train/pg_clipfrac": 0.2, + } + + model = types.ModuleType("miles.backends.megatron_utils.model") + model.log_train_step = log_train_step + megatron_utils = types.ModuleType("miles.backends.megatron_utils") + megatron_utils.model = model + backends = types.ModuleType("miles.backends") + backends.megatron_utils = megatron_utils + package = types.ModuleType("miles") + package.backends = backends + monkeypatch.setitem(sys.modules, "miles", package) + monkeypatch.setitem(sys.modules, "miles.backends", backends) + monkeypatch.setitem(sys.modules, "miles.backends.megatron_utils", megatron_utils) + monkeypatch.setitem(sys.modules, "miles.backends.megatron_utils.model", model) + + miles._install_train_metric_capture() + model.log_train_step(args=args) + actor = SimpleNamespace(args=args) + assert miles._actor_train_metrics(actor) == { + "train/train_rollout_kl": 0.1, + "train/ess_ratio": 0.9, + "train/pg_clipfrac": 0.2, + } + assert miles._actor_train_metrics(actor) is None + + +def test_miles_keeps_colocated_lora_ipc_storage_alive_until_transfer(monkeypatch): + events = [] + update = types.ModuleType( + "miles.backends.megatron_utils.update_weight.update_weight_from_tensor" + ) + + def send(*_values, **_kwargs): + events.append("send") + return ["ref"], "storage" + + update._send_to_colocated_engine = send + common = types.ModuleType("miles.backends.megatron_utils.update_weight.common") + + def check(results, *, is_lora): + assert results == ["done"] and is_lora + events.append("check") + + common._check_weight_sync_results = check + weight = types.ModuleType("miles.backends.megatron_utils.update_weight") + weight.update_weight_from_tensor = update + megatron_utils = types.ModuleType("miles.backends.megatron_utils") + megatron_utils.update_weight = weight + backends = types.ModuleType("miles.backends") + backends.megatron_utils = megatron_utils + package = types.ModuleType("miles") + package.backends = backends + ray = types.ModuleType("ray") + + def get(refs): + assert refs == ["ref"] + events.append("get") + return ["done"] + + ray.get = get + for name, module in { + "miles": package, + "miles.backends": backends, + "miles.backends.megatron_utils": megatron_utils, + "miles.backends.megatron_utils.update_weight": weight, + "miles.backends.megatron_utils.update_weight.update_weight_from_tensor": update, + "miles.backends.megatron_utils.update_weight.common": common, + "ray": ray, + }.items(): + monkeypatch.setitem(sys.modules, name, module) + monkeypatch.setattr( + torch.distributed, + "barrier", + lambda *, group: events.append(("barrier", group)), + ) + + miles._install_colocated_lora_ipc_sync() + refs, storage = update._send_to_colocated_engine( + [], ipc_gather_group="group", lora_config={} + ) + + assert refs == ["ref"] and storage == "storage" + assert events == ["send", "get", "check", ("barrier", "group")] + + +def test_miles_round_stats_use_checkpoint_rollout_and_train_metrics( + tmp_path, monkeypatch +): + constants = types.ModuleType("sglang.srt.constants") + constants.GPU_MEMORY_TYPE_CUDA_GRAPH = "graph" + constants.GPU_MEMORY_TYPE_KV_CACHE = "kv" + constants.GPU_MEMORY_TYPE_WEIGHTS = "weights" + sglang = types.ModuleType("sglang") + srt = types.ModuleType("sglang.srt") + sglang.srt = srt + srt.constants = constants + monkeypatch.setitem(sys.modules, "sglang", sglang) + monkeypatch.setitem(sys.modules, "sglang.srt", srt) + monkeypatch.setitem(sys.modules, "sglang.srt.constants", constants) + + checkpoint = tmp_path / "island.pt" + args = SimpleNamespace( + actor_num_gpus_per_node=1, + actor_num_nodes=1, + advantage_estimator="grpo", + yeto_rl_model="org/model", + yeto_rl_data="org/data", + yeto_rl_base_model_revision=MODEL_REVISION, + yeto_rl_data_revision="d" * 40, + expert_model_parallel_size=1, + yeto_rl_layout_hash="c" * 64, + lr=1e-4, + yeto_rl_lora_config_hash=LORA_CONFIG_HASH, + n_samples_per_prompt=2, + num_steps_per_rollout=1, + over_sampling_batch_size=2, + yeto_rl_reward_sha256="e" * 64, + rollout_batch_size=2, + seq_length=128, + seed=7, + rollout_max_response_len=16, + custom_generate_function_path=None, + use_session_server=False, + tito_model="default", + yeto_rl_completed_groups_path=str(checkpoint), + yeto_rl_learner_id=0, + ) + torch.save( + { + "schema_version": miles._ISLAND_CHECKPOINT_SCHEMA, + "config": miles._island_checkpoint_config(args), + "policy_version": 3, + "rollout_metrics": { + "active_groups": 3.0, + "cancelled_groups": 1.0, + "tool_wait_seconds": 4.0, + "group_p50_seconds": 5.0, + "group_p95_seconds": 6.0, + "group_p99_seconds": 7.0, + }, + }, + checkpoint, + ) + + class Remote: + def __init__(self, result=None): + self.result = result + + async def remote(self, *_args, **_kwargs): + return self.result + + runtime = object.__new__(MilesIslandRuntime) + runtime.args = args + runtime._policy_version = 3 + runtime._rollout_id = 3 + runtime._rollout_offloaded = False + runtime.rollout_manager = SimpleNamespace( + generate=Remote(object()), + offload=Remote(), + ) + + async def train(*_args): + pass + + runtime.actor_model = SimpleNamespace(train=train) + + async def no_op(): + pass + + runtime._pause_rollout = no_op + runtime._rollout_batches = lambda _pack: [ + { + "weight_versions": [["yeto:3"]] * 4, + "response_lengths": [1, 2, 3, 4], + "sample_indices": [0, 1, 2, 3], + "raw_reward": [1.0, 1.0, 0.0, 2.0], + } + ] + optimizer_steps = iter((10, 11)) + + async def actor_call(method, *_args, **_kwargs): + if method == "yeto_rl_optimizer_steps": + return next(optimizer_steps) + assert method == "yeto_rl_train_metrics" + return { + "train/train_rollout_kl": 0.1, + "train/ess_ratio": 0.8, + "train/pg_clipfrac": 0.25, + } + + runtime._actor_call = actor_call + stats = asyncio.run(runtime._run_local_round(3, 2, 2, 1)) + assert stats.active_groups == 3 + assert stats.cancelled_groups == 1 + assert stats.tool_wait_seconds == 4.0 + assert stats.zero_variance_group_ratio == 0.5 + assert stats.mean_kl == 0.1 + assert stats.ess_ratio == 0.8 + assert stats.clip_fraction == 0.25 + + +def test_island_checkpoint_restores_only_complete_same_policy_groups( + tmp_path, monkeypatch +): + checkpoint = tmp_path / "island.pt" + args = SimpleNamespace( + actor_num_gpus_per_node=1, + actor_num_nodes=1, + advantage_estimator="grpo", + yeto_rl_model="org/model", + yeto_rl_data="org/data", + yeto_rl_base_model_revision=MODEL_REVISION, + yeto_rl_data_revision="d" * 40, + expert_model_parallel_size=1, + yeto_rl_layout_hash="c" * 64, + lr=1e-4, + yeto_rl_lora_config_hash=LORA_CONFIG_HASH, + n_samples_per_prompt=2, + num_steps_per_rollout=1, + over_sampling_batch_size=2, + yeto_rl_reward_sha256="e" * 64, + rollout_batch_size=1, + seq_length=128, + seed=7, + rollout_max_response_len=16, + custom_generate_function_path=None, + use_session_server=False, + tito_model="default", + yeto_rl_completed_groups_path=str(checkpoint), + ) + + def sample(status, version, index=0): + value = SimpleNamespace( + status=SimpleNamespace(value=status), + weight_versions=[f"yeto:{version}"], + index=index, + ) + value.to_dict = lambda: { + "status": status, + "weight_versions": [f"yeto:{version}"], + "index": index, + } + return value + + class Sample: + @staticmethod + def from_dict(value): + version = int(value["weight_versions"][0].split(":", 1)[1]) + return sample(value["status"], version, value["index"]) + + package = types.ModuleType("miles") + utils = types.ModuleType("miles.utils") + sample_types = types.ModuleType("miles.utils.types") + sample_types.Sample = Sample + package.utils = utils + utils.types = sample_types + monkeypatch.setitem(sys.modules, "miles", package) + monkeypatch.setitem(sys.modules, "miles.utils", utils) + monkeypatch.setitem(sys.modules, "miles.utils.types", sample_types) + + complete = [sample("completed", 3), sample("truncated", 3)] + incomplete = [sample("aborted", 3), sample("completed", 3)] + + class DataSource: + def __init__(self, buffer=None): + self.buffer = list(buffer or []) + + def get_samples(self, _count): + return [] + + def add_samples(self, groups): + self.buffer.extend(groups) + + source = DataSource([complete, incomplete]) + miles._save_completed_groups(args, 3, 4, source, {"reward": 1.5}) + assert source.buffer == [complete] + payload = torch.load(checkpoint, weights_only=True) + for name, value in { + "model": "org/model", + "dataset": "org/data", + "seq_length": 128, + "seed": 7, + }.items(): + assert payload["config"][name] == value + assert payload["completed_groups"][0][0]["status"] == "completed" + assert payload["completed_groups"][0][1]["status"] == "truncated" + assert "tensors" not in payload and "local_lora" not in payload + assert not list(tmp_path.glob(".island.pt.tmp-*")) + + restored = DataSource() + miles._restore_completed_groups(args, 3, restored) + assert len(restored.buffer) == 1 + args.seed = 8 + incompatible = DataSource() + miles._restore_completed_groups(args, 3, incompatible) + assert incompatible.buffer == [] + args.seed = 7 + stale = DataSource() + miles._restore_completed_groups(args, 4, stale) + assert stale.buffer == [] + + used = [sample("completed", 3, 10), sample("completed", 3, 11)] + unused = [sample("completed", 3, 20), sample("truncated", 3, 21)] + incomplete = [sample("aborted", 3, 30), sample("completed", 3, 31)] + upstream = types.ModuleType("miles.rollout.sglang_rollout") + + def generate(_args, _rollout_id, data_source, evaluation=False): + assert not evaluation + miles.queue_completed_groups( + _args, [used, unused, incomplete], data_source.get_samples + ) + return SimpleNamespace(samples=[used], metrics={"reward": 2.0}) + + upstream.generate_rollout = generate + + class GenerateState: + def submit_generate_tasks(self, _samples): + pass + + upstream.GenerateState = GenerateState + rollout = types.ModuleType("miles.rollout") + rollout.sglang_rollout = upstream + package = types.ModuleType("miles") + package.rollout = rollout + monkeypatch.setitem(sys.modules, "miles", package) + monkeypatch.setitem(sys.modules, "miles.rollout", rollout) + monkeypatch.setitem(sys.modules, "miles.rollout.sglang_rollout", upstream) + checkpoint.unlink() + source = DataSource() + + miles.generate_rollout(args, 3, source) + + assert source.buffer == [unused] + payload = torch.load(checkpoint, weights_only=True) + assert [sample["index"] for sample in payload["completed_groups"][0]] == [20, 21] + assert payload["rollout_metrics"] == { + "reward": 2.0, + "active_groups": 0.0, + "cancelled_groups": 0.0, + "tool_wait_seconds": 0.0, + "group_p50_seconds": 0.0, + "group_p95_seconds": 0.0, + "group_p99_seconds": 0.0, + } + + +def test_miles_apply_resets_optimizer_and_restores_scheduler_progress(monkeypatch): + name = "base_model.model.layer.lora_A.weight" + parameter = torch.nn.Parameter(torch.zeros(1, 2)) + parameter.main_param = torch.zeros(1, 2, dtype=torch.float32) + + class Mapping: + def hf_to_megatron(self, value, _module): + return value + + def megatron_to_hf(self, value, _module): + return {name: value} + + side = SimpleNamespace( + mapping=Mapping(), + megatron_module=None, + param_weight=parameter, + ) + optimizer_state = { + "exp_avg": torch.ones_like(parameter.main_param), + "exp_avg_sq": torch.ones_like(parameter.main_param), + "step": torch.tensor(9.0), + } + inner = SimpleNamespace( + param_groups=[{"step": 9}], + state={parameter.main_param: optimizer_state}, + ) + unrelated = torch.nn.Parameter(torch.ones(1)) + unrelated_state = {"step": torch.tensor(4.0)} + inner.state[unrelated] = unrelated_state + + def copy_main_to_model(): + parameter.copy_(parameter.main_param) + + child = SimpleNamespace( + optimizer=inner, + _copy_main_params_to_model_params=copy_main_to_model, + ) + optimizer = SimpleNamespace(chained_optimizers=[child]) + scheduler_steps = [] + + class Scheduler: + num_steps = 0 + + def step(self, increment): + scheduler_steps.append(increment) + self.num_steps += increment + + scheduler = Scheduler() + backups = [] + applied = torch.tensor([[3.0, 4.0]]) + state_fn = state(2, {name: applied}) + actor = SimpleNamespace( + args=SimpleNamespace( + yeto_rl_base_model_revision=MODEL_REVISION, + yeto_rl_lora_config_hash=LORA_CONFIG_HASH, + yeto_rl_layout_hash=state_fn.layout_hash, + global_batch_size=64, + num_steps_per_rollout=1, + ), + model=[SimpleNamespace(start_param_sync=lambda **_kwargs: pytest.fail( + "replicated LoRA apply must not start distributed optimizer sync" + ))], + optimizer=optimizer, + opt_param_scheduler=scheduler, + weights_backuper=SimpleNamespace(backup=backups.append), + ) + monkeypatch.setattr(miles, "_adapter_sides", lambda _actor: [(name, side)]) + cache_releases = [] + monkeypatch.setattr(torch.cuda, "empty_cache", lambda: cache_releases.append(True)) + + reset_count, applied_hash = miles._actor_apply_policy(actor, {name: applied}, 2) + assert actor.optimizer is optimizer + assert actor.opt_param_scheduler is scheduler + assert scheduler.num_steps == 128 + assert scheduler_steps == [128] + assert inner.param_groups[0]["step"] == 9 + assert parameter.main_param not in inner.state + assert inner.state[unrelated] is unrelated_state + assert reset_count == 1 + assert applied_hash == miles.policy_hash(state_fn) + assert torch.equal(parameter.main_param, applied) + assert torch.equal(parameter, applied) + assert backups == ["actor"] + assert cache_releases == [True] + + scheduler.num_steps = 192 + with pytest.raises(RuntimeError, match="ahead of the committed policy"): + miles._actor_apply_policy(actor, {name: applied}, 2) + + +def test_miles_apply_hash_mismatch_is_a_strict_failure(monkeypatch): + constants = types.ModuleType("sglang.srt.constants") + constants.GPU_MEMORY_TYPE_CUDA_GRAPH = "graph" + constants.GPU_MEMORY_TYPE_KV_CACHE = "kv" + constants.GPU_MEMORY_TYPE_WEIGHTS = "weights" + sglang = types.ModuleType("sglang") + srt = types.ModuleType("sglang.srt") + sglang.srt = srt + srt.constants = constants + monkeypatch.setitem(sys.modules, "sglang", sglang) + monkeypatch.setitem(sys.modules, "sglang.srt", srt) + monkeypatch.setitem(sys.modules, "sglang.srt.constants", constants) + + runtime = object.__new__(MilesIslandRuntime) + runtime._trainer_awake = False + runtime._rollout_offloaded = True + + async def no_op(*_args, **_kwargs): + pass + + runtime._pause_rollout = no_op + runtime.actor_model = SimpleNamespace(onload=no_op) + + async def actor_call(*_args, **_kwargs): + return 1, "wrong-policy-hash" + + runtime._actor_call = actor_call + with pytest.raises(StrictRlInvariantError) as failure: + asyncio.run(runtime._apply_global_policy(state(0, tensors()))) + assert failure.value.metric == "policy_hash_mismatch_after_apply" + + +def test_miles_policy_apply_event_has_only_namespaced_policy_hash(monkeypatch): + constants = types.ModuleType("sglang.srt.constants") + constants.GPU_MEMORY_TYPE_CUDA_GRAPH = "graph" + constants.GPU_MEMORY_TYPE_KV_CACHE = "kv" + constants.GPU_MEMORY_TYPE_WEIGHTS = "weights" + sglang = types.ModuleType("sglang") + srt = types.ModuleType("sglang.srt") + sglang.srt = srt + srt.constants = constants + monkeypatch.setitem(sys.modules, "sglang", sglang) + monkeypatch.setitem(sys.modules, "sglang.srt", srt) + monkeypatch.setitem(sys.modules, "sglang.srt.constants", constants) + + async def no_op(*_args, **_kwargs): + pass + + class Remote: + async def remote(self, *_args, **_kwargs): + pass + + state_fn = state(1, tensors()) + events = [] + runtime = object.__new__(MilesIslandRuntime) + runtime.args = SimpleNamespace(yeto_rl_learner_id=0, offload_train=True) + runtime._trainer_awake = False + runtime._rollout_offloaded = True + runtime._optimizer_reset_count = 0 + runtime._pause_rollout = no_op + runtime._resume_rollout = no_op + runtime._set_rollout_version = no_op + runtime.actor_model = SimpleNamespace( + onload=no_op, + offload=no_op, + update_weights=no_op, + ) + runtime.rollout_manager = SimpleNamespace( + onload_weights=Remote(), + onload_kv=Remote(), + ) + + async def actor_call(*_args, **_kwargs): + return 2, miles.policy_hash(state_fn) + + runtime._actor_call = actor_call + runtime._append_event = events.append + asyncio.run(runtime._apply_global_policy(state_fn)) + + assert events[0]["sync/global_policy_hash"] == miles.policy_hash(state_fn) + assert "policy_hash" not in events[0] + assert "trainer_ranks" not in events[0] + + +@pytest.mark.parametrize( + "versions", + [["not-a-policy-token"], ["yeto:3", "yeto:4"]], +) +def test_miles_rejects_invalid_or_mixed_rollout_versions(monkeypatch, versions): + constants = types.ModuleType("sglang.srt.constants") + constants.GPU_MEMORY_TYPE_CUDA_GRAPH = "graph" + constants.GPU_MEMORY_TYPE_KV_CACHE = "kv" + constants.GPU_MEMORY_TYPE_WEIGHTS = "weights" + sglang = types.ModuleType("sglang") + srt = types.ModuleType("sglang.srt") + sglang.srt = srt + srt.constants = constants + monkeypatch.setitem(sys.modules, "sglang", sglang) + monkeypatch.setitem(sys.modules, "sglang.srt", srt) + monkeypatch.setitem(sys.modules, "sglang.srt.constants", constants) + + class Generate: + async def remote(self, _rollout_id): + return object() + + runtime = object.__new__(MilesIslandRuntime) + runtime._policy_version = 3 + runtime._rollout_id = 3 + runtime.rollout_manager = SimpleNamespace(generate=Generate()) + + async def no_op(): + pass + + runtime._pause_rollout = no_op + runtime._rollout_batches = lambda _pack: [ + {"weight_versions": [versions]} + ] + with pytest.raises(StrictRlInvariantError) as failure: + asyncio.run(runtime._run_local_round(3, 1, 1, 1)) + assert failure.value.metric == "mixed_version_group_count" diff --git a/tests/test_rl_export.py b/tests/test_rl_export.py new file mode 100644 index 0000000..678da82 --- /dev/null +++ b/tests/test_rl_export.py @@ -0,0 +1,171 @@ +from __future__ import annotations + +import struct + +import pytest +import torch + +from yeto.export import CKPT_MAGIC +from yeto.rl.core import ( + canonical_layout_hash, + canonical_lora_config_hash, + canonical_state, + flat_tensor, + tensors_from_flat, +) +from yeto.rl.export import derive_peft_lora_specs, export_rl_checkpoint + + +def _model(tmp_path): + transformers = pytest.importorskip("transformers") + config = transformers.LlamaConfig( + vocab_size=32, + hidden_size=8, + intermediate_size=16, + num_hidden_layers=1, + num_attention_heads=2, + num_key_value_heads=2, + ) + path = tmp_path / "model" + config.save_pretrained(path) + return path, config + + +def test_attention_regex_is_resolved_before_peft_moe_conversion(tmp_path): + transformers = pytest.importorskip("transformers") + config = transformers.OlmoeConfig( + vocab_size=32, + hidden_size=8, + intermediate_size=16, + num_hidden_layers=1, + num_attention_heads=2, + num_key_value_heads=2, + num_experts=2, + num_experts_per_tok=1, + ) + path = tmp_path / "olmoe" + config.save_pretrained(path) + + specs = derive_peft_lora_specs(str(path), None, rank=2, targets="attention") + + assert specs + assert all("self_attn" in spec.name for spec in specs) + + +MODEL_REVISION = "a" * 40 + + +def _write_checkpoint( + path, values: torch.Tensor, layout_hash: str | None, version=2, ledger_size=2 +): + body = bytearray() + body += struct.pack(" None: help="training loop selector; auto infers diffusion for diffusion aliases, " "otherwise uses the causal-LM learner", ) + rl = p.add_argument_group("Miles RL") + rl.add_argument( + "--training-mode", + choices=["sft", "rl"], + default="sft", + help="training workflow (default: sft)", + ) + rl.add_argument("--rl-runtime", choices=["miles"], default="miles") + rl.add_argument("--rl-image", default=MILES_IMAGE) + rl.add_argument( + "--reward-function", + default=None, + help="RL reward callable as package.module:function", + ) + rl.add_argument( + "--advantage-estimator", choices=["grpo"], default="grpo" + ) + rl.add_argument("--n-samples-per-prompt", type=int, default=4) + rl.add_argument("--rollout-batch-size", type=int, default=32) + rl.add_argument("--over-sampling-batch-size", type=int, default=None) + rl.add_argument("--rollout-max-response-len", type=int, default=32768) + rl.add_argument("--custom-generate-function-path", default=None) + rl.add_argument("--use-session-server", action="store_true") + rl.add_argument("--session-server-ip", default=None) + rl.add_argument("--session-server-port", type=int, nargs="+", default=None) + rl.add_argument("--tito-model", default=None) + rl.add_argument("--local-rl-rounds-per-sync", type=int, default=1) + rl.add_argument( + "--rl-sync-preset", choices=["strict-avg"], default="strict-avg" + ) + rl.add_argument( + "--rl-policy-version", choices=["strict"], default="strict" + ) + rl.add_argument( + "--rl-completed-groups-path", + default="~/yeto-rl/island-checkpoint.pt", + ) + rl.add_argument("--experimental-rl-sync", action="store_true") p.add_argument( "--output", default=None, diff --git a/yeto/export.py b/yeto/export.py index d02c9e1..9ac9e81 100644 --- a/yeto/export.py +++ b/yeto/export.py @@ -20,6 +20,7 @@ numel x f32 momentum ledger_count u32 per entry: learner_id u32, merges u64, steps u64, tokens u64 + optional: 32-byte semantic layout hash (new checkpoints) The outer momentum is parsed (and validated) but not needed for export; it only matters when the syncer itself resumes from the checkpoint. @@ -49,6 +50,8 @@ class Checkpoint: fragments: list[tuple[int, torch.Tensor, torch.Tensor]] # learner_id -> (merges, steps, tokens) ledger: dict[int, tuple[int, int, int]] + # HELLO's semantic layout fingerprint, persisted by current syncers. + layout_hash: str | None # Digest of the exact byte buffer decoded into this checkpoint. sha256: str @@ -94,13 +97,22 @@ def take(n: int, what: str) -> bytes: ) ledger[learner_id] = (merges, steps, tokens) + layout_hash = None + if len(data) - off == 32: + layout_hash = take(32, "layout_hash").hex() if off != len(data): raise ValueError( f"{path}: {len(data) - off} trailing bytes after the checkpoint " f"payload (parsed {off} of {len(data)}); file corrupt or from an " "incompatible syncer version" ) - return Checkpoint(global_step, fragments, ledger, hashlib.sha256(data).hexdigest()) + return Checkpoint( + global_step, + fragments, + ledger, + layout_hash, + hashlib.sha256(data).hexdigest(), + ) def _read_f32(raw: bytes) -> torch.Tensor: @@ -129,6 +141,14 @@ def validate_against_layout(ckpt: Checkpoint, layout: FragmentLayout) -> None: f"fragment {fid}: checkpoint numel {params.numel()} " f"!= layout numel {frag.numel}" ) + if ckpt.layout_hash is not None: + from .protocol import layout_fingerprint + + rebuilt = layout_fingerprint(layout).hex() + if ckpt.layout_hash != rebuilt: + problems.append( + f"checkpoint layout hash {ckpt.layout_hash} != rebuilt {rebuilt}" + ) if problems: raise ValueError( "checkpoint does not match the rebuilt fragment layout; make sure " diff --git a/yeto/launcher.py b/yeto/launcher.py index 92337a2..82d858b 100644 --- a/yeto/launcher.py +++ b/yeto/launcher.py @@ -20,14 +20,18 @@ from __future__ import annotations +import hashlib +import json import os +import re import shlex import signal import subprocess import sys import threading import time -from pathlib import Path +from dataclasses import dataclass +from pathlib import Path, PurePosixPath from . import delivery from .gpu_spec import ClusterSpec, parse_gpu_spec @@ -225,6 +229,27 @@ def syncer_command(args, num_learners: int, binary: str = "~/yeto-syncer") -> st """The syncer invocation shared by the syncer-cluster task (local controller mode) and the head-node subprocess (head controller mode). --resume makes any restart pick up from the on-disk checkpoint.""" + if getattr(args, "training_mode", "sft") == "rl": + return ( + "mkdir -p ~/yeto-output && " + f"{binary}" + f" --port {SYNCER_PORT}" + f" --learners {num_learners}" + f" --quorum {args.quorum}" + f" --grace-ms {args.grace_ms}" + f" --grace-gamma {args.grace_gamma}" + f" --grace-tau {args.grace_tau}" + f" --pipeline {args.pipeline}" + f" --sync-interval-steps {args.sync_interval_steps}" + f" --delta-correction {args.delta_correction}" + f" --total-steps {args.total_steps}" + f" --outer-lr {args.outer_lr}" + f" --outer-momentum {args.outer_momentum}" + " --max-base-lag 0 --learner-weight equal" + " --checkpoint-path ~/yeto-output/yeto-state.ckpt" + " --checkpoint-every 1 --resume" + " --event-tape ~/yeto-output/yeto-tape.jsonl" + ) return ( f"{binary}" f" --port {SYNCER_PORT}" @@ -477,6 +502,162 @@ def prepare_launch_args(args) -> None: f"{expected_adapter_sha256.lower()}, got {adapter_sha256}" ) args.diffusion_adapter_sha256 = adapter_sha256 + _prepare_rl_args(args) + + +def _rl_callable(value: str | None, flag: str, *, required: bool) -> None: + if value is None and not required: + return + module, separator, function = (value or "").partition(":") + if ( + not separator + or not module + or module.endswith(".py") + or not function.isidentifier() + ): + raise ValueError(f"{flag} must be package.module:function") + + +def _rl_miles_function(value: str | None) -> None: + if value is None: + return + parts = value.split(".") + if len(parts) < 2 or any(not part.isidentifier() for part in parts): + raise ValueError( + "--custom-generate-function-path must be package.module.function" + ) + + +def _prepare_rl_args(args) -> None: + if getattr(args, "training_mode", "sft") != "rl": + return + + from .models import resolve_model_kind + + if resolve_model_kind(args.model, args.model_kind) != "causal-lm": + raise ValueError("RL v0 supports only causal language models") + if args.tuning != "lora": + raise ValueError("RL v0 requires --tuning lora") + if args.lora_r <= 0: + raise ValueError("RL v0 requires a positive LoRA rank") + if args.total_steps <= 0: + raise ValueError("RL v0 requires --total-steps > 0") + if args.rollout_batch_size <= 0 or args.n_samples_per_prompt <= 0: + raise ValueError("RL v0 requires positive rollout batch and sample counts") + if args.rollout_max_response_len <= 0: + raise ValueError("RL v0 requires --rollout-max-response-len > 0") + if args.local_rl_rounds_per_sync != 1: + raise ValueError("RL v0 requires --local-rl-rounds-per-sync 1") + if args.seq_len < 2: + raise ValueError("RL v0 requires --seq-len >= 2") + args.seq_len = max(args.seq_len, args.rollout_max_response_len) + if args.over_sampling_batch_size is None: + args.over_sampling_batch_size = args.rollout_batch_size + elif args.over_sampling_batch_size < args.rollout_batch_size: + raise ValueError( + "RL --over-sampling-batch-size must be at least --rollout-batch-size" + ) + _rl_miles_function(args.custom_generate_function_path) + if not args.use_session_server and ( + args.session_server_ip is not None + or args.session_server_port is not None + or args.tito_model is not None + ): + raise ValueError( + "--session-server-ip/--session-server-port/--tito-model requires " + "--use-session-server" + ) + if args.session_server_port is not None and ( + len(args.session_server_port) not in {1, 2} + or any(port <= 0 or port > 65535 for port in args.session_server_port) + or ( + len(args.session_server_port) == 2 + and args.session_server_port[1] <= args.session_server_port[0] + ) + ): + raise ValueError( + "--session-server-port requires one positive port or an increasing range" + ) + + specs = parse_gpu_spec(args.gpu) + if getattr(args, "external_learners", 0): + raise ValueError("RL v0 does not support external learner slots") + if args.tensor_parallel != 1 or args.pipeline_parallel != 1: + raise ValueError("RL v0 requires TP=PP=1") + if args.expert_parallel is not None: + if args.expert_parallel <= 0: + raise ValueError("RL expert parallelism must be positive") + if any(spec.total_gpus % args.expert_parallel for spec in specs): + raise ValueError("RL expert parallelism must divide every island") + for spec in specs: + dp = spec.total_gpus # TP=PP=CP=1 in the fixed Miles/Megatron path. + if args.rollout_batch_size * args.n_samples_per_prompt % dp: + raise ValueError( + "RL rollout_batch_size*n_samples_per_prompt must be divisible " + "by every island data-parallel size" + ) + + _rl_callable(args.reward_function, "--reward-function", required=True) + if not re.fullmatch( + r"docker:[^\s@]+@sha256:[0-9a-fA-F]{64}", args.rl_image or "" + ): + raise ValueError( + "--rl-image must be docker:@sha256:<64 hex digest>" + ) + if getattr(args, "learner_image", None) is not None: + raise ValueError("RL uses digest-pinned --rl-image, not --learner-image") + if not args.experimental_rl_sync: + args.fragments = 1 + args.quorum = len(specs) + args.grace_ms = 0 + args.pipeline = 1 + args.sync_interval_steps = 0.0 + args.delta_correction = "none" + args.outer_lr = 1.0 + args.outer_momentum = 0.0 + args.merge_alpha = 0.0 + args.wire_dtype = "f32" + else: + for name, expected, flag in ( + ("fragments", 1, "--fragments 1"), + ("pipeline", 1, "--pipeline 1"), + ("merge_alpha", 0.0, "--merge-alpha 0"), + ("wire_dtype", "f32", "--wire-dtype f32"), + ): + if getattr(args, name) != expected: + raise ValueError( + f"RL's one-fragment f32 bridge still requires {flag}" + ) + if args.spot: + _rl_checkpoint_mount(args.rl_completed_groups_path) + + provenance = getattr(args, "_provenance", None) + if provenance: + for name in ("model", "dataset"): + source = provenance.get(name) or {} + if source.get("source") != "huggingface" or not source.get( + "resolved_revision" + ): + raise ValueError( + f"RL v0 requires a revision-pinned Hugging Face {name}" + ) + from .provenance import python_spec_path, python_spec_sha256 + + reward_path = python_spec_path(args.reward_function, base_dir=REPO_ROOT) + try: + reward_path.relative_to(REPO_ROOT.resolve()) + except ValueError as exc: + raise ValueError( + "RL reward source must be inside the synced Yeto workdir" + ) from exc + args.reward_sha256 = python_spec_sha256( + args.reward_function, base_dir=REPO_ROOT + ) + # The fixed Miles commit enables remote model code in its internal + # Megatron and SGLang loaders. Keep that trust decision explicit through + # Yeto's existing flag. + if not args.trust_remote_code: + raise ValueError("pinned Miles requires explicit --trust-remote-code") # AWS keeps this SSM parameter pointing at the CURRENT Deep Learning Base @@ -640,6 +821,185 @@ def causal_kernel_setup_steps(args) -> list[str]: DIFFUSION_SAMPLE_OUTPUT_DIR = "~/yeto-output" +def _rl_checkpoint_mount(value: str) -> str: + path = PurePosixPath(value) + if ( + not (value.startswith("~/") or value.startswith("/")) + or value.endswith("/") + or ".." in path.parts + or str(path.parent) in {".", "~", "/"} + ): + raise ValueError( + "Spot RL --rl-completed-groups-path must be a file in an " + "absolute or ~/ subdirectory" + ) + return str(path.parent) + + +def _rl_checkpoint_storage_name(cluster_prefix: str, learner_id: int) -> str: + stem = re.sub(r"[^a-z0-9-]+", "-", cluster_prefix.lower()).strip("-") or "yeto" + suffix = f"-{hashlib.sha256(cluster_prefix.encode()).hexdigest()[:8]}-rl-{learner_id}" + return stem[: 63 - len(suffix)].rstrip("-") + suffix + + +def make_miles_island_task( + args, + spec: ClusterSpec, + learner_id: int, + num_learners: int, + syncer_addr: str, +): + """Create one Ray/Miles island from the pinned Miles checkout.""" + + import sky + + from .datasource import learner_data_arg, learner_file_mounts + from .models import resolve + from .provenance import is_local_reference + from .rl import MILES_COMMIT, MILES_PEFT_VERSION, MILES_REPOSITORY + + if not getattr(args, "source_sha256", None) or not getattr( + args, "reward_sha256", None + ): + raise ValueError("RL task requires prepared source and reward provenance") + + flags = ( + f" --model {shlex.quote(args.model)}" + f" --data {shlex.quote(learner_data_arg(args.data))}" + " --syncer $SYNCER_ADDR" + " --learner-id $LEARNER_ID" + f" --reward-function {shlex.quote(args.reward_function)}" + f" --reward-sha256 {shlex.quote(args.reward_sha256)}" + f" --source-sha256 {shlex.quote(args.source_sha256)}" + f" --global-rounds {args.total_steps}" + f" --groups-per-round {args.rollout_batch_size}" + f" --samples-per-group {args.n_samples_per_prompt}" + f" --over-sampling-batch-size {args.over_sampling_batch_size}" + " --optimizer-steps 1" + f" --rollout-max-response-len {args.rollout_max_response_len}" + f" --completed-groups-path {shlex.quote(args.rl_completed_groups_path)}" + f" --event-tape ~/yeto-output/rl-island-{learner_id}.jsonl" + f" --actor-num-nodes {spec.num_nodes}" + f" --actor-num-gpus-per-node {spec.gpus_per_node}" + f" --lora-r {args.lora_r}" + f" --lora-targets {args.lora_targets}" + f" --inner-lr {args.inner_lr}" + f" --seq-len {args.seq_len}" + f" --seed {args.seed}" + f" --wan-streams {args.wan_streams}" + " --miles-root ~/miles" + ) + if args.expert_parallel is not None: + flags += f" --expert-parallel {args.expert_parallel}" + if args.custom_generate_function_path: + flags += ( + " --custom-generate-function-path " + f"{shlex.quote(args.custom_generate_function_path)}" + ) + if args.use_session_server: + flags += " --use-session-server" + if args.session_server_ip: + flags += f" --session-server-ip {shlex.quote(args.session_server_ip)}" + if args.session_server_port: + ports = " ".join(str(port) for port in args.session_server_port) + flags += f" --session-server-port {ports}" + if args.tito_model: + flags += f" --tito-model {shlex.quote(args.tito_model)}" + if args.model_revision: + flags += f" --model-revision {shlex.quote(args.model_revision)}" + if args.data_revision: + flags += f" --data-revision {shlex.quote(args.data_revision)}" + if args.trust_remote_code: + flags += " --trust-remote-code" + miles_setup = ( + "set -e\n" + f"if [ ! -d ~/miles/.git ]; then git clone --no-checkout " + f"{shlex.quote(MILES_REPOSITORY)} ~/miles; fi\n" + f"git -C ~/miles fetch --depth 1 origin {MILES_COMMIT}\n" + f"git -C ~/miles checkout --detach {MILES_COMMIT}\n" + f"python3 -m pip install -q --no-deps -e ~/miles " + f"'peft=={MILES_PEFT_VERSION}'" + ) + model = resolve(args.model) + if is_local_reference(model): + prefetch = ": # local model; no Hub prefetch" + else: + revision = ( + f" --revision {shlex.quote(args.model_revision)}" + if args.model_revision + else "" + ) + prefetch = ( + f"(nohup huggingface-cli download {shlex.quote(model)}{revision} " + ">/tmp/hf-prefetch.log 2>&1 &) || true" + ) + file_mounts = dict(learner_file_mounts(args.data)) + local_token = os.path.expanduser(HF_TOKEN_PATH) + if os.path.isfile(local_token): + file_mounts[HF_TOKEN_PATH] = local_token + envs = { + "SYNCER_ADDR": syncer_addr, + "LEARNER_ID": str(learner_id), + "HF_HUB_ENABLE_HF_TRANSFER": "1", + "NVTE_FLASH_ATTN": "0", + "NVTE_FUSED_ATTN": "0", + "NVTE_UNFUSED_ATTN": "1", + } + if os.environ.get("HF_TOKEN"): + envs["HF_TOKEN"] = os.environ["HF_TOKEN"] + task = sky.Task( + name=f"yeto-rl-island-{learner_id}", + setup="\n".join((WAN_TUNING, HF_TOKEN_ENV, miles_setup, prefetch)), + run=( + f"{HF_TOKEN_ENV}\n" + "set -e\n" + "cd ~/sky_workdir\n" + 'MASTER_ADDR=$(echo "$SKYPILOT_NODE_IPS" | head -n1)\n' + "ray stop --force >/dev/null 2>&1 || true\n" + 'if [ "$SKYPILOT_NODE_RANK" = "0" ]; then\n' + " ray start --head --node-ip-address=\"$MASTER_ADDR\" " + "--port=6379 --include-dashboard=false\n" + " trap 'ray stop --force >/dev/null 2>&1 || true' EXIT\n" + " PYTHONPATH=~/sky_workdir " + f"python3 -m yeto.rl.learner{flags}\n" + "else\n" + " until ray start --address=\"$MASTER_ADDR:6379\"; do sleep 2; done\n" + " while ray status --address=\"$MASTER_ADDR:6379\" " + ">/dev/null 2>&1; do sleep 5; done\n" + "fi" + ), + envs=envs, + num_nodes=spec.num_nodes, + workdir=str(REPO_ROOT), + file_mounts=file_mounts or None, + ) + resources = { + "infra": f"{spec.cloud}/{spec.region}" if spec.region else spec.cloud, + "accelerators": spec.accelerators, + "cpus": args.learner_cpus, + "instance_type": args.learner_instance_type, + "use_spot": args.spot, + "disk_size": args.disk_size, + } + resources["image_id"] = args.rl_image + if spec.num_nodes > 1: + resources["network_tier"] = "best" + task.set_resources(sky.Resources(**resources)) + if args.spot: + checkpoint_mount = _rl_checkpoint_mount(args.rl_completed_groups_path) + task.set_storage_mounts( + { + checkpoint_mount: sky.Storage( + name=_rl_checkpoint_storage_name(args.cluster_prefix, learner_id), + persistent=False, + mode=sky.StorageMode.MOUNT, + sync_on_reconstruction=True, + ) + } + ) + return task + + def make_learner_task(args, spec: ClusterSpec, learner_id: int, num_learners: int, syncer_addr: str): import sky @@ -1254,6 +1614,27 @@ def cluster_up(self, cluster: str) -> bool: ) return status == sky.ClusterStatus.UP + def rl_strict_failure(self, cluster: str, job_id: int) -> str | None: + """Return a strict RL event from a failed job's existing log.""" + + import sky + + try: + lines = sky.tail_logs( + cluster, job_id, follow=False, preload_content=False + ) + for line in lines: + text = str(line).strip() + if ( + "[yeto-rl-strict-failure]" in text + or "RL strict failure " in text + or "StrictRlInvariantError:" in text + ): + return text + except Exception: + return None + return None + def relaunch(self, task, cluster: str): """Re-provision `cluster` (same spec) and submit `task` as a new job. @@ -1303,10 +1684,14 @@ def __init__( self.command = syncer_command(args, num_learners, binary=binary) self.binary = os.path.expanduser(binary) self.log_file = os.path.expanduser(log_file) + self.event_tape = os.path.expanduser("~/yeto-output/yeto-tape.jsonl") self.proc: subprocess.Popen | None = None # The log file persists across controller jobs on a reused head; # forward only what this controller's syncer writes, not history. self._log_offset = os.path.getsize(self.log_file) if os.path.exists(self.log_file) else 0 + self._event_offset = ( + os.path.getsize(self.event_tape) if os.path.exists(self.event_tape) else 0 + ) def start(self) -> None: if os.path.exists(self.binary): @@ -1330,7 +1715,7 @@ def start(self) -> None: log_f.close() # Popen holds its own duplicate of the fd print(f"[launcher] syncer subprocess started (pid {self.proc.pid})", flush=True) - def probe(self) -> str | None: + def probe(self) -> str | _RlStrictFailure | None: """None if the subprocess is healthy, else a reason string. Exit code 0 means the syncer completed its total steps — terminal @@ -1342,8 +1727,26 @@ def probe(self) -> str | None: code = self.proc.poll() if code is None or code == 0: return None + strict_failure = self._strict_failure() + if strict_failure is not None: + return strict_failure return f"syncer subprocess exited with code {code}" + def _strict_failure(self) -> "_RlStrictFailure | None": + try: + with open(self.event_tape, encoding="utf-8") as handle: + handle.seek(self._event_offset) + for line in handle: + event = json.loads(line) + if event.get("event") == "rl_strict_failure": + return _RlStrictFailure( + f"{event.get('metric', 'strict_failure')}: " + f"{event.get('error', 'syncer strict invariant failed')}" + ) + except (OSError, ValueError, TypeError): + return None + return None + def restart(self) -> None: self.start() @@ -1391,6 +1794,11 @@ def _forward(): ABANDONED = "abandoned" +@dataclass(frozen=True) +class _RlStrictFailure: + reason: str + + class _RelaunchAttempt: """Result slot for one background relaunch attempt.""" @@ -1447,6 +1855,7 @@ def __init__( thread_cls=threading.Thread, syncer_probe=None, syncer_restart=None, + fixed_roster: bool = False, ): """`learners` maps cluster name -> (task, job_id); `syncer` is (name, task, job_id) for a cluster syncer, or None with @@ -1458,6 +1867,7 @@ def __init__( self.recover_timeout = recover_timeout self.on_relaunch = on_relaunch self.thread_cls = thread_cls + self.fixed_roster = fixed_roster self.learners = { name: self._make_record(name, task, job_id) for name, (task, job_id) in learners.items() @@ -1525,6 +1935,8 @@ def _poll_local_syncer(self) -> None: reason = self.syncer_probe() if reason is None: return + if isinstance(reason, _RlStrictFailure): + raise RuntimeError(f"strict RL syncer failed: {reason.reason}") print( f"[launcher] syncer: {reason}; restarting the local syncer " "(resumes from its checkpoint)", @@ -1548,6 +1960,11 @@ def _poll(self, rec, is_syncer: bool) -> None: rec["exit"] = str(status) print(f"[launcher] {rec['name']} job finished: {status}") else: + strict_failure = self._strict_failure(rec) + if strict_failure is not None: + raise RuntimeError( + f"strict RL job {rec['name']} failed: {strict_failure}" + ) self._enter_recovering(rec, verdict, is_syncer) elif rec["state"] == RECOVERING: self._drive_recovery(rec, is_syncer) @@ -1572,6 +1989,17 @@ def _probe(self, rec): return "cluster is not UP (preempted or deleted)", status return None, status + def _strict_failure(self, rec) -> str | None: + if not self.fixed_roster: + return None + probe = getattr(self.ops, "rl_strict_failure", None) + if probe is None: + return None + try: + return probe(rec["name"], rec["job_id"]) + except Exception: + return None + def _enter_recovering(self, rec, reason: str, is_syncer: bool) -> None: rec["state"] = RECOVERING rec["failed_at"] = self.ops.now() @@ -1647,6 +2075,13 @@ def _abandon(self, rec, elapsed: float) -> None: rec["state"] = ABANDONED rec["exit"] = f"ABANDONED after {elapsed:.0f}s" self._down(rec["name"]) + if self.fixed_roster: + message = ( + f"fixed-roster learner {rec['name']} could not recover within " + f"{self.recover_timeout}s" + ) + print(f"[launcher] ERROR: {message}", file=sys.stderr) + raise RuntimeError(message) remaining = sum(1 for r in self.learners.values() if r["state"] != ABANDONED) print( f"[launcher] LEARNER {rec['name']} ABANDONED after {elapsed:.0f}s " @@ -1854,9 +2289,14 @@ def run(args, on_clusters=None, local_syncer=None) -> int: # 2. Learners, in parallel. tasks = {} rids = {} + task_factory = ( + make_miles_island_task + if getattr(args, "training_mode", "sft") == "rl" + else make_learner_task + ) for m, spec in enumerate(specs): name = learner_names[m] - task = make_learner_task(args, spec, m, num_learners, syncer_addr) + task = task_factory(args, spec, m, num_learners, syncer_addr) tasks[name] = task print(f"[launcher] launching learner {m} on {spec} as {name}") rids[name] = ( @@ -1910,6 +2350,7 @@ def spawn_tail(name: str, job_id: int) -> None: on_relaunch=spawn_tail, syncer_probe=local_syncer.probe if head_mode else None, syncer_restart=local_syncer.restart if head_mode else None, + fixed_roster=getattr(args, "training_mode", "sft") == "rl", ) exit_codes = controller.run() failed = [n for n, s in exit_codes.items() if "SUCCEEDED" not in s] @@ -1922,7 +2363,12 @@ def spawn_tail(name: str, job_id: int) -> None: print("[launcher] no learner succeeded; recover from the syncer " "checkpoint with yeto-export", file=sys.stderr) return 1 - source = next((n for n in done if "-l0-" in n), done[0]) + rl_mode = getattr(args, "training_mode", "sft") == "rl" + source = ( + syncer_cluster + if rl_mode and not head_mode + else next((n for n in done if "-l0-" in n), done[0]) + ) output = getattr(args, "output", None) local_dest = ( os.path.expanduser(output) @@ -1930,16 +2376,20 @@ def spawn_tail(name: str, job_id: int) -> None: else os.path.expanduser("~/yeto-output") ) os.makedirs(local_dest, exist_ok=True) - try: - subprocess.run(delivery.fetch_cmd(source, local_dest), check=True) - print(f"[launcher] fine-tuned model fetched to {local_dest}") - except subprocess.CalledProcessError as e: - print( - f"[launcher] fetching {source}:~/yeto-output failed ({e}); " - "recover from the syncer checkpoint with yeto-export", - file=sys.stderr, - ) - return 2 + if rl_mode and head_mode: + print(f"[launcher] committed RL checkpoint retained at {local_dest}") + else: + try: + subprocess.run(delivery.fetch_cmd(source, local_dest), check=True) + artifact = "committed RL checkpoint" if rl_mode else "fine-tuned model" + print(f"[launcher] {artifact} fetched to {local_dest}") + except subprocess.CalledProcessError as e: + print( + f"[launcher] fetching {source}:~/yeto-output failed ({e}); " + "recover from the syncer checkpoint with yeto-export", + file=sys.stderr, + ) + return 2 if delivery.is_remote(output): try: delivery.deliver(output, local_dest) diff --git a/yeto/rl/__init__.py b/yeto/rl/__init__.py new file mode 100644 index 0000000..e7e86aa --- /dev/null +++ b/yeto/rl/__init__.py @@ -0,0 +1,9 @@ +"""Pinned Miles reinforcement-learning integration.""" + +MILES_REPOSITORY = "https://github.com/radixark/miles" +MILES_COMMIT = "dfc66ff38752bfa2c5d325e0037ebc4b537c06de" +MILES_PEFT_VERSION = "0.20.0" +MILES_IMAGE = ( + "docker:radixark/miles@sha256:" + "95b3afa9ee4313f5633e6ed3779c8276353cc8e24a2462e4f54ec0d5978fbae7" +) diff --git a/yeto/rl/bridge.py b/yeto/rl/bridge.py new file mode 100644 index 0000000..1f83556 --- /dev/null +++ b/yeto/rl/bridge.py @@ -0,0 +1,327 @@ +"""One fixed-roster RL island's PULL/local/PUSH/BCAST loop.""" + +from __future__ import annotations + +import json +import sys +import time +from dataclasses import asdict, dataclass, replace +from pathlib import Path +from typing import Protocol + +from ..protocol import DTYPE_F32, PullRequest, SyncerClient +from ..tensor_io import pack_tensor, unpack_fragment +from .core import ( + CanonicalLoraState, + CanonicalTensorSpec, + LocalRoundStats, + StrictRlInvariantError, + build_avg_layout, + canonical_state, + flat_tensor, + policy_delta, + tensors_from_flat, +) + + +class IslandRuntime(Protocol): + def initialize(self) -> CanonicalLoraState: ... + + def apply_global_policy(self, state: CanonicalLoraState) -> None: ... + + def run_local_round( + self, + *, + expected_policy_version: int, + groups: int, + samples_per_group: int, + optimizer_steps: int, + ) -> LocalRoundStats: ... + + def export_local_policy(self) -> CanonicalLoraState: ... + + def record_local_round(self, stats: LocalRoundStats) -> None: ... + + def shutdown(self) -> None: ... + + +@dataclass(frozen=True) +class BridgeConfig: + syncer_addr: tuple[str, int] + learner_id: int + global_rounds: int + groups_per_round: int + samples_per_group: int + local_optimizer_steps: int + expected_specs: tuple[CanonicalTensorSpec, ...] + base_model_revision: str + lora_config_hash: str + layout_hash: str + event_tape: str + wan_streams: int = 4 + + +class StrictRlBridge: + def __init__(self, runtime: IslandRuntime, config: BridgeConfig) -> None: + self.runtime = runtime + self.config = config + initialized = runtime.initialize() + if ( + initialized.base_model_revision, + initialized.lora_config_hash, + initialized.layout_hash, + ) != ( + config.base_model_revision, + config.lora_config_hash, + config.layout_hash, + ): + raise StrictRlInvariantError( + "layout_hash_mismatch", + "Miles initialized a different canonical LoRA identity", + ) + self.initial = canonical_state( + 0, + initialized.tensors, + base_model_revision=initialized.base_model_revision, + lora_config_hash=initialized.lora_config_hash, + layout_hash=initialized.layout_hash, + expected_specs=config.expected_specs, + ) + self.specs = self.initial.specs + self.layout = build_avg_layout(self.specs) + self.client = SyncerClient( + config.syncer_addr, + config.learner_id, + self.layout, + dtype=DTYPE_F32, + num_streams=config.wan_streams, + # A dead syncer connection makes this island exit. The launcher + # restarts the same logical ID, which reapplies the committed cut + # and recomputes any uncommitted local result. + max_reconnects=0, + ) + self.current: CanonicalLoraState | None = None + self.permits: dict[int, PullRequest] = {} + self.pushed_step: int | None = None + + def run(self) -> CanonicalLoraState: + try: + self.client.start() + if self.config.learner_id == 0: + self.client.send_init( + 0, + pack_tensor( + flat_tensor(self.initial.tensors, self.specs), + DTYPE_F32, + ), + ) + while True: + self.client.check_health() + if self.client.finalizing.is_set(): + return self._finalize() + progressed = self._drain_messages() + permit = self._ready_permit() + if permit is not None: + self._run_round(permit) + progressed = True + if not progressed: + time.sleep(0.05) + except StrictRlInvariantError as error: + self._append_event( + { + "event": "rl_strict_failure", + "metric": error.metric, + "value": 1, + "error": f"{type(error).__name__}: {error}", + } + ) + print( + f"[yeto-rl-strict-failure] {error.metric}: {error}", + file=sys.stderr, + flush=True, + ) + raise + finally: + try: + self.runtime.shutdown() + finally: + self.client.close() + + def _drain_messages(self) -> bool: + progressed = False + for update in self.client.drain_updates(): + progressed = True + if update.fragment_id != 0: + raise RuntimeError("RL received a nonzero fragment") + if self.current is not None: + if update.version < self.current.policy_version: + continue + if update.version == self.current.policy_version: + continue + if update.version != self.current.policy_version + 1: + raise RuntimeError( + f"RL policy jumped from {self.current.policy_version} " + f"to {update.version}" + ) + state = self._state_from_payload(update.version, update.data) + self.runtime.apply_global_policy(state) + self.current = state + self.pushed_step = None + self.permits = { + step: permit + for step, permit in self.permits.items() + if step > state.policy_version + } + + for permit in self.client.drain_pulls(): + progressed = True + if permit.fragment_id != 0 or permit.round_attempt != 1: + raise RuntimeError("RL received an invalid PULL permit") + if permit.global_step > self.config.global_rounds: + raise RuntimeError("RL received a PULL beyond configured rounds") + current_version = self.current.policy_version if self.current else -1 + if permit.global_step <= current_version: + continue + previous = self.permits.get(permit.global_step) + if previous is not None and previous != permit: + raise RuntimeError("RL received conflicting PULL permits") + self.permits[permit.global_step] = permit + return progressed + + def _state_from_payload(self, version: int, data: bytes) -> CanonicalLoraState: + try: + flat = unpack_fragment(self.layout.fragments[0], data, DTYPE_F32) + return canonical_state( + version, + tensors_from_flat(flat, self.specs), + base_model_revision=self.config.base_model_revision, + lora_config_hash=self.config.lora_config_hash, + layout_hash=self.config.layout_hash, + ) + except (TypeError, ValueError) as error: + metric = ( + "nonfinite_delta_count" + if "NaN or Inf" in str(error) + else "layout_hash_mismatch" + ) + raise StrictRlInvariantError(metric, str(error)) from error + + def _ready_permit(self) -> PullRequest | None: + if self.current is None: + return None + target = self.current.policy_version + 1 + if self.pushed_step == target: + self.permits.pop(target, None) + return None + return self.permits.pop(target, None) + + def _run_round(self, permit: PullRequest) -> None: + base = self.current + if base is None or permit.global_step != base.policy_version + 1: + raise RuntimeError("RL attempted a local round without its exact base") + stats = self.runtime.run_local_round( + expected_policy_version=base.policy_version, + groups=self.config.groups_per_round, + samples_per_group=self.config.samples_per_group, + optimizer_steps=self.config.local_optimizer_steps, + ) + if ( + stats.island_id != self.config.learner_id + or stats.base_policy_version != base.policy_version + or stats.local_round_id != permit.global_step + ): + raise RuntimeError("Miles returned LocalRoundStats for a different round") + + try: + exported = self.runtime.export_local_policy() + local = canonical_state( + exported.policy_version, + exported.tensors, + base_model_revision=exported.base_model_revision, + lora_config_hash=exported.lora_config_hash, + layout_hash=exported.layout_hash, + expected_specs=self.specs, + ) + delta = policy_delta(local, base) + except ValueError as error: + message = str(error) + if "NaN or Inf" in message: + metric = "nonfinite_delta_count" + elif any( + value in message + for value in ("layout", "names, shapes, or dtypes") + ): + metric = "layout_hash_mismatch" + else: + raise + raise StrictRlInvariantError(metric, message) from error + stats = replace(stats, delta_l2_norm=float(delta.norm().item())) + self.runtime.record_local_round(stats) + payload = pack_tensor(delta, DTYPE_F32) + self._append_event( + { + "event": "rl_local_round", + **asdict(stats), + "rl/active_groups": stats.active_groups, + "rl/completed_groups": stats.completed_groups, + "rl/cancelled_groups": stats.cancelled_groups, + "rl/completed_trajectories": stats.completed_trajectories, + "rl/action_tokens": stats.action_tokens, + "rl/tool_wait_seconds": stats.tool_wait_seconds, + "rl/reward_mean": stats.reward_mean, + "rl/reward_std": stats.reward_std, + "rl/rollout_seconds": stats.rollout_seconds, + "rl/group_p50_seconds": stats.group_p50_seconds, + "rl/group_p95_seconds": stats.group_p95_seconds, + "rl/group_p99_seconds": stats.group_p99_seconds, + "rl/zero_variance_group_ratio": stats.zero_variance_group_ratio, + "rl/global_policy_version": stats.base_policy_version, + "rl/rollout_policy_version": stats.base_policy_version, + "rl/mixed_version_group_count": 0, + "rl/local_delta_norm": stats.delta_l2_norm, + "rl/current_vs_rollout_kl": stats.mean_kl, + "rl/ess_ratio": stats.ess_ratio, + "rl/clip_fraction": stats.clip_fraction, + "sync/bytes_sent": 48 + len(payload), + } + ) + self.client.push_fragment( + 0, + permit.global_step, + permit.round_attempt, + base.policy_version, + permit.global_step, + 1, + 1, + payload, + ) + self.pushed_step = permit.global_step + + def _finalize(self) -> CanonicalLoraState: + manifest, fragments = self.client.wait_for_final_fragments() + if ( + manifest.global_step != self.config.global_rounds + or manifest.versions != (manifest.global_step,) + or len(fragments) != 1 + ): + raise RuntimeError("RL received an inconsistent final checkpoint cut") + final = self._state_from_payload( + manifest.global_step, + fragments[0].data, + ) + self.runtime.apply_global_policy(final) + self.client.acknowledge_finalization(manifest) + return final + + def _append_event(self, event: dict) -> None: + path = Path(self.config.event_tape).expanduser() + path.parent.mkdir(parents=True, exist_ok=True) + event = { + "island_id": self.config.learner_id, + "time_unix": time.time(), + **event, + } + with path.open("a", encoding="utf-8") as handle: + handle.write( + json.dumps(event, sort_keys=True, separators=(",", ":")) + "\n" + ) diff --git a/yeto/rl/core.py b/yeto/rl/core.py new file mode 100644 index 0000000..3686bca --- /dev/null +++ b/yeto/rl/core.py @@ -0,0 +1,296 @@ +"""Canonical PEFT LoRA values at the Yeto/Miles synchronization boundary.""" + +from __future__ import annotations + +import math +import re +import hashlib +import json +from collections.abc import Mapping, Sequence +from dataclasses import dataclass + +import torch + +from ..fragments import MERGE_AVG, Fragment, FragmentLayout +from ..protocol import layout_fingerprint + +_PEFT_LORA_NAME = re.compile(r"\.lora_(?:A|B)\.weight\Z") + + +class StrictRlInvariantError(RuntimeError): + """A deterministic INIT strict-run invariant violation.""" + + def __init__(self, metric: str, message: str) -> None: + super().__init__(message) + self.metric = metric + + +@dataclass(frozen=True, order=True) +class CanonicalTensorSpec: + name: str + shape: tuple[int, ...] + dtype: str + numel: int + + def __post_init__(self) -> None: + if not self.name or not _PEFT_LORA_NAME.search(self.name): + raise ValueError(f"not a canonical PEFT LoRA tensor name: {self.name!r}") + if not self.shape or any(dim <= 0 for dim in self.shape): + raise ValueError(f"invalid shape for {self.name!r}: {self.shape}") + if self.dtype != "float32": + raise ValueError(f"canonical LoRA dtype must be float32, got {self.dtype!r}") + if math.prod(self.shape) != self.numel: + raise ValueError(f"shape/numel mismatch for {self.name!r}") + + +@dataclass(frozen=True) +class CanonicalLoraState: + base_model_revision: str + lora_config_hash: str + layout_hash: str + policy_version: int + tensors: Mapping[str, torch.Tensor] + + def __post_init__(self) -> None: + if not re.fullmatch(r"[0-9a-f]{40}", self.base_model_revision): + raise ValueError("base model revision must be an immutable commit") + for name, value in ( + ("LoRA config", self.lora_config_hash), + ("layout", self.layout_hash), + ): + if not re.fullmatch(r"[0-9a-f]{64}", value): + raise ValueError(f"{name} hash must be a lowercase SHA256") + if self.policy_version < 0: + raise ValueError("policy version must be non-negative") + + @property + def specs(self) -> tuple[CanonicalTensorSpec, ...]: + return canonical_specs(self.tensors) + + +def canonical_specs( + tensors: Mapping[str, torch.Tensor], +) -> tuple[CanonicalTensorSpec, ...]: + if not tensors: + raise ValueError("canonical LoRA state is empty") + return tuple( + CanonicalTensorSpec( + name, + tuple(int(dim) for dim in tensor.shape), + "float32", + tensor.numel(), + ) + for name, tensor in sorted(tensors.items()) + ) + + +def build_avg_layout(specs: Sequence[CanonicalTensorSpec]) -> FragmentLayout: + ordered = tuple(sorted(specs)) + if not ordered: + raise ValueError("canonical LoRA layout is empty") + if len({spec.name for spec in ordered}) != len(ordered): + raise ValueError("canonical LoRA tensor names must be unique") + return FragmentLayout( + [ + Fragment( + merge_mode=MERGE_AVG, + tensors=[(spec.name, spec.numel) for spec in ordered], + identity_shapes={spec.name: spec.shape for spec in ordered}, + ) + ] + ) + + +def canonical_layout_hash(specs: Sequence[CanonicalTensorSpec]) -> str: + """Semantic hash persisted by the syncer and rebuilt by the exporter.""" + + return layout_fingerprint(build_avg_layout(specs)).hex() + + +def canonical_lora_config_hash( + *, rank: int, target_modules: Sequence[str] +) -> str: + if rank <= 0 or not target_modules: + raise ValueError("canonical LoRA config requires rank and target modules") + payload = { + "bias": "none", + "lora_alpha": rank, + "lora_dropout": 0.0, + "r": rank, + "target_modules": sorted(set(target_modules)), + "task_type": "CAUSAL_LM", + } + encoded = json.dumps( + payload, sort_keys=True, separators=(",", ":") + ).encode("utf-8") + return hashlib.sha256(encoded).hexdigest() + + +def _canonical_tensor(tensor: torch.Tensor, name: str) -> torch.Tensor: + if not isinstance(tensor, torch.Tensor) or not tensor.is_floating_point(): + raise TypeError(f"{name!r} must be a floating-point torch.Tensor") + value = tensor.detach().to(device="cpu", dtype=torch.float32).contiguous() + if not torch.isfinite(value).all().item(): + raise ValueError(f"{name!r} contains NaN or Inf") + return value.clone() + + +def canonical_state( + policy_version: int, + tensors: Mapping[str, torch.Tensor], + *, + base_model_revision: str, + lora_config_hash: str, + layout_hash: str | None = None, + expected_specs: Sequence[CanonicalTensorSpec] | None = None, +) -> CanonicalLoraState: + normalized = { + name: _canonical_tensor(tensor, name) + for name, tensor in sorted(tensors.items()) + } + specs = canonical_specs(normalized) + if expected_specs is not None and specs != tuple(expected_specs): + raise ValueError("canonical LoRA names, shapes, or dtypes changed") + actual_layout_hash = canonical_layout_hash(specs) + if layout_hash is not None and layout_hash != actual_layout_hash: + raise ValueError("canonical LoRA layout hash changed") + return CanonicalLoraState( + base_model_revision, + lora_config_hash, + actual_layout_hash, + policy_version, + normalized, + ) + + +def flat_tensor( + tensors: Mapping[str, torch.Tensor], + specs: Sequence[CanonicalTensorSpec] | None = None, +) -> torch.Tensor: + specs = tuple(sorted(specs or canonical_specs(tensors))) + if set(tensors) != {spec.name for spec in specs}: + raise ValueError("tensor names do not match canonical specs") + values = [] + for spec in specs: + tensor = _canonical_tensor(tensors[spec.name], spec.name) + if tuple(tensor.shape) != spec.shape: + raise ValueError(f"shape mismatch for {spec.name!r}") + values.append(tensor.reshape(-1)) + return torch.cat(values) + + +def tensors_from_flat( + flat: torch.Tensor, + specs: Sequence[CanonicalTensorSpec], +) -> dict[str, torch.Tensor]: + specs = tuple(specs) + if specs != tuple(sorted(specs)): + raise ValueError("canonical LoRA specs are not sorted") + flat = _canonical_tensor(flat.reshape(-1), "flat LoRA policy") + expected = sum(spec.numel for spec in specs) + if flat.numel() != expected: + raise ValueError( + f"flat LoRA policy has {flat.numel()} values, expected {expected}" + ) + tensors = {} + offset = 0 + for spec in specs: + tensors[spec.name] = flat[ + offset : offset + spec.numel + ].reshape(spec.shape).clone() + offset += spec.numel + return tensors + + +def policy_delta(local: CanonicalLoraState, base: CanonicalLoraState) -> torch.Tensor: + if local.policy_version != base.policy_version: + raise ValueError("local and base policy versions differ") + if local.specs != base.specs: + raise ValueError("local and base LoRA layouts differ") + if ( + local.base_model_revision, + local.lora_config_hash, + local.layout_hash, + ) != ( + base.base_model_revision, + base.lora_config_hash, + base.layout_hash, + ): + raise ValueError("local and base canonical LoRA identities differ") + delta = flat_tensor(local.tensors, local.specs) - flat_tensor( + base.tensors, base.specs + ) + if not torch.isfinite(delta).all().item(): + raise ValueError("local LoRA delta contains NaN or Inf") + return delta.contiguous() + + +def policy_hash(state: CanonicalLoraState) -> str: + digest = hashlib.sha256() + digest.update(b"yeto-rl-policy-v1\0") + digest.update(state.base_model_revision.encode("ascii")) + digest.update(state.lora_config_hash.encode("ascii")) + digest.update(state.layout_hash.encode("ascii")) + digest.update(state.policy_version.to_bytes(8, "little")) + for spec in state.specs: + digest.update(spec.name.encode("utf-8")) + digest.update(_canonical_tensor(state.tensors[spec.name], spec.name).numpy().tobytes()) + return digest.hexdigest() + + +@dataclass(frozen=True) +class LocalRoundStats: + island_id: int + local_round_id: int + base_policy_version: int + active_groups: int + completed_groups: int + cancelled_groups: int + completed_trajectories: int + action_tokens: int + tool_wait_seconds: float + group_p50_seconds: float + group_p95_seconds: float + group_p99_seconds: float + reward_mean: float + reward_std: float + zero_variance_group_ratio: float + mean_kl: float | None + ess_ratio: float | None + clip_fraction: float + delta_l2_norm: float + rollout_seconds: float + train_seconds: float + + def __post_init__(self) -> None: + for name in ( + "island_id", + "local_round_id", + "base_policy_version", + "active_groups", + "completed_groups", + "cancelled_groups", + "completed_trajectories", + "action_tokens", + ): + if getattr(self, name) < 0: + raise ValueError(f"{name} must be non-negative") + for name in ( + "reward_mean", + "reward_std", + "tool_wait_seconds", + "group_p50_seconds", + "group_p95_seconds", + "group_p99_seconds", + "zero_variance_group_ratio", + "clip_fraction", + "delta_l2_norm", + "rollout_seconds", + "train_seconds", + ): + if not math.isfinite(getattr(self, name)): + raise ValueError(f"{name} must be finite") + for name in ("mean_kl", "ess_ratio"): + value = getattr(self, name) + if value is not None and not math.isfinite(value): + raise ValueError(f"{name} must be finite when present") diff --git a/yeto/rl/export.py b/yeto/rl/export.py new file mode 100644 index 0000000..fe22df1 --- /dev/null +++ b/yeto/rl/export.py @@ -0,0 +1,223 @@ +"""Export a committed RL-AVG checkpoint as a standard PEFT adapter.""" + +from __future__ import annotations + +import argparse +import os +import re +import tempfile +from collections.abc import Sequence +from pathlib import Path + +from ..export import parse_checkpoint +from .core import ( + CanonicalLoraState, + CanonicalTensorSpec, + canonical_layout_hash, + canonical_lora_config_hash, + canonical_state, + tensors_from_flat, +) + + +def target_modules(choice: str, config) -> str: + """Reuse Yeto's model-driven public LoRA target semantics.""" + + from ..learner import resolve_lora_targets + + return resolve_lora_targets(choice, config) + + +def derive_peft_lora_specs( + model: str, + revision: str | None, + *, + rank: int, + targets: str | Sequence[str], + trust_remote_code: bool = False, +) -> tuple[CanonicalTensorSpec, ...]: + """Build PEFT's actual LoRA tensor contract without allocating weights.""" + + from accelerate import init_empty_weights + from peft import LoraConfig, get_peft_model, get_peft_model_state_dict + from transformers import AutoConfig, AutoModelForCausalLM + + config = AutoConfig.from_pretrained( + model, + revision=revision, + trust_remote_code=trust_remote_code, + ) + if isinstance(targets, str) and targets in { + "auto", + "attention", + "all-linear", + }: + targets = target_modules(targets, config) + with init_empty_weights(): + base = AutoModelForCausalLM.from_config( + config, + trust_remote_code=trust_remote_code, + ) + if isinstance(targets, str) and targets != "all-linear": + pattern = re.compile(targets) + targets = sorted( + name for name, _ in base.named_modules() if pattern.fullmatch(name) + ) + adapter = get_peft_model( + base, + LoraConfig( + r=rank, + lora_alpha=rank, + target_modules=targets, + lora_dropout=0.0, + bias="none", + task_type="CAUSAL_LM", + ), + ) + state = get_peft_model_state_dict(adapter) + if not state: + raise ValueError("PEFT found no LoRA tensors for the requested model") + return tuple( + CanonicalTensorSpec( + name, + tuple(int(dim) for dim in tensor.shape), + "float32", + tensor.numel(), + ) + for name, tensor in sorted(state.items()) + ) + + +def adapter_targets(specs: Sequence[CanonicalTensorSpec]) -> list[str]: + targets = { + spec.name.rsplit(".lora_", 1)[0].rsplit(".", 1)[-1] + for spec in specs + } + return sorted(targets) + + +def write_peft_adapter( + state: CanonicalLoraState, + output_dir: str | Path, + *, + base_model: str, + model_revision: str | None, + rank: int, +) -> None: + """Write the two standard PEFT adapter files.""" + + from peft import LoraConfig + from peft.utils import CONFIG_NAME + from safetensors.torch import save_file + + output = Path(output_dir).expanduser() + output.mkdir(parents=True, exist_ok=True) + weights = output / "adapter_model.safetensors" + temporary = output / "adapter_model.safetensors.tmp" + save_file(dict(state.tensors), temporary) + os.replace(temporary, weights) + + config = LoraConfig( + r=rank, + lora_alpha=rank, + target_modules=adapter_targets(state.specs), + lora_dropout=0.0, + bias="none", + task_type="CAUSAL_LM", + inference_mode=True, + base_model_name_or_path=base_model, + revision=model_revision, + ) + with tempfile.TemporaryDirectory(dir=output) as temporary_dir: + config.save_pretrained(temporary_dir) + os.replace(Path(temporary_dir) / CONFIG_NAME, output / CONFIG_NAME) + + +def export_rl_checkpoint( + checkpoint_path: str | Path, + output_dir: str | Path, + *, + model: str, + model_revision: str, + rank: int, + lora_targets: str = "auto", + trust_remote_code: bool = False, +) -> CanonicalLoraState: + from ..models import resolve + + model = resolve(model) + checkpoint = parse_checkpoint(checkpoint_path) + if len(checkpoint.fragments) != 1: + raise ValueError("RL checkpoint must contain exactly one fragment") + version, params, _ = checkpoint.fragments[0] + if version != checkpoint.global_step: + raise ValueError("RL checkpoint fragment version is not committed") + + specs = derive_peft_lora_specs( + model, + model_revision, + rank=rank, + targets=lora_targets, + trust_remote_code=trust_remote_code, + ) + layout_hash = canonical_layout_hash(specs) + if checkpoint.layout_hash is None: + raise ValueError("RL checkpoint does not contain a canonical layout hash") + if checkpoint.layout_hash != layout_hash: + raise ValueError("RL checkpoint canonical layout hash does not match exporter") + targets = adapter_targets(specs) + state = canonical_state( + version, + tensors_from_flat(params, specs), + base_model_revision=model_revision, + lora_config_hash=canonical_lora_config_hash( + rank=rank, + target_modules=targets, + ), + layout_hash=layout_hash, + ) + write_peft_adapter( + state, + output_dir, + base_model=model, + model_revision=model_revision, + rank=rank, + ) + return state + + +def parse_args(argv=None): + parser = argparse.ArgumentParser( + prog="yeto-rl-export", + description="Export a committed Yeto RL checkpoint as a PEFT adapter.", + ) + parser.add_argument("--checkpoint", required=True) + parser.add_argument("--model", required=True) + parser.add_argument("--model-revision", required=True) + parser.add_argument("--lora-r", type=int, required=True) + parser.add_argument( + "--lora-targets", + choices=["auto", "attention", "all-linear"], + default="auto", + ) + parser.add_argument("--trust-remote-code", action="store_true") + parser.add_argument("--output-dir", required=True) + return parser.parse_args(argv) + + +def main(argv=None) -> None: + args = parse_args(argv) + state = export_rl_checkpoint( + args.checkpoint, + args.output_dir, + model=args.model, + model_revision=args.model_revision, + rank=args.lora_r, + lora_targets=args.lora_targets, + trust_remote_code=args.trust_remote_code, + ) + print(f"exported committed RL policy v{state.policy_version} to {args.output_dir}") + + +if __name__ == "__main__": + main() diff --git a/yeto/rl/learner.py b/yeto/rl/learner.py new file mode 100644 index 0000000..ce2386a --- /dev/null +++ b/yeto/rl/learner.py @@ -0,0 +1,466 @@ +"""Entrypoint for one pinned-Miles RL island.""" + +from __future__ import annotations + +import argparse +import json +import os +import sys +from pathlib import Path +from typing import Any + +from .bridge import BridgeConfig, StrictRlBridge +from .export import adapter_targets, derive_peft_lora_specs +from .miles import MilesIslandRuntime, verify_miles_revision + + +def parse_args(argv=None): + parser = argparse.ArgumentParser(prog="python3 -m yeto.rl.learner") + parser.add_argument("--model", required=True) + parser.add_argument("--model-revision", required=True) + parser.add_argument("--data", required=True) + parser.add_argument("--data-revision", required=True) + parser.add_argument("--syncer", required=True) + parser.add_argument("--learner-id", type=int, required=True) + parser.add_argument("--reward-function", required=True) + parser.add_argument("--reward-sha256", required=True) + parser.add_argument("--source-sha256", required=True) + parser.add_argument("--global-rounds", type=int, required=True) + parser.add_argument("--groups-per-round", type=int, required=True) + parser.add_argument("--samples-per-group", type=int, required=True) + parser.add_argument("--over-sampling-batch-size", type=int, required=True) + parser.add_argument("--optimizer-steps", type=int, required=True) + parser.add_argument("--rollout-max-response-len", type=int, required=True) + parser.add_argument("--custom-generate-function-path", default=None) + parser.add_argument("--use-session-server", action="store_true") + parser.add_argument("--session-server-ip", default=None) + parser.add_argument("--session-server-port", type=int, nargs="+", default=None) + parser.add_argument("--tito-model", default=None) + parser.add_argument("--completed-groups-path", required=True) + parser.add_argument("--event-tape", required=True) + parser.add_argument("--actor-num-nodes", type=int, required=True) + parser.add_argument("--actor-num-gpus-per-node", type=int, required=True) + parser.add_argument("--expert-parallel", type=int, default=None) + parser.add_argument("--lora-r", type=int, required=True) + parser.add_argument( + "--lora-targets", + choices=["auto", "attention", "all-linear"], + required=True, + ) + parser.add_argument("--inner-lr", type=float, required=True) + parser.add_argument("--seq-len", type=int, required=True) + parser.add_argument("--seed", type=int, required=True) + parser.add_argument("--wan-streams", type=int, default=4) + parser.add_argument("--miles-root", required=True) + parser.add_argument("--trust-remote-code", action="store_true") + return parser.parse_args(argv) + + +def _provider_value(provider, *names: str): + for name in names: + value = getattr(provider, name, None) + if value is not None: + return value + raise ValueError( + f"Megatron-Bridge provider {type(provider).__name__} lacks {names[0]}" + ) + + +def _positive_int(provider, *names: str) -> int: + value = _provider_value(provider, *names) + if not isinstance(value, int) or value <= 0: + raise ValueError(f"Megatron-Bridge provider has invalid {names[0]}={value!r}") + return value + + +def _text(value: Any) -> str: + return str(getattr(value, "value", value)) + + +def _miles_callable(spec: str) -> str: + module, separator, function = spec.partition(":") + if not separator or not module or not function.isidentifier(): + raise ValueError("RL callable must be package.module:function") + return f"{module}.{function}" + + +def build_miles_argv( + args, + *, + model_path: str | Path, + prompt_path: str | Path, + provider, + target_modules: list[str], +) -> list[str]: + """Construct Miles arguments from Bridge's actual model provider.""" + + hidden = _positive_int(provider, "hidden_size") + heads = _positive_int(provider, "num_attention_heads") + layers = _positive_int(provider, "num_layers") + ffn = _positive_int(provider, "ffn_hidden_size") + query_groups = int(getattr(provider, "num_query_groups", heads)) + kv_channels = int(getattr(provider, "kv_channels", hidden // heads)) + max_positions = int( + _provider_value( + provider, + "seq_length", + "max_sequence_length", + "max_position_embeddings", + ) + ) + if args.seq_len > max_positions: + raise ValueError( + f"--seq-len {args.seq_len} exceeds model context limit {max_positions}" + ) + vocab_size = _positive_int(provider, "vocab_size", "padded_vocab_size") + normalization = _text(getattr(provider, "normalization", "RMSNorm")) + epsilon = float( + _provider_value(provider, "layernorm_epsilon", "norm_epsilon") + ) + position_type = _text(getattr(provider, "position_embedding_type", "rope")) + rope_type = None + if position_type == "yarn": + position_type, rope_type = "rope", "yarn" + rotary_base = int(getattr(provider, "rotary_base", 10000)) + actor_gpus = args.actor_num_nodes * args.actor_num_gpus_per_node + is_moe = getattr(provider, "num_moe_experts", None) is not None + expert_parallel = getattr(args, "expert_parallel", None) or ( + actor_gpus if is_moe else 1 + ) + if not is_moe and expert_parallel != 1: + raise ValueError("EP>1 requires a MoE model") + if actor_gpus % expert_parallel: + raise ValueError("expert parallelism must divide Miles actor world size") + if is_moe and expert_parallel > 1 and args.lora_targets == "all-linear": + raise ValueError( + "EP>1 requires replicated attention LoRA, not expert-sharded all-linear LoRA" + ) + # Megatron DP includes the ranks rearranged into EP groups. With the v0 + # TP=PP=CP=1 contract, Miles reports one rollout shard per actor rank. + data_parallel = actor_gpus + global_batch = ( + args.groups_per_round * args.samples_per_group // args.optimizer_steps + ) + if global_batch % data_parallel: + raise ValueError("Miles global batch must divide evenly across DP ranks") + if not target_modules: + raise ValueError("PEFT selected no LoRA target modules") + target_value = ",".join(target_modules) + + values = [ + "train.py", + "--train-backend", "megatron", + "--hf-checkpoint", str(model_path), + "--load", str(model_path), + "--megatron-to-hf-mode", "bridge", + "--model-name", type(provider).__name__, + "--num-layers", str(layers), + "--hidden-size", str(hidden), + "--num-attention-heads", str(heads), + "--num-query-groups", str(query_groups), + "--kv-channels", str(kv_channels), + "--ffn-hidden-size", str(ffn), + "--max-position-embeddings", str(max_positions), + "--seq-length", str(args.seq_len), + "--normalization", normalization, + "--norm-epsilon", str(epsilon), + "--position-embedding-type", position_type, + "--rotary-base", str(rotary_base), + "--vocab-size", str(vocab_size), + "--lora-rank", str(args.lora_r), + "--lora-alpha", str(args.lora_r), + "--lora-dropout", "0", + "--lora-type", "canonical_lora", + "--target-modules", target_value, + "--lora-base-cpu-backup", + "--actor-num-nodes", str(args.actor_num_nodes), + "--actor-num-gpus-per-node", str(args.actor_num_gpus_per_node), + "--num-gpus-per-node", str(args.actor_num_gpus_per_node), + "--rollout-num-gpus-per-engine", "1", + "--colocate", + "--no-offload-train", + "--sglang-mem-fraction-static", "0.4", + "--tensor-model-parallel-size", "1", + "--pipeline-model-parallel-size", "1", + "--context-parallel-size", "1", + "--expert-model-parallel-size", str(expert_parallel), + "--expert-tensor-parallel-size", "1", + "--prompt-data", str(prompt_path), + "--input-key", "messages", + "--label-key", "label", + "--metadata-key", "metadata", + "--apply-chat-template", + "--rollout-seed", str(args.seed + args.learner_id), + "--sglang-enable-deterministic-inference", + "--num-rollout", str(args.global_rounds), + "--rollout-batch-size", str(args.groups_per_round), + "--n-samples-per-prompt", str(args.samples_per_group), + "--over-sampling-batch-size", str(args.over_sampling_batch_size), + "--num-steps-per-rollout", str(args.optimizer_steps), + "--global-batch-size", str(global_batch), + "--balance-data", + "--rollout-max-context-len", str(args.seq_len), + "--rollout-max-response-len", str(args.rollout_max_response_len), + "--rollout-function-path", "yeto.rl.miles.generate_rollout", + "--rollout-all-samples-process-path", "yeto.rl.miles.queue_completed_groups", + "--custom-rm-path", _miles_callable(args.reward_function), + "--advantage-estimator", "grpo", + "--lr", str(args.inner_lr), + "--custom-megatron-init-path", "yeto.rl.miles.configure_miles_bridge", + "--accumulate-allreduce-grads-in-fp32", + "--attention-softmax-in-fp32", + "--attention-backend", "unfused", + "--no-gradient-accumulation-fusion", + "--bf16", + "--no-load-optim", + "--no-load-rng", + "--no-save-optim", + "--no-save-rng", + "--finetune", + "--seed", str(args.seed), + "--sglang-max-lora-rank", str(args.lora_r), + "--pin-rollout-manager-to-head", + ] + if args.custom_generate_function_path: + values.extend( + ( + "--custom-generate-function-path", + args.custom_generate_function_path, + ) + ) + if args.use_session_server: + values.append("--use-session-server") + if args.session_server_ip: + values.extend(("--session-server-ip", args.session_server_ip)) + if args.session_server_port: + values.append("--session-server-port") + values.extend(str(port) for port in args.session_server_port) + if args.tito_model: + values.extend(("--tito-model", args.tito_model)) + if query_groups < heads: + values.append("--group-query-attention") + if rope_type is not None: + values.extend(("--rope-type", rope_type)) + if bool(getattr(provider, "gated_linear_unit", False)): + values.append("--swiglu") + if not bool(getattr(provider, "share_embeddings_and_output_weights", True)): + values.append("--untie-embeddings-and-output-weights") + if getattr(provider, "add_bias_linear", None) is False: + values.append("--disable-bias-linear") + if bool(getattr(provider, "add_qkv_bias", False)): + values.append("--add-qkv-bias") + if bool(getattr(provider, "qk_layernorm", False)): + values.append("--qk-layernorm") + if getattr(provider, "num_moe_experts", None) is not None: + values.extend( + ( + "--num-experts", + str(_positive_int(provider, "num_moe_experts")), + "--moe-ffn-hidden-size", + str(_positive_int(provider, "moe_ffn_hidden_size")), + "--moe-router-topk", + str(_positive_int(provider, "moe_router_topk")), + "--moe-layer-freq", + str(_provider_value(provider, "moe_layer_freq")), + ) + ) + shared = getattr(provider, "moe_shared_expert_intermediate_size", None) + if shared is not None: + values.extend( + ("--moe-shared-expert-intermediate-size", str(int(shared))) + ) + if bool(getattr(provider, "multi_latent_attention", False)): + values.append("--multi-latent-attention") + for name in ( + "q_lora_rank", + "kv_lora_rank", + "qk_head_dim", + "qk_pos_emb_head_dim", + "v_head_dim", + ): + value = getattr(provider, name, None) + if value is not None: + values.extend((f"--{name.replace('_', '-')}", str(int(value)))) + return values + + +def _messages(row: dict[str, Any]) -> list[dict[str, Any]]: + value = row.get("messages", row.get("prompt", row.get("input"))) + if isinstance(value, str): + value = [{"role": "user", "content": value}] + if ( + not isinstance(value, list) + or not value + or any(not isinstance(item, dict) for item in value) + ): + raise ValueError("RL rows must contain messages or a string prompt/input") + return value + + +def prepare_prompt_data( + source: str, + revision: str | None, + output_path: str | Path, +) -> Path: + from ..data import load_rows + + rows = load_rows(source, revision=revision) + output = Path(output_path).expanduser() + output.parent.mkdir(parents=True, exist_ok=True) + temporary = output.with_name(output.name + ".tmp") + count = 0 + with temporary.open("w", encoding="utf-8") as handle: + for raw in rows: + row = dict(raw) + normalized = { + "messages": _messages(row), + "label": row.get("label"), + "metadata": { + key: value for key, value in row.items() if key != "messages" + }, + } + if "tools" in row: + normalized["tools"] = row["tools"] + handle.write( + json.dumps( + normalized, + ensure_ascii=False, + separators=(",", ":"), + ) + + "\n" + ) + count += 1 + if count == 0: + temporary.unlink(missing_ok=True) + raise ValueError("RL prompt dataset is empty") + os.replace(temporary, output) + return output + + +def _parse_miles_args(argv: list[str]): + previous = sys.argv + try: + sys.argv = argv + from miles.utils.arguments import parse_args as parse_miles_args + + return parse_miles_args() + finally: + sys.argv = previous + + +def _syncer_address(value: str) -> tuple[str, int]: + host, separator, port = value.rpartition(":") + if not separator or not host: + raise ValueError("--syncer must be HOST:PORT") + return host, int(port) + + +def main(argv=None) -> None: + args = parse_args(argv) + from ..provenance import ( + is_immutable_commit, + python_spec_sha256, + verify_source_tree_sha256, + ) + + verify_source_tree_sha256(args.source_sha256) + if not is_immutable_commit(args.model_revision) or not is_immutable_commit( + args.data_revision + ): + raise ValueError("RL model and dataset revisions must be immutable commits") + reward_sha256 = python_spec_sha256(args.reward_function) + if reward_sha256 != args.reward_sha256.lower(): + raise ValueError( + f"reward source SHA256 mismatch: expected {args.reward_sha256.lower()}, " + f"got {reward_sha256}" + ) + verify_miles_revision(args.miles_root) + + from miles.utils.misc import load_function + + load_function(_miles_callable(args.reward_function)) + if args.custom_generate_function_path: + load_function(args.custom_generate_function_path) + from huggingface_hub import snapshot_download + from megatron.bridge import AutoBridge + + from ..models import resolve + from ..provenance import is_local_reference + + model = resolve(args.model) + if is_local_reference(model): + model_path = str(Path(model).expanduser().resolve()) + else: + model_path = snapshot_download(repo_id=model, revision=args.model_revision) + prompt_path = prepare_prompt_data( + args.data, + args.data_revision, + "~/yeto-rl/prompts.jsonl", + ) + + model_bridge = AutoBridge.from_hf_pretrained( + model_path, + trust_remote_code=args.trust_remote_code, + ) + provider = model_bridge.to_megatron_provider(load_weights=False) + provider.finalize() + specs = derive_peft_lora_specs( + model_path, + None, + rank=args.lora_r, + targets=args.lora_targets, + trust_remote_code=args.trust_remote_code, + ) + miles_targets = adapter_targets(specs) + from .core import canonical_layout_hash, canonical_lora_config_hash + + layout_hash = canonical_layout_hash(specs) + lora_config_hash = canonical_lora_config_hash( + rank=args.lora_r, + target_modules=miles_targets, + ) + miles_args = _parse_miles_args( + build_miles_argv( + args, + model_path=model_path, + prompt_path=prompt_path, + provider=provider, + target_modules=miles_targets, + ) + ) + miles_args.yeto_rl_trust_remote_code = args.trust_remote_code + miles_args.yeto_rl_model = args.model + miles_args.yeto_rl_data = args.data + miles_args.yeto_rl_base_model_revision = args.model_revision + miles_args.yeto_rl_data_revision = args.data_revision + miles_args.yeto_rl_lora_config_hash = lora_config_hash + miles_args.yeto_rl_layout_hash = layout_hash + miles_args.yeto_rl_reward_sha256 = args.reward_sha256 + miles_args.yeto_rl_completed_groups_path = args.completed_groups_path + miles_args.yeto_rl_event_tape = args.event_tape + miles_args.yeto_rl_learner_id = args.learner_id + + runtime = MilesIslandRuntime(miles_args) + bridge = StrictRlBridge( + runtime, + BridgeConfig( + syncer_addr=_syncer_address(args.syncer), + learner_id=args.learner_id, + global_rounds=args.global_rounds, + groups_per_round=args.groups_per_round, + samples_per_group=args.samples_per_group, + local_optimizer_steps=args.optimizer_steps, + wan_streams=args.wan_streams, + expected_specs=specs, + base_model_revision=args.model_revision, + lora_config_hash=lora_config_hash, + layout_hash=layout_hash, + event_tape=args.event_tape, + ), + ) + final = bridge.run() + print(f"[rl] learner {args.learner_id} finalized policy v{final.policy_version}") + + +if __name__ == "__main__": + main() diff --git a/yeto/rl/miles.py b/yeto/rl/miles.py new file mode 100644 index 0000000..a9130bb --- /dev/null +++ b/yeto/rl/miles.py @@ -0,0 +1,1136 @@ +"""Adapter for the pinned Miles commit's Megatron/SGLang runtime.""" + +from __future__ import annotations + +import asyncio +import importlib +import json +import os +import statistics +import subprocess +import time +from collections.abc import Mapping +from dataclasses import asdict +from pathlib import Path +from typing import Any + +import torch + +from . import MILES_COMMIT, MILES_REPOSITORY +from .core import ( + CanonicalLoraState, + LocalRoundStats, + StrictRlInvariantError, + canonical_state, + policy_hash, +) + +_CANONICAL_PREFIX = "base_model.model." + + +def verify_miles_revision(root: str | Path) -> Path: + """Verify repository, commit, tracked files, and the imported package.""" + + root = Path(root).expanduser().resolve() + + def git(*args: str) -> str: + try: + return subprocess.run( + ["git", "-C", str(root), *args], + check=True, + capture_output=True, + text=True, + ).stdout.strip() + except (OSError, subprocess.CalledProcessError) as exc: + raise RuntimeError(f"cannot verify Miles checkout at {root}") from exc + + commit = git("rev-parse", "HEAD") + branch = git("rev-parse", "--abbrev-ref", "HEAD") + origin = git("config", "--get", "remote.origin.url") + if commit != MILES_COMMIT: + raise RuntimeError( + f"Miles revision mismatch: expected {MILES_COMMIT}, got {commit}" + ) + if origin.removesuffix(".git").rstrip("/") != MILES_REPOSITORY: + raise RuntimeError( + f"Miles origin mismatch: expected {MILES_REPOSITORY}, got {origin}" + ) + if branch != "HEAD": + raise RuntimeError("Miles checkout is not detached") + if git("status", "--porcelain", "--untracked-files=all"): + raise RuntimeError("Miles checkout is not clean") + + miles = importlib.import_module("miles") + package_path = Path(miles.__file__).resolve() + if not package_path.is_relative_to(root): + raise RuntimeError(f"imported Miles package {package_path} is outside {root}") + return root + + +def _adapter_sides(actor) -> list[tuple[str, Any]]: + """Return PEFT names paired with Bridge's actual conversion tasks.""" + + from megatron.bridge import AutoBridge + + bridge = AutoBridge.from_hf_pretrained( + actor.args.hf_checkpoint, + trust_remote_code=bool(actor.args.yeto_rl_trust_remote_code), + ) + model_bridge = getattr(bridge, "_model_bridge", None) + build_tasks = getattr(model_bridge, "build_adapter_conversion_tasks", None) + if build_tasks is None: + raise RuntimeError("pinned Megatron-Bridge lacks adapter conversion tasks") + tasks_by_base = build_tasks(actor.model) + sides: list[tuple[str, Any]] = [] + for base_name in sorted(tasks_by_base): + tasks = sorted( + tasks_by_base[base_name], + key=lambda task: task.adapter_key or "", + ) + for task in tasks: + for side in (task.linear_in_task, task.linear_out_task): + parameter = side.param_weight + main = getattr(parameter, "main_param", None) + if ( + main is None + or main.dtype != torch.float32 + or main.numel() != parameter.numel() + ): + raise RuntimeError( + f"LoRA parameter {side.param_name!r} has no complete " + "FP32 optimizer master" + ) + converted = side.mapping.megatron_to_hf( + main.view(parameter.shape), + side.megatron_module, + ) + if len(converted) != 1: + raise RuntimeError( + f"ambiguous LoRA mapping for {side.param_name!r}" + ) + raw_name = next(iter(converted)) + name = ( + raw_name + if raw_name.startswith(_CANONICAL_PREFIX) + else _CANONICAL_PREFIX + raw_name + ) + if not name.endswith((".lora_A.weight", ".lora_B.weight")): + raise RuntimeError(f"non-PEFT LoRA mapping {name!r}") + sides.append((name, side)) + names = [name for name, _ in sides] + if not names or len(names) != len(set(names)): + raise RuntimeError("Miles produced an empty or duplicate LoRA mapping") + mapped = {id(side.param_weight) for _, side in sides} + trainable = { + id(parameter) + for chunk in actor.model + for parameter in chunk.parameters() + if parameter.requires_grad + } + if mapped != trainable: + raise RuntimeError( + "Miles adapter conversion does not cover every trainable parameter" + ) + return sorted(sides) + + +@torch.no_grad() +def _export_fp32_policy(actor) -> dict[str, torch.Tensor]: + tensors = {} + for name, side in _adapter_sides(actor): + parameter = side.param_weight + converted = side.mapping.megatron_to_hf( + parameter.main_param.view(parameter.shape), + side.megatron_module, + ) + value = next(iter(converted.values())) + tensors[name] = value.detach().to( + device="cpu", dtype=torch.float32 + ).contiguous().clone() + return tensors + + +def _optimizer_children(optimizer) -> list[Any]: + return list(getattr(optimizer, "chained_optimizers", (optimizer,))) + + +def _reset_optimizer_state(actor, parameters: list[torch.Tensor]) -> int: + parameter_ids = {id(parameter) for parameter in parameters} + for child in _optimizer_children(actor.optimizer): + optimizer = getattr(child, "optimizer", child) + for parameter in list(optimizer.state): + if id(parameter) in parameter_ids: + optimizer.state.pop(parameter, None) + return len(parameter_ids) + + +@torch.no_grad() +def _copy_masters_to_model(actor) -> None: + for child in _optimizer_children(actor.optimizer): + copy = getattr(child, "_copy_main_params_to_model_params", None) + if copy is None: + raise RuntimeError("pinned Megatron optimizer lacks main-to-model copy") + copy() + + +def _restore_scheduler_progress(actor, policy_version: int) -> None: + scheduler = actor.opt_param_scheduler + batch_size = actor.args.global_batch_size + target = policy_version * actor.args.num_steps_per_rollout * batch_size + if scheduler is None or batch_size <= 0 or scheduler.num_steps % batch_size: + raise RuntimeError("Miles scheduler progress is not an integral optimizer step") + if scheduler.num_steps > target: + raise RuntimeError("Miles scheduler is ahead of the committed policy") + if scheduler.num_steps < target: + scheduler.step(increment=target - scheduler.num_steps) + + +def _actor_export_policy(self): + if torch.distributed.is_initialized() and torch.distributed.get_rank() != 0: + return None + return _export_fp32_policy(self) + + +@torch.no_grad() +def _actor_apply_policy( + self, + tensors: Mapping[str, torch.Tensor], + policy_version: int, +): + sides = dict(_adapter_sides(self)) + if set(tensors) != set(sides): + missing = sorted(set(sides) - set(tensors)) + extra = sorted(set(tensors) - set(sides)) + raise RuntimeError( + f"global LoRA mapping mismatch: missing={missing}, extra={extra}" + ) + + mapped = {} + for name, side in sides.items(): + value = tensors[name].detach().to( + device=side.param_weight.device, + dtype=torch.float32, + ) + target = side.mapping.hf_to_megatron(value, side.megatron_module) + if target.numel() != side.param_weight.numel(): + raise RuntimeError(f"global LoRA shape mismatch for {name!r}") + mapped[name] = target.reshape(side.param_weight.shape).contiguous() + + _restore_scheduler_progress(self, policy_version) + reset_parameter_count = _reset_optimizer_state( + self, + [side.param_weight.main_param for side in sides.values()], + ) + for name, side in sides.items(): + side.param_weight.main_param.view(side.param_weight.shape).copy_(mapped[name]) + _copy_masters_to_model(self) + if torch.distributed.is_initialized(): + torch.distributed.barrier() + self.weights_backuper.backup("actor") + torch.cuda.empty_cache() + + identity = { + "base_model_revision": self.args.yeto_rl_base_model_revision, + "lora_config_hash": self.args.yeto_rl_lora_config_hash, + "layout_hash": self.args.yeto_rl_layout_hash, + } + applied = canonical_state(policy_version, _export_fp32_policy(self), **identity) + canonical_state( + policy_version, + tensors, + expected_specs=applied.specs, + **identity, + ) + return reset_parameter_count, policy_hash(applied) + + +def _actor_optimizer_steps(self) -> int: + scheduler = self.opt_param_scheduler + batch = self.args.global_batch_size + if scheduler is None or batch <= 0 or scheduler.num_steps % batch: + raise RuntimeError("Miles optimizer step counter is not integral") + return int(scheduler.num_steps // batch) + + +def _actor_train_metrics(self): + if torch.distributed.is_initialized() and torch.distributed.get_rank() != 0: + return None + metrics = getattr(self.args, "_yeto_rl_train_metrics", None) + if hasattr(self.args, "_yeto_rl_train_metrics"): + del self.args._yeto_rl_train_metrics + return metrics + + +def _install_train_metric_capture() -> None: + from miles.backends.megatron_utils import model + + original = model.log_train_step + if getattr(original, "_yeto_rl_capture", False): + return + + def log_train_step(*values, **kwargs): + metrics = original(*values, **kwargs) + if not torch.distributed.is_initialized() or torch.distributed.get_rank() == 0: + args = kwargs["args"] + args._yeto_rl_train_metrics = { + key: float(metrics[key]) + for key in ( + "train/train_rollout_kl", + "train/ess_ratio", + "train/pg_clipfrac", + ) + if key in metrics + } + return metrics + + log_train_step._yeto_rl_capture = True + model.log_train_step = log_train_step + + +def _install_colocated_lora_ipc_sync() -> None: + """Keep CUDA IPC producers alive until SGLang finishes the transfer.""" + + import ray + from miles.backends.megatron_utils.update_weight import ( + update_weight_from_tensor, + ) + from miles.backends.megatron_utils.update_weight.common import ( + _check_weight_sync_results, + ) + + original = update_weight_from_tensor._send_to_colocated_engine + if getattr(original, "_yeto_rl_synchronized", False): + return + + def synchronized_send(*values, **kwargs): + refs, long_lived_tensors = original(*values, **kwargs) + if kwargs.get("lora_config") is not None: + results = ray.get(refs) + _check_weight_sync_results(results, is_lora=True) + group = kwargs.get("ipc_gather_group") + if group is not None: + torch.distributed.barrier(group=group) + return refs, long_lived_tensors + + synchronized_send._yeto_rl_synchronized = True + update_weight_from_tensor._send_to_colocated_engine = synchronized_send + + +def configure_miles_bridge(args) -> None: + """Install Yeto actor methods inside each Miles Ray worker.""" + + from megatron.bridge import AutoBridge + from megatron.bridge.training import config as bridge_config + + original = AutoBridge.to_megatron_provider + + def configured_provider(self, *values, **kwargs): + provider = original(self, *values, **kwargs) + provider.attention_backend = args.attention_backend + return provider + + AutoBridge.to_megatron_provider = configured_provider + # The INIT adapter is replicated. Keep complete fp32 masters and Adam + # state on every DP/EP rank; no sharded gather path is part of v0. + args.use_distributed_optimizer = False + original_ddp_config = bridge_config.DistributedDataParallelConfig + + def replicated_lora_ddp_config(*values, **kwargs): + kwargs["use_distributed_optimizer"] = False + return original_ddp_config(*values, **kwargs) + + bridge_config.DistributedDataParallelConfig = replicated_lora_ddp_config + _install_train_metric_capture() + _install_colocated_lora_ipc_sync() + install_miles_actor_adapter() + + +def install_miles_actor_adapter() -> None: + """Install methods in both the driver and each Miles Ray worker.""" + + from miles.backends.megatron_utils.actor import MegatronTrainRayActor + + methods = { + "yeto_rl_export_policy": _actor_export_policy, + "yeto_rl_apply_policy": _actor_apply_policy, + "yeto_rl_optimizer_steps": _actor_optimizer_steps, + "yeto_rl_train_metrics": _actor_train_metrics, + } + for name, method in methods.items(): + existing = getattr(MegatronTrainRayActor, name, None) + if existing is not None and ( + getattr(existing, "__module__", None), + getattr(existing, "__qualname__", None), + ) != (method.__module__, method.__qualname__): + raise RuntimeError(f"Miles actor already defines incompatible {name}") + setattr(MegatronTrainRayActor, name, method) + + +def _policy_token(version: int) -> str: + return f"yeto:{version}" + + +def _version_from_token(token: object) -> int: + prefix, separator, value = str(token).partition(":") + if prefix != "yeto" or not separator: + raise RuntimeError(f"invalid rollout policy token {token!r}") + try: + version = int(value) + except ValueError as exc: + raise RuntimeError(f"invalid rollout policy token {token!r}") from exc + if version < 0: + raise RuntimeError(f"invalid rollout policy token {token!r}") + return version + + +def _validate_rollout_groups(data: object, groups: int, samples: int) -> None: + if not isinstance(data, list) or len(data) != groups: + raise RuntimeError( + f"Miles produced {len(data) if isinstance(data, list) else 0} " + f"groups, expected {groups}" + ) + for index, group in enumerate(data): + if not isinstance(group, list) or len(group) != samples: + raise RuntimeError( + f"Miles group {index} contains " + f"{len(group) if isinstance(group, list) else 0} trajectories, " + f"expected {samples}" + ) + for sample in group: + status = getattr(getattr(sample, "status", None), "value", None) + if isinstance(sample, list) or status not in {"completed", "truncated"}: + raise RuntimeError("Miles returned an incomplete trajectory") + + +_ISLAND_CHECKPOINT_SCHEMA = 2 + + +def _island_checkpoint_config(args) -> dict[str, Any]: + return { + "actor_num_gpus_per_node": args.actor_num_gpus_per_node, + "actor_num_nodes": args.actor_num_nodes, + "advantage_estimator": args.advantage_estimator, + "model": args.yeto_rl_model, + "dataset": args.yeto_rl_data, + "base_model_revision": args.yeto_rl_base_model_revision, + "data_revision": args.yeto_rl_data_revision, + "seq_length": args.seq_length, + "seed": args.seed, + "expert_model_parallel_size": args.expert_model_parallel_size, + "layout_hash": args.yeto_rl_layout_hash, + "lr": args.lr, + "lora_config_hash": args.yeto_rl_lora_config_hash, + "n_samples_per_prompt": args.n_samples_per_prompt, + "num_steps_per_rollout": args.num_steps_per_rollout, + "over_sampling_batch_size": args.over_sampling_batch_size, + "reward_sha256": args.yeto_rl_reward_sha256, + "rollout_batch_size": args.rollout_batch_size, + "rollout_max_response_len": args.rollout_max_response_len, + "custom_generate_function_path": args.custom_generate_function_path, + "use_session_server": args.use_session_server, + "tito_model": args.tito_model, + } + + +def _complete_group_for_policy(group: object, policy_version: int, size: int) -> bool: + if not isinstance(group, list) or len(group) != size: + return False + for sample in group: + status = getattr(getattr(sample, "status", None), "value", None) + versions = getattr(sample, "weight_versions", None) + if status not in {"completed", "truncated"} or not versions: + return False + try: + observed = {_version_from_token(token) for token in versions} + except RuntimeError: + return False + if observed != {policy_version}: + return False + return True + + +def _group_indices(group: object) -> tuple[int, ...] | None: + if not isinstance(group, list): + return None + try: + return tuple(int(sample.index) for sample in group) + except (AttributeError, TypeError, ValueError): + return None + + +def queue_completed_groups(args, all_samples, data_source) -> None: + """Retain complete oversampling results until this round selects its batch.""" + + source = getattr(data_source, "__self__", None) + if source is None or not callable(getattr(source, "add_samples", None)): + raise RuntimeError("pinned Miles did not provide a bound data source method") + source.add_samples( + [ + group + for group in all_samples + if _complete_group_for_policy( + group, + args.yeto_rl_policy_version, + args.n_samples_per_prompt, + ) + ] + ) + + +def _atomic_save_island_checkpoint(path: Path, payload: dict[str, Any]) -> None: + path.parent.mkdir(parents=True, exist_ok=True) + temporary = path.with_name(f".{path.name}.tmp-{os.getpid()}") + try: + torch.save(payload, temporary) + os.replace(temporary, path) + finally: + temporary.unlink(missing_ok=True) + + +def _serialize_completed_groups(groups: list[list[Any]]) -> list[list[dict[str, Any]]]: + serialized = [] + for group in groups: + serialized_group = [] + for sample in group: + to_dict = getattr(sample, "to_dict", None) + if not callable(to_dict): + raise RuntimeError("pinned Miles Sample lacks checkpoint serialization") + serialized_group.append(to_dict()) + serialized.append(serialized_group) + return serialized + + +def _deserialize_completed_groups(groups: object) -> list[list[Any]]: + from miles.utils.types import Sample + + if not isinstance(groups, list): + return [] + try: + return [ + [Sample.from_dict(sample) for sample in group] + for group in groups + if isinstance(group, list) + ] + except (TypeError, ValueError) as error: + raise RuntimeError("invalid Miles samples in island checkpoint") from error + + +def _restore_completed_groups(args, policy_version: int, data_source) -> None: + if getattr(data_source, "_yeto_rl_checkpoint_loaded", False): + return + data_source._yeto_rl_checkpoint_loaded = True + path = Path(args.yeto_rl_completed_groups_path).expanduser() + if not path.is_file(): + return + try: + payload = torch.load(path, map_location="cpu", weights_only=True) + except Exception as error: # local queue is disposable, global state is not + print(f"[rl] discarded unreadable island checkpoint {path}: {error}") + return + expected_config = _island_checkpoint_config(args) + if ( + not isinstance(payload, dict) + or payload.get("schema_version") != _ISLAND_CHECKPOINT_SCHEMA + or payload.get("policy_version") != policy_version + or payload.get("config") != expected_config + ): + print(f"[rl] discarded island checkpoint outside policy/config contract: {path}") + return + groups = [ + group + for group in _deserialize_completed_groups(payload.get("completed_groups")) + if _complete_group_for_policy( + group, policy_version, args.n_samples_per_prompt + ) + ] + if groups: + data_source.add_samples(groups) + + +def _save_completed_groups( + args, + policy_version: int, + local_round_id: int, + data_source, + metrics: Mapping[str, Any] | None, +) -> None: + buffer = getattr(data_source, "buffer", None) + if not isinstance(buffer, list): + raise RuntimeError("Miles data source lacks a completed-group queue") + completed = [ + group + for group in buffer + if _complete_group_for_policy( + group, policy_version, args.n_samples_per_prompt + ) + ] + buffer[:] = completed + numeric_metrics = { + str(name): float(value) + for name, value in (metrics or {}).items() + if isinstance(value, (int, float)) + } + _atomic_save_island_checkpoint( + Path(args.yeto_rl_completed_groups_path).expanduser(), + { + "schema_version": _ISLAND_CHECKPOINT_SCHEMA, + "config": _island_checkpoint_config(args), + "local_round_id": local_round_id, + "policy_version": policy_version, + "rollout_metrics": numeric_metrics, + "local_round_stats": None, + "completed_groups": _serialize_completed_groups(completed), + }, + ) + + +def _percentile(values: list[float], fraction: float) -> float: + if not values: + return 0.0 + ordered = sorted(values) + position = (len(ordered) - 1) * fraction + lower = int(position) + upper = min(lower + 1, len(ordered) - 1) + return ordered[lower] + (ordered[upper] - ordered[lower]) * (position - lower) + + +def _run_rollout_with_metrics(generate, args, rollout_id, data_source, evaluation): + from miles.rollout import sglang_rollout + + state = { + "active": 0, + "peak_active": 0, + "cancelled": 0, + "durations": [], + } + generate_state = getattr(sglang_rollout, "GenerateState", None) + original = getattr(generate_state, "submit_generate_tasks", None) + if original is None: + raise RuntimeError("pinned Miles lacks GenerateState.submit_generate_tasks") + + def submit_generate_tasks(self, samples): + before = set(self.pendings) + original(self, samples) + started = time.monotonic() + tasks = self.pendings - before + state["active"] += len(tasks) + state["peak_active"] = max(state["peak_active"], state["active"]) + + def finished(task): + state["active"] -= 1 + state["durations"].append(time.monotonic() - started) + if task.cancelled(): + state["cancelled"] += 1 + return + try: + group = task.result() + except Exception: + state["cancelled"] += 1 + return + statuses = [ + getattr(getattr(sample, "status", None), "value", None) + for sample in group + ] + if any(status not in {"completed", "truncated"} for status in statuses): + state["cancelled"] += 1 + + for task in tasks: + task.add_done_callback(finished) + + generate_state.submit_generate_tasks = submit_generate_tasks + try: + return generate(args, rollout_id, data_source, evaluation=evaluation), state + finally: + generate_state.submit_generate_tasks = original + + +def generate_rollout(args, rollout_id: int, data_source, evaluation: bool = False): + """Keep Miles rollout, adding group/version queue contracts.""" + + from miles.rollout.sglang_rollout import generate_rollout as miles_generate + + if not evaluation: + args.yeto_rl_policy_version = rollout_id + _restore_completed_groups(args, rollout_id, data_source) + buffer = getattr(data_source, "buffer", None) + if not isinstance(buffer, list): + raise RuntimeError("Miles data source lacks a completed-group queue") + buffer[:] = [ + group + for group in buffer + if _complete_group_for_policy( + group, rollout_id, args.n_samples_per_prompt + ) + ] + output, lifecycle = _run_rollout_with_metrics( + miles_generate, + args, + rollout_id, + data_source, + evaluation, + ) + if not evaluation: + if lifecycle["active"] != 0: + raise RuntimeError("Miles returned with rollout groups still active") + _validate_rollout_groups( + output.samples, + args.rollout_batch_size, + args.n_samples_per_prompt, + ) + consumed = {_group_indices(group) for group in output.samples} + data_source.buffer[:] = [ + group + for group in data_source.buffer + if _group_indices(group) not in consumed + ] + samples = [sample for group in output.samples for sample in group] + tool_wait_seconds = sum( + float(getattr(sample, "non_generation_time", 0.0)) + for sample in samples + ) + round_metrics = { + "active_groups": lifecycle["peak_active"], + "cancelled_groups": lifecycle["cancelled"], + "tool_wait_seconds": tool_wait_seconds, + "group_p50_seconds": _percentile(lifecycle["durations"], 0.50), + "group_p95_seconds": _percentile(lifecycle["durations"], 0.95), + "group_p99_seconds": _percentile(lifecycle["durations"], 0.99), + } + _save_completed_groups( + args, + rollout_id, + rollout_id + 1, + data_source, + {**(output.metrics or {}), **round_metrics}, + ) + return output + + +class MilesIslandRuntime: + """Synchronous wrapper over the pinned Miles Ray APIs.""" + + def __init__(self, args) -> None: + self.args = args + self.loop = asyncio.new_event_loop() + self.rollout_manager = None + self.actor_model = None + self._owns_ray = False + self._trainer_awake = False + self._rollout_offloaded = True + self._policy_version: int | None = None + self._rollout_id = 0 + self._optimizer_reset_count = 0 + + async def _onload_trainer(self) -> None: + if self._trainer_awake: + return + await self.actor_model.onload() + self._trainer_awake = True + + async def _offload_trainer(self) -> None: + if self._trainer_awake and self.args.offload_train: + await self.actor_model.offload() + self._trainer_awake = False + + def _run(self, coroutine): + return self.loop.run_until_complete(coroutine) + + async def _actor_call(self, method: str, *args, rank0: bool = False): + results = await self.actor_model._broadcast(method, *args) + expected = self.args.actor_num_nodes * self.args.actor_num_gpus_per_node + if not isinstance(results, list) or len(results) != expected: + raise RuntimeError( + f"Miles returned {len(results) if isinstance(results, list) else 0} " + f"actor results, expected {expected}" + ) + if rank0: + exported = [result for result in results if result is not None] + if len(exported) != 1: + raise RuntimeError("Miles must export policy only on global rank 0") + return exported[0] + if not results or any(result != results[0] for result in results[1:]): + raise RuntimeError(f"Miles actor ranks disagree on {method}") + return results[0] + + async def _initialize(self) -> CanonicalLoraState: + import ray + from miles.ray.placement_group import ( + create_placement_groups, + create_rollout_manager, + create_training_models, + ) + + install_miles_actor_adapter() + if not ray.is_initialized(): + ray.init(address="auto") + self._owns_ray = True + expected_gpus = ( + self.args.actor_num_nodes * self.args.actor_num_gpus_per_node + ) + deadline = time.monotonic() + 300 + while True: + visible_gpus = int(ray.cluster_resources().get("GPU", 0)) + if visible_gpus >= expected_gpus or time.monotonic() >= deadline: + break + await asyncio.sleep(2) + if visible_gpus != expected_gpus: + raise RuntimeError( + f"Miles Ray cluster has {visible_gpus} GPUs, expected {expected_gpus}" + ) + groups = create_placement_groups(self.args) + self.rollout_manager, _ = create_rollout_manager( + self.args, groups["rollout"] + ) + self.actor_model, critic = await create_training_models( + self.args, groups, self.rollout_manager + ) + if critic is not None: + raise RuntimeError("RL v0 does not support a Miles critic") + self._trainer_awake = not self.args.offload_train + await self._onload_trainer() + tensors = await self._actor_call("yeto_rl_export_policy", rank0=True) + await self._offload_trainer() + return canonical_state( + 0, + tensors, + base_model_revision=self.args.yeto_rl_base_model_revision, + lora_config_hash=self.args.yeto_rl_lora_config_hash, + layout_hash=self.args.yeto_rl_layout_hash, + ) + + def initialize(self) -> CanonicalLoraState: + return self._run(self._initialize()) + + async def _engines(self) -> list[Any]: + info = await self.rollout_manager.get_updatable_engines_and_lock.remote() + engines = list(info.rollout_engines) + if not engines: + raise RuntimeError("Miles created no updatable SGLang engine") + return engines + + async def _pause_rollout(self) -> None: + engines = await self._engines() + # Retracted requests resume after a weight update; abort them so one + # trajectory can never cross the global policy boundary. + await asyncio.gather( + *(engine.pause_generation.remote("abort") for engine in engines) + ) + + async def _resume_rollout(self) -> None: + engines = await self._engines() + await asyncio.gather( + *(engine.continue_generation.remote() for engine in engines) + ) + + async def _set_rollout_version(self, version: int) -> None: + token = _policy_token(version) + engines = await self._engines() + await asyncio.gather( + *(engine.update_weight_version.remote(token) for engine in engines) + ) + + async def _apply_global_policy(self, state: CanonicalLoraState) -> None: + from sglang.srt.constants import ( + GPU_MEMORY_TYPE_CUDA_GRAPH, + GPU_MEMORY_TYPE_KV_CACHE, + GPU_MEMORY_TYPE_WEIGHTS, + ) + + await self._pause_rollout() + if not self._rollout_offloaded: + await self.rollout_manager.offload.remote( + tags=[ + GPU_MEMORY_TYPE_CUDA_GRAPH, + GPU_MEMORY_TYPE_KV_CACHE, + GPU_MEMORY_TYPE_WEIGHTS, + ] + ) + self._rollout_offloaded = True + await self._onload_trainer() + reset_parameter_count, applied_hash = await self._actor_call( + "yeto_rl_apply_policy", + dict(state.tensors), + state.policy_version, + ) + expected_hash = policy_hash(state) + if applied_hash != expected_hash: + raise StrictRlInvariantError( + "policy_hash_mismatch_after_apply", + "policy hash mismatch after trainer apply", + ) + await self._offload_trainer() + await self.rollout_manager.onload_weights.remote() + await self.actor_model.update_weights() + # update_weights resumes generation; close the admission boundary + # until KV/weights and the explicit version are all installed. + await self._pause_rollout() + await self.rollout_manager.onload_kv.remote() + self._rollout_offloaded = False + await self._set_rollout_version(state.policy_version) + self._policy_version = state.policy_version + self._rollout_id = state.policy_version + await self._resume_rollout() + self._optimizer_reset_count += 1 + self._append_event( + { + "event": "rl_policy_apply", + "policy_version": state.policy_version, + "optimizer_reset_count": self._optimizer_reset_count, + "reset_parameter_count": reset_parameter_count, + "rl/global_policy_version": state.policy_version, + "rl/optimizer_reset_count": self._optimizer_reset_count, + "sync/global_policy_hash": applied_hash, + } + ) + + def apply_global_policy(self, state: CanonicalLoraState) -> None: + self._run(self._apply_global_policy(state)) + + def _rollout_batches(self, data_pack) -> list[Mapping[str, Any]]: + import ray + + references = data_pack.get("data_ref") + if not isinstance(references, list): + raise RuntimeError("Miles returned an invalid rollout shard list") + expected_dp = ( + self.args.actor_num_nodes * self.args.actor_num_gpus_per_node + ) + if len(references) != expected_dp: + raise RuntimeError( + f"Miles returned {len(references)} DP shards, expected {expected_dp}" + ) + return [ray.get(reference.inner) for reference in references] + + def _rollout_metrics(self, policy_version: int) -> dict[str, float]: + path = Path(self.args.yeto_rl_completed_groups_path).expanduser() + try: + payload = torch.load(path, map_location="cpu", weights_only=True) + if ( + not isinstance(payload, dict) + or payload.get("schema_version") != _ISLAND_CHECKPOINT_SCHEMA + or payload.get("policy_version") != policy_version + or payload.get("config") != _island_checkpoint_config(self.args) + or not isinstance(payload.get("rollout_metrics"), Mapping) + ): + raise RuntimeError("Miles island checkpoint lacks rollout metrics") + return { + name: float(value) + for name, value in payload["rollout_metrics"].items() + } + except (TypeError, ValueError) as error: + raise RuntimeError("Miles returned invalid Yeto group metrics") from error + + async def _run_local_round( + self, + expected_policy_version: int, + groups: int, + samples_per_group: int, + optimizer_steps: int, + ) -> LocalRoundStats: + from sglang.srt.constants import ( + GPU_MEMORY_TYPE_CUDA_GRAPH, + GPU_MEMORY_TYPE_KV_CACHE, + GPU_MEMORY_TYPE_WEIGHTS, + ) + + if expected_policy_version != self._policy_version: + raise RuntimeError("Miles round requested from an unapplied global policy") + rollout_started = time.monotonic() + data_pack = await self.rollout_manager.generate.remote(self._rollout_id) + rollout_seconds = time.monotonic() - rollout_started + await self._pause_rollout() + batches = self._rollout_batches(data_pack) + versions = [ + sample_versions + for batch in batches + for sample_versions in batch.get("weight_versions", []) + ] + expected_samples = groups * samples_per_group + if not isinstance(versions, list) or len(versions) != expected_samples: + raise RuntimeError( + f"Miles produced {len(versions) if isinstance(versions, list) else 0} " + f"versioned samples, expected {expected_samples}" + ) + try: + observed = { + _version_from_token(token) + for sample_versions in versions + for token in sample_versions + } + except RuntimeError as error: + raise StrictRlInvariantError( + "mixed_version_group_count", + str(error), + ) from error + if any(not sample_versions for sample_versions in versions) or observed != { + expected_policy_version + }: + raise StrictRlInvariantError( + "mixed_version_group_count", + f"Miles rollout mixed policy versions: {observed}", + ) + rollout_metrics = self._rollout_metrics(expected_policy_version) + + await self.rollout_manager.offload.remote( + tags=[ + GPU_MEMORY_TYPE_CUDA_GRAPH, + GPU_MEMORY_TYPE_KV_CACHE, + GPU_MEMORY_TYPE_WEIGHTS, + ] + ) + self._rollout_offloaded = True + before = await self._actor_call("yeto_rl_optimizer_steps") + train_started = time.monotonic() + await self.actor_model.train(self._rollout_id, data_pack) + train_seconds = time.monotonic() - train_started + self._trainer_awake = True + after = await self._actor_call("yeto_rl_optimizer_steps") + if after - before != optimizer_steps: + raise RuntimeError( + f"Miles performed {after - before} optimizer steps, " + f"expected {optimizer_steps}" + ) + train_metrics = await self._actor_call( + "yeto_rl_train_metrics", + rank0=True, + ) + try: + mean_kl = float(train_metrics["train/train_rollout_kl"]) + ess_ratio = float(train_metrics["train/ess_ratio"]) + clip_fraction = float(train_metrics["train/pg_clipfrac"]) + except (KeyError, TypeError, ValueError) as error: + raise RuntimeError("Miles did not return required GRPO train metrics") from error + self._rollout_id += 1 + response_lengths = [ + int(value) + for batch in batches + for value in batch.get("response_lengths", []) + ] + sample_indices = [ + int(value) + for batch in batches + for value in batch.get("sample_indices", []) + ] + if ( + len(response_lengths) != expected_samples + or len(sample_indices) != expected_samples + or len(set(sample_indices)) != expected_samples + ): + raise RuntimeError("Miles DP rollout shards do not form one complete batch") + raw_rewards = batches[0].get("raw_reward") + if not isinstance(raw_rewards, list) or len(raw_rewards) != expected_samples: + raise RuntimeError("Miles rollout lacks scalar raw rewards") + if any(batch.get("raw_reward") != raw_rewards for batch in batches[1:]): + raise RuntimeError("Miles DP rollout shards disagree on raw rewards") + try: + rewards = [float(value) for value in raw_rewards] + except (TypeError, ValueError) as error: + raise RuntimeError("Miles RL v0 requires scalar rewards") from error + return LocalRoundStats( + island_id=int(self.args.yeto_rl_learner_id), + local_round_id=expected_policy_version + 1, + base_policy_version=expected_policy_version, + active_groups=int(rollout_metrics["active_groups"]), + completed_groups=groups, + cancelled_groups=int(rollout_metrics["cancelled_groups"]), + completed_trajectories=expected_samples, + action_tokens=sum(response_lengths), + tool_wait_seconds=rollout_metrics["tool_wait_seconds"], + group_p50_seconds=rollout_metrics["group_p50_seconds"], + group_p95_seconds=rollout_metrics["group_p95_seconds"], + group_p99_seconds=rollout_metrics["group_p99_seconds"], + reward_mean=statistics.fmean(rewards), + reward_std=statistics.pstdev(rewards), + zero_variance_group_ratio=sum( + len(set(rewards[index : index + samples_per_group])) == 1 + for index in range(0, expected_samples, samples_per_group) + ) + / groups, + mean_kl=mean_kl, + ess_ratio=ess_ratio, + clip_fraction=clip_fraction, + delta_l2_norm=0.0, + rollout_seconds=rollout_seconds, + train_seconds=train_seconds, + ) + + def run_local_round( + self, + *, + expected_policy_version: int, + groups: int, + samples_per_group: int, + optimizer_steps: int, + ) -> LocalRoundStats: + return self._run( + self._run_local_round( + expected_policy_version, + groups, + samples_per_group, + optimizer_steps, + ) + ) + + async def _export_local_policy(self) -> CanonicalLoraState: + if not self._trainer_awake: + await self._onload_trainer() + tensors = await self._actor_call("yeto_rl_export_policy", rank0=True) + await self._offload_trainer() + if self._policy_version is None: + raise RuntimeError("Miles has no applied global policy") + return canonical_state( + self._policy_version, + tensors, + base_model_revision=self.args.yeto_rl_base_model_revision, + lora_config_hash=self.args.yeto_rl_lora_config_hash, + layout_hash=self.args.yeto_rl_layout_hash, + ) + + def export_local_policy(self) -> CanonicalLoraState: + return self._run(self._export_local_policy()) + + def record_local_round(self, stats: LocalRoundStats) -> None: + path = Path(self.args.yeto_rl_completed_groups_path).expanduser() + try: + payload = torch.load(path, map_location="cpu", weights_only=True) + except Exception as error: + raise RuntimeError("cannot update Miles island checkpoint") from error + if ( + not isinstance(payload, dict) + or payload.get("schema_version") != _ISLAND_CHECKPOINT_SCHEMA + or payload.get("policy_version") != stats.base_policy_version + or payload.get("config") != _island_checkpoint_config(self.args) + ): + raise RuntimeError("Miles island checkpoint changed before round commit") + payload["local_round_id"] = stats.local_round_id + payload["local_round_stats"] = asdict(stats) + _atomic_save_island_checkpoint(path, payload) + + def _append_event(self, event: dict[str, Any]) -> None: + path = Path(self.args.yeto_rl_event_tape).expanduser() + path.parent.mkdir(parents=True, exist_ok=True) + event = { + "island_id": int(self.args.yeto_rl_learner_id), + "time_unix": time.time(), + **event, + } + with path.open("a", encoding="utf-8") as handle: + handle.write( + json.dumps(event, sort_keys=True, separators=(",", ":")) + "\n" + ) + + async def _shutdown(self) -> None: + if self.actor_model is not None: + await self._offload_trainer() + if self.rollout_manager is not None: + await self.rollout_manager.dispose.remote() + + def shutdown(self) -> None: + try: + self._run(self._shutdown()) + finally: + if self._owns_ray: + import ray + + ray.shutdown() + self.loop.close() From 3caa18a5f62be0410a092a47c6b5735f133eecd4 Mon Sep 17 00:00:00 2001 From: AlexEisie <1987460907@qq.com> Date: Thu, 30 Jul 2026 10:10:21 +0800 Subject: [PATCH 2/6] docs(rl): clarify Miles RL implementation --- docs/MILES_RL.md | 50 ++++++++++++++++++++++++++++++++++-------------- 1 file changed, 36 insertions(+), 14 deletions(-) diff --git a/docs/MILES_RL.md b/docs/MILES_RL.md index febc978..8f9a574 100644 --- a/docs/MILES_RL.md +++ b/docs/MILES_RL.md @@ -84,6 +84,9 @@ MoE `auto`/`attention` path), not LoRA on sharded expert weights. The replicated path disables Megatron's distributed optimizer so every trainer rank retains the complete fp32 LoRA masters and optimizer state needed by the export/apply contract. +When `--expert-parallel` is omitted, dense models use EP=1 and MoE models use +the full island world size. An explicit MoE EP value must divide that world +size. The RL path consumes `--lora-r` and `--lora-targets`. Its effective PEFT configuration fixes alpha to the rank, dropout to zero, and bias to `none`; @@ -105,14 +108,16 @@ The ordinary Yeto SFT and diffusion modes remain separate. Passing ## Current integration seams -The Miles checkout remains a clean detached checkout. The current adapter -installs Yeto export, apply, and optimizer-step methods on the pinned -`MegatronTrainRayActor` at process startup and invokes them through that -commit's non-FT `RayTrainGroup._broadcast` path. Yeto drives Miles rollout and -training primitives directly so the island-local result is not published to -SGLang before the global merge. There is currently no maintained Miles patch -or upstream train-loop synchronization hook. This private compatibility seam -is supported only for the pinned commit and must be revalidated if Miles is +The Miles checkout remains a clean detached checkout. At process startup the +current adapter installs Yeto export, apply, optimizer-step, and train-metric +methods on the pinned `MegatronTrainRayActor`, and invokes them through that +commit's non-FT `RayTrainGroup._broadcast` path. It also wraps the pinned +Megatron-Bridge provider/DDP construction, Miles train logging, and colocated +LoRA IPC completion in process memory. Yeto drives Miles rollout and training +primitives directly so the island-local result is not published to SGLang +before the global merge. There is currently no maintained Miles patch or +upstream train-loop synchronization hook. This private compatibility seam is +supported only for the pinned commit and must be revalidated if Miles is changed. The launcher selects strict RL behavior through the existing syncer's general @@ -134,6 +139,8 @@ The RL flags describe the work at the Miles boundary: - `G`: `--rollout-batch-size`, the complete GRPO groups per island round; - `K`: `--n-samples-per-prompt`, the trajectories per group; - local work: `--local-rl-rounds-per-sync 1` in v0. +- optimizer work: one Miles optimizer step per island round; v0 exposes no + separate optimizer-step control. `N`, `G`, and `K` must be positive. `G × K` must be divisible by every island's Miles data-parallel size. `--over-sampling-batch-size` defaults to @@ -193,12 +200,17 @@ Ray head and joins the remaining island nodes as workers. A single entry is a supported parity path; multiple entries enable fixed-roster averaging. External learner slots are not supported. -Use a Miles-compatible learner image. The image used for v0 validation is: +The default production base image is pinned by digest: ```text docker:radixark/miles@sha256:95b3afa9ee4313f5633e6ed3779c8276353cc8e24a2462e4f54ec0d5978fbae7 ``` +GPU validation used a local derivative of that exact base image with the +pinned Miles checkout and PEFT version preinstalled. The public launcher +performs those same source and PEFT setup steps on the digest-pinned base +image. + The Miles source itself is independently pinned to: ```text @@ -398,11 +410,13 @@ Recovery always starts from the most recent committed global checkpoint: | syncer VM or disk is lost | the current launcher has no durable mount for the global checkpoint, so automatic recovery is not available | | learner exits while applying a global policy | replacement reapplies the complete committed policy | -A dead syncer connection makes an island process exit rather than silently -continue local work. The launcher can restart the syncer through its existing -recovery path and restart affected logical learners individually; it does not -need to restart every healthy roster member. Duplicate computation is allowed, -but duplicate merge is not. +A dead syncer connection makes an island exit at the next bridge health check. +If it drops while a synchronous Miles rollout/train call is in progress, that +local round may finish redundant computation, but it cannot begin another +round or become authoritative after restart. The launcher can restart the +syncer through its existing recovery path; each learner whose bridge exits is +restarted under the same logical ID rather than through a roster-wide restart +command. Duplicate computation is allowed, but duplicate merge is not. Each island writes policy-apply, optimizer-reset, and `LocalRoundStats` JSONL; the syncer writes roster, base-version, layout, merge, and responder metrics to @@ -446,6 +460,14 @@ Common startup and progress failures have distinct meanings: - **optimizer steps but negligible LoRA change:** first inspect reward variance, advantages, and the configured LR schedule. +## Intentional differences from INIT + +| INIT plan | Current implementation | Assessment | +| --- | --- | --- | +| **Miles integration:** maintain a thin Miles branch with stable policy export, apply, and post-train synchronization hooks. | Keep the pinned upstream Miles checkout unchanged and adapt that exact version at runtime. | The required training and synchronization semantics are implemented and validated. A Miles upgrade still requires explicit compatibility revalidation. | +| **Global checkpoint recovery:** keep the authoritative policy recoverable across syncer failures, including replacement of its machine or disk. | Automatic recovery works while the syncer's disk is retained; losing that VM or disk also loses the checkpoint. | This does not change the RL algorithm, but it remains a production disaster-recovery gap. | +| **Monitoring:** connect the planned RL and synchronization metrics to a dashboard. | Emit the metrics to JSONL without enabling a dashboard. | The records are sufficient for validation and offline diagnosis, but not centralized live monitoring. | + ## Validation status The current source has automated coverage for multi-node task construction, From 9ebd42060414fecfca2cabd89f669e91f93d1041 Mon Sep 17 00:00:00 2001 From: AlexEisie <1987460907@qq.com> Date: Thu, 30 Jul 2026 12:28:51 +0800 Subject: [PATCH 3/6] feat(rl): adopt maintained Miles sync hooks Pin agentenv/miles yeto-sync and move replicated LoRA export/apply, optimizer reset, native post-train synchronization, and round metrics behind its public boundary. Keep Miles in control of rollout, GRPO training, offload, and SGLang publication while removing the old actor, train-loop, provider, logging, and IPC runtime injection. Validate rollout policy versions before optimizer work, preserve exact-base fixed-roster averaging and completed-group recovery, and document the maintained integration. The remaining INIT differences are syncer checkpoint loss with its VM/disk and JSONL-only monitoring without a dashboard. --- docs/MILES_RL.md | 111 +++-- tests/test_rl_core.py | 776 ++++++++----------------------- tests/test_rl_integration.py | 223 ++++++++- tests/test_rl_launcher.py | 5 + yeto/rl/__init__.py | 4 +- yeto/rl/bridge.py | 93 +++- yeto/rl/core.py | 9 +- yeto/rl/learner.py | 43 +- yeto/rl/miles.py | 858 ++++++++++------------------------- 9 files changed, 832 insertions(+), 1290 deletions(-) diff --git a/docs/MILES_RL.md b/docs/MILES_RL.md index 8f9a574..daf7dad 100644 --- a/docs/MILES_RL.md +++ b/docs/MILES_RL.md @@ -1,7 +1,7 @@ # Miles RL v0 Miles RL runs reinforcement learning in several independent -[Miles](https://github.com/radixark/miles) islands and uses Yeto to turn their +[Miles](https://github.com/agentenv/miles) islands and uses Yeto to turn their complete local LoRA updates into one committed global policy. Miles owns the rollout and GRPO loop inside each island; Yeto owns fleet lifecycle, the cross-island synchronization boundary, recovery, and the final adapter. @@ -10,10 +10,12 @@ This mode is not Decoupled DiLoCo. It is synchronous, fixed-roster LoRA FedAvg: every global round waits for one exact-base result from every logical island and averages them equally. -> **Status:** the core path is implemented and has passed real multi-GPU dense, -> MoE EP, two-island averaging, recovery, long-task, and 20-merge validation on -> eight A100 GPUs. The 24-hour soak and the operational gaps listed in -> [Validation status](#validation-status) remain outside that evidence. +> **Status:** the core path and the pinned Miles synchronization branch are +> implemented. The algorithm and model path previously passed real multi-GPU +> dense, MoE EP, two-island averaging, recovery, long-task, and 20-merge +> validation on eight A100 GPUs. The thin-branch boundary has automated +> integration coverage but has not yet repeated that GPU matrix; see +> [Validation status](#validation-status). ## Why Miles is the RL runtime @@ -108,17 +110,27 @@ The ordinary Yeto SFT and diffusion modes remain separate. Passing ## Current integration seams -The Miles checkout remains a clean detached checkout. At process startup the -current adapter installs Yeto export, apply, optimizer-step, and train-metric -methods on the pinned `MegatronTrainRayActor`, and invokes them through that -commit's non-FT `RayTrainGroup._broadcast` path. It also wraps the pinned -Megatron-Bridge provider/DDP construction, Miles train logging, and colocated -LoRA IPC completion in process memory. Yeto drives Miles rollout and training -primitives directly so the island-local result is not published to SGLang -before the global merge. There is currently no maintained Miles patch or -upstream train-loop synchronization hook. This private compatibility seam is -supported only for the pinned commit and must be revalidated if Miles is -changed. +Yeto pins the `agentenv/miles` `yeto-sync` branch by commit. That branch adds +the narrow boundary required by this mode: + +- `RayTrainGroup.export_trainable_state()` exports the complete replicated + LoRA from Megatron global rank zero in canonical PEFT form; +- `RayTrainGroup.apply_trainable_state()` validates and writes the complete + global LoRA on every trainer rank, clears LoRA optimizer state, preserves LR + scheduler progress, and refreshes the actor backup; +- the native Miles train loop loads one external policy-sync hook, calls it + after local training and before the normal trainer-to-SGLang publication, + and finalizes it only after the final global policy has been published. + +Miles still owns rollout generation, reward and advantage computation, GRPO +training, offload, checkpoint calls, and SGLang weight transport. Yeto does +not call private actor-group broadcast methods or replace Miles' local train +loop. The hook applies the committed global LoRA before Miles publishes +weights, so an island-local LoRA is never exposed to the next rollout. + +The pinned source is checked out as a clean detached commit. The public +boundary is deliberately limited to the non-FT Megatron actor group used by +v0; Miles' experimental fault-tolerant trainer is not adapted. The launcher selects strict RL behavior through the existing syncer's general controls: `--max-base-lag 0`, `--learner-weight equal`, `--quorum M`, and @@ -158,8 +170,8 @@ At committed version `v`, each island follows the same sequence: 1. Apply the complete global LoRA `theta_v` to the Megatron trainer and SGLang, then mark rollout policy version `v` active. 2. Generate exactly `G` groups of `K` terminal trajectories (completed or - truncated at the configured limit). Every recorded rollout weight version - must be `v`. + truncated at the configured limit). Before training starts, every recorded + rollout weight version must be exactly `v`. 3. Run the one configured Miles GRPO training cycle without publishing the local result to SGLang. 4. Export the complete local LoRA `theta_i_v` and send @@ -214,15 +226,15 @@ image. The Miles source itself is independently pinned to: ```text -https://github.com/radixark/miles -dfc66ff38752bfa2c5d325e0037ebc4b537c06de +https://github.com/agentenv/miles +a91bd34e50416aeb1da111f74d52b296e8216b96 ``` The launcher checks out that commit as a detached HEAD and installs the project's pinned PEFT version. At learner startup Yeto verifies the repository -origin, commit, clean worktree, and imported package path. Runtime -adaptation uses the compatibility seam described above without modifying the -checkout. +origin, commit, clean worktree, and imported package path. The checkout itself +contains the thin policy-sync boundary described above; Yeto does not modify +it at process startup. An illustrative two-island run is: @@ -423,9 +435,11 @@ the syncer writes roster, base-version, layout, merge, and responder metrics to `~/yeto-output/yeto-tape.jsonl`. Miles group-task completion supplies peak active groups, cancellations, and duration percentiles; `Sample.non_generation_time` supplies tool wait; grouped raw rewards supply the -zero-variance ratio; and Miles' rank-zero train log supplies KL, ESS, and clip -fraction. The bridge records the actual protocol payload bytes, while each -canonical global apply records its policy hash. +zero-variance ratio. The thin Miles branch attaches the rank-zero native train +log's KL, ESS, and clip fraction, plus the measured optimizer-train duration, +to the same local-state export consumed by the hook. The bridge records those +values, the actual protocol payload bytes, and each canonical global apply's +policy hash. The launcher does not currently enable a Miles or Yeto dashboard for these records. @@ -438,18 +452,20 @@ recoverable through the committed-checkpoint paths above. ## Runtime compatibility and diagnostics -The validated Miles/Transformer Engine stack cannot reliably use its -FlashAttention CUTE GQA kernel on A100. RL islands therefore select unfused -attention before Transformer Engine imports and propagate that choice into -the actual Megatron provider. This is a runtime-wide choice, not a model -family workaround. +Model loading, attention kernels, LoRA conversion, and trainer-to-SGLang +transport remain Miles responsibilities. Yeto passes the selected Miles +arguments and does not patch Megatron providers or transport functions at +runtime. A change to the pinned Miles, Megatron-Bridge, Transformer Engine, or +SGLang stack therefore requires a real GPU compatibility run in addition to +the tensor-contract tests. Common startup and progress failures have distinct meanings: - **PEFT import failure:** the learner did not run the current Miles setup, which installs the pinned `peft==0.20.0`. -- **`Operation creation failed` under `flash_attn/cute/pack_gqa.py`:** the - process did not receive the RL launch environment or provider backend. +- **Attention-kernel startup failure:** check the pinned Miles image and its + Megatron/Transformer Engine compatibility; do not add a model-name branch in + Yeto. - **Miles revision/origin/dirty-tree error:** the runtime is not using the supported checkout; do not bypass this check or patch that tree in place. - **PEFT/Megatron mapping mismatch:** the model is outside the currently @@ -464,25 +480,24 @@ Common startup and progress failures have distinct meanings: | INIT plan | Current implementation | Assessment | | --- | --- | --- | -| **Miles integration:** maintain a thin Miles branch with stable policy export, apply, and post-train synchronization hooks. | Keep the pinned upstream Miles checkout unchanged and adapt that exact version at runtime. | The required training and synchronization semantics are implemented and validated. A Miles upgrade still requires explicit compatibility revalidation. | | **Global checkpoint recovery:** keep the authoritative policy recoverable across syncer failures, including replacement of its machine or disk. | Automatic recovery works while the syncer's disk is retained; losing that VM or disk also loses the checkpoint. | This does not change the RL algorithm, but it remains a production disaster-recovery gap. | | **Monitoring:** connect the planned RL and synchronization metrics to a dashboard. | Emit the metrics to JSONL without enabling a dashboard. | The records are sufficient for validation and offline diagnosis, but not centralized live monitoring. | ## Validation status -The current source has automated coverage for multi-node task construction, -multi-rank actor results, DP rollout-shard collection, EP validation, +The current source has automated coverage for the Miles policy hook and its +ordering around native SGLang publication, multi-node task construction, +multi-rank apply/export, DP rollout-shard collection, EP validation, single-island and multi-island sync, canonical identity, completed-group recovery, strict failures, provenance, checkpoint export, and the unchanged -SFT/diffusion defaults. The 2026-07-30 regression passed all 73 focused RL -tests, the full Python suite (`766 passed, 4 skipped`), `cargo fmt --check`, and -all 58 Rust tests. +SFT/diffusion defaults. Exact current test counts are recorded in the pull +request rather than frozen in this document. -Real validation ran on one GCP Spot VM with eight NVIDIA A100-SXM4-40GB GPUs. -It used the pinned Miles commit, PEFT 0.20.0, immutable model and dataset -revisions, real model generations, real rewards, and production Yeto learner -and Rust syncer paths. Observation hooks only captured tensors, tokens, -metrics, and fault windows. +The following real validation ran on one GCP Spot VM with eight NVIDIA +A100-SXM4-40GB GPUs before the integration moved from runtime adaptation to +the maintained Miles hook. It used the same policy, merge, model, reward, +learner, and Rust syncer semantics, but it is evidence for the RL algorithm +and model path rather than GPU validation of the new hook implementation. - A one-island Qwen3-4B DP=8 run completed two global rounds. Each round trained on 16 trajectories and 8192 action tokens with nonconstant rewards. @@ -525,6 +540,7 @@ metrics, and fault windows. The following boundaries remain unvalidated or intentionally excluded: +- the maintained Miles hook has not yet repeated the real GPU matrix above; - the requested 24-hour soak was not run; the 20 consecutive merges are the bounded-duration stability evidence; - no physical multi-node island or end-to-end SkyPilot provisioning and Spot @@ -533,9 +549,6 @@ The following boundaries remain unvalidated or intentionally excluded: - the syncer checkpoint still has no durable mount for syncer VM or disk loss; - metrics remain JSONL-only and the launcher enables no dashboard. -The runtime-injection Miles seam described above remains an implementation -constraint, not a stable upstream interface. - ## Extending v0 safely The useful extension boundary is the observable contract, not a model-name @@ -563,10 +576,12 @@ For implementation navigation: | area | responsibility | | --- | --- | | `yeto/rl/core.py` | canonical LoRA state and the single AVG layout | -| `yeto/rl/miles.py` | pinned Miles runtime, trainer/SGLang policy boundary | +| `yeto/rl/miles.py` | external sync hook, rollout-version checks, and Yeto bridge adapter | | `yeto/rl/bridge.py` | exact-base island loop and protocol interaction | | `yeto/rl/export.py` | committed checkpoint to PEFT adapter | | `yeto/rl/learner.py` | island entry point and Miles argument mapping | +| `agentenv/miles:miles/backends/megatron_utils/trainable_state.py` | replicated LoRA export/apply and optimizer reset | +| `agentenv/miles:train.py` | native post-train, pre-publication hook ordering | | `syncer/src/server.rs` | fixed roster and checkpoint-before-broadcast commit | Run the focused and full regressions before a GPU shakedown: diff --git a/tests/test_rl_core.py b/tests/test_rl_core.py index 49b2fed..44c0a72 100644 --- a/tests/test_rl_core.py +++ b/tests/test_rl_core.py @@ -9,8 +9,10 @@ import pytest import torch -import yeto.rl.miles as miles from yeto.fragments import MERGE_AVG +from yeto.protocol import PullRequest +from yeto.rl import miles +from yeto.rl.bridge import BridgeConfig, StrictRlBridge from yeto.rl.core import ( CanonicalTensorSpec, LocalRoundStats, @@ -22,10 +24,8 @@ policy_delta, tensors_from_flat, ) -from yeto.rl.bridge import BridgeConfig, StrictRlBridge -from yeto.protocol import PullRequest from yeto.rl.export import adapter_targets, derive_peft_lora_specs -from yeto.rl.miles import MilesIslandRuntime +from yeto.rl.miles import MilesPolicySync def tensors(): @@ -281,187 +281,6 @@ def test_local_round_event_contains_every_init_metric(tmp_path): assert event["sync/bytes_sent"] == 48 + len(pushed[0][-1]) -def test_miles_admission_pause_aborts_inflight_rollouts(): - modes = [] - - class PauseGeneration: - async def remote(self, mode): - modes.append(mode) - - class Engine: - pause_generation = PauseGeneration() - - runtime = object.__new__(MilesIslandRuntime) - - async def engines(): - return [Engine()] - - runtime._engines = engines - asyncio.run(runtime._pause_rollout()) - assert modes == ["abort"] - - -def test_miles_bridge_propagates_attention_backend(monkeypatch): - class AutoBridge: - def to_megatron_provider(self): - return SimpleNamespace(attention_backend="auto") - - megatron = types.ModuleType("megatron") - bridge = types.ModuleType("megatron.bridge") - training = types.ModuleType("megatron.bridge.training") - config = types.ModuleType("megatron.bridge.training.config") - - class DistributedDataParallelConfig: - def __init__(self, **kwargs): - self.use_distributed_optimizer = kwargs["use_distributed_optimizer"] - - bridge.AutoBridge = AutoBridge - config.DistributedDataParallelConfig = DistributedDataParallelConfig - training.config = config - monkeypatch.setitem(sys.modules, "megatron", megatron) - monkeypatch.setitem(sys.modules, "megatron.bridge", bridge) - monkeypatch.setitem(sys.modules, "megatron.bridge.training", training) - monkeypatch.setitem(sys.modules, "megatron.bridge.training.config", config) - installed = [] - monkeypatch.setattr(miles, "_install_train_metric_capture", lambda: None) - monkeypatch.setattr(miles, "_install_colocated_lora_ipc_sync", lambda: None) - monkeypatch.setattr(miles, "install_miles_actor_adapter", lambda: installed.append(True)) - - args = SimpleNamespace( - attention_backend="unfused", use_distributed_optimizer=True - ) - miles.configure_miles_bridge(args) - - assert AutoBridge().to_megatron_provider().attention_backend == "unfused" - assert not args.use_distributed_optimizer - assert not config.DistributedDataParallelConfig( - use_distributed_optimizer=True - ).use_distributed_optimizer - assert installed == [True] - - -def test_miles_actor_calls_handle_rank0_export_and_all_rank_results(): - runtime = object.__new__(MilesIslandRuntime) - runtime.args = SimpleNamespace(actor_num_nodes=1, actor_num_gpus_per_node=2) - - class Actors: - def __init__(self, results): - self.results = results - - async def _broadcast(self, method, *args): - return self.results - - runtime.actor_model = Actors([{"policy": 1}, None]) - assert asyncio.run(runtime._actor_call("export", rank0=True)) == {"policy": 1} - runtime.actor_model = Actors([(4, "hash"), (4, "hash")]) - assert asyncio.run(runtime._actor_call("apply")) == (4, "hash") - runtime.actor_model = Actors([3, 4]) - with pytest.raises(RuntimeError, match="ranks disagree"): - asyncio.run(runtime._actor_call("steps")) - - -def test_miles_waits_for_the_complete_multinode_ray_island(monkeypatch): - canonical = state(0, tensors()) - args = SimpleNamespace( - actor_num_nodes=2, - actor_num_gpus_per_node=4, - offload_train=True, - yeto_rl_base_model_revision=canonical.base_model_revision, - yeto_rl_lora_config_hash=canonical.lora_config_hash, - yeto_rl_layout_hash=canonical.layout_hash, - ) - resource_counts = iter((4, 8)) - ray = types.ModuleType("ray") - ray.is_initialized = lambda: True - ray.cluster_resources = lambda: {"GPU": next(resource_counts)} - actor = SimpleNamespace() - - async def broadcast(method, *values): - assert method == "yeto_rl_export_policy" - return [canonical.tensors] + [None] * 7 - - async def onload(): - pass - - async def offload(): - pass - - actor._broadcast = broadcast - actor.onload = onload - actor.offload = offload - - async def create_training_models(*values): - return actor, None - - placement = types.ModuleType("miles.ray.placement_group") - placement.create_placement_groups = lambda _args: {"rollout": object()} - placement.create_rollout_manager = lambda *values: (object(), None) - placement.create_training_models = create_training_models - external_miles = types.ModuleType("miles") - ray_package = types.ModuleType("miles.ray") - ray_package.placement_group = placement - external_miles.ray = ray_package - monkeypatch.setitem(sys.modules, "ray", ray) - monkeypatch.setitem(sys.modules, "miles", external_miles) - monkeypatch.setitem(sys.modules, "miles.ray", ray_package) - monkeypatch.setitem(sys.modules, "miles.ray.placement_group", placement) - monkeypatch.setattr(miles, "install_miles_actor_adapter", lambda: None) - - async def no_wait(_seconds): - pass - - monkeypatch.setattr(miles.asyncio, "sleep", no_wait) - runtime = MilesIslandRuntime(args) - try: - initialized = asyncio.run(runtime._initialize()) - finally: - runtime.loop.close() - - assert initialized.layout_hash == canonical.layout_hash - - -def test_miles_keeps_non_offloaded_trainer_resident(): - calls = [] - - async def onload(): - calls.append("onload") - - async def offload(): - calls.append("offload") - - runtime = object.__new__(MilesIslandRuntime) - runtime.args = SimpleNamespace(offload_train=False) - runtime.actor_model = SimpleNamespace(onload=onload, offload=offload) - runtime._trainer_awake = True - - asyncio.run(runtime._onload_trainer()) - asyncio.run(runtime._offload_trainer()) - assert runtime._trainer_awake - assert calls == [] - - -def test_miles_collects_every_data_parallel_rollout_shard(monkeypatch): - monkeypatch.setitem(sys.modules, "ray", SimpleNamespace(get=lambda value: value)) - runtime = object.__new__(MilesIslandRuntime) - runtime.args = SimpleNamespace( - actor_num_nodes=1, - actor_num_gpus_per_node=4, - expert_model_parallel_size=2, - ) - shards = [ - SimpleNamespace(inner={"sample_indices": [0]}), - SimpleNamespace(inner={"sample_indices": [1]}), - SimpleNamespace(inner={"sample_indices": [2]}), - SimpleNamespace(inner={"sample_indices": [3]}), - ] - assert runtime._rollout_batches({"data_ref": shards}) == [ - {"sample_indices": [0]}, - {"sample_indices": [1]}, - {"sample_indices": [2]}, - {"sample_indices": [3]}, - ] - - def test_miles_rollout_lifecycle_metrics_use_real_task_completion(monkeypatch): class Sample: def __init__(self, status): @@ -508,213 +327,73 @@ async def run(): assert len(lifecycle["durations"]) == 2 -def test_miles_train_metric_capture_returns_rank_zero_values(monkeypatch): - args = SimpleNamespace() - - def log_train_step(*_values, **_kwargs): - return { - "train/train_rollout_kl": 0.1, - "train/ess_ratio": 0.9, - "train/pg_clipfrac": 0.2, - } - - model = types.ModuleType("miles.backends.megatron_utils.model") - model.log_train_step = log_train_step - megatron_utils = types.ModuleType("miles.backends.megatron_utils") - megatron_utils.model = model - backends = types.ModuleType("miles.backends") - backends.megatron_utils = megatron_utils +@pytest.mark.parametrize( + "versions", + [ + ["yeto:2"], + ["yeto:3", "yeto:2"], + [], + ["invalid:3"], + 3, + ], + ids=["stale", "mixed", "empty", "invalid", "malformed"], +) +def test_generate_rollout_rejects_invalid_policy_versions_before_train( + tmp_path, monkeypatch, capsys, versions +): + upstream = types.ModuleType("miles.rollout.sglang_rollout") + upstream.generate_rollout = object() + rollout = types.ModuleType("miles.rollout") + rollout.sglang_rollout = upstream package = types.ModuleType("miles") - package.backends = backends + package.rollout = rollout monkeypatch.setitem(sys.modules, "miles", package) - monkeypatch.setitem(sys.modules, "miles.backends", backends) - monkeypatch.setitem(sys.modules, "miles.backends.megatron_utils", megatron_utils) - monkeypatch.setitem(sys.modules, "miles.backends.megatron_utils.model", model) - - miles._install_train_metric_capture() - model.log_train_step(args=args) - actor = SimpleNamespace(args=args) - assert miles._actor_train_metrics(actor) == { - "train/train_rollout_kl": 0.1, - "train/ess_ratio": 0.9, - "train/pg_clipfrac": 0.2, - } - assert miles._actor_train_metrics(actor) is None - - -def test_miles_keeps_colocated_lora_ipc_storage_alive_until_transfer(monkeypatch): - events = [] - update = types.ModuleType( - "miles.backends.megatron_utils.update_weight.update_weight_from_tensor" - ) - - def send(*_values, **_kwargs): - events.append("send") - return ["ref"], "storage" - - update._send_to_colocated_engine = send - common = types.ModuleType("miles.backends.megatron_utils.update_weight.common") - - def check(results, *, is_lora): - assert results == ["done"] and is_lora - events.append("check") + monkeypatch.setitem(sys.modules, "miles.rollout", rollout) + monkeypatch.setitem(sys.modules, "miles.rollout.sglang_rollout", upstream) - common._check_weight_sync_results = check - weight = types.ModuleType("miles.backends.megatron_utils.update_weight") - weight.update_weight_from_tensor = update - megatron_utils = types.ModuleType("miles.backends.megatron_utils") - megatron_utils.update_weight = weight - backends = types.ModuleType("miles.backends") - backends.megatron_utils = megatron_utils - package = types.ModuleType("miles") - package.backends = backends - ray = types.ModuleType("ray") - - def get(refs): - assert refs == ["ref"] - events.append("get") - return ["done"] - - ray.get = get - for name, module in { - "miles": package, - "miles.backends": backends, - "miles.backends.megatron_utils": megatron_utils, - "miles.backends.megatron_utils.update_weight": weight, - "miles.backends.megatron_utils.update_weight.update_weight_from_tensor": update, - "miles.backends.megatron_utils.update_weight.common": common, - "ray": ray, - }.items(): - monkeypatch.setitem(sys.modules, name, module) + samples = [ + SimpleNamespace( + status=SimpleNamespace(value="completed"), + weight_versions=versions, + index=0, + ), + SimpleNamespace( + status=SimpleNamespace(value="completed"), + weight_versions=["yeto:3"], + index=1, + ), + ] + output = SimpleNamespace(samples=[samples], metrics={}) monkeypatch.setattr( - torch.distributed, - "barrier", - lambda *, group: events.append(("barrier", group)), - ) - - miles._install_colocated_lora_ipc_sync() - refs, storage = update._send_to_colocated_engine( - [], ipc_gather_group="group", lora_config={} + miles, + "_run_rollout_with_metrics", + lambda *_args, **_kwargs: ( + output, + {"active": 0, "peak_active": 0, "cancelled": 0, "durations": []}, + ), ) - - assert refs == ["ref"] and storage == "storage" - assert events == ["send", "get", "check", ("barrier", "group")] - - -def test_miles_round_stats_use_checkpoint_rollout_and_train_metrics( - tmp_path, monkeypatch -): - constants = types.ModuleType("sglang.srt.constants") - constants.GPU_MEMORY_TYPE_CUDA_GRAPH = "graph" - constants.GPU_MEMORY_TYPE_KV_CACHE = "kv" - constants.GPU_MEMORY_TYPE_WEIGHTS = "weights" - sglang = types.ModuleType("sglang") - srt = types.ModuleType("sglang.srt") - sglang.srt = srt - srt.constants = constants - monkeypatch.setitem(sys.modules, "sglang", sglang) - monkeypatch.setitem(sys.modules, "sglang.srt", srt) - monkeypatch.setitem(sys.modules, "sglang.srt.constants", constants) - - checkpoint = tmp_path / "island.pt" + monkeypatch.setattr(miles, "_save_completed_groups", lambda *_args: None) + event_tape = tmp_path / "events.jsonl" args = SimpleNamespace( - actor_num_gpus_per_node=1, - actor_num_nodes=1, - advantage_estimator="grpo", - yeto_rl_model="org/model", - yeto_rl_data="org/data", - yeto_rl_base_model_revision=MODEL_REVISION, - yeto_rl_data_revision="d" * 40, - expert_model_parallel_size=1, - yeto_rl_layout_hash="c" * 64, - lr=1e-4, - yeto_rl_lora_config_hash=LORA_CONFIG_HASH, n_samples_per_prompt=2, - num_steps_per_rollout=1, - over_sampling_batch_size=2, - yeto_rl_reward_sha256="e" * 64, - rollout_batch_size=2, - seq_length=128, - seed=7, - rollout_max_response_len=16, - custom_generate_function_path=None, - use_session_server=False, - tito_model="default", - yeto_rl_completed_groups_path=str(checkpoint), - yeto_rl_learner_id=0, - ) - torch.save( - { - "schema_version": miles._ISLAND_CHECKPOINT_SCHEMA, - "config": miles._island_checkpoint_config(args), - "policy_version": 3, - "rollout_metrics": { - "active_groups": 3.0, - "cancelled_groups": 1.0, - "tool_wait_seconds": 4.0, - "group_p50_seconds": 5.0, - "group_p95_seconds": 6.0, - "group_p99_seconds": 7.0, - }, - }, - checkpoint, - ) - - class Remote: - def __init__(self, result=None): - self.result = result - - async def remote(self, *_args, **_kwargs): - return self.result - - runtime = object.__new__(MilesIslandRuntime) - runtime.args = args - runtime._policy_version = 3 - runtime._rollout_id = 3 - runtime._rollout_offloaded = False - runtime.rollout_manager = SimpleNamespace( - generate=Remote(object()), - offload=Remote(), + rollout_batch_size=1, + yeto_rl_completed_groups_path=str(tmp_path / "island.pt"), + yeto_rl_event_tape=str(event_tape), + yeto_rl_learner_id=7, ) - async def train(*_args): - pass - - runtime.actor_model = SimpleNamespace(train=train) - - async def no_op(): - pass - - runtime._pause_rollout = no_op - runtime._rollout_batches = lambda _pack: [ - { - "weight_versions": [["yeto:3"]] * 4, - "response_lengths": [1, 2, 3, 4], - "sample_indices": [0, 1, 2, 3], - "raw_reward": [1.0, 1.0, 0.0, 2.0], - } - ] - optimizer_steps = iter((10, 11)) - - async def actor_call(method, *_args, **_kwargs): - if method == "yeto_rl_optimizer_steps": - return next(optimizer_steps) - assert method == "yeto_rl_train_metrics" - return { - "train/train_rollout_kl": 0.1, - "train/ess_ratio": 0.8, - "train/pg_clipfrac": 0.25, - } + with pytest.raises(StrictRlInvariantError) as failure: + miles.generate_rollout(args, 3, SimpleNamespace(buffer=[])) - runtime._actor_call = actor_call - stats = asyncio.run(runtime._run_local_round(3, 2, 2, 1)) - assert stats.active_groups == 3 - assert stats.cancelled_groups == 1 - assert stats.tool_wait_seconds == 4.0 - assert stats.zero_variance_group_ratio == 0.5 - assert stats.mean_kl == 0.1 - assert stats.ess_ratio == 0.8 - assert stats.clip_fraction == 0.25 + assert failure.value.metric == "mixed_version_group_count" + assert ( + "[yeto-rl-strict-failure] mixed_version_group_count" + in capsys.readouterr().err + ) + event = json.loads(event_tape.read_text()) + assert event["event"] == "rl_strict_failure" + assert event["metric"] == "mixed_version_group_count" + assert event["island_id"] == 7 def test_island_checkpoint_restores_only_complete_same_policy_groups( @@ -851,6 +530,7 @@ def submit_generate_tasks(self, _samples): assert source.buffer == [unused] payload = torch.load(checkpoint, weights_only=True) assert [sample["index"] for sample in payload["completed_groups"][0]] == [20, 21] + assert payload["rollout_metrics"].pop("rollout_seconds") >= 0 assert payload["rollout_metrics"] == { "reward": 2.0, "active_groups": 0.0, @@ -862,213 +542,157 @@ def submit_generate_tasks(self, _samples): } -def test_miles_apply_resets_optimizer_and_restores_scheduler_progress(monkeypatch): - name = "base_model.model.layer.lora_A.weight" - parameter = torch.nn.Parameter(torch.zeros(1, 2)) - parameter.main_param = torch.zeros(1, 2, dtype=torch.float32) - - class Mapping: - def hf_to_megatron(self, value, _module): - return value - - def megatron_to_hf(self, value, _module): - return {name: value} - - side = SimpleNamespace( - mapping=Mapping(), - megatron_module=None, - param_weight=parameter, - ) - optimizer_state = { - "exp_avg": torch.ones_like(parameter.main_param), - "exp_avg_sq": torch.ones_like(parameter.main_param), - "step": torch.tensor(9.0), - } - inner = SimpleNamespace( - param_groups=[{"step": 9}], - state={parameter.main_param: optimizer_state}, +def test_miles_policy_hook_uses_public_trainable_state_api(tmp_path, monkeypatch): + canonical = state(1, tensors()) + trainable_state = types.ModuleType( + "miles.backends.megatron_utils.trainable_state" ) - unrelated = torch.nn.Parameter(torch.ones(1)) - unrelated_state = {"step": torch.tensor(4.0)} - inner.state[unrelated] = unrelated_state - - def copy_main_to_model(): - parameter.copy_(parameter.main_param) - - child = SimpleNamespace( - optimizer=inner, - _copy_main_params_to_model_params=copy_main_to_model, + trainable_state.make_trainable_state = lambda version, values: SimpleNamespace( + policy_version=version, + layout_hash=canonical.layout_hash, + tensors=values, ) - optimizer = SimpleNamespace(chained_optimizers=[child]) - scheduler_steps = [] - - class Scheduler: - num_steps = 0 - - def step(self, increment): - scheduler_steps.append(increment) - self.num_steps += increment - - scheduler = Scheduler() - backups = [] - applied = torch.tensor([[3.0, 4.0]]) - state_fn = state(2, {name: applied}) - actor = SimpleNamespace( - args=SimpleNamespace( - yeto_rl_base_model_revision=MODEL_REVISION, - yeto_rl_lora_config_hash=LORA_CONFIG_HASH, - yeto_rl_layout_hash=state_fn.layout_hash, - global_batch_size=64, - num_steps_per_rollout=1, - ), - model=[SimpleNamespace(start_param_sync=lambda **_kwargs: pytest.fail( - "replicated LoRA apply must not start distributed optimizer sync" - ))], - optimizer=optimizer, - opt_param_scheduler=scheduler, - weights_backuper=SimpleNamespace(backup=backups.append), + for name in ( + "miles", + "miles.backends", + "miles.backends.megatron_utils", + ): + monkeypatch.setitem(sys.modules, name, types.ModuleType(name)) + monkeypatch.setitem( + sys.modules, + "miles.backends.megatron_utils.trainable_state", + trainable_state, ) - monkeypatch.setattr(miles, "_adapter_sides", lambda _actor: [(name, side)]) - cache_releases = [] - monkeypatch.setattr(torch.cuda, "empty_cache", lambda: cache_releases.append(True)) - - reset_count, applied_hash = miles._actor_apply_policy(actor, {name: applied}, 2) - assert actor.optimizer is optimizer - assert actor.opt_param_scheduler is scheduler - assert scheduler.num_steps == 128 - assert scheduler_steps == [128] - assert inner.param_groups[0]["step"] == 9 - assert parameter.main_param not in inner.state - assert inner.state[unrelated] is unrelated_state - assert reset_count == 1 - assert applied_hash == miles.policy_hash(state_fn) - assert torch.equal(parameter.main_param, applied) - assert torch.equal(parameter, applied) - assert backups == ["actor"] - assert cache_releases == [True] - - scheduler.num_steps = 192 - with pytest.raises(RuntimeError, match="ahead of the committed policy"): - miles._actor_apply_policy(actor, {name: applied}, 2) - - -def test_miles_apply_hash_mismatch_is_a_strict_failure(monkeypatch): - constants = types.ModuleType("sglang.srt.constants") - constants.GPU_MEMORY_TYPE_CUDA_GRAPH = "graph" - constants.GPU_MEMORY_TYPE_KV_CACHE = "kv" - constants.GPU_MEMORY_TYPE_WEIGHTS = "weights" - sglang = types.ModuleType("sglang") - srt = types.ModuleType("sglang.srt") - sglang.srt = srt - srt.constants = constants - monkeypatch.setitem(sys.modules, "sglang", sglang) - monkeypatch.setitem(sys.modules, "sglang.srt", srt) - monkeypatch.setitem(sys.modules, "sglang.srt.constants", constants) - - runtime = object.__new__(MilesIslandRuntime) - runtime._trainer_awake = False - runtime._rollout_offloaded = True - - async def no_op(*_args, **_kwargs): - pass - - runtime._pause_rollout = no_op - runtime.actor_model = SimpleNamespace(onload=no_op) - async def actor_call(*_args, **_kwargs): - return 1, "wrong-policy-hash" - - runtime._actor_call = actor_call - with pytest.raises(StrictRlInvariantError) as failure: - asyncio.run(runtime._apply_global_policy(state(0, tensors()))) - assert failure.value.metric == "policy_hash_mismatch_after_apply" - - -def test_miles_policy_apply_event_has_only_namespaced_policy_hash(monkeypatch): - constants = types.ModuleType("sglang.srt.constants") - constants.GPU_MEMORY_TYPE_CUDA_GRAPH = "graph" - constants.GPU_MEMORY_TYPE_KV_CACHE = "kv" - constants.GPU_MEMORY_TYPE_WEIGHTS = "weights" - sglang = types.ModuleType("sglang") - srt = types.ModuleType("sglang.srt") - sglang.srt = srt - srt.constants = constants - monkeypatch.setitem(sys.modules, "sglang", sglang) - monkeypatch.setitem(sys.modules, "sglang.srt", srt) - monkeypatch.setitem(sys.modules, "sglang.srt.constants", constants) - - async def no_op(*_args, **_kwargs): - pass + versions = [] class Remote: - async def remote(self, *_args, **_kwargs): - pass + def __init__(self, function): + self.function = function - state_fn = state(1, tensors()) - events = [] - runtime = object.__new__(MilesIslandRuntime) - runtime.args = SimpleNamespace(yeto_rl_learner_id=0, offload_train=True) - runtime._trainer_awake = False - runtime._rollout_offloaded = True - runtime._optimizer_reset_count = 0 - runtime._pause_rollout = no_op - runtime._resume_rollout = no_op - runtime._set_rollout_version = no_op - runtime.actor_model = SimpleNamespace( - onload=no_op, - offload=no_op, - update_weights=no_op, + async def remote(self, *args, **kwargs): + return self.function(*args, **kwargs) + + engine = SimpleNamespace( + update_weight_version=Remote(lambda version: versions.append(version)) ) - runtime.rollout_manager = SimpleNamespace( - onload_weights=Remote(), - onload_kv=Remote(), + rollout_manager = SimpleNamespace( + get_updatable_engines_and_lock=Remote( + lambda: SimpleNamespace(rollout_engines=[engine]) + ) ) - async def actor_call(*_args, **_kwargs): - return 2, miles.policy_hash(state_fn) + class Actor: + async def apply_trainable_state(self, value, *, reset_optimizer): + assert reset_optimizer + self.value = value + return 2 - runtime._actor_call = actor_call - runtime._append_event = events.append - asyncio.run(runtime._apply_global_policy(state_fn)) + async def export_trainable_state(self): + return self.value - assert events[0]["sync/global_policy_hash"] == miles.policy_hash(state_fn) - assert "policy_hash" not in events[0] - assert "trainer_ranks" not in events[0] + actor = Actor() + args = SimpleNamespace( + yeto_rl_base_model_revision=canonical.base_model_revision, + yeto_rl_lora_config_hash=canonical.lora_config_hash, + yeto_rl_layout_hash=canonical.layout_hash, + yeto_rl_event_tape=str(tmp_path / "events.jsonl"), + yeto_rl_learner_id=0, + ) + hook = MilesPolicySync(args) + hook.actor_model = actor + hook.rollout_manager = rollout_manager + asyncio.run(hook._apply_global_policy(canonical)) -@pytest.mark.parametrize( - "versions", - [["not-a-policy-token"], ["yeto:3", "yeto:4"]], -) -def test_miles_rejects_invalid_or_mixed_rollout_versions(monkeypatch, versions): - constants = types.ModuleType("sglang.srt.constants") - constants.GPU_MEMORY_TYPE_CUDA_GRAPH = "graph" - constants.GPU_MEMORY_TYPE_KV_CACHE = "kv" - constants.GPU_MEMORY_TYPE_WEIGHTS = "weights" - sglang = types.ModuleType("sglang") - srt = types.ModuleType("sglang.srt") - sglang.srt = srt - srt.constants = constants - monkeypatch.setitem(sys.modules, "sglang", sglang) - monkeypatch.setitem(sys.modules, "sglang.srt", srt) - monkeypatch.setitem(sys.modules, "sglang.srt.constants", constants) - - class Generate: - async def remote(self, _rollout_id): - return object() - - runtime = object.__new__(MilesIslandRuntime) - runtime._policy_version = 3 - runtime._rollout_id = 3 - runtime.rollout_manager = SimpleNamespace(generate=Generate()) - - async def no_op(): - pass + assert versions == ["yeto:1"] + event = json.loads((tmp_path / "events.jsonl").read_text()) + assert event["reset_parameter_count"] == 2 + assert event["sync/global_policy_hash"] - runtime._pause_rollout = no_op - runtime._rollout_batches = lambda _pack: [ - {"weight_versions": [versions]} + +def test_miles_policy_hook_builds_round_stats_without_revalidating_versions( + tmp_path, monkeypatch +): + checkpoint = tmp_path / "island.pt" + args = SimpleNamespace( + actor_num_gpus_per_node=2, + actor_num_nodes=1, + advantage_estimator="grpo", + yeto_rl_model="org/model", + yeto_rl_data="org/data", + yeto_rl_base_model_revision=MODEL_REVISION, + yeto_rl_data_revision="d" * 40, + expert_model_parallel_size=1, + yeto_rl_layout_hash="c" * 64, + lr=1e-4, + yeto_rl_lora_config_hash=LORA_CONFIG_HASH, + n_samples_per_prompt=2, + num_steps_per_rollout=1, + over_sampling_batch_size=2, + yeto_rl_reward_sha256="e" * 64, + rollout_batch_size=2, + seq_length=128, + seed=7, + rollout_max_response_len=16, + custom_generate_function_path=None, + use_session_server=False, + tito_model="default", + yeto_rl_completed_groups_path=str(checkpoint), + yeto_rl_learner_id=0, + ) + torch.save( + { + "schema_version": miles._ISLAND_CHECKPOINT_SCHEMA, + "config": miles._island_checkpoint_config(args), + "policy_version": 3, + "rollout_metrics": { + "active_groups": 2, + "cancelled_groups": 0, + "tool_wait_seconds": 1, + "group_p50_seconds": 2, + "group_p95_seconds": 3, + "group_p99_seconds": 4, + "rollout_seconds": 5, + }, + }, + checkpoint, + ) + batches = [ + { + "response_lengths": [1, 2], + "sample_indices": [0, 1], + "raw_reward": [0.0, 1.0, 2.0, 2.0], + }, + { + "response_lengths": [3, 4], + "sample_indices": [2, 3], + "raw_reward": [0.0, 1.0, 2.0, 2.0], + }, ] - with pytest.raises(StrictRlInvariantError) as failure: - asyncio.run(runtime._run_local_round(3, 1, 1, 1)) - assert failure.value.metric == "mixed_version_group_count" + monkeypatch.setitem( + sys.modules, + "ray", + SimpleNamespace(get=lambda reference: reference), + ) + hook = MilesPolicySync(args) + data_pack = { + "data_ref": [SimpleNamespace(inner=batch) for batch in batches] + } + + train_state = SimpleNamespace( + train_rollout_kl=0.1, + ess_ratio=0.8, + pg_clipfrac=0.25, + train_seconds=1.5, + ) + stats = hook._round_stats(3, data_pack, train_state) + + assert stats.action_tokens == 10 + assert stats.reward_mean == 1.25 + assert stats.zero_variance_group_ratio == 0.5 + assert stats.rollout_seconds == 5 + assert stats.mean_kl == 0.1 + assert stats.ess_ratio == 0.8 + assert stats.clip_fraction == 0.25 + assert stats.train_seconds == 1.5 diff --git a/tests/test_rl_integration.py b/tests/test_rl_integration.py index 949a4d7..566182a 100644 --- a/tests/test_rl_integration.py +++ b/tests/test_rl_integration.py @@ -2,20 +2,30 @@ from __future__ import annotations +import asyncio +import json import socket import subprocess +import sys import threading import time +import types from pathlib import Path +from types import SimpleNamespace import pytest import torch from yeto.export import parse_checkpoint from yeto.protocol import DTYPE_F32, SyncerClient -from yeto.rl.core import CanonicalTensorSpec, build_avg_layout from yeto.rl.bridge import BridgeConfig, StrictRlBridge -from yeto.rl.core import LocalRoundStats, canonical_state +from yeto.rl.core import ( + CanonicalTensorSpec, + LocalRoundStats, + build_avg_layout, + canonical_state, +) +from yeto.rl.miles import MilesPolicySync, _island_checkpoint_config from yeto.tensor_io import pack_tensor, unpack_fragment ROOT = Path(__file__).resolve().parent.parent @@ -336,6 +346,215 @@ def test_single_island_runs_the_real_syncer_parity_path(syncer_binary, tmp_path) process.wait() +def test_miles_public_hook_runs_against_real_syncer( + syncer_binary, tmp_path, monkeypatch +): + checkpoint_path = tmp_path / "state.ckpt" + island_checkpoint = tmp_path / "island.pt" + event_tape = tmp_path / "island.jsonl" + port = _port() + process = _start( + syncer_binary, + port, + checkpoint_path, + rounds=1, + learners=1, + ) + initial = _state( + 0, + {"base_model.model.layer.lora_A.weight": torch.zeros(1, 2)}, + ) + trainable_module = types.ModuleType( + "miles.backends.megatron_utils.trainable_state" + ) + def make_trainable_state( + version, + tensors, + *, + train_rollout_kl=None, + ess_ratio=None, + pg_clipfrac=None, + train_seconds=None, + ): + return SimpleNamespace( + policy_version=version, + layout_hash=initial.layout_hash, + tensors=tensors, + train_rollout_kl=train_rollout_kl, + ess_ratio=ess_ratio, + pg_clipfrac=pg_clipfrac, + train_seconds=train_seconds, + ) + + trainable_module.make_trainable_state = make_trainable_state + for name in ("miles", "miles.backends", "miles.backends.megatron_utils"): + monkeypatch.setitem(sys.modules, name, types.ModuleType(name)) + monkeypatch.setitem( + sys.modules, + "miles.backends.megatron_utils.trainable_state", + trainable_module, + ) + monkeypatch.setitem(sys.modules, "ray", SimpleNamespace(get=lambda value: value)) + + class Remote: + def __init__(self, function): + self.function = function + + async def remote(self, *args, **kwargs): + return self.function(*args, **kwargs) + + events = [] + engine = SimpleNamespace( + update_weight_version=Remote( + lambda version: events.append(("version", version)) + ) + ) + rollout_manager = SimpleNamespace( + get_updatable_engines_and_lock=Remote( + lambda: SimpleNamespace(rollout_engines=[engine]) + ) + ) + + class Actor: + def __init__(self): + self.state = SimpleNamespace( + policy_version=0, + layout_hash=initial.layout_hash, + tensors=initial.tensors, + ) + + async def export_trainable_state(self): + return self.state + + async def apply_trainable_state(self, state, *, reset_optimizer): + assert reset_optimizer + self.state = state + return len(state.tensors) + + actor = Actor() + args = SimpleNamespace( + actor_num_gpus_per_node=1, + actor_num_nodes=1, + advantage_estimator="grpo", + yeto_rl_model="org/model", + yeto_rl_data="org/data", + yeto_rl_base_model_revision=MODEL_REVISION, + yeto_rl_data_revision="d" * 40, + expert_model_parallel_size=1, + yeto_rl_layout_hash=initial.layout_hash, + lr=1e-4, + yeto_rl_lora_config_hash=LORA_CONFIG_HASH, + n_samples_per_prompt=2, + num_steps_per_rollout=1, + over_sampling_batch_size=1, + rollout_batch_size=1, + seq_length=128, + seed=7, + rollout_max_response_len=16, + custom_generate_function_path=None, + use_session_server=False, + tito_model=None, + yeto_rl_reward_sha256="e" * 64, + yeto_rl_completed_groups_path=str(island_checkpoint), + yeto_rl_event_tape=str(event_tape), + yeto_rl_learner_id=0, + num_rollout=1, + start_rollout_id=0, + ) + args.yeto_rl_bridge_config = BridgeConfig( + syncer_addr=("127.0.0.1", port), + learner_id=0, + global_rounds=1, + groups_per_round=1, + samples_per_group=2, + local_optimizer_steps=1, + wan_streams=0, + expected_specs=initial.specs, + base_model_revision=MODEL_REVISION, + lora_config_hash=LORA_CONFIG_HASH, + layout_hash=initial.layout_hash, + event_tape=str(event_tape), + ) + + async def run_hook(): + hook = MilesPolicySync(args) + await hook.initialize(actor_model=actor, rollout_manager=rollout_manager) + actor.state = trainable_module.make_trainable_state( + 0, + { + "base_model.model.layer.lora_A.weight": torch.tensor( + [[1.0, 3.0]] + ) + }, + train_rollout_kl=0.1, + ess_ratio=0.8, + pg_clipfrac=0.25, + train_seconds=1.5, + ) + torch.save( + { + "schema_version": 2, + "config": _island_checkpoint_config(args), + "policy_version": 0, + "rollout_metrics": { + "active_groups": 1, + "cancelled_groups": 0, + "tool_wait_seconds": 0, + "group_p50_seconds": 1, + "group_p95_seconds": 1, + "group_p99_seconds": 1, + "rollout_seconds": 1, + }, + }, + island_checkpoint, + ) + rollout_data = { + "data_ref": [ + SimpleNamespace( + inner={ + "weight_versions": [["yeto:0"], ["yeto:0"]], + "response_lengths": [2, 3], + "sample_indices": [0, 1], + "raw_reward": [0.0, 1.0], + } + ) + ] + } + await hook.after_local_train( + rollout_id=0, + actor_model=actor, + rollout_data=rollout_data, + ) + events.append("miles_weight_publish") + await hook.finalize() + + try: + asyncio.run(run_hook()) + assert process.wait(timeout=10) == 0 + assert events == [ + ("version", "yeto:0"), + ("version", "yeto:1"), + "miles_weight_publish", + ] + assert torch.equal( + next(iter(actor.state.tensors.values())), + torch.tensor([[1.0, 3.0]]), + ) + round_event = next( + event + for event in map(json.loads, event_tape.read_text().splitlines()) + if event.get("event") == "rl_local_round" + ) + assert round_event["rl/current_vs_rollout_kl"] == 0.1 + assert round_event["rl/ess_ratio"] == 0.8 + assert round_event["rl/clip_fraction"] == 0.25 + assert round_event["train_seconds"] == 1.5 + finally: + if process.poll() is None: + process.kill() + process.wait() + + def test_fixed_roster_exact_base_duplicate_disconnect_and_manual_average( syncer_binary, tmp_path ): diff --git a/tests/test_rl_launcher.py b/tests/test_rl_launcher.py index 55745b2..66d5661 100644 --- a/tests/test_rl_launcher.py +++ b/tests/test_rl_launcher.py @@ -541,6 +541,11 @@ def test_miles_argv_uses_provider_capabilities_without_model_family_branches(): argv[argv.index("--rollout-all-samples-process-path") + 1] == "yeto.rl.miles.queue_completed_groups" ) + assert ( + argv[argv.index("--external-policy-sync-path") + 1] + == "yeto.rl.miles.create_policy_sync" + ) + assert "--custom-megatron-init-path" not in argv assert "--use-distributed-optimizer" not in argv assert "--no-offload-train" in argv assert argv[argv.index("--sglang-mem-fraction-static") + 1] == "0.4" diff --git a/yeto/rl/__init__.py b/yeto/rl/__init__.py index e7e86aa..61cc6c9 100644 --- a/yeto/rl/__init__.py +++ b/yeto/rl/__init__.py @@ -1,7 +1,7 @@ """Pinned Miles reinforcement-learning integration.""" -MILES_REPOSITORY = "https://github.com/radixark/miles" -MILES_COMMIT = "dfc66ff38752bfa2c5d325e0037ebc4b537c06de" +MILES_REPOSITORY = "https://github.com/agentenv/miles" +MILES_COMMIT = "a91bd34e50416aeb1da111f74d52b296e8216b96" MILES_PEFT_VERSION = "0.20.0" MILES_IMAGE = ( "docker:radixark/miles@sha256:" diff --git a/yeto/rl/bridge.py b/yeto/rl/bridge.py index 1f83556..21d7e9a 100644 --- a/yeto/rl/bridge.py +++ b/yeto/rl/bridge.py @@ -106,15 +106,7 @@ def __init__(self, runtime: IslandRuntime, config: BridgeConfig) -> None: def run(self) -> CanonicalLoraState: try: - self.client.start() - if self.config.learner_id == 0: - self.client.send_init( - 0, - pack_tensor( - flat_tensor(self.initial.tensors, self.specs), - DTYPE_F32, - ), - ) + self.start() while True: self.client.check_health() if self.client.finalizing.is_set(): @@ -147,6 +139,46 @@ def run(self) -> CanonicalLoraState: finally: self.client.close() + def start(self) -> None: + self.client.start() + if self.config.learner_id == 0: + self.client.send_init( + 0, + pack_tensor( + flat_tensor(self.initial.tensors, self.specs), + DTYPE_F32, + ), + ) + + def wait_for_global_policy(self, version: int) -> CanonicalLoraState: + while self.current is None or self.current.policy_version < version: + self.client.check_health() + self._drain_messages() + if self.current is None or self.current.policy_version < version: + time.sleep(0.05) + if self.current.policy_version != version: + raise RuntimeError( + f"RL policy jumped past expected version {version}" + ) + return self.current + + def wait_for_initial_policy(self) -> CanonicalLoraState: + while self.current is None: + self.client.check_health() + self._drain_messages() + if self.current is None: + time.sleep(0.05) + return self.current + + def wait_for_round(self) -> PullRequest: + while True: + self.client.check_health() + self._drain_messages() + permit = self._ready_permit() + if permit is not None: + return permit + time.sleep(0.05) + def _drain_messages(self) -> bool: progressed = False for update in self.client.drain_updates(): @@ -232,14 +264,36 @@ def _run_round(self, permit: PullRequest) -> None: ): raise RuntimeError("Miles returned LocalRoundStats for a different round") + self.submit_local_state( + permit, + base, + self.runtime.export_local_policy(), + stats, + ) + + def submit_local_state( + self, + permit: PullRequest, + base: CanonicalLoraState, + local: CanonicalLoraState, + stats: LocalRoundStats, + ) -> LocalRoundStats: + if self.current is None or base.policy_version != self.current.policy_version: + raise RuntimeError("RL attempted to submit without its exact base") + if ( + permit.global_step != base.policy_version + 1 + or stats.island_id != self.config.learner_id + or stats.base_policy_version != base.policy_version + or stats.local_round_id != permit.global_step + ): + raise RuntimeError("Miles returned LocalRoundStats for a different round") try: - exported = self.runtime.export_local_policy() local = canonical_state( - exported.policy_version, - exported.tensors, - base_model_revision=exported.base_model_revision, - lora_config_hash=exported.lora_config_hash, - layout_hash=exported.layout_hash, + local.policy_version, + local.tensors, + base_model_revision=local.base_model_revision, + lora_config_hash=local.lora_config_hash, + layout_hash=local.layout_hash, expected_specs=self.specs, ) delta = policy_delta(local, base) @@ -296,6 +350,15 @@ def _run_round(self, permit: PullRequest) -> None: payload, ) self.pushed_step = permit.global_step + return stats + + def finalize(self) -> CanonicalLoraState: + while not self.client.finalizing.is_set(): + self.client.check_health() + self._drain_messages() + if not self.client.finalizing.is_set(): + time.sleep(0.05) + return self._finalize() def _finalize(self) -> CanonicalLoraState: manifest, fragments = self.client.wait_for_final_fragments() diff --git a/yeto/rl/core.py b/yeto/rl/core.py index 3686bca..f3c8ac5 100644 --- a/yeto/rl/core.py +++ b/yeto/rl/core.py @@ -2,10 +2,10 @@ from __future__ import annotations -import math -import re import hashlib import json +import math +import re from collections.abc import Mapping, Sequence from dataclasses import dataclass @@ -257,7 +257,7 @@ class LocalRoundStats: zero_variance_group_ratio: float mean_kl: float | None ess_ratio: float | None - clip_fraction: float + clip_fraction: float | None delta_l2_norm: float rollout_seconds: float train_seconds: float @@ -283,14 +283,13 @@ def __post_init__(self) -> None: "group_p95_seconds", "group_p99_seconds", "zero_variance_group_ratio", - "clip_fraction", "delta_l2_norm", "rollout_seconds", "train_seconds", ): if not math.isfinite(getattr(self, name)): raise ValueError(f"{name} must be finite") - for name in ("mean_kl", "ess_ratio"): + for name in ("mean_kl", "ess_ratio", "clip_fraction"): value = getattr(self, name) if value is not None and not math.isfinite(value): raise ValueError(f"{name} must be finite when present") diff --git a/yeto/rl/learner.py b/yeto/rl/learner.py index ce2386a..fdec189 100644 --- a/yeto/rl/learner.py +++ b/yeto/rl/learner.py @@ -3,15 +3,16 @@ from __future__ import annotations import argparse +import asyncio import json import os import sys from pathlib import Path from typing import Any -from .bridge import BridgeConfig, StrictRlBridge +from .bridge import BridgeConfig from .export import adapter_targets, derive_peft_lora_specs -from .miles import MilesIslandRuntime, verify_miles_revision +from .miles import verify_miles_revision def parse_args(argv=None): @@ -206,7 +207,7 @@ def build_miles_argv( "--custom-rm-path", _miles_callable(args.reward_function), "--advantage-estimator", "grpo", "--lr", str(args.inner_lr), - "--custom-megatron-init-path", "yeto.rl.miles.configure_miles_bridge", + "--external-policy-sync-path", "yeto.rl.miles.create_policy_sync", "--accumulate-allreduce-grads-in-fp32", "--attention-softmax-in-fp32", "--attention-backend", "unfused", @@ -440,26 +441,24 @@ def main(argv=None) -> None: miles_args.yeto_rl_event_tape = args.event_tape miles_args.yeto_rl_learner_id = args.learner_id - runtime = MilesIslandRuntime(miles_args) - bridge = StrictRlBridge( - runtime, - BridgeConfig( - syncer_addr=_syncer_address(args.syncer), - learner_id=args.learner_id, - global_rounds=args.global_rounds, - groups_per_round=args.groups_per_round, - samples_per_group=args.samples_per_group, - local_optimizer_steps=args.optimizer_steps, - wan_streams=args.wan_streams, - expected_specs=specs, - base_model_revision=args.model_revision, - lora_config_hash=lora_config_hash, - layout_hash=layout_hash, - event_tape=args.event_tape, - ), + miles_args.yeto_rl_bridge_config = BridgeConfig( + syncer_addr=_syncer_address(args.syncer), + learner_id=args.learner_id, + global_rounds=args.global_rounds, + groups_per_round=args.groups_per_round, + samples_per_group=args.samples_per_group, + local_optimizer_steps=args.optimizer_steps, + wan_streams=args.wan_streams, + expected_specs=specs, + base_model_revision=args.model_revision, + lora_config_hash=lora_config_hash, + layout_hash=layout_hash, + event_tape=args.event_tape, ) - final = bridge.run() - print(f"[rl] learner {args.learner_id} finalized policy v{final.policy_version}") + from train import train as miles_train + + asyncio.run(miles_train(miles_args)) + print(f"[rl] learner {args.learner_id} finalized") if __name__ == "__main__": diff --git a/yeto/rl/miles.py b/yeto/rl/miles.py index a9130bb..8790315 100644 --- a/yeto/rl/miles.py +++ b/yeto/rl/miles.py @@ -8,6 +8,7 @@ import os import statistics import subprocess +import sys import time from collections.abc import Mapping from dataclasses import asdict @@ -25,8 +26,6 @@ policy_hash, ) -_CANONICAL_PREFIX = "base_model.model." - def verify_miles_revision(root: str | Path) -> Path: """Verify repository, commit, tracked files, and the imported package.""" @@ -67,305 +66,6 @@ def git(*args: str) -> str: return root -def _adapter_sides(actor) -> list[tuple[str, Any]]: - """Return PEFT names paired with Bridge's actual conversion tasks.""" - - from megatron.bridge import AutoBridge - - bridge = AutoBridge.from_hf_pretrained( - actor.args.hf_checkpoint, - trust_remote_code=bool(actor.args.yeto_rl_trust_remote_code), - ) - model_bridge = getattr(bridge, "_model_bridge", None) - build_tasks = getattr(model_bridge, "build_adapter_conversion_tasks", None) - if build_tasks is None: - raise RuntimeError("pinned Megatron-Bridge lacks adapter conversion tasks") - tasks_by_base = build_tasks(actor.model) - sides: list[tuple[str, Any]] = [] - for base_name in sorted(tasks_by_base): - tasks = sorted( - tasks_by_base[base_name], - key=lambda task: task.adapter_key or "", - ) - for task in tasks: - for side in (task.linear_in_task, task.linear_out_task): - parameter = side.param_weight - main = getattr(parameter, "main_param", None) - if ( - main is None - or main.dtype != torch.float32 - or main.numel() != parameter.numel() - ): - raise RuntimeError( - f"LoRA parameter {side.param_name!r} has no complete " - "FP32 optimizer master" - ) - converted = side.mapping.megatron_to_hf( - main.view(parameter.shape), - side.megatron_module, - ) - if len(converted) != 1: - raise RuntimeError( - f"ambiguous LoRA mapping for {side.param_name!r}" - ) - raw_name = next(iter(converted)) - name = ( - raw_name - if raw_name.startswith(_CANONICAL_PREFIX) - else _CANONICAL_PREFIX + raw_name - ) - if not name.endswith((".lora_A.weight", ".lora_B.weight")): - raise RuntimeError(f"non-PEFT LoRA mapping {name!r}") - sides.append((name, side)) - names = [name for name, _ in sides] - if not names or len(names) != len(set(names)): - raise RuntimeError("Miles produced an empty or duplicate LoRA mapping") - mapped = {id(side.param_weight) for _, side in sides} - trainable = { - id(parameter) - for chunk in actor.model - for parameter in chunk.parameters() - if parameter.requires_grad - } - if mapped != trainable: - raise RuntimeError( - "Miles adapter conversion does not cover every trainable parameter" - ) - return sorted(sides) - - -@torch.no_grad() -def _export_fp32_policy(actor) -> dict[str, torch.Tensor]: - tensors = {} - for name, side in _adapter_sides(actor): - parameter = side.param_weight - converted = side.mapping.megatron_to_hf( - parameter.main_param.view(parameter.shape), - side.megatron_module, - ) - value = next(iter(converted.values())) - tensors[name] = value.detach().to( - device="cpu", dtype=torch.float32 - ).contiguous().clone() - return tensors - - -def _optimizer_children(optimizer) -> list[Any]: - return list(getattr(optimizer, "chained_optimizers", (optimizer,))) - - -def _reset_optimizer_state(actor, parameters: list[torch.Tensor]) -> int: - parameter_ids = {id(parameter) for parameter in parameters} - for child in _optimizer_children(actor.optimizer): - optimizer = getattr(child, "optimizer", child) - for parameter in list(optimizer.state): - if id(parameter) in parameter_ids: - optimizer.state.pop(parameter, None) - return len(parameter_ids) - - -@torch.no_grad() -def _copy_masters_to_model(actor) -> None: - for child in _optimizer_children(actor.optimizer): - copy = getattr(child, "_copy_main_params_to_model_params", None) - if copy is None: - raise RuntimeError("pinned Megatron optimizer lacks main-to-model copy") - copy() - - -def _restore_scheduler_progress(actor, policy_version: int) -> None: - scheduler = actor.opt_param_scheduler - batch_size = actor.args.global_batch_size - target = policy_version * actor.args.num_steps_per_rollout * batch_size - if scheduler is None or batch_size <= 0 or scheduler.num_steps % batch_size: - raise RuntimeError("Miles scheduler progress is not an integral optimizer step") - if scheduler.num_steps > target: - raise RuntimeError("Miles scheduler is ahead of the committed policy") - if scheduler.num_steps < target: - scheduler.step(increment=target - scheduler.num_steps) - - -def _actor_export_policy(self): - if torch.distributed.is_initialized() and torch.distributed.get_rank() != 0: - return None - return _export_fp32_policy(self) - - -@torch.no_grad() -def _actor_apply_policy( - self, - tensors: Mapping[str, torch.Tensor], - policy_version: int, -): - sides = dict(_adapter_sides(self)) - if set(tensors) != set(sides): - missing = sorted(set(sides) - set(tensors)) - extra = sorted(set(tensors) - set(sides)) - raise RuntimeError( - f"global LoRA mapping mismatch: missing={missing}, extra={extra}" - ) - - mapped = {} - for name, side in sides.items(): - value = tensors[name].detach().to( - device=side.param_weight.device, - dtype=torch.float32, - ) - target = side.mapping.hf_to_megatron(value, side.megatron_module) - if target.numel() != side.param_weight.numel(): - raise RuntimeError(f"global LoRA shape mismatch for {name!r}") - mapped[name] = target.reshape(side.param_weight.shape).contiguous() - - _restore_scheduler_progress(self, policy_version) - reset_parameter_count = _reset_optimizer_state( - self, - [side.param_weight.main_param for side in sides.values()], - ) - for name, side in sides.items(): - side.param_weight.main_param.view(side.param_weight.shape).copy_(mapped[name]) - _copy_masters_to_model(self) - if torch.distributed.is_initialized(): - torch.distributed.barrier() - self.weights_backuper.backup("actor") - torch.cuda.empty_cache() - - identity = { - "base_model_revision": self.args.yeto_rl_base_model_revision, - "lora_config_hash": self.args.yeto_rl_lora_config_hash, - "layout_hash": self.args.yeto_rl_layout_hash, - } - applied = canonical_state(policy_version, _export_fp32_policy(self), **identity) - canonical_state( - policy_version, - tensors, - expected_specs=applied.specs, - **identity, - ) - return reset_parameter_count, policy_hash(applied) - - -def _actor_optimizer_steps(self) -> int: - scheduler = self.opt_param_scheduler - batch = self.args.global_batch_size - if scheduler is None or batch <= 0 or scheduler.num_steps % batch: - raise RuntimeError("Miles optimizer step counter is not integral") - return int(scheduler.num_steps // batch) - - -def _actor_train_metrics(self): - if torch.distributed.is_initialized() and torch.distributed.get_rank() != 0: - return None - metrics = getattr(self.args, "_yeto_rl_train_metrics", None) - if hasattr(self.args, "_yeto_rl_train_metrics"): - del self.args._yeto_rl_train_metrics - return metrics - - -def _install_train_metric_capture() -> None: - from miles.backends.megatron_utils import model - - original = model.log_train_step - if getattr(original, "_yeto_rl_capture", False): - return - - def log_train_step(*values, **kwargs): - metrics = original(*values, **kwargs) - if not torch.distributed.is_initialized() or torch.distributed.get_rank() == 0: - args = kwargs["args"] - args._yeto_rl_train_metrics = { - key: float(metrics[key]) - for key in ( - "train/train_rollout_kl", - "train/ess_ratio", - "train/pg_clipfrac", - ) - if key in metrics - } - return metrics - - log_train_step._yeto_rl_capture = True - model.log_train_step = log_train_step - - -def _install_colocated_lora_ipc_sync() -> None: - """Keep CUDA IPC producers alive until SGLang finishes the transfer.""" - - import ray - from miles.backends.megatron_utils.update_weight import ( - update_weight_from_tensor, - ) - from miles.backends.megatron_utils.update_weight.common import ( - _check_weight_sync_results, - ) - - original = update_weight_from_tensor._send_to_colocated_engine - if getattr(original, "_yeto_rl_synchronized", False): - return - - def synchronized_send(*values, **kwargs): - refs, long_lived_tensors = original(*values, **kwargs) - if kwargs.get("lora_config") is not None: - results = ray.get(refs) - _check_weight_sync_results(results, is_lora=True) - group = kwargs.get("ipc_gather_group") - if group is not None: - torch.distributed.barrier(group=group) - return refs, long_lived_tensors - - synchronized_send._yeto_rl_synchronized = True - update_weight_from_tensor._send_to_colocated_engine = synchronized_send - - -def configure_miles_bridge(args) -> None: - """Install Yeto actor methods inside each Miles Ray worker.""" - - from megatron.bridge import AutoBridge - from megatron.bridge.training import config as bridge_config - - original = AutoBridge.to_megatron_provider - - def configured_provider(self, *values, **kwargs): - provider = original(self, *values, **kwargs) - provider.attention_backend = args.attention_backend - return provider - - AutoBridge.to_megatron_provider = configured_provider - # The INIT adapter is replicated. Keep complete fp32 masters and Adam - # state on every DP/EP rank; no sharded gather path is part of v0. - args.use_distributed_optimizer = False - original_ddp_config = bridge_config.DistributedDataParallelConfig - - def replicated_lora_ddp_config(*values, **kwargs): - kwargs["use_distributed_optimizer"] = False - return original_ddp_config(*values, **kwargs) - - bridge_config.DistributedDataParallelConfig = replicated_lora_ddp_config - _install_train_metric_capture() - _install_colocated_lora_ipc_sync() - install_miles_actor_adapter() - - -def install_miles_actor_adapter() -> None: - """Install methods in both the driver and each Miles Ray worker.""" - - from miles.backends.megatron_utils.actor import MegatronTrainRayActor - - methods = { - "yeto_rl_export_policy": _actor_export_policy, - "yeto_rl_apply_policy": _actor_apply_policy, - "yeto_rl_optimizer_steps": _actor_optimizer_steps, - "yeto_rl_train_metrics": _actor_train_metrics, - } - for name, method in methods.items(): - existing = getattr(MegatronTrainRayActor, name, None) - if existing is not None and ( - getattr(existing, "__module__", None), - getattr(existing, "__qualname__", None), - ) != (method.__module__, method.__qualname__): - raise RuntimeError(f"Miles actor already defines incompatible {name}") - setattr(MegatronTrainRayActor, name, method) - - def _policy_token(version: int) -> str: return f"yeto:{version}" @@ -402,6 +102,52 @@ def _validate_rollout_groups(data: object, groups: int, samples: int) -> None: raise RuntimeError("Miles returned an incomplete trajectory") +def _validate_rollout_policy_versions(data: list[list[Any]], rollout_id: int) -> None: + expected = _policy_token(rollout_id) + for group in data: + for sample in group: + versions = getattr(sample, "weight_versions", None) + if not isinstance(versions, list) or not versions or any( + not isinstance(version, str) or version != expected + for version in versions + ): + raise StrictRlInvariantError( + "mixed_version_group_count", + f"Miles rollout did not use only policy {expected}", + ) + + +def _append_rl_event(args, event: dict[str, Any]) -> None: + path = Path(args.yeto_rl_event_tape).expanduser() + path.parent.mkdir(parents=True, exist_ok=True) + event = { + "island_id": int(args.yeto_rl_learner_id), + "time_unix": time.time(), + **event, + } + with path.open("a", encoding="utf-8") as handle: + handle.write(json.dumps(event, sort_keys=True, separators=(",", ":")) + "\n") + + +def _record_strict_failure(args, error: StrictRlInvariantError, bridge=None) -> None: + _append_rl_event( + args, + { + "event": "rl_strict_failure", + "metric": error.metric, + "value": 1, + "error": f"{type(error).__name__}: {error}", + }, + ) + print( + f"[yeto-rl-strict-failure] {error.metric}: {error}", + file=sys.stderr, + flush=True, + ) + if bridge is not None: + bridge.client.close() + + _ISLAND_CHECKPOINT_SCHEMA = 2 @@ -662,6 +408,7 @@ def generate_rollout(args, rollout_id: int, data_source, evaluation: bool = Fals group, rollout_id, args.n_samples_per_prompt ) ] + rollout_started = time.monotonic() output, lifecycle = _run_rollout_with_metrics( miles_generate, args, @@ -677,6 +424,11 @@ def generate_rollout(args, rollout_id: int, data_source, evaluation: bool = Fals args.rollout_batch_size, args.n_samples_per_prompt, ) + try: + _validate_rollout_policy_versions(output.samples, rollout_id) + except StrictRlInvariantError as error: + _record_strict_failure(args, error) + raise consumed = {_group_indices(group) for group in output.samples} data_source.buffer[:] = [ group @@ -689,6 +441,7 @@ def generate_rollout(args, rollout_id: int, data_source, evaluation: bool = Fals for sample in samples ) round_metrics = { + "rollout_seconds": time.monotonic() - rollout_started, "active_groups": lifecycle["peak_active"], "cancelled_groups": lifecycle["cancelled"], "tool_wait_seconds": tool_wait_seconds, @@ -706,100 +459,64 @@ def generate_rollout(args, rollout_id: int, data_source, evaluation: bool = Fals return output -class MilesIslandRuntime: - """Synchronous wrapper over the pinned Miles Ray APIs.""" +class _BridgeRuntime: + def __init__(self, initial: CanonicalLoraState, args) -> None: + self.initial = initial + self.args = args + + def initialize(self) -> CanonicalLoraState: + return self.initial + + def apply_global_policy(self, _state: CanonicalLoraState) -> None: + pass + + def record_local_round(self, stats: LocalRoundStats) -> None: + path = Path(self.args.yeto_rl_completed_groups_path).expanduser() + try: + payload = torch.load(path, map_location="cpu", weights_only=True) + except Exception as error: + raise RuntimeError("cannot update Miles island checkpoint") from error + if ( + not isinstance(payload, dict) + or payload.get("schema_version") != _ISLAND_CHECKPOINT_SCHEMA + or payload.get("policy_version") != stats.base_policy_version + or payload.get("config") != _island_checkpoint_config(self.args) + ): + raise RuntimeError("Miles island checkpoint changed before round commit") + payload["local_round_id"] = stats.local_round_id + payload["local_round_stats"] = asdict(stats) + _atomic_save_island_checkpoint(path, payload) + + def shutdown(self) -> None: + pass + + +class MilesPolicySync: + """Yeto synchronization hook called by the Miles training loop.""" def __init__(self, args) -> None: self.args = args - self.loop = asyncio.new_event_loop() - self.rollout_manager = None self.actor_model = None - self._owns_ray = False - self._trainer_awake = False - self._rollout_offloaded = True - self._policy_version: int | None = None - self._rollout_id = 0 - self._optimizer_reset_count = 0 - - async def _onload_trainer(self) -> None: - if self._trainer_awake: - return - await self.actor_model.onload() - self._trainer_awake = True - - async def _offload_trainer(self) -> None: - if self._trainer_awake and self.args.offload_train: - await self.actor_model.offload() - self._trainer_awake = False - - def _run(self, coroutine): - return self.loop.run_until_complete(coroutine) - - async def _actor_call(self, method: str, *args, rank0: bool = False): - results = await self.actor_model._broadcast(method, *args) - expected = self.args.actor_num_nodes * self.args.actor_num_gpus_per_node - if not isinstance(results, list) or len(results) != expected: - raise RuntimeError( - f"Miles returned {len(results) if isinstance(results, list) else 0} " - f"actor results, expected {expected}" - ) - if rank0: - exported = [result for result in results if result is not None] - if len(exported) != 1: - raise RuntimeError("Miles must export policy only on global rank 0") - return exported[0] - if not results or any(result != results[0] for result in results[1:]): - raise RuntimeError(f"Miles actor ranks disagree on {method}") - return results[0] - - async def _initialize(self) -> CanonicalLoraState: - import ray - from miles.ray.placement_group import ( - create_placement_groups, - create_rollout_manager, - create_training_models, - ) + self.rollout_manager = None + self.bridge = None + self.current = None + self.permit = None + self.optimizer_reset_count = 0 - install_miles_actor_adapter() - if not ray.is_initialized(): - ray.init(address="auto") - self._owns_ray = True - expected_gpus = ( - self.args.actor_num_nodes * self.args.actor_num_gpus_per_node - ) - deadline = time.monotonic() + 300 - while True: - visible_gpus = int(ray.cluster_resources().get("GPU", 0)) - if visible_gpus >= expected_gpus or time.monotonic() >= deadline: - break - await asyncio.sleep(2) - if visible_gpus != expected_gpus: - raise RuntimeError( - f"Miles Ray cluster has {visible_gpus} GPUs, expected {expected_gpus}" - ) - groups = create_placement_groups(self.args) - self.rollout_manager, _ = create_rollout_manager( - self.args, groups["rollout"] - ) - self.actor_model, critic = await create_training_models( - self.args, groups, self.rollout_manager - ) - if critic is not None: - raise RuntimeError("RL v0 does not support a Miles critic") - self._trainer_awake = not self.args.offload_train - await self._onload_trainer() - tensors = await self._actor_call("yeto_rl_export_policy", rank0=True) - await self._offload_trainer() - return canonical_state( - 0, - tensors, + def _canonical_state(self, state) -> CanonicalLoraState: + canonical = canonical_state( + state.policy_version, + state.tensors, base_model_revision=self.args.yeto_rl_base_model_revision, lora_config_hash=self.args.yeto_rl_lora_config_hash, layout_hash=self.args.yeto_rl_layout_hash, ) - - def initialize(self) -> CanonicalLoraState: - return self._run(self._initialize()) + if state.layout_hash != canonical.layout_hash: + raise StrictRlInvariantError( + "layout_hash_mismatch", + "Miles and Yeto computed different LoRA layouts", + ) + return canonical async def _engines(self) -> list[Any]: info = await self.rollout_manager.get_updatable_engines_and_lock.remote() @@ -808,111 +525,68 @@ async def _engines(self) -> list[Any]: raise RuntimeError("Miles created no updatable SGLang engine") return engines - async def _pause_rollout(self) -> None: - engines = await self._engines() - # Retracted requests resume after a weight update; abort them so one - # trajectory can never cross the global policy boundary. - await asyncio.gather( - *(engine.pause_generation.remote("abort") for engine in engines) - ) - - async def _resume_rollout(self) -> None: - engines = await self._engines() - await asyncio.gather( - *(engine.continue_generation.remote() for engine in engines) - ) - async def _set_rollout_version(self, version: int) -> None: token = _policy_token(version) - engines = await self._engines() await asyncio.gather( - *(engine.update_weight_version.remote(token) for engine in engines) + *( + engine.update_weight_version.remote(token) + for engine in await self._engines() + ) ) async def _apply_global_policy(self, state: CanonicalLoraState) -> None: - from sglang.srt.constants import ( - GPU_MEMORY_TYPE_CUDA_GRAPH, - GPU_MEMORY_TYPE_KV_CACHE, - GPU_MEMORY_TYPE_WEIGHTS, - ) + from miles.backends.megatron_utils.trainable_state import make_trainable_state - await self._pause_rollout() - if not self._rollout_offloaded: - await self.rollout_manager.offload.remote( - tags=[ - GPU_MEMORY_TYPE_CUDA_GRAPH, - GPU_MEMORY_TYPE_KV_CACHE, - GPU_MEMORY_TYPE_WEIGHTS, - ] - ) - self._rollout_offloaded = True - await self._onload_trainer() - reset_parameter_count, applied_hash = await self._actor_call( - "yeto_rl_apply_policy", - dict(state.tensors), - state.policy_version, + reset_count = await self.actor_model.apply_trainable_state( + make_trainable_state(state.policy_version, state.tensors), + reset_optimizer=True, ) - expected_hash = policy_hash(state) - if applied_hash != expected_hash: + applied = self._canonical_state(await self.actor_model.export_trainable_state()) + applied_hash = policy_hash(applied) + if applied_hash != policy_hash(state): raise StrictRlInvariantError( "policy_hash_mismatch_after_apply", "policy hash mismatch after trainer apply", ) - await self._offload_trainer() - await self.rollout_manager.onload_weights.remote() - await self.actor_model.update_weights() - # update_weights resumes generation; close the admission boundary - # until KV/weights and the explicit version are all installed. - await self._pause_rollout() - await self.rollout_manager.onload_kv.remote() - self._rollout_offloaded = False await self._set_rollout_version(state.policy_version) - self._policy_version = state.policy_version - self._rollout_id = state.policy_version - await self._resume_rollout() - self._optimizer_reset_count += 1 + self.optimizer_reset_count += 1 self._append_event( { "event": "rl_policy_apply", "policy_version": state.policy_version, - "optimizer_reset_count": self._optimizer_reset_count, - "reset_parameter_count": reset_parameter_count, + "optimizer_reset_count": self.optimizer_reset_count, + "reset_parameter_count": reset_count, "rl/global_policy_version": state.policy_version, - "rl/optimizer_reset_count": self._optimizer_reset_count, + "rl/optimizer_reset_count": self.optimizer_reset_count, "sync/global_policy_hash": applied_hash, } ) - def apply_global_policy(self, state: CanonicalLoraState) -> None: - self._run(self._apply_global_policy(state)) - def _rollout_batches(self, data_pack) -> list[Mapping[str, Any]]: import ray references = data_pack.get("data_ref") if not isinstance(references, list): raise RuntimeError("Miles returned an invalid rollout shard list") - expected_dp = ( - self.args.actor_num_nodes * self.args.actor_num_gpus_per_node - ) - if len(references) != expected_dp: + expected = self.args.actor_num_nodes * self.args.actor_num_gpus_per_node + if len(references) != expected: raise RuntimeError( - f"Miles returned {len(references)} DP shards, expected {expected_dp}" + f"Miles returned {len(references)} DP shards, expected {expected}" ) return [ray.get(reference.inner) for reference in references] def _rollout_metrics(self, policy_version: int) -> dict[str, float]: path = Path(self.args.yeto_rl_completed_groups_path).expanduser() + payload = torch.load(path, map_location="cpu", weights_only=True) + if ( + not isinstance(payload, dict) + or payload.get("schema_version") != _ISLAND_CHECKPOINT_SCHEMA + or payload.get("policy_version") != policy_version + or payload.get("config") != _island_checkpoint_config(self.args) + or not isinstance(payload.get("rollout_metrics"), Mapping) + ): + raise RuntimeError("Miles island checkpoint lacks rollout metrics") try: - payload = torch.load(path, map_location="cpu", weights_only=True) - if ( - not isinstance(payload, dict) - or payload.get("schema_version") != _ISLAND_CHECKPOINT_SCHEMA - or payload.get("policy_version") != policy_version - or payload.get("config") != _island_checkpoint_config(self.args) - or not isinstance(payload.get("rollout_metrics"), Mapping) - ): - raise RuntimeError("Miles island checkpoint lacks rollout metrics") return { name: float(value) for name, value in payload["rollout_metrics"].items() @@ -920,87 +594,9 @@ def _rollout_metrics(self, policy_version: int) -> dict[str, float]: except (TypeError, ValueError) as error: raise RuntimeError("Miles returned invalid Yeto group metrics") from error - async def _run_local_round( - self, - expected_policy_version: int, - groups: int, - samples_per_group: int, - optimizer_steps: int, - ) -> LocalRoundStats: - from sglang.srt.constants import ( - GPU_MEMORY_TYPE_CUDA_GRAPH, - GPU_MEMORY_TYPE_KV_CACHE, - GPU_MEMORY_TYPE_WEIGHTS, - ) - - if expected_policy_version != self._policy_version: - raise RuntimeError("Miles round requested from an unapplied global policy") - rollout_started = time.monotonic() - data_pack = await self.rollout_manager.generate.remote(self._rollout_id) - rollout_seconds = time.monotonic() - rollout_started - await self._pause_rollout() + def _round_stats(self, rollout_id: int, data_pack, train_state) -> LocalRoundStats: batches = self._rollout_batches(data_pack) - versions = [ - sample_versions - for batch in batches - for sample_versions in batch.get("weight_versions", []) - ] - expected_samples = groups * samples_per_group - if not isinstance(versions, list) or len(versions) != expected_samples: - raise RuntimeError( - f"Miles produced {len(versions) if isinstance(versions, list) else 0} " - f"versioned samples, expected {expected_samples}" - ) - try: - observed = { - _version_from_token(token) - for sample_versions in versions - for token in sample_versions - } - except RuntimeError as error: - raise StrictRlInvariantError( - "mixed_version_group_count", - str(error), - ) from error - if any(not sample_versions for sample_versions in versions) or observed != { - expected_policy_version - }: - raise StrictRlInvariantError( - "mixed_version_group_count", - f"Miles rollout mixed policy versions: {observed}", - ) - rollout_metrics = self._rollout_metrics(expected_policy_version) - - await self.rollout_manager.offload.remote( - tags=[ - GPU_MEMORY_TYPE_CUDA_GRAPH, - GPU_MEMORY_TYPE_KV_CACHE, - GPU_MEMORY_TYPE_WEIGHTS, - ] - ) - self._rollout_offloaded = True - before = await self._actor_call("yeto_rl_optimizer_steps") - train_started = time.monotonic() - await self.actor_model.train(self._rollout_id, data_pack) - train_seconds = time.monotonic() - train_started - self._trainer_awake = True - after = await self._actor_call("yeto_rl_optimizer_steps") - if after - before != optimizer_steps: - raise RuntimeError( - f"Miles performed {after - before} optimizer steps, " - f"expected {optimizer_steps}" - ) - train_metrics = await self._actor_call( - "yeto_rl_train_metrics", - rank0=True, - ) - try: - mean_kl = float(train_metrics["train/train_rollout_kl"]) - ess_ratio = float(train_metrics["train/ess_ratio"]) - clip_fraction = float(train_metrics["train/pg_clipfrac"]) - except (KeyError, TypeError, ValueError) as error: - raise RuntimeError("Miles did not return required GRPO train metrics") from error - self._rollout_id += 1 + expected_samples = self.args.rollout_batch_size * self.args.n_samples_per_prompt response_lengths = [ int(value) for batch in batches @@ -1026,111 +622,133 @@ async def _run_local_round( rewards = [float(value) for value in raw_rewards] except (TypeError, ValueError) as error: raise RuntimeError("Miles RL v0 requires scalar rewards") from error + try: + mean_kl = ( + None + if train_state.train_rollout_kl is None + else float(train_state.train_rollout_kl) + ) + ess_ratio = ( + None if train_state.ess_ratio is None else float(train_state.ess_ratio) + ) + clip_fraction = ( + None + if train_state.pg_clipfrac is None + else float(train_state.pg_clipfrac) + ) + train_seconds = float(train_state.train_seconds) + except (AttributeError, TypeError, ValueError) as error: + raise RuntimeError( + "Miles did not export valid round train statistics" + ) from error + if train_seconds < 0: + raise RuntimeError("Miles exported a negative train duration") + metrics = self._rollout_metrics(rollout_id) + group_size = self.args.n_samples_per_prompt return LocalRoundStats( island_id=int(self.args.yeto_rl_learner_id), - local_round_id=expected_policy_version + 1, - base_policy_version=expected_policy_version, - active_groups=int(rollout_metrics["active_groups"]), - completed_groups=groups, - cancelled_groups=int(rollout_metrics["cancelled_groups"]), + local_round_id=rollout_id + 1, + base_policy_version=rollout_id, + active_groups=int(metrics["active_groups"]), + completed_groups=self.args.rollout_batch_size, + cancelled_groups=int(metrics["cancelled_groups"]), completed_trajectories=expected_samples, action_tokens=sum(response_lengths), - tool_wait_seconds=rollout_metrics["tool_wait_seconds"], - group_p50_seconds=rollout_metrics["group_p50_seconds"], - group_p95_seconds=rollout_metrics["group_p95_seconds"], - group_p99_seconds=rollout_metrics["group_p99_seconds"], + tool_wait_seconds=metrics["tool_wait_seconds"], + group_p50_seconds=metrics["group_p50_seconds"], + group_p95_seconds=metrics["group_p95_seconds"], + group_p99_seconds=metrics["group_p99_seconds"], reward_mean=statistics.fmean(rewards), reward_std=statistics.pstdev(rewards), zero_variance_group_ratio=sum( - len(set(rewards[index : index + samples_per_group])) == 1 - for index in range(0, expected_samples, samples_per_group) + len(set(rewards[index : index + group_size])) == 1 + for index in range(0, expected_samples, group_size) ) - / groups, + / self.args.rollout_batch_size, mean_kl=mean_kl, ess_ratio=ess_ratio, clip_fraction=clip_fraction, delta_l2_norm=0.0, - rollout_seconds=rollout_seconds, + rollout_seconds=metrics["rollout_seconds"], train_seconds=train_seconds, ) - def run_local_round( - self, - *, - expected_policy_version: int, - groups: int, - samples_per_group: int, - optimizer_steps: int, - ) -> LocalRoundStats: - return self._run( - self._run_local_round( - expected_policy_version, - groups, - samples_per_group, - optimizer_steps, + async def _initialize(self, *, actor_model, rollout_manager) -> None: + from .bridge import StrictRlBridge + + self.actor_model = actor_model + self.rollout_manager = rollout_manager + initial = self._canonical_state(await actor_model.export_trainable_state()) + runtime = _BridgeRuntime(initial, self.args) + self.bridge = StrictRlBridge(runtime, self.args.yeto_rl_bridge_config) + self.bridge.start() + self.current = self.bridge.wait_for_initial_policy() + await self._apply_global_policy(self.current) + self.args.start_rollout_id = self.current.policy_version + if self.current.policy_version < self.args.num_rollout: + self.permit = self.bridge.wait_for_round() + + async def initialize(self, *, actor_model, rollout_manager) -> None: + try: + await self._initialize( + actor_model=actor_model, + rollout_manager=rollout_manager, ) - ) - - async def _export_local_policy(self) -> CanonicalLoraState: - if not self._trainer_awake: - await self._onload_trainer() - tensors = await self._actor_call("yeto_rl_export_policy", rank0=True) - await self._offload_trainer() - if self._policy_version is None: - raise RuntimeError("Miles has no applied global policy") - return canonical_state( - self._policy_version, - tensors, - base_model_revision=self.args.yeto_rl_base_model_revision, - lora_config_hash=self.args.yeto_rl_lora_config_hash, - layout_hash=self.args.yeto_rl_layout_hash, - ) - - def export_local_policy(self) -> CanonicalLoraState: - return self._run(self._export_local_policy()) + except StrictRlInvariantError as error: + self._record_strict_failure(error) + raise - def record_local_round(self, stats: LocalRoundStats) -> None: - path = Path(self.args.yeto_rl_completed_groups_path).expanduser() - try: - payload = torch.load(path, map_location="cpu", weights_only=True) - except Exception as error: - raise RuntimeError("cannot update Miles island checkpoint") from error + async def _after_local_train( + self, *, rollout_id, actor_model, rollout_data + ) -> None: if ( - not isinstance(payload, dict) - or payload.get("schema_version") != _ISLAND_CHECKPOINT_SCHEMA - or payload.get("policy_version") != stats.base_policy_version - or payload.get("config") != _island_checkpoint_config(self.args) + actor_model is not self.actor_model + or self.current is None + or self.permit is None ): - raise RuntimeError("Miles island checkpoint changed before round commit") - payload["local_round_id"] = stats.local_round_id - payload["local_round_stats"] = asdict(stats) - _atomic_save_island_checkpoint(path, payload) + raise RuntimeError( + "Miles called policy synchronization outside an active round" + ) + if rollout_id != self.current.policy_version: + raise RuntimeError( + "Miles rollout ID differs from the global policy version" + ) + train_state = await actor_model.export_trainable_state() + local = self._canonical_state(train_state) + stats = self._round_stats(rollout_id, rollout_data, train_state) + self.bridge.submit_local_state(self.permit, self.current, local, stats) + self.current = self.bridge.wait_for_global_policy(rollout_id + 1) + await self._apply_global_policy(self.current) + if self.current.policy_version < self.args.num_rollout: + self.permit = self.bridge.wait_for_round() + else: + self.permit = None + + async def after_local_train(self, *, rollout_id, actor_model, rollout_data) -> None: + try: + await self._after_local_train( + rollout_id=rollout_id, + actor_model=actor_model, + rollout_data=rollout_data, + ) + except StrictRlInvariantError as error: + self._record_strict_failure(error) + raise + + async def finalize(self) -> None: + if self.bridge is None or self.current is None: + raise RuntimeError("Miles finalized an uninitialized policy synchronizer") + final = self.bridge.finalize() + if policy_hash(final) != policy_hash(self.current): + raise RuntimeError("final policy differs from the committed global policy") + self.bridge.client.close() def _append_event(self, event: dict[str, Any]) -> None: - path = Path(self.args.yeto_rl_event_tape).expanduser() - path.parent.mkdir(parents=True, exist_ok=True) - event = { - "island_id": int(self.args.yeto_rl_learner_id), - "time_unix": time.time(), - **event, - } - with path.open("a", encoding="utf-8") as handle: - handle.write( - json.dumps(event, sort_keys=True, separators=(",", ":")) + "\n" - ) + _append_rl_event(self.args, event) - async def _shutdown(self) -> None: - if self.actor_model is not None: - await self._offload_trainer() - if self.rollout_manager is not None: - await self.rollout_manager.dispose.remote() + def _record_strict_failure(self, error: StrictRlInvariantError) -> None: + _record_strict_failure(self.args, error, self.bridge) - def shutdown(self) -> None: - try: - self._run(self._shutdown()) - finally: - if self._owns_ray: - import ray - ray.shutdown() - self.loop.close() +def create_policy_sync(args) -> MilesPolicySync: + return MilesPolicySync(args) From 4f2b9636d65634ea05d2bad1233f07987be20a1f Mon Sep 17 00:00:00 2001 From: AlexEisie <1987460907@qq.com> Date: Thu, 30 Jul 2026 23:54:33 +0800 Subject: [PATCH 4/6] feat(rl): add reproducible Miles LM benchmark Compare native Miles, one strict Yeto island, and fixed-roster federated Yeto on identical GPU, prompt-group, trajectory, optimizer, and expert-parallel budgets. Run real Miles rollout and GRPO jobs, retain paired rollout captures, validate prompt identity and completed work, evaluate standard PEFT artifacts on a held-out set, and write atomic resumable reward, pass@k, throughput, synchronization, and cost reports. Reuse the public Miles runner without the Yeto hook for the native reference, isolate concurrent island ports, derive canonical adapter targets through Megatron-Bridge mappings, preserve declared Hugging Face model architectures and provider rotary settings, and pin the maintained gated-attention Miles revision. Reject incomplete native adapter layouts and missing PEFT keys, preserve adapter tensor values during PEFT namespace normalization, and allow a replacement learner joining at terminal state to apply and acknowledge the final committed policy. Add focused coverage for workload fairness, provenance, artifact validation, report aggregation, generic model mappings, native-hook isolation, and terminal replacement recovery. --- README.md | 7 +- docs/RL_BENCHMARK.md | 138 +++ scripts/benchmark_rl.py | 2097 ++++++++++++++++++++++++++++++++++ tests/test_rl_benchmark.py | 921 +++++++++++++++ tests/test_rl_export.py | 53 + tests/test_rl_integration.py | 70 ++ tests/test_rl_launcher.py | 246 +++- yeto/rl/__init__.py | 2 +- yeto/rl/bridge.py | 15 +- yeto/rl/export.py | 28 +- yeto/rl/learner.py | 232 ++-- 11 files changed, 3732 insertions(+), 77 deletions(-) create mode 100644 docs/RL_BENCHMARK.md create mode 100644 scripts/benchmark_rl.py create mode 100644 tests/test_rl_benchmark.py diff --git a/README.md b/README.md index 46e9a6e..290798f 100644 --- a/README.md +++ b/README.md @@ -198,6 +198,8 @@ per-seed results for the completed Qwen3.6, LTX-Video, and Wan2.2 benchmarks. for 1T-class MoE). [docs/MILES_RL.md](docs/MILES_RL.md) — fixed-roster Miles RL across causal-LM LoRA islands: rollout/training boundaries, recovery, export, and limitations. +[docs/RL_BENCHMARK.md](docs/RL_BENCHMARK.md) — equal-hardware native Miles, +single-island Yeto, and federated Yeto RL benchmark contract and runner. [docs/MLX.md](docs/MLX.md) — the Apple-silicon island backend: Macs as learner islands (`yeto launch --external-learners`, cross Mac↔NVIDIA runs). @@ -206,7 +208,7 @@ learner islands (`yeto launch --external-learners`, cross Mac↔NVIDIA runs). python3 -m pytest tests/ # includes a real syncer+learner loop (cd syncer && cargo test) -Three heavier harnesses (all support `--dry-run`): +Four heavier harnesses (all support `--dry-run`): # smoke every supported model with the auto fleet planner, tiered by # size; sequential, self-cleaning, writes a pass/fail report @@ -215,6 +217,9 @@ Three heavier harnesses (all support `--dry-run`): # causal-LM quality check against equal-hardware synchronous baselines python scripts/compare_diloco.py --data --settings all --dry-run + # Miles RL quality check: native, one Yeto island, and federated Yeto + python scripts/benchmark_rl.py --dry-run + # diffusion quality check with an explicit, fixed media shape python scripts/benchmark_diffusion_diloco.py \ --model Lightricks/LTX-Video --data \ diff --git a/docs/RL_BENCHMARK.md b/docs/RL_BENCHMARK.md new file mode 100644 index 0000000..de25ba3 --- /dev/null +++ b/docs/RL_BENCHMARK.md @@ -0,0 +1,138 @@ +# Miles RL LM Benchmark + +`scripts/benchmark_rl.py` compares native Miles with the strict Yeto RL path +on one GPU host. It trains real causal language models with real Miles +rollouts and the selected reward callable; it does not inject trajectories, +use synthetic optimizer steps, or provision cloud resources. + +## Comparison + +For a benchmark size `M` and `G` GPUs per federated island, the harness runs +three arms: + +| arm | topology | purpose | +| --- | --- | --- | +| `native-miles-mM` | one native Miles island on `M*G` GPUs | reference with optimizer state preserved across rounds and no Yeto policy hook | +| `yeto-single-mM` | one Yeto+Miles island on `M*G` GPUs | isolates the Yeto hook, strict checkpoint/apply contract, and LoRA optimizer reset | +| `yeto-federated-mM` | `M` Yeto+Miles islands on `G` GPUs each | measures fixed-roster exact-base averaging and the multi-island split | + +`native -> yeto-single` measures the synchronization contract itself. +`yeto-single -> yeto-federated` measures island partitioning and averaging. +`native -> yeto-federated` is the end-to-end product comparison. + +The single-island Yeto average is an identity operation on LoRA weights, but +it is not equivalent to native Miles: every committed global apply clears the +LoRA optimizer state while preserving LR scheduler progress. + +## Work Accounting + +Let `K` be `--groups-per-island`, `N` be `--samples-per-group`, and `R` be +`--global-rounds`. + +- Each federated island processes `K` prompt groups per round; the native and + Yeto-single arms process `M*K` groups per round. +- Every arm owns `M*G` GPUs and processes `R*M*K*N` trajectories. +- `K*N` must divide `optimizer_steps*G`, which gives every rank the same Miles + batch for all three arms. +- Expert parallelism is one by default. An explicit `--expert-parallel` must + divide `G` and is held fixed across all three arms, including MoE models. +- The maximum action-token budget is identical. Actual response tokens are + recorded because learned policies can terminate at different lengths. + +Training prompts are held in a round-major stream. The single-island arms see +the combined stream; federated island `i` sees its fixed slice from every +round. Rollout capture files are checked against those prompt identities +before a result is accepted. Every island also receives the same deterministic +sampling seed. Independent Miles/SGLang arms may still produce non-bitwise- +identical trajectories; the retained captures record the actual outputs. + +The final rows of the source dataset are held out before any training stream +is built. Every artifact is evaluated on those rows with the same per-sample +generation seeds and the same Miles reward callable. The report includes mean +reward and standard pass@k estimates. A sample passes when its reward is +strictly greater than `--pass-threshold`. + +## Running + +Run inside the pinned Miles RL environment with a clean detached checkout at +the commit recorded by `yeto.rl`. The largest requested `M` needs `M*G` visible +GPUs. The harness starts one local Ray cluster per arm; concurrent federated +islands use independent placement groups and disjoint host-port ranges within +that cluster. `--miles-port-base` moves those ranges when the defaults conflict +with another local service. + +```bash +python scripts/benchmark_rl.py \ + --model Qwen/Qwen3-4B \ + --model-revision \ + --data \ + --data-revision \ + --reward-function project.rewards:score \ + --islands 2,4 \ + --gpus-per-island 2 \ + --global-rounds 8 \ + --groups-per-island 4 \ + --samples-per-group 4 \ + --eval-prompts 64 \ + --eval-samples-per-prompt 4 \ + --pass-k 1,4 \ + --trust-remote-code +``` + +With `G=2`, this example uses `4` total GPUs for every `M=2` arm and `8` +total GPUs for every `M=4` arm. Use `--dry-run` to inspect every topology and +work budget without importing Ray, loading data, or touching a model: + +```bash +python scripts/benchmark_rl.py --dry-run +``` + +Models must be Hugging Face repositories selected by an immutable commit; +mutable local model directories are rejected. Remote datasets likewise require +an immutable Hub revision. A local JSON/JSONL, Parquet, or +`datasets.save_to_disk` input is accepted without `--data-revision`; its +materialized files are content-hashed in the run manifest. + +## Evaluation Scope + +Held-out evaluation uses the pinned base model plus the standard PEFT adapter +in an independent process and, on CUDA, one visible GPU. The base model must +fit that evaluation device. It implements ordinary single-turn LM generation +with the model chat template. This is appropriate for the v0 LM benchmark, +but it does not claim to reproduce a custom multi-turn generate function or a +session-server environment. Such workloads need a benchmark built on the +corresponding Miles environment evaluator rather than this Transformers +generation path. + +Yeto artifacts are always exported from the authoritative syncer checkpoint. +They are never taken from one island's local adapter. Native Miles uses its +own final LoRA save. The harness validates its complete tensor layout against +the model's PEFT contract and adds the PEFT wrapper prefix omitted by Miles; +tensor values are not changed. Miles performs the native save inside the +measured job, while Yeto checkpoint export follows the measured training job. +Results report both `train_wall_s` and the comparable `artifact_ready_s`; each +arm's adapter preparation duration is also retained as `artifact_s`. + +## Outputs And Resume + +The default output directories are `rl-benchmark-work/` and +`rl-benchmark-report/`. + +- `config.json` contains the immutable workload, complete Yeto-source and + built-syncer implementation fingerprint, Miles commit, reward and data + hashes, and fairness contract. +- `results.jsonl` is updated atomically after each completed arm and seed. +- `summary.json` contains aggregates across training seeds. +- `report.md` contains quality deltas and systems measurements. +- Each arm directory retains Miles logs, real rollout captures, evaluation + samples, and either the native adapter or authoritative Yeto export. + +Use `--resume` with unchanged inputs and arguments to skip completed records. +The harness refuses changed data, implementation, arms, or workload settings. +Use `--overwrite` only when intentionally starting a new result set. + +Systems fields include training and artifact-ready wall time, trajectories/s, +action tokens/s, GPU-hours, estimated cost, Yeto synchronization time and +bytes, and KL when emitted by the Yeto hook. Native Miles does not pass through +that hook, so hook-only diagnostics are intentionally absent for the native +arm. diff --git a/scripts/benchmark_rl.py b/scripts/benchmark_rl.py new file mode 100644 index 0000000..7656c78 --- /dev/null +++ b/scripts/benchmark_rl.py @@ -0,0 +1,2097 @@ +#!/usr/bin/env python3 +"""Benchmark native Miles, one Yeto island, and federated Yeto RL. + +For a requested M and G, every arm owns M*G GPUs and processes the same +round-major prompt groups. The harness runs on one host and does not create +cloud resources. +""" + +from __future__ import annotations + +import argparse +import asyncio +import json +import math +import os +import shutil +import signal +import socket +import statistics +import subprocess +import sys +import time +import warnings +from contextlib import contextmanager +from dataclasses import asdict, dataclass +from pathlib import Path +from types import SimpleNamespace +from typing import Any + +REPO_ROOT = Path(__file__).resolve().parent.parent +sys.path.insert(0, str(REPO_ROOT)) +from yeto.benchmark_resume import write_json_atomic +from yeto.provenance import file_sha256 + +SYNCER_BIN = REPO_ROOT / "syncer/target/release/yeto-syncer" +_MILES_PORT_STRIDE = 1000 +_RESUME_EXCLUDES = { + "_active_seed", + "_pass_ks", + "dry_run", + "overwrite", + "report_dir", + "resume", + "work_dir", +} +_IMPLEMENTATION_PATHS = ( + Path(__file__), + REPO_ROOT / "yeto", + REPO_ROOT / "syncer/src", + REPO_ROOT / "syncer/Cargo.toml", + REPO_ROOT / "syncer/Cargo.lock", + SYNCER_BIN, +) + + +@dataclass(frozen=True) +class Arm: + name: str + kind: str + benchmark_islands: int + islands: int + gpus_per_island: int + groups_per_round: int + + +@dataclass(frozen=True) +class PromptStreams: + combined_rows: tuple[dict[str, Any], ...] + island_rows: tuple[tuple[dict[str, Any], ...], ...] + combined_ids: tuple[int, ...] + island_ids: tuple[tuple[int, ...], ...] + + +@dataclass(frozen=True) +class WorkerSpec: + learner_id: int + gpus: int + groups_per_round: int + prompt_path: Path + policy_sync: bool + + +def _positive_csv(spec: str, flag: str) -> list[int]: + try: + values = [int(value.strip()) for value in spec.split(",") if value.strip()] + except ValueError as exc: + raise ValueError(f"{flag} must be a comma-separated list of integers") from exc + if not values: + raise ValueError(f"{flag} must contain at least one value") + if any(value <= 0 for value in values): + raise ValueError(f"{flag} values must be positive") + if len(values) != len(set(values)): + raise ValueError(f"{flag} contains duplicates") + return values + + +def select_arms( + spec: str, + gpus_per_island: int, + groups_per_island: int, +) -> list[Arm]: + arms = [] + for islands in _positive_csv(spec, "--islands"): + total_gpus = islands * gpus_per_island + total_groups = islands * groups_per_island + arms.extend( + ( + Arm( + f"native-miles-m{islands}", + "native", + islands, + 1, + total_gpus, + total_groups, + ), + Arm( + f"yeto-single-m{islands}", + "single", + islands, + 1, + total_gpus, + total_groups, + ), + Arm( + f"yeto-federated-m{islands}", + "federated", + islands, + islands, + gpus_per_island, + groups_per_island, + ), + ) + ) + return arms + + +def workload( + arm: Arm, *, rounds: int, samples_per_group: int +) -> dict[str, int | float]: + prompt_groups = arm.islands * arm.groups_per_round * rounds + total_gpus = arm.islands * arm.gpus_per_island + return { + "total_gpus": total_gpus, + "prompt_groups": prompt_groups, + "trajectories": prompt_groups * samples_per_group, + "groups_per_gpu_per_round": (arm.islands * arm.groups_per_round / total_gpus), + } + + +def validate_workload(args) -> None: + for name in ( + "global_rounds", + "groups_per_island", + "samples_per_group", + "optimizer_steps", + "gpus_per_island", + ): + if getattr(args, name) <= 0: + raise ValueError(f"--{name.replace('_', '-')} must be positive") + samples = args.groups_per_island * args.samples_per_group + divisor = args.optimizer_steps * args.gpus_per_island + if samples % divisor: + raise ValueError( + "groups-per-island*samples-per-group must be divisible by " + "optimizer-steps*gpus-per-island" + ) + + +def paired_prompt_streams( + rows: list[dict[str, Any]], + *, + islands: int, + groups: int, + rounds: int, +) -> PromptStreams: + if not rows: + raise ValueError("RL benchmark training split is empty") + count = islands * groups * rounds + ids = tuple(index % len(rows) for index in range(count)) + combined = tuple(dict(rows[index]) for index in ids) + island_ids = tuple( + tuple( + ids[round_id * islands * groups + island_id * groups + offset] + for round_id in range(rounds) + for offset in range(groups) + ) + for island_id in range(islands) + ) + island_rows = tuple( + tuple(dict(rows[index]) for index in indices) for indices in island_ids + ) + return PromptStreams(combined, island_rows, ids, island_ids) + + +def _normalized_prompt(row: dict[str, Any], prompt_id: int | str) -> dict[str, Any]: + value = row.get("messages", row.get("prompt", row.get("input"))) + if isinstance(value, str): + value = [{"role": "user", "content": value}] + if ( + not isinstance(value, list) + or not value + or any(not isinstance(message, dict) for message in value) + ): + raise ValueError("RL rows must contain messages or a string prompt/input") + metadata = dict(row.get("metadata") or {}) + for key, item in row.items(): + if key not in {"messages", "prompt", "input", "label", "metadata", "tools"}: + metadata.setdefault(key, item) + metadata["benchmark_prompt_id"] = prompt_id + output = { + "messages": value, + "label": row.get("label"), + "metadata": metadata, + } + if "tools" in row: + output["tools"] = row["tools"] + return output + + +def _write_jsonl(path: Path, rows) -> None: + path.parent.mkdir(parents=True, exist_ok=True) + temporary = path.with_name(path.name + ".tmp") + with temporary.open("w", encoding="utf-8") as handle: + for row in rows: + handle.write( + json.dumps(row, ensure_ascii=False, separators=(",", ":")) + "\n" + ) + temporary.replace(path) + + +def write_prompt_files( + streams: PromptStreams, + evaluation_rows: list[dict[str, Any]], + directory: Path, +) -> tuple[Path, tuple[Path, ...], Path]: + combined = directory / "combined.jsonl" + _write_jsonl( + combined, + ( + _normalized_prompt(row, prompt_id) + for row, prompt_id in zip( + streams.combined_rows, streams.combined_ids, strict=True + ) + ), + ) + island_paths = [] + for island_id, (rows, prompt_ids) in enumerate( + zip(streams.island_rows, streams.island_ids, strict=True) + ): + path = directory / f"island-{island_id}.jsonl" + _write_jsonl( + path, + ( + _normalized_prompt(row, prompt_id) + for row, prompt_id in zip(rows, prompt_ids, strict=True) + ), + ) + island_paths.append(path) + evaluation = directory / "eval.jsonl" + _write_jsonl( + evaluation, + ( + _normalized_prompt(row, f"eval:{index}") + for index, row in enumerate(evaluation_rows) + ), + ) + return combined, tuple(island_paths), evaluation + + +def verify_worker_inputs( + *, + prompt_path: Path, + prompt_sha256: str, + reward_function: str, + reward_sha256: str, +) -> None: + from yeto.provenance import python_spec_sha256 + + actual_prompt = file_sha256(prompt_path) + if actual_prompt != prompt_sha256.lower(): + raise RuntimeError( + "prompt source SHA256 mismatch: " + f"expected {prompt_sha256.lower()}, got {actual_prompt}" + ) + actual_reward = python_spec_sha256(reward_function, base_dir=REPO_ROOT) + if actual_reward != reward_sha256.lower(): + raise RuntimeError( + "reward source SHA256 mismatch: " + f"expected {reward_sha256.lower()}, got {actual_reward}" + ) + + +def worker_specs( + arm: Arm, + combined_prompt_path: Path, + island_prompt_paths: tuple[Path, ...], +) -> list[WorkerSpec]: + if arm.kind != "federated": + return [ + WorkerSpec( + 0, + arm.gpus_per_island, + arm.groups_per_round, + combined_prompt_path, + arm.kind != "native", + ) + ] + if len(island_prompt_paths) != arm.islands: + raise ValueError("federated prompt file count does not match its roster") + return [ + WorkerSpec( + learner_id, + arm.gpus_per_island, + arm.groups_per_round, + island_prompt_paths[learner_id], + True, + ) + for learner_id in range(arm.islands) + ] + + +def syncer_command(arm: Arm, port: int, run_dir: Path, *, rounds: int) -> list[str]: + return [ + str(SYNCER_BIN), + "--port", + str(port), + "--learners", + str(arm.islands), + "--quorum", + str(arm.islands), + "--grace-ms", + "0", + "--pipeline", + "1", + "--sync-interval-steps", + "0", + "--delta-correction", + "none", + "--total-steps", + str(rounds), + "--outer-lr", + "1", + "--outer-momentum", + "0", + "--max-base-lag", + "0", + "--learner-weight", + "equal", + "--checkpoint-path", + str(run_dir / "state.ckpt"), + "--checkpoint-every", + "1", + "--resume", + "--event-tape", + str(run_dir / "syncer.jsonl"), + ] + + +def miles_extra_argv(worker: WorkerSpec, run_dir: Path, rounds: int) -> list[str]: + values = [ + "--save-debug-rollout-data", + str(run_dir / "rollouts" / f"island-{worker.learner_id}" / "{rollout_id}.pt"), + ] + if not worker.policy_sync: + values.extend( + ( + "--save", + str(run_dir / "native-checkpoint"), + "--save-interval", + str(rounds), + ) + ) + return values + + +def canonical_native_adapter_tensors(tensors, specs) -> dict[str, Any]: + expected = {spec.name: tuple(spec.shape) for spec in specs} + mapped = {} + for name, tensor in tensors.items(): + canonical_name = name if name in expected else f"base_model.model.{name}" + if ( + canonical_name not in expected + or canonical_name in mapped + or tuple(tensor.shape) != expected[canonical_name] + ): + raise RuntimeError("native Miles adapter does not match the PEFT contract") + mapped[canonical_name] = tensor + if set(mapped) != set(expected): + raise RuntimeError("native Miles adapter does not match the PEFT contract") + return {name: mapped[name] for name in expected} + + +def standardize_native_adapter(args, source: Path, output: Path) -> Path: + import torch + + from yeto.rl.core import ( + canonical_layout_hash, + canonical_lora_config_hash, + canonical_state, + ) + from yeto.rl.export import ( + adapter_targets, + derive_peft_lora_specs, + write_peft_adapter, + ) + + specs = derive_peft_lora_specs( + args.model, + args.model_revision, + rank=args.lora_r, + targets=args.lora_targets, + trust_remote_code=args.trust_remote_code, + ) + raw = torch.load( + source / "adapter_model.bin", + map_location="cpu", + weights_only=True, + ) + if not isinstance(raw, dict): + raise TypeError("native Miles adapter does not contain a tensor mapping") + targets = adapter_targets(specs) + state = canonical_state( + args.global_rounds, + canonical_native_adapter_tensors(raw, specs), + base_model_revision=args.model_revision, + lora_config_hash=canonical_lora_config_hash( + rank=args.lora_r, + target_modules=targets, + ), + layout_hash=canonical_layout_hash(specs), + expected_specs=specs, + ) + write_peft_adapter( + state, + output, + base_model=args.model, + model_revision=args.model_revision, + rank=args.lora_r, + ) + return output + + +def worker_payload( + args, + worker: WorkerSpec, + *, + run_dir: Path, + model_path: Path, + syncer: str | None, + reward_sha256: str, +) -> dict[str, Any]: + worker_dir = run_dir / f"island-{worker.learner_id}" + miles_port_base = args.miles_port_base + worker.learner_id * _MILES_PORT_STRIDE + values = { + "model": args.model, + "model_revision": args.model_revision, + "data": args.data, + "data_revision": file_sha256(worker.prompt_path), + "syncer": syncer, + "learner_id": worker.learner_id, + "reward_function": args.reward_function, + "reward_sha256": reward_sha256, + "global_rounds": args.global_rounds, + "groups_per_round": worker.groups_per_round, + "samples_per_group": args.samples_per_group, + "over_sampling_batch_size": worker.groups_per_round, + "optimizer_steps": args.optimizer_steps, + "rollout_max_response_len": args.rollout_max_response_len, + "custom_generate_function_path": None, + "use_session_server": False, + "session_server_ip": None, + "session_server_port": None, + "tito_model": None, + "completed_groups_path": str(worker_dir / "completed-groups.pt"), + "event_tape": str(worker_dir / "events.jsonl"), + "actor_num_nodes": 1, + "actor_num_gpus_per_node": worker.gpus, + "expert_parallel": args.expert_parallel, + "lora_r": args.lora_r, + "lora_targets": args.lora_targets, + "inner_lr": args.inner_lr, + "seq_len": args.seq_len, + "seed": args._active_seed, + "rollout_seed": args._active_seed, + "rollout_engine_base_port": miles_port_base + 100, + "sglang_router_port": miles_port_base, + "sglang_router_prometheus_port": miles_port_base + 1, + "train_master_base_port": miles_port_base + 2, + "wan_streams": args.wan_streams, + "miles_root": str(args.miles_root.expanduser().resolve()), + "trust_remote_code": args.trust_remote_code, + } + return { + "arguments": values, + "model_path": str(model_path), + "prompt_path": str(worker.prompt_path), + "policy_sync": worker.policy_sync, + "extra_argv": miles_extra_argv(worker, run_dir, args.global_rounds), + } + + +def run_training_worker(config_path: Path) -> int: + payload = json.loads(config_path.read_text(encoding="utf-8")) + args = SimpleNamespace(**payload["arguments"]) + miles_root = str(Path(args.miles_root).expanduser().resolve()) + for path in (miles_root, str(REPO_ROOT)): + if path not in sys.path: + sys.path.insert(0, path) + + verify_worker_inputs( + prompt_path=Path(payload["prompt_path"]), + prompt_sha256=args.data_revision, + reward_function=args.reward_function, + reward_sha256=args.reward_sha256, + ) + + from miles.utils.misc import load_function + + from yeto.rl.learner import _miles_callable, run_miles + + load_function(_miles_callable(args.reward_function)) + run_miles( + args, + model_path=payload["model_path"], + prompt_path=payload["prompt_path"], + yeto_policy_sync=bool(payload["policy_sync"]), + extra_argv=tuple(payload["extra_argv"]), + ) + return 0 + + +def summarize_rollouts( + rollout_paths: tuple[tuple[Path, ...], ...], + *, + expected_prompt_ids: tuple[tuple[int, ...], ...], + samples_per_group: int, +) -> dict[str, Any]: + import torch + + if len(rollout_paths) != len(expected_prompt_ids): + raise RuntimeError("rollout island count does not match prompt manifest") + rewards = [] + action_tokens = 0 + trajectories = 0 + truncated = 0 + for island_paths, expected in zip(rollout_paths, expected_prompt_ids, strict=True): + observed = [] + for path in island_paths: + # Miles' own capture contains NumPy routing arrays and is generated + # inside this run, so use its trusted debug-data load semantics. + payload = torch.load(path, map_location="cpu", weights_only=False) + samples = payload.get("samples") if isinstance(payload, dict) else None + if not isinstance(samples, list) or len(samples) % samples_per_group: + raise RuntimeError(f"invalid Miles rollout capture: {path}") + for start in range(0, len(samples), samples_per_group): + group = samples[start : start + samples_per_group] + prompt_ids = { + (sample.get("metadata") or {}).get("benchmark_prompt_id") + for sample in group + if isinstance(sample, dict) + } + if len(prompt_ids) != 1: + raise RuntimeError("Miles rollout group lost prompt identity") + observed.append(prompt_ids.pop()) + for sample in group: + try: + reward = float(sample["reward"]) + length = int(sample["response_length"]) + except (KeyError, TypeError, ValueError) as exc: + raise RuntimeError( + "Miles rollout capture lacks scalar work metrics" + ) from exc + if length < 0 or not math.isfinite(reward): + raise RuntimeError( + "Miles rollout capture has invalid work metrics" + ) + status = sample.get("status") + if status not in {"completed", "truncated"}: + raise RuntimeError( + f"Miles rollout capture has invalid status {status!r}" + ) + rewards.append(reward) + action_tokens += length + trajectories += 1 + truncated += status == "truncated" + if tuple(observed) != tuple(expected): + raise RuntimeError( + f"prompt stream mismatch: expected {tuple(expected)}, got {tuple(observed)}" + ) + return { + "prompt_groups": trajectories // samples_per_group, + "trajectories": trajectories, + "action_tokens": action_tokens, + "reward_mean": statistics.fmean(rewards), + "reward_std": statistics.stdev(rewards) if len(rewards) > 1 else 0.0, + "truncated_trajectories": truncated, + } + + +def _read_jsonl(path: Path) -> list[dict[str, Any]]: + rows = [] + for line in path.read_text(encoding="utf-8").splitlines(): + if line.strip(): + value = json.loads(line) + if not isinstance(value, dict): + raise ValueError(f"expected JSON objects in {path}") + rows.append(value) + return rows + + +def prepare_evaluation_prompt( + tokenizer, + row: dict[str, Any], + *, + max_prompt_tokens: int, + device: str, +) -> tuple[str, dict[str, Any]]: + import torch + from miles.utils.chat_template_utils import apply_chat_template + + prompt = apply_chat_template( + row["messages"], + tokenizer=tokenizer, + tools=row.get("tools"), + tokenize=False, + add_generation_prompt=True, + ) + encoded = dict( + tokenizer( + prompt, + add_special_tokens=False, + return_tensors="pt", + ) + ) + input_ids = encoded.get("input_ids") + if not isinstance(input_ids, torch.Tensor) or input_ids.ndim != 2: + raise RuntimeError("tokenizer did not return batched input_ids") + prompt_tokens = input_ids.shape[-1] + for name, value in encoded.items(): + if not isinstance(value, torch.Tensor): + raise TypeError(f"tokenizer returned non-tensor field {name!r}") + if value.ndim >= 2 and value.shape[-1] == prompt_tokens: + value = value[..., -max_prompt_tokens:] + encoded[name] = value.to(device) + if "attention_mask" not in encoded: + encoded["attention_mask"] = torch.ones_like(encoded["input_ids"]) + return prompt, encoded + + +def generation_pad_token_id(tokenizer) -> int | None: + if tokenizer.pad_token_id is not None: + return int(tokenizer.pad_token_id) + eos = tokenizer.eos_token_id + if isinstance(eos, (list, tuple)): + return int(eos[0]) if eos else None + return int(eos) if eos is not None else None + + +async def evaluate_rewards(args, samples: list[Any]) -> list[Any]: + from miles.rollout.rm_hub import async_rm + + return list(await asyncio.gather(*(async_rm(args, sample) for sample in samples))) + + +def load_peft_adapter(peft_model, model, adapter_path): + message = "Found missing adapter keys while loading the checkpoint:" + with warnings.catch_warnings(): + warnings.filterwarnings("error", message=message) + try: + return peft_model.from_pretrained(model, adapter_path) + except UserWarning as exc: + if not str(exc).startswith(message): + raise + raise RuntimeError(str(exc)) from exc + + +def run_evaluation_worker(config_path: Path) -> int: + started = time.monotonic() + payload = json.loads(config_path.read_text(encoding="utf-8")) + miles_root = str(Path(payload["miles_root"]).expanduser().resolve()) + for path in (miles_root, str(REPO_ROOT)): + if path not in sys.path: + sys.path.insert(0, path) + + import torch + from miles.utils.types import Sample + from peft import PeftModel + from transformers import AutoConfig, AutoModelForCausalLM, AutoTokenizer + + from yeto.rl.export import _rl_model_factory + from yeto.rl.learner import _miles_callable + + verify_worker_inputs( + prompt_path=Path(payload["eval_path"]), + prompt_sha256=payload["eval_sha256"], + reward_function=payload["reward_function"], + reward_sha256=payload["reward_sha256"], + ) + + device = payload["device"] + dtype = torch.float32 if device == "cpu" else torch.bfloat16 + tokenizer = AutoTokenizer.from_pretrained( + payload["model_path"], + trust_remote_code=payload["trust_remote_code"], + ) + config = AutoConfig.from_pretrained( + payload["model_path"], + trust_remote_code=payload["trust_remote_code"], + ) + model_factory = _rl_model_factory(config) + model_kwargs = { + "config": config, + "dtype": dtype, + "low_cpu_mem_usage": True, + } + if model_factory is AutoModelForCausalLM: + model_kwargs["trust_remote_code"] = payload["trust_remote_code"] + model = model_factory.from_pretrained( + payload["model_path"], + **model_kwargs, + ) + model = load_peft_adapter(PeftModel, model, payload["adapter_path"]) + model.to(device) + model.eval() + + rows = _read_jsonl(Path(payload["eval_path"])) + max_response = int(payload["max_response_len"]) + max_prompt = int(payload["seq_len"]) - max_response + reward_args = SimpleNamespace( + **payload["reward_arguments"], + custom_rm_path=_miles_callable(payload["reward_function"]), + multi_lora=False, + ) + samples = [] + sample_records = [] + with torch.no_grad(): + for prompt_index, row in enumerate(rows): + prompt, encoded = prepare_evaluation_prompt( + tokenizer, + row, + max_prompt_tokens=max_prompt, + device=device, + ) + prompt_tokens = encoded["input_ids"][0].tolist() + for sample_index in range(payload["samples_per_prompt"]): + generation_seed = ( + int(payload["seed"]) + + prompt_index * payload["samples_per_prompt"] + + sample_index + ) + torch.manual_seed(generation_seed) + if device != "cpu": + torch.cuda.manual_seed_all(generation_seed) + generation = { + "max_new_tokens": max_response, + "do_sample": payload["temperature"] > 0, + "pad_token_id": generation_pad_token_id(tokenizer), + } + if generation["do_sample"]: + generation.update( + temperature=payload["temperature"], + top_p=payload["top_p"], + ) + output = model.generate(**encoded, **generation) + response_tokens = output[0, encoded["input_ids"].shape[1] :].tolist() + eos_ids = tokenizer.eos_token_id + eos_ids = set( + eos_ids if isinstance(eos_ids, (list, tuple)) else [eos_ids] + ) + eos_ids.discard(None) + truncated = len(response_tokens) == max_response and ( + not response_tokens or response_tokens[-1] not in eos_ids + ) + metadata = dict(row.get("metadata") or {}) + if row.get("tools") is not None: + metadata["tools"] = row["tools"] + sample = Sample( + group_index=prompt_index, + index=len(samples), + prompt=prompt, + tokens=prompt_tokens + response_tokens, + response=tokenizer.decode( + response_tokens, skip_special_tokens=True + ), + response_length=len(response_tokens), + label=row.get("label"), + metadata=metadata, + status=( + Sample.Status.TRUNCATED + if truncated + else Sample.Status.COMPLETED + ), + ) + samples.append(sample) + sample_records.append( + { + "prompt_index": prompt_index, + "sample_index": sample_index, + "generation_seed": generation_seed, + "response": sample.response, + "response_tokens": sample.response_length, + "truncated": truncated, + } + ) + + rewards = asyncio.run(evaluate_rewards(reward_args, samples)) + if not isinstance(rewards, list) or len(rewards) != len(samples): + raise RuntimeError("Miles reward callable did not return one reward per sample") + scalar_rewards = [] + for record, sample, value in zip(sample_records, samples, rewards, strict=True): + try: + reward = float(value) + except (TypeError, ValueError) as exc: + raise RuntimeError( + "RL benchmark requires scalar evaluation rewards" + ) from exc + if not math.isfinite(reward): + raise RuntimeError("evaluation reward contains NaN or Inf") + sample.reward = reward + record["reward"] = reward + scalar_rewards.append(reward) + + grouped = [ + scalar_rewards[index : index + payload["samples_per_prompt"]] + for index in range(0, len(scalar_rewards), payload["samples_per_prompt"]) + ] + summary = summarize_rewards( + grouped, + pass_ks=tuple(payload["pass_ks"]), + threshold=float(payload["pass_threshold"]), + ) + summary.update( + { + "prompts": len(rows), + "samples": len(samples), + "response_tokens": sum(sample.response_length for sample in samples), + "truncated_samples": sum( + sample.status == Sample.Status.TRUNCATED for sample in samples + ), + "wall_s": time.monotonic() - started, + } + ) + result_path = Path(payload["result_path"]) + write_json_atomic(result_path, summary) + _write_jsonl(Path(payload["samples_path"]), sample_records) + return 0 + + +def _free_port() -> int: + with socket.socket() as sock: + sock.bind(("127.0.0.1", 0)) + return int(sock.getsockname()[1]) + + +def _tail(path: Path, lines: int = 40) -> str: + if not path.exists(): + return "" + return "\n".join( + path.read_text(encoding="utf-8", errors="replace").splitlines()[-lines:] + ) + + +def _stop_process(process: subprocess.Popen, timeout: int = 20) -> None: + if process.poll() is not None: + return + try: + os.killpg(process.pid, signal.SIGTERM) + except (OSError, ProcessLookupError): + process.terminate() + try: + process.wait(timeout=timeout) + except subprocess.TimeoutExpired: + try: + os.killpg(process.pid, signal.SIGKILL) + except (OSError, ProcessLookupError): + process.kill() + process.wait(timeout=10) + + +def _visible_cuda_devices() -> list[str] | None: + raw = os.environ.get("CUDA_VISIBLE_DEVICES") + if raw is None or not raw.strip(): + return None + return [value.strip() for value in raw.split(",") if value.strip()] + + +def _visible_gpu_uuids() -> set[str] | None: + visible = _visible_cuda_devices() + if visible is None: + return None + result = subprocess.run( + [ + "nvidia-smi", + "--query-gpu=index,uuid", + "--format=csv,noheader,nounits", + ], + capture_output=True, + text=True, + check=False, + ) + if result.returncode != 0: + raise RuntimeError(f"nvidia-smi failed: {result.stderr.strip()}") + index_to_uuid = {} + for line in result.stdout.splitlines(): + parts = [value.strip() for value in line.split(",")] + if len(parts) >= 2: + index_to_uuid[parts[0]] = parts[1] + return {index_to_uuid.get(value, value) for value in visible} + + +def wait_for_free_gpus(limit_mb: int = 2000, timeout_s: int = 300) -> None: + visible_uuids = _visible_gpu_uuids() + deadline = time.monotonic() + timeout_s + last = "" + while True: + result = subprocess.run( + [ + "nvidia-smi", + "--query-compute-apps=gpu_uuid,pid,process_name,used_gpu_memory", + "--format=csv,noheader,nounits", + ], + capture_output=True, + text=True, + check=False, + ) + if result.returncode != 0: + raise RuntimeError(f"nvidia-smi failed: {result.stderr.strip()}") + holders = [] + for line in result.stdout.splitlines(): + parts = [value.strip() for value in line.split(",")] + if len(parts) < 4: + continue + gpu_uuid, pid, name, memory = parts[0], parts[1], parts[2], parts[-1] + if visible_uuids is not None and gpu_uuid not in visible_uuids: + continue + if not memory.isdigit() or int(memory) > limit_mb: + holders.append(f"pid {pid} ({name}): {memory} MiB") + if not holders: + return + current = "; ".join(holders) + if current != last: + print(f"[rl-benchmark] waiting for GPUs: {current}", flush=True) + last = current + if time.monotonic() >= deadline: + raise RuntimeError(f"GPUs still occupied after {timeout_s}s: {current}") + time.sleep(3) + + +@contextmanager +def local_ray_cluster(total_gpus: int): + if os.environ.get("RAY_ADDRESS"): + raise RuntimeError("unset RAY_ADDRESS before running the local RL benchmark") + import ray + + if ray.is_initialized(): + raise RuntimeError("RL benchmark requires no pre-existing Ray connection") + context = ray.init( + num_gpus=total_gpus, + include_dashboard=True, + logging_level="ERROR", + ) + try: + yield context.address_info["address"] + finally: + ray.shutdown() + + +def _wait_for_port( + port: int, process: subprocess.Popen, log: Path, timeout_s: int = 30 +) -> None: + deadline = time.monotonic() + timeout_s + while time.monotonic() < deadline: + if process.poll() is not None: + raise RuntimeError( + f"syncer exited before listening with code {process.returncode}:\n{_tail(log)}" + ) + with socket.socket() as sock: + sock.settimeout(0.2) + if sock.connect_ex(("127.0.0.1", port)) == 0: + return + time.sleep(0.1) + raise RuntimeError(f"syncer did not listen on port {port}:\n{_tail(log)}") + + +def _wait_for_training( + processes: list[subprocess.Popen], + logs: list[Path], + *, + syncer: subprocess.Popen | None, + syncer_log: Path | None, + timeout_s: int, +) -> None: + deadline = time.monotonic() + timeout_s + pending = set(range(len(processes))) + while pending: + for index in list(pending): + returncode = processes[index].poll() + if returncode is None: + continue + pending.remove(index) + if returncode != 0: + raise RuntimeError( + f"Miles island {index} failed with code {returncode}:\n{_tail(logs[index])}" + ) + if syncer is not None and syncer.poll() not in (None, 0): + raise RuntimeError( + f"syncer failed with code {syncer.returncode}:\n" + f"{_tail(syncer_log) if syncer_log else ''}" + ) + if time.monotonic() >= deadline: + raise RuntimeError(f"RL benchmark arm timed out after {timeout_s}s") + if pending: + time.sleep(1) + if syncer is not None: + try: + returncode = syncer.wait(timeout=60) + except subprocess.TimeoutExpired as exc: + raise RuntimeError("syncer did not finish after all Miles islands") from exc + if returncode != 0: + raise RuntimeError( + f"syncer failed with code {returncode}:\n" + f"{_tail(syncer_log) if syncer_log else ''}" + ) + + +def _training_environment(ray_address: str, miles_root: Path) -> dict[str, str]: + env = dict(os.environ) + env["RAY_ADDRESS"] = ray_address + existing = env.get("PYTHONPATH") + paths = [str(miles_root), str(REPO_ROOT)] + if existing: + paths.append(existing) + env["PYTHONPATH"] = os.pathsep.join(paths) + env["NVTE_FLASH_ATTN"] = "0" + env["NVTE_FUSED_ATTN"] = "0" + env["NVTE_UNFUSED_ATTN"] = "1" + return env + + +def _run_training_processes( + args, + arm: Arm, + workers: list[WorkerSpec], + *, + run_dir: Path, + model_path: Path, + reward_sha256: str, +) -> float: + port = _free_port() if arm.kind != "native" else None + syncer_address = f"127.0.0.1:{port}" if port is not None else None + payload_paths = [] + for worker in workers: + path = run_dir / f"island-{worker.learner_id}" / "worker.json" + write_json_atomic( + path, + worker_payload( + args, + worker, + run_dir=run_dir, + model_path=model_path, + syncer=syncer_address, + reward_sha256=reward_sha256, + ), + ) + payload_paths.append(path) + + syncer = None + syncer_handle = None + syncer_log = run_dir / "syncer.log" + processes = [] + handles = [] + logs = [] + started = time.monotonic() + try: + with local_ray_cluster( + arm.islands * arm.gpus_per_island, + ) as ray_address: + if port is not None: + syncer_handle = syncer_log.open("w", encoding="utf-8") + syncer = subprocess.Popen( + syncer_command( + arm, + port, + run_dir, + rounds=args.global_rounds, + ), + cwd=REPO_ROOT, + stdout=syncer_handle, + stderr=subprocess.STDOUT, + start_new_session=True, + ) + _wait_for_port(port, syncer, syncer_log) + env = _training_environment(ray_address, args.miles_root) + for worker, payload_path in zip(workers, payload_paths, strict=True): + log = run_dir / f"island-{worker.learner_id}" / "miles.log" + log.parent.mkdir(parents=True, exist_ok=True) + handle = log.open("w", encoding="utf-8") + handles.append(handle) + logs.append(log) + processes.append( + subprocess.Popen( + [ + sys.executable, + str(Path(__file__).resolve()), + "_train-worker", + str(payload_path), + ], + cwd=REPO_ROOT, + stdout=handle, + stderr=subprocess.STDOUT, + env=env, + start_new_session=True, + ) + ) + _wait_for_training( + processes, + logs, + syncer=syncer, + syncer_log=syncer_log, + timeout_s=args.arm_timeout_min * 60, + ) + finally: + for process in processes: + _stop_process(process) + if syncer is not None: + _stop_process(syncer) + for handle in handles: + handle.close() + if syncer_handle is not None: + syncer_handle.close() + return time.monotonic() - started + + +def _mean(values: list[float]) -> float | None: + return statistics.fmean(values) if values else None + + +def summarize_yeto_events(run_dir: Path, islands: int) -> dict[str, Any]: + local_rounds = [] + apply_events = [] + for island_id in range(islands): + path = run_dir / f"island-{island_id}" / "events.jsonl" + if not path.exists(): + continue + for event in _read_jsonl(path): + if event.get("event") == "rl_local_round": + local_rounds.append(event) + elif event.get("event") == "rl_policy_apply": + apply_events.append(event) + sync_records = ( + _read_jsonl(run_dir / "syncer.jsonl") + if (run_dir / "syncer.jsonl").exists() + else [] + ) + + def values(name: str) -> list[float]: + return [ + float(event[name]) for event in local_rounds if event.get(name) is not None + ] + + return { + "local_rounds": len(local_rounds), + "policy_applies": len(apply_events), + "rollout_s": sum(values("rollout_seconds")), + "optimizer_train_s": sum(values("train_seconds")), + "mean_kl": _mean(values("mean_kl")), + "mean_ess_ratio": _mean(values("ess_ratio")), + "mean_clip_fraction": _mean(values("clip_fraction")), + "sync_bytes_sent": int(sum(values("sync/bytes_sent"))), + "mean_sync_ms": _mean( + [ + float(event["ms"]) + for event in sync_records + if event.get("ms") is not None + ] + ), + } + + +def ensure_syncer() -> None: + subprocess.run( + ["cargo", "build", "--release", "--locked", "--quiet"], + cwd=REPO_ROOT / "syncer", + check=True, + ) + + +def resolve_model_path(args) -> Path: + from huggingface_hub import snapshot_download + + from yeto.models import resolve + + model = resolve(args.model) + args.model = model + return Path(snapshot_download(repo_id=model, revision=args.model_revision)) + + +def _run_checked( + command: list[str], + log: Path, + *, + env: dict[str, str] | None = None, + timeout_s: int, +) -> None: + log.parent.mkdir(parents=True, exist_ok=True) + with log.open("w", encoding="utf-8") as handle: + process = subprocess.Popen( + command, + cwd=REPO_ROOT, + stdout=handle, + stderr=subprocess.STDOUT, + env=env, + start_new_session=True, + ) + try: + returncode = process.wait(timeout=timeout_s) + except subprocess.TimeoutExpired as exc: + _stop_process(process) + raise RuntimeError( + f"command timed out: {' '.join(command)}\n{_tail(log)}" + ) from exc + except BaseException: + _stop_process(process) + raise + if returncode != 0: + raise RuntimeError( + f"command failed with code {returncode}: {' '.join(command)}\n{_tail(log)}" + ) + + +def evaluate_artifact( + args, + *, + adapter_path: Path, + model_path: Path, + eval_path: Path, + reward_sha256: str, + run_dir: Path, + seed: int, +) -> dict[str, Any]: + result_path = run_dir / "eval.json" + payload = { + "model_path": str(model_path), + "adapter_path": str(adapter_path), + "eval_path": str(eval_path), + "eval_sha256": file_sha256(eval_path), + "reward_function": args.reward_function, + "reward_sha256": reward_sha256, + "reward_arguments": { + "reward_function": args.reward_function, + "n_samples_per_prompt": args.samples_per_group, + "rollout_batch_size": args.groups_per_island, + "rollout_max_response_len": args.rollout_max_response_len, + "seed": seed, + "rm_type": None, + "group_rm": False, + "reward_key": None, + "eval_reward_key": None, + }, + "seq_len": args.seq_len, + "max_response_len": args.rollout_max_response_len, + "samples_per_prompt": args.eval_samples_per_prompt, + "pass_ks": list(args._pass_ks), + "pass_threshold": args.pass_threshold, + "temperature": args.eval_temperature, + "top_p": args.eval_top_p, + "seed": args.eval_seed + seed, + "device": args.eval_device, + "trust_remote_code": args.trust_remote_code, + "miles_root": str(args.miles_root), + "result_path": str(result_path), + "samples_path": str(run_dir / "eval-samples.jsonl"), + } + config_path = run_dir / "eval-worker.json" + write_json_atomic(config_path, payload) + env = dict(os.environ) + paths = [str(args.miles_root), str(REPO_ROOT)] + if env.get("PYTHONPATH"): + paths.append(env["PYTHONPATH"]) + env["PYTHONPATH"] = os.pathsep.join(paths) + if args.eval_device == "cuda": + visible = _visible_cuda_devices() + env["CUDA_VISIBLE_DEVICES"] = (visible or ["0"])[0] + _run_checked( + [ + sys.executable, + str(Path(__file__).resolve()), + "_eval-worker", + str(config_path), + ], + run_dir / "eval.log", + env=env, + timeout_s=args.arm_timeout_min * 60, + ) + if not result_path.is_file(): + raise RuntimeError("evaluation worker produced no result") + return json.loads(result_path.read_text(encoding="utf-8")) + + +def run_arm( + args, + arm: Arm, + *, + seed: int, + model_path: Path, + prompt_paths: tuple[Path, tuple[Path, ...], Path], + streams: PromptStreams, + reward_sha256: str, +) -> dict[str, Any]: + run_dir = args.work_dir / f"seed-{seed}" / arm.name + if run_dir.exists(): + shutil.rmtree(run_dir) + run_dir.mkdir(parents=True) + combined_path, island_paths, eval_path = prompt_paths + workers = worker_specs(arm, combined_path, island_paths) + args._active_seed = seed + wait_for_free_gpus() + train_wall_s = _run_training_processes( + args, + arm, + workers, + run_dir=run_dir, + model_path=model_path, + reward_sha256=reward_sha256, + ) + + artifact_s = None + if arm.kind == "native": + native_adapter = ( + run_dir + / "native-checkpoint" + / f"iter_{args.global_rounds - 1:07d}" + / "adapter" + ) + export_started = time.monotonic() + adapter_path = standardize_native_adapter( + args, + native_adapter, + run_dir / "adapter", + ) + artifact_s = time.monotonic() - export_started + else: + checkpoint = run_dir / "state.ckpt" + adapter_path = run_dir / "adapter" + export_started = time.monotonic() + from yeto.rl.export import export_rl_checkpoint + + state = export_rl_checkpoint( + checkpoint, + adapter_path, + model=args.model, + model_revision=args.model_revision, + rank=args.lora_r, + lora_targets=args.lora_targets, + trust_remote_code=args.trust_remote_code, + ) + artifact_s = time.monotonic() - export_started + if state.policy_version != args.global_rounds: + raise RuntimeError( + f"authoritative RL checkpoint ended at v{state.policy_version}, " + f"expected v{args.global_rounds}" + ) + if not (adapter_path / "adapter_config.json").is_file() or not any( + (adapter_path / name).is_file() + for name in ("adapter_model.safetensors", "adapter_model.bin") + ): + raise RuntimeError(f"arm produced no standard PEFT adapter: {adapter_path}") + + rollout_paths = tuple( + tuple( + run_dir / "rollouts" / f"island-{worker.learner_id}" / f"{round_id}.pt" + for round_id in range(args.global_rounds) + ) + for worker in workers + ) + expected_ids = ( + streams.island_ids if arm.kind == "federated" else (streams.combined_ids,) + ) + training = summarize_rollouts( + rollout_paths, + expected_prompt_ids=expected_ids, + samples_per_group=args.samples_per_group, + ) + expected_work = workload( + arm, + rounds=args.global_rounds, + samples_per_group=args.samples_per_group, + ) + if ( + training["prompt_groups"] != expected_work["prompt_groups"] + or training["trajectories"] != expected_work["trajectories"] + ): + raise RuntimeError("Miles completed work does not match the paired budget") + sync = None if arm.kind == "native" else summarize_yeto_events(run_dir, arm.islands) + if sync is not None and sync["local_rounds"] != arm.islands * args.global_rounds: + raise RuntimeError("Yeto event tape is missing a local RL round") + + wait_for_free_gpus() + evaluation = evaluate_artifact( + args, + adapter_path=adapter_path, + model_path=model_path, + eval_path=eval_path, + reward_sha256=reward_sha256, + run_dir=run_dir, + seed=seed, + ) + total_gpus = arm.islands * arm.gpus_per_island + gpu_hours = ( + total_gpus * train_wall_s + + (evaluation["wall_s"] if args.eval_device == "cuda" else 0.0) + ) / 3600.0 + return { + "arm": arm.name, + "kind": arm.kind, + "m": arm.benchmark_islands, + "seed": seed, + "islands": arm.islands, + "gpus_per_island": arm.gpus_per_island, + "total_gpus": total_gpus, + "train_wall_s": train_wall_s, + "artifact_s": artifact_s, + "artifact_ready_s": train_wall_s + (artifact_s or 0.0), + "gpu_hours": gpu_hours, + "estimated_cost": ( + gpu_hours * args.gpu_hour_cost if args.gpu_hour_cost is not None else None + ), + "training": training, + "sync": sync, + "eval": evaluation, + } + + +def summarize_rewards( + grouped_rewards: list[list[float]], + *, + pass_ks: tuple[int, ...], + threshold: float, +) -> dict[str, Any]: + if not grouped_rewards or any(not group for group in grouped_rewards): + raise ValueError("reward groups must be non-empty") + group_size = len(grouped_rewards[0]) + if any(len(group) != group_size for group in grouped_rewards): + raise ValueError("reward groups must have one fixed sample count") + if any(k <= 0 or k > group_size for k in pass_ks): + raise ValueError("pass@k must be positive and cannot exceed samples per prompt") + values = [float(value) for group in grouped_rewards for value in group] + pass_at_k = {} + for k in pass_ks: + estimates = [] + for group in grouped_rewards: + successes = sum(value > threshold for value in group) + failures = group_size - successes + miss = ( + 0.0 + if failures < k + else math.comb(failures, k) / math.comb(group_size, k) + ) + estimates.append(1.0 - miss) + pass_at_k[str(k)] = statistics.fmean(estimates) + return { + "reward_mean": statistics.fmean(values), + "reward_std": statistics.stdev(values) if len(values) > 1 else 0.0, + "pass_at_k": pass_at_k, + } + + +def expected_record_keys( + seeds: tuple[int, ...], arms: list[Arm] +) -> set[tuple[str, int, int]]: + return {(arm.name, arm.benchmark_islands, seed) for seed in seeds for arm in arms} + + +def validate_result_records(records: list[dict], expected: set[tuple]) -> None: + seen = set() + for record in records: + key = (record.get("arm"), record.get("m"), record.get("seed")) + if key not in expected: + raise ValueError(f"results contain an unexpected record: {key}") + if key in seen: + raise ValueError(f"results contain a duplicate record: {key}") + seen.add(key) + + +def annotate_deltas(records: list[dict]) -> list[dict]: + rewards = { + (record["m"], record["seed"], record["arm"]): record["eval"]["reward_mean"] + for record in records + } + output = [] + for record in records: + m, seed = record["m"], record["seed"] + reward = record["eval"]["reward_mean"] + native = rewards.get((m, seed, f"native-miles-m{m}")) + single = rewards.get((m, seed, f"yeto-single-m{m}")) + row = dict(record) + row["delta_vs_native"] = ( + None + if record["arm"].startswith("native-") or native is None + else reward - native + ) + row["delta_vs_single"] = ( + reward - single + if record["arm"].startswith("yeto-federated-") and single is not None + else None + ) + output.append(row) + return output + + +def parse_seeds(spec: str) -> tuple[int, ...]: + try: + seeds = tuple(int(value.strip()) for value in spec.split(",") if value.strip()) + except ValueError as exc: + raise ValueError("--seeds must be a comma-separated list of integers") from exc + if not seeds: + raise ValueError("--seeds must contain at least one value") + if len(seeds) != len(set(seeds)): + raise ValueError("--seeds contains duplicates") + return seeds + + +def validate_args(args, arms: list[Arm], *, check_runtime: bool) -> None: + from yeto.models import resolve + from yeto.provenance import is_immutable_commit, is_local_reference + + validate_workload(args) + parse_seeds(args.seeds) + if is_local_reference(resolve(args.model)): + raise ValueError("RL benchmark does not accept a mutable local model") + if not is_immutable_commit(args.model_revision): + raise ValueError("--model-revision must be an immutable commit") + data_path = Path(args.data).expanduser() + if data_path.exists(): + if args.data_revision is not None: + raise ValueError("--data-revision cannot be used with a local dataset") + elif not is_immutable_commit(args.data_revision or ""): + raise ValueError("remote --data requires an immutable --data-revision") + module, separator, function = args.reward_function.partition(":") + if not separator or not module or not function.isidentifier(): + raise ValueError("--reward-function must be package.module:function") + if args.eval_prompts <= 0 or args.eval_samples_per_prompt <= 0: + raise ValueError("evaluation prompt and sample counts must be positive") + if args.seq_len <= args.rollout_max_response_len: + raise ValueError("--seq-len must exceed --rollout-max-response-len") + if not 0 <= args.eval_temperature or not 0 < args.eval_top_p <= 1: + raise ValueError( + "evaluation temperature must be non-negative and top-p in (0, 1]" + ) + if args.inner_lr <= 0 or args.lora_r <= 0 or args.wan_streams <= 0: + raise ValueError("LoRA rank, learning rate, and WAN streams must be positive") + if args.arm_timeout_min <= 0: + raise ValueError("--arm-timeout-min must be positive") + if args.gpu_hour_cost is not None and args.gpu_hour_cost < 0: + raise ValueError("--gpu-hour-cost must be non-negative") + args._pass_ks = tuple(_positive_csv(args.pass_k, "--pass-k")) + if max(args._pass_ks) > args.eval_samples_per_prompt: + raise ValueError("pass@k cannot exceed --eval-samples-per-prompt") + if args.expert_parallel is not None and ( + args.expert_parallel <= 0 or args.gpus_per_island % args.expert_parallel + ): + raise ValueError("--expert-parallel must divide --gpus-per-island") + port_range_end = ( + args.miles_port_base + + max(arm.benchmark_islands for arm in arms) * _MILES_PORT_STRIDE + - 1 + ) + if args.miles_port_base <= 0 or port_range_end > 65535: + raise ValueError("Miles host port ranges must fit in host port space") + if args.overwrite and args.resume: + raise ValueError("--overwrite and --resume are mutually exclusive") + if args.work_dir.expanduser().resolve() == args.report_dir.expanduser().resolve(): + raise ValueError("--work-dir and --report-dir must differ") + if check_runtime: + if not args.trust_remote_code: + raise ValueError("pinned Miles requires explicit --trust-remote-code") + if not args.miles_root.expanduser().is_dir(): + raise ValueError(f"Miles checkout does not exist: {args.miles_root}") + import torch + + required = max(arm.islands * arm.gpus_per_island for arm in arms) + if torch.cuda.device_count() < required: + raise ValueError( + f"largest arm needs {required} GPUs, but torch sees " + f"{torch.cuda.device_count()}" + ) + + +def materialize_prompt_matrix( + args, islands: tuple[int, ...] +) -> tuple[dict[int, tuple[tuple[Path, tuple[Path, ...], Path], PromptStreams]], int]: + from yeto.data import load_rows + + dataset = load_rows(args.data, revision=args.data_revision) + total_rows = len(dataset) + if total_rows <= args.eval_prompts: + raise ValueError( + f"prompt dataset has {total_rows} rows; need more than {args.eval_prompts}" + ) + train_rows = total_rows - args.eval_prompts + required = max(islands) * args.groups_per_island * args.global_rounds + source_train = [dict(dataset[index]) for index in range(min(train_rows, required))] + evaluation = [dict(dataset[index]) for index in range(train_rows, total_rows)] + output = {} + for m in islands: + streams = paired_prompt_streams( + source_train, + islands=m, + groups=args.groups_per_island, + rounds=args.global_rounds, + ) + paths = write_prompt_files( + streams, + evaluation, + args.work_dir / "data" / f"m{m}", + ) + output[m] = (paths, streams) + return output, train_rows + + +def load_prompt_matrix( + data_root: Path, islands: tuple[int, ...] +) -> dict[int, tuple[tuple[Path, tuple[Path, ...], Path], PromptStreams]]: + output = {} + for m in islands: + directory = data_root / f"m{m}" + combined_path = directory / "combined.jsonl" + island_paths = tuple(directory / f"island-{index}.jsonl" for index in range(m)) + eval_path = directory / "eval.jsonl" + combined_rows = tuple(_read_jsonl(combined_path)) + island_rows = tuple(tuple(_read_jsonl(path)) for path in island_paths) + + def ids(rows) -> tuple[int, ...]: + return tuple(int(row["metadata"]["benchmark_prompt_id"]) for row in rows) + + streams = PromptStreams( + combined_rows, + island_rows, + ids(combined_rows), + tuple(ids(rows) for rows in island_rows), + ) + output[m] = ((combined_path, island_paths, eval_path), streams) + return output + + +def _resume_identity(args, arms: list[Arm]) -> dict[str, Any]: + from yeto.benchmark_resume import implementation_fingerprint, jsonable_arguments + from yeto.rl import MILES_COMMIT + + return { + "format_version": 1, + "benchmark": "miles-rl-lm", + "arguments": jsonable_arguments(args, exclude=_RESUME_EXCLUDES), + "arms": [asdict(arm) for arm in arms], + "miles_commit": MILES_COMMIT, + "implementation_sha256": implementation_fingerprint( + REPO_ROOT, + _IMPLEMENTATION_PATHS, + ), + } + + +def write_run_config( + args, + arms: list[Arm], + *, + data_manifest: dict[str, Any], +) -> None: + from yeto.benchmark_resume import jsonable_arguments + from yeto.rl import MILES_COMMIT + + config = { + "format_version": 1, + "arguments": jsonable_arguments( + args, + exclude={"_active_seed", "_pass_ks"}, + ), + "arms": [asdict(arm) for arm in arms], + "miles_commit": MILES_COMMIT, + "resume_identity": _resume_identity(args, arms), + "data_manifest": data_manifest, + "fairness_contract": { + "same_model_reward_recipe_and_training_seed": True, + "same_total_gpus": True, + "same_global_rounds": True, + "same_prompt_groups_and_trajectory_budget": True, + "same_per_rank_batch": True, + "same_expert_parallel": True, + "paired_round_major_prompt_streams": True, + "paired_held_out_generation_seeds": True, + "yeto_artifact": "authoritative syncer checkpoint export", + "native_artifact": "Miles final LoRA save, PEFT-normalized by the harness", + }, + } + args.report_dir.mkdir(parents=True, exist_ok=True) + write_json_atomic(args.report_dir / "config.json", config) + + +def write_results(report_dir: Path, records: list[dict[str, Any]]) -> None: + report_dir.mkdir(parents=True, exist_ok=True) + path = report_dir / "results.jsonl" + temporary = path.with_name(path.name + ".tmp") + with temporary.open("w", encoding="utf-8") as handle: + for record in records: + handle.write(json.dumps(record, sort_keys=True) + "\n") + temporary.replace(path) + + +def load_results(report_dir: Path) -> list[dict[str, Any]]: + path = report_dir / "results.jsonl" + return _read_jsonl(path) if path.exists() else [] + + +def _mean_std(values: list[float]) -> tuple[float | None, float | None]: + if not values: + return None, None + return statistics.fmean(values), statistics.stdev(values) if len( + values + ) > 1 else 0.0 + + +def aggregate_records(records: list[dict[str, Any]]) -> list[dict[str, Any]]: + annotated = annotate_deltas(records) + grouped: dict[tuple[str, int], list[dict[str, Any]]] = {} + for record in annotated: + grouped.setdefault((record["arm"], record["m"]), []).append(record) + output = [] + for (arm, m), group in grouped.items(): + rewards = [record["eval"]["reward_mean"] for record in group] + reward_mean, reward_std = _mean_std(rewards) + pass_keys = sorted(group[0]["eval"]["pass_at_k"], key=int) + sync_records = [record["sync"] for record in group if record["sync"]] + output.append( + { + "arm": arm, + "m": m, + "runs": len(group), + "total_gpus": group[0]["total_gpus"], + "reward_mean": reward_mean, + "reward_std": reward_std, + "pass_at_k": { + key: statistics.fmean( + record["eval"]["pass_at_k"][key] for record in group + ) + for key in pass_keys + }, + "delta_vs_native": _mean( + [ + record["delta_vs_native"] + for record in group + if record["delta_vs_native"] is not None + ] + ), + "delta_vs_single": _mean( + [ + record["delta_vs_single"] + for record in group + if record["delta_vs_single"] is not None + ] + ), + "train_wall_s": statistics.fmean( + record["train_wall_s"] for record in group + ), + "artifact_s": _mean( + [ + record["artifact_s"] + for record in group + if record["artifact_s"] is not None + ] + ), + "artifact_ready_s": statistics.fmean( + record["artifact_ready_s"] for record in group + ), + "eval_wall_s": statistics.fmean( + record["eval"]["wall_s"] for record in group + ), + "trajectories_per_s": statistics.fmean( + record["training"]["trajectories"] / record["train_wall_s"] + for record in group + ), + "action_tokens_per_s": statistics.fmean( + record["training"]["action_tokens"] / record["train_wall_s"] + for record in group + ), + "gpu_hours": statistics.fmean(record["gpu_hours"] for record in group), + "estimated_cost": _mean( + [ + record["estimated_cost"] + for record in group + if record["estimated_cost"] is not None + ] + ), + "mean_kl": _mean( + [ + record["mean_kl"] + for record in sync_records + if record["mean_kl"] is not None + ] + ), + "mean_sync_ms": _mean( + [ + record["mean_sync_ms"] + for record in sync_records + if record["mean_sync_ms"] is not None + ] + ), + "sync_bytes_sent": _mean( + [float(record["sync_bytes_sent"]) for record in sync_records] + ), + } + ) + order = {"native-miles": 0, "yeto-single": 1, "yeto-federated": 2} + return sorted( + output, + key=lambda row: ( + row["m"], + next( + rank for prefix, rank in order.items() if row["arm"].startswith(prefix) + ), + ), + ) + + +def write_report(args, records: list[dict[str, Any]]) -> None: + aggregates = aggregate_records(records) + (args.report_dir / "summary.json").write_text( + json.dumps(aggregates, indent=2, sort_keys=True) + "\n", + encoding="utf-8", + ) + + def fmt(value, digits=3): + return "-" if value is None else f"{value:.{digits}f}" + + pass_keys = sorted({key for row in aggregates for key in row["pass_at_k"]}, key=int) + lines = [ + f"# Miles RL LM benchmark: {args.model}", + "", + ( + f"Rounds: {args.global_rounds}; seeds: {args.seeds}; held-out prompts: " + f"{args.eval_prompts}; samples/prompt: {args.eval_samples_per_prompt}" + ), + "", + "## Quality", + "", + "| arm | M | runs | reward | " + + " | ".join(f"pass@{key}" for key in pass_keys) + + " | delta vs native | delta vs single |", + "|---|---:|---:|---:|" + "---:|" * len(pass_keys) + "---:|---:|", + ] + for row in aggregates: + reward = f"{row['reward_mean']:.4f} +/- {row['reward_std']:.4f}" + passes = " | ".join(fmt(row["pass_at_k"].get(key), 4) for key in pass_keys) + lines.append( + f"| {row['arm']} | {row['m']} | {row['runs']} | {reward} | {passes} | " + f"{fmt(row['delta_vs_native'], 4)} | {fmt(row['delta_vs_single'], 4)} |" + ) + lines.extend( + [ + "", + "## Systems", + "", + "| arm | GPUs | train s | artifact-ready s | eval s | traj/s | action tok/s | GPU-h | cost | sync ms | sent MB | KL |", + "|---|---:|---:|---:|---:|---:|---:|---:|---:|---:|---:|---:|", + ] + ) + for row in aggregates: + sent_mb = ( + None if row["sync_bytes_sent"] is None else row["sync_bytes_sent"] / 1e6 + ) + lines.append( + f"| {row['arm']} | {row['total_gpus']} | {row['train_wall_s']:.1f} | " + f"{row['artifact_ready_s']:.1f} | {row['eval_wall_s']:.1f} | " + f"{row['trajectories_per_s']:.2f} | " + f"{row['action_tokens_per_s']:.1f} | {row['gpu_hours']:.3f} | " + f"{fmt(row['estimated_cost'], 2)} | {fmt(row['mean_sync_ms'], 2)} | " + f"{fmt(sent_mb, 3)} | {fmt(row['mean_kl'], 5)} |" + ) + report = "\n".join(lines) + "\n" + (args.report_dir / "report.md").write_text(report, encoding="utf-8") + print(report) + + +def build_parser() -> argparse.ArgumentParser: + parser = argparse.ArgumentParser( + description=__doc__, + formatter_class=argparse.RawDescriptionHelpFormatter, + ) + parser.add_argument("--model", required=True) + parser.add_argument("--model-revision", required=True) + parser.add_argument("--data", required=True) + parser.add_argument("--data-revision", default=None) + parser.add_argument("--reward-function", required=True) + parser.add_argument("--islands", default="2") + parser.add_argument("--seeds", default="17,29,43") + parser.add_argument("--global-rounds", type=int, default=8) + parser.add_argument("--groups-per-island", type=int, default=4) + parser.add_argument("--samples-per-group", type=int, default=4) + parser.add_argument("--optimizer-steps", type=int, default=1) + parser.add_argument("--gpus-per-island", type=int, default=1) + parser.add_argument("--expert-parallel", type=int, default=1) + parser.add_argument("--miles-port-base", type=int, default=21000) + parser.add_argument("--rollout-max-response-len", type=int, default=256) + parser.add_argument("--seq-len", type=int, default=1024) + parser.add_argument("--inner-lr", type=float, default=1e-4) + parser.add_argument("--lora-r", type=int, default=16) + parser.add_argument( + "--lora-targets", + choices=["auto", "attention", "all-linear"], + default="auto", + ) + parser.add_argument("--wan-streams", type=int, default=4) + parser.add_argument("--eval-prompts", type=int, default=64) + parser.add_argument("--eval-samples-per-prompt", type=int, default=None) + parser.add_argument("--pass-k", default="1,4") + parser.add_argument("--pass-threshold", type=float, default=0.0) + parser.add_argument("--eval-temperature", type=float, default=1.0) + parser.add_argument("--eval-top-p", type=float, default=1.0) + parser.add_argument("--eval-seed", type=int, default=100000) + parser.add_argument("--eval-device", choices=["cpu", "cuda"], default="cuda") + parser.add_argument("--gpu-hour-cost", type=float, default=None) + parser.add_argument("--arm-timeout-min", type=int, default=240) + parser.add_argument( + "--miles-root", + type=Path, + default=Path.home() / "miles", + ) + parser.add_argument( + "--work-dir", + type=Path, + default=REPO_ROOT / "rl-benchmark-work", + ) + parser.add_argument( + "--report-dir", + type=Path, + default=REPO_ROOT / "rl-benchmark-report", + ) + parser.add_argument("--trust-remote-code", action="store_true") + parser.add_argument("--overwrite", action="store_true") + parser.add_argument("--resume", action="store_true") + parser.add_argument("--dry-run", action="store_true") + return parser + + +def print_plan(args, arms: list[Arm]) -> None: + plan = { + "model": args.model, + "expert_parallel": args.expert_parallel, + "seeds": parse_seeds(args.seeds), + "arms": [ + { + **asdict(arm), + **workload( + arm, + rounds=args.global_rounds, + samples_per_group=args.samples_per_group, + ), + } + for arm in arms + ], + "fairness": { + "same_total_gpus": True, + "same_prompt_groups": True, + "same_trajectories": True, + "same_per_rank_batch": True, + "same_expert_parallel": True, + "paired_prompt_streams": True, + }, + } + for arm in arms: + print( + f"{arm.name}: {arm.islands} island(s) x {arm.gpus_per_island} GPU, " + f"{arm.groups_per_round} groups/round" + ) + print("PLAN_JSON " + json.dumps(plan, sort_keys=True)) + + +def main(argv=None) -> int: + raw_argv = list(sys.argv[1:] if argv is None else argv) + if raw_argv and raw_argv[0] in {"_train-worker", "_eval-worker"}: + if len(raw_argv) != 2: + raise SystemExit(f"{raw_argv[0]} requires one config path") + worker = ( + run_training_worker + if raw_argv[0] == "_train-worker" + else run_evaluation_worker + ) + return worker(Path(raw_argv[1])) + + args = build_parser().parse_args(raw_argv) + args.miles_root = args.miles_root.expanduser().resolve() + args.work_dir = args.work_dir.expanduser().resolve() + args.report_dir = args.report_dir.expanduser().resolve() + if args.eval_samples_per_prompt is None: + args.eval_samples_per_prompt = args.samples_per_group + try: + arms = select_arms( + args.islands, + args.gpus_per_island, + args.groups_per_island, + ) + validate_args(args, arms, check_runtime=not args.dry_run) + except ValueError as exc: + raise SystemExit(str(exc)) from exc + print_plan(args, arms) + if args.dry_run: + return 0 + + from yeto.benchmark_resume import ( + build_data_manifest, + load_resume_config, + validate_data_manifest, + ) + from yeto.provenance import python_spec_sha256 + from yeto.rl.miles import verify_miles_revision + + miles_root = str(args.miles_root) + if miles_root not in sys.path: + sys.path.insert(0, miles_root) + reward_sha256 = python_spec_sha256(args.reward_function, base_dir=REPO_ROOT) + args.reward_sha256 = reward_sha256 + verify_miles_revision(args.miles_root) + ensure_syncer() + model_path = resolve_model_path(args) + distinct_m = tuple(sorted({arm.benchmark_islands for arm in arms})) + + if args.resume: + if not args.work_dir.is_dir() or not args.report_dir.is_dir(): + raise SystemExit("--resume requires existing work and report directories") + try: + manifest = load_resume_config( + args.report_dir / "config.json", + _resume_identity(args, arms), + ) + data_root, _, _ = validate_data_manifest(args.work_dir, manifest) + prompt_matrix = load_prompt_matrix(data_root, distinct_m) + except ValueError as exc: + raise SystemExit(str(exc)) from exc + else: + for path in (args.work_dir, args.report_dir): + if path.exists(): + if not args.overwrite: + raise SystemExit( + f"{path} already exists; pass --overwrite to replace it" + ) + shutil.rmtree(path) + args.work_dir.mkdir(parents=True) + args.report_dir.mkdir(parents=True) + try: + prompt_matrix, train_rows = materialize_prompt_matrix(args, distinct_m) + except (RuntimeError, ValueError) as exc: + raise SystemExit(str(exc)) from exc + first_eval = prompt_matrix[distinct_m[0]][0][2] + source = args.data if Path(args.data).expanduser().exists() else None + data_manifest = build_data_manifest( + args.work_dir, + args.work_dir / "data", + first_eval, + train_rows=train_rows, + eval_rows=args.eval_prompts, + source=source, + ) + write_run_config(args, arms, data_manifest=data_manifest) + + records = load_results(args.report_dir) if args.resume else [] + expected = expected_record_keys(parse_seeds(args.seeds), arms) + try: + validate_result_records(records, expected) + except ValueError as exc: + raise SystemExit(str(exc)) from exc + completed = {(record["arm"], record["m"], record["seed"]) for record in records} + for seed in parse_seeds(args.seeds): + for arm in arms: + key = (arm.name, arm.benchmark_islands, seed) + if key in completed: + print(f"[rl-benchmark] resume: skipping {arm.name} seed={seed}") + continue + print(f"[rl-benchmark] running {arm.name} seed={seed}", flush=True) + paths, streams = prompt_matrix[arm.benchmark_islands] + record = run_arm( + args, + arm, + seed=seed, + model_path=model_path, + prompt_paths=paths, + streams=streams, + reward_sha256=reward_sha256, + ) + records.append(record) + completed.add(key) + write_results(args.report_dir, records) + + write_report(args, records) + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/tests/test_rl_benchmark.py b/tests/test_rl_benchmark.py new file mode 100644 index 0000000..d2d70b9 --- /dev/null +++ b/tests/test_rl_benchmark.py @@ -0,0 +1,921 @@ +"""Pure-logic tests for the Miles RL benchmark harness.""" + +from __future__ import annotations + +import asyncio +import importlib.util +import json +import os +import subprocess +import sys +import warnings +from pathlib import Path +from types import SimpleNamespace + +import numpy as np +import pytest +import torch + +ROOT = Path(__file__).resolve().parent.parent +SPEC = importlib.util.spec_from_file_location( + "benchmark_miles_rl", ROOT / "scripts" / "benchmark_rl.py" +) +benchmark = importlib.util.module_from_spec(SPEC) +sys.modules[SPEC.name] = benchmark +SPEC.loader.exec_module(benchmark) + + +def _args(**overrides): + values = { + "global_rounds": 3, + "groups_per_island": 4, + "samples_per_group": 2, + "optimizer_steps": 1, + "gpus_per_island": 2, + } + values.update(overrides) + return SimpleNamespace(**values) + + +def test_three_arms_have_equal_hardware_and_trajectory_budgets(): + arms = benchmark.select_arms("2", gpus_per_island=2, groups_per_island=4) + + assert [arm.name for arm in arms] == [ + "native-miles-m2", + "yeto-single-m2", + "yeto-federated-m2", + ] + assert [ + (arm.islands, arm.gpus_per_island, arm.groups_per_round) for arm in arms + ] == [ + (1, 4, 8), + (1, 4, 8), + (2, 2, 4), + ] + + budgets = [benchmark.workload(arm, rounds=3, samples_per_group=2) for arm in arms] + assert {budget["total_gpus"] for budget in budgets} == {4} + assert {budget["prompt_groups"] for budget in budgets} == {24} + assert {budget["trajectories"] for budget in budgets} == {48} + assert {budget["groups_per_gpu_per_round"] for budget in budgets} == {2.0} + + +def test_arm_selection_rejects_duplicates_and_nonpositive_values(): + with pytest.raises(ValueError, match="at least one"): + benchmark.select_arms("", 1, 1) + with pytest.raises(ValueError, match="duplicates"): + benchmark.select_arms("2,2", 1, 1) + with pytest.raises(ValueError, match="positive"): + benchmark.select_arms("0", 1, 1) + + +def test_workload_validation_requires_equal_per_rank_batch(): + benchmark.validate_workload(_args()) + with pytest.raises(ValueError, match="divisible"): + benchmark.validate_workload(_args(groups_per_island=3, samples_per_group=3)) + + +def test_local_ray_cluster_uses_rays_short_default_temp_root(monkeypatch): + calls = {} + + def init(**kwargs): + calls["init"] = kwargs + return SimpleNamespace(address_info={"address": "local"}) + + ray = SimpleNamespace( + is_initialized=lambda: False, + init=init, + shutdown=lambda: calls.setdefault("shutdown", True), + ) + monkeypatch.delenv("RAY_ADDRESS", raising=False) + monkeypatch.setitem(sys.modules, "ray", ray) + + with benchmark.local_ray_cluster(8) as address: + assert address == "local" + + assert "_temp_dir" not in calls["init"] + assert calls["init"]["include_dashboard"] is True + assert calls["shutdown"] is True + + +def test_expert_parallel_default_is_fixed_across_all_arms(): + args = benchmark.build_parser().parse_args( + [ + "--model", + "org/model", + "--model-revision", + "a" * 40, + "--data", + "org/data", + "--data-revision", + "b" * 40, + "--reward-function", + "pkg.reward:score", + ] + ) + + assert args.expert_parallel == 1 + + +def test_prompt_streams_are_round_major_and_exactly_paired(): + rows = [ + { + "messages": [{"role": "user", "content": f"prompt {index}"}], + "label": str(index), + } + for index in range(8) + ] + + streams = benchmark.paired_prompt_streams(rows, islands=2, groups=2, rounds=2) + + assert streams.combined_ids == tuple(range(8)) + assert streams.island_ids == ((0, 1, 4, 5), (2, 3, 6, 7)) + for round_id in range(2): + combined = streams.combined_ids[round_id * 4 : (round_id + 1) * 4] + federated = tuple( + prompt_id + for island in streams.island_ids + for prompt_id in island[round_id * 2 : (round_id + 1) * 2] + ) + assert federated == combined + + +def test_prompt_streams_cycle_deterministically_without_touching_eval_rows(): + rows = [ + {"messages": [{"role": "user", "content": value}]} for value in ("a", "b", "c") + ] + streams = benchmark.paired_prompt_streams(rows, islands=2, groups=2, rounds=2) + assert streams.combined_ids == (0, 1, 2, 0, 1, 2, 0, 1) + + +def test_worker_specs_keep_native_outside_yeto_and_partition_federated_prompts( + tmp_path, +): + native, single, federated = benchmark.select_arms("2", 2, 4) + combined = tmp_path / "combined.jsonl" + islands = (tmp_path / "island-0.jsonl", tmp_path / "island-1.jsonl") + + native_specs = benchmark.worker_specs(native, combined, islands) + single_specs = benchmark.worker_specs(single, combined, islands) + federated_specs = benchmark.worker_specs(federated, combined, islands) + + assert len(native_specs) == len(single_specs) == 1 + assert native_specs[0].policy_sync is False + assert single_specs[0].policy_sync is True + assert native_specs[0].gpus == single_specs[0].gpus == 4 + assert native_specs[0].groups_per_round == single_specs[0].groups_per_round == 8 + assert [spec.prompt_path for spec in federated_specs] == list(islands) + assert [spec.learner_id for spec in federated_specs] == [0, 1] + assert all(spec.policy_sync for spec in federated_specs) + assert all( + spec.gpus == 2 and spec.groups_per_round == 4 for spec in federated_specs + ) + + +def test_federated_workers_use_disjoint_miles_host_ports(tmp_path): + args = benchmark.build_parser().parse_args( + [ + "--model", + "org/model", + "--model-revision", + "a" * 40, + "--data", + "org/data", + "--data-revision", + "b" * 40, + "--reward-function", + "pkg.reward:score", + ] + ) + args._active_seed = 17 + arm = benchmark.select_arms("2", 2, 4)[2] + combined = tmp_path / "combined.jsonl" + island_paths = (tmp_path / "island-0.jsonl", tmp_path / "island-1.jsonl") + for path in island_paths: + path.write_text("{}\n", encoding="utf-8") + workers = benchmark.worker_specs(arm, combined, island_paths) + + payloads = [ + benchmark.worker_payload( + args, + worker, + run_dir=tmp_path, + model_path=tmp_path / "model", + syncer="127.0.0.1:30000", + reward_sha256="c" * 64, + ) + for worker in workers + ] + + assert [ + payload["arguments"]["rollout_engine_base_port"] for payload in payloads + ] == [21100, 22100] + assert [payload["arguments"]["sglang_router_port"] for payload in payloads] == [ + 21000, + 22000, + ] + assert [ + payload["arguments"]["sglang_router_prometheus_port"] for payload in payloads + ] == [21001, 22001] + assert [payload["arguments"]["train_master_base_port"] for payload in payloads] == [ + 21002, + 22002, + ] + assert {payload["arguments"]["expert_parallel"] for payload in payloads} == {1} + + +def test_reward_summary_uses_standard_pass_at_k_estimator(): + result = benchmark.summarize_rewards( + [[1.0, 0.0, 0.0, 0.0], [0.0, 0.0, 0.0, 0.0]], + pass_ks=(1, 4), + threshold=0.0, + ) + + assert result["reward_mean"] == pytest.approx(0.125) + assert result["pass_at_k"] == {"1": pytest.approx(0.125), "4": 0.5} + with pytest.raises(ValueError, match="cannot exceed"): + benchmark.summarize_rewards([[1.0, 0.0]], pass_ks=(3,), threshold=0.0) + + +def test_resume_keys_cover_every_seed_m_and_arm(): + arms = benchmark.select_arms("2,4", 1, 2) + keys = benchmark.expected_record_keys((17, 29), arms) + + assert len(keys) == 12 + assert ("native-miles-m2", 2, 17) in keys + assert ("yeto-federated-m4", 4, 29) in keys + benchmark.validate_result_records( + [ + {"arm": name, "m": m, "seed": seed} + for name, m, seed in sorted(keys, key=str) + ], + keys, + ) + with pytest.raises(ValueError, match="duplicate"): + benchmark.validate_result_records( + [ + {"arm": "native-miles-m2", "m": 2, "seed": 17}, + {"arm": "native-miles-m2", "m": 2, "seed": 17}, + ], + keys, + ) + + +def test_resume_identity_survives_json_round_trip(): + args = benchmark.build_parser().parse_args( + [ + "--model", + "org/model", + "--model-revision", + "a" * 40, + "--data", + "org/data", + "--data-revision", + "b" * 40, + "--reward-function", + "pkg.reward:score", + ] + ) + args.eval_samples_per_prompt = args.samples_per_group + args.reward_sha256 = "c" * 64 + arms = benchmark.select_arms( + args.islands, args.gpus_per_island, args.groups_per_island + ) + benchmark.validate_args(args, arms, check_runtime=False) + + identity = benchmark._resume_identity(args, arms) + + assert json.loads(json.dumps(identity)) == identity + assert identity["arguments"]["reward_sha256"] == "c" * 64 + + +def test_resume_identity_fingerprints_all_yeto_sources_and_syncer_binary( + monkeypatch, +): + args = benchmark.build_parser().parse_args( + [ + "--model", + "org/model", + "--model-revision", + "a" * 40, + "--data", + "org/data", + "--data-revision", + "b" * 40, + "--reward-function", + "pkg.reward:score", + ] + ) + args.eval_samples_per_prompt = args.samples_per_group + args.reward_sha256 = "c" * 64 + arms = benchmark.select_arms( + args.islands, + args.gpus_per_island, + args.groups_per_island, + ) + observed = {} + + def fingerprint(root, paths): + observed["root"] = root + observed["paths"] = tuple(paths) + return "d" * 64 + + monkeypatch.setattr( + "yeto.benchmark_resume.implementation_fingerprint", + fingerprint, + ) + + identity = benchmark._resume_identity(args, arms) + + assert identity["implementation_sha256"] == "d" * 64 + assert benchmark.REPO_ROOT / "yeto" in observed["paths"] + assert benchmark.SYNCER_BIN in observed["paths"] + + +def test_report_deltas_distinguish_yeto_contract_from_federation(): + records = [ + {"arm": "native-miles-m2", "m": 2, "seed": 17, "eval": {"reward_mean": 0.5}}, + {"arm": "yeto-single-m2", "m": 2, "seed": 17, "eval": {"reward_mean": 0.4}}, + {"arm": "yeto-federated-m2", "m": 2, "seed": 17, "eval": {"reward_mean": 0.3}}, + ] + + annotated = {row["arm"]: row for row in benchmark.annotate_deltas(records)} + + assert annotated["native-miles-m2"]["delta_vs_native"] is None + assert annotated["yeto-single-m2"]["delta_vs_native"] == pytest.approx(-0.1) + assert annotated["yeto-federated-m2"]["delta_vs_native"] == pytest.approx(-0.2) + assert annotated["yeto-federated-m2"]["delta_vs_single"] == pytest.approx(-0.1) + + +def test_aggregate_report_keeps_the_three_comparisons_in_order(tmp_path): + records = [] + for arm, kind, reward, sync in ( + ("native-miles-m2", "native", 0.5, None), + ( + "yeto-single-m2", + "single", + 0.4, + {"mean_kl": 0.1, "mean_sync_ms": 2.0, "sync_bytes_sent": 100}, + ), + ( + "yeto-federated-m2", + "federated", + 0.3, + {"mean_kl": 0.2, "mean_sync_ms": 3.0, "sync_bytes_sent": 200}, + ), + ): + records.append( + { + "arm": arm, + "kind": kind, + "m": 2, + "seed": 17, + "total_gpus": 4, + "train_wall_s": 10.0, + "artifact_s": None if kind == "native" else 2.0, + "artifact_ready_s": 10.0 if kind == "native" else 12.0, + "gpu_hours": 1.0, + "estimated_cost": None, + "training": {"trajectories": 8, "action_tokens": 32}, + "sync": sync, + "eval": { + "reward_mean": reward, + "pass_at_k": {"1": reward}, + "wall_s": 1.0, + }, + } + ) + + aggregates = benchmark.aggregate_records(records) + + assert [row["arm"] for row in aggregates] == [ + "native-miles-m2", + "yeto-single-m2", + "yeto-federated-m2", + ] + assert aggregates[1]["delta_vs_native"] == pytest.approx(-0.1) + assert aggregates[2]["delta_vs_single"] == pytest.approx(-0.1) + assert aggregates[0]["artifact_ready_s"] == 10.0 + assert aggregates[1]["artifact_ready_s"] == 12.0 + + args = SimpleNamespace( + model="org/model", + global_rounds=1, + seeds="17", + eval_prompts=1, + eval_samples_per_prompt=1, + report_dir=tmp_path, + ) + benchmark.write_report(args, records) + report = (tmp_path / "report.md").read_text() + assert "native-miles-m2" in report + assert "artifact-ready s" in report + + +def test_dry_run_does_not_import_ray_or_materialize_data(monkeypatch, capsys): + monkeypatch.setitem(sys.modules, "ray", None) + result = benchmark.main( + [ + "--model", + "org/model", + "--model-revision", + "a" * 40, + "--data", + "org/data", + "--data-revision", + "b" * 40, + "--reward-function", + "pkg.reward:score", + "--islands", + "2", + "--dry-run", + ] + ) + + assert result == 0 + output = capsys.readouterr().out + assert "native-miles-m2" in output + assert "yeto-single-m2" in output + assert "yeto-federated-m2" in output + plan = json.loads(output.split("PLAN_JSON ", 1)[1]) + assert plan["fairness"]["same_total_gpus"] + assert plan["fairness"]["same_expert_parallel"] + assert plan["expert_parallel"] == 1 + + +def test_dry_run_works_when_script_is_invoked_outside_repo(tmp_path): + command = [ + sys.executable, + str(ROOT / "scripts" / "benchmark_rl.py"), + "--model", + "org/model", + "--model-revision", + "a" * 40, + "--data", + "org/data", + "--data-revision", + "b" * 40, + "--reward-function", + "pkg.reward:score", + "--dry-run", + ] + env = dict(os.environ) + env.pop("PYTHONPATH", None) + + result = subprocess.run( + command, + cwd=tmp_path, + env=env, + capture_output=True, + text=True, + check=False, + ) + + assert result.returncode == 0, result.stderr + assert "PLAN_JSON" in result.stdout + + +def test_benchmark_rejects_mutable_local_model_directory(tmp_path): + model = tmp_path / "model" + model.mkdir() + args = benchmark.build_parser().parse_args( + [ + "--model", + str(model), + "--model-revision", + "a" * 40, + "--data", + "org/data", + "--data-revision", + "b" * 40, + "--reward-function", + "pkg.reward:score", + ] + ) + args.eval_samples_per_prompt = args.samples_per_group + arms = benchmark.select_arms( + args.islands, + args.gpus_per_island, + args.groups_per_island, + ) + + with pytest.raises(ValueError, match="local model"): + benchmark.validate_args(args, arms, check_runtime=False) + + +def test_benchmark_rejects_rollout_port_ranges_outside_host_port_space(): + args = benchmark.build_parser().parse_args( + [ + "--model", + "org/model", + "--model-revision", + "a" * 40, + "--data", + "org/data", + "--data-revision", + "b" * 40, + "--reward-function", + "pkg.reward:score", + "--islands", + "2", + "--miles-port-base", + "65000", + ] + ) + args.eval_samples_per_prompt = args.samples_per_group + arms = benchmark.select_arms( + args.islands, + args.gpus_per_island, + args.groups_per_island, + ) + + with pytest.raises(ValueError, match="Miles host port ranges"): + benchmark.validate_args(args, arms, check_runtime=False) + + +def test_materialized_prompt_files_preserve_holdout_and_prompt_identity(tmp_path): + rows = [ + { + "messages": [{"role": "user", "content": f"prompt {index}"}], + "label": str(index), + "metadata": {"source_id": index}, + } + for index in range(6) + ] + train, evaluation = rows[:-2], rows[-2:] + streams = benchmark.paired_prompt_streams(train, islands=2, groups=2, rounds=1) + combined, islands, evaluation_path = benchmark.write_prompt_files( + streams, + evaluation, + tmp_path, + ) + + combined_rows = [json.loads(line) for line in combined.read_text().splitlines()] + eval_rows = [json.loads(line) for line in evaluation_path.read_text().splitlines()] + assert [row["metadata"]["benchmark_prompt_id"] for row in combined_rows] == [ + 0, + 1, + 2, + 3, + ] + assert [row["label"] for row in eval_rows] == ["4", "5"] + assert [row["metadata"]["source_id"] for row in eval_rows] == [4, 5] + assert len(islands) == 2 + + +def test_strict_syncer_command_is_one_fragment_exact_base_avg(tmp_path): + arm = benchmark.select_arms("2", 2, 4)[2] + command = benchmark.syncer_command(arm, 29400, tmp_path, rounds=3) + + def value(flag): + return command[command.index(flag) + 1] + + assert value("--learners") == "2" + assert value("--quorum") == "2" + assert value("--total-steps") == "3" + assert value("--pipeline") == "1" + assert value("--sync-interval-steps") == "0" + assert value("--delta-correction") == "none" + assert value("--outer-lr") == "1" + assert value("--outer-momentum") == "0" + assert value("--max-base-lag") == "0" + assert value("--learner-weight") == "equal" + assert "--mark-final-checkpoint" not in command + + +def test_worker_miles_extras_capture_real_rollouts_and_only_native_saves(tmp_path): + native, single, _ = benchmark.select_arms("2", 2, 4) + native_worker = benchmark.worker_specs(native, tmp_path / "all", ())[0] + single_worker = benchmark.worker_specs(single, tmp_path / "all", ())[0] + + native_extra = benchmark.miles_extra_argv(native_worker, tmp_path / "native", 3) + single_extra = benchmark.miles_extra_argv(single_worker, tmp_path / "single", 3) + + assert "--save-debug-rollout-data" in native_extra + assert "--save-debug-rollout-data" in single_extra + assert "--save" in native_extra + assert native_extra[native_extra.index("--save-interval") + 1] == "3" + assert "--save-hf" not in native_extra + assert "--save-hf" not in single_extra + + +def test_native_miles_adapter_names_are_mapped_to_the_exact_peft_contract(): + specs = ( + SimpleNamespace( + name="base_model.model.model.layers.0.q_proj.lora_A.weight", + shape=(2, 4), + ), + SimpleNamespace( + name="base_model.model.model.layers.0.q_proj.lora_B.weight", + shape=(4, 2), + ), + ) + raw = { + "model.layers.0.q_proj.lora_A.weight": torch.ones(2, 4), + "model.layers.0.q_proj.lora_B.weight": torch.ones(4, 2), + } + + mapped = benchmark.canonical_native_adapter_tensors(raw, specs) + + assert tuple(mapped) == tuple(spec.name for spec in specs) + assert all( + torch.equal(tensor, torch.ones(spec.shape)) + for tensor, spec in zip(mapped.values(), specs) + ) + + raw.pop("model.layers.0.q_proj.lora_B.weight") + with pytest.raises(RuntimeError, match="does not match the PEFT contract"): + benchmark.canonical_native_adapter_tensors(raw, specs) + + +def test_native_miles_adapter_is_rewritten_as_standard_peft(tmp_path, monkeypatch): + from yeto.rl import export as rl_export + from yeto.rl.core import CanonicalTensorSpec + + specs = ( + CanonicalTensorSpec( + "base_model.model.model.layers.0.q_proj.lora_A.weight", + (2, 4), + "float32", + 8, + ), + CanonicalTensorSpec( + "base_model.model.model.layers.0.q_proj.lora_B.weight", + (4, 2), + "float32", + 8, + ), + ) + source = tmp_path / "miles" + source.mkdir() + torch.save( + { + "model.layers.0.q_proj.lora_A.weight": torch.ones(2, 4), + "model.layers.0.q_proj.lora_B.weight": torch.ones(4, 2), + }, + source / "adapter_model.bin", + ) + written = {} + monkeypatch.setattr( + rl_export, "derive_peft_lora_specs", lambda *args, **kwargs: specs + ) + monkeypatch.setattr( + rl_export, + "write_peft_adapter", + lambda state, output, **kwargs: written.update( + state=state, + output=output, + kwargs=kwargs, + ), + ) + args = SimpleNamespace( + model="org/model", + model_revision="a" * 40, + lora_r=2, + lora_targets="attention", + global_rounds=3, + trust_remote_code=True, + ) + + output = benchmark.standardize_native_adapter(args, source, tmp_path / "adapter") + + assert output == tmp_path / "adapter" + assert written["state"].policy_version == 3 + assert written["state"].specs == specs + assert written["kwargs"] == { + "base_model": "org/model", + "model_revision": "a" * 40, + "rank": 2, + } + + +def test_rollout_summary_verifies_prompt_pairing_and_counts_real_work(tmp_path): + paths = [] + for island_id, prompt_ids in enumerate(((0, 1), (2, 3))): + path = tmp_path / f"island-{island_id}.pt" + samples = [] + for prompt_id in prompt_ids: + for sample_id in range(2): + samples.append( + { + "metadata": {"benchmark_prompt_id": prompt_id}, + "reward": float((prompt_id + sample_id) % 2), + "response_length": prompt_id + sample_id + 1, + "rollout_routed_experts": np.array([prompt_id]), + "status": "completed", + } + ) + torch.save({"rollout_id": 0, "samples": samples}, path) + paths.append((path,)) + + summary = benchmark.summarize_rollouts( + tuple(paths), + expected_prompt_ids=((0, 1), (2, 3)), + samples_per_group=2, + ) + + assert summary["prompt_groups"] == 4 + assert summary["trajectories"] == 8 + assert summary["action_tokens"] == 24 + assert summary["reward_mean"] == 0.5 + assert summary["truncated_trajectories"] == 0 + + payload = torch.load(paths[1][0], weights_only=False) + for sample in payload["samples"][:2]: + sample["metadata"]["benchmark_prompt_id"] = 99 + torch.save(payload, paths[1][0]) + with pytest.raises(RuntimeError, match="prompt stream mismatch"): + benchmark.summarize_rollouts( + tuple(paths), + expected_prompt_ids=((0, 1), (2, 3)), + samples_per_group=2, + ) + + payload = torch.load(paths[1][0], weights_only=False) + for sample in payload["samples"][:2]: + sample["metadata"]["benchmark_prompt_id"] = 2 + payload["samples"][0]["status"] = "aborted" + torch.save(payload, paths[1][0]) + with pytest.raises(RuntimeError, match="status"): + benchmark.summarize_rollouts( + tuple(paths), + expected_prompt_ids=((0, 1), (2, 3)), + samples_per_group=2, + ) + + +def test_evaluation_prompt_matches_miles_rendering_and_truncates_all_token_fields( + monkeypatch, +): + calls = {} + + def render(messages, *, tokenizer, tools, tokenize, add_generation_prompt): + calls.update( + messages=messages, + tokenizer=tokenizer, + tools=tools, + tokenize=tokenize, + add_generation_prompt=add_generation_prompt, + ) + return "rendered prompt" + + monkeypatch.setitem( + sys.modules, + "miles.utils.chat_template_utils", + SimpleNamespace(apply_chat_template=render), + ) + + class Tokenizer: + def __call__(self, text, *, add_special_tokens, return_tensors): + assert text == "rendered prompt" + assert add_special_tokens is False + assert return_tensors == "pt" + return { + "input_ids": torch.tensor([[1, 2, 3, 4]]), + "attention_mask": torch.tensor([[1, 1, 1, 1]]), + "token_type_ids": torch.tensor([[0, 0, 1, 1]]), + } + + tokenizer = Tokenizer() + row = { + "messages": [{"role": "user", "content": "question"}], + "tools": [{"type": "function", "function": {"name": "lookup"}}], + } + + prompt, encoded = benchmark.prepare_evaluation_prompt( + tokenizer, + row, + max_prompt_tokens=3, + device="cpu", + ) + + assert prompt == "rendered prompt" + assert encoded["input_ids"].tolist() == [[2, 3, 4]] + assert encoded["attention_mask"].tolist() == [[1, 1, 1]] + assert encoded["token_type_ids"].tolist() == [[0, 1, 1]] + assert calls == { + "messages": row["messages"], + "tokenizer": tokenizer, + "tools": row["tools"], + "tokenize": False, + "add_generation_prompt": True, + } + + +def test_evaluation_rewards_use_miles_single_sample_contract(monkeypatch): + observed = [] + + async def score(_args, sample): + observed.append(sample) + return float(sample) + + monkeypatch.setitem( + sys.modules, + "miles.rollout.rm_hub", + SimpleNamespace(async_rm=score), + ) + + rewards = asyncio.run(benchmark.evaluate_rewards(SimpleNamespace(), [1, 2, 3])) + + assert rewards == [1.0, 2.0, 3.0] + assert observed == [1, 2, 3] + + +def test_evaluation_rejects_missing_peft_adapter_keys(): + class PeftModel: + @classmethod + def from_pretrained(cls, model, adapter_path): + warnings.warn( + "Found missing adapter keys while loading the checkpoint: ['lora_A']." + ) + return model + + with pytest.raises(RuntimeError, match="missing adapter keys"): + benchmark.load_peft_adapter(PeftModel, object(), "adapter") + + +def test_generation_pad_token_keeps_valid_zero_id(): + assert ( + benchmark.generation_pad_token_id( + SimpleNamespace(pad_token_id=0, eos_token_id=2) + ) + == 0 + ) + assert ( + benchmark.generation_pad_token_id( + SimpleNamespace(pad_token_id=None, eos_token_id=[2, 3]) + ) + == 2 + ) + + +def test_worker_input_verification_rejects_prompt_or_reward_drift( + tmp_path, + monkeypatch, +): + prompt = tmp_path / "prompts.jsonl" + prompt.write_text("{}\n", encoding="utf-8") + expected_prompt = benchmark.file_sha256(prompt) + monkeypatch.setattr( + "yeto.provenance.python_spec_sha256", + lambda spec, base_dir=None: "a" * 64, + ) + + benchmark.verify_worker_inputs( + prompt_path=prompt, + prompt_sha256=expected_prompt, + reward_function="pkg.reward:score", + reward_sha256="a" * 64, + ) + prompt.write_text('{"changed":true}\n', encoding="utf-8") + with pytest.raises(RuntimeError, match="prompt source SHA256 mismatch"): + benchmark.verify_worker_inputs( + prompt_path=prompt, + prompt_sha256=expected_prompt, + reward_function="pkg.reward:score", + reward_sha256="a" * 64, + ) + + with pytest.raises(RuntimeError, match="reward source SHA256 mismatch"): + benchmark.verify_worker_inputs( + prompt_path=prompt, + prompt_sha256=benchmark.file_sha256(prompt), + reward_function="pkg.reward:score", + reward_sha256="b" * 64, + ) + + +def test_syncer_is_always_built_from_the_locked_source(monkeypatch): + calls = [] + + def run(command, **kwargs): + calls.append((command, kwargs)) + return SimpleNamespace(returncode=0) + + monkeypatch.setattr(benchmark.subprocess, "run", run) + + benchmark.ensure_syncer() + + assert calls == [ + ( + ["cargo", "build", "--release", "--locked", "--quiet"], + {"cwd": benchmark.REPO_ROOT / "syncer", "check": True}, + ) + ] + + +def test_gpu_drain_check_ignores_compute_apps_on_hidden_devices(monkeypatch): + monkeypatch.setenv("CUDA_VISIBLE_DEVICES", "0") + + def run(command, **_kwargs): + query = next(value for value in command if value.startswith("--query-")) + if query.startswith("--query-gpu"): + stdout = "0, GPU-visible\n1, GPU-hidden\n" + else: + stdout = "GPU-hidden, 123, python, 30000\n" + return SimpleNamespace(returncode=0, stdout=stdout, stderr="") + + monkeypatch.setattr(benchmark.subprocess, "run", run) + + assert benchmark._visible_gpu_uuids() == {"GPU-visible"} + benchmark.wait_for_free_gpus(timeout_s=0) diff --git a/tests/test_rl_export.py b/tests/test_rl_export.py index 678da82..7b0e212 100644 --- a/tests/test_rl_export.py +++ b/tests/test_rl_export.py @@ -6,6 +6,7 @@ import torch from yeto.export import CKPT_MAGIC +from yeto.rl import export as rl_export from yeto.rl.core import ( canonical_layout_hash, canonical_lora_config_hash, @@ -31,6 +32,58 @@ def _model(tmp_path): return path, config +def test_rl_model_factory_preserves_the_checkpoint_architecture(monkeypatch): + transformers = pytest.importorskip("transformers") + + class DeclaredConditionalGeneration: + @classmethod + def from_config(cls, config, **kwargs): + return cls() + + monkeypatch.setattr( + transformers, + "DeclaredConditionalGeneration", + DeclaredConditionalGeneration, + raising=False, + ) + + config = type( + "Config", + (), + {"architectures": ["DeclaredConditionalGeneration"]}, + )() + assert rl_export._rl_model_factory(config) is DeclaredConditionalGeneration + + config.architectures = None + assert rl_export._rl_model_factory(config) is transformers.AutoModelForCausalLM + + +def test_declared_rl_architecture_does_not_receive_auto_factory_kwargs(monkeypatch): + transformers = pytest.importorskip("transformers") + + class DeclaredConditionalGeneration: + @classmethod + def _from_config(cls, config): + return cls() + + monkeypatch.setattr( + transformers, + "DeclaredConditionalGeneration", + DeclaredConditionalGeneration, + raising=False, + ) + config = type( + "Config", + (), + {"architectures": ["DeclaredConditionalGeneration"]}, + )() + + assert isinstance( + rl_export._rl_model_from_config(config, trust_remote_code=True), + DeclaredConditionalGeneration, + ) + + def test_attention_regex_is_resolved_before_peft_moe_conversion(tmp_path): transformers = pytest.importorskip("transformers") config = transformers.OlmoeConfig( diff --git a/tests/test_rl_integration.py b/tests/test_rl_integration.py index 566182a..6071daa 100644 --- a/tests/test_rl_integration.py +++ b/tests/test_rl_integration.py @@ -346,6 +346,76 @@ def test_single_island_runs_the_real_syncer_parity_path(syncer_binary, tmp_path) process.wait() +def test_terminal_replacement_receives_final_policy(syncer_binary, tmp_path): + layout = _layout() + checkpoint_path = tmp_path / "state.ckpt" + port = _port() + process = _start( + syncer_binary, + port, + checkpoint_path, + rounds=1, + learners=1, + ) + original = _client(port, 0, layout) + replacement = None + thread = None + try: + original.send_init(0, pack_tensor(torch.zeros(2), DTYPE_F32)) + assert _wait_item(original.drain_updates).version == 0 + assert _wait_item(original.drain_pulls).global_step == 1 + _push(original, 1, 0, [1, 3]) + manifest, _ = original.wait_for_final_fragments(timeout=10) + assert manifest.global_step == 1 + original.close() + + runtime = _FakeMiles([0, 0], 0) + replacement = StrictRlBridge( + runtime, + BridgeConfig( + syncer_addr=("127.0.0.1", port), + learner_id=0, + global_rounds=1, + groups_per_round=1, + samples_per_group=2, + local_optimizer_steps=1, + wan_streams=0, + expected_specs=runtime.current.specs, + base_model_revision=runtime.current.base_model_revision, + lora_config_hash=runtime.current.lora_config_hash, + layout_hash=runtime.current.layout_hash, + event_tape=str(tmp_path / "replacement.jsonl"), + ), + ) + replacement.start() + result = {} + + def receive(): + result["state"] = replacement.wait_for_initial_policy() + + thread = threading.Thread(target=receive, daemon=True) + thread.start() + thread.join(timeout=3) + assert not thread.is_alive(), "replacement ignored the terminal policy" + assert result["state"].policy_version == 1 + assert process.poll() is None + + final = replacement.finalize() + assert final.policy_version == 1 + assert torch.equal( + next(iter(final.tensors.values())), + torch.tensor([[1.0, 3.0]]), + ) + assert process.wait(timeout=10) == 0 + finally: + original.close() + if replacement is not None: + replacement.client.close() + if process.poll() is None: + process.kill() + process.wait() + + def test_miles_public_hook_runs_against_real_syncer( syncer_binary, tmp_path, monkeypatch ): diff --git a/tests/test_rl_launcher.py b/tests/test_rl_launcher.py index 66d5661..d7e0a22 100644 --- a/tests/test_rl_launcher.py +++ b/tests/test_rl_launcher.py @@ -7,7 +7,7 @@ import pytest -import yeto.launcher as launcher +from yeto import launcher from yeto.cli import parse_args from yeto.launcher import ( FleetController, @@ -15,8 +15,8 @@ make_miles_island_task, syncer_command, ) -from yeto.rl import MILES_PEFT_VERSION -from yeto.rl import MILES_COMMIT, MILES_REPOSITORY +from yeto.rl import MILES_COMMIT, MILES_PEFT_VERSION, MILES_REPOSITORY +from yeto.rl import learner as rl_learner from yeto.rl.learner import build_miles_argv from yeto.rl.miles import verify_miles_revision @@ -550,6 +550,7 @@ def test_miles_argv_uses_provider_capabilities_without_model_family_branches(): assert "--no-offload-train" in argv assert argv[argv.index("--sglang-mem-fraction-static") + 1] == "0.4" assert "--sglang-enable-deterministic-inference" in argv + assert argv[argv.index("--rollout-seed") + 1] == "1" assert "--pin-rollout-manager-to-head" in argv for recipe_flag in ( "--rollout-shuffle", @@ -572,6 +573,48 @@ def test_miles_argv_uses_provider_capabilities_without_model_family_branches(): ): assert recipe_flag not in argv + native_argv = build_miles_argv( + args, + model_path="/model", + prompt_path="/prompts.jsonl", + provider=provider, + target_modules=["qkv_proj", "out_proj"], + yeto_policy_sync=False, + ) + assert "--external-policy-sync-path" not in native_argv + assert "--rollout-all-samples-process-path" not in native_argv + assert native_argv[native_argv.index("--rollout-function-path") + 1] == ( + "miles.rollout.sglang_rollout.generate_rollout" + ) + assert not any(value.startswith("yeto.rl.") for value in native_argv) + + paired_args = argparse.Namespace( + **vars(args), + rollout_seed=91, + rollout_engine_base_port=22000, + sglang_router_port=21900, + sglang_router_prometheus_port=21901, + train_master_base_port=21902, + ) + paired_argv = build_miles_argv( + paired_args, + model_path="/model", + prompt_path="/prompts.jsonl", + provider=provider, + target_modules=["qkv_proj", "out_proj"], + ) + assert paired_argv[paired_argv.index("--rollout-seed") + 1] == "91" + assert paired_argv[ + paired_argv.index("--rollout-engine-base-port") + 1 + ] == "22000" + assert paired_argv[paired_argv.index("--sglang-router-port") + 1] == "21900" + assert paired_argv[ + paired_argv.index("--sglang-router-prometheus-port") + 1 + ] == "21901" + assert paired_argv[ + paired_argv.index("--train-master-base-port") + 1 + ] == "21902" + dense_values = vars(provider).copy() dense_values.pop("num_moe_experts") dense_provider = argparse.Namespace(**dense_values) @@ -585,6 +628,19 @@ def test_miles_argv_uses_provider_capabilities_without_model_family_branches(): target_modules=["qkv_proj", "out_proj"], ) assert dense_argv[dense_argv.index("--expert-model-parallel-size") + 1] == "1" + assert "--qkv-format" not in dense_argv + + gdn_values = vars(dense_provider).copy() + gdn_values["experimental_attention_variant"] = "gated_delta_net" + gdn_argv = build_miles_argv( + dense_args, + model_path="/model", + prompt_path="/prompts.jsonl", + provider=argparse.Namespace(**gdn_values), + target_modules=["qkv_proj", "out_proj"], + ) + assert gdn_argv[gdn_argv.index("--qkv-format") + 1] == "bshd" + dense_args.expert_parallel = 2 with pytest.raises(ValueError, match="EP>1 requires a MoE"): build_miles_argv( @@ -607,6 +663,190 @@ def test_miles_argv_uses_provider_capabilities_without_model_family_branches(): ) +def test_miles_argv_preserves_mrope_provider_configuration(): + args = argparse.Namespace( + seq_len=128, + groups_per_round=4, + samples_per_group=2, + optimizer_steps=1, + lora_targets="attention", + lora_r=4, + seed=1, + learner_id=0, + global_rounds=1, + inner_lr=1e-5, + reward_function="pkg.reward:score", + over_sampling_batch_size=4, + rollout_max_response_len=64, + custom_generate_function_path=None, + use_session_server=False, + actor_num_nodes=1, + actor_num_gpus_per_node=2, + expert_parallel=None, + ) + provider = argparse.Namespace( + hidden_size=16, + num_attention_heads=4, + num_layers=2, + ffn_hidden_size=32, + num_query_groups=2, + kv_channels=4, + seq_length=256, + vocab_size=64, + layernorm_epsilon=1e-6, + position_embedding_type="mrope", + rotary_base=1000000, + rotary_percent=0.25, + mrope_section=[11, 11, 10], + ) + + argv = build_miles_argv( + args, + model_path="/model", + prompt_path="/prompts.jsonl", + provider=provider, + target_modules=["q_proj"], + ) + + assert argv[argv.index("--position-embedding-type") + 1] == "mrope" + assert argv[argv.index("--rotary-percent") + 1] == "0.25" + section = argv.index("--mrope-section") + assert argv[section + 1 : section + 4] == ["11", "11", "10"] + + +def test_megatron_targets_follow_the_exact_peft_bridge_mappings(): + qkv = types.SimpleNamespace( + megatron_param=( + "language_model.decoder.layers.3.self_attention.linear_qkv.weight" + ), + hf_param={ + "q": "model.language_model.layers.3.self_attn.q_proj.weight", + "k": "model.language_model.layers.3.self_attn.k_proj.weight", + "v": "model.language_model.layers.3.self_attn.v_proj.weight", + }, + ) + projection = types.SimpleNamespace( + megatron_param=( + "language_model.decoder.layers.3.self_attention.linear_proj.weight" + ), + hf_param="model.language_model.layers.3.self_attn.o_proj.weight", + ) + mappings = { + value: qkv for value in qkv.hf_param.values() + } | {projection.hf_param: projection} + registry = types.SimpleNamespace( + hf_to_megatron_lookup=lambda name: mappings.get(name) + ) + bridge = types.SimpleNamespace( + _model_bridge=types.SimpleNamespace(mapping_registry=lambda: registry) + ) + specs = tuple( + types.SimpleNamespace(name=f"base_model.model.{module}.lora_{side}.weight") + for module in ( + "model.language_model.layers.3.self_attn.q_proj", + "model.language_model.layers.3.self_attn.k_proj", + "model.language_model.layers.3.self_attn.v_proj", + "model.language_model.layers.3.self_attn.o_proj", + ) + for side in ("A", "B") + ) + + assert rl_learner.megatron_adapter_targets(specs, bridge) == [ + "language_model.decoder.layers.3.self_attention.linear_k", + "language_model.decoder.layers.3.self_attention.linear_proj", + "language_model.decoder.layers.3.self_attention.linear_q", + "language_model.decoder.layers.3.self_attention.linear_v", + ] + + +def test_megatron_targets_reject_a_peft_module_without_a_bridge_mapping(): + bridge = types.SimpleNamespace( + _model_bridge=types.SimpleNamespace( + mapping_registry=lambda: types.SimpleNamespace( + hf_to_megatron_lookup=lambda _name: None + ) + ) + ) + specs = ( + types.SimpleNamespace( + name="base_model.model.model.layers.0.self_attn.q_proj.lora_A.weight" + ), + ) + + with pytest.raises(ValueError, match="has no Megatron-Bridge mapping"): + rl_learner.megatron_adapter_targets(specs, bridge) + + +def test_miles_runner_keeps_native_arm_outside_yeto_policy_sync(monkeypatch): + captured = {} + + class Provider: + def finalize(self): + captured["finalized"] = True + + class Bridge: + def to_megatron_provider(self, load_weights): + assert load_weights is False + return Provider() + + class AutoBridge: + @staticmethod + def from_hf_pretrained(path, trust_remote_code): + assert path == "/model" + assert trust_remote_code is True + return Bridge() + + async def train(args): + captured["miles_args"] = args + + def build(*args, **kwargs): + captured["policy_sync"] = kwargs["yeto_policy_sync"] + return ["train.py"] + + monkeypatch.setitem( + sys.modules, + "megatron.bridge", + types.SimpleNamespace(AutoBridge=AutoBridge), + ) + monkeypatch.setitem(sys.modules, "train", types.SimpleNamespace(train=train)) + monkeypatch.setattr(rl_learner, "derive_peft_lora_specs", lambda *a, **k: ()) + monkeypatch.setattr(rl_learner, "adapter_targets", lambda specs: []) + monkeypatch.setattr( + rl_learner, + "megatron_adapter_targets", + lambda specs, bridge: [], + ) + monkeypatch.setattr(rl_learner, "build_miles_argv", build) + monkeypatch.setattr( + rl_learner, + "_parse_miles_args", + lambda argv: argparse.Namespace(argv=argv), + ) + args = argparse.Namespace( + trust_remote_code=True, + lora_r=4, + lora_targets="auto", + learner_id=0, + ) + + rl_learner.run_miles( + args, + model_path="/model", + prompt_path="/prompts.jsonl", + yeto_policy_sync=False, + extra_argv=("--save-debug-rollout-data", "/rollouts/{rollout_id}.pt"), + ) + + assert captured["finalized"] + assert captured["policy_sync"] is False + assert captured["miles_args"].argv == [ + "train.py", + "--save-debug-rollout-data", + "/rollouts/{rollout_id}.pt", + ] + assert not hasattr(captured["miles_args"], "yeto_rl_bridge_config") + + class _Status: def __init__(self, value): self.value = value diff --git a/yeto/rl/__init__.py b/yeto/rl/__init__.py index 61cc6c9..580e272 100644 --- a/yeto/rl/__init__.py +++ b/yeto/rl/__init__.py @@ -1,7 +1,7 @@ """Pinned Miles reinforcement-learning integration.""" MILES_REPOSITORY = "https://github.com/agentenv/miles" -MILES_COMMIT = "a91bd34e50416aeb1da111f74d52b296e8216b96" +MILES_COMMIT = "c951c667c2b754cf244e1787845c05b41b50d4df" MILES_PEFT_VERSION = "0.20.0" MILES_IMAGE = ( "docker:radixark/miles@sha256:" diff --git a/yeto/rl/bridge.py b/yeto/rl/bridge.py index 21d7e9a..40cb0a1 100644 --- a/yeto/rl/bridge.py +++ b/yeto/rl/bridge.py @@ -9,7 +9,7 @@ from pathlib import Path from typing import Protocol -from ..protocol import DTYPE_F32, PullRequest, SyncerClient +from ..protocol import DTYPE_F32, FinalManifest, PullRequest, SyncerClient from ..tensor_io import pack_tensor, unpack_fragment from .core import ( CanonicalLoraState, @@ -165,6 +165,9 @@ def wait_for_global_policy(self, version: int) -> CanonicalLoraState: def wait_for_initial_policy(self) -> CanonicalLoraState: while self.current is None: self.client.check_health() + if self.client.finalizing.is_set(): + _, self.current = self._terminal_state() + return self.current self._drain_messages() if self.current is None: time.sleep(0.05) @@ -361,6 +364,12 @@ def finalize(self) -> CanonicalLoraState: return self._finalize() def _finalize(self) -> CanonicalLoraState: + manifest, final = self._terminal_state() + self.runtime.apply_global_policy(final) + self.client.acknowledge_finalization(manifest) + return final + + def _terminal_state(self) -> tuple[FinalManifest, CanonicalLoraState]: manifest, fragments = self.client.wait_for_final_fragments() if ( manifest.global_step != self.config.global_rounds @@ -372,9 +381,7 @@ def _finalize(self) -> CanonicalLoraState: manifest.global_step, fragments[0].data, ) - self.runtime.apply_global_policy(final) - self.client.acknowledge_finalization(manifest) - return final + return manifest, final def _append_event(self, event: dict) -> None: path = Path(self.config.event_tape).expanduser() diff --git a/yeto/rl/export.py b/yeto/rl/export.py index fe22df1..8328202 100644 --- a/yeto/rl/export.py +++ b/yeto/rl/export.py @@ -20,6 +20,30 @@ ) +def _rl_model_factory(config): + """Use the model class declared by the pinned checkpoint when available.""" + + import transformers + + for architecture in getattr(config, "architectures", None) or (): + factory = getattr(transformers, architecture, None) + if factory is not None: + return factory + return transformers.AutoModelForCausalLM + + +def _rl_model_from_config(config, *, trust_remote_code: bool): + import transformers + + factory = _rl_model_factory(config) + if factory is transformers.AutoModelForCausalLM: + return factory.from_config( + config, + trust_remote_code=trust_remote_code, + ) + return factory._from_config(config) + + def target_modules(choice: str, config) -> str: """Reuse Yeto's model-driven public LoRA target semantics.""" @@ -40,7 +64,7 @@ def derive_peft_lora_specs( from accelerate import init_empty_weights from peft import LoraConfig, get_peft_model, get_peft_model_state_dict - from transformers import AutoConfig, AutoModelForCausalLM + from transformers import AutoConfig config = AutoConfig.from_pretrained( model, @@ -54,7 +78,7 @@ def derive_peft_lora_specs( }: targets = target_modules(targets, config) with init_empty_weights(): - base = AutoModelForCausalLM.from_config( + base = _rl_model_from_config( config, trust_remote_code=trust_remote_code, ) diff --git a/yeto/rl/learner.py b/yeto/rl/learner.py index fdec189..63c3652 100644 --- a/yeto/rl/learner.py +++ b/yeto/rl/learner.py @@ -7,6 +7,7 @@ import json import os import sys +from collections.abc import Sequence from pathlib import Path from typing import Any @@ -85,6 +86,51 @@ def _miles_callable(spec: str) -> str: return f"{module}.{function}" +def megatron_adapter_targets(specs, bridge) -> list[str]: + """Map the exact PEFT contract onto Bridge's Megatron module paths.""" + + model_bridge = getattr(bridge, "_model_bridge", None) + if model_bridge is None: + raise ValueError("Megatron-Bridge does not expose its model mapping") + registry = model_bridge.mapping_registry() + modules = { + spec.name.removeprefix("base_model.model.").rsplit(".lora_", 1)[0] + for spec in specs + } + targets = set() + for module in modules: + hf_weight = f"{module}.weight" + mapping = registry.hf_to_megatron_lookup(hf_weight) + if mapping is None: + raise ValueError(f"PEFT module {module!r} has no Megatron-Bridge mapping") + megatron_module = mapping.megatron_param.removesuffix(".weight") + prefix, separator, leaf = megatron_module.rpartition(".") + if not separator: + raise ValueError(f"invalid Megatron adapter module {megatron_module!r}") + if leaf in {"linear_qkv", "linear_fc1"}: + hf_params = mapping.hf_param + component = next( + ( + name + for name, value in hf_params.items() + if value == hf_weight + ), + None, + ) if isinstance(hf_params, dict) else None + allowed = {"q", "k", "v"} if leaf == "linear_qkv" else {"gate", "up"} + if component not in allowed: + raise ValueError( + f"PEFT module {module!r} cannot use canonical Megatron LoRA" + ) + leaf = ( + f"linear_{component}" + if leaf == "linear_qkv" + else f"linear_fc1_{component}" + ) + targets.add(f"{prefix}.{leaf}") + return sorted(targets) + + def build_miles_argv( args, *, @@ -92,6 +138,7 @@ def build_miles_argv( prompt_path: str | Path, provider, target_modules: list[str], + yeto_policy_sync: bool = True, ) -> list[str]: """Construct Miles arguments from Bridge's actual model provider.""" @@ -123,6 +170,7 @@ def build_miles_argv( if position_type == "yarn": position_type, rope_type = "rope", "yarn" rotary_base = int(getattr(provider, "rotary_base", 10000)) + rotary_percent = float(getattr(provider, "rotary_percent", 1.0)) actor_gpus = args.actor_num_nodes * args.actor_num_gpus_per_node is_moe = getattr(provider, "num_moe_experts", None) is not None expert_parallel = getattr(args, "expert_parallel", None) or ( @@ -167,6 +215,7 @@ def build_miles_argv( "--norm-epsilon", str(epsilon), "--position-embedding-type", position_type, "--rotary-base", str(rotary_base), + "--rotary-percent", str(rotary_percent), "--vocab-size", str(vocab_size), "--lora-rank", str(args.lora_r), "--lora-alpha", str(args.lora_r), @@ -191,7 +240,8 @@ def build_miles_argv( "--label-key", "label", "--metadata-key", "metadata", "--apply-chat-template", - "--rollout-seed", str(args.seed + args.learner_id), + "--rollout-seed", + str(getattr(args, "rollout_seed", args.seed + args.learner_id)), "--sglang-enable-deterministic-inference", "--num-rollout", str(args.global_rounds), "--rollout-batch-size", str(args.groups_per_round), @@ -202,12 +252,15 @@ def build_miles_argv( "--balance-data", "--rollout-max-context-len", str(args.seq_len), "--rollout-max-response-len", str(args.rollout_max_response_len), - "--rollout-function-path", "yeto.rl.miles.generate_rollout", - "--rollout-all-samples-process-path", "yeto.rl.miles.queue_completed_groups", + "--rollout-function-path", + ( + "yeto.rl.miles.generate_rollout" + if yeto_policy_sync + else "miles.rollout.sglang_rollout.generate_rollout" + ), "--custom-rm-path", _miles_callable(args.reward_function), "--advantage-estimator", "grpo", "--lr", str(args.inner_lr), - "--external-policy-sync-path", "yeto.rl.miles.create_policy_sync", "--accumulate-allreduce-grads-in-fp32", "--attention-softmax-in-fp32", "--attention-backend", "unfused", @@ -222,6 +275,24 @@ def build_miles_argv( "--sglang-max-lora-rank", str(args.lora_r), "--pin-rollout-manager-to-head", ] + if yeto_policy_sync: + values.extend( + ( + "--rollout-all-samples-process-path", + "yeto.rl.miles.queue_completed_groups", + "--external-policy-sync-path", + "yeto.rl.miles.create_policy_sync", + ) + ) + for flag, name in ( + ("--rollout-engine-base-port", "rollout_engine_base_port"), + ("--sglang-router-port", "sglang_router_port"), + ("--sglang-router-prometheus-port", "sglang_router_prometheus_port"), + ("--train-master-base-port", "train_master_base_port"), + ): + value = getattr(args, name, None) + if value is not None: + values.extend((flag, str(value))) if args.custom_generate_function_path: values.extend( ( @@ -242,6 +313,10 @@ def build_miles_argv( values.append("--group-query-attention") if rope_type is not None: values.extend(("--rope-type", rope_type)) + if position_type == "mrope": + section = _provider_value(provider, "mrope_section") + values.append("--mrope-section") + values.extend(str(int(value)) for value in section) if bool(getattr(provider, "gated_linear_unit", False)): values.append("--swiglu") if not bool(getattr(provider, "share_embeddings_and_output_weights", True)): @@ -252,6 +327,11 @@ def build_miles_argv( values.append("--add-qkv-bias") if bool(getattr(provider, "qk_layernorm", False)): values.append("--qk-layernorm") + if ( + _text(getattr(provider, "experimental_attention_variant", None)) + == "gated_delta_net" + ): + values.extend(("--qkv-format", "bshd")) if getattr(provider, "num_moe_experts", None) is not None: values.extend( ( @@ -356,6 +436,84 @@ def _syncer_address(value: str) -> tuple[str, int]: return host, int(port) +def run_miles( + args, + *, + model_path: str | Path, + prompt_path: str | Path, + yeto_policy_sync: bool = True, + extra_argv: Sequence[str] = (), +) -> None: + """Run one Miles job, optionally with Yeto's external policy boundary.""" + + from megatron.bridge import AutoBridge + + model_bridge = AutoBridge.from_hf_pretrained( + model_path, + trust_remote_code=args.trust_remote_code, + ) + provider = model_bridge.to_megatron_provider(load_weights=False) + provider.finalize() + specs = derive_peft_lora_specs( + model_path, + None, + rank=args.lora_r, + targets=args.lora_targets, + trust_remote_code=args.trust_remote_code, + ) + canonical_targets = adapter_targets(specs) + miles_targets = megatron_adapter_targets(specs, model_bridge) + miles_argv = build_miles_argv( + args, + model_path=model_path, + prompt_path=prompt_path, + provider=provider, + target_modules=miles_targets, + yeto_policy_sync=yeto_policy_sync, + ) + miles_argv.extend(extra_argv) + miles_args = _parse_miles_args(miles_argv) + + if yeto_policy_sync: + from .core import canonical_layout_hash, canonical_lora_config_hash + + layout_hash = canonical_layout_hash(specs) + lora_config_hash = canonical_lora_config_hash( + rank=args.lora_r, + target_modules=canonical_targets, + ) + miles_args.yeto_rl_trust_remote_code = args.trust_remote_code + miles_args.yeto_rl_model = args.model + miles_args.yeto_rl_data = args.data + miles_args.yeto_rl_base_model_revision = args.model_revision + miles_args.yeto_rl_data_revision = args.data_revision + miles_args.yeto_rl_lora_config_hash = lora_config_hash + miles_args.yeto_rl_layout_hash = layout_hash + miles_args.yeto_rl_reward_sha256 = args.reward_sha256 + miles_args.yeto_rl_completed_groups_path = args.completed_groups_path + miles_args.yeto_rl_event_tape = args.event_tape + miles_args.yeto_rl_learner_id = args.learner_id + miles_args.yeto_rl_bridge_config = BridgeConfig( + syncer_addr=_syncer_address(args.syncer), + learner_id=args.learner_id, + global_rounds=args.global_rounds, + groups_per_round=args.groups_per_round, + samples_per_group=args.samples_per_group, + local_optimizer_steps=args.optimizer_steps, + wan_streams=args.wan_streams, + expected_specs=specs, + base_model_revision=args.model_revision, + lora_config_hash=lora_config_hash, + layout_hash=layout_hash, + event_tape=args.event_tape, + ) + + from train import train as miles_train + + asyncio.run(miles_train(miles_args)) + print(f"[rl] learner {args.learner_id} finalized") + + def main(argv=None) -> None: args = parse_args(argv) from ..provenance import ( @@ -376,6 +534,9 @@ def main(argv=None) -> None: f"got {reward_sha256}" ) verify_miles_revision(args.miles_root) + miles_root = str(Path(args.miles_root).expanduser().resolve()) + if miles_root not in sys.path: + sys.path.insert(0, miles_root) from miles.utils.misc import load_function @@ -383,7 +544,6 @@ def main(argv=None) -> None: if args.custom_generate_function_path: load_function(args.custom_generate_function_path) from huggingface_hub import snapshot_download - from megatron.bridge import AutoBridge from ..models import resolve from ..provenance import is_local_reference @@ -398,67 +558,7 @@ def main(argv=None) -> None: args.data_revision, "~/yeto-rl/prompts.jsonl", ) - - model_bridge = AutoBridge.from_hf_pretrained( - model_path, - trust_remote_code=args.trust_remote_code, - ) - provider = model_bridge.to_megatron_provider(load_weights=False) - provider.finalize() - specs = derive_peft_lora_specs( - model_path, - None, - rank=args.lora_r, - targets=args.lora_targets, - trust_remote_code=args.trust_remote_code, - ) - miles_targets = adapter_targets(specs) - from .core import canonical_layout_hash, canonical_lora_config_hash - - layout_hash = canonical_layout_hash(specs) - lora_config_hash = canonical_lora_config_hash( - rank=args.lora_r, - target_modules=miles_targets, - ) - miles_args = _parse_miles_args( - build_miles_argv( - args, - model_path=model_path, - prompt_path=prompt_path, - provider=provider, - target_modules=miles_targets, - ) - ) - miles_args.yeto_rl_trust_remote_code = args.trust_remote_code - miles_args.yeto_rl_model = args.model - miles_args.yeto_rl_data = args.data - miles_args.yeto_rl_base_model_revision = args.model_revision - miles_args.yeto_rl_data_revision = args.data_revision - miles_args.yeto_rl_lora_config_hash = lora_config_hash - miles_args.yeto_rl_layout_hash = layout_hash - miles_args.yeto_rl_reward_sha256 = args.reward_sha256 - miles_args.yeto_rl_completed_groups_path = args.completed_groups_path - miles_args.yeto_rl_event_tape = args.event_tape - miles_args.yeto_rl_learner_id = args.learner_id - - miles_args.yeto_rl_bridge_config = BridgeConfig( - syncer_addr=_syncer_address(args.syncer), - learner_id=args.learner_id, - global_rounds=args.global_rounds, - groups_per_round=args.groups_per_round, - samples_per_group=args.samples_per_group, - local_optimizer_steps=args.optimizer_steps, - wan_streams=args.wan_streams, - expected_specs=specs, - base_model_revision=args.model_revision, - lora_config_hash=lora_config_hash, - layout_hash=layout_hash, - event_tape=args.event_tape, - ) - from train import train as miles_train - - asyncio.run(miles_train(miles_args)) - print(f"[rl] learner {args.learner_id} finalized") + run_miles(args, model_path=model_path, prompt_path=prompt_path) if __name__ == "__main__": From 8e1eef98af8ba85a5b7a56590fcf49f43a37b803 Mon Sep 17 00:00:00 2001 From: AlexEisie <1987460907@qq.com> Date: Fri, 31 Jul 2026 08:24:05 +0800 Subject: [PATCH 5/6] docs(rl): record current Miles validation Describe the maintained external policy hook, the gated-attention Megatron-Bridge compatibility in the pinned c951c667 Miles revision, and the equal-hardware native, single-island, and federated benchmark workflow. Record the three-seed Qwen3.6-27B run on eight H200 GPUs: nine eight-round jobs, 2,304 real training trajectories, 576 held-out generations, aggregate reward and pass@k results, and the shared 768-token response-cap limitation. Keep the broader four-A100 MoE, recovery, parity, export, session/tool, and 20-merge campaign distinct from the current-pin benchmark, update automated-test evidence, and state the remaining multi-node, Spot, durable-checkpoint, dashboard, and soak boundaries precisely. --- docs/MILES_RL.md | 151 +++++++++++++++++++++++++++++------------------ 1 file changed, 93 insertions(+), 58 deletions(-) diff --git a/docs/MILES_RL.md b/docs/MILES_RL.md index daf7dad..0aaa36a 100644 --- a/docs/MILES_RL.md +++ b/docs/MILES_RL.md @@ -11,10 +11,10 @@ FedAvg: every global round waits for one exact-base result from every logical island and averages them equally. > **Status:** the core path and the pinned Miles synchronization branch are -> implemented. The algorithm and model path previously passed real multi-GPU -> dense, MoE EP, two-island averaging, recovery, long-task, and 20-merge -> validation on eight A100 GPUs. The thin-branch boundary has automated -> integration coverage but has not yet repeated that GPU matrix; see +> implemented. The current pin completed a real equal-hardware Qwen3.6-27B +> benchmark across native Miles, one Yeto island, and two Yeto islands on eight +> H200 GPUs. A separate four-A100 campaign covers multi-GPU dense, MoE EP, +> recovery, export, parity, session/tool, and 20-merge behavior; see > [Validation status](#validation-status). ## Why Miles is the RL runtime @@ -121,6 +121,9 @@ the narrow boundary required by this mode: - the native Miles train loop loads one external policy-sync hook, calls it after local training and before the normal trainer-to-SGLang publication, and finalizes it only after the final global policy has been published. +- its Megatron-Bridge compatibility preserves canonical LoRA query and gate + rows for architectures that declare gated attention, while leaving ordinary + attention unchanged. Miles still owns rollout generation, reward and advantage computation, GRPO training, offload, checkpoint calls, and SGLang weight transport. Yeto does @@ -227,7 +230,7 @@ The Miles source itself is independently pinned to: ```text https://github.com/agentenv/miles -a91bd34e50416aeb1da111f74d52b296e8216b96 +c951c667c2b754cf244e1787845c05b41b50d4df ``` The launcher checks out that commit as a detached HEAD and installs the @@ -476,6 +479,16 @@ Common startup and progress failures have distinct meanings: - **optimizer steps but negligible LoRA change:** first inspect reward variance, advantages, and the configured LR schedule. +## Benchmark + +[`RL_BENCHMARK.md`](RL_BENCHMARK.md) describes the local equal-hardware LM +benchmark. It compares native Miles, one strict Yeto island, and strict +federated Yeto with paired prompt budgets and held-out reward/pass@k +evaluation. The runner uses real Miles generation and training, and evaluates +Yeto only from the authoritative syncer checkpoint export. All three arms use +one explicit expert-parallel setting, while same-host federated islands receive +disjoint Miles port ranges. + ## Intentional differences from INIT | INIT plan | Current implementation | Assessment | @@ -490,64 +503,86 @@ ordering around native SGLang publication, multi-node task construction, multi-rank apply/export, DP rollout-shard collection, EP validation, single-island and multi-island sync, canonical identity, completed-group recovery, strict failures, provenance, checkpoint export, and the unchanged -SFT/diffusion defaults. Exact current test counts are recorded in the pull -request rather than frozen in this document. - -The following real validation ran on one GCP Spot VM with eight NVIDIA -A100-SXM4-40GB GPUs before the integration moved from runtime adaptation to -the maintained Miles hook. It used the same policy, merge, model, reward, -learner, and Rust syncer semantics, but it is evidence for the RL algorithm -and model path rather than GPU validation of the new hook implementation. - -- A one-island Qwen3-4B DP=8 run completed two global rounds. Each round - trained on 16 trajectories and 8192 action tokens with nonconstant rewards. - Local delta norms were `0.22248` and `0.10623`; both M=1 merges matched the - saved local policy within `3.64e-12`. -- Two concurrent Qwen3-4B DP=4 islands used different prompt-token hashes. - Their local delta norms were `0.223218` and `0.222651`; the committed f32 - policy exactly matched the offline mean, and both islands applied identical - initial and final policies. -- `allenai/OLMoE-1B-7B-0125-Instruct` ran one EP=8 island with replicated - attention LoRA, 16 trajectories, 4096 action tokens, nonconstant rewards, - and a `0.10110` local delta. The merge error was zero. Standard PEFT loaded - the exported adapter, produced finite logits with a nonzero adapter effect, - and completed real generation. -- A direct pinned-Miles round and the production Yeto+Miles M=1 path produced - identical sampled tokens and rewards; their LoRA tensors had max error - `0.0`. The exact trainer-to-SGLang LoRA checksum checker passed throughout. - Trainer-versus-rollout KL stayed around `6e-4` to `1.3e-3` before and after - global applies rather than showing a material token-path mismatch. -- A 20-round Qwen run committed versions `1..20` exactly once, completing 160 - trajectories and 20,480 action tokens. Every round had a nonzero finite - local delta, no mixed-version group was observed, and current-versus-rollout - KL remained between `0.000329` and `0.001067`. -- Learners were killed during rollout, after local train, before push, after - push, before broadcast, and after global apply. Every replacement completed, - and each case produced one committed step. A separate syncer restart resumed - the committed version and completed the next round. -- Completed-group recovery retained one real four-trajectory oversampling - group across learner replacement and selected that group after restart. A - separately replayed real trained delta was accepted once, while a stale - exact-base update caused the strict connection to close. -- A Qwen3-4B session-server run completed two rounds and 16 real environment - tasks. Ten trajectories made actual calculator calls and consumed their - returned values; the second round had nonconstant rewards and a `0.10059` - local update. Fourteen session traces had exact TITO reconstruction. The two - remaining traces reached the 512-token response cap before a terminal token, - were correctly marked truncated, and accounted for the reported 25% TITO - structural mismatch in that round. Trainer-versus-rollout absolute logprob - error remained below `0.01`. +SFT/diffusion defaults. The current source passed 85 focused Yeto tests, +the full Yeto suite with 799 passed and 4 skipped, 58 Rust syncer tests, and 10 +focused Miles tests. + +The current Miles pin was exercised with the immutable `Qwen/Qwen3.6-27B` +revision on one eight-H200 host. The equal-hardware benchmark used three seeds +and, for each seed, compared native Miles on eight GPUs, one Yeto island on +eight GPUs, and two four-GPU Yeto islands. Every arm completed eight real RL +rounds, 64 prompt groups, and 256 trajectories, for 2,304 training trajectories +in total. The nine resulting adapters were evaluated on 16 held-out prompts +with four samples each, for 576 evaluation generations. The complete run +exited successfully in about 11 hours 35 minutes. + +| arm | mean reward | pass@1 | pass@4 | +| --- | ---: | ---: | ---: | +| native Miles | `0.1094 +/- 0.0312` | `0.1094` | `0.2083` | +| Yeto single island | `0.0938 +/- 0.0541` | `0.0938` | `0.1875` | +| Yeto two-island federation | `0.1146 +/- 0.0592` | `0.1146` | `0.2500` | + +Nearly every response reached the shared 768-token cap. These bounded-run +quality values therefore show that the real training, synchronization, export, +and evaluation paths work together; they are not a claim of converged policy +quality. The small arm differences are also within the variation across three +seeds. + +The broader compatibility and recovery campaign ran on one GCP Spot VM with +four NVIDIA A100-SXM4-40GB GPUs. It directly exercised the maintained external +policy hook. Yeto was `9ebd42060414fecfca2cabd89f669e91f93d1041` +(tested source digest +`80dfdbb7c2b3d8bc18cee51ac55e63a28b93cc0811c660ac625837b80baeab6e`) and +Miles `a91bd34e50416aeb1da111f74d52b296e8216b96`, plus propagation of +`provider.attention_backend` in the bridge model provider. PEFT was `0.20.0`. + +- Qwen3-0.6B completed multiple real DeepMath/reward/GRPO updates. Qwen3-4B + completed two rounds in one DP=4 island. +- Two concurrent Qwen3-4B islands, each DP=2 and using different prompts, + committed a policy whose maximum error from the offline f32 mean was `0`; + both islands finished on the same policy. +- `allenai/OLMoE-1B-7B-0125-Instruct` completed EP=4 training with attention + LoRA. Its local delta norm was `0.1417596`, and the committed policy's maximum + error from the expected merge was `0`. +- Learners were replaced during rollout, before push, after push, and after + global apply. Each case recovered and committed exactly once, with no stale + update accepted. The after-push and after-apply cases first exposed a final + replacement hang; after the fix, both real reruns exited successfully. +- A syncer restart recovered checkpoint version 1 and committed version 2. + Learners observed base versions `[0, 1]`, with no duplicate merge or stale + update. +- A 20-round Qwen run committed versions `1..20` in order and applied versions + `0..20`. It completed 160 trajectories and 20,480 action tokens with no + mixed-version or stale group; the maximum merge error was `7.28e-12`. +- Native Miles and the Yeto path produced identical rollout tokens and rewards, + with maximum LoRA tensor error `0.0`. +- The committed version-20 policy exported as standard PEFT, loaded normally, + and completed real GPU inference and generation. Its maximum logit difference + from the base policy was `0.25`. +- A Qwen3-4B session/tool run completed two rounds and 16 real tasks, including + 10 actual calculator calls. Its local delta norms were `0.0001186` and + `0.1004914`, with no mixed-version or stale group. + +The evidence archive is +`gs://yeto-exp2-52-model-training-497007/miles-rl-validation/20260730-thin-a100-4g/` +with SHA-256 +`d0945b0077308bda368ae9cc45d45d1be12b5e6fcafb12cba3dbbbd95128e142`. The following boundaries remain unvalidated or intentionally excluded: -- the maintained Miles hook has not yet repeated the real GPU matrix above; +- the H200 benchmark covered the current pin's concurrent same-host drivers + and a real gated-attention model; the A100 campaign's MoE EP, + fault-injection, and session-server cases have not all been repeated on that + pin; - the requested 24-hour soak was not run; the 20 consecutive merges are the bounded-duration stability evidence; -- no physical multi-node island or end-to-end SkyPilot provisioning and Spot - VM replacement was exercised, although multi-node task construction is - covered automatically; -- the syncer checkpoint still has no durable mount for syncer VM or disk loss; -- metrics remain JSONL-only and the launcher enables no dashboard. +- no physical multi-node island was exercised; +- neither end-to-end SkyPilot provisioning nor a real Spot preemption was + exercised; learner recovery used process termination and replacement; +- syncer restart was verified only with its existing checkpoint disk retained, + so syncer VM or disk loss remains unrecovered; +- metrics remain JSONL-only and the launcher enables no dashboard; +- Diffusion RL is outside the v0 contract and is therefore not a missing test. ## Extending v0 safely From bb3ecdaed4229f53f19c26527d444df0c1c40c0f Mon Sep 17 00:00:00 2001 From: AlexEisie <1987460907@qq.com> Date: Fri, 31 Jul 2026 11:23:32 +0800 Subject: [PATCH 6/6] docs(benchmarks): publish Miles RL results Archive the three-seed equal-hardware Qwen3.6-27B comparison across native Miles, one Yeto island, and two Yeto islands. Record configuration provenance, aggregate and per-seed quality, execution and synchronization metrics, and the verified adapter contract. Document the response-cap limitation and keep the result scoped to strict v0 averaging rather than treating the small observed reward differences as convergence or an algorithmic win. --- docs/BENCHMARK_RESULTS.md | 152 ++++++++++++++++++++++++++++++++++---- 1 file changed, 137 insertions(+), 15 deletions(-) diff --git a/docs/BENCHMARK_RESULTS.md b/docs/BENCHMARK_RESULTS.md index f36e1d9..7ba7d44 100644 --- a/docs/BENCHMARK_RESULTS.md +++ b/docs/BENCHMARK_RESULTS.md @@ -1,20 +1,22 @@ # Benchmark Results -This document is the result archive for Yeto's completed LM and diffusion -benchmarks. Benchmark definitions, fairness contracts, arm semantics, and -execution instructions remain in [LM_BENCHMARK.md](LM_BENCHMARK.md) and +This document is the result archive for Yeto's completed LM, Miles RL, and +diffusion benchmarks. Benchmark definitions, fairness contracts, arm +semantics, and execution instructions remain in +[LM_BENCHMARK.md](LM_BENCHMARK.md), [RL_BENCHMARK.md](RL_BENCHMARK.md), and [DIFFUSION_BENCHMARK.md](DIFFUSION_BENCHMARK.md). Aggregate values are mean +/- sample standard deviation across training -seeds. Per-seed deltas compare a DiLoCo artifact only with the synchronous -baseline having the same topology and seed. The per-seed tables enumerate +seeds. SFT and diffusion per-seed deltas compare a DiLoCo artifact only with +the synchronous baseline having the same topology and seed. Miles RL deltas +use the paired native Miles or single-island arm. The per-seed tables enumerate every durable result record and its quality, execution, and synchronization metrics. Aggregate cost is the mean cost of one result item; explicitly reported full-run cost also includes setup and idle overhead. These runs are historical evidence rather than current release gates. The -Qwen run predates current resume-manifest validation; the LTX and Wan runs -predate current logical-rank training-stream pairing. +first Qwen SFT run predates current resume-manifest validation; the LTX and +Wan runs predate current logical-rank training-stream pairing. ## Qwen3.6-27B LM @@ -327,6 +329,119 @@ topology because this run also changed to NF4, attention-only LoRA, 512-token sequences, and gradient checkpointing. The best DiLoCo artifact was `m8`, seed 17, at 0.920297 CE/token (-0.21% versus paired sync). +## Qwen3.6-27B Miles RL + +### Configuration + +| item | value | +|---|---| +| completed | 2026-07-31 00:11:20 UTC | +| model | `Qwen/Qwen3.6-27B` at revision `6a9e13bd6fc8f0983b9b99948120bc37f49c13e9` | +| data | local MATH-500 JSONL, SHA-256 `af174db82e27aa94b3f492083c30b07282b546104fe9d7a5d184823d35e7daf9`; 484 training candidates and 16 held-out prompts | +| reward | `benchmark_runtime.rewards:math_score`, SHA-256 `1c9328a9d35d021f6e43fb3c1714d2866a088b16b6013a401c2fec29d0d5ed08` | +| hardware | one host with 8 x NVIDIA H200 143771 MiB GPUs | +| runtime | Miles `c951c667c2b754cf244e1787845c05b41b50d4df`; Yeto implementation fingerprint `45020cc2c1842b481f3c674e3bd143034ad6db3e2d2711a873dbb160b33f0c7b` | +| container | `sha256:95b3afa9ee4313f5633e6ed3779c8276353cc8e24a2462e4f54ec0d5978fbae7` | +| arms | native Miles: 1 island x 8 GPUs; Yeto single: 1 island x 8 GPUs; Yeto federated: 2 islands x 4 GPUs | +| training | attention LoRA r16, inner LR `1e-5`, sequence 2048, 8 global rounds, 1 optimizer step per local round, 4 samples per prompt group, 768-token response cap | +| paired budget | 64 prompt groups and 256 trajectories per arm and seed; all arms used 8 GPUs and identical model, reward, prompt, trajectory, and maximum action-token budgets | +| evaluation | 16 held-out prompts, 4 samples per prompt, temperature 1, top-p 1, paired generation seeds; pass when reward is greater than 0 | +| seeds | 17, 29, and 43 | +| full run | 9 durable records, 2,304 training trajectories, 576 held-out generations, about 11 hours 35 minutes | +| result archive | `results.jsonl` SHA-256 `1258c857e2e57051ff48784b12fc7d7061eb18baa4a0570d836a15b28b5a4053` | + +Native Miles is the equal-hardware RL reference. `yeto-single-m2` adds the +strict Yeto checkpoint/apply contract and LoRA optimizer reset without +splitting the island. `yeto-federated-m2` then isolates the two-island split +and exact equal-weight global average. + +### Aggregate Quality Results + +| arm | M | runs | reward | pass@1 | pass@4 | delta vs native | delta vs single | +|---|---:|---:|---:|---:|---:|---:|---:| +| `native-miles-m2` | 2 | 3 | 0.1094 +/- 0.0312 | 0.1094 | 0.2083 | - | - | +| `yeto-single-m2` | 2 | 3 | 0.0938 +/- 0.0541 | 0.0938 | 0.1875 | -0.0156 | - | +| `yeto-federated-m2` | 2 | 3 | 0.1146 +/- 0.0592 | 0.1146 | 0.2500 | +0.0052 | +0.0208 | + +### Aggregate Systems Results + +| arm | GPUs | train s | artifact-ready s | eval s | traj/s | action tok/s | GPU-h | mean sync s | sent MB | mean KL | +|---|---:|---:|---:|---:|---:|---:|---:|---:|---:|---:| +| `native-miles-m2` | 8 | 1583.9 | 1599.7 | 3096.7 | 0.164 | 125.7 | 4.380 | - | - | - | +| `yeto-single-m2` | 8 | 1568.6 | 1570.4 | 3071.3 | 0.163 | 125.4 | 4.339 | 150.624 | 612.369 | 0.000152 | +| `yeto-federated-m2` | 8 | 1456.8 | 1459.8 | 2922.1 | 0.176 | 135.0 | 4.049 | 138.427 | 1224.738 | 0.000155 | + +### Per-Seed Quality Results + +| arm | seed | reward | pass@1 | pass@4 | delta vs native | delta vs single | eval tokens | capped samples | +|---|---:|---:|---:|---:|---:|---:|---:|---:| +| `native-miles-m2` | 17 | 0.140625 | 0.140625 | 0.2500 | - | - | 48882 | 63/64 | +| `yeto-single-m2` | 17 | 0.125000 | 0.125000 | 0.2500 | -0.015625 | - | 48879 | 63/64 | +| `yeto-federated-m2` | 17 | 0.140625 | 0.140625 | 0.2500 | +0.000000 | +0.015625 | 48880 | 63/64 | +| `native-miles-m2` | 29 | 0.109375 | 0.109375 | 0.1875 | - | - | 49152 | 64/64 | +| `yeto-single-m2` | 29 | 0.125000 | 0.125000 | 0.1875 | +0.015625 | - | 49152 | 64/64 | +| `yeto-federated-m2` | 29 | 0.156250 | 0.156250 | 0.3125 | +0.046875 | +0.031250 | 49152 | 64/64 | +| `native-miles-m2` | 43 | 0.078125 | 0.078125 | 0.1875 | - | - | 49152 | 64/64 | +| `yeto-single-m2` | 43 | 0.031250 | 0.031250 | 0.1250 | -0.046875 | - | 49152 | 64/64 | +| `yeto-federated-m2` | 43 | 0.046875 | 0.046875 | 0.1875 | -0.031250 | +0.015625 | 49152 | 64/64 | + +### Per-Seed Execution Results + +| arm | seed | topology | groups | trajectories | capped trajectories | action tokens | train s | artifact s | artifact-ready s | eval s | GPU-h | +|---|---:|---|---:|---:|---:|---:|---:|---:|---:|---:|---:| +| `native-miles-m2` | 17 | 1 x 8 | 64 | 256 | 255/256 | 196531 | 1843.9 | 45.634 | 1889.5 | 2896.3 | 4.902 | +| `yeto-single-m2` | 17 | 1 x 8 | 64 | 256 | 255/256 | 196554 | 1624.2 | 2.580 | 1626.8 | 2845.7 | 4.400 | +| `yeto-federated-m2` | 17 | 2 x 4 | 64 | 256 | 255/256 | 196554 | 1411.0 | 1.421 | 1412.4 | 2914.5 | 3.945 | +| `native-miles-m2` | 29 | 1 x 8 | 64 | 256 | 256/256 | 196608 | 1443.0 | 0.776 | 1443.7 | 3428.1 | 4.159 | +| `yeto-single-m2` | 29 | 1 x 8 | 64 | 256 | 256/256 | 196608 | 1506.1 | 1.537 | 1507.6 | 3389.0 | 4.288 | +| `yeto-federated-m2` | 29 | 2 x 4 | 64 | 256 | 255/256 | 196576 | 1472.9 | 6.281 | 1479.1 | 2936.7 | 4.089 | +| `native-miles-m2` | 43 | 1 x 8 | 64 | 256 | 255/256 | 196531 | 1464.8 | 0.977 | 1465.8 | 2965.6 | 4.079 | +| `yeto-single-m2` | 43 | 1 x 8 | 64 | 256 | 255/256 | 196531 | 1575.5 | 1.431 | 1576.9 | 2979.3 | 4.329 | +| `yeto-federated-m2` | 43 | 2 x 4 | 64 | 256 | 255/256 | 196531 | 1486.5 | 1.448 | 1487.9 | 2915.1 | 4.113 | + +### Per-Seed Synchronization Results + +| arm | seed | merges | final version | local rounds | policy applies | roster per merge | mean sync s | sent MB | mean KL | ESS | clip fraction | +|---|---:|---:|---:|---:|---:|---|---:|---:|---:|---:|---:| +| `yeto-single-m2` | 17 | 8 | 8 | 8 | 9 | 1/1 | 159.448 | 612.369 | 0.000146 | 1.0 | 0.0 | +| `yeto-federated-m2` | 17 | 8 | 8 | 16 | 18 | 2/2 | 135.047 | 1224.738 | 0.000149 | 1.0 | 0.0 | +| `yeto-single-m2` | 29 | 8 | 8 | 8 | 9 | 1/1 | 143.062 | 612.369 | 0.000154 | 1.0 | 0.0 | +| `yeto-federated-m2` | 29 | 8 | 8 | 16 | 18 | 2/2 | 139.269 | 1224.738 | 0.000156 | 1.0 | 0.0 | +| `yeto-single-m2` | 43 | 8 | 8 | 8 | 9 | 1/1 | 149.361 | 612.369 | 0.000156 | 1.0 | 0.0 | +| `yeto-federated-m2` | 43 | 8 | 8 | 16 | 18 | 2/2 | 140.964 | 1224.738 | 0.000159 | 1.0 | 0.0 | + +### Interpretation + +All nine arms completed the paired workload and evaluation. Every Yeto run +reached global version 8. Each federated merge received both fixed-roster +islands, and both islands ended with the same global policy hash. All nine +standard PEFT adapters contained 224 finite tensors. + +Federation had an observed mean reward delta of +0.0052 versus native Miles +and +0.0208 versus Yeto single, while single-island Yeto was -0.0156 versus +native. These differences are smaller than the variation across three seeds. +With only 16 held-out prompts, they do not establish a quality winner. + +The response cap dominated this run: 2,297 of 2,304 training trajectories and +573 of 576 evaluation samples reached the 768-token maximum. The reward and +pass@k values therefore demonstrate a complete real RL, synchronization, +export, and held-out evaluation path; they are not evidence of converged +policy quality. Federated training was about 8% faster than native Miles in +this single-host run, but one host and three seeds are not a scaling result. + +An initial native calibration was discarded after the harness found that +Miles' native save omitted PEFT's `base_model.model.` key prefix. Before the +formal nine records, the harness was changed to verify all 224 names, shapes, +and tensors and to fail on missing adapter keys. The normalization adds only +the wrapper prefix; all three formal native adapters were bitwise equal to the +standardized tensors after their bfloat16-to-float32 conversion. + +This benchmark evaluates the strict v0 contract rather than a general RL +algorithm sweep. Native Miles supplies the same-runtime, equal-hardware +reference; the global Yeto path uses one complete LoRA fragment, full-roster +synchronization, and exact equal-weight averaging. It does not include a +separate FSDP2 RL runtime or outer-optimizer ablations. + ## LTX-Video Diffusion ### Configuration @@ -735,16 +850,18 @@ at 0.305034. ## Cross-Benchmark Summary -Absolute losses are not comparable across LM, LTX, and Wan; only paired -deltas against each workload's synchronous baseline are comparable. All -three baselines improved substantially over the untrained model, so the +Absolute SFT and diffusion losses are not comparable across LM, LTX, and Wan; +only paired deltas within a workload are comparable. Miles RL reward and +pass@k are also not comparable with those losses. The three SFT/diffusion +baselines improved substantially over the untrained model, so those benchmarks measured real fine-tuning rather than a failed training recipe. -| workload | default `m2` vs sync | strongest DiLoCo result | main signal | -|---|---:|---|---| -| Qwen3.6-27B | +9.91% | `direct-rda` +0.56% | Changing outer update/application fixed most of the gap; slightly faster synchronization did not. | -| LTX-Video | +1.44% | `m2` seed 17 at parity | The default path was already stable; `q4` cut payload by about 33% with a modest quality tradeoff. | -| Wan2.2 short / long | +13.77% / +17.07% | `unthrottled` -19.25% / -5.48% | Quality depended strongly on synchronization cadence; short `m2` was seed-stable, while the gain cost 2.9x / 5.0x payload. | +| workload | primary comparison | strongest observed result | main signal | +|---|---|---|---| +| Qwen3.6-27B SFT | default `m2` +9.91% vs sync | `direct-rda` +0.56% | Changing outer update/application fixed most of the gap; slightly faster synchronization did not. | +| Qwen3.6-27B Miles RL | federated reward +0.0052 vs native | federated pass@4 0.2500 vs native 0.2083 | All paths completed, but the differences were below seed variation and almost every response hit the token cap. | +| LTX-Video | default `m2` +1.44% vs sync | `m2` seed 17 at parity | The default path was already stable; `q4` cut payload by about 33% with a modest quality tradeoff. | +| Wan2.2 short / long | default `m2` +13.77% / +17.07% vs sync | `unthrottled` -19.25% / -5.48% | Quality depended strongly on synchronization cadence; short `m2` was seed-stable, while the gain cost 2.9x / 5.0x payload. | The central result is that one synchronization preset is not optimal for every model. Qwen primarily exposed an outer Nesterov/LR/blending problem; @@ -753,6 +870,11 @@ seed-stable short `m2` result but substantial long-run seed variance. `q4` reliably reduced bytes, but compression was only useful when the underlying optimization recipe was already sound. +The Miles RL run answers a narrower question: native Miles, the Yeto +single-island contract, and two-island strict averaging all completed the same +real workload and produced valid adapters. Its limited held-out set and +response cap do not support selecting one arm on quality. + For LM, `direct-rda` is the strongest production candidate. For LTX, default `m2` is acceptable and `q4` is a reasonable bandwidth tradeoff. Wan should not adopt `unthrottled` directly: first repeat long `direct-rda` for all three