mlxrl

Fast on-policy MLX RL for Apple Silicon, including single-turn GRPO-family RL and multi-turn GiGPO; not a general RL framework, not preference tuning, and not distributed training.

mlxrl is a small, single-process RL post-training library for LLMs on Apple Silicon. It is built around one idea: GRPO on MLX should be a fast batched rollout path with a thin loss and optimizer step on top, not a framework.

The current implementation targets QLoRA GRPO/GiGPO on local 4-bit MLX models. It reuses mlx-lm model loading, LoRA layers, KV caches, and sampling utilities, and keeps generation and training in one Python process with one model object.

mlxrl is pre-1.0. The correctness gates are stable, but import APIs and config fields may change before a 1.0 release.

Quickstart

Requires Apple Silicon, Python 3.11–3.14, MLX 0.32.2+, and MLX-LM 0.32.0+. See SECURITY.md for model trust boundaries and private reporting.

git clone https://github.com/inductiveML/mlxrl.git
cd mlxrl
UV_CACHE_DIR=.uv-cache uv sync --all-groups
UV_CACHE_DIR=.uv-cache uv run mlxrl train \
  --config examples/qwen3_0_6b_grpo.toml \
  --available-memory-gb 48

For the measured 9B-on-48GB shape, use the checkpointed G=2 config:

UV_CACHE_DIR=.uv-cache uv run mlxrl train \
  --config examples/qwen35_9b_g2_checkpoint.toml \
  --available-memory-gb 48 \
  --dry-run

What Works

• Batched group rollouts with MLX-LM KV caches and sampling.
• Full-forward old-policy logprob recompute for training-time pi_old.
• Adapter-disabled reference policy on the same model object.
• GRPO, Dr. GRPO, DAPO, and GSPO loss variants.
• RLOO (REINFORCE Leave-One-Out) as a critic-free rollout objective.
• Multi-turn agentic rollout through an Environment / EnvFactory protocol.
• GiGPO with episode-level and anchor-state step-level advantages.
• Full-forward trajectory action-logprob recompute for 4-bit correctness.
• QLoRA injection on dense and heterogeneous/hybrid layer stacks.
• Qwen3.5-style hybrid support via MLX-LM auto LoRA targeting, including DeltaNet linear_attn.in_proj_* and dense attention q/k/v/o_proj.
• Per-layer gradient checkpointing through mlx_lm.tuner.trainer.grad_checkpoint for linear-attention/DeltaNet backward memory.
• Micro-batched gradient accumulation for token-mean policy losses.
• beta == 0 reference-forward skip.
• Phase 4 benchmark harness for mlxrl, mlx-tune, mlx-lm-lora, and mlx-lm.

Install

Install the PyPI distribution:

pip install inductive-mlxrl

The Python import package and CLI command are still mlxrl:

mlxrl --help

Source installs are also supported:

UV_CACHE_DIR=.uv-cache uv sync --all-groups

Run commands through the local environment:

UV_CACHE_DIR=.uv-cache uv run mlxrl --help

Python 3.11+ is required. The MLX runtime path is supported on Apple Silicon with Python 3.11 through 3.14, matching the current MLX wheel support window. Runtime dependencies are intentionally small: mlx and mlx-lm on macOS arm64. Development dependencies include pytest, ruff, pyright, mlx-tune, and mlx-lm-lora for comparison benchmarks. The PyPI distribution name is inductive-mlxrl; the import package and console script remain mlxrl.

Quick Smoke Tests

Dense Qwen3 0.6B:

UV_CACHE_DIR=.uv-cache uv run mlxrl phase0-smoke \
  --model mlx-community/Qwen3-0.6B-4bit \
  --prompt "What is 2+2?"

Hybrid Qwen3.5 9B with rank-16 LoRA:

UV_CACHE_DIR=.uv-cache uv run mlxrl phase0-smoke \
  --model mlx-community/Qwen3.5-9B-MLX-4bit \
  --rank 16 \
  --scale 2.0 \
  --prompt "What is 2+2?"

