gdiamos/amx-moe-e256-4day

A causal language model trained end to end on one CPU core --- a single Intel Emerald Rapids core, bf16 through AMX, OMP_NUM_THREADS=1. 3,316,320 active parameters per token, 100,956,515 stored.

The point of the project is not that a small model runs on a CPU. It is that the architecture is derived from a single-core roofline, that training is confined to the same core, and that at this scale the interesting behaviours show up much earlier in the token budget than we expected.

Paper: Outrageously Small Neural Networks: Emergent Basic Reasoning at 6,616 tok/sec on One Intel AMX Core, shipped here as paper.pdf. It reports the single-core roofline this architecture is derived from, the budget law that bounds how many experts a token budget can support, and the failures that had to be diagnosed to get here -- expert collapse to 1 among them.

What it does

This is a base model: it predicts the next token and has had no instruction tuning. It is well calibrated under teacher forcing and it cannot generate --- free-running, it enters a repetition basin within about five tokens. That is expected of this checkpoint and is what the instruction-tuned sibling exists to fix. Use it as a starting point for fine-tuning, not as a generator.

Running it

AutoModelForCausalLM.from_pretrained will not work: the architecture is not one transformers knows --- chunked sliding-window attention interleaved with log-decay linear attention, and a tied readout and block-routed experts. The model's own source ships here under m2r/, unmodified from the repository that trained it.

hf download gdiamos/amx-moe-e256-4day --local-dir amx-moe-e256-4day
cd amx-moe-e256-4day && pip install -r requirements.txt && python example.py

import sys, torch
from safetensors.torch import load_file
from tokenizers import Tokenizer
sys.path.insert(0, ".")                 # the folder you downloaded

from m2r.config import load
from m2r.model.torch_model import Model, swa_mask

cfg = load("training_config.yaml")
model = Model(cfg.model).to(torch.bfloat16)
model.load_state_dict(load_file("model.safetensors"))
model.eval()
tok = Tokenizer.from_file("tokenizer.json")
mask = swa_mask(cfg.model, dtype=torch.bfloat16)

# Pad to a whole number of blocks and read the last REAL position. Right-padding
# is safe -- attention is causal and the MLP is position-wise.
PAD_TO = max(cfg.model.window, cfg.model.route_block or 1, 256)
ids = [1] + tok.encode("The capital of France is").ids    # 1 is BOS
for _ in range(60):
    n = len(ids)
    x = torch.tensor([ids + [0] * ((-n) % PAD_TO)])
    with torch.no_grad():
        h = model.body(x, mask)[:, n - 1]
    logits = (h @ model.emb.t().to(h.dtype)).float()[0] / 0.8
    v, i = logits.topk(40)
    ids.append(int(i[torch.multinomial(v.softmax(-1), 1)]))
print(tok.decode(ids[1:]))

Sample rather than take the argmax: greedy decoding loops within a few tokens. BOS (id 1) matters --- the model was trained with attention confined to document boundaries keyed on that token, so a prompt without it is unlike anything it saw in training.

What is in this repo

file
model.safetensors	the weights, bf16
generation.json	decode settings, and the vocabulary mask described below
config.json	every architecture field, machine readable
training_config.yaml	the run's config, and what example.py loads
tokenizer.json	a tokenizers BPE; Tokenizer.from_file loads it alone
m2r/	the model source, imported by example.py
example.py	load and generate, correctly
paper.pdf	the write-up, when shipped with this export
LICENSE	Apache 2.0

Architecture

d_model	256
layers	6
layer types	lin, swa, swa, lin, swa, lin
mixers	sliding-window attention (window 256), log-decay linear attention (d_state 32)
MLP width	640
vocabulary	16384
readout	tied to the embedding
parameters	100,956,515 stored, 3,316,320 active per token
experts	256, top-4, d_ff_e 160, route_block 256
MoE layers	[0, 3, 5]

Attention is confined to document boundaries: a training window packs many documents, and without isolation sliding-window attention reaches into its neighbours while linear attention carries state across the whole window.

The shape is deliberate. One AMX core sustains roughly 2,231 GF/s of bf16 matrix multiply at these dimensions but pays a 1.4--1.5 microsecond floor per GEMM dispatch, so every design choice here is about issuing few large matrix multiplies rather than many small ones.

This is a BASE model

It has had no instruction tuning and no QA tuning, and it will not answer a question put to it. Scored on the held-out extractive QA set this project uses for its tuned models, it gets 0.0% exact match and 0.1% F1 -- not a defect, just the wrong eval for this checkpoint. The instruction-tuned sibling gdiamos/amx-reasoning-v1-instruct reaches 18.2% EM and 23.2% F1 on the same rows after two further training stages, and is a dense model.

What this checkpoint is for: a starting point for fine-tuning, and the evidence for the architecture claim below.

Against the dense model it costs the same to run

This is the point of the architecture, so here is the whole comparison rather than the flattering part of it. dense-4day-sft is the same body, same data, same schedule, same seed, with one dense MLP where this model has expert banks on its three linear-attention layers --- matched to 0.02% on active parameters (3,315,552 against 3,316,320) while storing 18x more.

