DeepSpeed - 优化器

优化和调度是深度学习中大型模型获得更好性能的基础。DeepSpeed 是一个开源的深度学习优化库，它使用各种支持的技术更有效地帮助模型训练:内存优化、梯度累积和混合精度训练。

DeepSpeed 的两个关键组件是 DeepSpeed 优化器和 DeepSpeed 调度器。它们协同工作，可以有效地管理系统资源、加速训练并减少普通硬件设置上的内存占用，以训练可能具有数十亿个参数的模型。

让我们通过代码中的示例详细了解 DeepSpeed 优化器的工作原理。我们将在下一章中介绍 DeepSpeed 调度程序。

什么是 DeepSpeed 优化器?

DeepSpeed 优化器通过高效分配内存来管理模型优化。它支持与任何流行的深度学习框架(例如 PyTorch)进行本地接口的优化，因此，它可以处理包括动量和梯度累积在内的优化器状态。这是一个深度速度优化器，其主要功能包括零冗余优化器、ZeRO、混合精度训练和梯度检查点。

DeepSpeed Optimizer 的主要功能

以下是 DeepSpeed Optimizer 的主要功能 −

1. 零冗余优化器 (ZeRO)

通过将优化器、梯度和模型参数的状态划分到多个设备，可以减少内存消耗。

这使得能够在容量有限的设备上训练巨型模型。

2.混合精度训练

通过同时使用 16 位和 32 位浮点表示，混合精度训练可以最小化内存消耗，同时不会降低模型准确性。

3. 梯度检查点

它将模型分片成块，并在前向传递期间仅存储激活的子集；因此，它可能会在后向传递期间计算中间值以节省内存。

使用 DeepSpeed 优化器的示例

以下是基于 PyTorch 的示例，使用带有 ZeRO 的 DeepSpeed 优化器 −

import deepspeed
import torch
import torch.nn as nn
import torch.optim as optim

# 示例模型定义
class SampleModel(nn.Module):
    def __init__(self):
        super(SampleModel, self).__init__()
        self.fc = nn.Linear(10, 1)

    def forward(self, x):
        return self.fc(x)

# 初始化模型和优化器
model = SampleModel()
optimizer = optim.Adam(model.parameters(), lr=0.001)

# DeepSpeed 配置
ds_config = {
    "train_batch_size": 8,
    "gradient_accumulation_steps": 2,
    "optimizer": {
        "type": "Adam",
        "params": {
            "lr": 0.001,
        }
    },
    "zero_optimization": {
        "stage": 1
    }
}

# 初始化 DeepSpeed
model_engine, optimizer, _, _ = deepspeed.initialize(model=model, optimizer=optimizer, config_params=ds_config)

# 样本输入和正向传递
inputs = torch.randn(8, 10)
outputs = model_engine(inputs)
loss = output.mean()

# 反向传递和优化
model_engine.backward(loss)
model_engine.step()

输出

在 PyCharm 或 VSCode 等 IDE 环境中执行时，它将看起来像 −

Deepspeed is initiated
Input tensor: torch.Size([8, 10])
Forward pass completed
Loss: -0.015
Backward pass and optimizer step complete

以上是示例IDE(例如 PyCharm 或 VSCode)显示已应用优化器的代码片段，终端输出将显示此优化器的成功执行。

应用本章中显示的示例和输出将使将这些工具应用于您的深度学习工作流程变得更加容易。