mojolearn

mojolearn

Reproduce the model. Reproduce the prediction.

mojolearn makes machine-learning results reproducible down to the bits, across GPU vendors. One Mojo codebase runs on Apple, NVIDIA and AMD GPUs and covers over 100 algorithms and model components in 13 families, from tree ensembles to transformers and hybrid attention and state space models.

Open source, Apache 2.0. Created by Andrew Hendel. The identity contract covers supported NVIDIA and AMD workloads within one release. Scope.

Try it Paper (PDF) Measurements GitHub Model and evidence

A language model trained twice, with the same bits

A 162-million-parameter language model was trained twice from the same initialization, on the same 2.62 billion FineWeb-Edu tokens per run. The two routes used different hardware at every segment and together spanned NVIDIA, AMD and Apple. One segment trained jointly on an NVIDIA H100 and an AMD MI325X in separate machines for 400 steps.

The runs matched at all 5,000 optimizer steps in training state, summed gradients, all 64 per-shard losses and the learning rate. They wrote identical final checkpoint files, one on an Apple M3 Ultra and the other on an NVIDIA L40S.

Route A · Apple M3 Ultra4129921e99ed404db1ca4b7b94ebc43c0bee18e7711deb326b31ec89e0b402c8
Route B · NVIDIA L40S4129921e99ed404db1ca4b7b94ebc43c0bee18e7711deb326b31ec89e0b402c8

The training experiment

Two hardware routes. The same bits at every step.

162,147,840 parameters and 2.62 billion tokens per route. Each route resumes its checkpoint on the next segment’s hardware.

  1. Segment 1

    Steps 1 to 1,000

    Route A

    2 × NVIDIA H100

    1,000 steps
    equal bits

    Route B

    1 × AMD MI325X

  2. Segment 2

    Steps 1,001 to 2,000

    Route A

    2 × NVIDIA H100

    1,000 steps
    equal bits

    Route B

    2 × AMD MI300X

  3. Segment 3

    Steps 2,001 to 2,400

    Route A

    1 × NVIDIA H100 + 1 × AMD MI325X

    Joint training, separate machines

    400 steps
    equal bits

    Route B

    1 × NVIDIA H100

  4. Segment 4

    Steps 2,401 to 3,900

    Route A

    2 × AMD MI300X

    1,500 steps
    equal bits

    Route B

    2 × NVIDIA H100

  5. Segment 5

    Steps 3,901 to 4,900

    Route A

    2 × NVIDIA H100

    1,000 steps
    equal bits

    Route B

    2 × AMD MI300X

  6. Segment 6

    Steps 4,901 to 5,000

    Route A

    1 × Apple M3 Ultra

    100 steps
    equal bits

    Route B

    1 × NVIDIA L40S

Same initialization. 5,000 matching updates. Identical final checkpoint files.

Show the routes as a table
Hardware and matching steps for every segment of the two training routes
SegmentStepsRoute ARoute BMatching steps
11 to 1,0002 × NVIDIA H1001 × AMD MI325X1,000
21,001 to 2,0002 × NVIDIA H1002 × AMD MI300X1,000
32,001 to 2,4001 × NVIDIA H100 + 1 × AMD MI325X (joint training)1 × NVIDIA H100400
42,401 to 3,9002 × AMD MI300X2 × NVIDIA H1001,500
53,901 to 4,9002 × NVIDIA H1002 × AMD MI300X1,000
64,901 to 5,0001 × Apple M3 Ultra1 × NVIDIA L40S100
Adapted from the paper’s training-route figure. Segment widths do not represent steps or elapsed time. Source Released model, checkpoints and evidence.

The model has GPT-3 Small’s depth and width with modern Llama-style blocks, and uses FP32 throughout. The model card has its architecture, quality evaluations and limitations, the weights and both routes’ evidence.

What exact agreement makes possible

A small rounding difference can change a tree’s winning split or a training update, and every later step inherits it. With the bits fixed, one question gets a clean answer: did the model change, or did the hardware?

Start with an open-source model. Reproduce the fine-tuning.

A bank or a healthcare research team adapts a pretrained open-source model to its own task. Given the same imported weights, data order, optimizer state, settings and randomness, a supported fine-tuning run in IDENTICAL mode produces the same adapted model on NVIDIA and on AMD. An independent team can repeat the adaptation on its own hardware and compare the result exactly, without repeating pretraining.

Replay a regulated model’s output

A risk model’s prediction is reviewed months later, or a clinical group reproduces an output at another institution. Keep the model, inputs, preprocessing and release, and the reviewer has an exact computation to repeat on whatever supported GPU they have.

Run a model trained elsewhere on any GPU

Load a model trained in another framework, develop on a MacBook and serve on NVIDIA or AMD. SmolLM2-360M logits have matched on an M2 Pro, RTX 4090, H100, MI325X and CPU (release history).

Resume training on whatever GPUs are available

NVIDIA capacity today, AMD tomorrow. A checkpoint with optimizer state resumes on the other vendor and follows the same numerical trajectory, as the two-route experiment above did at every handoff.

Check a fleet before a rollout

Keep a set of inputs and expected output bits, and run them across every NVIDIA and AMD server before deploying. Differences too small to flip a label, but large enough to matter downstream, become visible.

Try it on your hardware

Install verification support, check your device, then run representative checks against recorded references.

python -m pip install "mojolearn[verify]"
mojolearn doctor
python -m mojolearn verify --quick

