AI模型并行训练：成就AI应用架构师的技术基石

禅与计算机程序设计艺术

79人浏览 · 2026-03-30 19:21:10

禅与计算机程序设计艺术 · 2026-03-30 19:21:10 发布

AI模型并行训练：成就AI应用架构师的技术基石

一、引言 (Introduction)

钩子 (The Hook)

当OpenAI在2023年发布GPT-4时，一个惊人的细节引发了AI社区的广泛讨论：其训练过程动用了超过10000块GPU，形成了一个庞大的分布式计算集群。然而，如果你深入探究会发现，即使拥有如此规模的硬件，若没有精妙的模型并行策略，GPT-4的训练周期可能会延长10倍以上，甚至根本无法完成。

这并非孤例。Meta的LLaMA 2（700亿参数）、Google的PaLM（5400亿参数）、Anthropic的Claude（1.3万亿参数）——这些推动AI革命的大模型背后，都离不开“模型并行训练”这一核心技术。一个残酷的现实是：当今主流的大语言模型（LLM）参数规模已突破万亿，单块GPU的显存（即使是40GB的A100）早已无法容纳其完整权重，更不用说训练过程中的梯度和激活值。

作为AI应用架构师，你是否曾面临以下困境：

训练一个中等规模的Transformer模型（如10亿参数）时，单卡显存频繁溢出，即使调小batch size也无济于事？
尝试使用数据并行扩展时，发现通信开销随GPU数量增加而急剧上升，性能不升反降？
面对客户提出的“如何在有限硬件资源下训练自定义大模型”的需求，却不知从何下手设计并行策略？

模型并行训练，正是解决这些问题的技术基石。它不仅是突破硬件限制的“钥匙”，更是AI架构师设计高效、可扩展训练系统的“必修课”。

定义问题/阐述背景 (The “Why”)

1. 模型规模的爆炸式增长：从“单卡能容”到“分布式刚需”

回顾AI模型的发展历程，参数规模呈现指数级增长：

2018年，BERT（3.4亿参数）可在单块V100 GPU上训练；
2020年，GPT-3（1750亿参数）需要数千块GPU协同；
2023年，PaLM 2（未公开具体参数，但推测远超万亿）依赖专用TPU集群。

这种增长背后的核心驱动力是“规模定律”（Scaling Law）：模型性能随参数、数据和计算量的增加而单调提升。但随之而来的是“内存墙”问题：

模型权重：1750亿参数的GPT-3（FP16精度）需要约350GB内存（1750e9 × 2字节）；
激活值：训练时的中间激活值通常是权重的2-4倍（取决于序列长度和batch size）；
梯度和优化器状态：Adam优化器的状态是权重的4倍（参数+动量+方差）。

单块GPU的显存（目前最高为H100的80GB）远无法满足需求，必须通过并行训练将模型拆分到多块设备上。

2. 模型并行的独特价值：突破数据并行的瓶颈

分布式训练的主流方法分为“数据并行”和“模型并行”两类：

数据并行：将数据拆分到多设备，每个设备存储完整模型副本，通过梯度同步更新参数（如PyTorch DDP）。
- 优点：实现简单，适用于中小模型；
- 缺点：通信开销随设备数量线性增长，且无法解决单设备内存不足问题（模型必须完整放入单卡）。
模型并行：将模型结构拆分到多设备，每个设备仅存储部分模型参数，协同完成前向/反向传播。
- 优点：突破单卡内存限制，支持超大规模模型；
- 缺点：实现复杂，需深入理解模型结构和硬件特性。

随着模型规模超越单卡容量，模型并行从“可选优化”变为“必选方案”。例如，训练1000亿参数的模型时，即使使用数据并行，单卡仍需存储完整模型（约2000GB FP16权重），这显然不可能。此时，模型并行是唯一可行的路径。

3. AI应用架构师的核心能力：从“模型使用者”到“系统设计者”

传统AI工程师关注模型调优和应用开发，而AI应用架构师需要从“系统视角”设计端到端的训练与部署方案。模型并行能力直接决定了架构师能否：

