[RFC] Add a single-host MLX path for Apple Silicon serving and training

[BobbyZhouZijian]

Follows #134, which proposed the feature; this issue is the RFC and its decision record. I ran the experiments below on an M4 Pro / 48 GB / macOS 26.6.1 with mlx 0.32.2, mlx-lm 0.31.3 and mlx-tune 0.6.0 on Python 3.14.7; every number quoted is measured, not estimated.

Primary area

Training and runtime (also touches Service/API and Artifact management).

Summary and decision to make

Reef should gain one Apple Silicon path that covers both serving and training, built directly on mlx and mlx-lm, behind the existing TrainingRuntime and TrainingBackend contracts.

The smallest durable decision that needs agreement is: Reef's service assembly stops hardcoding a single training runtime kind. reef/service/assembly.py:83-121 builds {"type": "ray_training", ...} for every WeightTrainingRecipe and requires reef.ray_address at line 91. Until that is a selection rather than a constant, no non-Ray weight-training deployment can be assembled at all, on any hardware.

Two smaller decisions follow from it:

1. SlimeTrainingBackend is already runtime-generic — reef/recipe/base.py:335 uses it for every weight-training recipe, and its only Slime-specific line is the "runtime": "slime" label it reports in experiment_config (reef/train/slime_backend/backend.py:47). Rename it RuntimeTrainingBackend (keeping an alias) so a second runtime is not forced to import a package named after the first.
2. The MLX path consumes Reef's existing externally-supplied batches and does not generate rollouts inside the training process. This differs from what [Feature] Add a native MLX training backend for Apple Silicon #134 proposed; the reasoning is below.

Motivation and evidence

Mac-local development is currently impossible, and the blocker is not MLX. It is the service layer. _connect_training_runtime accepts no runtime kind; the only path from a WeightTrainingRecipe to a runtime goes through Ray. A contributor on Apple Silicon cannot boot the model-evolution loop even with a working local trainer in hand.

The serving half is already solved by mlx-lm, better than I expected. mlx_lm.server is an OpenAI-compatible server (/v1/chat/completions, /v1/completions, /v1/models, /health), and it supplies exactly the two things Reef's weight surface needs. Measured:

• Token-native rollout log-probabilities. With {"logprobs": true} the response carries token ids, not strings:

  logprobs.content: [{"id": 19, "logprob": -0.4375}, {"id": 17, "logprob": 0.0},
                     {"id": 16, "logprob": 0.0}, {"id": 151645, "logprob": -0.03125}]
  

  That is PolicySample.tokens and PolicySample.rollout_log_probs directly, with no re-tokenization. top_logprobs: N additionally returns {id, token, logprob} triples, matching topk_indices / topk_log_probs.
• Per-request LoRA addressing. The request body's adapters field selects the adapter, and the server keys its model cache on (model_path, adapter_path, draft_model_path). It is genuinely wired: a nonexistent path returns HTTP 404 {"error": "The adapter path does not exist: ..."}, a real one loads and serves. This is what WeightInferenceHooks.prepare_request needs to guarantee no harness silently samples the frozen base.

This changes the shape of the training half. #134 proposed adopting mlx-tune, whose GRPO trainer samples its own completions inside the training process, and acknowledged that this conflicts with how Reef models topology. It does conflict — tests/plugin_contracts/test_slime_backend_scope.py pins the invariant that "Reef supplies batches externally and never calls the runtime's rollout generation package". Since mlx_lm.server can hand back token-native log-probs from real served traffic, that conflict is avoidable rather than something to argue about: the MLX backend can consume the same externally-reserved batches every other Reef backend consumes.

recipes/tttd is the precedent. It already assembles served rollouts into GroupedPolicyBatch comparison sets, drops constant-reward groups for lack of relative signal (recipes/tttd/processor.py:158-165), and runs a group-relative objective. A group-relative MLX objective needs no new data path — only a new execution backend.

On mlx-tune. #134 recommends it for "native MLX LoRA/QLoRA and GRPO support, tests, tagged releases, checkpointing". I evaluated it against a from-scratch mlx+mlx-lm implementation on an identical workload (Qwen2.5-1.5B-Instruct-4bit, LoRA on all 28 blocks, group size 8, 256 max new tokens, chat template applied, same reward function, 4 steps):

	s/step	peak unified memory	installed packages
mlx-tune GRPOTrainer	5.7	11.22 GB	75
mlx + mlx-lm only	3.0	10.84 GB	35

The 40 extra packages include mlx-vlm, mlx-audio, mlx-embeddings, opencv-python, fastapi, uvicorn, sounddevice, miniaudio, pandas, pyarrow, datasets and scipy. The speed gap comes from batching: mlx-lm's BatchGenerator generates the group in one continuous batch, mlx-tune loops the group sequentially over forked KV caches.

