Gala-598M-MLX / PAPER.md
junafinity's picture
Gala-598M: weights, code, logs, report
76bbe95 verified
|
Raw History Blame Contribute Delete
21.9 kB

Gala: a million tokens on a laptop, and what it cost

Single-machine research report · code, checkpoints, and raw logs released · 27 Aug 2026

Abstract

We trained a family of small non-Transformer language models, Gala, on one Apple M3 Max in stock MLX, then streamed a million tokens of real FineWeb text through the 188M-parameter member. Ingest took 36.7 seconds. Decode afterwards ran at 377 tok/s out of 3.07 MB of recurrent state and 4.59 GB peak memory, and none of those numbers moved between 32k and 1M tokens of context. The matched Transformer baseline at that length would hold roughly 33 GB of KV cache and, extrapolating its measured curve, decode at 4 to 5 tok/s. That's the deployment story, and it's the headline.

The quality story is smaller and we tell it straight: at 120M matched FineWeb-Edu tokens, every Gala variant lands within 0.8% of the Transformer's cross-entropy. The dense-MLP variant matches or slightly beats it (3.8762 vs 3.8887, single seed); the sparse variants lose by 0.016 and 0.030 nats. We also document how an algebraically exact formula for the chunked delta rule's triangular inverse blew up a provably bounded recurrence ten steps into training, with the measurements that make the failure legible and the blocked substitution that fixed it for a 9% throughput fee. We close with the unified-memory practicalities (the 128 GB ceiling is a cliff, not a slope) and a gated roadmap toward a few billion parameters.

1. Introduction

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 unified memory shared zero-copy between CPU and GPU, and a framework, MLX [11], whose mixture-of-experts primitives (gather_mm, gather_qmm) make very wide, sparsely touched parameter spaces cheap.

Gala is what happens if you take that inversion seriously: hoard parameters and state, the cheap resources, and ration the scarce one, FLOPs per token. (Gala is a variety of apple. The model targets Apple Silicon. That's the whole joke, and we won't make it again.) This report is not a victory lap. At the budget we could afford, a plain Transformer still edges the sparse variants on validation loss, and we print the numbers. What Gala wins, measured rather than argued, is deployment: its per-context state is three megabytes, and it stays three megabytes.

Four things here should transfer even if you never run our code: the million-token demo (§3) and the decode-scaling measurement behind it (§4); a matched-token comparison of four architectures on one machine, caveats stated instead of buried (§5); a numerics post-mortem of the chunked delta rule that any from-scratch port outside the fused-CUDA ecosystem can hit (§6); and operational findings for training under a unified-memory ceiling, ending in a checklist (§7). Section 8 is the roadmap.

2. The model, and what things are called

One pre-norm residual cell, weight-tied, applied 4 times per token with input re-injection and a learned per-loop offset [8,9]:

h ← h + e + loop_embed[ℓ]
h ← h + GDNMixer(RMSNorm h)     Gated DeltaNet; fp32 fast-weight state
h ← h + Hoard_A(RMSNorm h)      1,024–4,096 blocks × 64 neurons, top-32, ReLU²
h ← h + WindowAttn(RMSNorm h)   RoPE, exact attention over a 256-token band
h ← h + Hoard_B(RMSNorm h)

The mixer is Gated DeltaNet [3,5]: short conv, L2-normed queries and keys, data-dependent decay and write strength, constant-size fast-weight state per head. The sparse MLP keeps its working name, the hoard: BDH's wide, positive, sparsely active neuron space [1,2] realized as product-key-routed [6] blocks applied with mx.gather_mm, tokens sorted by block id so each block's weights are read contiguously. A token touches 32 of 1,024 blocks in Gala-S and 32 of 4,096 in Gala-M3. The attention window supplies exact local retrieval, the known weak spot of fixed-state models. Baselines live in the same file and trainer.

We renamed the architecture for this report; the released code keeps the working name hoard in configs, flags, and filenames. The mapping, for diffing against the release:

paper name code name total params active per token per loop
Gala-S hoard_small 188.2M 58.2M
Gala-D gdn_hybrid_small 77.8M 77.8M
Gala-M3 hoard_m3 888.2M ~89M
Transformer transformer_small 76.7M 76.7M

