细粒度张量并行

概述

细粒度张量并行通过为不同模型组件启用独立的张量并行大小，扩展了标准张量并行功能。细粒度TP不将单一的全局tensor_parallel_size应用于所有层，而是允许用户通过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模式	混合模式	预填充	解码
embedding	✅	✅	✅	✅	✅
o_proj	❌	✅	❌	❌	✅
mlp	✅	✅	✅	✅	✅
LMhead	✅	✅	✅	✅	✅

⚠️ 注意：

• o_proj TP仅在解码阶段的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模式部署。在此设置中，对权重密集型层进行分片可减少冗余存储和内存压力。