I-BERT基础模型

ibert-roberta-base模型是RoBERTa的纯整数量化版本，相关论文详见此处。该模型采用INT8格式存储所有参数，并完全使用整数运算进行推理。特别地，I-BERT将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-base \
         --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-base",       
  "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}
}