1. 为什么要在单卡 A100 上折腾一个循环思考的小模型循环思考这个词听起来有点玄但落到工程上其实很朴素让模型在生成答案之前先在内部把同一段隐状态反复过几遍用额外的计算深度换取推理质量。这跟人类遇到难题时会在脑子里多转几圈是一个道理。我这次的目标很明确——用一张 A100 80G在 8 小时以内从零训练一个参数量不大、但具备这种迭代推理能力的小模型并且全程不依赖任何预训练权重。先说清楚这件事的定位。它不是要跟动辄几百亿参数的大模型比通用能力而是想验证一个假设在固定参数量的前提下把计算花在深度循环上是否比单纯堆层数更划算。这个思路和 Transformer 的经典设计有直接关系——标准 Transformer 每一层参数只被用一次而循环结构让同一组参数被复用多次等于用时间换空间。对显存吃紧、又想探索推理深度的场景来说这是个性价比很高的方向。适合读这篇的人大概有三类一是手里有单卡或少量卡、想认真做点小规模训练实验的工程师二是对 Transformer、MoE、LoRA 这些概念已经看过不少科普、但没真正从零跑通过一次训练的人三是想理解循环深度这类结构改动到底怎么落地、坑在哪里的研究者或学生。我会把数据构造、模型结构、训练配置、显存控制、以及中途踩的坑全部摊开讲参数给到能直接抄的程度。需要提前说明的是8 小时这个预算不是拍脑袋定的。A100 80G 的 FP16/BF16 算力大约在 300 TFLOPS 量级实际训练受限于显存带宽和 kernel 效率能稳定跑到的有效算力通常只有峰值的三到四成。按这个折算8 小时能提供的有效计算量是有限的所以模型规模、序列长度、循环次数这三者必须精打细算任何一个放大都会直接吃掉预算。后面每一节我都会围绕这个约束来展开。2. 循环深度到底改了什么结构设计与参数量账本2.1 标准 Transformer 的一次性问题先把基线摆出来。一个标准的 decoder-only Transformer输入 token 经过 embedding 后依次穿过 N 个结构相同的 block每个 block 里是自注意力加前馈网络最后接一个输出头。关键在于第 1 层的参数和第 N 层的参数是各自独立的信息只沿着层数方向单向流动一次。这意味着模型的思考深度被层数硬性锁死了想更深就得加层加层就加参数、加显存。循环结构的改动点很小但很关键不再堆 N 个不同的 block而是只保留一个或少数几个共享 block让隐状态在它里面反复迭代 K 次。数学上可以写成 h_{t1} Block(h_t, x)其中 x 是固定的输入条件t 从 1 到 K。这样参数量只跟 block 本身的大小有关跟循环次数 K 完全解耦。K 从 4 提到 8参数量一个不涨只多花计算时间。2.2 参数量与计算量的账要分开算很多人第一次接触循环结构会混淆两个量参数量和计算量FLOPs。我列个表把关系讲清楚这也是我设计时的核心依据。配置项标准 12 层 Transformer循环结构1 个 block 循环 12 次参数量约 12 倍单层约 1 倍单层前向 FLOPs约 12 倍单层约 12 倍单层显存占用激活高需存 12 层中间态低可只存循环态推理深度固定 12可动态调整 K从表里能看出循环结构省的是参数和激活显存不省计算。这正好契合单卡场景A100 的算力相对充裕但显存是硬约束。用循环换参数等于把瓶颈从显存挪到了算力上而算力恰好是我们相对富裕的资源。2.3 我最终选定的结构参数经过几轮试跑我定下来的配置是隐藏维度 512注意力头数 8每个头的维度 64前馈网络中间层维度 2048也就是 4 倍扩展共享 block 数量为 2两个 block 交替循环比单 block 表达能力强一些循环次数 K 设为 6。词表大小控制在 32000用字节级 BPE。这样算下来总参数量大约在 4000 万到 5000 万之间属于小模型范畴单卡训练毫无压力。这里有个经验点共享 block 数量不要设成 1。我最早用单 block 循环发现模型很难同时学好底层特征提取和高层语义整合这两件性质不同的事loss 下降明显偏慢。改成 2 个 block 交替后收敛速度肉眼可见地变好。这其实符合直觉——循环复用同一组参数如果这组参数还要兼顾差异很大的功能就会互相打架。2.4 循环次数 K 不是越大越好K 的选择有个反直觉的地方不是越大越好。我做过 K2、4、6、8、12 的对比在固定 8 小时预算下K6 的综合表现最好。K 太小循环带来的深度收益不明显跟普通浅层模型差不多K 太大单步训练时间线性增长同样的 8 小时里能过的 token 数大幅减少模型反而欠拟合。这本质上是在每步算得更深和总共见更多数据之间做权衡。提示循环次数 K 和训练步数是此消彼长的关系。调大 K 之前先确认你的数据量是否足够支撑更少的训练步数否则容易陷入结构很 fancy 但没训透的尴尬。3. 数据从哪来小模型的数据构造与配比策略3.1 从零训练意味着数据要自己攒没有预训练权重意味着模型对世界的全部认知都来自我喂给它的数据。这一步的投入产出比往往被低估——很多人把精力全花在调结构上结果数据一塌糊涂loss 曲线再漂亮也没用。我的数据来源分三块公开的中文通用语料、一部分英文技术文本、以及我针对推理类任务专门构造的合成数据。通用语料负责让模型学会语言的基本规律技术文本让它对结构化表达更敏感合成数据则是为了强化多步推理这个我们真正想要的能力。三者的配比大概是 6:2:2。这个比例不是固定的训练前期可以多用通用语料打基础后期逐步提高合成数据的采样权重。3.2 合成推理数据怎么造合成数据是这次实验里最花心思的部分。我的做法是构造一批需要多步才能得出答案的样本比如简单的算术链、逻辑排序、以及带中间步骤的问答。关键不在于题目多难而在于答案的生成过程必须显式包含中间推理步骤让模型在训练时能学到先想再答的模式。举个具体例子我会生成这样的样本问题是一个数加上 7 再乘以 3 等于 45这个数是多少目标输出不是直接给答案而是45 除以 3 等于 1515 减去 7 等于 8所以答案是 8。这种带步骤的监督信号配合循环结构能让模型在内部迭代时逐渐对齐到逐步逼近的行为。我大概造了 50 万条这类样本覆盖算术、逻辑、简单规划三类。3.3 分词器的坑别用现成的从零训练有个容易被忽略的细节——分词器也得自己训。我一开始图省事想直接套用一个现成的中文分词器结果发现词表和我的数据分布不匹配很多高频词被切得七零八落直接拖累了训练效率。后来我用自己攒的语料重新训了一个 BPE 分词器词表 32000效果立刻不一样。训分词器时有个参数值得注意词表大小。太小会导致序列变长、计算量上升太大则 embedding 层参数膨胀对小模型不划算。32000 这个量级对中文为主、夹杂英文的语料是比较平衡的选择。训完后一定要检查一下压缩率也就是平均每个 token 覆盖多少字符中文语料下这个值在 1.5 到 2 之间比较健康。3.4 数据清洗的几条硬标准清洗这块我踩过坑总结几条硬标准。第一去重必须做而且要做得狠重复数据会让模型过拟合到特定片段loss 看着降得快其实是假象。第二长度过滤太短的碎片和超长的文档都要处理我一般把长度控制在 64 到 1024 个 token 之间。第三敏感和低质内容过滤这个不用多说是底线。第四格式统一把各种奇怪的空白、控制字符清干净否则分词器会产出大量无意义的 token。注意清洗完的数据一定要抽样人工看几十条。我见过太多人清洗脚本跑完就直接开训结果模型学出一堆乱码回头排查才发现是清洗环节把正常文本也误伤了。4. 训练配置8 小时预算怎么分配才不浪费4.1 精度选择BF16 是单卡首选A100 对 BF16 有原生支持这是单卡训练的首选精度。相比 FP16BF16 的动态范围大得多不容易出现梯度溢出基本不需要 loss scaling 那套麻烦事。实测下来BF16 下训练非常稳定我全程没遇到过一次 NaN。FP16 虽然理论算力一样但为了防溢出要额外处理反而增加调试成本不划算。显存方面BF16 下模型参数、梯度、优化器状态加起来大概是参数量的 12 到 16 倍如果用 Adam 类优化器。5000 万参数的话这部分占用不到 1G完全不是瓶颈。真正的显存大头是激活值和注意力矩阵尤其是序列长度拉长之后。4.2 批次与序列长度的组合批次大小和序列长度是一对需要联调的参数。我的策略是先把序列长度定在 512然后用梯度累积把有效批次撑到 256。为什么是 512 而不是更长因为循环结构下每个 token 要过 K 次 block激活显存本来就比普通结构高序列再拉长很容易 OOM。512 是个稳妥的起点等训练稳定后再考虑逐步加长。梯度累积步数设为 8微批次大小 32这样 32 乘 8 等于 256 的有效批次。这个规模对小模型来说足够稳定梯度噪声不会太大。如果你显存更紧张可以把微批次降到 16累积步数提到 16效果基本等价只是训练速度略慢。4.3 学习率与调度学习率我用的是带预热的余弦退火。峰值学习率定在 3e-4预热步数 2000 步然后余弦衰减到峰值的十分之一。这个配置对小模型比较友好既不会因为学习率太大而发散也不会因为太小而收敛慢。从零训练时预热特别重要因为初始参数是随机的一上来就用大学习率很容易把模型带偏。优化器用 AdamW权重衰减 0.1beta 取默认的 0.9 和 0.95。梯度裁剪阈值设 1.0防止偶发的梯度尖峰。这几个参数我基本没怎么调属于比较通用的配置直接抄问题不大。4.4 8 小时的时间账现在算总账。序列长度 512有效批次 256那么每个训练步处理的 token 数是 512 乘 256 约等于 13 万。循环结构下每个 token 的计算量是普通结构的 K 倍K6所以单步实际计算量相当于普通结构的 78 万 token。按 A100 的有效算力估算单步耗时大概在 1.5 到 2 秒之间。8 小时等于 28800 秒按每步 1.8 秒算能跑大约 16000 步。16000 步乘以每步 13 万 token总共见过约 20 亿 token。对 5000 万参数的小模型来说这个数据量是相当充足的甚至有点过训练的意思。这也说明 8 小时预算和模型规模是匹配的没有浪费。配置项取值说明精度BF16A100 原生支持稳定序列长度512兼顾显存与效率微批次32显存友好梯度累积8有效批次 256峰值学习率3e-4余弦退火预热步数2000从零训练必备循环次数 K6深度与速度的平衡点预计总步数约 160008 小时预算内5. 训练过程中的真实踩坑记录5.1 第一个坑loss 不降反升开训后大概 500 步loss 突然从 3.2 反弹到 4.5我当时第一反应是学习率太大。但把学习率砍半后问题依旧说明不是这个原因。逐步排查后发现是循环结构里的残差连接出了问题——我在每次循环迭代时都加了残差导致隐状态的数值随着迭代次数累积不断放大K6 时已经溢出到不稳定区间。修复方法是在循环迭代之间做归一化具体是在每次进入 block 前加一层 LayerNorm把隐状态重新拉回稳定范围。改完之后 loss 曲线立刻变得平滑。这个坑的教训是循环结构里数值稳定性比普通结构更敏感因为同一组运算被反复叠加任何微小的放大效应都会被迭代次数放大。5.2 第二个坑显存碎片导致的间歇性 OOM训练到中途偶尔会报 OOM但重跑同样的步数又没事。这种间歇性 OOM 最烦人因为它不是配置问题而是显存碎片。原因是循环结构下每次迭代都要分配和释放中间张量长时间运行后显存被切得七零八落某一步需要一块较大的连续显存时就分配失败了。解决办法有两个一是设置环境变量让 PyTorch 使用更激进的显存分配策略减少碎片二是把循环迭代中能复用的张量提前分配好避免反复申请释放。我两个都做了之后连续跑了几个小时再没出现过 OOM。这个经验对任何长时间训练都适用不只是循环结构。5.3 第三个坑合成数据比例过高导致语言能力退化中期我为了提高推理能力把合成数据的采样权重从 0.2 提到了 0.5结果发现模型在通用文本上的表现明显变差生成的句子开始变得机械、重复。这是典型的灾难性遗忘——合成数据格式高度统一模型过度拟合到这种格式后把通用语言能力给挤掉了。调整方案是把合成数据权重回调到 0.3并且在合成数据里混入一定比例的多样化表达避免格式过于单一。同时我在验证集里同时监控通用语言 loss 和推理任务准确率两个指标确保两者都不掉队。这个教训是任何一类数据占比过高都会带来偏科训练时一定要有多个维度的监控指标。5.4 第四个坑验证集泄漏这个坑比较隐蔽。我构造合成数据时用了一个随机种子生成题目结果验证集也用同一个种子生成导致验证集和训练集有大量重复样本。模型在验证集上的准确率虚高我一度以为训练很成功直到换了一批全新生成的验证样本才发现真实表现差了一大截。修复很简单训练集和验证集用完全独立的生成流程和种子。但这个坑提醒我合成数据的划分不能想当然必须确保验证集是真正没见过的。后来我干脆把验证集全部换成人工构造的样本彻底杜绝泄漏可能。6. 循环结构 vs 常规结构实测对比与选型建议6.1 同等参数量下的对比为了验证循环结构到底值不值我做了一组对照实验一个是我最终的循环模型2 个共享 blockK6另一个是参数量相近的常规 6 层 Transformer。两者用完全相同的数据、相同的训练步数、相同的优化器配置。结果在推理类任务上循环模型的准确率高出约 8 个百分点在通用语言任务上两者基本持平。这个结果说明循环结构确实在需要多步推理的任务上有优势而且没有牺牲通用能力。原因我分析是循环迭代给了模型一个内部 scratchpad它可以在隐空间里反复修正自己的表示这种能力是固定层数的模型不具备的。6.2 训练速度的代价当然代价也有。循环模型单步训练时间比常规模型长约 40%因为同样的参数量要跑 6 次。这意味着在固定时间预算下循环模型能过的数据更少。所以它适合的场景是数据量充足、任务偏推理、且你愿意用训练时间换质量。如果你的任务只是简单的分类或生成常规结构可能更划算。6.3 什么情况下该选循环结构我总结了几条判断标准。第一任务需要多步推理比如数学、逻辑、规划类。第二显存是瓶颈而算力相对充裕。第三你希望推理时能动态调整计算深度简单问题少循环几次、难题多循环几次。如果这三条里有两条符合循环结构就值得一试。反之如果任务简单、追求极致训练速度那就老老实实用常规结构。对比维度循环结构常规结构参数量效率高低单步训练速度慢约 40%快推理任务表现优一般通用任务表现持平持平推理深度可调支持不支持数值稳定性需额外处理较稳7. 训练完之后推理阶段的动态深度玩法7.1 用置信度决定循环几次循环结构训练完之后最大的红利在推理阶段。因为循环次数 K 是运行时可调的我可以根据模型输出的置信度动态决定循环几次。具体做法是先跑较少的循环次数比如 3 次看输出分布的熵如果熵很低说明模型很确定直接输出如果熵很高说明模型拿不准就继续增加循环次数到 6 甚至 8再重新判断。这个策略在实测中能省下不少计算。简单问题平均只需要 3 到 4 次循环难题才用满 8 次整体推理成本比固定 K8 低了约三成而质量几乎没损失。这就是循环结构相对固定深度模型的独特优势——计算可以按需分配。7.2 早停与稳定性动态循环要配一个早停机制。我的做法是监控相邻两次循环的输出差异如果差异小于某个阈值说明模型已经收敛再循环也没意义直接停。这个阈值我设在输出分布的 KL 散度小于 0.01。加上早停后既避免了无谓计算也防止了过度循环导致的输出漂移。提示动态循环次数虽然灵活但一定要设上限。我见过有人不设上限结果个别样本陷入循环出不来单条推理耗时爆炸。上限设 8 到 10 次比较稳妥。7.3 和 LoRA 微调的结合思路训练好的基础模型如果要适配具体任务可以接 LoRA 微调。LoRA 的好处是只训练少量低秩矩阵不动原模型参数显存占用极低单卡就能搞定。对循环结构来说LoRA 可以挂在共享 block 的注意力层上微调时循环次数保持不变。这样既保留了循环推理的能力又能快速适配新任务是很实用的组合。具体操作上LoRA 的秩设 8 到 16 就够学习率比预训练时高一个量级比如 1e-3训练步数几百到几千步即可。因为基础模型已经学到了通用表示LoRA 只需要做小幅调整。这套流程我在几个下游任务上试过效果稳定而且每次微调只要几十分钟。8. 一些关于 MoE 和扩展性的延伸想法8.1 循环结构能不能和 MoE 结合MoE混合专家的核心思想是让不同的 token 走不同的专家网络从而在参数量很大时保持计算量可控。它和循环结构其实是可以叠加的把共享 block 里的前馈网络换成 MoE 层每次循环时根据 token 动态选择专家。这样既有了循环带来的深度又有了 MoE 带来的参数容量。不过要提醒的是MoE 会引入负载均衡的额外复杂度训练时容易出现专家利用不均的问题需要加辅助损失来约束。对单卡小模型来说MoE 的收益可能不明显因为专家数量上不去反而增加了工程复杂度。我的建议是先把循环结构跑通有余力再考虑 MoE。8.2 扩展到更大模型时的注意事项如果将来要把这套方案扩展到更大的模型有几个点要提前想清楚。第一循环次数 K 在大模型上要重新调因为大模型单步计算更重K 太大时间成本吃不消。第二数值稳定性问题会更突出归一化的位置和方式要仔细设计。第三动态推理的调度逻辑要更精细否则省下的计算可能被调度开销吃掉。8.3 这套方案适合谁、不适合谁最后说点实在的。这套方案适合想认真做小规模训练实验、对模型结构有探索兴趣、手里有单卡或少量卡的人。它不适合追求开箱即用、只想调 API 的人也不适合任务本身很简单、用现成模型就够的场景。训练一个自己的小模型最大的价值不在于它多强而在于你对它的每一个细节都了如指掌这种掌控感是调 API 永远给不了的。我在整个实验里最大的体会是约束反而催生创造力。正是因为只有一张 A100、只有 8 小时我才被迫去思考哪些计算是必要的、哪些结构改动是真正有效的。如果算力无限我可能就直接堆参数了反而学不到这些东西。循环结构这个方向我觉得还有不少可以挖的空间尤其是动态深度和推理效率的结合值得继续折腾。