TORCHTITAN: One-Stop PyTorch Native Solution for Production Ready LLM Pretraining

文章标题：TORCHTITAN：用于生产就绪的LLM预训练的一站式PyTorch原生解决方案
作者/机构：Wanchao Liang1, Tianyu Liu1∗, Less Wright1, Will Constable1, Andrew Gu1 Chien-Chin Huang1, Iris Zhang1, Wei Feng1, Howard Huang1, Junjie Wang1 Sanket Purandare2†, Gokul Nadathur1, Stratos Idreos2
机构：1Meta, 2Harvard University

A1 主要贡献

大规模语言模型（LLM）的训练是一项艰巨的任务，需要在并行性、计算和通信之间取得精妙平衡，同时解决复杂的内存和计算权衡问题。现有系统在集成和堆叠先进的分布式训练技术方面存在困难，导致效率低下、架构僵化、硬件利用率不足，并且缺乏对生产级训练的充分支持。这些问题的根源在于缺少统一的张量和设备抽象。

本文介绍了TORCHTITAN，这是一个开源的、PyTorch原生的分布式训练系统，旨在统一和推进最先进的技术，简化集成并减少工程开销。TORCHTITAN的核心研究贡献在于识别并将并行化与优化技术的核心原则统一到一个内聚的框架中。通过利用和扩展PyTorch的分布式张量（DTensor）和DeviceMesh，TORCHTITAN提供了一个统一的抽象，简化了并行策略的组合。它不仅仅是单个技术的集合，而是一个完整的分布式训练系统。

本文的主要贡献总结如下：
1. 增强DTensor：通过扩展其分片以支持n-D并行，增加与torch.compile的兼容性以进行编译器优化，并通过状态字典支持实现n-D模型的高效检查点，我们推进了DTensor的发展。同时，我们解决了关键错误，以增强DTensor的生产就绪性。
2. 组合并行技术：我们展示了如何组合各种并行技术，促进了在大型语言模型训练中对多维并行的探索（§2.1）。
3. 创新的软硬件协同设计：我们利用先进的硬件特性，如Float8训练和SymmetricMemory，以提高GPU效率；提供可定制的激活检查点策略以权衡内存与计算；并利用torch.compile进一步优化内存、计算和通信（§2.2）。
4. 提供生产级训练支持：通过集成可扩展且高效的分布式检查点以促进快速故障恢复，集成Flight Recorder等调试工具来调试崩溃/卡住的作业，并提供广泛的日志记录指标（§2.3）。
5. 全面的性能评估：我们在Llama 3.1系列模型上广泛评估了TORCHTITAN，通过堆叠1D到4D的并行性，在8到512个GPU的规模上展示了其弹性可扩展性、效率、收敛性和准确性。与优化基线相比，在NVIDIA H100 GPU上，Llama 3.1 8B在128 GPU规模（1D）上实现了65.08%的加速，Llama 3.1 70B在256 GPU规模（2D）上实现了12.59%的加速，Llama 3.1 405B在512 GPU规模（3D）上实现了30%的加速，并展示了4D并行在实现长上下文训练中的有效性（§3.2）。
6. 提供系统的训练配方和指南：我们提供系统的训练配方和指导方针，帮助用户应对分布式训练的复杂性，针对不同的模型大小和集群配置优化训练效率（§3.3）。

A2 方法细节

2 弹性源于可组合性

图 1: 可组合和模块化的TORCHTITAN初始化工作流。

模块化设计支持多维分片组合。TORCHTITAN以模块化的方式集成了各种并行策略，使用户能够轻松地选择和组合多维分片。这种可组合性通过简化对优化训练效率的探索，从而能够应对大规模扩展所带来的挑战。

正交的代码库组织。TORCHTITAN的代码库经过精心组织，以实现可组合性和可扩展性。我们刻意将三个主要组件分开，并使其尽可能正交：（1）与并行性无关且为可读性而设计的模型定义；（2）将并行性和训练优化应用于特定模型的并行助手；以及（3）一个通用的训练循环。所有这些组件都可以通过TOML文件进行配置，并支持命令行覆盖，从而可以轻松地在现有代码库之上添加新的模型和并行技术。

2.1 可组合的N-D并行训练

在本节中，我们将详细介绍在大型集群上扩展模型训练的整个过程，包括元设备初始化和核心的可组合多维并行技术，以展示如何在TORCHTITAN中组合这些技术，从而在不断增大的规模上高效地训练LLM。TORCHTITAN中相应的代码片段可以在附录A中找到。

2.1.1 使用元设备进行大规模模型初始化

解决模型实例化内存瓶颈。随着LLM的指数级增长，甚至在训练开始之前就出现了扩展挑战，特别是在实例化大型模型以进行分片时，很容易超出CPU或GPU的内存限制。为了解决这个问题，TORCHTITAN启用了元设备初始化，即模型首先在一个只存储元数据的元设备上创建，这使得初始化速度极快。然后，模型被分片为分布式张量（DTensors），每个参数的本地分片驻留在元设备上。最后，使用用户定义的函数执行参数初始化，确保正确的DTensor分片布局和适当的RNG种子使用。

2.1.2 完全分片数据并行

从FSDP1到FSDP2的演进。原始的完全分片数据并行（FSDP）【索引46，PyTorch FSDP: Experiences on Scaling Fully Sharded Data Parallel，2023，Proc. VLDB Endow.，https://doi.org/10.14778/3611540.3611569】是ZeRO的一种有效实现，在PyTorch中提供了大规模模型训练的能力。然而，PyTorch中的原始实现（FSDP1）由于其FlatParameter实现而存在各种限制。鉴于这些限制，TORCHTITAN集成了一个新版本的完全分片数据并行（FSDP2），该版本使用基于每个参数的分布式张量分片表示，从而提供了与模型并行技术和其他需要操作单个参数的功能更好的可组合性。

