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

Summary

This paper presents a novel scheme for reducing the communication cost of training LLMs using FSDP approach. To this end, the authors suggest to quantize the weights differences (instead of the weights themselves) and applying Hadamard transformations for making the gradients smooth. They show up to ~4x training speedup on 128 GPUs.

Strengths

1. The paper studies a highly important problem (communication cost of llm training). 2. Although the ideas of quantizing the weights difference and Hadamard smoothing are not new, using them in this direction is interesting. 3. The scheme is evaluated on a solid setting (16x4 A100 and 16x8 H100). 4. The theoretical convergence analysis of difference quantization is nice. 5. The paper is well-written and easy to follow.

Weaknesses

1. Although, the idea is interesting, I think the authors should provide more ablations on various aspects of their design choice (see questions section) to elaborate the approach better. 2. To motivate the problem more, I would like to see some practical analysis of the communication volume based on the number of nodes. For example, how much data (in GB) will be sent between gpus in a single node (and between different nodes) when we scale to a high number of gpus.

Questions

I have some questions/ideas about the paper. I would be happy if the authors answer them. 1. The paper applies symmetric quantization for all schemes. However, in the communication reduction techniques, it is natural to use asymmetric quantization (which is a bit more accurate). Also, stochastic rounding could be another option to have unbiased gradients. Have you tried those schemes to see the accuracy changes? Maybe you don't need to apply hadamard in that case and the accuracy will be recovered! 2. The main motivation behind the hadamard smoothing is to remove the outliers. This is studied in several works [1, 2, 3, 4] (which are not cited here also). My main question is why we do not see an analysis of the gradient distribution in the paper so we can make sure that the hadamard is important. 3. There isn't any study on the overhead of applying Hadamard scheme. I agree that there are fast hadamard kernels when the input is power-of-two. However, it is not clear for me if your input is not power-of-two. Also, there isn't an ablation on the size of the Hadamard in that case. References: 1. Training Transformers with 4-bit Integers (https://arxiv.org/abs/2306.11987) 2. QuaRot: Outlier-Free 4-Bit Inference in Rotated LLMs (https://arxiv.org/abs/2404.00456) 3. QuIP#: Even Better LLM Quantization with Hadamard Incoherence and Lattice Codebooks (https://arxiv.org/pdf/2402.04396) 4. SpinQuant: LLM quantization with learned rotations (https://arxiv.org/abs/2405.16406)

Rating

7

Confidence

4

Soundness

3

Presentation

3

Contribution

3

Limitations

Yes. The author discuss that the work focuses only on communication reduction.

Summary

Thank you for submitting your paper to Neurips 2024. The paper proposes SDP4Bit, a communication quantization method that mitigates accuracy loss by weight difference quantization and 8-bit intra-node and 4-bit internode quantization. The authors provide convergence analysis to show that SDP4Bit is compatible with both biased and unbiased compressors and experimental results to show that SDP4Bit outperforms existing solutions.

Strengths

1. The authors provide convergence analysis to show that weight difference quantization is compatible with both biased and unbiased compressors. 2. The two-level quantization method is suitable for the existing GPU cluster network topology. 3. The authors' implementation uses several system optimizations.

Weaknesses

1. Missing key ablation studies. The two proposed strategies in this paper are weight difference quantization and two-level int8-int4 all-to-all gradient averaging. So I suggest the authors add the following two experiments to compare the accuracies. - Running ZeRO++ with two-level int8-int4 all-to-all gradient averaging and its original quantization method - Running ZeRO++ with its original int4-int4 all-to-all gradient averaging and the weight difference quantization. 2. The figures are too small to read. 3. Related work could add AG-SGD, which compresses changes in activations rather than the activations themselves, so it does not rely on the assumption that gradients are unbiased. [1] Fine-tuning Language Models over Slow Networks using Activation Quantization with Guarantees

Questions

1. Under what assumptions are weight differences generally easier to quantify? I've read several papers based on this intuition, but none that provide a comprehensive analysis of this intuition. 2. If you only need to assume that the magnitude of the weight difference is smaller than the weight, how do you ensure this? 3. If you ensure this via the learning rate, I think it will potentially lead to higher wall clock convergence time.