Quick verification covers representative inference and classical training checks; neural training has a separate opt-in check. The verification guideexplains the scopes, and the support matrixlists the GPU wheels in each release. See the model example.

One source can be portable without producing the same bits

Mojo provides compilation across GPU backends. Getting their arithmetic to agree requires more work. Floating-point addition depends on order, and instruction selection and math libraries can change the final bits.

  1. Fix the order of arithmetic

    Adding the same floating-point values in a different order can change the answer. Reduction trees and accumulation rules keep that order stable across devices.

  2. Control rounding and elementary functions

    Fused operations, square root, exp and log can take different paths on different backends. The library uses explicit arithmetic rules and its own portable routines where needed.

  3. Keep hardware scheduling separate

    Thread layout, memory placement and kernel scheduling can vary while preserving the operations that determine the output bits. Mojo provides portability; mojolearn implements the numerical contract.

The numerical-pathway ledgerexplains the implementation choices. Neural training and inference use full FP32 arithmetic. Some other algorithms, including tree histograms, use integer accumulation where appropriate.

Measured runtime

Scroll the table to compare both GPUs →

WorkloadNVIDIA L40SAMD MI325X
mojolearnOpponentmojolearnOpponent
Language-model training step8 layers, 2,048 tokens35.4 ms21.6 msPyTorchcompiled FP3242.8 ms15.2 msPyTorchcompiled FP32
Transformer block forward2,048 tokens3.99 ms3.59 msPyTorcheager FP323.23 ms1.58 msPyTorchcompiled FP32
Gradient boosting fit †Istella, 2.04M rows6.14 s8.18 sCatBoostGPU6.00 s44.4 sCatBoostEPYC CPU
PCA fitIstella, 2.04M rows102 ms93.5 msPyTorchGPU + upload51.8 ms207 msPyTorchGPU + upload

mojolearn in IDENTICAL mode; neural comparisons in FP32. One timed sample after warmup per cell. October 11, 2026 snapshot of unreleased builds. Neural rows time operations, not full training runs.

Neural opponents use the faster recorded PyTorch FP32 arm, eager or compiled. PCA opponent totals add the separately measured upload. † CatBoost on the AMD system ran on its EPYC CPU. Boosting leaf counts differ or are unverified, so those fit times are contextual comparisons.

Comparable recorded quality for language-model loss, boosting AUC and PCA explained variance. Transformer quality is unreported in this snapshot.

Quality and comparison details
Language-model loss
Final-step loss is 8.41845 for mojolearn and both selected PyTorch arms, rounded to five decimal places. This is a two-step benchmark, not a full training-quality evaluation.
Transformer output
The selected timing cells do not record an opponent quality metric in this snapshot; equivalence is not established here.
Gradient boosting
AUC is 0.980129 for mojolearn, 0.980018 for CatBoost on NVIDIA and 0.979899 for CatBoost on the AMD host CPU. Log loss is 0.186686, 0.187292 and 0.188093, respectively. The board marks NVIDIA fit comparability unknown and AMD fit work not comparable because leaf counts differ.
PCA explained variance
The explained-variance ratio sums to approximately 1.000000 for mojolearn and both PyTorch opponents. Matching this metric does not establish identical components.

Quality agreement does not mean bitwise agreement with opponents. Opponent results are retained measurements from earlier runs.

Full benchmarks, comparisons and per-cell identity status NVIDIA · AMD. Every race with its opponents.

The contract, and what it costs

Compare the same release, input bytes, model state, settings and supported execution configuration. The release policy requires bitwise agreement between NVIDIA and AMD GPUs. Apple GPU and CPU execution are also available, and the paper records broader agreement for the configurations it tested. Results may change between releases. See the release notes and support matrix.

Identity has a cost: a fixed arithmetic order and portable numerical routines give up some throughput that vendor-specific implementations use. The measurements record timings, quality and scope for complete implementations, so a ratio does not isolate the cost of identity alone. Agreement is between mojolearn executions; it does not establish accuracy, fairness or fitness for a particular use, which need their own evaluation.

mojolearn is an alpha. Reproductions, bug reports and workloads where exact agreement would be useful are welcome.

From the paper

Bitwise-Identical Machine Learning Across Metal, CUDA, and HIP

Read the full paper (PDF)October 9, 2026 preprint · 28 pages

The abstract below describes the configurations verified in the paper, including its Apple and CPU results. The current release contract is described above.

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. It supports training and inference for over 100 machine-learning algorithms in 13 families, from decision trees to neural networks. Neural models trained in other frameworks can be served with identical outputs across vendors, and training workloads can be handed off between GPU vendors mid-run, or shared by GPUs of different vendors at once, while maintaining bitwise identity. Two training runs of a 162-million-parameter GPT-3 Small-shaped model follow different hardware routes collectively spanning Apple, NVIDIA and AMD GPUs and remain bitwise identical at all 5,000 optimizer steps, each processing 2.62 billion tokens. Neural training and inference run in full FP32 floating-point arithmetic, without reducing the computation to integers. To my knowledge, I am the first to enable these capabilities.

One Mojo codebase implements these algorithms across Metal, CUDA and HIP under a shared numerical contract. Mojo compiles the code for different GPUs, but the numerical contract, introduced in this work, fixes arithmetic order, rounding and elementary functions while allowing backend-specific scheduling.

PyPIArchived software releasesModel and experiment evidence