torch-npu-comm-test

通过 PyTorch torch.distributed 接口测试昇腾 NPU 通信算子性能。支持指定任意 tensor shape、dtype，使用 torchrun 启动，贴近真实训练场景的通信算子测试与性能分析。Use for testing collective communication operators (AllReduce, AllGather, ReduceScatter, etc.) with specific tensor shapes via torch.distributed on Ascend NPU.

Zur Installation springen

Repository: ascend-ai-coding/awesome-ascend-skills
Letzte Quellaktivität: 18. Mai 2026 um 12:39
Erkannte Sprache von SKILL.md: Chinesisch
Sterne: 166
Forks: 56

Installationsoptionen

Standardmäßig ist der Prompt ausgewählt, der zuerst die Quelle prüft. Sie können zu einem direkten Befehl wechseln oder eine lokale Kopie herunterladen.

Mit Codex oder Claude installieren Kopieren Sie diesen Prompt, fügen Sie ihn in Codex, Claude oder einen anderen Assistant ein und lassen Sie die Skill-Seite prüfen und installieren.

Ein direkter Befehl überspringt den Prüf-Prompt. Prüfen Sie die Quelle, bevor Sie ihn ausführen.

npx skills add https://github.com/ascend-ai-coding/awesome-ascend-skills --skill torch-npu-comm-test

Der Befehl bleibt in einer Zeile. Scrollen Sie horizontal, um ihn vor dem Kopieren vollständig zu prüfen.

Sie bevorzugen eine lokale Kopie? Laden Sie die Dateien herunter, die SkillsMP derzeit vorliegen.

Quelldateien prüfen

Lesen Sie SKILL.md und alle von SkillsMP angezeigten Begleitdateien, bevor Sie sich für eine Installation entscheiden.

Datei-Explorer

5 Dateien

SKILL.md

Quellanweisungen · Schreibgeschützte Vorschau

Zurück zur Suche Zurück zur Installation

name	torch-npu-comm-test
description	通过 PyTorch torch.distributed 接口测试昇腾 NPU 通信算子性能。支持指定任意 tensor shape、dtype，使用 torchrun 启动，贴近真实训练场景的通信算子测试与性能分析。Use for testing collective communication operators (AllReduce, AllGather, ReduceScatter, etc.) with specific tensor shapes via torch.distributed on Ascend NPU.
keywords	["torch.distributed","通信算子","collective communication","allreduce","allgather","reduce_scatter","torchrun","HCCL","性能测试","shape"]

Torch Communication Operator Test

通过 PyTorch torch.distributed 接口，使用 HCCL 后端在昇腾 NPU 上测试通信算子的功能与性能。

Overview

何时使用本 Skill（vs hccl-test）

场景	推荐工具
验证 HCCL 库基础连通性和带宽	hccl-test（mpirun）
测试特定 tensor shape 下的通信性能	torch-npu-comm-test
复现训练中某一层梯度的通信耗时	torch-npu-comm-test
测试 bf16 / fp16 等训练常用 dtype	torch-npu-comm-test
测试进程子组（subgroup）通信	torch-npu-comm-test
新集群交付验收、大规模打流	hccl-test（mpirun）

核心优势

任意 Shape：直接指定 --shape 4096,12288，而非仅数据量大小
贴近业务：通过 torch.distributed 调用，与真实训练框架一致
灵活 Dtype：支持 fp32、fp16、bf16、int32
详细统计：avg / min / max / p50 / p95 / p99 延迟 + 算法带宽 / 总线带宽
子组测试：可指定 rank 子集创建 process group 测试

Quick Reference

# 前置：确保已 source CANN 环境
source /usr/local/Ascend/ascend-toolkit/set_env.sh

# 单机 8 卡 AllReduce，shape [4096, 12288]，fp16
torchrun --nproc_per_node=8 scripts/comm-bench.py \
    --op all_reduce --shape 4096,12288 --dtype fp16

# 单机 8 卡 ReduceScatter，shape [32768, 4096]，bf16，100 次迭代
torchrun --nproc_per_node=8 scripts/comm-bench.py \
    --op reduce_scatter --shape 32768,4096 --dtype bf16 --iters 100

# 双机 16 卡 AllGather
torchrun --nnodes=2 --nproc_per_node=8 \
    --master_addr=175.99.1.2 --master_port=29500 \
    scripts/comm-bench.py --op all_gather --shape 4096,4096 --dtype fp16

# 批量测试多个算子和 shape
./scripts/batch-bench.sh --npus 8 \
    --ops "all_reduce,all_gather,reduce_scatter" \
    --shapes "4096,4096;4096,12288;32768,4096" --dtype fp16

1. Prerequisites

1.1 软件依赖

组件	要求
Python	>= 3.8
PyTorch	>= 2.1
torch_npu	与 PyTorch 版本配套
CANN	>= 8.0

1.2 环境配置

# 1. Source CANN 环境（必须）
source /usr/local/Ascend/ascend-toolkit/set_env.sh

