gdiamos/amx-reasoning-moe-4day-v2

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,315,744 active parameters per token, 30,029,417 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.

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-reasoning-moe-4day-v2 --local-dir amx-reasoning-moe-4day-v2
cd amx-reasoning-moe-4day-v2 && 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	30,029,417 stored, 3,315,744 active per token
experts	64, 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.

Training data

The run consumed 1,749,999,616 tokens --- 12.8% of the 13,651,264,478-token corpus prepared for it, taken in schedule order. The table below is the PREPARED corpus, not the consumed slice; per-source shares are of the corpus.

source	tokens	share
reasoning_anneal.code_reasoning	3,031,714,478	22.2%
foundation.github_code	2,926,000,000	21.4%
foundation.web	2,217,600,000	16.2%
foundation.code_reasoning	1,062,600,000	7.8%
foundation.math	924,000,000	6.8%
reasoning_anneal.github_code	582,000,000	4.3%
instruction.tulu_sft	290,000,000	2.1%
instruction.web	290,000,000	2.1%
reasoning_anneal.web	257,050,000	1.9%
reasoning_anneal.tulu_sft	218,250,000	1.6%
instruction.code_reasoning	217,500,000	1.6%
instruction.github_code	217,500,000	1.6%
foundation.tulu_sft	215,600,000	1.6%
instruction.qa	174,000,000	1.3%
instruction.math	168,200,000	1.2%
reasoning_anneal.math	145,500,000	1.1%
foundation.qa	77,000,000	0.6%
foundation.task_induct	77,000,000	0.6%
foundation.task_shift	77,000,000	0.6%
foundation.task_add	77,000,000	0.6%
reasoning_anneal.qa	72,750,000	0.5%
reasoning_anneal.task_induct	48,500,000	0.4%
reasoning_anneal.task_shift	48,500,000	0.4%
reasoning_anneal.task_add	48,500,000	0.4%
instruction.ultrachat	24,650,000	0.2%
instruction.instruct_sft	24,650,000	0.2%
reasoning_anneal.ultrachat	24,250,000	0.2%
reasoning_anneal.instruct_sft	24,250,000	0.2%
foundation.ultrachat	23,100,000	0.2%
foundation.instruct_sft	23,100,000	0.2%
instruction.task_induct	14,500,000	0.1%
instruction.task_shift	14,500,000	0.1%
instruction.task_add	14,500,000	0.1%
total prepared	13,651,264,478
consumed by this run	1,749,999,616	12.8%

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.109375,
      3.09375,
      3.109375
    ],
    "mean": 3.1041666666666665
  },
  "val_reasoning_anneal": {
    "last_3": [
      2.296875,
      2.296875,
      2.296875
    ],
    "mean": 2.296875
  },
  "val_flat": {
    "last_3": [
      5.036965370178223,
      5.066086769104004,
      5.095897674560547
    ],
    "mean": 5.066316604614258
  },
  "val_loss": {
    "last_3": [
      2.84375,
      2.796875,
      2.765625
    ],
    "mean": 2.8020833333333335
  },
  "steps": 427245,
  "tokens": 1749999616
}

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-best-4day-v2-20260919T044803 at step 427246.