Gala: an Apple-Silicon-first model, trained on one

I spent five days treating my MacBook (M3 Max, 40-core GPU, 128 GB unified memory) as a training machine instead of an inference toy. The result is a small language-model family, working name Gala, designed around what this chip is actually good at: enormous unified memory and modest FLOPs. The design rule that falls out is simple. Hoard parameters and state, because memory is cheap here. Ration FLOPs per token, because compute is not.

Here is the exit exam. Ten and a half million tokens of real FineWeb text streamed through the trained model, batch 1, on the laptop:

Nothing grows: not the state, not the latency, and (the part I wasn't sure about going in) not the loss. Ingest is linear at ~28k tok/s, so the full 10M took about six minutes. The matched Transformer we trained on identical data decodes at 134 tok/s by 32k context and would need roughly 33 GB of KV cache at 1M. At 10M it stops being a comparison and becomes a category difference.

Everything below is measured, the logs ship with the weights, and the model is honest about what it cannot do. That last part turns out to be the most interesting section.

Why design for a Mac at all

The Transformer is the child of its hardware: dense matmuls near 1,000 TFLOPS, scarce high-bandwidth memory, FLOPs treated as free. A Mac is the opposite machine. An M3 Max has maybe 1/70th the matmul throughput of an H100, but 128 GB of memory shared zero-copy between CPU and GPU, and a framework, MLX, whose mixture-of-experts primitives (gather_mm, gather_qmm) make very wide, sparsely touched parameter spaces cheap.

Nobody trains here. The public record before this week: toy GPT-2 replications, a 2M-param documented ceiling for on-Mac pretraining, and the most serious MLX-pretraining author renting an H100 for his 440M runs. Meanwhile the community's loudest recurring complaint about Macs is long-context pain: nine-minute prompt processing, KV caches that outgrow RAM, whole products built to page KV to SSD. So the experiment was: design a model for this chip, train it on this chip, and measure everything.

The architecture, in one paragraph

Four cells, each: a Gated DeltaNet mixer (a fast-weight recurrence with a fixed-size fp32 state), a hoard MLP (product-key-routed sparse blocks: 1,024 blocks of 64 neurons, top-32 per token, applied with mx.gather_mm, so each token touches about 3% of a 537M-parameter pool), and a 256-token sliding-window attention for exact local recall. 598M total parameters, ~78M active per token, tied embeddings, GPT-2 BPE. Pure MLX, no custom kernels.

Does it actually work? Three seeds say yes

Matched everything: 120M identical FineWeb-Edu tokens, identical batch order, identical optimizer (Muon on hidden matrices, AdamW elsewhere), against a Transformer baseline matched on active FLOPs per token.

Gala: 3.816 ± 0.008 across three seeds. Transformer: 3.887 ± 0.003. The gap is about ten times the seed noise. Anticipating the obvious objection, we also gave the Transformer tied embeddings, which helped it (3.860) and it still lost.

Then we asked whether the lead survives more data, and trained both on 300M tokens:

The lead grew, from 0.071 to 0.103 nats. This is the capacity thesis doing what it promised: the 537M-parameter hoard keeps converting data after the 77M dense model starts flattening. Two diverging curves are the strongest chart a scaling argument can produce, and this one came off a laptop.

It's worth saying what the first version of this model did: it lost. Our original design reused one weight-tied cell four times per token, and it trailed the Transformer by 0.030. The single biggest improvement of the whole project was unrolling that loop into four untied cells at identical FLOPs. The elegant trick was the mistake.

What its memory really is

The 10M table invites a fair objection: reading is not remembering. So we measured the difference instead of hand-waving it, with passkey retrieval done properly: 30 trials per cell, six depths, two context lengths, and chance at 20%.

Two results. With the needle inside the 256-token attention window: 100%. Thirty out of thirty, at both 8k and 32k context. With the needle anywhere beyond the window: chance, at every depth.

The in-window control is what makes this a finding rather than an excuse. The model aces the task whenever the answer is in reach, so the boundary is the architecture, not model stupidity. Gala has photographic recall up close and a lossy 3 MB gist of everything else. The gist is real and measurable (loss at position 2-4k into a document runs 0.36 better than at position 0-256), but it cannot fetch an exact fact from a million tokens back. If your workload streams a lot of text cheaply, this is your shape. If it hunts needles, you want retrieval or full attention on top.

The bug that almost ended the project on night one

Our first bf16 launch went NaN by step 10. The second one too. Bisection cleared the optimizer; per-tensor probes pointed into the mixer; and a chunk-by-chunk trace of the recurrent state ended the argument:

The chunked delta rule solves a small triangular system every 64 tokens. We had computed its inverse with the exact nilpotent identity (I−A)(I+A²)(I+A⁴)…, which is exact in real arithmetic, differentiable, and fast. It is also a trap: it explicitly forms A^(2^k), and once training correlates the keys, those intermediates blow up and their fp32 rounding corrupts the inverse. The corrupted factor feeds the cross-chunk recurrence, the error compounds about 500× per chunk, and a provably bounded recurrence hits infinity inside one forward pass. Everything passes at init because at init the keys aren't correlated. The bug waits for learning to start.

The fix is boring on purpose: blocked forward substitution, the same structure production CUDA kernels use, for about a 9% throughput fee. Theory caught this instability class in the spring (arXiv:2605.21325); as far as we can tell, no public kernel ever shipped the bad form. We may be its only field casualty, which is exactly why the write-up, the failing configuration, and the 150-line probe that found it (debug_nan.py) are all in the repo. If you reimplement chunked delta rules outside the fused-CUDA ecosystem, this section is for you.

The serving story, quantified

Batch-1 decode on the same checkpoints: the Transformer starts faster (643 tok/s at 1k, both models are launch-bound when the cache is small) and decays 4.8× by 32k as every token rereads a growing cache. Gala holds 364-394 tok/s from 1k to 32k, and, per the first table, to 10M. Crossover lands near 8k. Per-session state is ~3 MB against a cache that grows forever, which is also the difference between serving a thousand sessions in 3 GB and serving a few.

What we learned about training on Apple Silicon

Things we measured that transfer to any MLX training project:

1. The 128 GB ceiling is a cliff, not a slope. Doubling the batch to fill memory pushed peak usage to 125.8 GB and halved throughput. Size the batch under the ceiling, never to it. We hit the same inversion twice on two different models.
2. bf16 needs fp32 master weights and optimizer state. bf16-only AdamW moments diverged within 10 steps, reliably. The fp32-master wrapper costs a little memory and buys correctness.
3. Keep the recurrent core in fp32. Apple GPUs run fp32 and fp16 at the same ALU rate, so the safety is nearly free.
4. mx.compile is free money (+7.5% on our training step) once your step avoids Python control flow. Gradient checkpointing passed parity at a 34% cost; worth it only at long sequence.
5. Big-lab recipes don't automatically transfer down. Meta's memory-layer tricks (dedicated value LR, silu output gating) both hurt at our scale. We ran the probes so you don't have to.

What this is not

A chatbot. It's a 598M research model trained on 420M tokens of FineWeb-Edu; it writes plausible English and knows nothing. The Transformer still trains about 1.5× faster in wall-clock, because its attention kernels have had years of optimization and our sparse path is first-generation unfused MLX (a fused Metal kernel via mx.fast.metal_kernel with a custom VJP is the top roadmap item). The 128k passkey evaluation timed out before finishing. One machine, one dataset, no downstream evals yet, and every scale claim beyond what's in the tables is roadmap, not result.

Reproduce it

Everything is one command on any Apple Silicon Mac with MLX. The 10M demo is python bench_1m.py --total 10485760, about seven minutes on an M3 Max. The full training campaign is python campaign3.py: roughly 36 hours, fully unattended, restartable, with a live HTML dashboard. Code, both checkpoints, every training log, and the five-day worklog including the failures:

junafinity/Gala-598M-MLX

Total electricity for everything in this article: about 2 kWh. Nothing left the laptop.

What's next

In order: a fused Metal kernel for the sparse path (the wall-clock gap is a kernel gap, not an architecture gap), standard downstream evals, a bigger constant state (300 MB is still constant, and constant is the property that matters), and the 0.5B → 3B ladder that 128 GB of unified memory makes possible. Apple shipped RDMA over Thunderbolt and demoed distributed training at WWDC26, so the multi-Mac chapter has an obvious opening. Nobody has published multi-Mac pretraining yet. That flag is still on the table.

Key references: Gated DeltaNet (arXiv:2412.06464), chunked delta-rule parallelization (arXiv:2406.06484), product-key memories (arXiv:1907.05242), triangular inversion stability (arXiv:2605.21325), MLX, flash-linear-attention. Full bibliography and the formal write-up (PAPER.md) in the model repo.