FSDP2的优势与集成。TORCHTITAN集成并利用FSDP2作为其默认的1D并行策略，得益于其改进的内存管理（通常比FSDP1的每GPU内存需求低7%）和轻微的性能提升（平均比FSDP1快1.5%）。关于FSDP2的更多细节和使用示例见附录B.1。TORCHTITAN通过嵌入适当的默认设置，包括根据您的全局大小自动分片，使得使用FSDP2变得简单。

HSDP支持更大规模扩展。为了扩展到更大的全局规模，TORCHTITAN还集成了混合分片数据并行（HSDP），它通过创建带有副本组的2D DeviceMesh来扩展FSDP2。详情见附录B.2。

2.1.3 张量并行

TP/SP作为关键模型并行技术。张量并行（TP）【索引30，Efficient large-scale language model training on gpu clusters using megatron-lm，2021，SC ’21，https://doi.org/10.1145/3458817.3476209】与序列并行（SP）【索引18 ，Reducing activation recomputation in large transformer models，2023，Proceedings of Machine Learning and Systems】是实现大规模模型训练的关键模型并行技术。TP在TORCHTITAN中使用PyTorch的RowwiseParallel和ColwiseParallel API实现，其中模型参数被分区为DTensor并进行分片计算（图3）。通过利用DTensor，TP的实现无需修改模型代码，这使得在不同模型上能够更快地启用，并提供了与本文提到的其他功能更好的可组合性。

张量与序列并行（TP/SP）的协同。虽然TP对计算最密集的部分进行分区，但序列并行（SP）对序列维度上的归一化或dropout层执行分片计算，这些层否则会产生大的复制激活张量，从而对每个GPU的内存限制构成挑战。有关TP和FSDP+TP的更多细节、图示和用法，请参见附录B.3。由于TP和SP之间的协同关系，TORCHTITAN原生将这两者捆绑在一起，它们由TP度设置共同控制。

损失并行以优化内存。在计算损失函数时，模型输出通常很大，特别是在使用TP/SP时，它们会在词汇表维度上被分片。简单地计算交叉熵损失需要收集所有分片，导致高内存使用。损失并行能够在不完全收集模型输出的情况下高效计算损失，通过最小化通信开销和启用并行的分片计算，显著减少内存消耗并提高训练速度。由于这些优势，TORCHTITAN默认实现了损失并行。

2.1.4 流水线并行

PP通过P2P通信减少开销。对于大规模预训练，TORCHTITAN采用流水线并行（PP），通过利用点对点（P2P）通信来最小化通信开销。PP将模型划分为S个阶段，每个阶段在一组独立的设备上运行。通常，每个阶段代表一个模型层或一组相邻层，但也可以包含部分层。在前向传播过程中，每个阶段接收输入激活（阶段0除外），在本地进行计算，并发送输出激活（阶段S-1除外）。最后一个阶段计算损失并发起反向传播，按相反顺序发送梯度。为提高效率，输入批次被分割成微批次，流水线调度在微批次之间重叠计算和通信。TORCHTITAN支持多种流水线调度【索引29，PipeDream: generalized pipeline parallelism for DNN training，2019，SOSP ’19，https://doi.org/10.1145/3341301.3359646；索引15 ，GPipe: efficient training of giant neural networks using pipeline parallelism，2019，Curran Associates Inc.；索引30，Efficient large-scale language model training on gpu clusters using megatron-lm，2021，SC ’21，https://doi.org/10.1145/3458817.3476209；索引37 ，Zero bubble pipeline parallelism，2023，https://arxiv.org/abs/2401.10241】。最近，TORCHTITAN增加了对新调度的支持，包括ZeroBubble 和'Flexible-Interleaved-1F1B'，利用流水线IR快速表达新调度为一系列计算操作，并依赖编译器传递来插入和优化通信操作【索引35，Training with zero-bubble Pipeline Parallelism，2024，PyTorch Forum Post】。

PP训练循环的实现。PP的训练循环与标准训练不同，它创建流水线阶段并执行调度，而不是直接调用model.forward()。由于损失是按微批次计算的，TORCHTITAN引入了一个共享的loss_fn来统一流水线和非流水线工作流，减少了代码分歧。torch.distributed.pipelining还简化了与数据并行的交互，确保归约仅在最后一个微批次之后发生，并处理分片/取消分片操作（例如，与ZeRO-3一起），以及在流水线调度执行器内透明地应用梯度缩放。有关TORCHTITAN中PP实现的更多详细信息，请参见附录B.4。

2.1.5 上下文并行

CP实现4D并行以支持长上下文。TORCHTITAN已扩展以包含上下文并行（CP）【索引22，Ring attention with blockwise Transformers for near-infinite context，2023，https://arxiv.org/abs/2310.01889；索引21 ，Blockwise parallel Transformers for large context models，2024，Advances in Neural Information Processing Systems；索引31，Megatron Core API Guide: Context Parallel，2023，https://docs.nvidia.com/megatron-core/developer-guide/latest/apiguide/context_parallel.html】，通过将CP作为现有DP、TP和PP的附加维度，实现了4D并行。CP通过在GPU之间拆分上下文维度来扩展模型训练，显著增加了最大可训练上下文长度，而不会导致内存不足（OOM）错误。例如，在8 个H100 GPU上对Llama 3.1 8B模型，使用CP能够训练上下文长度高达262,144个token，并且随着CP度的增加，MFU（模型浮点运算利用率）仅有轻微下降【索引36，Breaking barriers: Training long context llms with 1M sequence length in PyTorch using Context Parallel，2025，PyTorch Forum Post】。有关CP集成的更多详细信息，请参阅附录B.5。

