Minimind 模型训练

我的配置

四卡 4090 服务器，160G 内存，VPN

Download Datasets

我并没有找到在服务器环境下只下载单个数据集的方法，只能通过 huggingface（或modelscope） 全部下载下来，代码如下：

from huggingface_hub import snapshot_download
repo_id = "jingyaogong/minimind_dataset"  # 数据集仓库 ID
local_dir = "./dataset"  # 本地下载目录
snapshot_download(repo_id=repo_id, repo_type="dataset", local_dir=local_dir)

由于我的资源是足够的，所以打算训练完整的 MiniMind2 模型，参数如下：

Model Name	params	len_vocab	rope_theta	n_layers	d_model	share	选择
MiniMind2-Small	26M	6400	1e6	8	512	-
MiniMind2-MoE	145M	6400	1e6	8	640	1+4
MiniMind2	104M	6400	1e6	16	768	-	✓

• rope_theta:
• share+route:

Training Processure

两张 4090 可提供 48G 显存，所以我可以取比较大的 batch_size。以下是我的训练配置：

参数	epochs	batch_size	dim	n_layers	max_seq_len	len_vocab
配置	6	48	768	16	512	6400

参数总量计算：$Params\approx 12Lh^2=16\times12\times 768^2=113M$，与实际的 104M 略有差异，之后会回来找原因。

训练命令

1. 注意，由于我的计算节点上是没有网络的，所以需要设置 wandb 为离线模式；
2. SFT 用多卡服务器跑的时候需要大量的内存（> 100G），不然会报错；

wandb.init(project=args.wandb_project, name=args.wandb_run_name, mode="offline") # 需要修改

torchrun --nproc_per_node 2 train_pretrain.py --use_wandb
torchrun --nproc_per_node 4 train_full_sft.py --epochs 4 --use_wandb
torchrun --nproc_per_node 4 train_full_sft.py --epochs 3 --use_wandb --max_seq_len 2048 --batch_size 8

测试 LLM 命令

这里选择的测试框架是 lm-evaluation-harness，可以参考 lm-evaluation-harness 的文档。
在开始测试之前，需要把我们训练好的 pytorch 模型转化为 huggingface 的模型。

python convert_model.py

在按照文档配置好环境后，运行以下命令：

lm_eval --model hf --model_args pretrained="you/hf/model/path",device=cuda,dtype=auto --task arc_challenge,arc_easy,boolq,ceval-valid,cmmlu,translation --batch_size 16 --trust_remote_code --output_path "your/output/path"

即可获得模型的评估结果。我这里选择了以下几个评估任务：

任务名称	类型	描述
arc_challenge	科学推理	AI2 Reasoning Challenge 的挑战集，包含需要复杂推理的科学问题，主要是验证模型英文能力
arc_easy	科学推理	AI2 Reasoning Challenge 的简单集，包含基础科学知识问题，主要是验证模型英文能力
boolq	阅读理解	自然语言的是/否问题回答任务，主要是验证模型英文能力
ceval-valid	中文评估	中文基础模型评估基准，覆盖多个学科领域，主要是验证模型中文能力
cmmlu	中文评估	中文多任务语言理解基准，测试模型在专业领域的表现，主要是验证模型中文能力
translation	翻译	主要是测试模型的中英互译能力