GPU 实例训练 LLaMA2-7B 模型

创建 GPU 实例

1. 登录潞晨云网站
   1. 登录潞晨云网站后，进入“GPU Instances”页面。
2. 创建实例
   1. （可选）配置共享存储：进入共享存储页面，点击创建存储，根据需要创建合适的共享存储空间。
   2. 点击“Create GPU Instance”按钮，开始创建新的 GPU 实例。
3. 选择配置
   1. 选择 GPU 类型：H100-NVLINK-80GB。
   2. 选择服务器地点。
   3. （可选）挂载之前的共享存储。
4. 配置实例信息
   1. 实例名称：colossal
   2. 镜像：选择 CUDA (12.1)ubuntu==20.04, cuda==12.1, python==3.11, conda。这里提供了多种镜像供选择，用户可根据自身需求选择合适的镜像或上传自定义镜像。
   3. 外部存储：本文未使用外部存储，因为数据量较小，可以直接存放在机器内部。
5. 初始化实例
   1. 完成上述配置后，等待实例初始化完成。
   2. 初始化完成后，可以通过 SSH 连接到实例。

环境配置

安装 Colossal-AI

• 克隆官方仓库

git clone https://github.com/hpcaitech/ColossalAI.git
cd ColossalAI

• 编译 Kernels

BUILD_EXT=1 pip install .

• 安装 Flash Attention

pip install flash-attn --no-build-isolation

• 安装 NVIDIA Apex

git clone https://github.com/NVIDIA/apex
cd apex
git checkout 22.04-dev

# 修改 setup.py 文件 36行
#if (bare_metal_major != torch_binary_major) or (bare_metal_minor != torch_binary_minor):
#    pass
# 安装依赖
pip install -r requirements.txt

# 编译安装
pip install -v --disable-pip-version-check --no-cache-dir --global-option="--cpp_ext" --global-option="--cuda_ext" ./
python setup.py install

由于 torch 的新版本不支持 torch._six 模块，所以需要修改以下代码 在文件/root/.local/lib/python3.11/site-packages/apex/amp/_initialize.py 中,注释掉第 2 行，即

from torch._six import string_classes

在第 42 行，string_classes 换为 str

• 安装 TensorNVME

pip install camke
git clone https://github.com/hpcaitech/TensorNVMe.git
pip install -v --no-cache-dir .

注：由于该教程默认使用 Colossal-AI main 分支代码（2024.12.8），因此，若由于主分支升级而出现部分包缺失缺失或安装问题，可自行通过 pip 进行安装。

训练 LLaMA2-7B 模型

1. 准备训练脚本

   1. 修改 /root/my_env/ColossalAI/examples/language/llama/scripts/benchmark_7B/gemini.sh 脚本：

   #!/bin/bash
   cd ../../
   export OMP_NUM_THREADS=64
   colossalai run --nproc_per_node 8 benchmark.py -g -x -b 6 -s 10 --shard_param_frac 0

2. 执行训练脚本

bash /root/my_env/ColossalAI/examples/language/llama/scripts/benchmark_7B/gemini.sh

结果图：

• 尝试 3D 并行
• dp8+tp1+pp1+zero2

  colossalai run --nproc_per_node 8 benchmark.py -p 3d -g -x -b 16 -s 10 --tp 4 --zero 2

结果图：

• dp8+tp1+pp1+zero2

  colossalai run --nproc_per_node 8 benchmark.py -p 3d -g -x -b 128 -s 10 --pp 2 --zero 1

结果图：

多机训练指南

以下多机启动指南同样适用于使用 torch.distributed 进行分布式通信初始化的框架如 megatron,fsdp,nanotron 等框架。

1. 环境准备

确保所有节点上都安装了相同版本的 Python 和 PyTorch。推荐使用 Anaconda 来管理环境，以避免版本冲突。

• 如果使用 IB 网络，需要安装必要的软件perftest。
• 为了在不同节点之间无密码执行命令，需要配置 SSH 无密码登录：
  • 生成 SSH 密钥对（在主节点上）。
  • 将公钥复制到所有节点。

2. 配置环境变量

在脚本中设置必要的环境变量，这些变量用于配置分布式训练的参数。

• MASTER_ADDR: 主节点的 IP 地址，所有其他节点将连接到这个地址。
• MASTER_PORT: 主节点的端口号，用于通信。
• NUM_GPUS_PER_NODE: 每个节点上可用的 GPU 数量。
• NNODES: 参与训练的节点总数，由node_ip_list.txt中的 IP 数量确定。
• NCCL_DEBUG (可选): 设置为INFO或更高级别可以帮助调试 NCCL 相关的通信问题。

export MASTER_ADDR="10.42.128.230"
export MASTER_PORT="12345"
export NUM_GPUS_PER_NODE=8
export NNODES=$(grep -c '^' /root/ColossalAI/examples/language/llama/node_ip_list.txt)
export NCCL_DEBUG=INFO

