mlx-train-perf

Train on longer sequences, and get through short ones faster, on the Mac you already have.

mlx-train-perf is a set of drop-in parts for MLX LoRA and QLoRA fine-tuning: Metal kernels that cut what a single training step allocates, sequence packing for datasets made of short examples, and a planner that answers "will this fit in my RAM?" before you download the weights. It is not a trainer and does not want to be one. You keep mlx_lm's training loop and swap in the pieces you need.

Measured on one M1 Max (32 GB), Qwen3-8B-4bit QLoRA: the longest sequence you can train goes from 7,936 tokens to 23,040, and an Alpaca-shaped instruction dataset trains 3.0× faster per real token (2 to 3× depending on how short your examples are).

If a fine-tune does not fit in memory, mlx-lm's LoRA guide offers five remedies. The fourth is:

Longer examples require more memory. If it makes sense for your data, one thing you can do is break your examples into smaller sequences when making the {train, valid, test}.jsonl files.

And if you leave max_seq_length at its default of 2048 while your data is longer than that, the trainer does it for you:

[WARNING] Some sequences are longer than 2048 tokens. The longest sentence 6144 will be
truncated to 2048. Consider pre-splitting your data to save memory.

That advice is correct, and it means training on the first 2048 tokens of every contract, transcript, or source file and discarding the rest. The advice exists because MLX's attention has a memory-light forward pass and a backward pass that rebuilds the full (N, N) score matrix. An MLX maintainer put it plainly while closing an out-of-memory report:

The RAM needed for training grows quadratically as sequence length increases, so I'm afraid the OOM is not something we can simply solve.

It is solvable one layer up. This library replaces that backward pass with one that keeps O(N) state, for the architectures it supports, so context costs what it should. You raise max_seq_length instead of cutting up your data.

Including the cases where it does not.

Every number here has a committed script under scripts/ that reproduces it, all measured on one M1 Max (32 GB, macOS 26.5). The loss-layer figures were taken on mlx 0.31.2 and reproduce on the pinned 0.32.0; the flash-attention memory figures were taken on 0.32.0 in 0.2.0, and the 0.3.0 context-ceiling figures on 0.32.0.

pip install mlx-train-perf            # the loss kernel + planner
pip install "mlx-train-perf[mlx-lm]"  # plus the mlx-lm training adapter
pip install "mlx-train-perf[guard]"   # plus external supervision for benchmark runs

Apple Silicon only. Requires mlx >=0.32.0,<0.33, the version the kernels' JIT contract is verified against. The mlx-lm adapter and the flash-attention wrapper need the optional mlx-lm extra.

There is no flag to bolt onto mlx_lm.lora. These parts attach to a loaded model object, so you drive mlx_lm's train() from a short Python script instead of the CLI. That script is the whole difference, and it runs about six lines longer than the one you would have written anyway.

What you run today:

mlx_lm.lora --model mlx-community/Qwen3-8B-4bit --train --data ./data \
    --batch-size 1 --grad-checkpoint

[WARNING] Some sequences are longer than 2048 tokens. The longest sentence 6144 will be
truncated to 2048. Consider pre-splitting your data to save memory.

Passing --max-seq-length 8192 trades the truncation for a crash: the step peaks at 25.68 GiB, and a 32 GB Mac has roughly 24.5 GiB to give it. With the flash path that same step peaks at 12.75 GiB, and the longest sequence you can train moves from 7,936 tokens to 23,040 (scripts/northstar_context_sweep.py).

import mlx.core as mx
from mlx_lm import load
from mlx_lm.tuner.trainer import TrainingArgs, train
from mlx_train_perf.adapters.mlx_lm import make_loss_fn
from mlx_train_perf.attention import enable_flash_attention

model, tokenizer = load("mlx-community/Qwen3-8B-4bit")
model.set_dtype(mx.bfloat16)  # 4-bit checkpoints compute in fp16; the kernels need bf16/fp32
# ... freeze the base model and apply linear_to_lora_layers as in a normal mlx-lm LoRA run ...

args = TrainingArgs(batch_size=1, max_seq_length=8192, grad_checkpoint=True, iters=600)
enable_flash_attention(model, seq_len=8192, batch_size=1)

train(model=model, optimizer=opt, train_dataset=ds, args=args,
      loss=make_loss_fn(model, impl="auto"))

Two independent levers sit in those last three lines. enable_flash_attention swaps each layer's attention for the O(N) path, which is what moves the ceiling. make_loss_fn routes the loss through the fused kernel and frees the logit buffer on top of that. Either one works without the other.

The elisions above are the parts of a stock mlx-lm LoRA run that do not change. examples/finetune_long_context.py is the same thing with nothing left out: argument parsing, dataset loading, the freeze and adapter setup, and adapter saving. Run it as is.

