使用新执行算法以恒定内存训练大型神经网络
机器学习
2020-06-08 v5 分布式、并行与集群计算
机器学习
摘要
广泛流行的基于 transformer 的 NLP 模型（如 BERT 和 Turing-NLG）具有趋向数十亿参数的巨大容量。当前执行方法需要暴力资源，如 HBM 设备和用于数据并行的高速互连。本文中，我们引入一种称为 L2L（layer-to-layer，层到层）的新型中继式执行技术，在任何时刻，设备内存主要仅容纳正在执行层的占用空间。模型驻留在附属于 CPU 或 FPGA 的 DRAM 内存中，作为我们称为 eager param-server（EPS，急切参数服务器）的实体。为克服与 EPS 来回传输参数的带宽问题，模型以许多微批次逐层执行，而非传统的整个模型小批次方法。L2L 使用 16GB V100 设备实现，运行 BERT-Large 时设备批大小可达 256。我们的结果显示，与最先进基线相比内存减少 45%、吞吐量提升 40%。L2L 还能在单张 16GB V100 和 512GB CPU 内存的机器上容纳多达 500 亿参数的模型，且无需任何模型划分。L2L 可扩展至任意深度，使研究者能在负担得起的设备上开发，这是迈向 AI 民主化的重要一步。通过在主机 EPS 中运行优化器，我们展示了一种用于更快吞吐和收敛的新型混合精度形式。此外，EPS 通过跨迭代变化层来支持动态神经架构方法。最后，我们还提出并演示了 L2L 的恒定内存变体，并提出未来增强。本工作首先在 GPU 上完成，但也面向所有高 TFLOPS/Watt 加速器。
引用
@article{arxiv.2002.05645,
title = {Training Large Neural Networks with Constant Memory using a New Execution Algorithm},
author = {Bharadwaj Pudipeddi and Maral Mesmakhosroshahi and Jinwen Xi and Sujeeth Bharadwaj},
journal= {arXiv preprint arXiv:2002.05645},
year = {2020}
}