The smoke gate prints the model id, layer count, LoRA target keys, per-layer LoRA module counts, total/trainable parameter counts, and logits shape. It fails if any trainable leaf is not lora_a or lora_b.

Training Commands

Toy hand-computed GRPO math gate:

UV_CACHE_DIR=.uv-cache uv run mlxrl phase1-toy-gate

Small built-in GSM8K-style run:

UV_CACHE_DIR=.uv-cache uv run mlxrl phase1-gsm8k \
  --model mlx-community/Qwen3-0.6B-4bit \
  --steps 20 \
  --group-size 4 \
  --max-tokens 64

Config-driven run:

UV_CACHE_DIR=.uv-cache uv run mlxrl train \
  --config examples/qwen3_0_6b_grpo.toml \
  --available-memory-gb 48

Reference multi-turn GiGPO run against the built-in recurring-state env:

UV_CACHE_DIR=.uv-cache uv run mlxrl train \
  --config examples/qwen3_0_6b_gigpo_recurring.toml \
  --available-memory-gb 48

The config schema validates model id, quant bits, group size, completion/prompt lengths, checkpointing granularity, iogpu.wired_limit_mb, optimizer settings, algorithm hyperparameters, KL beta, and seed before a model is loaded. CLI overrides such as --steps, --group-size, --max-tokens, --algorithm, --beta, and --seed apply on top of the file.

For DeltaNet / linear-attention models, enable per-layer checkpointing:

UV_CACHE_DIR=.uv-cache uv run mlxrl phase1-gsm8k \
  --model mlx-community/Qwen3.5-9B-MLX-4bit \
  --rank 16 \
  --scale 2.0 \
  --checkpoint-completion-forward \
  --steps 1 \
  --group-size 2 \
  --max-tokens 256

Despite the historical CLI name, --checkpoint-completion-forward now enables per-transformer-block checkpointing at model setup. The old whole-model mx.checkpoint(...) wrapper was removed because it does not cap DeltaNet's per-layer scan memory.

Phase 2 rollout equivalence check:

UV_CACHE_DIR=.uv-cache uv run mlxrl phase2-equivalence \
  --model mlx-community/Qwen3-0.6B-4bit \
  --group-size 4 \
  --max-tokens 32 \
  --compile-decode-step \
  --batch-groups

Import API

Minimal model setup:

from mlxrl.policy import LoRAConfig, load_policy_with_lora

model, tokenizer, report = load_policy_with_lora(
    model_id="mlx-community/Qwen3.5-9B-MLX-4bit",
    config=LoRAConfig(
        rank=16,
        scale=2.0,
        dropout=0.0,
        grad_checkpoint=True,
    ),
)

One optimizer step:

import mlx.optimizers as optim

from mlxrl.algo import GRPOAlgorithm
from mlxrl.train import batch_from_rollouts, optimizer_step

optimizer = optim.Adam(learning_rate=1e-5)
algorithm = GRPOAlgorithm()
batch = batch_from_rollouts(
    model=model,
    completions=completions,
    rewards=rewards,
    group_size=4,
    pad_token_id=pad_token_id,
    algorithm=algorithm,
    compute_reference=beta != 0.0,
)
metrics = optimizer_step(
    model=model,
    optimizer=optimizer,
    batch=batch,
    beta=beta,
    pad_token_id=pad_token_id,
    algorithm=algorithm,
    use_checkpoint=True,
    micro_batch_size=2,
)

Multi-turn GiGPO uses the additive trajectory path:

from mlxrl.algo import GiGPOAlgorithm
from mlxrl.env import RecurringStateTextEnv
from mlxrl.rollout import generate_agentic_trajectories
from mlxrl.train import batch_from_trajectories, optimizer_step_trajectory

