使用 DeepSpeed 进行内存优化

随着深度学习模型复杂性和大规模计算的不断增长，内存优化在训练过程中至关重要。DeepSpeed 提供了多种节省内存的技术，例如卸载、梯度检查点和 ZeRO。开发人员可以基于这些节省内存的技术，在标准硬件上训练非常庞大的模型。此外，这些技术还可以训练以前受硬件限制的模型。

DeepSpeed 不仅在研究领域取得了令人瞩目的进步，在工业领域也同样取得了成功，因此已成为深度学习从业者不可或缺的工具。这意味着通过这些策略，您可以使模型消耗更少的内存，并真正突破硬件的限制。

为什么需要内存优化？

内存优化是训练深度学习模型中最关键的组成部分之一。对于像 GPT 和 BERT 这样拥有数十亿参数的模型，在可用硬件上进行训练时，必须有效地管理内存。DeepSpeed 是一个开源库，用于训练深度学习模型，它具有 ZeRO 优化器、卸载技术和梯度检查点等功能，以避免训练过程中的主要内存占用。

深度学习中的内存问题

深度学习模型的规模和复杂性近年来显著增加。这些大型模型需要大量的内存进行训练。微软开发的深度学习优化库 DeepSpeed 为这些挑战提供了强大的解决方案。

深度学习模型有点像艺术，它们代表着新兴事物，随着模型规模的增长，内存相关的问题也会随之出现。一些最常见的内存相关问题包括：

模型参数 - 大型模型，如 GPT-3，拥有数千亿个参数，因此需要大量的内存来存储。
梯度 - 在训练过程中计算每个参数的梯度也必须计算并保存在内存中，这会消耗更多的内存。
激活映射 - 前向传递过程中产生的所有中间值都需要存储，直到反向传递仅用于梯度计算，这称为激活映射。
批大小 - 更大的批大小可以提高收敛速度，但会消耗更多的内存。
数据并行 - 将数据在多个 GPU 之间进行分发是减少训练时间的好策略，但毫无疑问，它确实会消耗大量的内存，除非得到控制。

除非识别出这些陷阱，否则即使在消费级硬件上训练大型模型也变得不可能。DeepSpeed 通过使用创新的节省内存技术克服了这些挑战。

DeepSpeed 的内存优化技术

DeepSpeed 在训练模型时有多种方法可以优化内存使用。一些方法包括 ZeRO（零冗余优化器）、梯度检查点和激活重计算。

零冗余优化器 (ZeRO)
梯度检查点
卸载技术

1. 零冗余优化器 (ZeRO)

ZeRO 主要关注在优化器状态、梯度和模型参数的冗余副本被移除的地方进行内存优化。ZeRO 经历以下三个阶段：

阶段 1 - 将优化器状态跨 GPU 分片，每个 GPU 存储一部分优化器状态。
阶段 2 - 进一步减少内存，因为梯度跨 GPU 分片。
阶段 3 - 模型参数被分片，现在可以训练高达万亿参数的模型。

示例

import deepspeed

model = MyModel()  # your dl model
optimizer = torch.optim.Adam(model.parameters(), lr=1e-4)

# DeepSpeed configuration for ZeRO
ds_config = {
 "train_batch_size": 8,
 "zero_optimization": {
    "stage": 2,   # adjust the stage of ZeRO here
    "allgather_partitions": True,
    "reduce_scatter": True,
    "allgather_bucket_size": 5e8,
    "overlap_comm": True,
    "contiguous_gradients": True
 }
}
# Initialize DeepSpeed
model_engine, optimizer, _, _ = deepspeed.initialize(
  model=model,
  optimizer=optimizer,
  config_params=ds_config
)

# Training loop
for batch in train_dataloader:
    loss = model_engine(batch)
model_engine.backward(loss)
model_engine.step()

您会注意到内存使用量大大降低，特别是对于大型模型。然后，内存分析器可以突出显示 ZeRO 优化启动的位置。

2. 梯度检查点

梯度检查点通过不在缓冲区中存储前向传递期间的激活来减少内存。相反，它们在反向传递中被重建，牺牲一点计算来节省一些内存。

示例

import torch
from torch.utils.checkpoint import checkpoint

def custom_forward(*inputs):
    return model(*inputs)

# Gradient checkpointing
outputs = checkpoint(custom_forward, input_data)
loss = criterion(outputs, labels)
loss.backward()

在这种情况下，节省的内存将取决于中间激活的大小。

3. 卸载技术

DeepSpeed 还提供另一种形式的卸载。它允许您将模型的部分内容（如优化器状态和梯度）移动到 CPU 甚至 NVMe 存储器，从而释放 GPU 内存以供其他用途使用。

CPU 卸载

DeepSpeed 允许我们将优化器状态和梯度卸载到 CPU。这也释放了宝贵的 GPU 内存。如果 GPU 上的内存有限，但 CPU 上的内存相当多，这将非常有用。

示例

ds_config = {
    "train_batch_size": 8,
    "zero_optimization": {
        "stage": 2,
        "offload_optimizer": {
            "device": "cpu",
            "pin_memory": True
        },
        "offload_param": {
            "device": "cpu",
            "pin_memory": True
        }
    }
}
model_engine, optimizer, _, _ = deepspeed.initialize(
   model=model,
   optimizer=optimizer,
   config_params=ds_config
)

由于卸载到 CPU 的传输涉及设备间通信成本，因此训练速度相对较慢，但在那些原本无法适应内存受限 GPU 的模型大小上，它仍然很有用。

NVMe 卸载

对于大型模型，这还不够。DeepSpeed 还将优化器状态和梯度卸载到 NVMe 存储器。这将进一步提高模型的规模，即使在训练中也能实现。

示例

ds_config = {
    "train_batch_size": 8,
    "zero_optimization": {
        "stage": 2,
        "offload_optimizer": {
            "device": "nvme",
            "nvme_path": "/local_nvme"
        },
        "offload_param": {
            "device": "nvme",
            "nvme_path": "/local_nvme"
        }
    }
}
model_engine, optimizer, _, _ = deepspeed.init(
   model=model,
   optimizer=optimizer,
   config_params=ds_config
)

使用 NVMe 进行卸载将能够训练海量模型，尽管速度将在很大程度上取决于 NVMe 驱动器的 I/O 速度。

内存优化案例研究

让我们讨论一些使用 DeepSpeed 进行内存优化的真实案例研究 -

案例研究 1：使用 ZeRO 优化训练 GPT-2

使用 DeepSpeed，一个研究团队将这个拥有 15 亿个参数的 GPT-2 模型的训练扩展到了消费级 GPU。使用 ZeRO 阶段 3，可以在 4 个 NVIDIA RTX 3090 GPU 上对其进行训练，每个 GPU 的总内存为 24 GB。如果没有使用 ZeRO，训练将是不可能的，因为该模型每个 GPU 需要超过 50 GB 的内存。

案例研究 2：使用 NVMe 卸载 1750 亿参数模型

微软利用 DeepSpeed 的卸载功能，在一个内存有限的 GPU 集群上训练了一个 1750 亿参数的模型。在模型训练期间，几乎没有内存瓶颈来卸载优化器状态和参数，这表明即使 GPU 资源有限，卸载也可以为超大型模型铺平道路。

打印页面