Off NVIDIA, you fix the platform yourself

One model (Gemma 4 31B, bf16), one long-sequence LoRA job (Lab 27's), one Mac Studio (M5 Ultra, 256 GB), one training run per arm, Metal only. The patch sits on an unmerged MLX change and hasn't been reviewed by anyone who maintains MLX. Quality is held-out loss and next-token agreement; Lab 18's coding loop wasn't rerun. The NVIDIA comparison is Lab 27's DGX Spark run, not rerun here, on different hardware, framework and date. The open question is the carrying cost of a fork nobody maintains.

Self-funded editorial. The Mac Studio was bought by the practice. Disclosure: Google Cloud is a client of this practice, and the model under test, Gemma 4 31B, is a Google model. Google did not commission, fund, preview or have any say in this lab. Apple's MLX is the framework under test. Neither Apple, NVIDIA nor Anthropic is a client. I wrote and tested the patch using Claude Code, an Anthropic product. No vendor paid for this answer.

In Lab 27 the Mac Studio trained a dense Gemma 4 31B fine-tune 3.8 times faster than a DGX Spark. But only after I rewrote part of the model's code. Stock MLX, Apple's machine-learning framework, ran out of memory on the longest training example. MLX had a fix in review. So I tested it against the same job. Would the pending fix let the Studio train it on stock software, and if not, what would it take to fix it myself?

The pending fix didn't do it. Fixing it myself took a patch to MLX's attention code, a 30-line patch to the training library, a source build, and a fork that nobody maintains. The fixed Studio then trained the job 15% faster than my workaround, with the same held-out loss.

What was broken?

Attention is the part of a language model where every word looks back at earlier words to decide what matters. It builds a table of scores, one row per word and one column per earlier word. For a 30,553-token transcript that table is huge, and it grows with the square of the length.

Using a model only needs that table for a moment. Training needs more. It runs a backward pass afterward, working out how each weight should change, and done the plain way that means holding the whole table. MLX already had a shortcut for the backward pass. It works through the table one block at a time and never holds all of it. But it only understood one rule: each word sees every word before it.

Gemma 4 uses a second rule in 50 of its 60 layers, a sliding window. Each word sees only the 1,024 words before it. MLX's shortcut didn't handle the window, so those layers fell back to the full table. At 32,768 tokens, one layer's backward pass needed 198 GB. That's why the Studio ran out of memory at 224 GB in Lab 27.

Did MLX's pending change close it?

No. The pending change, MLX pull request 4734, comes from an Apple engineer on the MLX team. It adds the sliding window, and a reviewer approved it for merge on October 21. I built MLX from that branch and measured it. The backward pass still needed 12.9 GB at 8,192 tokens and 50.5 GB at 16,384, the same as before. The change teaches the window to the part of the code that uses the model. The part training needs is marked not yet implemented. There was no issue open for it and no date.

So test a pending fix against your job before you wait for it. This one never promised the backward pass, and it didn't train this model.

What did fixing it take?

The patch teaches MLX's existing shortcut the window. It skips blocks of the table that sit entirely outside the window and masks the cells at its edge. Two assumptions in the old code broke with a window, and both would have produced wrong numbers without an error. The first block each word visits is no longer block zero, and some blocks are never visited at all. The core change is 107 added and 15 removed lines across six files, plus a test. I wrote it using Claude Code, Anthropic's coding agent. The patch and its first checks took about an hour. The full training runs went overnight.

Memory now grows about in step with length. One sliding-window layer's backward pass dropped to 4.74 GB at 32,768 tokens. Against an fp32 reference, gradients held across 867 small test shapes, except one case that missed by 0.0004 and a one-token window, where the true gradient is zero and only rounding is left. A whole-model gradient check couldn't tell any of the training paths apart, mine included. So the kernel checks and the trained result carry the evidence.

Then the test that matters. Lab 27's full job, 86 steps from the same data and starting point, trained through mlx-lm's own attention in 2 hours 17 minutes. My rewritten attention had taken 2 hours 41. Held-out loss was 0.1496 against 0.1510, one run each. I didn't rerun Lab 18's coding loop on the new adapter, and the baseline was Lab 27's logged run, not a fresh one beside it. Within those limits, my workaround wasn't flattering the Studio. It was costing it about 15%.

So what's the catch?

Somebody has to carry it. The fix sits on an unmerged change, builds from source with Xcode's Metal toolchain, and has to be retested, and usually rebased, on every MLX release you move to. I published it as a fork. I'm not going to maintain it, and nobody else does yet.

That's the overhead of not going with NVIDIA's ecosystem. In Lab 27 the Spark trained this same job without any framework changes, after I pinned an older release of its training library. It took 10.25 hours against the fixed Studio's 2.29. That's a different box, framework and date, so it isn't a clean price for the ecosystem. But nobody had to patch the Spark's framework. On MLX, the speed was there. The platform work was mine.

Is that worth it? For one owner with one model, an evening of work bought a faster, cleaner venue. For a team that can't own a fork, it's a reason to stay on the stack someone else maintains. Off NVIDIA, budget for fixing the platform yourself.

The numbers

Where each layer belongs

Sources

Method and disclosure

Self-funded. The Mac Studio was bought by the practice. No vendor was briefed, reviewed this, or was given advance sight. Google Cloud is a client of this practice, and the model under test is a Google model; Google had no part in this lab.

Build: ml-explore/mlx at PR #4734's head (fdf1d256a, 2026-10-20), editable install from source into its own virtual environment, Xcode 27.1's Metal toolchain via DEVELOPER_DIR, about 2 minutes per build. A baseline build of #4734 alone and a patched build, side by side; Lab 27's MLX 0.32.4 environment left untouched.

The patch: the Metal forward may emit the logsumexp under a window; the backward primitive carries window_size, including in its cache key; the tiled backward skips key tiles outside each query tile's window and masks edge tiles; dQ is initialized on each query tile's first visited key tile; dK and dV are zeroed when a window with fewer queries than keys leaves key tiles unvisited; CUDA throws for a windowed backward; a new test, test_sdpa_window_grad. An opt-in experiment, MLX_SDPA_VJP_FP32_TILES=1, computes the backward's score tiles in fp32.

Checks: backward peak memory with mx.grad at 8K to 32K tokens; gradient parity on 867 shapes against an fp32 reference, scored against MLX's own causal backward on the same shape; per-layer checks on captured Gemma 4 31B activations against an fp32 CPU reference; MLX's test_fast_sdpa (36 tests) with the fp32 flag off and on.

Training: Lab 27's trainer copied with one flag to install the mlx-lm runtime patch, Lab 27's shared inputs, order and seed-42 starting adapter, 86 steps. Validation with Lab 27's own script in Lab 27's runtime. The rewritten-attention arm reproduced Lab 27's longest-example step to every digit on the patched build, so Lab 27's logged run is the baseline.

Model-level gradient checks against Lab 27's fp32 CPU truth (5,085-token example) were near noise for every bf16 path, as Lab 27 found, and didn't separate them; the per-layer checks and the trained outcome carry the evidence.

Peer check: the draft was attacked by one outside model, ChatGPT (gpt-5.6-sol), before revision; Google Antigravity was over its quota, and Keith ruled one reviewer enough. The findings and my adjudication are kept with the working files.

The patch, probe, checks, runtime patch and trainer are the instrument and are published (the fork and the lab detail). The verdict and the trained adapters stay proprietary.

I wrote the patch using Claude Code, which also ran the bench, so the drafting model's vendor also appears on the bench. The fork's commits and README say it was generated with Claude.

The essay on this page was drafted by Claude Opus 5.5, an Anthropic model, from the lab’s frozen structured record under this practice’s voice specification, checked by a deterministic gate against every figure in the record, and validated by the author before publication. The record is the canonical surface: where prose and record disagree, the record rules.