想象一下,你拥有一间专为大规模人群烹制美食而设计的宏大高端厨房。这间厨房就是你的大型语言模型(LLM)。
传统上,如果你要举办小型家庭晚餐、中型聚会和大型宴会,你需要建造三间完全独立的厨房。你必须购买三套烤箱,雇佣三组不同的厨师团队,并支付三个独立的建设项目费用。这极其昂贵且浪费。
Star Elastic 是一项改变游戏规则的新发明。它不再建造三间厨房,而是建造一间超级灵活的厨房,能够瞬间自我调整以适应任何需求。
以下是其工作原理,分解为简单的概念:
1. “一间厨房,多种尺寸”的诀窍
通常,为了获得更小、更便宜的 AI 模型,研究人员必须先训练一个巨型模型,然后尝试通过切除部分(如修剪树木)来“缩小”它。这既缓慢又昂贵,而且往往导致模型的智能程度不如从头训练出来的模型。
Star Elastic 的做法不同。它采用一个巨大的“父”模型(Nemotron Nano v3),并教会它一项特殊技能:如何同时具备多种不同的尺寸。
- 这就像俄罗斯套娃(Matryoshka)。在这个 300 亿参数的巨大娃娃内部,有一个完美的 230 亿参数娃娃,而在这个娃娃内部,又有一个完美的 120 亿参数娃娃。
- 它们都生活在同一个“房子”里(同一个计算机文件)。你不需要下载三个不同的文件;只需打开那个大文件,然后告诉它:“我今天只需要小版本。”
2. “智能路由器”(厨房经理)
模型如何知道该使用哪些部分?它使用一个可学习的路由器。
- 想象一位厨房经理正在查看你的订单。
- 如果你点一份简单的沙拉(简单的问题),经理会说:“好的,我们只用小型搅拌机和基础刀具套装。”
- 如果你点一份复杂的舒芙蕾(难的数学问题),经理会说:“我们需要大烤箱、重型搅拌机以及所有厨师!”
- 论文表明,这位经理经过训练,能够确切知道“厨房”的哪些部分最重要。它为小版本保留最佳工具,仅在需要大版本时才添加额外的重型工具。
3. “思考与回答”策略
这是论文中最巧妙的技巧,称为弹性预算控制。
- 当 AI 解决难题时,它通常做两件事:思考(推理步骤)和回答(撰写最终结果)。
- 旧方法:AI 在思考和回答时使用相同的“大脑尺寸”。这就像用一辆巨大的、耗油的卡车去杂货店,然后再去邮局,即使邮局就在街对面。
- Star Elastic 方法:AI 可以切换档位!
- 第一阶段(思考):它使用更小、更快的模型来头脑风暴并思考问题。这既便宜又快速。
- 第二阶段(回答):一旦思考完成,它就切换到更大、更智能的模型来撰写最终答案。这确保了答案的完美。
- 结果:你获得了巨型大脑的准确性,同时享受了小型大脑的速度和成本。论文声称,与使用单一固定尺寸相比,这能使 AI 的准确率提高16%,速度提升1.9 倍。
4. “魔法切片”(零样本提取)
通常,如果你想要一个更小的模型,你必须训练它数月。有了 Star Elastic,“切片”是瞬间完成的。
- 因为该模型从一开始就被训练为具有灵活性,你可以**零样本(zero-shot)**地“切下”小版本或中版本。
- 这意味着你不需要重新训练它们。你只需取出大文件,应用“切割”,砰——你就立即拥有了一个可用的、高质量的小模型。
- 论文声称,与从头构建模型相比,这节省了360 倍的训练成本;与之前的压缩方法相比,节省了7 倍的成本。
5. “小行李箱”(量化)
最后,论文讨论了如何将这些模型进一步缩小,以便在手机或小型计算机上运行。
- 他们使用了一种称为量化感知蒸馏的技术。想象一下,将一幅沉重的高清画作转化为一个占用空间极小但仍完美呈现的数字文件。
- 他们创建了适配4 位或8 位格式(非常小的数字足迹)的模型版本。
- 即使在这些微小的格式中,“俄罗斯套娃”技巧仍然有效。你仍然可以从一个微小的文件中切分出小型、中型和大型版本。
胜利总结
- 成本:你只需训练一次,就能获得多个模型。这就像购买一张主题公园门票,可以乘坐所有过山车,而不是为每个游乐设施单独购票。
- 速度:通过使用小大脑进行思考,大大脑进行回答,你节省了时间和金钱。
- 存储:你只需存储一个文件,即可访问三种不同尺寸的模型。
简而言之,Star Elastic 阻止了我们为每项工作构建独立、僵化的 AI 模型。相反,它构建了一个变色龙般的 AI,能够瞬间改变其尺寸和算力以适应任务,在节省巨额资金和时间的同时,实际上变得更智能。
技术摘要:Star Elastic
问题陈述
训练具有不同参数规模的大语言模型(LLM)家族目前成本过高且效率低下。传统方法需要为每个模型规模进行独立的训练运行,从而成倍增加了计算和存储成本。尽管最近的压缩方法(例如结构化剪枝和知识蒸馏)通过从大型父模型衍生出小型模型来降低成本,但它们仍需要每个压缩模型数百亿个 token,并且往往无法支持异构架构。
此外,现有的高效推理框架存在架构僵化的问题。它们无论 token 的难度如何都分配静态资源,强制在推理的“思考”(推理)和“回答”阶段使用单一模型规模。当这两个阶段的复杂度差异较大时,这会导致计算效率低下。最后,当前的弹性嵌套模型(例如 Nemotron Elastic)仅限于 Mamba-Attention 架构,缺乏对日益流行的混合 Mamba-Attention-专家混合(MoE)设计的支持。
方法论:Star Elastic
Star Elastic 是一种后训练方法,旨在利用单次训练运行的计算资源,从单个父推理模型生成 N 个嵌套子模型。它针对混合 Mamba-Attention-MoE 架构,并引入了用于动态推理的弹性预算控制。
1. 弹性模型构建
该流程通过三个阶段将现有的 LLM 转换为弹性模型:
- 重要性估计:根据重要性对组件进行排序以指导剪枝。
- 宽度:将基于激活的评分扩展到混合架构。计算嵌入通道、FFN/MoE 通道、Mamba 头/通道以及注意力头的分数。对于 MoE 层,使用路由加权专家激活剪枝(REAP),根据路由权重和激活范数衡量专家对输出幅度的贡献。
- 深度:使用完整模型 logits 与移除特定层后的 logits 之间的归一化均方误差(MSE)迭代估计层重要性。
- 弹性公式化:模型支持沿宽度和深度轴进行嵌套。
- 宽度:子网络根据重要性排名选择组件子集(嵌入维度、头、专家、FFN 通道)。异构配置允许每层独立选择。
- 深度:二进制选择向量确定活跃层,首先移除重要性较低的层。
- 弹性训练:一个可学习的路由器将用户提供的预算规范映射到最优架构配置。
- 路由器使用 Gumbel-Softmax 生成用于训练的可微掩码。
- 损失函数:结合来自非弹性化父模型的知识蒸馏(KD)与路由器损失,后者使所选配置与硬件约束(延迟、内存)保持一致。
- 两阶段课程:
- 阶段 1(短上下文):在所有目标规模上进行均匀预算采样,以防止梯度竞争。
- 阶段 2(扩展上下文):非均匀采样优先处理较大预算(例如,50% 用于完整模型),以确保在较大变体中保留推理能力,同时使较小模型适应长上下文推理链。
2. 弹性预算控制
Star Elastic 在推理阶段启用动态模型选择:
- 思考阶段:使用较小的子模型(MS)以最小开销生成广泛的推理轨迹。
- 回答阶段:切换到较大的子模型(ML)以更高的保真度合成最终答案。
- 缓存兼容性:该架构在嵌套模型之间保留了 Mamba 和注意力层结构,允许 KV 和 SSM 缓存状态在子模型之间移植(尽管当前框架对无缝复用的支持有限,但该架构使其成为可能)。
3. 量化感知蒸馏(QAD)
为了支持在量化环境(FP8 和 NVFP4)中的部署,Star Elastic 应用了 QAD。
- FP8:通过对 BF16 弹性检查点进行训练后量化(PTQ)实现。
- NVFP4:需要一个简短的 QAD 阶段,其中使用相同的嵌套掩码层次结构,从 BF16 教师模型蒸馏出量化学生模型。这恢复了 PTQ 期间损失的精度,同时保留了零样本切片能力。
主要贡献
- 首个混合 MoE 的弹性后训练:将 Nemotron Elastic 框架从 Mamba-Attention 扩展到混合 Mamba-Attention-MoE 架构,引入了用于专家排名的 REAP 和每层异构选择。
- 弹性预算控制:引入了一种动态推理方案,将不同的子模型分配给思考和回答阶段,优化了精度 - 延迟帕累托前沿。
- 可学习路由器:一个通过端到端训练自动确定嵌套架构的路由器,在不过度解耦配置与训练目标的情况下,同时优化精度和资源约束。
- 显著的成本降低:与从头预训练相比,训练 token 减少了360 倍;与生成模型家族的最先进压缩方法相比,减少了7 倍。
- 量化弹性:生成了保留零样本切片能力的嵌套 NVFP4 和 FP8 检查点,使得可以从单个量化文件中部署多种模型规模。
实验结果
该方法应用于Nemotron Nano v3 (30B/3.6A),使用160B 训练 token生成了23B (2.8A) 和 12B (2.0A) 变体。
- 精度:所有嵌套模型均达到或超过同等规模独立训练基线的水平。例如,23B 弹性变体在 LiveCodeBench 和 IFBench 上优于独立训练的 Qwen3-30B-A3B。
- 效率:
- 训练:与预训练相比 token 减少 360 倍;与顺序压缩相比减少 7 倍。
- 部署内存:将 12B、23B 和 30B 变体存储在单个弹性检查点中需要58.9 GB(BF16),而单独的检查点需要126.1 GB。
- 吞吐量:嵌套模型在 H100 GPU 上的吞吐量最高提升2.4 倍。
- 弹性预算控制:MS→ML 策略(小思考 → 大回答)与静态预算控制相比,实现了高达16% 的更高精度和1.9 倍的更低延迟,显著推进了精度 - 速度帕累托前沿。
- 量化:FP8 和 NVFP4 变体恢复了97–100% 的 BF16 精度,同时支持在消费级 GPU(例如 RTX 5080)上部署,而 BF16 变体在这些设备上会出现内存溢出(OOM)。
意义与主张
该论文声称,Star Elastic 通过以下方式从根本上改变了推理 LLM 的部署范式:
- 消除冗余训练:使模型家族的创建成为单次过程,而非一系列昂贵且独立的运行。
- 解锁动态推理:超越静态架构,允许资源分配匹配推理与答案合成的特定计算需求。
- 弥合与量化的差距:证明弹性嵌套结构可以在低精度格式(NVFP4/FP8)中保留,从而在资源受限的硬件上实现高性能推理,同时不牺牲多模型规模的灵活性。
作者将这项工作定位为迈向“多合一”推理 LLM 的一步,这些模型在精度上处于最先进水平,并且在各种部署场景中具有极高的成本效益。
每周获取最佳 machine learning 论文。
受到斯坦福、剑桥和法国科学院研究人员的信赖。
请查收邮箱确认订阅。
出了点问题,再试一次?
无垃圾邮件,随时退订。