去年年中我们团队把模型训练的预算砍了一半产品端的预期反而提了两倍新模型要能在手机上跑得动要接实时对话场景还得保留代码解释和长文档归纳的能力。这一矛盾逼着我们把思路从“追参数量”切换到“怎么把每个参数用到位”于是就有了xLLM这个项目。代号里的x在我们内部一直读作“叉”——既是实验性的意思也是极致压缩的意思用尽量少的参数、尽量省的算力训练出超过同尺寸平均水平的语言模型。眼下第一轮预训练跑完结果基本达到预期1.5B参数量级五个标准评测集上比同尺寸基线平均高3.2个点INT4量化后端侧推理内存压到4GB以内。这篇博客就是把架构选型、数据配比、训练策略到部署量化踩过的坑完整梳理一遍给正在折腾轻量级预训练架构的团队做个参照。1. 项目缘起轻量级不该只是“蒸馏一下”1.1 从大模型直接裁小的那条路真的走不通最早我们想过最省事的方案——直接把开源7B模型的权重拿来蒸馏到1.5B。做了两轮实验结果都不理想模型确实能快速收敛到差不多能用的水平但一上评测集就露馅尤其在多轮对话的一致性和代码长上下文追踪上几乎是一迁移就碎。后来复盘出的原因是容量瓶颈。小模型和大模型之间的能力差距不只是“知识量”的差别更多是中间层表征的复杂度和特征复用方式不同。蒸馏只能让教师模型“教”给学生模型知识却教不了容量不足导致的表征结构塌陷。换句话说对一个1.5B的模型你不能指望它像7B模型那样用一个很宽的FFN和一个很深的特征金字塔来记忆纷繁复杂的模式。它的容量天然有限必须在架构层面就需要把“参数效率”这个指标放在第一位。所以我们很快调整方向不走蒸馏后期补救的路子而是从模型结构、训练数据、目标函数三个维度一起重新设计。这也是xLLM和其他“小模型”最大的不同——它是一个“原生设计”的轻量级架构不是一个“阉割版”的大模型。1.2 轻量级预训练的几个硬性指标项目启动时我们立了几条硬指标后续所有架构取舍都围绕这些数字展开指标目标值备注模型参数量≤1.5B实际激活参数约1.1B预训练算力成本≤320 GPU·day以A100-80G为标准计端侧推理内存≤4GBINT4量化后含KV Cache目标效果追平同规模3B模型以5个公开评测集为参照最长序列长度8192训练阶段动态扩展到8K有几个指标需要特别解释一下。参数量定在1.5B不是因为1.5B有什么魔法而是端侧内存和延迟的硬约束下能接受的上限。我们做过内存测算1.5B模型FP16权重占3GB再加激活值和KV Cache直接落手机是不可能的所以必须配合量化。4GB内存这条线对应的是市面上主流旗舰机型App的可用内存水位超过这个数字系统容易杀后台没有实际部署价值。算力成本那条更有意思。320 GPU·day意味着如果我们用64张A100并行理想情况下5天训完一轮实际加上反复实验和回滚项目周期被控制在三周以内。这个预算不是拍脑袋定的而是我们用小规模代理实验估算出来的——先把数据量、模型参数量、训练步数和loss收益之间的关系摸了个大概再倒推出这个体量的模型大约需要多少tokens然后反推算力。后面我会详细说这个过程。1.3 为什么还要重造轮子开源小模型不够用可能有人会问市面上不是有大量2B、3B的开源模型吗为什么不直接用我们评估了一圈主要卡在两个点上。第一绝大多数开源小模型的训练数据以英文为主中文能力的衰减非常明显而我们产品的核心用户大量使用中文第二这些模型的架构基本都是大模型的等比例缩小参数量下去了但计算量和内存占用并没有按我们希望的幅度下降。换句话说它们在“轻量”这个维度上做得还不够彻底。我们自己训xLLM本质是想要一个“可定制、可控、可裁剪”的底座。架构和数据都是自己的后续如果想针对特定行业做持续预训练或者裁剪掉不必要的模块换推理速度都有操作空间。对一个内部平台型团队来说这个自主性比直接调API或者用开源权重更值钱。2. 架构拆解xLLM轻量化的三个关键决策2.1 注意力机制滑窗加全局稀疏别让平方复杂度毁掉一切标准Transformer的全局注意力复杂度随序列长度平方增长。放在7B、13B这种大模型上靠着算力和优化库还能硬扛但放在小模型上就非常不划算——小模型参数少省下的算力全被注意力的平方复杂度吃掉了。我们在8K序列长度上做过profiling标准MHA的注意力计算占整个forward的45%以上这显然不可接受。xLLM采用了一种混合注意力路由方案交替使用局部滑动窗口注意力和全局稀疏注意力。具体来说每个Transformer块按以下模式排列奇数层局部滑动窗口注意力窗口大小设为512。偶数层全局稀疏注意力通过Top-K路由挑选与当前token最相关的128个历史tokenK128。这样设计的好处立竿见影。计算复杂度从标准的O(n²)降到了O(n×w n×K)w是窗口大小512K是稀疏路由数量128。在8K序列下单层注意力的显存消耗比标准MHA下降了约60%forward速度提升了将近一倍。当然这种设计不是免费的午餐消融实验显示纯全局注意力在长文本依赖类任务上能比混合方案高1.1个点但训练速度会慢45%。我们在这个项目里选的是速度优先因为预训练预算就卡在那里后期如果要做长文档专项增强可以在fine-tuning阶段把部分层切换回全局注意力这个灵活性是架构设计时预留的。2.2 前馈网络想要MoE的收益又不想踩MoE的坑FFN层在Transformer里占据了大约三分之二的参数量是轻量化躲不开的点位。最粗暴的办法是直接砍宽度但这个做法在消融实验里掉点非常严重原因也不难理解FFN承担的是特征变换和信息写入的工作宽度不够模型的“记忆体”就太小了。我们试过把隐藏维度从标准比例缩小一半效果直接掉4个点以上这条路基本堵死。后来我们把目光投向MoEMixture of Experts思路。MoE的核心特点是参数共享、每个token只激活部分专家于是我们在xLLM里设计了一种“混合共享专家”结构大部分层使用一个共享FFN它承载通用的非线性变换和大部分知识存储每隔3层插入一组包含8个专家的路由FFN每个token通过Top-K路由选择2个专家激活。这里要特别说一下负载均衡loss的处理。MoE大家最怕的两个问题一是专家坍缩即所有token都涌向同一两个专家其他专家形同虚设二是通信开销专家分布在多卡上时会引入all-to-all通信。我们在xLLM里做了两件事控制风险路由专家只在同一张卡内激活从根上避开通信开销同时在loss里加入负载均衡正则项系数设为0.01。路由表也加了噪声让路由选择在训练初期不要过早确定给所有专家一个被充分训练的机会。这个设计的最终效果是同等参数量下xLLM的理论FLOPs比密集FFN版本低了约35%而效果损失控制在0.5个点以内对比直接砍宽度要划算得多。2.3 参数复用把省下的钱花在刀刃上模型架构里的另一个隐蔽浪费是各层参数完全独立。语言模型相邻层之间的表征往往有较高相似度如果每层都完整保存一套独立的投影矩阵等于每层都在从零开始学习一套相似度极高的映射冗余很高。xLLM采用了一种周期性的下投影共享策略每3个Transformer块共享一套下投影矩阵和层归一化参数只有注意力模块和FFN的核心权重保持独立。这个设计当时团队内部争论了很久反对者说参数共享会限制每层特征的差异化程度。消融实验给了我们信心在1.5B尺度下这种共享只造成0.3个点的效果损失却节省了大约2亿参数。省出来的空间被我们用来加宽FFN的中间层以及增加注意力头的数量——这两个改动带来的收益远大于0.3个点。这个取舍背后其实是一种架构哲学在轻量模型中参数预算就应该“集中砸在关键模块”而不是均匀铺开。如果你每层都塞一点参数最后的结果往往是每层都差一口气。我们验证下来把共享省下的预算投入到FFN宽度上效果立竿见影。2.4 最终架构参数一览xLLM第一轮版本的完整配置如下配置项数值隐藏层维度1536注意力头数16每个头64维FFN隐藏维度4096层数24其中每3层插一组路由专家共6组路由专家数每组8个Top-K2滑窗大小512全局稀疏K128参数量约1.5B这个组合是我们在多次代理实验里试出来的核心感受是1.5B尺度下宽度比深度更值钱。把隐藏维度从1280提到1536再配合理的FFN宽度比单纯加两层Transformer更能提升最终loss的下限。如果你也想做类似规模的模型建议先固定FLOPs预算然后用小规模实验扫一遍“深度-宽度”的等高线不要拍脑袋定结构。3. 数据管线与预训练策略轻量模型的半条命在数据3.1 数据配比不是所有语料都值得同样对待架构是模型的上限数据是模型的下限。在小模型上数据的影响甚至比架构更大因为模型容量有限能吸收的信息量本来就少喂什么就显得格外重要。xLLM预训练采用了动态配比策略而不是从头到尾用同一个混合比例。前期前40%的steps我们重点喂代码和数学数据比例分别到25%和15%。这个时期的目标是让模型建立“严格逻辑”和“结构化推理”的能力基底这类能力一旦早期没形成后期很难补。中期40%-80%加入大量通用英文和中文文本让模型把语言流利度和知识广度提上来。最后20%的steps则把书籍和高质量长文的比例调高目的是拉长上下文的建模能力因为长文档依赖需要模型在训练后期接触更多长序列样本才能稳住。几个数据配比的关键经验代码和数学数据的比例加起来不要低于30%。这不是我拍脑袋的结论而是消融实验里发现低于这个比例模型的推理类评测指标会明显塌。中文数据单独给权。通用开源数据里中文占比通常只有百分之几如果不做针对性补充中文能力会非常差。数据去重做两轮。第一轮用MinHash去近似重复第二轮用embedding相似度去语义重复。我们在第一轮实验里偷懒只做了第一轮去重结果模型反复“背诵”同一段语料损失曲线出现了周期性尖峰。3.2 训练目标语言模型损失加蒸馏再加对比学习xLLM没有只跑单纯的next token prediction而是用了组合目标函数。损失由三部分组成L L_lm 0.3 × L_distill 0.05 × L_ctrL_lm是标准交叉熵损失L_distill是蒸馏损失那是用同尺寸的合成教师模型一个训练更充分的1.5B模型做软标签监督L_ctr是对比学习项用来拉近同一段落内不同句子的表征距离增强语义连贯性。蒸馏项权重设为0.3是一个相对保守的值。我们试过0.5甚至0.8结果模型过于模仿教师分布自我表达能力反而下降而设成0.3可以让蒸馏起到“平滑标签”的效果又不至于喧宾夺主。对比学习部分我们放在最后20%的steps才开启因为太早加入会让表征空间过度拥挤影响语言建模的收敛速度。3.3 训练工程BF16混合精度、动态Batch与学习率调度训练效率方面我们用的方案是BF16混合精度加梯度检查点gradient checkpointing优化器使用AdamW并开了ZeRO Stage 2把优化器状态分布到多卡上。说实话这套组合在现在的大模型训练里已经不算新鲜但xLLM有一个自己的不同点我们对Batch Size做了动态调度。训练初期Batch Size只设128因为这个时候模型还在“快速认字”阶段太小的Batch会导致梯度噪声大不稳定中期逐步涨到256最后10%涨到512。配合Batch Size的上涨学习率采用余弦退火调度峰值设为2e-4warmup步数2000。这套动态策略比固定大Batch从头训到尾在最终loss上低了约0.08别小看这个数字在1.5B模型上已经是非常明显的提升。监控方面我们重点盯着两个指标一是“梯度范数”如果超过5就要警惕超过10基本意味着loss即将尖峰二是“历史困惑度差值”我们管它叫“perplexity gap”如果当前batch的PPL和同分布验证集PPL差值异常拉大说明模型正在过拟合训练集中的某个局部簇需要立刻检查是否数据重复。4. 训练稳定与效果调优那些模型没告诉你的细节4.1 显存优化让小模型也吃上大模型的成熟红利在1.5B这个体量下显存问题没有7B那么致命但依然存在尤其是在开启长序列训练后。我们沿用了一整套优化组合每一步都有明确目的FlashAttention V2把注意力算子的显存占用压到最低。序列打包sequence packing把不同样本拼接到同一条序列里减少短样本造成的填充浪费。Paged Optimizer States这是借鉴推理引擎的分页思想把偶尔访问的优化器状态换到CPU内存给激活值让路。这几项叠加下来我们能在单张A100-80G上塞下8K序列、Batch Size为8的配置这在标准实现里是做不到的。如果你在同等规模模型上老是OOM优先排查这几项是否都打开了很多时候不是显存不够而是显存被白白浪费。4.2 训练不稳定loss尖峰和梯度爆炸的现场xLLM预训练过程中遇到过几次loss尖峰最严重的一次发生在第4200步左右loss从2.7突然跳到3.4随后又慢慢降回来。起初我们以为只是随机波动直到发现同样的尖峰隔了几百步又出现一次才意识到不对。排查之后定位到两个原因。第一个是数据重复污染——某个批次里混入了大量未去重的高重复片段模型对这些片段过度学习导致梯度方向突然偏离第二个是学习率在warmup结束后有个陡峭的上升期接近峰值的梯度范数偏大遇到数据异常就会推得更高。解决方案很直接把学习率峰值从2.5e-4微调回2e-4并给梯度裁剪设了双重保险——先按全局范数裁剪到5再按单参数梯度裁剪到1。之后loss尖峰再没出现。有一个细节想提醒大家不要在loss尖峰出现的第一时间就停掉训练去改超参。很多时候模型自己会恢复你一动反而破坏了它在正常轨迹上的优化。我们的惯例是连续三次尖峰才介入排查给模型一点自我修复的时间。4.3 下游效果评估不能只看几个公开榜预训练快结束时我们在5个公开评测集上做了测试平均分确实不错但单独拆开看就有意思了代码生成类任务分数很高多轮对话一致性也不错但数学推理类任务明显拖后腿。这和数据配比里数学比例偏低有关我们及时在最后5%的steps里上调了数学数据的权重最终把数学评测拉回了正常水平。这里想分享的经验是不要只盯公开评测集的平均分一定要按任务类型拆开看。每个任务类型背后对应的是不同的能力维度平均分高可能只是长板和短板互相抵消后的假象。对小模型来说能力短板比平均分波动更值得关注因为部署后的用户往往只会盯住最差的那一档体验。5. 从预训练到端侧部署轻量架构的最后考验5.1 量化实践PTQ在稀疏激活面前会翻车预训练模型出来后部署环节第一个绕不开的就是量化。我们最初尝试了最简单的PTQ训练后量化把权重直接映射到INT4。结果大跌眼镜xLLM量化后掉点严重尤其在代码和数学任务上有些case直接从正确变成完全乱写。后来分析原因发现是稀疏专家模块导致激活分布出现明显的高峰尖尾PTQ的校准数据对这个分布拟合不足量化误差被显著放大。解决方案是把量化方案从PTQ换成QAT量化感知训练。我们在预训练的最后阶段对模型插入伪量化算子用很小的学习率继续训练几千步让模型自己去适应低比特表示的扰动。QAT之后的模型在INT4下效果损失从2.1个点缩小到0.7个点。这个经验推广到一般场景就是如果你的模型里有MoE或者稀疏路由之类的结构不要轻易相信PTQ哪怕PTQ在标准Transformer上效果不错遇到分布尖峰的结构照样会翻车。5.2 推理加速算子融合和KV Cache打磨端侧部署用的推理引擎是自研的轻量runtime核心优化做三件事。第一是算子融合把LayerNorm加QKV投影合并成一个kernel减少kernel launch开销第二是Key-Value Cache的内存精简我们用INT8存储KV Cache滑动窗口层只保留窗口内最近512个token的KV稀疏层则固定缓存Top-K选中的那128个历史token突破了8K序列下KV Cache内存随长度线性增长的瓶颈第三是动态batch和PagedAttention的引入虽然这是服务端推理常见的方案但端侧runtime里也能用它来管理碎片化内存尤其是长时间对话场景里效果明显。这三项优化合在一起xLLM在骁龙8 Gen2上的首token延迟大约是180ms续生token延迟约25ms/token基本满足实时对话的体验要求。如果你要在端侧部署类似的模型建议优先动手算子融合和KV Cache优化它们对体验的改善最直接不需要动模型结构成本最低。5.3 持续预训练与领域适配的一个建议模型部署上线之后一定会遇到的一个事情是领域适配需求。有人会觉得直接拿业务数据做SFT就行但我们的经验是如果领域数据量足够大比如超过几亿token先做一步持续预训练continual pretraining效果更稳。持续预训练时要把学习率压到预训练的十分之一大概2e-5级别同时保持领域数据和通用数据的配比在1比1左右否则很容易出现“灾难性遗忘”通用能力掉得飞快。在xLLM上我们用这个方案做过一次金融领域的适配持续预训练了50亿token领域专用指标提升明显通用评测集的下降控制在0.8个点以内。这个结果验证了xLLM的架构可塑性也说明轻量模型在行业落地上是可以走通“通用底座加行业持续预训练”这条路的。6. 排障速查表与实战心得6.1 常见问题速查表问题可能原因排查建议训练loss不下数据配比失衡、学习率过低先做小规模代理实验确认baseline可收敛检查数据配比中代码数学占比是否过低loss周期性尖峰数据重复或批次相似度过高检查去重是否完整暂停训练观察两步loss是否自动恢复梯度范数持续高于10学习率过高或数据噪声过大调低峰值学习率加强梯度裁剪到5MoE专家坍缩负载均衡loss权重过低把负载均衡权重从0.01提到0.03给路由加噪声量化后掉点严重激活分布尖峰PTQ校准不足改用QAT校准集覆盖目标任务类型端侧首token延迟过高模型过大或算子未融合优先合并LayerNorm和QKV投影检查是否开启了KV Cache量化领域适配后通用能力下降持续预训练学习率过高或通用数据不足学习率降到预训练十分之一通用数据比例保持在50%左右6.2 两条最想分享的实战心得第一小模型的训练日志一定要留足信息。我们每次实验都会记录包括梯度范数、激活分布的分位数、各层参数更新幅度在内的十几个指标。很多人只盯loss曲线这是不够的——loss是结果的投影梯度范数和激活分布才是过程的原因。xLLM项目的几次关键调优全部都是先通过梯度范数的异常发现端倪再顺着去查数据或结构问题而不是对着loss曲线盲猜。第二架构设计一定要留“后门”。这里的“后门”指的是可切换、可配置的模块。比如xLLM的注意力机制虽然预训练用的是滑窗加稀疏但代码里保留了切换到全局注意力的接口MoE的路由专家数量也是可配置的。为什么这么做因为预训练阶段结束之后下游任务很可能会对某些能力有特殊要求你不可能为了一个任务重新预训练整个模型但可以微调数据量不大的情况下通过结构调整来适配。我们在长文档场景里就把部分稀疏层切换成了全局注意力效果立竿见影。如果当时把架构写死就得推倒重来了。6.3 后续可以扩展的方向xLLM第一轮预训练完整体验下来最让我兴奋的不是最终的评测分数而是它验证了一条以前不太被重视的路径小模型可以通过架构的精细设计和数据的动态配比在有限的算力预算内逼近大模型的水平。当然它离真正的通用人工智能还差得很远很多边角case依然会翻车推理类的极限任务也兜不住。接下来团队打算在两个方向上继续试一是把滑窗加稀疏注意力的比例再做一版更激进的配置试探参数量和效果悬崖的边界在哪里二是给路由专家加更细粒度的领域标签让模型在处理特定领域数据时能优先激活对应的专家路径。后一个方向如果跑通小模型做行业适配的成本还能再降一截。如果后续有进展我会再来分享。
