在8GB消费级GPU上从零开始训练BERT的注意事项
原文作者：sidneyradcliffe
我在自己的台式电脑（配备 Nvidia 3060 Ti 8GB GPU）上从零开始训练了一个 BERT 模型（Devlin 等人，2019 年）。模型架构、分词器和训练器都来自 Hugging Face 库，我的主要工作是编写代码、设置数据（大约 20GB 的未压缩文本），然后让我的这台电脑运行。(并确保其工作正常，GPU 利用率高）。
- 代码以 Jupyter notebook的形式提供，点击这里获取。
- 这些数据可以在 Hugging Face 的数据集中找到。
通常情况下，训练大语言模型需要使用 GPU 或 TPU 集群，并非仅仅一台台式电脑可以进行。下图说明了我用于训练该模型的计算资源与训练原始 BERT-base 模型的资源之间的差异。
尽管 BERT-base 和这个模型的训练时间相同，但 BERT-base 训练 token 数量大约是这个模型的 30 倍（BERT-base 训练了大约 40 个epoch，而这个模型只训练了一个epoch）。
上图显示了 GLUE dev-set 分数，以说明该模型在自然语言任务中的表现。对 GLUE 的微调总共花费了约 12 个小时（在 4 天/约 100 小时的预训练基础上）。下表列出了 GLUE-dev 的详细结果：
| Model | MNLI (m/mm) | SST-2 | STSB | RTE | QNLI | QQP | MRPC | CoLA | Average |
|---|---|---|---|---|---|---|---|---|---|
| This model | 79.3/80.1 | 89.1 | 61.9 | 55.9 | 86.3 | 86.4 | 74.8 | 41.0 | 72.7 |
| BERT-Base* | 83.2/83.4 | 91.9 | 86.7 | 59.2 | 90.6 | 87.7 | 89.3 | 56.5 | 80.9 |
*BERT-Base 指的是经过 fully trained 的 BERT 模型，其结果来自 Cramming（Geiping 等人，2022）。
虽然我们可以看到，BERT-Base 在每个任务中的表现都更好；但 “This model” 的结果在 2018 年初可能会非常出色（可能在某些任务上达到了最先进水平）。
没有进行超参数调整（ hyperparameter tuning ）。也没有使用特殊技术来改进训练。优化器（Optimizer）和学习率调度（learning rate schedule）以 Cramming（Geiping 等人，2022 年）为指导，但没有使用 Cramming 中的模型架构变化和其他建议。我先做了几次规模较小的训练（约 1-12 小时）。
通过 Weights & Biases，我能够进行远程监控训练。
这项工作受到了 Cramming（[Geiping et al，2022 年]((arxiv.org/abs/2212.14…)）的启发，该论文介绍了如何在有限的计算资源下（仅用 24 小时）训练出性能良好的 BERT 模型。
100 小时训练运行的曲线图
预训练过程中的loss值变化
Cramming（Geiping 等人，2022）推荐的学习率调度
GPU 利用率约为 98%
GPU 内存使用率约为 98%，这是通过调整 batch size 实现的
GPU 的温度保持在 76 – 80 摄氏度之间，高温天的温度更高
References:
- Geiping, Jonas, and Tom Goldstein. “Cramming: Training a Language Model on a Single GPU in One Day.” arXiv preprint arXiv:2212.14034 (2022). URL arxiv.org/abs/2212.14….
- Jacob Devlin, Ming-Wei Chang, Kenton Lee, and Kristina Toutanova. BERT: Pre-training of Deep Bidirectional Transformers for Language Understanding. arXiv:1810.04805 [cs], May 2019. URL arxiv.org/abs/1810.04….
- Vaswani et al. (2017) Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N. Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention Is All You Need. arXiv:1706.03762 [cs], December 2017. URL arxiv.org/abs/1706.03….
- Alec Radford, Karthik Narasimhan, Tim Salimans, and Ilya Sutskever. 2018. Improving language understanding with unsupervised learning. Technical report, OpenAI, s3-us-west-2.amazonaws.com/openai-asse…