Gala-D is the same mixer stack with dense MLPs in place of the hoard, in four untied cells. Gala-M3 takes its widths from the chip's datasheet (d=768 on 64-lane tile boundaries, 6 mixer heads, 12 attention heads, 4,096 blocks). The Transformer is a plain 8-layer full-attention baseline; it and Gala-D are matched to Gala-S on non-embedding active FLOPs per token.

Setup, once, for everything below: one M3 Max (16-core CPU, 40-core GPU, 128 GB unified memory), macOS 15, MLX 0.32.2 on Metal, bf16 parameters with fp32 master weights. Data is FineWeb-Edu sample-10BT [12], GPT-2 BPE, 588.0M train / 12.0M validation tokens. The trainer runs Muon [10] at 0.02 on hidden 2-D matrices and AdamW at 3e-3 elsewhere, betas 0.9/0.95, 100 warmup steps, 30% linear cooldown, grad clip 1.0, seed 0. Before any long run we make the stack memorize Shakespeare: 300 steps to a train cross-entropy of 0.054, a cheap check that everything can actually fit data.

3. A million tokens on a laptop

bench_1m.py streams real FineWeb text through the trained Gala-S checkpoint in 16,384-token segments at batch 1, then decodes greedily. No tricks: same checkpoint the quality numbers come from, same unfused MLX graph.

context cumulative prefill prefill tok/s decode tok/s recurrent state peak GB
32,768 1.2 s 27,741 394.8 3.07 MB 4.59
131,072 4.8 s 27,329 373.1 3.07 MB 4.59
524,288 18.5 s 28,294 377.4 3.07 MB 4.59
1,048,576 36.7 s 28,540 377.0 3.07 MB 4.59

One million tokens ingested in 37 seconds, on a laptop, after which the model decodes at the same speed it decoded at 32k. Prefill throughput holds near 28k tok/s regardless of how much context came before, so ingest time is linear. State is 3.07 MB in every row. Peak memory is 4.59 GB in every row; it's set by the 16k-segment prefill, not by the context, so it would read the same at ten million.

The contrast, honestly: we didn't run the Transformer at 1M, because on this machine there'd be no point. Its KV cache alone would be about 33 GB for this 8-layer model, its measured decode curve (§4) extrapolates to 4 to 5 tok/s at that depth, and quadratic prefill puts ingest somewhere between tens of minutes and hours. The baseline isn't a strawman, it's the matched model from §5. The claim is about scaling shape: at some context length, every growing-cache model falls off this laptop, and Gala doesn't.

The caveat, in full, because this is the section people will quote: it measures capability, not recall. Gala-S was trained at sequence length 1,024, and everything more than 256 tokens behind the cursor lives in a lossy 3 MB state; we have no evidence yet that the model can fetch a fact from 900k tokens back. Needle-in-a-haystack is the obvious next experiment, and a matched wall-clock study at sequence length 8,192 with two seeds is running as we write this; no result from it appears here.

4. Decode: constant state against a growing cache

The controlled version of the same story. Batch-1 greedy decode of 96 tokens after prefills of increasing length, same checkpoints as §5:

decode

context Gala-S tok/s Transformer tok/s Gala-S peak GB Transformer peak GB
1,024 364.5 642.9 0.65 0.35
4,096 393.5 536.2 1.53 0.93
16,384 384.0 253.1 4.47 4.00
32,768 373.9 134.0 10.38 11.52

Gala-S is flat. Its decode-resident state is about 3 MB at any context: per loop, 4 heads of 128×128 fp32 fast weights (256 KB), a window KV of at most 256 tokens, and a 3-token conv tail. The peak-GB column grows for both models because it includes prefill transients (chunked prefill over a 32k prompt is itself memory-hungry), not decode residency. After prefill, Gala holds its 3 MB while the Transformer holds a cache that grows without bound.

The Transformer starts 1.8× faster at 1k. Both models are kernel-launch-bound at batch 1, and its 8 layers launch fewer kernels than Gala's 4 loops of 4 sub-blocks. It decays 4.8× by 32k as every new token re-reads the whole cache. Crossover sits near 8k; at 32k, Gala decodes 2.8× faster in 0.9× the peak memory. Prefill is mixed: the Transformer wins short contexts (0.05 s vs 0.10 s at 4k), Gala wins long ones (1.8 s vs 2.4 s at 32k) as O(L²) overtakes chunked-linear-plus-window. Absolute numbers are unfused MLX graphs, not an inference engine; the scaling shapes are the result.

