[论文解读] Efficient Large-Scale Language Model Training on GPU Clusters Using Megatron-LM

Deepak Narayanan, Mohammad Shoeybi+10 more|arXiv (Cornell University)|Apr 9, 2021

Topic Modeling被引用 7

一句话总结

本文提出 PTD-P，一种在 Megatron-LM 中结合流水线并行、张量并行和数据并行的新方法，可在 3072 块 GPU 上高效训练万亿参数语言模型。通过引入交错流水线调度，其在 52% 的理论峰值 GPU 利用率下实现了 502 petaFLOP/s 的吞吐量，显著提升了吞吐量，同时保持了内存效率。

ABSTRACT

Large language models have led to state-of-the-art accuracies across a range of tasks. However, training these models efficiently is challenging for two reasons: a) GPU memory capacity is limited, making it impossible to fit large models on even a multi-GPU server, and b) the number of compute operations required to train these models can result in unrealistically long training times. Consequently, new methods of model parallelism such as tensor and pipeline parallelism have been proposed. Unfortunately, naive usage of these methods leads to fundamental scaling issues at thousands of GPUs, e.g., due to expensive cross-node communication or devices spending significant time waiting on other devices to make progress. In this paper, we show how different types of parallelism methods (tensor, pipeline, and data parallelism) can be composed to scale to thousands of GPUs and models with trillions of parameters. We survey techniques for pipeline parallelism and propose a novel interleaved pipeline parallelism schedule that can improve throughput by 10+% with memory footprint comparable to existing approaches. We quantitatively study the trade-offs between tensor, pipeline, and data parallelism, and provide intuition as to how to configure distributed training of a large model. Our approach allows us to perform training iterations on a model with 1 trillion parameters at 502 petaFLOP/s on 3072 GPUs with achieved per-GPU throughput of 52% of theoretical peak. Our code is open sourced at https://github.com/nvidia/megatron-lm.

研究动机与目标

• 解决大规模语言模型因超出 GPU 内存容量而需要极长训练时间的挑战。
• 实现在数千块 GPU 上训练参数超过一万亿的模型的可扩展性。
• 在保持严格优化器语义的前提下，优化流水线、张量和数据并行的组合。
• 通过一种新型调度策略，减少流水线冲洗开销，提升小批量大小下的吞吐量。
• 在大规模 GPU 集群上实现高计算效率（达到理论峰值的 52%）。

提出的方法

• 结合节点间流水线并行、节点内张量并行和数据并行（PTD-P），实现跨数千块 GPU 的训练扩展。
• 提出一种交错流水线调度，与基线方法相比，将流水线冲洗时间减少并使吞吐量提升超过 10%。
• 使用激活重计算技术减少内存占用，使更大模型能够适配 GPU 内存限制。
• 应用算子融合和优化的数据布局，以最小化内存绑定内核和通信开销。
• 采用成本模型指导模型在设备间的划分，以最小化通信量并最大化设备利用率。
• 利用高带宽互连（如 NVLink）实现服务器内部优化，以及优化的跨服务器通信，以减少多节点训练中的延迟。

实验结果

研究问题

• RQ1 如何有效组合流水线、张量和数据并行，以实现跨数千块 GPU 的训练扩展？
• RQ2 何种调度策略能最小化流水线冲洗开销并最大化吞吐量，尤其是在小批量大小下？
• RQ3 在 GPU 集群上训练万亿参数模型时，如何克服内存限制？
• RQ4 激活重计算和通信优化对训练效率有何影响？
• RQ5 所提方法在大规模 GPU 集群上能实现多大程度的高利用率（例如达到峰值的 52%）？

主要发现

• 所提出的交错流水线调度相比标准流水线方法，将训练吞吐量提升了 10% 以上。
• 该方法在 3072 块 GPU 上实现了 502 petaFLOP/s 的峰值训练吞吐量，达到理论 GPU 峰值性能的 52%。
• 使用所提出的 PTD-P 框架，训练一个 1 万亿参数模型大约需要 3 个月时间。
• 张量、流水线和数据并行的组合使训练可扩展至超越单块或多块 GPU 服务器的内存和计算极限。
• 激活重计算在不显著增加 FLOP 数量的前提下减少了内存使用，从而支持更大规模模型的训练。
• 该方法在最小化流水线同步导致的空闲时间的同时，保持了严格的优化器语义。