amx-reasoning-v1

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

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

Training data

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.