keeeeenw/MicroLlama2-base

MicroLlama V2

Thank you for the support and warm reception of MicroLlama V1! I still want to answer the same question: can an individual with consumer GPUs and a small budget build a meaningful language model?

For V2, I went small again: a 388M-parameter base language model, trained from scratch on 50B valid next-token targets, with a Qwen3-based backbone and Engram-style conditional memory.

My original compute budget was $500. I went over it: I have spent $880 on compute so far. This is my reported compute spend, not an all-inclusive accounting of storage, electricity or personal hardware. The NVIDIA 50B run is complete; the model in this repository is that checkpoint.

At 5.76B training targets, V2 already exceeded V1's reported-50B seven-task mean by 3.37 percentage points. This is an encouraging result for small-model experimentation. It is a comparison between different model sizes, datasets and recipes, rather than a controlled claim of training efficiency.

Model Details

V1 followed the TinyLlama lineage. V2 uses a new, randomly initialized Qwen3-based architecture and a lightweight Engram-inspired memory implementation. It does not start from pretrained Qwen weights and is not an instruction-tuned or Qwen thinking model.

• Developed by: Zixiao Ken Wang (keeeeenw)
• Funded by: myself; $880 in compute spent so far, above the original $500 goal
• Model type: autoregressive base language model, 387,898,305 parameters
• Backbone / memory: 337,694,720 / 50,203,585 parameters
• Language: primarily English
• Context: 8,192 tokens
• License: Apache 2.0 for this model and inference implementation; upstream dataset terms still apply
• Final NVIDIA checkpoint: step 381,824, 50,000,054,177 valid targets (a final-update overshoot of 54,177)

24 decoder layers
hidden size 1024; intermediate size 2816; SiLU
16 query heads; 4 key/value heads; head dimension 64
65,536 model vocabulary rows; tied input/output embeddings
RoPE theta 10,000; RMSNorm epsilon 1e-6
Engram after layers 3 and 11 (zero based)
normalized multi-head bigram/trigram lookup; four hash heads per order
128 memory dimensions per order; shared gated key/value projections

The memory is inspired by Engram, not a complete reproduction of its architecture: this version omits convolution, multi-branch integration and memory offloading. The tokenizer is the pinned LiquidAI/LFM2.5-230M-Base tokenizer, reused as a tokenizer only. It has 64,402 IDs; unused model-vocabulary rows are suppressed during generation. Training and the examples below use plain text without a chat template or automatic BOS insertion.

Training data and hardware

1. First 10B targets: FineWeb-Edu.
2. 10B to 23.774B: 60% additional FineWeb-Edu and 40% DCLM-Edu.
3. 23.774B to 50B: 48% FineWeb-Edu, 32% DCLM-Edu, 15% Cosmopedia v2 and 5% FineMath 4+, applied to future valid targets.

Training used plain PyTorch on rented NVIDIA RTX 4090/5090 hardware, including a four-RTX-5090 continuation on Vast.ai. The tokenizer, data manifests and full restart states are preserved. Cosmopedia is an existing publicly released synthetic corpus; unlike V1's original goals, this recipe does include public model-generated training text. I did not commission a new synthetic corpus for this run.

Model Sources

• Inference model: MicroLlama2-base
• Full checkpoints and token archives: MicroLlama2-388M-checkpoints
• Inference implementation: configuration_microllama2.py, modeling_microllama2.py, and infer.py
• Technical report: will be linked shortly.
• Training code and reproduction instructions: MicroLlama2-training

Uses

This is a research base completion model. Useful directions include continued pretraining, task-specific fine-tuning, small-model experiments and studying how factual answers and hallucinations change across checkpoints. It has not been aligned for helpful conversation, safety-critical advice or tool use.

1. Install dependencies

Use Python 3.10 or newer and an appropriate PyTorch build for your hardware:

pip install -U "torch>=2.9.1,<3" "transformers==5.20.2" "huggingface-hub==2.2.0" safetensors