algorithm = GiGPOAlgorithm(omega=1.0, gamma=1.0)
trajectories = generate_agentic_trajectories(
    model=model,
    tokenizer=tokenizer,
    env_factory=lambda task, seed, group: RecurringStateTextEnv(
        task=str(task),
        max_turns=3,
    ),
    tasks=["finish the task"],
    group_size=4,
    sampling=sampling,
    rollout_mode="parallel_per_turn",
)
batch = batch_from_trajectories(
    model=model,
    trajectories=trajectories,
    group_size=4,
    pad_token_id=pad_token_id,
    algorithm=algorithm,
    compute_reference=beta != 0.0,
)
metrics = optimizer_step_trajectory(
    model=model,
    optimizer=optimizer,
    batch=batch,
    beta=beta,
    pad_token_id=pad_token_id,
    algorithm=algorithm,
)

micro_batch_size=0 keeps the original whole-batch path. Micro-batching is currently exact for token-mean policy losses: base GRPO, DAPO, GSPO token mode, RLOO, and Dr. GRPO with loss_reduction="token_mean". Sequence-reduced losses should keep micro_batch_size=0, except GRPO and score centering's explicit sequence_mean mode, which accumulates by completion count. CONF token credits support both of these reductions.

ECHO World Modeling

ECHO is an additive GiGPO trainer feature for tagged multi-turn trajectories. Producer code may attach token_roles to each Trajectory, aligned to full_token_ids: 0 = MASKED, 1 = ACTION, 2 = ECHO. ACTION tokens keep the normal GiGPO policy-gradient loss. ECHO tokens add an independently normalized SFT term, alpha * NLL, in the same forward/backward pass.

echo_alpha defaults to 0.0, which leaves current GiGPO behavior untouched. For early world-modeling, start small, roughly 0.005 to 0.05, and consider echo_schedule = "linear_taper_to_zero" with echo_taper_steps = N so training returns to pure RL. alpha is model-size-sensitive; too-high values can overfit tool responses or collapse RL.

ECHO is content-agnostic by design. mlxrl never inspects token text to decide what is an environment or tool response. The rollout producer owns the tags: MASKED tokens have zero gradient, ACTION tokens receive the RL advantage, and ECHO tokens receive only the SFT loss with no KL penalty or importance ratio. The trainer logs loss_echo and token-level echo_acc separately from the ACTION loss.

Policy Semantics

• The base model is frozen before LoRA injection.
• Only LoRA adapter leaves are trainable.
• Reference logprobs are computed by temporarily disabling adapters on the same model object; there is no second reference model in memory.
• Old-policy logprobs are recomputed with a full forward for the training batch. Rollout-time logprobs are captured for inspection, but 4-bit sequential decode and full-forward prefill are not numerically identical on hybrid/quantized models, so recompute remains the default training semantics.
• When beta == 0, the reference forward is skipped and the policy logprobs are used as a zero-KL placeholder.
• PPO, DPO, and ORPO are intentionally out of scope. PPO needs a separate critic and value forward; DPO/ORPO are offline preference objectives with no rollout phase. mlxrl is critic-free, on-policy, and rollout-based by design.

Algorithms

Concrete algorithms implement the small Algorithm protocol: compute advantages, optionally filter a prepared batch, then compute a loss from policy, old-policy, and reference logprobs. rollout/, policy/, and train/ do not import concrete algorithm implementations.

algorithm	defining behavior
GRPO	group-normalized rewards, token-level importance ratio
Dr. GRPO	centered or normalized rewards with decoupled length reduction
DAPO	asymmetric low/high clipping plus optional dynamic zero-advantage group filtering
GSPO	sequence-level, length-normalized importance ratio and clipping
RLOO	leave-one-out group baseline, no critic, no std-normalized advantage
GiGPO	multi-turn episode advantage plus anchor-state step advantage
Score centering (experimental)	centered rewards, token-mean REINFORCE with sampler-score correction

Experimental score centering

Score centering is an explicit single-turn opt-in based on Marek and Ryabinin, 2026. Existing algorithm defaults and old-policy recompute semantics stay unchanged.

UV_CACHE_DIR=.uv-cache uv run mlxrl train \
  --config examples/qwen3_0_6b_score_centering.toml \
  --available-memory-gb 48

The config selects the objective and the number of sampler candidates to retain:

[algorithm]
name = "score-centering"
score_centering_top_k = 128

