PaperlayerView exact revision on arXiv

Source-linked AI summary

Towards Efficient Post-training Quantization of Pre-trained Language Models

Haoli Bai, Lu Hou, Lifeng Shang, Xin Jiang, Irwin King, Michael R. Lyu

arXiv:2109.15082v1cs.CL

TL;DR

Quantization-aware training for pre-trained language models requires end-to-end training and full training-set access, with costs in time, memory, and data security. The paper proposes module-wise reconstruction error minimization and parallel module training for post-training quantization. Experiments report improved performance and reduced training time and memory use, including over QAT.

  • Problem

    Most PLM quantization methods use QAT, which requires slow end-to-end training, substantial memory, and access to the full training set.

    Source: §1 · p. 1

  • Method

    MREM minimizes quantization reconstruction error across modules, with a parallel strategy that trains modules locally on separate devices.

    Source: §1 · p. 2

  • Results

    Parallel MREM reports 85.5% accuracy on 4K training samples, one-third of QAT’s memory per GPU, and more than 150× faster training on BERT-large.

    Source: §1 · p. 2

  • Takeaways & Limitations

    On GLUE and SQuAD, MREM improves PTQ performance while reducing training time, memory overhead, and data consumption relative to QAT.

    Source: §1 · p. 2

    The current parallel strategy trains modules separately; caching module-wise gradients for the backward pass is proposed to reduce its discrepancy from conventional end-to-end training.

    Source: §6 · p. 11

Abstract

from arXiv · show

Network quantization has gained increasing attention with the rapid growth of large pre-trained language models~(PLMs). However, most existing quantization methods for PLMs follow quantization-aware training~(QAT) that requires end-to-end training with full access to the entire dataset. Therefore, they suffer from slow training, large memory overhead, and data security issues. In this paper, we study post-training quantization~(PTQ) of PLMs, and propose module-wise quantization error minimization~(MREM), an efficient solution to mitigate these issues. By partitioning the PLM into multiple modules, we minimize the reconstruction error incurred by quantization for each module. In addition, we design a new model parallel training strategy such that each module can be trained locally on separate computing devices without waiting for preceding modules, which brings nearly the theoretical training speed-up (e.g., $4\times$ on $4$ GPUs). Experiments on GLUE and SQuAD benchmarks show that our proposed PTQ solution not only performs close to QAT, but also enjoys significant reductions in training time, memory overhead, and data consumption.

1 INTRODUCTION

The introduction frames QAT’s time, memory, and data-access costs against PTQ’s promise, while noting that layer-wise PTQ can miss correlations among Transformer layers. It presents MREM, parallel module training, and annealed teacher forcing as a route to better PLM quantization efficiency and performance.

1 INTRODUCTION

  • QAT requires extensive training, end-to-end memory, and full training-set access, creating time, resource, and data-security challenges.The introduction notes that distillation can add memory overhead because of the teacher model.Source: §1 · p. 1
  • Existing PTQ often calibrates layer-wise with limited data, but this can perform poorly on PLMs because Transformer attention and feed-forward layers are coupled.The paper says layer-wise training ignores underlying correlations among layers.Source: §1 · p. 1
  • MREM jointly minimizes reconstruction error across multiple Transformer layers in each module, with module size adjustable to balance layer correlation and memory.The paper contrasts this with prior block-wise objectives requiring costly second-order Hessian computation.Source: §1 · p. 2
  • Parallel MREM assigns modules to separate devices for local training, while annealed teacher forcing uses full-precision modules to provide clean signals and limit error propagation.The introduction reports near-theoretical speedup, including 4× on 4 GPUs.Source: §1 · p. 2
  • 85.5% accuracy on 4K training samples, one-third of QAT’s memory per GPU, and more than 150× faster training are reported for parallel MREM on BERT-large.The paper also reports results on GLUE and SQuAD, with faster training, less memory overhead, and improved data security over QAT.Source: §1 · p. 2

2 RELATED WORK

Related work reviews PLM compression methods, with particular attention to quantization and parallel training. It situates the paper’s PTQ approach among compression alternatives and distinguishes its parallel strategy from pipeline parallelism.

