☰
从零手写64M小模型:3块钱2小时跑通大语言模型训练全流程
2026/10/16 11:08:41 网站建设 项目流程

1. 为什么我要花时间折腾一个 64M 的小模型

大语言模型这两年卷得厉害,动不动就是几百亿参数、上千张卡、几百万美元的训练成本。作为一个普通开发者,我每次看到这类新闻都有一种"热闹是他们的,我什么也没有"的感觉。直到我刷到 MiniMind 这个项目,标题里那句"花3块钱、2小时,从零训练一个 64M 参数的大语言模型"直接把我钉在了椅子上。

先说清楚这个项目是什么。MiniMind 是一个从零手写实现的大语言模型训练项目,它把模型结构、分词器、预训练、监督微调这几个环节全部用原生 PyTorch 实现了一遍,没有依赖任何高层封装框架。整个模型只有约 64M 参数,训练数据用的是公开的小规模中文语料,在一张消费级显卡上跑两小时左右就能得到一个能对话的模型,电费折算下来大概三块钱。

它能做什么?能让你亲手跑通"数据准备→分词器训练→预训练→SFT微调→对话推理"这条完整链路,而不是像调 API 那样只看到输入输出。它解决的核心问题是:绝大多数人学大模型都停留在"用"的层面,对"训"这件事完全没有体感,而这个项目把门槛压到了一个人一台机器一个下午就能跨过去。

适合谁看?三类人。第一类是刚入门大模型、想搞懂训练流程到底长什么样的学生和转行者;第二类是想快速验证某个想法、但不想一上来就烧大钱的独立开发者;第三类是老工程师,想找个干净的小项目当教学素材或者二次开发的起点。如果你已经能熟练训 7B 模型,这篇对你价值有限,但如果你想搞清楚"最小可运行的大模型长什么样",那接着往下看。

我自己的背景是做过几年后端,大模型这块属于半路出家,读过不少论文但一直没真正从零训过。这次拿 MiniMind 练手,踩了不少坑,也攒了一些文档里不会写的经验,下面全部摊开讲。

2. 项目整体设计与思路拆解

2.1 为什么是 64M 这个量级

很多人第一反应是:64M 也太小了吧,能叫大语言模型吗?这个问题得拆开看。参数量决定的是模型的"容量上限",也就是它能记住多少东西、能建模多复杂的规律。64M 参数确实做不了知识问答,你问它"某年某事件发生了什么"它大概率胡扯,因为它根本没在足够大的语料上见过这些知识。

但训练流程本身和 7B、70B 模型是完全一样的。预训练阶段的自回归语言建模目标、SFT 阶段的指令跟随微调、损失函数、优化器、学习率调度、梯度累积、混合精度——这些环节一个都不少。换句话说,64M 是一个"麻雀虽小五脏俱全"的规模,它让你用最低的成本把整条工程链路走通,把每个环节的坑都踩一遍,而这些经验在放大到 7B 时依然适用。

从成本角度算一笔账。64M 参数的模型,用 FP16 存权重也就 128MB 左右,加上优化器状态(Adam 的动量和方差各一份)和梯度,显存占用大概在 1GB 出头。这意味着什么?一张 8GB 显存的消费级显卡绰绰有余,甚至 6GB 的老卡也能跑。训练数据如果用几 GB 的中文语料,单卡跑两小时左右就能过几轮。电费按 300W 功耗算,两小时 0.6 度电,按商业电价也就几毛到一块钱,加上显卡折旧和机器损耗,说"三块钱"是个很实在的估算。

2.2 为什么坚持从零手写而不调库

这是 MiniMind 最核心的设计取舍。市面上训练框架很多,有封装好的 Trainer,有各种自动化工具,为什么还要手写?

我的理解是,封装框架帮你屏蔽了细节,但代价是你不知道发生了什么。比如学习率调度,你调一个参数就完事了,但它是 warmup 多少步、按 cosine 还是 linear 衰减、衰减到多少,这些直接影响训练稳定性。再比如梯度累积,框架帮你做了,但你知道它是在 loss 上除累积步数还是在梯度上除吗?这两种做法在混合精度下结果是不一样的。

