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

Architecture

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.

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.