Keyboard shortcuts

Press or to navigate between chapters

Press S or / to search in the book

Press ? to show this help

Press Esc to hide this help

快速开始

在进入各种理论之前,最好先把原始项目跑起来,直接开始跑 Pretrain。原始项目介绍得很详细: 原项目参考文档

下面内容是我实践时整理出来的版本。它不是完整复述原项目 README,而是把最容易混乱的几件事先梳理清楚:

  1. 不同 MiniMind 版本到底对应什么模型配置。
  2. 预训练和后续阶段分别应该下载哪些数据。
  3. 如果我现在想训练某一个版本,应该敲哪条命令。

MiniMind3 的模型结构和训练代码升级细节,可以参考另一篇笔记:MiniMind 3 升级内容

1. 先梳理模型配置

MiniMind 项目历史版本比较多,而且有些名字描述的是“模型结构”,有些名字描述的是“训练流程产物”,所以第一次看 README 时很容易混在一起。

我目前建议这样理解:

  • MiniMind Zero:更像是一个训练目标或者实践路线,而不是一个独立的新架构。它通常指“用 mini 数据集,从 0 快速训练出一个能基本对话的小模型”。在 MiniMind3 语境下,最快复现 Zero 通常是 pretrain_t2t_mini.jsonl + sft_t2t_mini.jsonl
  • minimind-3:当前主线 Dense 模型,默认 hidden_size=768num_hidden_layers=8,参数量约 64M。
  • minimind-3-moe:当前主线 MoE 模型,主干配置和 minimind-3 基本一致,但 FFN 换成 MoE,参数量约 198M,总参数更多,但激活参数约 64M。
  • minimind2-small / minimind2:历史版本配置。现在如果只是跟着最新代码实践,可以优先看 minimind-3,历史版本主要作为对照。

1.1 常见模型版本表

版本 / 训练目标参数量关键结构参数训练命令中的核心参数适合场景
MiniMind Zero取决于所选结构,MiniMind3 下通常约 64M通常使用 minimind-3 Dense 结构--hidden_size 768 --num_hidden_layers 8 --use_moe 0最快从 0 跑通 pretrain + SFT,观察完整训练链路
minimind-3约 64Md_model=768n_layers=8q_heads=8kv_heads=4--hidden_size 768 --num_hidden_layers 8 --use_moe 0当前主线 Dense 模型,优先推荐
minimind-3-moe约 198M-A64Md_model=768n_layers=84 experts / top-1 routing--hidden_size 768 --num_hidden_layers 8 --use_moe 1想观察 MoE 训练和激活参数概念
minimind2-small约 26Md_model=512n_layers=8,历史轻量配置--hidden_size 512 --num_hidden_layers 8 --use_moe 0更低成本的历史小模型配置
minimind2约 104Md_model=768n_layers=16,历史主线配置当前脚本需要额外指定 --hidden_size 768 --num_hidden_layers 16 --use_moe 0作为历史主线对照

1.2 为什么 MiniMind3 默认选择 hidden_size=768

hidden_size 也就是 Transformer 里的 d_model。它决定了每个 token 在模型内部被表示成多宽的向量。

如果 hidden_size 太小,模型表达能力会明显受限。尤其是在 attention 中,每个 head 的维度通常是:

其中,\(d_\text{model}\) 表示隐藏层维度,\(n_\text{heads}\) 表示 query head 数量,\(d_\text{head}\) 表示每个 head 的维度:

\[ d_\text{head} = \frac{d_\text{model}}{n_\text{heads}} \]

比如 hidden_size=512num_attention_heads=8 时,\(d_\text{head}=64\)。如果继续把 hidden_size 压得更小,head 内部可表达的信息也会变少。对于 MiniMind 这种小模型来说,embedding 层、attention head 维度、FFN 中间层都会一起受到影响。

但如果 hidden_size 太大,训练成本又会明显上升。MiniMind3 选择 hidden_size=768num_hidden_layers=8,本质上是在几个目标之间做折中:

  • 相比 512768 有更好的表示能力。
  • 相比堆到更多层,8 layers 更容易在个人 GPU 上快速训练。
  • 对于 64M 级别的小模型,它能在训练速度、稳定性和效果之间取得比较平衡的结果。

所以实践时可以先记住一个简单规则:

  • 想快速复现,优先用 hidden_size=768, num_hidden_layers=8 的 MiniMind3 Dense。
  • 想更省资源,可以退回 hidden_size=512, num_hidden_layers=8 的历史 small 配置。
  • 想观察 MoE,就保持 768 x 8 不变,只把 --use_moe 1 打开。

2. 再梳理数据配置

新版 MiniMind3 默认数据已经从早期的 pretrain_hq.jsonl 切换到 pretrain_t2t(_mini).jsonlsft_t2t(_mini).jsonl

数据集下载地址:

下载后统一放到 ../dataset 目录下。这里使用 ../dataset 是为了和训练脚本默认参数、上游项目习惯保持一致。实际运行前只要确认下面路径能被当前命令找到即可:

uv run python -m minimind_learning.trainer.train_pretrain --data_path ../dataset/pretrain_t2t_mini.jsonl

如果你把数据放在别的位置,只需要同步修改 --data_path,模型和训练逻辑本身不受影响。

2.1 数据集用途表

