LLM.int8() Mixed-Precision Decompositions

论文上下文

Paper title: Large Transformer Model Inference Optimization | Lil'Log Abstract: The figure illustrates the LLM.int8() method by decomposing matrix multiplication into an 8-bit quantization path for regular values and a 16-bit path for outliers. Source body: Table of Contents Methods Overview Distillation Quantization Challenges for Transformer Quantization Post-training quantization (PTQ) Mixed-precision quantization Quantization at fine-grained granularity Second order information for quantization Outlier smoothing Quantization-aware training (QAT) Pruning How to prune? How to retrain? Sparsity N:M Sparsity via Pruning Sparsified Transformer Mixture-of-Experts Routing Strategy Improvement Kernel Improvement Architectural Optimization Sparse Attention Patterns Recurrence Memory Saving Designs Adaptive Attention Citation References [Updated on 2023-01-24: add a small section on Distillation .] Large transformer models are mainstream nowadays, creating SoTA results for a variety of tasks. They are powerful but very expensive to train and use. The extremely high inference cost, in both time and memory, is a big bottleneck for adopting a powerful transformer for solving real-world tasks at scale. Why is it hard to run inference for large transformer models? Besides the increasing size of SoTA models, there are two main factors contributing to the inference challenge ( Pope et al. 2022 ): Large memory footprint . Both model parameters and intermediate states are needed in memory at inference time. For example, The KV cache should be stored in memory during decoding time; E.g. For a batch size of 512 and context length of 2048, the KV cache totals 3TB, that is 3x the model size (!). Inference cost from the attention mechanism scales quadratically with input sequence length. Low parallelizability. Inference generation is executed in an autoregressive fashion, making the decoding process hard to parallel. In this post, we will look into several approaches for making transformer inference more efficient. Some are general network compression methods, while others are specific to transformer architecture. Methods Overview # We in general consider the following as goals for model inference optimization: Reduce the memory footprint of the model by using fewer GPU devices and less GPU memory; Reduce the desired computation complexity by lowering the number of FLOPs needed; Reduce the inference latency and make things run faster. Several methods can be used to make inference cheaper in memory or/and faster in time. Apply various parallelism to scale up the model across a large number of GPUs. Smart parallelism of model components and data makes it possible to run a model of trillions of parameters. Memory offloading to offload temporarily unused data to the CPU and read them back when needed later. This helps with memory usage but causes higher latency. Smart batching strategy; E.g. EffectiveTransformer packs consecutive sequences together to remove padding within one batch. Network compression techniques, such as pruning, quantization, distillation . A model of smaller size, in terms of parameter count or bitwidth, should demand less memory and run faster. Improvement specific to a target model architecture. Many architectural changes , especially those for attention layers, help with transformer decoding speed. Check the previous post on large model training on different types of training parallelism and memory saving designs including CPU memory offloading. This post focuses on network compression techniques and architecture-specific improvement for transformer models. Distillation # Knowledge Distillation ( KD ; Hinton et al. 2015 , Gou et al. 2020 ) is a straightforward way to build a smaller, cheaper model ( “student model” ) to speed up inference by transferring skills from a pre-trained expensive model ( “teacher model” ) into the student. There is no much restriction on how the student architecture should be constructed, except for a matched output space with the teacher in order to construct a proper learning objective. The generic framework of teacher-student knowledge distillation training. (Image source: Gou et al. 2020 ) Given a dataset, a student model is trained to mimic outputs of a teacher via distillation loss. Usually a neural network has a softmax layer; For example, a LLM outputs a probability distribution over tokens. Let’s denote the logits layer right before softmax as $\mathbf{z}_t$ and $\mathbf{z}_s$ for teacher and student models, respectively. The distillation loss minimizes the difference between two softmax outputs with a high temperature $T$. When ground truth labels $\mathbf{y}$ are known, we can combine it with a supervised learning objective between ground truth and the student’s soft logits using e.g. cross-entropy. $$ \mathcal{L}_\text{KD} = \mathcal{L}_\text{distll}(\text{softmax}(\mathbf{z}_t, T), \text{softmax}(\mathbf{z}_s, T)) + \lambda\mathcal{L}_\text{CE}(\mathbf{y}, \mathbf{z}_s) $$ where $\lambda$ is a hyperparameter to balance between soft and hard learning objectives. A common choice for $\mathcal{L}_\text{distll}$ is KL divergence / cross entropy. A successful early trial is DistilBERT ( Sanh et al. 2019 ) that is able to reduce the parameters of a BERT by 40% while maintaining 97% performance of BERT on fine-tuned downstream tasks and running 71% faster. The loss of pre-training DistilBERT is a combination of soft distillation loss, supervised training loss (i.e. Masked language modeling loss $\mathcal{L}_\text{MLM}$ in the case of BERT) and a special cosine embedding loss to align the hidden state vectors between teacher and student. Distillation can be easily combined with quantization , pruning or sparsification techniques, where the teacher model is the original full-precision, dense model and the student is quantized, pruned, or trimmed to have higher sparsity level. Quantization # There are two common approaches for applying quantization on a deep neural network: Post-Training Quantization (PTQ) : A model is first trained to convergence and then we convert its weights to lower precision without more training. It is usually quite cheap to implement, in comparison to