MicroLlama V2

Thank you for the support for MicroLlama V1! V2 continues the same experiment: seeing how far I can take a small language model with consumer hardware and a modest budget.

MicroLlama V2 is a 388M-parameter base model, trained from scratch with a Qwen3-based backbone and Engram-style memory. The checkpoint available here was trained on about 56.6B tokens.

I budgeted $500 for compute and have spent $880 so far.

October 30, 2026 update

This release adds about 6.6B tokens of continued training on Apple M5 Ultra through the Mac pipeline. Here’s how the 56.6B checkpoint compares with earlier runs and other small models. Scores are percentages.

The current release averages 51.35%, up from 48.93% at 49.2B, and is 1.81 percentage points behind Liquid. Training is still running on the M5 Ultra, with more checkpoint updates to come. Stay tuned!

¹ TinyLlama’s row comes from its published 3T results. Those seven scores average 52.97%; the published average is 52.99%. Liquid uses its native BOS token. The models differ in size, training data, token budgets and context; these comparisons show where V2 stands, rather than measuring training efficiency.

Model details

V1 followed the TinyLlama lineage. For V2, I used a randomly initialized Qwen3-based architecture and added a lightweight memory module inspired by Engram. The model was trained from scratch; it does not use pretrained Qwen weights.

• Developer: Zixiao Ken Wang (keeeeenw)
• Funding: self-funded
• Model: autoregressive base model, about 388M parameters
• Backbone / memory: about 338M / 50M parameters
• Language: primarily English
• Context: 8K model limit; the Mac continuation uses 2K training sequences
• Current checkpoint: step 432,299, about 56.6B training tokens, released October 30, 2026
• License: Apache 2.0 for the model and inference code; upstream dataset terms also apply

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 Engram-inspired module uses normalized bigram and trigram lookups with a learned gate. It omits Engram’s convolution, multi-branch integration and memory offloading.

V2 uses the LiquidAI/LFM2.5-230M-Base tokenizer, with 64,402 token IDs. The model has a padded vocabulary of 65,536 rows; unused rows are suppressed during generation. Inputs are plain text, with no chat template or automatic BOS token.

Training data and hardware

I trained the model for its first 50B tokens in PyTorch on rented NVIDIA RTX 4090/5090 GPUs, including a four-RTX-5090 run on Vast.ai:

1. First 10B tokens: FineWeb-Edu.
2. 10B to 23.8B: 60% FineWeb-Edu and 40% DCLM-Edu.
3. 23.8B to 50B: 48% FineWeb-Edu, 32% DCLM-Edu, 15% Cosmopedia v2 and 5% FineMath 4+.

The released Mac continuation adds about 6.6B tokens from the pinned Dolmino mixture, using 2K training sequences on Apple M5 Ultra. Source manifests and checks against previously used training rows are retained locally.

The recipe includes existing public synthetic data, including Cosmopedia. I did not generate a new synthetic corpus for this run.

Project links

• Inference model: MicroLlama2-base
• Original NVIDIA checkpoints and token archives: MicroLlama2-388M-checkpoints
• Inference code: configuration_microllama2.py, modeling_microllama2.py, and infer.py
• Training code: MicroLlama2-training
• Technical report: in preparation

Using the model

MicroLlama V2 is a base text-completion model. It can be used for continued pretraining, fine-tuning and small-model research. It has not been tuned as a chat assistant.

1. Install dependencies

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

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

The inference adapter was tested with Transformers 5.20.2 and PyTorch 2.9.1 on CPU in FP32. CUDA and MPS are also supported; hardware and precision can affect outputs. Loading the custom architecture requires trust_remote_code=True, so review the model code before running it.

2. Generate text

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. Greedy decoding and sampling carry both the attention cache and Engram history; beam search, speculative decoding and static caches are not supported. On MPS, the integer memory hash runs on CPU and floating-point model computation runs on GPU.

A normal from_pretrained call downloads about 1.55 GB of FP32 weights, plus the tokenizer, configuration and model code. The weights match the saved Mac checkpoint exactly. Training state and datasets are kept in separate archives; avoid downloading an entire archive unless you need it.

Evaluation

Seven-task comparison

The table in the October 30 update follows the original V1 task setup, based on TinyLlama’s GPT4All evaluation guide. Our saved evaluations use a pinned lm-evaluation-harness, zero-shot prompts and CPU FP32. The evaluation limit is 8K tokens, including for checkpoints trained with 2K sequences. We report token-normalized accuracy (acc_norm) except ordinary accuracy (acc) for WinoGrande and BoolQ. The average gives each of the seven tasks equal weight.

The 56.6B row matches the weights in this repository. All seven tasks finished and both independent evaluation audits passed. The exported inference model was checked against the saved training checkpoint, including its weights, outputs and cached generation.

The 49.2B row shows the late NVIDIA checkpoint. The original final 50B NVIDIA checkpoint has no full seven-task result. TinyLlama is a published reference; the other rows use our saved CPU evaluations.

At 5.8B tokens, V2 was already 3.37 points above V1’s reported 50B average. More recent results fluctuate between checkpoints, so a small change in the mean should be read alongside the individual tasks.

Earlier expanded evaluation (38.1B)

We also ran a broader evaluation on an earlier V2 checkpoint, following the SmolLM2-360M base definitions. It uses 2K context, character-normalized likelihood, generated exact-match answers for TriviaQA and GSM8K, five-shot GSM8K, and no BoolQ. ARC is the average of ARC-E and ARC-C. This is a separate protocol, so its scores and mean should not be mixed with the seven-task comparison.

The V2 results here are from 38.1B tokens. V1, V2 and Liquid were evaluated on CPU in FP32; SmolLM2 is the published reference. Liquid uses its native BOS token and a different loader version. TriviaQA includes all task rows, including repeated question IDs. V2 trails Liquid by 6.77 points on the nine-row average.

Limitations

These results have not been repeated across multiple seeds. The public benchmarks have been checked throughout training. The expanded tasks were outside the initial seven-task exclusion checks, so benchmark contamination has not been ruled out.

V2 can invent facts, repeat itself and leave answers unfinished. Generated mathematics and factual question answering remain weak. These benchmark scores alone do not establish reliable chat or reasoning ability.

Checkpoints and further work

The original NVIDIA run is preserved in MicroLlama2-388M-checkpoints. Its final 50B checkpoint includes model weights, Adam state, RNG state and data-reader state, about 4.66 GB in total. Some intermediate checkpoints are missing from the archive.

This repository serves the October 30 Mac inference checkpoint. Its full training state is retained locally. The previous NVIDIA inference release remains available at revision e78d8650, and checkpoint_source.json records the current checkpoint’s origin and weight hash.

The original token archives are available under datasets/tokenized-v1. Download only what you need; the upstream dataset terms apply.

I’d also like to compare factual completions across checkpoints: where the model improves, where it forgets, and where it confidently gets things wrong. That study is still to come.

Citation

If you use MicroLlama V2, please cite the model release. I’ll add the technical report when it is ready.

@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.