一文详解张量并行Tensor parallel的概念和原理应用

张量并行概念

张量并行（Tensor Parallelism）是一种模型并行技术，其核心思想是将模型的张量操作（如矩阵乘法、注意力计算等）拆分成多个子任务，分配到不同设备（如GPU）上并行执行。以下从概念、区别与联系三个方面展开分析：

---

一、张量并行的概念

1. 核心思想：

   • 将模型中的大张量（如权重矩阵）沿特定维度（行或列）切分，分配到多个设备上。
   • 每个设备仅持有部分参数，执行局部计算，再通过设备间通信（如All-Reduce）合并结果。
   • 适用于参数量超出单设备显存的大模型（如Transformer）。

2. 典型场景：

   • 训练阶段：拆分模型参数，同步梯度（如Megatron-LM的层内并行）。
   • 推理阶段：拆分计算图，合并前向结果（如vLLM的注意力头并行）。

3. 优势：

   • 降低单设备显存占用，支持更大模型。
   • 提升计算吞吐量（若通信开销可控）。

---

二、PyTorch分布式训练中的张量并行

1. 目标与场景：

   • 训练优化：解决显存不足问题，加速梯度计算。
   • 常与数据并行（Data Parallelism）结合，形成混合并行策略。

2. 实现方式：

   • 参数切分：将权重矩阵按行或列拆分到多设备（如线性层的分块计算）。
   • 通信模式：
     • 前向传播：需要All-Gather操作合并中间结果。
     • 反向传播：需同步梯度（All-Reduce）。
   • 工具支持：
     • torch.distributed + NCCL通信库。
     • 框架级支持（如DeepSpeed、FairScale）。

3. 挑战：

   • 通信开销与计算效率的权衡。
   • 动态计算图的灵活性受限。

---

三、vLLM推理框架中的张量并行

1. 目标与场景：

   • 推理优化：降低延迟，提升吞吐量，支持长序列生成。
   • 针对Transformer结构的生成任务（如LLM的Token生成）。

2. 实现方式：

   • 注意力并行：
     • 拆分注意力头（Heads）到不同设备，独立计算后合并结果。
     • 优化KV缓存管理（如PagedAttention减少显存碎片）。
   • 轻量通信：
     • 仅需前向传播中的All-Gather，无需反向通信。
     • 采用异步通信与计算重叠（如CUDA Stream优化）。