用途:在昇腾 NPU 上通过 PyTorch torch.distributed(HCCL 后端)测试通信算子的功能与性能,适合验证特定 tensor shape 下的通信性能、复现训练中某层梯度的通信耗时、测试 bf16/fp16 等训练常用 dtype 以及进程子组通信。若只需验证 HCCL 基础连通性和带宽或做新集群验收,建议改用 hccl-test。
主要能力:支持 AllReduce、AllGather、ReduceScatter、Broadcast、AllToAll、Reduce、P2P Send/Recv;可指定任意 shape(如 --shape 4096,12288)与 dtype(fp32、fp16、bf16、int32);输出 avg/min/max/p50/p95/p99 延迟及算法带宽、总线带宽;支持 --group-ranks 子组测试、--check 结果校验、--output json 等。
使用方式:先 source CANN 环境(source /usr/local/Ascend/ascend-toolkit/set_env.sh),单机示例:torchrun --nproc_per_node=8 scripts/comm-bench.py --op all_reduce --shape 4096,12288 --dtype fp16;多机需指定 --nnodes、--node_rank、--master_addr、--master_port;批量测试用 ./scripts/batch-bench.sh。依赖 Python>=3.8、PyTorch>=2.1、配套 torch_npu、CANN>=8.0。
限制:多机要求各节点 CANN 版本一致且网络互通,建议配置 HCCL_IF_IP 或 HCCL_SOCKET_IFNAME;shape 过大会 OOM;fp16 结果校验可能因精度失败。具体脚本与参数细节以资料为准,未提供的收费、平台等信息暂未提供。