More decisive than the numbers, four behaviours are incompatible with the lifecycle #134 itself asks for:

1. A degenerate step reports success. When a group's reward standard deviation is below 1e-8 the optimizer update is skipped, but the trainer still writes an adapter and returns {"status": "success"}. Observed on a 4-step run: all four steps skipped, adapter L2 delta 0.000000 across 0/112 tensors, success returned. An unchanged adapter would reach Reef as a trained candidate.
2. No optimizer state between calls. AdamW and the cosine schedule are constructed inside train(). Driving it one Reef step at a time resets the moments and the schedule every step. Three consecutive single-step calls produced losses of +1.87, -23.59 and -7.05 but parameter deltas of 0.232135 and 0.232173 — the signature of a fresh optimizer's first step each time.
3. Checkpointing is not restart-safe. _save_adapters_and_config overwrites one directory in place with no temporary file, no rename, and no step-tagged output; it wraps everything in try/except, prints a warning and returns False on failure, and no caller checks the return value. [Feature] Add a native MLX training backend for Apple Silicon #134's "cannot expose a partially written adapter as active" cannot be met by delegating to it.
4. Silent degradation. use_native = use_native and HAS_NATIVE_TRAINING; if native training is unavailable it falls back to a subprocess mlx_lm.lora SFT run without failing.

Separately, the objective is not the algorithm it advertises. The native loss is -advantage * log_prob with no KL term, no clipping and no reference model; beta is stored and never used, and loss_type (grpo / dapo / dr_grpo / bnpo) is not read anywhere in the native path, so all four settings compute the same thing. In fairness to upstream, one thing that looks like a defect is not: the log-prob includes the prompt tokens, but because group advantages are standardized and group members share a prompt, the prompt term cancels — I measured cos(mlx-tune gradient, completion-only gradient) = 0.999999. The real deviation is that it sums rather than averages per token, i.e. the usual length bias.

Finally, mlx-tune's adapter output is adapters.safetensors + an adapter_config.json holding fine_tune_type / num_layers / lora_parameters. This is not a mlx-tune quirk — it is mlx-lm's adapter format, which mlx_lm.tuner.utils.load_adapters reads. It is not Hugging Face PEFT, so Reef's existing admission policy rejects it: running reef.artifact.peft.PEFTValidator against a produced adapter gives AdapterArtifactError: adapter_config.json must declare a peft_type. Whatever we build on, Reef needs an MLX adapter validator or a conversion step.

Goals and non-goals

Goals:

• A clean checkout on Apple Silicon installs an optional MLX extra, with no CUDA, Ray or Slime, and runs the full model-evolution loop: serve, record, batch, train, evaluate, select, publish, activate, roll back, recover.
• Serving and training share one base model and one adapter format on one host.
• The MLX backend stays behind TrainingBackend / TrainingRuntime; no MLX concept enters backend-neutral recipe, processor or scenario APIs.
• Every published adapter carries base-model, tokenizer, quantization, dtype and library-version provenance, and cannot be activated unless its bytes are complete.

Non-goals:

• Feature parity with Slime; multi-node or multi-Mac distribution; asynchronous or disaggregated rollout.
• Making SGLang mandatory, or depending on SGLang's MLX backend in this iteration.
• Supporting every mlx-lm architecture or every quantization mode in the first release.
• Changing anything about the CUDA/Slime path's behaviour.

Proposal

Serving. A new runtime kind mlx registered through reef.runtime.registry.register_runtime_kind, implementing TrainingRuntime and composing an InferenceBackend that speaks to a local mlx_lm.server. It implements serving_adapter_name() / serving_adapter_runtime_load_id() so the existing weight surface addresses every request to the scenario's published adapter, and restore_checkpoint() so rollback is a real capability rather than the base class's refusal.

A companion inference backend (selected via the existing reef.inference_backend_factory hook, the same seam SGLangChatTrainingInferenceBackend uses) requests logprobs and translates the response into PolicySample.tokens and rollout_log_probs. Reef never re-tokenizes a completion.

Training. The backend consumes PolicyBatch / GroupedPolicyBatch reserved by Reef, exactly as the Slime path does. One prepare_step is one optimizer update over one reserved batch. The optimizer lives with the backend across steps, so Reef's step N+1 continues the moments and schedule step N left behind. Rollout generation stays out of the training process; test_slime_backend_scope.py's invariant holds for the MLX path too.

The initial objective is group-relative policy gradient over GroupedPolicyBatch comparison sets, with per-token length normalization and an explicit importance-ratio correction against rollout_log_probs for stale samples. A degenerate group (no reward variance) produces PreparedStep.skipped(), never a candidate.

