企业官网建设流程全解析

• 示例工程

导读

本文围绕 DeepSpeedExamples 仓库中的benchmarks/communication通信基准测试套件展开，该套件用于在 Python 层直接测量 DeepSpeed 与 PyTorch 分布式通信原语（all_reduce、all_gather、all_to_all、pt2pt、broadcast）的延迟与带宽。通过阅读本文，你将掌握单操作与全操作基准的运行方法、全部命令行参数的语义、纯 PyTorch 与 DeepSpeed 两种运行模式的区别、Intel CPU（oneCCL）支持方式，以及如何基于模板扩展新的通信基准操作。

一、套件定位：为什么需要 Python 层通信基准

分布式训练的性能瓶颈往往藏在通信环节。与 OSU Micro-Benchmarks、NCCL Tests、oneCCL Benchmark 这类 C 级（语言级）基准不同，本套件的核心意图是在Python 层度量 DeepSpeed 或 PyTorch 分布式通信操作的延迟与带宽，其价值体现在两个方面：

• 快速定位问题层：当训练出现挂起（hang）或性能劣化时，可以便捷地判断问题究竟源自通信软件栈的哪一层，而不必逐层排查底层库；
• 度量预期通信性能：既能量化 DeepSpeed 通信封装（deepspeed.comm）的性能，也能度量纯 PyTorchtorch.distributed的通信性能，为上层训练框架的性能基线提供参照。

从源码结构看，该套件由 5 个独立的单操作基准脚本（all_reduce.py、all_gather.py、all_to_all.py、pt2pt.py、broadcast.py）、统一入口 run_all.py、共享工具 utils.py 与默认常量 constants.py 组成。

二、快速开始：两种运行方式

2.1 运行单个通信操作

以 all_reduce 为例，直接使用 DeepSpeed 启动即可。默认情况下（不带--scan），基准会计算一个恰好能装入 GPU 显存的大消息尺寸（受--mem-factor控制，默认 0.3），用于测量单一大消息下的性能：

deepspeed all_reduce.py

跨消息尺寸扫描（--scan），测量从 2 的幂次方小消息到大消息的完整延迟/带宽曲线：

deepspeed all_reduce.py --scan

2.2 运行纯 PyTorch 分布式通信（不引入 DeepSpeed）

如果只想度量原生torch.distributed的性能，通过--dist="torch"切换分布式框架即可。该模式下基准完全跳过 DeepSpeed 的导入与初始化。MPI 启动方式：

mpirun -np 16 --hostfile ${HOSTFILE} -x LD_LIBRARY_PATH -x PATH -x LD_PRELOAD python all_reduce.py --scan --dist="torch"

Slurm 启动方式：

srun -n 16 python all_reduce.py --scan --dist="torch"

从 utils.py 的init_torch_distributed实现可以看到，纯 PyTorch 模式下基准会从多种环境变量（如OMPI_COMM_WORLD_RANK、SLURM_PROCID、MPI_LOCALRANKID等）自动发现 rank、local_rank 与 world_size；若未设置MASTER_ADDR，则会尝试通过 mpi4py 获取主节点的地址并广播给所有 rank，默认端口为29500（见 constants.py 中的TORCH_DISTRIBUTED_DEFAULT_PORT）。

2.3 运行全部通信基准

deepspeed run_all.py

run_all.py与单个基准脚本共享同一套参数，扫描参数（如最大消息尺寸、带宽单位等）会被透传给每个通信操作。从 run_all.py 的main()可以看出，它按参数依次调用run_all_reduce、run_all_gather、run_all_to_all、run_pt2pt、run_broadcast，未指定任何操作时默认运行全部 5 种。

若希望有选择地运行部分操作，将其作为参数传入即可，例如只跑 all_reduce、all_to_all 与 broadcast：

deepspeed run_all.py --scan --all-reduce --all-to-all --broadcast

三、命令行参数全解

run_all.py与ds_bench的完整参数如下（默认值取自 constants.py 与 utils.py 中的benchmark_parser）：