2.2 优化训练效率

2.2.1 使用激活检查点导航计算-内存权衡

AC和SAC作为标准内存优化技术。激活检查点（AC）【索引7，Training Deep Nets with Sublinear Memory Cost，2016，https://arxiv.org/abs/1604.06174；索引14 ，Transcending runtime-memory tradeoffs in checkpointing by being fusion aware，2023，Proceedings of Machine Learning and Systems；索引33，µ-TWO: 3 Faster Multi-Model Training with Orchestration and Memory Optimization，2023，https://proceedings.mlsys.org/paper_files/paper/2023/file/a72071d84c001596e97a2c7e1e880559-Paper-mlsys2023.pdf】和选择性激活检查点（SAC）【索引18 ，Reducing activation recomputation in large transformer models，2023，Proceedings of Machine Learning and Systems】是标准的训练技术，通过在反向传播期间重新计算激活来换取内存节省，从而减少峰值GPU内存使用。即使在应用了多维并行之后，这些技术通常也是必需的。

TORCHTITAN提供灵活的AC/SAC选项。TORCHTITAN利用torch.utils.checkpoint在TransformerBlock级别提供灵活的AC和SAC选项。AC策略包括“完全”AC、操作级SAC和层级SAC。在TransformerBlock内，完全AC通过重新计算反向传播所需的所有激活张量来工作，而操作级SAC则保存计算密集型PyTorch操作的结果，只重新计算其他操作。层级SAC的工作方式与完全AC类似，但是包装应用于每x个TransformerBlock（x由用户指定），以实现内存和重新计算之间的可配置权衡。（详情见附录B.6。）

利用torch.compile提升效率。torch.compile在PyTorch 2中发布【索引4，PyTorch 2: Faster machine learning through dynamic python bytecode transformation and graph compilation，2024，ASPLOS ’24，https://doi.org/10.1145/3620665.3640366】，它使用TorchDynamo作为前端将PyTorch操作提取到FX图中，并使用TorchInductor作为后端将FX图编译为融合的Triton代码以提高性能。在TORCHTITAN中，我们使用区域编译，将torch.compile应用于Transformer模型中的每个独立的TransformerBlock。这有两个主要好处：（1）我们为每个区域获得一个完整的图（没有图中断），与FSDP2和TP（以及更普遍的torch.Tensor子类，如DTensor）以及其他PyTorch分布式训练技术兼容；（2）由于Llama模型将相同的TransformerBlock层一个接一个地堆叠，torch.compile可以识别出相同的结构正在被重复编译，并只编译一次，从而大大减少了编译时间。

torch.compile的可组合性优势。torch.compile通过计算融合和计算-通信重排，以一种模型无关的方式和简单的用户界面，带来了吞吐量和内存方面的效率（见第3.2节）。下面我们进一步阐述torch.compile的可组合性如何帮助TORCHTITAN通过集成诸如异步TP和Float8等高级功能，以简单的用户界面解锁硬件优化的性能增益。

2.2.3 使用异步张量并行最大化通信重叠

AsyncTP实现计算与通信重叠。默认情况下，TP在分片计算之前/之后会产生阻塞性通信，导致计算资源无法被有效利用。异步TP（AsyncTP）【索引42，Overlap communication with dependent computation via decomposition in large deep learning models，2022，ASPLOS】通过将注意力（attention）和前馈（feed-forward）模块内的TP矩阵乘法分解为更小的块，并在每个部分之间重叠通信集合，从而实现计算与通信的重叠。这种重叠是通过微流水线优化实现的，即在计算其他矩阵乘法块的同时，结果正在被通信。

基于SymmetricMemory的实现。PyTorch AsyncTP基于SymmetricMemory抽象，该抽象创建节点内缓冲区以实现更快的通信集合。这是通过在每个GPU上分配一个共享内存缓冲区来实现的，以便提供直接的点对点（P2P）访问【索引34，Introducing Async Tensor Parallelism in PyTorch，2024，PyTorch Forum Post】。

与torch.compile集成。通过TORCHTITAN与torch.compile的集成，AsyncTP可以轻松配置，以在较新的硬件（H100或更新的、节点内配备NVSwitch的GPU）上实现显著的端到端加速（详见第3.2节）。使用细节见附录B.7。

2.2.4 通过混合精度训练和Float8支持提升吞吐量

FSDP2内置混合精度支持。混合精度训练【索引27，Mixed precision training，2018，https://arxiv.org/abs/1710.03740】在确保训练稳定性的同时，节省了内存和计算资源。FSDP2内置了对基本torch.dtype的混合精度训练支持。这涵盖了以低精度（例如torch.bfloat16）执行FSDP的all-gather和计算，并以高精度（例如torch.float32）执行无损的FSDP reduce-scatter（梯度）以获得更好的数值结果的常用方法。使用细节见附录B.8。

