`hccl_checkpoint_engine.py` — 这是基于 **HCCL（Huawei Collective Communication Librar...¶

文件路径: verl/checkpoint_engine/hccl_checkpoint_engine.py

文件概述¶

这是基于 HCCL（Huawei Collective Communication Library） 的权重同步后端实现，专门适配华为 Ascend NPU 硬件。HCCL 之于华为 NPU，就像 NCCL 之于 NVIDIA GPU。

本文件的整体结构和逻辑与 nccl_checkpoint_engine.py 高度相似，主要区别在于： 1. 使用 torch.npu 代替 torch.cuda 2. 使用 StatelessProcessGroup + stateless_init_process_group 代替 ray.util.collective 3. BroadcastOperation 是同步执行（不使用 run_in_executor） 4. 需要额外的 dist_port 用于初始化分布式进程组

关键代码讲解¶

1. MasterMetadata -- 主节点元数据¶

@dataclass
class MasterMetadata:
    zmq_ip: str
    zmq_port: int
    dist_ip: str      # 新增：分布式通信地址
    dist_port: int     # 新增：分布式通信端口

与 NCCL 版本相比，多了 dist_ip 和 dist_port，用于初始化 HCCL 的分布式进程组。

2. BroadcastOperation -- 同步广播¶

class BroadcastOperation:
    def __init__(self, rank, process_group, bucket, metadata, socket, topic):
        ...
        self._run()  # 注意：这里直接同步执行！

    def _run(self):
        # 1. ZMQ 传元数据
        if self.rank == 0:
            self.socket.send_string(self.topic, flags=zmq.SNDMORE)
            self.socket.send_pyobj(self.metadata)
        else:
            self.socket.recv_string()
            self.metadata = self.socket.recv_pyobj()

        # 2. HCCL broadcast 传张量
        self.pyhccl.broadcast(self.bucket, src=0)

    async def wait_for_complete(self):
        return self.metadata  # 已经同步完成，直接返回

与 NCCL 版本的关键差异： - NCCL 版本使用 asyncio.get_running_loop().run_in_executor() 异步执行 - HCCL 版本在 __init__ 中直接调用 _run()，同步阻塞执行 - wait_for_complete() 不需要 await 任何 task，直接返回

3. HCCLCheckpointEngine -- 主体类¶

初始化差异¶

@CheckpointEngineRegistry.register("hccl")
class HCCLCheckpointEngine(CheckpointEngine):
    def __init__(self, bucket_size, group_name="default", rebuild_group=False,
                 is_master=False, rollout_dtype=torch.bfloat16):
        ...
        self.pyhccl = None
        self.device = torch.npu.current_device()  # 使用 NPU 设备
        if self.is_master:
            self._start_zmq_server()
            self.dist_port, _ = get_free_port(self.ip)  # 额外分配分布式端口

prepare -- 使用 NPU 内存¶

def prepare(self):
    self.send_buf = torch.zeros(self.bucket_size, dtype=torch.uint8, device="npu")
    self.recv_buf = torch.zeros(self.bucket_size, dtype=torch.uint8, device="npu")
    return MasterMetadata(zmq_ip=self.ip, zmq_port=self.zmq_port,
                          dist_ip=self.ip, dist_port=self.dist_port) if self.is_master else None

注意使用 device="npu" 而非 device="cuda"。

init_process_group -- 使用 stateless 方式¶

def init_process_group(self, rank, world_size, master_metadata):
    if rank < 0:
        self.rank = rank
        return

    if self.rebuild_group or self.pyhccl is None:
        # 使用 stateless_init_process_group 初始化
        self.pyhccl = stateless_init_process_group(
            master_metadata.dist_ip, master_metadata.dist_port,
            rank, world_size, self.device
        )
        self.rank = rank
        self.world_size = world_size

    if self.rank > 0:
        self._connect_zmq_client(master_metadata)

    # 使用 all_reduce 代替 barrier（NPU 适配）
    signal = torch.tensor([1], dtype=torch.int8, device=torch.npu.current_device())
    self.pyhccl.all_reduce(signal)

与 NCCL 版本的差异： - 使用 stateless_init_process_group 而非 collective.init_collective_group - 使用 all_reduce 代替 barrier 进行同步

finalize -- 清理 NPU 资源¶

def finalize(self):
    if self.rebuild_group:
        if self.rank >= 0:
            self.pyhccl.destroyComm(self.pyhccl.comm)  # 销毁 HCCL 通信句柄
            self.pyhccl = None
        self.rank = None
        self.world_size = None

    self.send_buf = None
    self.recv_buf = None
    torch.npu.empty_cache()  # 清理 NPU 缓存

send_weights / receive_weights¶

这两个方法的逻辑与 NCCL 版本完全一致，仅在细节上有所不同： - 使用 torch.npu.synchronize() 代替 torch.cuda.synchronize() - 直接使用 weight.view(-1).view(torch.uint8) 而非 cp.asarray(...) - BroadcastOperation 传入 process_group=self.pyhccl 而非 group_name

核心类/函数列表¶

名称	类型	作用
`MasterMetadata`	dataclass	master 的 ZMQ 和分布式通信地址
`BroadcastOperation`	类	封装一次同步 HCCL broadcast
`HCCLCheckpointEngine`	类	HCCL 后端的完整实现

与 NCCL 版本的对比¶

方面	NCCL	HCCL
硬件	NVIDIA GPU	华为 Ascend NPU
设备 API	`torch.cuda`	`torch.npu`
内存分配	cupy (master) + torch	torch (npu)
集合通信库	`ray.util.collective`	`stateless_init_process_group`
广播执行	异步 (`run_in_executor`)	同步 (直接调用)
同步原语	`collective.barrier`	`pyhccl.all_reduce`
清理	`torch.cuda.empty_cache()`	`torch.npu.empty_cache()`

与其他模块的关系¶

base.py：继承 CheckpointEngine，通过 @CheckpointEngineRegistry.register("hccl") 注册。
vllm.distributed.utils：使用 StatelessProcessGroup。
verl.utils.distributed：使用 stateless_init_process_group 初始化进程组。

小结¶

HCCL 后端是 NCCL 后端在华为 Ascend NPU 上的适配版本。代码结构和通信逻辑几乎完全相同，主要差异在于底层设备 API 和通信库的替换。如果理解了 NCCL 版本，这个文件可以快速理解。一个值得注意的设计差异是 BroadcastOperation 改为同步执行，可能与 HCCL 的异步支持情况有关。

hccl_checkpoint_engine.py — 这是基于 **HCCL（Huawei Collective Communication Librar...¶