注意：上面MASTER_ADDR只是一个示例，用户需要根据实际的 IP 地址进行替换。在服务器上可以通过以下命令查看 IP 地址（需先安装相应的工具iproute2, 安装方式：apt install iproute2）：

ip addr show

3. 创建节点列表文件

创建一个文本文件（如node_ip_list.txt），每行记录一个参与训练的节点的 IP 地址。这用于确定参与训练的节点数量。例如：

10.42.128.230
10.42.128.231
10.42.128.232

4. 编写启动脚本（一种简单的拉起多机运行的脚本）

定义启动训练的函数，该函数负责在每个节点上启动训练进程。如果是主节点，直接运行训练脚本；如果不是主节点，通过 SSH 远程执行。

start_training_on_node() {
    local ip=$1
    local rank=$2
    if [ "$ip" == "$MASTER_ADDR" ]; then
        # 如果是主节点，则直接运行
        torchrun \
            --nnodes=$NNODES \
            --node_rank=$rank \
            --nproc_per_node=$NUM_GPUS_PER_NODE \
            --master_addr=$MASTER_ADDR \
            --master_port=$MASTER_PORT \
            /root/ColossalAI/examples/language/llama/benchmark.py -c 7b -g -p 3d -x -b 64 -s 10 --pp 2 --sp 2 --tp 1 --sp_mode "all_to_all" --zero 1 -l 32768
    else
        # 如果不是主节点，通过SSH远程执行
        ssh root@$ip "
            export PATH=/root/.local/bin:\$PATH
            export OMP_NUM_THREADS=8
            torchrun \
                --nnodes=$NNODES \
                --node_rank=$rank \
                --nproc_per_node=$NUM_GPUS_PER_NODE \
                --master_addr=$MASTER_ADDR \
                --master_port=$MASTER_PORT \
                /root/ColossalAI/examples/language/llama/benchmark.py -c 7b -g -p 3d -x -b 64 -s 10 --pp 2 --sp 2 --tp 1 --sp_mode \"all_to_all\" --zero 1 -l 32768
        "
    fi
}

• 将函数导出，以便在子 shell 中调用。

export -f start_training_on_node

• 获取节点排名（rank），主节点的 rank 为 0。

ranks=$(seq 0 $((NNODES - 1)))

• 启动训练。

combined=$(paste -d ' ' /root/ColossalAI/examples/language/llama/node_ip_list.txt <(echo "$ranks"))
echo "$combined" | xargs -P $NNODES -n 2 bash -c 'start_training_on_node "$1" "$2"' _

5. 运行脚本

确保所有节点之间可以通过 SSH 无密码访问。然后，在主节点上执行脚本：

#!/bin/bash
export MASTER_ADDR="10.42.128.230"
export MASTER_PORT="12345"
export NUM_GPUS_PER_NODE=8
# export NCCL_DEBUG=INFO
export NNODES=$(grep -c '^' /root/ColossalAI/examples/language/llama/node_ip_list.txt)
start_training_on_node() {
    local ip=$1
    local rank=$2
    if [ "$ip" == "$MASTER_ADDR" ]; then
        torchrun \
            --nnodes=$NNODES \
            --node_rank=$rank \
            --nproc_per_node=$NUM_GPUS_PER_NODE \
            --master_addr=$MASTER_ADDR \
            --master_port=$MASTER_PORT \
            /root/ColossalAI/examples/language/llama/benchmark.py -c 7b -g -p 3d -x -b 64 -s 10 --pp 2 --sp 2 --tp 1 --sp_mode "all_to_all" --zero 1 -l 32768
    else
        ssh root@$ip "
            export PATH=/root/.local/bin:\$PATH
            export OMP_NUM_THREADS=8
            torchrun \
                --nnodes=$NNODES \
                --node_rank=$rank \
                --nproc_per_node=$NUM_GPUS_PER_NODE \
                --master_addr=$MASTER_ADDR \
                --master_port=$MASTER_PORT \
                /root/ColossalAI/examples/language/llama/benchmark.py -c 7b -g -p 3d -x -b 64 -s 10 --pp 2 --sp 2 --tp 1 --sp_mode \"all_to_all\" --zero 1 -l 32768
        "
    fi
}

export -f start_training_on_node
ranks=$(seq 0 $((NNODES - 1)))
combined=$(paste -d ' ' /root/ColossalAI/examples/language/llama/node_ip_list.txt <(echo "$ranks"))
echo "$combined" | xargs -P $NNODES -n 2 bash -c 'start_training_on_node "$1" "$2"' _

以上即为多机训练的完整流程及配置。希望这能帮助您顺利使用 colossalai 进行分布式训练。