斯坦福 CS336 从零构建大模型 (2025 春) - 第七讲:并行计算基础(Parallelism 1)
斯坦福 CS336 从零构建大模型 (2025 春) - 第七讲:并行计算基础(Parallelism 1)
文章目录
斯坦福 CS336 第七讲(Parallelism 1)是一节极其硬核的系统架构课。随着模型参数量和数据量的爆炸,单张 GPU 已经无法装下现代大模型,训练被迫从单卡优化转向“数据中心级别”的大规模分布式并行。
讲师系统性地拆解了多机并行的底层通信基础,以及数据并行(DP)、模型并行(MP)和激活值并行等核心策略。以下是本讲不遗漏任何知识点的详细总结:
一、硬件通信基础与集群原语 (Hardware & Communications)
要想做好并行,首先必须理解底层硬件的物理限制。
- 物理拓扑的“快慢交织”:在标准的 GPU 节点(如包含 8 张 H100 的服务器)内部,GPU 之间通过 NVLink/NVSwitch 连接,通信速度极快;但跨机器的 GPU 通信必须经过慢得多的网络交换机(如 Infiniband),而当集群规模超过 256 张 GPU 时,往往还需要跨越更慢的脊叶(Spine-leaf)交换机。这决定了不同的并行策略必须被放置在对应的通信层级上。
- 集群通信原语的数学等价性:并行算法底层依赖于 All-reduce(全规约)、Broadcast(广播)、Reduce-scatter(规约散播)和 All-gather(全收集)等操作。讲师特别强调了一个将在后续大幅优化显存的关键恒等式:All-reduce = Reduce-scatter + All-gather。
二、数据并行 (DP) 与 ZeRO / FSDP 显存优化
朴素的数据并行(Naive DP)是将整个模型复制到每一张 GPU 上,然后把数据批次(Batch)切分给不同 GPU 计算,最后通过 All-reduce 同步梯度。它在算力扩展上很好,但在显存利用上是灾难性的。
- 16字节的显存账本:讲师算了一笔账,在混合精度训练中,每个参数实际上需要占用高达 16 字节的显存。其中参数本身和梯度只占极小部分,绝大部分显存被 Adam 优化器状态(主权重、一阶动量、二阶动量)吃掉了。
- ZeRO 系列优化:为了打破显存瓶颈,微软提出了 ZeRO(零冗余优化),通过巧妙的通信操作消除冗余:
- ZeRO-1(优化器状态切分):所有卡仍保留完整的参数和梯度,但在参数更新时,每张卡只负责更新自己分到的那部分参数(利用属于它的优化器状态)。它先用 Reduce-scatter 汇总自己负责的梯度,更新后,再用 All-gather 将更新后的参数广播给所有人。这与朴素 DP 的通信成本完全相同,但极大节省了优化器显存。
- ZeRO-2(梯度切分):进一步切分梯度。在反向传播中,一旦某层的梯度算完,立刻用 Reduce-scatter 发给负责该层的 GPU,并释放本地显存。
- ZeRO-3(即 FSDP,完全参数切分):连模型权重也切碎。前向和反向传播时,模型会逐层“按需动态拉取(All-gather)”所需的参数,算完立刻丢弃(Free)。这带来了 3 倍于参数量的通信开销,但通过“计算与通信重叠(Overlapping)”(即在计算当前层时,后台预取下一层的参数),极大地掩盖了通信延迟,效率惊人地高。
- DP 的核心限制:数据并行的扩展极限受限于全局批次大小(Batch Size)。你不能切分出比 Batch Size 更多的并行份数,因此批次大小被视为一种极其稀缺的调度资源。
三、模型并行 (Model Parallelism)
ZeRO-3 虽然切分了参数,但没有切分前向传播产生的激活值(Activations),且需要极高的动态通信。因此必须引入模型并行:
- 流水线并行 (Pipeline Parallelism, PP):
- 按神经网络的深度(层数)切分,把不同层放到不同机器上。
- 痛点:严重的“气泡(Bubble)”问题,即大部分 GPU 处于空闲状态等待前置节点的计算结果。
- 优化:通过微批次(Micro-batches)来缩小气泡。进阶玩法是 Zero Bubble / DualPipe:将反向传播拆分为“计算激活值梯度 (B)”和“计算权重梯度 (W)”,并把无串行依赖的 W 计算强行塞进流水线的空闲气泡中,完美榨干算力,但工程实现极其复杂。
- 张量并行 (Tensor Parallelism, TP):
- 按矩阵乘法的宽度切分(例如将一个巨大的 MLP 矩阵垂直和水平切分给不同的 GPU 算,然后求和)。
- 它没有流水线的气泡问题,且不消耗 Batch Size 资源,但在前向和反向传播的每一层都需要极其频繁的 All-reduce 同步激活值。
- 黄金法则:由于极度渴求通信带宽,张量并行只能被限制在单一物理节点(通常最多 8 张 GPU)的内部使用。跨节点使用会导致性能断崖式下跌(暴降 40%~65%)。
四、激活值显存与序列并行 (Sequence Parallelism)
- 激活值刺客:随着序列变长,即便用了 TP 和 PP,激活值显存仍会爆炸。因为 TP 只能均分那些进行了矩阵乘法的操作,而像 LayerNorm 和 Dropout 这样的非矩阵运算操作依然会在每个 GPU 上产生冗余的激活值。
- 序列并行 (SP):将非矩阵乘法的操作沿着序列维度(Sequence dimension)切分给各个 GPU,配合 All-gather 和 Reduce-scatter 进行通信,从而彻底消除最后的显存冗余。
- 配合激活值重计算(Activation Recomputation),可以用廉价的算力换取极其昂贵的显存,从而支持更大的 Batch Size。
五、3D 并行组合法则与真实大模型案例
讲师最后给出了工业界组合这三种并行的 “3D 并行黄金法则”:
- 第一步:优先拉满张量并行(TP),直到塞满单个机器的 8 张 GPU,以容纳模型和激活值。
- 第二步:如果模型依然塞不下,跨机器使用流水线并行(PP)或 ZeRO-3(FSDP)继续切分。
- 第三步:在满足显存后,用剩余所有的 GPU 规模来做数据并行(DP),以横向扩展总算力。
前沿案例应用:
- DeepSeek V3:使用了 16路流水线并行、64路专家并行(类似于处理 MoE 的 TP)和 ZeRO-1 数据并行。
- Llama 3:严格遵循 TP=8 + 序列并行 + 流水线并行 + 数据并行的堆叠策略。有趣的是,Meta 在训练 Llama 3 时经历了 148 次物理 GPU 损坏中断,这表明在超大规模训练中,容错系统甚至比并行算法本身更加重要,有时 GPU 甚至会发生静默的数据损坏(Silent data corruption)产生垃圾数据。
六、核心概念问答 (Q&A)
Q1:什么是 All-reduce(全规约)、Broadcast(广播)、Reduce-scatter(规约散播)和 All-gather(全收集)?
在大型大语言模型的分布式训练中,我们需要极其频繁地在多个 GPU(甚至跨服务器节点)之间传递参数、梯度或激活值。为了高效管理这些复杂的数据流动,底层硬件和软件系统(如 NVIDIA 的 NCCL 库以及 PyTorch 的 torch.distributed)提供了一系列经典的集合通信原语(Collective Operations)。
这四个概念代表了集群节点间协同交换数据的四种最核心的模式:
-
Broadcast(广播)
- 运行机制:一对多的通信。将某一个特定节点(Rank)上掌握的完整张量(Tensor),一模一样地复制并广播给集群中的所有其他节点。
- 效果:操作前只有 Rank 0 有数据 A;操作后所有 Rank 手里都有了数据 A。
-
All-gather(全收集)
- 运行机制:多对多的拼接通信。起初,每个节点手里都只掌握着数据的一个“切片”(例如把一个大矩阵按节点数切分成了多份,每人拿一份)。All-gather 会把所有节点上的这些切片收集起来并拼接完整,然后把这份完整的数据发送给每一个人。
- 效果:操作前每个人只有自己的局部片段;操作结束后,所有节点都拥有了完整的全局数据。
-
Reduce-scatter(规约散播)
- 运行机制:多对多的“聚合后打散”通信。它包含两步逻辑上的动作:首先,它对所有节点提供的张量执行“规约(Reduce)”操作(通常是将各个节点上张量的对应位置进行相加求和)。但是,它不会把完整的求和结果发给所有人,而是将最终的结果切碎(Scatter),让每个节点只分到汇总结果的其中一小块。
- 效果:操作前每个人都有一个完整尺寸的输入;操作后每个人只获得了一部分求和结果。
-
All-reduce(全规约)
- 运行机制:多对多的“聚合后广播”通信。系统会将所有节点上的输入张量进行规约计算(例如求和汇总),并将完整的最终聚合结果复制给每一个节点。
- 应用场景:这是朴素数据并行(DDP)中最核心的操作。每张卡在反向传播中独立算出了自己的梯度后,系统会强行插入一步 All-reduce(求和取平均),从而确保下一轮更新前,所有卡拿到的梯度是完全一致的全局梯度。
Q2:ZeRO 优化是如何利用等价公式(All-reduce = Reduce-scatter + All-gather)节省显存的?
在这节课中,讲师着重强调了一个非常核心的底层等价关系:All-reduce = Reduce-scatter + All-gather。
这个数学和通信层面的等价性,正是目前大模型训练中节省显存的终极武器——ZeRO 算法(包括 FSDP)的基石:
在传统的朴素数据并行中,更新参数需要执行庞大的 All-reduce,这不仅耗费极高的通信带宽,还要求每一张显卡都必须在显存里保留一份完整的、极其庞大的“优化器状态(Optimizer State)”。
为了打破显存瓶颈,ZeRO 算法利用了上述等价性,将原本的 All-reduce 拆成了两步:
- 第一步(Reduce-scatter 汇总梯度):各个 GPU 算出梯度后,不再需要全局完整的梯度。通过 Reduce-scatter,每个 GPU 只收集属于它负责更新的那一小部分参数的梯度之和。拿到局部梯度后,每个 GPU 就利用自身显存里保存的局部优化器状态,完成这部分参数的更新。
- 第二步(All-gather 拼回完整模型):等大家都各自把自己负责的参数切片更新完毕后,再执行一次 All-gather,将全网更新好的局部参数广播拼接起来,组成一个完整的最新模型,以供下一次前向传播使用。
总结:借由原语的拆分替换,集群在完全不增加额外总通信带宽成本(两者产生的通信字节量一致)的前提下,成功将优化器状态和梯度的显存占用从“每张卡存一份(冗余)”降到了“全网只存一份(切片分摊)”,从而释放出了巨大的显存空间去容纳百亿甚至千亿参数的大模型。
我们可以用一个“4个合伙人算账”的生活例子,来非常直观地理解 All-gather、All-reduce 以及那个神奇的等价公式到底是怎么运作的。
假设有 4 个合伙人(分别代表 4 张 GPU),他们共同经营一家公司,公司有 4 个部门的开销需要核算:研发、市场、销售、行政(分别代表模型的 4 块参数)。
每个合伙人手里都有一本自己记录的账本(代表不同的数据批次算出的不同梯度/误差),上面记着这 4 个部门的开销。目前的终极目标是:算出每个部门的总开销,并让每个人手里都拿到一份完整的最终总账单。
- 简单粗暴的方法:All-reduce(全规约)
这就像是开了一个大会:
- 动作:4 个人同时把各自账本上的 4 个部门开销大声念出来,然后大家各自在小本子上把所有的数加起来(执行数学求和)。
- 结果:开完会,每个人手里都有了一份完整的、加总过后的总账单。
- 代价:通信量非常大,大家喊来喊去,而且每个人都要做所有的加法计算。最关键的是,每个人在开会时,桌子上都必须摆下整个公司庞大的总账册(占用极大的显存)。
- 只拼接不计算的方法:All-gather(全收集)
假设现在情况变了,不需要做加法了。假设 4 个合伙人每个人本来就只负责保管最终总账单的四分之一:
- 合伙人 0 手里只有【研发】的最终账单;
- 合伙人 1 手里只有【市场】的最终账单;以此类推。
- 动作:此时调用 All-gather。合伙人 0 说“这是研发的账”,把它复印 3 份发给其他人。合伙人 1 说“这是市场的账”,也复印发给其他人。大家只是互相交换和拼接数据,不做任何加减法。
- 结果:交换拼接完毕后,大家都集齐了 4 个部门的完整账单。
- 见证奇迹的时刻:ZeRO 优化是如何拆分 All-reduce 的?
实际训练大模型时,如果直接用 All-reduce,每张显卡都要保留完整的优化器状态(厚重的总账册),显存会直接爆掉。所以系统底层利用了公式:All-reduce = Reduce-scatter + All-gather。他们是这样巧妙合作的:
-
第一步:Reduce-scatter(规约散播 —— 分工算账)
大家不再同时算所有的账了,而是分工合作。
合伙人 0 说:“我专门负责算【研发】的账,你们把账本里关于研发的开销都发给我,我来求和。”
合伙人 1 说:“我负责算【市场】的账,你们把市场的开销都发给我。”
动作结果:这就是 Reduce-scatter。做完这一步,求和(Reduce)完成了,但结果打散(Scatter)在了不同人手里。现在合伙人 0 算出了研发的总开销,合伙人 1 算出了市场的总开销。 -
第二步:更新参数(ZeRO 节省显存的魔法)
既然合伙人 0 只负责更新【研发】的数据,那么他的显存里就只需要保存【研发】这一小部分的优化器状态(历史账本),其他三个部门的历史账单他直接扔掉不管了。这瞬间帮每张显卡省下了四分之三的显存!他拿着刚才汇总的研发总开销,更新好研发的最终参数。 -
第三步:All-gather(全收集 —— 拼回完整结果)
现在,每个人都把自己负责的那部分参数更新完毕了。
大家再执行一次 All-gather,互相交换拼接一下刚才算好的最终结果。
最终结果:大家都拿到了一份完整的最新模型总账单。
总结:All-reduce 是一步到位,大家一起算、一起拿结果,但要求每个人显存里都装下整个世界。而在底层,通过先让大家分工各自求和一小块(Reduce-scatter),各自用极小的显存更新完这一小块后,再互相交换拼图(All-gather),最终达到了和 All-reduce 一模一样的效果,却把显存消耗降到了最低。
Q3:优化器状态(Optimizer States)为什么会占据如此庞大的显存?
在大型语言模型的混合精度训练中,Adam 优化器之所以占用极其庞大的显存,是因为它不仅需要维持模型本身的参数,还需要为每个参数额外存储高精度的“主权重”历史梯度的统计信息(一阶和二阶动量)。
我们可以通过一笔精确的“显存账”来直观理解。在典型的混合精度训练中,平均每个模型参数大约需要消耗 16 字节的显存,其中优化器状态占据了绝对的主导地位:
- 模型参数(Parameters):通常使用 BF16 或 FP16 格式保存,占用 2 字节。
- 梯度(Gradients):同样使用 BF16/FP16 格式进行反向传播,占用 2 字节。
- Adam 优化器状态(Optimizer States):独占了剩下的 12 字节,它具体由以下三个部分组成:
- FP32 主权重(Master weights,4 字节):因为 BF16/FP16 精度较低,在累加极小的梯度更新步时容易发生数值丢失。因此,Adam 必须在显存中单独保留一份 32 位单精度(FP32)的模型权重副本来进行高精度的参数更新。
- 一阶动量估计(First moment estimates,4 字节):Adam 会追踪历史梯度的滑动平均值(即动量),这部分状态通常也需要以 32 位(FP32)或至少较高的精度存储。
- 二阶动量估计(Second moment estimates,4 字节):Adam 还需要追踪过去梯度的平方的滑动平均值(即方差信息),用以自适应地调整每个参数的学习率,同样占用额外的空间。
Q4:既然参数和优化器被切分了,为什么激活值依然是显存刺客?
在大型语言模型的训练中,激活值(Activations)之所以占据极大的显存,主要由以下几个核心因素决定:
-
反向传播的数学必然性与层层累加
在神经网络的运行机制中,前向传播时每一层计算出的激活值,必须原封不动地保留在显存中,直到反向传播阶段用来计算梯度(根据链式法则,计算某层权重的梯度必须依赖该层的输入激活值)。因此,在训练过程中,随着前向传播不断往深层推进,激活值的显存占用会不断累加(“越来越大”),直到反向传播开始并逐层使用完它们之后,这些显存才会被逐渐释放。模型的层数(Layers)越多,囤积在显存里的激活值总量就越庞大。 -
庞大的维度乘数与注意力机制的“平方级爆炸”
根据课程中推导的 Transformer 单层激活值显存公式: s b h × 34 + 5 a s 2 b sbh \times 34 + 5as^2b sbh×34+5as2b(其中 s 为序列长度,b 为批次大小,h 为隐藏层维度,a 为注意力头数):
- 线性增长部分 ( s b h sbh sbh):前馈网络(MLP)和其他逐点运算(Point-wise operations)产生的激活值,会随着批次大小和隐藏层维度的增大而呈线性增长。
- 平方级增长部分 ( s 2 b s^2b s2b):这是最致命的部分。在计算自注意力机制(如 Softmax 操作)时,需要计算任意两个 Token 之间的关系,因此其生成的激活值大小与序列长度 (s) 的平方成正比。这意味着当处理长文本时,这部分激活值的显存占用会呈现爆炸式飙升。
- 并行切分中的“顽固分子”
虽然我们可以将极其占显存的“模型参数”和“优化器状态”完美切碎分摊给不同的 GPU,但激活值却很难被完全切分。即使使用了张量并行(Tensor Parallelism),像 LayerNorm 和 Dropout 这种非矩阵乘法的逐点操作依然会在每张显卡上产生冗余的激活值内存占用。
如何解决?正是因为激活值对显存的消耗极其恐怖,工程师们不得不引入激活值重计算(Activation Recomputation / Checkpointing)技术——宁愿在反向传播时消耗算力把激活值重新算一遍,也不把它们存进显存里。此外,结合 Flash Attention 技术,也能通过在线分块计算(Tiling)直接在局部极速缓存中完成 Softmax,从而避免将庞大的 s 2 s^2 s2 注意力矩阵写入全局显存。
Q5:流水线并行 (Pipeline Parallelism) 和张量并行 (Tensor Parallelism) 的原理与优缺点?
一、流水线并行 (Pipeline Parallelism, PP) —— “接力跑工厂”
流水线并行是按模型的“层(Layers)”或者说“深度(Depth)”来切分的。假设你的模型有 4 层,你有两台机器。你可以把第 1、2 层放在机器 A 上,把第 3、4 层放在机器 B 上。机器 A 算完前面的层后,把中间结果(激活值)传给机器 B,机器 B 接着往下算。这就像工厂流水线的接力作业。
- 痛点:严重的“气泡(Bubble)”。流水线的致命弱点是干等。当机器 A 在处理第一批数据时,机器 B 没有任何数据可算,只能处于闲置状态(Idle);反之,当机器 A 传完数据给 B,在 B 算完并把误差反向传回来之前,A 又得干等。这种算力处于闲置、白白流失的时间片段,在图表上看起来就像一个个空白的泡泡,因此被称为“气泡(Bubble)”。
- 基础优化:微批次(Micro-batches)。为了不让机器闲着,我们可以把原本很大的一批数据(Batch)切分成好几块更小的“微批次”。机器 A 只要算完“微批次1”,就赶紧甩给机器 B,然后机器 A 立刻开始算“微批次2”。这样两台机器就能尽量交叠工作,大幅缩小了等待的“气泡”时间。
- 进阶魔法:Zero Bubble / DualPipe(零气泡)。这是一种将算力压榨到极限的黑科技。在反向传播(计算误差梯度)时,其实包含两部分计算:计算激活值的梯度(必须按顺序等前面的传过来)和计算权重(W)的梯度(随时可以算,没有先后依赖)。工程师们写出了极其复杂的调度代码,把计算权重 W 梯度的任务,精准地“塞进”机器原本在干等的那些气泡时间里。这样显卡就永远在全速运转,但因为需要干预底层计算图的执行顺序,工程实现简直是一场噩梦。
二、张量并行 (Tensor Parallelism, TP) —— “大家一起算同一道大题”
张量并行不按层切,而是直接把一层内部极其庞大的矩阵乘法,顺着宽度(Width)切成几块。假设这层有一个巨大的矩阵,机器 A 和 B 各自只拿这个大矩阵的一半。数据进来后,A 和 B 同时开始算自己负责的那一半。
- 优点:没有气泡,不吃 Batch Size。大家是同时干活的,所以完全没有流水线那种“你等我、我等你”的气泡问题。而且它不需要像流水线那样把批次(Batch)切碎,所以不消耗批次大小这一稀缺资源。
- 痛点:频繁的 All-reduce 同步通信。既然一道大题被拆给了两台机器算,那么算完之后,A 和 B 必须把各自的局部结果加起来,才能得到这一层最终完整的正确答案(特征输出)。这个“汇总加和并广播给所有人”的操作,就是 All-reduce。在前向和反向传播中,每一层的计算都需要这样停下来强制同步一次,导致极其频繁的通信。
- 黄金法则:为什么被锁死在单机 8 卡?正是因为每一层都要极度频繁地交换庞大的矩阵结果,张量并行对通信带宽的渴求极高。在物理硬件上,同一台机器(Node)内部的 8 张 GPU 是通过 NVLink 这种超高速专用通道直接相连的,速度极快(可达 900 GB/s),勉强能承受这种频繁的 All-reduce。但如果你想跨机器使用张量并行,数据就必须经过慢得多的机房网络交换机(比如 Infiniband)。一旦跨机器(例如扩展到 16 卡或 32 卡),GPU 把大量时间花在了等待网络传输上,吞吐量性能会遭遇断崖式暴跌(下跌高达 42% 甚至 65%)。因此,业界铁律是:张量并行只能在单台物理机(最多 8 张 GPU)的内部使用。
七、第七讲复习题
基于斯坦福 CS336 第七讲关于**大规模分布式并行(Parallelism 1)**的核心内容,为您精心整理了对应的复习题。
📝 第七讲复习题 (Lecture 7: Parallelism 1)
一、集合通信与底层规律 (Collective Communications)
- 通信的数学等价性:在大模型的分布式通信中,All-reduce(全归约)操作在数学和通信成本上,完美等价于哪两个底层集合通信操作的组合?为什么这个等价关系在后来的 ZeRO 优化中极其重要?
- 内存刺客 (The Memory Hog):在传统的朴素数据并行(Naive Data Parallelism)中,如果我们在 GPU 上训练模型,存储模型的总显存占用其实远超模型参数本身的大小。导致显存暴涨的“罪魁祸首”是哪一部分?
二、数据并行与 ZeRO 优化 (Data Parallelism & ZeRO/FSDP)
- ZeRO 的三个切分阶段:微软提出的 ZeRO 优化(零冗余优化器)分为 Stage 1、2、3。请问这三个阶段分别切分(Shard)了哪些模型状态变量?PyTorch 中大名鼎鼎的 FSDP 对应的是哪一个阶段?
- FSDP 的“重叠魔法” (Communication Overlap):既然 ZeRO-3 (FSDP) 连模型参数都切碎并分布到不同机器上了,这意味着在每一层的前向和反向传播时都要立刻发起通信去要回参数。为什么这种疯狂的通信操作没有严重拖慢 GPU 的训练速度?
- 数据并行的终极瓶颈:为什么我们不能单纯靠无脑堆叠数据并行(Data Parallelism)来无限扩展系统规模?在数据并行中,什么参数是受限的“核心系统资源”?
三、模型并行 (Model Parallelism: Pipeline & Tensor)
- 流水线并行的致命伤 (Pipeline Bubble):如果我们简单地按层切割模型并分配给不同的 GPU(流水线并行),系统会面临什么最致命的效率问题?业界通常用哪两种策略来缓解这个问题?
- TP 与 PP 的网络拓扑部署:在真实的数据中心物理拓扑中,张量并行(TP)和流水线并行(PP)的部署原则有什么根本不同?分别适合部署在什么样的硬件连接链路上?
- 张量并行 (Tensor Parallelism) 的通信开销:当使用 TP 切分一个大型 MLP 的矩阵乘法时,在前向传播和反向传播中,GPU 之间通常需要触发什么同步操作?
四、激活值显存与 3D 并行综合 (Activation & 3D Parallelism)
- 序列并行的救场 (Sequence Parallelism):即使我们使用了张量并行(TP)切分了矩阵乘法,系统依然会遗留很大一部分冗余的激活值显存占用。这部分顽固的显存是由什么操作产生的?序列并行(SP)是如何解决它的?
- 大厂的 3D 并行经验法则 (Rule of Thumb):假设你现在拥有几千张卡,要训练一个超大模型,结合系统通信带宽的物理限制,讲师给出的标准 3D 并行“部署顺序(Rule of Thumb)”是怎样的?
八、参考答案与知识点解析
-
通信的数学等价性?
答案:All-reduce 完全等价于 Reduce-scatter 加上 All-gather。这个身份等价性极其重要,因为在传统的梯度同步中我们需要一次 All-reduce;而在 ZeRO 优化中,我们可以先用 Reduce-scatter 聚合梯度,各自用分片的优化器状态更新分片的参数,最后再用 All-gather 把更新后的参数广播出去。这样做的总通信带宽成本和原来完全一样,但却免费获得了显存的极大节省。 -
内存刺客 (The Memory Hog)?
答案:罪魁祸首是优化器状态(Optimizer States)。以 Adam 优化器为例,除了保存 2 字节(BF16/FP16)的权重和 2 字节的梯度外,还需要保存 4 字节的 FP32 原始主权重(Master weights)、4 字节的动量(First moment)以及 4 字节的方差(Second moment)。这导致每个参数需要占用高达 16 字节的显存,优化器状态占据了绝大部分空间。 -
ZeRO 的三个切分阶段?
答案:
- Stage 1:仅切分优化器状态(Optimizer States)。
- Stage 2:切分优化器状态和梯度(Gradients)。
- Stage 3:切分优化器状态、梯度以及模型参数(Parameters)。
FSDP(Fully Sharded Data Parallel)就是完全等价于 ZeRO Stage 3 的 PyTorch 实现。
-
FSDP 的“重叠魔法” (Communication Overlap)?
答案:系统利用了通信与计算的重叠(Overlapping communication and computation)。当 GPU 正在计算当前层的矩阵乘法时,系统会在后台利用网络带宽提前预取(Prefetching)下一层所需的参数数据。当计算完成进入下一层时,参数已经加载完毕,从而极大地隐藏了通信的延迟泡沫。 -
数据并行的终极瓶颈?
答案:数据并行高度消耗批次大小(Batch Size)这一核心资源。你无法将机器数量扩展到超过全局批次大小,因为每张卡至少需要分配 1 个样本。此外,随着批次大小增加,它会面临强烈的边际收益递减(达到“临界批次大小” Critical Batch Size 后,额外的样本对优化梯度的帮助微乎其微)。 -
流水线并行的致命伤 (Pipeline Bubble)?
答案:致命伤是流水线气泡(Pipeline Bubble),因为后续的 GPU 必须等待前面的 GPU 计算完才能开工,导致绝大部分时间 GPU 处于空闲闲置状态(利用率极低)。
缓解策略 1:将大批次切分为多个微批次(Micro-batches),使得不同层可以交替流水线作业以缩小气泡比例。
缓解策略 2:采用如 1F1B 或 Zero Bubble (DualPipe) 这样的极致调度策略,把反向传播中用于计算权重梯度的独立任务,塞进激活值反向传播产生的空闲气泡中进行计算。 -
TP 与 PP 的网络拓扑部署?
答案:
- 张量并行 (TP):需要极高的通信带宽和极低的延迟,因此通常严格限制在单台物理机/单个节点内部(通过最高速的 NVLink 连接,通常不超过 8 张 GPU)。
- 流水线并行 (PP):只需在层与层之间进行点对点(Point-to-point)传递激活值,通信量相对较小。因此非常适合部署在跨机器、跨机架甚至跨数据中心的慢速网络链路(如以太网/Infiniband)上。
-
张量并行 (Tensor Parallelism) 的通信开销?
答案:在切分了矩阵后,无论是前向传播的最后,还是反向传播的最后,所有参与张量并行的 GPU 都必须执行一次 All-reduce 来对部分和(Partial sums)进行汇总或对梯度进行同步。这也是为什么它极其消耗通信带宽的原因。 -
序列并行的救场 (Sequence Parallelism)?
答案:张量并行只能切分矩阵乘法运算。而在模型中,像 LayerNorm、Dropout 等逐点运算(Point-wise operations)产生的激活值依然会在所有并行 GPU 上重复存储一份冗余的完整拷贝,这部分冗余激活值显存会随规模线性暴增。序列并行 (SP) 的思路就是将这些不需要跨 Token 交互的逐点运算,直接沿着序列维度(Sequence dimension)切分开交给各个 GPU 独立计算,算完后再用集合通信拼起来,从而完美消灭这最后的显存冗余。 -
大厂的 3D 并行经验法则 (Rule of Thumb)?
答案:
- 首先使用 张量并行 (TP) 直到单机的极限(通常单节点 8 张卡),确保把庞大的模型和激活值装进显存,因为这是最高效的。
- 如果单机依然装不下模型,就在跨节点机器间引入 流水线并行 (PP) 或 ZeRO-3 (FSDP)。
- 最后,一旦你的单体模型能在上述组合中跑起来,就用 数据并行 (DP) 横向复制这套组合,填满集群中剩下的所有 GPU,以此来提升整体的算力吞吐和训练速度。
AtomGit 是由开放原子开源基金会联合 CSDN 等生态伙伴共同推出的新一代开源与人工智能协作平台。平台坚持“开放、中立、公益”的理念,把代码托管、模型共享、数据集托管、智能体开发体验和算力服务整合在一起,为开发者提供从开发、训练到部署的一站式体验。
更多推荐

所有评论(0)