手写的好处是每一行你都得自己想清楚。模型结构里 attention 怎么算、位置编码用哪种、归一化放在哪、残差怎么连,全都要自己写。写完之后你对"一个大模型到底由什么组成"这件事的理解会完全不一样。我自己的感受是,读完十篇论文不如自己手撸一遍 forward 函数。

当然手写也有代价,就是代码量大、容易出 bug。所以 MiniMind 的定位很明确:它是一个教学和实验项目,不是生产级框架。你要拿它上生产那肯定不合适,但你要拿它学原理、做小规模实验,它非常合适。

2.3 训练流程的分段设计

整个项目把训练拆成了几个清晰的阶段,这个拆分逻辑值得说一下。

第一阶段是分词器训练。大模型不能直接处理文字,得先把文本切成 token。分词器决定了词表大小,词表大小又直接影响 embedding 层的参数量和模型的表达能力。MiniMind 用的是自己训练的 BPE 分词器,词表规模控制在几万这个量级。为什么不直接用现成的分词器?因为现成的分词器往往是为大语料、大词表设计的,小模型用小词表更划算,而且自己训一遍能让你理解 BPE 到底在干什么。

第二阶段是预训练。这是最耗资源的一步,目标是让模型学会"预测下一个 token"这个基本能力。预训练完之后模型其实已经能生成通顺的文本了,但它不会听指令,你问它问题它可能接着你的问题往下编。

第三阶段是 SFT,也就是监督微调。用一批"指令-回答"格式的数据继续训练,让模型学会"看到问题要回答"这个模式。这一步数据量不用很大,几千到几万条就有效果,但数据质量很关键。

第四阶段是推理和对话。把训练好的权重加载进来,实现一个简单的对话循环,支持多轮上下文。

这个分段设计的好处是每一步都能单独验证。预训练完先看看生成效果,不行就回去调;SFT 完再看看指令跟随能力,不行就换数据。如果全揉在一起,出了问题你根本不知道是哪一环的锅。

3. 核心细节解析与实操要点

3.1 模型结构里几个关键选择

先看模型主体。MiniMind 用的是标准的 decoder-only Transformer,也就是 GPT 那一套。这个选择没什么悬念,因为现在主流的大语言模型基本都是这个架构,学它最有迁移价值。

具体到几个细节。归一化用的是 RMSNorm 而不是 LayerNorm,这是近几年的主流做法,好处是省掉了均值计算,速度略快,效果基本持平。激活函数用的是 SwiGLU,比 ReLU 和 GELU 效果更好,代价是参数量会多一点,因为 SwiGLU 有三个投影矩阵而不是两个。位置编码用的是 RoPE,旋转位置编码,它的优势是能更好地处理长序列,而且实现上比较优雅。

注意力这块,小模型其实用不着分组查询注意力或者多头潜在注意力这些省显存的花活,标准的多头注意力就够了。但如果你想练手,可以自己改成 GQA 试试,看看显存和速度有什么变化。

注意:RMSNorm 和 LayerNorm 的一个关键区别是 RMSNorm 不做去均值,只做缩放。如果你从 LayerNorm 改过来,记得把 bias 参数去掉,RMSNorm 一般不带 bias。

3.2 分词器训练的实操细节

分词器这块我踩的坑最多。BPE 训练的核心参数是词表大小,这个值怎么定?太小了,一个汉字可能被切成好几个 token,序列变长,训练和推理都变慢;太大了,embedding 层参数暴涨,小模型扛不住。

我的经验是,纯中文语料的话,词表设在 3 万到 5 万之间比较合适。MiniMind 默认用的是这个区间。你可以做个简单测试:拿一段典型文本,用不同词表的分词器切一下,看看平均每个汉字对应多少 token。理想情况是接近 1,如果超过 1.5 说明词表偏小。