One trap worth knowing: enable_flash_attention replaces each layer's attention in place and there is no undo. Load a fresh copy of the model for inference, because the training-configured object raises AttentionInputError as soon as a KV cache appears.

On a 16 GB machine this shape does not fit. Ask the planner what does, rather than finding out three minutes into a run.

Alpaca averages 84 tokens per example under Qwen3's chat template, and mlx-lm's trainer runs one step per batch of them. On an 8B model at batch 1 that step takes 2.5 s to carry 84 real tokens. Packing fills the row to 4,096 tokens with whole examples instead, and the step then takes 40.4 s to carry about 4,000 — roughly 48 times the tokens for 16 times the wall clock. The difference is fixed per-step cost that a short batch pays in full, and a packed row pays once.

import functools
from mlx_train_perf.adapters.mlx_lm import make_packed_loss_fn
from mlx_train_perf.data.packing import packed_iterate_batches

enable_flash_attention(model, seq_len=4096, batch_size=1, packed=True)
args = TrainingArgs(batch_size=1, max_seq_length=4096, grad_checkpoint=True, iters=600)

train(model=model, optimizer=opt, train_dataset=ds, args=args,
      loss=make_packed_loss_fn(model),
      iterate_batches=functools.partial(
          packed_iterate_batches,
          max_position_embeddings=model.args.max_position_embeddings,
      ))

Measured on Qwen3-8B-4bit: 33.1 real tokens per second unpacked against 99.2 packed, a factor of 3.00 (scripts/bench_packed_training.py). Dataset items are (tokens, offset) pairs, where the offset is the prompt length. The gain comes from amortizing the fixed step cost, so it shrinks as your examples get longer and disappears once they already fill a row. Sequence packing has the conservative steady-state range and the full recipe.

mlx-train-perf plan --config ./Qwen3-8B-4bit/config.json --batch 1 --lora-rank 8 \
    --attention flash --max-seq

The alternative is a 16 GB download and an out-of-memory crash three minutes into training. This loads no weights and spends no GPU time. It reads the config, prices the run against your machine's memory, and hands back the longest sequence that fits. Ask about one specific config with --seq-len instead and it answers fits or does not fit, with the peak it predicted. The estimate leans toward over-predicting, which is the safe direction for a tool whose job is keeping you off the cliff.

These pieces are usable without mlx_lm and without this project's adapter. mlx is the only runtime dependency; mlx-lm is an optional extra that exists solely for the adapter and the model-instance wrappers.

from mlx_train_perf import linear_cross_entropy, DenseHead, QuantizedHead
from mlx_train_perf.attention import flash_attention
from mlx_train_perf.recurrent import chunked_gated_delta

# Loss: hidden states in, scalar out, no (N, V) tensor in between.
loss = linear_cross_entropy(hidden, head, targets, impl="auto", reduction="mean")

# Attention: a drop-in for mx.fast.scaled_dot_product_attention on the training path.
out = flash_attention(q, k, v, scale=scale, causal=True)

# GatedDelta: Qwen 3.5's linear-attention recurrence, argument-compatible with mlx-lm's own op.
out, state = chunked_gated_delta(q, k, v, g, beta, state, mask=mask)

At masked positions, chunked_gated_delta's y is unspecified and differs from gated_delta_ops's own output there — the two agree on state and on every valid position's y, but not on what a masked position's output looks like. Mask your own loss the same way mlx_train_perf's adapter does, and don't rely on y at those positions.

head is a DenseHead, a QuantizedHead, or a tied embedding via tied_head(...). For flash_attention, q/k/v are (B, H, N, D) with head_dim in {64, 96, 128} and grouped-query heads mapped contiguously, matching mx.fast.scaled_dot_product_attention's own convention; pass segments=PackedMask(...) for block-diagonal packing. chunked_gated_delta follows mlx-lm's own gated_delta_ops convention instead: (B, T, H, D) with the sequence axis second, and whatever key/value head dims the checkpoint carries.

What you are signing up for, stated plainly:

• In-place mutation. enable_flash_attention(model) and enable_gated_delta_training(model) both swap layers on a live model object and have no undo. flash_attention and chunked_gated_delta are themselves pure functions and mutate nothing, so if you own your model code, call them directly and skip the wrappers.
• Training only. enable_flash_attention and enable_gated_delta_training both refuse a KV cache. Reload the model for inference.
• Typed refusals, never silent fallbacks. An unsupported architecture, head dim, dtype or mask raises at enable time or on the first call, naming the reason.
• The mlx pin is a policy, not neglect. mlx>=0.32.0,<0.33 is narrow because the kernels' JIT contract is re-verified against each mlx release before the range widens, rather than assumed forward-compatible.
• Calibration is one-time and host-synced. Warm it at your training shape before a compiled step traces, or accept a single in-trace stall on the first call.

