分布式训练完全指南：数据并行、模型并行、流水并行与混合并行（so-large-lm 第8章）

文档教程大模型人工智能【免费下载链接】so-large-lm大模型基础: 一文了解大模型基础知识项目地址https://gitcode.com/datawhalechina/so-large-lm点击查看免费下载导读本文是 Datawhale 开源项目 so-large-lm大模型基础教程第 8 章「分布式训练」的完整技术解读。文章以 GPT-3 这类千亿级大模型的真实训练诉求为背景系统拆解数据并行、模型并行、流水并行与混合并行四种策略的切分方式、通信代价与适用场景并结合仓库第 4 章MoE 架构、第 6 章训练优化与第 14 章Llama 家族实践的源码级资料交叉印证。读完本文你将能依据模型规模与集群拓扑为训练任务选择并组合合适的并行策略并理解 AllReduce、梯度同步等关键通信原语的作用。1. 为什么分布式训练越来越流行近年来深度学习被广泛应用到计算机视觉、语言理解、语音识别、广告推荐等各个领域。这些领域有一个共同特点模型规模越来越大。以 GPT-3 为例其参数量达到1750 亿。即便使用1024 张 80 GB 的 A100完整训练 GPT-3 也需要约1 个月的时间。模型规模的扩大对硬件算力、内存的发展提出了严苛要求。然而由于 内存墙Memory Wall 的存在单一设备的算力及容量受限于物理定律持续提高芯片集成度越来越困难难以跟上模型规模扩大的需求。所谓内存墙是指处理器运算速度与内存访问速度之间的差距持续拉大导致算力再高也会被内存带宽与容量瓶颈所限制。为了解决算力增速不足的问题人们开始考虑用多节点集群进行分布式训练以聚合多台设备的算力与显存分布式训练由此势在必行。2. 常见的并行策略从矩阵乘法说起需要强调简单的机器堆叠并不一定会带来算力的增长。神经网络的训练并不是单纯的把原来一个设备做的事情现在分给多个设备各自做——它不仅需要多个设备进行计算还涉及设备之间的数据传输。只有协调好集群中的计算与通信才能实现高效的分布式训练。为了直观理解不同并行策略的本质区别我们以矩阵乘法为例。假设神经网络中某一层做矩阵乘法输入 $x$ 的形状为 $4 \times 5$模型参数 $w$ 的形状为 $5 \times 8$输出 $out$ 的形状为 $4 \times 8$。单机单卡训练中先计算得到 $out$将 $out$ 传递给下一层并最终计算得到 $loss$然后在反向传播过程中得到 $\frac{\partial loss}{\partial w}$用于更新 $w$。在分布式训练中依据切分 $x$ 还是切分 $w$的不同可以划分为数据并行与模型并行两类基本策略并由此衍生出流水并行与混合并行。3. 数据并行Data Parallelism数据并行将数据 $x$ 进行切分而每个设备上的模型 $w$ 是相同的。如下图所示$x$ 被按照第 0 维度平均切分到 2 个设备上两个设备上都持有完整的 $w$。两台设备分别得到的输出都只是逻辑上输出的一半形状为 $2 \times 8$将两个设备上的输出拼接到一起才能得到逻辑上完整的输出。关键问题——梯度同步由于数据被分发到了 2 个设备上反向传播过程中各设备得到的 $\frac{\partial loss}{\partial w}$ 会不一样。如果直接使用各自设备上的梯度更新各自的模型会造成 2 个设备上的模型不一致训练就失去了意义到底该用哪个模型。因此数据并行策略下在反向传播过程中需要对各个设备上的梯度执行AllReduce集合通信如 NVIDIA NCCL 提供的 AllReduce 原语将各设备梯度求和/平均后广播回所有设备从而确保各设备上的模型始终保持一致。适用场景当数据集较大、模型较小时反向过程中为同步梯度产生的通信代价较小此时选择数据并行一般更有优势。常见的视觉分类模型如 ResNet50比较适合采用数据并行。4. 模型并行Model Parallelism当神经网络非常巨大时数据并行同步梯度的代价会变得很大甚至网络可能巨大到无法存放到单一计算设备中这时可以采用模型并行策略。模型并行每个设备上的数据是完整的、一致的而模型 $w$ 被切分到各个设备上每个设备只拥有模型的一部分所有计算设备上的模型拼在一起才是完整的模型。如下图所示$w$ 被按照第 1 维度平均切分到 2 个设备上两个设备上都有完整的 $x$两个设备上的输出也需要通过拼接才能得到逻辑上完整的输出。优势与代价优势省去了多个设备之间的梯度 AllReduce代价由于每个设备都需要完整的数据输入数据需要在多个设备之间进行广播产生通信代价数据不会复制多份而是通过广播传递输入数据。例如上图中最终得到的 $out$$4 \times 8$如果作为下一层网络的输入就需要被广播发送到两个设备上。适用场景语言模型如 BERT常采用模型并行。5. 流水并行Pipeline Parallelism当神经网络过于巨大、无法在一个设备上存放时除了模型并行还可以选择流水并行。流水并行将网络切分为多个阶段stage分发到不同的计算设备上各计算设备之间以接力的方式完成训练。以一个逻辑上的 4 层网络T1T4为例4 层网络被切分到 2 个计算设备上其中 GPU0 上执行 T1 与 T2 的运算GPU1 上执行 T3 与 T4 的计算GPU0 完成前两层的计算后其输出被当作 GPU1 的输入继续完成后两层的计算。流水并行的本质是按层深度切分模型。它与第 4 章中提到的模型拆到多台机器、网络带宽成为瓶颈的直觉一致——docs/content/ch04.md 中给出的模型并行示例正是GPU1[layer1, layer2]、GPU2[layer3, layer4]、GPU3[layer5, layer6]这种按层切分的方式。6. 混合并行Hybrid Parallelism以 GPT-3 为例真实的大模型训练通常混用多种并行策略。以 GPT-3 为例其训练时的设备并行方案如下模型首先被分为64 个阶段进行流水并行每个阶段都运行在6 台 DGX-A100 主机上在6 台主机之间进行的是数据并行训练每台主机有8 张 GPU 显卡同一台机器上的 8 张 GPU 显卡之间进行模型并行训练。这套方案的要点在于在通信带宽充裕的尺度上做数据并行跨主机在通信带宽紧张、延迟敏感的尺度上做模型并行机内 NVLink整体再用流水并行把超大规模模型化整为零。并行策略的选择直接影响训练效率框架对并行训练的接口支持程度则决定了算法工程师的开发效率。7. 仓库纵深大模型实践中的并行策略佐证在 docs/content/ch14.mdLlama 开源家族中可以找到与本章策略一一对应的工程实践证据Llama-1650 亿参数采用模型并行与序列并行sequence parallelism并结合 xformers 的高效因果多头注意力、手动实现反向传播、激活检查点checkpointing等技巧同时优化 GPU 之间的通信——这正是本章所说协调好计算与通信在真实训练中的体现。Llama-3最大规模模型Meta结合了数据并行、模型并行和流水并行三种策略当同时在 16K GPU 上训练时最高效的策略实现了每个 GPU 超过 400 TFLOPS的计算利用率最终在两个定制的 24K GPU 集群上完成训练。这些改进使 Llama-3 的训练效率比 Llama-2 提高了约三倍总有效训练时间超过 95%。此外仓库其他章节还揭示了与分布式训练直接相关的两个底层事实优化器的显存开销docs/content/ch06.mdAdam 优化器需要同时存储 $\theta_t, g_t, m_t, v_t$ 四组量将存储从 2 倍模型参数增加到4 倍模型参数。对于动辄几十上百 GB 的模型这直接决定了单卡显存是否放得下、是否需要切分模型——是模型并行的重要动因。混合精度训练docs/content/ch06.md将主权重保存在 FP32其余运算在 FP16 中执行并通过损失缩放避免梯度下溢小于 $2^{-24}$ 的 FP16 值会变为 0从而显著降低显存与通信带宽压力是现代分布式训练中与并行策略配套使用的标准手段。专家并行Expert Parallelismdocs/content/ch04.md混合专家模型MoE天然有利于并行化——每个输入只激活一小部分专家专家可以被切分到不同设备上与数据并行结合后能进一步提升大规模稀疏模型的训练规模上限。8. 总结与选型建议并行策略切分对象同步/通信方式主要代价典型适用场景数据并行切分数据 $x$反向传播时对梯度做 AllReduce梯度同步通信开销数据集大、模型小如 ResNet50模型并行切分参数 $w$前向传播时广播输入数据数据广播通信开销模型无法单卡存放如 BERT流水并行按层切分网络为多个阶段阶段间接力传递激活流水线气泡bubble导致设备空闲超大规模模型按层拆分混合并行多策略组合分层使用不同通信原语需要精细的调度设计千亿级模型如 GPT-3、Llama-3核心选型逻辑可以概括为三点模型放不下单卡→ 优先考虑模型并行或流水并行先把模型拆开数据规模大、梯度同步成本低→ 优先数据并行把吞吐量铺开模型极大且集群规模大→ 采用混合并行在带宽充裕的层级做数据并行、在带宽紧张的层级做模型并行整体用流水并行组织阶段。进一步阅读仓库中的 docs/content/ch04.mdMoE 与检索增强等新架构如何影响并行、docs/content/ch06.md优化器显存分析与混合精度训练以及 docs/content/ch14.mdLlama-1/2/3 的真实训练并行方案可与本章内容互为补充构建完整的大模型如何被高效训练出来的知识闭环。赞分享文档教程大模型人工智能【免费下载链接】so-large-lm大模型基础: 一文了解大模型基础知识项目地址https://gitcode.com/datawhalechina/so-large-lm点击查看免费下载相关推荐Transformers 多 GPU 分布式训练并行方法指南数据并行、模型并行、流水线与 3D 并行Transformers 多 GPU 分布式训练并行方法指南数据并行、模型并行、流水线与 3D 并行 在大模型训练场景中单张 GPU 既受显存上限约束也难人工智能深度学习机器学习预训练微调NLP计算机视觉语音多模态三步完成小爱音箱AI改造从人工智障到智能学霸的终极指南三步完成小爱音箱AI改造从人工智障到智能学霸的终极指南 还在为小爱音箱只能执行简单指令而烦恼吗想让你的智能音箱真正理解你的需求成为贴心的AI助手吗今天人工智能AI 应用语音智能家居交互助手Drawio桌面版终极指南如何在离线环境下高效绘制专业图表Drawio桌面版终极指南如何在离线环境下高效绘制专业图表 Drawio desktop是基于Electron框架构建的开源图表绘制桌面应用它完美融合了dr桌面应用图形学上一篇ik_llama.cpp 对 LlaMA-4纯文本的支持从移植落地到 1 位量化配方全解析下一篇cargo-vet快速入门如何在10分钟内为你的Rust项目添加供应链安全检查创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考