mojolearn

Bitwise-identical machine learning across Apple, NVIDIA and AMD GPUs.

The same machine-learning workload can produce different bits on different GPUs, changing predictions, learned models and subsequent training updates. mojolearn is a GPU machine-learning library whose default identical mode produces bitwise-identical results across verified Apple, NVIDIA and AMD GPUs and x86-64 and Arm CPUs. Results agree bit for bit, not merely within a numerical tolerance.

• Training and inference for transformers and hybrids. Decoder language models, Mamba-1/2/3 state-space models and Samba attention/state-space hybrids train and serve under the same contract, alongside CNNs, recurrent and graph networks.
• 100+ machine-learning algorithms and model components across 13 families, from neural networks to tree ensembles and classical models.
• Serve models trained elsewhere. Neural models trained in other frameworks can be served with identical outputs across vendors.
• Move training between vendors. A training run can be handed off from one GPU vendor to another mid-run, or shared by Apple, NVIDIA and AMD GPUs working on one model at once, and the training state after every step is bitwise identical to a single GPU's.
• GPT-3 Small, trained twice, the same bits. Two 5,000-step runs of a 162-million-parameter GPT-3 Small-shaped decoder, each on 2.62 billion tokens, took different hardware routes across Apple, NVIDIA and AMD GPUs, with checkpoints handed between vendors and one segment trained jointly on an NVIDIA H100 and an AMD MI325X. The runs agree at every step and write the same final checkpoint.
• Full FP32. Neural training and inference run in full FP32 floating-point arithmetic, without reducing the computation to integers.
• Every algorithm runs on the Mac's own GPU, where XGBoost, LightGBM, CatBoost and cuML have no supported GPU path.

To the author's knowledge, mojolearn is the first system to enable these capabilities.

One Mojo codebase implements these algorithms across Metal, CUDA and HIP under a shared numerical contract. Mojo compiles code for different GPUs, but maintaining that contract and implementing the algorithms is original work.

Why it matters

A model may serve predictions on NVIDIA GPUs, run locally on an Apple laptop and later be evaluated on an AMD server. With identical weights and inputs, different GPUs normally return different output bits, so hardware becomes another variable in serving, regression testing and auditing. In a tree learner one rounding difference can change the winning split and every node beneath it. In a training loop it perturbs a gradient, then the optimizer state, then every step after that.

mojolearn removes hardware from that list. A team can test a deployed model on different hardware, replay a recorded decision during an audit, replace serving hardware without numerical change, and resume a checkpoint on another vendor's GPU on the same trajectory as uninterrupted training.

How it works

Floating-point addition is not associative, and GPU vendors differ in reduction order, fused multiply-add contraction, denormal handling, square root, tie-breaking and elementary functions. Recompiling does not fix any of that. mojolearn's numerical contract fixes the operations that determine model bits, including reduction order, rounding and elementary functions, and leaves tiling, thread layout, memory placement and kernel scheduling free, so each backend can be tuned for its own hardware without changing a bit. The kernels are mojolearn's own, written in Mojo, and replace vendor math routines wherever their results would depend on the backend.

Algorithms

The table shows representative methods, with related estimator variants grouped together.

Also provided are evaluation metrics, cross-validation, CPU training and inference, Hugging Face checkpoint loading and decoding, BF16 and INT8 weight storage, tokenized corpora, multi-GPU execution plans and a built-in verifier.

Numeric modes

identical is the default and, outside the tree learners, the only mode. Gradient boosting, random forests and Extra Trees also offer two opt-in modes.

The identical contract is across hardware within one release. mojolearn is alpha: a release may change identical output bits (a faster or more accurate route), provided the new bits agree on every supported vendor and the CPU. Compare runs made with the same release.

Select a mode per estimator with numeric_mode=, per process with mojolearn.set_numeric_mode(...), or before import with MOJOLEARN_NUMERIC_MODE. A configuration that cannot meet its mode's contract raises a named error rather than silently returning something weaker.

Install

pip install mojolearn   # Mac: Apple GPU; Linux: NVIDIA or AMD GPU, or CPU

macOS arm64 wheels include Apple Metal support. In 0.8.25 on Linux x86-64, pip install mojolearn automatically installs the mojolearn-nvidia and mojolearn-amd GPU packages and selects the backend for the GPU it finds. Both platforms also support CPU training and inference.

On NVIDIA, native kernels are used first; this release ships them for the sm_89 architecture. Every other NVIDIA GPU with compute capability 8.0 or newer (A100, H100/H200, Blackwell, ...) uses the PTX build that ships in the same package, which the driver compiles on your machine. It serves every numeric mode, identical included: its results are recorded in the reference table beside the NVIDIA and AMD results and must equal them. See NVIDIA PTX.

