I-BERT大型模型

该模型ibert-roberta-large是RoBERTa的纯整数量化版本，由这篇论文首次提出。I-BERT采用INT8格式存储所有参数，并完全使用整数运算执行推理过程。其创新点在于将Transformer架构中的所有浮点运算（如矩阵乘法、GELU激活、Softmax和层归一化）替换为高度近似的整数运算。在Nvidia T4 GPU上的测试表明，相较于浮点运算版本，该模型能实现最高4倍的推理加速。通过量化感知微调得到的最佳模型参数可导出（例如到TensorRT），实现模型的纯整数部署。

微调流程

I-BERT微调包含三个阶段：(1)基于预训练模型在下游任务上进行全精度微调；(2)模型量化；(3)量化模型的纯整数微调（即量化感知训练）。

全精度微调

I-BERT的全精度微调与RoBERTa微调流程相似。以MRPC文本分类任务为例，可执行以下命令：

python examples/text-classification/run_glue.py \
         --model_name_or_path kssteven/ibert-roberta-large \
         --task_name MRPC \
         --do_eval \
         --do_train \
         --evaluation_strategy epoch \
         --max_seq_length 128 \
         --per_device_train_batch_size 32 \
         --save_steps 115 \
         --learning_rate 2e-5 \
         --num_train_epochs 10 \
         --output_dir $OUTPUT_DIR

模型量化

完成全精度微调后，需在检查点目录的config.json文件中将quantize属性设为true：

{                                  
  "_name_or_path": "kssteven/ibert-roberta-large",       
  "architectures": [               
    "IBertForSequenceClassification"                    
  ],                               
  "attention_probs_dropout_prob": 0.1,                  
  "bos_token_id": 0,               
  "eos_token_id": 2,               
  "finetuning_task": "mrpc",       
  "force_dequant": "none",         
  "hidden_act": "gelu",            
  "hidden_dropout_prob": 0.1,      
  "hidden_size": 768,              
  "initializer_range": 0.02,       
  "intermediate_size": 3072,       
  "layer_norm_eps": 1e-05,         
  "max_position_embeddings": 514,  
  "model_type": "ibert",           
  "num_attention_heads": 12,       
  "num_hidden_layers": 12,         
  "pad_token_id": 1,               
  "position_embedding_type": "absolute",                
  "quant_mode": true,             
  "tokenizer_class": "RobertaTokenizer",                
  "transformers_version": "4.4.0.dev0",                 
  "type_vocab_size": 1,            
  "vocab_size": 50265              
}

修改后加载检查点时，模型将自动切换至纯整数运算模式。注意需同步删除目录中的optimizer.pt、scheduler.pt和trainer_state.json文件，否则HuggingFace框架不会为后续纯整数微调重置优化器、调度器或训练状态。

纯整数微调（量化感知训练）

加载修改后的检查点即可进行纯整数微调。以下示例命令仅model_name_or_path参数有所变化：

python examples/text-classification/run_glue.py \
         --model_name_or_path $CHECKPOINT_DIR
         --task_name MRPC \
         --do_eval \
         --do_train \
         --evaluation_strategy epoch \
         --max_seq_length 128 \
         --per_device_train_batch_size 32 \
         --save_steps 115 \
         --learning_rate 1e-6 \
         --num_train_epochs 10 \
         --output_dir $OUTPUT_DIR

引用信息

若使用I-BERT，请引用我们的论文：

@article{kim2021bert,
  title={I-BERT: Integer-only BERT Quantization},
  author={Kim, Sehoon and Gholami, Amir and Yao, Zhewei and Mahoney, Michael W and Keutzer, Kurt},
  journal={arXiv preprint arXiv:2101.01321},
  year={2021}
}