MiniMind 学习笔记(十二):Pretrain 实操——从版本梳理到 8GB 显卡上的真实训练

MiniMind 学习笔记(十二):Pretrain 实操------从版本梳理到 8GB 显卡上的真实训练

在进入各种理论之前,最好先把原始项目跑起来。这一节是 Pretrain 的实操篇 :不讲公式,讲的是第一次跑 MiniMind 时最容易混乱的三件事------版本对应什么配置、每个阶段下载哪些数据、想训练某个版本应该敲哪条命令;然后是在一块 RTX 3070 Laptop(8GB)上真实训练时踩过的工程坑:float16 vs bfloat16、num_workers、诊断轮、显存实测,以及 pretrain 刚结束的模型能生成出什么。

文章目录

  • [MiniMind 学习笔记(十二):Pretrain 实操------从版本梳理到 8GB 显卡上的真实训练](#MiniMind 学习笔记(十二):Pretrain 实操——从版本梳理到 8GB 显卡上的真实训练)
    • 前言:先跑起来,再谈理论
    • 一、先梳理模型配置
      • [1.1 常见模型版本表](#1.1 常见模型版本表)
      • [1.2 为什么 MiniMind3 默认 hidden_size=768](#1.2 为什么 MiniMind3 默认 hidden_size=768)
    • 二、再梳理数据配置
      • [2.1 数据集用途表](#2.1 数据集用途表)
      • [2.2 原项目推荐的数据组合和训练参数](#2.2 原项目推荐的数据组合和训练参数)
      • [2.3 mini 和主线数据怎么选](#2.3 mini 和主线数据怎么选)
      • [2.4 max_seq_len 是 token 长度,不是字符长度](#2.4 max_seq_len 是 token 长度,不是字符长度)
    • 三、运行命令
      • [3.1 快速训练 MiniMind Zero 的 pretrain 阶段](#3.1 快速训练 MiniMind Zero 的 pretrain 阶段)
      • [3.2-3.4 其他三条命令的差异点](#3.2-3.4 其他三条命令的差异点)
      • [3.5 多卡训练](#3.5 多卡训练)
      • [3.6 可选参数](#3.6 可选参数)
      • [3.7 训练输入输出一览](#3.7 训练输入输出一览)
    • [四、实验结果:8GB 显卡上的工程现实](#四、实验结果:8GB 显卡上的工程现实)
      • [4.1 训练代码补充的诊断指标](#4.1 训练代码补充的诊断指标)
    • 五、测试的重要性:不要等完整训练跑起来才发现参数不合适
      • [5.1 第一层:test_gpu_diagnostics.py](#5.1 第一层:test_gpu_diagnostics.py)
      • [5.2 第二层:真实 pretrain 诊断短跑](#5.2 第二层:真实 pretrain 诊断短跑)
    • 六、显存占用的理论值
      • [6.1 RTX 3070 Laptop 8GB 实测表](#6.1 RTX 3070 Laptop 8GB 实测表)
      • [6.2 allocated 和 reserved 的区别](#6.2 allocated 和 reserved 的区别)
    • 七、实际执行的配置
    • [八、初步 Eval:pretrain 刚结束的模型长什么样](#八、初步 Eval:pretrain 刚结束的模型长什么样)
    • 总结
    • 参考

前言:先跑起来,再谈理论

MiniMind 项目历史版本比较多,而且有些名字描述的是"模型结构",有些名字描述的是"训练流程产物",第一次看 README 很容易混在一起。这一篇按实践顺序展开:
#mermaid-svg-fFFdrfunSGha2OEl{font-family:"trebuchet ms",verdana,arial,sans-serif;font-size:16px;fill:#333;}@keyframes edge-animation-frame{from{stroke-dashoffset:0;}}@keyframes dash{to{stroke-dashoffset:0;}}#mermaid-svg-fFFdrfunSGha2OEl .edge-animation-slow{stroke-dasharray:9,5!important;stroke-dashoffset:900;animation:dash 50s linear infinite;stroke-linecap:round;}#mermaid-svg-fFFdrfunSGha2OEl .edge-animation-fast{stroke-dasharray:9,5!important;stroke-dashoffset:900;animation:dash 20s linear infinite;stroke-linecap:round;}#mermaid-svg-fFFdrfunSGha2OEl .error-icon{fill:#552222;}#mermaid-svg-fFFdrfunSGha2OEl .error-text{fill:#552222;stroke:#552222;}#mermaid-svg-fFFdrfunSGha2OEl .edge-thickness-normal{stroke-width:1px;}#mermaid-svg-fFFdrfunSGha2OEl .edge-thickness-thick{stroke-width:3.5px;}#mermaid-svg-fFFdrfunSGha2OEl .edge-pattern-solid{stroke-dasharray:0;}#mermaid-svg-fFFdrfunSGha2OEl .edge-thickness-invisible{stroke-width:0;fill:none;}#mermaid-svg-fFFdrfunSGha2OEl .edge-pattern-dashed{stroke-dasharray:3;}#mermaid-svg-fFFdrfunSGha2OEl .edge-pattern-dotted{stroke-dasharray:2;}#mermaid-svg-fFFdrfunSGha2OEl .marker{fill:#333333;stroke:#333333;}#mermaid-svg-fFFdrfunSGha2OEl .marker.cross{stroke:#333333;}#mermaid-svg-fFFdrfunSGha2OEl svg{font-family:"trebuchet ms",verdana,arial,sans-serif;font-size:16px;}#mermaid-svg-fFFdrfunSGha2OEl p{margin:0;}#mermaid-svg-fFFdrfunSGha2OEl .label{font-family:"trebuchet ms",verdana,arial,sans-serif;color:#333;}#mermaid-svg-fFFdrfunSGha2OEl .cluster-label text{fill:#333;}#mermaid-svg-fFFdrfunSGha2OEl .cluster-label span{color:#333;}#mermaid-svg-fFFdrfunSGha2OEl .cluster-label span p{background-color:transparent;}#mermaid-svg-fFFdrfunSGha2OEl .label text,#mermaid-svg-fFFdrfunSGha2OEl span{fill:#333;color:#333;}#mermaid-svg-fFFdrfunSGha2OEl .node rect,#mermaid-svg-fFFdrfunSGha2OEl .node circle,#mermaid-svg-fFFdrfunSGha2OEl .node ellipse,#mermaid-svg-fFFdrfunSGha2OEl .node polygon,#mermaid-svg-fFFdrfunSGha2OEl .node path{fill:#ECECFF;stroke:#9370DB;stroke-width:1px;}#mermaid-svg-fFFdrfunSGha2OEl .rough-node .label text,#mermaid-svg-fFFdrfunSGha2OEl .node .label text,#mermaid-svg-fFFdrfunSGha2OEl .image-shape .label,#mermaid-svg-fFFdrfunSGha2OEl .icon-shape .label{text-anchor:middle;}#mermaid-svg-fFFdrfunSGha2OEl .node .katex path{fill:#000;stroke:#000;stroke-width:1px;}#mermaid-svg-fFFdrfunSGha2OEl .rough-node .label,#mermaid-svg-fFFdrfunSGha2OEl .node .label,#mermaid-svg-fFFdrfunSGha2OEl .image-shape .label,#mermaid-svg-fFFdrfunSGha2OEl .icon-shape .label{text-align:center;}#mermaid-svg-fFFdrfunSGha2OEl .node.clickable{cursor:pointer;}#mermaid-svg-fFFdrfunSGha2OEl .root .anchor path{fill:#333333!important;stroke-width:0;stroke:#333333;}#mermaid-svg-fFFdrfunSGha2OEl .arrowheadPath{fill:#333333;}#mermaid-svg-fFFdrfunSGha2OEl .edgePath .path{stroke:#333333;stroke-width:2.0px;}#mermaid-svg-fFFdrfunSGha2OEl .flowchart-link{stroke:#333333;fill:none;}#mermaid-svg-fFFdrfunSGha2OEl .edgeLabel{background-color:rgba(232,232,232, 0.8);text-align:center;}#mermaid-svg-fFFdrfunSGha2OEl .edgeLabel p{background-color:rgba(232,232,232, 0.8);}#mermaid-svg-fFFdrfunSGha2OEl .edgeLabel rect{opacity:0.5;background-color:rgba(232,232,232, 0.8);fill:rgba(232,232,232, 0.8);}#mermaid-svg-fFFdrfunSGha2OEl .labelBkg{background-color:rgba(232, 232, 232, 0.5);}#mermaid-svg-fFFdrfunSGha2OEl .cluster rect{fill:#ffffde;stroke:#aaaa33;stroke-width:1px;}#mermaid-svg-fFFdrfunSGha2OEl .cluster text{fill:#333;}#mermaid-svg-fFFdrfunSGha2OEl .cluster span{color:#333;}#mermaid-svg-fFFdrfunSGha2OEl div.mermaidTooltip{position:absolute;text-align:center;max-width:200px;padding:2px;font-family:"trebuchet ms",verdana,arial,sans-serif;font-size:12px;background:hsl(80, 100%, 96.2745098039%);border:1px solid #aaaa33;border-radius:2px;pointer-events:none;z-index:100;}#mermaid-svg-fFFdrfunSGha2OEl .flowchartTitleText{text-anchor:middle;font-size:18px;fill:#333;}#mermaid-svg-fFFdrfunSGha2OEl rect.text{fill:none;stroke-width:0;}#mermaid-svg-fFFdrfunSGha2OEl .icon-shape,#mermaid-svg-fFFdrfunSGha2OEl .image-shape{background-color:rgba(232,232,232, 0.8);text-align:center;}#mermaid-svg-fFFdrfunSGha2OEl .icon-shape p,#mermaid-svg-fFFdrfunSGha2OEl .image-shape p{background-color:rgba(232,232,232, 0.8);padding:2px;}#mermaid-svg-fFFdrfunSGha2OEl .icon-shape .label rect,#mermaid-svg-fFFdrfunSGha2OEl .image-shape .label rect{opacity:0.5;background-color:rgba(232,232,232, 0.8);fill:rgba(232,232,232, 0.8);}#mermaid-svg-fFFdrfunSGha2OEl .label-icon{display:inline-block;height:1em;overflow:visible;vertical-align:-0.125em;}#mermaid-svg-fFFdrfunSGha2OEl .node .label-icon path{fill:currentColor;stroke:revert;stroke-width:revert;}#mermaid-svg-fFFdrfunSGha2OEl :root{--mermaid-font-family:"trebuchet ms",verdana,arial,sans-serif;} ① 梳理模型配置

Zero / minimind-3 / MoE / 历史版本
② 梳理数据配置

pretrain_t2t(_mini) 等数据集
③ 选择运行命令

单卡 / 多卡 torchrun
④ 8GB 显卡上的工程现实

float16 / num_workers=0
⑤ 训练前诊断

test_gpu_diagnostics + 短跑
⑥ 显存构成分析

12N 参数 + 激活值
⑦ 实际执行配置

小 batch × 大累积
⑧ 初步 Eval

看 pretrain 模型的续写能力

一、先梳理模型配置

四个名字的准确定位:

  • MiniMind Zero :更像是一个训练目标或实践路线,而不是独立架构------"用 mini 数据集,从 0 快速训练出能基本对话的小模型"。MiniMind3 语境下最快复现 Zero 是 pretrain_t2t_mini.jsonl + sft_t2t_mini.jsonl;
  • minimind-3 :当前主线 Dense 模型,默认 hidden_size=768、num_hidden_layers=8,约 64M 参数;
  • minimind-3-moe :主干与 minimind-3 一致,FFN 换成 MoE,总参数约 198M,激活参数约 64M;
  • minimind2-small / minimind2:历史版本,主要作为对照。

1.1 常见模型版本表

版本 / 训练目标 参数量 关键结构参数 命令核心参数 适合场景
MiniMind Zero MiniMind3 下约 64M 通常用 minimind-3 Dense 结构 --hidden_size 768 --num_hidden_layers 8 --use_moe 0 最快从 0 跑通 pretrain + SFT
minimind-3 约 64M d_model=768,n_layers=8,q_heads=8,kv_heads=4 同上 当前主线,优先推荐
minimind-3-moe 约 198M-A64M 768×8,4 experts / top-1 routing --use_moe 1 观察 MoE 与激活参数概念
minimind2-small 约 26M d_model=512,n_layers=8 --hidden_size 512 --num_hidden_layers 8 更低成本的历史小配置
minimind2 约 104M d_model=768,n_layers=16 --hidden_size 768 --num_hidden_layers 16 历史主线对照

1.2 为什么 MiniMind3 默认 hidden_size=768

hidden_size 就是 Transformer 里的 d m o d e l d_{model} dmodel,决定每个 token 在模型内部被表示成多宽的向量。太小,表达能力受限------attention 中每个 head 的维度通常是:

d h e a d = d m o d e l n h e a d s d_{head} = \frac{d_{model}}{n_{heads}} dhead=nheadsdmodel

比如 hidden_size=512、num_attention_heads=8 时 d h e a d = 64 d_{head}=64 dhead=64;继续压小,head 内部可表达的信息就变少,embedding、attention head 维度、FFN 中间层会一起受影响。太大,训练成本明显上升。

MiniMind3 选 768×8 是折中:比 512 表示能力更好,比堆更多层更容易在个人 GPU 上快速训练。实践规则:
#mermaid-svg-Mme7vfHCmLprDjZq{font-family:"trebuchet ms",verdana,arial,sans-serif;font-size:16px;fill:#333;}@keyframes edge-animation-frame{from{stroke-dashoffset:0;}}@keyframes dash{to{stroke-dashoffset:0;}}#mermaid-svg-Mme7vfHCmLprDjZq .edge-animation-slow{stroke-dasharray:9,5!important;stroke-dashoffset:900;animation:dash 50s linear infinite;stroke-linecap:round;}#mermaid-svg-Mme7vfHCmLprDjZq .edge-animation-fast{stroke-dasharray:9,5!important;stroke-dashoffset:900;animation:dash 20s linear infinite;stroke-linecap:round;}#mermaid-svg-Mme7vfHCmLprDjZq .error-icon{fill:#552222;}#mermaid-svg-Mme7vfHCmLprDjZq .error-text{fill:#552222;stroke:#552222;}#mermaid-svg-Mme7vfHCmLprDjZq .edge-thickness-normal{stroke-width:1px;}#mermaid-svg-Mme7vfHCmLprDjZq .edge-thickness-thick{stroke-width:3.5px;}#mermaid-svg-Mme7vfHCmLprDjZq .edge-pattern-solid{stroke-dasharray:0;}#mermaid-svg-Mme7vfHCmLprDjZq .edge-thickness-invisible{stroke-width:0;fill:none;}#mermaid-svg-Mme7vfHCmLprDjZq .edge-pattern-dashed{stroke-dasharray:3;}#mermaid-svg-Mme7vfHCmLprDjZq .edge-pattern-dotted{stroke-dasharray:2;}#mermaid-svg-Mme7vfHCmLprDjZq .marker{fill:#333333;stroke:#333333;}#mermaid-svg-Mme7vfHCmLprDjZq .marker.cross{stroke:#333333;}#mermaid-svg-Mme7vfHCmLprDjZq svg{font-family:"trebuchet ms",verdana,arial,sans-serif;font-size:16px;}#mermaid-svg-Mme7vfHCmLprDjZq p{margin:0;}#mermaid-svg-Mme7vfHCmLprDjZq .label{font-family:"trebuchet ms",verdana,arial,sans-serif;color:#333;}#mermaid-svg-Mme7vfHCmLprDjZq .cluster-label text{fill:#333;}#mermaid-svg-Mme7vfHCmLprDjZq .cluster-label span{color:#333;}#mermaid-svg-Mme7vfHCmLprDjZq .cluster-label span p{background-color:transparent;}#mermaid-svg-Mme7vfHCmLprDjZq .label text,#mermaid-svg-Mme7vfHCmLprDjZq span{fill:#333;color:#333;}#mermaid-svg-Mme7vfHCmLprDjZq .node rect,#mermaid-svg-Mme7vfHCmLprDjZq .node circle,#mermaid-svg-Mme7vfHCmLprDjZq .node ellipse,#mermaid-svg-Mme7vfHCmLprDjZq .node polygon,#mermaid-svg-Mme7vfHCmLprDjZq .node path{fill:#ECECFF;stroke:#9370DB;stroke-width:1px;}#mermaid-svg-Mme7vfHCmLprDjZq .rough-node .label text,#mermaid-svg-Mme7vfHCmLprDjZq .node .label text,#mermaid-svg-Mme7vfHCmLprDjZq .image-shape .label,#mermaid-svg-Mme7vfHCmLprDjZq .icon-shape .label{text-anchor:middle;}#mermaid-svg-Mme7vfHCmLprDjZq .node .katex path{fill:#000;stroke:#000;stroke-width:1px;}#mermaid-svg-Mme7vfHCmLprDjZq .rough-node .label,#mermaid-svg-Mme7vfHCmLprDjZq .node .label,#mermaid-svg-Mme7vfHCmLprDjZq .image-shape .label,#mermaid-svg-Mme7vfHCmLprDjZq .icon-shape .label{text-align:center;}#mermaid-svg-Mme7vfHCmLprDjZq .node.clickable{cursor:pointer;}#mermaid-svg-Mme7vfHCmLprDjZq .root .anchor path{fill:#333333!important;stroke-width:0;stroke:#333333;}#mermaid-svg-Mme7vfHCmLprDjZq .arrowheadPath{fill:#333333;}#mermaid-svg-Mme7vfHCmLprDjZq .edgePath .path{stroke:#333333;stroke-width:2.0px;}#mermaid-svg-Mme7vfHCmLprDjZq .flowchart-link{stroke:#333333;fill:none;}#mermaid-svg-Mme7vfHCmLprDjZq .edgeLabel{background-color:rgba(232,232,232, 0.8);text-align:center;}#mermaid-svg-Mme7vfHCmLprDjZq .edgeLabel p{background-color:rgba(232,232,232, 0.8);}#mermaid-svg-Mme7vfHCmLprDjZq .edgeLabel rect{opacity:0.5;background-color:rgba(232,232,232, 0.8);fill:rgba(232,232,232, 0.8);}#mermaid-svg-Mme7vfHCmLprDjZq .labelBkg{background-color:rgba(232, 232, 232, 0.5);}#mermaid-svg-Mme7vfHCmLprDjZq .cluster rect{fill:#ffffde;stroke:#aaaa33;stroke-width:1px;}#mermaid-svg-Mme7vfHCmLprDjZq .cluster text{fill:#333;}#mermaid-svg-Mme7vfHCmLprDjZq .cluster span{color:#333;}#mermaid-svg-Mme7vfHCmLprDjZq div.mermaidTooltip{position:absolute;text-align:center;max-width:200px;padding:2px;font-family:"trebuchet ms",verdana,arial,sans-serif;font-size:12px;background:hsl(80, 100%, 96.2745098039%);border:1px solid #aaaa33;border-radius:2px;pointer-events:none;z-index:100;}#mermaid-svg-Mme7vfHCmLprDjZq .flowchartTitleText{text-anchor:middle;font-size:18px;fill:#333;}#mermaid-svg-Mme7vfHCmLprDjZq rect.text{fill:none;stroke-width:0;}#mermaid-svg-Mme7vfHCmLprDjZq .icon-shape,#mermaid-svg-Mme7vfHCmLprDjZq .image-shape{background-color:rgba(232,232,232, 0.8);text-align:center;}#mermaid-svg-Mme7vfHCmLprDjZq .icon-shape p,#mermaid-svg-Mme7vfHCmLprDjZq .image-shape p{background-color:rgba(232,232,232, 0.8);padding:2px;}#mermaid-svg-Mme7vfHCmLprDjZq .icon-shape .label rect,#mermaid-svg-Mme7vfHCmLprDjZq .image-shape .label rect{opacity:0.5;background-color:rgba(232,232,232, 0.8);fill:rgba(232,232,232, 0.8);}#mermaid-svg-Mme7vfHCmLprDjZq .label-icon{display:inline-block;height:1em;overflow:visible;vertical-align:-0.125em;}#mermaid-svg-Mme7vfHCmLprDjZq .node .label-icon path{fill:currentColor;stroke:revert;stroke-width:revert;}#mermaid-svg-Mme7vfHCmLprDjZq :root{--mermaid-font-family:"trebuchet ms",verdana,arial,sans-serif;} 快速复现
更省资源
观察 MoE
想训练哪个版本?
hidden_size=768, layers=8

Dense (minimind-3)
hidden_size=512, layers=8

历史 small 配置
768×8 不变

只加 --use_moe 1

二、再梳理数据配置

新版 MiniMind3 默认数据已从早期的 pretrain_hq.jsonl 切换到 pretrain_t2t(_mini).jsonl 和 sft_t2t(_mini).jsonl。下载地址:ModelScope / HuggingFace,下载后统一放到 ../dataset 目录(与训练脚本默认参数、上游习惯一致)。

2.1 数据集用途表

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

只想先把 pretrain 跑起来,最小只需要 pretrain_t2t_mini.jsonl;要完成 Zero 风格对话模型,再加 sft_t2t_mini.jsonl。

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

把 README 推荐与脚本默认值合在一起看:

目标 数据组合 阶段 epoch 关键训练参数 说明
Zero / 快速复现 pretrain_t2t_mini + sft_t2t_mini Pretrain 1 bs=32, lr=5e-4, accum=8, seq≈768 快速得到 base 权重
Zero / 快速复现 同上 Full SFT 1 bs=16, lr=1e-5, accum=1, seq≈768 在 pretrain 权重上继续训
Dense 主线 pretrain_t2t + sft_t2t Pretrain 2 bs=32, lr=5e-4, accum=8, seq≈380 接近完整主线训练
Dense 主线 同上 Full SFT 2 bs=16, lr=1e-5, accum=1, seq=768 主线 SFT
MoE 可沿用上述组合 Pretrain / SFT 同 Dense 加 --use_moe 1 总参数大、激活约同 Dense,更慢
DPO 对齐 dpo.jsonl DPO 1 bs=4, lr=4e-8, beta=0.15, seq=1024 SFT 之后的阶段,非必需

上游说单卡 3090 上 1 epoch 约 2.31h,指的是 mini 两阶段合计的 Zero 快速路线(MoE 约 3.23h)。这些数字只用于估算门槛------不同 GPU、数据加载速度、max_seq_len、batch size 都会影响耗时。

第一次实践的建议:先用 pretrain_t2t_mini, epochs=1, seq=768 + sft_t2t_mini, epochs=1, seq=768 把完整链路跑通,再切主线数据或开 MoE。

2.3 mini 和主线数据怎么选

  • 先学习代码、验证环境 → pretrain_t2t_mini.jsonl,目标是短时间内完整跑通 数据读取 → tokenizer → forward → loss → backward → checkpoint → 评估;
  • 想认真训练主线模型 → pretrain_t2t.jsonl,数据量更大,训练更久,适合和上游 README 的结果对齐;
  • 显存/时间有限 → 先不要急着上完整数据。

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

MiniMind tokenizer 在中文上大致 1 token ≈ 1.5~1.7 个中文字符,英文压缩比更高。max_seq_len 是在两个问题间做平衡:太短,长样本被截断、语义不完整;太长,短样本 padding 浪费算力。

上游经验:pretrain_t2t_mini 可设长一些(≈768);pretrain_t2t 主线可设均衡些(≈380)。当前脚本默认 --max_seq_len 340------快速跑通用默认即可,用 mini 数据且想减少截断就手动改 768。

三、运行命令

命令默认从项目根目录执行(Windows 下把 \ 换成 ^ 续行)。

3.1 快速训练 MiniMind Zero 的 pretrain 阶段

bash 复制代码
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-3.4 其他三条命令的差异点

命令 与 Zero 路线的差异 产出
Dense 主线(3.2) 数据换 pretrain_t2t.jsonl,seq=380, epochs=2 pretrain_768.pth
MoE(3.3) --use_moe 1 pretrain_768_moe.pth
minimind2-small 风格(3.4) hidden_size=512, seq=512 pretrain_512.pth

MoE 总参数更多、训练更慢,适合观察"总参数量 vs 激活参数量"的区别,而不是最快跑通流程。

3.5 多卡训练

bash 复制代码
uv run torchrun --nproc_per_node 2 -m minimind_learning.trainer.train_pretrain \
  --data_path ../dataset/pretrain_t2t_mini.jsonl ... # 其余参数同单卡

--nproc_per_node 2 启动 2 个进程(对应 2 张 GPU)。注意 batch_size 是每个进程 上的 batch size,总 batch 会随 GPU 数变大;要保持训练动态接近单卡,有时需同步调小 batch_size 或 accumulation_steps。

3.6 可选参数

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

3.7 训练输入输出一览

  • 输入:--data_path(数据)、tokenizer 默认从 ../model 加载;
  • 输出:../out/<save_weight>_<hidden>.pth(半精度推理权重)、../checkpoints/..._resume.pth(模型+优化器+进度的完整续训状态,见上一篇笔记 #11 的 checkpoint 分析)。

四、实验结果:8GB 显卡上的工程现实

实际训练使用一块 RTX 3070 Laptop GPU,显存只有 8GB。几轮测试后发现,除了模型结构,两个训练参数非常关键:

① 混合精度类型 。对这块消费级显卡,float16 比默认的 bfloat16 更合适------bfloat16 在新卡或数据中心卡上很好用,但 30 系消费级卡上不一定最优;实测 float16 训练路径更顺。

② num_workers 。很容易以为 worker 越多数据加载越快,但这台机器显存和内存都紧张、batch 不大,CPU→GPU 拷贝并不是瓶颈;反而 worker 数太大后,每个 worker 都要启动自己的进程并持有 Dataset/tokenizer 状态(Windows 上进程启动开销更明显 )。实测 num_workers=0 反而更稳定、启动更快。

所以这台机器上 pretrain 优先加:

bash 复制代码
--dtype float16
--num_workers 0

4.1 训练代码补充的诊断指标

训练不只是看 loss 是否下降,小显存 GPU 上还要观察吞吐、梯度和真实 token 数:

  • tokens_seen :实际参与 loss 计算的 token 数(排除 label 中 -100 的 padding/忽略位置)。比 step 更稳定,因为不同 batch 有效长度不同;
  • tokens/s:观察吞吐------batch size、seq_len、num_workers 或 dtype 改动后若明显下降,说明这组参数可能不是最优;
  • grad_norm :clip_grad_norm_ 返回的是裁剪前的全局梯度范数。长时间很大或突现尖峰,说明学习率/batch/数据分布/混合精度状态可能有问题;异常接近 0,可能梯度传播不充分。
python 复制代码
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)

wandb log 里除 loss 外还记录 logits_loss / aux_loss / learning_rate / tokens_seen / tokens_per_second / grad_norm / optimizer_step------这些指标不改变训练行为,只是让训练过程更透明,帮助在正式长跑前判断参数是否值得继续。

五、测试的重要性:不要等完整训练跑起来才发现参数不合适

训练一旦跑起来,可能几分钟后才因为显存不够、数据加载过慢、checkpoint 太大或 CUDA 状态异常而失败,每次从头试非常浪费时间。更好的方式是正式实验前先做一个很小的诊断轮,只回答几个工程问题:

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

理论上可以估算显存,但实际还有 CUDA cache、optimizer state、临时张量、DataLoader、后台进程占用------最后还是要在自己机器上实际跑一下,而且最好留出余量。
#mermaid-svg-laNOLnxm2z9Ss8ok{font-family:"trebuchet ms",verdana,arial,sans-serif;font-size:16px;fill:#333;}@keyframes edge-animation-frame{from{stroke-dashoffset:0;}}@keyframes dash{to{stroke-dashoffset:0;}}#mermaid-svg-laNOLnxm2z9Ss8ok .edge-animation-slow{stroke-dasharray:9,5!important;stroke-dashoffset:900;animation:dash 50s linear infinite;stroke-linecap:round;}#mermaid-svg-laNOLnxm2z9Ss8ok .edge-animation-fast{stroke-dasharray:9,5!important;stroke-dashoffset:900;animation:dash 20s linear infinite;stroke-linecap:round;}#mermaid-svg-laNOLnxm2z9Ss8ok .error-icon{fill:#552222;}#mermaid-svg-laNOLnxm2z9Ss8ok .error-text{fill:#552222;stroke:#552222;}#mermaid-svg-laNOLnxm2z9Ss8ok .edge-thickness-normal{stroke-width:1px;}#mermaid-svg-laNOLnxm2z9Ss8ok .edge-thickness-thick{stroke-width:3.5px;}#mermaid-svg-laNOLnxm2z9Ss8ok .edge-pattern-solid{stroke-dasharray:0;}#mermaid-svg-laNOLnxm2z9Ss8ok .edge-thickness-invisible{stroke-width:0;fill:none;}#mermaid-svg-laNOLnxm2z9Ss8ok .edge-pattern-dashed{stroke-dasharray:3;}#mermaid-svg-laNOLnxm2z9Ss8ok .edge-pattern-dotted{stroke-dasharray:2;}#mermaid-svg-laNOLnxm2z9Ss8ok .marker{fill:#333333;stroke:#333333;}#mermaid-svg-laNOLnxm2z9Ss8ok .marker.cross{stroke:#333333;}#mermaid-svg-laNOLnxm2z9Ss8ok svg{font-family:"trebuchet ms",verdana,arial,sans-serif;font-size:16px;}#mermaid-svg-laNOLnxm2z9Ss8ok p{margin:0;}#mermaid-svg-laNOLnxm2z9Ss8ok .label{font-family:"trebuchet ms",verdana,arial,sans-serif;color:#333;}#mermaid-svg-laNOLnxm2z9Ss8ok .cluster-label text{fill:#333;}#mermaid-svg-laNOLnxm2z9Ss8ok .cluster-label span{color:#333;}#mermaid-svg-laNOLnxm2z9Ss8ok .cluster-label span p{background-color:transparent;}#mermaid-svg-laNOLnxm2z9Ss8ok .label text,#mermaid-svg-laNOLnxm2z9Ss8ok span{fill:#333;color:#333;}#mermaid-svg-laNOLnxm2z9Ss8ok .node rect,#mermaid-svg-laNOLnxm2z9Ss8ok .node circle,#mermaid-svg-laNOLnxm2z9Ss8ok .node ellipse,#mermaid-svg-laNOLnxm2z9Ss8ok .node polygon,#mermaid-svg-laNOLnxm2z9Ss8ok .node path{fill:#ECECFF;stroke:#9370DB;stroke-width:1px;}#mermaid-svg-laNOLnxm2z9Ss8ok .rough-node .label text,#mermaid-svg-laNOLnxm2z9Ss8ok .node .label text,#mermaid-svg-laNOLnxm2z9Ss8ok .image-shape .label,#mermaid-svg-laNOLnxm2z9Ss8ok .icon-shape .label{text-anchor:middle;}#mermaid-svg-laNOLnxm2z9Ss8ok .node .katex path{fill:#000;stroke:#000;stroke-width:1px;}#mermaid-svg-laNOLnxm2z9Ss8ok .rough-node .label,#mermaid-svg-laNOLnxm2z9Ss8ok .node .label,#mermaid-svg-laNOLnxm2z9Ss8ok .image-shape .label,#mermaid-svg-laNOLnxm2z9Ss8ok .icon-shape .label{text-align:center;}#mermaid-svg-laNOLnxm2z9Ss8ok .node.clickable{cursor:pointer;}#mermaid-svg-laNOLnxm2z9Ss8ok .root .anchor path{fill:#333333!important;stroke-width:0;stroke:#333333;}#mermaid-svg-laNOLnxm2z9Ss8ok .arrowheadPath{fill:#333333;}#mermaid-svg-laNOLnxm2z9Ss8ok .edgePath .path{stroke:#333333;stroke-width:2.0px;}#mermaid-svg-laNOLnxm2z9Ss8ok .flowchart-link{stroke:#333333;fill:none;}#mermaid-svg-laNOLnxm2z9Ss8ok .edgeLabel{background-color:rgba(232,232,232, 0.8);text-align:center;}#mermaid-svg-laNOLnxm2z9Ss8ok .edgeLabel p{background-color:rgba(232,232,232, 0.8);}#mermaid-svg-laNOLnxm2z9Ss8ok .edgeLabel rect{opacity:0.5;background-color:rgba(232,232,232, 0.8);fill:rgba(232,232,232, 0.8);}#mermaid-svg-laNOLnxm2z9Ss8ok .labelBkg{background-color:rgba(232, 232, 232, 0.5);}#mermaid-svg-laNOLnxm2z9Ss8ok .cluster rect{fill:#ffffde;stroke:#aaaa33;stroke-width:1px;}#mermaid-svg-laNOLnxm2z9Ss8ok .cluster text{fill:#333;}#mermaid-svg-laNOLnxm2z9Ss8ok .cluster span{color:#333;}#mermaid-svg-laNOLnxm2z9Ss8ok div.mermaidTooltip{position:absolute;text-align:center;max-width:200px;padding:2px;font-family:"trebuchet ms",verdana,arial,sans-serif;font-size:12px;background:hsl(80, 100%, 96.2745098039%);border:1px solid #aaaa33;border-radius:2px;pointer-events:none;z-index:100;}#mermaid-svg-laNOLnxm2z9Ss8ok .flowchartTitleText{text-anchor:middle;font-size:18px;fill:#333;}#mermaid-svg-laNOLnxm2z9Ss8ok rect.text{fill:none;stroke-width:0;}#mermaid-svg-laNOLnxm2z9Ss8ok .icon-shape,#mermaid-svg-laNOLnxm2z9Ss8ok .image-shape{background-color:rgba(232,232,232, 0.8);text-align:center;}#mermaid-svg-laNOLnxm2z9Ss8ok .icon-shape p,#mermaid-svg-laNOLnxm2z9Ss8ok .image-shape p{background-color:rgba(232,232,232, 0.8);padding:2px;}#mermaid-svg-laNOLnxm2z9Ss8ok .icon-shape .label rect,#mermaid-svg-laNOLnxm2z9Ss8ok .image-shape .label rect{opacity:0.5;background-color:rgba(232,232,232, 0.8);fill:rgba(232,232,232, 0.8);}#mermaid-svg-laNOLnxm2z9Ss8ok .label-icon{display:inline-block;height:1em;overflow:visible;vertical-align:-0.125em;}#mermaid-svg-laNOLnxm2z9Ss8ok .node .label-icon path{fill:currentColor;stroke:revert;stroke-width:revert;}#mermaid-svg-laNOLnxm2z9Ss8ok :root{--mermaid-font-family:"trebuchet ms",verdana,arial,sans-serif;} 是
否
理论估算

(参数量 / 激活值)
test_gpu_diagnostics.py

构造小 jsonl, 跑若干 micro-step

输出耗时/显存/tokens-per-s
真实 pretrain 短跑

--max_train_steps 32 --profile_train

打印每步耗时和显存
显存余量足够?
正式长跑
降 batch / 升 accumulation

回到 B 重测

5.1 第一层:test_gpu_diagnostics.py

更像一个"训练前分析器"而非普通单元测试:默认跳过,显式设置 RUN_GPU_DIAGNOSTICS=1 才运行。核心循环与真实训练完全同构(autocast → scale → backward → unscale_ → clip → step → update):

python 复制代码
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)

运行方式(通过环境变量控制配置):

powershell 复制代码
$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

一次 768×8, seq=768, bs=16 的 summary 输出:

json 复制代码
{
  "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 或显存波动时,就可能接近边界。

5.2 第二层:真实 pretrain 诊断短跑

给 train_pretrain.py 加了三个只用于诊断的参数:--max_train_steps(只跑 N 个 micro-step)、--profile_train、--profile_interval。不把诊断结果存进 checkpoint,只在训练时打印每步耗时和显存:

python 复制代码
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
    step_time = time.time() - step_start_time
    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, tokens/s={tokens_per_second:.1f}, "
           f"optimizer_step={did_optimizer_step}, "
           f"cuda_alloc={allocated_mb:.1f}MB, cuda_reserved={reserved_mb:.1f}MB, "
           f"cuda_peak={peak_mb:.1f}MB")

短跑命令(--max_train_steps 32,配 accumulation_steps=16 即只做 2 次 optimizer step):

bash 复制代码
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

结果与诊断测试很接近:

text 复制代码
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, bs=16 能跑,但峰值 reserved 已到 6.8GB 左右------对 8GB 显卡属于"可以跑,但需要留意后台占用"。

六、显存占用的理论值

MiniMind3 Dense 约 64M 参数,fp16 下权重本身:

64 × 10 6 × 2 bytes ≈ 128 MB 64 \times 10^6 \times 2\ \text{bytes} \approx 128\ \text{MB} 64×106×2 bytes≈128 MB

离 8GB 很远,但训练显存不只存权重。以 AdamW + FP16 为例,参数量 N N N 至少还有:

组成 大小 说明
权重 2 N 2N 2N FP16
梯度 2 N 2N 2N 通常也是 FP16
AdamW 一阶动量 exp_avg 4 N 4N 4N 通常 FP32
AdamW 二阶动量 exp_sq_avg 4 N 4N 4N 通常 FP32

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

相对 FP16 权重的 2 N 2N 2N,大约是 6 倍 ------不能只拿权重大小判断训练显存。但这次实测中,参数相关的常驻显存并不是最大头,更明显的增长来自激活值 :激活随 batch size 和 seq_len 增长,单个隐藏状态张量 [B, L, D] 在 fp16 下约 B × L × D × 2 B \times L \times D \times 2 B×L×D×2 bytes,例如 B=16, L=768, D=768:

16 × 768 × 768 × 2 ≈ 18 MB 16 \times 768 \times 768 \times 2 \approx 18\ \text{MB} 16×768×768×2≈18 MB

单项不夸张,但每一层都要为反向传播保留中间激活 ,attention、MLP、loss 还会产生临时张量,最终远大于单项估算。尤其 batch size 是最直接的放大器:8→16 显存明显上台阶,16→32 在 8GB 卡上就很容易接近边界。
#mermaid-svg-ZGr9nbP0CQGqAshE{font-family:"trebuchet ms",verdana,arial,sans-serif;font-size:16px;fill:#333;}@keyframes edge-animation-frame{from{stroke-dashoffset:0;}}@keyframes dash{to{stroke-dashoffset:0;}}#mermaid-svg-ZGr9nbP0CQGqAshE .edge-animation-slow{stroke-dasharray:9,5!important;stroke-dashoffset:900;animation:dash 50s linear infinite;stroke-linecap:round;}#mermaid-svg-ZGr9nbP0CQGqAshE .edge-animation-fast{stroke-dasharray:9,5!important;stroke-dashoffset:900;animation:dash 20s linear infinite;stroke-linecap:round;}#mermaid-svg-ZGr9nbP0CQGqAshE .error-icon{fill:#552222;}#mermaid-svg-ZGr9nbP0CQGqAshE .error-text{fill:#552222;stroke:#552222;}#mermaid-svg-ZGr9nbP0CQGqAshE .edge-thickness-normal{stroke-width:1px;}#mermaid-svg-ZGr9nbP0CQGqAshE .edge-thickness-thick{stroke-width:3.5px;}#mermaid-svg-ZGr9nbP0CQGqAshE .edge-pattern-solid{stroke-dasharray:0;}#mermaid-svg-ZGr9nbP0CQGqAshE .edge-thickness-invisible{stroke-width:0;fill:none;}#mermaid-svg-ZGr9nbP0CQGqAshE .edge-pattern-dashed{stroke-dasharray:3;}#mermaid-svg-ZGr9nbP0CQGqAshE .edge-pattern-dotted{stroke-dasharray:2;}#mermaid-svg-ZGr9nbP0CQGqAshE .marker{fill:#333333;stroke:#333333;}#mermaid-svg-ZGr9nbP0CQGqAshE .marker.cross{stroke:#333333;}#mermaid-svg-ZGr9nbP0CQGqAshE svg{font-family:"trebuchet ms",verdana,arial,sans-serif;font-size:16px;}#mermaid-svg-ZGr9nbP0CQGqAshE p{margin:0;}#mermaid-svg-ZGr9nbP0CQGqAshE .label{font-family:"trebuchet ms",verdana,arial,sans-serif;color:#333;}#mermaid-svg-ZGr9nbP0CQGqAshE .cluster-label text{fill:#333;}#mermaid-svg-ZGr9nbP0CQGqAshE .cluster-label span{color:#333;}#mermaid-svg-ZGr9nbP0CQGqAshE .cluster-label span p{background-color:transparent;}#mermaid-svg-ZGr9nbP0CQGqAshE .label text,#mermaid-svg-ZGr9nbP0CQGqAshE span{fill:#333;color:#333;}#mermaid-svg-ZGr9nbP0CQGqAshE .node rect,#mermaid-svg-ZGr9nbP0CQGqAshE .node circle,#mermaid-svg-ZGr9nbP0CQGqAshE .node ellipse,#mermaid-svg-ZGr9nbP0CQGqAshE .node polygon,#mermaid-svg-ZGr9nbP0CQGqAshE .node path{fill:#ECECFF;stroke:#9370DB;stroke-width:1px;}#mermaid-svg-ZGr9nbP0CQGqAshE .rough-node .label text,#mermaid-svg-ZGr9nbP0CQGqAshE .node .label text,#mermaid-svg-ZGr9nbP0CQGqAshE .image-shape .label,#mermaid-svg-ZGr9nbP0CQGqAshE .icon-shape .label{text-anchor:middle;}#mermaid-svg-ZGr9nbP0CQGqAshE .node .katex path{fill:#000;stroke:#000;stroke-width:1px;}#mermaid-svg-ZGr9nbP0CQGqAshE .rough-node .label,#mermaid-svg-ZGr9nbP0CQGqAshE .node .label,#mermaid-svg-ZGr9nbP0CQGqAshE .image-shape .label,#mermaid-svg-ZGr9nbP0CQGqAshE .icon-shape .label{text-align:center;}#mermaid-svg-ZGr9nbP0CQGqAshE .node.clickable{cursor:pointer;}#mermaid-svg-ZGr9nbP0CQGqAshE .root .anchor path{fill:#333333!important;stroke-width:0;stroke:#333333;}#mermaid-svg-ZGr9nbP0CQGqAshE .arrowheadPath{fill:#333333;}#mermaid-svg-ZGr9nbP0CQGqAshE .edgePath .path{stroke:#333333;stroke-width:2.0px;}#mermaid-svg-ZGr9nbP0CQGqAshE .flowchart-link{stroke:#333333;fill:none;}#mermaid-svg-ZGr9nbP0CQGqAshE .edgeLabel{background-color:rgba(232,232,232, 0.8);text-align:center;}#mermaid-svg-ZGr9nbP0CQGqAshE .edgeLabel p{background-color:rgba(232,232,232, 0.8);}#mermaid-svg-ZGr9nbP0CQGqAshE .edgeLabel rect{opacity:0.5;background-color:rgba(232,232,232, 0.8);fill:rgba(232,232,232, 0.8);}#mermaid-svg-ZGr9nbP0CQGqAshE .labelBkg{background-color:rgba(232, 232, 232, 0.5);}#mermaid-svg-ZGr9nbP0CQGqAshE .cluster rect{fill:#ffffde;stroke:#aaaa33;stroke-width:1px;}#mermaid-svg-ZGr9nbP0CQGqAshE .cluster text{fill:#333;}#mermaid-svg-ZGr9nbP0CQGqAshE .cluster span{color:#333;}#mermaid-svg-ZGr9nbP0CQGqAshE div.mermaidTooltip{position:absolute;text-align:center;max-width:200px;padding:2px;font-family:"trebuchet ms",verdana,arial,sans-serif;font-size:12px;background:hsl(80, 100%, 96.2745098039%);border:1px solid #aaaa33;border-radius:2px;pointer-events:none;z-index:100;}#mermaid-svg-ZGr9nbP0CQGqAshE .flowchartTitleText{text-anchor:middle;font-size:18px;fill:#333;}#mermaid-svg-ZGr9nbP0CQGqAshE rect.text{fill:none;stroke-width:0;}#mermaid-svg-ZGr9nbP0CQGqAshE .icon-shape,#mermaid-svg-ZGr9nbP0CQGqAshE .image-shape{background-color:rgba(232,232,232, 0.8);text-align:center;}#mermaid-svg-ZGr9nbP0CQGqAshE .icon-shape p,#mermaid-svg-ZGr9nbP0CQGqAshE .image-shape p{background-color:rgba(232,232,232, 0.8);padding:2px;}#mermaid-svg-ZGr9nbP0CQGqAshE .icon-shape .label rect,#mermaid-svg-ZGr9nbP0CQGqAshE .image-shape .label rect{opacity:0.5;background-color:rgba(232,232,232, 0.8);fill:rgba(232,232,232, 0.8);}#mermaid-svg-ZGr9nbP0CQGqAshE .label-icon{display:inline-block;height:1em;overflow:visible;vertical-align:-0.125em;}#mermaid-svg-ZGr9nbP0CQGqAshE .node .label-icon path{fill:currentColor;stroke:revert;stroke-width:revert;}#mermaid-svg-ZGr9nbP0CQGqAshE :root{--mermaid-font-family:"trebuchet ms",verdana,arial,sans-serif;} 合计 12N ≈ 权重6倍
batch 是最直接放大器
激活值 (随 B×L×层数)
每层保留中间激活
attention/MLP/loss 临时张量
参数相关常驻 (随 N)
权重 2N
梯度 2N
exp_avg 4N
exp_sq 4N
训练总显存

8GB 卡上激活才是大头

6.1 RTX 3070 Laptop 8GB 实测表

模型均为 MiniMind3 Dense(768×8, dtype=float16):

max_seq_len batch_size 诊断峰值显存 结论
512 1 约 0.62GB 很轻
512 4 约 1.30GB 很稳
512 8 约 2.19GB 很稳
512 16 约 3.90GB 稳
512 32 接近 7.96GB 不推荐,接近上限
768 8 约 3.03GB 推荐,比较稳
768 16 约 5.67~6.43GB 可用,但需留余量

6.2 allocated 和 reserved 的区别

seq=768, bs=16 在诊断测试中峰值 allocated 约 5.95GB,真实短跑峰值 reserved 约 6.79GB,两个数不一样是因为:

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

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

七、实际执行的配置

8GB 卡上训练 MiniMind3 Dense,推荐两个配置。

更稳的版本:

cmd 复制代码
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

稍微激进的版本 (bs=16, accum=16)------注意作者的实际教训:这个版本实际跑不了:因为同时开着浏览器刷 B 站,第一个 batch 还能跑,第二个 batch 显存到了 7.9GB,开始来回 swap,直接卡死了。

两个配置的 effective batch size 都和上游推荐保持一致:

8 × 768 × 32 = 196608 , 16 × 768 × 16 = 196608 8 \times 768 \times 32 = 196608, \qquad 16 \times 768 \times 16 = 196608 8×768×32=196608,16×768×16=196608

(单步 batch × seq_len × 累积步数,即每次参数更新的有效 token 数,与上游 32×8 的默认组合等价。)

结论:显存紧张时,可以把单步 batch size 降低,再用更大的 accumulation_steps 把有效 batch size 补回来------这是小显存 GPU 上很常用的做法。(梯度累积的原理见笔记 #11。)

八、初步 Eval:pretrain 刚结束的模型长什么样

用刚训练完的 pretrain_768.pth 做粗糙的生成测试。手动对话:

cmd 复制代码
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

批量评测则加 --eval_file .\eval_prompts.jsonl --eval-output-file .\eval_result.jsonl(每行一个 {"prompt": ...})。

因为这是 pretrain 阶段的模型,还没经过 SFT,它更适合观察"续写能力"和局部语言模式,而不是期待稳定的指令跟随。几个有代表性的输出:

text 复制代码
Input: '<|im_start|>北京市是'
Output: 中国的首都,位于中国的东部。北京市是中国的首都,位于中国的东部。
它的面积约为 1200 万平方公里,人口约为 120 万。北京市是一个著名的历史文化名城...

------续写是通顺的(数字纯属幻觉),说明模型确实学到了局部语言模式。

text 复制代码
💬: 推荐一些中国的美食
🤖️: 中国美食有许多,包括寿司、刺身、炸鸡、糖果、火锅、糯米饭等等...

------"中国美食包括寿司刺身"是典型的 pretrain 模型行为:语言流畅、事实混乱。

比较神奇的是,即使 pretrain 数据中没有对话性数据,SFT 之前模型已经能进行一些"勉强前言搭后语"的成语接龙。但这种能力很不稳定------更多说明模型学到了局部语言模式,而不是已经具备真正稳定的指令跟随能力;推理速度约 60+ tokens/s。

总结

主题 一句话要点
版本梳理 Zero 是路线不是架构;主线 = minimind-3(768×8, 64M);MoE 198M 总参/A64M 激活
hidden_size d h e a d = d m o d e l / n h e a d s d_{head}=d_{model}/n_{heads} dhead=dmodel/nheads;768 是表示能力 vs 训练成本的折中
数据 mini 跑流程,t2t 训主线;max_seq_len 是 token 数(中文 1 token≈1.5~1.7 字)
命令 结构参数 + 数据路径 + bs/accum/lr;多卡 torchrun 注意 per-process batch
8GB 实战 30 系消费卡选 float16;Windows 上 num_workers=0 反而更稳
诊断指标 tokens_seen/tokens/s 比 step 更稳;grad_norm 是裁剪前的范数
诊断轮 test_gpu_diagnostics + --max_train_steps 短跑,能较准确预测真实显存
显存构成 参数相关 12N(权重 6 倍);激活随 B×L 增长,小模型上才是大头;batch 是最直接放大器
allocated vs reserved 关注 reserved 与总显存余量;开着浏览器刷 B 站真的会把训练卡死
小显存策略 降单步 batch、用 accumulation_steps 补回有效 batch(8×32 = 16×16 = 196608)
pretrain 模型表现 续写通顺、事实幻觉、无稳定指令跟随------正常,等 SFT

Pretrain 篇章还剩两章:评估(4.eval ------train/val loss、perplexity、生成样例的系统分析方法**)和实践总结(5.practice)**,后续继续更新。

参考

相关推荐
YOLO数据集集合1 小时前
无人机视角行人与车辆检测数据集 | 无人机航拍 行人检测 车辆检测 智慧城市 公共安全9140期
人工智能·目标检测·无人机·智慧城市·车辆识别·无人机视角
Joecien1 小时前
【2026实测】百炼 CLI 托管 Agent 教程:bl managed-agent 配置校验、版本回滚与变更预演(附完整命令)
人工智能·git·阿里云·知识图谱·agi
jimmyleeee1 小时前
大模型安全之三十六:大模型数据管理----从投毒防御到偏见治理的完整框架
人工智能·安全
IT古董2 小时前
《FDE前沿部署工程师实战教程》29 - Enterprise AI Security:Agent安全体系设计
大数据·数据库·人工智能
m4Rk_2 小时前
【论文阅读】Agent 记忆机制(87):VizoMem——把文本历史转化为可检索的视觉记忆
论文阅读·人工智能·学习·开源·github
让学习成为一种生活方式2 小时前
啤酒花基因组与重测序--Nature Communications
人工智能·算法
2601_949950632 小时前
错题总是反复错?用练题簿在线刷题,把复习重点找出来
人工智能·小程序·刷题·练习·小程序推荐
lpfasd1232 小时前
GitHub非AI开源方向调研报告
人工智能·开源·github
知几蜗牛2 小时前
OLMo-core 3的token gerrymandering提醒:MoE路由要按时间窗验收
人工智能