If your model family is not in the support list, open an issue and name it. The refusal list is a statement about what has been verified, not about what the kernels could cover.

The idea is the same one behind Cut Cross-Entropy and Liger-Kernel on the CUDA side, ported to a Metal kernel: compute the cross-entropy loss and its gradient without ever building the full (N, V) logits tensor. For a large vocabulary that tensor is the single biggest allocation in the training step, and it is pure waste. You only need the per-token loss and a gradient back into the hidden states.

Standard cross-entropy in a trainer materializes logits of shape (batch·seq, vocab). At Qwen3-8B's vocabulary (151,936) and a 2048-token sequence, that is a 0.6 GB tensor in bf16, plus another for the softmax gradient in the backward pass. The fused kernel never allocates it: the forward regenerates logits in registers tile-by-tile over the vocabulary and returns three N-length arrays (the per-token NLL, the log-sum-exp, and the target logit); the backward recomputes the needed tiles instead of reading a stored matrix.

Measured in isolation, at n=8192, V=151936, D=4096, bf16 (scripts/bench_loss_layer.py):

About 3900× less memory for the loss layer, at a 1.64× cost on the forward pass.

The fused loss is exact to bf16 tolerance against the stock trainer (per-step loss curves match to about 2e-3), and the throughput cost is small: roughly 8–12% slower per step at bf16 (scripts/bench_train_step.py, Qwen3-8B-4bit, LoRA r=8, gradient checkpointing on).

New in 0.2.0 and opt-in. Removing the logit tensor frees real memory, but on its own it barely moves the training peak at long context. The reason is attention. mx.fast.scaled_dot_product_attention has a memory-light forward and an O(N²) backward that rebuilds the (N, N) score matrix one layer at a time. Once the logits are gone, that backward is what sets the peak.

0.2.0 adds a flash-attention path with a Metal forward and a Metal backward, neither of which materializes the score matrix. It keeps O(N) saved state — the attention output and the log-sum-exp — and recomputes the tiles it needs. You switch it on per model with enable_flash_attention and train exactly as before.

On Qwen3-8B-4bit (LoRA rank 8, batch 1, gradient checkpointing on, bf16) the two attention paths are close at a 2048-token sequence. At 8192 they are not: the flash path halves the whole step's peak memory.

