当前位置：首页 > news >正文

多卡并行策略：Tensor Parallelism（TP）在昇腾上的通信机制

news 2026/3/27 7:24:58

当DeepSeek模型的规模膨胀到67B甚至更大时，单张加速卡——哪怕是拥有64GB HBM的昇腾910B——也显得捉襟见肘。67B参数的模型仅权重本身就需要约130GB（FP16），这还不算推理时KV Cache和中间激活值（Activation）的巨大开销。

面对物理显存的“硬墙”，我们唯一的出路是并行化（Parallelism）。在众多并行策略中，Tensor Parallelism (TP，张量并行)凭借其极低的延迟特性，成为了大模型在线推理的绝对主流。本篇将深入解构TP在昇腾平台上的实现原理、通信机制以及针对DeepSeek模型的特定优化策略。

1. 并行策略的江湖：为什么是TP？

在分布式训练和推理的江湖中，主要有三大门派：

Pipeline Parallelism (PP，流水线并行)：按层切分。Card 0算前10层，Card 1算后10层。
- 缺点：存在严重的“气泡”（Pipeline Bubble）。当Card 0在计算时，Card 1只能干等，硬件利用率低，且首字延迟（Latency）随着卡数线性增加。
Data Parallelism (DP，数据并行)：每张卡复制一份完整的模型，处理不同的请求。
- 缺点：前提是单卡能装下整个模型。对于67B模型，这显然不可能。
Tensor Parallelism (TP，张量并行)：按矩阵切分。每一层的矩阵乘法被拆分到多张卡上同时计算。
- 优点：所有卡时刻都在忙碌，延迟极低（接近单卡速度），显存占用被均匀分摊。

对于追求极致响应速度的在线推理服务，TP是当之无愧的王者。

2. TP的数学魔法：行列切分

TP的核心思想是将一个大的矩阵乘法Y=X×WY = X \times WY=X×W拆解为多个小的矩阵乘法。以DeepSeek模型中最常见的MLP（多层感知机）和Self-Attention层为例，TP采用了经典的“列切分-行切分”交替策略，以最小化通信次数。

2.1 列切分（Column Parallel）

假设输入XXX的维度是[B,H][B, H][B,H]，权重WWW的维度是[H,4H][H, 4H][H,4H]。我们将WWW按列切分为两半W1,W2W_1, W_2W1,W2，分别放在Card 0和Card 1上。

Card 0 计算Y1=X×W1Y_1 = X \times W_1Y1=X×W1
Card 1 计算Y2=X×W2Y_2 = X \times W_2Y2=X×W2
此时，每张卡得到的输出Y1,Y2Y_1, Y_2Y1,Y2是完整输出的一部分通道。这一步完全独立，不需要任何通信。

2.2 行切分（Row Parallel）

紧接着下一层，假设权重VVV的维度是[4H,H][4H, H][4H,H]。我们将VVV按行切分为V1,V2V_1, V_2V1,V2。

Card 0 计算Z1=Y1×V1Z_1 = Y_1 \times V_1Z1=Y1×V1
Card 1 计算Z2=Y2×V2Z_2 = Y_2 \times V_2Z2=Y2×V2
注意，根据矩阵乘法结合律，Z=Y×V=Y1×V1+Y2×V2Z = Y \times V = Y_1 \times V_1 + Y_2 \times V_2Z=Y×V=Y1×V1+Y2×V2。
因此，最终结果ZZZ需要将两张卡的计算结果Z1Z_1Z1和Z2Z_2Z2相加。

2.3 AllReduce：通信的代价

这就引入了TP中唯一的通信操作：AllReduce。
Card 0和Card 1必须交换各自持有的Z1Z_1Z1和Z2Z_2Z2，并计算和。由于这一步发生在每一层Transformer Block的末尾，DeepSeek-67B有几十层，意味着一次前向传播需要几十次AllReduce。通信带宽成为了系统的生死线。

3. 昇腾上的实现：HCCL与HCCS

在昇腾系统中，多卡通信由HCCL (Huawei Collective Communication Library)负责，它是对标NVIDIA NCCL的通信库。

3.1 物理通道：HCCS互联

在昇腾服务器内部，8张NPU卡并非通过PCIe互联，而是通过华为自研的HCCS (Huawei Cache Coherent System)协议全互联。

带宽：HCCS的总双向带宽高达392GB/s（不同型号略有差异），远超PCIe 4.0的64GB/s。
拓扑：8张卡通常组成一个Hypercube Mesh或全互联拓扑，任意两张卡之间都可以直通，无需CPU中转。

这意味着，只要你的TP部署在同一台服务器内（Intra-server），通信开销是非常小的。切记：千万不要跨服务器做TP！跨服务器走的是以太网（RoCE），延迟比HCCS高两个数量级，会导致推理速度断崖式下跌。

3.2 算法实现：Ring vs Tree

HCCL内部实现了多种AllReduce算法：

Ring AllReduce：卡之间围成一个环，数据像接力棒一样传递。适合带宽受限但延迟不敏感的场景。
Tree/Mesh AllReduce：适合昇腾这种全互联拓扑，延迟更低。

在CANN 7.0中，HCCL会自动根据当前的物理拓扑选择最优算法，通常无需人工干预。

4. 实战代码：DeepSpeed与MindSpore

4.1 使用DeepSpeed (PyTorch)

在昇腾上，deepspeed_npu插件已经完美适配了HCCL。以下是启动TP推理的标准姿势：

importtorchimportdeepspeedimporttorch_npu# 1. 初始化分布式环境# DeepSpeed会自动读取 rank, world_size 等环境变量deepspeed.init_distributed(dist_backend="hccl")# 2. 加载模型model=AutoModelForCausalLM.from_pretrained(...)# 3. 注入TP逻辑# DeepSpeed会自动识别Transformer层，用支持TP的HCCL算子替换原生Linear算子model=deepspeed.init_inference(model,mp_size=2,# 设置TP并行度为2dtype=torch.bfloat16,replace_with_kernel_inject=True# 使用高性能Kernel)# 4. 推理# 代码与单卡完全一致，通信细节被封装在底层output=model.generate(...)