揭秘大模型的训练方法：使用PyTorch进行超大规模深度学习模型训练

作者

我相信我们每个人都已经快要被 LLM（大语言模型）研究界的节奏“卷”麻了。几乎每天都会有一款新的 SOTA（最先进）模型横空出世，打破现有基准。如果你曾好奇，这样的技术爆发背后到底是什么驱动力——其实很大程度上，是因为研究人员能够在超大规模集群上进行训练和验证，而这一切都要归功于并行化（parallelism）。

你听说过“5D并行”吗？这个术语最早由 Meta AI 在其论文《The Llama 3 Herd of Models》（《Llama 3 模型群》）中提出。传统意义上，它指的是将数据并行（data parallelism）、张量并行（tensor parallelism）、上下文并行（context parallelism）、流水线并行（pipeline parallelism）以及专家并行（expert parallelism）结合使用的技术体系。而最近，一种新的范式——ZeRO（Zero Redundancy Optimizer）——的出现彻底改变了游戏规则：通过减少分布式计算中的冗余来优化内存使用。每种并行方式都针对训练过程中的不同挑战点，而当它们结合在一起使用时，就可以支撑起拥有数十亿、甚至万亿参数的超大模型。

我此前分享过一些更底层的技巧（公众号首页的 PyTorch 合集），来帮助大家更快地训练和部署 PyTorch 模型。这些小技巧当然有帮助，但归根结底，它们只是“锦上添花”的工具。如果你没有对它们的适用场景和原理有一个系统的理解，那极有可能会用错地方，甚至适得其反。

这篇文章的重点，是帮助你从概念上理解模型操作的高阶组织方式，并结合 PyTorch 的实际示例进行讲解。正是这些并行训练的基本原则，才是如今支撑超大规模迭代与部署（比如拥有数千万日活用户系统）背后的真正动力。

数据并行

数据并行（Data Parallelism） 是最直接、也是应用最广泛的一种并行训练技术。它的基本思想是：创建多个相同的模型副本，并在不同的数据子集上分别进行训练。每个副本在本地计算梯度后，通过一种名为 all-reduce 的操作将所有梯度进行聚合，然后再用于更新每一个模型副本。

当你的模型本身可以放进一张 GPU 的显存，但训练数据量过大、无法顺序处理时，数据并行就是一个非常高效的解决方案。

PyTorch 主要通过以下两个模块内置了对数据并行的支持：

  • torch.nn.DataParalle

  • torch.nn.parallel.DistributedDataParallel（DDP）

其中，DDP 更受推荐，因为它在多节点环境下具有更好的可扩展性和效率。

NVIDIA 的 NeMo 框架就对这种并行方式提供了一个非常直观且酷炫的示意图和演示：

下面是一个简单的代码示例：

核心要点（Key Takeaways）

  • 小模型 / 大数据集：仅当模型本身可以完全放入一张 GPU 的显存，但数据集过大无法一次处理时，该方法才最有效。

  • 模型复制：每张 GPU 上都会存储一份完全相同的模型参数副本。

  • 小批量拆分：输入数据会被自动分割到多个 GPU 上，每张卡负责处理一个 mini-batch。

  • 梯度同步：前向和反向传播之后，所有 GPU 会进行梯度同步，以保证模型副本的一致性。

优势与注意事项（Benefits & Considerations）

✅ 简单高效：实现简单，能很好地集成到现有 PyTorch 代码中，对于大规模数据训练效果极佳，扩展性强。

⚠️ 通信开销：在梯度同步过程中存在通信开销，这可能在大规模系统中成为性能瓶颈。

张量并行

与数据并行通过划分数据不同，张量并行（Tensor Parallelism，也称为模型并行，Model Parallelism）的核心思想是将模型本身在多个设备间切分。这种方式会将大型的权重矩阵和中间张量进行分片，每个设备只处理整个模型计算的一部分。与数据并行中每张 GPU 上复制整个模型不同，Tensor 并行是将模型的不同层或张量分布在多个设备上，每个设备负责部分前向与反向传播计算。