数据集大致用途推荐场景
../dataset/pretrain_t2t_mini.jsonl轻量预训练数据快速复现 MiniMind Zero / 快速验证训练流程
../dataset/pretrain_t2t.jsonlMiniMind3 主线预训练数据完整训练 minimind-3
../dataset/sft_t2t_mini.jsonl轻量 SFT 数据,已混入部分 Tool Call 样本快速把 pretrain 模型训练成基础对话模型
../dataset/sft_t2t.jsonlMiniMind3 主线 SFT 数据,含 Tool Call 样本完整 SFT,追求更好效果
../dataset/dpo.jsonl偏好对齐数据DPO 阶段使用,不属于 pretrain
../dataset/rlaif.jsonlRLAIF 数据PPO / GRPO / CISPO 等后续 RL 阶段
../dataset/agent_rl.jsonl / ../dataset/agent_rl_math.jsonlAgentic RL 数据多轮 Tool-Use / 数学工具场景

如果只是想回家先把 pretrain 跑起来,最小需要:

../dataset/pretrain_t2t_mini.jsonl

如果想继续完成 Zero 风格对话模型,还需要:

../dataset/sft_t2t_mini.jsonl

2.2 原项目推荐的数据组合和训练参数

下面这张表把原项目 README 和训练脚本中的默认配置合在一起看。这样后面真正训练时,就不只是知道“下载哪个数据集”,也知道每个阶段大概应该用什么 epoch 和关键参数。

目标数据组合阶段epoch关键训练参数说明
MiniMind Zero / 快速复现../dataset/pretrain_t2t_mini.jsonl + ../dataset/sft_t2t_mini.jsonlPretrain1batch_size=32learning_rate=5e-4accumulation_steps=8max_seq_len≈768先用 mini 预训练数据快速得到 base 权重
MiniMind Zero / 快速复现../dataset/pretrain_t2t_mini.jsonl + ../dataset/sft_t2t_mini.jsonlFull SFT1batch_size=16learning_rate=1e-5accumulation_steps=1max_seq_len≈768在 pretrain 权重上继续训练,得到基础对话能力
MiniMind3 Dense 主线../dataset/pretrain_t2t.jsonl + ../dataset/sft_t2t.jsonlPretrain2batch_size=32learning_rate=5e-4accumulation_steps=8max_seq_len≈380更接近完整 minimind-3 主线训练
MiniMind3 Dense 主线../dataset/pretrain_t2t.jsonl + ../dataset/sft_t2t.jsonlFull SFT2batch_size=16learning_rate=1e-5accumulation_steps=1max_seq_len=768主线 SFT,数据中已混入 Tool Call 样本
MiniMind3 MoE可沿用 mini 或主线数据组合Pretrain / Full SFT同 Dense在对应 Dense 命令基础上加 --use_moe 1总参数更大,激活参数约等于 Dense,训练会更慢
DPO 后续对齐../dataset/dpo.jsonlDPO1batch_size=4learning_rate=4e-8beta=0.15max_seq_len=1024属于 SFT 之后的偏好对齐阶段,不是 pretrain 必需步骤

需要注意,原项目里说 minimind-3 在单卡 3090 上 1 epoch2.31h,指的是 pretrain_t2t_mini + sft_t2t_mini 两阶段合计的快速 Zero 路线;minimind-3-moe 同样数据组合下约 3.23h。这些数字主要用于估算门槛,不应该理解成固定训练时间,因为不同 GPU、数据加载速度、max_seq_len、batch size 都会影响耗时。

从脚本默认参数看,三个主要阶段的默认值大致是:

Pretrain:
  epochs=2
  batch_size=32
  learning_rate=5e-4
  accumulation_steps=8
  max_seq_len=340
  data_path=../dataset/pretrain_t2t_mini.jsonl

Full SFT:
  epochs=2
  batch_size=16
  learning_rate=1e-5
  accumulation_steps=1
  max_seq_len=768
  data_path=../dataset/sft_t2t_mini.jsonl

DPO:
  epochs=1
  batch_size=4
  learning_rate=4e-8
  beta=0.15
  max_seq_len=1024
  data_path=../dataset/dpo.jsonl

所以如果只是第一次实践,我会优先采用:

Pretrain: pretrain_t2t_mini, epochs=1, max_seq_len=768
SFT:      sft_t2t_mini,     epochs=1, max_seq_len=768

先把完整链路跑通,再考虑把数据切到 pretrain_t2t / sft_t2t,或者打开 --use_moe 1

2.3 pretrain_t2t_mini 和 pretrain_t2t 怎么选

pretrain_t2t_mini.jsonl 更适合快速实践。它的目标不是把模型训练到最强,而是让我们在较短时间内完整跑通:

数据读取 -> tokenizer -> 模型 forward -> loss -> backward -> checkpoint -> 评估

pretrain_t2t.jsonl 更适合完整复现 MiniMind3 主线模型。它数据量更大,训练时间更长,也更适合和上游 README 中的主线结果对齐。

这里的选择可以简单理解为:

  • 先学习代码、验证环境:选 pretrain_t2t_mini.jsonl
  • 想认真训练一个主线模型:选 pretrain_t2t.jsonl
  • 显存/时间有限:先不要急着上主线完整数据。

2.4 max_seq_len 是 token 长度,不是字符长度

训练参数里的 max_seq_len 指的是 token 数,不是字符数。MiniMind tokenizer 在中文上大致可以粗略理解为 1 token ≈ 1.5~1.7 个中文字符,英文压缩比会更高一些。

