细粒度张量并行 (Fine-grained TP)

细粒度张量并行 (Fine-grained TP)#

概述#

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

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

细粒度 TP 的优势#

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

降低单设备内存占用：细粒度 TP 将大型权重矩阵（例如 LM Head、o_proj）分片到多个设备上，降低了峰值内存使用量，从而支持更大的批次或在内存受限的硬件上进行部署——无需量化。
加速 GEMM 中的内存访问：在解码密集型工作负载中，GEMM 性能通常受内存带宽限制。权重分片减少了每个设备需要获取的权重数据量，从而降低了 DRAM 流量并提高了带宽效率——对于 LM Head 和 o_proj 等延迟敏感层尤其如此。

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

支持场景#

模型#

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

组件与执行模式支持#

TP 配置	Eager	Graph	Hybrid	Prefill	Decode
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 通过 --additional-config 内的 finegrained_tp_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 分离的解码实例，并行配置为 DP32+EP32，细粒度 TP 规模为 8；性能数据如下。

模块	内存节省	TPOT 影响 (batch=24)
o_proj TP = 8	5.8 GB	+1.5 ms (性能下降)
LM head TP = 8	1.51 GB	−1.2 ms (性能提升)
FFN TP = 8	0.9 GB	−1.0 ms (性能提升)
Embedding TP = 8	1.51 GB	−1.0 ms (性能提升)
总计	9.72 GB	—

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

✅ 部署建议#

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