还有一个坑是特殊 token。分词器需要预留一些特殊符号,比如 padding、unknown、bos、eos,还有对话场景下的角色标记。这些 token 的 id 一定要固定好,训练和推理时必须一致,否则模型会完全懵掉。我见过有人训练时用一套 id、推理时用另一套,结果模型输出全是乱码。

训练分词器的数据量不用太大,几百 MB 到 1GB 的干净文本就够了。关键是干净,要去掉 HTML 标签、乱码、重复内容。数据脏的话分词器会学出一堆奇怪的 token。

3.3 预训练的数据和超参

预训练数据我用的是公开的中文语料,清洗之后大概几个 GB。这里要强调一下数据清洗的重要性。原始语料里往往混着大量重复、广告、乱码,直接拿来训会让模型学到很多垃圾模式。我的做法是先做去重,用简单的哈希或者 MinHash 都行;然后过滤掉长度过短和过长的样本;最后做一遍基本的字符集过滤,把非中文非英文的乱码清掉。

超参方面,学习率是重中之重。小模型的学习率可以设得比大模型高一些,因为它的损失曲面相对简单。MiniMind 默认的学习率在 1e-3 到 3e-3 这个量级,配合 warmup 和 cosine 衰减。warmup 步数一般占总步数的 1% 到 5%,太少容易初期震荡,太多浪费训练时间。

批次大小受显存限制,小卡上可能只能设到 8 或者 16。这时候就要用梯度累积,比如累积 8 步,等效批次就是 128。梯度累积的实现要注意,loss 要除以累积步数再 backward,否则梯度会偏大。

混合精度训练基本是标配,用 FP16 或者 BF16 能省一半显存、快不少。BF16 的动态范围更大,不容易溢出,如果显卡支持优先用 BF16。

3.4 SFT 数据的构造要点

SFT 阶段的数据格式很关键。每条样本一般是一个"指令-回答"对,拼成模型能理解的格式。常见做法是用特殊 token 把不同角色隔开,比如用户说的话前面加个标记,模型回答前面加另一个标记,结尾加 eos。

数据质量比数量重要得多。我试过用几千条精心构造的数据,效果比几万条随便爬的数据好。什么叫精心构造?回答要准确、完整、格式统一,不能有前后矛盾,不能有事实错误。小模型本来就容易学歪,数据里有点噪声它就跟着学。

还有一个细节是 loss 的计算范围。SFT 时通常只对"回答"部分算 loss,不对"指令"部分算。因为指令是给定的输入,模型不需要学怎么生成指令。这个通过 mask 实现,把指令部分的 label 设成忽略值就行。

提示:SFT 数据里如果混入了预训练阶段的纯文本,会让模型困惑。建议 SFT 数据格式高度统一,不要一会儿是问答、一会儿是续写。

4. 实操过程与核心环节实现

4.1 环境准备与依赖安装

环境这块不复杂,核心就是 PyTorch 加几个辅助库。我的建议是直接用 conda 建一个干净的环境,避免和系统里的其他包冲突。

conda create -n minimind python=3.10 conda activate minimind pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu118 pip install transformers datasets tokenizers numpy tqdm

PyTorch 版本根据你的显卡驱动选,CUDA 11.8 或者 12.1 都行。transformers 主要用来加载分词器做对比,datasets 用来处理数据,tokenizers 用来训 BPE。

装完之后跑一句torch.cuda.is_available()确认显卡能用。如果返回 False,检查驱动和 CUDA 版本是否匹配。这一步别偷懒,我见过有人训了半天发现一直在用 CPU,速度慢了几十倍。

4.2 数据准备与分词器训练

数据准备分两步。第一步是把原始语料整理成纯文本,一行一条或者一段一条都行。第二步是训分词器。

