张量并行负载均衡策略：让每位 GPU 厨师都不空转

可视化演示

动态演示张量并行负载均衡策略的核心机制，通过入门/进阶/实战/优化四种模式展示关键概念与交互流程。

负载不均成因

厨房里，八位厨师领到了同样大小的蛋糕块。乍一看很公平，但每一块内部奶油、水果和坚果的分布天差地别。有的层夹着厚厚的水果，要一颗颗摆整齐；有的层只有薄薄一层奶油，抹两下就完事。于是工作量大的厨师忙得满头大汗，工作量小的厨师早早站在一边喝咖啡。整个蛋糕上桌的时间，永远由最慢的那位决定——这就是后厨的“木桶效应”。

把镜头切到 GPU 集群，道理完全相同。张量并行把大矩阵切成若干子块分给多张卡。矩阵不同区域的非零元素密度、算子复杂度都不一样，天然就有“厚层”和“薄层”；再加上集群里可能混着 A100、A800、4090 等不同算力的卡，快卡等慢卡的现象被进一步放大。通信也会掺和一脚：如果某张卡总是转发别人的数据，它就成了新的瓶颈，形成长尾延迟。

为了量化这种忙闲不均，我们引入 GPU 利用率的概念：让 N 张卡并行工作，每张卡实际有效计算时间为 \(T_i\)，真正能跑完一轮的总耗时却是最慢卡的完成时间 \(T_{\max}\)。那么利用率可以写成：

主厨要救场，第一步不是急着调整蛋糕配方，而是先盯住监控屏找出“谁最晚放下锅铲”。GPU 也是一样，定位到瓶颈卡，才能谈如何均衡。

计算量均衡

发现平均分不公平后，主厨换了思路：给手速快、工具好的厨师多分几层奶油厚的，给动作慢的厨师少分几层薄的。任务单不再追求“每个人拿一样多”，而是“每个人都在差不多的时间收工”。这个朴素的逻辑，正是负载均衡的核心。

在张量并行中，对应动作是“按算力加权分配工作量”。每张 GPU 有自己实测出来的算力，比如相对速度 \(P_i\)。总工作量 \(W_{\text{total}}\) 不该被 N 整除，而应该按每张卡算力占全体算力的比例来分：

主厨每周都会重新测一次厨师们的工作速度，再更新任务单。训练集群也一样，每次跑大任务前做一次微基准测试，算力分配表就能常保持新鲜，这一步是所有后续调优的地基。

通信均衡策略

蛋糕块分得再均匀，最后总要把各块拼成一个整体。拼的时候需要传递奶油、水果、托盘，如果所有厨师都挤同一个窗口，传菜口瞬间堵死。你可能会想：那多开几个窗口不就好了？但每多一个窗口，就要多派人手，而且高峰期依然有人闲着等料。后厨的瓶颈，有时候不在做蛋糕，而在传蛋糕。

张量并行里也有一模一样的“传菜口”：每步计算完，各 GPU 要执行 AllReduce 等集合通信，把各自的中间结果汇总再分发。通信开销由三部分决定：要传的数据量 \(V\)、网络带宽 \(B\)、以及发送次数 \(r\)。每次通信都有固定的启动延迟 \(L\)，于是通信耗时可以写成：

主厨后来发现，只要让不同厨师错峰去传菜窗口，别一窝蜂挤在同一个时间段，传菜口的拥堵就消失了。通信均衡的本质，不是“不传菜”，而是“大家都有菜可传、又都不堵车”。

切分策略调优

蛋糕应该横切成大块，还是竖切成细条？横切的话块数少，厨师之间要传递的切面也少，但每块依旧很大，并行度有限；竖切的话能分给更多人，可每一条都要和其他条交换奶油和水果，传菜次数急剧增加。主厨揉着太阳穴感叹：切法没有绝对好坏，全看这块蛋糕“耐不耐烦等传菜”。

张量并行中的行切分与列切分就是这个选择。权重矩阵可以按行切（类似横切层）也可以按列切（类似竖切条），切分粒度决定了并行度与通信量的比值。为了科学决策，我们定义一个计算通信比：

主厨的经验是：先尝一口蛋糕，判断它是“耐放型”还是“易化型”。耐放型可以大卸八块慢慢做；易化型必须少切快拼。你的模型是哪种类型，跑一次 profile 就能心中有数。