Rating

5

Confidence

4

Soundness

3

Presentation

4

Contribution

3

Limitations

Yes.

Summary

This paper proposed a 4bit quantization framework for sharded data parallelism training. It proposed to quantize the weight differences between iterations as the first method to reduce the accuracy degradation. It also proposed to mixed-precision quantization for intra- and inter- node communication. Experiment results show that the proposed system can achieve 2x to 4x speedup on training GPT models with size from millions to a few billions.

Strengths

1. The problem of communication bottleneck in FSDP, especially multi-node training is important. The paper aims to solve this problem and the results seems to encourage this approach. 2. The paper implemented the framework in real hardware environment and achieves 2-4x speedup by reducing the communication.

Weaknesses

1. The experiment results shows training/validation loss on Pile for the proposed method and baseline (FP) and QSDP and other method. The loss is an auxiliary variable to the performance of the model. A direct comparison on real dataset will be more convincing. For example, what is the accuracy of the model training by proposed methods and ZeRO/QSDP on mmlu and hellaswag/piqa/winogrande/ARC/ARE etc reasoning tasks, and gsk8k etc math tasks? This is critical to show the accuracy/precision improvement over other 4-bit communication methods. 2. The paper shows the speedup of the proposed methods, but I would like to know which results in Table 4 is the training throughput for ZeRO++ and QSDP (if applicable)? I would like to see a comparison for accuracy-training_speed between the proposed method and ZeRO /QSDP. Furthermore, is there any results on the training throughput (tokens/sec) at certain sequence length (instead of training TFLOPs)? 3. The paper introduces a few method, including qWD, two-level quantization, Hadamard transforms, etc. I would like to know which ones are the invention of the paper? For example, using Hadamard transform to handle the outlier has been known for a while.

Questions

Please refer to the session "Weakness"

Rating

6

Confidence

4

Soundness

3

Presentation

3

Contribution

3

Limitations

n/a

Summary

The paper introduces a novel approach to reduce communication overhead in Sharded Data Parallelism (ShardedDP) for training large language models (LLMs). It proposes two key techniques: quantization on weight differences and two-level gradient smooth quantization, which effectively compress weights and gradients to nearly 4 bits without compromising training accuracy. The method is implemented within the Megatron-LM framework and achieves up to 4.08× speedup in end-to-end throughput on 128 GPUs. The paper also provides theoretical guarantees of convergence and demonstrates negligible impact on training loss for models with up to 6.7 billion parameters.

Strengths

1. Weight difference and two-level gradient quantization are more effective at retaining accuracy and significantly improving throughput than baseline approaches 2. The authors present a complete story with a new quantization algorithm, accompanying GPU kernel implementations, and detailed evaluation/ablation over a range of models and baselines which together clearly illustrate all the benefits of each component of their approach.

Weaknesses

1. The connection between the convergence analysis and the quantization schemes is not entirely clear. Specifically, it is not clear if the quantizers correspond to the biased/unbiased compressors in Algorithm 4 and if so which of them is biased and which is unbiased. It appears from lines 5 and 7 of algorithm 4 that TLq is unbiased and qWD is biased but I couldn't find any mention/proof of that, so I am not sure if that is the case or if I am missing the point entirely. 2. Most of the empirical analysis is limited to the older generation of GPT models, and it is possible that may not extend as well to modern models like Mistral series, Phi series, Orca series, Llama series etc as well as the newer generation of GPT models

Questions

1. Why is 4-bit quantization needed? Is 8-bit quantization not sufficient in practice, especially since we don't expect foundation models to be trained too frequently? 2. Why is inter node communication quantized to 4 bits while intra node communication is quantized to 8 bits and not the other way around? 3. I think the Hadamard transform matrix might need to be post multiplied by a diagonal matrix with randomly chosen +1/-1 values on its diagonal for it to be effective at removing outliers and flattening the data (https://arxiv.org/abs/1011.1595). Is that what has been done here? If not, why?

Rating

6

Confidence

4

Soundness

3

Presentation

3

Contribution

3

Limitations

N/A

Response to rebuttal

Thank you for addressing my concerns. As I had already recommended accepting the paper, I will keep my score.