Repository navigation
feat: MLX training backend for Apple Silicon + OpenClaw-RL end to end - #325
Open
BobbyZhouZijian wants to merge 23 commits into
Open
BobbyZhouZijian wants to merge 23 commits into
BobbyZhouZijian wants to merge 23 commits into
Conversation
|
Merge Oncall: @BobbyZhouZijian The merge oncall will route reviews and keep this pull request moving toward a clear decision. |
github-actions
Bot
added
area: core
BobbyZhouZijian
marked this pull request as ready for review
September 11, 2026 09:30
BobbyZhouZijian
requested review from
Benjamin-eecs and
hanfeiyu
as code owners
September 11, 2026 09:30
…OpenClaw-RL on it Reef's training path assumed a GPU cluster: Ray, Megatron, a separate rollout engine. This adds a runtime kind that serves and trains one model in a single process on Apple Silicon, and runs OpenClaw-RL's GSM8K homework stream on it — a benchmark whose reference stack needs seven GPUs. `reef/train/mlx_backend/` holds the backend: - `engine.py` — the model, the LoRA adapter and the optimizer in one process; serving reads those parameters and training writes them, kept apart by the engine lock. Serving drives the engine as a single sequence: mlx-lm's continuous batcher cannot merge Qwen3.5's hybrid cache, which splices a GatedDeltaNet conv-state onto an attention KV and Metal-OOMs the decode. - `gated_delta.py` — a chunkwise, differentiable GatedDeltaNet recurrence with per-layer checkpointing. mlx-lm implements the recurrence as a Metal kernel with no vjp, so in eval mode a gradient cannot cross one; running it in training mode for the span of each backward is what lets every layer train rather than the last attention block alone. - `objective.py` / `rows.py` — OpenClaw-RL's top-K OPD surrogate, verified against the torch implementation, plus a frozen-base KL term. The head's log-probs are scored in chunks so a long response's backward fits. - `inference.py` / `messages.py` / `runtime.py` — an OpenAI-compatible serving path: chat templates, a request's toolset, streamed completions, and reasoning and tool calls read back out of a reply. `recipes/openclawrl/` carries the method: the recipe, the PRM judge speaking an OpenAI dialect, and a candidate-evaluation plugin that gates publication on a pinned style probe, so a step that makes the policy answer worse is held out of serving instead of compounding into it. `results/2026-09-10-gsm8k-stream-qwen3.5-9b-mlx/` records a 72-session run on one 48 GB machine: all 32 layers trained at rank 64, `sessions_to_adaptation` 14 (the 7-GPU reference's is also 14), zero OOM across the stream. The record states what it does not establish — no held-out control was run for that configuration — and `mlx-runtime-notes.md` carries the runtime-general measurements: topology, how to read a step's metrics, the adapter artifact format, and the capacity envelope. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01SuEa3FyrkW1xc3sjmAJSEB
BobbyZhouZijian
force-pushed
the
worktree-issue-134
branch
from
September 12, 2026 08:59
670f13a to
a1192cb
Compare
This was referenced Sep 12, 2026
…s speed The reference harness kills a `hermes chat` turn at `TURN_TIMEOUT_S`, 600s, a ceiling set by the reference's GPU-speed serving. A killed turn is not reported as an error: the benchmark records the dead turn as a `no-reply` failure, so on a slower runtime the constant quietly converts serving latency into policy failures, and those failures then feed the training signal. Measured over a 72-session Hermes + Harbor stream (Qwen3.5-9B-4bit, all 32 layers, single-sequence serving on an M4 Pro): a turn takes 397s at the median and 1141s at the worst. 44 sessions replied and **28 died at the ceiling**, every one between 605s and 610s. They are not the policy falling silent — at each of the run's 36 training steps the same weights answered 8 of 8 probe prompts (`answered_rate` 1.0). The margin over the default is about 1.5x, so a longer reply, a tool round-trip, or a training step landing mid-turn crosses it. `TURN_TIMEOUT_S` now reads `OPENCLAWRL_TURN_TIMEOUT_S`, defaulting to 600 so the reference path is unchanged. The runtime notes carry the measurement and the two costs that share that headroom — a training backward, and a candidate gate's probe, which runs on every step including the ones it rejects, since the probe is what produces the decision. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01SuEa3FyrkW1xc3sjmAJSEB
The first MLX record here that runs the benchmark the way it is defined: the real Hermes agent in a Harbor container, a judge service beside it, reef-eval sequencing 72 sessions, against a host-native Reef+MLX service — no Mac-native driver standing in for Hermes — with the candidate gate on. What it establishes is a runtime result. The reference harness kills a turn at `TURN_TIMEOUT_S` and the benchmark records the dead turn as a `no-reply` failure, so on a runtime whose turns take 397s at the median the 600s default converts serving latency into policy failures. Two runs differing in that one constant: 28 of 72 sessions empty at 600s against 3 of 72 at 1200s, one-sided Fisher p = 1.6e-07. The gate's own probe is the independent check — the same weights answered 8 of 8 pinned prompts at every step of both runs, so the empty sessions were generation outrunning the clock, not a silent policy. What it does not establish is learning, and the record says so: 2 accepts in 72 is not distinguishable from the previous run's 0, `sessions_to_adaptation` was not reached, and the probe's `gold_rate` never moved. The gate selected 20 steps and held out 47, which is the gate working, but no ungated Hermes run of this configuration exists to compare against. Also recorded: the probe scores style only while a session needs style, shown work and the answer together, which is why a rising clean rate does not carry the other two; and a gap in the criterion, where an accepted reply opens with a tool's diff output and still passes. The three operational traps that each cost a discarded run are in the README: the relative artifact-repository paths, the stream identity covering more than `--name`, and the two timeouts being coupled. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01SuEa3FyrkW1xc3sjmAJSEB
main reworked the ground this branch stands on: #432 gave Reef ownership of inference/training coordination, #446 unified training data as ATIF items, and #412 moved the scenario and runtime modules. The MLX backend is re-integrated rather than merely merged. Conflicts and how they resolved: - `reef/runtime/adapters/` and `reef/service/deploy/settings.py` are gone on main; the deletions are accepted. The MLX kind registers from `reef/train/__init__.py`, which is the direction the boundary allows. - `reef/surface/base.py` and `reef/inference/sglang/chat.py` take main's restructured versions; `RecoveryRestorer` is re-added in main's ABC idiom. - `reef/service/assembly.py` takes main's version: the function this branch had changed built a runtime from `ServiceSettings`, which no longer exists. Re-integration: - The batch path reads ATIF items. The captured tensors move from a `PolicySample`'s attributes to the item's `extra.reef.training`, where the openclawrl processor already writes them; the tests use main's own `policy_trajectory` helper so both sides agree on that layout. - `MLXServingRuntime` (new) owns the parameters, the admission gate and the served version; `MLXRuntime` keeps candidates, the objective and checkpoints, and reaches the engine through it. The factory returns the pair #432's contract names. The seam follows the state: three methods needed both halves, and all three have the same shape — hold serving, touch the engine, release — so serving exposes an explicit `hold`/`release` rather than a context manager, because activation deliberately stays held until Reef's commit. - `MLXInferenceBackend` implements both `InferenceHandler` and `InferenceBackend`; a receiver that runs as a separate service splits those, an in-process engine does not. The receiver operations that would mean moving memory Reef does not own (`offload`, `onload_*`, `unload_adapter`) raise instead of pretending, which is what the base class does for `load_adapter_files`. - `assistant_message.py` moves under the MLX backend: main's boundary test counts it as an unowned module in `reef/runtime/`, and the MLX backend is its only consumer. One judgement worth flagging: the shipped-config layout test now skips `results/` directories. A results directory records the deployment a past run used; migrating it to the current schema would make it a config that was never executed, which is worse than an old one, because the record's whole value is that it is what ran. Verified: 3528 passed, 88 skipped (slime/ray/inference modules excluded — those dependencies are absent here). mypy clean, ruff 0.14.7, isort 5.13.2, black 24.3.0, and the Python design and statement policies pass. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01SuEa3FyrkW1xc3sjmAJSEB
… activates A live run on the merged tree found what 3528 passing tests did not. **Nothing could start.** #432 replaced the public layout, so the MLX runtime's config keys were no longer consumed: the deployment selects a runtime through `training.backend` now, and its options arrive as `training.options`. This adds `MLXDeployment(InProcessTrainingDeployment)` and registers `mlx` as a built-in backend; the base class already refuses the Ray, executor and inference-engine settings an in-process integration has no use for. **Candidate evaluation is a factory class.** `CandidateEvaluationPluginFactory` replaced the plain `build` function, and it receives both runtimes. That correction matters here: the probe applies a candidate's weights and generates from them, so it belongs to the *training* half — the serving half deliberately never sees a candidate before Reef selects it. **Activation belongs to the half that serves.** `activate_candidate` is declared on `InferenceRuntime`, so Reef calls it there, and the split had left it on the training runtime holding the staged snapshots. Every test passed, because they called it directly on the training object rather than through the caller Reef uses; the first live training step failed with "MLXServingRuntime does not support candidate activation". The serving half now owns the staged snapshots — it is the one that adopts them — and training stages and discards through it. Two regression tests assert the seam as Reef holds it. Verified by running it: four optimizer steps on Qwen3.5-9B, `rows=8` each, `adapter_tensors_changed=304` (every LoRA tensor moving, so the gradient path is intact), `adapter_delta_l2` 0.84-1.02 (the pre-merge run's first step was 0.837), the gate deciding as `regression-gate`, the probe's `answered_rate` at 1.0, and the served version advancing on every activation. Suite: 3530 passed, 88 skipped. mypy, ruff, isort, black and the design and statement policies pass. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01SuEa3FyrkW1xc3sjmAJSEB
…struction The repository had no MLX deployment anyone could run. The reference `serve.yaml` is the seven-GPU stack, and the two configs beside the run records are archives in the pre-#432 schema, deliberately left as what those runs used. So the current shape had to be rediscovered from startup errors, one field at a time. `serve-mlx.yaml` is that shape, and it is the file the verification run used: `training.backend: mlx` with the runtime's options under `training.options`, the model named in `inference.model-path`, and the candidate gate in `evaluation`. The comments carry the reasons that are expensive to rediscover — why `in_proj_qkv` is adapted, why a cached prefill is ruled out, why the judge needs a large token budget, and that relative storage paths make two runs from one checkout collide on the artifact repository. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01SuEa3FyrkW1xc3sjmAJSEB
Three CI failures, one of which reverses an earlier call of mine. `isort` reformats `mlx_backend/runtime.py`: the hook runs with `--profile=black` rather than the pyproject settings I had been checking against, and the two disagree on collapsing one import. Running pre-commit itself, rather than the tools with my own arguments, is the check that matches CI. The 2026-09-14 record linked the candidate gate one directory too high. The archived run configs are migrated to the current schema. I had excluded results directories from the layout check instead, on the grounds that a record should stay the config that ran. The config tests then failed for the better reason: the repository requires every config under `recipes/` to load with the current code, which is exactly what stops a file from looking authoritative while being unrunnable. Weakening that for two archives costs more than translating them, and the translation changes no value - only the spelling and the section each setting sits in. Both READMEs now say the run used the schema of its date. The layout check covers results directories again. Verified with `pre-commit run --all-files` (all hooks pass) and the suite: 3530 passed, 88 skipped. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01SuEa3FyrkW1xc3sjmAJSEB
`serve-mlx.yaml` went into the cookbook *training* config set, which selects
slime deployments by `training.backend == "slime"`; this one names `mlx`. It
belongs in the example-deployment set, which selects on carrying a `reef`
section.
The mistake survived my own run because `test_sao_configs.py` opens with
`importorskip("torch")` and this machine had no torch, so the file I was
editing never executed. Installing torch runs it — and 55 further tests that
were quietly skipping with it.
Suite with torch present: 3585 passed, 67 skipped. The 30 remaining failures
and errors are all `No module named 'slime'`, which this machine does not have
and CI does; none has another cause.
Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01SuEa3FyrkW1xc3sjmAJSEB
13 tasks
…oaching its own score Two defects found by running the gated OpenClaw-RL stream and then re-running the pinned probe offline against the base weights and every saved candidate. **The deployment's `kl_coef` never reached this recipe's objective.** `_run_training` spent it only on the generic loss path, shaping advantages the way TTT-Discover's `incorporate_kl_penalty` does. OpenClaw-RL's recipe always names the `openclawrl` family, whose objective took its coefficient from a separate default of zero — so a config setting `kl_coef: 0.05` trained with no KL term at all, and the run's own metrics reported `kl_coef 0.0` while the deployment asked for 0.05. The key now defaults the objective's coefficient, so it prices drift through whichever mechanism the active family carries: the openclawrl loss adds its own KL term, as the slime reference does. An explicit `openclawrl.kl_coef` still overrides. Both recorded MLX runs trained without the term their configs asked for; their READMEs now say so. **The probe stated the style rule it was scoring.** Its instruction ended "do not use bold, headings, bullet points, or numbered lists", which the student persona never says up front — so the probe measured compliance with its own instruction rather than what the policy had internalised, and a model could score well on it while opening every session with markdown. Uncoached, the baseline falls from 0.125 to 0.0, which is the number that agrees with the 67-of-72 markdown rejection rate. **At `max_tokens: 96` the probe scored truncation.** The model opens with a preamble, so the budget ran out before the arithmetic: `no-shown-work` fired on 50 of 50 replies at 96 against 13 of 50 at 320, where replies reach a median 1056 characters against the 770 a real session's first reply runs to. The corrected probe shows a hard floor rather than a resolution problem — all three style markers appear in 100 of 100 replies, base and candidates alike — so the evaluator now also reports `mean_violations`, the distance to the criterion, which moves while `clean_rate` is pinned at zero. It is reported, not gated on. Verification: `pre-commit run --all-files` passes. 3310 tests pass; the 9 failures and every collection error are `No module named 'slime'`, which this environment does not have. The two new MLX runtime tests were confirmed to fail on the unpatched code (`assert 0.0 == 0.05`). mypy's configured 3.10 run is blocked in this venv by a numpy stub before it reaches any Reef file; forced to 3.13 it reports 10 pre-existing errors, none in the files changed here. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01SuEa3FyrkW1xc3sjmAJSEB
…raining pairs OpenClaw-RL trains on a reply bound to the student reaction that judges it. `SessionIndex` drops a session after its idle TTL, and an expired session can never bind the two -- `processor.ingest` runs expiry before observing, so the next request opens a fresh session with nothing to bind back to. Hermes runs each turn as a fresh process starting from `[system, user]`, so no cross-turn request extends the previous one and the session tag is the only thing carrying that link. Once the session is gone, nothing recovers it: the reply is never judged and never trains. The MLX example set `session-ttl-s: 45` against a documented default of 900. That holds on the reference's timing, where a reply comes back in seconds. On this runtime a single generation takes over two minutes and a training step blocks the engine for longer still, so the session routinely expired between a reply and the reaction to it. Measured against a live run of this configuration, two independent readings agree: 15 of 31 consecutive inference requests were further apart than the 45s window (48%), and the processor reported 12 sessions expiring unbound against 9 binding (57%). About half the method's training pairs were being discarded before they were judged. Raised to 1800, which covers a generation, the student's round trip, and a training step in between. The cost is that a finished session waits that long before its last turn retires, which delays a step; the short value silently dropped the data instead. Both recorded MLX runs carry this setting, so their results are what the method does on roughly half its intended signal. Their READMEs now say so -- this is the first thing to re-check before reading either run as a property of the method. Verification: `pre-commit run --all-files` passes, and `test_user_facing_example_deployment_resolves[serve-mlx.yaml]` still resolves the config. The archived run configs keep their original values; only their READMEs change. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01SuEa3FyrkW1xc3sjmAJSEB
… that invalidated them Both recorded MLX runs trained with no KL term and on roughly half their intended training pairs. Neither number was a property of the method, so they are withdrawn rather than left standing with a correction note: a reader reaching for a result should not find one that measured a defect. Replaced by a defect record covering the three faults running the path turned up, each with the measurement that established it: - `kl_coef` never reached the objective this recipe selects. The runs' own step metrics reported `kl_coef 0.0` against a configured 0.05. - The style probe stated the rule it was scoring, so it measured compliance with its own instruction; uncoached, the baseline falls from 0.125 to 0.0, which agrees with the 67-of-72 session-level rejection rate. At `max_tokens: 96` it also scored truncation -- `no-shown-work` on 50 of 50 offline probe replies against 13 of 50 at 320. - `session-ttl-s: 45` against a documented default of 900 retired sessions before the student reaction that judges a reply arrived. 15 of 31 consecutive requests were further apart than that window, and sessions expiring unbound ran level with those binding. The record also carries the first 21 sessions of a run with the first two fixes active and the TTL still short, because that run is the evidence they hold and its flat probe is what the third defect predicts. It is explicitly not a learning result: 25 optimizer steps move `||lora_b||` from 0.84 to 8.43 and move the probe not at all, and the three accepts are all 2-turn sessions, which fits an easy-problem subset as well as a changed policy. The control that would separate them has not been run and the record says so. The seven-GPU reference run is untouched -- it does not take the MLX path, and it remains the recipe's reference result. Verification: `pre-commit run --all-files` passes; `test_user_facing_example_deployments_are_discovered` updated for the two removed configs and passing with the rest of that suite's 16 discovery/resolve cases. The 9 remaining failures in `test_sao_configs.py` are `No module named 'slime'`, unchanged by this commit. Link targets and the README anchor checked by hand. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01SuEa3FyrkW1xc3sjmAJSEB
…t change what the policy learns Adds the controlled comparison the earlier record was missing: two runs differing only in `session-ttl-s`, same KL term, same corrected probe, same timeouts, same base weights, same stream in the same order. The fix works on its own terms, on two measurements that do not share a failure mode. No request gap comes near the 1800s window against 22 of 54 exceeding the old 45s one, and the same agent work yields roughly five times the optimizer steps (14 by session 8, against 3). It changes nothing about learning. Over the first 26 steps of each arm the probe averages 3.306 on half the training pairs and 3.293 on all of them -- a difference of 0.013 against a step-to-step jitter of ±0.13, with `clean_rate` 0.000 on every step of both. So the TTL was a real defect and was not the reason the policy does not learn, which is worth recording as plainly as the fix itself. Also records what the 1800s arm rules out: the objective is not starved of direction (139 positive advantages against 76 negative), and the newly active KL term is not pinning the weights back (`||lora_b||` grows 0.84 to 8.43, read from the safetensors). Two readings this file previously carried are withdrawn. The "easy problem subset" account of the accepts is disproved by s006, which accepted on one arm and failed on the other from the same base weights -- sampling noise, and the 2-turn correlation runs the other way, the session being short because the reply came out clean. And the 86/14 tool-versus-student split of judged turns, though measured, cannot explain the failure: the seven-GPU reference run drives the same stream through the same harness and so carries the same split, and it adapted at session 14. What differs from the run that learned is left stated and unisolated: a 4B policy trained full-parameter against a 4-bit 9B through a rank-64 LoRA, 16 judged turns per step against 8, lr 1e-5 against 3e-5, local judges against hosted ones. Verification: `pre-commit run --all-files` passes. Data regenerated for both arms from the runs' own commit records and lab verdicts. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01SuEa3FyrkW1xc3sjmAJSEB
9 tasks done
2 of 4 tasks
BobbyZhouZijian
requested review from
MikeZheng777 and
quao627
as code owners
October 14, 2026 11:05
…terfaces - MLXRuntime takes a TrainingMethod and the recipe's StepScheduling, resolves the objective, and refuses a learning-rate schedule it cannot apply. - The durable-commit handshake moves to commit_candidate; reject_candidate no longer takes a selection decision. - The OpenAI-dialect judge settings (prm_api, prm_model, prm_api_key_env) move to OpenClawRLJudgingSettings and prm_clients. - The candidate evaluator becomes one ConfiguredEvaluationPlugin over RegressionCheckMixin, and requires the MLX training runtime. - Recovery restores the head through Surface.restore_recovered. - serve-mlx.yaml uses the pipeline section; MLX tests move to tests/mlx_backend. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
github-actions
Bot
added
the
area: ci
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
- Add docs/user-guide/mlx.rst: install, configuration, training options, loss families, serving, the published adapter, sizing and the example. - Link it from configuration, installation, evolve-your-model, the OpenClaw-RL recipe page and the example README; document the OpenAI judge dialect settings. - Replace the runtime notes and the defects page with one result page for the session TTL comparison; rename the CSVs away from "gate". Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
- Fix the gated-delta fallback recursion, the inexact cached-prompt gradient check, and the non-atomic adapter save. - Losses resolve per family with backend "mlx"; the TTT-Discover and OpenClaw-RL losses move to recipes/tttd/mlx.py and recipes/openclawrl/mlx.py. - Split the engine into config, contracts, training and engine modules; remove the unused continuous batcher and inference backend half. - Serving uses the base admission and publication contract; recovery reloads the adapter through activate_checkpoint, and the unused RecoveryRestorer interface is gone. The mlx kind is a lazy builtin. - training.options is a typed config that refuses unknown keys; method settings move to loss_options. - Share the assistant-message split and tool-call types with reef/inference/chat.py without changing SGLang or vLLM behaviour. - The OpenClaw-RL judge takes a transport (SGLang or OpenAI chat) and reads the credential from the deployment environment; the candidate evaluator validates its settings at the entry point. - Comments state contracts; tests share one fake engine. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
reef/train may not import reef/inference, so the MLX backend could not reuse the assistant-message split, the reasoning stream splitter and the SSE writers from reef/inference/chat.py. Move these engine-independent helpers to reef/core/chat_protocol.py; the token-native chat handler, SGLang, vLLM and MLX import them from there. Behaviour does not change. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
4 of 6 tasks
Resolve two conflicts and follow #819 in the MLX backend: - OpenClawRLRecord2Record no longer takes the pipeline context (#819); keep this branch's environ argument. - Keep both import sets in tests/artifact/test_reef_git_lfs.py. - Rename the prepare_training_step keyword serving_runtime_load_id to parent_runtime_load_id, as TrainingRuntime now declares it. - Remove MLXConfig.max_staleness. #819 removed the runtime staleness option; the record-to-record stage applies recipe.config.max_staleness. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
The harness used the reef_url host for two things: the shim's upstream on the host, and the shim address that hermes calls from the task container. On Docker runtimes such as colima, a container cannot reach the host's LAN address, and the host cannot resolve host.docker.internal. No single reef_url worked, and every hermes call failed with a connection error. Add the optional shim_host agent argument. Hermes calls the shim at that host; the shim keeps reef_url as its upstream. Without shim_host, the behavior does not change. The README describes the setup for these runtimes. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
An OpenClaw-RL stream on a 48 GB Mac lost the Reef process at step 17. Hermes asks for 8192 completion tokens, and its prompts grow during a session; one 12288-token row made the training backward pass exceed the machine's memory. The cost of one row grows faster than linearly with its length (13.5, 26.2 and 35.9 GiB at 4096, 8192 and 10240 tokens on the 9B config), and fewer adapted layers do not remove it (45.5 GiB at 12288 tokens with 16 layers). training.options.max_context_tokens bounds the prompt and completion tokens of one request, as rollout-max-context-len does in the GPU reference. A request's completion limit is lowered to fit; a prompt that reaches the limit fails with HTTP 400. Rows that Reef trains come from these requests, so the limit also bounds a training row. 0 keeps no limit. serve-mlx.yaml sets 10240 and raises inference.timeout-s to 1800: requests wait while a long training step holds the engine, and the 300-second default failed them with HTTP 503. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
On the MLX GSM8K stream, the policy stopped answering in text after about 20 steps and ran tool calls that did nothing for the student: background terminal commands such as "ymd", mkdir, and rm of the homework file. Hermes turns then hit their timeout. The upstream judge prompt scores any tool output without an error as +1. Re-judging recorded turns with the run's own PRM path gave +1 to 28 of 34 such tool calls and to 7 of 20 text replies, which the student judges. The reward term therefore moved the policy from text replies to tool calls. - pipeline.record2record.config.tool_turn_reward: 'prm' keeps the upstream judging. 'neutral' gives a turn whose next state is a tool result no reward, so only an accepted hint trains it. - The candidate evaluation gets a reply probe. reply_probe_file names hermes conversations where the policy has read the homework and should answer. decide() refuses a candidate whose reply_rate drops more than reply_margin below the best selected one, before the clean_rate check. The style probe has no tools and could not see this collapse. On the run's checkpoints, the base weights answered 8 of 8 and the step-20 adapter 0 of 8. - serve-mlx.yaml sets tool_turn_reward: neutral and the reply probe; reply_probe.json holds stream problems 64-71 in a recorded hermes conversation. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
This branch has not been deployed
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Motivation
Reef's weight-training recipes need a GPU cluster today. The OpenClaw-RL reference stack uses seven GPUs: Megatron on four, a rollout engine, a PRM and a 32B student. This pull request adds an in-process MLX backend, so one Apple Silicon Mac can serve and train a LoRA adapter on the traffic that Reef records. The PRM and the student only need text, so they can use a hosted model.
Changes
reef/train/mlx_backend/).training.backend: mlxselects it. One process holds the base model, the adapter and the optimizer. Serving and training take turns on the same weights.config.py:training.optionsis a typed config. Unknown keys and wrong types fail at startup.engine.py,training.py: model loading, generation, adapter files, and the generic training mechanisms (chunked output head, gradient checkpointing, prompt prefill, optimizer step).gated_delta.py: a chunkwise GatedDeltaNet recurrence with a backward pass, so the gradient crosses hybrid layers.serving.py,inference.py:/v1/chat/completionswith tools, reasoning and streaming, one sequence at a time. Serving uses the base admission and publication contract.activate_checkpointreloads the committed adapter after a restart.runtime.py: candidates are trained, exported and staged without changing serving.probe_candidatelets an evaluator generate from an unselected candidate.mlxruntime kind is a lazy built-in reference, like the other bundled kinds.register_loss_family_ref(name, reference, backend="mlx"), as Tinker does.recipes/tttd/mlx.pyandrecipes/openclawrl/mlx.pyhold the two losses. The OpenClaw-RL loss is checked against the torch kernel.training.options.loss_optionsholds the method settings.PRMJudgetakes a transport:SGLangTransport(the upstream payload, unchanged) orOpenAIChatTransport(any OpenAI-compatible endpoint). New settings:prm_api,prm_model,prm_api_key_env. The credential comes from the deployment environment.recipes/openclawrl/candidate_evaluator.pyprobes each candidate on fixed GSM8K problems and appliesRegressionCheckMixintoclean_rate. It needs the MLX training runtime.reef/trainmay not importreef/inference. The assistant-message split, the streamed reasoning split, the SSE writers andParsedToolCallmove fromreef/inference/chat.pytoreef/core/chat_protocol.py. The token-native handler, SGLang, vLLM and MLX use them from there. Their behaviour does not change.recipes/openclawrl/examples/openclawrl/serve-mlx.yamlruns the GSM8K stream on a Mac.results/2026-09-16-mlx-session-ttl/comparessession_ttl_s45 s against 1800 s.Defects found and fixed while the backend was built. These are not in the docs:
kl_coefdid not reach the OpenClaw-RL loss, so step metrics showedkl_coef 0.0against a configured 0.05. The loss now readstraining.options.kl_coefunlessloss_options.kl_coefis set.max_tokens: 96it also scored truncation: the shown-work check failed on 50 of 50 replies at 96 and on 13 of 50 at 320. The probe no longer states the rule, and the default is 320.session_ttl_s: 45was below the gap between a reply and its reaction on MLX (15 of 31 gaps were longer). About half of the training pairs expired unjudged. The example now uses 1800.prefill_step_sizenow needslora_layers: 1and no key, value orlinear_attnadapter. A test checks gradient equality.MLXServingRuntime.activate_checkpointnow does the reload through the existing hook.Documentation impact
docs/user-guide/mlx.rst(in the site navigation);docs/reference/configuration.rst,docs/getting-started/installation.rst,docs/user-guide/evolve-your-model.rst,docs/developer-guide/write-a-recipe.rst,docs/contributing/codebase-structure.rst;docs/user-guide/recipes/openclawrl.rst(Apple Silicon section, judge dialect settings);recipes/openclawrl/examples/openclawrl/README.md; the session TTL result page.npm ci,npm run check:docs,npm run lintandnpm run buildpass indocs/sitewith Node 26.8.1. CI uses Node 22, which I did not run locally. The option names and defaults on the MLX page were checked againstreef/train/mlx_backend/config.py. Thereef serve -c serve-mlx.yamlrun and the harness commands in the example README were not run after this change.Compatibility and operational impact
mlx(mlx>=0.32,<0.33,mlx-lm>=0.31,<0.32), with environment markers for macOS arm64. A CPU or CUDA installation is unchanged; Reef imports without MLX.mlxand runtime kindmlx.pipeline.record2record.config:prm_api(defaultsglang, unchanged behaviour),prm_model,prm_api_key_env.reef.core.chat_protocol, andParsedToolCallis no longer inreef.inference.vllm.chat. No wire contract changes.Verification
pytest tests -q -p no:cacheprovider --ignore-glob='*sglang*'with torch, Ray and the pinned Slime source onPYTHONPATH, on macOS arm64 with MLX installed: 6357 passed, 146 skipped.pyteston everytests/*sglang*module: 117 passed.tests/mlx_backend/ran on this Mac, including the torch parity tests, the gradient-path tests and the cached-prefill gradient-equality test onmlx-community/Qwen2.5-0.5B-Instruct-4bit.mlxandmlx_lmimports blocked, as on Linux CI: Reef, the recipes and the MLX config and runtime modules import, and the MLX-only test modules skip.pre-commit run --all-files: pass.python -m mypy: no issues in 390 source files.AI assistance
Claude Code wrote most of this change with the author: the MLX backend, its refactor into the current structure, the review against
AGENTS.md, the documentation and the tests. The author remains responsible for every line, claim and result.Checklist
AGENTS.md.review checklist in
AGENTS.md. Documentation impact describes the changesand verification, or explains why no documentation update is needed.
README.i18n.yaml, or do not affect README content.Review and merge
Reviewers are requested automatically from
CODEOWNERS. See themaintenance model
for review and merge requirements.
🤖 Generated with Claude Code