细粒度张量并行

概述

细粒度张量并行通过允许为不同模型组件设置**独立的张量并行大小**，扩展了标准张量并行。细粒度 TP 并非对所有层应用单一的全局 tensor_parallel_size，而是允许用户通过 finegrained_tp_config 参数为关键模块（如嵌入层、语言模型头 (lm_head)、注意力输出投影 (o_proj) 和 MLP 块）配置独立的 TP 大小。

此功能支持在单个模型内采用异构并行策略，从而对跨设备的权重分布、内存布局和通信模式提供更精细的控制。该特性与标准稠密 Transformer 架构兼容，并能无缝集成到 vLLM 的服务流程中。

---

细粒度张量并行的优势

细粒度张量并行通过针对性的权重分片带来两大性能优势：

• 降低单设备内存占用：
  细粒度TP将大权重矩阵（如LM Head、o_proj）跨设备分片，降低峰值内存使用，支持更大批次或在内存受限硬件上部署——无需量化。

• 加速GEMM中的内存访问：
  在解码密集型负载中，GEMM性能常受内存限制。权重分片减少单设备权重获取量，降低DRAM流量并提升带宽效率——尤其适用于LM Head和o_proj等延迟敏感层。

综合来看，这些效果使实践者能够更好地平衡内存、通信与计算——尤其是在高并发服务场景中——同时保持与标准稠密Transformer模型的兼容性。

---

支持场景

模型

细粒度 TP 是**模型无关的**，支持所有标准稠密 Transformer 架构，包括 Llama、Qwen、DeepSeek（基础/稠密变体）等。

组件与执行模式支持

TP 配置	即时模式	图模式	混合模式	预填充	解码
embedding	✅	✅	✅	✅	✅
o_proj	❌	✅	❌	❌	✅
mlp	✅	✅	✅	✅	✅
LMhead	✅	✅	✅	✅	✅

⚠️ 注意：

• o_proj TP 仅在 Decode 阶段的 Graph 模式下支持，因为 eager 模式下的 dummy_run 不会触发 o_proj。
• mlp TP 支持稠密模型，或 MoE 模型中的稠密层。例如，DeepSeek-R1 的前三个稠密层。

配置限制

任何组件的细粒度 TP 大小必须满足：

• ≤ 数据并行 (DP) 大小，并且
• 能整除 DP 大小（即 dp_size % tp_size == 0），以确保有效的设备分配和通信分组。

⚠️ 违反这些约束将导致运行时错误或未定义行为。

---

如何使用细粒度 TP

配置格式

细粒度 TP 通过 finegrained_tp_config 中的 --additional-config 字段进行控制。

--additional-config '{
    "finegrained_tp_config": {
        "embedding_tensor_parallel_size": 8,
        "lmhead_tensor_parallel_size": 8,
        "oproj_tensor_parallel_size": 8,
        "mlp_tensor_parallel_size": 8
    }
}'

使用示例

vllm serve deepseek-ai/DeepSeek-R1 \
    --data-parallel-size 16 \
    --tensor-parallel-size 1 \
    --enable-expert-parallel \
    --additional-config '{
        "finegrained_tp_config": {
            "embedding_tensor_parallel_size": 8,
            "lmhead_tensor_parallel_size": 8,
            "mlp_tensor_parallel_size": 8
        }
    }'

---

实验结果

为了评估细粒度 TP 在大规模服务场景中的有效性，我们使用 DeepSeek-R1-W8A8 模型，在 32 卡 Ascend Atlas A2 推理产品*64G (A2) 环境中部署 PD 分离的 decode 实例，并行配置为 DP32+EP32，细粒度 TP 大小为 8；性能数据如下。

模块	内存节省	TPOT 影响（批次=24）
o_proj TP = 8	5.8 GB	+1.5 毫秒（性能下降）
LM head TP = 8	1.51 GB	−1.2 毫秒（性能提升）
FFN TP = 8	0.9 GB	−1.0 毫秒（性能提升）
Embedding TP = 8	1.51 GB	−1.0 毫秒（性能提升）
总计	9.72 GB	—

• 我们在单卡高内存容量以及 TPOT 收益方面取得了显著成果。

---

✅ 部署建议

细粒度 TP 在 PD 分离的 decode 实例**中**最为有效，因为模型通常以全 DP 模式部署。在此设置下，对权重较大的层进行分片可以减少冗余存储和内存压力。