5. Quality: matched tokens, stated caveats

The protocol: 120M identical FineWeb-Edu tokens per model, same data order (seed 0), same batch shape (32 × 1024, one flagged exception), same optimizer and schedule, sequential solo runs. Matched tokens isolates architecture from kernel maturity: the attention path inherits years of fused-kernel work (mx.fast.scaled_dot_product_attention); the hoard path is first-generation unfused MLX. The wall-clock view, stated plainly: the Transformer finished its 120M tokens in 1.7 h against Gala-S's 3.0 h. At equal time it sees about 1.8× the tokens and, on these curves, would likely widen its lead. If you care about time-to-quality on today's software, weight that view.

validation loss

model total (active/token/loop) val loss @ 120M Δ vs Transformer train tok/s wall
Gala-D 78M (78M) 3.8762 −0.013 14,589 2.3 h
Transformer 77M (77M) 3.8887 — 19,602 1.7 h
Gala-M3 888M (~89M) 3.9046 +0.016 6,686 5.0 h
Gala-S 188M (58M) 3.9188 +0.030 10,990 3.0 h

Read it with the error bars we don't have. Every run is single-seed, and deltas under about 0.02 nats could plausibly be seed noise. So: Gala-D matches or slightly beats full attention at matched tokens, which is directionally consistent with the production-scale hybrid consensus [4,7]. Gala-S and Gala-M3 lose, by 0.030 and 0.016 nats — which is to say every variant lands within 0.8% of the Transformer's cross-entropy. Gala-M3 carries an extra confound: memory forced it to 16,384 batch tokens per step against everyone else's 32,768, so its trajectory differs in more than architecture; we report it as the chip-derived result, not a controlled comparison. It led Gala-S by 0.10 nats at 6.6M tokens (the wide hoard helps early) and finished 0.014 ahead. And 120M tokens is far short of compute-optimal for any of these models; the ordering at convergence is not knowable from this data.

6. Exact algebra, unstable arithmetic

We got nan by step 10. The second launch too. This section is for anyone reimplementing chunked delta rules outside the fused-CUDA ecosystem; we think it's the part that earns your time.

Bisection first: fp32 was clean, bf16 with AdamW active was broken, and fp32 master weights with fp32 optimizer state — the standard remedy — did not help. So the failure lived in the forward/backward computation, not the update arithmetic. Per-tensor finiteness probes (debug_nan.py, 150 lines, ships in the repo) pointed at the GDN mixer, and a per-chunk trace of the delta-rule state S ended the argument:

chunk  0: |S|max 3.9        chunk  8: |S|max 1.2e23
chunk  2: |S|max 1.4e6      chunk 12: |S|max 5.8e32
chunk  4: |S|max 3.3e11     chunk 15: |S|max nan   (inf − inf)

That trace came from feeding finite inputs through the fp32 chunk core. Roughly 500× growth per 64-token chunk, in fp32, from finite inputs. The instability is a property of the algorithm, not of bf16; the bf16 trajectory just reached the triggering key statistics first.

It can't be the math. Per token, S ← exp(g)(I − βkkᵀ)S + βkvᵀ with β ∈ (0,1), unit keys, and decay ≤ 1, so the exact state obeys ‖S_T‖ ≤ ‖S_0‖ + Σβ‖v‖ — with our activations, a bound in the tens after 1,024 tokens. Growth to 10²³ is impossible in ℝ.

