pretrain_kaiyuan2b/scripts/kaiyuan2b-training/MULTINODE_TRAINING.md

多机多卡训练说明

本文档说明当前仓库如何从单机 8 卡训练扩展到 g0033-g0036 四机 32 卡训练，并解释相关原理、通信配置、ZeRO 模式和本次脚本改动。

1. 当前仓库结构

本仓库的训练主体是 Megatron-LM。外层脚本负责拼接模型参数、数据参数、优化器参数，然后通过 torchrun 启动 Megatron 的 pretrain_gpt.py。

关键文件：

• training_smoke_qwen3_1p7b.sh：Qwen3 1.7B 训练入口。
• training_smoke_gpt2.sh：小模型 smoke test 训练入口。
• start_training.sh：原有单机后台启动脚本。
• stop_training.sh：原有单机停止脚本。
• start_multinode_training.sh：新增四机启动脚本。
• stop_multinode_training.sh：新增四机停止脚本。
• params/qwen3_1p7b/*.sh：模型结构、数据路径、训练超参数。

当前 qwen3_1p7b 的并行方式是纯数据并行：

--tensor-model-parallel-size 1
--pipeline-model-parallel-size 1

也就是说，每张 GPU 都持有完整模型，拿不同数据算梯度，然后所有 GPU 同步梯度。

2. 多机多卡训练原理

先用一个很土的比喻：32 张 GPU 就像 32 个学生一起做同一道题。

每个学生拿到一小份不同的数据，独立算出自己的答案，也就是梯度。算完之后，大家把答案平均一下。平均后的答案就是这一步真正用来更新模型的梯度。因为每个人都使用同一个平均梯度更新，所以 32 张 GPU 上的模型参数仍然保持一致。

这个过程叫数据并行，英文是 data parallel。

2.1 单卡训练

单卡训练时只有一个进程：

数据 -> 模型 -> loss -> backward -> optimizer step

没有跨 GPU 通信，最简单。

2.2 单机多卡训练

单机 8 卡训练时，一般是一张 GPU 对应一个训练进程：

GPU0: batch shard 0 -> gradient 0
GPU1: batch shard 1 -> gradient 1
...
GPU7: batch shard 7 -> gradient 7

反向传播后，8 张卡做一次梯度同步。常见通信方式是 NCCL 的 all-reduce。

all-reduce 可以理解为：

每张卡都有一个梯度
大家把梯度加起来
再除以 GPU 数量
最后每张卡都拿到同一个平均梯度

2.3 多机多卡训练

四机 32 卡时，逻辑没有变，只是通信范围从一台机器里的 8 张卡扩展到四台机器里的 32 张卡：

g0033: GPU0-GPU7
g0034: GPU0-GPU7
g0035: GPU0-GPU7
g0036: GPU0-GPU7

每张 GPU 还是一个训练进程，总共 32 个进程。

区别是：

• 同一台机器内部通信通常走 NVLink/NVSwitch。
• 不同机器之间通信走网卡，比如 IB、RoCE 或以太网。
• torchrun 负责把 32 个进程组成一个训练集群。
• NCCL 负责 GPU 之间真正的数据通信。

3. torchrun 的几个核心参数

多机训练最关键的是这几个参数：

--nproc_per_node 8
--nnodes 4
--node_rank 0
--master_addr g0033
--master_port 6000

含义：

• nproc_per_node：每台机器启动几个训练进程。这里每台 8 张 H200，所以是 8。
• nnodes：总共有几台机器。这里是 4。
• node_rank：当前机器编号。g0033 是 0，g0034 是 1，g0035 是 2，g0036 是 3。
• master_addr：主节点地址。这里用 g0033。
• master_port：主节点监听端口。四台机器都要能访问这个端口。

四台机器上的参数应该是：

g0033: --node_rank 0 --master_addr g0033 --master_port 6000
g0034: --node_rank 1 --master_addr g0033 --master_port 6000
g0035: --node_rank 2 --master_addr g0033 --master_port 6000
g0036: --node_rank 3 --master_addr g0033 --master_port 6000

torchrun 会自动给每个进程设置这些环境变量：

• RANK：全局进程编号，范围 0-31。
• LOCAL_RANK：本机进程编号，范围 0-7。
• WORLD_SIZE：总进程数，这里是 32。
• LOCAL_WORLD_SIZE：本机进程数，这里是 8。

Megatron-LM 会读取这些信息，然后建立 data parallel group。

4. batch size 怎么变

当前 qwen3_1p7b 参数里有：

--micro-batch-size 16
--global-batch-size 2048

Megatron 里的关系是：

global_batch_size = micro_batch_size * data_parallel_size * gradient_accumulation_steps

单机 8 卡时：

2048 = 16 * 8 * 16
gradient_accumulation_steps = 16

四机 32 卡时：

2048 = 16 * 32 * 4
gradient_accumulation_steps = 4

所以扩到 32 卡后，如果 global batch 不变，每个 global step 需要累积的次数会减少。通常这会提升吞吐，因为更多 GPU 同时处理数据。

如果想进一步提高吞吐，也可以尝试增大 global-batch-size，但这会改变训练动力学，需要关注 loss 曲线、学习率设置和收敛行为。

5. 网络通信配置

多机训练最常见的问题不是模型代码，而是网络通信。

需要确认：

• g0033 能 ssh 到 g0034-g0036。
• g0034-g0036 能访问 g0033 的 MASTER_PORT，默认是 6000。
• 四台机器能互相解析主机名，至少都能解析 g0033。
• 四台机器上的代码路径一致。
• 四台机器上的数据路径一致。
• 四台机器上的 Python、CUDA、NCCL、Transformer Engine、Megatron 版本一致。

常用 NCCL 环境变量：

NCCL_DEBUG=INFO
NCCL_SOCKET_IFNAME=eth0
GLOO_SOCKET_IFNAME=eth0
NCCL_IB_HCA=mlx5_0,mlx5_1
NCCL_IB_DISABLE=0

解释：

• NCCL_DEBUG=INFO：打印 NCCL 日志，排查多机通信问题很有用。
• NCCL_SOCKET_IFNAME：指定 NCCL 用哪块网卡做 socket 通信。
• GLOO_SOCKET_IFNAME：指定 Gloo 用哪块网卡。
• NCCL_IB_HCA：指定使用哪些 IB/RoCE 设备。
• NCCL_IB_DISABLE=1：禁用 IB，只走 TCP。排查问题时可以临时用，但性能通常差很多。

具体网卡名需要在目标机器上查看：

ip addr
ibdev2netdev

如果不知道集群网卡，先用：

NCCL_DEBUG=INFO

跑 smoke test，看 NCCL 日志里选了什么网卡。

6. 本次脚本修改

6.1 训练脚本不再写死单机参数

原来 training_smoke_qwen3_1p7b.sh 和 training_smoke_gpt2.sh 写死为：

--nproc_per_node 8
--nnodes 1
--node_rank 0
--master_addr localhost
--master_port 6000

现在改成从环境变量读取：

NPROC_PER_NODE=${NPROC_PER_NODE:-8}
NNODES=${NNODES:-1}
NODE_RANK=${NODE_RANK:-0}
MASTER_ADDR=${MASTER_ADDR:-localhost}
MASTER_PORT=${MASTER_PORT:-6000}

也就是说，默认仍然是单机 8 卡，不影响原来的启动方式。但多机脚本可以给每台机器传不同的 NODE_RANK。

6.2 新增 start_multinode_training.sh

这个脚本默认使用四台机器：

g0033 g0034 g0035 g0036

它会在 g0033 上执行，然后通过 ssh 到其他机器启动训练。

每台机器会启动同一个训练脚本，但传入不同的环境变量：

NNODES=4
NPROC_PER_NODE=8
NODE_RANK=0/1/2/3
MASTER_ADDR=g0033
MASTER_PORT=6000

日志会写到：

/apps/yi/model_training/artifacts/logs/<train_name>_node0.log
/apps/yi/model_training/artifacts/logs/<train_name>_node1.log
/apps/yi/model_training/artifacts/logs/<train_name>_node2.log
/apps/yi/model_training/artifacts/logs/<train_name>_node3.log

运行状态会写到：

/apps/yi/model_training/artifacts/run_state/<train_name>_node0.pid
/apps/yi/model_training/artifacts/run_state/<train_name>_node1.pid
/apps/yi/model_training/artifacts/run_state/<train_name>_node2.pid
/apps/yi/model_training/artifacts/run_state/<train_name>_node3.pid

6.3 新增 stop_multinode_training.sh

停止多机任务时，不应该只杀 g0033 上的进程，还要杀 g0034-g0036 上的进程。

所以新增：

bash stop_multinode_training.sh <train_name>

它会逐台 ssh 上去调用原有的 stop_training.sh。

6.4 修正 stop_training.sh 的默认 ARTIFACT_ROOT

原来的 start_training.sh 默认是：

/apps/yi/model_training/artifacts

但 stop_training.sh 默认是：

/ssd1/yi/artifacts

这会导致默认情况下 stop 找不到 pid 文件。本次已改成一致：

/apps/yi/model_training/artifacts

6.5 新增 ZERO_STAGE 开关

训练脚本现在支持：

ZERO_STAGE=0
ZERO_STAGE=1
ZERO_STAGE=2

默认是 ZERO_STAGE=0，保持原行为。

7. 如何启动

进入训练脚本目录：

cd /apps/yi/model_training/scripts/kaiyuan2b-training

如果训练环境在 Docker 容器内，四台机器需要先启动同名容器，并且使用 host network：

docker run -dit \
  --gpus all \
  --network=host \
  --ipc=host \
  --shm-size=64g \
  -v /ssd1/yi:/ssd1/yi \
  -w /ssd1/yi/pretrain_kaiyuan2b \
  --name megatron-ngc25-training \
  base-mirror.tencentcloudcr.com/mode-optimization/megatron-env:ngc25.10 \
  bash

检查容器网络模式：

docker inspect megatron-ngc25-training --format '{{.HostConfig.NetworkMode}}'

应该输出：

host

如果要从 host 上一条命令启动四台机器容器内训练，给多机脚本加：

CONTAINER_NAME=megatron-ngc25-training

7.1 单机 8 卡，原方式不变

bash start_training.sh qwen3_1p7b qwen3_1p7b_smoke_yi qwen3_single_node

停止：

bash stop_training.sh qwen3_single_node

7.2 四机 32 卡，默认 zero0

bash start_multinode_training.sh \
  qwen3_1p7b \
  qwen3_1p7b_smoke_yi \
  qwen3_32gpu_z0

如果训练环境在容器里：

CONTAINER_NAME=megatron-ngc25-training \
bash start_multinode_training.sh \
  qwen3_1p7b \
  qwen3_1p7b_smoke_yi \
  qwen3_32gpu_z0

停止：

bash stop_multinode_training.sh qwen3_32gpu_z0

容器内训练对应的停止方式：

CONTAINER_NAME=megatron-ngc25-training \
bash stop_multinode_training.sh qwen3_32gpu_z0

7.3 四机 32 卡，指定 NCCL 网卡

示例：

NCCL_DEBUG=INFO \
NCCL_SOCKET_IFNAME=eth0 \
GLOO_SOCKET_IFNAME=eth0 \
bash start_multinode_training.sh \
  qwen3_1p7b \
  qwen3_1p7b_smoke_yi \
  qwen3_32gpu_eth0

如果集群有 IB/RoCE，可能是：

NCCL_DEBUG=INFO \
NCCL_SOCKET_IFNAME=bond0 \
GLOO_SOCKET_IFNAME=bond0 \
NCCL_IB_HCA=mlx5_0,mlx5_1 \
bash start_multinode_training.sh \
  qwen3_1p7b \
  qwen3_1p7b_smoke_yi \
  qwen3_32gpu_ib

网卡名需要按实际机器修改。

7.4 四机 32 卡，ZeRO-1

ZERO_STAGE=1 \
bash start_multinode_training.sh \
  qwen3_1p7b \
  qwen3_1p7b_smoke_yi \
  qwen3_32gpu_z1

7.5 四机 32 卡，ZeRO-2

ZERO_STAGE=2 \
bash start_multinode_training.sh \
  qwen3_1p7b \
  qwen3_1p7b_smoke_yi \
  qwen3_32gpu_z2

8. ZeRO 模式回顾

严格说，DeepSpeed ZeRO 和 Megatron-LM distributed optimizer 不是同一套实现。但概念上可以对应理解。

训练中主要有三类大对象：

参数 parameter
梯度 gradient
优化器状态 optimizer state

以 Adam 为例，optimizer state 通常包括一阶动量和二阶动量。它们会占不少显存。

8.1 ZeRO-0

每张卡都有完整的：

• 参数
• 梯度
• optimizer state

优点：

• 最简单。
• 通信模式最常规。
• 如果显存够，通常吞吐最好。

缺点：

• 最吃显存。

本仓库中：

ZERO_STAGE=0

不会额外加 distributed optimizer 参数。

8.2 ZeRO-1

切分 optimizer state。

每张卡仍有完整参数和梯度，但 optimizer state 在 data parallel 组内分片存储。

优点：

• 明显节省 optimizer state 显存。
• 通信开销相对 ZeRO-2/3 小。
• 很适合显存有压力但还没到必须切梯度/参数的情况。

本仓库中：

ZERO_STAGE=1

会加：

--use-distributed-optimizer
--data-parallel-sharding-strategy optim

8.3 ZeRO-2

切分 optimizer state 和 gradient。

每张卡仍有完整参数，但 optimizer state 和梯度会分片。

优点：

• 比 ZeRO-1 更省显存。
• 可以支持更大的 batch 或更大的模型。

缺点：

• 通信更复杂。
• 如果模型本来就能轻松放下，ZeRO-2 不一定比 ZeRO-0/1 更快。

本仓库中：

ZERO_STAGE=2

会加：

--use-distributed-optimizer
--data-parallel-sharding-strategy optim_grads

8.4 ZeRO-3

切分参数、梯度、optimizer state。

优点：

• 最省显存。

缺点：

• 通信开销最大。
• 配置和 checkpoint 复杂度更高。

本次没有加 ZeRO-3，因为当前目标是兼容 zero0/1/2，而且 1.7B 模型在 32 张 H200 上通常不需要 ZeRO-3 来解决显存问题。

9. 推荐实验顺序

建议不要一上来就跑正式训练。按下面顺序验证：

1. 单机 smoke test。

bash start_training.sh gpt_smoke smoke smoke_single

2. 四机小模型 smoke test。

NCCL_DEBUG=INFO \
bash start_multinode_training.sh gpt_smoke smoke smoke_32gpu

3. 四机 Qwen3 smoke test，zero0。

NCCL_DEBUG=INFO \
bash start_multinode_training.sh qwen3_1p7b qwen3_1p7b_smoke_yi qwen3_32gpu_z0

4. 对比 ZeRO-1。

NCCL_DEBUG=INFO \
ZERO_STAGE=1 \
bash start_multinode_training.sh qwen3_1p7b qwen3_1p7b_smoke_yi qwen3_32gpu_z1

5. 对比 ZeRO-2。

NCCL_DEBUG=INFO \
ZERO_STAGE=2 \
bash start_multinode_training.sh qwen3_1p7b qwen3_1p7b_smoke_yi qwen3_32gpu_z2

主要观察：

• 是否能正常建立 NCCL 通信。
• 日志中的 world size 是否是 32。
• step time 是否稳定。
• tokens/s 或 samples/s 是否提升。
• GPU util 是否接近满载。
• 是否出现 NCCL timeout。
• checkpoint 是否能正常保存和加载。

10. 常见问题

10.1 卡在初始化

常见原因：

• MASTER_ADDR 解析不到。
• MASTER_PORT 不通。
• 某台机器没有成功启动。
• node_rank 重复或缺失。
• 防火墙阻止通信。

排查：

ssh g0034 hostname
ssh g0035 hostname
ssh g0036 hostname

在 g0034-g0036 上测试：

nc -vz g0033 6000

10.2 NCCL 选错网卡

现象：

• 初始化很慢。
• NCCL timeout。
• 跨机吞吐极低。

解决：

NCCL_DEBUG=INFO
NCCL_SOCKET_IFNAME=<正确网卡>
GLOO_SOCKET_IFNAME=<正确网卡>

10.3 数据路径不存在

四台机器都需要能访问：

/ssd/yi/converted_data/megatron_phase1
/apps/yi/model_training/data/tokenizer

如果这些不是共享存储，就要保证每台机器本地都有同样路径和文件。

10.4 checkpoint 路径问题

当前 checkpoint 路径是：

/apps/yi/model_training/artifacts/checkpoints/<train_name>

多机训练最好使用共享存储。否则不同节点各写各的 checkpoint，恢复时很容易出问题。

10.5 zero1/zero2 启动失败

可能原因：

• 当前部署的 Megatron-LM 版本不支持 --data-parallel-sharding-strategy optim 或 optim_grads。
• Megatron-LM 子模块版本和脚本预期不一致。

当前本地仓库里的 Megatron-LM 子模块目录是空的，所以我无法在本地直接检查目标集群上的 Megatron 参数解析。若目标机器上的 Megatron 版本不支持这些参数，需要在目标机上用：

python /apps/yi/model_training/Megatron-LM/pretrain_gpt.py --help | grep -E "distributed-optimizer|data-parallel-sharding"

如果参数名不同，需要按目标 Megatron 版本调整。

11. 为什么这样改

这次改法的原则是：尽量小改动，保留原来的单机行为，把多机能力做成环境变量和独立编排脚本。

具体原因：

• 原训练脚本逻辑已经能训练，不需要重写。
• torchrun 原生支持多机，只需要把 nnodes/node_rank/master_addr 参数暴露出来。
• 单机和多机共用同一个训练脚本，避免两套参数漂移。
• start_multinode_training.sh 只负责编排，不改 Megatron 训练逻辑。
• ZERO_STAGE 用环境变量控制，方便跑吞吐对比实验。
• 默认 ZERO_STAGE=0，默认 NNODES=1，所以老命令不受影响。

最终目标是：

原来：
g0033 上 8 个 torchrun worker

现在：
g0033 上 8 个 torchrun worker
g0034 上 8 个 torchrun worker
g0035 上 8 个 torchrun worker
g0036 上 8 个 torchrun worker

总共 32 个 worker 组成一个 WORLD_SIZE=32 的训练任务