迈向资源高效的大语言模型:蒸馏管道的端到端能耗核算

arXiv cs.LG 论文

摘要

本文提出了一种针对大语言模型蒸馏管道的端到端能耗核算框架,通过分阶段测量能耗并构建能耗-质量帕累托前沿,揭示了此前被忽视的教师模型侧成本。

arXiv:2605.13981v1 公告类型:新 摘要:大规模语言模型的部署增加驱动了GPU需求和数据中心规模的激增,引发了对电力使用、电网压力以及现代AI工作负载影响的担忧。蒸馏通常被视为获得更便宜、更高效模型的最有效途径之一,但这些说法很少考虑完整的端到端能耗和资源成本,包括关键的教师模型侧工作负载,如数据生成、logit缓存和评估。我们提出了一个全面的能耗核算框架,通过详细的分阶段GPU设备功耗跟踪来测量蒸馏管道的完整计算成本。在我们的实验中,我们分离并记录了不同阶段的实证能耗使用,并系统测量了两种常见蒸馏方法的能耗和排放:经典的基于logit的知识蒸馏和合成数据监督微调,构建了能耗-质量帕累托前沿,揭示了此前被忽视的成本。基于这些测量和分析,我们得出了在能耗和预算约束下选择蒸馏方法和超参数的实用设计规则,并发布了一个开源测量工具和核算协议,为可比较、可复现的蒸馏研究提供标准化基础,并明确计及完整管道能耗影响。
查看原文
查看缓存全文

缓存时间: 2026/05/15 06:26

# 迈向资源高效的大语言模型:蒸馏流程的全流程能量核算

来源:https://arxiv.org/html/2605.13981

###### 摘要

大型语言模型的广泛部署推动了GPU需求和数据中心规模的激增,引发了人们对用电量、电网压力以及现代AI工作负载影响的担忧。蒸馏通常被认为是获取更便宜、更高效模型的最有效途径之一,但这些说法很少考虑完整的端到端能量和资源成本,包括关键的教学侧工作负载,如数据生成、logit缓存和评估。我们提出了一个全面的能量核算框架,通过详细的阶段级GPU设备功耗跟踪来测量蒸馏流程的完整计算成本。在我们的实验中,我们分离并记录了不同阶段的实验能量使用情况,并系统地测量了两种常见蒸馏方法的能量和排放:经典的基于logit的知识蒸馏和合成数据监督微调,构建了能量-质量帕累托前沿,揭示了以前被忽视的成本。基于这些测量和分析,我们推导出在能量和预算约束下选择蒸馏方法和超参数的实际设计规则,并发布了一个开源测量工具和核算协议,为可比较、可复现的蒸馏研究提供标准化基础,明确核算完整流程的能量影响。

机器学习,能源,人工智能,深度学习,ICML,效率

## 1 引言

大型语言模型(LLM)的快速部署加速了GPU需求和数据中心扩张,加剧了人们对用电量、电网压力和环境影响的担忧。从“绿色AI”的角度来看,能量应与准确性和性能一起评估。蒸馏通常被视为减少计算量的主要手段,通过产生更小的学生模型来保留教师模型的许多质量,同时实现更便宜的推理。然而,它需要大量的教师侧工作——合成数据生成、logit缓存、过滤和评估——这些工作可能堪比甚至超过学生训练,尤其是在超参数搜索(例如,KD温度、解码设置、过滤阈值)下。先前的工作通常报告学生侧的节省(FLOPs、运行时间或训练能量),同时忽略这些上游成本,使得可持续性和成本主张难以证实。

我们通过一个端到端的、蒸馏特定的能量核算框架来弥补这一差距,该框架 (i) 描绘了流程阶段,(ii) 在一致协议下测量每阶段的能量、质量和吞吐量,以及 (iii) 在跨流程比较中纳入教师侧成本。我们的分析围绕三个问题展开:

1. 1. 蒸馏流程(KD或合成监督微调(SFT))在固定硬件和训练预算下,何时能比强SFT基线提供更好的能量-质量权衡?
2. 2. 教师侧成本(生成、logit缓存、评估)与学生训练相比如何,以及它们何时主导整体能量预算?
3. 3. 在什么条件下——学生规模、序列长度、教师复用和目标质量——蒸馏相对于替代方案实际上能减少端到端能量使用和排放?

为了回答这些问题,我们做出了三个贡献:

- • 蒸馏能量协议和工具。我们正式化了蒸馏阶段、记录规则和指标,以产生KD和合成SFT的可比能量图表,使用基于NVML的GPU能量作为基准,并结合经验CPU/CO2e估计。
- • 具有阶段级核算的受控基准测试。在固定硬件、软件和训练预算下,我们对1B/7B/13B OLMo-2学生进行基准测试,并报告阶段级能量、运行时间和J/token,构建在能量-质量空间中分离不同机制的帕累托前沿。
- • 设计规则和盈亏平衡条件。我们量化了教师复用、序列长度和关键超参数(KD温度/损失权重;解码和合成数据复用)如何改变这些前沿,并推导出蒸馏在能量/排放方面真正更便宜的条件——以及何时并非如此。

综合来看,我们的结果为经济且可持续的AI提供了蒸馏特定视角:我们不将更小的学生模型视为自动高效,而是展示如何核算整个流程,测量其能量成本,并决定何时蒸馏是一种能源和资源高效的选择。

## 2 背景与相关工作