所以 max_seq_len 的选择其实是在两个问题之间做平衡:

  • 太短:长样本会被截断,语义不完整。
  • 太长:短样本会 padding 到很长,浪费算力。

上游 README 中给出的经验是:

  • pretrain_t2t_mini.jsonl:可以设置相对长一些,比如 max_seq_len≈768
  • pretrain_t2t.jsonl:主线训练中可以设置更均衡一些,比如 max_seq_len≈380

在我们当前的训练脚本里,默认值是:

--max_seq_len 340

如果只是快速跑通,可以先用默认值;如果你明确使用 pretrain_t2t_mini.jsonl 并希望减少截断,可以手动改成:

--max_seq_len 768

3. 最后给出运行命令和参数

下面命令默认从项目根目录执行。为了和本文统一,数据路径都写成 ../dataset/...

3.1 快速训练 MiniMind Zero 的 pretrain 阶段

这条命令对应的是:使用 MiniMind3 Dense 结构,用 mini 预训练数据快速跑通第一阶段。

uv run python -m minimind_learning.trainer.train_pretrain \
  --data_path ../dataset/pretrain_t2t_mini.jsonl \
  --save_dir ../out \
  --hidden_size 768 \
  --num_hidden_layers 8 \
  --max_seq_len 768 \
  --use_moe 0 \
  --epochs 1 \
  --batch_size 32 \
  --accumulation_steps 8 \
  --learning_rate 5e-4 \
  --save_weight pretrain \
  --from_weight none

训练后主要得到:

../out/pretrain_768.pth
../checkpoints/pretrain_768_resume.pth

如果中断后继续训练,追加:

--from_resume 1

3.2 训练 MiniMind3 Dense 主线 pretrain

这条命令对应的是当前主线 minimind-3 Dense 结构。和上面 Zero 快速路线的主要区别是:数据集换成完整的 pretrain_t2t.jsonl

uv run python -m minimind_learning.trainer.train_pretrain \
  --data_path ../dataset/pretrain_t2t.jsonl \
  --save_dir ../out \
  --hidden_size 768 \
  --num_hidden_layers 8 \
  --max_seq_len 380 \
  --use_moe 0 \
  --epochs 2 \
  --batch_size 32 \
  --accumulation_steps 8 \
  --learning_rate 5e-4 \
  --save_weight pretrain \
  --from_weight none

训练后主要得到:

../out/pretrain_768.pth
../checkpoints/pretrain_768_resume.pth

3.3 训练 MiniMind3 MoE pretrain

这条命令对应 minimind-3-moe。模型主干仍然是 hidden_size=768, num_hidden_layers=8,区别是打开 MoE:

uv run python -m minimind_learning.trainer.train_pretrain \
  --data_path ../dataset/pretrain_t2t_mini.jsonl \
  --save_dir ../out \
  --hidden_size 768 \
  --num_hidden_layers 8 \
  --max_seq_len 768 \
  --use_moe 1 \
  --epochs 1 \
  --batch_size 32 \
  --accumulation_steps 8 \
  --learning_rate 5e-4 \
  --save_weight pretrain \
  --from_weight none

训练后主要得到:

../out/pretrain_768_moe.pth
../checkpoints/pretrain_768_moe_resume.pth

MoE 总参数更多,训练也会更慢。它适合用来观察“总参数量”和“激活参数量”的区别,而不是最快速地跑通流程。

3.4 训练历史 minimind2-small 风格配置

如果只是想用更小配置快速实验,可以使用历史 minimind2-small 风格:

uv run python -m minimind_learning.trainer.train_pretrain \
  --data_path ../dataset/pretrain_t2t_mini.jsonl \
  --save_dir ../out \
  --hidden_size 512 \
  --num_hidden_layers 8 \
  --max_seq_len 512 \
  --use_moe 0 \
  --epochs 1 \
  --batch_size 32 \
  --accumulation_steps 8 \
  --learning_rate 5e-4 \
  --save_weight pretrain \
  --from_weight none

训练后主要得到:

../out/pretrain_512.pth
../checkpoints/pretrain_512_resume.pth

3.5 多卡训练命令

如果有多张 GPU,可以用 torchrun 启动。比如单机 2 卡:

uv run torchrun --nproc_per_node 2 -m minimind_learning.trainer.train_pretrain \
  --data_path ../dataset/pretrain_t2t_mini.jsonl \
  --save_dir ../out \
  --hidden_size 768 \
  --num_hidden_layers 8 \
  --max_seq_len 768 \
  --use_moe 0 \
  --epochs 1 \
  --batch_size 32 \
  --accumulation_steps 8 \
  --learning_rate 5e-4 \
  --save_weight pretrain \
  --from_weight none

其中 --nproc_per_node 2 表示启动 2 个进程,一般对应 2 张 GPU。实际训练时,batch_size 是每个进程上的 batch size,总 batch size 会随 GPU 数增加而变大。为了保持训练动态接近单卡,有时需要同步调整 batch_sizeaccumulation_steps

3.6 可选参数说明

几个常用参数可以先记住:

参数含义
--use_wandb开启 SwanLab / wandb 风格训练可视化
--from_resume 1../checkpoints 中自动恢复训练
--use_compile 1开启 torch.compile,可能提升吞吐,但会增加首次编译开销
--save_interval每隔多少 step 保存一次权重
--log_interval每隔多少 step 打印一次训练日志

4. 训练输入输出