# 2. 验证 NPU 可用
python3 -c "import torch; import torch_npu; print(f'NPU available: {torch.npu.is_available()}, count: {torch.npu.device_count()}')"

# 3. 验证 HCCL 后端
python3 -c "import torch.distributed as dist; print(f'HCCL available: {dist.is_hccl_available()}')"

1.3 多机额外要求

所有节点 CANN 版本一致
所有节点间网络互通（通信网卡可达）
建议配置 HCCL_IF_IP（指定 HCCL 通信网卡 IP）或 HCCL_SOCKET_IFNAME（指定网卡名）

2. Supported Operations

算子	`--op` 参数值	torch.distributed API	通信模式	适用场景
AllReduce	`all_reduce`	`dist.all_reduce(tensor)`	多对多	梯度聚合
AllGather	`all_gather`	`dist.all_gather_into_tensor(out, inp)`	多对多	参数收集、TP 前向
ReduceScatter	`reduce_scatter`	`dist.reduce_scatter_tensor(out, inp)`	多对多	TP 反向梯度
Broadcast	`broadcast`	`dist.broadcast(tensor, src)`	一对多	参数分发
AllToAll	`all_to_all`	`dist.all_to_all_single(out, inp)`	多对多	MoE 路由
Reduce	`reduce`	`dist.reduce(tensor, dst)`	多对一	指标汇总
P2P Send/Recv	`send_recv`	`dist.send()` / `dist.recv()`	点对点	PP 流水线

详细 API 参数与 shape 语义见 references/supported-ops.md

3. Usage

3.1 单机测试

torchrun --nproc_per_node=<NPU数> scripts/comm-bench.py \
    --op <算子> --shape <shape> --dtype <dtype> \
    [--iters <迭代次数>] [--warmup <预热次数>]

示例：

# 测试 AllReduce，模拟 LLaMA-65B 某层梯度 shape
torchrun --nproc_per_node=8 scripts/comm-bench.py \
    --op all_reduce --shape 8192,8192 --dtype bf16 --iters 200 --warmup 20

3.2 多机测试

# 在每台机器上执行（或通过 pdsh/fabric 批量下发）
torchrun --nnodes=<机器数> --nproc_per_node=<每机NPU数> \
    --node_rank=<当前机器编号> \
    --master_addr=<主节点IP> --master_port=29500 \
    scripts/comm-bench.py --op all_reduce --shape 4096,12288 --dtype fp16

3.3 子组测试

测试部分 rank 之间的通信（模拟 TP/PP 子组）：

torchrun --nproc_per_node=8 scripts/comm-bench.py \
    --op all_reduce --shape 4096,4096 --dtype fp16 \
    --group-ranks "0,1,2,3"

3.4 批量测试

./scripts/batch-bench.sh --npus 8 \
    --ops "all_reduce,all_gather,reduce_scatter" \
    --shapes "4096,4096;4096,12288;32768,4096" \
    --dtype fp16 --iters 100

4. Parameters

4.1 comm-bench.py 参数

参数	说明	默认值	示例
`--op`	通信算子类型	`all_reduce`	`--op reduce_scatter`
`--shape`	Tensor shape（逗号分隔）	`1024,1024`	`--shape 4096,12288`
`--dtype`	数据类型	`fp16`	`--dtype bf16`
`--iters`	测量迭代次数	`50`	`--iters 200`
`--warmup`	预热迭代次数	`10`	`--warmup 20`
`--group-ranks`	参与通信的 rank 列表	全部 rank	`--group-ranks "0,1,2,3"`
`--output`	输出格式	`table`	`--output json`
`--reduce-op`	Reduce 操作类型	`sum`	`--reduce-op max`
`--src-rank`	Broadcast/Reduce 的源/目标 rank	`0`	`--src-rank 0`
`--async-op`	是否使用异步操作	`false`	`--async-op`
`--check`	是否校验结果正确性	`false`	`--check`

4.2 数据类型映射

`--dtype` 值	PyTorch dtype	元素大小
`fp32`	`torch.float32`	4 bytes
`fp16`	`torch.float16`	2 bytes
`bf16`	`torch.bfloat16`	2 bytes
`int32`	`torch.int32`	4 bytes

5. Understanding Results

5.1 输出格式

================================================================
Op: all_reduce | Shape: [4096, 12288] | Dtype: fp16 | Ranks: 8
================================================================
data_size(MB)  avg(us)    min(us)    max(us)    p99(us)    algbw(GB/s)  busbw(GB/s)
96.00          1234.5     1200.1     1300.2     1298.5     77.76        135.98
================================================================

5.2 字段说明

字段	含义
`data_size(MB)`	单个 rank 参与通信的数据量
`avg(us)`	平均延迟（微秒）
`min(us)` / `max(us)`	最小 / 最大延迟
`p99(us)`	99 百分位延迟
`algbw(GB/s)`	算法带宽：`data_size / avg_time`
`busbw(GB/s)`	总线带宽：考虑算子数据搬运量的有效带宽

