Gala-598M: weights, code, logs, report
76bbe95 verified 2 months ago
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](paper_assets/decode.png)

| 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](paper_assets/loss_curves.png)

| 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.