深度学习专题：模型训练的张量并行（一）

张量并行的基本原理

（一）张量并行的定义

张量并行（Tensor Parallelism）是一种将单个张量分割到多个 GPU 上进行并行计算的技术，尤其在 Transformer 架构中广泛用于大模型训练和推理。

• 将权重矩阵按行或列拆分到不同设备
• 每个设备持有部分参数，计算部分结果，通过通信 如 all-reduce 聚合完整输出
• 目标是减少单个设备的显存占用，同时利用多设备并行计算

（二）张量并行的切分方式

输入 X: [b, s, h] (batch_size, sequence_length, hidden_size)

权重 W: [h, o] (hidden_size, output_size)

并行度 n: 设备数量

1. 行并行（Row Parallelism）

   • 将权重矩阵按行拆分到不同设备，输入矩阵随之按列拆分
   • 每个设备计算拆分后的输入矩阵与拆分后的权重矩阵的乘积
   • 输出结果通过 all-reduce 聚合到所有设备

   (1) 切分权重矩阵W
   权重W形状: [h, o]
   按行切分: W被切成n块，每块形状 [h/n, o]
   
   (2) 切分输入张量X
   输入X形状: [b, s, h]
   对应切分: X的最后一维h也被切成n块，每块形状 [b, s, h/n]
   
   (3) 计算每个 GPU 的输出
   GPU0: Y0 = X0 @ W0    
   # X0: [b, s, h/n], W0: [h/n, o] → Y0: [b, s, o]
   
   GPU1: Y1 = X1 @ W1    
   # X1: [b, s, h/n], W1: [h/n, o] → Y1: [b, s, o]
   
   ...
   
   GPUn-1: Yn-1 = Xn-1 @ Wn-1    
   # Xn-1: [b, s, h/n], Wn-1: [h/n, o] → Yn-1: [b, s, o]
   
   (4) 聚合输出
   每个 GPU 计算完成后，将结果通过 all-reduce 聚合到所有设备
   数学上：Y = X @ W = (X0 @ W0) + (X1 @ W1) + ... + (Xn-1 @ Wn-1)
   所以要得到最终Y，需要把所有Yi相加：Y = ∑ Yi
   

2. 列并行（Column Parallelism）

   • 将权重矩阵按列拆分到不同设备
   • 每个设备计算输入矩阵与拆分后的权重矩阵的乘积
   • 输出结果通过 all-gather 聚合到所有设备

   (1) 切分权重矩阵W
   权重W形状: [h, o]
   按列切分: W被切成n块，每块形状 [h, o/n]
   
   (2) 输入张量X保持不变
   输入X形状: [b, s, h]
   保持不变: X完整广播到所有GPU，形状 [b, s, h]
   
   (3) 计算每个 GPU 的输出
   GPU0: Y0 = X @ W0    
   # X: [b, s, h], W0: [h, o/n] → Y0: [b, s, o/n]
   
   GPU1: Y1 = X @ W1    
   # X: [b, s, h], W1: [h, o/n] → Y1: [b, s, o/n]
   
   ...
   
   GPUn-1: Yn-1 = X @ Wn-1    
   # X: [b, s, h], Wn-1: [h, o/n] → Yn-1: [b, s, o/n]
   
   (4) 聚合输出
   每个 GPU 计算完成后，将结果通过 all-gather 聚合到所有设备
   数学上：Y = X @ W = [X @ W0 | X @ W1 | ... | X @ Wn-1]
   所以要得到最终Y，需要把所有Yi拼接起来：Y = concat(Y0, Y1, ..., Yn-1)