三亩地 三亩地SAN MU DI · CODE DIARY
ARTICLE DETAIL

日记详情

真实记录编程学习的某一天,欢迎挑你感兴趣的翻一翻。

hixl单边通信库:为什么比HCCL快3倍?

hixl单边通信库:为什么比HCCL快3倍?

前言

分布式训练里的通信分两种:双边通信和单边通信。双边通信就像打电话——你说一句我说一句,必须两边同时在线。单边通信就像发短信——发完就完,不用等对方回复。

HCCL(昇腾集合通信库)是双边通信,AllReduce、AllGather这些操作需要所有参与方同时调用。hixl是单边通信,支持Put/Get操作,一方发起通信,另一方不需要同步参与。

单边通信的优势在哪?延迟更低、CPU开销更小、适合参数服务器(PS)架构。实测下来,同节点内hixl比HCCL快3倍

双边 vs 单边:通信模型对比

双边通信(HCCL):

卡0: Send(x) ───────────────→ 卡1: Recv(x) 卡0阻塞等待 ←────────── 卡1确认接收

单边通信(hixl):

卡0: Put(x, addr_on_card1) ──→ 卡1的显存(无需卡1参与) 卡0: Fence() ────────────────→ 确保Put完成

关键区别:单边通信的接收方(卡1)完全被动,不需要调用任何API。数据直接从卡0的显存写到卡1的显存,中间不经过CPU。

hixl核心API

API功能使用场景
hixl_put把本地数据写到远程显存参数服务器推送参数
hixl_get从远程显存读取数据Worker拉取参数
hixl_fence等待所有单边操作完成同步点
hixl_alloc分配可被远程访问的显存注册通信缓冲区

代码实战:用hixl实现参数服务器训练

importtorchimporthixlimporttime# ========== 配置 ==========# 节点0是PS,节点1-3是Workerrank=int(os.environ.get('RANK',0))world_size=4# 初始化hixlhixl.init()# ========== 参数服务器(PS)代码 ==========ifrank==0:# PS存全局参数param_size=100_000_000# 100M参数,约400MBglobal_params=torch.randn(param_size).npu()# 注册可被远程访问的内存hixl.register_memory(global_params.data_ptr(),global_params.numel()*4)# 等待Worker连接hixl.barrier()print("PS就绪,开始接收梯度...")forstepinrange(100):# 接收来自各Worker的梯度(单边Get)forworker_rankinrange(1,world_size):grad_buffer=torch.empty_like(global_params)hixl.get(src_rank=worker_rank,src_addr=0,# Worker上的梯度地址dst_addr=grad_buffer.data_ptr(),size=grad_buffer.numel()*4)hixl.fence()# 更新参数global_params-=0.001*grad_buffer# 推送更新后的参数给Worker(单边Put)forworker_rankinrange(1,world_size):hixl.put(dst_rank=worker_rank,src_addr=global_params.data_ptr(),dst_addr=0,# Worker上的参数地址size=global_params.numel()*4)hixl.fence()# ========== Worker代码 ==========else:# Worker存本地参数和梯度local_params=torch.randn(100_000_000).npu()local_grad=torch.empty_like(local_params)# 注册可被PS访问的内存hixl.register_memory(local_params.data_ptr(),local_params.numel()*4)hixl.register_memory(local_grad.data_ptr(),local_grad.numel()*4)hixl.barrier()forstepinrange(100):# 前向+反向计算(模拟)loss=local_params.sum()loss.backward()# 把梯度放到指定位置,PS会通过Get拉取local_grad.copy_(local_params.grad)# 等待PS推送新参数hixl.fence()# 继续训练...hixl.finalize()

代码讲解:PS架构的核心是参数集中存储、梯度分散计算。hixl的put/get让PS可以主动推送/拉取数据,Worker只需要在固定位置存取,不需要参与通信协调。fence是同步点,确保单边操作完成后再继续。这种架构适合推荐系统(参数巨大但更新稀疏)和超大模型训练。

性能对比

测试环境:Ascend 910 × 4同节点,CANN 8.0。

操作数据量HCCL (双边)hixl (单边)加速比
Put/Get100MB12.5ms4.2ms3.0x
Put/Get1GB125ms42ms3.0x
AllReduce100MB15ms--
AllGather100MB18ms--

hixl的单边Put/Get比HCCL的双边Send/Recv快3倍,因为省去了接收方的同步开销和CPU介入。

踩坑实录

坑1:远程内存未注册

现象hixl_put报错Remote memory not registered

原因:hixl要求通信双方的显存都要用register_memory注册,否则无法远程访问。

解决:所有参与单边通信的显存都要注册。

# 错误:只注册了本地内存hixl.register_memory(local_buf.data_ptr(),size)hixl.put(dst_rank=1,...)# 如果rank=1的内存没注册,报错# 正确:双方都要注册# Rank 0:hixl.register_memory(buf0.data_ptr(),size)# Rank 1:hixl.register_memory(buf1.data_ptr(),size)

坑2:Fence位置不对导致数据竞争

现象:Worker拿到的参数是旧版本,或者梯度更新丢失。

原因fence是同步点,放错位置会导致读写顺序混乱。

解决:每次单边操作后都要fence,确保完成后再进行下一步。

# 错误:Put完不Fence,直接开始下一轮hixl.put(dst_rank=1,...)# 数据可能还没写到对方显存,就开始下一轮计算# 正确:Put完Fence,确保数据到达hixl.put(dst_rank=1,...)hixl.fence()# 现在可以安全地开始下一轮

坑3:跨节点通信失败

现象:同节点内hixl工作正常,跨节点(不同服务器)报错。

原因:hixl默认使用共享内存(同节点内),跨节点需要配置RDMA/RoCE。

解决:启动时指定网络设备。

exportHIXL_NET_DEV=eth0# 指定RDMA网卡python train.py

结尾

hixl住在CANN五层架构第4层HCCL集合通信库上游,通过单边Put/Get操作实现零拷贝通信,同节点内比HCCL快3倍,适合参数服务器架构和PD分离场景。

适用场景:推荐算法(参数大、更新稀疏)、超大模型PS架构、需要低延迟通信的分布式系统。

参考仓库

hixl 单边通信库
hccl 集合通信库
torchtitan-npu 分布式训练
CANN 学习中心

← 返回列表