支持先进的Float8训练。TORCHTITAN还支持更先进的Float8混合精度训练，这是一种派生数据类型，选择性地应用于线性层（在NVIDIA H100等较新硬件上可用），在确保训练稳定性的同时实现了显著的性能提升（见第3.2节报告）。来自torchao.float8的Float8功能支持多种张量级缩放策略，包括动态、延迟和静态（详见【索引28，FP8 formats for deep learning，2022，https://arxiv.org/abs/2209.05433；索引33 ，Float8 in PyTorch 1.x，2023，PyTorch Discussion Thread】），同时与其他关键的PyTorch原生系统（如autograd、torch.compile、FSDP2和TP，具备Float8 all-gather能力）兼容【索引35，Enabling Float8 all-gather in FSDP2，2024，PyTorch Forum Post】。

2.3 生产就绪的训练

为了实现生产级的训练，TORCHTITAN开箱即用地提供了与关键功能的无缝集成。这些功能包括（1）使用PyTorch分布式检查点（DCP）进行高效的检查点操作，以及（2）通过与Flight Recorder集成来调试卡住或崩溃的作业。

2.3.1 可扩展且高效的分布式检查点

检查点的重要性与挑战。检查点在训练大型语言模型中至关重要，原因有二：它们便于模型在推理和评估等应用中复用，并在发生故障时提供恢复机制。一个理想的检查点工作流程应确保在不同并行策略下易于复用，并保持高性能而不拖慢训练。典型的检查点方法有两种。第一种是将状态（模型参数和优化器状态）聚合成一个非分片的版本，该版本与并行策略无关，便于复用但需要昂贵的通信。第二种方法是每个训练器保存其本地的分片状态，这加快了过程但因嵌入了并行信息而使复用变得复杂。

DCP的解决方案。DCP通过使用DTensor解决了这些挑战，DTensor封装了全局和本地张量信息，且与并行策略无关。DCP将此信息转换为内部格式进行存储。在加载时，DCP将存储的分片与当前基于DTensor的模型参数和优化器状态进行匹配，从存储中获取所需的分片。TORCHTITAN有效地使用DCP来平衡效率和可用性。此外，DCP通过异步检查点提高了效率，它在一个单独的线程中处理存储持久化，允许此操作与后续的训练迭代重叠。与同步分布式检查点相比，TORCHTITAN利用DCP的异步检查点将Llama 3.1 8B模型的检查点开销减少了5-15倍【索引35，Optimizing checkpointing efficiency with PyTorch DCP，2024，PyTorch Forum Post】。

2.3.2 使用Flight Recorder调试作业崩溃

大规模调试的挑战。由于通信内核的异步性，大规模调试NCCL集合操作超时具有挑战性。PyTorch的Flight Recorder通过记录所有集合和p2p操作的开始、结束和入队时间，以及进程组、源/目标排名、张量大小和堆栈跟踪等元数据来解决这个问题。

Flight Recorder的诊断价值。这些数据对于诊断并行代码中的挂起问题非常宝贵。对于PP，它可以精确定位GPU上最后完成的发送或接收操作，帮助调试调度错误。对于FSDP和TP，它可以识别未能调用集合操作的排名，有助于发现PP调度或TP逻辑中的问题。

A3 设计原则

3.3 使用TORCHTITAN 4D并行进行扩展

扩展LLM的并行策略。扩展大型语言模型（LLM）需要并行策略来处理在数千个GPU上不断增长的模型大小和数据。TORCHTITAN通过可组合的4D并行实现了高效扩展。本节重点介绍使用TORCHTITAN 4D并行的一些关键观察和动机，重点关注图2中所示的特定组合。

图 2: 使用4D并行进行扩展

3.3.1 使用FSDP进行扩展

FSDP的适用性与局限。FSDP (ZeRO) 是一种适用于任何模型架构的通用技术，当通信速度快于计算时（例如，最多512个GPU），它通常作为第一级并行就足够了。然而，随着规模的扩大，集合通信的延迟会随着设备数量的增加而线性增长，从而限制了效率。为了克服这个问题，可以将像TP和PP这样的模型并行与FSDP结合起来。

3.3.2 2D并行：TP与FSDP结合

TP+FSDP的优势。张量并行（TP）通过在多个GPU之间分配工作来减少集合通信延迟，从而允许使用更小的有效批量大小，并减少大型模型或长序列的峰值内存使用。将FSDP和TP结合起来，可以在固定的问题/批量大小下实现强扩展（详情见图4）。TP还通过优化矩阵乘法的形状来提高FLOP利用率。然而，TP引入了阻塞性集合通信，并且通常仅限于节点内扩展（例如，通过NVLink），其并行度通常上限为8。扩展到超过4192个GPU需要将TP与PP结合。

3.3.3 3D并行：PP与2D并行结合

PP的通信优势。流水线并行（PP）通过以点对点的方式在各阶段之间仅传输激活和梯度，从而降低了通信带宽需求。PP在更大规模或带宽受限的集群中，对于缓解FSDP通信延迟尤其有效。PP的效率取决于流水线调度和微批量大小，这些因素会影响流水线“气泡”的大小。

3.3.4 长上下文训练与4D并行

CP支持超长上下文。上下文并行（CP）通过在GPU之间拆分上下文（序列）维度来避免OOM错误，从而实现超长上下文训练。CP主要用于长上下文训练，使模型能够捕捉更长的token相关性，从而提升整体模型质量。为了扩展序列长度，CP可以单独使用，也可以与DP一起使用。在训练大型模型或使用大量GPU时，我们可以将CP与3D并行结合，其中TP通常保持在最内层的DeviceMesh维度，而CP应用于次外层的DeviceMesh维度。

A4 实验环境