Transformers 5.20.2 is the current stable version used to validate this inference adapter. The reference tests use PyTorch 2.9.1, CPU and FP32. CUDA/MPS can be selected below; lower precision and other hardware can produce different results. This custom architecture requires trust_remote_code=True: review the two model files before enabling it. Beam search and speculative/static-cache decoding are not supported by this adapter; ordinary greedy decoding and sampling carry both KV-cache and Engram suffix state.

2. Run code!

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer

model_id = "keeeeenw/MicroLlama2-base"
# For reproducible experiments, replace main with a commit from this repo's history.
revision = "main"
# Set this to a large/external drive if you do not want the default local cache.
cache_dir = None

device = (
    "cuda" if torch.cuda.is_available()
    else "mps" if torch.backends.mps.is_available()
    else "cpu"
)
tokenizer = AutoTokenizer.from_pretrained(model_id, revision=revision, cache_dir=cache_dir, trust_remote_code=True)
model = AutoModelForCausalLM.from_pretrained(
    model_id,
    revision=revision,
    cache_dir=cache_dir,
    trust_remote_code=True,
    dtype=torch.float32,
    attn_implementation="eager",
).to(device).eval()

# A plain continuation prompt, not a chat/instruction template.
prompt = "The water cycle describes"
inputs = tokenizer(prompt, add_special_tokens=False, return_tensors="pt").to(device)
with torch.inference_mode():
    output = model.generate(**inputs, max_new_tokens=64, do_sample=False, use_cache=True)
print(tokenizer.decode(output[0], skip_special_tokens=True))

For sampling, use do_sample=True, temperature=0.8, top_p=0.9, top_k=40. For batches, set tokenizer.padding_side = "left". Keep the prompt plus generation within 8,192 tokens. MPS runs the signed-int64 memory hash on CPU while keeping floating-point model computation on the GPU.

The root inference files reuse the same FP32 weights as the final checkpoint. A normal from_pretrained call downloads roughly 1.55 GB of weights, plus small tokenizer/config/code files. It does not download Adam states, the checkpoint history or the tokenized dataset. This inference repository contains no Adam states, training datasets or checkpoint history. Avoid an unrestricted snapshot_download of the separate archive repository unless you intend to fetch those files.

Evaluation

Original seven-task setup (same task family as V1)

V1's original evaluation followed TinyLlama's GPT4All evaluation guide. For our V2 comparison, we ran the seven full tasks using a version-pinned lm-evaluation-harness setup, native tokenizers, CPU FP32 and 8K context. We use token-normalized acc_norm except raw acc for WinoGrande and BoolQ. Scores below are percentages; the average is an unweighted seven-task mean.

The late V2 benchmark column is the independently reconciled 49.237B checkpoint, not an evaluation of the exact final 50B weights. Final-checkpoint training and held-out validation completed, but we do not claim a full benchmark suite for those exact weights. All reference columns in this table are our saved full CPU evaluations under this seven-task protocol.

Model	Training targets/tokens	HellaSwag	OBQA	WinoGrande	ARC-C	ARC-E	BoolQ	PIQA	Average
MicroLlama V1, ~300M	50B reported	34.30	30.60	51.54	23.29	39.06	53.15	64.58	42.36
MicroLlama V2, ~388M	5.762B	36.29	30.20	51.14	28.41	51.47	58.53	64.09	45.73
MicroLlama V2, ~388M	49.237B	43.89	32.40	53.43	30.12	54.92	60.31	67.46	48.93
Liquid LFM2.5-350M-Base	28T reported	45.15	36.20	55.25	40.96	68.81	58.53	67.19	53.16
SmolLM2-360M	4T reported	56.36	38.20	58.96	38.40	68.01	61.74	71.93	56.23

V2's late mean is 4.22 percentage points behind Liquid, using about 0.176% of its reported training tokens (roughly 1/569). Some task scores are close to Liquid's; ARC remains substantially behind. Different architectures, tokenizers, data and exposure budgets mean this ratio is descriptive, not a controlled efficiency or SOTA claim.