2.1 Network Compression for Pre-trained Language Models

  • Quantization reduces parameter and activation precision; prior PLM methods largely use QAT, while PTQ is an alternative for limited training resources.The review also describes pruning, distillation, and neural architecture search as other compression directions.Source: §2.1 · pp. 2–3
  • Distillation transfers knowledge from a large teacher to a smaller student, but teacher models or precomputed representations can consume substantial memory.PTQ is related when the full-precision model provides layer-wise supervision to the quantized student.Source: §2.1 · pp. 2–3

2.2 Parallel Training for Pre-trained Language Models

  • Pipeline parallelism assigns modules to devices but incurs intermediate-tensor communication; the proposed strategy avoids real-time worker communication, though it can be suboptimal.The review notes that teacher forcing is used to alleviate this issue.Source: §2.2 · p. 3
  • Data-, operator-level, and optimizer model parallelism partition training data, computations, or model states in different ways, and could be combined with the paper’s strategy.The review identifies scaling to very large PLMs as a potential application of such combinations.Source: §2.2 · p. 3

3 MOTIVATION

The section introduces post-training quantization of PLMs as important but challenging, then reviews background on network quantization.

3 MOTIVATION

  • The section frames PLM post-training quantization as important yet challenging and first reviews network-quantization background.Source: §3 · p. 3

3.1 Quantization Background

Quantization replaces full-precision weights or activations with lower-bit representations, with this section contrasting quantization-aware training (QAT) and post-training quantization (PTQ). PTQ includes reconstruction-error minimization (REM) using a small calibration set, while QAT uses quantized forward passes during training.

3.1 Quantization Background

  • Network quantization replaces full-precision weights or activations with lower-bit representations using symmetric uniform quantization and step sizes.The quantization setup also uses asymmetric quantization for activations after self-attention and GeLU, and skips quantization for layer normalization, skip connections, biases, and the final classification head.Source: §3.1 · p. 3

3.1.1 Quantization-aware Training (QAT)

  • QAT trains with a quantized network over the training set, using either classification loss or distillation and a straight-through estimator for back-propagation.The forward pass uses quantized weights and activations, although the quantization function is non-differentiable.Source: §3.1 · p. 3, §3.1.1 · p. 3

3.1.2 Post-training Quantization (PTQ)

  • PTQ can either adjust quantization without training data or use a small unlabeled calibration set to minimize reconstruction error.In REM, the greedy procedure minimizes the distance between quantized and full-precision multiplication outputs and can be solved quickly with calibration data.Source: §3.1.2 · pp. 3–4

3.2 Why Post-training Quantization?

The comparison weighs PTQ's efficiency and limited data access against its quantized-performance gap with QAT. The figures and accompanying text illustrate these trade-offs using BERT-large on MNLI.

3.2.1 Training Time

  • QAT can take nearly four times longer than full-precision fine-tuning.The comparison attributes this to QAT iterating over the full training set for multiple epochs, with knowledge distillation further prolonging training.Source: §3.2.1 · p. 4

3.2.2 Memory Overhead

  • QAT with knowledge distillation consumes 8.3 GB more memory than full-precision fine-tuning, while PTQ can fit on a single GTX 1080 Ti.PTQ caches intermediate results during layer-wise REM.Source: §3.2.2 · p. 4

3.2.3 Data Accessibility

  • PTQ can use 1K–4K randomly sampled calibration instances, whereas QAT requires access to the entire training set.The text presents limited calibration data as reducing exposure of original training instances in third-party quantization services.Source: §3.2.3 · p. 4

3.2.4 Performance

  • QAT performance remains close to full precision across the illustrated bit-widths, while PTQ performance drops significantly, especially at low bit-widths.The passage identifies this performance drop as the main concern for REM-based PTQ.Source: §3.2.4 · p. 4

4 METHODOLOGY

The methodology extends reconstruction-error minimization to modules and trains modules in parallel to accelerate PTQ. It also uses annealed teacher forcing to address reconstruction-error propagation while distinguishing its parallel strategy from pipeline parallelism.

4.1 Module-wise Reconstruction Error Minimization

  • MREM jointly optimizes coupled linear layers within each module to reduce reconstruction error that can propagate across Transformer layers.The number of modules can be adjusted to memory constraints, trading module granularity against memory overhead.Source: §4.1 · pp. 4–5

