细粒度张量并行

细粒度张量并行#

概述#

细粒度张量并行 (Fine-Grained TP) 通过为不同的模型组件启用独立的张量并行大小，扩展了标准的张量并行。Fine-Grained TP 不再对所有层应用统一的全局 tensor_parallel_size，而是允许用户通过 finegrained_tp_config 参数，为关键模块（如 Embedding、语言模型头 (LM Head)、Attention 输出投影 (o_proj) 和 MLP 块）单独配置 TP 大小。

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

细粒度 TP 的优势#

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

降低单设备显存占用：细粒度 TP 将巨大的权重矩阵（如 LM Head、o_proj）跨设备分片，从而降低峰值显存使用量。这使得在不使用量化的情况下，也能支持更大的 Batch Size 或在显存受限的硬件上进行部署。
加速 GEMM 内存访问：在以解码 (Decode) 为主的工作负载中，GEMM 性能通常受限于内存带宽。权重分片减少了单个设备需要获取的权重数据量，降低了 DRAM 流量并提高了带宽效率——对于像 LM Head 和 o_proj 这样对延迟敏感的层尤其有效。

综上所述，这些优化使得开发者能够更好地平衡内存、通信和计算资源（特别是在高并发推理场景下），同时保持对标准 Dense Transformer 模型的兼容性。

支持的场景#

模型支持#

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

组件与执行模式支持情况#

TP 配置项目	Eager 模式	图模式 (Graph)	混合 (Hybrid)	预填充 (Prefill)	解码 (Decode)
embedding	✅	✅	✅	✅	✅
o_proj	❌	✅	❌	❌	✅
mlp	✅	✅	✅	✅	✅
LMhead	✅	✅	✅	✅	✅

⚠️ 注意：

o_proj 的 TP 仅在 Decode 阶段的图模式 (Graph mode) 下支持，因为 Eager 模式下的 dummy_run 不会触发 o_proj。

mlp TP 支持 Dense 模型，或者 MoE 模型中的 Dense 层（例如 DeepSeek-R1 的前三层 Dense 层）。

配置限制：#

任何组件的细粒度 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 910B*64G (A2) 环境下部署 PD 分离的 Decode 实例。并行配置为 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 分离的解码 (Decode) 实例中最为有效。在该场景下，模型通常以全 DP 模式部署，通过对权重密集的层进行分片，可以有效减少冗余存储并缓解显存压力。