由崔在元著

Zeus logo

本文作者为崔在元，密歇根大学博士生，ML.ENERGY 创新计划负责人。

深度学习消耗大量能源。例如，在 AWS p4d 实例上训练单个 200B LLM 模型消耗了约 11.9 GWh（来源：CIDR 2024 大会），这个数字足以为一千多户美国家庭供电一年。

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

宙斯主要提供两种类型的工具：

  1. 程序化和命令行 GPU 能耗测量工具
  2. 几种寻找最佳机器学习和/或 GPU 配置的能耗优化工具

宙斯可以造福那些想要的人

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

第一部分：测量能量

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

为了使能源测量尽可能简单和透明，核心工具宙斯提供的功能是 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 在各种应用中使用。例如：

  • 机器学习.能源领导者排行榜：首个开源基准，衡量LLM文本生成消耗的能源量。
  • 机器学习.能源竞技场：一个在线服务，允许用户根据响应质量和能源消耗将LLM响应并排比较。

查看我们的博客文章，深入了解准确 GPU 能源测量的技术细节。

第二部分：优化能源

让我来向您介绍宙斯提供的两个能量优化器。

GlobalPowerLimitOptimizer

GPU 允许用户配置其最大功耗，称为功率限制。通常，当您将 GPU 的功率限制从默认最大值降低时，计算可能会略微变慢，但您将节省不成比例的更多能源。宙斯中的 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）。用户也可以根据需要创建自己的最佳选择器。

管道频率优化器

基于我们的研究论文《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 在集群中运行，寻找稳定、经济、可持续的电力来源来为数据中心供电正变得越来越具有挑战性。找到比减速更能降低能源消耗的方法，可以降低平均功耗，从而有助于解决电力供应问题。

与宙斯一起，我们希望迈出深度学习能量测量和优化的第一步。

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