深度学习能耗测量与优化

作者：Jae-Won Chung

Zeus logo

本文作者是Jae-Won Chung，密歇根大学博士生，ML.ENERGY Initiative 的负责人。

深度学习消耗相当多的能源。例如，在 AWS p4d 实例上训练一个 2000 亿参数的大模型 (LLM) 消耗了大约 11.9 GWh 的电量（来源：CIDR 2024 主题演讲），这个电量足以独自为一个以上一千个美国家庭提供一年的电力。

Zeus 是一个开源工具箱，用于测量和优化深度学习工作负载的能耗。我们的目标是通过提供具有最小假设的可组合工具，使基于准确测量的能耗优化对于各种深度学习工作负载和设置尽可能简单。

Zeus 主要提供两种类型的工具

可通过编程方式和命令行使用的 GPU 能耗测量工具
若干能耗优化工具，用于寻找最佳的 ML 和/或 GPU 配置

Zeus 可以使希望以下的用户受益

测量和优化他们的电费成本
减少 GPU 的散热（通过降低功耗）
报告研发过程中的能源使用情况
减少用电产生的碳足迹

第一部分：测量能源

就像性能优化一样，准确测量是有效能耗优化的基础。像硬件最大功耗这样的常用功耗估算代理有时与实际测量值相差甚远。

为了使能耗测量尽可能简单和透明，Zeus 提供的核心实用工具是 ZeusMonitor 类。让我们看看实际的代码片段

from zeus.monitor import ZeusMonitor

# All four GPUs are measured simultaneously.
monitor = ZeusMonitor(gpu_indices=[0,1,2,3])

# Measure total time and energy within the window.
monitor.begin_window("training")
for e in range(100):

    # Measurement windows can arbitrarily be overlapped.
    monitor.begin_window("epoch")
    for x, y in train_dataloader:
        y_hat = model(x)
        loss = criterion(y, y_hat)
        loss.backward()
        optim.step()
    measurement = monitor.end_window("epoch")
    print(f"Epoch {e}: {measurement.time} s, {measurement.total_energy} J")

measurement = monitor.end_window("training")
print(f"Entire training: {measurement.time} s, {measurement.total_energy} J")

您在上面看到的是一个典型的 PyTorch 训练循环，它使用四个 GPU 进行数据并行训练。在循环内部，我们创建了一个 ZeusMonitor 实例，并传入要监控的 GPU 索引列表。然后，使用该监控器，我们可以通过配对调用 begin_window 和 end_window 来测量训练脚本中任意执行窗口的时间和能耗。只要它们的名称不同，多个窗口可以任意重叠和嵌套，而不会影响彼此的测量。

ZeusMonitor 在窗口周围增加的开销非常小——通常是几毫秒。这使得 ZeusMonitor 可以用于各种应用。例如

ML.ENERGY 排行榜：第一个关于大模型 (LLM) 文本生成消耗多少能量的开源基准测试。
ML.ENERGY 竞技场：一项在线服务，允许用户并排比较大模型 (LLM) 的响应，基于响应质量和能耗。

请参阅我们的博客文章，深入了解准确的 GPU 能耗测量技术细节。

第二部分：优化能源

让我向您介绍 Zeus 提供的两种能耗优化器。

全局功耗限制优化器

GPU 允许用户配置其最大功耗，称为功耗限制。通常，当您将 GPU 的功耗限制从默认最大值降低时，计算速度可能会略微变慢，但您会节省不成比例的更多能源。Zeus 中的 GlobalPowerLimitOptimizer 会自动全局寻找所有 GPU 的最佳功耗限制。

from zeus.monitor import ZeusMonitor
from zeus.optimizer.power_limit import GlobalPowerLimitOptimizer

# The optimizer measures time and energy through the ZeusMonitor.
monitor = ZeusMonitor(gpu_indices=[0,1,2,3])
plo = GlobalPowerLimitOptimizer(monitor)

for e in range(100):
    plo.on_epoch_begin()
    for x, y in train_dataloader:
        plo.on_step_begin()

        y_hat = model(x)
        loss = criterion(y, y_hat)
        loss.backward()
        optim.step()

        plo.on_step_end()
    plo.on_epoch_end()