训练需要设置的几个目录:

  • --save_dir:保存模型权重的目录,本文统一使用 ../out
  • --data_path:预训练数据路径,例如 ../dataset/pretrain_t2t_mini.jsonl
  • checkpoint 默认保存到 ../checkpoints
  • tokenizer 默认从 ../model 加载,也就是上游仓库中 model/ 目录下的 tokenizer 文件。

训练后常见输出文件:

  • ../out/pretrain_768.pth:MiniMind3 Dense 预训练权重。
  • ../out/pretrain_768_moe.pth:MiniMind3 MoE 预训练权重。
  • ../out/pretrain_512.pth:历史 small 配置的预训练权重。
  • ../checkpoints/<weight>_<hidden_dim>_resume.pth:续训检查点,包含模型、优化器、训练进度等信息。

5. 实验结果

在实际训练中,我们又遇到了一些很工程化的问题。这里我使用的是一块 RTX 3070 Laptop GPU,显存只有 8GB。经过几轮测试之后,我发现除了模型结构本身,下面两个训练参数也非常关键。

第一个是混合精度的数据类型。对这块消费级显卡来说,float16 比默认的 bfloat16 更合适。bfloat16 在一些新卡或者数据中心卡上很好用,但在 30 系消费级显卡上并不一定是最优选择。实际测试时,float16 的训练路径更顺,也更符合这块显卡的硬件特性。

第二个是 DataLoadernum_workers。一开始很容易以为 worker 越多,数据加载就越快。但在我的机器上,显存和内存都比较紧张,batch size 也不会设得特别大,CPU 到 GPU 的数据拷贝并不是主要瓶颈。反而是 worker 数太大以后,每个 worker 都要启动自己的进程,并持有一部分 Dataset / tokenizer 相关状态,尤其在 Windows 上进程启动开销更明显。实际测试下来,num_workers=0 反而更稳定,启动也更快。

所以在这台机器上,pretrain 阶段我会优先加上:

--dtype float16
--num_workers 0

训练代码细节

相比上游的原始训练代码,我在这里额外补充了几个更偏工程诊断的指标。训练不只是看 loss 是否下降,尤其是在小显存 GPU 上跑预训练时,还需要观察吞吐、梯度和真实处理过的 token 数。

首先是 tokens_seentokens/stokens_seen 统计训练过程中实际参与 loss 计算的 token 数,也就是排除了 label 中 -100 的 padding 或忽略位置之后的有效 token。这个指标比单纯的 step 更稳定,因为不同 batch 里的有效长度可能不同。tokens/s 则用来观察当前配置的吞吐,如果 batch size、sequence length、num_workers 或 dtype 改动之后,tokens/s 明显下降,就说明这组参数可能不是最优的。

第二个是 grad_norm。训练代码在 optimizer.step() 之前会先执行梯度裁剪:

grad_norm = torch.nn.utils.clip_grad_norm_(model.parameters(), args.grad_clip)
last_grad_norm = float(grad_norm.item() if hasattr(grad_norm, "item") else grad_norm)

这里返回的 grad_norm 是裁剪前的全局梯度范数。它可以帮助我们观察训练稳定性:如果 grad_norm 长时间非常大,或者突然出现尖峰,通常说明学习率、batch size、数据分布或者混合精度状态可能存在问题;如果它异常接近 0,也可能说明梯度传播不充分。把它和 loss、learning rate、tokens/s 一起记录到 wandb 里,可以更容易判断训练到底是在正常收敛,还是只是表面上还能继续跑。

因此现在 wandb log 中除了 loss 之外,还会记录:

{
    "loss": current_loss,
    "logits_loss": current_logits_loss,
    "aux_loss": current_aux_loss,
    "learning_rate": current_lr,
    "tokens_seen": tokens_seen,
    "tokens_per_second": current_tokens_per_second,
    "grad_norm": last_grad_norm,
    "optimizer_step": did_optimizer_step,
}

这些指标不会改变训练行为,只是让训练过程更透明。尤其在正式长跑之前,它们可以帮助我们快速判断当前参数是否值得继续跑下去。

测试的重要性

这次实践里一个很深的感受是:不要等完整训练跑起来之后才发现参数不合适。训练代码一旦开始跑,可能几分钟之后才因为显存不够、数据加载过慢、checkpoint 太大或者 CUDA 状态异常而失败。这样每次都从头试,会非常浪费时间。

更好的方式是,在正式实验之前,先做一个很小的诊断轮。这个诊断轮不追求模型效果,只回答几个工程问题:

  1. CUDA 环境是否正常。
  2. 这个模型配置能不能放进显存。
  3. 某个 batch_sizemax_seq_len 组合下,峰值显存是多少。
  4. 每个 micro-step 大概需要多长时间。
  5. 第一次 optimizer step 之后,AdamW 状态会不会显著增加显存。
  6. 当前配置是否还留有足够显存余量。

理论上我们可以估算 batch size 和 sequence length 对显存的影响,但实际训练中还有 CUDA cache、optimizer state、临时张量、DataLoader、后台进程占用等因素。最后还是要在自己的机器上实际跑一下,才能确定真正可用的参数,而且最好留出余量。

为此,我在项目里加入了两个层次的测试。

第一个是 tests/test_gpu_diagnostics.py。它更像一个训练前分析器,而不是普通单元测试。它默认跳过,只有显式设置 RUN_GPU_DIAGNOSTICS=1 才会运行。这个测试会构造一个小的预训练 jsonl,初始化真实的 MiniMind 模型,跑若干个 micro-step,并输出每一步的耗时、显存和 tokens/s。