硬件配置：实验在NVIDIA H100 GPU上进行，每个GPU配备95 GiB内存。每台主机装备8个GPU，并通过NVSwitch连接。两台主机组成一个机架，连接到一个TOR交换机。后端RDMA网络连接各个TOR交换机。
数据集：所有实验均使用C4数据集（英文变体），这是一个经过清洗的大规模Common Crawl网络爬虫语料库【索引39，Exploring the limits of transfer learning with a unified text-to-text Transformer，2020，J. Mach. Learn. Res.】。
软件配置：TORCHTITAN集成了一个可检查点的数据加载器。分词器使用与Llama 3.1一同发布的官方版本（tiktoken）。

A4 实验结果

本节通过在Llama 3.1 8B、70B和405B模型上，从8个GPU扩展到512个GPU，并采用1D到4D的并行策略，展示了TORCHTITAN的弹性和可扩展性。实验展示了叠加各项优化技术（§2.2）后训练吞吐量的提升。在更高维度的并行实验中，基线总是包含之前所有已应用的技术。

实验期间内存读数稳定，吞吐量（token/秒/GPU）每10个迭代计算一次，并总是在第90个迭代时读取。由于启用Float8后，训练同时涉及BFLOAT16和FP8 Tensor Core，它们的峰值FLOPS不同，导致模型FLOPS利用率（MFU）定义不明确，因此未报告MFU。值得注意的是，在8或128个H100 GPU上，不使用Float8的1D Llama 3.1 8B模型训练实现了33%到42%的MFU。

1D并行（FSDP）在Llama 3.1 8B上的性能（8 GPUs）
- 实验设置：混合精度训练，选择性激活检查点，局部批量大小2，全局批量大小16。
- 实验结果：如表1所示，基线FSDP吞吐量为6,258 Tok/Sec。添加torch.compile后提升6.64%。再添加Float8后，总吞吐量达到9,409 Tok/Sec，相比基线提升50.35%，同时内存占用略有下降。

表 1: Llama 3.1 8B模型上的1D并行（FSDP），8个GPU。混合精度训练。选择性激活检查点。本地批量大小2，全局批量大小16。（每个GPU的统计数据）

1D并行（FSDP）在Llama 3.1 8B上的性能（128 GPUs）
- 实验设置：混合精度训练，选择性激活检查点，局部批量大小2，全局批量大小256。
- 实验结果：如表2所示，基线FSDP吞吐量为5,645 Tok/Sec。添加torch.compile后提升14.82%。再添加Float8后，总吞吐量达到9,319 Tok/Sec，相比基线提升65.08%。

表 2: Llama 3.1 8B模型上的1D并行（FSDP），128个GPU。混合精度训练。选择性激活检查点。本地批量大小2，全局批量大小256。（每个GPU的统计数据）

2D并行（FSDP+TP）在Llama 3.1 70B上的性能（256 GPUs）
- 实验设置：torch.compile + Float8，混合精度训练，完全激活检查点，FSDP度32，TP度8，局部批量大小16，全局批量大小512。
- 实验结果：如表3所示，2D并行基线吞吐量为897 Tok/Sec。添加AsyncTP后，吞吐量提升至1,010 Tok/Sec，加速12.59%，内存占用降低。

表 3: Llama 3.1 70B模型上的2D并行（FSDP + TP）+ torch.compile + Float8，256个GPU。混合精度训练。完全激活检查点。FSDP度32，TP度8。本地批量大小16，全局批量大小512。（每个GPU的统计数据）

3D并行（FSDP+TP+PP）在Llama 3.1 405B上的性能（512 GPUs）
- 实验设置：torch.compile + Float8 + AsyncTP，混合精度训练，完全激活检查点，FSDP度4，TP度8，PP度16，局部批量大小32，全局批量大小128。
- 实验结果：如表4所示，使用1F1B流水线调度时吞吐量为100 Tok/Sec。切换到Interleaved 1F1B调度后，吞吐量提升至130 Tok/Sec，加速30.00%。

表 4: Llama 3.1 405B模型上的3D并行（FSDP + TP + PP）+ torch.compile + Float8 + AsyncTP，512个GPU。混合精度训练。完全激活检查点。FSDP度4，TP度8，PP度16。本地批量大小32，全局批量大小128。（每个GPU的统计数据）

上下文并行（CP）在Llama 3.1 8B上的性能（8 GPUs）
- 实验设置：FSDP + CP + torch.compile + Float8，混合精度训练，完全激活检查点，局部批量大小1。
- 实验结果：如表5所示，通过调整FSDP和CP的并行度，成功将序列长度从32,768扩展到262,144，展示了CP在长上下文训练中的有效性。

表 5: Llama 3.1 8B模型上的FSDP + CP + torch.compile + Float8，8个GPU。混合精度训练。完全激活检查点。本地批量大小1。（每个GPU的统计数据）

4D并行在Llama 3.1 405B上的性能（512 GPUs）
- 实验设置：FSDP + TP + PP + CP + torch.compile + Float8 + AsyncTP + 1F1B，混合精度训练，完全激活检查点，TP度8，PP度8，局部批量大小8。
- 实验结果：如表6所示，通过4D并行的组合，同样成功将序列长度从32,768扩展到262,144，验证了4D并行在大型模型长上下文训练中的可组合性和有效性。

表 6: Llama 3.1 405B模型上的4D并行（FSDP + TP + PP + CP）+ torch.compile + Float8 + AsyncTP + 1F1B，512个GPU。混合精度训练。完全激活检查点。TP度8，PP度8。本地批量大小8。（每个GPU的统计数据）

关于额外的实验细节和用于验证正确性的损失收敛性测试，请参见附录B.10。

