大模型流水线并行训练优化

在华为昇腾集群上训练千亿参数大模型时，单个 NPU 的 HBM 显存装不下整个模型，需要用流水线并行（Pipeline Parallelism）将模型按序切分为多个连续阶段，分配到不同 NPU 上。

给定 $N$ 层神经网络，需要划分到 $K$ 个 NPU 上，每个 NPU 至少分到 $1$ 层，且分配给同一个 NPU 的层必须在原始顺序中连续。第 $i$ 层的计算耗时为 $C[i]$，显存需求为 $W[i]$，每个 NPU 的显存上限为 $M$。

这道题的本质是：在额外约束（每段显存不超过 $M$）下，将数组切分成 $K$ 段，最小化最大段和。这是经典的二分答案问题。

算法框架：

1. 预处理：检查 $K > N$ 或任意 $W[i] > M$，直接返回 $-1$。
2. 二分答案：$L = \max(C)$，$R = \sum(C)$，每次取 $mid$ 进行贪心验证。
3. 贪心验证：从第 1 层开始，累计计算耗时和显存。当加入下一层会超过 $mid$ 或 $M$ 时，断开新段。
4. 若最终段数 $\le K$，说明 $mid$ 可行，尝试更小；否则需要更大。

给定阈值 $T$，贪心过程如下：

维护当前段的计算耗时和 $t$、显存和 $s$、已用段数 $cnt$。初始 $t = s = 0$，$cnt = 1$。

对每一层 $i$，判断能否加入当前段：

如果任何时刻 $C[i] > T$ 或 $W[i] > M$，说明单个层就超过了限制，$T$ 不可行。实际上 $C[i] > T$ 在我们的二分下界 $L = \max(C)$ 的设定下不会发生，但 $W[i] > M$ 需要在预处理中捕获。

最终如果 $cnt \le K$，则 $T$ 是一个可行的最大耗时上限。

算法步骤：

1. 读入 $N, K, M$ 及数组 $C, W$。
2. 若 $K > N$ 或存在 $W[i] > M$，输出 $-1$。
3. 二分 $L = \max(C)$，$R = \sum(C) + 1$。
4. 贪心验证每个 $mid$。
5. 输出 $L$。

def solve():
    import sys
    data = sys.stdin.read().split()
    idx = 0
    N, K, M = int(data[idx]), int(data[idx+1]), int(data[idx+2])
    idx += 3
    C = [int(data[idx+i]) for i in range(N)]
    idx += N
    W = [int(data[idx+i]) for i in range(N)]

    # 无解判定
    if K > N:
        print(-1)
        return
    for i in range(N):
        if W[i] > M:
            print(-1)
            return

    # 贪心验证：给定阈值 T，最少需要几段
    def check(T):
        cnt = 1
        t = 0  # 当前段计算耗时和
        s = 0  # 当前段显存和
        for i in range(N):
            if t + C[i] <= T and s + W[i] <= M:
                t += C[i]
                s += W[i]
            else:
                cnt += 1
                t = C[i]
                s = W[i]
        return cnt <= K

    # 二分答案：左闭右开 [L, R)
    L = max(C)
    R = sum(C) + 1
    while L < R:
        mid = (L + R) // 2
        if check(mid):
            R = mid
        else:
            L = mid + 1

    print(L)

solve()
  

对于本题的数据规模（$N$ 可达 $10^5$ 级别，$C[i]$ 可达 $10^4$ 级别），$\log(\sum C) \approx 30$ 左右，总操作量约 $3 \times 10^6$，在 2 秒时限内绰绰有余。