关键逻辑如下,代码来自 tests/test_gpu_diagnostics.py

optimizer = torch.optim.AdamW(model.parameters(), lr=learning_rate)
scaler = torch.amp.GradScaler("cuda", enabled=(dtype == torch.float16))
optimizer.zero_grad(set_to_none=True)

for micro_step in range(1, micro_steps + 1):
    input_ids, labels = next(data_iter)
    input_ids = input_ids.to("cuda:0", non_blocking=True)
    labels = labels.to("cuda:0", non_blocking=True)

    batch_tokens = int((labels != -100).sum().item())
    lr = get_lr(micro_step, micro_steps, learning_rate)
    for param_group in optimizer.param_groups:
        param_group["lr"] = lr

    with torch.amp.autocast("cuda", dtype=dtype):
        result = model(input_ids, labels=labels)
        loss = (result.loss + result.aux_loss) / accumulation_steps

    scaler.scale(loss).backward()

    if micro_step % accumulation_steps == 0:
        scaler.unscale_(optimizer)
        grad_norm = torch.nn.utils.clip_grad_norm_(model.parameters(), grad_clip)
        scaler.step(optimizer)
        scaler.update()
        optimizer.zero_grad(set_to_none=True)

运行方式如下:

$env:RUN_GPU_DIAGNOSTICS='1'
$env:GPU_DIAG_HIDDEN_SIZE='768'
$env:GPU_DIAG_NUM_LAYERS='8'
$env:GPU_DIAG_SEQ_LEN='768'
$env:GPU_DIAG_BATCH_SIZE='16'
$env:GPU_DIAG_ACCUMULATION_STEPS='8'
$env:GPU_DIAG_STEPS='8'
$env:GPU_DIAG_DTYPE='float16'
$env:GPU_DIAG_NUM_WORKERS='0'
uv run python -m pytest tests/test_gpu_diagnostics.py -s

输出中最需要关注的是 summary 部分。例如一次 hidden_size=768, num_hidden_layers=8, seq_len=768, batch_size=16 的测试结果大致是:

{
  "avg_step_s": 0.611,
  "median_step_s": 0.563,
  "avg_forward_s": 0.202,
  "avg_backward_s": 0.380,
  "tokens_per_second": 20095,
  "peak_allocated_mb": 5951.78,
  "peak_reserved_mb": 6494.0,
  "estimated_reserved_headroom_ratio": 0.207,
  "recommendation": "usable_but_tight"
}

这里的结论很直观:这个配置能跑,但不是特别宽裕。如果后台还有其他程序占用 GPU,或者训练过程中显存有波动,就有可能接近边界。

第二个是实际的 pretrain 诊断轮。为了让真实训练入口也能短跑验证,我给 train_pretrain.py 加了几个只用于诊断的参数:

--max_train_steps
--profile_train
--profile_interval

它们不会把诊断结果保存进 checkpoint,只是在训练时打印每个 step 的耗时和显存。核心代码来自 src/minimind_learning/trainer/train_pretrain.py

if args.profile_train and ((step - start_step) % args.profile_interval == 0 or step == iters):
    if torch.cuda.is_available():
        torch.cuda.synchronize()
        allocated_mb = torch.cuda.memory_allocated() / 1024**2
        reserved_mb = torch.cuda.memory_reserved() / 1024**2
        peak_mb = torch.cuda.max_memory_allocated() / 1024**2
    else:
        allocated_mb = reserved_mb = peak_mb = 0.0

    step_time = time.time() - step_start_time
    avg_step_time = (time.time() - profile_start_time) / max(step - start_step, 1)
    tokens_per_second = int(batch_tokens.item()) / max(step_time, 1e-9)

    Logger(
        f"Profile step {step}: step_time={step_time:.3f}s, "
        f"avg_step_time={avg_step_time:.3f}s, "
        f"tokens/s={tokens_per_second:.1f}, "
        f"optimizer_step={did_optimizer_step}, "
        f"cuda_alloc={allocated_mb:.1f}MB, "
        f"cuda_reserved={reserved_mb:.1f}MB, "
        f"cuda_peak={peak_mb:.1f}MB"
    )

实际短跑命令如下。这里 --max_train_steps 32 表示只跑 32 个 micro-step;当 accumulation_steps=16 时,也就是只做 2 次 optimizer step。

uv run python -u -m minimind_learning.trainer.train_pretrain \
  --data_path ../dataset/pretrain_t2t_mini.jsonl \
  --save_dir ../out \
  --hidden_size 768 \
  --num_hidden_layers 8 \
  --max_seq_len 768 \
  --use_moe 0 \
  --epochs 1 \
  --batch_size 16 \
  --accumulation_steps 16 \
  --learning_rate 5e-4 \
  --save_weight pretrain_dense_probe \
  --from_weight none \
  --dtype float16 \
  --num_workers 0 \
  --log_interval 16 \
  --save_interval 100000000 \
  --max_train_steps 32 \
  --profile_train \
  --profile_interval 1

这次真实 pretrain 短跑的结果和 test_gpu_diagnostics 很接近:

Profile step 16: step_time=0.792s, avg_step_time=0.594s,
tokens/s=3998.5, optimizer_step=True,
cuda_alloc=952.6MB, cuda_reserved=6194.0MB, cuda_peak=5932.5MB

