Mojolearn 162M: one model, two hardware routes

DOI

A 162M-parameter decoder trained twice from the same initialization on the same 2.62 billion FineWeb-Edu tokens, once along each of two different routes through NVIDIA, AMD and Apple GPUs. The two runs are bitwise identical at every one of the 5,000 optimizer steps: training state, summed gradients, all 64 per-shard losses and the learning rate. Their final checkpoint files have the same SHA-256. This repository distributes that one final model with the evidence from both routes.

The model is GPT-3 Small-sized and trained with GPT-3 Small's hyperparameters, but its blocks are modern Llama-style blocks; see the comparison below. It is a research artifact for reproducible training, not a production language model.

Compared with GPT-3 Small

GPT-3 Small (Brown et al., 2020) This model
Layers / width / heads × head dim 12 / 768 / 12 × 64 same
Context 2,048 same
Batch 0.5M tokens 524,288 tokens (64 shards × 4 × 2,048)
Peak learning rate 6.0e-4 same
Optimizer Adam β=(0.9, 0.95), weight decay 0.1 AdamW, same β and decay, eps 1e-8
Schedule warmup, cosine decay to 10% 250-step warmup, cosine to 10%
Positions learned table rotary (θ = 10,000)
Norm LayerNorm RMSNorm (eps 1e-6), pre-norm, no final norm
MLP GELU, 3,072 wide SwiGLU, 2,048 wide (same parameter count)
Biases yes none
Input / output embeddings tied untied
Parameters 125M 162,147,840 (the extra 38.6M are the untied input table)
Tokenizer GPT-2 BPE Mojolearn BPE, 50,257 ids, trained on FineWeb-Edu
Training data 300B tokens, GPT-3 mix 2.62B tokens, FineWeb-Edu
Arithmetic mixed precision FP32 throughout

Per token, the matrix-multiply work is the same as GPT-3 Small's: SwiGLU's three 768 × 2,048 matrices hold exactly as many weights as GELU's two 768 × 3,072 ones, attention is identical, and both models multiply by a 50,257 × 768 output matrix. The untied input table is a lookup with no arithmetic, but it adds 38.6M parameters to every optimizer update and to every checkpoint whose bits must agree.

The token budget is Chinchilla-scale: about 21 tokens per parameter of a 125M-parameter GPT-3 Small, or about 16 per parameter of this 162M model. It is about 1/114 of GPT-3 Small's 300B tokens, so this model is not expected to match GPT-3 Small's quality, and the comparison below is for orientation only.

Results

All results are from the single distributed checkpoint (step 5,000).

