There was a time when you could comfortably train state-of-the-art vision and language models at home on your workstation. The first convolutional neural net to beat ImageNet (AlexNet) was trained for 5-6 days on two gamer-grade GPUs. In contrast, today's Top-1 ImageNet model (CoAtNet) takes 20,000 TPU-v3 days. And things are even worse in the NLP world: training GPT‑3 on a top-tier server with 8x A100 would take decades.

So, can individual researchers and small labs still train state-of-the-art models? Yes we can! All it takes is for a bunch of us to come together. In fact, we're doing it right now and you are invited to join!

In this demo, we train a model similar to OpenAI DALL-E — a Transformer model that generates images from text descriptions. It is trained on LAION-400M, the world's largest openly available image-text-pair dataset with 400 million samples. Our model is based on the dalle‑pytorch implementation by Phil Wang with a few tweaks to make it communication-efficient.

Modern distributed training algorithms are designed for HPC clusters with a 10-100 gigabit per second bandwidth. In turn, a typical Internet connection runs at 10-100 megabits per second: that’s three orders of magnitude slower. To make distributed training efficient, you need to win back these three orders of magnitude. This may seem daunting at first, but in reality, DL researchers have already made all the necessary pieces for solving this puzzle:

Speed‑up How to achieve
4-16x Large-batch training: You et al. (2019) proposed a way for training neural networks efficiently with larger batches, and hence, fewer communication rounds.
4-32x Gradient compression: from simple 8-bit quantization to advanced techniques such as Deep Gradient Compression, PowerSGD, 1-bit Adam, and many others. As a rule of thumb, these techniques can safely reduce communication by 16-32x. More extreme compression is often possible, but it may affect stability or final quality.
4-24x Parameter sharing: reusing parameters between model layers results in a model with fewer parameters, and hence, fewer gradients to communicate. Lan et al. (2019) and Xue et al. (2021) propose efficient parameter sharing architectures for NLP and computer vision.
1.5-2x Overlapping computation with communication: running network communication in background while computing the next portion of gradients. This is a long-standing trick from HPC that was recently adapted for DL training. Ren et al. (2021) show that updating parameters in background while computing the next batch of gradients does not harm convergence.

These techniques are already more than enough to cover 1000x slower communication. This means that in practice you can pick and choose which of them you want in your training run. For this demo, we use 8x larger batches, 4x compression, 12x parameter sharing and partial overlapping. If you don’t want parameter sharing, you can trade it for more advanced gradient compression or larger batches.

Most distributed DL frameworks assume that the computation is performed by a fleet of identical devices, typically GPU servers or TPU cores. Under this assumption, each device can be assigned an equal part of computation, such as processing a fixed batch size of training samples. However, this quickly breaks down if workers use different device types. If one participant uses a GPU (e.g. P100) and another runs on TPU-v2-8, it is difficult to find a regime where both devices will be fully utilized.

To make the best use of all available devices, we let each device accumulate gradients at its own pace with individually tuned batch size and some other features (e.g. gradient checkpointing or using XLA). Once workers collectively aggregate some predefined global batch size, they average their gradients with weights proportional to each worker's individual contribution (i.e. number of samples processed).

This technique allows the "swarm" to automatically adjust its behavior as peers join, leave or fail. For instance, if several high-performance peers join the experiment, other peers will need to process a smaller number of samples per optimizer step, and hence, the collaboration will train faster with the same hyperparameters. In turn, if one of the workers fails and loses its progress (e.g. due to a fp16 overflow), others will make up for that by processing slightly more. For more details on how this works, please refer to "Deep Learning In Open Collaborations" paper or the corresponding blog post.

How do I join?

