1. 项目缘起与整体设计思路1.1 为什么我想用单卡A100折腾一个会循环思考的小模型先说清楚这个项目到底在干什么。标题里的循环思考不是玄学指的是让模型在推理阶段对同一段输入做多轮内部迭代——每一轮把上一轮的隐状态或中间结果重新喂回去逐步修正自己的表示最终再输出答案。这类结构在文献里常被叫做循环深度recurrent depth或者迭代推理本质上是把想一遍变成想好几遍。我手上只有一张A100 80G租的按小时计费所以整个项目的约束非常明确8小时内必须跑完一轮从零开始的训练并且要能看到循环思考带来的实际收益而不是只跑通一个玩具。这个约束直接决定了后面所有的技术选型——不能上太大的模型不能做全参数SFT数据规模要克制训练流程要能在一个晚上跑完。适合谁来参考这篇内容如果你已经写过Transformer、跑过LoRA微调但没试过把循环结构真正塞进训练流程里那这篇正好。如果你是完全的新手建议先把Transformer的前向传播手写一遍再回来不然中间很多设计取舍你会看不懂为什么这么选。1.2 核心思路拆解循环到底循环在哪里很多人一听循环思考就以为是RNN那种时间步循环其实不是。我这里的做法是层间共享多轮迭代模型主体还是标准Transformer的Decoder结构但取其中若干层作为一个思考块thinking block推理时把这个块重复调用N次每次的输入是上一次的输出加上原始输入的残差。为什么这么设计因为如果直接堆层数参数量会线性增长A100单卡8小时根本训不动。而共享权重的循环块参数量固定但有效深度可以随迭代次数变化——这正好是用时间换深度的思路。代价是训练时梯度要穿过多次迭代显存和计算量都会上去所以迭代次数在训练阶段我控制在3次推理阶段可以放到5到8次看效果。另一个关键取舍是要不要用MoE。热搜里MoE架构很火但MoE的问题是路由网络本身要训练而且专家并行在单卡上收益有限反而增加显存碎片。我最后的方案是主体用Dense只在FFN部分做一个轻量的专家分组类似MoE但共享大部分参数这样既蹭到了MoE的稀疏激活思路又不会让单卡训练崩掉。1.3 技术栈选型与理由组件选型理由基座结构Decoder-only Transformer自回归生成任务最成熟社区工具链全循环机制共享权重块重复调用参数量可控有效深度可调微调方式LoRA 部分层全参全参SFT在8小时内训不收敛LoRA能快速适配精度bf16 gradient checkpointingA100对bf16友好checkpointing省显存优化器AdamW cosine schedule小模型训练标配warmup 500步数据自建指令数据集约50万条规模适中8小时内能过2到3个epoch这里重点说LoRA。热搜里lora微调、lora训练、秋叶lora训练器这些词很热但很多人把LoRA当成万能药。我的经验是LoRA适合适配不适合从零学新能力。所以这个项目里LoRA只用在注意力层的Q、V投影上FFN和循环块的核心参数还是全参训练否则循环思考这个新结构根本学不出来。2. 核心细节解析与实操要点2.1 循环块的具体结构设计循环块我放了4层Transformer Decoder Layer隐藏维度768注意力头数12FFN中间维度3072。为什么是4层因为迭代3次相当于有效深度12层这个深度对小模型来说已经能学到一定的多步推理能力再深的话单卡8小时训不完。每一轮迭代的输入构造是这样的# 伪代码示意 h embed(input_ids) for i in range(num_iterations): h_prev h h thinking_block(h, attention_mask) h h h_prev # 残差连接防止迭代发散 h layer_norm(h) logits lm_head(h)注意那个残差连接这是整个设计里最容易被忽略但最关键的细节。如果不加残差多轮迭代后隐状态会爆炸或者坍缩训练直接不收敛。我第一版就是忘了加loss在前200步就变成nan了。2.2 训练阶段的迭代次数怎么定训练时迭代次数不能太多原因是梯度要穿过所有迭代步显存占用和计算量都是线性增长的。我实测下来迭代1次显存约28G单步约0.4秒迭代3次显存约52G单步约1.1秒迭代5次显存约78G单步约1.8秒接近OOM边缘所以训练阶段定在3次推理阶段可以动态调整。这里有个技巧训练时可以用随机迭代次数比如在2到4之间随机采样这样模型对迭代次数更鲁棒推理时想用几次都行。这个技巧我是从dropout的思路迁移过来的实测确实有效。2.3 LoRA配置与参数选择LoRA的rank我设的是16alpha设32dropout 0.05。为什么rank不设大因为循环块本身参数量就有限rank太大反而过拟合。alpha是rank的2倍这是社区常用的经验值能让LoRA更新的幅度和全参更新匹配。目标模块只选q_proj和v_proj不选k_proj和o_proj。原因是Q和V决定了关注什么和取出什么是注意力里最需要适配的部分K和O相对稳定全参训练时已经学得差不多了。这个选择让LoRA参数量控制在总参数的2%左右训练速度提升明显。注意LoRA的权重初始化要用高斯分布标准差设0.02不要用全零。全零初始化会让训练初期梯度消失我踩过这个坑loss前100步几乎不动。2.4 数据准备的关键细节数据我用了三个来源混合通用指令数据30万条、多步推理数据15万条、自构造的循环任务数据5万条。循环任务数据是自己写的模板生成的比如计算((35)*2-4)/3这种需要多步才能算出来的题专门用来训练循环思考能力。数据格式统一成instruction-input-output三段式长度截断到512。为什么是512因为循环3次后有效序列长度相当于1536再长显存扛不住。截断时优先保留output部分input可以截这个顺序不能反。3. 实操过程与核心环节实现3.1 环境搭建与依赖安装环境是Ubuntu 22.04 CUDA 12.1 PyTorch 2.1。依赖清单如下pip install torch2.1.0 --index-url https://download.pytorch.org/whl/cu121 pip install transformers4.36.0 pip install peft0.7.0 pip install datasets2.16.0 pip install accelerate0.25.0 pip install bitsandbytes0.41.0版本一定要锁死尤其是transformers和peft的版本匹配我试过用最新版结果LoRA注入失败排查了两小时才发现是版本问题。3.2 模型定义与循环块实现核心代码结构如下我简化了非关键部分class ThinkingBlock(nn.Module): def __init__(self, config, num_layers4): super().__init__() self.layers nn.ModuleList([ DecoderLayer(config) for _ in range(num_layers) ]) self.norm nn.LayerNorm(config.hidden_size) def forward(self, h, mask): for layer in self.layers: h layer(h, mask) return self.norm(h) class RecurrentModel(nn.Module): def __init__(self, config, num_iter3): super().__init__() self.embed nn.Embedding(config.vocab_size, config.hidden_size) self.block ThinkingBlock(config) self.lm_head nn.Linear(config.hidden_size, config.vocab_size) self.num_iter num_iter def forward(self, input_ids, mask): h self.embed(input_ids) for _ in range(self.num_iter): h_prev h h self.block(h, mask) h h h_prev return self.lm_head(h)注意h h h_prev这行就是前面说的残差少了它整个训练会崩。3.3 训练脚本与关键参数训练用accelerate做单卡bf16batch size设8梯度累积4步等效batch 32。学习率3e-4warmup 500步cosine衰减到1e-5。总步数按数据量算50万条数据batch 32一个epoch约15600步跑2个epoch约31200步。8小时能不能跑完实测单步1.1秒31200步约9.5小时超了。所以我把数据砍到35万条跑2个epoch约22000步约6.7小时留出1小时做验证和推理测试。这个时间预算是整个项目的硬约束所有决策都要围绕它。accelerate launch train.py \ --model_name recurrent-small \ --data_path ./data/mixed.jsonl \ --batch_size 8 \ --grad_accum 4 \ --lr 3e-4 \ --warmup 500 \ --epochs 2 \ --max_len 512 \ --num_iter 3 \ --lora_rank 16 \ --lora_alpha 32 \ --output_dir ./ckpt3.4 训练过程监控与关键节点训练过程中我重点盯三个指标loss、梯度范数、显存占用。loss从初始的10.8降到2.3左右收敛梯度范数稳定在0.5到1.5之间显存峰值52G。如果梯度范数超过5说明迭代次数太多或者学习率太高要立刻调。第5000步左右loss会有一个平台期这时候不要慌是循环结构在学怎么迭代过了这个平台loss会继续降。我在这个阶段差点以为模型训废了后来发现是正常现象。3.5 推理阶段的迭代次数调优训练完的模型在推理时可以动态调迭代次数。我做了个对比测试迭代次数简单任务准确率多步推理准确率单条推理耗时178%42%0.08s381%58%0.19s582%67%0.31s881%69%0.48s可以看到迭代次数从1增加到5多步推理准确率提升明显但到8次后收益递减。所以实际部署时我建议按任务难度动态选简单任务1到2次复杂任务5次左右。4. 常见问题与排查技巧实录4.1 训练不收敛的几种典型情况情况一loss直接变nan。最常见原因是残差没加或者layer norm位置不对。检查循环块里是不是有h h h_prev以及norm是不是放在残差之后。情况二loss降不下去卡在5左右。大概率是LoRA rank太小或者学习率太低。把rank从8提到16学习率从1e-4提到3e-4一般能解决。情况三loss震荡严重。迭代次数太多导致梯度不稳定。把训练迭代次数从5降到3或者加梯度裁剪max_grad_norm设1.0。4.2 显存不够用的排查顺序显存OOM的排查我总结了一个顺序按这个顺序查基本能定位先看batch size是不是太大从8降到4试试再看序列长度512降到384然后开gradient checkpointing最后才考虑降迭代次数为什么这个顺序因为前三个对训练效果影响小降迭代次数会直接影响模型能力是最后手段。4.3 LoRA相关的坑热搜里lora微调实战教程qwen、safetensors lora这些词很热但实际用起来有几个坑LoRA权重保存格式默认保存成adapter_model.bin但如果你想用safetensors格式要在save_pretrained时加safe_serializationTrue。LoRA合并推理时可以把LoRA权重合并回基座用merge_and_unload()合并后推理速度提升约15%。LoRA和全参混合训练如果同时训LoRA和全参要确保优化器对两组参数用不同的学习率否则全参会把LoRA的更新淹没。4.4 循环结构的特有问题循环结构有两个特有问题普通Transformer不会遇到问题一迭代发散。表现是推理时迭代次数越多输出越乱。原因是残差累积导致数值爆炸。解决办法是每轮迭代后做一次layer norm或者把残差系数设成0.5而不是1.0。问题二迭代次数和训练不一致。训练用3次推理用8次效果反而下降。这是因为模型没见过8次迭代的分布。解决办法是训练时随机迭代次数让模型适应不同次数。4.5 常见问题速查表问题现象可能原因解决办法loss变nan残差缺失/norm位置错检查循环块结构loss卡住不降LoRA rank小/学习率低rank提到16lr提到3e-4显存OOMbatch大/序列长降batch开checkpointing推理输出乱迭代发散加norm降残差系数迭代多了效果差训练推理次数不一致训练时随机迭代次数LoRA不生效版本不匹配锁死transformers和peft版本5. 迭代效果验证与扩展方向5.1 怎么验证循环思考真的有用验证不能只看loss要做消融实验。我做了三组对比迭代1次、3次、5次在同一个测试集上跑。测试集分两类单步任务如直接问答和多步任务如数学计算、逻辑推理。结果显示单步任务上迭代次数影响不大准确率差异在3%以内但多步任务上迭代5次比迭代1次准确率高了25个百分点。这就证明了循环结构确实在多步推理上起作用而不是单纯增加了参数量。5.2 后续可以怎么扩展这个项目跑通后有几个方向可以继续做迭代次数自适应让模型自己决定要想几轮简单问题少想复杂问题多想。可以用一个小的分类头预测需要的迭代次数。循环块和MoE结合把循环块里的FFN换成MoE结构进一步增加容量而不增加计算量。更长序列现在限制在512如果能上2048循环思考在长文本推理上的优势会更明显。5.3 一些实操心得最后分享几个我在这个项目里踩坑总结的经验。第一8小时的时间预算要留20%的buffer我第一版没留结果训练到7.5小时的时候发现验证脚本有bug只能重新跑。第二循环结构的调试要从迭代1次开始先确保基础结构能训再逐步加迭代次数一上来就3次会很难定位问题。第三LoRA和全参混合训练时先把全参训到收敛再加LoRA同时训容易互相干扰。这个项目最大的收获是让我理解了有效深度和参数深度的区别。一个共享权重的循环块参数量只有4层但迭代5次后有效深度20层在多步推理上的表现接近一个12层的独立模型而参数量只有后者的三分之一。对于算力有限的场景这种用时间换深度的思路值得一试。