NVIDIA NeMo Framework 是一款面向云的端到端框架，用于训练和部署具有数十亿甚至万亿参数的生成式 AI 模型。本文将介绍如何在 Amazon EKS 集群上运行分布式训练工作负载的步骤，并展示如何使用 NVIDIA NeMo Framework 提高训练效率和性能。

关键要点

• 🚀 NVIDIA NeMo Framework 的优势

  NVIDIA NeMo Framework 提供了一套完整的工具、脚本和配方，以支持 LLM 训练的每个阶段，从数据准备到训练和部署。它提供了各种定制技术，并针对语言和图像应用程序的规模化推理进行了优化，使用多 GPU 和多节点配置。NVIDIA NeMo 简化了生成式 AI 模型的开发，使其对企业更具成本效益和效率。通过提供端到端管道、高级并行技术、节省内存的策略和分布式检查点，NVIDIA NeMo 确保 AI 模型训练是简化的、可扩展的和高性能的。

  NVIDIA NeMo 的主要优势包括：

  • 端到端管道：支持数据准备、训练等不同阶段，允许对自定义数据进行即插即用。
  • 并行技术：包括数据并行、张量并行、管道并行、序列并行、专家并行和上下文并行。
  • 节省内存的技术：包括选择性激活重新计算、CPU 卸载（激活、权重）、注意力机制（Flash Attention、分组查询注意力、多查询注意力和滑动窗口注意力）和分布式优化器（Torch FSDP、分布式优化器）。
  • 数据加载器：支持不同的架构。
  • 分布式检查点：支持模型训练过程中的检查点保存和恢复。

• 🪐 Amazon EKS 作为分布式训练平台

  Amazon EKS 是一款托管的 Kubernetes 服务，使在 AWS 上运行 Kubernetes 集群变得简单直接。它管理 Kubernetes 控制平面的可用性和可扩展性，并提供计算节点自动扩展和生命周期管理支持，以帮助您运行高可用性容器应用程序。

  Amazon EKS 是运行分布式训练工作负载的理想平台，因为它与 AWS 服务的强大集成和性能特性。它与 Amazon FSx for Lustre（一种高吞吐量文件系统）无缝集成，使用 FSx CSI 驱动程序通过持久卷声明实现快速数据访问和管理。Amazon EKS 还与 Amazon CloudWatch 集成，以提供全面的日志记录和监控，提供对集群性能和资源利用率的洞察。它支持 Amazon Simple Storage Service (Amazon S3) 用于可扩展且持久的数据存储和管理，为大型数据集提供可访问性。通过 Elastic Fabric Adapter (EFA) 实现增强的网络性能，它在节点之间提供低延迟、高吞吐量连接。这些特性共同使 Amazon EKS 成为优化 AI 和机器学习 (ML) 训练工作流程的强大而高效的选择。

• 🌌 解决方案概述

  本文介绍了在 EKS 集群上运行分布式训练工作负载的步骤。主要步骤如下：

  1. 设置启用 EFA 的 2 节点 24xlarge 集群。
  2. 设置 FSx for Lustre 文件系统，以便您拥有一个共享的数据存储库来存储训练数据集和模型检查点。
  3. 设置 NVIDIA NeMo 环境。
  4. 修改 NVIDIA NeMo Kubernetes 清单以准备数据集并训练模型。

  本指南详细介绍了每个步骤，并提供代码示例和配置文件，帮助您快速开始使用 NVIDIA NeMo 和 Amazon EKS 进行分布式训练。

• 🚀 先决条件

  您需要能够启动一个基于 CPU 的 Amazon Elastic Compute Cloud (Amazon EC2) 实例，您将使用它来创建 EKS 集群。当您的实例启动并运行后，SSH 到您的 EC2 实例并安装以下 CLI：

  • 最新版本的 AWS 命令行界面 (AWS CLI)
  • kubectl
  • eksctl
  • helm

  这些步骤可能会因您使用的平台而异。请相应地查阅上述文档以在其他平台上安装 CLI。我们还要求您拥有 p4de.24xlarge 实例的容量预留，并拥有 capacityReservationID。