Pier: Efficient Large Language Model pretraining with Relaxed Global Communication

TL;DR

Pier introduces a scalable optimizer with relaxed global communication to significantly accelerate large language model pretraining while maintaining performance, leveraging innovative techniques and system architecture for efficient parallelization.

Contribution

The paper proposes Pier, a novel optimizer with relaxed global communication, enabling faster LLM pretraining without sacrificing model quality, and demonstrates its effectiveness on various GPT models.

Findings

Speeds up GPT-2 XL training by up to 3.7x on 256 GPUs

Reduces GPT-2 7B training time by 54.5% with combined parallel strategies

Maintains validation loss and downstream performance despite acceleration

Abstract

Global communication, such as all-reduce and allgather, is the prominent performance bottleneck in large language model (LLM) pretraining. To address this issue, we present Pier, an efficient and scalable optimizer with relaxed global communication. Pier is built upon DiLoCo, which leverages an inner optimizer within groups of processors and an outer optimizer that requires global communication. To preserve the convergence and model performance, Pier incorporates two key techniques for the outer optimizer: momentum warmup and momentum decay. Pier employs an efficient and scalable system architecture to enable complex parallelization strategies in LLM pretraining. We examine the model performance and runtime reduction of Pier using the GPT model family (e.g., small, medium, XL, and 7B) and the OpenWebText dataset with a suite of thirteen downstream tasks. With data parallel strategy, Pier speeds up GPT-2 XL training by up to 2.7x-3.7x on 256 NVIDIA A100 GPUs and 1.2x-1.9x on 64 GH200 Superchips, respectively, without degradation of validation loss or downstream task performance. With data parallel and tensor parallel, Pier reduces the time cost GPT-2 7B model training by 54.5% on 128 A100s.