Evaluation Result
Held-out FineWeb-Edu loss 3.265 nats/token (perplexity 26.2)
Held-out bits per UTF-8 byte 0.976
HellaSwag, 0-shot, validation (10,042) acc 28.0%, acc_norm 29.6%
Final training loss (step 5,000, mean of 64 shards) 3.140 nats/token
  • Held-out loss is computed by Mojolearn's own CPU forward pass (LanguageModelInference.loss) on 128 windows of 2,049 tokens (262,144 scored tokens), one every 721,773 tokens from the start of the corpus's held-out shard (FineWeb-Edu sample-10BT file 013_00000.parquet, token range 2,926,502,182–3,110,556,447, which no training step read); the windows span its first half, tokens 2,926,502,182–3,018,169,402. Bits per byte divide the summed loss by the UTF-8 length of the scored tokens and do not depend on the tokenizer. tools/fetch_heldout.py (which reads the token stream from the project's R2 store; it is not public) and tools/heldout_loss.py reproduce it; details in results/heldout-loss.json.

  • HellaSwag uses the EleutherAI lm-evaluation-harness formulation (acc_norm normalizes by the choice's UTF-8 length) through reference_torch.py. GPT-3 Small reported 33.7% zero-shot after 300B tokens. Details: results/hellaswag.json.

  • Training loss comes from the chain records both routes share. It is not a held-out metric.

  • Generations (greedy, 48 new tokens, Mojolearn CPU forward pass) are in results/generations.json. They are fluent and on topic, repeat themselves as greedy decoding at this scale does, and state facts unreliably. For example:

    Photosynthesis is the process by which plants convert sunlight into energy. The process of photosynthesis is the process by which plants convert sunlight into energy. …

Using the model

With Mojolearn (the training library)

pip install mojolearn tokenizers numpy
import numpy as np
from tokenizers import Tokenizer
from mojolearn import ByteLanguageModelConfig, LanguageModelInference
from mojolearn import _byte_lm_checkpoint as ck

state = ck.load("checkpoint.blm")            # verifies every array digest
shape = ByteLanguageModelConfig(**{**state["model_shape"], "batch": 1, "length": 128})
model = LanguageModelInference(state["parameters"], shape=shape)
tok = Tokenizer.from_file("tokenizer.json")

ids = tok.encode("The water cycle describes how water").ids
for _ in range(24):
    ids.append(int(model.next_bytes(np.array([ids[-128:]], dtype=np.int32))[0]))
print(tok.decode(ids))

LanguageModelInference runs Mojolearn's FP32 forward pass on the CPU, with logits, loss and greedy next_bytes. It was tested with mojolearn 0.8.35 from PyPI on macOS arm64 (tools/generate_mojolearn.py, results/generations.json). LanguageModelTrainer.from_checkpoint_binary("checkpoint.blm") restores the full training state, AdamW moments included, for continued training (not exercised in this package).

Anywhere else: safetensors + PyTorch

model.safetensors holds the same weights as checkpoint.blm's parameter array, byte for byte (tools/export_safetensors.py writes and checks it), under Mojolearn's tensor names. Linear weights are stored [out_features, in_features]. reference_torch.py is a plain PyTorch implementation of the forward pass:

pip install torch safetensors tokenizers
python3 reference_torch.py "The water cycle describes how water" --tokens 40

On held-out tokens its logits agree with Mojolearn's to within 3.8e-05 (largest absolute difference) and it picks the same greedy token at every position checked (tools/compare_reference.py, results/reference-agreement.json). It is a convenience, not part of the bitwise claim: PyTorch sums in its own order.

The weights do not load into Transformers' LlamaForCausalLM unchanged, because this model has no final norm before the output head.

Tokenizer

tokenizer.json is a Hugging Face tokenizers file (byte-level BPE with GPT-2's pre-tokenizer pattern). Its vocabulary is identical to ranks.tsv, the file whose SHA-256 the training recipe pins. Do not substitute the GPT-2 or GPT-3 tokenizer. Training documents were concatenated without <|endoftext|> separators, so the model never learned to emit id 50,256 and will not stop on its own.

Training

Initialization NumPy default_rng(93261), N(0, 0.02); norm weights 1 + N(0, 0.02). Both routes start from the same step-0 file (SHA-256 2b03554d…)
Steps 5,000 per route
Tokens 524,288 per step; 2,621,440,000 per route
Gradient 64 logical shards, each a mean over 4 × 2,048 tokens; shard gradients summed in a fixed order
Optimizer AdamW, β = (0.9, 0.95), eps 1e-8, weight decay 0.1, no gradient clipping
Learning rate 250-step linear warmup to 6e-4, cosine decay to 6e-5; an exact FP32 table of all 5,000 values is in the recipe
Arithmetic FP32 everywhere; Mojolearn IDENTICAL mode
Data FineWeb-Edu sample-10BT files 000–003 for training (2,926,502,182 tokens); file 013 held out

The complete recipe is supplement/gpt3-six-segment-2026-09-28/recipe.json.

Two routes, one checkpoint

Steps Route A Route B
1–1,000 2 × NVIDIA H100 1 × AMD MI325X
1,001–2,000 2 × NVIDIA H100 2 × AMD MI300X
2,001–2,400 H100 + MI325X jointly, in separate machines 1 × NVIDIA H100
2,401–3,900 2 × AMD MI300X 2 × NVIDIA H100
3,901–4,900 2 × NVIDIA H100 2 × AMD MI300X
4,901–5,000 Apple M3 Ultra NVIDIA L40S

Each route continues from its own checkpoints. The segments also ran different Mojolearn releases: route A used 0.8.15, 0.8.15, 0.8.17, 0.8.18, 0.8.19 and 0.8.24; route B used 0.8.18 and then 0.8.22. The bits still agree at every step. All 56 common saved checkpoints match. The final checkpoint SHA-256, recorded by both routes and re-hashed from both stored copies before packaging, is:

4129921e99ed404db1ca4b7b94ebc43c0bee18e7711deb326b31ec89e0b402c8

checkpoint.blm is that original file: parameters, both AdamW moment arrays, initialization flags and the training metadata (format in ARTIFACTS.md).

Verifying

python3 tools/verify_release.py        # Python 3.11+; file hashes, checkpoint header and array digests
python3 tools/check_gpt3_evidence.py   # both routes' training records

Neither needs a GPU or a network connection. The evidence checker confirms that the retained records are complete and agree; it does not recompute training. Its exact scope is in the supplement README. Matching bits establish reproducibility, not model quality; the quality results are reported separately above.

Limitations

  • A small research model trained on 2.6B tokens: it is not instruction tuned, it hallucinates, repeats itself under greedy decoding and does not stop on its own.
  • English educational web text only (FineWeb-Edu); it inherits that corpus's biases.
  • The Mojolearn forward pass used here runs on the CPU in FP32 and is slow; it is a reference path, not a serving stack.

License, data and citation

The weights, tokenizer and code in this repository are released under the Apache License 2.0, the license of Mojolearn. The training data is FineWeb-Edu (HuggingFaceFW/fineweb-edu), released under ODC-By 1.0; its terms and the terms of the underlying Common Crawl data apply to the data, not to this repository's files.

Source: github.com/mojolearn/mojolearn, where this model card, its tools and evidence live in release/gpt3-small-162m/.

Cite this model with its DOI, 10.57967/hf/10832:

@misc{hendel2026mojolearn162m,
  author    = {Hendel, Andrew},
  title     = {Mojolearn 162M: One Model, Two Hardware Routes},
  year      = {2026},
  publisher = {Hugging Face},
  doi       = {10.57967/hf/10832},
  url       = {https://huggingface.co/mojolearn-ai/mojolearn-162m}
}

Cite the Mojolearn library with its concept DOI 10.5281/zenodo.22068632.

Downloads last month
33
Safetensors
Model size
0.2B params
Tensor type
F32
·
Inference Providers NEW
This model isn't deployed by any Inference Provider. 🙋 Ask for provider support

Dataset used to train mojolearn-ai/mojolearn-162m