大模型分布式训练并行策略详解(1) 2.2 张量并行(Tensor Parallelism, TP) 核心思想 :模型纵向切分,层内分割。把单层网络的参数矩阵切分到多卡,每张卡只存部分权重。 以Transformer的MLP层为例(简化表示): plaintext 原始计算: Y GeLU(X @ A) @ B A: [hiddensize, 4hiddensize] B: [4hiddensize, hiddensize] TP2时的切分方式(按列切分A,按行切分B): ┌──────────────────────────────────────────────────────────┐ │ GPU 0 持有 A[:, :2h] 和 B[:2h, :] │ │ GPU 1 持有 A[:, 2h:] 和 B[2h:, :] │ │ │ │ 前向传播: │ │ GPU 0: Z0 X @ A[:, :2h] → GeLU(Z0) @ B[:2h, :] Y0 │ │ GPU 1: Z1 X @ A[:, 2h:] → GeLU(Z1) @ B[2h:, :] Y1 │ │ │ │ 需要AllGather通信合并 Y Y0 + Y1 │ └──────────────────────────────────────────────────────────┘ 2.3 流水线并行(Pipeline Parallelism, PP) 核心思想 :模型横向切分,层间分割。把不同层放到不同卡,数据像流水线一样流动。 气泡问题与GPipe改进: 朴素流水线(如GPipe)把batch切成更多微批次(microbatches),让GPU尽可能填满: plaintext 微批次数量 4, PP度数 4 时的填充效果: GPU 0: [F0][F1][F2][F3][B3][B2][B1][B0] GPU 1: [F0][F1][F2][F3][B3][B2][B1][B0] GPU 2: [F0][F1][F2][F3][B3][B2][B1][B0] GPU 3: [F0][F1][F2][F3][B3][B2][B1][B0] 气泡率 ≈ (PP1)/(PP+M1) M为微批次数量,M越大气泡越小 1F1B(One Forward One Backward)调度:当前主流方案,显存更优,但气泡稍大。 关键特征: 显存占用:每卡只存部分层,显存与层数成正比减少 加速效果:接近线性加速(通信极少),但受气泡影响,通常8095%效率 硬性限制 :PP度数 ≤ 总层数,且最好整除;层间通信量小(激活值),可跨节点
来自: