Gala-598M: weights, code, logs, report

76bbe95 verified 2 months ago

5 kB

---
license: apache-2.0
language:
- en
library_name: mlx
pipeline_tag: text-generation
datasets:
- HuggingFaceFW/fineweb-edu
tags:
- mlx
- apple-silicon
- pretrained-from-scratch
- gated-deltanet
- linear-attention
- product-key-memory
- long-context
---
# Gala-598M — a Mac-first language model, trained from scratch on one M3 Max

Gala is a non-transformer LM designed around Apple Silicon's actual profile: lots of
unified memory, modest FLOPs. **Hoard parameters and state (cheap on a Mac), ration FLOPs
per token (scarce).** Everything here — pretraining, three-seed baselines, a 300M-token
scale study, long-context finetunes, memory characterization, and a 10.5M-token context
demo — ran on a single MacBook (M3 Max, 40-core GPU, 128 GB) in stock MLX. No CUDA, no
custom kernels, no cluster.

![10M tokens](https://huggingface.co/datasets/junafinity/gala-article-assets/resolve/main/fig_10m.png)

## Architecture

Four untied cells, each: **Gated DeltaNet mixer** (fixed-size fp32 fast-weight state) →
**hoard MLP** (product-key-routed sparse blocks via `mx.gather_mm`: 1,024 blocks × 64
neurons, top-32 — a token touches 3% of a 537M-parameter pool) → **256-token sliding-window
attention** → hoard MLP. 598M total parameters, ~78M active per token. Tied embeddings.
GPT-2 BPE.

## Headline results (all measured, logs in `results/`)

**Quality, matched data** — 120M identical FineWeb-Edu tokens, same batch order, same
optimizer, 3 seeds each:

| | seeds | mean ± spread |
|---|---|---|
| Gala-598M | 3.8145 / 3.8128 / 3.8203 | **3.816 ± 0.008** |
| Transformer-77M (matched active FLOPs) | 3.8887 / 3.8856 / 3.8863 | 3.887 ± 0.003 |

Gap ≈ 10× seed noise. A tied-embeddings Transformer variant (3.860) also loses. At **300M
tokens** the lead **grows**: 3.581 vs 3.684.

**Context scaling** — 10,485,760 real FineWeb tokens streamed through this checkpoint,
batch 1, on the laptop: ingest ~6 minutes, decode **382–388 tok/s flat at every depth**,
loss-on-next-2k **flat from 32k to 10.5M**, recurrent state **3.07 MB constant**, peak RAM
5.3 GB. The matched Transformer decays 643 → 134 tok/s by 32k and would need ~33 GB of KV
at 1M.

**Memory, characterized honestly** — passkey retrieval, 30 trials × 6 depths, chance = 20%:
**100% (30/30) when the needle is inside the 256-token window, at both 8k and 32k context;
chance at every depth beyond it.** The state carries a measurable gist (loss at depth
2–4k into a document is 0.36 better than at depth 0–256) but not verbatim facts. The
in-window control shows the boundary is architectural, not model capability.

## Files

| file | what |
|---|---|
| `gala_u_300M.safetensors` | main checkpoint: 300M FineWeb-Edu tokens, val 3.581 |
| `gala_u_32k_ft.safetensors` | + curriculum finetune to seq 32k (1k→8k→32k), val 3.564 @32k |
| `model.py` … `campaign4.py` | complete training/eval/orchestration code (pure MLX) |
| `results/` | raw training logs and campaign verdicts for every number above |
| `PAPER.md`, `CHRONICLE.md`, `NOVELTY.md` | the report, the five-day worklog, the novelty sweep |

## Use

```bash
pip install "mlx>=0.32" numpy tiktoken
python sample.py --run <dir-with-checkpoint-and-config> --prompt "The history of" --tokens 150
python bench_1m.py --run <dir> --total 10485760     # the 10M demo, ~7 min on an M3 Max
```

(To use these weights with the scripts: put `model.safetensors` + `config.json` in one
directory — e.g. rename `gala_u_300M.safetensors` and `gala_u_300M.config.json`.)

## Training

Muon (0.02, NorMuon-style row normalization) on hidden matrices + AdamW (3e-3) elsewhere,
fp32 master weights and optimizer state, bf16 compute, fully-fp32 GDN mixer, `mx.compile`
(+7.5%), seq 1024, batch 32,768 tokens, warmup 100 / 30% linear cooldown, grad clip 1.0.
~11.3k tok/s training throughput on the M3 Max.

One numerics note worth reading if you reimplement chunked delta rules (`PAPER.md` §6):
computing the chunk's triangular inverse via the exact nilpotent identity
(I−A)(I+A²)(I+A⁴)… detonates in fp32 once training correlates the keys (~500×-per-chunk
state growth into inf, from finite inputs, through a provably bounded recurrence). Use
blocked forward substitution. Instability class independently established by
arXiv:2605.21325; the training-side account and the probe that found it (`debug_nan.py`)
are in this repo.

## Limitations

A research model, not an assistant: 420M training tokens total, FineWeb-Edu only, no
instruction tuning, no downstream evals yet. Verbatim recall is window-local by
construction (see above). Training wall-clock is ~1.5× slower than the matched Transformer
on today's unfused MLX sparse path. Single machine, single dataset; scale conclusions
beyond 598M/300M tokens are roadmap, not claims.

## Citation

```
@misc{gala2026,
  title  = {Gala: a million tokens on a laptop, and what it cost},
  author = {junafinity},
  year   = {2026},
  note   = {https://huggingface.co/junafinity/Gala-598M-MLX}
}
```

	---
	license: apache-2.0
	language:
	- en
	library_name: mlx
	pipeline_tag: text-generation
	datasets:
	- HuggingFaceFW/fineweb-edu
	tags:
	- mlx
	- apple-silicon
	- pretrained-from-scratch
	- gated-deltanet
	- linear-attention
	- product-key-memory
	- long-context
	---

	# Gala-598M — a Mac-first language model, trained from scratch on one M3 Max

	Gala is a non-transformer LM designed around Apple Silicon's actual profile: lots of
	unified memory, modest FLOPs. **Hoard parameters and state (cheap on a Mac), ration FLOPs
	per token (scarce).** Everything here — pretraining, three-seed baselines, a 300M-token
	scale study, long-context finetunes, memory characterization, and a 10.5M-token context
	demo — ran on a single MacBook (M3 Max, 40-core GPU, 128 GB) in stock MLX. No CUDA, no
	custom kernels, no cluster.

	![10M tokens](https://huggingface.co/datasets/junafinity/gala-article-assets/resolve/main/fig_10m.png)

	## Architecture

	Four untied cells, each: Gated DeltaNet mixer (fixed-size fp32 fast-weight state) →
	hoard MLP (product-key-routed sparse blocks via `mx.gather_mm`: 1,024 blocks × 64
	neurons, top-32 — a token touches 3% of a 537M-parameter pool) → **256-token sliding-window
	attention** → hoard MLP. 598M total parameters, ~78M active per token. Tied embeddings.
	GPT-2 BPE.

	## Headline results (all measured, logs in `results/`)

	Quality, matched data — 120M identical FineWeb-Edu tokens, same batch order, same
	optimizer, 3 seeds each:

	\| \| seeds \| mean ± spread \|
	\|---\|---\|---\|
	\| Gala-598M \| 3.8145 / 3.8128 / 3.8203 \| 3.816 ± 0.008 \|
	\| Transformer-77M (matched active FLOPs) \| 3.8887 / 3.8856 / 3.8863 \| 3.887 ± 0.003 \|

	Gap ≈ 10× seed noise. A tied-embeddings Transformer variant (3.860) also loses. At **300M
	tokens the lead grows**: 3.581 vs 3.684.

	Context scaling — 10,485,760 real FineWeb tokens streamed through this checkpoint,
	batch 1, on the laptop: ingest ~6 minutes, decode 382–388 tok/s flat at every depth,
	loss-on-next-2k flat from 32k to 10.5M, recurrent state 3.07 MB constant, peak RAM
	5.3 GB. The matched Transformer decays 643 → 134 tok/s by 32k and would need ~33 GB of KV
	at 1M.

	Memory, characterized honestly — passkey retrieval, 30 trials × 6 depths, chance = 20%:
	**100% (30/30) when the needle is inside the 256-token window, at both 8k and 32k context;
	chance at every depth beyond it.** The state carries a measurable gist (loss at depth
	2–4k into a document is 0.36 better than at depth 0–256) but not verbatim facts. The
	in-window control shows the boundary is architectural, not model capability.

	## Files

	\| file \| what \|
	\|---\|---\|
	\| `gala_u_300M.safetensors` \| main checkpoint: 300M FineWeb-Edu tokens, val 3.581 \|
	\| `gala_u_32k_ft.safetensors` \| + curriculum finetune to seq 32k (1k→8k→32k), val 3.564 @32k \|
	\| `model.py` … `campaign4.py` \| complete training/eval/orchestration code (pure MLX) \|
	\| `results/` \| raw training logs and campaign verdicts for every number above \|
	\| `PAPER.md`, `CHRONICLE.md`, `NOVELTY.md` \| the report, the five-day worklog, the novelty sweep \|

	## Use

	```bash
	pip install "mlx>=0.32" numpy tiktoken
	python sample.py --run <dir-with-checkpoint-and-config> --prompt "The history of" --tokens 150
	python bench_1m.py --run <dir> --total 10485760 # the 10M demo, ~7 min on an M3 Max
	```

	(To use these weights with the scripts: put `model.safetensors` + `config.json` in one
	directory — e.g. rename `gala_u_300M.safetensors` and `gala_u_300M.config.json`.)

	## Training

	Muon (0.02, NorMuon-style row normalization) on hidden matrices + AdamW (3e-3) elsewhere,
	fp32 master weights and optimizer state, bf16 compute, fully-fp32 GDN mixer, `mx.compile`
	(+7.5%), seq 1024, batch 32,768 tokens, warmup 100 / 30% linear cooldown, grad clip 1.0.
	~11.3k tok/s training throughput on the M3 Max.

	One numerics note worth reading if you reimplement chunked delta rules (`PAPER.md` §6):
	computing the chunk's triangular inverse via the exact nilpotent identity
	(I−A)(I+A²)(I+A⁴)… detonates in fp32 once training correlates the keys (~500×-per-chunk
	state growth into inf, from finite inputs, through a provably bounded recurrence). Use
	blocked forward substitution. Instability class independently established by
	arXiv:2605.21325; the training-side account and the probe that found it (`debug_nan.py`)
	are in this repo.

	## Limitations

	A research model, not an assistant: 420M training tokens total, FineWeb-Edu only, no
	instruction tuning, no downstream evals yet. Verbatim recall is window-local by
	construction (see above). Training wall-clock is ~1.5× slower than the matched Transformer
	on today's unfused MLX sparse path. Single machine, single dataset; scale conclusions
	beyond 598M/300M tokens are roadmap, not claims.

	## Citation

	```
	@misc{gala2026,
	title = {Gala: a million tokens on a laptop, and what it cost},
	author = {junafinity},
	year = {2026},
	note = {https://huggingface.co/junafinity/Gala-598M-MLX}
	}
	```