amx-diffusion-pretrain-4day

A 126M-parameter (10.1M active) block-diffusion mixture-of-experts language model. It was pretrained from scratch for 95 hours on 4 CPU cores (Intel Xeon Silver 4514Y with AMX, bf16), with no GPU.

This is a research checkpoint from a CPU-training study. It is pretrained only. It is not instruction-tuned and has no benchmark evaluations.

Model

Parameters	125,981,216 total, 10,146,656 active per token
Layers	6 hybrid: lin swa swa lin swa lin (linear attention / sliding-window attention, window 2048)
Width	d_model 320, 2 query heads + 1 KV head of 32, d_state 32
FFN	dense d_ff 800 (ReLU²)
MoE (layers 0, 3, 5)	64 experts: 2 shared + 62 routed, top-2 routed, product-key router, d_ff_e 640, one routing decision per 256-token block
Vocab	16,384 (BPE, tokenizer.json)
Context	4,096
Readout	tied: logits = norm(h) · embᵀ, full softmax

Objective: block diffusion. The sequence is split into 64-token blocks. Each block gets a mask rate t ~ U[0.125, 1] and its masked tokens are replaced by [MASK] (id 8). The model sees a clean stream and a noisy stream, and it predicts the masked tokens of each block given the clean prefix before it. The loss is cross-entropy on the masked positions with uniform weighting. Decoding is planned as 8 denoising steps per block.

Routing.

• The shared experts always run at weight 1.
• The two routed experts' gates are a softmax over just the two selected logits (shared_gate: one, gate_norm: selected).
• Expert selection adds a DeepSeek-V3-style balancing bias. It affects only which experts are picked, not their gates. The bias moves by a sign step of 1e-3 each step on the global counts and is frozen for the last 5% of training.
• A 1e-4 auxiliary balance loss acts as a guard.

Training

Data	~2.7B tokens of a 3-phase mix: foundation (code, web, math, reasoning traces), reasoning anneal, instruction
Steps	41,198 × 65,536 tokens (16 windows of 4,096)
Optimizer	AdamW β (0.9, 0.95), weight decay 0.01, grad clip 1.0
LR	peak 1.6e-3 (from a grid at this batch size), 5% warmup, cosine to 0.1×
Precision	bf16 autocast, fp32 master weights
Hardware	4 cores of one Xeon Silver 4514Y socket; ZeRO-1 over shared memory; ~7,900 tok/s mean
Wall time	95.4 h

Results

Validation loss (nats per masked token) on held-out windows of each phase, at three mask rates:

phase	t = 0.125	t = 0.5	t = 1.0	mean
foundation	1.750	2.913	6.213	3.625
reasoning_anneal	1.338	2.323	5.475	3.046
instruction	1.689	2.730	5.854	3.424

A sibling run trained the same model with a softmax gate over all 64 experts. It ended at 1.771 on foundation t = 0.125 (3.637 mean). In that run the routed experts' gates collapsed to ~0.01, so most routed picks had almost no effect on the output. With the gating of this model, the routed gates stay at 0.2–0.4. eval_results.json holds the full validation history.

Files

• model.safetensors: fp32 weights
• modeling_amx_diffusion.py: standalone model and block sampler (plain PyTorch)
• config.yaml: the training config (model, objective, data, optimizer, schedule)
• tokenizer.json: the 16k BPE tokenizer (shared with gdiamos/amx-moe-4day-w2048)
• eval_results.json: final and per-step validation losses

Usage

This is not a transformers model. modeling_amx_diffusion.py is a standalone, plain-PyTorch implementation. It needs torch, safetensors, pyyaml and huggingface_hub, plus tokenizers for text. It matches the training code's forward to 1e-13 in float64 on these weights.

import os, sys
from huggingface_hub import hf_hub_download
from tokenizers import Tokenizer

repo = "gdiamos/amx-diffusion-pretrain-4day"
sys.path.insert(0, os.path.dirname(hf_hub_download(repo, "modeling_amx_diffusion.py")))
from modeling_amx_diffusion import AmxDiffusion

model = AmxDiffusion.from_pretrained(repo)
tok = Tokenizer.from_file(hf_hub_download(repo, "tokenizer.json"))

prompt = [1] + tok.encode("def fibonacci(n):\n").ids        # 1 = BOS
print(tok.decode(model.generate(prompt, max_new_tokens=64)))

generate decodes greedily, one 64-token block at a time, in 8 passes per block. Each pass commits the masked positions with the largest top-1 − top-2 probability margin. Each pass re-runs the whole sequence, so it is simple rather than fast.

For scoring, model(clean, noisy) returns the final-normed hidden states of both streams. Put [MASK] (id 8) at the positions to predict in noisy. model.logits(hidden) then gives the vocabulary logits. Ids 0, 7 and 8 are never predicted (id 8 doubles as <eos> in the tokenizer). Sequence lengths must be a multiple of 64 and start with BOS.

Limitations

• Pretraining only: no instruction tuning, alignment, or safety work.
• The sampler is a reference implementation with no KV/state cache, and there is no int8 serving export.
• Generations are repetitive: it is a 10M-active-parameter base model.
• Small: it is a study of CPU training, not a capable assistant.
• 10% of the planned code-reasoning data for the anneal phase was missing.