启动 Colossal-AI

作者: Chuanrui Wang, Shenggui Li, Siqi Mai

预备知识:

• 分布式训练
• Colossal-AI 总览

简介

正如我们在前面的教程中所提到的，在您的配置文件准备好后，您需要为 Colossal-AI 初始化分布式环境。我们把这个过程称为 launch。在本教程中，您将学习如何在您的服务器上启动 Colossal-AI，不管是小型的还是大型的。

在 Colossal-AI 中，我们提供了几种启动方法来初始化分布式后端。 在大多数情况下，您可以使用 colossalai.launch 和 colossalai.get_default_parser 来通过命令行传递参数。如果您想使用 SLURM、OpenMPI 和 PyTorch 等启动工具，我们也提供了几个启动的辅助方法以便您的使用。您可以直接从这些启动工具设置的环境变量中访问 rank 和 world size 大小。

在本教程中，我们将介绍如何启动 Colossal-AI 来初始化分布式后端：

• 用 colossalai.launch 启动
• 用 Colossal-AI命令行 启动
• 用 SLURM 启动
• 用 OpenMPI 启动

启动分布式环境

为了启动 Colossal-AI，我们需要两类参数:

1. 配置文件
2. 分布式设置

无论我们使用何种启动方式，配置文件是必须要求的，而分布式设置有可能依情况而定。配置文件可以是配置文件的路径或 Python dictionary 的形式。分布式设置可以通过命令行或多进程启动器传递。

命令行解析器

在使用 launch 之前, 我们首先需要了解我们需要哪些参数来进行初始化。 如分布式训练 中 基本概念 一节所述 ，涉及的重要参数是:

1. host
2. port
3. rank
4. world_size
5. backend

在 Colossal-AI 中，我们提供了一个命令行解析器，它已经提前添加了这些参数。您可以通过调用 colossalai.get_default_parser() 来获得这个解析器。这个解析器通常与 colossalai.launch 一起使用。

# add these lines in your train.py
import colossalai

# get default parser
parser = colossalai.get_default_parser()

# if you want to add your own arguments
parser.add_argument(...)

# parse arguments
args = parser.parse_args()

您可以在您的终端传入以下这些参数。

python train.py --host <host> --rank <rank> --world_size <world_size> --port <port> --backend <backend>

backend 是用户可选的，默认值是 nccl。

本地启动

为了初始化分布式环境，我们提供了一个通用的 colossalai.launch API。colossalai.launch 函数接收上面列出的参数，并在通信网络中创建一个默认的进程组。方便起见，这个函数通常与默认解析器一起使用。

import colossalai

# parse arguments
args = colossalai.get_default_parser().parse_args()

# launch distributed environment
colossalai.launch(config=args.config,
                  rank=args.rank,
                  world_size=args.world_size,
                  host=args.host,
                  port=args.port,
                  backend=args.backend
)

用 Colossal-AI命令行工具 启动

为了更好地支持单节点以及多节点的训练，我们通过封装PyTorch的启动器实现了一个更加方便的启动器。 PyTorch自带的启动器需要在每个节点上都启动命令才能启动多节点训练，而我们的启动器只需要一次调用即可启动训练。

首先，我们需要在代码里指定我们的启动方式。由于这个启动器是PyTorch启动器的封装，那么我们自然而然应该使用colossalai.launch_from_torch。 分布式环境所需的参数，如 rank, world size, host 和 port 都是由 PyTorch 启动器设置的，可以直接从环境变量中读取。

config.py

BATCH_SIZE = 512
LEARNING_RATE = 3e-3
WEIGHT_DECAY = 0.3
NUM_EPOCHS = 2

train.py

import colossalai

colossalai.launch_from_torch(
    config="./config.py",
)
...

接下来，我们可以轻松地在终端使用colossalai run来启动训练。下面的命令可以在当前机器上启动一个4卡的训练任务。 你可以通过设置nproc_per_node来调整使用的GPU的数量，也可以改变master_port的参数来选择通信的端口。

# 在当前节点上启动4卡训练 （默认使用29500端口）
colossalai run --nproc_per_node 4 train.py

# 在当前节点上启动4卡训练，并使用一个不同的端口
colossalai run --nproc_per_node 4 --master_port 29505 test.py

