FasterTransformer项目中的BERT模型量化与稀疏化实践指南-优快云博客

本文链接：https://blog.youkuaiyun.com/gitblog_00263/article/details/148487440

FasterTransformer项目中的BERT模型量化与稀疏化实践指南

FasterTransformer Transformer related optimization, including BERT, GPT 项目地址: https://gitcode.com/gh_mirrors/fa/FasterTransformer

前言

在深度学习领域，模型优化一直是提升推理效率的关键。本文将详细介绍如何利用FasterTransformer项目对BERT模型进行量化(Quantization)和稀疏化(Sparsity)处理，这些技术可以显著减少模型大小并提高推理速度，同时尽可能保持模型精度。

环境准备

基础环境配置

首先需要设置几个关键环境变量：

export ROOT_DIR=</项目根目录路径>
export DATA_DIR=</数据存储目录路径>
export MODEL_DIR=</模型检查点存储路径>

量化工具安装

需要安装PyTorch量化工具包：

pip install pytorch-quantization

数据准备

下载SQuAD数据集：

bash $ROOT_DIR/data/squad/squad_download.sh

下载预训练模型及相关文件(以bert-base-uncased为例)：

wget https://s3.amazonaws.com/models.huggingface.co/bert/google/bert_uncased_L-12_H-768_A-12/pytorch_model.bin
wget https://s3.amazonaws.com/models.huggingface.co/bert/google/bert_uncased_L-12_H-768_A-12/config.json
wget https://s3.amazonaws.com/models.huggingface.co/bert/google/bert_uncased_L-12_H-768_A-12/vocab.txt

稀疏化训练

稀疏化基础

稀疏化训练通过在训练过程中强制部分权重为零来减少模型参数，特别适合NVIDIA Ampere架构GPU的Tensor Core加速。

稀疏化训练策略

推荐的三阶段训练流程：

密集预训练阶段1(序列长度128)
密集预训练阶段2(序列长度512)
稀疏预训练阶段2(使用与密集阶段2相同的超参数)

启动稀疏训练只需添加--sparse标志，并使用--dense_checkpoint进行初始化。通常，稀疏阶段训练更多步数可以获得更好的精度。

稀疏化与量化的结合

稀疏化可以与量化感知训练(QAT)结合使用，实现更高效的模型压缩。

训练后量化(PTQ)

完整流程

首先微调一个浮点密集模型：

python -m torch.distributed.launch --nproc_per_node=8 run_squad.py \
  --init_checkpoint=$MODEL_DIR/bert-base-uncased/pytorch_model.bin \
  --do_train \
  --train_file=$DATA_DIR/v1.1/train-v1.1.json \
  --train_batch_size=4 \
  --learning_rate=3e-5 \
  --num_train_epochs=2 \
  --do_predict \
  --predict_file=$DATA_DIR/v1.1/dev-v1.1.json \
  --eval_script=$DATA_DIR/v1.1/evaluate-v1.1.py \
  --do_eval \
  --do_lower_case \
  --bert_model=bert-base-uncased \
  --max_seq_length=384 \
  --doc_stride=128 \
  --vocab_file=$MODEL_DIR/bert-base-uncased/vocab.txt \
  --config_file=$MODEL_DIR/bert-base-uncased/config.json \
  --output_dir=$MODEL_DIR/bert-base-uncased-finetuned \
  --fp16 \
  --quant-disable

执行PTQ量化：

python run_squad.py \
  --init_checkpoint=$MODEL_DIR/bert-base-uncased-finetuned/pytorch_model.bin \
  --do_calib \
  --train_file=$DATA_DIR/v1.1/train-v1.1.json \
  --train_batch_size=16 \
  --num-calib-batch=16 \
  --do_predict \
  --predict_file=$DATA_DIR/v1.1/dev-v1.1.json \
  --eval_script=$DATA_DIR/v1.1/evaluate-v1.1.py \
  --do_eval \
  --do_lower_case \
  --bert_model=bert-base-uncased \
  --max_seq_length=384 \
  --doc_stride=128 \
  --output_dir=$MODEL_DIR/bert-base-uncased-PTQ-mode-2 \
  --fp16 \
  --calibrator percentile \
  --percentile 99.999 \
  --quant_mode ft2

量化模式说明

quant_mode参数与FasterTransformer中的int8_mode统一，可选值：

ft1: 模式1量化
ft2: 模式2量化
ft3: 模式3量化

不同模式在精度和性能上有不同权衡，需要根据实际需求选择。

量化感知微调(QAT)

当PTQ结果不理想时，可以采用QAT进一步恢复精度。

标准QAT流程

校准预训练模型：

python run_squad.py \
  --init_checkpoint=$MODEL_DIR/bert-base-uncased/pytorch_model.bin \
  --do_calib \
  --train_file=$DATA_DIR/v1.1/train-v1.1.json \
  --train_batch_size=16 \
  --num-calib-batch=16 \
  --output_dir=$MODEL_DIR/bert-base-uncased-calib-mode-2 \
  --fp16 \
  --calibrator percentile \
  --percentile 99.99 \
  --quant_mode ft2

执行量化感知微调：

python -m torch.distributed.launch --nproc_per_node=8 run_squad.py \
  --init_checkpoint=$MODEL_DIR/bert-base-uncased-calib-mode-2/pytorch_model.bin \
  --do_train \
  --train_file=$DATA_DIR/v1.1/train-v1.1.json \
  --train_batch_size=4 \
  --learning_rate=3e-5 \
  --num_train_epochs=2 \
  --output_dir=$MODEL_DIR/bert-base-uncased-QAT-mode-2 \
  --fp16 \
  --quant_mode ft2

结合知识蒸馏的QAT

知识蒸馏可以进一步提升量化模型的精度，通常从PTQ检查点开始：

python -m torch.distributed.launch --nproc_per_node=8 run_squad.py \
  --init_checkpoint=$MODEL_DIR/bert-base-uncased-PTQ-mode-2-for-KD/pytorch_model.bin \
  --do_train \
  --train_file=$DATA_DIR/v1.1/train-v1.1.json \
  --train_batch_size=4 \
  --learning_rate=3e-5 \
  --num_train_epochs=10 \
  --output_dir=$MODEL_DIR/bert-base-uncased-QAT-mode-2 \
  --fp16 \
  --quant_mode ft2 \
  --distillation \
  --teacher=$MODEL_DIR/bert-base-uncased-finetuned/pytorch_model.bin

结果分析

不同方法在SQuAD v1.1开发集上的典型结果：

| 方法 | Exact Match | F1 Score | |------|-------------|----------| | 浮点模型 | 82.63 | 89.53 | | PTQ模式1 | 81.92 | 89.09 | | PTQ模式2 | 80.36 | 88.09 | | QAT模式1 | 82.17 | 89.37 | | QAT模式2 | 82.02 | 89.30 | | QAT+蒸馏 | 83.67 | 90.37 |