如何用学术预算训练BERT模型

学术预算下的BERT模型训练

在自然语言处理领域,像BERT这样的大型语言模型已经被广泛应用。然而,预训练这些模型通常被认为是只有少数资金充足的工业实验室才能承担得起的奢侈。那么,如何在更有限的预算下训练这样的模型呢?本文将介绍一种在24小时内使用单个低端深度学习服务器预训练掩码语言模型的方法。

项目背景

本文介绍的方法基于Peter Izsak、Moshe Berchansky和Omer Levy在EMNLP 2021会议上发表的论文《How to Train BERT with an Academic Budget》。该项目的代码已在GitHub上开源,可以在IntelLabs/academic-budget-bert仓库中找到。

核心理念

这种训练方法的核心在于通过软件优化、设计选择和超参数调优的组合,在原始预训练成本的一小部分下生成与BERT-base在GLUE任务上具有竞争力的模型。主要包括以下几个方面:

时间预算训练
基于时间的学习率调度
混合精度训练
验证调度优化
检查点设置

实现细节

预训练脚本

预训练使用run_pretraining.py脚本进行。以下是使用论文中最佳配置进行训练的示例命令:

deepspeed run_pretraining.py \
  --model_type bert-mlm --tokenizer_name bert-large-uncased \
  --hidden_act gelu \
  --hidden_size 1024 \
  --num_hidden_layers 24 \
  --num_attention_heads 16 \
  --intermediate_size 4096 \
  --hidden_dropout_prob 0.1 \
  --attention_probs_dropout_prob 0.1 \
  --encoder_ln_mode pre-ln \
  --lr 1e-3 \
  --train_batch_size 4096 \
  --train_micro_batch_size_per_gpu 32 \
  --lr_schedule time \
  --curve linear \
  --warmup_proportion 0.06 \
  --gradient_clipping 0.0 \
  --optimizer_type adamw \
  --weight_decay 0.01 \
  --adam_beta1 0.9 \
  --adam_beta2 0.98 \
  --adam_eps 1e-6 \
  --total_training_time 24.0 \
  --early_exit_time_marker 24.0 \
  --dataset_path <dataset path> \
  --output_dir /tmp/training-out \
  --print_steps 100 \
  --num_epochs_between_checkpoints 10000 \
  --job_name pretraining_experiment \
  --project_name budget-bert-pretraining \
  --validation_epochs 3 \
  --validation_epochs_begin 1 \
  --validation_epochs_end 1 \
  --validation_begin_proportion 0.05 \
  --validation_end_proportion 0.01 \
  --validation_micro_batch 16 \
  --deepspeed \
  --data_loader_type dist \
  --do_validation \
  --use_early_stopping \
  --early_stop_time 180 \
  --early_stop_eval_loss 6 \
  --seed 42 \
  --fp16

时间预算训练

通过设置--total_training_time=24.0可以将预训练限制在24小时内完成。这种基于时间的训练方法允许研究人员更好地控制训练过程,特别是在计算资源有限的情况下。

基于时间的学习率调度

学习率调度示意图

学习率可以根据配置的总训练时间进行调度变化。使用--lr_schedule time来选择基于时间的学习率调度,并通过--curve=linear定义退火曲线的形状。学习率的预热阶段由--warmup_proportion指定,它表示训练会话中可用时间预算的比例。

例如,对于24小时的训练会话,warmup_proportion=0.1将占用24小时的10%,即2.4小时(或144分钟)来达到峰值学习率。之后,学习率将按计划在时间预算结束时降至0。

检查点设置

系统支持两种类型的检查点:

训练检查点:保存模型权重、优化器状态和训练参数。通过--num_epochs_between_checkpoints定义。
微调检查点:保存模型权重和配置,以便后续进行微调。通过--finetune_time_markers定义。

finetune_time_markers可以通过提供整体训练进度的时间标记列表来分配训练时间预算中的多个点。例如,--finetune_time_markers=0.5将在达到50%的训练时间预算时保存一个微调检查点。

验证调度

通过--do_validation启用预训练期间的验证。使用--validation_epochs=<num>控制验证运行之间的周期数。为了控制开始和结束时的验证运行量,可以使用validation_begin_proportion和validation_end_proportion来指定时间比例,并使用validation_epochs_begin和validation_epochs_end来相应地控制自定义值。

混合精度训练

通过添加--fp16支持混合精度训练。使用--fp16_backend=ds来使用Deepspeed的混合精度后端,或使用--fp16_backend=apex来使用apex(--fp16_opt控制优化级别)。

微调过程

使用run_glue.py对保存的检查点进行GLUE任务的微调。以下是在MRPC任务上进行微调的示例命令:

python run_glue.py \
  --model_name_or_path <path to model> \
  --task_name MRPC \
  --max_seq_length 128 \
  --output_dir /tmp/finetuning \
  --overwrite_output_dir \
  --do_train --do_eval \
  --evaluation_strategy steps \
  --per_device_train_batch_size 32 --gradient_accumulation_steps 1 \
  --per_device_eval_batch_size 32 \
  --learning_rate 5e-5 \
  --weight_decay 0.01 \
  --eval_steps 50 --evaluation_strategy steps \
  --max_grad_norm 1.0 \
  --num_train_epochs 5 \
  --lr_scheduler_type polynomial \
  --warmup_steps 50

结论

通过这种方法,研究人员可以在有限的计算资源和时间预算下训练出与BERT-base性能相当的模型。这不仅大大降低了预训练的成本,也为更多的研究团队提供了探索和改进大型语言模型的机会。

该方法的成功证明,通过精心的优化和设计,即使在资源受限的情况下,也能够实现高质量的模型训练。这对于推动自然语言处理技术在学术界和小型研究团队中的发展具有重要意义。

引用建议

如果您发现这篇论文或相关代码对您的研究有帮助,请考虑引用以下论文:

@inproceedings{izsak-etal-2021-train,
    title = "How to Train {BERT} with an Academic Budget",
    author = "Izsak, Peter  and
      Berchansky, Moshe  and
      Levy, Omer",
    booktitle = "Proceedings of the 2021 Conference on Empirical Methods in Natural Language Processing",
    month = nov,
    year = "2021",
    publisher = "Association for Computational Linguistics",
    url = "https://aclanthology.org/2021.emnlp-main.831",
}

通过这种方法,研究人员可以在有限的资源下探索和改进大型语言模型,为自然语言处理领域的进步做出贡献。🚀💡