from tokenizers import Tokenizer, models, trainers, pre_tokenizers tokenizer = Tokenizer(models.BPE()) tokenizer.pre_tokenizer = pre_tokenizers.ByteLevel(add_prefix_space=False) trainer = trainers.BpeTrainer( vocab_size=32000, special_tokens=["<pad>", "<unk>", "<bos>", "<eos>", "<user>", "<assistant>"], min_frequency=2 ) tokenizer.train(files=["corpus.txt"], trainer=trainer) tokenizer.save("tokenizer.json")

词表大小我设的 32000,特殊 token 里加了对话用的角色标记。min_frequency 设 2 是为了过滤掉只出现一次的词,减少词表里的噪声。

训完之后一定要验证一下。拿几句话编码再解码,看看能不能还原。再统计一下平均每条文本的 token 数,判断词表大小是否合适。

4.3 预训练的关键代码与参数

预训练的核心循环其实不长,但每个细节都要对。

model.train() optimizer = torch.optim.AdamW(model.parameters(), lr=3e-3, betas=(0.9, 0.95), weight_decay=0.1) scheduler = get_cosine_schedule_with_warmup(optimizer, warmup_steps=500, total_steps=total_steps) scaler = torch.cuda.amp.GradScaler() for step, batch in enumerate(dataloader): input_ids = batch["input_ids"].cuda() labels = input_ids.clone() with torch.cuda.amp.autocast(dtype=torch.bfloat16): outputs = model(input_ids, labels=labels) loss = outputs.loss / accumulation_steps scaler.scale(loss).backward() if (step + 1) % accumulation_steps == 0: scaler.unscale_(optimizer) torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0) scaler.step(optimizer) scaler.update() scheduler.step() optimizer.zero_grad()

几个关键点。梯度裁剪的 max_norm 设 1.0,防止梯度爆炸,小模型训练时这个很重要。AdamW 的 betas 第二个值设 0.95 而不是默认的 0.999,这是大模型训练的常见做法,能让优化器更快响应最近的梯度。weight_decay 设 0.1 做正则化。

学习率 3e-3 对小模型来说不算激进,但如果你发现 loss 一直震荡,可以降到 1e-3 试试。warmup 500 步是经验值,数据量大的话可以按比例增加。

训练过程中要盯着 loss 曲线。正常情况下 loss 应该平滑下降,如果出现尖刺说明学习率太大或者数据有问题,如果一直不降说明学习率太小或者模型结构有 bug。

4.4 SFT 微调与对话推理

SFT 的代码和预训练几乎一样,区别在数据格式和 loss mask。

def build_sft_sample(instruction, response, tokenizer, max_len=512): prompt = f"<user>{instruction}<assistant>" prompt_ids = tokenizer.encode(prompt) response_ids = tokenizer.encode(response) + [tokenizer.eos_token_id] input_ids = prompt_ids + response_ids labels = [-100] * len(prompt_ids) + response_ids if len(input_ids) > max_len: input_ids = input_ids[:max_len] labels = labels[:max_len] return {"input_ids": input_ids, "labels": labels}

-100 是 PyTorch 里 CrossEntropyLoss 默认忽略的标签值,这样指令部分就不参与 loss 计算。

SFT 的学习率要比预训练小,一般设 1e-4 到 5e-5,因为模型已经有一定基础了,太大的学习率会破坏预训练学到的知识。训练轮数也不用多,几轮就够,多了容易过拟合。

推理的时候实现一个简单的对话循环:

def chat(model, tokenizer, history, max_new_tokens=128): prompt = "" for user, assistant in history: prompt += f"<user>{user}<assistant>{assistant}" prompt += f"<user>{history[-1][0]}<assistant>" input_ids = torch.tensor([tokenizer.encode(prompt)]).cuda() output = model.generate(input_ids, max_new_tokens=max_new_tokens, temperature=0.7, top_p=0.9, do_sample=True) response = tokenizer.decode(output[0][input_ids.shape[1]:], skip_special_tokens=True) return response

temperature 和 top_p 控制生成的随机性。temperature 越高越随机,0.7 是个比较平衡的值。top_p 做核采样,0.9 表示只从累积概率前 90% 的 token 里采样。

5. 常见问题与排查技巧实录