基于业务需求选择合适的并行策略（如低延迟场景选张量并行，高吞吐量场景选流水线并行）；
优化硬件资源利用率（如GPU内存、网络带宽），降低训练成本；
应对大规模模型的工程挑战（如容错、负载均衡、跨节点通信）；
评估和选型分布式框架（如Megatron-LM、DeepSpeed、Colossal-AI）。

可以说，不懂模型并行的AI架构师，将无法驾驭下一代大模型技术。

亮明观点/文章目标 (The “What” & “How”)

本文将系统梳理AI模型并行训练的技术体系，帮助你从“理论认知”到“实战落地”，构建成为顶级AI应用架构师的核心竞争力。读完本文后，你将掌握：

核心知识

模型并行的底层原理：张量并行、流水线并行、混合并行的数学本质与实现逻辑；
关键技术对比：数据并行vs模型并行，不同模型并行策略的优缺点与适用场景；
硬件与框架协同：GPU/TPU架构对并行训练的影响，主流框架（PyTorch/TensorFlow）的模型并行API。

实战能力

从零实现：用PyTorch手动拆分Transformer层，构建基础张量并行模型；
框架应用：基于Megatron-LM实现GPT-2的张量并行训练，基于DeepSpeed实现BERT的流水线并行；
性能调优：定位通信瓶颈，优化显存占用，实现计算与通信重叠。

架构设计

场景化方案：根据模型类型（CNN/Transformer/GNN）和硬件条件（单节点多卡/多节点集群）选择并行策略；
大规模部署：设计混合并行系统（模型并行+数据并行+专家并行），支撑万亿参数模型训练；
最佳实践：规避负载不均衡、死锁、精度损失等常见陷阱，保障训练稳定性与效率。

本文结构：

基础知识：从分布式训练基础到模型并行核心概念，构建理论框架；
核心技术：详解张量并行、流水线并行、混合并行的原理与实现；
实战演练：三大案例覆盖从基础到工业级的模型并行落地；
进阶优化：性能调优、容错机制、大规模部署的最佳实践；
未来趋势：自动并行、硬件感知优化、量子加速等前沿方向。

无论你是初涉分布式训练的AI工程师，还是需要设计大规模训练系统的架构师，本文都将为你提供体系化的技术指导。让我们开始这场模型并行的深度探索之旅吧！

二、基础知识/背景铺垫 (Foundational Concepts)

核心概念定义：从“分布式训练”到“模型并行”

1. 分布式训练的本质：打破“内存墙”与“算力墙”

分布式训练的核心目标是通过多设备协同，突破单设备的内存和算力限制，实现超大规模模型的训练与推理。其数学本质可抽象为：

问题：求解优化问题 $ \theta^* = \arg\min_\theta \mathcal{L}(\theta; D) $，其中 $ D $ 为训练数据，$ \mathcal{L} $ 为损失函数，$ \theta $ 为模型参数。
挑战：当 $ \theta $ 规模过大（如万亿参数）或 $ D $ 过多（如万亿样本），单设备无法存储 $ \theta $ 或高效计算梯度 $ \nabla_\theta \mathcal{L} $。
解决方案：通过“拆分”策略将问题分解到多设备，协同完成计算与通信。

分布式训练的“拆分”维度主要有三：

数据拆分（数据并行）：拆分 $ D $，每个设备处理部分数据；
模型拆分（模型并行）：拆分 $ \theta $，每个设备存储部分参数；
混合拆分：同时拆分 $ D $ 和 $ \theta $（如数据并行+模型并行）。

2. 模型并行的定义与分类

模型并行（Model Parallelism）：将模型的计算图（Computational Graph）或参数（Parameters）拆分到多个设备（GPU/TPU/CPU），通过设备间通信协同完成前向传播（Forward Pass）和反向传播（Backward Pass）。

根据拆分粒度和方式，模型并行可分为三大类（如图2-1所示）：

张量并行（Tensor Parallelism）：拆分单个张量（如权重矩阵）到多设备，并行执行矩阵运算（如矩阵乘法）；
流水线并行（Pipeline Parallelism）：按层拆分模型（如将Transformer的Encoder层拆分到不同设备），通过流水线执行前向/反向传播；
混合并行（Hybrid Parallelism）：结合张量并行、流水线并行和数据并行，适用于超大规模模型（如GPT-3、PaLM）。