To participate in the demo training run, you need to follow these steps:

  1. Log in to your Hugging Face account or create a new one,
  2. Join our organization using this invite link,
  3. Create a new User Access token at huggingface.co/settings/token and copy it to the clipboard,
  4. Run the starter notebook locally if you have a GPU, or join using Google Colab and/or Kaggle (please don't use multiple accounts at once),
  5. Keep it running for at least 15 minutes if you'd like your contribution to be recorded.

  • The notebook requires an NVIDIA GPU and Linux. It's also possible to run it inside WSL (Windows Subsystem for Linux) on Windows 11.
  • You need to install Anaconda, PyTorch, the CUDA toolkit, and other dependencies. In case of Ubuntu, you can run:
    sudo apt install -y git tmux
    curl https://repo.anaconda.com/archive/Anaconda3-2021.11-Linux-x86_64.sh > Anaconda3-2021.11-Linux-x86_64.sh
    bash Anaconda3-2021.11-Linux-x86_64.sh -b -p ~/anaconda3
    source ~/anaconda3/bin/activate
    conda install -y pytorch torchvision torchaudio cudatoolkit=11.3 -c pytorch
    
    git clone https://github.com/learning-at-home/dalle-hivemind
    cd dalle-hivemind && python3 -m pip install -r requirements.txt
    
    # Re-install bitsandbytes for the actual CUDA version
    pip uninstall -y bitsandbytes-cuda111
    pip install -y bitsandbytes-cuda113==0.26.0
  • Download the notebook, remove the first 4 lines (they install dependencies but we've already did that), and run it in Jupyter.
  • If it fails to find libcudart.so.11.0, add the line below to ~/.bashrc:
    export LD_LIBRARY_PATH="/usr/local/cuda/lib64:/usr/local/cuda/extras/CUPTI/lib64:$LD_LIBRARY_PATH"
  • As reported by @Umbreon, the notebook uses 2 CPU cores, 13 GB RAM, 50 GB upload, and 50GB download per 6 hours.

If you have any questions or would like to follow news related to this training run, please join the #neurips‑demo channel in our Discord:

If everything went fine, you will see that the notebook begins to process minibatches and communicate with peers. You can view the global collaboration progress in the interactive dashboard. This dashboard also visualizes individual contributions of each peer and statistics, such as swarm-wide GPU performance. More materials (including a notebook for inference and the latest model checkpoint) are available on the organization page.

Sometimes it will also print communication error messages: this is fine as long as the peer continues training after the error. These error messages are typically caused by other peers disconnecting abruptly (or, rarely, network faults). Our code is designed to recover from these errors, but it still shows them for clarity.

Practical aspects

Our aim is to train a large model in a decentralized fashion on consumer hardware or low-end cloud instances. This means we need to make the model, dataset, and other memory buffers fit onto a few GB of disk, 12-16 GB of CPU RAM, and 8-12 GB of GPU memory. Unfortunately, this rules out many popular techniques such as ZeRO-Offload: there is simply not enough RAM for that. Instead, we must make better use of what limited memory we have. To do this, we use two techniques: 8-bit Optimizers for GPU memory and dataset streaming for RAM & HDD.

8-bit optimizers: Using optimizers such as LAMB or Adam requires four times as much GPU memory as simply storing model parameters (8 bytes vs 2 bytes) because of additional gradient statistics. As such, for training large models with many parameters, the optimizer state takes the largest amount of memory. With 8-bit optimizers, this amount is reduced by 75% (2 bytes), making it much easier to fit large models onto consumer GPUs.

Naturally, we can combine this technique with offloading and store 8-bit optimizer states in the CPU memory rather than in the GPU memory (0 bytes GPU, 2 bytes CPU). To perform an optimizer update, we transfer the GPU gradients to the CPU, update the model parameters, and then copy the new weights to the GPU. We can do this for each weight one-by-one so that the additional CPU memory required for the optimizer update is minimal. This combination of offloading and 8-bit optimizers means that we conserve GPU memory (0 bytes per parameter) and also use only a limited amount of CPU memory (2 bytes per parameter).

We've created an interactive "calculator" to help you get a better "feel" of which techniques to use and how they fit together. It covers 8-bit optimizers, offloading, as well as other popular tricks that emphasize memory efficiency. For each model configuration, it computes GPU and RAM usage based on the shapes of tensors that will be allocated. Of course, the real-world memory usage depends on the implementation details and technical inefficiencies. A good rule of thumb for PyTorch is to increase both CPU and GPU memory by 20-30%.

Dataset streaming: Usually data is stored on disk and needs to be fully or partially loaded into RAM for training. Large datasets used for pretraining measure in hundreds of gigabytes or even terabytes. This can pose a significant problem, as most desktop and cheap cloud instances simply do not have that much free space. Furthermore, downloading the data over the Internet would take up hours before one can even begin training.

To circumvent these problems, it is possible to stream the data in the same way as you stream online videos. Participants download a small random portion of the training dataset and immediately begin training on it, while additional data is loaded in the background. As such, we can train a model with virtually no storage overhead from the dataset, and switching to a new dataset is as simple as changing an argument of the dataset class.

Here's our tutorial covering these methods:

In this section, we discuss common concerns related to security of collaborative training:

Q: If I join a collaborative experiment, do I allow other people to execute code on my computer?

A: During the training, participants only exchange data (gradients, statistics, model weights) and never send code to each other. No other peer can execute arbitrary code on your computer.

To join the experiment, you typically need to run the code (implementing the model, data streaming, training loop, etc.) from a repository or a Colab notebook provided by the authors of the experiment. This is no different from running any other open source project/Colab notebook.

Q: Can a malicious participant influence the training outcome?

A: It is indeed possible unless we use some defense mechanisms. For instance, a malicious participant can damage model weights by sending large numbers instead of correct gradients. The same can happen due to broken hardware or misconfiguration.

  • One possible defense is using authentication combined with model checkpointing. In this case, participants should log in (e.g. with their Hugging Face account) to interact with the rest of the collaboration. In turn, moderators can screen potential participants and add them to an allowlist. If something goes wrong (e.g. a participant sends invalid gradients and the model diverges), the moderators remove them from the list and revert the model to the latest checkpoint unaffected by the attack.

    Nice bonus: using this data, the moderators can acknowledge the personal contribution of each participant.

  • Another defense is replacing the naive averaging of the peers' gradients with an aggregation technique that is robust to outliers. Karimireddy et al. (2020) suggested such a technique (named CenteredClip) and proved that it does not significantly affect the model's convergence.

    In our case, CenteredClip is useful but not enough to protect from malicious participants, since it implies that the CenteredClip procedure itself is performed by a trusted server. By contrast, in our decentralized system, all participants can aggregate a part of the gradients, and we cannot assume any of them to be trusted.

    Recently, Gorbunov et al. (2021) proposed a robust aggregation protocol for decentralized systems that does not require this assumption. This protocol uses CenteredClip as a subroutine but is able to detect and ban participants who performed it incorrectly.

In this section, we provide a recipe for you to run a collaborative training experiment yourself.

Got confused? Feel free to ask any questions in our Discord!

  1. Set up dataset streaming:
    • Upload your dataset to the Hugging Face Hub in a streaming-friendly format (example).
    • Set up dataset streaming (see the "Memory-Efficient Training" section).
  2. Write the code of training peers (example):
    • Implement your model, set up dataset streaming, and write the training loop.
    • Get familiar with the hivemind library (quickstart).
    • In the training loop, wrap up your PyTorch optimizer with hivemind.Optimizer (example).
  3. (optional) Write the code of auxiliary peers (example):
    • Auxiliary peers are a special kind of peers responsible for logging experiment progress (e.g., to Weights & Biases) and uploading model checkpoints (e.g., to Hugging Face Hub).
    • Such peers don't need to calculate gradients and may be launched on cheap machines without GPUs.
    • They can serve as a convenient entry point to hivemind.DHT (i.e., their address can be specified as initial_peers).
    • It is useful to fix their address by providing host_maddrs and identity_path arguments to hivemind.DHT (these are forwarded to the underlying libp2p daemon).
  4. (optional) Make it easier for other people to join:
    • Create notebooks for free GPU providers (Google Colab, Kaggle, AWS SageMaker, etc.). People may run them online and/or download and run them on their own hardware.
    • Create a Hugging Face organization with all resources related to the training (dataset, model, inference demo, how-to-join walkthrough, links to a dashboard with loss and other metrics, etc.). Look at ours for an example.
    • Set up an authentication system (see the "Security" section). For example, you can ask people to join your organization with their Hugging Face accounts (the website allows either sharing a link for joining or manually approving new participants). This allows you to screen the peers, acknowledge their contributions (e.g., make a leaderboard), and ban accounts who behave maliciously. You can use our authentication system or deploy your own (our server implementation might be a good start).
    • Set up an inference demo for your model (e.g., using Spaces) or a script that periodically uploads the inference results to show the training progress.

Organizers

This demonstration was created by Alexander Borzunov*, Max Ryabinin*, Tim Dettmers*, Quentin Lhoest*, Lucile Saulnier*, Michael Diskin, Yacine Jernite, and Thomas Wolf.

Learn more