The backend runs in the dispatcher's existing local worker (dispatched = False, reef/dispatcher.py:354), the same execution position CordisBackend uses. No Ray, no second process.

Publication. A candidate is written to a staging directory, fsynced, and renamed into place; the adapter directory a reader sees is either the previous complete one or the new complete one. The directory holds adapters.safetensors and adapter_config.json in mlx-lm's format so the stock loader can read it, plus reef_provenance.json carrying base model, tokenizer, quantization, dtype, LoRA parameters, library versions, and the rollout implementation and policy version that produced it. Before publication the backend measures the adapter's parameter delta against the pre-step snapshot and refuses to emit a candidate whose weights did not move.

Failure behaviour. Unsupported model architectures, quantization modes, algorithm settings and any distributed configuration fail during preparation with actionable errors, before generation starts. An interrupted step leaves no adapter directory in a partial state; recover_pending_step discards an unreferenced staging directory and reconciles against Reef's durable commit.

Smallest useful version. One recipe, one model family, LoRA over a 4-bit quantized base, single scenario, single host, group-relative policy gradient, no KL term.

Public interfaces and configuration

• Service config (new). A reef.runtime section selecting the runtime kind; absent, it defaults to ray_training so every existing deployment is unaffected. reef.ray_address becomes required only for that kind.
• Packaging (new). An optional mlx extra pinning a tested mlx / mlx-lm matrix. The dependency-boundary assertion in .github/workflows/ci.yml gains a corresponding exact-set check. reef/train/mlx_backend/ must import no MLX dependency at module scope.
• Artifacts (new). An MLX adapter validator alongside PEFTValidator, and the reef_provenance.json schema (reef.mlx.adapter/1) as a persisted format.
• Python API (rename). SlimeTrainingBackend → RuntimeTrainingBackend, with the old name kept as an alias.
• Unchanged: recipe extension contract, processor contract, scenario lifecycle, HTTP API, commit log format, CUDA path behaviour.

Compatibility and migration

Additive. The runtime kind defaults to today's behaviour, the extra is opt-in, and the rename keeps an alias. No existing artifact, commit log, config file or deployment needs migration. Rollout is by stage (below); rollback of any stage is removing the new kind from config. No deprecation is proposed for the Slime path — it stays the production path, and the MLX path is explicitly a development and small-scale one.

Security, privacy, and trust boundaries

mlx_lm.server should bind to loopback; the RFC does not propose exposing it. No new credentials or network egress beyond the model download the deployment already performs.

One trust boundary needs care and is a change from the CUDA path: mlx_lm.server accepts an adapter filesystem path in the request body. Reef must set that field itself in prepare_request and must never forward a client-supplied adapters value, or a caller could make the server load arbitrary local directories. The MLX inference backend should strip the field from incoming payloads before adding its own.

Adapter artifacts contain trained weights derived from recorded traffic; retention is governed by the existing artifact repository, unchanged.

Operations and observability

Measured cost on an M4 Pro, Qwen2.5-1.5B-Instruct-4bit, LoRA over 28 blocks, group size 8, 256 max new tokens: 3.0 s per training step, 10.84 GB peak unified memory, with generation dominating. A 0.5B 4-bit base with two prompts × a group of 8 at 256 tokens runs at 2.8 s/step and 7.72 GB; memory scales with the number of sequences held in the batched forward, not with model size. Unified memory is the binding constraint, so generation length, group size, adapted layer count and checkpoint cadence must all be configurable and bounded.

Per-step metrics to emit: reward mean and standard deviation, fraction of degenerate groups skipped, completion tokens generated, adapter parameter delta, peak memory, and step wall clock. The adapter delta is not decoration — it is how an operator distinguishes a real update from the silent no-op described above.

Known operator-visible failure modes: a degenerate batch producing no candidate (expected, logged); an unsupported architecture failing at preparation; unified memory exhaustion under a group size or generation length set too high.

Testing and documentation

Without Apple hardware (ordinary CI):

• Backend lifecycle contract tests against a fake runtime: initial_state → prepare_step → evaluate → settle_step / abort_step, restart recovery, and the rule that a partially written adapter can never be activated.
• Atomic publication tests: interrupt between staging and rename, assert the reader sees the previous complete adapter.
• A degenerate-reward batch yields a skipped step and no candidate.
• Provenance round-trip and adapter-validator admission and rejection cases.
• The packaging boundary test: reef.train.mlx_backend is importable without mlx installed, and importing it does not import MLX.

