[论文解读] Cramming: Training a Language Model on a Single GPU in One Day

Jonas Geiping, Tom Goldstein|arXiv (Cornell University)|Dec 28, 2022

Topic Modeling被引用 30

一句话总结

本文在单个消费级 GPU 上从头完全训练一个基于 Transformer 的语言模型 24 小时，并评估下游 GLUE 性能，通过精心设计的流水线和数据整理实现了类似 BERT 的结果。

ABSTRACT

Recent trends in language modeling have focused on increasing performance through scaling, and have resulted in an environment where training language models is out of reach for most researchers and practitioners. While most in the community are asking how to push the limits of extreme computation, we ask the opposite question: How far can we get with a single GPU in just one day? We investigate the downstream performance achievable with a transformer-based language model trained completely from scratch with masked language modeling for a single day on a single consumer GPU. Aside from re-analyzing nearly all components of the pretraining pipeline for this scenario and providing a modified pipeline with performance close to BERT, we investigate why scaling down is hard, and which modifications actually improve performance in this scenario. We provide evidence that even in this constrained setting, performance closely follows scaling laws observed in large-compute settings. Through the lens of scaling laws, we categorize a range of recent improvements to training and architecture and discuss their merit and practical applicability (or lack thereof) for the limited compute setting.

研究动机与目标

• 研究在单个 GPU 上在 24 小时内从零开始训练一个 Transformer 语言模型能达到多高的程度。
• 在严格的计算约束下重新审视并优化预训练流水线的各个组成部分。
• 评估下游 GLUE 性能，以在极端资源条件下对比 cram（挤压式训练）与 BERT 及相关基线。
• 分析数据选择与整理在低资源 regime 下对性能的影响。

提出的方法

• 实现一个基于 PyTorch 的预训练流水线，具备自动混合精度，并仅在计算约束有正当性时使用标准组件。
• 使用 128-token 打包序列设置，采用 32768 词的 WordPiece 词汇表，以及来自 Wikipedia 与 Books Corpus 的小写、仅 ASCII 的语料库。
• 系统性地消融架构选择（如注意力偏置、嵌入方案、归一化等）和训练超参数，以确定在 24 小时预算下能带来收益的因素。
• 利用扩展定律来解释结果并指导可能的收益点，关注数据吞吐量和梯度计算效率，而非大幅度的架构改动。
• 尝试数据获取与处理（The Pile、C4 子集）以及数据过滤（基于压缩的过滤、排序），以提升下游性能。

实验结果

研究问题

• RQ1 在单个 GPU 上从头预训练一个语言模型一天内，可以达到怎样的下游 GLUE 性能？
• RQ2 在极端计算资源减少的情况下，哪些架构、训练和数据选择能带来显著提升？
• RQ3 在挤压式（低资源）情景下，大型规模训练中观察到的扩展定律如何体现？
• RQ4 在 24 小时预训练场景中，数据整理与数据集预处理对下游性能的影响程度如何？

主要发现

• 在 cramming 设置下的性能遵循与大规模计算条件相似的扩展定律，较大模型每梯度学习更多，而较小模型提供吞吐量上的优势。
• 在 24 小时预算内，重大架构重新设计带来的收益有限，而可针对性地加速梯度计算的优化在不增大模型规模的前提下带来改进。
• 仔细的数据整理与处理（例如过滤不可压缩内容、对序列进行排序、增大端批大小）可以提升下游性能，特别是数据源如 C4 与 The Pile 时。
• 在预训练阶段跳过 dropout 以在单轮预算内最大化参数更新，且在监督微调阶段保留 dropout 以实现正则化。
• 与 Izsak 等人（2021）相比，cramming 配方在 GLUE 任务上取得显著提升，在不同 GPU（RTX 2080 Ti、A4000、A6000）上实现了有竞争力的 MNLI、QQP、QNLI 和 SST-2 结果。