usage: ds_bench [-h] [--local_rank LOCAL_RANK] [--trials TRIALS] [--warmups WARMUPS] [--maxsize MAXSIZE] [--async-op] [--bw-unit {Gbps,GBps}] [--backend {nccl}] [--dist {deepspeed,torch}] [--scan] [--raw] [--all-reduce] [--all-gather] [--all-to-all] [--pt2pt] [--broadcast] [--dtype DTYPE] [--mem-factor MEM_FACTOR] [--debug] optional arguments: -h, --help show this help message and exit --local_rank LOCAL_RANK --trials TRIALS Number of timed iterations（计时迭代次数，默认 50） --warmups WARMUPS Number of warmup (non-timed) iterations（预热迭代次数，默认 5） --maxsize MAXSIZE Max message size as a power of 2（最大消息尺寸，以 2 的幂表示，默认 24，即 2^24 个元素） --async-op Enables non-blocking communication（启用非阻塞通信） --bw-unit {Gbps,GBps} 带宽单位（默认 Gbps） --backend {nccl} Communication library to use（通信后端，支持 nccl / ccl / mpi） --dist {deepspeed,torch} Distributed DL framework to use（分布式框架，默认 deepspeed） --scan Enables scanning all message sizes（扫描全部消息尺寸） --raw Print the message size and latency without units（原始输出，不带单位） --all-reduce Run all_reduce --all-gather Run all_gather --all-to-all Run all_to_all --pt2pt Run pt2pt --broadcast Run broadcast --dtype DTYPE PyTorch tensor dtype（默认 float） --mem-factor MEM_FACTOR Proportion of max available GPU memory to use for single-size evals（单尺寸评估占用可用显存的比例，默认 0.3） --debug Enables all_to_all debug prints（启用 all_to_all 调试打印）

几个关键参数的底层影响：

• --trials/--warmups：以 all_reduce.py 的timed_all_reduce为例，先执行warmups次通信用于建立连接、预热 CUDA context，再用 CUDA Event（torch.cuda.Event(enable_timing=True)）记录trials次通信的累计耗时并取平均。预热迭代数太少会导致计时包含连接建立开销，测量值偏高。
• --maxsize：扫描模式下消息元素数从2^1到2^(maxsize-1)递增（见 all_reduce.py），默认--maxsize 24对应约 2^23 个元素的最大扫描尺寸。
• --mem-factor：决定单尺寸评估时张量占用 GPU 显存的比例。对于无需输出张量的操作（all_reduce、pt2pt、broadcast），源码会将该参数翻倍使用（注释为 "Don't need output tensor, so we double mem_factor"），见 all_reduce.py。all_gather 由于输出张量随世界大小线性放大，且支持all_gather_into_tensor时额外增加 0.2 的余量（见 all_gather.py）。遇到 OOM 时按脚本提示降低该值即可。
• --async-op：将阻塞式通信替换为async_op=True的非阻塞调用；pt2pt 场景下对应isend/irecv配对（见 pt2pt.py）。
• --backend：通信库后端。源码层面 DeepSpeed 模式通过deepspeed.init_distributed(dist_backend=backend)初始化（见 utils.py），实际可用后端由加速器决定，默认取get_accelerator().communication_backend_name()（如 NVIDIA 环境下为 nccl）。

ds_bench是预打包在 DeepSpeed 安装目录中的run_all.py封装，参数完全一致：

<path to deepspeed>/bin/ds_bench --scan --trials=10

四、输出指标解读：Throughput 与 BusBW 的计算口径

每次计时结束后，基准会打印一行形如Size (Bytes) / Description / Duration / Throughput / BusBW的结果。三个指标的计算逻辑集中在 utils.py 的get_bw中，且不同通信操作的吞吐与总线带宽公式不同：

其中size为单 rank 的消息字节数，n为 world size。选择Gbps单位时两个指标统一乘以 8（字节转比特），见 utils.py。BusBW 反映了集合通信算法在总线上的实际有效带宽，比原始吞吐更能衡量硬件互连的利用效率。

需要留意的是，all_gather 的size在公式中先乘以了n（每个 rank 最终持有 n 份数据的拼接结果）；而 all_reduce 假设了"两两通信"模型，因此乘以 2。--raw参数可输出不带单位的原始数值，便于脚本化处理或与 C 级基准结果对比。

五、Intel CPU 支持（oneCCL）

除 GPU 外，该套件还支持通过 oneCCL 在 Intel CPU 等设备上运行，只需在所有 Python 脚本后追加--device cpu参数。例如在 Intel CPU 上运行单一大消息尺寸的 all_reduce：

deepspeed all_reduce.py --device cpu

