devbox
码创造者
这个屌丝很懒，什么也没留下！
关注作者
热门标签
热门文章
当前位置:   article > 正文

训练BERT，我只花了一半的时间

作者：码创造者 | 2024-08-15 12:48:35

赞

踩

bert 训练慢

相信很多人都知道Hugging Face，也都用过它的Transformers预训练语言模型，但你们有没有觉得它训练的有点太慢了呢？

这时候，字节第二快的男人要站出来了（第一快是我mentor），手把手教你怎么让训练时间缩短一半。

训练BERT

首先我们要安装Transformers库，这很简单：

pip install transformers

然后我们直接把官方的例子拷贝下来，这里我们用的是GLUE任务，地址是github.com/huggingface/。因为代码太长了，这里就不放了，拷贝下来后文件名是run_glue.py。

接着我们就可以直接运行这个代码了，我们采用mrpc数据集，开启FP16训练，命令如下：

  1. python run_glue.py \
  2.   --model_name_or_path bert-base-cased \
  3.   --task_name mrpc \
  4.   --do_train \
  5.   --do_eval \
  6.   --max_seq_length 128 \
  7.   --per_device_train_batch_size 32 \
  8.   --num_train_epochs 3 \
  9.   --output_dir /tmp/mrpc/ \
  10.   --overwrite_output_dir \
  11.   --fp16

我这里是单卡训练的，训练完后输出如下：

  1. ***** train metrics *****
  2.   epoch                    =        3.0
  3.   train_loss               =     0.3921
  4.   train_runtime            = 0:00:45.06
  5.   train_samples            =       3668
  6.   train_samples_per_second =    244.166
  7.   train_steps_per_second   =      7.655

可以看出，训练总共耗时「45秒」，是不是有点等不及了呢？

加速训练

首先我们需要安装训练加速库，这里我们用到的是LightSeq，项目地址是github.com/bytedance/li。不过我们还是直接pip安装：

pip install lightseq

然后我们需要做的就是将Hugging Face的BERT替换成LightSeq的BERT，代码如下，放在文件replace_module.py中。

  1. from lightseq.training.ops.pytorch.transformer_encoder_layer import (
  2.     LSTransformerEncoderLayer,
  3. )
  4.  
  5. class LSHFTransformerEncoderLayer(LSTransformerEncoderLayer):
  6.     def __init__(self, *args, **kwargs):
  7.         super(LSHFTransformerEncoderLayer, self).__init__(*args, **kwargs)
  8.  
  9.     def forward(self, hidden_states, encoder_padding_mask, *args, **kwargs):
  10.         encoder_padding_mask /= -10000.0
  11.         output = super().forward(hidden_states, encoder_padding_mask)
  12.         return (output, None, None, None)
  13.  
  14. def gen_ls_bert_config(training_args, config):
  15.     bert_config = LSTransformerEncoderLayer.get_config(
  16.         max_batch_tokens=4096,
  17.         max_seq_len=config.max_position_embeddings,
  18.         hidden_size=config.hidden_size,
  19.         intermediate_size=config.intermediate_size,
  20.         nhead=config.num_attention_heads,
  21.         attn_prob_dropout_ratio=config.attention_probs_dropout_prob,
  22.         activation_dropout_ratio=0.1,
  23.         hidden_dropout_ratio=config.hidden_dropout_prob,
  24.         pre_layer_norm=False,
  25.         fp16=training_args.fp16,
  26.         local_rank=training_args.local_rank,
  27.     )
  28.     return bert_config
  29.  
  30. def inject_ls_enc_layer(model, training_args, config):
  31.     for i in range(config.num_hidden_layers):
  32.         bert_config = gen_ls_bert_config(training_args, config)
  33.         model.bert.encoder.layer[i] = LSHFTransformerEncoderLayer(bert_config)

这里LSHFTransformerEncoderLayer是继承的LightSeq中的LSTransformerEncoderLayer类，然后重写了forward函数。原因是Hugging Face的输入格式和LightSeq略有不同，需要在forward之前转换一下。

gen_ls_bert_config函数是用来定义LightSeq的encoder参数配置，这里直接从Hugging Face的主函数入口获取即可。

inject_ls_enc_layer函数就是用来替换BERT中的每一层encoder的，首先定义每一层的参数配置，然后用LSHFTransformerEncoderLayer类去替换原始的encoder层即可。

然后我们打开run_glue.py，在头文件处加上inject_ls_enc_layer的引用：

from replace_module import inject_ls_enc_layer