val_flat (full-vocabulary), mean of window	dense	MoE E=256	delta
steps 400k-660k	5.0972	4.5904	-0.507
steps 800k-1.0M	4.0050	3.3181	-0.687
steps 1.1M-1.2M (instruction phase)	4.0259	4.0601	+0.034
final (mean of last 5)	3.9227	4.0147	+0.092

So the MoE is clearly ahead through the middle of training and the two CONVERGE by the end. It does not finish ahead.

Two honest caveats on top of that:

THE TWO VALIDATION METRICS DISAGREE. val_flat is full-vocabulary cross-entropy; val_loss is the sampled training objective, a (1+n_negatives)-way discrimination against unigram-drawn negatives. On identical held-out tokens they rank these two runs in OPPOSITE orders for most of training. We report val_flat because it is perplexity and the sampled number depends on the proposal distribution, but a reader should know the other exists and says something different.

THE TASK LADDER SAYS TIE. At matched step 750,000, in-context induction, positional shift and two-digit addition average 37% for this model against 35% for dense, n=40 per cell. MoE is ahead on retrieval (induct 85% vs 68%, shift 88% vs 82%) and behind on arithmetic (48% vs 60%).

The defensible claim is therefore not that block-routed MoE beats dense here. It is that it MATCHES dense at equal inference cost while storing 18x the parameters, on one CPU core, which is what makes the capacity worth having.

Training data

source	tokens	share
reasoning_anneal.code_reasoning	1,184,900,000	24.6%
foundation.github_code	1,026,000,000	21.3%
foundation.web	756,000,000	15.7%
foundation.code_reasoning	372,600,000	7.7%
foundation.math	324,000,000	6.7%
reasoning_anneal.github_code	204,000,000	4.2%
instruction.tulu_sft	117,300,000	2.4%
instruction.code_reasoning	76,500,000	1.6%
instruction.github_code	76,500,000	1.6%
instruction.web	68,340,000	1.4%
reasoning_anneal.web	68,000,000	1.4%
reasoning_anneal.tulu_sft	68,000,000	1.4%
foundation.tulu_sft	67,500,000	1.4%
reasoning_anneal.math	51,000,000	1.1%
instruction.math	38,250,000	0.8%
foundation.ultrachat	31,199,214	0.6%
reasoning_anneal.ultrachat	31,199,214	0.6%
instruction.ultrachat	31,199,214	0.6%
foundation.task_induct	27,000,000	0.6%
foundation.task_shift	27,000,000	0.6%
foundation.task_add	27,000,000	0.6%
instruction.instruct_sft	25,500,000	0.5%
foundation.instruct_sft	21,600,000	0.4%
reasoning_anneal.instruct_sft	17,000,000	0.4%
reasoning_anneal.task_induct	17,000,000	0.4%
reasoning_anneal.task_shift	17,000,000	0.4%
reasoning_anneal.task_add	17,000,000	0.4%
instruction.task_induct	5,100,000	0.1%
instruction.task_shift	5,100,000	0.1%
instruction.task_add	5,100,000	0.1%
instruction.short_sft2	4,080,000	0.1%
foundation.short_sft2	4,050,000	0.1%
reasoning_anneal.short_sft2	3,400,000	0.1%
reasoning_anneal.short_sft	1,700,000	0.0%
instruction.short_sft	1,530,000	0.0%
foundation.short_sft	1,350,000	0.0%
total	4,819,997,642

Every natural-language and code source above is a curated artefact built with the help of large models --- quality classification, rephrasing, model-assisted extraction, and in the case of the reasoning corpus, traces that are themselves generated output. Training a model this small on them is a form of distillation, with no teacher present at training time. This is worth stating plainly, because it means results at this scale depend on corpora that did not exist when models of this size were last studied seriously.

Validation loss

The run's own numbers on its fixed held-out set, as logged. These are a sampled loss --- a (1 + n_negatives)-way discrimination, not a full-vocabulary one --- except val_flat, which is full-vocabulary.

{
  "val_foundation": {
    "last_3": [
      3.125,
      3.125,
      3.171875
    ],
    "mean": 3.140625
  },
  "val_reasoning_anneal": {
    "last_3": [
      2.84375,
      2.859375,
      2.90625
    ],
    "mean": 2.8697916666666665
  },
  "val_flat": {
    "last_3": [
      3.9145889282226562,
      4.108060836791992,
      4.163447856903076
    ],
    "mean": 4.062032540639241
  },
  "val_loss": {
    "last_3": [
      2.75,
      2.765625,
      2.765625
    ],
    "mean": 2.7604166666666665
  },
  "steps": 1198729,
  "tokens": 4910010368
}

Limitations

A research artifact, and the honest summary is that the failures are specific rather than diffuse.

It degenerates into repetition. Free-running, it enters an absorbing state within about five tokens, in every domain. No decode-time patch fixes this --- truncation sampling has nothing to reshape in a distribution that concentrated. Instruction tuning does fix it, which is what the instruct sibling of this repo is.

It has had no alignment, safety, or preference training of any kind, and will reproduce the biases and errors of its training corpus.

License and provenance

Apache 2.0, for the weights and for the source in m2r/; full text in LICENSE.

The training data is a mixture of public code, web, math and instruction corpora, with per-source token counts above. Those corpora carry their own terms, which the Apache licence on this model does not alter and does not extend to them.

Produced by tools/export_hf.py from run moe-grow-e256-4day-sft-window-20260907T054932 at step 1198730.