如果你在使用一个集群，并且想进行多节点的训练，你需要使用Colossal-AI的命令行工具进行一键启动。我们提供了两种方式来启动多节点任务

• 通过--hosts来启动

这个方式适合节点数不多的情况。假设我们有两个节点，分别为host和host2。我们可以用以下命令进行多节点训练。 比起单节点训练，多节点训练需要手动设置--master_addr （在单节点训练中master_addr默认为127.0.0.1）。

警告

多节点训练时，master_addr不能为localhost或者127.0.0.1，它应该是一个节点的名字或者IP地址。

# 在两个节点上训练
colossalai run --nproc_per_node 4 --host host1,host2 --master_addr host1 test.py

• 通过--hostfile来启动

这个方式适用于节点数很大的情况。host file是一个简单的文本文件，这个文件里列出了可以使用的节点的名字。 在一个集群中，可用节点的列表一般由SLURM或者PBS Pro这样的集群资源管理器来提供。比如，在SLURM中， 你可以从SLURM_NODELIST这个环境变量中获取到当前分配列表。在PBS Pro中，这个环境变量为PBS_NODEFILE。 可以通过echo $SLURM_NODELIST 或者 cat $PBS_NODEFILE 来尝试一下。如果你没有这样的集群管理器， 那么你可以自己手动写一个这样的文本文件即可。

提供给Colossal-AI的host file需要遵循以下格式，每一行都是一个节点的名字。

host1
host2

如果host file准备好了，那么我们就可以用以下命令开始多节点训练了。和使用--host一样，你也需要指定一个master_addr。 当使用host file时，我们可以使用一些额外的参数：

• --include: 设置你想要启动训练的节点。比如，你的host file里有8个节点，但是你只想用其中的6个节点进行训练， 你可以添加--include host1,host2,host3,...,host6，这样训练任务只会在这6个节点上启动。

• --exclude: 设置你想排除在训练之外的节点。当你的某一些节点坏掉时，这个参数会比较有用。比如假如host1的GPU有一些问题，无法正常使用， 那么你就可以使用--exclude host1来将其排除在外，这样你就可以训练任务就只会在剩余的节点上启动。

# 使用hostfile启动
colossalai run --nproc_per_node 4 --hostfile ./hostfile --master_addr host1  test.py

# 只使用部分节点进行训练
colossalai run --nproc_per_node 4 --hostfile ./hostfile --master_addr host1  --include host1 test.py

# 不使用某些节点进行训练
colossalai run --nproc_per_node 4 --hostfile ./hostfile --master_addr host1  --exclude host2 test.py

用 SLURM 启动

如果您是在一个由 SLURM 调度器管理的系统上， 您也可以使用 srun 启动器来启动您的 Colossal-AI 脚本。我们提供了辅助函数 launch_from_slurm 来与 SLURM 调度器兼容。 launch_from_slurm 会自动从环境变量 SLURM_PROCID 和 SLURM_NPROCS 中分别读取 rank 和 world size ，并使用它们来启动分布式后端。

您可以在您的训练脚本中尝试以下操作。

import colossalai

colossalai.launch_from_slurm(
    config=<CONFIG>,
    host=args.host,
    port=args.port
)

您可以通过在终端使用这个命令来初始化分布式环境。

srun python train.py --host <master_node> --port 29500

用 OpenMPI 启动

如果您对OpenMPI比较熟悉，您也可以使用 launch_from_openmpi 。 launch_from_openmpi 会自动从环境变量 OMPI_COMM_WORLD_LOCAL_RANK， MPI_COMM_WORLD_RANK 和 OMPI_COMM_WORLD_SIZE 中分别读取local rank、global rank 和 world size，并利用它们来启动分布式后端。

您可以在您的训练脚本中尝试以下操作。

colossalai.launch_from_openmpi(
    config=<CONFIG>,
    host=args.host,
    port=args.port
)

以下是用 OpenMPI 启动多个进程的示例命令。

mpirun --hostfile <my_hostfile> -np <num_process> python train.py --host <node name or ip> --port 29500

• --hostfile: 指定一个要运行的主机列表。
• --np: 设置总共要启动的进程（GPU）的数量。例如，如果 --np 4，4个 python 进程将被初始化以运行 train.py。