小白入门的DeepSpeed原理和代码实现（一）训练的内存需求

最新推荐文章于 2025-10-18 13:49:34 发布

原创最新推荐文章于 2025-10-18 13:49:34 发布 · 1.8k 阅读

47 ·

CC 4.0 BY-SA版权

文章标签：

#人工智能 #AI大模型 #大模型 #LLM #ai #agi

这篇文章探讨了在多个GPU上分布式训练神经网络的技术，并测试了微软的DeepSpeed库提供的多种分布式训练优化。

一、训练的内存需求

本页的大部分内容将集中在减少训练神经网络所需的内存的技术上。因此，为了更好地理解和欣赏这些技术，我将首先分解训练过程中内存消耗的所有方式。

根据Sohoni等人的说法，我将使用以下术语来描述训练期间内存消耗的来源（Sohoni等人只描述了三个来源，但我将梯度内存和优化器内存分开，因为DeepSpeed中也做了同样的区分，稍后将展示）：

模型内存：存储模型权重所需的内存。
梯度内存：存储每个模型权重梯度所需的内存。
优化器内存：存储优化器所需的任何额外状态所需的内存。
激活内存：存储在前向传递期间计算的中间值所需的内存。

测量四种内存消耗源

让我们现在看看如何测量这四种来源消耗的内存量，从一个简单的全连接神经网络开始：

model = torch.nn.Sequential(
    torch.nn.Linear(512, 1024),
    torch.nn.Linear(1024, 1024),
    torch.nn.Linear(1024, 1024),
    torch.nn.Linear(1024, 512)
)

要获得总的模型内存，我们计算模型中的参数数量，并乘以每个参数用来表示的字节数（在这种情况下是4个字节，因为每个权重是一个32位浮点数）：

def get_model_memory(model: torch.nn.Module):
    """
    返回给定模型的内存使用情况
    """
    total_memory = 0
    for param in model.parameters():
        total_memory += param.numel() * param.element_size()
    return total_memory

print("模型内存：{:,} 字节".format(get_model_memory(model)))

模型内存：12,597,248 字节

梯度内存与模型内存相同，因为在后向传递期间我们需要为每个权重存储梯度。

消耗的总优化器内存取决于训练期间使用的优化器类型。下表显示了PyTorch中一些流行的优化器所需的优化器内存，其中M_model代表模型内存[2]：

在这里插入图片描述

最后，要计算总的激活内存，我们需要计算在前向传递期间计算的隐藏层输出。有几种不同的方法可以做到这一点；这里，我们将递归地注入PyTorch前向Hook到每个子模块中以捕获该子模块的输出大小。

首先，我们定义了Hook以及一个简单的辅助类ActivationCounter，用于跟踪模型隐藏层输出的总大小：

class ActivationCounter:

    def __init__(self):
        self.activation_bytes = 0

    def add_activations(self, tensor):
        self.activation_bytes += tensor.numel() * tensor.element_size()

def activation_counter_hook(counter: ActivationCounter):

    def hook(self, input, output):
        counter.add_activations(output.data)

    return hook

接下来，我们定义一个函数，它将递归地将Hook注入到模型的每个子模块中：

def register_hooks_recursive(model, counter: ActivationCounter):
  for module in model.children():
      module.register_forward_hook(activation_counter_hook(counter))
      register_hooks_recursive(module, counter)

activation_counter = ActivationCounter()
register_hooks_recursive(model, activation_counter)

现在要找到激活内存消耗的总量，我们执行模型的一个前向传递，然后打印activation_counter.activation_bytes的值：

inputs = torch.randn(4, 512)
outputs = model(inputs)

# 因为Hook只捕获层输出，我们需要单独添加原始输入张量的大小
activation_counter.add_activations(inputs)

print("激活内存：{:,} 字节".format(
  activation_counter.activation_bytes
))