4.2 Accelerated Parallel Training

  • Modules on separate devices sample queued inputs and train locally without waiting for their predecessors, yielding nearly the theoretical N× speed-up.Training proceeds sequentially for the first t0 steps to fill queues, then modules train in parallel with local gradient propagation and FIFO queue updates.Source: Figure 3 · p. 5, §4.2 · p. 5

4.2.1 Annealed Teaching Forcing

  • Annealed teacher forcing mixes full-precision and quantized module outputs to limit propagated reconstruction error while transitioning toward normal quantized training.The full-precision output can stop error propagation but risks forward inconsistency; λ decays from a larger early value, followed by training without teacher forcing.Source: §4.2.1 · pp. 5–6

4.2.2 Comparison with Pipeline Parallelism

  • Unlike pipeline parallelism, MREM uses local training with stale synchronous updates, giving negligible bubble time when stragglers stay within the staleness threshold.The passage says this condition can be met with balanced module partitions or a larger threshold.Source: §4.2.2 · p. 6

5 EXPERIMENTS

The experiments section outlines evaluation of MREM for PTQ and previews benchmark comparisons and analyses of design factors.

5 EXPERIMENTS

  • The experiments evaluate the proposed MREM for post-training quantization of pre-trained language models.Source: §5 · p. 6
  • The planned main results compare MREM with QAT, REM, and other quantization baselines on GLUE and SQuAD.Source: §5 · p. 6
  • Further discussions examine teacher forcing, the number of model partitions, and calibration-data size.The passage says code will be released upon acceptance.Source: §5 · p. 6

5.1 Experimental Setup

The experiments evaluate post-training quantization on GLUE and SQuAD using fine-tuned BERT models. They compare sequential and parallel MREM with QAT and REM under shared quantization settings.

5.1.1 Datasets and Metrics

  • GLUE and SQuAD evaluation uses 4,096 randomly sampled calibration instances by default, except RTE and MRPC, which use their full training sets.Experiments repeat ten times with different calibration sets and report means and standard deviations.Source: §5.1.1 · p. 6
  • Evaluation reports task-specific GLUE and SQuAD metrics, plus training time, memory overhead, calibration-set size, and average GLUE performance.MNLI includes matched and mismatched accuracy; SQuAD reports exact match and F1.Source: §5.1.1 · p. 6, Table 1 · p. 7, Table 2 · p. 7

5.1.2 Implementation

  • BERT-base and BERT-large are tested with sequential MREM-S and parallel MREM-P, using four modules across four NVIDIA V100 GPUs by default.Training runs for 2,000 steps on GLUE and 4,000 on SQuAD, with linearly decaying learning rates.Source: §5.1.2 · p. 6
  • QAT and REM are the main baselines, and all methods use the same quantization scheme; MREM is tested only once after training to protect development-set data.The study uses TWN or LAQ for 2- and 4-bit weights, and LSQ for activation quantization.Source: §5.1.2 · pp. 6–7

5.2 Main Results: Comparison with QAT and REM

MREM improves PTQ performance over REM while approaching QAT, and parallel training preserves similar performance while reducing training time. It also uses less memory than QAT and relies on a small calibration set.

5.2.1 Performance

  • MREM-S improves on REM and is closer to QAT; MREM-P remains close to or slightly below MREM-S, and can surpass it on one SQuAD v1.1 configuration.On MNLI matched accuracy with 4-bit weights, MREM-S reaches 83.5%±0.1 for BERT-base and 86.1%±0.1 for BERT-large, respectively 10.2% and 16.1% above REM and 1.1% and 0.8% below QAT.Source: §5.2.1 · p. 7

5.2.2 Training Time

  • 84 minutes for 4-bit BERT-large training on MNLI is 38× faster than QAT and 7× faster than full-precision fine-tuning; MREM-P is another 4× faster than MREM-S.Together, parallel MREM is reported as more than 150× faster than QAT.Source: §5.2.2 · pp. 7–8

5.2.3 Memory Overhead

  • MREM uses 10.8 GB on BERT-large versus QAT's 29.8 GB, while module count trades lower memory for potentially worse performance.The authors report that MREM uses around one-third of QAT's memory; longer SQuAD sequences can cause QAT memory overflow even on a 32 GB V100.Source: §5.2.3 · p. 8