最后在定义完model后，将model中的encoder替换即可，利用上面引用的替换函数：

  1. model = AutoModelForSequenceClassification.from_pretrained(
  2.     model_args.model_name_or_path,
  3.     from_tf=bool(".ckpt" in model_args.model_name_or_path),
  4.     config=config,
  5.     cache_dir=model_args.cache_dir,
  6.     revision=model_args.model_revision,
  7.     use_auth_token=True if model_args.use_auth_token else None,
  8. )
  9.  
  10. # 在model定义后立刻替换
  11. inject_ls_enc_layer(model, training_args, config)

我们重新运行上一次运行的命令：

  1. python run_glue.py \
  2.   --model_name_or_path bert-base-cased \
  3.   --task_name mrpc \
  4.   --do_train \
  5.   --do_eval \
  6.   --max_seq_length 128 \
  7.   --per_device_train_batch_size 32 \
  8.   --num_train_epochs 3 \
  9.   --output_dir /tmp/mrpc/ \
  10.   --overwrite_output_dir \
  11.   --fp16

最终输出如下：

  1. ***** train metrics *****
  2.   epoch                    =        3.0
  3.   train_loss               =     0.6077
  4.   train_runtime            = 0:00:25.08
  5.   train_samples            =       3668
  6.   train_samples_per_second =    438.603
  7.   train_steps_per_second   =     13.751

这次运行时间只有「25秒」！不愧是字节最快的男人。

加载预训练参数

有眼尖的小伙伴可能发现了，上面加速后效果变差了呀。没错，因为新建了encoder类之后，参数都是随机初始化的了，所以要重新加载一下预训练参数。

LightSeq的encoder类初始化的时候提供了预训练参数初始化的选项，我们只需要将预训练参数从Hugging Face的BERT中提取出来即可：

  1. def get_hf_bert_enc_layer_params(layer):
  2.     init_ws = []
  3.     init_bs = []
  4.  
  5.     init_ws.append(layer.attention.self.query.weight.detach().clone())
  6.     init_bs.append(layer.attention.self.query.bias.detach().clone())
  7.     init_ws.append(layer.attention.self.key.weight.detach().clone())
  8.     init_bs.append(layer.attention.self.key.bias.detach().clone())
  9.     init_ws.append(layer.attention.self.value.weight.detach().clone())
  10.     init_bs.append(layer.attention.self.value.bias.detach().clone())
  11.     init_ws.append(layer.attention.output.dense.weight.detach().clone())
  12.     init_bs.append(layer.attention.output.dense.bias.detach().clone())
  13.     init_ws.append(layer.attention.output.LayerNorm.weight.detach().clone())
  14.     init_bs.append(layer.attention.output.LayerNorm.bias.detach().clone())
  15.  
  16.     init_ws.append(layer.intermediate.dense.weight.detach().clone())
  17.     init_bs.append(layer.intermediate.dense.bias.detach().clone())
  18.     init_ws.append(layer.output.dense.weight.detach().clone())
  19.     init_bs.append(layer.output.dense.bias.detach().clone())
  20.     init_ws.append(layer.output.LayerNorm.weight.detach().clone())
  21.     init_bs.append(layer.output.LayerNorm.bias.detach().clone())
  22.  
  23.     return init_ws, init_bs

注意参数在列表中的顺序不能错了，然后将这两个列表加入到LSHFTransformerEncoderLayer类的初始化参数中去：

  1. def inject_ls_enc_layer(model, training_args, config):
  2.     for i in range(config.num_hidden_layers):
  3.         bert_config = gen_ls_bert_config(training_args, config)
  4.         # 提取预训练参数
  5.         init_ws, init_bs = get_hf_bert_enc_layer_params(model.bert.encoder.layer[i])
  6.         # 利用预训练参数进行初始化
  7.         model.bert.encoder.layer[i] = LSHFTransformerEncoderLayer(
  8.             bert_config, init_ws, init_bs
  9.         )

接着运行命令不变，效果就上来啦。

和竞品比如何？

另一款知名的训练加速库DeepSpeed你们可能也听过，那和它比速度怎么样呢？

Hugging Face已经内置了DeepSpeed，可以直接开启。不过它并没有替换掉encoder，所以模型还是用PyTorch写的，速度依然很慢。因此我们需要手动替换一下encoder。