这种技术在训练超大规模模型（尤其是基于 Transformer 的架构）时尤为有用，因为它们通常无法完整装入一张 GPU 的显存。

虽然 PyTorch 本身不直接提供对张量并行的内置支持，但通过 PyTorch 强大的张量操作能力和分布式通信原语，可以比较容易地实现自定义版本。如果你希望使用更健壮的解决方案，可以考虑使用像 DeepSpeed（可以参考此前的一篇文章《使用PyTorch Lightning优化AutoEncoder：多GPU训练DeepSpeed增强指南》）或 Megatron-LM 这样的框架，它们扩展了 PyTorch 并封装了 Tensor 并行的实现。

下面是一个简化版张量并行矩阵乘法的实现代码：

核心要点（Key Takeaways）

  • 适用于更大模型：当模型无法完整加载进单个 GPU 时，张量并行非常有效。

  • 权重分片：不是每张 GPU 都保存完整模型，而是对模型参数进行切分。

  • 联合计算：前向与反向传播由多个设备共同完成，需精心协调确保每部分计算正确。

  • 自定义操作：通常会用到定制的 CUDA kernel 或第三方库来提升计算效率。

优势与注意事项（Benefits & Considerations）

✅ 显存效率高：通过切分大张量，可以训练超过单卡显存限制的模型，同时也大幅降低矩阵运算的延迟。

⚠️ 实现复杂度高：设备间需要高度协调，特别是当扩展到两个以上的 GPU 时，必须处理好同步、负载均衡和通信等问题，否则可能会出现 GPU 空转等现象。

⚠️ 框架增强支持：如 Megatron-LM 等框架为 Tensor 并行设立了行业标准，并且可与 PyTorch 无缝集成，但集成过程本身仍可能不够“开箱即用”。

上下文并行

上下文并行（Context Parallelism）是一种完全不同的思路：它关注的是输入数据的“上下文维度（context dimension）”，在处理序列类模型（如 transformers）时尤其强大。它的核心理念是将长序列或上下文信息拆分开来，让不同部分并行处理，从而在不压垮内存或算力的情况下，让模型处理更长的上下文。

这种方式在多任务学习（例如多任务 NLP）中尤为适用，特别是在多个任务需要同时训练的场景下。

和张量并行类似，PyTorch 目前也没有原生支持上下文并行，但我们可以通过灵活的数据重构手动实现。例如：对于一个 transformer 模型，如果需要处理长文本，我们可以将其拆成若干片段，每个片段并行处理，最后再合并起来。

下面是一个简单的自定义 transformer 模块，它将长序列拆分成多个段落分别处理，最后合并输出：

核心要点（Key Takeaways）

  • 序列维度切分：将长序列的上下文维度进行拆分，使得不同段落可以并行处理。

  • 长序列可扩展性：尤其适合超长序列的处理场景，避免一次性加载整段文本带来的内存压力。

  • Attention 机制分担：在 Transformer 中，可将 self-attention 的计算负载切分到多个 GPU 上分别处理。

优势与注意事项（Benefits & Considerations）

✅ 长序列处理更高效：上下文被拆分后，可在不增加显存开销的前提下处理长文本或多任务。

⚠️ 边界依赖问题：上下文段之间可能存在依赖关系，因此可能需要使用重叠分段或结果融合机制来缓解跨段信息损失的问题。

⚠️ 新兴研究方向：目前上下文并行还属于较新领域，未来可能会有更多框架和工具专门支持这种训练方式。

流水线并行
流水线并行（Pipeline Parallelism）的核心思想是：将一个神经网络划分为多个连续的阶段（stages），每个阶段由不同的 GPU 来处理。随着数据在网络中的流动，中间结果依次传递至下一个阶段，类似于装配线的工作方式。这种错峰执行允许计算与通信并行进行，从而提高整体吞吐量。

幸运的是，PyTorch 提供了开箱即用的 API：Pipe，可以非常方便地创建这种分段模型。该 API 会自动将顺序模型划分为多个微批次（micro-batches），并让它们流经指定的多个 GPU。

下面是一个简单的使用 Pipe 实现流水线并行的示例：

