PCCL: Megatron 5D 并行 Device Order Search 工具指南（v2.1）-真武 PPU 云服务(ppu)-阿里云帮助中心

1. 概述

目前 LLM 大规模模型训练框架 Megatron-LM 并行策略越来越多，也越来越复杂，Tensor Parallel, Context Parallel，Pipeline Parallel, Data Parallel 等等，这些并行策略的使用，会选择不同的 PPUs 组合作为通信 Group，这在目前不支持 ICN Switch 架构的平台上，比如对真武810E 来说，经常会选择不是完全 ICN 互联的 PPUs 作为通信 Group，这样的选择会导致通信性能下降，最终影响端到端的模型训练性能。

Device Order Search 工具就是为了解决这个问题而开发的，它会通过搜索的方式，尽力去保证所有的并行通信 Group 内的 PPUs 都是 ICN 互联的。

用户只需要将 Search 出来的 PPU device order 设置给 CUDA_VISIBLE_DEVICES 即可。

1.1 支持的并行类型

通信 Group	备注
TP	Used for Tensor Parallel
CP	Used for Context Parallel
EP	Used for Expert Parallel
DP	Used for Data Parallel
PP	Used for Pipeline Parallel
TP * CP	Used for loading balance to calculate the auxiliary loss
TP * EP	Used for token dispatcher
DP * CP	Used for Data Parallel when Context Parallel also enabled
(DP * CP) Modulo EP	Used for MoE layer Data Parallel when both Expert Parallel and Context Parallel also enabled
DP Modulo EP	Used for MoE layer Data Parallel when Expert Parallel also enabled

1.2 重点优化

在上述这些并行策略下，我们通常只需要满足 Group 内的 PPUs 能组成 ring graph 即可满足 ICN 互联的需求，但是依然有些场景需要特别注意：

真武810E 上面存在部分 PPUs 之间有 2 根 ICN link 的情况，对于 TP 2 场景，我们会优先考虑带有 2 根 ICN link 的切分策略，保证 TP 性能最大化。
对于 Context Parallel + TE 场景，计算和通信存在 Overlap 情况，Ring Attention 采用了 P2P Send/Recv 的通信 Pattern，我们会优先保证 Rank 顺序的 PPUs 是完全 ICN 互联的，避免转发造成的性能损失。

当然如果符合优化场景的并行策略最终没有能 Search 出来 ICN 互联的设置，我们也会最终 fallback 到普通的 ring search。

1.3 发布形式

该 Search 工具脚本，会随着 SDK 发布出去，具体路径：PPU_SDK/pccl_tools/DeviceOrderSearch.py

2. 使用介绍

可以通过 --help 参考工具的使用方法

usage: DeviceOrderSearch.py [-h] [--tensor-model-parallel-size TENSOR_MODEL_PARALLEL_SIZE] [--pipeline-model-parallel-size PIPELINE_MODEL_PARALLEL_SIZE] [--context-parallel-size CONTEXT_PARALLEL_SIZE] [--expert-model-parallel-size EXPERT_MODEL_PARALLEL_SIZE]
                            [--use-tp-pp-dp-mapping] [--world-size WORLD_SIZE] [--rank RANK] [--ppu-num-per-node] [--devices-order  [...]]

optional arguments:
  -h, --help            show this help message and exit
  --ppu-num-per-node    configure ppu number for host by hand
  --devices-order  [ ...]
                        specified device order to search

distributed:
  --tensor-model-parallel-size TENSOR_MODEL_PARALLEL_SIZE, --tp TENSOR_MODEL_PARALLEL_SIZE
                        Degree of tensor model parallelism.
  --pipeline-model-parallel-size PIPELINE_MODEL_PARALLEL_SIZE, --pp PIPELINE_MODEL_PARALLEL_SIZE
                        Degree of pipeline model parallelism.
  --context-parallel-size CONTEXT_PARALLEL_SIZE, --cp CONTEXT_PARALLEL_SIZE
                        Degree of context parallelism.
  --expert-model-parallel-size EXPERT_MODEL_PARALLEL_SIZE, --ep EXPERT_MODEL_PARALLEL_SIZE
                        Degree of expert model parallelism.
  --use-tp-pp-dp-mapping
                        If set, distributed ranks initialize order is changed from tp-dp-pp to tp-pp-dp. Make sure EP and CP aren't used with this option enabled
  --world-size WORLD_SIZE
                        configure whole world_size.
  --rank RANK           configure current rank.

为了方便使用，我们对不同并行策略参数名字也做了优化，比如：

你可以使用 --tensor-model-parallel-size 指定 TP 的维度，也可以简化使用 --tp 来使用, 其他 --cp、--pp、 --ep、 --dp 同理。
你可以在带有 PPU 的物理机上面进行运行，脚本会自动检测 ICN Topology，你也可以在 CPU 机器上面，通过指定 --ppu-num-per-node 8 or 16, 脚本会自动生成对应的 ICN topology。
脚本默认会打印所有并行维度大于 1 的 Rank 组合，方便对比，检测和分析。

使用示范如下：

拿到上面 Device Order Search 出来的配置之后，在训练脚本中，加入 CUDA_VISIBLE_DEVICES 的配置命令：

export CUDA_VISIBLE_DEVICES="0,3,1,2,4,7,5,6,9,10,8,11,13,14,12,15"