hccl 做的是底层集合通信——AllReduce、AllGather、ReduceScatter 等基础通信算法。直接调 hccl 的 API 写分布式训练逻辑,需要在 Python 代码里手动管理通信域、同步点、数据切分。张量并行、流水线并行这类策略下,通信模式远不止一个 AllReduce——层间传激活值用 send/recv,MoE 里 token 路由用 all-to-all,专家参数聚合用 reduce——每个通信模式都要手动调 hccl 的底层 API。

hcomm 在 hccl 之上封了一层高层通信原语(Communication Primitive),把分布式训练中常见的通信模式——TP 通信、PP 通信、MoE 通信——打包成函数。底层还是调 hccl,但调用方式从「管理通信域 + 手动拆分 + 逐卡同步」变成了「调一个函数」。

hcomm 和 hccl 的分工

hccl 是通信库——它知道怎么在 8 张 NPU 之间用 Ring 算法做完一次 AllReduce。hccl 不管张量在模型里是什么角色、分片规则是什么、和其他通信操作怎么配合。

hcomm 是通信原语库——它知道张量并行里每个 attention head 对应哪张卡、流水线并行里 send/recv 的配对关系、MoE 里 token 该怎么按专家分组路由。

模型层(用户代码)
    ↓ 调 hcomm 的高层原语
hcomm(通信原语层)
    ├─ tp_all_reduce()   → 内部调 hcclAllReduce
    ├─ tp_all_gather()   → 内部调 hcclAllGather  
    ├─ pp_send_recv()    → 内部调 hcclSend + hcclRecv
    └─ moe_all_to_all()  → 内部调 hcclGroupAlltoAll
    ↓ 调 hccl 底层 API
hccl(集合通信库)
    ↓ 硬件直通
昇腾NPU 的 RoCE 网卡 / HBM 互连

三种典型通信原语

TP 通信:张量并行的激活值同步

张量并行把 attention 的 QKV 投影和 MLP 的中间层矩阵按列或按行切到不同 NPU 上。切完之后,不同 NPU 算出来的激活值需要同步才能继续往下算。

标准办法:手动调 dist.all_reduce。hcomm 的做法:封装成 tp_all_reduce,自动选择通信组和最优算法:

# 用 hcomm 做张量并行的激活值同步
import hcomm
from hcomm import tp

# 初始化 TP 通信组
tp_group = hcomm.init_tensor_parallel_group(tp_rank=0, tp_size=4)

# 前向:attention 输出需要跨 TP 组做 all-reduce
# hcomm 内部自动选 Ring 或 Mesh 算法(取决于 NPU 拓扑)
attention_output = model.attention(hidden_states)
synced_output = tp.all_reduce(attention_output, group=tp_group)

# MLP 层的列并行:激活值用 all-gather 拼接
# 每张卡只算了 1/4 的中间维度,需要拼成完整的
column_output = model.ffn(attention_output)
full_output = tp.all_gather(column_output, dim=-1, group=tp_group)

tp.all_reduce 不只是 dist.all_reduce 换了个名字。它内部做了拓扑感知的算法选择:

// hcomm 内部:TP all-reduce 的拓扑感知算法选择
CommResult TpAllReduceOp::Execute(Tensor& activation, const TpGroup& group) {
    // 每张 NPU 算一部分 activation[part_i],需要所有 NPU 拿到完整的 activation
    switch (group.npu_topology) {
        case TOPOLOGY_NVLINK_FULL_MESH:
            // 如果 NPU 之间是 NVLink 全互联
            // 用 halving-doubling 算法——对数步数完成,比 Ring 快
            return HcclAllReduce_HalvingDoubling(activation, group);
        
        case TOPOLOGY_RING:
            // 环形拓扑
            // 用 Ring 算法,N-1 步完成
            return HcclAllReduce_Ring(activation, group);
        
        case TOPOLOGY_CUSTOM:
            // 用户自定义拓扑
            // 回退到 naive all-reduce(Broadcast + Reduce)
            return HcclAllReduce_Naive(activation, group);
    }
}

PP 通信:流水线并行的层间传输

流水线并行把模型的不同层放到不同 NPU 上。前向传播时,第 1-8 层在卡 0 上跑,第 9-16 层在卡 1 上跑——卡 0 的最后一层激活值要发到卡 1 的第一层。

hcomm 的 pp_send_recv 把 hccl 的点对点 send/recv 封装成一对匹配的原语:

# 流水线并行:用 hcomm 的 pp_send_recv 传激活值
from hcomm import pp

# 卡0:第1-8层算完,把最后的激活值发给卡1
if pp_rank == 0:
    output = pipeline_stage_0(input_data)
    pp.send(output, dst=1, tag=PP_TAG_FORWARD)
    
    # 反向传播时收卡1的梯度
    grad = pp.recv(src=1, tag=PP_TAG_BACKWARD)
    pipeline_stage_0.backward(grad)

# 卡1:收卡0的激活值,接着算第9-16层
if pp_rank == 1:
    input_data = pp.recv(src=0, tag=PP_TAG_FORWARD)
    output = pipeline_stage_1(input_data)
    
    # 反向传播时把梯度发回卡0
    pp.send(output.grad, dst=0, tag=PP_TAG_BACKWARD)

PP 通信的关键瓶颈是气泡时间——卡 0 在等卡 1 算完才能收到梯度,这段时间卡 0 闲置。hcomm 不做气泡消除(那是调度策略的事),但它优化了 send/recv 的延迟——底层的 hcclSend / hcclRecv 走的是 NPU 之间的 SDMA 直连路径,数据不经过 CPU 内存,延迟在微秒级。

MoE 通信:专家并行的 token 路由

MoE(Mixture of Experts)模型里,每个 token 只激活少数几个专家,需要通过 all-to-all 通信把 token 路由到对应的专家 NPU 上:

# MoE 的 all-to-all token 路由
from hcomm import moe

# 输入:每张 NPU 上有一部分 token,每个 token 有一个 expert_id
# expert_id 决定了这个 token 要发给哪张 NPU
tokens_per_npu = [batch // num_experts] * num_experts

# hcomm 的 moe.all_to_all 封装了 hccl 的 Grouped All-to-All
# 自动处理 token 的拆分、路由、拼接
routed_tokens = moe.all_to_all(
    local_tokens,           # 本卡上的 token [num_local_tokens, hidden_dim]
    expert_ids,             # 每个 token 要发到哪个 expert (0~num_experts-1)
    num_experts=8,
    tokens_per_expert=tokens_per_npu
)

# 路由完成后,每张 NPU 拿到属于自己专家组的全部 token
# 跑专家计算
expert_output = moe_layer(routed_tokens)

# 反向 all-to-all:专家结果路由回原始位置
final_output = moe.all_to_all(
    expert_output,
    reverse_expert_ids,    # 反向路由表
    num_experts=8,
    tokens_per_expert=tokens_per_npu
)

通信域管理

hcomm 最大的便利不在单个通信操作,而在通信域(Communicator)的自动管理。

分布式训练里一张 NPU 可能同时属于多个通信组:TP 组、PP 组、DP 组。每个组有独立的 hccl communicator、独立的同步点、独立的通信拓扑。手动管理这些组的生命周期——创建、同步、销毁——容易出错。

hcomm 用 init_context 一次搞定所有通信域的初始化:

# hcomm 通信域初始化——一次性创建所有并行组
import hcomm

ctx = hcomm.init_context(
    tp_size=4,     # 张量并行:4 张 NPU 切模型参数
    pp_size=2,     # 流水线并行:2 组串行
    dp_size=4,     # 数据并行:4 个副本(总共 32 张 NPU = 4×2×4)
    backend='hccl'
)

# ctx 自动创建了三个通信组的 communicator
# tp_group: 4 张 NPU 的 Ring
# pp_group: 2 组的 send/recv 配对
# dp_group: 4 张 NPU 的 AllReduce

# 训练循环中直接用
for batch in dataloader:
    # TP:每层算完自动调 tp.all_reduce
    loss = model(batch)
    
    # DP:梯度同步走 ctx 的 dp_group
    ctx.dp.all_reduce(grads)
    
    # 更新参数(TP 组内已经自动同步过激活值)
    optimizer.step()

分布式训练的通信复杂度,根本不在单个 AllReduce 或 AllGather 怎么调——hccl 已经把这些做透了。真正的复杂度在编排——TP 的激活值同步、PP 的层间传输、MoE 的 token 路由,三套通信逻辑要在训练循环里正确交错执行。hcomm 的价值就是把这三套逻辑封装成命名清晰的原语,让训练代码的通信层从几百行 hccl 调用变成几行函数调用。

Logo

作为“人工智能6S店”的官方数字引擎,为AI开发者与企业提供一个覆盖软硬件全栈、一站式门户。

更多推荐