This objective uses group-centered rewards without standard-deviation scaling. The sampler records probabilities after top-p, min-p, top-k, and temperature, including the selected token's probability even when it falls outside the recorded head. Greedy sampling is recorded as a point mass. Candidate probabilities retain their original mass; the trainer approximates the remaining tail as described in the paper. The correction coefficients are detached, and trainer logprobs are normalized in fp32. kl_beta remains the independent adapter-disabled reference-model penalty.

For the import API, construct ScoreCenteringAlgorithm(sampler_top_k=128) from mlxrl.algo and use SamplingConfig(record_sampler_top_k=128) for generation. Pass the resulting completions and that algorithm to batch_from_rollouts and optimizer_step. Missing sampler records are rejected. Completion.sampler_trace contains the selected-token and candidate probabilities; config-driven runs also save sampler_logprobs, sampler_topk_ids, and sampler_topk_logprobs in the output NPZ. Ordinary old_policy_logprobs remain separate.

Capture works with naive, prefix-cached, batched, and compiled single-turn rollouts and preserves seeded token draws. The objective supports token-weighted micro-batching. Multi-turn/ECHO training is outside this initial experiment. The memory preflight estimates have not been calibrated for this objective. Its benefit and overhead on a particular model must be measured; the paper's large-model results do not establish a benefit for MLX QLoRA.

The reproducible paired benchmark is benchmarks/run_score_centering.py; its protocol and measured results are in the 4-bit benchmark report.

Experimental confidence warm start (CONF)

CONF adds detached, centered step-confidence credit during the first eight optimizer updates, then continues with the selected base loss and optimizer. It is experimental and disabled by default. GRPO stays the default algorithm; score centering remains a separate opt-in.

UV_CACHE_DIR=.uv-cache uv run mlxrl train \
  --config examples/qwen3_0_6b_confidence.toml

The example uses explicit step: expression = numeric_result reasoning lines, four prompt groups per update, and per-completion loss averaging. To compare against its matched baseline, change only [confidence] enabled = false. The prompt format and loss reduction are independent settings and stay fixed.

Confidence is the model's temperature-1 probability of the written result, computed with fp32 normalization before optimization. Each step receives lambda * (confidence - token_weighted_mean_confidence) on eligible reasoning tokens. Result tokens, final answers and padding receive no extra credit. Unscored or unalignable text contributes zero. This arithmetic reader does not validate correctness or infer steps from arbitrary prose.

The default controller targets a credit share of 0.875, caps its measured relative p99 at 10, and holds the previous scale when fewer than four mixed groups carry credit. The cap is reapplied on held scales whenever measurable. With no initial calibration signal, credit stays zero. Logs expose the scale, actual share, scored steps, credited tokens and handoff; the output NPZ retains the credit history. Held scales can produce shares above or below the target.

For custom tasks, supply completion-token spans through ConfidenceStep and TokenSpan; see the Python integration and validation notes. The initial integration supports single-turn GRPO and score centering. It does not support DAPO filtering, GiGPO, ECHO or multi-turn credit.

RAMP's reported speedup is specific to its research setup. The MLX arithmetic screening harness is benchmarks/run_confidence.py; it compares all four combinations across seeds, including sampling and confidence overhead. The 12-run 4-bit screening found mixed outcomes: CONF raised final GRPO accuracy on two seeds and lowered it on one; with score centering it lowered accuracy on two seeds and tied on one. These short runs do not establish a reliable gain or a 2× speedup.

Memory Preflight

mlxrl train can estimate memory before loading the model:

UV_CACHE_DIR=.uv-cache uv run mlxrl train \
  --config examples/qwen3_0_6b_grpo.toml \
  --available-memory-gb 48 \
  --dry-run