从源码实现看，--device参数默认值为cuda（见 utils.py 与 constants.py 的DEFAULT_DEVICE），并可通过--backend ccl指定 oneCCL 后端。需要说明的是，CPU 模式下当前计时依赖torch.cpu.Event，而各timed_*函数在device == "cpu"时仅打印 "No Event support on CPU to measure time for now" 后直接返回（见 all_reduce.py），因此 CPU 路径目前主要面向通信逻辑验证，尚不具备与 GPU 同等的计时能力。

六、深入理解：基准的底层实现机制

6.1 计时流程

所有操作遵循统一的三段式计时流程（以 broadcast.py 的timed_broadcast为例）：

1. sync_all()：先同步设备（get_accelerator().synchronize()+dist.barrier()，见 utils.py），确保所有 rank 就绪；
2. 执行warmups次通信并再次同步，完成连接建立与预热；
3. 记录起始 CUDA Event，循环执行trials次通信，记录结束 Event，取平均单次耗时。

6.2 显存安全的消息尺寸计算

max_numel（utils.py）针对不同操作做了差异化处理，避免 OOM：

• all_reduce / pt2pt / broadcast：元素数 = 可用显存 × mem_factor ÷ 单元素字节数；
• all_gather：由于输出缓冲随 world size 放大，先除以 world size，再向下取整到最近的 2 的幂（非 2 的幂尺寸下 all_gather 性能会下降）；
• all_to_all：元素数必须能被 world size 整除（分块语义），同样向下取整到 2 的幂。

扫描模式下，张量按torch.ones(world_size, M)构造并以global_rank填充，便于调试时区分各 rank 的数据；分配失败捕获RuntimeError，若为 OOM 则打印警告并提前退出（见 all_reduce.py）。

6.3 两种分布式框架的初始化差异

• DeepSpeed 模式（--dist=deepspeed，默认）：调用deepspeed.init_distributed(dist_backend=backend)，后续通过import deepspeed.comm as dist使用 DeepSpeed 的通信封装；
• PyTorch 模式（--dist=torch）：调用torch.distributed.init_process_group(backend)，并自动从 MPI/Slurm 环境变量推断分布式参数（见 utils.py）。

值得一提的实现细节：all_gather 在两种模式下分别调用TorchBackend.get_all_gather_function()与dist.allgather_fn（见 all_gather.py），确保统一走底层all_gather_into_tensor以节省内存。

七、扩展指南：添加新的通信基准

本套件设计上便于横向扩展，添加新通信操作的通用流程（以添加reduce_scatter为例）如下：

1. 复制模板：拷贝一个结构相似的基准文件作为起点，例如复制all_reduce.py得到reduce_scatter.py；
2. 三处配套修改：
   • 在 utils.py 的get_bw中为新操作添加带宽公式；
   • 在max_numel（utils.py）中为新操作添加最大张量元素数公式（注意其显存与整除约束，如参考 all_gather 的 world size 缩放处理）；
   • 在benchmark_parser（utils.py）中为新操作添加--<op>布尔参数；
3. 替换通信调用：用新脚本中查找替换（find-replace）的方式，将模板里的dist.all_reduce(...)替换为新的通信原语调用；
4. 确定默认 mem_factor：为run_<collective>_single()找到合适的默认显存占用比例，避免默认运行即 OOM；
5. 接入统一入口：在 run_all.py 中导入新操作并加入ops_to_run的调度逻辑。

完成上述五步后，新操作即可通过deepspeed reduce_scatter.py单独运行，也能通过deepspeed run_all.py --reduce-scatter与其他操作联合扫描。

八、典型使用场景小结

• 训练前摸底：用deepspeed run_all.py --scan一次性获取全部 5 种操作在各消息尺寸下的延迟与 BusBW 曲线，建立集群通信性能基线；
• 问题定位：分别运行 DeepSpeed 模式与--dist="torch"模式，对比结果即可判断性能劣化来自 DeepSpeed 通信封装层还是底层库；
• 与 C 级基准对照：使用--raw输出无单位原始数据，便于与 NCCL Tests / OSU 等 C 级基准结果交叉验证；
• 显存受限场景：通过--mem-factor缩放消息尺寸，在有限显存下完成大消息性能评估。

该套件全部源码与配置均位于仓库 benchmarks/communication 目录下，读者可结合各timed_*函数与get_bw公式深入理解每个指标的精确口径，也可按第七节的扩展流程定制自己的通信基准。

• 示例工程