232 lines
6.2 KiB
Markdown
232 lines
6.2 KiB
Markdown
# pretrain_kaiyuan2b
|
||
|
||
本仓库用于基于 Megatron-LM 复现 Kaiyuan-2B 的预训练流程。
|
||
|
||
## 0. 环境配置：Docker 与 submodule
|
||
|
||
### 0.1 Docker 环境
|
||
|
||
训练和推理建议直接使用准备好的 Megatron-LM 容器环境。示例：
|
||
|
||
```bash
|
||
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：
|
||
|
||
```bash
|
||
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`：定义模型结构
|
||
|
||
它们分别回答三个问题：
|
||
|
||
1. 用什么数据训练
|
||
2. 训练的超参数如何设置
|
||
3. 模型结构如何定义
|
||
|
||
## 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 行，把：
|
||
|
||
```python
|
||
"tokens": [x.prompt_tokens + x.generated_tokens.tolist() for x in result],
|
||
```
|
||
|
||
改成：
|
||
|
||
```python
|
||
"tokens": [x.prompt_tokens.tolist() + x.generated_tokens for x in result],
|
||
```
|
||
|
||
否则会报错：
|
||
|
||
```text
|
||
AttributeError: 'list' object has no attribute 'tolist'
|
||
```
|
||
|
||
## 8. 常用脚本
|
||
|
||
### 8.1 数据下载
|
||
|
||
```bash
|
||
bash scripts/download_kaiyuan.sh
|
||
```
|
||
|
||
### 8.2 训练
|
||
|
||
```bash
|
||
cd scripts/kaiyuan2b-training
|
||
bash training_smoke_gpt2.sh
|
||
```
|
||
|
||
### 8.3 Profiling
|
||
|
||
```bash
|
||
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 脚本
|