外链图片转存失败,源站可能有防盗链机制,建议将图片保存下来直接上传
图2-1：模型并行的三大分类及其适用场景

3. 关键术语解析：从“通信原语”到“并行粒度”

术语	定义	核心作用
张量（Tensor）	多维数组，模型参数、输入数据、中间激活值的基本存储形式（如权重矩阵 $ W \in \mathbb{R}^{n \times m} $）	模型并行的基本操作单元
计算图（Computational Graph）	模型的数学表达形式，由节点（操作）和边（张量流动）组成	模型拆分的逻辑基础（按节点或边拆分）
通信原语（Communication Primitives）	设备间数据传输的基本操作，如all-gather、reduce-scatter、broadcast	模型并行中协同计算的核心机制
并行粒度（Parallelism Granularity）	模型拆分的精细程度，如算子级、层级、模块级	影响通信效率和负载均衡（细粒度vs粗粒度）
同步/异步（Synchronous/Asynchronous）	设备间是否等待彼此完成计算再同步数据	影响训练稳定性（同步收敛更稳定，异步速度更快）

4. 模型并行与数据并行的本质区别

为直观理解两者差异，我们以一个简单的两层MLP模型（$ y = \sigma(W_2 \sigma(W_1 x + b_1) + b_2) $）为例：

数据并行场景（图2-2左）：
- 设备0和设备1各存储完整模型（$ W_1, b_1, W_2, b_2 $）；
- 输入数据 $ x $ 拆分为 $ x_0 $ 和 $ x_1 $，分别输入两设备；
- 前向传播：设备0计算 $ y_0 $，设备1计算 $ y_1 $；
- 反向传播：两设备分别计算梯度 $ \nabla W_1^0, \nabla W_1^1 $，通过all-reduce求和后更新参数。
模型并行场景（图2-2右）：
- 设备0存储 $ W_1, b_1 $，设备1存储 $ W_2, b_2 $；
- 输入 $ x $ 完整传入设备0，计算 $ h = \sigma(W_1 x + b_1) $；
- 设备0将中间结果 $ h $ 发送给设备1；
- 设备1计算 $ y = \sigma(W_2 h + b_2) $，反向传播时梯度沿原路返回。

外链图片转存失败,源站可能有防盗链机制,建议将图片保存下来直接上传
图2-2：左：数据并行（模型副本+数据拆分）；右：模型并行（模型拆分+数据流动）

核心差异总结：

维度	数据并行	模型并行
存储内容	完整模型副本	部分模型参数
通信对象	梯度（反向传播）	中间激活值（前向）+ 梯度（反向）
通信量	$ O(P \times \text{模型大小}) $	$ O(\text{中间激活值大小}) $
适用场景	模型可放入单卡内存	模型无法放入单卡内存
实现复杂度	低（框架自动支持，如DDP）	高（需手动拆分模型结构）

关键结论：当模型大小超过单卡内存时，数据并行失效，必须使用模型并行；当模型和数据均超单卡容量时，需混合使用两者。

通信原语	功能描述	模型并行应用场景
Broadcast	从一个源设备向所有目标设备发送数据（1-to-N）	前向传播中输入数据分发、参数初始化同步
Gather/Scatter	Gather：多设备数据汇聚到一个设备（N-to-1）；Scatter：一个设备数据拆分到多设备（1-to-N）	张量并行中输入/输出拆分与合并
All-Gather	每个设备将本地数据发送给其他所有设备，最终每个设备拥有完整数据集合（N-to-N）	张量并行中特征拼接（如多头注意力输出合并）
Reduce-Scatter	先按维度拆分数据，对各分片执行reduce（如求和），再将结果分发到对应设备	张量并行中梯度计算（如矩阵乘法梯度拆分）
All-Reduce	所有设备对同一数据执行reduce操作，结果广播到所有设备（N-to-N）	数据并行梯度同步、混合并行中的参数更新

模型并行的分类与数学基础

1. 按拆分维度分类：张量并行、流水线并行与混合并行