5.1 训练不收敛的排查思路

这是最常见的问题,loss 不降或者降得很慢。排查顺序我总结成一张表。

现象可能原因排查方法
loss 完全不动学习率太小或为 0打印 optimizer 的 lr 确认
loss 剧烈震荡学习率太大降一个数量级重试
loss 先降后升过拟合或学习率衰减不当检查数据量、调整调度
loss 降但生成乱码分词器不匹配检查训练和推理的分词器是否一致
loss 正常但显存爆批次太大或没开混合精度减小批次、开启 AMP

我遇到过一次 loss 死活不降,查了半天发现是数据加载的时候把 input_ids 和 labels 搞反了,模型在学"根据答案预测问题",当然学不会。所以数据 pipeline 一定要单独验证,拿几条样本打印出来看看对不对。

5.2 显存不足的应对策略

小卡训练显存紧张是常态。按优先级排序,能省显存的手段有这些。

第一,开混合精度,BF16 或 FP16,能省将近一半。第二,减小批次大小,配合梯度累积保持等效批次。第三,用梯度检查点,牺牲一点速度换显存,能省不少。第四,缩短序列长度,但要注意别短到影响模型学上下文。第五,如果还不行,考虑用 LoRA 之类的参数高效微调方法,只训一小部分参数。

注意:梯度检查点会让反向传播重新计算前向,速度大概慢 20% 到 30%,但显存能省 50% 以上,显存紧张时非常值得。

5.3 生成质量差的调优方向

模型训完了但生成的东西不通顺或者答非所问,原因可能有很多。

如果是预训练阶段就不通顺,检查数据质量和训练步数。小模型需要足够的步数才能学会语言规律,步数不够就是没学会。如果预训练没问题但 SFT 后变差,多半是 SFT 数据质量不行或者学习率太大把预训练知识冲掉了。

如果是重复生成同一个词,这是小模型的常见毛病,可以试试调高 temperature 或者加重复惩罚。如果是答非所问,检查 SFT 数据的格式是否和推理时的 prompt 格式一致,格式不一致模型就不知道该怎么接。

还有一个容易被忽略的点是 eos token。如果模型不会在合适的时候输出 eos,它就会一直生成下去。检查 SFT 数据里每条回答结尾有没有加 eos,这个很关键。

5.4 几个文档里不会写的经验

第一,训练日志一定要记全。loss、学习率、梯度范数、吞吐量,这些都要打出来。出了问题回头看日志,比瞎猜快得多。

第二,小规模先跑通再放大。别一上来就用全量数据训,先用几千条数据跑几百步,确认流程没问题再上全量。这样能省大量时间。

第三,随机种子要固定。不然每次结果都不一样,你没法判断改动到底有没有效果。

第四,模型保存要定期做。训练中途崩了不至于从头再来。保存的时候把优化器状态也存上,方便断点续训。

第五,别迷信默认参数。默认参数是别人在别人的数据上试出来的,你的数据不一样,该调就得调。

6. 这个项目还能怎么玩

跑通基础流程之后,其实有很多可以扩展的方向。比如把模型结构改一改,试试 GQA 或者 MoE,看看效果和速度有什么变化。比如换不同的分词器,对比词表大小对训练的影响。比如在 SFT 数据上做文章,构造特定领域的指令数据,看看小模型能不能学会某个垂直场景。

我自己下一步打算试试把序列长度拉长,看看 RoPE 在长上下文下的表现,顺便验证一下位置编码的外推能力。另外想试试用不同的数据配比,看看中文和英文的比例对模型能力的影响。

这个项目最大的价值不在于训出来的模型有多强,而在于它把大模型训练这件事从"黑盒"变成了"白盒"。你亲手跑一遍,很多之前模糊的概念会一下子清晰起来。三块钱两小时,换来的是对整条链路的真实体感,这笔账怎么算都划算。

需要专业的网站建设服务?

联系我们获取免费的网站建设咨询和方案报价,让我们帮助您实现业务目标

立即咨询