It's the chunk algebra. The chunkwise-parallel form [5] solves T = (I+A)⁻¹ per chunk for a strictly-lower A = (βK)Kᵀ ⊙ D. We had used the nilpotent identity (I−A)(I+A²)(I+A⁴)…, which is exact in ℝ, differentiable, and about 2 log₂ C batched matmuls: precisely the clean trick a from-scratch port reaches for. But it explicitly forms A^(2^k). One optimizer step into training, keys within a chunk correlate, the intermediate powers grow combinatorially, and their fp32 rounding puts O(1) absolute error into T (the true inverse really is large, measured ‖T‖max ≈ 1.3×10³, and the rounding error arrives at the same order). The corrupted W = T(βK ⊙ e^G) then feeds the cross-chunk feedback v_new = U − W·S, which amplifies the error geometrically (the measured ~500× per chunk) into inf 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, which is the point: blocked forward substitution, row substitution inside 16×16 diagonal blocks, block substitution across them. It's the structure fused production kernels already use [13; Megatron's torch_chunk_gated_delta_rule], and it never forms a power of A. All six chunk-vs-recurrent parity tests pass unchanged, and the identical configuration then trained through 480M+ tokens with zero non-finite values. Cost: about 9% training throughput (11.0k against 12.2k tok/s on Gala-S). We don't claim the substitution is novel; kernel authors settled this long ago. After our run we found the theory side had settled it too: Sobczyk et al. [15] showed this spring that the repeated-squaring family of triangular inversions is unstable, with intermediate growth like 2ⁿ/√n, and recommended blocked-recursive schemes — kernel-level analysis, no training runs. The flash-linear-attention tracker holds a related field report of solve_tril explosions traced to unnormalized keys [16], and the state-precision bug in ollama's gated_delta_step [14] is a cousin. What we add to that record is the training-side account: the failure is data-dependent, absent at init, triggered in plain fp32 by keys that correlate a few optimizer steps into learning, and it detonates a provably bounded recurrence from inside — with the per-chunk trace to show it. As far as we can tell, no public kernel ever shipped the nilpotent-product form; we may be its only field casualty, which is exactly why the warning is worth printing. Two hardening changes ship with the fix: the mixer computes fully in fp32 behind boundary casts, and both optimizers keep fp32 masters and state for bf16 parameters.

7. Training on 128 GB: the ceiling is a cliff

Forward+backward throughput at seq 1024, 5-iteration mean:

config dtype batch tok/s peak GB
Gala-S bf16 32 12,240 64.1
Gala-S fp32 32 9,062 86.5
Gala-D bf16 32 15,630 45.7
Transformer bf16 32 19,080 40.6
Gala-S bf16 64 6,590 125.8
Gala-M3 bf16 32 6,375 79.6
Gala-M3 bf16 16 8,152 41.3

bf16 runs about 35% faster than fp32 throughout. The bolded row is the operational finding. Unified memory doesn't degrade gracefully: doubling Gala-S's batch lands at 125.8 GB, against the 128 GB ceiling, and throughput halves. Gala-M3 tells the same story from the other side, 28% faster at batch 16 than at batch 32. On this hardware you size the batch under the ceiling. Filling memory is exactly wrong.

What we'd tell anyone training a recurrent-state model on unified memory:

  1. Size the batch under the ceiling, never to fill it. The last few GB cost half your throughput.
  2. Never form powers of a triangular matrix to invert it. Use (blocked) forward substitution; exact-in-ℝ identities are not fp32-safe once the data correlates.
  3. Keep recurrent state, decay gates, and the chunk solve in fp32. Apple GPUs run fp32 at full rate, so this is nearly free.
  4. Give bf16 parameters fp32 masters and fp32 optimizer state. bf16-only AdamW moments diverged for us within 10 steps, independently of the solve bug.
  5. Gate every long run: parity tests, a short memorization run with a loss ceiling, and per-tensor finiteness probes wired up before you need them.

8. Roadmap

Everything below is gated the way this report's runs were: a tier starts when the one before it produces evidence. None of it is a promise.

Gala-0.5B fits the current M3 Max. Before it starts we want three things from the models we have: needle-in-a-haystack numbers at 4k and 16k (the recall question §3 left open), a seed spread on at least the S and D configs so deltas like −0.013 mean something, and router-health histograms showing the hoard isn't quietly half-collapsed. Scaling an unmeasured router is how you waste a month.

Gala-1B is a software problem before it's a hardware one. The hoard path needs a fused Metal kernel with a hand-written VJP, decode needs the 4-bit gather_qmm path, and the multi-token-prediction head that's designed but unimplemented should earn its keep here. Machine class: M4/M5-Ultra or a small Mac cluster. Gate: Gala-0.5B trains through its budget with the evidence above holding.

Gala-3B is roughly the trainable ceiling for 128 GB, with 16,384 hoard blocks, and it needs a 10B+ token data pipeline with a pinned dataset revision this time. Gate: the 1B fused kernels hit their throughput targets, because at 3B an unfused graph stops being quaint and starts being the schedule.

Gala-27B-class is the Qwen3.8-27B reference point, and we're not romantic about it: at that scale the sensible moves are probably hybrid attention ratios closer to production practice and distillation from open models rather than from-scratch pretraining. Gate: everything above, plus a reason to believe the hoard's parameters-for-FLOPs trade still pays at depth.

9. Limitations

Single seed per configuration, no measured spread. 120M tokens per model, several-fold short of compute-optimal. No downstream or recall evaluation yet; the million-token demo is speed and memory only, and the 256-token window may be hiding real retrieval failures. No ablations of loop count, window size, or hoard top-k. Router health is unmeasured, so partial router collapse can't be excluded. Gala-M3's batch size is a confound. FineWeb-Edu was streamed without a pinned revision, so byte-exact reproduction needs pinning. The MTP head, CPU-stream routing, and 4-bit decode remain designed but unimplemented. Decode and million-token numbers come from unfused MLX graphs; an optimized engine would move the absolutes, not the shapes.

10. Conclusion

On one M3 Max, in stock MLX: a trained 188M model ingested a million tokens in 37 seconds and decoded at an unchanging 377 tok/s from 3 MB of state; three variants landed within 0.8% of a matched Transformer's cross-entropy, the dense one edging ahead on a single seed; an exact shortcut was caught detonating a provably bounded recurrence; and the memory ceiling turned out to be a cliff. The machine was a co-author. Both of its defining constraints, the unified-memory ceiling and full-rate fp32, are visible in every table above.

Artifacts. All code (model.py, train.py, bench.py, bench_decode.py, bench_1m.py, debug_nan.py, run_chain.sh), configs, raw training logs, and checkpoints for every model in this report ship with the release, which keeps the working name hoard. Total wall clock behind every number here: about 13.5 hours, plus 37 seconds for the part people will remember.

References

[1] Kosowski, A., Uznański, P., Chorowski, J., Stamirowska, Z., Bartoszkiewicz, M. The Dragon Hatchling: The Missing Link between the Transformer and Models of the Brain. arXiv:2509.26507, 2025. [2] BDH-CQ: In-Context Learning with Recurrent Latent Reasoning. arXiv:2608.09888, 2026. [3] Yang, S., Kautz, J., Hatamizadeh, A. Gated Delta Networks: Improving Mamba2 with Delta Rule. arXiv:2412.06464, 2024. [4] Gated DeltaNet-2. arXiv:2605.22791, 2026. [5] Yang, S., et al. Parallelizing Linear Transformers with the Delta Rule over Sequence Length. arXiv:2406.06484, 2024. [6] Lample, G., et al. Large Memory Layers with Product Keys. arXiv:1907.05242, 2019. [7] Peng, B., et al. RWKV-7 "Goose" with Expressive Dynamic State Evolution. arXiv:2503.14456, 2025. [8] LT2: Linear-Time Looped Transformers. arXiv:2605.20670, 2026. [9] Saunshi, N., et al. Reasoning with Latent Thoughts: On the Power of Looped Transformers. arXiv:2502.17416, 2025. [10] Jordan, K., et al. modded-nanogpt and the Muon optimizer. github.com/KellerJordan/modded-nanogpt. [11] Hannun, A., et al. MLX: Efficient and flexible machine learning on Apple Silicon. github.com/ml-explore/mlx. [12] Penedo, G., et al. The FineWeb Datasets: Decanting the Web for the Finest Text Data at Scale. arXiv:2406.17557, 2024. [13] fla-org, flash-linear-attention. github.com/fla-org/flash-linear-attention. [14] ollama issue #15865, gated_delta_step state precision. github.com/ollama/ollama/issues/15865. [15] Sobczyk et al. Fast and Stable Triangular Inversion for Delta-Rule Linear Transformers. arXiv:2605.21325, 2026. [16] flash-linear-attention issue #537, Gated Delta Rule, solve_tril Numerical Explosion. github.com/fla-org/flash-linear-attention/issues/537.