# 多机多卡训练说明

本文档说明当前仓库如何从单机 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` 的并行方式是纯数据并行：

```bash
--tensor-model-parallel-size 1
--pipeline-model-parallel-size 1
```

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

## 2. 多机多卡训练原理

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

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

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

### 2.1 单卡训练

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

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

没有跨 GPU 通信，最简单。

### 2.2 单机多卡训练

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

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

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

all-reduce 可以理解为：

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

### 2.3 多机多卡训练

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

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

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

区别是：

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

## 3. torchrun 的几个核心参数

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

```bash
--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`：主节点监听端口。四台机器都要能访问这个端口。

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

```text
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` 参数里有：

```bash
--micro-batch-size 16
--global-batch-size 2048
```

Megatron 里的关系是：

```text
global_batch_size = micro_batch_size * data_parallel_size * gradient_accumulation_steps
```

单机 8 卡时：

```text
2048 = 16 * 8 * 16
gradient_accumulation_steps = 16
```

四机 32 卡时：

```text
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 环境变量：

```bash
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。排查问题时可以临时用，但性能通常差很多。

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

```bash
ip addr
ibdev2netdev
```

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

```bash
NCCL_DEBUG=INFO
```

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

## 6. 本次脚本修改

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

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

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

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

```bash
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

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

```bash
g0033 g0034 g0035 g0036
```

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

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

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

日志会写到：

```text
/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
```

运行状态会写到：

```text
/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
bash stop_multinode_training.sh <train_name>
```

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

### 6.4 修正 stop_training.sh 的默认 ARTIFACT_ROOT

原来的 `start_training.sh` 默认是：

```bash
/apps/yi/model_training/artifacts
```

但 `stop_training.sh` 默认是：

```bash
/ssd1/yi/artifacts
```

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

```bash
/apps/yi/model_training/artifacts
```

### 6.5 新增 ZERO_STAGE 开关

训练脚本现在支持：

```bash
ZERO_STAGE=0
ZERO_STAGE=1
ZERO_STAGE=2
```

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

## 7. 如何启动

进入训练脚本目录：

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

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

```bash
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
```

检查容器网络模式：

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

应该输出：

```text
host
```

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

```bash
CONTAINER_NAME=megatron-ngc25-training
```

### 7.1 单机 8 卡，原方式不变

```bash
bash start_training.sh qwen3_1p7b qwen3_1p7b_smoke_yi qwen3_single_node
```

停止：

```bash
bash stop_training.sh qwen3_single_node
```

### 7.2 四机 32 卡，默认 zero0

```bash
bash start_multinode_training.sh \
  qwen3_1p7b \
  qwen3_1p7b_smoke_yi \
  qwen3_32gpu_z0
```

如果训练环境在容器里：

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

停止：

```bash
bash stop_multinode_training.sh qwen3_32gpu_z0
```

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

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

### 7.3 四机 32 卡，指定 NCCL 网卡

示例：

```bash
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，可能是：

```bash
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

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

### 7.5 四机 32 卡，ZeRO-2

```bash
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 不是同一套实现。但概念上可以对应理解。

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

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

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

### 8.1 ZeRO-0

每张卡都有完整的：

- 参数
- 梯度
- optimizer state

优点：

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

缺点：

- 最吃显存。

本仓库中：

```bash
ZERO_STAGE=0
```

不会额外加 distributed optimizer 参数。

### 8.2 ZeRO-1

切分 optimizer state。

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

优点：

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

本仓库中：

```bash
ZERO_STAGE=1
```

会加：

```bash
--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 更快。

本仓库中：

```bash
ZERO_STAGE=2
```

会加：

```bash
--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
bash start_training.sh gpt_smoke smoke smoke_single
```

2. 四机小模型 smoke test。

```bash
NCCL_DEBUG=INFO \
bash start_multinode_training.sh gpt_smoke smoke smoke_32gpu
```

3. 四机 Qwen3 smoke test，zero0。

```bash
NCCL_DEBUG=INFO \
bash start_multinode_training.sh qwen3_1p7b qwen3_1p7b_smoke_yi qwen3_32gpu_z0
```

4. 对比 ZeRO-1。

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

5. 对比 ZeRO-2。

```bash
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` 重复或缺失。
- 防火墙阻止通信。

排查：

```bash
ssh g0034 hostname
ssh g0035 hostname
ssh g0036 hostname
```

在 g0034-g0036 上测试：

```bash
nc -vz g0033 6000
```

### 10.2 NCCL 选错网卡

现象：

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

解决：

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

### 10.3 数据路径不存在

四台机器都需要能访问：

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

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

### 10.4 checkpoint 路径问题

当前 checkpoint 路径是：

```text
/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 版本不支持这些参数，需要在目标机上用：

```bash
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`，所以老命令不受影响。

最终目标是：

```text
原来：
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 的训练任务
```