Fine-tuning Language Models over Slow Networks using Activation Compression with Guarantees

TL;DR

This work tackles the bottleneck of activations in pipeline-parallel training of large language models by introducing AQ-SGD, an activation-change compression algorithm that provides convergence guarantees for non-convex objectives without assuming gradient unbiasedness. The method quantizes activation deltas across epochs, stores per-sample activation history, and uses compressed gradients, achieving robust convergence even at 2-4 bit activation precision. The authors prove an $O(1/\sqrt{T})$ convergence rate under mild assumptions and demonstrate up to 4.3× end-to-end speedups on slow networks, with up to 4.9× when combined with state-of-the-art gradient compression. The approach is validated on DeBERTa-1.5B and GPT-2-1.5B benchmarks, showing that end-to-end communication compression can substantially accelerate distributed fine-tuning without sacrificing model quality.

Abstract

Communication compression is a crucial technique for modern distributed learning systems to alleviate their communication bottlenecks over slower networks. Despite recent intensive studies of gradient compression for data parallel-style training, compressing the activations for models trained with pipeline parallelism is still an open problem. In this paper, we propose AC-SGD, a novel activation compression algorithm for communication-efficient pipeline parallelism training over slow networks. Different from previous efforts in activation compression, instead of compressing activation values directly, AC-SGD compresses the changes of the activations. This allows us to show, to the best of our knowledge for the first time, that one can still achieve $O(1/\sqrt{T})$ convergence rate for non-convex objectives under activation compression, without making assumptions on gradient unbiasedness that do not hold for deep learning models with non-linear activation functions.We then show that AC-SGD can be optimized and implemented efficiently, without additional end-to-end runtime overhead.We evaluated AC-SGD to fine-tune language models with up to 1.5 billion parameters, compressing activations to 2-4 bits.AC-SGD provides up to 4.3X end-to-end speed-up in slower networks, without sacrificing model quality. Moreover, we also show that AC-SGD can be combined with state-of-the-art gradient compression algorithms to enable "end-to-end communication compression: All communications between machines, including model gradients, forward activations, and backward gradients are compressed into lower precision.This provides up to 4.9X end-to-end speed-up, without sacrificing model quality.

Figures (11)

• Figure 1: (a) Fine-tuning GPT2-1.5B with different activation precisions in communication; (b) Average absolute value of activations and their changes for GPT2-1.5B during training.
• Figure 2: The communication pattern of training large language models with both data parallelism and pipeline model parallelism. $C$ denotes a compression module. The goal of this paper is to understand the design of $C$ for forward activation and backward gradient.
• Figure 3: Convergence (loss vs. # steps) of different approaches. $\boldsymbol{\times}$ represents divergence.
• Figure 4: End-to-end training performance over different networks. $\boldsymbol{\times}$ represents divergence.
• Figure 5: Convergence and Throughput of AQ-SGD combined with gradient compression.

Theorems & Definitions (6)

• Theorem 3.1
• Lemma A.1
• Lemma A.2
• Theorem A.3
• Lemma A.4
• Lemma A.5