Gated Apple Silicon qualification, following #27's lane and #22's evidence requirements: real MLX generation and optimization, checkpoint resume, artifact round-trip and reload for inference, rollback, memory ceilings, and a reported learning curve with quality, throughput, peak unified memory and known limitations.

Documentation: a new Apple Silicon section in docs/user-guide/evolve-your-model.rst documenting the single-host topology honestly, including what it does not do, and the migration path to a Slime deployment.

Alternatives considered

• Adopt mlx-tune as [Feature] Add a native MLX training backend for Apple Silicon #134 proposes. Rejected on the evidence above. Concretely, it would save roughly 85 lines of the ~290 in my prototype — the generation loop, the log-prob helper and the advantage computation — while the adapter publication, provenance, optimizer ownership, delta verification and degeneracy guard all still have to be written, and the resulting wrapper bypasses GRPOTrainer.train() anyway. In exchange it costs 40 extra packages, an objective that is not the advertised algorithm, and the four lifecycle behaviours listed above. Its _perf.py (wired-limit configuration, gradient checkpointing, compiled steps, bucketed shapes) is genuinely good work and worth reimplementing; that is a reason to read it, not to depend on it.
• Depend on SGLang's MLX backend. This is the right long-term convergence — it would give Reef the same engine on both sides. But it is in flux: sgl-project/sglang has an open redesign RFC (#32321, "Apple Silicon serving redesign") and a run of recent startup-crash and cache fixes on the MLX path. Proposed instead: keep the inference-runtime boundary explicit so this is a later substitution, not a rewrite.
• MLX training with llama.cpp/GGUF serving. Rejected: two weight formats and a conversion step, and GGUF export from a quantized base is documented upstream as not working.
• Do nothing. Apple Silicon contributors keep running trainers outside Reef, losing batch reservation, candidate evaluation, artifact lineage, activation, rollback and recovery — which is the whole point of the framework.

Risks and unresolved questions

• Unverified: adapter cache growth in mlx_lm.server. It caches models keyed by adapter path, and each Reef publication is a new path. Whether activation and rollback reload correctly rather than hitting a stale cache entry, and how the cache is bounded, needs testing before this is trusted for long runs.
• Objective fidelity. The first version has no KL term against a reference policy. Whether that is acceptable for a development-scale path, or whether a reference-policy KL is required from the start, is a decision for the maintainers.
• Length normalization. Per-token mean versus summed log-prob is a real behavioural choice (measured: gradient cosine 0.9886 between them, 8× magnitude difference). Proposed default is per-token mean; worth an explicit decision since it is not what mlx-tune or every reference implementation does.
• Prompt formatting. Rollouts must carry the chat template. Without it generations do not terminate and every step pays the full token budget — I hit this and it silently tripled step cost.
• Upstream churn. mlx and mlx-lm move quickly. The pinned matrix must be tested, and upstream API changes treated as capability changes, not silent behaviour changes.
• Scope of the invariant. test_slime_backend_scope.py is written against the Slime package specifically. If the maintainers want it to be a project-wide rule, it should be restated generically — this proposal complies with it either way.

Implementation plan and ownership

Stage 1 — service layer, no MLX. Runtime-kind selection in reef/service/assembly.py and ServiceSettings; the RuntimeTrainingBackend rename with alias; backend lifecycle contract tests against a fake runtime. Reviewable and mergeable on CPU CI with no new dependency. This is the compatibility-commitment point: it changes a public config surface.

Stage 2 — MLX serving. The mlx extra, the mlx runtime kind over mlx_lm.server, the inference backend translating token-native log-probs, and the adapter validator plus provenance schema. At the end of this stage a Mac can serve and record trainable rollouts, with no training yet.

Stage 3 — MLX training. The backend, the group-relative objective, atomic publication with delta verification, recovery, one runnable example recipe, and the Apple Silicon qualification under #27 with a #22-compliant reproduction.

Coordination: adapter publication with #14, candidate gating with #5, accelerator coordination with #28, qualification with #27.

Ownership: I intend to implement Stages 1–3 and maintain them, subject to this RFC being accepted. Maintainer review is needed most on the Stage 1 config surface and on the objective-fidelity question, since those constrain later work.

Submission checks

• I searched existing issues, pull requests, and RFCs for this decision.
• I understand this issue is the RFC and that discussion is not acceptance.
• I will keep this issue updated with material design changes and implementation links.

---

Disclosure per CONTRIBUTING.md: this proposal was drafted with AI assistance (Claude Code). The experiments — the mlx-tune evaluation, the mlx_lm.server probes, the head-to-head benchmark, the gradient-cancellation check and the from-scratch prototype — were run locally on the hardware described, and I have reviewed and verified each factual claim and measurement above.