The estimator is calibrated to measured anchors: 6.245 GB for Qwen3-0.6B/G4/prompt≈19/T256, 25.9 GB for Qwen3.5-9B/G2/seq609/per-layer-checkpointed, 45.9 GB for Qwen3.5-9B/G4/seq609/per-layer-checkpointed, and 36 GB for Qwen3.5-9B/G2/seq128/no-checkpoint. For hybrid 9B no-checkpoint long-sequence configs, it reports an OOM-risk lower bound rather than a fake precise peak. Multi-turn configs include max_turns and max_observation_len in the estimated trajectory length and are labeled as estimated OOM-risk shapes until new measured anchors are added. For an obviously too-large Qwen3.5-9B/G8/prompt97/T512/no-checkpoint config on 48 GB, it flags the run and suggests the measured-boundary fallback around G4/T512/checkpointed.

Benchmarks

Local M4 Max Phase 4 snapshot:

• 454 rollout tok/s on Qwen3-0.6B GRPO with G=4 and 256-token completions.
• 0.283 end-to-end it/s with full mlxrl training semantics.
• 3.2x faster rollout and 2.2x higher end-to-end it/s than mlx-tune v0.5.1 on the same run shape.
• 1.3x faster rollout than sequential mlx-lm generation at G=4.

These are the two-pass means from benchmarks/results/gate5_full_reconciled.md, run with MLX 0.31.2, MLX-LM 0.31.3, mlx-community/Qwen3-0.6B-4bit, 100 measured steps with 5 warmup steps discarded:

target	comparison	rollout tok/s	grad s/step	samples/s	it/s	peak GB
mlxrl	apples-to-apples GRPO	454.1	1.282	1.133	0.283	6.25
mlx-lm	generation-only, G=1	347.0	-	1.355	-	0.52
mlx-lm-g4	generation-only, sequential G=4	349.7	-	1.366	-	0.52
mlx-tune	package-speed reference	142.2	0.502	0.519	0.130	6.16
mlx-lm-lora	package-speed reference	557.9	0.592	1.648	0.412	5.32

mlx-lm-lora reports higher raw package-speed throughput in this snapshot, but its benchmarked path is not the same training problem as mlxrl's live old-policy/reference semantics and completion-loss masking. That is the honest case where mlxrl is not faster; the apples-to-apples comparison label is reserved for mlxrl's own semantic path. On the 9B Noether real workload, the checkpointed MLX path measured about 6x faster than the previous torch-MPS path; that workload is separate from the public Phase 4 package-speed harness.

Run the Phase 4 harness:

UV_CACHE_DIR=.uv-cache uv run python benchmarks/run_phase4.py run \
  --targets mlxrl,mlx-lm,mlx-tune,mlx-lm-lora \
  --model mlx-community/Qwen3-0.6B-4bit \
  --steps 100 \
  --warmup-steps 5 \
  --group-size 4 \
  --max-tokens 256 \
  --passes 2 \
  --output benchmarks/results/phase4.jsonl \
  --summary benchmarks/results/phase4.md \
  --allow-missing-baselines

The harness reports synchronized rollout tok/s, gradient seconds per step, samples/s, it/s, and peak MLX memory. mlx-lm targets are generation-only; external package targets are useful speed references but may not match mlxrl training semantics.

Development

See CONTRIBUTING.md and DESIGN.md before adding algorithms or changing rollout/logprob semantics.

Run the quality gates:

UV_CACHE_DIR=.uv-cache uv run pytest
UV_CACHE_DIR=.uv-cache uv run ruff check .
UV_CACHE_DIR=.uv-cache uv run pyright

MLX lazy evaluation matters. Any mx.eval(...) or mx.synchronize() in this repo should mark a real boundary: sampled token append/EOS checks, logprob freezing before adapter mutation, per-micro-batch graph release, optimizer updates, or benchmark timing boundaries.

Layout

mlxrl/
  rollout/   # batched group generation
  env/       # agentic environment protocols and reference envs
  policy/    # model loading, LoRA setup, logprob passes
  algo/      # GRPO-family advantages and losses
  train/     # value_and_grad and optimizer integration
  data/      # toy GSM8K data and rewards
  cli.py
tests/
benchmarks/

Non-Goals

• No inference server or second model copy.
• No CUDA or torch fallback.
• No distributed training.
• No broad RL framework abstractions beyond the small algorithm interface.