EnergyLens: 面向多GPU大语言模型推理优化的预测性能耗感知探索

# 面向多GPU LLM推理优化的预测性能源感知探索
来源：https://arxiv.org/html/2605.14249
宋智烨 麻省理工学院 剑桥，马萨诸塞州 02139 zhiye@mit\.edu &李景美 麻省理工学院 剑桥，马萨诸塞州 02139 kyungmi@mit\.edu &李恩京 IBM研究院 约克镇高地，纽约州 10598 eunkyung\.lee@us\.ibm\.com &张欣 IBM研究院 约克镇高地，纽约州 10598 xzhang@us\.ibm\.com &Tamar Eilam IBM研究院 约克镇高地，纽约州 10598 eilamt@us\.ibm\.com &Anantha P\. Chandrakasan 麻省理工学院 剑桥，马萨诸塞州 02139

###### 摘要

我们提出 EnergyLens，一个面向能源感知的大型语言模型（LLM）推理优化的端到端框架。随着LLM规模的扩大，预测并降低其能源足迹对可持续性和数据中心运营变得至关重要，但现有方法要么需要生产级代码和高昂的剖析成本，要么无法准确捕捉多GPU的能源行为。因此，实践者在决定优先采用哪些优化方案，以及在穷举剖析不可行的情况下选择现有部署配置时，缺乏有效工具。EnergyLens 通过一个直观的、基于einsum的接口来弥补这一缺口，该接口能够捕捉包含融合、并行化以及计算-通信重叠在内的LLM规格，并结合了负载不均衡感知的MoE建模，以及针对多GPU场景的经验驱动的通信能量模型。我们针对 Llama3 和 Qwen3-MoE，在张量并行和专家并行的配置下验证了 EnergyLens，在多GPU预填充和解码能量方面实现了9.25%到13.19%的平均绝对百分比误差，在 Megatron 风格的重叠配置中，跨SM分配实现了12.97%的平均绝对百分比误差。我们的能量驱动探索揭示了，不同配置在预填充和解码效率上的能量差异可达1.47倍和52.9倍，这推动了分布式服务的需求。我们进一步表明，仅凭直觉很难优化计算-通信重叠，但 EnergyLens 能够正确识别帕累托最优的重叠配置。

## 1 引言

以LLM为核心的生成式AI的到来，已深刻改变了机器学习。随着LLM扩展至数千亿参数并为持续增长的用户群体服务，其计算需求转化为巨大的能源成本和碳足迹。因此，理解和优化推理能量对可持续性、数据中心供给以及服务效率变得至关重要。

应对这一挑战需要在LLM推理开发周期的各个阶段（而不仅仅是在部署后）都能发挥作用的能量模型。在实践中，实践者需要在两种场景下获得能量估计：一是在实现之前，决定优先采用哪些优化，例如融合、并行化或计算-通信重叠；二是在部署期间，当工作负载和延迟目标变化时，在穷举剖析成本过高的情况下选择配置。在帕累托前沿难以仅凭直觉推断的场景中，用于空间探索的模型尤为重要。现有方法各自解决了部分问题，但在能源感知的LLM优化方面仍有不足。直接测量是准确的，但需要GPU访问权以及实施每个配置的大量工程工作，使得广泛的探索既缓慢又昂贵。粗略的整体方法，如基于热设计功耗（TDP）的估算，虽然易于应用，但忽略了LLM推理各阶段之间巨大的功率差异；在对A100 GPU上Llama3-70B的测量中，基于TDP的方法可能高估解码能量高达60%。LLMCO2能够预测功率，但原生不支持融合实现或混合专家（MoE）架构，也无法提供理解优化机会所需的能量归因。内核级预测器（如EnergAIzer）能够提供计算内核的归因，但未能捕捉越来越主导实际LLM服务的多GPU效应，包括并行化和计算-通信重叠。

这些局限性尤其重要，因为现代服务堆栈越来越多地将分布式执行与架构多样性相结合。模型通常部署在多GPU上以满足内存和延迟约束，而MoE模型引入了额外的专家并行执行和路由不均衡，这同时改变了内核形状和每GPU的工作分布。随着多GPU并行化对网络提出更高需求，准确的通信能量建模变得至关重要。否则，重要的优化机会可能被错失。以 Llama3-70B 在 8 GPU 张量并行配置下为例，通信占总能量的23%。计算-通信重叠进一步使延迟-能量权衡复杂化，其程度是直觉和现有工具无法捕捉的。由此产生的配置空间涵盖融合策略、并行化选项、重叠设置、批量大小等维度，每个维度都与序列长度和延迟目标相互作用。这种组合爆炸使得穷举剖析变得不切实际。

EnergyLens 通过支持在无需现有代码或GPU访问权的情况下，快速探索能量-延迟权衡空间，填补了这一空白。开发者使用我们的简单接口，仅需几行代码即可指定LLM架构以及融合、并行化和重叠设置。该框架解析这些规范，并结合经验驱动的、感知重叠的通信建模和负载不均衡感知的MoE建模，生成详细的能量分解和能量-延迟权衡。例如，为比较 Llama3-70B 预填充在7个批量大小和32个输入长度下的两种融合设置，需要在一台8 GPU服务器上进行12小时的剖析，还需大量工程工作。而使用 EnergyLens，生成相应结果和可视化只需15分钟。

我们的主要贡献如下：