核心要点（Key Takeaways）

  • 分阶段计算（Staged Computation）：模型被分为多个阶段，每个阶段分配给不同的 GPU。

  • 微批处理（Micro-Batching）：输入的大批次会被自动拆分成多个微批次，在流水线上连续流动处理。

  • 吞吐量提升（Increased Throughput）：即便每个 GPU 正在处理不同的微批次，但由于彼此并行工作，总体吞吐量显著提升。

优势与注意事项（Benefits & Considerations）

✅ 资源利用率更高：各 GPU 计算任务重叠，设备利用效率提升。

⚠️ 延迟 vs 吞吐量：吞吐量上升的同时，延迟可能略微上升，尤其是微批数量过少时。

⚠️ 调度复杂性：合理划分 micro-batch 和负载均衡是性能最优的关键。

这种方式非常适合用于训练非常深的模型（如 50 层以上的 Transformer），或者希望充分利用多块 GPU 的情况。

专家并行

专家并行（expert parallelism）是一种受到 MoE（Mixture-of-Experts，混合专家）模型启发的技术，旨在在保持计算成本可控的前提下扩展模型容量。在这种范式中，整个模型由多个专门化的“专家”子网络组成，每个专家专注于处理某一类输入。

通过一个门控机制（gating mechanism），模型在每个输入样本到来时动态选择部分专家来参与计算。也就是说：

对于每一个输入，仅有一小部分专家会被激活用于处理，从而大大增加了模型容量，但计算成本不会随之线性增长。

这种结构允许我们训练出参数规模极大（例如数千亿甚至上万亿）的模型，同时依然可以保持推理或训练过程中的计算效率，是近年来大模型架构设计的关键突破之一。

目前 PyTorch 尚未提供开箱即用的专家并行支持，但由于其模块化设计良好，用户可以较为灵活地自定义实现这一策略。

典型的做法是：

  • 定义一组专家层（expert layers），

  • 配合一个门控机制（gate）来决定每个输入应激活哪些专家。

在实际部署中，专家并行常与其他并行策略结合使用。例如，可以同时结合数据并行与专家并行，一方面处理大规模数据集，另一方面管理庞大的模型参数——通过智能路由机制将计算任务仅分发给相关专家，实现训练和推理效率的统一。

下面是一个简化版的实现示例：

核心要点（Key Takeaways）

  • 混合专家（Mixture-of-Experts）：每个样本只激活部分专家进行处理，大幅降低单样本计算量，但整体模型容量却可扩展至非常大。

  • 动态路由（Dynamic Routing）：门控函数会根据输入内容，动态决定哪些专家被激活，从而实现任务感知的计算调度。

  • 专家层级的并行（Expert-Level Parallelism）：专家可以被分布在多个设备上，从而实现真正的并行计算，降低瓶颈。

优势与注意事项（Benefits & Considerations）

✅ 可扩展的模型容量：通过专家并行，可以训练拥有海量参数的模型，而不必对每个输入都进行全模型计算。

✅ 计算效率高：每次仅启用部分专家，减少了不必要的计算。

⚠️ 路由机制复杂：门控设计至关重要，不合理的路由可能导致专家负载不均、模型训练不稳定。

⚠️ 前沿研究热点：专家并行仍处于持续研究与优化阶段，尤其在门控机制、专家同步、负载均衡等方面仍有提升空间。

ZeRO: Zero Redundancy Optimizer

ZeRO（Zero Redundancy Optimizer） 是大规模训练中的一项内存优化突破，它作为 DeepSpeed 库的一部分，旨在解决分布式训练中的内存限制问题。通过对优化器状态、梯度和模型参数进行分割，ZeRO 消除了每个 GPU 存储完整副本所带来的冗余，从而实现了显著的内存节省。

ZeRO 的工作原理是将优化器状态和梯度的存储分配到所有参与设备，而不是在每个 GPU 上复制一份。这一策略不仅减少了内存使用量，还使得在单个 GPU 的内存容量无法承载的大型模型可以得以训练。

ZeRO 通常分为三个阶段，每个阶段解决内存冗余的不同方面：

