跳转到主要内容
社区

深度学习能耗测量与优化

作者： 2024 年 5 月 11 日2024 年 11 月 21 日暂无评论
Zeus logo

本文由Jae-Won Chung撰写，他是密歇根大学的博士生，也是ML.ENERGY Initiative的负责人。

深度学习消耗了相当多的能量。例如，在 AWS p4d 实例上训练单个 200B LLM 消耗了大约 11.9 GWh（来源：CIDR 2024 主题演讲），这个数量足以单独为一千多个普通美国家庭供电一年。

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

Zeus 主要提供两类工具：

  1. 程序化和命令行 GPU 能耗**测量**工具
  2. 多种能耗**优化**工具，用于寻找最佳的 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 提供的两种能耗优化器。

GlobalPowerLimitOptimizer

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

一些内置策略包括“最小化时间”（时间，这仍然可能将功耗限制从默认值降低，因为有些工作负载即使在较低的功耗限制下也几乎没有减速），“最小化能耗”（能耗），“介于两者之间”（ZeusCost），以及“在最大减速下最小化能耗”（MaxSlowdownConstraint）。用户还可以根据需要创建自己的最佳选择器。

PipelineFrequencyOptimizer

流水线频率优化器，基于我们的研究论文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，迈出深度学习能耗测量与优化的第一步。

想知道接下来去哪里？这里有一些有用的链接：