张量并行

定义

张量并行(Tensor Parallelism)是模型并行的一种,将单个层内的权重矩阵按行或列切分到多块 GPU,各卡计算部分结果后通过集合通信(all-reduce)合并,对上层逻辑表现为一个完整的层。

原理

以 Transformer 为例,Megatron-LM 的经典切法是:注意力按头切分,前馈层第一个矩阵按列切、第二个按行切,如此一个 FFN 块只需一次 all-reduce。由于每层前向反向都要通信,张量并行对带宽极其敏感,通常限制在单机 NVLink 互联的 4 到 8 卡内。

应用

张量并行解决「单卡放不下一层计算」的问题,与流水线并行(按层切分)、数据并行(按样本切分)组合成 3D 并行,是训练千亿级模型的标准配置;推理框架如 vLLM、TensorRT-LLM 也用它把大模型摊到多卡服务。