gdiamos/amx-reasoning-v1

A 7,492,448-parameter causal language model trained on 4,909,998,080 tokens, from the m2-reasoning project. Small enough to train on one CPU core -- this checkpoint was trained on a single Intel Emerald Rapids core using AMX through oneDNN.

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. The model's own source is included here instead, under m2r/, unmodified from the repository that trained it.

hf download gdiamos/amx-reasoning-v1 --local-dir amx-reasoning-v1
cd amx-reasoning-v1 && pip install -r requirements.txt && python example.py

or directly:

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: the
# training forward chunks by `window` / the linear-attention chunk, and a
# prompt is whatever length it is. 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. At this size greedy decoding falls into a repetition loop within a few tokens, which says more about argmax than about the model.

example.py is exactly that script. 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
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: config.py, model/torch_model.py, model/moe.py
example.py	load and generate
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	7,492,448 stored, 3,315,552 active per token

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 entire window.

Training data

source	tokens	share
reasoning_anneal.code_reasoning	1,475,753,528	30.1%
foundation.github_code	1,175,719,140	23.9%
foundation.web	870,198,318	17.7%
foundation.code_reasoning	417,689,772	8.5%
foundation.math	367,413,336	7.5%
reasoning_anneal.github_code	260,043,420	5.3%
reasoning_anneal.web	77,672,272	1.6%
reasoning_anneal.math	65,010,364	1.3%
foundation.task_induct	29,460,000	0.6%
foundation.task_shift	29,460,000	0.6%
foundation.task_add	29,460,000	0.6%
foundation.instruct_sft	21,520,530	0.4%
reasoning_anneal.instruct_sft	21,519,548	0.4%
reasoning_anneal.task_induct	19,640,000	0.4%
reasoning_anneal.task_shift	19,640,000	0.4%
reasoning_anneal.task_add	19,640,000	0.4%
reasoning_anneal.short_sft2	3,448,784	0.1%
foundation.short_sft2	3,446,820	0.1%
foundation.short_sft	1,632,084	0.0%
reasoning_anneal.short_sft	1,632,084	0.0%
total	4,910,000,000

Metrics

These are the run's own validation numbers, on its fixed held-out set, as logged. They 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.140625,
      3.28125,
      3.28125
    ],
    "mean": 3.234375
  },
  "val_reasoning_anneal": {
    "last_3": [
      2.765625,
      2.703125,
      2.765625
    ],
    "mean": 2.7447916666666665
  },
  "val_flat": {
    "last_3": [
      3.3155245780944824,
      3.2603909969329834,
      3.385425329208374
    ],
    "mean": 3.3204469680786133
  },
  "val_loss": {
    "last_3": [
      2.765625,
      2.703125,
      2.765625
    ],
    "mean": 2.7447916666666665
  },
  "steps": 1198729,
  "tokens": 4909998080
}

Intended use and limitations

A research artifact. At this scale the model produces locally fluent text and is not reliable for factual questions, instruction following, or reasoning.

It degenerates into repetition. example.py as shipped, sampling at temperature 0.8 with top-k 40, produces something like

The capital of France is commonly used in enhancing enhancing productivity practices worldwide worldwide worldwide worldwide worldwide worldwide...

Greedy decoding is worse -- it loops within a few tokens. No repetition penalty is applied anywhere; adding one would hide the behaviour rather than change it. Treat fluent-looking spans as coincidence rather than capability.

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/. The full text is in LICENSE.

The training data is a mixture of public code, web, math and instruction corpora, with per-source token counts in the table 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-20260831T204431-16k-4day at step 1198730.