A5 结论

TORCHTITAN是一个功能强大且灵活的LLM训练框架，它实现了并行技术（FSDP、TP、PP、CP）、内存优化（Float8、激活检查点）和PyTorch编译器集成的无缝组合，从而提高了效率。其模块化设计支持不断演变的架构和硬件，通过多轴指标促进创新。

TORCHTITAN专为可解释性和生产级训练而设计，提供弹性可扩展性、全面的训练配方以及关于分布式训练策略的专家指导。实验证明，与优化基线相比，它在Llama 3.1 8B（128 GPU，1D）上加速训练65.08%，在Llama 3.1 70B（256 GPU，2D）上加速12.59%，在Llama 3.1 405B（512 GPU，3D）上加速30%，同时通过4D可组合性实现了长上下文训练。凭借其强大的功能和高效率，TORCHTITAN是应对挑战性LLM训练任务的理想一站式解决方案。

A6 附录

A 可组合的4D并行演练

高级流程。我们已经讨论了使用TORCHTITAN 4D并行进行扩展以及应用不同并行策略以将训练扩展到数千个GPU的动机。本节将逐步介绍TORCHTITAN中的4D并行代码。第一步是在元设备上创建模型实例（例如Llama模型的Transformer）。然后，我们根据pipeline_parallel_split_points配置将模型拆分为多个PP阶段来应用PP。注意，对于循环调度的PP，我们可能会从PP拆分中获得多个model_parts，其中model_parts中的每个项都是一个阶段-模型-块。接下来，在GPU上实际初始化分片模型之前，我们为每个模型部分应用SPMD风格的分布式训练技术，包括TP、激活检查点、torch.compile、FSDP和混合精度训练。

流水线并行（PP）应用。要将PP应用于模型，我们在高层次上运行以下代码。pipeline_llama_manual_split根据手动给出的pipeline_parallel_split_points配置将模型拆分为多个阶段，方法是从一个完整的模型（在元设备上）中移除未使用的模型组件。然后，build_pipeline_schedule根据pipeline_parallel_schedule配置的指示，使用torch.distributed.pipelining中的各种选项（包括1F1B【索引29，PipeDream: generalized pipeline parallelism for DNN training，2019，SOSP ’19，https://doi.org/10.1145/3341301.3359646】、GPipe【索引15 ，GPipe: efficient training of giant neural networks using pipeline parallelism，2019，Curran Associates Inc.】、交错1F1B【索引30，Efficient large-scale language model training on gpu clusters using megatron-lm，2021，SC ’21，https://doi.org/10.1145/3458817.3476209】等）来创建流水线调度。

stages, models = pipeline_llama_manual_split(
    model, pp_mesh, parallel_dims, job_config, device, model_config
)
pp_schedule = build_pipeline_schedule(job_config, stages, loss_fn)
return pp_schedule, models

张量并行（TP）与完全分片数据并行（FSDP）应用。TP和FSDP在SPMD风格的models_parallelize_fns函数中应用。为了应用TP，我们利用DTensor的parallelize_module API，通过提供一个TP“计划”作为模型参数应如何分片的指令。在下面的示例中，我们展示了分片重复的TransformerBlock的（不完整）代码。然后，我们通过包装每个单独的TransformerBlock，然后是整个模型来应用FSDP。请注意，PyTorch中的FSDP2实现支持混合精度训练。默认情况下，我们在参数的all-gather和激活计算中使用torch.bfloat16，在梯度的reduce-scatter通信和优化器更新中使用torch.float32。

mp_policy = MixedPrecisionPolicy(param_dtype, reduce_dtype)
fsdp_config = {"mesh": dp_mesh, "mp_policy": mp_policy}
for layer_id, transformer_block in model.layers.items():
    # As an optimization, do not reshard_after_forward for the last
    # TransformerBlock since FSDP would prefetch it immediately
    reshard_after_forward = int(layer_id) < len(model.layers) - 1
    fully_shard(
        transformer_block, **fsdp_config, reshard_after_forward=reshard_after_forward,
    )
fully_shard(model, **fsdp_config)

上下文并行（CP）应用。独立地，我们可以通过在Python上下文管理器下运行每个训练迭代来应用CP。

optional_context_parallel_ctx = (
    utils.create_context_parallel_ctx(
        cp_mesh=world_mesh["cp"],
        cp_buffers=[input_ids, labels] + [m.freqs_cis for m in model_parts],
        cp_seq_dims=[1, 1] + [0 for _ in model_parts],
        cp_no_restore_buffers={input_ids, labels},
        cp_rotate_method=job_config.experimental.context_parallel_rotate_method,
    )
    if parallel_dims.cp_enabled
    else None
)
with train_context(optional_context_parallel_ctx):
    pred = model(input_ids)
    loss = loss_fn(pred, labels)

B 补充材料

B.1 完全分片数据并行

FSDP2的技术优势。FSDP2通过替换原始FSDP1的FlatParameter分片方法，推进了张量分片技术。具体来说，参数现在表示为在张量维度0上分片的DTensor。这提供了与模型并行技术和其他需要操作单个参数的功能更好的可组合性，允许分片的状态字典由DTensor表示而无需任何通信，并通过DTensor提供了一个更简单的元设备初始化流程。例如，FSDP2解锁了更细粒度的张量级量化，特别是Float8张量量化，我们将在结果部分展示这一点。

