Fine-Grained Tensor Parallelism (Finegrained TP)

细粒度张量并行（Finegrained TP）#

概述#

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

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

细粒度 TP 的优势#

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

  • 降低每设备内存占用：
    Finegrained TP 将大型权重矩阵（如 LM Head、o_proj）分片到设备上，降低峰值内存使用量，并允许在内存受限的硬件上部署更大的批次或模型——无需量化。

  • 加快 GEMM 中的内存访问：
    在以解码为主的工作负载中，GEMM 性能通常受内存限制。权重分片减少了每设备的权重获取量，减少了 DRAM 流量并提高了带宽效率——尤其对于像 LM Head 和 o_proj 这样的延迟敏感层。

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

支持的场景#

模型#

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

组件和执行模式支持#

TP 配置

Eager

Graph

Hybrid

Prefill

Decode

embedding

✅

✅

✅

✅

✅

o_proj

❌

✅

❌

❌

✅

mlp

✅

✅

✅

✅

✅

LMhead

✅

✅

✅

✅

✅

⚠️ 注意

  • o_proj TP 仅在 Graph 模式下的 Decode 阶段支持，因为 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 910B*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 模式部署。在此配置下，分片权重密集层可以减少冗余存储和内存压力。

目录