减少大型 Transformer 模型中的激活重计算
机器学习
2022-05-12 v1 计算与语言
摘要
训练大型 transformer 模型是现代 AI 最重要的计算挑战之一。在本文中，我们展示如何通过减少激活重计算来显著加速大型 transformer 模型的训练。激活重计算常用于应对内存容量限制。传统上不存储用于反向传播的激活，而是重新计算，这节省了内存但增加了冗余计算。在本工作中，我们展示大部分此类冗余计算是不必要的，因为我们可以在不依赖它的情况下充分降低内存消耗。我们提出两种新颖却非常简单的技术：序列并行与选择性激活重计算。结合张量并行，这些技术几乎消除了重计算激活的需求。我们在规模达一万亿参数的语言模型上评估了我们的方法，并表明我们的方法将激活内存减少 5 倍，同时将激活重计算带来的执行时间开销降低逾 90%。例如，当在 2240 块 NVIDIA A100 GPU 上训练 530B 参数的 GPT-3 风格模型时，我们实现了 54.2% 的模型浮点利用率（Model Flops Utilization），比使用重计算时所达 42.1% 快 29%。我们的实现将在 Megatron-LM 与 NeMo-Megatron 中提供。
引用
@article{arxiv.2205.05198,
title = {Reducing Activation Recomputation in Large Transformer Models},
author = {Vijay Korthikanti and Jared Casper and Sangkug Lym and Lawrence McAfee and Michael Andersch and Mohammad Shoeybi and Bryan Catanzaro},
journal= {arXiv preprint arXiv:2205.05198},
year = {2022}
}