1. 为什么要在单卡上折腾一个会循环思考的小模型1.1 从大就是好到小而会想的转向过去两年大家聊模型动辄就是千亿参数、万卡集群仿佛没有个几百张卡都不好意思说自己在训模型。但真正在一线做过落地的人心里都清楚绝大多数业务场景根本用不上那么大的模型反倒是能不能在一张卡上、一天之内跑出一个能用的东西这种需求更真实。我这次要聊的就是这么一个项目——一张 A100、8 小时从零训练一个会循环思考的小模型。先说清楚这个标题里的三个关键词。一张 A100是硬件约束80GB 显存算力天花板摆在那8 小时是时间约束意味着你不能搞那种跑三天的实验会循环思考是能力目标指的是模型在推理时能对同一个问题反复迭代、逐步修正自己的中间表示而不是一次性前向传播就出结果。这种能力在数学推理、代码生成、多步规划这类任务上特别有用因为这类问题的答案往往不是一眼看穿而是想几步才对。那为什么不用现成的大模型微调一下因为我想验证一个假设循环思考这种能力未必需要千亿参数才能涌现一个结构设计得当的小模型配合合适的训练策略也能学会多想几步。这个假设如果成立对边缘部署、对成本敏感的场景意义就很大了。1.2 这个项目适合谁来参考我把话说在前头这篇内容不是给纯小白看的手把手装环境教程也不是给大厂训练平台工程师看的分布式训练指南。它更适合这几类人一是手里有一张或几张卡、想认真做点模型结构实验的独立研究者二是想理解 Transformer、MoE、LoRA、SFT 这些概念怎么在一个真实项目里串起来的中级开发者三是想搞清楚循环思考到底是怎么回事、能不能抄到自己任务里的算法工程师。读完之后你应该能拿到这些东西一套完整的、可复现的训练流程对核心结构选型的取舍逻辑踩过的坑和对应的排查方法以及一份能直接改改就用的配置参考。我不会只告诉你怎么做更会告诉你为什么这么做以及这么做会踩什么坑。2. 整体设计思路小模型怎么才能会想2.1 循环思考的本质是什么先把概念掰开揉碎。循环思考这个词听起来玄乎其实核心思想很朴素让模型对同一个输入做多次处理每次处理都基于上一次的中间结果进行修正直到结果稳定或者达到预设的迭代次数。打个比方你解一道复杂的应用题不会看一眼就写答案而是先列个大概思路然后检查一遍发现某步算错了回头改再检查再改最后得出答案。模型也一样普通的前向传播相当于看一眼就答而循环思考相当于答完自己检查几遍。在结构上实现循环思考有几种常见路子。一种是把同一组 Transformer 层重复调用多次每次的输出作为下一次的输入这叫权重共享的循环另一种是显式地维护一个思考状态向量每一轮迭代都更新这个状态最后再解码成输出。我这次选的是第一种因为它实现简单、参数少而且天然适合小模型——毕竟你只有一张卡参数量得省着用。2.2 为什么选 Transformer MoE 的组合Transformer 是绕不开的这个没什么好纠结的。它的自注意力机制天生适合建模序列内的长距离依赖而且生态成熟各种优化实现随手就能拿到。真正需要动脑子的是怎么在有限参数下提升模型的思考容量。这里我引入了MoE混合专家架构。MoE 的核心思路是与其让所有参数对每个输入都参与计算不如把参数分成若干专家每个输入只激活其中一小部分。这样做的好处是总参数量可以做得比较大但每次前向传播的实际计算量激活参数量保持很小。对于单卡训练来说这意味着我可以在显存允许的范围内塞进更多参数同时不显著拖慢训练速度。具体到我的设计我用了 8 个专家每个 token 路由到 top-2 专家。这样总参数量大概是稠密模型的 4 倍左右但激活参数量只比稠密模型多一点点。路由用的是经典的 top-k gating加了一个负载均衡损失防止所有 token 都挤到同一个专家上——这个坑我后面会详细讲因为不加载荷均衡损失的话训练到一半你会发现 8 个专家里有 6 个基本没被激活过。2.3 LoRA 和 SFT 在流程里的位置有人可能会问你都是从零训练了还要 LoRA 干嘛这里要澄清一个常见误解从零训练和 LoRA 微调不是互斥的它们解决的是不同阶段的问题。我的流程是这样的先用从零训练的方式在一个较大的通用语料上把基础模型训出来让它具备基本的语言能力和初步的循环思考能力然后用SFT监督微调在高质量的指令数据上做对齐让模型学会按照思考-修正-输出的格式来回答问题最后如果你有特定领域的任务再用LoRA做轻量级适配。LoRA 的好处是只训练一小部分低秩矩阵显存占用小、训练快特别适合在已经训好的基础模型上做领域迁移。所以整个流程是从零预训练 → SFT 对齐 → LoRA 领域适配三段式每一段的目标和手段都不一样。8 小时的时间预算主要花在第一段后面两段加起来大概占 1 到 2 小时。3. 核心细节解析与实操要点3.1 模型结构的关键参数怎么定结构设计这块我踩过不少坑最后定下来的配置是经过反复权衡的。先看整体规模我最终选的配置是隐藏维度 512、12 层 Transformer、8 个注意力头、词表大小 32000。这个规模在 A100 上单卡训练完全 hold 得住而且 8 小时内能跑完我准备的数据量。为什么是 512 而不是 768 或 1024因为我要留出显存给 MoE 的专家参数和循环迭代带来的额外开销。循环思考意味着同一组层要被调用多次如果单层太宽迭代几次显存就爆了。512 是一个比较舒服的平衡点实测下来单次前向的激活显存大概在 20GB 左右留足了余量。循环迭代次数我设的是4 次。这个数字不是拍脑袋定的我做过消融实验迭代 1 次相当于普通模型在推理任务上准确率大概 42%2 次跳到 51%4 次到 58%8 次反而降到 55% 左右——过拟合了模型开始想太多把自己绕进去。所以 4 次是个甜点。MoE 那边专家数量 8、top-2 路由、专家隐藏维度跟主模型一致。这里有个细节专家的初始化很重要如果所有专家用同样的初始化训练初期路由会非常不稳定。我的做法是给每个专家的初始化加一点扰动让它们起点略有差异这样路由网络更容易分化出不同的专家分工。3.2 数据准备质量和多样性比数量更重要8 小时的时间预算决定了你不可能在数据量上堆太多。我准备的数据大概是20GB 左右的文本混合了通用语料、代码、数学题解和一部分多步推理的合成数据。这里的关键不是量大而是数据里得有思考过程的样本。什么意思如果你只喂模型问题-答案对它学不会循环思考因为它没见过思考过程长什么样。所以我专门构造了一批带中间步骤的数据格式大概是问题 → 第一步推理 → 第二步推理 → ... → 答案。模型在训练中会逐渐学会模仿这种多步模式推理时自然就会多迭代几轮。数据清洗这块我做了几件事去重用 MinHash 做近似去重阈值设 0.8、过滤过短和过长的样本保留 50 到 2048 token 之间的、以及用一个小分类器过滤掉质量明显偏低的内容。这些步骤听起来琐碎但实测下来对最终效果影响很大——脏数据训出来的模型循环思考时会胡思乱想迭代越多错得越离谱。3.3 训练超参的取舍逻辑超参这块我列个表把关键参数和选择理由说清楚参数取值选择理由学习率3e-4小模型常用值配合 warmup 稳定Warmup2000 步防止初期路由网络震荡Batch size512序列打包后单卡显存能承受的最大值序列长度2048兼顾长文本和显存优化器AdamW权重衰减 0.1稳定梯度裁剪1.0防止 MoE 路由梯度爆炸精度bf16A100 原生支持省显存这里重点说两个。学习率 3e-4 是配合 warmup 用的如果不加 warmup 直接上 3e-4MoE 的路由网络在前几百步会剧烈震荡专家分化不出来。梯度裁剪设 1.0 也是针对 MoE 的因为路由的 softmax 在某些 token 上会产生很大的梯度不裁剪的话训练很容易发散。还有一个容易被忽略的点权重衰减对 MoE 专家要单独处理。我试过对所有参数统一用 0.1 的权重衰减结果专家参数被压得太狠表达能力下降。后来改成专家参数用 0.01、其他参数用 0.1效果好很多。4. 实操过程与核心环节实现4.1 环境准备与依赖安装环境这块我不啰嗦直接给关键步骤。基础环境是 Python 3.10 PyTorch 2.1 CUDA 12.1这几个版本的组合在 A100 上最稳。装依赖的时候注意flash-attention 一定要装对应 CUDA 版本的预编译包自己编译容易出各种玄学问题。pip install torch2.1.0 --index-url https://download.pytorch.org/whl/cu121 pip install flash-attn --no-build-isolation pip install transformers datasets accelerate safetensors装完之后跑个简单的验证脚本确认 A100 被正确识别、bf16 可用、flash-attention 能正常调用。这一步别省我见过太多人训练跑了一半才发现精度没开对白白浪费几小时。4.2 模型代码的核心实现模型部分我拆成几个模块写。先是 MoE 层核心是路由和专家计算class MoELayer(nn.Module): def __init__(self, dim, num_experts8, top_k2): super().__init__() self.num_experts num_experts self.top_k top_k self.router nn.Linear(dim, num_experts, biasFalse) self.experts nn.ModuleList([ nn.Sequential( nn.Linear(dim, dim * 4), nn.GELU(), nn.Linear(dim * 4, dim) ) for _ in range(num_experts) ]) def forward(self, x): # x: [batch, seq, dim] gate_logits self.router(x) gate_probs F.softmax(gate_logits, dim-1) topk_probs, topk_idx gate_probs.topk(self.top_k, dim-1) topk_probs topk_probs / topk_probs.sum(dim-1, keepdimTrue) output torch.zeros_like(x) for i in range(self.num_experts): mask (topk_idx i).any(dim-1) if mask.any(): expert_out self.experts[i](x[mask]) weight topk_probs[mask][topk_idx[mask] i].unsqueeze(-1) output[mask] expert_out * weight return output这段代码看着简单但有几个坑。第一路由的 softmax 要在 float32 下算bf16 下数值精度不够会导致路由结果不稳定。第二负载均衡损失必须加否则专家会塌缩。负载均衡损失的计算方式是每个专家被激活的频率与均匀分布的 KL 散度乘以一个系数我用的是 0.01加到总损失里。循环思考的实现是在 Transformer 主干外面套一层循环class RecurrentTransformer(nn.Module): def __init__(self, num_layers12, dim512, num_iters4): super().__init__() self.layers nn.ModuleList([TransformerBlock(dim) for _ in range(num_layers)]) self.num_iters num_iters self.iter_embed nn.Embedding(num_iters, dim) # 迭代步数嵌入 def forward(self, x): for it in range(self.num_iters): x x self.iter_embed(torch.tensor(it, devicex.device)) for layer in self.layers: x layer(x) return x这里的关键设计是迭代步数嵌入。如果不加这个模型在每一轮迭代时看到的是完全相同的输入它没法区分我现在是第几轮思考。加上步数嵌入后模型能感知到迭代进度行为会更合理——早期迭代偏向探索后期迭代偏向收敛。4.3 训练循环与显存优化训练循环这块我用的是标准的 PyTorch 训练循环加梯度累积。因为单卡 batch size 有限我用了4 步梯度累积等效 batch size 到 512。梯度累积的代码很简单但要注意累积期间不要清零梯度只在累积满之后才 step 和 zero_grad。显存优化我做了几件事。一是激活检查点gradient checkpointing对每个 Transformer 块开启用计算换显存实测能省 40% 左右的激活显存。二是序列打包把多个短序列拼成一条长序列减少 padding 浪费。三是优化器状态用 8-bit 量化用 bitsandbytes 的 AdamW8bit优化器状态显存直接砍半。这几招下来原本需要 60GB 显存的配置压到了 35GB 左右A100 的 80GB 显存跑起来很从容还能留出空间做验证和 checkpoint 保存。4.4 SFT 和 LoRA 阶段的衔接预训练跑完之后模型已经具备了基本的循环思考能力但输出格式还很随意。SFT 阶段我用的是大概 5 万条高质量的指令数据格式统一成问题 → 思考过程 → 答案。训练 2 个 epoch学习率降到 1e-5其他超参跟预训练一致。LoRA 阶段是可选的。如果你有特定领域的需求比如想让模型专门处理法律文书或者医疗问答就在 SFT 模型上挂 LoRA 适配器。LoRA 的秩我设的是 16alpha 设 32只作用于注意力的 Q、V 投影和 MoE 的路由层。训练数据量不用太大几千条就够训练时间大概 20 分钟。5. 常见问题与排查技巧实录5.1 训练不收敛的几种典型表现训练不收敛是最常见的问题但表现不一样原因也不一样。我整理了一个速查表表现可能原因排查方法解决Loss 震荡不下降学习率过高打印每步 loss降到 1e-4 或加 warmupLoss 突然变 NaN梯度爆炸监控梯度范数梯度裁剪设 1.0专家激活不均缺负载均衡损失统计各专家激活率加负载均衡损失循环迭代无提升步数嵌入没生效对比不同迭代次数检查嵌入维度我重点说专家激活不均这个。训练到大概 5000 步的时候我发现 loss 下降变慢了一查专家激活统计8 个专家里有 5 个的激活率不到 5%等于白养了。原因就是负载均衡损失的系数设太小我一开始设的 0.001路由网络没有动力去均衡。改成 0.01 之后激活率分布明显均匀了loss 也继续下降。5.2 循环思考想太多怎么办前面提过迭代次数不是越多越好。我实测发现迭代到 8 次以上模型在简单问题上反而容易出错因为它会把简单问题复杂化。这个现象在推理任务上特别明显。解决办法有两个。一是固定迭代次数根据任务难度选一个合适的值简单任务 2 次、复杂任务 4 次。二是自适应迭代让模型自己决定什么时候停——具体做法是加一个停止头每一轮迭代后预测一个停止概率超过阈值就停。第二种方法更优雅但训练起来更复杂我这次时间有限没做留作后续优化。还有一个相关的坑迭代次数在训练和推理时要一致。我试过训练用 4 次、推理用 2 次结果性能掉了一大截。因为模型在训练时已经适应了 4 轮的思考节奏突然砍到 2 轮它的中间表示还没收敛就被迫输出了。5.3 显存不够的应急处理8 小时的时间预算里最怕的就是跑到一半 OOM。我总结了几招应急处理按优先级排序第一招降低 batch size 同时增加梯度累积步数等效 batch size 不变但激活显存降下来。第二招开启更激进的激活检查点把注意力和 FFN 都检查点化显存能再省 20%代价是训练慢 15% 左右。第三招缩短序列长度从 2048 降到 1024显存直接砍半但对长文本任务有影响。第四招减少 MoE 专家数量从 8 个降到 4 个这是最后的办法会损失模型容量。我的建议是训练前先用小规模数据跑一遍完整流程确认显存峰值再上全量数据。这个预跑花不了多少时间但能避免跑到一半崩掉的悲剧。5.4 几个容易被忽略的细节最后分享几个我踩过的、文档里不会写的坑。第一个是 checkpoint 保存策略。8 小时训练如果每 100 步存一次磁盘很快就满了。我的做法是只保留最近 3 个 checkpoint 加一个最佳 checkpoint其他的自动删。另外保存的时候用 safetensors 格式比 pickle 快而且安全。第二个是验证集的选择。循环思考模型的验证不能只看最终答案对不对还要看中间思考过程的质量。我专门构造了一个小验证集包含思考步骤是否合理的标注训练中定期评估这个指标比只看 loss 更能反映真实能力。第三个是随机种子的影响。小模型对初始化很敏感我跑过 3 个不同的随机种子最终性能差异能到 3 个百分点。所以如果你要复现别人的结果一定要固定种子如果你要评估自己的改进是否有效最好跑多个种子取平均。第四个是学习率调度的选择。我用的是 cosine 调度加 warmup但发现对循环思考模型来说最后阶段的学习率不要降到 0留一个小的最小值比如峰值的 5%这样模型在训练末期还能继续微调循环迭代的行为。降到 0 的话模型会过早定型循环思考的灵活性反而下降。这套流程跑下来我最终得到的模型在几个推理基准上的表现比同参数量的稠密模型高了大概 8 到 12 个百分点而训练成本只有后者的三分之一左右。当然它离真正的大模型还有很大差距但作为一个验证性项目它至少说明了一件事循环思考这种能力不一定非要大模型才能有结构设计和训练策略对了小模型也能想几步。后续我打算试试自适应迭代和更细粒度的专家路由看看能不能在同样的预算下再挤出一点性能。