{ "slug": "mlx-windowed-backward", "lab_number": 28, "title": "Off NVIDIA, you fix the platform yourself", "type": "editorial", "status": "published", "sponsor": null, "author": "Keith Townsend", "date": "October 22, 2026", "date_iso": "2026-10-22", "layers": [ "layer0a" ], "axes": [ { "instrument": "4plus1", "layers": [ "layer0a" ], "has_dapm_table": true } ], "supersedes": [], "superseded_by": null, "corrects": [], "corrected_by": [], "vendors": [ { "key": "apple", "role": "hw" }, { "key": "google", "role": "model" }, { "key": "nvidia", "role": "ref" }, { "key": "anthropic", "role": "bench" } ], "themes": [ "accelerator-lock-in", "operating-cost-of-ownership" ], "finding": "MLX's pending sliding-window change (PR #4734) didn't fix training: Gemma 4's sliding-window layers still built the full score matrix in the backward pass. An owner-written patch (a core MLX diff of +107/-15 lines) made it scale about linearly (4.7 GB against 198 GB at 32,768 tokens), and Lab 27's 31B fine-tune then trained on mlx-lm's own attention path 15% faster than the model-side workaround, at the same held-out loss in one run each.", "video": null, "lab_detail": "https://labs.layer2c.com/downloads/mlx-windowed-backward-lab-detail.md", "sources": [ { "label": "This lab's patch: kltownsend/mlx, branch sdpa-windowed-backward (GitHub)", "url": "https://github.com/kltownsend/mlx/tree/sdpa-windowed-backward" }, { "label": "MLX PR #4734: sliding window attention", "url": "https://github.com/ml-explore/mlx/pull/4734" }, { "label": "MLX PR #4563: VJP for scaled dot product attention", "url": "https://github.com/ml-explore/mlx/pull/4563" }, { "label": "MLX issue #3860: fp32 matmul defaults to TF32-class precision", "url": "https://github.com/ml-explore/mlx/issues/3860" } ], "question": "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 attention in the model code. Stock MLX, Apple's machine-learning framework, ran out of memory on the longest example. MLX had a fix in review, PR #4734, adding sliding-window attention. The Spark ran the job on NVIDIA's ecosystem without framework changes, after pinning an older transformers release. So would the pending fix let the Studio do the same, and if not, what does it take to fix it yourself? It didn't, and fixing it took a patch to MLX's attention backward, a source build, and a fork that nobody maintains.", "load": "Lab 27's job unchanged: Gemma 4 31B, dense, bf16, LoRA rank 32 on attention q, k, v and o, Lab 18's corpus of 171 agent transcripts (mean 8,663 tokens, longest 30,553), 2 epochs, 86 steps of 4, 21 held-out validation examples. One Mac Studio (M5 Ultra, 256 GB), owned. MLX built from source at PR #4734's head with this lab's backward patch; mlx-lm 0.32.2 with a 30-line runtime patch.", "verdict": { "scope": "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.", "independence": "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.", "calls": [ { "kind": "do", "text": "On a non-NVIDIA framework like MLX, budget to close platform gaps yourself. Stock MLX couldn't train this job: the attention backward for Gemma 4's sliding-window layers built the full score matrix, 198 GB for one layer at 32,768 tokens. The core fix was +107/-15 lines in MLX's attention backward, plus a 30-line mlx-lm patch and a source build. On NVIDIA, Lab 27's Spark ran the same job without framework changes, after pinning an older transformers release." }, { "kind": "do", "text": "Test the vendor's pending fix against your job before you wait for it. MLX PR #4734 adds sliding-window attention, and it was approved for merge. With it, training still went quadratic: 12.9 GB at 8,192 tokens, 50.5 GB at 16,384. The forward pass was fixed; the backward, which training needs, was marked not yet implemented." }, { "kind": "do", "text": "Count the platform fix as worth more than a model-side workaround. With the patch, Lab 27's 86-step job trained through mlx-lm's own attention path in 2 hours 17 minutes against 2 hours 41 for my rewritten attention, with held-out loss 0.1496 against 0.1510, one run each; Lab 18's coding loop wasn't rerun. The fix sits in MLX's shared Metal attention path, so other models that use its window_size form can use it; the workaround had to be written per model. Only Gemma 4 31B was trained." }, { "kind": "dont", "text": "Don't treat a working fork as free. It sits on an unmerged change, needs a source build with Xcode's Metal toolchain, and has to be retested, and usually rebased, on every MLX release you move to. I published it and I'm not maintaining it. That carrying cost is the overhead the NVIDIA ecosystem doesn't charge." } ] }, "objection": { "q": "Why patch a framework? Wait for upstream, or buy NVIDIA.", "a": [ "Waiting wasn't a plan. No issue existed for the backward. The one change in review, #4734, fixed the forward pass and marked the backward not yet implemented, and with it the job still ran out of memory. There was no date for the rest.", "Buying NVIDIA is a real answer, and it's the point of the title. In Lab 27 the DGX Spark trained this job without framework changes, after pinning an older transformers release. It took 10.25 hours. The Studio, fixed, took 2.29. That crosses hardware, frameworks and dates, so it isn't a clean price for the ecosystem, but it's the speed the fix bought on MLX.", "And the cost was bounded. A core MLX diff of +107/-15 lines, a 30-line mlx-lm patch, a test, and a two-minute source build, written using Claude Code. The patch and its first checks took about an hour of a session that started at 7:54 PM; the full training runs finished overnight. What isn't bounded is the maintenance, which is why the fork is public and unmaintained." ] }, "numbers": [ { "label": "Attention backward memory, one sliding-window layer", "value": "4.74 GB against 198 GB at 32,768 tokens", "note": "Patched against the boolean mask mlx-lm passes today. Peak above inputs, bf16, head_dim 256, 32 query and 16 KV heads, mx.grad through mx.fast.scaled_dot_product_attention. Patched: 1.24 / 2.49 / 4.74 GB at 8,192 / 16,384 / 32,768. Boolean mask: 12.94 / 50.25 / 198.0." }, { "label": "With PR #4734 alone", "value": "12.94 and 50.5 GB at 8,192 and 16,384 tokens", "note": "window_size=1024 on #4734's head (fdf1d256a) without this lab's patch. Still quadratic: the backward falls back to the full matrix." }, { "label": "Training time, 86 steps", "value": "8,228 s against 9,670 s", "note": "Patched MLX, mlx-lm's own attention path, against Lab 27's rewritten attention on MLX 0.32.4. Lab 27's DGX Spark took 36,885 s on its stock stack." }, { "label": "Held-out validation", "value": "loss 0.1496 against 0.1510", "note": "21 examples, 11,458 labeled tokens, Lab 18's loss, scored in Lab 27's runtime. Next-token match 94.77% against 94.81%; same prediction as Lab 27's Spark-trained adapter 98.46% against 98.31%. Lab 27's two Spark-trained adapters: 0.1536 and 0.1534." }, { "label": "Training loss, 86 steps", "value": "within 0.010 of Lab 27 on average", "note": "Mean absolute step-loss gap 0.010, largest 0.083. Final step 0.1758 against 0.1746." }, { "label": "Longest example, one step (30,553 tokens)", "value": "95.4 s, 182 GB against 114.8 s, 205 GB", "note": "Patched against rewritten attention, both on the patched build. Process footprint. The rewritten attention reproduced Lab 27's step to every digit on this build. Whole-run footprint peaked at 231 GB for both, set by a different step." }, { "label": "Forward loss, longest example", "value": "1.5508651853, identical to stock mlx-lm", "note": "Patched windowed path against mlx-lm's stock boolean mask, forward only, same build." }, { "label": "Gradient parity, small shapes", "value": "867 cases", "note": "fp32/fp16/bf16, head dims 64 to 512, grouped-query ratios, queries shorter than keys, windows 1 to 5,000, causal and not, tile sizes 256 and 1,024, against an fp32 reference. Every window from 7 up within twice MLX's own causal backward error, except one bf16 case that missed by 0.0004. At window 1 the true gradient is zero and only rounding is left." }, { "label": "Attention gradient error, real Gemma 4 activations", "value": "dQ 5.1 to 6.7% against 2.1 to 2.9%", "note": "Patched backward against MLX's dense fallback, sliding layers 0 to 4, relative to an fp32 CPU reference; cosine 0.998 or better. With an opt-in fp32 score tile, 0.7 to 1.5%, better than the fallback. That adapter scored 0.1530, no better." }, { "label": "Size of the fix", "value": "+107 / -15 lines, 6 files", "note": "Core MLX diff including a new test, on top of PR #4734, plus a 30-line mlx-lm runtime patch. The opt-in fp32 experiment adds about 45 more. Builds from source in about 2 minutes with Xcode's Metal toolchain." }, { "label": "Whole-model gradient check, 5,085 tokens", "value": "cosine -0.037 to fp32 truth", "note": "Patched path, LoRA gradient at the start of training, against Lab 27's fp32 CPU truth. Grad norm 2,668 against the truth's 536, the largest of the bf16 paths. Stock mlx-lm -0.056, Lab 27's rewrite 0.141: every bf16 path is near noise here, so this check couldn't validate any of them." }, { "label": "Metered spend", "value": "$0", "note": "Owned hardware throughout." } ], "opened": [ { "q": "Will MLX take a windowed backward upstream?", "a": "The patch is public on a fork, with a README and the tests. No pull request or issue went to MLX. Until upstream ships one, training Gemma 4 at this length on stock MLX means carrying the fork or rewriting attention in the model." }, { "q": "Why did the MLX-trained adapter run Lab 18's coding loop slower than the PyTorch-trained one?", "a": "In Lab 27, on the Studio, the PyTorch-trained adapter ran the 22 tasks in 1 hour 46 minutes and the MLX-trained one in about 2.5 hours, with two apparatus faults in the slower arm. Training with PyTorch on the Studio would separate the trainer software from the hardware, and PyTorch on Apple's GPU stops at a 32-bit indexing limit near 11,600 tokens. Open." } ], "not_proved": [ "One model, one job, one Mac Studio, one training run per arm. The patch passed MLX's own SDPA tests and this lab's checks; it hasn't been reviewed by anyone who maintains MLX.", "Metal only. MLX's CUDA backend has no sliding-window kernels, and the patch makes a windowed CUDA backward throw rather than run.", "It builds on #4734, which isn't merged and may change before it is. A window combined with an arbitrary mask still falls back to the full matrix in training.", "The new adapters weren't run through Lab 18's 22-task coding loop. Held-out loss and next-token agreement are the quality evidence here.", "The NVIDIA comparison is Lab 27's Spark run on its stock stack, not a fresh run; no faster Spark software was tried there either.", "End-to-end gradient correctness is unresolved. The whole-model gradient check against fp32 truth was near noise for every bf16 path, the patched path's grad norm was the largest, and only one training run per arm was made. The kernel-level checks and the trained outcome are the evidence, and they're encouraging, not production validation.", "The 15% speedup is against Lab 27's logged rewritten-attention run, not a side-by-side full rerun on the patched build (the longest-example step was rerun and reproduced Lab 27 exactly), and not against an NVIDIA box, and the whole-run peak footprint didn't drop: a different step sets it.", "I wrote the patch and the checks using Claude Code; the lab detail explains every change in plain language. Nobody outside this lab has reviewed the code." ], "executive_summary": [ "Lab 27 trained Lab 18's dense Gemma 4 31B fine-tune on a Mac Studio, but only after rewriting attention in the model code. Stock MLX ran out of memory at a 224 GB footprint on the 30,553-token example. The cause was narrow: MLX's memory-saving attention backward handled the plain causal mask and nothing else, and Gemma 4 uses a sliding window, where each token sees only the 1,024 before it, in 50 of its 60 layers. Those layers fell back to building the full score matrix, which grows with the square of the length.", "MLX PR #4734, from an MLX maintainer, adds sliding-window attention and was approved for merge on October 21. Built from source with it, the backward was still quadratic: 12.94 GB at 8,192 tokens and 50.5 GB at 16,384, the same as the boolean mask. Its own code marks the windowed backward as not yet implemented.", "The patch teaches the existing tiled backward the window: skip tiles that lie entirely outside it, mask the cells at its edge, and fix two bookkeeping assumptions the causal-only code never had to make. Backward memory scaled about linearly over the range measured: 1.24, 2.49 and 4.74 GB at 8K, 16K and 32K tokens, against 198 GB for the boolean mask at 32K.", "Gradients held against an fp32 reference on 867 small shapes, within twice the error of MLX's own causal backward, except one bf16 case that missed by 0.0004 and the one-token window, where the true gradient is zero and only rounding is left. A whole-model gradient check against fp32 truth couldn't separate any bf16 path, this one included. Lab 27's full 86-step job then trained through mlx-lm's own attention path in 8,228 seconds against 9,670 for the rewritten attention, with held-out loss 0.1496 against 0.1510. Its forward loss on the longest example matched stock mlx-lm to every digit." ], "dapm": [ { "instrument": "4plus1", "reading": "decision", "layer": "layer0a", "placement": "Ceded", "cells": { "p": "Ceded" }, "note": "Upstream MLX's feature decisions, on the Mac Studio training runtime as shipped (0.32.4) and with PR #4734: what the stock attention kernels support is the MLX project's decision, and its decision here was that the windowed backward wasn't implemented yet. The owner couldn't train the job on that runtime. Proposal, Keith's to rule." }, { "instrument": "4plus1", "reading": "decision", "layer": "layer0a", "placement": "Retained", "cells": { "p": "Retained" }, "note": "The deployed fork and its patch: the owner decided what the backward supports and made it so. Retaining it means carrying it: rebase and rebuild per MLX release, no upstream maintainer. Lab 27 placed the Studio's training runtime Retained under the portability reading; this table reads decision and supplements that placement rather than reversing it." } ], "seam_map": null, "writeup": [ "Lab 27 left a question open: would stock MLX train Gemma 4 31B once sliding-window attention landed? The change that adds it, MLX PR #4734, was approved for merge on October 21. So I built MLX from that branch on the Studio and measured the backward pass, the part training needs. It was still quadratic. #4734 teaches the forward kernels the window and leaves the backward marked not yet implemented.", "The gap was narrow. MLX already had a tiled backward, merged in September, that walks the score matrix one block at a time and never holds all of it. It skipped blocks above the causal diagonal. It didn't know about windows, so a windowed layer fell back to the full matrix. The patch adds the second skip, for blocks entirely outside the window, and masks the cells at the window's edge. Two assumptions in the causal-only code broke with a window, and both would have produced wrong gradients silently: the first block each query visits is no longer block zero, and some key blocks are never visited at all.", "With the patch, memory for one sliding-window layer's backward went from 198 GB at 32,768 tokens to 4.74. Gradients matched an fp32 reference across 867 small shapes. On real Gemma 4 activations they pointed the right way, cosine 0.998 or better, with about twice the dense fallback's bf16 error. The extra error comes from the existing tiled backward, which keeps its scratch scores in bf16. Keeping them in fp32 cut it fourfold and didn't train a better adapter.", "The test that matters is the job. Lab 27's 86 steps, same data, order and starting adapter, ran through mlx-lm's own attention path in 2 hours 17 minutes, against 2 hours 41 for Lab 27's rewritten attention. Held-out loss was 0.1496 against 0.1510, within 0.004 of all three of Lab 27's earlier adapters (0.1510 to 0.1536), one run each; Lab 18's coding loop wasn't rerun. The baseline is Lab 27's logged run, not a side-by-side rerun. Within those limits, the model-side rewrite wasn't flattering the Studio. It cost it about 15%.", "What it took: a source build of MLX with Xcode's Metal toolchain, an unmerged branch, a +107/-15 line diff, a test, and a 30-line patch to mlx-lm so Gemma passes the window instead of a boolean mask. The fork is public. I'm not maintaining it, and nobody else is yet. On NVIDIA, Lab 27's Spark ran this job on the shipped stack; that's what the ecosystem buys, and this is what it costs to do without it." ], "abstraction": [], "method": [ "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." ], "assessments": [], "publisher": "The Advisor Bench LLC", "program": "Layer2C Labs", "url": "https://labs.layer2c.com/labs/mlx-windowed-backward", "markdown_url": "https://labs.layer2c.com/labs/mlx-windowed-backward.md", "json_url": "https://labs.layer2c.com/labs/mlx-windowed-backward.json" }