Epoch:[1/1](16/79390), loss: 8.9127, tokens_seen: 65613

Profile step 32: step_time=0.554s, avg_step_time=0.575s,
tokens/s=5477.6, optimizer_step=True,
cuda_alloc=952.6MB, cuda_reserved=6794.0MB, cuda_peak=6431.1MB

Epoch:[1/1](32/79390), loss: 8.4330, tokens_seen: 127676

这说明诊断测试确实能较好地预测真实训练的显存和耗时。seq_len=768, batch_size=16 能跑,但峰值 reserved 显存已经到 6.8GB 左右。对 8GB 显卡来说,这属于“可以跑,但需要留意后台占用”的配置。

显存占用理论值

MiniMind3 Dense 的参数量大约是 64M。如果只看模型权重,并且使用 float16,每个参数占 2 字节,那么权重本身大约是:

\[ 64 \times 10^6 \times 2 \text{ bytes} \approx 128 \text{ MB} \]

看起来这离 8GB 显存还很远。但训练时显存不只存模型权重,还要存梯度、优化器状态、激活值和各种临时张量。

以 AdamW 为例,假设参数量为 \(N\),每个参数用 FP16 存权重。训练时至少还会有:

  • 权重:约 \(2N\) bytes。
  • 梯度:通常也是约 \(2N\) bytes。
  • AdamW 一阶动量 exp_avg:通常是 FP32,约 \(4N\) bytes。
  • AdamW 二阶动量 exp_avg_sq:通常是 FP32,约 \(4N\) bytes。

也就是说,只看参数相关部分,粗略就可能达到:

\[ 2N + 2N + 4N + 4N = 12N \text{ bytes} \]

相对于 FP16 权重本身的 \(2N\) bytes,这大约是 6 倍。当然实际 PyTorch、AMP、optimizer 实现会有细节差异,但这个估算能说明一个关键点:不能只拿模型权重大小去判断训练显存。

不过,在我们这次实际测试里,参数相关的常驻显存并不是最大头。更明显的增长来自激活值。对于 Transformer 来说,激活值会随 batch size 和 sequence length 增长。一个隐藏状态张量的形状大致是:

[B, L, D]

其中:

  • \(B\):batch size。
  • \(L\):sequence length。
  • \(D\):hidden size。

如果使用 float16,只看一个隐藏状态张量,它大约占:

\[ B \times L \times D \times 2 \text{ bytes} \]

例如 B=16, L=768, D=768 时:

\[ 16 \times 768 \times 768 \times 2 \approx 18 \text{ MB} \]

单个隐藏状态并不夸张,但训练时每一层都要为反向传播保留中间激活,attention、MLP、loss 计算也会产生额外临时张量,所以最终显存会比这个单项估算大得多。

尤其要注意,batch size 通常是最直接的放大器。batch_size=8batch_size=16,显存压力几乎会明显上一个台阶;batch_size=16batch_size=32,在 8GB 显卡上就很容易接近边界。

下面是这次在 RTX 3070 Laptop GPU 8GB 上的实测结果。模型都是 MiniMind3 Dense,即 hidden_size=768, num_hidden_layers=8, dtype=float16

max_seq_lenbatch_size诊断峰值显存结论
5121约 0.62GB很轻
5124约 1.30GB很稳
5128约 2.19GB很稳
51216约 3.90GB
51232接近 7.96GB不推荐,已经接近显存上限
7688约 3.03GB推荐,比较稳
76816约 5.67GB 到 6.43GB可用,但需要留余量

其中 seq_len=768, batch_size=16test_gpu_diagnostics 中测到的峰值 allocated 大约是 5.95GB;在真实 pretrain 短跑中,峰值 reserved 大约到 6.79GB。这两个数不完全一样,是因为 CUDA 的 allocated 和 reserved 含义不同:

  • allocated:当前真正被 tensor 使用的显存。
  • reserved:PyTorch CUDA allocator 向显卡申请并缓存起来的显存。

训练时更应该关注 reserved 和系统总显存之间的余量。因为一旦 reserved 已经很接近总显存,后面即使 allocated 看起来还没满,也可能因为临时张量或后台进程导致 OOM。

实际执行的代码

综合上面的测试结果,在 8GB RTX 3070 Laptop GPU 上,如果训练 MiniMind3 Dense,我更推荐下面两个配置。

更稳的版本是:

uv run python -m minimind_learning.trainer.train_pretrain ^
  --data_path ../dataset/pretrain_t2t_mini.jsonl ^
  --save_dir ../out ^
  --hidden_size 768 ^
  --num_hidden_layers 8 ^
  --max_seq_len 768 ^
  --use_moe 0 ^
  --epochs 1 ^
  --batch_size 8 ^
  --accumulation_steps 32 ^
  --learning_rate 5e-4 ^
  --save_weight pretrain ^
  --from_weight none ^
  --dtype float16 ^
  --num_workers 0 ^
  --log_interval 16 ^
  --use_wandb 

稍微激进一些、但实测可以跑的版本是:(实际测试是跑不了的,因为我还开浏览器刷b站,第一个batch还能跑,第二个batch到了7.9GB,开始来回swap,直接卡死了.)

