多GPU环境下模型并行与数据并行的混合部署要点

文章导读
混合部署的核心是先决定哪一层并行在解决当前问题。数据并行解决吞吐,模型并行解决单卡装不下。不要在一开始就把张量并行、流水线并行、数据并行全部打开;打开的顺序和资源分配顺序,由模型体积和机器拓扑共同决定。通常的建议是:单卡能稳定跑完一个最小批次,就用纯数据并行;跑不下去,再根据机器拓扑在TP和PP之间做出选择;单副本跑通后还不够吞吐,才叠加数据并行。
📋 目录
  1. A 先判断模型体积是否真的超出一张卡
  2. B 模型并行拆哪一维,先看卡间拓扑
  3. C 数据并行叠加到模型副本之外,而不是替换模型并行
  4. D 落地时按这个顺序检查
A A

混合部署的核心是先决定哪一层并行在解决当前问题。数据并行解决吞吐,模型并行解决单卡装不下。不要在一开始就把张量并行、流水线并行、数据并行全部打开;打开的顺序和资源分配顺序,由模型体积和机器拓扑共同决定。通常的建议是:单卡能稳定跑完一个最小批次,就用纯数据并行;跑不下去,再根据机器拓扑在TP和PP之间做出选择;单副本跑通后还不够吞吐,才叠加数据并行。

混合部署的价值是让模型完整跑起来并利用多卡吞吐。先用模型体积估算显存占用、再用机器拓扑选主并行维度,最后用数据并行扩展吞吐。张量并行要留在节点内、流水线并行可以跨节点,数据并行看数据加载速度。是否引入模型并行,以单卡能不能稳定跑一个最小批次为准。

先判断模型体积是否真的超出一张卡

用粗略方式估算:把参数量乘上当前精度的字节数,得到权重文件最低占用;训练时还要算上梯度、优化器状态和激活值,所以模型权重能放进卡,不代表训练流程能放进卡。最直接的判断是跑一个最小批次,观察显存峰值接近上限还是已经OOM。

在PyTorch中可以用这段代码在训练结束时输出本进程的显存峰值(单位是字节):

# 训练结束后读取本进程显存峰值
print(torch.cuda.max_memory_allocated())

如果这个值已经接近物理显存上限,数据并行方案就成了硬约束:再加卡也只能增加吞吐,不能解决单卡放不下的问题。此时才需要进入模型并行。

多GPU环境下模型并行与数据并行的混合部署要点

模型并行拆哪一维,先看卡间拓扑

张量并行和流水线并行都叫模型并行,但通信模式差别很大。张量并行会在每一步计算中同步拆分的矩阵运算,通信频率高;流水线并行只在层边界传递激活值,通信频率低但会引入气泡。因此张量并行要放在同一个高速互联域里,比如同一台机器或同一个NVSwitch/NVLink域;流水线并行可以放宽到跨节点。

查看拓扑的标准命令是:

nvidia-smi topo -m

矩阵中标出NVLink或较低延迟互连的卡组优先做TP,只能走PCIe甚至跨节点互连的卡组优先做PP或DP。

多GPU环境下模型并行与数据并行的混合部署要点

数据并行叠加到模型副本之外,而不是替换模型并行

当TP和PP把模型切成一份能跑的副本后,再在多份副本之间做数据并行。数据并行与模型并行不是互斥关系,而是先用模型并行解决单副本显存,再通过增加副本推进吞吐。

叠加时注意两点:

  • 数据并行副本数不是越多越好。每份副本都会重新占一份显存,batch size增大后整体吞吐是否上升,要看数据加载能否跟上。
  • 各副本在梯度同步时存在快的等慢的问题。如果某些卡利用率明显低于其他卡,先排查数据加载不均匀,再排查TP或PP维度上的通信过载。

落地时按这个顺序检查

  1. nvidia-smi -L 确认总GPU数和每张卡编号。
  2. nvidia-smi topo -m 确认NVLink/PCIe分组,把张量并行限制在高速域内。
  3. 先用最小批次跑一次,观察torch.cuda.max_memory_allocated()的峰值;接近上限再调整TP与PP的切分方式。
  4. 用监控工具看各卡利用率曲线。持续有空闲或低使用率卡,多半是流水线气泡或数据加载不均衡。
  5. 扩大batch size时同步调整学习率;具体规则不同框架不同,需要按实际框架确认。
  6. 不要只看单卡显存,还要看通信链路流量和GPU利用率,才能定位瓶颈。

最后留一条判断标准:模型能放进单卡时,混合部署带来的复杂度往往超过收益;不要为了用满多卡而强行引入模型并行。