FSDP2的内存与性能改进。作为从FSDP1到FSDP2重写的一部分，FSDP2通过避免使用记录流（record stream）实现了一个改进的内存管理系统。这使得内存释放具有确定性，因此相对于FSDP1，每个GPU的内存需求更低。例如，在Llama 2 7B上，FSDP2记录的GPU内存平均比FSDP1低7%。此外，通过编写高效的内核来执行多张量all-gather和reduce-scatter，FSDP2的性能与FSDP1相当，甚至略有提升——使用Llama 2 7B，FSDP2的吞吐量平均提高了1.5%。

性能提升来源。性能提升是两个小改进的结果。首先，FP32的reduce-scatter只运行一个除法内核（将本地FP32 reduce-scatter梯度预先除以全局大小，而不是分两步预除和后除以全局大小的平方根）。其次，在TORCHTITAN中，FSDP2集成的默认设置是在前向传播期间不重新分片transformer层的最后一个块，因为它将在反向传播开始时立即被重新收集。

使用方法。TORCHTITAN已将FSDP2完全集成为训练时的默认并行策略，data_parallel_shard_degree是命令行或TOML文件中的控制维度。注意，为方便使用，默认的data_parallel_shard_degree是-1，表示简单地使用所有可用的GPU，因此用户无需指定实际的全局大小。

B.2 混合分片数据并行

HSDP解决FSDP扩展瓶颈。混合分片数据并行（HSDP）是FSDP的扩展【索引45，DHEN: A deep and hierarchical ensemble network for largescale click-through rate prediction，2022，https://arxiv.org/abs/2203.11014】。在FSDP中，通信发生在FSDP组内的所有设备之间。然而，在某个点上，FSDP的通信开销会超过其相应的计算量，因为all-gather/reduce-scatter通信的延迟会随着设备数量的增加而线性增长。这导致了低的MFU，使得增加更多的GPU进行扩展变得没有价值。

HSDP的工作原理。HSDP通过创建一个包含一个维度上的副本组和另一个维度上的分片组的2-D DeviceMesh，在一定程度上解决了这个问题。每个分片组运行FSDP，而副本组运行普通的数据并行。这确保了FSDP通信发生在原始全局大小的一部分中，同时增加了跨副本组的反向梯度all-reduce。HSDP减少了FSDP的通信开销，并允许通过数据并行进一步扩展。

使用方法。TORCHTITAN通过使用两个可配置的设置：data_parallel_shard_degree和data_parallel_replicate_degree，使得实验HSDP变得容易。这两个设置控制我们创建的分片组和副本组的并行度。副本和分片度的乘积是实际的数据并行全局大小。

B.3 张量并行

TP/SP的工作机制。TP将transformer层的注意力（attention）和前馈网络（MLP）模块划分到多个设备上，使用的设备数量即为TP度。这允许多个GPU通过使用本地分片的模型参数协同处理同一个批次，代价是增加了all-reduce/all-gather/reduce-scatter操作来同步中间激活。

图 3: 张量并行详解（2个GPU，数据从左向右移动）。

TP与FSDP的结合。由于TP引入了额外的集合通信，它需要在快速网络（即NVLink）内进行。在训练LLM时，TP通常与FSDP结合使用，其中TP在节点内分片，FSDP在节点间分片，从而在不同的DeviceMesh维度上创建2D层次化分片。

图 4: FSDP2 + 张量并行（TP度为4）的分片布局，包含2个节点，每个节点4个GPU。

使用方法。由于TP和SP之间的协同关系，TORCHTITAN原生将这两者捆绑在一起，它们由命令行或TOML条目tensor_parallel_degree中的TP度设置共同控制。例如，将其设置为2意味着节点内的2个GPU将通过TP共享每个transformer层的注意力和MLP模块的计算负载，并通过序列并行共享归一化/dropout层的计算负载。损失并行是通过一个上下文管理器实现的，因为它需要控制模型前向计算之外的损失计算。它可以通过enable_loss_parallel来启用。

B.4 流水线并行

PP配置参数。我们公开了几个参数来配置PP。pipeline_parallel_degree控制参与PP的rank数量。pipeline_parallel_split_points接受一个字符串列表，表示进行拆分的层全限定名前缀。因此，流水线阶段的总数V将由这个列表的长度决定。pipeline_parallel_schedule接受要使用的调度名称。如果调度是多阶段的，那么应该有V > 1个阶段分配给每个流水线rank，否则V == 1。pipeline_parallel_microbatches控制将一个数据批次拆分成的微批次数量。

B.5 启用4D并行训练：上下文并行（CP）

CP的集成方式。为了解决上下文扩展问题，我们将上下文并行（CP）整合到了TORCHTITAN中。遵循TORCHTITAN的模块化设计原则，CP通过一个上下文管理器集成，该管理器动态地将对注意力算子（即缩放点积注意力）的调用替换为CP操作，确保不需要对模型代码进行任何更改。

CP的底层实现。在底层，CP沿着序列维度在CP设备网格上对DTensor进行分片。它扩展了DTensor调度器以处理CP特定的操作，例如环形注意力和因果注意力负载均衡，确保高效运行。通过扩展DTensor的功能以支持CP，TORCHTITAN确保CP与所有其他并行策略（FSDP、TP、PP）、优化（例如，激活检查点、torch.compile）和DCP完全兼容。这展示了TORCHTITAN模块化设计的可扩展性，它能够无缝地容纳未来的优化，同时保持性能和兼容性。

B.6 激活检查点

提供两种SAC选项。TORCHTITAN提供了两种类型的选择性激活检查点，允许在内存和重新计算之间进行更细致的权衡。具体来说，我们提供了“逐层”或“逐操作”选择性检查点的选项。逐操作的目标是释放那些重新计算速度较快的操作所使用的内存，并为那些重新计算速度较慢的操作保存中间结果（内存），从而实现更有效的吞吐量/内存权衡。

