Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

SDP4Bit: Toward 4-bit Communication Quantization in Sharded Data Parallelism for LLM Training

Jinda Jia, Cong Xie, Hanlin Lu, Daoce Wang, Hao Feng, Chengming Zhang, Baixi Sun, Haibin Lin, Zhi Zhang, Xin Liu, Dingwen Tao

TL;DR

This work proposes SDP4Bit (Toward 4Bit Communication Quantization in Sharded Data Parallelism for LLM Training), which effectively reduces the communication of weights and gradients to nearly 4 bits via two novel techniques: quantization on weight differences, and two-level gradient smooth quantization.

Abstract

Recent years have witnessed a clear trend towards language models with an ever-increasing number of parameters, as well as the growing training overhead and memory usage. Distributed training, particularly through Sharded Data Parallelism (ShardedDP) which partitions optimizer states among workers, has emerged as a crucial technique to mitigate training time and memory usage. Yet, a major challenge in the scalability of ShardedDP is the intensive communication of weights and gradients. While compression techniques can alleviate this issue, they often result in worse accuracy. Driven by this limitation, we propose SDP4Bit (Toward 4Bit Communication Quantization in Sharded Data Parallelism for LLM Training), which effectively reduces the communication of weights and gradients to nearly 4 bits via two novel techniques: quantization on weight differences, and two-level gradient smooth quantization. Furthermore, SDP4Bit presents an algorithm-system co-design with runtime optimization to minimize the computation overhead of compression. In addition to the theoretical guarantees of convergence, we empirically evaluate the accuracy of SDP4Bit on the pre-training of GPT models with up to 6.7 billion parameters, and the results demonstrate a negligible impact on training loss. Furthermore, speed experiments show that SDP4Bit achieves up to 4.08$\times$ speedup in end-to-end throughput on a scale of 128 GPUs.

SDP4Bit: Toward 4-bit Communication Quantization in Sharded Data Parallelism for LLM Training

TL;DR

This work proposes SDP4Bit (Toward 4Bit Communication Quantization in Sharded Data Parallelism for LLM Training), which effectively reduces the communication of weights and gradients to nearly 4 bits via two novel techniques: quantization on weight differences, and two-level gradient smooth quantization.

Abstract

Recent years have witnessed a clear trend towards language models with an ever-increasing number of parameters, as well as the growing training overhead and memory usage. Distributed training, particularly through Sharded Data Parallelism (ShardedDP) which partitions optimizer states among workers, has emerged as a crucial technique to mitigate training time and memory usage. Yet, a major challenge in the scalability of ShardedDP is the intensive communication of weights and gradients. While compression techniques can alleviate this issue, they often result in worse accuracy. Driven by this limitation, we propose SDP4Bit (Toward 4Bit Communication Quantization in Sharded Data Parallelism for LLM Training), which effectively reduces the communication of weights and gradients to nearly 4 bits via two novel techniques: quantization on weight differences, and two-level gradient smooth quantization. Furthermore, SDP4Bit presents an algorithm-system co-design with runtime optimization to minimize the computation overhead of compression. In addition to the theoretical guarantees of convergence, we empirically evaluate the accuracy of SDP4Bit on the pre-training of GPT models with up to 6.7 billion parameters, and the results demonstrate a negligible impact on training loss. Furthermore, speed experiments show that SDP4Bit achieves up to 4.08 speedup in end-to-end throughput on a scale of 128 GPUs.

Paper Structure

This paper contains 25 sections, 3 theorems, 20 equations, 11 figures, 9 tables, 4 algorithms.

Table of Contents

  1. Introduction
  2. Preliminaries
  3. Sharded Data Parallelism
  4. Quantization
  5. Collective Reduction Communication with Quantization
  6. Methodology
  7. Quantization on Weight Differences (qWD)
  8. Two-Level Gradient Smooth Quantization
  9. Two-level Gradient Quantization (TLq)
  10. TLq with Hadamard Smoother (TLq-HS)
  11. Performance Optimizations in Implementation
  12. Theoretical Analysis
  13. Counterexample of Biased Weight Compression
  14. Convergence Analysis
  15. Evaluation
  16. ...and 10 more sections

Key Result

Theorem 4.1

For arbitrary non-convex function under Assumption asm:smoothness and Assumption asm:gradient, taking learning rate $\eta \leq \frac{1}{10L \left(\frac{2}{\delta} + \rho \kappa + \rho + \kappa \right)}$, Algorithm alg:sgd_compress converges to a critical point with the following error bound:

Figures (11)

  • Figure 1: Training validation loss for GPT-6.7B; SDP4Bit is closely aligned with full precision training.
  • Figure 2: Communication of quantized weight differences.
  • Figure 3: Two-level gradient quantization: 8-bit intra-node and 4-bit inter-node quantization.
  • Figure 4: Histogram of (a) weights and (b) weight differences. Each vertical dashed line represents a quantization level corresponding to a 4-bit quantization lattice.
  • Figure 5: Validation loss comparison for the Baseline, ULq, TLq, and TLq-HS on the GPT-125M model. Uniformly applying 4-bit gradient quantization twice results in a noticeable gap compared to the baseline. In contrast, two-level quantization (8-bit for intra-node and 4-bit for inter-node) mitigates this gap. The Hadamard smoother further reduces the gap, making the loss nearly identical to the baseline.
  • ...and 6 more figures

Theorems & Definitions (9)

  • Definition 4.1: Unbiased $\kappa$-approximate compressor Alistarh2016QSGDCS
  • Definition 4.2: $\delta$-approximate compressor Karimireddy2019ErrorFF
  • Remark 4.1
  • Remark 4.2
  • Theorem 4.1: Convergence error bound
  • Remark 4.3
  • Lemma A.1
  • Theorem A.1: Convergence error bound
  • proof