feat: incorporate multi-device training scripts and README

2026-05-09 21:35:42 +08:00
parent 02868ec01a
commit 75eacf00c2
6 changed files with 1082 additions and 10 deletions
--- a/scripts/kaiyuan2b-training/MULTINODE_TRAINING.md
+++ b/scripts/kaiyuan2b-training/MULTINODE_TRAINING.md
@@ -0,0 +1,710 @@
+# 多机多卡训练说明
+
+本文档说明当前仓库如何从单机 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 的训练任务
+```
--- a/scripts/kaiyuan2b-training/start_multinode_training.sh
+++ b/scripts/kaiyuan2b-training/start_multinode_training.sh
@@ -0,0 +1,163 @@
+#!/usr/bin/env bash
+set -euo pipefail
+
+SCRIPT_DIR=$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)
+ARTIFACT_ROOT=${ARTIFACT_ROOT:-/apps/yi/model_training/artifacts}
+RUN_STATE_DIR="${ARTIFACT_ROOT}/run_state"
+LOG_DIR="${ARTIFACT_ROOT}/logs"
+
+usage() {
+  cat <<'EOF'
+Usage:
+  bash start_multinode_training.sh <model> [mode] [train_name]
+
+Default cluster:
+  g0033,g0034,g0035,g0036 with 8 GPUs per host.
+
+Environment overrides:
+  HOSTS="g0033 g0034 g0035 g0036"
+  MASTER_ADDR=g0033
+  MASTER_PORT=6000
+  NPROC_PER_NODE=8
+  ZERO_STAGE=0|1|2
+  CONTAINER_NAME=megatron-ngc25-training
+  NCCL_SOCKET_IFNAME=eth0
+  GLOO_SOCKET_IFNAME=eth0
+  NCCL_IB_HCA=mlx5_0,mlx5_1
+  NCCL_DEBUG=INFO
+  CHECKPOINT_KEEP_RECENT=3
+  CHECKPOINT_CLEANUP_INTERVAL_SECONDS=300
+  EXTRA_ARGS="--exit-duration-in-mins 120"
+
+Examples:
+  bash start_multinode_training.sh qwen3_1p7b qwen3_1p7b_smoke_yi qwen3_32gpu
+  CONTAINER_NAME=megatron-ngc25-training bash start_multinode_training.sh qwen3_1p7b qwen3_1p7b_smoke_yi qwen3_32gpu
+  ZERO_STAGE=1 bash start_multinode_training.sh qwen3_1p7b phase1 qwen3_phase1_zero1
+EOF
+}
+
+model=${1:-}
+mode=${2:-}
+train_name=${3:-}
+
+if [ -z "$model" ] || [ "$model" = "-h" ] || [ "$model" = "--help" ]; then
+  usage
+  exit 0
+fi
+
+case "$model" in
+  gpt_smoke)
+    train_script="${SCRIPT_DIR}/training_smoke_gpt2.sh"
+    mode=${mode:-smoke}
+    train_name=${train_name:-smoke_gpt_multinode}
+    ;;
+  qwen3_1p7b)
+    train_script="${SCRIPT_DIR}/training_smoke_qwen3_1p7b.sh"
+    mode=${mode:-qwen3_1p7b_smoke_yi}
+    train_name=${train_name:-qwen3_1p7b_multinode}
+    ;;
+  *)
+    echo "Unknown model: $model" >&2
+    usage >&2
+    exit 1
+    ;;
+esac
+
+read -r -a HOST_ARRAY <<< "${HOSTS:-g0033 g0034 g0035 g0036}"
+NNODES=${NNODES:-${#HOST_ARRAY[@]}}
+NPROC_PER_NODE=${NPROC_PER_NODE:-8}
+MASTER_ADDR=${MASTER_ADDR:-${HOST_ARRAY[0]}}
+MASTER_PORT=${MASTER_PORT:-6000}
+ZERO_STAGE=${ZERO_STAGE:-0}
+CONTAINER_NAME=${CONTAINER_NAME:-}
+NCCL_SOCKET_IFNAME=${NCCL_SOCKET_IFNAME:-}
+GLOO_SOCKET_IFNAME=${GLOO_SOCKET_IFNAME:-}
+NCCL_IB_HCA=${NCCL_IB_HCA:-}
+NCCL_DEBUG=${NCCL_DEBUG:-}
+NCCL_IB_DISABLE=${NCCL_IB_DISABLE:-}
+CHECKPOINT_KEEP_RECENT=${CHECKPOINT_KEEP_RECENT:-3}
+CHECKPOINT_CLEANUP_INTERVAL_SECONDS=${CHECKPOINT_CLEANUP_INTERVAL_SECONDS:-300}
+EXTRA_ARGS="--exit-signal-handler ${EXTRA_ARGS:-}"
+
+mkdir -p "$RUN_STATE_DIR" "$LOG_DIR"
+
+echo "Starting multinode training: model=${model}, mode=${mode}, train_name=${train_name}"
+echo "Hosts: ${HOST_ARRAY[*]}"
+echo "Distributed: nnodes=${NNODES}, nproc_per_node=${NPROC_PER_NODE}, master=${MASTER_ADDR}:${MASTER_PORT}, zero_stage=${ZERO_STAGE}"
+if [ -n "$CONTAINER_NAME" ]; then
+  echo "Container: ${CONTAINER_NAME}"
+fi
+
+for idx in "${!HOST_ARRAY[@]}"; do
+  host=${HOST_ARRAY[$idx]}
+  node_rank=$idx
+  node_train_name="${train_name}_node${node_rank}"
+  pid_file="${RUN_STATE_DIR}/${node_train_name}.pid"
+  meta_file="${RUN_STATE_DIR}/${node_train_name}.env"
+  log_file="${LOG_DIR}/${node_train_name}.log"
+
+  remote_cmd=$(cat <<EOF
+set -euo pipefail
+run_cmd=\$(cat <<'RUN_CMD'
+mkdir -p "$RUN_STATE_DIR" "$LOG_DIR"
+cd "$SCRIPT_DIR"
+ARTIFACT_ROOT="$ARTIFACT_ROOT" \\
+CHECKPOINT_KEEP_RECENT="$CHECKPOINT_KEEP_RECENT" \\
+CHECKPOINT_CLEANUP_INTERVAL_SECONDS="$CHECKPOINT_CLEANUP_INTERVAL_SECONDS" \\
+NPROC_PER_NODE="$NPROC_PER_NODE" \\
+NNODES="$NNODES" \\
+NODE_RANK="$node_rank" \\
+MASTER_ADDR="$MASTER_ADDR" \\
+MASTER_PORT="$MASTER_PORT" \\
+ZERO_STAGE="$ZERO_STAGE" \\
+NCCL_SOCKET_IFNAME="$NCCL_SOCKET_IFNAME" \\
+GLOO_SOCKET_IFNAME="$GLOO_SOCKET_IFNAME" \\
+NCCL_IB_HCA="$NCCL_IB_HCA" \\
+NCCL_DEBUG="$NCCL_DEBUG" \\
+NCCL_IB_DISABLE="$NCCL_IB_DISABLE" \\
+EXTRA_ARGS="$EXTRA_ARGS" \\
+setsid bash "$train_script" "$mode" "$train_name" > "$log_file" 2>&1 < /dev/null &
+pid=\$!
+pgid=\$(ps -o pgid= -p "\$pid" | tr -d ' ' || true)
+printf '%s\n' "\$pid" > "$pid_file"
+cat > "$meta_file" <<META
+MODEL=$model
+MODE=$mode
+TRAIN_NAME=$train_name
+NODE_TRAIN_NAME=$node_train_name
+HOST=$host
+NODE_RANK=$node_rank
+NNODES=$NNODES
+NPROC_PER_NODE=$NPROC_PER_NODE
+MASTER_ADDR=$MASTER_ADDR
+MASTER_PORT=$MASTER_PORT
+ZERO_STAGE=$ZERO_STAGE
+NCCL_SOCKET_IFNAME=$NCCL_SOCKET_IFNAME
+GLOO_SOCKET_IFNAME=$GLOO_SOCKET_IFNAME
+NCCL_IB_HCA=$NCCL_IB_HCA
+NCCL_DEBUG=$NCCL_DEBUG
+NCCL_IB_DISABLE=$NCCL_IB_DISABLE
+PID=\$pid
+PGID=\$pgid
+LOG_FILE=$log_file
+TRAIN_SCRIPT=$train_script
+META
+echo "host=$host node_rank=$node_rank pid=\$pid pgid=\${pgid:-unknown} log=$log_file"
+RUN_CMD
+)
+if [ -n "$CONTAINER_NAME" ]; then
+  docker exec "$CONTAINER_NAME" bash -lc "\$run_cmd"
+else
+  bash -lc "\$run_cmd"
+fi
+EOF
+)
+
+  if [ "$host" = "$(hostname -s)" ] || [ "$host" = "$(hostname)" ]; then
+    bash -lc "$remote_cmd"
+  else
+    ssh "$host" "bash -lc $(printf '%q' "$remote_cmd")"
+  fi
+done
+
+echo "Launched ${NNODES} nodes. Stop all nodes with: bash ${SCRIPT_DIR}/stop_multinode_training.sh ${train_name}"
--- a/scripts/kaiyuan2b-training/stop_multinode_training.sh
+++ b/scripts/kaiyuan2b-training/stop_multinode_training.sh
@@ -0,0 +1,44 @@
+#!/usr/bin/env bash
+set -euo pipefail
+
+SCRIPT_DIR=$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)
+
+usage() {
+  cat <<'EOF'
+Usage:
+  bash stop_multinode_training.sh <train_name>
+
+Environment overrides:
+  HOSTS="g0033 g0034 g0035 g0036"
+  CONTAINER_NAME=megatron-ngc25-training
+  ARTIFACT_ROOT=/apps/yi/model_training/artifacts
+  GRACE_SECONDS=600
+EOF
+}
+
+train_name=${1:-}
+if [ -z "$train_name" ] || [ "$train_name" = "-h" ] || [ "$train_name" = "--help" ]; then
+  usage
+  exit 0
+fi
+
+read -r -a HOST_ARRAY <<< "${HOSTS:-g0033 g0034 g0035 g0036}"
+CONTAINER_NAME=${CONTAINER_NAME:-}
+
+for idx in "${!HOST_ARRAY[@]}"; do
+  host=${HOST_ARRAY[$idx]}
+  node_train_name="${train_name}_node${idx}"
+  stop_cmd="cd \"$SCRIPT_DIR\" && ARTIFACT_ROOT=\"${ARTIFACT_ROOT:-/apps/yi/model_training/artifacts}\" GRACE_SECONDS=\"${GRACE_SECONDS:-600}\" bash stop_training.sh \"$node_train_name\""
+  if [ -n "$CONTAINER_NAME" ]; then
+    remote_cmd="docker exec \"$CONTAINER_NAME\" bash -lc $(printf '%q' "$stop_cmd")"
+  else
+    remote_cmd="$stop_cmd"
+  fi
+
+  echo "Stopping host=${host}, train_name=${node_train_name}"
+  if [ "$host" = "$(hostname -s)" ] || [ "$host" = "$(hostname)" ]; then
+    bash -lc "$remote_cmd" || true
+  else
+    ssh "$host" "bash -lc $(printf '%q' "$remote_cmd")" || true
+  fi
+done
--- a/scripts/kaiyuan2b-training/training_smoke_gpt2.sh
+++ b/scripts/kaiyuan2b-training/training_smoke_gpt2.sh
@@ -14,6 +14,12 @@ CKPT_DIR="${ARTIFACT_ROOT}/checkpoints/${TRAIN_NAME}"
 CHECKPOINT_KEEP_RECENT=${CHECKPOINT_KEEP_RECENT:-3}
 CHECKPOINT_CLEANUP_INTERVAL_SECONDS=${CHECKPOINT_CLEANUP_INTERVAL_SECONDS:-300}
 EXTRA_ARGS=${EXTRA_ARGS:-}
+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}
+ZERO_STAGE=${ZERO_STAGE:-0}

 source params/optim_common.sh
 source params/gpt_smoke/model.sh
@@ -46,6 +52,28 @@ PARALLEL_ARGS="
 #     --sequence-parallel
 # "

+case "$ZERO_STAGE" in
+  0)
+    ZERO_ARGS=""
+    ;;
+  1)
+    ZERO_ARGS="
+      --use-distributed-optimizer
+      --data-parallel-sharding-strategy optim
+    "
+    ;;
+  2)
+    ZERO_ARGS="
+      --use-distributed-optimizer
+      --data-parallel-sharding-strategy optim_grads
+    "
+    ;;
+  *)
+    echo "Unsupported ZERO_STAGE=${ZERO_STAGE}; expected 0, 1, or 2" >&2
+    exit 1
+    ;;
+esac
+
 mkdir -p "$CKPT_DIR" "$TB_DIR"

 cleanup_old_checkpoints_once() {
@@ -107,17 +135,18 @@ CHECKPOINT_CLEANUP_PID=$!
 trap 'kill "$CHECKPOINT_CLEANUP_PID" 2>/dev/null || true; cleanup_old_checkpoints_once "$CKPT_DIR" "$CHECKPOINT_KEEP_RECENT"' EXIT

 DISTRIBUTED_ARGS="
-  --nproc_per_node 8
-  --nnodes 1
-  --node_rank 0
-  --master_addr localhost
-  --master_port 6000
+  --nproc_per_node ${NPROC_PER_NODE}
+  --nnodes ${NNODES}
+  --node_rank ${NODE_RANK}
+  --master_addr ${MASTER_ADDR}
+  --master_port ${MASTER_PORT}
 "

 torchrun $DISTRIBUTED_ARGS \
  $MEGATRON_PATH/pretrain_gpt.py \
  $MODEL_ARGS \
  $OPTIM_ARGS \
+  $ZERO_ARGS \
  $PRECISION_ARGS \
  $PARALLEL_ARGS \
  $DATA_ARGS \
--- a/scripts/kaiyuan2b-training/training_smoke_qwen3_1p7b.sh
+++ b/scripts/kaiyuan2b-training/training_smoke_qwen3_1p7b.sh
@@ -16,6 +16,12 @@ CKPT_DIR="${ARTIFACT_ROOT}/checkpoints/${TRAIN_NAME}"
 CHECKPOINT_KEEP_RECENT=${CHECKPOINT_KEEP_RECENT:-3}
 CHECKPOINT_CLEANUP_INTERVAL_SECONDS=${CHECKPOINT_CLEANUP_INTERVAL_SECONDS:-300}
 EXTRA_ARGS=${EXTRA_ARGS:-}
+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}
+ZERO_STAGE=${ZERO_STAGE:-0}

 source "${PARAMS_DIR}/optim_common.sh"
 source "${PARAMS_DIR}/qwen3_1p7b/model.sh"
@@ -57,6 +63,28 @@ else
  exit 1
 fi

+case "$ZERO_STAGE" in
+  0)
+    ZERO_ARGS=""
+    ;;
+  1)
+    ZERO_ARGS="
+      --use-distributed-optimizer
+      --data-parallel-sharding-strategy optim
+    "
+    ;;
+  2)
+    ZERO_ARGS="
+      --use-distributed-optimizer
+      --data-parallel-sharding-strategy optim_grads
+    "
+    ;;
+  *)
+    echo "Unsupported ZERO_STAGE=${ZERO_STAGE}; expected 0, 1, or 2" >&2
+    exit 1
+    ;;
+esac
+
 mkdir -p "$CKPT_DIR" "$TB_DIR"

 cleanup_old_checkpoints_once() {
@@ -118,17 +146,18 @@ CHECKPOINT_CLEANUP_PID=$!
 trap 'kill "$CHECKPOINT_CLEANUP_PID" 2>/dev/null || true; cleanup_old_checkpoints_once "$CKPT_DIR" "$CHECKPOINT_KEEP_RECENT"' EXIT

 DISTRIBUTED_ARGS="
-  --nproc_per_node 8
-  --nnodes 1
-  --node_rank 0
-  --master_addr localhost
-  --master_port 6000
+  --nproc_per_node ${NPROC_PER_NODE}
+  --nnodes ${NNODES}
+  --node_rank ${NODE_RANK}
+  --master_addr ${MASTER_ADDR}
+  --master_port ${MASTER_PORT}
 "

 torchrun $DISTRIBUTED_ARGS \
  $MEGATRON_PATH/pretrain_gpt.py \
  $MODEL_ARGS \
  $OPTIM_ARGS \
+  $ZERO_ARGS \
  $PRECISION_ARGS \
  $PARALLEL_ARGS \
  $DATA_ARGS \