（1）张量并行（Tensor Parallelism）

核心思想：将单个张量（如权重矩阵）按维度拆分到多设备，并行执行张量运算（主要是矩阵乘法）。

例：权重矩阵 $ W \in \mathbb{R}^{n \times m} $ 的拆分

按行拆分（输入并行）：将 $ W $ 拆分为 $ W_1 \in \mathbb{R}^{n_1 \times m}, W_2 \in \mathbb{R}^{n_2 \times m} $，其中 $ n_1 + n_2 = n $；
按列拆分（输出并行）：将 $ W $ 拆分为 $ W_1 \in \mathbb{R}^{n \times m_1}, W_2 \in \mathbb{R}^{n \times m_2} $，其中 $ m_1 + m_2 = m $；
2D拆分：同时按行和列拆分（如2×2设备网格，$ W $ 拆为4个子矩阵）。

数学原理：矩阵乘法的分布式计算
设 $ Y = X W $，其中 $ X \in \mathbb{R}^{b \times n} $（输入批次， b 为 ba t c h s i z e ），$ W \in \mathbb{R}^{n \times m} $（权重），$ Y \in \mathbb{R}^{b \times m} $（输出）。

按列拆分 $ W $（图2-4左）：
- 将 $ W $ 拆分为 $ W_1 \in \mathbb{R}^{n \times m_1}, W_2 \in \mathbb{R}^{n \times m_2} $（$ m = m_1 + m_2 $）；
- 设备0计算 $ Y_1 = X W_1 $，设备1计算 $ Y_2 = X W_2 $；
- 通过All-Gather将 $ Y_1, Y_2 $ 拼接为 $ Y = [Y_1, Y_2] $。
按行拆分 $ W $（图2-4右）：
- 将 $ X $ 拆分为 $ X_1 \in \mathbb{R}^{b \times n_1}, X_2 \in \mathbb{R}^{b \times n_2} $（$ n = n_1 + n_2 $）；
- 设备0计算 $ Y_1 = X_1 W_1 $，设备1计算 $ Y_2 = X_2 W_2 $（$ W_1 \in \mathbb{R}^{n_1 \times m}, W_2 \in \mathbb{R}^{n_2 \times m} $）；
- 通过Reduce-Scatter将 $ Y_1, Y_2 $ 求和得到 $ Y = Y_1 + Y_2 $。

外链图片转存失败,源站可能有防盗链机制,建议将图片保存下来直接上传

优点：计算与通信高度重叠，并行粒度细，适合计算密集型算子（如Transformer中的QKV矩阵乘法）；
缺点：需模型算子级支持，对非张量运算（如激活函数）拆分困难。

（2）流水线并行（Pipeline Parallelism）

核心思想：将模型按层（Layer）拆分为多个阶段（Stage），每个阶段由一个或多个设备处理，通过流水线执行前向/反向传播，隐藏部分通信延迟。

例：3层MLP拆分为3个阶段（图2-5）

前向传播：
- 批次数据拆分为微批次（Micro-batch）$ B_1, B_2, B_3 $；
- 阶段1处理 $ B_1 $ → 输出发送给阶段2；阶段1处理 $ B_2 $ 时，阶段2处理 $ B_1 $（流水线重叠）；
- 最终阶段输出所有微批次结果，拼接为完整批次。
反向传播：按前向传播的逆序执行，梯度从最后阶段流向第一阶段（与前向形成“反向流水线”）。

外链图片转存失败,源站可能有防盗链机制,建议将图片保存下来直接上传

关键概念：

微批次（Micro-batch）：将批次数据拆分为更小单元，实现流水线重叠；
气泡（Bubble）：流水线启动和结束时的空闲时间（如图2-5中t0-t2和t6-t8的空白区域）；
阶段（Stage）：模型层的集合（如将Transformer的10层拆分为2个阶段，每个阶段5层）。

优点：实现简单（按层拆分），适合模型层次清晰的场景（如CNN、Transformer）；
缺点：存在流水线气泡，设备利用率随阶段数增加而降低（如n个阶段，理想利用率为1-1/n）。