现代ML系统的能量和碳排放成本推动了“绿色AI”视角,该视角认为效率应与预测性能一起评估(Schwartz等人的研究,2020 (https://arxiv.org/html/2605.13981#bib.bib8);Strubell等人的研究,2019 (https://arxiv.org/html/2605.13981#bib.bib19);Patterson等人的研究,2022 (https://arxiv.org/html/2605.13981#bib.bib2))。先前的工作还强调需要透明报告硬件、运行时间和位置假设,以及仔细处理测量不确定性(Henderson等人的研究,2020 (https://arxiv.org/html/2605.13981#bib.bib23))。我们的协议遵循这些建议,但侧重于蒸馏流程的*直接测量*能量(第4节 (https://arxiv.org/html/2605.13981#S4)),而不是GPU小时或FLOPs等代理指标。

##### 蒸馏的能量核算。

最近关于AI环境报告的生命周期审查将训练后适应确定为一个主要盲点。尽管微调、蒸馏、量化及相关方法的使用日益增多,但这些方法很少被单独测量,并且创建蒸馏或压缩模型所投入的能量也很少被报告或与下游推理节省进行比较(Lambert和Luccioni的研究,2026 (https://arxiv.org/html/2605.13981#bib.bib20))。

少量工作研究了蒸馏的环境成本和影响。Rafat等人的研究表明,用于CNN的KD可能碳排放密集,并强调了能量感知调优的重要性(Rafat等人的研究,2023 (https://arxiv.org/html/2605.13981#bib.bib24))。Yuan等人比较了蒸馏和非蒸馏的NLP模型,重点关注推理时的能量和运行时间,同时将教师训练和蒸馏流程视为沉没成本(Yuan等人的研究,2024 (https://arxiv.org/html/2605.13981#bib.bib22))。这些结果表明蒸馏并非天生“绿色”,但现有的研究通常没有提供一个阶段级、端到端的整体核算,该核算在一致预算下分离教师侧工作负载(例如,生成/logit缓存、过滤)与学生训练和评估,并提供如何减少这些成本的建议。

##### 测量工具和方法论。

能量报告通常依赖于估算工具链,如CodeCarbon和Experiment Impact Tracker(Courty等人的研究,2024 (https://arxiv.org/html/2605.13981#bib.bib16);Henderson等人的研究,2020 (https://arxiv.org/html/2605.13981#bib.bib23)),以及比较遥测测量与基于模型的估算并记录估算器误差模式的工作(Bouza等人的研究,2023 (https://arxiv.org/html/2605.13981#bib.bib10);Bannour等人的研究,2021 (https://arxiv.org/html/2605.13981#bib.bib13))。这些工具说明了保真度-便利性之间的权衡:仅使用估算器的方法易于部署,但可能错误估算设备能量,而基于遥测的记录则需要更紧密的集成。我们的框架将基于NVML的GPU遥测作为基准信号,结合CPU能量和CO2e的轻量级估算器,打包成一个具有明确阶段边界和记录规则的可复用工具。

## 3 蒸馏流程设置

我们研究三种标准机制:基线监督微调(SFT)、基于logit的知识蒸馏(KD)和在单个完全开放的LLM家族上的合成监督微调(合成SFT)。所有运行共享相同的硬件、软件栈、分词器和数据预处理。

### 3.1 硬件与环境

我们所有的实验都在单个NVIDIA H100 SXM 80 GB GPU和16个Intel Xeon Gold 6442Y CPU核心上运行,使用独占节点以避免来自相邻GPU上运行的作业的噪声。软件环境(Linux发行版、内核、NVIDIA驱动程序和CUDA/软件包版本)在所有运行中固定。每次实验的记录通过Git提交记录进行版本控制。完整的配置表见附录;测量工具和记录细节在第4节 (https://arxiv.org/html/2605.13981#S4) 中描述,控制环境漂移,使得观察到的能量差异纯粹来自流程结构、模型大小和超参数差异。

### 3.2 模型与任务

表1 (https://arxiv.org/html/2605.13981#S3.T1) 总结了教师/学生蒸馏中使用的模型。我们选择OLMo-2家族是因为其完全开放的权重和训练方法、跨大小的共享分词器和预训练血统,以及现有的环境和能量核算,使我们能够表示模型能量使用的完整生命周期。所有模型共享相同的分词器,我们以相同的方式对教师和学生的提示和输出进行分词,从而实现可比较的token计数和每token焦耳统计。

表1:实验中使用的教师和学生模型。
| 角色 | 模型(Hugging Face ID) | 参数量 |
|------|---------------------------|--------|
| 教师 | allenai/OLMo-2-0325-32B-SFT (https://huggingface.co/allenai/OLMo-2-0325-32B-SFT) | 32B |
| 学生 | allenai/OLMo-2-0425-1B (https://huggingface.co/allenai/OLMo-2-0425-1B) | 1B |
| 学生 | allenai/OLMo-2-1124-7B (https://huggingface.co/allenai/OLMo-2-1124-7B) | 7B |
| 学生 | allenai/OLMo-2-1124-13B (https://huggingface.co/allenai/OLMo-2-1124-13B) | 13B |

为避免特定领域的结论,我们在三个监督工作负载上实例化流程:指令遵循(TULU-3 SFT混合)(Lambert等人的研究,2024 (https://arxiv.org/html/2605.13981#bib.bib12))、数学推理(OpenR1-Math-220k)(Zhao等人的研究,2025 (https://arxiv.org/html/2605.13981#bib.bib14))和代码生成(Open-R1 Codeforces)(Penedo等人的研究,2025 (https://arxiv.org/html/2605.13981#bib.bib11))。所有数据集都通过一个统一的OLMo-2兼容聊天式提示流程进行格式化。TULU是我们大多数实验的主要设置,而数学和代码工作负载作为鲁棒性检查。在其它方面相同的运行中,我们观察到数据集之间的端到端能量差异相对于流程选择和学生规模的更大影响来说很小。因此,在第6节 (https://arxiv.org/html/2605.13981#S6) 中,我们报告主要能量指标(kWh和J/token)作为每个流程-大小配置在各数据集上的平均值。

### 3.3 蒸馏流程

我们考虑两种机制,它们反映了面向部署的LLM实际中使用的最常见的蒸馏方法:基于logit的KD和合成SFT。在所有机制中,优化器、调度、精度、有效批大小和早停法都是固定的(见第5节 (https://arxiv.org/html/2605.13981#S5))。下面的三个流程——直接SFT、基于logit的KD和合成序列蒸馏——在一个共享、受控的环境下提供了一组最小但有代表性的蒸馏模式,支持后续章节中的阶段级能量前沿和盈亏平衡分析。

#### 3.3.1 基于logit的知识蒸馏(KD)

这代表了Hinton等人(Hinton等人的研究,2015 (https://arxiv.org/html/2605.13981#bib.bib1))描述的传统的KD机制,其中学生通过离线蒸馏在两个阶段中训练以匹配教师的token级分布:

##### 教师logit缓存。

对于固定的训练语料库,教师(一个32B模型)以推理模式运行,从给定数据集的输入中生成样本。对于每个token位置,我们缓存前k个logit(k=100)和索引。

##### 学生KD训练。

给定缓存的教师分布p_t和硬标签y_hard,学生分布p_s训练如下:

L_KD(θ_s) = α CE(y_hard, p_s) + (1-α) T^2 KL(p_t^(T) ∥ p_s^(T)),  (1)

其中T是蒸馏温度,p^(T)表示由T软化的分布,α∈[0,1]权衡硬标签监督和软标签匹配。我们的核心网格使用默认选择α=0.5和T=1;两者都在敏感性实验中变化(见第6节 (https://arxiv.org/html/2605.13981#S6))。

#### 3.3.2 合成监督微调(SFT)

通常被称为“更便宜”的蒸馏方法,合成SFT实现了数据或序列蒸馏,即教师生成输出,然后将其视为学生监督微调的硬标签。

##### 教师数据生成。

对于每个数据集,我们保留原始提示(指令、数学问题、编程任务),并用教师生成的响应替换原始数据集标签,使用标准化的解码配置(具有固定top-p、温度和最大长度的核采样;见第5节 (https://arxiv.org/html/2605.13981#S5))。这个生成过程每个数据集运行一次,得到的合成语料库在学生大小和超参数设置之间复用。

##### 学生合成SFT。

学生使用标准自回归交叉熵目标在合成数据集上进行微调:

L_SFT(θ_s; x, y) = -∑_{t=1}^s log p_{θ_s}(y_t | x, y_{<t}),  (2)

其中x是输入提示,y = (y_1, ..., y_s)是目标序列,θ_s是学生参数,p_{θ_s}是学生条件token分布。在合成SFT机制中,目标y_t是教师对原始提示的延续。

#### 3.3.3 基线监督微调

为了提供一个非蒸馏的参考,我们还训练同样的学生在基线SFT机制中,模型直接使用原始数据集标签训练,目标函数与公式2 (https://arxiv.org/html/2605.13981#S3.E2) 相同,但使用原始数据集的目标标签y。我们使用一个全局训练预算和基于验证损失的早停法,并选择最佳检查点进行评估。这个基线锚定了质量和能量,并作为衡量引入教师的增量成本和潜在收益的对照。

## 4 能量核算

我们采用一个蒸馏特定的能量核算协议,该协议符合机器学习中系统性能量报告的最佳实践建议(Henderson等人的研究,2020 (https://arxiv.org/html/2605.13981#bib.bib23)),并针对我们的流程设计进行了定制。目标是测量*端到端*能量,使其具有可复现性,揭示蒸馏流程不同阶段的能量负载分布,并在不同模型大小和实验之间具有可比性。

### 4.1 阶段级协议

每次运行分解为具有明确开始/结束时间戳的不相交阶段。我们将总蒸馏能量定义为教师侧工作、学生训练和评估的总和,并将这些项映射到记录阶段:

- • E_prerun – 独立的、一次性环境稳定和烟雾

相似文章