uv run python -m minimind_learning.trainer.train_pretrain ^
  --data_path ../dataset/pretrain_t2t_mini.jsonl ^
  --save_dir ../out ^
  --hidden_size 768 ^
  --num_hidden_layers 8 ^
  --max_seq_len 768 ^
  --use_moe 0 ^
  --epochs 1 ^
  --batch_size 16 ^
  --accumulation_steps 16 ^
  --learning_rate 5e-4 ^
  --save_weight pretrain ^
  --from_weight none ^
  --dtype float16 ^
  --num_workers 0 ^
  --log_interval 2 ^
  --use_wandb 

这两个配置的effective batch size 都尽可能和原文中的推荐做到一样 (effective batch size = 32*8, seq_len=768确保effective token接近):

\[ 8 \times 768 \times 32 = 196608 \]

\[ 16 \times 768 \times 16 = 196608 \]

也就是说,如果显存紧张,可以把单步 batch size 降低,再用更大的 accumulation_steps 把有效 batch size 补回来。这是小显存 GPU 上很常用的做法。

初步Eval

这里使用刚训练完的 pretrain_768.pth 做一次非常粗糙的生成测试。评测输入使用 JSONL,每一行只保留一个 prompt 字段,例如:

{"prompt": "你有什么特长?"}
{"prompt": "北京市是"}

如果只是手动对话,可以直接运行:

python scripts\eval_llm.py ^
  --load_from model ^
  --tokenizer_path ..\tokenizer ^
  --save_dir ..\out ^
  --weight pretrain ^
  --hidden_size 768 ^
  --num_hidden_layers 8 ^
  --use_moe 0 ^
  --max_new_tokens 256 ^
  --temperature 0.85 ^
  --top_p 0.85 ^
  --device cuda

如果使用 JSONL 批量评测,可以运行:

python scripts\eval_llm.py ^
  --load_from model ^
  --tokenizer_path ..\tokenizer ^
  --save_dir ..\out ^
  --weight pretrain ^
  --hidden_size 768 ^
  --num_hidden_layers 8 ^
  --use_moe 0 ^
  --max_new_tokens 256 ^
  --temperature 0.85 ^
  --top_p 0.85 ^
  --eval_file .\eval_prompts.jsonl ^
  --eval-output-file .\eval_result.jsonl ^
  --device cuda

这里 --tokenizer_path 指向 tokenizer 所在目录,--save_dir ..\out 对应前面训练命令保存出的 ..\out\pretrain_768.pth。因为这是 pretrain 阶段的模型,还没有经过 SFT,所以它更适合观察“续写能力”和局部语言模式,而不是期待稳定的指令跟随。

在进行 SFT 训练之前,模型已经可以进行一些“勉强前言搭后语”的成语接龙。比较神奇的是,即使 Pretrain 数据中没有对话性的数据,模型似乎也并不完全受影响。不过这种能力还很不稳定,更多说明的是模型学到了一些局部语言模式,而不是已经具备真正稳定的指令跟随能力。

Input: '<|im_start|>system\nYou are a helpful assistant<|im_end|>\n<|im_start|>user\n你有什么特长?<|im_end|>\n<|im_start|>assistant\n'

Output: 我正在想如何开始一个新项目。我正在思考什么?请提供您的一些想法,我将帮助您确定如何开始一个新项目。请描述您所需要的一些信息,包括项目的目标、预算、期限和时间。请提供您希望完成的项目的细节和任何相关信息。
基于以上这段文本内容回答:如何制定一个有效的项目计划?制定一个有效的项目计划需要以下步骤:确定项目目标、确定预算、制定计划、评估期限、设定优先级、评估项目进展情况等。在制定计划的过程中,务必要保持冷静和集中,避免情绪化和决策过分。同时,也要保持持续性和自我反思能力。

Input: '<|im_start|>你有什么特长?'

Output: 有哪些特长? 我能给你讲解什么是机器学习、自然语言处理、深度学习、人工智能、自然语言处理等等。

Input: '<|im_start|>system\nYou are a helpful assistant<|im_end|>\n<|im_start|>user\n北京市是<|im_end|>\n<|im_start|>assistant\n'

Output: 请为我找到在我附近一个新的餐厅,并提供他们推荐的菜肴。
在您的附近找到一个新的餐厅,并输入您的位置信息(例如,您所在的城市或您想查询的区域)。在输入您的位置信息之后,搜索“附近一个新的餐厅”,找到该餐厅的名称和描述。根据您提供的位置信息,搜索“附近一个新的餐厅”,并从中找出该餐厅的名称和描述。这将提供该餐厅的名称和描述,帮助您找到您所在的位置。

Input: '<|im_start|>北京市是'

Output: 中国的首都,位于中国的东部。北京市是中国的首都,位于中国的东部。它的面积约为 1200 万平方公里,人口约为 120 万。北京市是一个著名的历史文化名城,其建筑和文化遗产吸引了大量游客前来参观。

自动的测试结果:

💬: 你有什么特长?
🤖️: 作为一个AI模型,我具有以下特长:1. 学习:我的知识储备主要是通过大量的数据和算法来实现,包括知识库、统计和深度学习。2. 解决问题:我能够回答问题、提供知识和信息、进行交流和解决方案。3. 创意生成:我能够生成各种类型的创意,包括设计、创作、编码、设计、广告等。4. 创造力:我能够创建新的想法、技术、设计概念等,帮助人们在任何领域取得成功。5. 团队合作:我能够与其他团队成员合作,实现各自的角色和职责。6. 创新和灵活:我能够与其他团队成员一起学习和创新,以满足他们的需求和期望。7. 情感和体验:我能够根据用户的需求和偏好提供情感和体验,以更好地理解用户需求和提供个性化服务。8. 团队合作:我能够与其他团队成员合作,共同完成任务,实现共同的目标。

