题目内容

在训练千亿参数的大模型时，单张 $NPU$ 的 $HBM$ 显存通常无法完整容纳模型。为此，常采用流水线并行（$Pipeline$ $Parallelism$）技术：将模型按网络结构顺序切分为若干连续阶段（$Stage$），并把每个阶段部署到不同的 $NPU$ 上。

现有一个包含 $N$ 层的神经网络，需要按顺序切分为 $K$ 个连续阶段，分配给 $K$ 张 $NPU$。每个 $NPU$ 恰好负责一个阶段，且每个阶段至少包含一层网络。

第 $i$ 层网络的前向与反向传播计算耗时为 $C[i]$，该层所需的显存容量为 $W[i]$。每张 $NPU$ 的显存物理上限为 $M$。一个阶段的总计算耗时等于该阶段内所有层 $C[i]$ 的累加和，总显存需求等于该阶段内所有层 $W[i]$ 的累加和。如果任意一个阶段的显存总和超过 $M$，就会触发 $OOM$，该切分方案无效。由于神经网络计算图存在顺序数据依赖，分配给同一个 $NPU$ 的层必须保持原始顺序连续，不能打乱或跨层组合。

整条流水线的吞吐效率受限于计算耗时最大的那个 $NPU$，即流水线瓶颈。目标是在保证不发生 $OOM$ 的前提下，找到一种合法切分方案，使所有 $NPU$ 中最大的计算耗时尽可能小。请输出这个最小的最大计算耗时。如果不存在合法方案，则输出 -1。

输入描述

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