在我们熟悉的 PyTorch 训练循环中，我们实例化了 GlobalPowerLimitOptimizer 并向其传递了一个 ZeusMonitor 实例，优化器通过它看到 GPU。然后，我们只需要让优化器知道训练进度（步数和 epoch 边界），优化器就会透明地执行所有必要的性能分析并收敛到最佳功耗限制。

如果您正在使用 HuggingFace 的 Trainer 或 SFTTrainer，集成甚至更容易

from zeus.monitor import ZeusMonitor
from zeus.optimizer.power_limit import HFGlobalPowerLimitOptimizer

# ZeusMonitor actually auto-detects CUDA_VISIBLE_DEVICES.
monitor = ZeusMonitor()
pl_optimizer = HFGlobalPowerLimitOptimizer(monitor)

# Pass in the optimizer as a Trainer callback. Also works for SFTTrainer.
trainer = Trainer(
    model=model,
    train_dataset=train_dataset,
    ...,
    callbacks=[pl_optimizer],
)

HFGlobalPowerLimitOptimizer 对 GlobalPowerLimitOptimizer 进行了封装，使其能够自动检测步数和 epoch 边界。我们在此处提供了集成示例此处，包括使用 QLoRA 运行 Gemma 7B 有监督微调。

现在，我们知道如何集成优化器了，但什么是最佳功耗限制呢？我们知道不同的用户在权衡时间和能源方面可能有不同的偏好，因此我们允许用户指定一个 OptimumSelector（本质上是策略模式）来表达他们的需求。

# Built-in strategies for selecting the optimal power limit.
from zeus.optimizer.power_limit import (
    GlobalPowerLimitOptimizer,
    Time,
    Energy,
    MaxSlowdownConstraint,
)

# Minimize energy while tolerating at most 10% slowdown.
plo = GlobalPowerLimitOptimizer(
    monitor,
    MaxSlowdownConstraint(factor=1.1),
)

一些内置策略包括“最小化时间” (Time，即使在较低功耗限制下，某些工作负载几乎没有减速，这仍然可能降低功耗限制)、“最小化能源” (Energy)、“介于两者之间” (ZeusCost) 和“在最大减速限制下最小化能源” (MaxSlowdownConstraint)。用户还可以根据需要创建自己的最佳选择器。

流水线频率优化器

流水线频率优化器，基于我们的研究论文 Perseus，是我们关于大型模型训练（如 GPT-3）能耗优化的最新工作。Perseus 可以在训练吞吐量没有或可忽略的下降的情况下降低大型模型训练的能耗。我们将简要介绍其工作原理。

one iteration of training with four stage pipeline parallelism

上图展示了使用四阶段流水线并行以 1F1B 调度运行的一次训练迭代的可视化。每个框代表一次前向或后向计算，并根据其功耗着色。

这里的关键观察是，当模型被划分为流水线阶段时，很难将其完全等大小地切分。这导致前向/后向计算框的宽度各不相同，因此在框之间存在计算空闲时间。您会注意到，那些较小的框即使比宽的框运行得稍慢，整体关键路径（蓝色线）也不会发生变化。

one iteration of training with four stage pipeline parallelism

这就是 Perseus 自动完成的工作。基于性能分析，它识别出不在关键路径上的计算框，并确定每个框精确的减慢量，从而最小化能耗。如果操作正确，我们减慢的计算会消耗更少的功率和能源，但流水线的整体迭代时间不会改变。

请参阅我们的指南，开始使用 Perseus！

总结

对于那些运行自己的本地计算设施的用户来说，能源消耗和由此产生的电费是不能轻易忽视的问题。从更广阔的视角来看，能源消耗不仅关乎电费，也关乎数据中心的电力供应。当数千个 GPU 在集群中运行时，寻找稳定、负担得起且可持续的电力来源来为数据中心供电变得越来越具挑战性。找到在不造成不成比例的减速的情况下大幅减少能源消耗的方法，可以降低平均功耗，从而有助于解决电力供应挑战。

借助于 Zeus，我们希望迈出深度学习能耗测量与优化的第一步。

想知道接下来做什么？这里有一些有用的链接

第一部分：测量能源

第二部分：优化能源

全局功耗限制优化器

流水线频率优化器

总结

文档

教程

资源