（3）混合并行（Hybrid Parallelism）

当模型规模极大（如万亿参数）时，单一并行策略无法满足需求，需组合使用张量并行、流水线并行和数据并行：

张量+数据并行：同一模型层内用张量并行拆分权重，不同数据批次用数据并行；
流水线+数据并行：不同模型阶段用流水线并行，每个阶段内部用数据并行；
张量+流水线+数据并行：如GPT-3训练中，使用张量并行拆分注意力层，流水线并行拆分Transformer层，数据并行扩展训练数据（图2-6）。

外链图片转存失败,源站可能有防盗链机制,建议将图片保存下来直接上传

2. 按通信模式分类：同步并行与异步并行

（1）同步并行（Synchronous Parallelism）

所有设备完成当前步骤后才进入下一步（如前向传播需等待所有设备完成计算），数学上等价于单设备训练，收敛性有保证。

优点：训练稳定，梯度更新无偏差；
缺点：慢设备拖慢整体（“木桶效应”），通信延迟直接影响训练速度。

（2）异步并行（Asynchronous Parallelism）

设备独立执行计算和更新，无需等待其他设备（如参数服务器架构中，设备计算完梯度后立即更新参数，不等待其他设备）。

优点：无木桶效应，设备利用率高；
缺点：梯度更新存在“陈旧性”（Staleness），可能导致收敛不稳定（需通过算法优化，如异步SGD的学习率衰减策略）。

工业界选择：绝大多数大模型训练采用同步并行（如GPT-3、LLaMA），因收敛性更可靠；异步并行多用于数据量极大且实时性要求高的场景（如推荐系统训练）。

关键挑战与评估指标

1. 模型并行的核心挑战

通信开销：设备间数据传输（如中间激活值、梯度）占用带宽，可能成为性能瓶颈；
负载均衡：模型拆分不均导致部分设备空闲（如流水线并行的气泡）；
内存优化：除参数外，中间激活值和梯度也占用大量内存，需通过Checkpointing（激活值检查点）等技术优化；
实现复杂度：需深入理解模型结构和硬件特性，手动拆分时易出错（如梯度计算方向错误）。

2. 评估指标：从“效率”到“性价比”

吞吐量（Throughput）：单位时间内处理的样本数（samples/sec），衡量训练速度；
内存利用率：GPU显存的实际使用比例（避免过度空闲或溢出）；
并行效率（Parallel Efficiency）：$ \text{效率} = \frac{\text{分布式训练速度}}{\text{单设备训练速度} \times \text{设备数}} $，理想值为1（线性加速）；
通信量（Communication Volume）：设备间传输的数据总量（GB），直接影响训练时间；
收敛速度：达到目标精度所需的迭代次数（并行策略不应显著增加迭代次数）。

示例：比较张量并行（TP）和流水线并行（PP）的效率

小模型（单卡可容纳）：TP/PP效率低于数据并行（因通信开销）；
大模型（单卡不可容纳）：TP/PP是唯一可行方案，效率取决于通信优化程度（如使用NVLink可提升TP效率至80%以上）。

本章小结

本章系统梳理了模型并行的基础知识，核心要点包括：

本质差异：模型并行通过拆分模型结构突破单卡内存限制，是超大规模模型训练的必选方案；
技术体系：按拆分方式分为张量并行（算子级拆分）、流水线并行（层级拆分）和混合并行（多维度拆分）；
硬件与框架：GPU的显存、计算能力和互连带宽是基础，Megatron-LM/DeepSpeed等框架简化了工业级实现；
核心挑战：通信开销、负载均衡、内存优化和实现复杂度是设计并行策略时需重点解决的问题。

理解这些基础知识后，我们将在第三章深入探讨张量并行、流水线并行和混合并行的实现原理与代码实战，从理论走向实践。

三、核心内容/实战演练 (The Core - “How-To”)

张量并行（Tensor Parallelism）：算子级拆分的艺术

1. 数学原理：从“矩阵拆分”到“梯度计算”

（1）前向传播：矩阵乘法的并行化

