docker run -it --rm \
  --gpus all \
  --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 需自行准备。

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

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

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

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

4. 模型定义与训练脚本

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

4.1 `training_*.sh`

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

4.2 `params/<model_name>`

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

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

它们分别回答三个问题：

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

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 或对应的推理脚本进行模型推理。

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'

8. 常用脚本

8.1 数据下载

bash scripts/download_kaiyuan.sh

8.2 训练

cd scripts/kaiyuan2b-training
bash training_smoke_gpt2.sh

8.3 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 脚本

README.md Unescape Escape

pretrain_kaiyuan2b

0. 环境配置：Docker 与 submodule

0.1 Docker 环境