本题要求把 N 层网络划分成 K 段连续区间,每段分给一个 NPU。
每个区间需要同时满足:
在华为昇腾(Ascend)集群上训练千亿参数的大模型时,由于单个昇腾 NPU 的 HBM 显存无法装下整个模型,我们通常会采用流水线并行(Pipeline Parallelism)技术。该技术会将模型按序切分为多个连续的阶段(Stage),并将每个阶段分配到集群中不同的 NPU上。
现有一个包含 N 层的神经网络模型,需要按顺序切分并部署到 K 个昇腾NPU上,即划分为 K 个连续的流水线阶段。每个 NPU 负责一个阶段,且每个NPU至少需要分配 1 层网络。
已知:
第 i 层网络前向与反向传播的计算耗时为 C[i]。
第 i 层网络所需的显存容量为 W[i]。
每个昇腾NPU的显存物理上限为 M。
在流水线并行中,整个集群的吞吐效率受限于最慢的那个 NPU(即计算耗时最大的阶段,被称为流水线瓶颈)。一个阶段的计算耗时等于分配给它的所有网络层 C[i] 的总和,其所需显存等于这些网络层 W[i] 的总和。
如果任意一个阶段的显存总和超过了单卡上限 M,则会导致 OOM(内存溢出),该切分方案无效。
注意:由于神经网络计算图的数据流向依赖,分配给同一个 NPU 的多个网络层必须在原始顺序中是连续的,绝对不可打乱或跨层自由组合。
你的任务:
请设计一个算法,在保证不发生 OOM 的前提下,找到一种切分方案,使得所有 NPU 中最大的计算耗时尽可能小。请输出这个最小的“最大计算耗时”。如果无法找到满足显存限制的切分方案,或者 NPU 数量多于网络层数(无法保证每卡至少 1 层),请输出 −1。
输入共三行,所有数字均通过空格分隔。
第一行:三个正整数 N K M,分别代表网络总层数、NPU数量(阶段数)、单卡显存上限。
第二行:N 个正整数,代表每一层的计算耗时数组 C。
第三行:N 个正整数,代表每一层的显存需求数组 W。
一个整数,表示在满足显存约束下,最小化的最大计算耗时。如果无解,输出 −1。
输入
3 4 100
1 1 1
1 1 1
输出
-1
说明
NPU 数量(4)大于模型层数(3),由于必须保证每个 NPU 至少分配一层网络,无法满足条件,返回 −1。
输入
4 2 10
2 2 2 2
6 6 6 6
输出
-1
说明
单卡显存上限为 10,但每层都需要 6 的显存,如果将任意两层放在同一个 NPU 上,显存需求就是 12,会发生 OOM。因此,这 4 层网络必须每层单独占用一个 NPU(需要 4 个 NPU),但集群只有 2 个 NPU,故无法完成切分,返回 −1。
输入
5 3 20
5 1 2 3 4
10 5 5 5 10
输出
6
说明 模型共有 5 层,需要分配到 3 个 NPU 上,单卡显存上限为 20。
最佳切分方案为:[层 0, 层 1] 分给 NPU1,[层 2, 层 3] 分给 NPU2,[层 4] 分给 NPU3。
NPU1:耗时 5+1=6,显存 10+5=15(≤20,合法)
NPU2:耗时 2+3=5,显存 5+5=10(≤20,合法)
NPU3:耗时 4,显存 10(≤20,合法)
最大计算耗时为 6,没有比 6 更优的合法方案。