gdiamos/amx-moe-4day-w2048

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,429 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-moe-4day-w2048 --local-dir amx-moe-4day-w2048
cd amx-moe-4day-w2048 && 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.0625,
      3.015625,
      3.03125
    ],
    "mean": 3.0364583333333335
  },
  "val_reasoning_anneal": {
    "last_3": [
      2.28125,
      2.28125,
      2.28125
    ],
    "mean": 2.28125
  },
  "val_flat": {
    "last_3": [
      4.885426998138428,
      4.837184429168701,
      4.90023946762085
    ],
    "mean": 4.874283631642659
  },
  "val_loss": {
    "last_3": [
      2.90625,
      2.890625,
      2.859375
    ],
    "mean": 2.8854166666666665
  },
  "steps": 427245,
  "tokens": 1749999616
}

Which run this is, and how its experts actually behave

This is the control arm of a router-normalisation A/B at window 2048: E=64 experts, top-4 (two of them shared), 1.75 B tokens, 427,246 steps. Its twin differs only by M2R_MOE_ROUTER_NORM. It is not the project's best model on any benchmark; it is a clean, completed baseline.

The three MoE layers behave completely differently from each other, which is worth knowing before fine-tuning from it. Measured on these weights with a real batch from the training mixture:

For reference the dense blocks' RMSNorm gains are 0.0726 (block 1), 0.7861 (block 2) and 1.3187 (block 4), so all three MoE layers sit well below them.

Read the three columns together, because the gain alone is misleading. Every MoE layer damped its input gain hard, but the expert banks grew enough to more than compensate in blocks 3 and 5 -- so block 0's branch is effectively inert while block 5's carries a fifth of the residual stream. The enormous middle column is the signature of that trade: a tiny input scaled up by very large weights.

Two consequences, stated plainly:

• Depth matters more than the MoE/dense distinction here. A fine-tune that assumes the expert banks are uniformly useful, or uniformly dead, will be wrong about two of the three layers either way.
• This is not a clean demonstration that block-routed MoE works. The gain collapse is real, it is what this project calls expert collapse, and block 0 is what it looks like when it goes all the way.

The run discarded 367 gradients in 427,245 steps (0.086%) as non-finite and finished. Sibling runs on nearby configurations did worse, so the stability here should not be assumed to transfer.

Validation loss, its own

Full-vocabulary held-out loss, 42 draws over the run:

Mean of the last six draws 4.5648; best single draw 3.2218. The spread between adjacent draws is large -- 3.22 to 4.90 within 130 k steps -- so single draws should not be compared across models.

The machine-generated block below reports 4.8743 for val_flat, which is the mean of the last THREE draws rather than six. Both are correct and the difference is entirely draw noise, which is the point: at this variance a three-draw window is not a stable estimate. Use the six-draw figure, and do not compare either against another model's single draw.

Limitations

A research artifact, and the honest summary is that the failures are specific rather than diffuse. The one above is architectural: the three MoE layers damped their input gains hard and only two of the three compensated, so block 0's experts are close to inert while block 5's carry a fifth of the residual.

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 torch-20260925T081457-4day-w2048-ctrl at step 427246.