Through the PyTorch torch.distributed interface, using the HCCL backend to test the functionality and performance of communication operators on Ascend NPUs.
Overview
When to use this Skill (vs hccl-test)
| Scenario | Recommended tool |
|---|---|
| Verify basic HCCL library connectivity and bandwidth | hccl-test（mpirun） |
| Testspecific tensor shape communication performance under | torch-npu-comm-test |
| Reproduce the communication time of a certain layer's gradient during training | torch-npu-comm-test |
| Test commonly used training dtypes such as bf16 / fp16 | torch-npu-comm-test |
| Test process subgroup communication | torch-npu-comm-test |
| New cluster delivery acceptance, large-scale traffic testing | hccl-test（mpirun） |
Core advantages
- Arbitrary Shape: directly specify
--shape 4096,12288, rather than only the data size - Close to business: through
torch.distributedcalls, consistent with real training frameworks - Flexible Dtype: supports fp32, fp16, bf16, int32
- Detailed statistics: avg / min / max / p50 / p95 / p99 latency + algorithmic bandwidth / bus bandwidth
- Subgroup testing: can specify a subset of ranks to create a process group for testing
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 Software dependencies
| Component | Requirement |
|---|---|
| Python | >= 3.8 |
| PyTorch | >= 2.1 |
| torch_npu | Matching the PyTorch version |
| CANN | >= 8.0 |
1.2 Environment configuration
# 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 Additional requirements for multi-machine
- Consistent CANN version across all nodes
- Network connectivity between all nodes (communication NICs reachable)
- Recommended to configure
HCCL_IF_IP(specify HCCL communication NIC IP) orHCCL_SOCKET_IFNAME(specify NIC name)
2. Supported Operations
| Operator | --op Parameter value |
torch.distributed API | Communication mode | Applicable scenario |
|---|---|---|---|---|
| AllReduce | all_reduce |
dist.all_reduce(tensor) |
Many-to-many | Gradient aggregation |
| AllGather | all_gather |
dist.all_gather_into_tensor(out, inp) |
Many-to-many | Parameter gathering, TP forward |
| ReduceScatter | reduce_scatter |
dist.reduce_scatter_tensor(out, inp) |
Many-to-many | TP backward gradient |
| Broadcast | broadcast |
dist.broadcast(tensor, src) |
One-to-many | Parameter distribution |
| AllToAll | all_to_all |
dist.all_to_all_single(out, inp) |
Many-to-many | MoE routing |
| Reduce | reduce |
dist.reduce(tensor, dst) |
Many-to-one | Metric aggregation |
| P2P Send/Recv | send_recv |
dist.send() / dist.recv() |
Point-to-point | PP pipeline |
For detailed API parameters and shape semantics, see references/supported-ops.md
3. Usage
3.1 Single-machine test
torchrun --nproc_per_node=<NPU数> scripts/comm-bench.py \
--op <算子> --shape <shape> --dtype <dtype> \
[--iters <迭代次数>] [--warmup <预热次数>]
Example：
# 测试 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 Multi-machine test
# 在每台机器上执行（或通过 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 Subgroup test
Test communication between some ranks (simulating TP/PP subgroups):
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 Batch test
./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 parameters
| Parameter | Description | Default value | Example |
|---|---|---|---|
--op |
Communication operator type | all_reduce |
--op reduce_scatter |
--shape |
Tensor shape (comma-separated) | 1024,1024 |
--shape 4096,12288 |
--dtype |
Data type | fp16 |
--dtype bf16 |
--iters |
Number of measurement iterations | 50 |
--iters 200 |
--warmup |
Number of warmup iterations | 10 |
--warmup 20 |
--group-ranks |
List of ranks participating in communication | All ranks | --group-ranks "0,1,2,3" |
--output |
Output format | table |
--output json |
--reduce-op |
Reduce operation type | sum |
--reduce-op max |
--src-rank |
Source/destination rank of Broadcast/Reduce | 0 |
--src-rank 0 |
--async-op |
Whether to use asynchronous operations | false |
--async-op |
--check |
Whether to verify the correctness of the results | false |
--check |
4.2 Data Type Mapping
--dtype Value |
PyTorch dtype | Element size |
|---|---|---|
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 Output Format
================================================================
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 Field Description
| Field | Meaning |
|---|---|
data_size(MB) |
Amount of data involved in communication per rank |
avg(us) |
Average latency (microseconds) |
min(us) / max(us) |
Minimum / maximum latency |
p99(us) |
99th percentile latency |
algbw(GB/s) |
Algorithm bandwidth:data_size / avg_time |
busbw(GB/s) |
Bus bandwidth: effective bandwidth considering the operator's data transfer volume |
5.3 Bandwidth Calculation Formula
Different operators have different data transfer volumes, and the bus bandwidth calculation method also differs (n = world_size):
| Operator | Algorithm bandwidth (algbw) | Bus bandwidth (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 |
Where S = data volume (bytes),t = elapsed time (seconds),n = number of ranks participating in communication.
6. Real-World Shape Examples
Common communication shape references in large model training:
| Model | Communication operator | Typical Shape | Dtype | Scenario |
|---|---|---|---|---|
| LLaMA-7B | AllReduce | [4096, 4096] |
bf16 | Attention weight gradient |
| LLaMA-7B | AllReduce | [4096, 11008] |
bf16 | FFN weight gradient |
| LLaMA-65B | AllReduce | [8192, 8192] |
bf16 | Attention weight gradient |
| GPT-3 175B | AllGather | [12288, 12288] |
fp16 | TP forward parameter aggregation |
| GPT-3 175B | ReduceScatter | [12288, 12288] |
fp16 | TP backward gradient |
| MoE model | AllToAll | [4096, 4096] |
bf16 | Expert routing |
| PP scenario | Send/Recv | [32, 2048, 4096] |
fp16 | Pipeline activation |
Usage example：
# 复现 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
| Problem | Cause | Solution |
|---|---|---|
HCCL is not available |
torch_npu is not installed correctly or the CANN environment has not been sourced | Execute source set_env.sh, reinstall torch_npu |
RuntimeError: HCCL error |
HCCL initialization failed (NIC/IP configuration) | Check HCCL_IF_IP or HCCL_SOCKET_IFNAME |
| OOM (Out of Memory) | Shape is too large and exceeds NPU memory | Reduce the shape or use fewer ranks |
| Multi-machine timeout | Network unreachable or port unreachable | Check firewall,MASTER_ADDR/MASTER_PORT |
| Result verification failed | Numerical precision issue (fp16) | Try fp32 or increase atol |
For detailed troubleshooting, see references/common-issues.md
8. Scripts
8.1 comm-bench.py
Main test script, started via 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
Batch test script, supports running multiple operator and shape combinations at once:
./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: See torch_npu SKILL distributed chapter in
- HCCL documentation: https://www.hiascend.com/document/detail/zh/CANNCommunityEdition/83RC1/devaids/hccltool/HCCLpertest_16_0001.html