5.3 带宽计算公式

不同算子的数据搬运量不同，总线带宽计算方式也不同（n = world_size）：

算子	算法带宽 (algbw)	总线带宽 (busbw)
AllReduce	`S / t`	`S × 2(n-1)/n / t`
AllGather	`S×n / t`	`S × (n-1) / t`
ReduceScatter	`S / t`	`S × (n-1)/n / t`
Broadcast	`S / t`	`S / t`
Reduce	`S / t`	`S × (n-1)/n / t`

其中 S = 数据量（bytes），t = 耗时（秒），n = 参与通信的 rank 数。

6. Real-World Shape Examples

常见大模型训练中的通信 shape 参考：

模型	通信算子	典型 Shape	Dtype	场景
LLaMA-7B	AllReduce	`[4096, 4096]`	bf16	Attention 权重梯度
LLaMA-7B	AllReduce	`[4096, 11008]`	bf16	FFN 权重梯度
LLaMA-65B	AllReduce	`[8192, 8192]`	bf16	Attention 权重梯度
GPT-3 175B	AllGather	`[12288, 12288]`	fp16	TP 前向参数聚合
GPT-3 175B	ReduceScatter	`[12288, 12288]`	fp16	TP 反向梯度
MoE 模型	AllToAll	`[4096, 4096]`	bf16	Expert 路由
PP 场景	Send/Recv	`[32, 2048, 4096]`	fp16	流水线 activation

使用示例：

# 复现 LLaMA-7B FFN 梯度通信
torchrun --nproc_per_node=8 scripts/comm-bench.py \
    --op all_reduce --shape 4096,11008 --dtype bf16

# 复现 GPT-3 TP AllGather
torchrun --nproc_per_node=8 scripts/comm-bench.py \
    --op all_gather --shape 12288,12288 --dtype fp16

7. Common Issues

问题	原因	解决方法
`HCCL is not available`	torch_npu 未正确安装或 CANN 环境未 source	执行 `source set_env.sh`，重新安装 torch_npu
`RuntimeError: HCCL error`	HCCL 初始化失败（网卡/IP 配置）	检查 `HCCL_IF_IP` 或 `HCCL_SOCKET_IFNAME`
OOM (Out of Memory)	Shape 过大超出 NPU 显存	减小 shape 或使用更少 rank
多机超时	网络不通或端口不可达	检查防火墙、`MASTER_ADDR`/`MASTER_PORT`
结果校验失败	数值精度问题（fp16）	尝试 fp32 或增大 atol

详细故障排查见 references/common-issues.md

8. Scripts

8.1 comm-bench.py

主测试脚本，通过 torchrun 启动：

torchrun --nproc_per_node=8 scripts/comm-bench.py \
    --op all_reduce --shape 4096,12288 --dtype fp16 --iters 100 --warmup 10

8.2 batch-bench.sh

批量测试脚本，支持一次运行多个算子和 shape 组合：

./scripts/batch-bench.sh --npus 8 \
    --ops "all_reduce,all_gather,reduce_scatter" \
    --shapes "4096,4096;4096,12288;32768,4096" \
    --dtype fp16 --iters 100

Official References

PyTorch Distributed: https://pytorch.org/docs/stable/distributed.html
torch_npu Distributed: 见 torch_npu SKILL 中的分布式章节
HCCL 文档: https://www.hiascend.com/document/detail/zh/CANNCommunityEdition/83RC1/devaids/hccltool/HCCLpertest_16_0001.html

torch-npu-comm-test

Quellinformationen

Installationsoptionen

Quelldateien prüfen

Mehr aus diesem Repository

Torch Communication Operator Test

Overview

何时使用本 Skill（vs hccl-test）

核心优势

Quick Reference

1. Prerequisites

1.1 软件依赖

1.2 环境配置

1.3 多机额外要求

2. Supported Operations

3. Usage

3.1 单机测试

3.2 多机测试

3.3 子组测试

3.4 批量测试

4. Parameters

4.1 comm-bench.py 参数

4.2 数据类型映射

5. Understanding Results

5.1 输出格式

5.2 字段说明

5.3 带宽计算公式

6. Real-World Shape Examples

7. Common Issues

8. Scripts

8.1 comm-bench.py

8.2 batch-bench.sh

Official References

Mehr aus diesem Repository

Torch Communication Operator Test

Overview

何时使用本 Skill（vs hccl-test）

核心优势

Quick Reference

1. Prerequisites

1.1 软件依赖

1.2 环境配置

1.3 多机额外要求

2. Supported Operations

3. Usage

3.1 单机测试

3.2 多机测试

3.3 子组测试

3.4 批量测试

4. Parameters

4.1 comm-bench.py 参数

4.2 数据类型映射

5. Understanding Results

5.1 输出格式

5.2 字段说明

5.3 带宽计算公式

6. Real-World Shape Examples

7. Common Issues

8. Scripts

8.1 comm-bench.py

8.2 batch-bench.sh

Official References