代码和上面类似，也是定义参数配置和encoder类：

  1. from deepspeed.ops.transformer import (
  2.     DeepSpeedTransformerConfig,
  3.     DeepSpeedTransformerLayer
  4. )
  5.  
  6. def gen_ds_bert_config(training_args, config):
  7.     bert_config = DeepSpeedTransformerConfig(
  8.         batch_size=4096,
  9.         hidden_size=config.hidden_size,
  10.         intermediate_size=config.intermediate_size,
  11.         heads=config.num_attention_heads,
  12.         attn_dropout_ratio=config.attention_probs_dropout_prob,
  13.         hidden_dropout_ratio=config.hidden_dropout_prob,
  14.         num_hidden_layers=config.num_hidden_layers,
  15.         initializer_range=0.02,
  16.         layer_norm_eps=1e-8,
  17.         local_rank=training_args.local_rank,
  18.         fp16=training_args.fp16,
  19.         pre_layer_norm=False,
  20.         huggingface=True,
  21.         training=True
  22.     )
  23.     return bert_config
  24.  
  25. def inject_ds_enc_layer(model, training_args, config):
  26.     for i in range(config.num_hidden_layers):
  27.         bert_config = gen_ds_bert_config(training_args, config)
  28.         model.bert.encoder.layer[i] = DeepSpeedTransformerLayer(bert_config)

然后在run_glue.py里引用inject_ds_enc_layer替换函数，并对model进行替换：

  1. from replace_module import inject_ds_enc_layer
  2.  
  3. model = AutoModelForSequenceClassification.from_pretrained(
  4.     model_args.model_name_or_path,
  5.     from_tf=bool(".ckpt" in model_args.model_name_or_path),
  6.     config=config,
  7.     cache_dir=model_args.cache_dir,
  8.     revision=model_args.model_revision,
  9.     use_auth_token=True if model_args.use_auth_token else None,
  10. )
  11.  
  12. # 在model定义后立刻替换
  13. inject_ds_enc_layer(model, training_args, config)

最后我们还需要定义一个DeepSpeed需要用到的运行参数配置ds_config.json：

  1. {
  2.   "train_micro_batch_size_per_gpu": "auto",
  3.   "optimizer": {
  4.     "type": "AdamW",
  5.     "params": {
  6.       "lr": "auto",
  7.       "betas": [
  8.         0.9,
  9.         0.999
  10.       ],
  11.       "eps": 1e-8,
  12.       "weight_decay": "auto",
  13.       "torch_adam": true
  14.     }
  15.   },
  16.   "scheduler": {
  17.     "type": "WarmupDecayLR",
  18.     "params": {
  19.       "warmup_num_steps": "auto",
  20.       "warmup_min_lr": "auto",
  21.       "warmup_max_lr": "auto",
  22.       "total_num_steps": "auto"
  23.     }
  24.   },
  25.   "gradient_clipping": "auto",
  26.   "fp16": {
  27.     "enabled": "auto",
  28.     "loss_scale": 0,
  29.     "initial_scale_power": 7
  30.   }
  31. }

运行命令需要稍稍修改，采用DeepSpeed的启动器：

  1. deepspeed --num_gpus=1 run_glue.py \
  2.   --model_name_or_path bert-base-cased \
  3.   --task_name mrpc \
  4.   --do_train \
  5.   --do_eval \
  6.   --max_seq_length 128 \
  7.   --per_device_train_batch_size 32 \
  8.   --num_train_epochs 3 \
  9.   --output_dir /tmp/mrpc/ \
  10.   --overwrite_output_dir \
  11.   --fp16 \
  12.   --deepspeed ds_config.json

输出结果如下：

  1. ***** train metrics *****
  2.   epoch                    =        3.0
  3.   train_loss               =     0.5865
  4.   train_runtime            = 0:00:37.17
  5.   train_samples            =       3668
  6.   train_samples_per_second =    296.032
  7.   train_steps_per_second   =      9.281

发现DeepSpeed用了整整「37秒」才训练完，和LightSeq的「25秒」相比还是有差距的。

总结

最终对比下来，Hugging Face花了「45秒」训练完成，DeepSpeed花了「37秒」，而LightSeq只花了「25秒」。

「项目地址：」

bytedance/lightseq

「技术原理：」

godweiyang：训练加速3倍！字节跳动推出业界首个NLP模型全流程加速引擎

「其它使用例子：」

godweiyang：只用几行代码，我让模型『训练』加速了3倍以上！

如果你对字节的技术比较感兴趣，欢迎加入我们，一起开发牛X的项目，做最快的男人。

「我的内推码：」
A7FSJMK
「内推链接：」

加入字节跳动

声明：本文内容由网友自发贡献，不代表【wpsshop博客】立场，版权归原作者所有，本站不承担相应法律责任。如您发现有侵权的内容，请联系我们。转载请注明出处：https://www.wpsshop.cn/w/码创造者/article/detail/983656
推荐阅读
相关标签

Copyright © 2003-2013 www.wpsshop.cn 版权所有，并保留所有权利。

  

闽ICP备14008679号