Memory-Efficient Pipeline-Parallel DNN Training
Memory-Efficient Pipeline-Parallel DNN Training
复制标题
DOI：
--
复制
发表时间：
2020-06
期刊：
影响因子：
--
通讯作者：
D. Narayanan;Amar Phanishayee;Kaiyu Shi;Xie Chen;M. Zaharia
中科院分区：
文献类型：
--
作者：
D. Narayanan;Amar Phanishayee;Kaiyu Shi;Xie Chen;M. Zaharia
NLP 和计算机视觉等领域的许多最先进的成果都是通过扩大现有模型中的参数数量获得的。然而，这些大型模型的权重参数和中间输出通常无法容纳在单个加速器设备的主存储器中；这意味着需要使用多个加速器来训练大型模型，而以高效的方式实现这一点具有挑战性。在这项工作中，我们提出了 PipeDream-2BW，这是一个执行内存高效管道并行性的系统，这是一种将数据和模型并行性与输入管道相结合的混合形式的并行性。我们的系统使用新颖的流水线和权重梯度合并策略，结合权重的双缓冲，以确保高吞吐量、低内存占用和类似于数据并行性的权重更新语义。此外，PipeDream-2BW 自动在可用的硬件资源上对模型进行分区，同时了解计算能力、内存容量和互连拓扑等约束，并确定何时采用现有的内存节省技术，例如激活重新计算，以牺牲额外的计算来降低内存占用。 PipeDream-2BW 能够将具有多达 25 亿个参数的大型语言模型的训练速度比优化基线提高高达 6.9 倍。
Many state-of-the-art results in domains such as NLP and computer vision have been obtained by scaling up the number of parameters in existing models. However, the weight parameters and intermediate outputs of these large models often do not fit in the main memory of a single accelerator device; this means that it is necessary to use multiple accelerators to train large models, which is challenging to do in a time-efficient way. In this work, we propose PipeDream-2BW, a system that performs memory-efficient pipeline parallelism, a hybrid form of parallelism that combines data and model parallelism with input pipelining. Our system uses a novel pipelining and weight gradient coalescing strategy, combined with the double buffering of weights, to ensure high throughput, low memory footprint, and weight update semantics similar to data parallelism. In addition, PipeDream-2BW automatically partitions the model over the available hardware resources, while being cognizant of constraints such as compute capabilities, memory capacities, and interconnect topologies, and determines when to employ existing memory-savings techniques, such as activation recomputation, that trade off extra computation for lower memory footprint. PipeDream-2BW is able to accelerate the training of large language models with up to 2.5 billion parameters by up to 6.9x compared to optimized baselines.