5.2.4 Data Accessibility

  • MREM and REM use only 4,096 randomly sampled training instances on MNLI and SQuAD, a tiny fraction of the training data used by QAT.The paper's summary says parallel MREM further strengthens PTQ's advantages without an apparent performance drop.Source: §5.2.4 · p. 8

5.3 Main Results: Comparison with Existing Methods

The comparison finds that MREM-S and MREM-P generally outperform existing PTQ approaches and achieve results close to QAT on GLUE.

5.3 Main Results: Comparison with Existing Methods

  • The comparison includes Q-BERT, Quant-Noise, TernaryBERT, and GOBO, with results taken from their original papers.Source: §5.3 · p. 8
  • 83.5% and 83.4% are the averaged MNLI accuracies of W4-E4-A8 MREM-S and MREM-P, respectively, on par with W2/4-E8-A8 Q-BERT.Table 4 reports GLUE development-set results and average performance across tasks.Source: §5.3 · pp. 8–9
  • On MNLI-m, W2-E2-A8 MREM-S and MREM-P surpass GOBO by 11.7% and 11.3%, respectively.Source: §5.3 · p. 9

5.4 Discussions

The discussions analyze teacher forcing, REM comparisons, partitioning, calibration-set size, reconstruction-error propagation, and per-channel quantization. Unless otherwise noted, experiments use BERT-base on MNLI.

5.4.1 Teacher Forcing

  • At 250 steps, teacher forcing improves MNLI-m by 3.4% for W2-E2-A4 BERT-base and 2.8% for BERT-large; its benefit is greater with fewer steps or lower bit-width.The authors set teacher forcing to 40% of iterations; tuning it between 20% and 80% yields little difference in final performance.Source: §5.4.1 · p. 9

5.4.2 Further Comparison with REM

  • With matched training steps, REM remains inferior to MREM-S across bit-widths and takes around 9× more training time.REM's recursive caching limits it to 250 MNLI steps and 500 SQuAD steps in the reported comparison, versus 2,000 and 4,000 for MREM.Source: §5.4.2 · p. 10

5.4.3 Number of Modules and Memory Overhead

  • Fewer modules slightly improve performance but increase memory: partitions {1, 2, 3, 4, 6} use {11.9, 6.7, 4.7, 3.7, 2.7} GB, respectively.The paper chooses four modules as a trade-off.Source: §5.4.3 · p. 10, Figure 5 · p. 10

5.4.4 Size of Calibration Data

  • REM leads MREM-S/P below 128 calibration samples, but its accuracy saturates around 60%; MREM-S/P better exploit larger calibration sets, with diminishing gains beyond 4,096 samples.The default calibration-set size is 4,096.Source: §5.4.4 · p. 10, Figure 5 · p. 10

5.4.5 Reconstruction Error Propagation

  • Across layers, MREM has lower and more slowly rising reconstruction error than REM for both W2-E2-A8 and W2-E2-A4 BERT-base models.The error generally grows over the first ten layers and then decreases; the authors speculate that the classification head contributes to the later decrease.Source: §5.4.5 · p. 10, Figure 5 · p. 10

5.4.6 Per-channel Quantization

  • Per-channel quantization improves REM by around 1.0% to 2.5%, but adds only incremental gain on MREM and requires storing more full-precision step sizes.The authors therefore do not use it by default.Source: §5.4.6 · p. 10, Table 7 · p. 10

6 CONCLUSION

The paper presents MREM as an efficient approach to post-training quantization of pre-trained language models, with sequential or parallel training. Experiments report improved performance and reduced training time and memory overhead using only thousands of instances; future work includes scaling and reducing discrepancies in parallel training.

6 CONCLUSION

  • Experiments show MREM improves quantization performance while reducing training time and memory overhead with only thousands of training instances.The conclusion reports these as experimental outcomes without specifying benchmark values here.Source: §6 · p. 11
  • Parallel MREM achieves speedup close to the theoretical limit without apparent performance degradation.MREM can be conducted sequentially or in parallel.Source: §6 · p. 11
  • The paper identifies slow training, high memory overhead, and data privacy issues from QAT's access to the full training set as problems for PLM quantization.Source: §6 · p. 11
  • Future directions include scaling MREM to larger PLMs, applying parallel training to pre-training, and caching module-wise gradients to reduce discrepancy with end-to-end training.Source: §6 · p. 11
Loading 2109.15082v1…