Appearance
将模型权重矩阵按维度切分到多个 GPU 上并行计算,也称 intra-layer parallelism。
为什么需要 Tensor Parallelism
大模型(如 LLaMA-70B、DeepSeek-V3)的单层权重远超单个 GPU 显存。Tensor Parallelism(TP)将每层的权重矩阵按行或列切分到多个 GPU,每个 GPU 仅持有部分权重和中间激活,使单 GPU 无法容纳的模型得以运行。
核心原理
- Column Parallel:权重按列切分,每个 GPU 计算 partial output,通过 AllReduce 聚合。用于 MLP 的 up/gate 投影和 Attention 的 QKV 投影。
- Row Parallel:权重按行切分,输入先切分,每个 GPU 计算部分结果后 AllReduce 求和。用于 MLP 的 down 投影和 Attention 的 output 投影。
- 通信模式:每层需要 2 次 AllReduce(Attention 后 + MLP 后),通信量与 hidden_dim 成正比。
- NVLink 依赖:TP 对通信带宽敏感,通常要求 GPU 间有 NVLink/NVSwitch 连接(节点内使用)。
在源码中的实现
vllm/distributed/tensor_parallel/— TP 通信原语(AllReduce、Send/Recv)。vllm/model_executor/layers/linear.py—ColumnParallelLinear和RowParallelLinear封装切分逻辑。vllm/model_executor/layers/parallel_embedding.py— Embedding 层的并行化。vllm/config.py—ParallelConfig.tensor_parallel_size控制 TP 度。
相关概念
- pipeline-parallelism — 另一种维度的模型并行,跨层切分
- moe — 专家并行通常与 TP 结合使用
- cuda-graph — TP 场景下 CUDA Graph 需处理分布式通信节点
- flash-attention — 注意力 kernel 需感知 TP 的 QKV 切分