prime: training a language model across unreliable machines over the internet

The problem is training on machines that are not in the same datacenter

prime is a framework for efficient, globally distributed training of AI models over the internet. The name changed: the project was previously called ZeroBand, and the packaging still carries that name, with the project metadata declaring the package zeroband at version 0.1.0 and describing ZeroBand as a production ready codebase for decentralized training of an LLM.

The distributed training method it implements is DiLoCo, and the reference experiment is INTELLECT-1, a 10B parameter model trained across nodes that are not in one datacenter. A research paper covering both the framework and that experiment is linked from the documentation.

The constraint that shapes every design decision is that the participants are remote machines rented from different providers. They are not under a single administrator, they are not on the same fast interconnect, and any of them can vanish. That is a different problem from scaling inside a datacenter, and it is the problem the framework is built around.

The audience is therefore research and infrastructure engineers who have access to many smaller GPU boxes and want to pool them into one training run, plus people who want to reproduce the INTELLECT-1 setup or build on top of its components.

ElasticDeviceMesh resizes the process group instead of crashing the run

The central abstraction is ElasticDeviceMesh, a distributed group that wraps two kinds of communication. It handles dynamic global process groups for fault-tolerant communication across the internet, and local process groups for communication within a node or a datacenter. Keeping those two scopes separate is what makes the global group cheap enough to churn.

The comparison the documentation draws is with the standard DeviceMesh in torch distributed, which crashes and needs a cold restart when the process group has to be resized. ElasticDeviceMesh instead manages the resizing itself when nodes join or leave.

The mechanism for knowing when to resize is a heartbeat. Dead nodes are discovered through the heartbeat and removed from the process group. There is one refinement on top of that: a node that is crashing attempts a best effort deathrattle to fail its own heartbeat quickly, which saves its comrades the wait for a timeout. In other words, a node that knows it is dying says so, and the rest of the group does not have to wait for the silence to become suspicious.

That combination is the piece worth understanding before anything else in the codebase. Without it, the only responses to a machine failure are a full restart or tolerating a stalled collective. With it, a run over the internet can absorb a machine loss and continue.

Checkpointing goes through /dev/shm so the training loop never blocks

Checkpointing is expensive at the model sizes involved, and the documentation reports it taking up to 20 minutes on the nodes tested. Blocking the training process for that long would waste compute, so the checkpoint path is staged.

The first stage writes into /dev/shm, a RAM backed filesystem. That operation is much faster, and once the checkpoint exists there the main training process is unblocked. The second stage uses two subprocesses to copy the checkpoint out of /dev/shm into the checkpoint directory on disk and to upload it to the remote, both asynchronously.

So at any moment there are potentially three copies of a full checkpoint in play: one in RAM, one being written to disk, and one in flight to remote storage. The RAM copy is also what makes the next feature possible.

Live checkpoint recovery handles the case of a node joining a run in progress. A joining node needs the most recent model and optimiser state before it can contribute, and it must get that within the window between two outer steps or the checkpoint it receives is stale. To make that fast, joining nodes request checkpoints from their peers, and each peer hosts a sidecar HTTP server serving the latest checkpoint straight out of /dev/shm.

The joining node then skips the inner steps entirely and joins the outer step with zero pseudo-gradients. The reasoning is about not stalling the others: if the newcomer also performed the inner steps, it would arrive late to the outer step by the time it downloaded and loaded the checkpoint, which would reduce the cluster's compute utilisation. Contributing zeros is a deliberate way to join without costing the existing nodes time.

A hand written int8 all-reduce kernel because the collectives do not support it

The communication path is the other place where the framework departs from stock PyTorch. In the experiments described, int8 quantization was applied to the pseudo gradients without any impact on the loss curves, which reduces the payload of each outer step all-reduce by 4x compared to fp32.

The catch is in the reduction itself. The accumulate step has to happen in fp32, meaning intermediate results are dequantized and re-quantized during the all-reduce, and no collective communication library supports that pattern. So the project implements a fully pipelined ring-reduce kernel in C++, JIT compiled as a custom operator using the torch library.

A second optimisation followed. Doing the quantization and dequantization with torch ops such as quantize_per_tensor, scatter_add and index was too slow and left the target network bandwidth of 4 Gbps underused. The response was custom multithreaded uint8 ops in C++, which improved quantization speed by more than 60x.

Bandwidth is also treated as a tuning problem rather than a given. Sharding the DiLoCo pseudo gradients within a node lets several connections open at once during the all-reduce, which the documentation reports as an 8x transfer speed improvement on some nodes. Separately, relying on the public IP forward gave poor or unstable peer-to-peer bandwidth on some compute providers, so the framework employs VPN technology to alter packet routing between nodes and use the available internet bandwidth better. In similar datacenter settings the documentation reports bandwidth utilization improvements of up to 40x over the earlier OpenDiLoCo release, reaching up to 4Gb/s connections between datacenters across the United States.

Memory sharding and CPU offload make a 10B model fit the available GPUs

Two more pieces of the implementation sit behind the training loop. The first is a PyTorch FSDP2 with DTensor ZeRO-3 implementation, used to fit 10B model training inside the available memory. It uses the fully_shard API from FSDP2, which wraps model parameters as DTensors and registers hooks to schedule all-gather and reduce-scatter when the tensors are used.