并行度权衡

主厨后来又犯过愁：是不是厨师越多越好？他把厨房从 4 人加到 10 人，结果发现效率反而下降了——灶台只有四个，新来的厨师挤来挤去，传菜通道也乱了，互相绊手，出餐比之前还慢。厨房的空间和传菜通道是有限的，这就是协作的代价。

张量并行同样有“最佳人数”。增加并行度 N，每张卡分到的计算量确实变少了，但通信开销会随着 N 增大而上升；负载不均造成的空闲等待通常也会更明显。因此总训练时间可以建模成三个部分之和：

真正的主厨会综合考虑灶台数量、人手和传菜口，而不是盲目堆人。工程上，张量并行度也常与数据并行度、流水并行度搭配，例如 64 卡下选 TP=4、PP=4、DP=4 的组合，而不是 64 卡全做张量并行。

工程实践要点

开业之后主厨依然不轻松，因为每位厨师的效率会波动：今天有人手快，明天有人状态差；奶油用完要补货，传菜口偶尔还会堵。所以他装了摄像头，盯着后厨的实时画面：谁在玩手机、谁在拼命赶工，发现问题立刻临时调拨人手。

真实的训练集群也需要这样一双眼睛。模型不同层的计算密度不同、序列长度动态变化、特殊算子只在某张卡上出现，这些都会让负载漂移。因此工程实践的第一步是定义负载均衡度指标，用最快卡与最慢卡的完成时间比来衡量忙闲差距：

主厨每晚收工后都要看“谁最晚放下锅铲”的复盘记录，然后修改第二天的任务单。训练工程师也应在每个训练阶段后检查 profiler 报告，把负载均衡度纳入常规迭代节奏，而不是出了问题才想起它。

常见误区辨析

后厨来了一批新厨师，他们有三个常见误解：第一，“每人分一样多就是公平”；第二，“人多一定更快”；第三，“只要把传菜速度提上去，菜就不会等”。结果每次考核都栽跟头。这些误区在 GPU 世界里也同样遍布。

理解这些误区，最好用一个同时惩罚“计算不均、通信开销、扩展损耗”的并行效率公式。假设单卡完成一轮耗时 \(T_1\)，N 卡理想并行时每卡应耗时 \(T_1/N\)，而实际由最慢卡决定的总耗时为 \(T_{\max}\)，那么：

这组公式就像后厨的“总出餐效率”：不看每人分到多少盘子，只看最后一桌菜几点上齐。平均分配、盲目加人、无视通信，这些直觉的坑，都要靠数据指标来填平。

AI 应用场景

Python 演示

# -*- coding: utf-8 -*-def calc_work(W_total, capacities):    total = sum(capacities)    return [W_total * c / total for c in capacities]  def sim_time(work_loads, speeds):    return [w / s for w, s in zip(work_loads, speeds)]  def lb_ratio(times):    return min(times) / max(times)  def report(name, times):    print(name, "各卡耗时:", [round(t, 2) for t in times])    print(name, "负载均衡度:", round(lb_ratio(times), 3))    print(name, "总耗时:", round(max(times), 2))  def main():    W_total = 1000.0    capacities = [1.0, 0.8, 1.2, 1.0]    n = len(capacities)    even_work = [W_total / n] * n    weighted_work = calc_work(W_total, capacities)    report("平均分配", sim_time(even_work, capacities))    report("按算力分配", sim_time(weighted_work, capacities))  if __name__ == "__main__":    main()

1. def calc_work(W_total, capacities):实现按算力比例分配工作量的核心函数，对应正文中的加权公式。
2. total = sum(capacities)先求出所有 GPU 的算力总和，作为分配基准。
3. return [W_total * c / total for c in capacities]每张卡的算力占比乘以总工作量，得到各自的任务量。
4. def sim_time(work_loads, speeds):计算每张卡完成自己任务所需的时间，即工作量除以算力。
5. def lb_ratio(times):用最快卡与最慢卡耗时之比计算负载均衡度 LB，越接近 1 越好。
6. even_work = [W_total / n] * n构造一个“平均分配”的对照组，用来凸显差异。
7. if __name__ == "__main__":标准入口，保证代码可直接运行，也可以作为模块被导入。

常见误区

面试题