PrimeIntellect-ai/prime-diloco

prime is a framework for efficient, globally distributed training of AI models over the internet.

What it solves

Prime is a framework designed for the efficient, globally distributed training of AI models over the internet. It addresses the challenges of training large models across geographically dispersed nodes, such as network instability, high latency, and the high cost of checkpointing and recovery.

How it works

Prime utilizes the DiLoCo (Distributed Low-Communication) approach to minimize communication overhead. It implements several key optimizations:

  • ElasticDeviceMesh: A fault-tolerant abstraction that manages dynamic process groups, allowing nodes to join or leave the training run without requiring a cold restart.
  • Asynchronous Checkpointing: To prevent blocking the main training process, checkpoints are first saved to a RAM-backed filesystem (/dev/shm) and then asynchronously uploaded to disk or remote storage.
  • Live Checkpoint Recovery: Joining nodes can quickly retrieve the latest model state from peers via a sidecar HTTP server, allowing them to join the training process mid-run.
  • Custom Int8 All-Reduce Kernel: A C++ implementation of a ring-reduce kernel that quantizes pseudo-gradients to Int8 to reduce network payload size by 4x without impacting loss curves.
  • Bandwidth Optimization: It maximizes network utilization through sharded pseudo-gradients (multiple simultaneous connections) and VPN technology to optimize peer-to-peer routing.
  • Memory Management: It uses PyTorch FSDP2/DTensor ZeRO-3 to shard model weights, gradients, and optimizer states across intra-node GPUs, and offloads the DiLoCo optimizer tensors to CPU memory.

Who it’s for

Researchers and engineers training large-scale AI models who need to distribute training across multiple data centers or over the public internet rather than within a single high-speed cluster.

Highlights

  • Fault-tolerant training via ElasticDeviceMesh and heartbeat mechanisms.
  • 4x reduction in communication payload using a custom Int8 all-reduce kernel.
  • Zero-blocking checkpointing through RAM-backed asynchronous uploads.
  • High-bandwidth internet training achieving up to 4Gb/s between data centers across the US.
  • FSDP2 integration for efficient memory sharding (ZeRO-3).

Related

  • Project
    hpcaitech/ColossalAI

    A distributed deep learning framework that makes training and inference for large AI models faster and cheaper through advanced parallelism and memory management.

  • Project
    PrimeIntellect-ai/OpenDiloco

    An open-source framework for globally distributed training of large models that reduces communication overhead using distributed weight averaging.

  • Project
    deepspeedai/DeepSpeed

    A deep learning optimization library that enables the training of models with trillions of parameters by optimizing memory and compute efficiency across distributed systems.

  • Project
    EvolvingLMMs-Lab/lmms-engine

    A unified, high-performance training engine for multimodal models that integrates distributed training, kernel fusion, and advanced optimizers to enable scalable pretraining and fine-tuning.

  • Project
    PrimeIntellect-ai/prime-rl

    A framework for large-scale asynchronous reinforcement learning designed to train frontier models with 1T+ parameters on up to 1000+ GPUs.