The optimisation detail worth knowing is bucketing. FSDP2 groups parameters into FSDPParamGroups so collectives run on larger tensors, which improves the protocol-to-payload ratio and improves overlap from pipelining. The same trick is applied to the pseudo gradients, bucketing them by layer.

The second piece is CPU offloading, and the claim attached to it is unusually specific: the DiLoCo optimizer adds no GPU overhead, because all the tensors it requires are offloaded to CPU memory. The justification is arithmetic. A global sync happens only every hundreds of steps, so the slower speed of copying and calculating the pseudo gradient on the CPU is negligible next to the time taken by the inner steps and the all-reduce.

That argument is the reason the optimizer can live off the GPU at all. If the global sync were frequent, offloading would be a bottleneck instead of a free win. The dependency list shows the implementation side: torch-shampoo is pulled from a git reference, along with liger-kernel-nightly, torchdata, transformers, datasets, pyarrow, fsspec with GCS support, zstandard, ninja, toposolve and psutil. Python 3.10 or newer is required.

Setting up locally means uv, a git submodule and a Hugging Face login

The install path is documented in five steps, and the first one is a convenience script.

curl -sSL https://raw.githubusercontent.com/PrimeIntellect-ai/prime/main/scripts/install/install.sh | bash

The step by step version starts with cloning the repository, then installing uv and sourcing its environment.

git clone [email protected]:PrimeIntellect-ai/prime.git

curl -LsSf https://astral.sh/uv/install.sh | sh
source $HOME/.local/bin/env

Then the environment is created, dependencies synced with the all extra, and the submodules initialised. The iperf package is installed because it is the tool used to measure the bandwidth this framework is trying to maximise.

sudo apt install iperf -y
uv venv
source .venv/bin/activate
uv sync --extra all
git submodule update --init --recursive

Authentication to Hugging Face follows, since the dataset step downloads from there.

huggingface-cli login

The data download itself writes a directory for the dataset, runs the subset script against the PrimeIntellect/fineweb-edu dataset with a single data worker and 32 shards, and then moves the result into place.

mkdir -p datasets
uv run python scripts/subset_data.py --dataset_name PrimeIntellect/fineweb-edu --data_world_size 1 --data_rank 0 --max_shards 32
mv fineweb-edu/ datasets/fineweb-edu/

The Quick Check is the smoke test. It launches two local processes with torchrun under a debug configuration file, and it sets the global address, rank, unique id, world size and port as environment variables, which is how the mesh learns the topology.

GLOO_SOCKET_IFNAME=lo GLOBAL_ADDR=localhost GLOBAL_RANK=0 GLOBAL_UNIQUE_ID=0 GLOBAL_WORLD_SIZE=1 GLOBAL_PORT=8989  uv run torchrun --nproc_per_node=2 src/zeroband/train.py  @configs/debug/diloco.toml

The optional all extra installs wandb and lm-eval, so experiment tracking and evaluation are available but not required.

Where prime is the wrong tool

The first wrong-tool case is a team with one good datacenter and no interest in the internet. If your nodes sit on one fast interconnect, a standard torch distributed setup has lower overhead, fewer moving parts and a much better documented failure model than a framework whose central feature is surviving remote node loss.

The second case is anyone who needs the research prototype to be a supported platform. The published tags are v0.5, v0.3 and v0.2, all from October 2024, and the last push to the repository was on 2026-07-17. The documentation also states plainly that the project was renamed from ZeroBand to prime, and the package metadata still uses the old name and a version of 0.1.0, which tells you the distribution metadata and the release tags have not been reconciled. Plan for the API to have moved, and expect to read the code rather than a versioned changelog.

The third case is a run where the nodes cannot reach each other well. The framework's bandwidth work assumes peer-to-peer paths that can be improved with VPN routing. If your compute providers only expose unstable public IP paths and you cannot change that, the all-reduce becomes the bottleneck no matter how the kernel performs.

On licensing, prime is Apache-2.0, which is a permissive grant with an explicit patent clause. That is a low-friction licence for both research and commercial use. The dependencies deserve their own look, though: torch-shampoo is pinned to a git reference on a main branch rather than to a released version, and pydantic_config is likewise referenced by commit hash, so a rebuild months from now can pick up different code than the run you validated today.

Editorial conclusion

prime suits a research team that has spare GPUs in several places and cannot justify a single large cluster, and who wants to reproduce or extend the INTELLECT-1 10B setup described in the paper. It does not suit a team that needs a supported training platform, because the published tags are v0.5, v0.3 and v0.2 from October 2024 while the last push was on 2026-07-17, a gap that leaves the API surface uncertain. Before adopting, read the paper linked in the documentation, run the Quick Check torchrun command on two local processes to see the mesh come up, and confirm your nodes can install iperf and reach each other over stable peer-to-peer links before planning a multi-site run.

Frequently asked questions

Official sources

1. Issues
2. License: Apache-2.0
3. PrimeIntellect-ai/prime-diloco on GitHub
4. README
5. Releases