pretrain_kaiyuan2b/README.md

pretrain_kaiyuan2b

本仓库用于基于 Megatron-LM 复现 Kaiyuan-2B 的预训练流程。

0. 环境配置：Docker 与 submodule

0.1 Docker 环境

训练和推理建议直接使用准备好的 Megatron-LM 容器环境。示例：

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

说明：拉取该镜像需要先登录对应的私有 Docker Registry，相关 credentials 需自行准备。

多机多卡训练时建议每台机器都用同名镜像和同名容器启动，并保持容器内代码、数据、checkpoint 路径一致。--network=host 用于避免 Docker bridge/NAT 影响 torchrun rendezvous 和 NCCL 跨机通信。

0.2 Megatron-LM submodule

Megatron-LM 作为 git submodule 放在仓库根目录，用于实际的模型训练与推理。

执行下面的命令初始化 submodule：

git submodule update --init --recursive Megatron-LM

1. 仓库功能

本仓库的目标为：

• 基于 Megatron-LM 进行 Kaiyuan-2B 模型的复现训练
• 提供数据下载、数据转换、模型定义、训练参数构建、训练与推理脚本
• 提供可复用的 profiling / training 两套脚本组织方式

2. 仓库组织

2.1 Megatron-LM

Megatron-LM 是训练与推理的核心代码，作为 submodule 放在仓库根目录。

2.2 scripts

scripts 目录存储所有外部脚本，包括：

