很久没有 @character_ai 的消息了,今天突然看到他们发布了新模型「Kaiju」。
Kaiju 不是追求学术基准的通用模型,而是专注于实际生产环境中的对话性能和部署效率。这篇博客分享了从架构设计、训练优化到数据策略和安全对齐的全过程,值得看看。
模型概述
Kaiju系列包括三个规模:小(13B 参数)、中(34B 参数)和大(110B 参数)。这些模型基于稠密 Transformer 架构,采用自回归生成方式,专为对话场景优化。不同于追求高基准分数的模型,Kaiju 优先考虑推理效率和用户互动质量,例如快速响应和生动对话。这反映了团队的观点:生产性能比学术指标更重要。
架构创新
文章的核心在于多项效率优化技术,确保模型在规模化部署时保持高性能:
· 多查询注意力(MQA):减少键-值(KV)缓存大小,提高对话推理效率。尽管在某些基准测试中略有质量损失,但对对话任务影响最小。
· 滑动窗口注意力:限制注意力范围为1024个 token,并交替使用全局层(比例 6:1)。这显著降低长上下文计算量,同时不牺牲检索准确性,避免了传统方法如注意力沉没(attention sinks)的复杂性。
· 跨层 KV 共享:每2-3层共享一个 KV 缓存,进一步压缩内存,无明显准确性下降。
· Int8 量化:权重和矩阵乘法使用8位整数存储和计算,比16位浮点数快20-30%,通过量化感知训练(QAT)确保精度接近原生水平。
· 其他优化:预层归一化(RMSNorm)和动态钳位(clamping)增强稳定性,避免训练中的数值问题。
模型训练
训练在 Google Cloud H100 GPU 集群上进行,采用模型并行策略(节点内张量和序列并行,节点间全分片数据并行)。关键效率技巧包括:
· 低精度计算:权重和 KV 使用 Int8,前向激活和局部梯度用 BF16,梯度累积和权重用 FP32。
· 梯度压缩:引入 Squinch 算法,将梯度压缩至6位(块状、对数均匀建模),减少通信开销。
· 稳定性增强:对于小模型,使用 Bungee 虚拟标量避免 Int8 溢出;还实验了三元权重更新(每参数1.6位),进一步压缩存储。
数据策略
数据混合是 Kaiju 成功的关键。作者将数据分为两类:
· MMLU Max:针对 AGI 基准,包含网络规模文本、代码和合成数据,使用 T5 嵌入计算相似度。
· Production Max:聚焦用户互动,强调指令跟随。
训练后期采用退火策略,逐步增加指令数据和 MMLU 相关内容,以平衡基准性能和实际应用。这避免了过度优化单一指标,确保模型在对话中更自然。
安全与对齐
安全是文章强调的重点,采用多阶段方法:
· 监督微调:使用高质量数据调整模型行为。
· 强化学习:基于用户反馈(如“滑动”或偏好)进行修改版在线直接偏好优化(DPO),提升互动质量。
· 分类器训练:集成可选分类器头,提供标记级安全指标。
· 推理时控制:使用分类器引导的束搜索,确保生成内容安全可靠。
挑战与解决方案
文章客观讨论了权衡取舍:例如,MQA 虽高效但可能影响基准分数,作者通过专注非 AGI 任务(如对话)化解;长上下文计算昂贵,则用滑动窗口和 KV 共享应对;低精度训练易不稳定,则引入 QAT 和 Bungee 等创新。整体上,这些解决方案证明了效率优化不一定牺牲质量,尤其在生产环境中。
博客地址:
邵猛,中年失业程序员 😂
专注 - Context Engineering, AI Agents.
分享 - AI papers, apps and OSS.
ex Microsoft MVP
合作 - 私信/邮箱:shaomeng@outlook.com
📢 公众号/小红书: AI 启蒙小伙伴