ATOM: Asynchronous Training of Massive Models for Deep Learning in a Decentralized Environment

TL;DR

ATOM introduces a decentralized, asynchronous training framework for large language models that leverages cost-effective hardware, enabling scalable training without central bottlenecks and outperforming traditional methods in slow network environments.

Contribution

The paper presents a novel decentralized training framework, ATOM, which allows training large models on commodity hardware with asynchronous model swapping and optimized partitioning strategies.

Findings

Up to 20x training efficiency improvement over existing decentralized pipeline methods.

Successfully trains GPT-3 configurations on consumer-grade hardware.

Avoids central failure points inherent in pipeline parallelism.

Abstract

The advent of the Transformer architecture has propelled the growth of natural language processing (NLP) models, leading to remarkable achievements in numerous NLP tasks. Yet, the absence of specialized hardware like expansive GPU memory and high-speed interconnects poses challenges for training large-scale models. This makes it daunting for many users to experiment with pre-training and fine-tuning large language models (LLMs). In this study, we introduce \atom, a resilient distributed training framework designed for asynchronous training of vast models in a decentralized setting using cost-effective hardware, including consumer-grade GPUs and Ethernet. Unlike conventional model partitioning methods that distribute sub-models across GPUs, \atom aims to accommodate a complete LLM on one host (peer) through seamlessly model swapping and concurrently trains multiple copies across various peers to optimize training throughput. Through static analysis, \atom identifies the best model partitioning strategy and flawlessly merges model execution with swapping. Key benefits of \atom include: Avoiding the central point of failure found in pipeline parallelism methods. Demonstrating superior performance and scalability compared to closely-integrated pipeline parallelism in slower networks. Our experiments using different GPT-3 model configurations reveal that, in scenarios with suboptimal network connections, \atom can enhance training efficiency up to $20 \times$ when juxtaposed with the state-of-the-art decentralized pipeline parallelism approaches.