- • 我们提出了 EnergyLens，一个框架，允许实践者通过几行基于einsum的接口（包含并行化和重叠注释）来描述密集型和MoE的LLM推理，并预测详细的能量和延迟分解，而无需执行跟踪、生产实现或GPU访问权。
- • 我们开发了一个分布式能量建模栈，结合了经验驱动的通信能量建模、感知重叠的聚合以及负载不均衡感知的MoE建模，从而支持张量并行、专家并行和 Megatron 风格的计算-通信重叠。
- • 我们在 Llama3 和 Qwen3-MoE 上验证了 EnergyLens，在多GPU能量方面实现了9.25%到13.19%的平均绝对百分比误差（MAPE），在重叠评估中实现了12.97%的MAPE。我们表明，这些误差足够小，足以恢复最佳部署配置，包括非直观的与重叠相关的权衡。

## 2 相关工作

### 2.1 内核级性能建模

一些内核性能模型提供了比周期精确的GPU模拟更快的替代方案。我们将 EnergAIzer 作为主要的计算内核后端，因为它能同时提供延迟和功率估计；我们框架的模块化设计也支持其他延迟后端，包括 NeuSight 和 Li 等人的模型。

### 2.2 LLM性能建模

现有的LLM性能建模涵盖了训练和推理，但这些工具共享两个关键局限性：它们专注于延迟和吞吐量而忽略了能量，且大多数需要现有实现提供的执行跟踪，从而排除了早期设计探索。LLMCO2 优先考虑能量，但不能生成促进优化机会所需的能量分解。此外，其图神经网络假设的是标准Transformer（如 Llama2）的未融合版本。应用操作融合会改变这种拓扑结构，需要对融合实现进行剖析并重新训练模型，但论文中并未解决这些问题。LLMServingSim v2.0 在其以性能为中心的模型基础上添加了一个粗略的三态功率模型，但对活跃GPU使用单一功率数值，忽略了诸如预填充和解码阶段之间多样的功率分布。EnergyLens 解决了所有这些缺口：以能量为中心的建模、分解可见性以及无需执行跟踪的输入。

### 2.3 计算-通信重叠

近期工作通过将张量并行集合与通用矩阵乘法重叠，以隐藏分布式训练中的通信延迟；这也在分布式推理的预填充阶段得到越来越多探索。然而，对其能量影响的预测建模仍未被探索。Lee 等人通过实验描述了重叠的特征但未做预测。Hong 等人使用一个假设固定SM分配用于通信且GEMM分区线性缩放的延迟模型来优化重叠策略。该延迟模型在变化的SM分配或较小工作负载下表现不佳。

## 3 EnergyLens 框架

图1：EnergyLens 框架能够在模型实现和部署的早期阶段，在高维探索空间中实现能量-延迟优化。图̃1 展示了 EnergyLens 框架。*EnergyLens 接口*（第 3.1 节）提供了一个基于 einsum 的规范，用于描述 LLM 架构、融合策略和并行化设置。一个*轻量级解释器*（第 3.2 节）解析这些规范，以提取 GEMM 维度、算术强度和所需的通信内核。它还支持 MoE 模型中的负载不均衡建模。由此产生的内核信息输入到*重叠感知通信模型*（第 3.3 节），该模型捕捉多GPU通信能量，以及一个*计算内核模型*，该模型利用 EnergAIzer 进行单GPU内核延迟和功率估计，并支持替代的计算延迟后端，如 NeuSight 和 Li 的模型。内核级结果经过一个考虑计算-通信重叠（第 3.4 节）的聚合阶段，以生成端到端的能量分解和能量-延迟帕累托前沿。

### 3.1 EnergyLens 接口

EnergyLens 提供了一个基于 einsum 的接口，用于指定 LLM 推理操作，并添加了用于并行化和重叠的轻量级注释。我们选择 einsum 符号，因为它在 PyTorch 等框架中广泛用于表达张量操作，使得能够直观地描述并轻松扩展到分布式实现。

图̃2 展示了仅用几行代码完成的融合 MoE 完整规格，符号定义见表̃2。融合和非融合密集模型的示例见附录̃B。例如，融合 QKV 投影写作 `bsm,miKh->bsiKh`，张量并行通过在切分维度上注释（如 `parallel="K"`）来指定。同一规格覆盖预填充和解码：`s` 在预填充中表示序列长度，在解码中 `s=1`。模型维度在运行时提供，因此一个规格可以表示不同规模的模型。

```
llm_MoE.EQS=[
  op("bsm,miKh->bsiKh",parallel="K",label="QKV ␣ Projection"),\|\\label\{line:qkv\}\|
  op("attention",parallel="H",attn_eqs=attn_eqs,label="Attention"),\|\\label\{line:attn\}\|
  op("bsHh,Hhm->bsm",parallel="H",label="Output ␣ Projection"),\|\\label\{line:out\}\|
  op("bsm,Em->bsE",label="Router"),\|\\label\{line:router\}\|
  op("bsm->bsmA",label="Scatter"),\|\\label\{line:scatter\}\|
  op("ETm,EmF->ETF",parallel="E",label="Gate ␣&␣ Up ␣ Projection"),\|\\label\{line:gateup\}\|
  op("ETf,Efm->ETm",parallel="E",label="Down ␣ Projection"),\|\\label\{line:down\}\|
  op("bsAm,bsA->bsm",parallel="A",label="Reduction")\|\\label\{line:reduce\}\|
]
```

（注：由于原文后续章节未提供，翻译至此结束。）