使用多块AMD GPU通过Megatron-DeepSpeed进行大型语言模型的预训练

Pre-training a large language model with Megatron-DeepSpeed on multiple AMD GPUs

2024年1月24日，作者：Douglas Jia

在这篇博客中，我们将向你展示如何使用Megatron-DeepSpeed框架在多块AMD GPU上预训练GPT-3模型。我们还将展示如何使用你预训练的模型执行文本生成任务的推理。

什么是Megatron-DeepSpeed？

微软通过将其DeepSpeed库整合到NVIDIA的Megatron-LM框架中开发了Megatron-DeepSpeed。DeepSpeed是微软的优化库。它旨在简化和增强分布式训练和推理。DeepSpeed引入了一系列优化，使这些过程更高效、更有效。

Megatron-LM是NVIDIA的大型且强大的变换器。它能够处理庞大的模型和复杂的深度学习任务，使其成为DeepSpeed带来的进步的理想起点。

Megatron-DeepSpeed脱颖而出的是其对丰富功能的全面支持，从专家混合训练到课程学习。这使其成为应对深度学习领域各种挑战的多功能工具。使用Megatron-DeepSpeed，你可以以前所未有的效率和规模训练更大的模型。

3D 并行

Megatron-DeepSpeed 的亮点在于其 3D 并行性的实现。这种方法结合了零冗余优化器（ZeRO）分片、来自 DeepSpeed 的流水线并行性以及来自 Megatron-LM 的张量并行性。该组合使您能够高效地训练庞大的模型，从而在模型可扩展性方面开辟了新的前沿。

与 TensorParallel 一样，ZeRO 执行张量分片。ZeRO 的独特之处在于它能够在计算时及时重建整个张量，而无需对模型进行任何修改。这种创新方法还支持各种卸载技术，以应对 GPU 内存限制。

Megatron-DeepSpeed 引入了 3D 并行性的三个关键组件：

• 数据并行（DataParallel）：复制设置并并行处理数据切片，在每一步结束时同步。

• 张量并行（TensorParallel）：将张量切片分布到多个 GPU 上进行独立的并行处理，允许横向切割。

• 流水线并行（PipelineParallel）：在层级上垂直拆分模型，将其分布在多个 GPU 上，以启用不同阶段的并行处理。

为什么使用 AMD GPU？

AMD GPUs 提供了强大的开源支持，具备如 ROCm 和 HIP 等工具，使其能够轻松适应 AI 工作流。我们的高性价比适用于任何寻求经济高效解决方案以进行 AI 和深度学习任务的个人。随着 AMD 在市场上影响力的增长，越来越多的机器学习库和框架添加了对 AMD GPU 的支持。

硬件和软件要求

为了实现此任务所需的计算能力，我们使用了 AMD 加速器云 (AAC)，这是一个提供按需云计算资源和 API 的平台。在 AAC 上，我们使用了一个 PyTorch Docker 容器（版本：rocm5.7_ubuntu22.04_py3.10_pytorch_2.0.1；我们也测试过版本：rocm6.1_ubuntu22.04_py3.10_pytorch_2.1.2）并配备了 8 个 GPU。

我们的方法是硬件无关的，这意味着访问 AAC 并不是成功运行我们的代码示例的必要条件。只要您能够访问加速设备，如 GPU 或张量处理单元 (TPUs)，您应该能够在最小修改的情况下运行这些代码示例。如果您使用的是 AMD GPU，请确保正确安装了 ROCm 及其兼容版本的 PyTorch。请参考以下两个教程进行安装：

• ROCm 安装

• PyTorch 安装

GPT-3 模型的预训练代码示例

首先，安装 DeepSpeed（以及其他所需的软件包）并克隆 Megatron-DeepSpeed GitHub 仓库到本地（或服务器上）。然后，您需要下