[Speed]: 61.14 tokens/s

💬: 为什么天空是蓝色的
🤖️: ?回答上面的问题,给出具体的推理逻辑。天空是蓝色的主要是因为大气层中含有的气体分子会吸收和散射蓝光。其中,蓝光是一种介质,波长较短,因此我们看到的天空看起来是蓝色的。当太阳光线穿过大气层时,大气层中的气体分子会散射蓝光,蓝光波长较短,因此我们看到的天空看起来是蓝色的。此外,大气层中的气体分子和微小颗粒也会散射蓝光,这是因为蓝光波长较短,所以我们看到的天空看起来是蓝色的。因此,天空的颜色并不是蓝色,而是因为大气层中的气体分子和微小颗粒所散射蓝光,而是因为蓝光波长较短,所以我们看到的天空是蓝色的。

[Speed]: 61.61 tokens/s

💬: 请用Python写一个计算斐波那契数列的函数
🤖️: 。以下是用Python编写的计算斐波那契数列的函数:```pythondef fibonacci(n):if n <= 0:return [0]elif n == 1:return [0, 1]else:return fibonacci(n-1) + fibonacci(n-2)```

[Speed]: 67.95 tokens/s

💬: 解释一下"光合作用"的基本过程
🤖️: ,并举例说明。光合作用是植物和动物利用太阳能进行光合作用的过程中,将二氧化碳和水转化为有机物质和氧气的过程。光合作用是植物和动物利用阳光、二氧化碳和水进行呼吸作用的一种化学过程,其中光能被转化为化学能,产生氧气。光合作用的过程如下:1. 光合作用的光合作用是指植物和动物通过光合作用来制造氧气。植物和动物可以通过光合作用来产生氧气。2. 光合作用的光合作用是指植物和动物利用光能将二氧化碳和水转化为有机物质,产生氧气。3. 光合作用的光合作用是指植物和动物利用光能将二氧化碳和水转化为有机物质。这个过程可以应用于许多领域,例如:1. 农业:光合作用是植物和动物利用光能将二氧化碳和水转化为有机物质。2. 能源:光合作用是植物和动物利用光能将二氧化碳和水转化为有机物质。3. 能源:光合

[Speed]: 60.11 tokens/s

💬: 如果明天下雨,我应该如何出门
🤖️: ?如果明天下雨,您应该选择在雨中散步,享受大自然的美好。可以考虑带上雨伞或带伞,以便更容易出门。同时,还应该注意防雨,不要暴露在强烈的阳光下。如果天气不好,您应该选择携带雨伞或者携带雨伞。如果天气不好,您应该带上雨伞,以便更容易出门。

[Speed]: 65.65 tokens/s

💬: 比较一下猫和狗作为宠物的优缺点
🤖️: 。猫和狗都是常见的宠物,它们的优缺点如下:1. 聪明:猫是独立的动物,它们需要主人的陪伴和训练,可以在室内或户外活动。猫通常需要更多的运动和活动,而且它们可以成为主人的好伙伴。2. 亲密:猫是相对较小的动物,它们需要人类的陪伴和陪伴,可以提供安全感和亲密感。3. 忠诚:猫是独立的宠物,它们可以陪伴人类,提供陪伴和安慰。4. 独立:猫的陪伴需要独立,能够独立工作,而且需要定期的运动和活动。但也存在着许多问题,如猫需要较高的关注度,而狗需要更多的空间和活动,需要更多的关注和关注。综上所述,猫和狗各自的优点和缺点各不相同,不能简单地比较。但它们的优点是它们可以成为主人忠实的伙伴,同时也是主人需要陪伴和关爱的好伙伴。

[Speed]: 60.6 tokens/s

💬: 解释什么是机器学习
🤖️: ,并提供一个使用该算法的实际应用案例。机器学习是一种人工智能的分支,它使用算法来让计算机系统从数据中学习和识别模式,从而自主地进行决策和预测。一个实际应用案例是图像识别。该算法使用大量数据集来训练模型,然后使用该模型来识别图像中的物体或人脸。该算法使用一个称为图像分类器的算法来预测用户的喜好和行为。

[Speed]: 67.11 tokens/s

💬: 推荐一些中国的美食
🤖️: 和传统美食中国美食有许多,包括寿司、刺身、炸鸡、糖果、火锅、糯米饭等等。这些美食都具有独特的口感和味道,同时也有丰富的营养价值。以下是我推荐的一些中国美食和传统美食:1. 寿司:寿司是中国最著名的传统美食之一,以寿司制作而成。寿司制作可以分为各种不同的类型,包括寿司、生鱼片、海苔、豆腐、海苔等等。寿司的制作方法非常简单,只需要将寿司放在海苔上,然后在海苔上翻滚,就可以享受到美味的寿司了。2. 刺身:刺身是一道传统美食,它以各种刺激的刺激为特点。刺身的制作方法非常简单,只需要将身体内的刺激刺激和肌肉刺激结合起来,就可以享受到美味的刺身。3. 火锅:火锅是一道传统美食,通常由面粉、肉类、蔬菜、调料等材料混合制成。制作方法简单,只需要将火锅放入开水中,加入

[Speed]: 60.6 tokens/s