ZeRO-1：优化器状态分区

  • 将优化器状态（例如动量缓存）在多个 GPU 之间进行分区。

  • 每个 GPU 只存储其对应参数部分的优化器状态。

  • 模型参数和梯度仍然在所有 GPU 上复制。

ZeRO-2：梯度分区

  • 包含 ZeRO-1 的所有功能。

  • 进一步将梯度在多个 GPU 之间进行分区。

  • 每个 GPU 只计算并存储其对应参数部分的梯度。

  • 模型参数仍然在所有 GPU 上复制。

ZeRO-3：参数分区

  • 包含 ZeRO-1 和 ZeRO-2 的所有功能。

  • 进一步将模型参数在多个 GPU 之间进行分区。

  • 每个 GPU 只存储模型参数的部分。

  • 在前向传播和反向传播过程中，需要聚集所有参数以完成计算。


ZeRO 提供了最大的灵活性，结合了数据并行性和模型并行性的优点，如上所示。

虽然 ZeRO 是 DeepSpeed 的一部分，但它与 PyTorch 的集成使其成为训练优化工具中的一个关键组成部分，能够高效管理内存，并使以前无法训练的超大模型在现代硬件上变得可行。

以下是一个 ZeRO 的简单实现示例：

核心要点（Key Takeaways）

    • 阶段选择：ZeRO 通常分多个阶段实现，每个阶段在内存节省和通信开销之间提供不同的平衡。根据模型大小、网络能力以及可接受的通信开销，选择合适的阶段至关重要。

    • 与其他技术的集成：ZeRO 可以无缝地融入到包含其他并行化策略的生态系统中，如数据并行性和模型并行性策略的结合。

    优势与注意事项（Benefits & Considerations）

    ✅ 通信开销：该策略的固有挑战是，减少内存冗余通常会增加 GPU 之间的数据交换量。因此，高效利用高速互联技术（如 NVLink 或 InfiniBand）变得尤为关键。。

    ⚠️ 配置复杂性：与传统优化器相比，ZeRO 引入了更多的配置参数。这些设置需要通过精心的实验和性能分析来匹配硬件的优势，确保优化器高效运行。配置包括但不限于梯度聚合的适当桶大小、各个状态（优化器状态、梯度、参数）的分区策略。

    ⚠️ 强大的监控工具：在启用了 ZeRO 的训练中，调试问题非常具有挑战性。因此，监控工具（例如 GPU 内存使用情况、网络延迟和总体吞吐量的监控工具）变得尤为重要。

    混合并行策略在大规模训练中的应用

    在大规模训练深度学习模型时，通常需要一种混合方法，即结合多种前述技术。例如，一个最先进的 LLM（大语言模型）可能会使用：

    • 数据并行性（Data Parallelism）来分发批次到各个节点；

    • 张量并行性（Tensor Parallelism）来拆分超大的权重矩阵；

    • 上下文并行性（Context Parallelism）来处理长序列；

    • 流水线并行性（Pipeline Parallelism）来链接模型的顺序阶段；

    • 专家并行性（Expert Parallelism）来动态分配计算资源；

    • 最后使用 ZeRO 来优化内存使用。

    这种协同作用确保了即便是拥有庞大参数量的模型也能保持高效和可训练。

    关键要点

    • 灵活应用并行技术：理解何时、在哪里以及如何应用这些并行技术是推动技术进步的关键。只有掌握这些技术的使用场景，才能充分发挥其优势。

    • PyTorch 的模块化与插件库：PyTorch 的模块化设计和即插即用的库，使得构建强大且可扩展的训练管道变得越来越易于实现，突破传统硬件限制成为可能。这为更广泛的研究人员和开发者提供了前所未有的机会。

    通过结合这些先进的并行化技术，深度学习模型的训练不仅可以跨越硬件限制，还能实现对更大规模模型的处理，让大型 AI 模型的训练成为现实。

    https://levelup.gitconnected.com/training-deep-learning-models-at-ultra-scale-using-pytorch-74c6cbaa814b

      （文：PyTorch研习社）

      发表评论