社区

深度学习能耗测量与优化

作者： 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. 多种能够寻找最佳机器学习和/或 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 Colosseum：一项在线服务，允许用户根据响应质量和能耗对比 LLM 的输出结果。

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

第二部分：能源优化

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

GlobalPowerLimitOptimizer

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）。用户也可以根据需要创建自己的最佳选择器。

PipelineFrequencyOptimizer

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

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