以Transformer中的多头注意力（Multi-Head Attention）为例，其核心操作为：
$W_Q, \quad K = X W_K, \quad V = X W_V$
其中 $ X \in \mathbb{R}^{b \times d} $（输入序列， b 为 ba t c h s i z e ， d 为隐藏层维度），$ W_Q, W_K, W_V \in \mathbb{R}^{d \times d_k \times h} $（ h 为头数，$ d_k = d/h $ 为单头维度）。

按列拆分权重矩阵（以 $ W_Q $ 为例，图3-1）：

将 $ W_Q $ 按输出维度（列）拆分为 $ W_Q^1, W_Q^2, …, W_Q^p $，其中 $ p $ 为张量并行度（设备数），每个 $ W_Q^i \in \mathbb{R}^{d \times (d_k h / p)} $；
设备 $ i $ 存储 $ W_Q^i, W_K^i, W_V^i $，输入 $ X $ 完整广播到所有设备；
设备 $ i $ 计算 $ Q^i = X W_Q^i, K^i = X W_K^i, V^i = X W_V^i $（此时 $ Q^i \in \mathbb{R}^{b \times (d_k h / p)} $，即拆分后的查询矩阵）；
注意力分数计算：$ \text{Attn}^i = \text{softmax}(\frac{Q^i (Kⁱ⁾T}{\sqrt{d_k}}) V^i $；
通过 All-Gather 操作将所有设备的 $ \text{Attn}^i $ 拼接为完整注意力输出 $ \text{Attn} \in \mathbb{R}^{b \times d} $。

外链图片转存失败,源站可能有防盗链机制,建议将图片保存下来直接上传

数学验证：
拆分后 $ W_Q = [W_Q^1, W_Q^2, …, W_Q^p]T $（列拼接），则 $ Q = X W_Q = X [W_Q^1, …, W_Q^p] = [X W_Q^1, …, X W_Q^p] = [Q^1, …, Q^p] $，与All-Gather结果一致。

（2）反向传播：梯度计算与Reduce-Scatter

反向传播中，需计算权重梯度 $ \nabla W_Q = X^T \nabla Q $。由于前向传播中 $ Q = [Q^1, …, Q^p] $，则 $ \nabla Q = [\nabla Q^1, …, \nabla Q^p] $（梯度按输入维度拆分）。

设备 $ i $ 计算 $ \nabla W_Q^i = X^T \nabla Q^i $（因 $ Q^i = X W_Q^i $，故 $ \nabla W_Q^i = X^T \nabla Q^i $）；
通过 Reduce-Scatter 操作，各设备将本地梯度 $ \nabla W_Q^i $ 发送到对应设备，完成梯度聚合。

注：反向传播的通信原语与前向传播“对偶”：前向用All-Gather，反向用Reduce-Scatter；前向用Broadcast，反向用Gather。

2. 代码实现：PyTorch手动实现张量并行

（1）环境准备

硬件：2张GPU（如A100），支持NVLink；
软件：PyTorch 2.0+，CUDA 11.7+，nccl通信后端；
初始化分布式环境：

import torch
import torch.distributed as dist
import torch.nn as nn

# 初始化分布式进程组（假设2设备）
dist.init_process_group(backend='nccl', init_method='env://', world_size=2, rank=0)  # 设备0
# dist.init_process_group(backend='nccl', init_method='env://', world_size=2, rank=1)  # 设备1（实际通过launch脚本启动）
rank = dist.get_rank()
device = torch.device(f'cuda:{rank}')

（2）张量并行线性层（TensorParallelLinear）

实现按列拆分的线性层，支持前向/反向传播：

class TensorParallelLinear(nn.Module):
    def __init__(self, in_features, out_features, parallel_size=2):
        super().__init__()
        self.parallel_size = parallel_size
        self.out_features = out_features
        self.in_features = in_features
        
        # 按列拆分权重：每个设备存储 out_features / parallel_size 列
        self.out_features_per_device = out_features // parallel_size
        self.weight = nn.Parameter(torch.randn(
            in_features, self.out_features_per_device, device=device
        ))
        self.bias = nn.Parameter(torch.randn(
            self.out_features_per_device, device=device
        ))
        
    def forward(self, x):
        # x: [b, in_features]，需广播到所有设备（实际中通过上游自动广播）
        out = torch.matmul(x, self.weight) + self.bias  # [b, out_features_per_device]
        
        # All-Gather：收集所有设备的输出并拼接
        # 创建输出缓冲区（所有设备大小相同）
        gathered_out = [torch.empty_like(out) for _ in range(self.parallel_size)]
        dist.all_gather(gathered_out, out)  # gathered_out[i] 为设备i的输出
        out = torch.cat(gathered_out, dim=-1)  # [b, out_features]
        return out

（3）测试与验证

单设备线性层：

# 设备0：单卡线性层
single_linear = nn.Linear(1024, 2048).to(device)
x = torch.randn(32, 1024).to(device)  # batch size=32，输入维度1024
y_single = single_linear(x)  # [32, 2048]

张量并行线性层（2设备）：

# 设备0和1分别初始化张量并行线性层
tp_linear = TensorParallelLinear(1024, 2048, parallel_size=2).to(device)
# 手动同步权重（实际中通过参数广播初始化）
if rank == 0:
    # 将单卡权重按列拆分并发送给设备1
    w0 = single_linear.weight[:, :1024].clone()  # 设备0权重（前1024列）
    w1 = single_linear.weight[:, 1024:].clone()   # 设备1权重（后1024列）
    b0 = single_linear.bias[:1024].clone()
    b1 = single_linear.bias[1024:].clone()
    dist.send(w1, dst=1)
    dist.send(b1, dst=1)
    tp_linear.weight.data.copy_(w0)
    tp_linear.bias.data.copy_(b0)
else:
    w1 = torch.empty_like(tp_linear.weight)
    b1 = torch.empty_like(tp_linear.bias)
    dist.recv(w1, src=0)
    dist.recv(b1, src=0)
    tp_linear.weight.data.copy_(w1)
    tp_linear.bias.data.copy_(b1)

# 前向传播
y_tp = tp_linear(x)
# 验证结果一致性（设备0上比较）
if rank == 0:
    print(torch.allclose(y_single, y_tp, atol=1e-6))  # 应输出True

关键结论：张量并行的核心是权重拆分+局部计算+通信拼接，需确保前向和反向传播的通信操作对称（如前向All-Gather对应反向Reduce-Scatter）。

流水线并行（Pipeline Parallelism）：层间拆分与流水线调度

1. 原理：从“顺序执行”到“流水线重叠”

以包含L层的Transformer模型为例，传统单设备执行流程为：
$x_0 \rightarrow \text{Layer}_1 \rightarrow x_1 \rightarrow \text{Layer}_2 \rightarrow ... \rightarrow \text{Layer}_L \rightarrow x_L$
其中 $ x_0 $ 为输入，$ x_L $ 为输出，每层执行时间为 $ t_1, t_2, …, t_L $，总时间 $ T = \sum_{i=1}^L t_i $。

流水线并行拆分（图3-2）：

将L层拆分为P个阶段（Stage），每个阶段包含连续的 $ L/P $ 层（如L=12，P=3，每个阶段4层）；
阶段 $ i $ 由设备 $ i $ 处理，输入数据拆分为M个微批次（Micro-batch）$ B_1, B_2, …, B_M $；
前向传播：
- $ t_0 $：设备0处理 $ B_1 $ 的阶段1 → 输出发送给设备1；
- $ t_1 $：设备0处理 $ B_2 $ 的阶段1，设备1处理 $ B_1 $ 的阶段2；
- $ t_2 $：设备0处理 $ B_3 $ 的阶段1，设备1处理 $ B_2 $ 的阶段2，设备2处理 $ B_1 $ 的阶段3；
- … 以此类推，形成流水线重叠。

外链图片转存失败,源站可能有防盗链机制,建议将图片保存下来直接上传

流水线效率分析：

理想加速比：$ P $（设备数），但实际因气泡（Bubble）存在而降低；
气泡占比：启动阶段（前P-1步）和结束阶段（后P-1步）存在气泡，总时间 $ T_{\text{pp}} = (M + P - 1) \times T_{\text{stage}} $，其中 $ T_{\text{stage}} $ 为单阶段处理一个微批次的时间；
优化方向：增加微批次数量M（气泡占比 $ (2P-2)/(M+P-1) \rightarrow 0 $，当 $ M \gg P $ 时）。

2. 代码实现：基于PyTorch实现简单流水线并行

（1）定义模型阶段（Stage）

class PipelineStage(nn.Module):
    def __init__(self, num_layers, hidden_dim=1024):
        super().__init__()
        self.layers = nn.ModuleList([
            nn.Sequential(
                nn.Linear(hidden_dim, hidden_dim),
                nn.GELU()
            ) for _ in range(num_layers)
        ])
    
    def forward(self, x):
        for layer in self.layers:
            x = layer(x)
        return x

（2）流水线调度器（Pipeline Scheduler）

实现微批次拆分、前向/反向传播调度：

class PipelineScheduler:
    def __init__(self, stages, num_microbatches=4):
        self.stages = stages  # 阶段列表，每个阶段在不同设备上
        self.num_stages = len(stages)
        self.num_microbatches = num_microbatches
        self.rank = dist.get_rank()
    
    def split_microbatches(self, x):
        # 将批次拆分为微批次（按batch维度）
        return torch.split(x, x.size(0) // self.num_microbatches, dim=0)
    
    def forward_backward(self, x, labels, criterion):
        microbatches = self.split_microbatches(x)
        loss = torch.tensor(0.0, device=x.device)
        
        # 前向传播：按阶段和微批次调度
        forward_outputs = [[] for _ in range(self.num_stages)]  # 存储各阶段输出
        for m in range(self.num_microbatches):
            mb = microbatches[m].to(self.stages[0].device)  # 微批次输入到第一阶段
            for s in range(self.num_stages):
                with torch.no_grad():  # 前向传播不保留梯度（反向时重新计算）
                    mb = self.stages[s](mb)
                forward_outputs[s].append(mb)
                if s < self.num_stages - 1:
                    # 发送到下一阶段（简化版：直接使用dist.send/recv）
                    next_device = self.stages[s+1].device
                    dist.send(mb.contiguous(), dst=s+1)
                    mb = torch.empty_like(mb).to(next_device)
                    dist.recv(mb, src=s)
        
        # 反向传播：按逆序处理微批次和阶段
        for m in reversed(range(self.num_microbatches)):
            # 最后阶段输出计算损失
            if self.rank == self.num_stages - 1:
                last_output = forward_outputs[-1][m]
                loss_mb = criterion(last_output, labels[m])
                loss_mb.backward()
                grad = last_output.grad.clone()
            else:
                grad = None
            
            # 梯度回传各阶段
            for s in reversed(range(self.num_stages - 1)):
                if self.rank == s + 1:
                    # 发送梯度给前一阶段
                    dist.send(grad, dst=s)
                    grad = torch.empty_like(forward_outputs[s][m]).to(self.stages[s].device)
                    dist.recv(grad, src=s+1)
                elif self.rank == s:
                    # 重新计算前向传播（带梯度）
                    mb = microbatches[m].to(self.stages[s].device)
                    for layer in self.stages[s].layers:
                        mb = layer(mb)
                    mb.backward(grad)
        
        return loss / self.num_microbatches

（3）测试与验证

单设备模型：

# 单卡12层模型
single_model = nn.Sequential(*[PipelineStage(1) for _ in range(12)]).to(device)
x = torch.randn(32, 1024).to(device)  # batch size=32
labels = torch.randint(0, 10, (32,)).to(device)  # 分类标签
criterion = nn.CrossEntropyLoss()
y_single = single_model(x)  # [32, 1024]
loss_single = criterion(y_single, labels)
loss_single.backward()

流水线并行模型（3阶段，每阶段4层）：

# 3个阶段（设备0、1、2）
stages = [
    PipelineStage(4).to(device=0),
    PipelineStage(4).to(device=1),
    PipelineStage(4).to(device=2)
]
scheduler = PipelineScheduler(stages, num_microbatches=4)
loss_pp = scheduler.forward_backward(x, labels, criterion)
# 验证损失一致性（略，需同步各阶段参数后比较）