使用方法。AC通过命令行或TOML文件中的两行设置启用。具体来说，mode可以是none、selective或full。当设置为selective时，会使用下一个配置selective_ac_type，它可以是一个正整数以启用选择性层检查点，也可以是op以启用选择性操作检查点。逐层检查点接受一个整数输入来指导检查点策略，其中1 = 检查点每一层（与full相同），2 = 每隔一层检查点，3 = 每隔两层检查点，依此类推。逐操作检查点由parallelize_llama.py中的_save_list策略驱动，该策略标记了高算术强度的操作，如matmul（矩阵乘法）和SPDA（缩放点积注意力）以保存中间结果，同时允许其他较低强度的操作被重新计算。注意，为了平衡总吞吐量，只有每隔一个matmul被标记为保存。

B.7 异步张量并行（AsyncTP）

SymmetricMemory机制。AsyncTP中使用的SymmetricMemory集合通信比标准的NCCL集合通信更快，其工作原理是让每个GPU分配一个相同的内存缓冲区，以提供直接的点对点（P2P）访问。SymmetricMemory依赖于节点内的NVSwitch，因此通常只适用于H100或更新的GPU。

使用方法。AsyncTP在TORCHTITAN TOML配置文件的实验部分中启用，通过enable_async_tensor_parallel布尔设置来开启或关闭。

B.8 在TORCHTITAN中定制FSDP2混合精度

混合精度策略配置。混合精度由apply_fsdp函数中的MixedPrecisionPolicy类控制，该类在TORCHTITAN中默认使用BF16作为param_dtype，并默认使用FP32作为reduce_dtype。reduce_dtype为FP32意味着在反向传播中用于梯度计算的reduce-scatter操作将在FP32中进行，以帮助最大化梯度更新的稳定性和精度。

B.9 TORCHTITAN：全面的功能集和降低的复杂性

TORCHTITAN实现新设计。TORCHTITAN广泛的功能集和宽泛的设计空间覆盖是由其统一的设计原则驱动的，即模块化、可组合性和可扩展性。利用这些原则，TORCHTITAN无缝集成了多种并行策略（FSDP、TP、PP和CP）和优化（例如，SAC、Float8训练）。这个统一的框架不仅支持先进的流水线调度和多维并行，还简化了新技术的集成，使其非常适用于前沿研究和生产级部署。

功能对比。下表重点介绍了TORCHTITAN在并行性、检查点和编译器支持方面与Megatron-LM、DeepSpeed和veScale的对比：

表 7: TORCHTITAN与Megatron-LM、DeepSpeed和veScale在并行性、编译器支持、激活检查点和模型检查点方面的比较。

代码复杂性与可维护性。TORCHTITAN的设计原则也使其代码复杂性显著降低。尽管提供了丰富的功能集，TORCHTITAN仍保持了一个紧凑和模块化的代码库，使其更易于扩展、维护和演进，同时确保高性能。下表比较了TORCHTITAN与Megatron-LM和DeepSpeed的代码行数（LOC）：

表 8: 各系统代码行数（LOC）比较。

B.10 扩展实验分析：性能和损失收敛

性能实验目标。我们在3.2节中的实验服务于多个目标：
* 建立可组合性和模块性：TORCHTITAN展示了各种并行和优化技术的无缝集成。
* 展示性能提升：在各种并行和优化中观察到显著的加速。
* 验证弹性可扩展性：TORCHTITAN能有效地随模型大小和GPU数量扩展。
* 消融研究：呈现了单个技术的详细性能增益。

各表格的具体分析：
* 表1：突出了编译器支持相对于eager执行的改进，以及Float8训练带来的进一步增益。
* 表2：展示了随着GPU数量增加，早期获得的增益如何扩展。
* 表3：显示了AsyncTP（一种软硬件协同设计技术）在2D训练结合torch.compile和Float8训练的基础上实现的加速。
* 表4：量化了在AsyncTP、torch.compile和Float8训练的基础上，Interleaved 1F1B调度相对于1F1B的优势。
* 表5：展示了CP在实现长上下文训练方面的有效性，即使在小规模下也是如此。
* 表6：展示了4D并行的可组合性，以及CP在实现大规模长上下文训练方面的有效性。

ZeRO变体选择。对于FSDP，除了涉及PP的实验使用ZeRO-2变体外，所有实验都使用ZeRO-3变体。这种区别是由于ZeRO-3在PP中效率低下，它会为每个微批次产生额外的all-gather调用。相比之下，ZeRO-2只在第一个微批次时收集一次参数，并在最后一个微批次的反向传播后重新分片。

损失收敛性验证。TORCHTITAN的设计原则影响了PyTorch分布式库中高级分布式训练功能（如FSDP2、AsyncTP、PP和CP）的开发。在这些贡献中，我们确保了单个技术以及它们各种并行和优化组合的损失收敛性。例如，下面是一系列覆盖并行和训练优化的损失收敛测试。我们使用“FSDP 8”表示FSDP度为8的实验，“FSDP 8, CP 8”表示在64个GPU上FSDP度为8、CP度为8的实验，以此类推。我们假设FSDP的正确性，这可以通过与DDP甚至单设备作业进行比较来进一步验证。

表 9: 损失收敛测试设置。

图 5: Llama 3.1 8B上的损失收敛测试。C4数据集。本地批量大小4，全局批量大小32。3000步，600个预热步。