Expanded SmolLM2-style base-model setup

We also evaluated a separate, broader setup following the official SmolLM2-360M base definitions: 2K context, character-normalized likelihood, generated exact-match answers for TriviaQA/GSM8K, five-shot GSM8K and no BoolQ. ARC is the unweighted average of ARC-E and ARC-C. These numbers and means cannot be mixed with the seven-task table above.

Task	V1 (50B reported)	V2 (38.107B)	Liquid (28T reported)	SmolLM2 published reference
HellaSwag	33.47	41.36	44.24	54.50
ARC average	31.18	42.46	54.88	53.00
PIQA	64.58	66.38	67.19	71.70
MMLU cloze	26.51	30.32	34.67	35.80
CommonsenseQA	30.38	32.19	32.27	38.00
TriviaQA	0.00	0.13	15.00	16.90
WinoGrande	50.12	48.22	51.93	52.50
OpenBookQA	27.00	34.00	33.00	37.40
GSM8K (five-shot)	0.45	1.06	23.88	3.20
Derived nine-row mean	29.30	32.90	39.67	40.33

The expanded V2 column is 38.107B, not the final checkpoint. V1/V2/Liquid are our CPU FP32 runs; the SmolLM2 column is its published reference under a different runtime. Liquid uses its required native BOS and a different loader/library version. TriviaQA retains all 17,944 task rows, including repeated question IDs. The V2–Liquid mean gap here is 6.77 pp. The report will provide the exact pins, per-task counts, paired exploratory comparisons and full checkpoint history.

Limitations

This is one training trajectory, not a multi-seed study. Public tasks were repeatedly inspected, and the expanded tasks were outside the initial seven-task exclusion ledger; we do not claim a sealed or fully uncontaminated test set. V2 can produce false facts, repetitive text, unsupported numerical examples and unfinished answers. It is especially weak at generated mathematics and factual QA. Benchmark similarity on selected tasks does not imply reliable chat or reasoning ability.

Full checkpoints and further research

The full restart archive is a separate repository, MicroLlama2-388M-checkpoints. Restart checkpoints are available under checkpoint/, including weights, Adam state, RNG state, source cursors, tokenizer and manifests. There are known historical gaps; not every update is archived.

The immutable final NVIDIA checkpoint is step-00381824 at revision 527ab410. The full restart payload is about 4.66 GB. Root-level files are a modern inference adapter for the same weights, not a new training run. See checkpoint_source.json for the immutable source and changes to inference metadata.

One direction I am interested in is continuation analysis across training: give the same factual prompt to multiple checkpoints, retain the complete continuations, and examine where correct statements, hallucinations, repetition or forgetting appear. Saved examples suggest these behaviors can coexist in one answer; a systematic factuality study remains future work.

Token archives and provenance are also preserved under datasets/tokenized-v1/. Download only the portions needed. Their upstream terms apply. The training code/reproduction guide is linked above; it is initially private while the author reviews the release.

Citation

If you use MicroLlama V2, please cite the model release; the technical report link will be added shortly.

@misc{wang2026microllama2,
  author = {Zixiao Ken Wang},
  title = {MicroLlama V2: A 388M-Parameter Language Model with Full Training Checkpoints},
  year = {2026},
  howpublished = {Hugging Face model repository},
  url = {https://huggingface.co/keeeeenw/MicroLlama2-base}
}

For V1, use Ken Wang. MicroLlama-300M. 2024. See CITATIONS.bib for the author-prescribed references for Qwen3, Engram, TinyLlama, SmolLM2, LiquidAI and the data/evaluation projects used here.

Acknowledgements

Thank you to the MicroLlama community and to the authors of Qwen3, Engram, TinyLlama, LitGPT, FlashAttention, FineWeb-Edu, DCLM-Edu, Cosmopedia, FineMath, SmolLM2, LiquidAI, Transformers, lm-evaluation-harness and LightEval. Their open work makes experiments at this scale possible.