(scripts/bench_train_step.py; M1 Max 32 GB, macOS 26.5, mlx 0.32.0. These memory and throughput figures are the 0.2.0 measurements, carried into 0.3.0 unchanged: 0.3.0 changed how the backward splits its kernel launches, not what it allocates, and the 0.3.0 context sweep below — measured fresh — confirms the flash path's memory still scales linearly in sequence length.)

The 32 GB machine that peaked near its ceiling with stock attention now runs the same step at half the memory. That headroom buys a longer sequence.

The attention op timed alone at the flagship shape (batch 1, 32 query / 8 KV heads, 8192 tokens, head_dim 128) is 0.186 s forward and 0.576 s backward (scripts/bench_attention_op.py). Its peak grows 2.00× from 2,048 to 4,096 tokens and 3.06× from 4,096 to 8,192; the second step is above 2× because the chained backward split adds a bounded constant of at most ~0.3 GB, not because the O(N) growth law changed. Stock attention grows 3.76× then 3.05× over the same doublings, from a far higher base. Flash is not universally cheaper: below about 2,100 tokens the stock op's simpler bookkeeping wins, and the two curves cross there. When the bottleneck moved has the full measurement, including what stock attention does on a 32 GB machine once it starts paging.

Turning flash attention on is not free. On the stock-loss path at 8192 it costs 5.3% of tokens/sec (74.0 vs 78.1); at 2048 the cost is 5.5% on the fused-loss path (86.4 vs 91.5) and 5.9% on the stock-loss path (92.1 vs 97.8). The fused-loss comparison at 8192 has no stock-attention number to pair with: on this 32 GB machine that baseline condition crosses the memory safety net's ceiling and records an abort instead of a number. That baseline running out of room is the problem flash attention exists to remove. Under flash attention the fused cross-entropy and mlx-lm's stock cross-entropy stay close: 0.94× at 2048 (86.4 vs 92.1 tok/s) and 0.99× at 8192 on Qwen3-8B, 0.92× and 0.97× on Llama-3.2-3B. The loss values match to bf16 tolerance throughout — the worst per-step difference across every measured pair is 2.4e-3. The worst attention-arm throughput ratio measured is 0.94× stock.

0.2.0 shipped this path with a launch-safety budget that capped context before memory did. That budget turned out to be guarding the wrong unit: the GPU watchdog acts on a single Metal command buffer, not a chain of them, and 0.3.0 measures the margin against the right thing (scripts/probe_command_buffer_packing.py; the reasoning is in How MLX packs Metal command buffers).

With that cap gone, the flash path is bound by memory, the same thing that bounds stock attention — and it needs far less of it. Measured the same day with the same search (scripts/northstar_context_sweep.py, Qwen3-8B-4bit QLoRA, gradient checkpointing, bf16):

Both arms stop at the same ~24.5 GiB, the effective memory ceiling on this machine at run time. Under that one budget the flash path trains 2.9× the context, because it holds O(N) saved state where stock holds the O(N²) score matrix. The ratio is the part that travels: raise the available memory and both ceilings rise together (a freshly booted or larger machine lets both climb toward the 28 GiB static ceiling), but the flash path keeps its roughly threefold reach. On the same machine in 0.2.0 this path was launch-capped near 10k tokens; removing that cap is what moved it. (These figures are this release's measurement; the two arms are comparable to each other, taken together, not to 0.2.0's numbers.)

enable_flash_attention is causal-only and training-only, and it refuses anything outside that up front rather than failing mid-run:

from mlx_train_perf.attention import enable_flash_attention

enable_flash_attention(model, seq_len=8192, batch_size=1)

Call it in place on a loaded model, after you set the compute dtype and before you build the loss and call train. mlx-lm's train wraps the step in mx.compile, and the kernel calibrates itself with a one-time host-synced timing probe. Passing seq_len (and batch_size) runs that calibration up front, at your training shape, so the compiled step traces with warm caches. Match them to the shape you actually train: batch_size defaults to 1 and must equal your training batch. If a compiled train traces at a shape the caches were not warmed for, the calibration runs once inside the traced region instead — the run completes, but the timing probe executes on a machine mid-trace rather than in the controlled up-front window (measured on mlx 0.32.0: a one-time stall, not a crash). Omit the hints and the call still succeeds — eager and mx.grad callers calibrate lazily on the first attention call — but a compiled train run should always pass them for calibration fidelity.

New in 0.4.0 and opt-in. Packing concatenates many short sequences into fixed 4,096-token rows so every step runs at full-context efficiency, as the worked example above describes. A block-diagonal attention mask keeps the sequences independent: a token attends another only when both belong to the same original sequence, enforced inside the flash Metal kernels by a per-token segment id rather than a materialized mask (the mask tensor an (N, N) approach would need is exactly the quadratic allocation this library exists to avoid).

Loss masking reproduces mlx-lm's unpacked semantics segment by segment, so the supervised token set is identical to an unpacked run. Three sequences packed into one row produce the same token count and a loss within measured bf16 tolerance of the same three run unpacked: worst difference 5.0e-4 against a 2e-2 pin sized from measured RoPE offset drift (tests/test_adapter_packed.py). Cross-sequence contamination is tested by construction: deliberately dropping the segment mask in the test suite moves the loss by 0.11, well past the pin.

Measured on Alpaca (pinned revision, 4,000-example sample, seed 42), LoRA rank 8, batch 1, gradient checkpointing, bf16, pack length 4,096 (scripts/bench_packed_training.py):

Samples per hour move the same way: 1,415 → 4,246 on Qwen3-8B and 2,617 → 7,726 on Llama-3.2-3B. "Real tokens" counts sequence content only, never padding or separators. Both arms of each pair were measured in one session at this release's code state.

Where the win comes from matters for whether you will see it too. At batch size 1, stock batching loses little to padding (17% on this dataset, mostly round-to-32 alignment) — the win comes from amortization. A packed row carries roughly 40–50 Alpaca sequences (47.6 on average under Qwen3's tokenizer, 38.1 under Llama's), so the fixed step cost is paid once per ~4,000 real tokens instead of once per 84, and attention runs at its 4,096-token efficiency instead of a ~100-token shape. A dataset of long sequences packs fewer per row and gains less; one that already fills the context gains nothing. The stock arm's per-step median also includes mx.compile's first trace of each batch width (stock widths vary; packed rows are one constant shape, which is itself part of the win), and a long training run amortizes those traces away. Reading the stock arm at its fastest repeated warm step instead of its median gives a conservative bound of about 2.0–2.3×, so the honest range is 2–2.7× on this dataset. The packed arm's own walls are flat to within 5%.

0.5.0 tightens the packed backward: the dK/dV kernel now bounds its query walk at each key block's segment end instead of masking cross-segment work after computing it. Timed with identical dispatch ranges on both arms (scripts/bench_packed_dkv.py), the dK/dV pass on an Alpaca-like row runs 6.2× faster at 4,096 tokens and 8.3× at 8,192; a single-segment row is unchanged. That pass is one slice of the training step, so the win on the full step is smaller: on Qwen3-8B-4bit the packed arm's median step drops from 44.7 s to 40.4 s (+10.7% tokens/sec, with the unpacked arm within half a percent of its prior measurement — the control that pins the gain to the packed backward). The table above is this release's measurement of both models.

The parts drop into the stock trainer the same way the loss does — a batch iterator, a loss function, and the flash-attention switch. Packing requires the flash path (the stock attention cannot express a block-diagonal mask):

import functools
import mlx.core as mx
from mlx_lm import load
from mlx_lm.tuner.trainer import train
from mlx_train_perf.adapters.mlx_lm import make_packed_loss_fn
from mlx_train_perf.attention import enable_flash_attention
from mlx_train_perf.data.packing import packed_iterate_batches

model, tokenizer = load("mlx-community/Qwen3-8B-4bit")
model.set_dtype(mx.bfloat16)  # 4-bit checkpoints compute in fp16; the kernels need bf16/fp32
enable_flash_attention(model, seq_len=4096, batch_size=1, packed=True)
# ... freeze the base model and apply linear_to_lora_layers as usual ...

train(
    model=model, optimizer=opt,
    train_dataset=dataset,          # items are (tokens, offset) pairs; offset = prompt length
    args=args,                      # args.max_seq_length is the pack length
    loss=make_packed_loss_fn(model),
    iterate_batches=functools.partial(
        packed_iterate_batches,
        max_position_embeddings=model.args.max_position_embeddings,
    ),
)

packed_iterate_batches re-packs each epoch with a fresh shuffle and hands the trainer fixed-shape batches; make_packed_loss_fn walks the model's layers itself to thread the segment mask (the stock model call hardcodes a causal mask) and refuses at construction if enable_flash_attention has not run. Pass packed=True with seq_len equal to your pack length and batch_size equal to your training batch: the calibration caches key on the exact batch size and sequence bucket, so matching hints keep the one-time kernel timing probes in the controlled window before mx.compile traces the step. The pack length must not exceed the model's trained context — packed sequences keep their relative positions, and the row as a whole runs at absolute positions up to the pack length.

New in 0.7.0 and opt-in. Qwen 3.5 mixes full-attention layers with GatedDeltaNet (linear-attention) ones in the same model, an architecture enable_flash_attention was never built to reach: most of its layers have no full-attention block at all. enable_gated_delta_training gives those layers a training path of their own, an in-tree chunk-parallel implementation of the same recurrence, switched on per loaded model the way the flash-attention wrapper already is.

from mlx_train_perf.recurrent import enable_gated_delta_training
from mlx_train_perf.adapters.mlx_lm import make_loss_fn

enable_gated_delta_training(model, impl="chunked")
loss = make_loss_fn(model, impl="auto")

Call it in place, on a loaded model, before you build the loss and call train, the same order enable_flash_attention uses. Every structurally linear-attention layer's GatedDeltaNet gets replaced by a proxy that routes through the chunk-parallel op; the model's full-attention layers are untouched and keep running on stock attention. One throughput note: each distinct padded sequence length traces its own compute graph (the chunk loop unrolls at trace time). mlx-lm pads each batch to a 32-token bucket capped at max_seq_length, so ragged data pays a small retrace for each new bucket it encounters, bounded and quickly amortized; fixed-length rows keep the trace count at one. Right-padded batches work as you'd expect: the recurrence and its depthwise convolution are both causal, so a padded position can only influence positions after it, and with right padding those are all padding too. The loss already masks out padded targets, so nothing downstream ever depends on what a pad position's own output was.

What this release covers, and what it does not:

impl="sequential" is a benchmarking control, not a second production path. It swaps in the same proxy but routes it through mlx-lm's own sequential op instead of this project's chunk-parallel one, which is what lets a throughput comparison attribute a difference to the op itself rather than to the proxy's block rewrite. Use the default impl="chunked" unless you're running that comparison yourself.

A loaded 4-bit checkpoint needs its floating parameters cast to bf16 before training, and the obvious way to do that is unsafe for this family:

import mlx.core as mx
from mlx.utils import tree_map_with_path

cast_predicate = model.cast_predicate

def cast_bf16(path, value):
    if cast_predicate(path) and mx.issubdtype(value.dtype, mx.floating):
        return value.astype(mx.bfloat16)
    return value

model.update(tree_map_with_path(cast_bf16, model.parameters()))

Do this before calling enable_gated_delta_training. model.set_dtype(mx.bfloat16) calls its predicate with a dtype, not a path, so it can never run qwen3_5's own cast_predicate, the rule that keeps A_log (the GatedDelta gating log-rates) in fp32, and it downcasts A_log along with everything else without telling you. enable_gated_delta_training refuses the forward outright when it sees a non-fp32 A_log, so the failure shows up as an error instead of a quietly wrong model.

impl picks how the loss is computed. "auto" is the default and the one to use.

• kernel — the fused Metal kernel. "auto" resolves here when the mlx version is verified and the head/dtype are supported (dense or tied fp32/bf16 head; 4-bit group-size-64 quantized head; hidden states in fp32 or bf16). It never materializes (N, V).
• chunked — a pure-MLX fallback that processes the vocabulary in fixed tiles. No Metal kernel, works anywhere MLX does, uses more memory than kernel but far less than naive. This is also the backward path the kernel forward pairs with today.
• naive — materializes the full logits. It is the correctness oracle the other two are tested against, not something to train with.

"auto" never silently downgrades. If it cannot use the kernel (unverified mlx, an unsupported head, fp16 hidden states) it raises a typed error naming the reason and the alternatives, so you always know which path ran.

Before a run, the planner estimates the peak training memory for a config and tells you whether it fits, or suggests a smaller batch or sequence length that would:

mlx-train-perf plan --config path/to/config.json --batch 1 --seq-len 4096 --lora-rank 8

The memory model is fit to measured Qwen3-8B train-step peaks and cross-model validated on Llama-3.2-3B to within about 9%. It accounts for the O(N²) attention backward described above, so it does not under-predict at long context the way a linear model would. It is an estimate, and it errs toward over-predicting, which is the safe direction for a tool whose job is to keep you off the OOM cliff.

Pass --attention flash to price the flash-attention path instead of the stock backward:

mlx-train-perf plan --config path/to/config.json --batch 1 --seq-len 8192 --lora-rank 8 --attention flash

The flash model is an analytic saved-state term plus a measured linear coefficient per loss implementation, each fit on its own arm's anchors so neither under-predicts the quantity the fit models (the weights term plus the measured training-loop marginal) at any of them, even before the safety cushion. Plans for the fused loss read at 1.09–1.19× the recorded peak on the fitted model and 1.22–1.24× cross-model (ranges rounded outward); before 0.6.0 a single worst-case coefficient pushed these to about 1.4× and 1.6×. The chunked loss uses the stock-loss arm's coefficient and reads at 1.13–1.26×; the naive loss reads higher still because its own loss term over-predicts away from its calibration shape. The validated range is 2,048 to 12,288 tokens; past that the fit extrapolates. The estimate models MLX active memory — if you are budgeting full resident footprint, bound the allocator's retained cache with mx.set_cache_limit(...) in your training process.

Instead of checking one config at a time, ask the planner for the largest sequence length or batch size that fits your budget:

mlx-train-perf plan --config path/to/config.json --batch 1 --lora-rank 8 --attention flash --max-seq
mlx-train-perf plan --config path/to/config.json --seq-len 8192 --lora-rank 8 --attention flash --max-batch

--max-seq searches for the largest --seq-len and still needs --batch; --max-batch searches for the largest --batch and still needs --seq-len. A budget that nothing fits, even at the smallest value searched, is refused with a typed error.

• Architectures: Llama, Qwen2 (the Qwen2.5 family), and Qwen3, for both the loss adapter and the flash-attention wrapper. Qwen 3.5 is supported by the loss adapter and by its own GatedDelta training path; its full-attention layers are not on the flash-attention wrapper yet, and sequence packing and the RAM-fit planner don't cover it. The adapter's model splitter handles all of these; other families raise a typed error.
• Quantization: 4-bit group-size-64 (the mlx-community QLoRA default), or a dense fp32/bf16 head.
• Training: LoRA / QLoRA. Full fine-tuning is estimated by the planner but is not the case this is tuned for.
• Hardware: Apple Silicon.

Each claim above has one script. They run on the GPU, take real wall-clock time, and print the artifacts they measured:

python scripts/bench_loss_layer.py        # the ~3900x loss-layer memory number
python scripts/bench_attention_op.py      # the single-op flash vs stock memory + timing
# the 12.75 vs 25.68 GiB training table: run each attention arm into its own --out dir
python scripts/bench_train_step.py --model mlx-community/Qwen3-8B-4bit --seq-len 8192 \
    --attention flash --impl kernel --compute-dtype bfloat16 --grad-checkpoint --out _artifacts/flash
python scripts/bench_train_step.py --model mlx-community/Qwen3-8B-4bit --seq-len 8192 \
    --attention stock --impl kernel --compute-dtype bfloat16 --grad-checkpoint --out _artifacts/stock
python scripts/northstar_context_sweep.py # the max-context sweep (1-2 h; heavy)
# the packed dK/dV block-skip ratios (6.2x / 8.3x): one invocation per layout and length
python scripts/bench_packed_dkv.py --n 4096 --layout alpaca --out _artifacts/packed_dkv
python scripts/bench_packed_dkv.py --n 8192 --layout alpaca --out _artifacts/packed_dkv
# the planner's flash-fit anchors and refit (per-loss-arm fit over the committed manifest)
python scripts/fit_calibration.py --manifest _artifacts/calib_050/refit_manifest.json --dry-run
# the packing table (3.00x / 2.95x): prep the dataset once per model, then run each arm
# into its own --out dir (30 timed steps per arm, the script default, matching the
# committed artifacts)
python scripts/prep_alpaca.py --model mlx-community/Qwen3-8B-4bit \
    --out _artifacts/packed_bench/alpaca_qwen3.jsonl --batch-size 1 --pack-len 4096 --max-samples 4000 --seed 42
python scripts/bench_packed_training.py --model mlx-community/Qwen3-8B-4bit \
    --data _artifacts/packed_bench/alpaca_qwen3.jsonl --arm stock --pack-len 4096 --batch-size 1 \
    --grad-checkpoint --compute-dtype bfloat16 --out _artifacts/packed_bench_050
python scripts/bench_packed_training.py --model mlx-community/Qwen3-8B-4bit \
    --data _artifacts/packed_bench/alpaca_qwen3.jsonl --arm packed --pack-len 4096 --batch-size 1 \
    --grad-checkpoint --compute-dtype bfloat16 --out _artifacts/packed_bench_050

Every benchmark and contribution run is fenced by a device-relative memory guard. A GPU over-allocation on Apple Silicon does not always fail cleanly: mx.set_memory_limit is advisory, so an allocation past the soft cap pages instead of raising, and a hard enough paging storm can panic the machine rather than kill the process.

The guard sets an active-memory ceiling from the machine's own RAM. It is anchored at 28 GiB on a 32 GB Mac — above the largest legitimate peak measured here (25.68 GiB) and below physical RAM — and scales from that anchor across the range from 16 GB up to a 1 TB machine. At start it takes the smaller of that static ceiling and what the machine actually has free right now, minus a 2 GiB cushion. A daemon thread samples active memory throughout the run and aborts the moment it reaches the ceiling, writing an honest aborted-status artifact instead of letting the storm build. If the machine is already too loaded to start safely — less than a quarter of RAM effectively available — the run refuses up front with a typed error. Between those two points it proceeds but prints a warning naming how much memory it expected free for the machine's class against how much it measured, so a crowded machine is visible rather than silent.

The guard is rank-local: every input it reads is this node's own RAM, availability, and process memory. On a multi-node mx.distributed job each rank sizes its own ceiling and flags its own crowding, and a breach hard-exits that rank — so run distributed training under a launcher (mpirun or mlx.launch) that propagates a rank failure to the whole job.

The incident that motivated this guard, and what the watchdog does and does not cover, is documented in When an MLX memory cap is not a safety boundary.

The guard above lives inside the worker, so it sees what MLX reports: active memory. It cannot see what macOS charges the process as a whole, which also includes MLX's retained buffer cache and everything the interpreter allocates outside MLX, and nothing outside the worker owns its process group. mlx-guard covers that side. It is a small native supervisor that launches the worker, samples its OS-accounted footprint from outside, and stops the whole process group against a limit you set. The benchmark runner can put every condition under it:

pip install "mlx-train-perf[guard]"

from mlx_train_perf.bench.runner import ExternalGuardConfig, run_conditions

paths = run_conditions(
    conditions,
    out_dir,
    session_id=session_id,
    guard=ExternalGuardConfig(
        max_footprint_bytes=28 * 1024**3,  # sized for a 32 GB Mac; scale it to your RAM
        wall_time_ms=60 * 60 * 1000,
        checkpoint_timeout_ms=10_000,
    ),
)

Supervision is off unless you pass guard=, and nothing in the package imports mlx-guard until you do. For now it is a Python API on run_conditions; the scripts/bench_*.py drivers and the mlx-train-perf CLI do not expose it yet. The footprint limit is yours to choose. It counts more than MLX's active memory, so set it at or above the in-process ceiling for your machine, not below it, and well under physical RAM.

When a limit trips, the supervisor first asks the worker to checkpoint, then sends TERM. The worker can only answer between repetitions or training steps, so checkpoint_timeout_ms should cover one full step. The supervisor's default is one second, which a multi-second training step will always miss, and the most it accepts is 60 seconds. A condition whose steps run longer than that cannot checkpoint in time and will always end as an abort. A worker that answers in time writes and syncs a checkpointed_partial artifact with its progress before it acknowledges. When the intervention ends without a completed checkpoint, because the worker never reached a safe point, answered too late, or could not write, the condition is recorded as aborted_external_guard. Either way the condition counts as unfinished, and the next run retries it from the start.

Each launch writes its own supervisor report under out_dir/_mlx_guard/, next to a small supervision record that says how the run ended, whether a checkpoint was acknowledged, and how the worker itself exited. Reports are never deleted or reused, because the supervisor keeps a journal beside each one and refuses a path it has already written to. The directory grows by one report per launch, and clearing it out is up to you. The runner refuses to start if _mlx_guard is a symlink or belongs to another user.

The runner falls back to a direct, unsupervised launch in one situation only: the supervisor could not be brought up. That covers mlx-guard not being installed, a broken install, a binary that fails its version or integrity check, and a version check that times out on a busy machine. Each fallback is written as a guard_fallback record and announced on stderr, because the condition's own artifact looks the same with or without a supervisor. Once a supervisor has started, the condition is never launched a second time, whatever goes wrong afterwards. A failure in the client is written to a guard_client_error record, and any artifact the worker produced is kept. If the client never learned how the worker ended and there is no artifact, the condition is recorded as an error of type GuardClientError. A failure the supervisor reports about itself or the launch becomes an error of type SupervisorReportedFailure.

A runner that dies takes its supervised worker down with it, unless the runner was started under nohup, which turns the supervisor's parent watch off. Pass on_parent_exit="detach" if a condition should outlive the runner on purpose. A later sweep cannot see a detached worker that is still running, so it would launch that condition again.

The wired limit, the memory ceiling watchdog and the wall-time backstop described above stay on inside a supervised worker. The two layers watch different numbers, and either one can stop a run. If both fire at once, the watchdog's breach record wins: the checkpoint stands down, and a result that finishes during a breach does not replace the record.

The extra pins mlx-guard==0.2.0 exactly. mlx-guard is pre-1.0, where a minor release may change behavior a consumer can see, and its reports can only be read back by the version that wrote them. To roll back, drop guard= or uninstall the extra: runs return to the unsupervised path with nothing else to change.

Four write-ups cover the work behind this library in more depth than a README can, including the measurements that went the wrong way. They are published at ineshin.space alongside the rest of my Apple Silicon work, and the source Markdown lives under docs/papers/.

• Fused linear cross-entropy on Apple GPUs explains how vocabulary chunking and a fused Metal kernel avoid materializing logits. It covers memory costs, the optimization ladder, failed performance models, and the limits of the evidence.
• When the bottleneck moved: from fused cross-entropy to FlashAttention on MLX explains why removing the logits matrix did not extend context once attention backward set the peak. It also covers the command-buffer correction that removed a false launch limit, while separating source-reported measurements from claims the available controls cannot support.
• How MLX packs Metal command buffers explains the operation and element thresholds that MLX 0.32.0 uses to commit Metal work. It applies them to tiled attention, then explains why a whole-chain launch budget rejected valid work. The macOS watchdog mechanism remains an inference.
• When an MLX memory cap is not a safety boundary reports the kernel-panic incident behind the memory guard: a wired limit that caps residency without rejecting allocation, an advisory soft limit, and the active-memory watchdog added as a third layer. It separates the observed record from reconstruction and keeps the panic-trigger mechanism labeled as an unverified hypothesis.

Every number above is from an M1 Max (32 GB), the machine this is developed on. Whether the kernel and the flash-attention path scale the way the memory model expects on larger machines is a question only other people's hardware can answer, so there is a one-command way to measure it and send the numbers back:

mlx-train-perf contribute --tier quick   # ~10-15 min; --tier full loads a model, ~1-2 h

It detects your machine, picks shapes for your RAM, prints a time estimate, runs the committed benches with the same memory guardrails the project uses, and writes one provenance-complete file plus a ready-to-paste PR. The three-step submission flow is in community-benchmarks/README.md.

Submitted results are folded into the table below (python scripts/aggregate_community.py). Each row is measured on that contributor's own hardware and reported as-is — nothing is extrapolated to machines no one has run, and a row does not imply a trainable-context ceiling beyond what that machine measured. The stock-attention baseline comparison is run on reference hardware by the maintainer, not asked of contributors.

The "Attn flash 2x ratio" column is how the flash forward+backward peak grows per sequence doubling. The O(N) target is about 2×. On the reference machine the largest measured pair (4096→8192) reads 3.06× rather than 2×, because the chained backward split adds a small, budget-bounded constant of at most ~0.3 GB; the growth law is still linear, so a reading a little above 2× on a given machine is expected, not a regression.

MIT. See LICENSE.