• 数据下载脚本（download_kaiyuan.sh）
• 数据转换脚本（convert_phase_to_megatron.py）
• 数据格式检查脚本（inspect_parquet.py）
• 权重转换脚本（convert_megatron_weight_to_hf.sh）
• toy 训练脚本（toy_model_training.sh）
• 模型定义与训练参数构建脚本（params/*）
• 模型训练脚本（training_*.sh / eval_*.sh）
• Docker 相关脚本（Dockerfile）

2.3 kaiyuan2b-profiling 与 kaiyuan2b-training

这两个目录都放在 scripts/ 下：

• scripts/kaiyuan2b-profiling：用于启动训练过程前的性能调优、profiling、实验验证脚本
• scripts/kaiyuan2b-training：结构和 profiling 基本一致，但用于实际执行模型训练

两者的脚本结构大体相同，通常都包含：

• training_*.sh：拉起具体训练任务
• eval_*.sh：执行模型推理/验证
• params/<model_name>/：模型参数定义，通常拆成 data、hparams、model 三部分

3. 数据下载与准备

3.1 数据下载

需要使用 scripts/download_kaiyuan.sh 下载 phase1 到 phase5 的数据。

数据分布与背景建议参考：

• 论文：https://arxiv.org/abs/2512.07612
• Hugging Face 数据集：https://huggingface.co/thu-pacman/PCMind-2.1-Kaiyuan-2B

3.2 数据转换

下载后的数据通常是 parquet 格式，需要结合 tokenizer 转换成 Megatron 支持的训练数据格式：.bin 和 .idx。

仓库中提供的转换脚本：

• scripts/convert_phase_to_megatron.py

该脚本直接读取 parquet，并使用 Megatron 的 tokenizer 与 IndexedDatasetBuilder 写出 .bin / .idx，不再生成 JSONL 中间文件。

转换前需要准备 tokenizer 的 4 个定义文件：

• merges.txt
• tokenizer.json
• vocab.json
• tokenizer_config.json

这些文件需要从 Kaiyuan-2B 的模型权重或对应 tokenizer 配置中手动提取出来，再交给 Megatron 的数据预处理流程使用:

wget https://hf-mirror.com/thu-pacman/PCMind-2.1-Kaiyuan-2B/resolve/refs%2Fpr%2F1/tokenizer.json
wget https://hf-mirror.com/thu-pacman/PCMind-2.1-Kaiyuan-2B/resolve/refs%2Fpr%2F1/tokenizer_config.json
wget https://hf-mirror.com/thu-pacman/PCMind-2.1-Kaiyuan-2B/resolve/refs%2Fpr%2F1/vocab.json
wget https://hf-mirror.com/thu-pacman/PCMind-2.1-Kaiyuan-2B/resolve/refs%2Fpr%2F1/merges.txt

转换示例：

python scripts/convert_phase_to_megatron.py \
  --input-dir /apps/yi/model_training/data/phase1 \
  --output-dir /ssd/yi/converted_data/megatron_phase1 \
  --megatron-dir /apps/yi/model_training/Megatron-LM \
  --tokenizer-model /apps/yi/model_training/data/tokenizer \
  --text-key text \
  --output-prefix-prefix phase1 \
  --num-shards 4 \
  --workers-per-shard 16 \
  --batch-size 8192 \
  --chunksize 64

并发参数建议：

• --num-shards：同时处理多少个 parquet 文件
• --workers-per-shard：每个 parquet 的 tokenizer worker 数
• 总 tokenizer worker 数约等于 num_shards * workers_per_shard
• 建议让总 worker 数接近机器物理 CPU core 数，再根据磁盘 I/O 和 tokenizer 吞吐调整

输出文件名保持与训练脚本兼容，例如：

/ssd/yi/converted_data/megatron_phase1/phase1_part-00000_text_document.bin
/ssd/yi/converted_data/megatron_phase1/phase1_part-00000_text_document.idx

4. 模型定义与训练脚本

模型定义主要放在 scripts/kaiyuan2b-training 中。

4.1 training_*.sh

training_*.sh 负责拉起具体的模型训练任务。

4.2 params/<model_name>

params/<model_name> 一般包含三部分：

• data：定义用什么数据训练
• hparams：定义训练超参数
• model：定义模型结构

它们分别回答三个问题：

1. 用什么数据训练
2. 训练的超参数如何设置
3. 模型结构如何定义

4.3 多机 32 卡训练速查

完整说明见：

• scripts/kaiyuan2b-training/MULTINODE_TRAINING.md

当前多机训练目标是 g0033-g0036 四台机器，每台 8 张 H200，总计 32 张 GPU。脚本仍使用 Megatron-LM + torchrun，默认并行方式是数据并行：

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

四机训练时，torchrun 的核心参数是：

--nproc_per_node 8
--nnodes 4
--node_rank 0/1/2/3
--master_addr <g0033 address>
--master_port 6000

training_smoke_qwen3_1p7b.sh 和 training_smoke_gpt2.sh 已支持通过环境变量覆盖这些参数：

NPROC_PER_NODE=8
NNODES=4
NODE_RANK=0
MASTER_ADDR=10.20.32.33
MASTER_PORT=6000

从 g0033 编排四机训练时使用：

cd /ssd1/yi/pretrain_kaiyuan2b/scripts/kaiyuan2b-training

CONTAINER_NAME=megatron-ngc25-training \
MASTER_ADDR=10.20.32.33 \
NCCL_DEBUG=INFO \
NCCL_SOCKET_IFNAME=eth00,eth01,eth02,eth03,eth04,eth05,eth06,eth07 \
GLOO_SOCKET_IFNAME=bond1.1032 \
bash start_multinode_training.sh gpt_smoke smoke smoke_32gpu

停止四机任务：

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

如果不使用 Docker 容器执行训练，可以不设置 CONTAINER_NAME。设置后，脚本会在每台机器上通过 docker exec <CONTAINER_NAME> 进入容器启动训练。

四台机器上的容器建议这样启动，路径按实际部署调整：

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

多机训练前需要确认：

• 四台机器都能访问同一套代码、tokenizer、训练数据路径。
• 四台机器都已导入同一个 Docker image。
• 容器网络模式是 host。
• g0033 能通过 ssh 访问 g0034-g0036。
• MASTER_ADDR:MASTER_PORT 能从其他节点访问。
• NCCL 使用正确网卡，例如 eth00-eth07。

ZeRO/分布式优化器通过 ZERO_STAGE 控制：

ZERO_STAGE=0  # 默认，不启用 distributed optimizer
ZERO_STAGE=1  # --use-distributed-optimizer --data-parallel-sharding-strategy optim
ZERO_STAGE=2  # --use-distributed-optimizer --data-parallel-sharding-strategy optim_grads

建议验证顺序：

1. 单机 gpt_smoke
2. 四机 gpt_smoke
3. 四机 qwen3_1p7b + ZERO_STAGE=0
4. 对比 ZERO_STAGE=1/2 的吞吐和显存

5. 模型说明

5.1 gpt_smoke

gpt_smoke 是一个用于验证数据链路、训练脚本和 eval 脚本是否正常工作的简单测试模型，不是最终目标模型。

其配置是一个大约 330M 参数量的 Toy Model，核心结构为：

• 12 层 Transformer Layer
• hidden size: 768
• FFN hidden size: 3072
• attention heads: 8
• query groups: 4
• 使用 GQA（Group Query Attention）
• 上下文窗口 / 最大位置长度：4096
• 位置编码：RoPE
• 归一化：RMSNorm
• 激活：SwiGLU

5.2 qwen3_1p7b

qwen3_1p7b 是我们实际要训练的主模型，一个约 1.7B 参数量的大模型。

其结构配置为：

• 28 层 Transformer Layer
• hidden size: 2048
• FFN hidden size: 6144
• attention heads: 16
• query groups: 8
• 使用 GQA（Group Query Attention）
• 上下文窗口 / 最大位置长度：4096
• 位置编码：RoPE
• 归一化：RMSNorm
• 激活：SwiGLU
• 不使用 bias linear
• 输出与输入 embedding 不共享权重

6. 训练超参数与 paper 对齐说明

qwen3_1p7b 的模型定义和训练超参数已经尽可能对齐原论文配置。

在实际执行时，我们优先保持：

• 模型结构与论文一致
• 训练超参数与论文一致或尽量接近
• 数据配比、训练步数、batch 设定尽量沿用论文设定

相关对照信息请参考论文附录中的 Table 11。

7. 推理与验证

训练完成后，可以使用 eval_<model_name>.sh 或对应的推理脚本进行模型推理。

注意: 执行推理前需在docker环境中安装 flask:

pip install flask-restful

7.1 推理前的必要修改

在执行推理前，需要手动修改 Megatron-LM/megatron/core/inference/text_generation_server/run_mcore_engine.py 的第 89 行，把：

"tokens": [x.prompt_tokens + x.generated_tokens.tolist() for x in result],

改成：

"tokens": [x.prompt_tokens.tolist() + x.generated_tokens for x in result],

否则会报错：

AttributeError: 'list' object has no attribute 'tolist'

并且需要将 Megatron-LM/tools/run_text_generation_server.py 的第64行, 将

inference_context = StaticInferenceContext(args.inference_max_requests, args.inference_max_sequence_length)

改为

inference_context = StaticInferenceContext(args.inference_max_requests, <any integer, such as 4096>)

避免报错:

[rank0]: AttributeError: 'Namespace' object has no attribute 'inference_max_sequence_length'

推理服务部署成功后会显示:

INFO:werkzeug:WARNING: This is a development server. Do not use it in a production deployment. Use a production WSGI server instead.
 * Running on all addresses (0.0.0.0)
 * Running on http://127.0.0.1:5000
 * Running on http://172.17.0.2:5000
INFO:werkzeug:Press CTRL+C to quit

切换到同Docker脚本下的另一个 bash terminal, 执行如下的命令即可测试模型推理:

curl -X PUT http://127.0.0.1:5000/api \
  -H "Content-Type: application/json" \
  -d '{
    "prompts": ["The capital of France is"],
    "tokens_to_generate": 50,
    "temperature": 0.8,
    "top_k": 0,
    "top_p": 0.9
  }'

8. 常用脚本

8.1 数据下载

bash scripts/download_kaiyuan.sh

8.2 训练

cd scripts/kaiyuan2b-training
bash training_smoke_gpt2.sh

8.3 动态启动训练

推荐使用 scripts/kaiyuan2b-training/start_training.sh 动态拉起训练任务。该脚本只负责选择并后台启动已有的 training_*.sh，不改变 data、hparams、model 分文件定义的组织方式。

当前支持的模型入口：

• gpt_smoke：对应 training_smoke_gpt2.sh
• qwen3_1p7b：对应 training_smoke_qwen3_1p7b.sh

示例：

cd scripts/kaiyuan2b-training

# 启动 gpt_smoke smoke 训练
bash start_training.sh gpt_smoke smoke smoke_gpt

# 启动 qwen3_1p7b smoke 训练
bash start_training.sh qwen3_1p7b qwen3_1p7b_smoke_yi qwen3_1p7b_smoke_yi

启动后会写入：

• PID 状态：/apps/yi/model_training/artifacts/run_state/<train_name>.pid
• 任务元信息：/apps/yi/model_training/artifacts/run_state/<train_name>.env
• 训练日志：/apps/yi/model_training/artifacts/logs/<train_name>.log
• TensorBoard 日志：/apps/yi/model_training/artifacts/tb_logs/<train_name>
• checkpoint：/apps/yi/model_training/artifacts/checkpoints/<train_name>

start_training.sh 会自动给训练命令追加 Megatron 的 --exit-signal-handler，用于支持收到 SIGTERM 后保存 checkpoint 并退出。

可以通过环境变量传入额外 Megatron 参数：

EXTRA_ARGS="--exit-duration-in-mins 120" \
  bash start_training.sh gpt_smoke smoke smoke_gpt_2h

8.4 停止训练

使用 stop_training.sh 按 train_name 停止任务：

cd scripts/kaiyuan2b-training
bash stop_training.sh smoke_gpt

停止脚本会向训练进程组发送 SIGTERM。由于训练启动时已经启用 --exit-signal-handler，Megatron 会在训练循环中保存 checkpoint 后退出。默认等待 300 秒，可通过 GRACE_SECONDS 调整：

GRACE_SECONDS=600 bash stop_training.sh smoke_gpt

不建议直接 kill -9，除非确认 checkpoint 保存已经卡死且必须释放 GPU。

8.5 保留最近 K 个 checkpoint

两个正式训练脚本都会自动清理旧 checkpoint，默认保留最近 3 个 iter_* 目录，避免长期训练炸盘：

• scripts/kaiyuan2b-training/training_smoke_gpt2.sh
• scripts/kaiyuan2b-training/training_smoke_qwen3_1p7b.sh

可以通过环境变量修改：

# 保留最近 5 个 checkpoint
CHECKPOINT_KEEP_RECENT=5 bash start_training.sh qwen3_1p7b

# 每 60 秒检查一次旧 checkpoint
CHECKPOINT_CLEANUP_INTERVAL_SECONDS=60 bash start_training.sh gpt_smoke

# 关闭自动清理
CHECKPOINT_KEEP_RECENT=0 bash start_training.sh gpt_smoke

清理逻辑只删除 checkpoint 目录下形如 iter_0001000 的旧目录，不删除 latest_checkpointed_iteration.txt、TensorBoard 日志或其他 artifact。

8.6 Profiling

cd scripts/kaiyuan2b-profiling
bash training_smoke_gpt2.sh

9. 目录快速索引

• Megatron-LM/：Megatron-LM 子模块
• scripts/download_kaiyuan.sh：数据下载
• scripts/convert_phase_to_megatron.py：数据转换
• scripts/inspect_parquet.py：数据格式检查
• scripts/convert_megatron_weight_to_hf.sh：权重导出/转换
• scripts/toy_model_training.sh：小规模训练验证
• scripts/kaiyuan2b-training/：实际训练脚本
• scripts/kaiyuan2b-profiling/：性能调优与 profiling 脚本