使用 torchtitan 在 Amazon SageMaker 上高效预训练类似 Llama 3

• 2026-01-27 13:05:29

高效的 Llama 3 类模型架构预训练：使用 Torchtitan 在 Amazon SageMaker 上的实践

Roy Allela、Trevor Harvey、Kanwaljit Khurmi、Less Wright 和 Wei Feng 于 2024 年 10 月 8 日发布于 Amazon SageMaker， 生成 AI， 中级200 永久链接评论 分享

本文与来自 Meta 的 Less Wright 和 Wei Feng 共同撰写

关键要点

通过预训练大规模语言模型 (LLM)，可以开发出理解和生成类人文本的强大 AI 系统。本文介绍了如何在 Amazon SageMaker 上使用 Torchtitan 库加速和简化 Meta Llama 3 类模型架构的预训练。我们展示了诸如 FSDP2、torchcompile 集成和 FP8 支持等关键特性，以优化训练效率，实现高达 3823 的训练速度提升。

预训练大型语言模型LLMs是开发强大 AI 系统的第一步，这些系统能够理解和生成类人文本。通过将模型暴露于大量多样化的数据中，预训练为 LLMs 学习一般语言模式、世界知识和推理能力奠定了基础。这个基础过程使 LLMs 能够在没有特定任务训练的情况下执行各种任务，从而使其高度多才多艺和适应力强。预训练对于建立强大的知识基础至关重要，这可以通过微调、迁移学习或少量学习方法进行进一步完善和专业化。

在这篇文章中，我们与 Meta 的 PyTorch 团队合作，展示了 Torchtitan 库如何加速和简化 Meta Llama 3 类模型架构的预训练过程。我们展示了 Torchtitan 的关键特性和能力，如 FSDP2、torchcompile 集成和 FP8 支持，这些都优化了训练效率。我们使用 Torchtitan 在 Amazon SageMaker 上的 p548xlarge 实例每个实例配备 8 个 Nvidia H100 GPU中预训练了一个 Meta Llama 3 8B 模型架构。我们展示了与未应用优化的基线相比，训练吞吐量提升了 3823。Amazon SageMaker 模型训练 无需管理基础设施即可减少大规模训练和调优机器学习 (ML) 模型所需的时间和成本。您可以利用当前可用的最高性能的机器学习计算基础设施，SageMaker 可以自动根据需要进行基础设施的扩展。

GitHub 上的完整代码示例 供您进一步学习。

Torchtitan 介绍

Torchtitan 是一个用于大规模 LLM 训练的参考架构，基于原生 PyTorch。其目的在于通过干净、简约的代码库展示 PyTorch 最新的分布式训练功能。该库旨在易于理解、使用和扩展，以便用于不同的训练目的，且在应用各种并行处理技术时，模型代码所需的更改最小。

Torchtitan 提供多个关键特性，包括支持 FSDP2 的每参数分片、张量并行处理、选择性层和运算符激活检查点和分布式检查点。它支持不同规模的 Meta Llama 3 类和 Llama 2 类模型架构的预训练，并包含多个数据集的配置。该库通过 TOML 文件提供简单的配置，并通过 TensorBoard 提供性能监控。接下来的部分中，我们将重点介绍 Torchtitan 的一些关键特性。

从 FSDP1 到 FSDP2 的过渡

FSDP1 和 FSDP2 是完全分片数据并行训练的两种方法。FSDP1 使用平面参数分片，将所有参数展平为 1D，然后将其连接为一个张量、填充并在工作节点间划分。这种方法提供了有界填充和高效的未分片存储，但可能并不总是允许对单个参数进行最佳分片。而 FSDP2 则将分片参数表示为在维度 0 上分片的 DTensors，单独处理每个参数。这种方法简化了参数的操作，例如每权重学习率、无通信的分片状态字典和更简单的元设备初始化。从 FSDP1 到 FSDP2 的过渡反映了在分布式训练中朝着更灵活和高效的参数处理方向的转变，解决了平面参数方法的局限性，同时可能引入新的优化机会。

Torchtitan 对 torchcompile 的支持

torchcompile 是 PyTorch 的一项关键特性，它通过最小的代码更改显著提升模型性能。通过即时编译 (JIT)，它分析和转换 PyTorch 代码为更高效的内核。Torchtitan 支持 torchcompile，在大型模型和复杂架构中提供显著的加速，使用的技术包括操作融合、内存规划和自动内核选择。这可以通过在模型的 TOML 配置文件中设置 compile = true 来实现。

Torchtitan 对 FP8 线性操作的支持

Torchtitan 提供 FP8 8 位浮点计算的支持，这显著降低了内存占用并提升了 LLM 训练的性能。FP8 拥有两种格式，E4M3 和 E5M2，各自针对训练的不同方面进行了优化。E4M3 提供更高的精度，非常适合正向传播，而 E5M2 则因其更大的动态范围，更适合反向传播。在较低精度下操作时，FP8 对模型准确性没有影响，我们通过 Meta Llama 3 8B 在 2000 步的收敛比较证明了这一点。Torchtitan 中的 FP8 支持基于 torchao 库，我们通过在模型的 TOML 配置文件中设置 enablefloat8linear = true 来启用 FP8。

银河npv免费加速器

Torchtitan 对 FP8 allgather 的支持

该特性使得 FP8 张量在多个 GPU 之间的高效通信成为可能，显著减少了相较于 bfloat16 allgather 操作的网络带宽。FP8 allgather 在 allgather 操作之前执行 float8 转换，进而减少消息体积。其效率的关键在于结合绝对最大值 (AMAX) AllReduce，这在优化器步骤之后对所有 float8 参数计算 AMAX，以避免多次小规模的 allreduce。与 FP8 支持类似，这也不会影响模型的准确性，我们通过 Meta Llama 3 8B 预训练的收敛比较再次验证了这一点。

