细粒度张量并行¶
概述¶
细粒度张量并行通过为**不同的模型组件启用独立的张量并行大小**,扩展了标准的张量并行。细粒度 TP 不再对所有层应用统一的全局tensor_parallel_size,而是允许用户通过finegrained_tp_config参数,为关键模块(如 Embedding、语言模型头、注意力输出投影和 MLP 块)单独配置 TP 大小。
此功能支持在单个模型内采用异构并行策略,从而对跨设备的权重分布、内存布局和通信模式进行更精细的控制。该特性与标准的密集 Transformer 架构兼容,并无缝集成到 vLLM 的推理流水线中。
细粒度 TP 的优势¶
细粒度张量并行通过有针对性的权重分片,带来两个主要的性能优势:
-
Reduced Per-Device Memory Footprint:
Fine-grained TP shards large weight matrices (e.g., LM Head, o_proj) across devices, lowering peak memory usage and enabling larger batches or deployment on memory-limited hardware—without quantization. -
Faster Memory Access in GEMMs:
In decode-heavy workloads, GEMM performance is often memory-bound. Weight sharding reduces per-device weight fetch volume, cutting DRAM traffic and improving bandwidth efficiency—especially for latency-sensitive layers like LM Head and o_proj.
综上所述,这些优化使得开发者能够更好地平衡内存、通信和计算资源,同时保持对标准密集 Transformer 模型的兼容性。
支持的场景¶
模型支持¶
细粒度 TP 是**模型无关**的,支持所有标准的密集 Transformer 架构,包括 Llama、Qwen、DeepSeek 等。
组件与执行模式支持情况¶
| TP 配置项目 | Eager 模式 | 图模式 | 混合模式 | 预填充 | 解码 |
|---|---|---|---|---|---|
| embedding | ✅ | ✅ | ✅ | ✅ | ✅ |
| o_proj | ❌ | ✅ | ❌ | ❌ | ✅ |
| mlp | ✅ | ✅ | ✅ | ✅ | ✅ |
| LMhead | ✅ | ✅ | ✅ | ✅ | ✅ |
⚠️ 注意:
o_proj的 TP 仅在解码阶段的图模式下支持,因为 Eager 模式 模式下的 dummy_run 不会触发 o_proj。mlpTP 支持密集模型,或者 MoE 模型中的密集层。例如,DeepSeek-R1 的前三个密集层。
配置限制¶
任何组件的细粒度 TP 大小必须满足:
- 必须 ≤ 数据并行的大小,且
- 必须**能整除 DP 大小**,以确保有效的设备分配和通信组划分。
⚠️ 违反这些约束将导致运行时错误或未定义的行为。
如何使用细粒度 TP¶
配置格式¶
细粒度 TP 通过 dp_size % tp_size == 0 中的 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 环境下部署 PD 分离的解码实例。并行配置为 DP32+EP32,细粒度 TP 大小设为 8。性能数据如下。
| 模块 | 内存节省 | TPOT 影响 |
|---|---|---|
| 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 模式部署,通过对权重密集的层进行分片,可以有效减少冗余存储并缓解内存压力。