mojolearn never substitutes the CPU for a GPU it cannot serve.

mojolearn doctor

reports what the installed wheel supports on this machine.

Verification combines inference and classical training in one routine scope. Neural training is a separate opt-in:

python -m mojolearn verify                    # every applicable routine lane, base fixture
python -m mojolearn verify --quick            # representative routine checks
python -m mojolearn verify --all              # entire routine scope, all nine fixtures
python -m mojolearn verify --neural-training  # neural training/backward/optimizer checks

Routine checks include bundled reference hashes, saved-model inference, and, on a GPU installation, direct GPU/CPU inference comparisons. Neural training is excluded from default, quick and all; combine --neural-training --all to exercise its complete fixture set. Reports name excluded and missing checks.

The advanced --cross-check quick|default|all option runs only direct GPU/CPU comparisons. Its default scope retains the 24-lane diagnostic sample; all has no lane cap and runs every eligible lane on every fixture using sequential child processes. Bare verify uses all eligible cross-check lanes on the base fixture, without that diagnostic sample cap.

Quick start

import numpy as np
import mojolearn

rng = np.random.default_rng(0)
X = rng.random((100_000, 20), dtype=np.float32)
y = (X[:, 0] + X[:, 1] > 1.0).astype(np.float32)

model = mojolearn.GradientBoosting(loss="Logloss", n_estimators=200, max_depth=6)
model.fit(X, y)
print(model.predict_proba(X[:5]))
print(model.numeric_mode_used(), mojolearn.vendor())

Run the same script on an Apple, NVIDIA or AMD GPU and the model and its predictions are the same bits.

Verify it yourself

The wheel ships the reference results recorded on Apple, NVIDIA, AMD and CPU, and a verifier that checks your machine against them.

In 0.8.25, install verification support with python -m pip install "mojolearn[verify]": the same library plus its optional verification dependency. On 0.8.24, use python -m pip install numpy. Optional ANN, CNN and sequence APIs that use NumPy are supported by mojolearn[numpy]; the verification extra includes that same dependency. The base install does not require NumPy. See verification support.

python -m mojolearn verify          # inference and classical training, base fixture
python -m mojolearn verify --quick  # representative checks from the same scope
python -m mojolearn verify --all    # the same scope on all nine fixtures

In 0.8.25, --all also includes applicable gradient, batch-size, ragged-batch and sampler/replay checks. On 0.8.24, add --batch-checks for those probes. --quick keeps its smaller scope.

Each part reports its comparison with the recorded references; missing references read OWED and inapplicable properties are named explicitly. docs/VERIFY.md describes the verifier and docs/VERIFY_EXTERNALLY.md shows how to check the claims from outside the project.

Scope

The guarantee holds for the same code revision, numerical profile, input bytes, hyperparameters and seed, in identical mode, on configurations the verification record covers. It is an FP32 guarantee. Weights trained in another framework reproduce under mojolearn's arithmetic, not that framework's. SUPPORT_MATRIX.md lists the verified devices and configurations.

Runtime rule

No Python compute in the runtime, ever. The Python package is only the API surface (argument checks, routing, calls into compiled Mojo); every computed value comes from Mojo, on the GPU for GPU paths.

Benchmark boards

These are the canonical pages. Update them in place; don't start new ones.

• Full-dataset A/B experiment board: NVIDIA/AMD IDENTICAL and Apple FAST attempts, recorded toggles, A/B times, quality metrics and decisions. Failed attempts and unrun experiments remain visible; combined measurements do not qualify individual switches.
• Apple M3 Ultra board: every lane, with our FAST and IDENTICAL arms next to the opponents. The IDENTICAL cells come from the 2026-10-04 sweep (tools/af_board_ident_update.py); FAST and opponent cells come from the 0.8.34 board.
• Apple FAST refresh: FAST before and after for each lane against the best opponent (tools/af_board_merge.py).
• Apple IDENTICAL refresh: the same view for IDENTICAL (tools/af_board_merge.py --mode identical).
• Board protocol and older boards.

Documentation

• Support matrix
• Verification
• Changelog
• Roadmap
• Getting started as a contributor
• Contributing and governance

Trademarks and affiliation

mojolearn is an independent project by Andrew Hendel, licensed under Apache-2.0. It is not affiliated with, sponsored by, or endorsed by Modular, Inc. MAX® and Mojo® are trademarks of Modular, Inc. Binary wheels include unmodified Modular runtime components redistributed under Modular's own license; see NOTICE.

Citation

To cite mojolearn, use CITATION.cff. The concept DOI is 10.5281/zenodo.22068632.