题目内容

在华为昇腾（$Ascend$）集群上训练千亿参数的大模型时，由于单个昇腾 $NPU$ 的 $HBM$ 显存无法装下整个模型，我们通常会采用流水线并行（$Pipeline$ $Parallelism$）技术。该技术会将模型按序切分为多个连续的阶段（$Stage$），并将每个阶段分配到集群中不同的 $NPU$上。

现有一个包含 $N$ 层的神经网络模型，需要按顺序切分并部署到 $K$ 个昇腾NPU上，即划分为 $K$ 个连续的流水线阶段。每个 $NPU$ 负责一个阶段，且每个NPU至少需要分配 $1$ 层网络。

已知：

第 $i$ 层网络前向与反向传播的计算耗时为 $C[i]$。

开通会员即可查看完整视频题解： 1.题目讲解 2.思路分析 3.逐行代码手写