在 Amazon SageMaker 上使用 Torchtitan 预训练 Meta Llama 3 8B

SageMaker 训练作业提供了多个关键优势，增强了使用 Torchtitan 预训练 Meta Llama 3 类模型架构的过程。它提供了一个完全托管的环境，简化了跨多个实例的大规模分布式训练，这对于高效地预训练 LLM 至关重要。SageMaker 支持自定义容器，便于无缝集成 Torchtitan 库及其依赖项，使所有必要组件均可用。

SageMaker 的内置分布式训练功能简化了多 GPU 和多节点作业的设置，减少了与这种配置通常相关的复杂性。此外，SageMaker 还与 TensorBoard 集成，实时监控和可视化训练指标，为预训练过程提供有价值的洞见。通过这些功能，研究人员和从业者可以更加专注于模型开发和优化，而不必管理基础设施，从而加速创建和完善自定义 LLM 的迭代过程。

解决方案概述

在接下来的部分中，我们将逐步引导您如何使用 Torchtitan 库准备一个自定义镜像，然后配置一个训练作业估算器函数，以在 SageMaker 上使用 c4 数据集大型清洗爬虫语料库启动 Meta Llama 3 8B 模型的预训练。c4 数据集是一个经过清洗和过滤，以删除低质量内容的大规模网络文本语料库，通常用于语言模型的预训练。

前置条件

在开始之前，请确保您具备以下条件：

一个 AWS 账户。SageMaker 域和 Amazon SageMaker Studio。有关如何创建这些的说明，请参阅 Amazon SageMaker 快速入门。一个 Hugging Face 访问令牌，以便您可以下载 Meta Llama 3 的模型和标记器以供后续使用。您需要 请求配额增加，以至少支持 SageMaker 上 1 个 mlp548xlarge 实例的训练作业。

构建 Torchtitan 自定义镜像

SageMaker BYOC自带容器 允许您使用自定义 Docker 容器来训练和部署 ML 模型。通常，SageMaker 提供自带算法和流行 ML 框架的预配置环境。但是，可能存在某些情况，您需要的独特或专有算法、依赖项或特定要求在内置选项中无法找到，需使用自定义容器。在这种情况下，我们需要使用 PyTorch、torchdata 和 torchao 包的夜间版本，以便使用 FP8 精度进行训练。

我们使用 Amazon SageMaker Studio 镜像构建便捷包，提供命令行界面 (CLI) 来简化直接从 SageMaker Studio 笔记本构建自定义容器镜像的过程。该工具免去了手动设置 Docker 构建环境的烦恼，简化了数据科学家和开发人员的工作流程。CLI 自动管理构建镜像所需的底层 AWS 服务，例如 Amazon S3Amazon S3、AWS CodeBuild 和 Amazon Elastic Container RegistryAmazon ECR，使您可以专注于 ML 任务，而无需关心基础设施设置。它提供简单的命令接口，处理 Dockerfile 和容器代码的打包，并提供用于 SageMaker 训练和托管的镜像 URI。

在开始之前，请确保您的 AWS 身份与访问管理IAM执行角色具有使用镜像构建 CLI 所需的 IAM 权限和策略。有关更多信息，请参阅 使用 Amazon SageMaker Studio 镜像构建 CLI 从 Studio 笔记本构建容器镜像。我们在 GitHub 仓库 中提供了用于构建自定义容器的 Jupyter 笔记本。

完成以下步骤以构建自定义镜像：

使用以下命令安装镜像构建包：

bash ! pip install sagemakerstudioimagebuild

为了扩展预构建镜像，您可以使用所附的深度学习库和设置，而无需从头创建镜像：

dockerfile FROM 763104351884dkrecr{REGION}amazonawscom/pytorchtraining230gpupy311cu121ubuntu2004sagemaker

接下来，指定要安装的库。您需要 PyTorch、torchdata 和 torchao 库的夜间版本：

dockerfile RUN pip install pre torch forcereinstall indexurl https//downloadpytorchorg/whl/nightly/cu121

RUN pip install pre torchdata indexurl https//downloadpytorchorg/whl/nightly

# 安装 Torchtitan 依赖 RUN pip install nocachedir datasetsgt=2190 tomligt=110 tensorboard sentencepiece tiktoken blobfile tabulate

# 安装 torchao 包以支持 FP8 RUN pip install pre torchao indexurl https//downloadpytorchorg/whl/nightly/cu121

# 显示已安装的包以供参考 RUN pip freeze

使用镜像构建 CLI 构建并推送镜像到 Amazon ECR：

bash !smdocker build repository torchtitanlatest

现在，您已准备好在 SageMaker 中使用此镜像进行 Torchtitan 模型预训练。

准备数据集可选

默认情况下，Torchtitan 库在其训练配置中使用 allenai/c4 “en” 数据集。该数据集在训练期间使用 HuggingFaceDataset 类直接流式传输。不过，您可能希望在 Amazon S3 中预训练 Meta Llama 3 模型的自定义数据集。为此，我们准备了一个 示例 Jupyter 笔记本，用于从 Hugging Face 数据集中心下载 allenai/c4 “en” 数据集到 S3 桶中。我们稍后使用 SageMaker 的 InputDataConfiguration 将数据集加载到训练实例中。您可以借助 [示例 Jupyter 笔记本](https//githubcom/awssamples